diff --git a/js/ai/src/generate.ts b/js/ai/src/generate.ts index 2cea52ebd1..3d55762992 100755 --- a/js/ai/src/generate.ts +++ b/js/ai/src/generate.ts @@ -57,13 +57,19 @@ import { type ToolResponsePart, } from './model.js'; import { isExecutablePrompt } from './prompt.js'; -import { isDynamicResourceAction, ResourceAction } from './resource.js'; +import { + isDynamicResourceAction, + resolveResources, + ResourceAction, + ResourceArgument, +} from './resource.js'; import { isDynamicTool, resolveTools, toToolDefinition, type ToolArgument, } from './tool.js'; + export { GenerateResponse, GenerateResponseChunk }; /** Specifies how tools should be called by the model. */ @@ -121,7 +127,7 @@ export interface GenerateOptions< /** List of registered tool names or actions to treat as a tool for this generation if supported by the underlying model. */ tools?: ToolArgument[]; /** List of dynamic resources to be made available to this generate request. */ - resources?: ResourceAction[]; + resources?: ResourceArgument[]; /** Specifies how tools should be called by the model. */ toolChoice?: ToolChoice; /** Configuration for the generation request. */ @@ -222,6 +228,10 @@ export async function toGenerateRequest( if (options.tools) { tools = await resolveTools(registry, options.tools); } + let resources: ResourceAction[] | undefined; + if (options.resources) { + resources = await resolveResources(registry, options.resources); + } const resolvedSchema = toJsonSchema({ schema: options.output?.schema, @@ -245,6 +255,7 @@ export async function toGenerateRequest( config: options.config, docs: options.docs, tools: tools?.map(toToolDefinition) || [], + resources: resources?.map((a) => a.__action) || [], output: { ...(resolvedFormat?.config || {}), ...options.output, @@ -285,7 +296,8 @@ async function toolsToActionRefs( for (const t of toolOpt) { if (typeof t === 'string') { - tools.push(await resolveFullToolName(registry, t)); + const names = await resolveFullToolNames(registry, t); + tools.push(...names); } else if (isAction(t) || isDynamicTool(t)) { tools.push(`/${t.__action.metadata?.type}/${t.__action.name}`); } else if (isExecutablePrompt(t)) { @@ -298,6 +310,27 @@ async function toolsToActionRefs( return tools; } +async function resourcesToActionRefs( + registry: Registry, + resOpt?: ResourceArgument[] +): Promise { + if (!resOpt) return; + + const resources: string[] = []; + + for (const r of resOpt) { + if (typeof r === 'string') { + const names = await resolveFullResourceNames(registry, r); + resources.push(...names); + } else if (isAction(r)) { + resources.push(`/resource/${r.__action.name}`); + } else { + throw new Error(`Unable to resolve resource: ${JSON.stringify(r)}`); + } + } + return resources; +} + function messagesFromOptions(options: GenerateOptions): MessageData[] { const messages: MessageData[] = []; if (options.system) { @@ -358,6 +391,10 @@ export async function generate< const params = await toGenerateActionOptions(registry, resolvedOptions); const tools = await toolsToActionRefs(registry, resolvedOptions.tools); + const resources = await resourcesToActionRefs( + registry, + resolvedOptions.resources + ); const streamingCallback = stripNoop( resolvedOptions.onChunk ?? resolvedOptions.streamingCallback ) as StreamingCallback; @@ -372,6 +409,7 @@ export async function generate< const request = await toGenerateRequest(registry, { ...resolvedOptions, tools, + resources, }); return new GenerateResponse(response, { request: response.request ?? request, @@ -458,6 +496,7 @@ export async function toGenerateActionOptions< ): Promise { const resolvedModel = await resolveModel(registry, options.model); const tools = await toolsToActionRefs(registry, options.tools); + const resources = await resourcesToActionRefs(registry, options.resources); const messages: MessageData[] = messagesFromOptions(options); const resolvedSchema = toJsonSchema({ @@ -478,6 +517,7 @@ export async function toGenerateActionOptions< docs: options.docs, messages: messages, tools, + resources, toolChoice: options.toolChoice, config: { version: resolvedModel.version, @@ -530,17 +570,49 @@ function stripUndefinedOptions(input?: any): any { return copy; } -async function resolveFullToolName( +async function resolveFullToolNames( registry: Registry, name: string -): Promise { +): Promise { + let names: string[]; + const parts = name.split(':'); + if (parts.length > 1) { + // Dynamic Action Provider + names = await registry.resolveActionNames( + `/dynamic-action-provider/${name}` + ); + if (names.length) { + return names; + } + } if (await registry.lookupAction(`/tool/${name}`)) { - return `/tool/${name}`; - } else if (await registry.lookupAction(`/prompt/${name}`)) { - return `/prompt/${name}`; - } else { - throw new Error(`Unable to determine type of of tool: ${name}`); + return [`/tool/${name}`]; + } + if (await registry.lookupAction(`/prompt/${name}`)) { + return [`/prompt/${name}`]; + } + throw new Error(`Unable to resolve tool: ${name}`); +} + +async function resolveFullResourceNames( + registry: Registry, + name: string +): Promise { + let names: string[]; + const parts = name.split(':'); + if (parts.length > 1) { + // Dynamic Action Provider + names = await registry.resolveActionNames( + `/dynamic-action-provider/${name}` + ); + if (names.length) { + return names; + } + } + if (await registry.lookupAction(`/resource/${name}`)) { + return [`/resource/${name}`]; } + throw new Error(`Unable to resolve resource: ${name}`); } export type GenerateStreamOptions< diff --git a/js/ai/src/generate/action.ts b/js/ai/src/generate/action.ts index 66229484e8..8161ed827d 100644 --- a/js/ai/src/generate/action.ts +++ b/js/ai/src/generate/action.ts @@ -56,7 +56,11 @@ import { type Part, type Role, } from '../model.js'; -import { findMatchingResource } from '../resource.js'; +import { + findMatchingResource, + resolveResources, + type ResourceAction, +} from '../resource.js'; import { resolveTools, toToolDefinition, type ToolAction } from '../tool.js'; import { assertValidToolNames, @@ -151,14 +155,15 @@ async function resolveParameters( registry: Registry, request: GenerateActionOptions ) { - const [model, tools, format] = await Promise.all([ + const [model, tools, resources, format] = await Promise.all([ resolveModel(registry, request.model, { warnDeprecated: true }).then( (r) => r.modelAction ), resolveTools(registry, request.tools), + resolveResources(registry, request.resources), resolveFormat(registry, request.output), ]); - return { model, tools, format }; + return { model, tools, resources, format }; } /** Given a raw request and a formatter, apply the formatter's logic and instructions to the request. */ @@ -246,12 +251,12 @@ async function generate( streamingCallback?: StreamingCallback; } ): Promise { - const { model, tools, format } = await resolveParameters( + const { model, tools, resources, format } = await resolveParameters( registry, rawRequest ); rawRequest = applyFormat(rawRequest, format); - rawRequest = await applyResources(registry, rawRequest); + rawRequest = await applyResources(registry, rawRequest, resources); // check to make sure we don't have overlapping tool names *before* generation await assertValidToolNames(tools); @@ -481,7 +486,8 @@ function getRoleFromPart(part: Part): Role { async function applyResources( registry: Registry, - rawRequest: GenerateActionOptions + rawRequest: GenerateActionOptions, + resources: ResourceAction[] ): Promise { // quick check, if no resources bail. if (!rawRequest.messages.find((m) => !!m.content.find((c) => c.resource))) { @@ -500,7 +506,12 @@ async function applyResources( updatedContent.push(p); continue; } - const resource = await findMatchingResource(registry, p.resource); + + const resource = await findMatchingResource( + registry, + resources, + p.resource + ); if (!resource) { throw new GenkitError({ status: 'NOT_FOUND', diff --git a/js/ai/src/model-types.ts b/js/ai/src/model-types.ts index 1e842d77ec..e64381345f 100644 --- a/js/ai/src/model-types.ts +++ b/js/ai/src/model-types.ts @@ -388,6 +388,8 @@ export const GenerateActionOptionsSchema = z.object({ messages: z.array(MessageSchema), /** List of registered tool names for this generation if supported by the underlying model. */ tools: z.array(z.string()).optional(), + /** List of registered resource names for this generation if supported by the underlying model. */ + resources: z.array(z.string()).optional(), /** Tool calling mode. `auto` lets the model decide whether to use tools, `required` forces the model to choose a tool, and `none` forces the model not to use any tools. Defaults to `auto`. */ toolChoice: z.enum(['auto', 'required', 'none']).optional(), /** Configuration for the generation request. */ diff --git a/js/ai/src/resource.ts b/js/ai/src/resource.ts index 1257a2ee75..1ffb8192d1 100644 --- a/js/ai/src/resource.ts +++ b/js/ai/src/resource.ts @@ -84,6 +84,45 @@ export interface ResourceAction matches(input: ResourceInput): boolean; } +/** + * A reference to a resource in the form of a name or a ResourceAction. + */ +export type ResourceArgument = ResourceAction | string; + +export async function resolveResources( + registry: Registry, + resources?: ResourceArgument[] +): Promise { + if (!resources || resources.length === 0) { + return []; + } + + return await Promise.all( + resources.map(async (ref): Promise => { + if (typeof ref === 'string') { + return await lookupResourceByName(registry, ref); + } else if (isAction(ref)) { + return ref; + } + throw new Error('Resources must be strings, or actions'); + }) + ); +} + +export async function lookupResourceByName( + registry: Registry, + name: string +): Promise { + const resource = + (await registry.lookupAction(name)) || + (await registry.lookupAction(`/resource/${name}`)) || + (await registry.lookupAction(`/dynamic-action-provider/${name}`)); + if (!resource) { + throw new Error(`Resource ${name} not found`); + } + return resource as ResourceAction; +} + /** * Defines a resource. * @@ -122,11 +161,28 @@ export type DynamicResourceAction = ResourceAction & { */ export async function findMatchingResource( registry: Registry, + resources: ResourceAction[], input: ResourceInput ): Promise { - for (const actKeys of Object.keys(await registry.listResolvableActions())) { - if (actKeys.startsWith('/resource/')) { - const resource = (await registry.lookupAction(actKeys)) as ResourceAction; + // First look in any resources explicitly listed in the generate request + for (const res of resources) { + if (res.matches(input)) { + return res; + } + } + + // Then search the registry + for (const registryKey of Object.keys( + await registry.listResolvableActions() + )) { + // We decided not to look in DAP actions because they might be slow. + // DAP actions with resources will only be found if they are listed in the + // resources section, and then they will be found above. + if (registryKey.startsWith('/resource/')) { + const resource = (await registry.lookupAction( + registryKey + )) as ResourceAction; + if (resource.matches(input)) { return resource; } diff --git a/js/ai/src/tool.ts b/js/ai/src/tool.ts index 6d6b4d3eeb..19c9665676 100644 --- a/js/ai/src/tool.ts +++ b/js/ai/src/tool.ts @@ -218,7 +218,8 @@ export async function lookupToolByName( const tool = (await registry.lookupAction(name)) || (await registry.lookupAction(`/tool/${name}`)) || - (await registry.lookupAction(`/prompt/${name}`)); + (await registry.lookupAction(`/prompt/${name}`)) || + (await registry.lookupAction(`/dynamic-action-provider/${name}`)); if (!tool) { throw new Error(`Tool ${name} not found`); } diff --git a/js/ai/tests/generate/generate_test.ts b/js/ai/tests/generate/generate_test.ts index 0522b8d817..4abd89038e 100644 --- a/js/ai/tests/generate/generate_test.ts +++ b/js/ai/tests/generate/generate_test.ts @@ -82,6 +82,7 @@ describe('toGenerateRequest', () => { ], config: undefined, docs: undefined, + resources: [], tools: [], output: {}, }, @@ -100,6 +101,7 @@ describe('toGenerateRequest', () => { ], config: undefined, docs: undefined, + resources: [], tools: [ { name: 'tellAFunnyJoke', @@ -132,6 +134,7 @@ describe('toGenerateRequest', () => { messages: [{ role: 'user', content: [{ text: 'Add 10 and 5.' }] }], config: undefined, docs: undefined, + resources: [], tools: [ { description: 'add two numbers together', @@ -167,6 +170,7 @@ describe('toGenerateRequest', () => { ], config: undefined, docs: undefined, + resources: [], tools: [ { name: 'tellAFunnyJoke', @@ -219,6 +223,7 @@ describe('toGenerateRequest', () => { ], config: undefined, docs: undefined, + resources: [], tools: [], output: {}, }, @@ -241,6 +246,7 @@ describe('toGenerateRequest', () => { ], config: undefined, docs: undefined, + resources: [], tools: [], output: {}, }, @@ -258,6 +264,7 @@ describe('toGenerateRequest', () => { ], config: undefined, docs: [{ content: [{ text: 'context here' }] }], + resources: [], tools: [], output: {}, }, @@ -303,6 +310,7 @@ describe('toGenerateRequest', () => { ], config: undefined, docs: undefined, + resources: [], tools: [], output: { constrained: true, @@ -434,7 +442,7 @@ describe('generate', () => { ); }); - it('applies resources', async () => { + it('applies resources in the registry', async () => { defineResource( registry, { name: 'testResource', template: 'test://resource/{param}' }, diff --git a/js/ai/tests/resource/resource_test.ts b/js/ai/tests/resource/resource_test.ts index 1829ed8a3e..fff8fa88fd 100644 --- a/js/ai/tests/resource/resource_test.ts +++ b/js/ai/tests/resource/resource_test.ts @@ -206,12 +206,19 @@ describe('resource', () => { return { content: [{ text: `bar` }] }; } ); + const resList = [ + dynamicResource({ uri: 'baz://qux' }, () => ({ + content: [{ text: `baz` }], + })), + ]; - const gotBar = await findMatchingResource(registry, { uri: 'bar://baz' }); + const gotBar = await findMatchingResource(registry, resList, { + uri: 'bar://baz', + }); assert.ok(gotBar); assert.strictEqual(gotBar.__action.name, 'testResource'); - const gotFoo = await findMatchingResource(registry, { + const gotFoo = await findMatchingResource(registry, resList, { uri: 'foo://bar/something', }); assert.ok(gotFoo); @@ -225,7 +232,21 @@ describe('resource', () => { dynamic: true, }); - const gotUnmatched = await findMatchingResource(registry, { + const gotBaz = await findMatchingResource(registry, resList, { + uri: 'baz://qux', + }); + assert.ok(gotBaz); + assert.strictEqual(gotBaz.__action.name, 'baz://qux'); + assert.deepStrictEqual(gotBaz.__action.metadata, { + resource: { + template: undefined, + uri: 'baz://qux', + }, + type: 'resource', + dynamic: true, + }); + + const gotUnmatched = await findMatchingResource(registry, resList, { uri: 'unknown://bar/something', }); assert.strictEqual(gotUnmatched, undefined); diff --git a/js/core/src/dynamic-action-provider.ts b/js/core/src/dynamic-action-provider.ts new file mode 100644 index 0000000000..af6182cdbe --- /dev/null +++ b/js/core/src/dynamic-action-provider.ts @@ -0,0 +1,213 @@ +/** + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type * as z from 'zod'; +import { Action, ActionMetadata, defineAction } from './action.js'; +import { ActionType, Registry } from './registry.js'; + +type DapValue = { + [K in ActionType]?: Action[]; +}; + +class SimpleCache { + private value: DapValue | undefined; + private expiresAt: number | undefined; + private ttlMillis: number; + private dap: DynamicActionProviderAction; + private dapFn: DapFn; + private fetchPromise: Promise | null = null; + + constructor( + dap: DynamicActionProviderAction, + config: DapConfig, + dapFn: DapFn + ) { + this.dap = dap; + this.dapFn = dapFn; + this.ttlMillis = !config.cacheConfig?.ttlMillis + ? 3 * 1000 + : config.cacheConfig?.ttlMillis; + } + + async getOrFetch(): Promise { + const isStale = + !this.value || + !this.expiresAt || + this.ttlMillis < 0 || + Date.now() > this.expiresAt; + if (!isStale) { + return this.value!; + } + + if (!this.fetchPromise) { + this.fetchPromise = (async () => { + try { + // Get a new value + this.value = await this.dapFn(); // this returns the actual actions + this.expiresAt = Date.now() + this.ttlMillis; + + // Also run the action + this.dap.run(this.value); // This returns metadata and shows up in dev UI + + return this.value; + } catch (error) { + console.error('Error fetching Dynamic Action Provider value:', error); + this.invalidate(); + throw error; // Rethrow to reject the fetchPromise + } finally { + // Allow new fetches after this one completes or fails. + this.fetchPromise = null; + } + })(); + } + return await this.fetchPromise; + } + + invalidate() { + this.value = undefined; + } +} + +export interface DynamicRegistry { + __cache: SimpleCache; + invalidateCache(): void; + getAction( + actionType: string, + actionName: string + ): Promise | undefined>; + listActionMetadata( + actionType: string, + actionName: string + ): Promise; +} + +export type DynamicActionProviderAction = Action< + z.ZodTypeAny, + z.ZodTypeAny, + z.ZodTypeAny +> & + DynamicRegistry & { + __action: { + metadata: { + type: 'dynamic-action-provider'; + }; + }; + }; + +export function isDynamicActionProvider( + obj: Action +): obj is DynamicActionProviderAction { + return obj.__action?.metadata?.type == 'dynamic-action-provider'; +} + +export interface DapConfig { + name: string; + description?: string; + cacheConfig?: { + // Negative = no caching + // Zero or undefined = default (3000 milliseconds) + // Positive number = how many milliseconds the cache is valid for + ttlMillis: number | undefined; + }; + metadata?: Record; +} + +export type DapFn = () => Promise; +export type DapMetadata = { + [K in ActionType]?: ActionMetadata[]; +}; + +function transformDapValue(value: DapValue): DapMetadata { + const metadata: DapMetadata = {}; + for (const key of Object.keys(value)) { + metadata[key] = value[key].map((a) => { + return a.__action; + }); + } + return metadata; +} + +export function defineDynamicActionProvider( + registry: Registry, + config: DapConfig | string, + fn: DapFn +): DynamicActionProviderAction { + let cfg: DapConfig; + if (typeof config == 'string') { + cfg = { name: config }; + } else { + cfg = { ...config }; + } + const a = defineAction( + registry, + { + ...cfg, + actionType: 'dynamic-action-provider', + metadata: { ...(cfg.metadata || {}), type: 'dynamic-action-provider' }, + }, + async (i, _options) => { + // The actions are retrieved and saved in a cache and then passed in here. + // We run this action to return the metadata for the actions only. + // We pass the actions in here to prevent duplicate calls to the mcp + // and also so we are guaranteed the same actions since there is only a + // single call to mcp client/host. + return transformDapValue(i); + } + ); + implementDap(a as DynamicActionProviderAction, cfg, fn); + return a as DynamicActionProviderAction; +} + +function implementDap( + dap: DynamicActionProviderAction, + config: DapConfig, + dapFn: DapFn +) { + dap.__cache = new SimpleCache(dap, config, dapFn); + dap.invalidateCache = () => { + dap.__cache.invalidate(); + }; + + dap.getAction = async (actionType: string, actionName: string) => { + const result = await dap.__cache.getOrFetch(); + if (result[actionType]) { + return result[actionType].find((t) => t.__action.name == actionName); + } + return undefined; + }; + + dap.listActionMetadata = async (actionType: string, actionName: string) => { + const result = await dap.__cache.getOrFetch(); + if (!result[actionType]) { + return []; + } + + // Match everything in the actionType + const metadata = result[actionType].map((a) => a.__action); + if (actionName == '*') { + return metadata; + } + + // Prefix matching + if (actionName.endsWith('*')) { + const prefix = actionName.slice(0, -1); + return metadata.filter((m) => m.name.startsWith(prefix)); + } + + // Single match or empty array + return metadata.filter((m) => m.name == actionName); + }; +} diff --git a/js/core/src/index.ts b/js/core/src/index.ts index 5d0b7dc3d9..632b9b30fc 100644 --- a/js/core/src/index.ts +++ b/js/core/src/index.ts @@ -51,6 +51,12 @@ export { type ContextProvider, type RequestData, } from './context.js'; +export { + defineDynamicActionProvider, + type DapConfig, + type DapFn, + type DynamicActionProviderAction, +} from './dynamic-action-provider.js'; export { GenkitError, UnstableApiError, diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts index 83b0dda6ff..90fbe91411 100644 --- a/js/core/src/registry.ts +++ b/js/core/src/registry.ts @@ -26,6 +26,7 @@ import { lookupBackgroundAction, } from './background-action.js'; import { ActionContext } from './context.js'; +import { isDynamicActionProvider } from './dynamic-action-provider.js'; import { GenkitError } from './error.js'; import { logger } from './logging.js'; import type { PluginProvider } from './plugin.js'; @@ -36,23 +37,30 @@ export type AsyncProvider = () => Promise; /** * Type of a runnable action. */ -export type ActionType = - | 'custom' - | 'embedder' - | 'evaluator' - | 'executable-prompt' - | 'flow' - | 'indexer' - | 'model' - | 'background-model' - | 'check-operation' - | 'cancel-operation' - | 'prompt' - | 'reranker' - | 'retriever' - | 'tool' - | 'util' - | 'resource'; +const ACTION_TYPES = [ + 'custom', + 'dynamic-action-provider', + 'embedder', + 'evaluator', + 'executable-prompt', + 'flow', + 'indexer', + 'model', + 'background-model', + 'check-operation', + 'cancel-operation', + 'prompt', + 'reranker', + 'retriever', + 'tool', + 'util', + 'resource', +] as const; +export type ActionType = (typeof ACTION_TYPES)[number]; + +export function isActionType(value: string): value is ActionType { + return (ACTION_TYPES as readonly string[]).includes(value); +} /** * A schema is either a Zod schema or a JSON schema. @@ -71,6 +79,7 @@ function parsePluginName(registryKey: string) { } interface ParsedRegistryKey { + dynamicActionHost?: string; actionType: ActionType; pluginName?: string; actionName: string; @@ -78,6 +87,7 @@ interface ParsedRegistryKey { /** * Parses the registry key into key parts as per the key format convention. Ex: + * - mcp-host:tool/my-tool * - /model/googleai/gemini-2.0-flash * - /prompt/my-plugin/folder/my-prompt * - /util/generate @@ -85,6 +95,30 @@ interface ParsedRegistryKey { export function parseRegistryKey( registryKey: string ): ParsedRegistryKey | undefined { + if (registryKey.startsWith('/dynamic-action-provider')) { + // Dynamic Action Provider format: 'dynamic-action-provider/mcp-host:tool/mytool' or 'mcp-host:tool/*' + const keyTokens = registryKey.split(':'); + const hostTokens = keyTokens[0].split('/'); + if (hostTokens.length < 3) { + return undefined; + } + if (keyTokens.length < 2) { + return { + actionType: 'dynamic-action-provider', + actionName: hostTokens[2], + }; + } + const tokens = keyTokens[1].split('/'); + if (tokens.length < 2 || !isActionType(tokens[0])) { + return undefined; + } + return { + dynamicActionHost: hostTokens[2], + actionType: tokens[0], + actionName: tokens.slice(1).join('/'), + }; + } + const tokens = registryKey.split('/'); if (tokens.length < 3) { // Invalid key format @@ -158,6 +192,24 @@ export class Registry { return new Registry(parent); } + async resolveActionNames(key: string): Promise { + const parsedKey = parseRegistryKey(key); + if (parsedKey?.dynamicActionHost) { + const hostId = `/dynamic-action-provider/${parsedKey.dynamicActionHost}`; + const dap = await this.actionsById[hostId]; + if (!dap || !isDynamicActionProvider(dap)) { + return []; + } + return ( + await dap.listActionMetadata(parsedKey.actionType, parsedKey.actionName) + ).map((m) => `${hostId}:${parsedKey.actionType}/${m.name}`); + } + if (await this.lookupAction(key)) { + return [key]; + } + return []; + } + /** * Looks up an action in the registry. * @param key The key of the action to lookup. @@ -168,8 +220,21 @@ export class Registry { O extends z.ZodTypeAny, R extends Action, >(key: string): Promise { - // We always try to initialize the plugin first. const parsedKey = parseRegistryKey(key); + if ( + parsedKey?.dynamicActionHost && + this.actionsById[ + `/dynamic-action-provider/${parsedKey.dynamicActionHost}` + ] + ) { + // If it's a dynamic action provider, get the dynamic action. + const action = await this.getDynamicAction(parsedKey); + if (action) { + return action as R; + } + } + + // We always try to initialize the plugin first. if (parsedKey?.pluginName && this.pluginsByName[parsedKey.pluginName]) { await this.initializePlugin(parsedKey.pluginName); @@ -184,6 +249,7 @@ export class Registry { ); } } + return ( ((await this.actionsById[key]) as R) || this.parent?.lookupAction(key) ); @@ -404,6 +470,21 @@ export class Registry { } } + async getDynamicAction( + key: ParsedRegistryKey + ): Promise | undefined> { + if (key.actionName.includes('*')) { + // * means multiple actions, this returns exactly one. + return undefined; + } + const actionId = `/dynamic-action-provider/${key.dynamicActionHost}`; + const dap = await this.actionsById[actionId]; + if (!dap || !isDynamicActionProvider(dap)) { + return undefined; + } + return await dap.getAction(key.actionType, key.actionName); + } + /** * Initializes a plugin already registered with {@link registerPluginProvider}. * @param name The name of the plugin to initialize. diff --git a/js/core/src/schema.ts b/js/core/src/schema.ts index f3fa32bb30..955cc1d9b9 100644 --- a/js/core/src/schema.ts +++ b/js/core/src/schema.ts @@ -130,7 +130,9 @@ export function parseSchema( options: ProvidedSchema ): T { const { valid, errors, schema } = validateSchema(data, options); - if (!valid) throw new ValidationError({ data, errors: errors!, schema }); + if (!valid) { + throw new ValidationError({ data, errors: errors!, schema }); + } return data as T; } diff --git a/js/core/tests/dynamic-action-provider_test.ts b/js/core/tests/dynamic-action-provider_test.ts new file mode 100644 index 0000000000..b1e6ff8ee5 --- /dev/null +++ b/js/core/tests/dynamic-action-provider_test.ts @@ -0,0 +1,240 @@ +/** + * Copyright 2025 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import * as assert from 'assert'; +import { beforeEach, describe, it } from 'node:test'; +import { setTimeout } from 'timers/promises'; +import { z } from 'zod'; +import { Action, defineAction } from '../src/action.js'; +import { + defineDynamicActionProvider, + isDynamicActionProvider, +} from '../src/dynamic-action-provider.js'; +import { initNodeFeatures } from '../src/node.js'; +import { Registry } from '../src/registry.js'; + +initNodeFeatures(); + +describe('dynamic action provider', () => { + let registry: Registry; + let tool1: Action; + let tool2: Action; + + beforeEach(() => { + registry = new Registry(); + tool1 = defineAction( + registry, + { name: 'tool1', actionType: 'tool' }, + async () => 'tool1' + ); + tool2 = defineAction( + registry, + { name: 'tool2', actionType: 'tool' }, + async () => 'tool2' + ); + }); + + it('gets a specific action', async () => { + let callCount = 0; + const dap = defineDynamicActionProvider(registry, 'my-dap', async () => { + callCount++; + return { + tool: [tool1, tool2], + }; + }); + + const action = await dap.getAction('tool', 'tool1'); + assert.strictEqual(action, tool1); + assert.strictEqual(callCount, 1); + }); + + it('lists action metadata', async () => { + let callCount = 0; + const dap = defineDynamicActionProvider(registry, 'my-dap', async () => { + callCount++; + return { + tool: [tool1, tool2], + }; + }); + + const metadata = await dap.listActionMetadata('tool', '*'); + assert.deepStrictEqual(metadata, [tool1.__action, tool2.__action]); + assert.strictEqual(callCount, 1); + }); + + it('caches the actions', async () => { + let callCount = 0; + const dap = defineDynamicActionProvider(registry, 'my-dap', async () => { + callCount++; + return { + tool: [tool1, tool2], + }; + }); + + let action = await dap.getAction('tool', 'tool1'); + assert.strictEqual(action, tool1); + assert.strictEqual(callCount, 1); + + // This should be cached + action = await dap.getAction('tool', 'tool2'); + assert.strictEqual(action, tool2); + assert.strictEqual(callCount, 1); + + const metadata = await dap.listActionMetadata('tool', '*'); + assert.deepStrictEqual(metadata, [tool1.__action, tool2.__action]); + assert.strictEqual(callCount, 1); + }); + + it('invalidates the cache', async () => { + let callCount = 0; + const dap = defineDynamicActionProvider(registry, 'my-dap', async () => { + callCount++; + return { + tool: [tool1, tool2], + }; + }); + + await dap.getAction('tool', 'tool1'); + assert.strictEqual(callCount, 1); + + dap.invalidateCache(); + + await dap.getAction('tool', 'tool2'); + assert.strictEqual(callCount, 2); + }); + + it('respects cache ttl', async () => { + let callCount = 0; + const dap = defineDynamicActionProvider( + registry, + { name: 'my-dap', cacheConfig: { ttlMillis: 10 } }, + async () => { + callCount++; + return { + tool: [tool1, tool2], + }; + } + ); + + await dap.getAction('tool', 'tool1'); + assert.strictEqual(callCount, 1); + + await setTimeout(20); + + await dap.getAction('tool', 'tool2'); + assert.strictEqual(callCount, 2); + }); + + it('lists actions with prefix', async () => { + let callCount = 0; + const tool3 = defineAction( + registry, + { name: 'other-tool', actionType: 'tool' }, + async () => 'other' + ); + const dap = defineDynamicActionProvider(registry, 'my-dap', async () => { + callCount++; + return { + tool: [tool1, tool2, tool3], + }; + }); + + const metadata = await dap.listActionMetadata('tool', 'tool*'); + assert.deepStrictEqual(metadata, [tool1.__action, tool2.__action]); + assert.strictEqual(callCount, 1); + }); + + it('lists actions with exact match', async () => { + let callCount = 0; + const dap = defineDynamicActionProvider(registry, 'my-dap', async () => { + callCount++; + return { + tool: [tool1, tool2], + }; + }); + + const metadata = await dap.listActionMetadata('tool', 'tool1'); + assert.deepStrictEqual(metadata, [tool1.__action]); + assert.strictEqual(callCount, 1); + }); + + it('handles concurrent requests', async () => { + let callCount = 0; + const dap = defineDynamicActionProvider(registry, 'my-dap', async () => { + callCount++; + await setTimeout(10); + return { + tool: [tool1, tool2], + }; + }); + + const [metadata1, metadata2] = await Promise.all([ + dap.listActionMetadata('tool', '*'), + dap.listActionMetadata('tool', '*'), + ]); + + assert.deepStrictEqual(metadata1, [tool1.__action, tool2.__action]); + assert.deepStrictEqual(metadata2, [tool1.__action, tool2.__action]); + assert.strictEqual(callCount, 1); + }); + + it('handles fetch errors', async () => { + let callCount = 0; + const dap = defineDynamicActionProvider(registry, 'my-dap', async () => { + callCount++; + if (callCount === 1) { + throw new Error('Fetch failed'); + } + return { + tool: [tool1, tool2], + }; + }); + + await assert.rejects(dap.listActionMetadata('tool', '*'), /Fetch failed/); + assert.strictEqual(callCount, 1); + + const metadata = await dap.listActionMetadata('tool', '*'); + assert.deepStrictEqual(metadata, [tool1.__action, tool2.__action]); + assert.strictEqual(callCount, 2); + }); + + it('returns metadata when run', async () => { + const dap = defineDynamicActionProvider(registry, 'my-dap', async () => { + return { + tool: [tool1, tool2], + }; + }); + + const result = await dap.run({ tool: [tool1, tool2] }); + assert.deepStrictEqual(result.result, { + tool: [tool1.__action, tool2.__action], + }); + }); + + it('identifies dynamic action providers', async () => { + const dap = defineDynamicActionProvider(registry, 'my-dap', async () => { + return {}; + }); + assert.ok(isDynamicActionProvider(dap)); + + const regularAction = defineAction( + registry, + { name: 'regular', actionType: 'tool' }, + async () => {} + ); + assert.ok(!isDynamicActionProvider(regularAction)); + }); +}); diff --git a/js/genkit/src/common.ts b/js/genkit/src/common.ts index 241e000d56..c1f1959521 100644 --- a/js/genkit/src/common.ts +++ b/js/genkit/src/common.ts @@ -138,6 +138,7 @@ export { type Action, type ActionContext, type ActionMetadata, + type DynamicActionProviderAction, type Flow, type FlowConfig, type FlowFn, diff --git a/js/genkit/src/genkit.ts b/js/genkit/src/genkit.ts index 77abd1e531..abe4f66d4a 100644 --- a/js/genkit/src/genkit.ts +++ b/js/genkit/src/genkit.ts @@ -103,6 +103,7 @@ import { GenkitError, Operation, ReflectionServer, + defineDynamicActionProvider, defineFlow, defineJsonSchema, defineSchema, @@ -115,6 +116,9 @@ import { setClientHeader, type Action, type ActionContext, + type DapConfig, + type DapFn, + type DynamicActionProviderAction, type FlowConfig, type FlowFn, type JSONSchema, @@ -241,6 +245,16 @@ export class Genkit implements HasRegistry { return dynamicTool(config, fn) as ToolAction; } + /** + * Defines and registers a dynamic action provider (e.g. mcp host) + */ + defineDynamicActionProvider( + config: DapConfig | string, + fn: DapFn + ): DynamicActionProviderAction { + return defineDynamicActionProvider(this.registry, config, fn); + } + /** * Defines and registers a schema from a Zod schema. * diff --git a/js/plugins/mcp/src/client/client.ts b/js/plugins/mcp/src/client/client.ts index 3a930de5bc..9f48739874 100644 --- a/js/plugins/mcp/src/client/client.ts +++ b/js/plugins/mcp/src/client/client.ts @@ -24,6 +24,7 @@ import { } from '@modelcontextprotocol/sdk/types.js'; import { GenkitError, + type DynamicActionProviderAction, type DynamicResourceAction, type ExecutablePrompt, type Genkit, @@ -113,6 +114,10 @@ export type McpClientOptions = { sessionId?: string; }; +export type McpClientOptionsWithCache = McpClientOptions & { + cacheTtlMillis?: number; +}; + /** * Represents a client connection to a single MCP (Model Context Protocol) server. * It handles the lifecycle of the connection (connect, disconnect, disable, re-enable, reconnect) @@ -123,6 +128,7 @@ export type McpClientOptions = { */ export class GenkitMcpClient { _server?: McpServerRef; + private _dynamicActionProvider: DynamicActionProviderAction | undefined; sessionId?: string; readonly name: string; @@ -152,6 +158,16 @@ export class GenkitMcpClient { this._initializeConnection(); } + set dynamicActionProvider(dap: DynamicActionProviderAction) { + this._dynamicActionProvider = dap; + } + + _invalidateDapCache(): void { + if (this._dynamicActionProvider) { + this._dynamicActionProvider.invalidateCache(); + } + } + get serverName(): string { return ( this.suppliedServerName ?? @@ -163,6 +179,7 @@ export class GenkitMcpClient { async updateRoots(roots: Root[]) { this.roots = roots; await this._server?.client.sendRootsListChanged(); + this._invalidateDapCache(); } /** @@ -188,6 +205,7 @@ export class GenkitMcpClient { if (this.roots) { await this.updateRoots(this.roots); } + this._invalidateDapCache(); } /** @@ -208,6 +226,7 @@ export class GenkitMcpClient { */ private async _connect() { if (this._server) await this._server.transport.close(); + this._invalidateDapCache(); logger.debug( `[MCP Client] Connecting MCP server '${this.serverName}' in client '${this.name}'.` ); @@ -249,6 +268,7 @@ export class GenkitMcpClient { transport, error, } as McpServerRef; + this._invalidateDapCache(); } /** @@ -262,6 +282,7 @@ export class GenkitMcpClient { ); await this._server.client.close(); this._server = undefined; + this._invalidateDapCache(); } } @@ -277,6 +298,7 @@ export class GenkitMcpClient { ); await this._disconnect(); this.disabled = true; + this._invalidateDapCache(); } } @@ -296,6 +318,7 @@ export class GenkitMcpClient { logger.debug(`[MCP Client] Reenabling MCP server in client '${this.name}'`); await this._initializeConnection(); this.disabled = !!this._server!.error; + this._invalidateDapCache(); } /** @@ -310,6 +333,7 @@ export class GenkitMcpClient { ); await this._disconnect(); await this._initializeConnection(); + this._invalidateDapCache(); } } diff --git a/js/plugins/mcp/src/client/host.ts b/js/plugins/mcp/src/client/host.ts index 298b0ee26f..2648db429d 100644 --- a/js/plugins/mcp/src/client/host.ts +++ b/js/plugins/mcp/src/client/host.ts @@ -16,6 +16,7 @@ import { Root } from '@modelcontextprotocol/sdk/types.js'; import { + type DynamicActionProviderAction, type DynamicResourceAction, type ExecutablePrompt, type Genkit, @@ -57,6 +58,24 @@ export interface McpHostOptions { roots?: Root[]; } +export type McpHostOptionsWithCache = Omit & { + /** + * A client name for this MCP host. This name is advertised to MCP Servers + * as the connecting client name. + */ + name: string; + + /** + * Cache TTL. The dynamic action provider has a cache for the available actions + * The default TTL is 3 seconds. + * The cache will automatically be invalidated if any connections change. + * Negative = no caching (expect noisy traces and slower resolution) + * Zero or undefined = use the default 3000 millis (3 seconds) + * Positive: The number of milliseconds to keep the cache. + */ + cacheTTLMillis?: number; +}; + /** Internal representation of client state for logging. */ interface ClientState { error?: { @@ -82,6 +101,7 @@ export class GenkitMcpHost { reject: (err: Error) => void; }[] = []; private _ready = false; + private _dynamicActionProvider: DynamicActionProviderAction | undefined; private roots: Root[] | undefined; rawToolResponses?: boolean; @@ -97,6 +117,16 @@ export class GenkitMcpHost { } } + set dynamicActionProvider(dap: DynamicActionProviderAction) { + this._dynamicActionProvider = dap; + } + + _invalidateCache(): void { + if (this._dynamicActionProvider) { + this._dynamicActionProvider.invalidateCache(); + } + } + /** * Returns a Promise that resolves when the host has attempted to connect * to all configured clients, or rejects if a critical error occurs during @@ -148,6 +178,7 @@ export class GenkitMcpHost { detail: `Details: ${e}`, }); } + this._invalidateCache(); } /** @@ -175,6 +206,7 @@ export class GenkitMcpHost { }); } delete this._clients[serverName]; + this._invalidateCache(); } /** @@ -198,6 +230,7 @@ export class GenkitMcpHost { `[MCP Host] Disabling MCP server '${serverName}' in host '${this.name}'` ); await client.disable(); + this._invalidateCache(); } /** @@ -224,6 +257,7 @@ export class GenkitMcpHost { detail: `Details: ${e}`, }); } + this._invalidateCache(); } /** @@ -251,6 +285,7 @@ export class GenkitMcpHost { detail: `Details: ${e}`, }); } + this._invalidateCache(); } /** @@ -290,6 +325,8 @@ export class GenkitMcpHost { this._readyListeners.pop()?.reject(err); } }); + + this._invalidateCache(); } /** @@ -446,6 +483,7 @@ export class GenkitMcpHost { for (const client of Object.values(this._clients)) { await client._disconnect(); } + this._invalidateCache(); } /** Helper method to track and log client errors. */ diff --git a/js/plugins/mcp/src/client/index.ts b/js/plugins/mcp/src/client/index.ts index 2032da820f..47e3557f69 100644 --- a/js/plugins/mcp/src/client/index.ts +++ b/js/plugins/mcp/src/client/index.ts @@ -18,11 +18,16 @@ import { SSEClientTransportOptions } from '@modelcontextprotocol/sdk/client/sse. import { StdioServerParameters } from '@modelcontextprotocol/sdk/client/stdio.js'; import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js'; import { GenkitMcpClient, McpClientOptions } from './client.js'; -import { GenkitMcpHost, McpHostOptions } from './host.js'; +import { + GenkitMcpHost, + McpHostOptions, + McpHostOptionsWithCache, +} from './host.js'; export { GenkitMcpClient, GenkitMcpHost }; export type { McpClientOptions, McpHostOptions, + McpHostOptionsWithCache, SSEClientTransportOptions, StdioServerParameters, Transport, diff --git a/js/plugins/mcp/src/index.ts b/js/plugins/mcp/src/index.ts index afad3a8cdd..90c8ab5e20 100644 --- a/js/plugins/mcp/src/index.ts +++ b/js/plugins/mcp/src/index.ts @@ -18,16 +18,23 @@ import type { Genkit } from 'genkit'; import { GenkitMcpClient, McpClientOptions, + McpClientOptionsWithCache, McpServerConfig, McpStdioServerConfig, } from './client/client.js'; -import { GenkitMcpHost, McpHostOptions } from './client/index.js'; +import { + GenkitMcpHost, + McpHostOptions, + McpHostOptionsWithCache, +} from './client/index.js'; import { GenkitMcpServer } from './server.js'; export { GenkitMcpClient, GenkitMcpHost, type McpClientOptions, + type McpClientOptionsWithCache, type McpHostOptions, + type McpHostOptionsWithCache, type McpServerConfig, type McpStdioServerConfig, }; @@ -62,11 +69,53 @@ export interface McpServerOptions { * * @param options Configuration for the MCP Client Host, including the definitions of MCP servers to connect to. * @returns A new instance of GenkitMcpHost. + * @deprecated Please use defineMcpHost instead. */ export function createMcpHost(options: McpHostOptions) { return new GenkitMcpHost(options); } +/** + * Creates an MCP Client Host that connects to one or more MCP servers. + * Each server is defined in the `mcpClients` option, where the key is a + * client-side name for the server and the value is the server's configuration. + * + * By default, all servers in the config will be attempted to connect unless + * their configuration includes `{disabled: true}`. + * + * ```ts + * const clientHost = defineMcpHost(ai, { + * name: "my-mcp-client-host", // Name for the host itself + * mcpServers: { + * // Each key is a name for this client/server configuration + * // Each value is an McpServerConfig object + * gitToolServer: { command: "uvx", args: ["mcp-server-git"] }, + * customApiServer: { url: "http://localhost:1234/mcp" } + * } + * }); + * ``` + * + * @param options Configuration for the MCP Client Host, including the definitions of MCP servers to connect to. + * @returns A new instance of GenkitMcpHost. + */ +export function defineMcpHost(ai: Genkit, options: McpHostOptionsWithCache) { + const mcpHost = new GenkitMcpHost(options); + const dap = ai.defineDynamicActionProvider( + { + name: options.name, + cacheConfig: { + ttlMillis: options.cacheTTLMillis, + }, + }, + async () => ({ + tool: await mcpHost.getActiveTools(ai), + resource: await mcpHost.getActiveResources(ai), + }) + ); + mcpHost.dynamicActionProvider = dap; + return mcpHost; +} + /** * Creates an MCP Client that connects to a single MCP server. * This is useful when you only need to interact with one MCP server, @@ -86,11 +135,61 @@ export function createMcpHost(options: McpHostOptions) { * @param options Configuration for the MCP Client, defining how it connects * to the MCP server and its behavior. * @returns A new instance of GenkitMcpClient. + * @deprecated Please use defineMcpClient instead. */ export function createMcpClient(options: McpClientOptions) { return new GenkitMcpClient(options); } +/** + * Defines an MCP Client that connects to a single MCP server. + * This is useful when you only need to interact with one MCP server, + * or if you want to manage client instances individually. + * + * ```ts + * const client = defineMcpClient(ai, { + * name: "mySingleMcpClient", // A name for this client instance + * command: "npx", // Example: Launching a local server + * args: ["-y", "@modelcontextprotocol/server-everything", "/path/to/allowed/dir"], + * }); + * + * // To get tools from this client: + * // const tools = await client.getActiveTools(ai); + * + * // Or in a generate call you can use: + * ai.generate({ + prompt: ``, + tools: ['mySingleMcpClient:tool/*'], + }); + * ``` + * + * @param options Configuration for the MCP Client, defining how it connects + * to the MCP server and its behavior. + * @returns A new instance of GenkitMcpClient. + */ +export function defineMcpClient( + ai: Genkit, + options: McpClientOptionsWithCache +) { + const mcpClient = new GenkitMcpClient(options); + const dap = ai.defineDynamicActionProvider( + { + name: options.name, + cacheConfig: { + ttlMillis: options.cacheTtlMillis, + }, + }, + async () => { + return { + tool: await mcpClient.getActiveTools(ai), + resource: await mcpClient.getActiveResources(ai), + }; + } + ); + mcpClient.dynamicActionProvider = dap; + return mcpClient; +} + /** * Creates an MCP server based on the supplied Genkit instance. All tools and prompts * will be automatically converted to MCP compatibility. diff --git a/js/plugins/mcp/src/server.ts b/js/plugins/mcp/src/server.ts index 31cfa9d001..ad454cab75 100644 --- a/js/plugins/mcp/src/server.ts +++ b/js/plugins/mcp/src/server.ts @@ -444,7 +444,7 @@ function toMcpResourceMessage( if (!url.startsWith('data:')) throw new GenkitError({ status: 'UNIMPLEMENTED', - message: `[MCP Server] MCP prompt messages only support base64 data images.`, + message: `[MCP Server] MCP resource messages only support base64 data images.`, }); const mimeType = contentType || url.substring(url.indexOf(':')! + 1, url.indexOf(';')); @@ -455,7 +455,7 @@ function toMcpResourceMessage( } else { throw new GenkitError({ status: 'UNIMPLEMENTED', - message: `[MCP Server] MCP prompt messages only support media and text parts.`, + message: `[MCP Server] MCP resource messages only support media and text parts.`, }); } }); diff --git a/js/testapps/mcp/src/index.ts b/js/testapps/mcp/src/index.ts index ca6ebebe6f..e4d19223c8 100644 --- a/js/testapps/mcp/src/index.ts +++ b/js/testapps/mcp/src/index.ts @@ -15,7 +15,7 @@ */ import { googleAI } from '@genkit-ai/googleai'; -import { createMcpHost } from '@genkit-ai/mcp'; +import { defineMcpHost } from '@genkit-ai/mcp'; import { genkit, z } from 'genkit'; import { logger } from 'genkit/logging'; import path from 'path'; @@ -50,7 +50,7 @@ export const ai = genkit({ logger.setLogLevel('debug'); // Set the logging level to debug for detailed output -export const mcpHost = createMcpHost({ +export const mcpHost = defineMcpHost(ai, { name: 'test-mcp-manager', mcpServers: { 'git-client': { @@ -81,6 +81,15 @@ ai.defineFlow('git-commits', async (q) => { return text; }); +ai.defineFlow('dynamic-git-commits', async (q) => { + const { text } = await ai.generate({ + prompt: `summarize last 5 commits in '${path.resolve(process.cwd(), '../../..')}'`, + tools: ['test-mcp-manager:tool/*'], // All the tools + }); + + return text; +}); + ai.defineFlow('get-file', async (q) => { const { text } = await ai.generate({ prompt: `summarize contexts of hello-world.txt (in '${process.cwd()}/test-workspace')`, @@ -90,6 +99,70 @@ ai.defineFlow('get-file', async (q) => { return text; }); +ai.defineFlow('dynamic-get-file', async (q) => { + const { text } = await ai.generate({ + prompt: `summarize contexts of hello-world.txt (in '${process.cwd()}/test-workspace')`, + tools: ['test-mcp-manager:tool/fs/read_file'], // Just this one tool + }); + + return text; +}); + +ai.defineFlow('dynamic-prefix-tool', async (q) => { + // When specifying tools you can use a prefix - so + // test-mcp-manager:tool/fs/read_* or + // test-mcp-manager:tool/fs/* will use only the tools whose + // names match the prefix. + const { text } = await ai.generate({ + prompt: `summarize contexts of hello-world.txt (in '${process.cwd()}/test-workspace')`, + tools: ['test-mcp-manager:tool/fs/read_*'], // Just read tools from the fs/ + }); + + return text; +}); + +ai.defineFlow('dynamic-disable-enable', async (q) => { + // This shows that the dap cache is invalidated any time + // we change something with the mcpHost config. + const { text: text1 } = await ai.generate({ + prompt: `summarize contexts of hello-world.txt (in '${process.cwd()}/test-workspace')`, + tools: ['test-mcp-manager:tool/fs/read_file'], // Just this one tool + }); + + // Now disable fs to show that we invalidate the dap cache + await mcpHost.disable('fs'); + let text2: string; + try { + // This should fail because the fs/read_file tool is not available + // after disabling the mcp client providing it. + const { text } = await ai.generate({ + prompt: `summarize contexts of hello-world.txt (in '${process.cwd()}/test-workspace')`, + tools: ['test-mcp-manager:tool/fs/read_file'], // Just this one tool + }); + text2 = + 'ERROR! This should have failed to find the tool but succeeded instead: ' + + text; + } catch (e: any) { + text2 = e.message; + } + + // If we re-enable the fs it will succeed. + await mcpHost.enable('fs'); + await mcpHost.reconnect('fs'); + const { text: text3 } = await ai.generate({ + prompt: `summarize contexts of hello-world.txt (in '${process.cwd()}/test-workspace')`, + tools: ['test-mcp-manager:tool/fs/read_file'], // Just this one tool + }); + return ( + 'Original:
' + + text1 + + '
After Disable:
' + + text2 + + '
After Enable:
' + + text3 + ); +}); + ai.defineFlow('test-resource', async (q) => { const { text } = await ai.generate({ prompt: [ @@ -102,6 +175,30 @@ ai.defineFlow('test-resource', async (q) => { return text; }); +ai.defineFlow('dynamic-test-resources', async (q) => { + const { text } = await ai.generate({ + prompt: [ + { text: 'analyze this: ' }, + { resource: { uri: 'test://static/resource/1' } }, + ], + resources: ['test-mcp-manager:resource/*'], + }); + + return text; +}); + +ai.defineFlow('dynamic-test-one-resource', async (q) => { + const { text } = await ai.generate({ + prompt: [ + { text: 'analyze this: ' }, + { resource: { uri: 'test://static/resource/1' } }, + ], + resources: ['test-mcp-manager:resource/everything/Resource 1'], + }); + + return text; +}); + ai.defineFlow('update-file', async (q) => { const { text } = await ai.generate({ prompt: `Improve hello-world.txt (in '${process.cwd()}/test-workspace') by rewriting the text, making it longer, just do it, use your imagination.`,