@mastra/mcp 0.5.0-alpha.3 → 0.5.0-alpha.5

package/src/client.test.ts

@@ -8,7 +8,7 @@ import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
 import { describe, it, expect, beforeEach, afterEach } from 'vitest';
 import { z } from 'zod';
 
-import { MastraMCPClient } from './client.js';
+import { InternalMastraMCPClient } from './client.js';
 
 async function setupTestServer(withSessionManagement: boolean) {
   const httpServer: HttpServer = createServer();
@@ -62,12 +62,12 @@ describe('MastraMCPClient with Streamable HTTP', () => {
     serverTransport: StreamableHTTPServerTransport;
     baseUrl: URL;
   };
-  let client: MastraMCPClient;
+  let client: InternalMastraMCPClient;
 
   describe('Stateless Mode', () => {
     beforeEach(async () => {
       testServer = await setupTestServer(false);
-      client = new MastraMCPClient({
+      client = new InternalMastraMCPClient({
         name: 'test-stateless-client',
         server: {
           url: testServer.baseUrl,
@@ -99,7 +99,7 @@ describe('MastraMCPClient with Streamable HTTP', () => {
   describe('Stateful Mode', () => {
     beforeEach(async () => {
       testServer = await setupTestServer(true);
-      client = new MastraMCPClient({
+      client = new InternalMastraMCPClient({
         name: 'test-stateful-client',
         server: {
           url: testServer.baseUrl,

package/src/client.ts

@@ -1,6 +1,7 @@
 import { MastraBase } from '@mastra/core/base';
+import type { RuntimeContext } from '@mastra/core/di';
 import { createTool } from '@mastra/core/tools';
-import { isZodType, resolveSerializedZodOutput } from '@mastra/core/utils';
+import { isZodType } from '@mastra/core/utils';
 import { Client } from '@modelcontextprotocol/sdk/client/index.js';
 import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
 import type { SSEClientTransportOptions } from '@modelcontextprotocol/sdk/client/sse.js';
@@ -14,9 +15,9 @@ import type { ClientCapabilities, LoggingLevel } from '@modelcontextprotocol/sdk
 import { CallToolResultSchema, ListResourcesResultSchema } from '@modelcontextprotocol/sdk/types.js';
 
 import { asyncExitHook, gracefulExit } from 'exit-hook';
-import jsonSchemaToZod from 'json-schema-to-zod';
-import type { JsonSchema } from 'json-schema-to-zod';
 import { z } from 'zod';
+import { jsonSchemaObjectToZodRawShape } from 'zod-from-json-schema';
+import type { JSONSchema } from 'zod-from-json-schema';
 
 // Re-export MCP SDK LoggingLevel for convenience
 export type { LoggingLevel } from '@modelcontextprotocol/sdk/types.js';
@@ -27,6 +28,7 @@ export interface LogMessage {
   timestamp: Date;
   serverName: string;
   details?: Record<string, any>;
+  runtimeContext?: RuntimeContext | null;
 }
 
 export type LogHandler = (logMessage: LogMessage) => void;
@@ -107,6 +109,7 @@ export class InternalMastraMCPClient extends MastraBase {
   private enableServerLogs?: boolean;
   private serverConfig: MastraMCPServerDefinition;
   private transport?: Transport;
+  private currentOperationContext: RuntimeContext | null = null;
 
   constructor({
     name,
@@ -159,6 +162,7 @@ export class InternalMastraMCPClient extends MastraBase {
         timestamp: new Date(),
         serverName: this.name,
         details,
+        runtimeContext: this.currentOperationContext,
       });
     }
   }
@@ -323,11 +327,35 @@ export class InternalMastraMCPClient extends MastraBase {
   }
 
   private convertInputSchema(
-    inputSchema: Awaited<ReturnType<Client['listTools']>>['tools'][0]['inputSchema'] | JsonSchema,
+    inputSchema: Awaited<ReturnType<Client['listTools']>>['tools'][0]['inputSchema'] | JSONSchema,
   ): z.ZodType {
-    return isZodType(inputSchema)
-      ? inputSchema
-      : resolveSerializedZodOutput(jsonSchemaToZod(inputSchema as JsonSchema));
+    if (isZodType(inputSchema)) {
+      return inputSchema;
+    }
+
+    try {
+      // Assuming inputSchema is a JSONSchema object type for tool inputs
+      const rawShape = jsonSchemaObjectToZodRawShape(inputSchema as JSONSchema);
+      return z.object(rawShape); // Wrap the raw shape to return a ZodType (object)
+    } catch (error: unknown) {
+      let errorDetails: string | undefined;
+      if (error instanceof Error) {
+        errorDetails = error.stack;
+      } else {
+        // Attempt to stringify, fallback to String()
+        try {
+          errorDetails = JSON.stringify(error);
+        } catch {
+          errorDetails = String(error);
+        }
+      }
+      this.log('error', 'Failed to convert JSON schema to Zod schema using zodFromJsonSchema', {
+        error: errorDetails,
+        originalJsonSchema: inputSchema,
+      });
+
+      throw new Error(errorDetails);
+    }
   }
 
   async tools() {
@@ -336,38 +364,49 @@ export class InternalMastraMCPClient extends MastraBase {
     const toolsRes: Record<string, any> = {};
     tools.forEach(tool => {
       this.log('debug', `Processing tool: ${tool.name}`);
+      try {
+        const mastraTool = createTool({
+          id: `${this.name}_${tool.name}`,
+          description: tool.description || '',
+          inputSchema: this.convertInputSchema(tool.inputSchema),
+          execute: async ({ context, runtimeContext }: { context: any; runtimeContext?: RuntimeContext | null }) => {
+            const previousContext = this.currentOperationContext;
+            this.currentOperationContext = runtimeContext || null; // Set current context
+            try {
+              this.log('debug', `Executing tool: ${tool.name}`, { toolArgs: context });
+              const res = await this.client.callTool(
+                {
+                  name: tool.name,
+                  arguments: context,
+                },
+                CallToolResultSchema,
+                {
+                  timeout: this.timeout,
+                },
+              );
+              this.log('debug', `Tool executed successfully: ${tool.name}`);
+              return res;
+            } catch (e) {
+              this.log('error', `Error calling tool: ${tool.name}`, {
+                error: e instanceof Error ? e.stack : JSON.stringify(e, null, 2),
+                toolArgs: context,
+              });
+              throw e;
+            } finally {
+              this.currentOperationContext = previousContext; // Restore previous context
+            }
+          },
+        });
 
-      const mastraTool = createTool({
-        id: `${this.name}_${tool.name}`,
-        description: tool.description || '',
-        inputSchema: this.convertInputSchema(tool.inputSchema),
-        execute: async ({ context }) => {
-          try {
-            this.log('debug', `Executing tool: ${tool.name}`, { toolArgs: context });
-            const res = await this.client.callTool(
-              {
-                name: tool.name,
-                arguments: context,
-              },
-              CallToolResultSchema,
-              {
-                timeout: this.timeout,
-              },
-            );
-            this.log('debug', `Tool executed successfully: ${tool.name}`);
-            return res;
-          } catch (e) {
-            this.log('error', `Error calling tool: ${tool.name}`, {
-              error: e instanceof Error ? e.stack : JSON.stringify(e, null, 2),
-              toolArgs: context,
-            });
-            throw e;
-          }
-        },
-      });
-
-      if (tool.name) {
-        toolsRes[tool.name] = mastraTool;
+        if (tool.name) {
+          toolsRes[tool.name] = mastraTool;
+        }
+      } catch (toolCreationError: unknown) {
+        // Catch errors during tool creation itself (e.g., if createTool has issues)
+        this.log('error', `Failed to create Mastra tool wrapper for MCP tool: ${tool.name}`, {
+          error: toolCreationError instanceof Error ? toolCreationError.stack : String(toolCreationError),
+          mcpToolDefinition: tool,
+        });
       }
     });
 

package/src/configuration.test.ts

@@ -1,6 +1,11 @@
 import { spawn } from 'child_process';
 import path from 'path';
+import { openai } from '@ai-sdk/openai';
+import { Agent } from '@mastra/core/agent';
+import { RuntimeContext } from '@mastra/core/di';
 import { describe, it, expect, beforeEach, afterEach, afterAll, beforeAll, vi } from 'vitest';
+import { allTools, mcpServerName } from './__fixtures__/fire-crawl-complex-schema';
+import type { LogHandler, LogMessage } from './client';
 import { MCPClient } from './configuration';
 
 vi.setConfig({ testTimeout: 80000, hookTimeout: 80000 });
@@ -8,6 +13,7 @@ vi.setConfig({ testTimeout: 80000, hookTimeout: 80000 });
 describe('MCPClient', () => {
   let mcp: MCPClient;
   let weatherProcess: ReturnType<typeof spawn>;
+  let clients: MCPClient[] = [];
 
   beforeAll(async () => {
     // Start the weather SSE server
@@ -50,11 +56,16 @@ describe('MCPClient', () => {
         },
       },
     });
+    clients.push(mcp);
   });
 
   afterEach(async () => {
     // Clean up any connected clients
     await mcp.disconnect();
+    const index = clients.indexOf(mcp);
+    if (index > -1) {
+      clients.splice(index, 1);
+    }
   });
 
   afterAll(async () => {
@@ -357,4 +368,265 @@ describe('MCPClient', () => {
       await mixedConfig.disconnect();
     });
   });
+
+  describe('Schema Handling', () => {
+    let complexClient: MCPClient;
+    let mockLogHandler: LogHandler & ReturnType<typeof vi.fn>;
+
+    beforeEach(async () => {
+      mockLogHandler = vi.fn();
+
+      complexClient = new MCPClient({
+        id: 'complex-schema-test-client-log-handler-firecrawl',
+        servers: {
+          'firecrawl-mcp': {
+            command: 'npx',
+            args: ['-y', 'tsx', path.join(__dirname, '__fixtures__/fire-crawl-complex-schema.ts')],
+            logger: mockLogHandler,
+          },
+        },
+      });
+    });
+
+    afterEach(async () => {
+      mockLogHandler.mockClear();
+      await complexClient?.disconnect().catch(() => {});
+    });
+
+    it('should process tools from firecrawl-mcp without crashing', async () => {
+      const tools = await complexClient.getTools();
+
+      Object.keys(allTools).forEach(toolName => {
+        expect(tools).toHaveProperty(`${mcpServerName.replace(`-fixture`, ``)}_${toolName}`);
+      });
+
+      expect(mockLogHandler.mock.calls.length).toBeGreaterThan(0);
+    });
+  });
+
+  describe('MCPClient Configuration', () => {
+    let clientsToCleanup: MCPClient[] = [];
+
+    afterEach(async () => {
+      await Promise.all(
+        clientsToCleanup.map(client =>
+          client.disconnect().catch(e => console.error(`Error disconnecting client during test cleanup: ${e}`)),
+        ),
+      );
+      clientsToCleanup = []; // Reset for the next test
+    });
+
+    it('should pass runtimeContext to the server logger function during tool execution', async () => {
+      type TestContext = { channel: string; userId: string };
+      const testContextInstance = new RuntimeContext<TestContext>();
+      testContextInstance.set('channel', 'test-channel-123');
+      testContextInstance.set('userId', 'user-abc-987');
+      const loggerFn = vi.fn();
+
+      const clientForTest = new MCPClient({
+        servers: {
+          stockPrice: {
+            command: 'npx',
+            args: ['tsx', path.join(__dirname, '__fixtures__/stock-price.ts')],
+            env: { FAKE_CREDS: 'test' },
+            logger: loggerFn,
+          },
+        },
+      });
+      clientsToCleanup.push(clientForTest);
+
+      const tools = await clientForTest.getTools();
+      const stockTool = tools['stockPrice_getStockPrice'];
+      expect(stockTool).toBeDefined();
+
+      await stockTool.execute({
+        context: { symbol: 'MSFT' },
+        runtimeContext: testContextInstance,
+      });
+
+      expect(loggerFn).toHaveBeenCalled();
+      const callWithContext = loggerFn.mock.calls.find(call => {
+        const logMessage = call[0] as LogMessage;
+        return (
+          logMessage.runtimeContext &&
+          typeof logMessage.runtimeContext.get === 'function' &&
+          logMessage.runtimeContext.get('channel') === 'test-channel-123' &&
+          logMessage.runtimeContext.get('userId') === 'user-abc-987'
+        );
+      });
+      expect(callWithContext).toBeDefined();
+      const capturedLogMessage = callWithContext?.[0] as LogMessage;
+      expect(capturedLogMessage?.serverName).toEqual('stockPrice');
+    }, 15000);
+
+    it('should pass runtimeContext to MCP logger when tool is called via an Agent', async () => {
+      type TestAgentContext = { traceId: string; tenant: string };
+      const agentTestContext = new RuntimeContext<TestAgentContext>();
+      agentTestContext.set('traceId', 'agent-trace-xyz');
+      agentTestContext.set('tenant', 'acme-corp');
+      const loggerFn = vi.fn();
+
+      const mcpClientForAgentTest = new MCPClient({
+        id: 'mcp-for-agent-test-suite',
+        servers: {
+          stockPriceServer: {
+            command: 'npx',
+            args: ['tsx', path.join(__dirname, '__fixtures__/stock-price.ts')],
+            env: { FAKE_CREDS: 'test' },
+            logger: loggerFn,
+          },
+        },
+      });
+      clientsToCleanup.push(mcpClientForAgentTest);
+
+      const agentName = 'stockAgentForContextTest';
+      const agent = new Agent({
+        name: agentName,
+        model: openai('gpt-4o'),
+        instructions: 'Use the getStockPrice tool to find the price of MSFT.',
+        tools: await mcpClientForAgentTest.getTools(),
+      });
+
+      await agent.generate('What is the price of MSFT?', { runtimeContext: agentTestContext });
+
+      expect(loggerFn).toHaveBeenCalled();
+      const callWithAgentContext = loggerFn.mock.calls.find(call => {
+        const logMessage = call[0] as LogMessage;
+        return (
+          logMessage.runtimeContext &&
+          typeof logMessage.runtimeContext.get === 'function' &&
+          logMessage.runtimeContext.get('traceId') === 'agent-trace-xyz' &&
+          logMessage.runtimeContext.get('tenant') === 'acme-corp'
+        );
+      });
+      expect(callWithAgentContext).toBeDefined();
+      if (callWithAgentContext) {
+        const capturedLogMessage = callWithAgentContext[0] as LogMessage;
+        expect(capturedLogMessage?.serverName).toEqual('stockPriceServer');
+      }
+    }, 20000);
+
+    it('should correctly use different runtimeContexts on sequential direct tool calls', async () => {
+      const loggerFn = vi.fn();
+      const clientForSeqTest = new MCPClient({
+        id: 'mcp-sequential-context-test',
+        servers: {
+          stockPriceServer: {
+            command: 'npx',
+            args: ['tsx', path.join(__dirname, '__fixtures__/stock-price.ts')],
+            env: { FAKE_CREDS: 'test' },
+            logger: loggerFn,
+          },
+        },
+      });
+      clientsToCleanup.push(clientForSeqTest);
+
+      const tools = await clientForSeqTest.getTools();
+      const stockTool = tools['stockPriceServer_getStockPrice'];
+      expect(stockTool).toBeDefined();
+
+      type ContextA = { callId: string };
+      const runtimeContextA = new RuntimeContext<ContextA>();
+      runtimeContextA.set('callId', 'call-A-111');
+      await stockTool.execute({ context: { symbol: 'MSFT' }, runtimeContext: runtimeContextA });
+
+      expect(loggerFn).toHaveBeenCalled();
+      let callsAfterA = [...loggerFn.mock.calls];
+      const logCallForA = callsAfterA.find(
+        call => (call[0] as LogMessage).runtimeContext?.get('callId') === 'call-A-111',
+      );
+      expect(logCallForA).toBeDefined();
+      expect((logCallForA?.[0] as LogMessage)?.runtimeContext?.get('callId')).toBe('call-A-111');
+
+      loggerFn.mockClear();
+
+      type ContextB = { sessionId: string };
+      const runtimeContextB = new RuntimeContext<ContextB>();
+      runtimeContextB.set('sessionId', 'session-B-222');
+      await stockTool.execute({ context: { symbol: 'GOOG' }, runtimeContext: runtimeContextB });
+
+      expect(loggerFn).toHaveBeenCalled();
+      let callsAfterB = [...loggerFn.mock.calls];
+      const logCallForB = callsAfterB.find(
+        call => (call[0] as LogMessage).runtimeContext?.get('sessionId') === 'session-B-222',
+      );
+      expect(logCallForB).toBeDefined();
+      expect((logCallForB?.[0] as LogMessage)?.runtimeContext?.get('sessionId')).toBe('session-B-222');
+
+      const contextALeak = callsAfterB.some(
+        call => (call[0] as LogMessage).runtimeContext?.get('callId') === 'call-A-111',
+      );
+      expect(contextALeak).toBe(false);
+    }, 20000);
+
+    it('should isolate runtimeContext between different servers on the same MCPClient', async () => {
+      const sharedLoggerFn = vi.fn();
+
+      const clientWithTwoServers = new MCPClient({
+        id: 'mcp-multi-server-context-isolation',
+        servers: {
+          serverX: {
+            command: 'npx',
+            args: ['tsx', path.join(__dirname, '__fixtures__/stock-price.ts')], // Re-use fixture, tool name will differ by server
+            logger: sharedLoggerFn,
+            env: { FAKE_CREDS: 'serverX-creds' }, // Make env slightly different for clarity if needed
+          },
+          serverY: {
+            command: 'npx',
+            args: ['tsx', path.join(__dirname, '__fixtures__/stock-price.ts')], // Re-use fixture
+            logger: sharedLoggerFn,
+            env: { FAKE_CREDS: 'serverY-creds' },
+          },
+        },
+      });
+      clientsToCleanup.push(clientWithTwoServers);
+
+      const tools = await clientWithTwoServers.getTools();
+      const toolX = tools['serverX_getStockPrice'];
+      const toolY = tools['serverY_getStockPrice'];
+      expect(toolX).toBeDefined();
+      expect(toolY).toBeDefined();
+
+      // --- Call tool on Server X with contextX ---
+      type ContextX = { requestId: string };
+      const runtimeContextX = new RuntimeContext<ContextX>();
+      runtimeContextX.set('requestId', 'req-X-001');
+
+      await toolX.execute({ context: { symbol: 'AAA' }, runtimeContext: runtimeContextX });
+
+      expect(sharedLoggerFn).toHaveBeenCalled();
+      let callsAfterToolX = [...sharedLoggerFn.mock.calls];
+      const logCallForX = callsAfterToolX.find(call => {
+        const logMessage = call[0] as LogMessage;
+        return logMessage.serverName === 'serverX' && logMessage.runtimeContext?.get('requestId') === 'req-X-001';
+      });
+      expect(logCallForX).toBeDefined();
+      expect((logCallForX?.[0] as LogMessage)?.runtimeContext?.get('requestId')).toBe('req-X-001');
+
+      sharedLoggerFn.mockClear(); // Clear for next distinct operation
+
+      // --- Call tool on Server Y with contextY ---
+      type ContextY = { customerId: string };
+      const runtimeContextY = new RuntimeContext<ContextY>();
+      runtimeContextY.set('customerId', 'cust-Y-002');
+
+      await toolY.execute({ context: { symbol: 'BBB' }, runtimeContext: runtimeContextY });
+
+      expect(sharedLoggerFn).toHaveBeenCalled();
+      let callsAfterToolY = [...sharedLoggerFn.mock.calls];
+      const logCallForY = callsAfterToolY.find(call => {
+        const logMessage = call[0] as LogMessage;
+        return logMessage.serverName === 'serverY' && logMessage.runtimeContext?.get('customerId') === 'cust-Y-002';
+      });
+      expect(logCallForY).toBeDefined();
+      expect((logCallForY?.[0] as LogMessage)?.runtimeContext?.get('customerId')).toBe('cust-Y-002');
+
+      // Ensure contextX did not leak into logs from serverY's operation
+      const contextXLeakInYLogs = callsAfterToolY.some(call => {
+        const logMessage = call[0] as LogMessage;
+        return logMessage.runtimeContext?.get('requestId') === 'req-X-001';
+      });
+      expect(contextXLeakInYLogs).toBe(false);
+    }, 25000); // Increased timeout for multiple server ops
+  });
 });