From 080f9fba1b5def1a83c59530c6c9cc4c8f7d0449 Mon Sep 17 00:00:00 2001
From: Gad Grandez <gadgrandez@gmail.com>
Date: Sat, 5 Apr 2025 15:40:12 -0500
Subject: [PATCH] feat(js/plugins/mcp): allow passing context data to the
 server for tool execution

---
 js/plugins/mcp/README.md           |  1 +
 js/plugins/mcp/src/client/tools.ts | 30 +++++++++++++++++++++++-------
 js/plugins/mcp/src/index.ts        |  2 ++
 js/plugins/mcp/src/server.ts       | 14 +++++++++++++-
 4 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/js/plugins/mcp/README.md b/js/plugins/mcp/README.md
index 248592c563..b7e3d36182 100644
--- a/js/plugins/mcp/README.md
+++ b/js/plugins/mcp/README.md
@@ -54,6 +54,7 @@ Most MCP servers are built to run as spawned processes on the same machine using
   - **`serverWebsocketUrl`: The URL of a remote server to connect to using the WebSocket MCP transport.
   - **`transport`**: An existing MCP transport object for connecting to the server.
 - **`rawToolResponses`**: (optional) A boolean flag. If `true`, tool responses are returned in their raw MCP format; otherwise, they are processed for Genkit compatibility.
+- **`sendGenkitContext`**: (optional) A boolean flag. If `true`, tool calls include context data from the current execution.
 
 ### Using MCP Actions
 
diff --git a/js/plugins/mcp/src/client/tools.ts b/js/plugins/mcp/src/client/tools.ts
index 04778bbdc8..f0774ce3f4 100644
--- a/js/plugins/mcp/src/client/tools.ts
+++ b/js/plugins/mcp/src/client/tools.ts
@@ -16,6 +16,7 @@
 
 import type { Client } from '@modelcontextprotocol/sdk/client/index.js' with { 'resolution-mode': 'import' };
 import type {
+  CallToolRequest,
   CallToolResult,
   Tool,
 } from '@modelcontextprotocol/sdk/types.js' with { 'resolution-mode': 'import' };
@@ -59,15 +60,30 @@ function registerTool(
       inputJsonSchema: tool.inputSchema as JSONSchema7,
       outputSchema: z.any(),
     },
-    async (args) => {
-      logger.debug(
-        `[@genkit-ai/mcp] Calling MCP tool ${params.name}/${tool.name} with arguments`,
-        JSON.stringify(args)
-      );
-      const result = await client.callTool({
+    async (args, { context }) => {
+      let callToolOptions: CallToolRequest['params'] = {
         name: tool.name,
         arguments: args,
-      });
+      };
+      if (params.sendGenkitContext) {
+        callToolOptions = {
+          ...callToolOptions,
+          _meta: {
+            context,
+          },
+        };
+        logger.debug(
+          `[@genkit-ai/mcp] Calling MCP tool ${params.name}/${tool.name} with arguments and Genkit context`,
+          JSON.stringify({ args, context })
+        );
+      } else {
+        logger.debug(
+          `[@genkit-ai/mcp] Calling MCP tool ${params.name}/${tool.name} with arguments`,
+          JSON.stringify(args)
+        );
+      }
+
+      const result = await client.callTool(callToolOptions);
       logger.debug(
         `MCP tool ${tool.name} result:`,
         JSON.stringify(result, null, 2)
diff --git a/js/plugins/mcp/src/index.ts b/js/plugins/mcp/src/index.ts
index c950cf037d..b362c671bf 100644
--- a/js/plugins/mcp/src/index.ts
+++ b/js/plugins/mcp/src/index.ts
@@ -38,6 +38,8 @@ export interface McpClientOptions {
   serverWebsocketUrl?: string | URL;
   /** Return tool responses in raw MCP form instead of processing them for Genkit compatibility. */
   rawToolResponses?: boolean;
+  /** Send the Genkit context to the server. */
+  sendGenkitContext?: boolean;
 }
 
 async function transportFrom(params: McpClientOptions): Promise<Transport> {
diff --git a/js/plugins/mcp/src/server.ts b/js/plugins/mcp/src/server.ts
index 55599f6170..7ee34ba87d 100644
--- a/js/plugins/mcp/src/server.ts
+++ b/js/plugins/mcp/src/server.ts
@@ -134,7 +134,19 @@ export class GenkitMcpServer {
         status: 'NOT_FOUND',
         message: `Tried to call tool '${req.params.name}' but it could not be found.`,
       });
-    const result = await tool(req.params.arguments);
+
+    const clientContext = req.params._meta?.context;
+
+    let result: any;
+
+    if (clientContext) {
+      const currentContext = this.ai.currentContext() || {};
+      result = await tool(req.params.arguments, {
+        context: { ...currentContext, ...clientContext },
+      });
+    } else {
+      result = await tool(req.params.arguments);
+    }
     return { content: [{ type: 'text', text: JSON.stringify(result) }] };
   }