From 99de0f07433190e17734742e8aaaeb20d3deb03f Mon Sep 17 00:00:00 2001 From: Michael Bleigh Date: Thu, 12 Jun 2025 17:54:36 -0700 Subject: [PATCH] feat(js/mcp): Adds getActivePrompts() to host. --- js/plugins/mcp/src/client/client.ts | 61 ++++++++++++---- js/plugins/mcp/src/client/host.ts | 26 +++++-- js/plugins/mcp/src/util/prompts.ts | 25 +++++++ js/plugins/mcp/src/util/resources.ts | 6 +- js/plugins/mcp/src/util/tools.ts | 2 +- .../tests/{manager_test.ts => host_test.ts} | 71 +++++++++++++++++-- 6 files changed, 162 insertions(+), 29 deletions(-) rename js/plugins/mcp/tests/{manager_test.ts => host_test.ts} (83%) diff --git a/js/plugins/mcp/src/client/client.ts b/js/plugins/mcp/src/client/client.ts index 71e7e00c46..296de70a9b 100644 --- a/js/plugins/mcp/src/client/client.ts +++ b/js/plugins/mcp/src/client/client.ts @@ -26,7 +26,12 @@ import { ToolAction, } from 'genkit'; import { logger } from 'genkit/logging'; -import { fetchDynamicTools, getExecutablePrompt, transportFrom } from '../util'; +import { + fetchAllPrompts, + fetchDynamicTools, + getExecutablePrompt, + transportFrom, +} from '../util'; export type { SSEClientTransportOptions, StdioServerParameters, Transport }; interface McpServerRef { @@ -39,12 +44,6 @@ export interface McpServerControls { /** when true, the server will be stopped and its registered components will * not appear in lists/plugins/etc */ disabled?: boolean; - /** - * If true, tool responses from the MCP server will be returned in their raw - * MCP format. Otherwise (default), they are processed and potentially - * simplified for better compatibility with Genkit's typical data structures. - */ - rawToolResponses?: boolean; } export type McpStdioServerConfig = StdioServerParameters; @@ -76,14 +75,25 @@ export type McpServerConfig = ( * Configuration options for an individual `GenkitMcpClient` instance. * This defines how the client connects to a single MCP server and how it behaves. */ -export type McpClientOptions = McpServerConfig & { - /** Name for this server configuration */ +export type McpClientOptions = { + /** Client name to advertise to the server. */ name: string; + /** Name for the server, defaults to the server's advertised name. */ + serverName?: string; + /** * An optional version number for this client. This is primarily for logging * and identification purposes. Defaults to '1.0.0'. */ version?: string; + /** + * If true, tool responses from the MCP server will be returned in their raw + * MCP format. Otherwise (default), they are processed and potentially + * simplified for better compatibility with Genkit's typical data structures. + */ + rawToolResponses?: boolean; + /** The server configuration to connect. */ + mcpServer: McpServerConfig; }; /** @@ -98,6 +108,7 @@ export class GenkitMcpClient { _server?: McpServerRef; readonly name: string; + readonly suppliedServerName?: string; private version: string; private serverConfig: McpServerConfig; private rawToolResponses?: boolean; @@ -111,14 +122,23 @@ export class GenkitMcpClient { constructor(options: McpClientOptions) { this.name = options.name; + this.suppliedServerName = options.serverName; this.version = options.version || '1.0.0'; - this.serverConfig = { ...options }; + this.serverConfig = options.mcpServer; this.rawToolResponses = !!options.rawToolResponses; - this.disabled = !!options.disabled; + this.disabled = !!options.mcpServer.disabled; this._initializeConnection(); } + get serverName(): string { + return ( + this.suppliedServerName ?? + this._server?.client.getServerVersion()?.name ?? + 'unknown-server' + ); + } + /** * Sets up a connection based on a provided map of server configurations. * - Reconnects existing servers if their configuration appears to have @@ -266,12 +286,12 @@ export class GenkitMcpClient { let tools: ToolAction[] = []; if (this._server) { - const capabilities = await this._server.client.getServerCapabilities(); + const capabilities = this._server.client.getServerCapabilities(); if (capabilities?.tools) tools.push( ...(await fetchDynamicTools(ai, this._server.client, { rawToolResponses: this.rawToolResponses, - serverName: this.name, + serverName: this.serverName, name: this.name, })) ); @@ -280,6 +300,21 @@ export class GenkitMcpClient { return tools; } + async getActivePrompts( + ai: Genkit, + options?: PromptGenerateOptions + ): Promise { + if (this._server?.client.getServerCapabilities()?.prompts) { + return fetchAllPrompts(this._server.client, { + ai, + serverName: this.serverName, + name: this.name, + options, + }); + } + return []; + } + /** * Get the specified prompt as an `ExecutablePrompt` available through this * client. If no such prompt is found, return undefined. diff --git a/js/plugins/mcp/src/client/host.ts b/js/plugins/mcp/src/client/host.ts index ac0eafe804..62fe418b36 100644 --- a/js/plugins/mcp/src/client/host.ts +++ b/js/plugins/mcp/src/client/host.ts @@ -26,9 +26,8 @@ import { GenkitMcpClient, McpServerConfig } from './client.js'; export interface McpHostOptions { /** - * An optional name for this MCP host. This name is primarily for - * logging and identification purposes within Genkit. - * Defaults to 'genkitx-mcp'. + * An optional client name for this MCP host. This name is advertised to MCP Servers + * as the connecting client name. Defaults to 'genkit-mcp'. */ name?: string; /** @@ -72,7 +71,7 @@ export class GenkitMcpHost { private _ready = false; constructor(options: McpHostOptions) { - this.name = options.name || 'genkitx-mcp'; + this.name = options.name || 'genkit-mcp'; if (options.mcpServers) { this.updateServers(options.mcpServers); @@ -119,7 +118,11 @@ export class GenkitMcpHost { `[MCP Host] Connecting to MCP server '${serverName}' in host '${this.name}'.` ); try { - const client = new GenkitMcpClient({ ...config, name: serverName }); + const client = new GenkitMcpClient({ + name: this.name, + serverName: serverName, + mcpServer: config, + }); this._clients[serverName] = client; } catch (e) { this.setError(serverName, { @@ -323,6 +326,19 @@ export class GenkitMcpHost { return allTools; } + async getActivePrompts(ai: Genkit): Promise { + await this.ready(); + let allPrompts: ExecutablePrompt[] = []; + + for (const serverName in this._clients) { + const client = this._clients[serverName]; + if (client.isEnabled() && !this.hasError(serverName)) { + allPrompts.push(...(await client.getActivePrompts(ai))); + } + } + return allPrompts; + } + /** * Get the specified prompt as an `ExecutablePrompt` available through the * specified server. If no such prompt is found, return undefined. diff --git a/js/plugins/mcp/src/util/prompts.ts b/js/plugins/mcp/src/util/prompts.ts index 95b68ebab2..98423c3a45 100644 --- a/js/plugins/mcp/src/util/prompts.ts +++ b/js/plugins/mcp/src/util/prompts.ts @@ -184,3 +184,28 @@ export async function getExecutablePrompt( } return undefined; } + +export async function fetchAllPrompts( + client: Client, + params: { + name: string; + serverName: string; + ai: Genkit; + options?: PromptGenerateOptions; + } +): Promise { + let cursor: string | undefined; + const out: ExecutablePrompt[] = []; + + while (true) { + const { nextCursor, prompts } = await client.listPrompts({ cursor }); + for (const p of prompts) { + out.push( + createExecutablePrompt(client, p, { ...params, promptName: p.name }) + ); + } + cursor = nextCursor; + if (!cursor) break; + } + return out; +} diff --git a/js/plugins/mcp/src/util/resources.ts b/js/plugins/mcp/src/util/resources.ts index 158db3e967..1b274b2183 100644 --- a/js/plugins/mcp/src/util/resources.ts +++ b/js/plugins/mcp/src/util/resources.ts @@ -25,7 +25,7 @@ function listResourcesTool( params: { asDynamicTool?: boolean } ): ToolAction { const actionMetadata = { - name: `mcp/list_resources`, + name: `${host.name}/list_resources`, description: `list all available resources`, inputSchema: z .object({ @@ -87,7 +87,7 @@ function readResourceTool( params: { asDynamicTool?: boolean } ) { const actionMetadata = { - name: `mcp/read_resource`, + name: `${host.name}/read_resource`, description: `this tool can read resources`, inputSchema: z.object({ server: z.string(), @@ -95,7 +95,7 @@ function readResourceTool( }), }; const fn = async ({ server, uri }) => { - const client = host.activeClients.find((c) => c.name === server); + const client = host.activeClients.find((c) => c.serverName === server); if (!client) { throw new GenkitError({ status: 'NOT_FOUND', diff --git a/js/plugins/mcp/src/util/tools.ts b/js/plugins/mcp/src/util/tools.ts index 6edb83817a..af6bb41552 100644 --- a/js/plugins/mcp/src/util/tools.ts +++ b/js/plugins/mcp/src/util/tools.ts @@ -106,7 +106,7 @@ function createDynamicTool( ); return ai.dynamicTool( { - name: `${params.name}/${tool.name}`, + name: `${params.serverName}/${tool.name}`, description: tool.description || '', inputJsonSchema: tool.inputSchema as JSONSchema7, outputSchema: z.any(), diff --git a/js/plugins/mcp/tests/manager_test.ts b/js/plugins/mcp/tests/host_test.ts similarity index 83% rename from js/plugins/mcp/tests/manager_test.ts rename to js/plugins/mcp/tests/host_test.ts index 6386466190..27fa612293 100644 --- a/js/plugins/mcp/tests/manager_test.ts +++ b/js/plugins/mcp/tests/host_test.ts @@ -210,9 +210,7 @@ describe('createMcpHost', () => { }, }); - fakeTransport.prompts.push({ - name: 'testPrompt', - }); + // Note: fakeTransport.prompts.push({ name: 'testPrompt' }); is moved to specific tests fakeTransport.getPromptResult = { messages: [ { @@ -230,7 +228,56 @@ describe('createMcpHost', () => { clientHost?.close(); }); + it('should list active prompts', async () => { + // Initially no prompts + assert.deepStrictEqual(await clientHost.getActivePrompts(ai), []); + + // Add a prompt to the first transport + fakeTransport.prompts.push({ + name: 'testPrompt1', + }); + let activePrompts = await clientHost.getActivePrompts(ai); + assert.strictEqual(activePrompts.length, 1); + assert.deepStrictEqual(await activePrompts[0].render(), { + messages: [ + { + role: 'user', + content: [ + { + text: 'prompt says: hello', + }, + ], + }, + ], + }); + + // Add a second transport with another prompt + const fakeTransport2 = new FakeTransport(); + fakeTransport2.prompts.push({ + name: 'testPrompt2', + }); + await clientHost.connect('test-server-2', { + transport: fakeTransport2, + }); + + activePrompts = await clientHost.getActivePrompts(ai); + assert.strictEqual(activePrompts.length, 2); + + // Disable the first server + await clientHost.disable('test-server'); + activePrompts = await clientHost.getActivePrompts(ai); + assert.strictEqual(activePrompts.length, 1); + + // Enable the first server again + await clientHost.enable('test-server'); + activePrompts = await clientHost.getActivePrompts(ai); + assert.strictEqual(activePrompts.length, 2); + }); + it('should execute prompt', async () => { + fakeTransport.prompts.push({ + name: 'testPrompt', + }); const prompt = await clientHost.getPrompt( ai, 'test-server', @@ -249,6 +296,9 @@ describe('createMcpHost', () => { }); it('should render prompt', async () => { + fakeTransport.prompts.push({ + name: 'testPrompt', + }); const prompt = await clientHost.getPrompt( ai, 'test-server', @@ -266,6 +316,9 @@ describe('createMcpHost', () => { }); it('should stream prompt', async () => { + fakeTransport.prompts.push({ + name: 'testPrompt', + }); const prompt = await clientHost.getPrompt( ai, 'test-server', @@ -343,20 +396,24 @@ describe('createMcpHost', () => { (await clientHost.getActiveTools(ai, { resourceTools: true })).map( (t) => t.__action.name ), - ['test-server/testTool', 'mcp/list_resources', 'mcp/read_resource'] + [ + 'test-server/testTool', + 'test-mcp-host/list_resources', + 'test-mcp-host/read_resource', + ] ); }); it('should list resources and templates', async () => { const listResourcesTool = ( await clientHost.getActiveTools(ai, { resourceTools: true }) - ).find((t) => t.__action.name === 'mcp/list_resources'); + ).find((t) => t.__action.name === 'test-mcp-host/list_resources'); assert.ok(listResourcesTool); const response = await listResourcesTool({}); assert.deepStrictEqual(response, { servers: { - 'test-server': { + 'test-mcp-host': { resources: [ { description: 'test resource', @@ -388,7 +445,7 @@ describe('createMcpHost', () => { const readResourceTool = ( await clientHost.getActiveTools(ai, { resourceTools: true }) - ).find((t) => t.__action.name === 'mcp/read_resource'); + ).find((t) => t.__action.name === 'test-mcp-host/read_resource'); assert.ok(readResourceTool); const response = await readResourceTool({