Only load capabilities advertised by the server. (#1425)

firebase · Nov 27, 2024 · 96957cd · 96957cd
1 parent 054f0af
commit 96957cd
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/js/plugins/mcp/src/client/resources.ts b/js/plugins/mcp/src/client/resources.ts
@@ -22,7 +22,7 @@ import type {
 import { Genkit, z } from 'genkit';
 import type { McpClientOptions } from '../index.js';
 
-export function registerResourceTools(
+export async function registerResourceTools(
   ai: Genkit,
   client: Client,
   params: McpClientOptions

diff --git a/js/plugins/mcp/src/index.ts b/js/plugins/mcp/src/index.ts
@@ -72,11 +72,14 @@ export function mcpClient(params: McpClientOptions) {
       { capabilities: {} }
     );
     await client.connect(transport);
-    await Promise.all([
-      registerAllTools(ai, client, params),
-      registerAllPrompts(ai, client, params),
-      registerResourceTools(ai, client, params),
-    ]);
+    const capabilties = await client.getServerCapabilities();
+    const promises: Promise<any>[] = [];
+    if (capabilties?.tools) promises.push(registerAllTools(ai, client, params));
+    if (capabilties?.prompts)
+      promises.push(registerAllPrompts(ai, client, params));
+    if (capabilties?.resources)
+      promises.push(registerResourceTools(ai, client, params));
+    await Promise.all(promises);
   });
 }