@mastra/mcp 0.4.3 → 0.5.0-alpha.1

package/src/server.test.ts

@@ -1,4 +1,4 @@
-import http from 'http';
+import http from 'node:http';
 import path from 'path';
 import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
 import { describe, it, expect, beforeAll, afterAll, afterEach, vi } from 'vitest';
@@ -31,11 +31,11 @@ describe('MCPServer', () => {
       });
     });
 
-    await new Promise(resolve => httpServer.listen(PORT, resolve));
+    await new Promise<void>(resolve => httpServer.listen(PORT, () => resolve()));
   });
 
   afterAll(async () => {
-    await new Promise(resolve => httpServer.close(resolve));
+    await new Promise<void>(resolve => httpServer.close(() => resolve()));
   });
 
   describe('MCPServer SSE transport', () => {
@@ -77,7 +77,6 @@ describe('MCPServer', () => {
     });
 
     it('should return 503 if message sent before SSE connection', async () => {
-      // Manually clear the SSE transport
       (server as any).sseTransport = undefined;
       const res = await fetch(`http://localhost:${PORT}/message`, {
         method: 'POST',

package/src/server.ts

@@ -1,29 +1,24 @@
-import { isVercelTool, isZodType, resolveSerializedZodOutput } from '@mastra/core';
+import { randomUUID } from 'crypto';
+import type * as http from 'node:http';
+import type { InternalCoreTool } from '@mastra/core';
+import { makeCoreTool } from '@mastra/core';
 import type { ToolsInput } from '@mastra/core/agent';
+import { MCPServerBase } from '@mastra/core/mcp';
+import type { MCPServerSSEOptions, ConvertedTool } from '@mastra/core/mcp';
+import { RuntimeContext } from '@mastra/core/runtime-context';
 import { Server } from '@modelcontextprotocol/sdk/server/index.js';
 import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js';
 import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
+import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
+import type { StreamableHTTPServerTransportOptions } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
 import { CallToolRequestSchema, ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js';
-import jsonSchemaToZod from 'json-schema-to-zod';
 import { z } from 'zod';
-import { zodToJsonSchema } from 'zod-to-json-schema';
-import { createLogger } from './logger';
 
-const logger = createLogger();
-
-type ConvertedTool = {
-  name: string;
-  description?: string;
-  inputSchema: any;
-  zodSchema: z.ZodTypeAny;
-  execute: any;
-};
-
-export class MCPServer {
+export class MCPServer extends MCPServerBase {
   private server: Server;
-  private convertedTools: Record<string, ConvertedTool>;
   private stdioTransport?: StdioServerTransport;
   private sseTransport?: SSEServerTransport;
+  private streamableHTTPTransport?: StreamableHTTPServerTransport;
 
   /**
    * Get the current stdio transport.
@@ -40,10 +35,10 @@ export class MCPServer {
   }
 
   /**
-   * Get a read-only view of the registered tools (for testing/introspection).
+   * Get the current streamable HTTP transport.
    */
-  tools(): Readonly<Record<string, ConvertedTool>> {
-    return this.convertedTools;
+  public getStreamableHTTPTransport(): StreamableHTTPServerTransport | undefined {
+    return this.streamableHTTPTransport;
   }
 
   /**
@@ -53,9 +48,11 @@ export class MCPServer {
    * @param opts.tools - Tool definitions to register
    */
   constructor({ name, version, tools }: { name: string; version: string; tools: ToolsInput }) {
+    super({ name, version, tools });
+
     this.server = new Server({ name, version }, { capabilities: { tools: {}, logging: { enabled: true } } });
-    this.convertedTools = this.convertTools(tools);
-    void logger.info(
+
+    this.logger.info(
       `Initialized MCPServer '${name}' v${version} with tools: ${Object.keys(this.convertedTools).join(', ')}`,
     );
 
@@ -68,55 +65,39 @@ export class MCPServer {
    * @param tools Tool definitions
    * @returns Converted tools registry
    */
-  private convertTools(tools: ToolsInput): Record<string, ConvertedTool> {
+  convertTools(tools: ToolsInput): Record<string, ConvertedTool> {
     const convertedTools: Record<string, ConvertedTool> = {};
     for (const toolName of Object.keys(tools)) {
-      let inputSchema: any;
-      let zodSchema: z.ZodTypeAny;
       const toolInstance = tools[toolName];
       if (!toolInstance) {
-        void logger.warning(`Tool instance for '${toolName}' is undefined. Skipping.`);
+        this.logger.warn(`Tool instance for '${toolName}' is undefined. Skipping.`);
         continue;
       }
+
       if (typeof toolInstance.execute !== 'function') {
-        void logger.warning(`Tool '${toolName}' does not have a valid execute function. Skipping.`);
+        this.logger.warn(`Tool '${toolName}' does not have a valid execute function. Skipping.`);
         continue;
       }
-      // Vercel tools: .parameters is either Zod or JSON schema
-      if (isVercelTool(toolInstance)) {
-        if (isZodType(toolInstance.parameters)) {
-          zodSchema = toolInstance.parameters;
-          inputSchema = zodToJsonSchema(zodSchema);
-        } else if (typeof toolInstance.parameters === 'object') {
-          zodSchema = resolveSerializedZodOutput(jsonSchemaToZod(toolInstance.parameters));
-          inputSchema = toolInstance.parameters;
-        } else {
-          zodSchema = z.object({});
-          inputSchema = zodToJsonSchema(zodSchema);
-        }
-      } else {
-        // Mastra tools: .inputSchema is always Zod
-        zodSchema = toolInstance?.inputSchema ?? z.object({});
-        inputSchema = zodToJsonSchema(zodSchema);
-      }
 
-      // Wrap execute to support both signatures (typed, returns Promise<any>)
-      const execute: (args: any, execOptions?: any) => Promise<any> = async (args, execOptions) => {
-        if (isVercelTool(toolInstance)) {
-          return (await toolInstance.execute?.(args, execOptions)) ?? undefined;
-        }
-        return (await toolInstance.execute?.({ context: args }, execOptions)) ?? undefined;
+      const options = {
+        name: toolName,
+        runtimeContext: new RuntimeContext(),
+        mastra: this.mastra,
+        logger: this.logger,
+        description: toolInstance?.description,
       };
+
+      const coreTool = makeCoreTool(toolInstance, options) as InternalCoreTool;
+
       convertedTools[toolName] = {
         name: toolName,
-        description: toolInstance?.description,
-        inputSchema,
-        zodSchema,
-        execute,
+        description: coreTool.description,
+        parameters: coreTool.parameters,
+        execute: coreTool.execute,
       };
-      void logger.info(`Registered tool: '${toolName}' [${toolInstance?.description || 'No description'}]`);
+      this.logger.info(`Registered tool: '${toolName}' [${toolInstance?.description || 'No description'}]`);
     }
-    void logger.info(`Total tools registered: ${Object.keys(convertedTools).length}`);
+    this.logger.info(`Total tools registered: ${Object.keys(convertedTools).length}`);
     return convertedTools;
   }
 
@@ -125,12 +106,12 @@ export class MCPServer {
    */
   private registerListToolsHandler() {
     this.server.setRequestHandler(ListToolsRequestSchema, async () => {
-      await logger.debug('Handling ListTools request');
+      this.logger.debug('Handling ListTools request');
       return {
         tools: Object.values(this.convertedTools).map(tool => ({
           name: tool.name,
           description: tool.description,
-          inputSchema: tool.inputSchema,
+          inputSchema: tool.parameters.jsonSchema,
         })),
       };
     });
@@ -145,17 +126,36 @@ export class MCPServer {
       try {
         const tool = this.convertedTools[request.params.name];
         if (!tool) {
-          await logger.warning(`CallTool: Unknown tool '${request.params.name}' requested.`);
+          this.logger.warn(`CallTool: Unknown tool '${request.params.name}' requested.`);
           return {
             content: [{ type: 'text', text: `Unknown tool: ${request.params.name}` }],
             isError: true,
           };
         }
-        await logger.debug(`CallTool: Invoking '${request.params.name}' with arguments:`, request.params.arguments);
-        const args = tool.zodSchema.parse(request.params.arguments ?? {});
-        const result = await tool.execute(args, request.params);
+
+        this.logger.debug(`CallTool: Invoking '${request.params.name}' with arguments:`, request.params.arguments);
+
+        const validation = tool.parameters.validate?.(request.params.arguments ?? {});
+        if (validation && !validation.success) {
+          this.logger.warn(`CallTool: Invalid tool arguments for '${request.params.name}'`, {
+            errors: validation.error,
+          });
+          return {
+            content: [{ type: 'text', text: `Invalid tool arguments: ${JSON.stringify(validation.error)}` }],
+            isError: true,
+          };
+        }
+        if (!tool.execute) {
+          this.logger.warn(`CallTool: Tool '${request.params.name}' does not have an execute function.`);
+          return {
+            content: [{ type: 'text', text: `Tool '${request.params.name}' does not have an execute function.` }],
+            isError: true,
+          };
+        }
+
+        const result = await tool.execute(validation?.value, { messages: [], toolCallId: '' });
         const duration = Date.now() - startTime;
-        await logger.info(`Tool '${request.params.name}' executed successfully in ${duration}ms.`);
+        this.logger.info(`Tool '${request.params.name}' executed successfully in ${duration}ms.`);
         return {
           content: [
             {
@@ -168,7 +168,7 @@ export class MCPServer {
       } catch (error) {
         const duration = Date.now() - startTime;
         if (error instanceof z.ZodError) {
-          await logger.warning('Invalid tool arguments', {
+          this.logger.warn('Invalid tool arguments', {
             tool: request.params.name,
             errors: error.errors,
             duration: `${duration}ms`,
@@ -183,7 +183,7 @@ export class MCPServer {
             isError: true,
           };
         }
-        await logger.error(`Tool execution failed: ${request.params.name}`, error);
+        this.logger.error(`Tool execution failed: ${request.params.name}`, { error });
         return {
           content: [{ type: 'text', text: `Error: ${error instanceof Error ? error.message : String(error)}` }],
           isError: true,
@@ -195,10 +195,10 @@ export class MCPServer {
   /**
    * Start the MCP server using stdio transport (for Windsurf integration).
    */
-  async startStdio() {
+  public async startStdio(): Promise<void> {
     this.stdioTransport = new StdioServerTransport();
     await this.server.connect(this.stdioTransport);
-    await logger.info('Started MCP Server (stdio)');
+    this.logger.info('Started MCP Server (stdio)');
   }
 
   /**
@@ -211,43 +211,135 @@ export class MCPServer {
    * @param req Incoming HTTP request
    * @param res HTTP response (must support .write/.end)
    */
-  async startSSE({
+  public async startSSE({ url, ssePath, messagePath, req, res }: MCPServerSSEOptions): Promise<void> {
+    if (url.pathname === ssePath) {
+      await this.connectSSE({
+        messagePath,
+        res,
+      });
+    } else if (url.pathname === messagePath) {
+      this.logger.debug('Received message');
+      if (!this.sseTransport) {
+        res.writeHead(503);
+        res.end('SSE connection not established');
+        return;
+      }
+      await this.sseTransport.handlePostMessage(req, res);
+    } else {
+      this.logger.debug('Unknown path:', { path: url.pathname });
+      res.writeHead(404);
+      res.end();
+    }
+  }
+
+  /**
+   * Handles MCP-over-StreamableHTTP protocol for user-provided HTTP servers.
+   * Call this from your HTTP server for the streamable HTTP endpoint.
+   *
+   * @param url Parsed URL of the incoming request
+   * @param httpPath Path for establishing the streamable HTTP connection (e.g. '/mcp')
+   * @param req Incoming HTTP request
+   * @param res HTTP response (must support .write/.end)
+   * @param options Optional options to pass to the transport (e.g. sessionIdGenerator)
+   */
+  public async startHTTP({
     url,
-    ssePath,
-    messagePath,
+    httpPath,
     req,
     res,
+    options = { sessionIdGenerator: () => randomUUID() },
   }: {
     url: URL;
-    ssePath: string;
-    messagePath: string;
-    req: any;
-    res: any;
+    httpPath: string;
+    req: http.IncomingMessage;
+    res: http.ServerResponse<http.IncomingMessage>;
+    options?: StreamableHTTPServerTransportOptions;
   }) {
-    if (url.pathname === ssePath) {
-      await logger.debug('Received SSE connection');
-      this.sseTransport = new SSEServerTransport(messagePath, res);
-      await this.server.connect(this.sseTransport);
+    if (url.pathname === httpPath) {
+      this.streamableHTTPTransport = new StreamableHTTPServerTransport(options);
+      try {
+        await this.server.connect(this.streamableHTTPTransport);
+      } catch (error) {
+        this.logger.error('Error connecting to MCP server', { error });
+        res.writeHead(500);
+        res.end('Error connecting to MCP server');
+        return;
+      }
+
+      try {
+        await this.streamableHTTPTransport.handleRequest(req, res);
+      } catch (error) {
+        this.logger.error('Error handling MCP connection', { error });
+        res.writeHead(500);
+        res.end('Error handling MCP connection');
+        return;
+      }
 
       this.server.onclose = async () => {
+        this.streamableHTTPTransport = undefined;
         await this.server.close();
-        this.sseTransport = undefined;
       };
+
       res.on('close', () => {
-        this.sseTransport = undefined;
+        this.streamableHTTPTransport = undefined;
       });
-    } else if (url.pathname === messagePath) {
-      await logger.debug('Received message');
-      if (!this.sseTransport) {
-        res.writeHead(503);
-        res.end('SSE connection not established');
-        return;
-      }
-      await this.sseTransport.handlePostMessage(req, res);
     } else {
-      await logger.debug('Unknown path:', url.pathname);
       res.writeHead(404);
       res.end();
     }
   }
+
+  public async handlePostMessage(req: http.IncomingMessage, res: http.ServerResponse<http.IncomingMessage>) {
+    if (!this.sseTransport) {
+      res.writeHead(503);
+      res.end('SSE connection not established');
+      return;
+    }
+    await this.sseTransport.handlePostMessage(req, res);
+  }
+
+  public async connectSSE({
+    messagePath,
+    res,
+  }: {
+    messagePath: string;
+    res: http.ServerResponse<http.IncomingMessage>;
+  }) {
+    this.logger.debug('Received SSE connection');
+    this.sseTransport = new SSEServerTransport(messagePath, res);
+    await this.server.connect(this.sseTransport);
+
+    this.server.onclose = async () => {
+      this.sseTransport = undefined;
+      await this.server.close();
+    };
+
+    res.on('close', () => {
+      this.sseTransport = undefined;
+    });
+  }
+
+  /**
+   * Close the MCP server and all its connections
+   */
+  async close() {
+    try {
+      if (this.stdioTransport) {
+        await this.stdioTransport.close?.();
+        this.stdioTransport = undefined;
+      }
+      if (this.sseTransport) {
+        await this.sseTransport.close?.();
+        this.sseTransport = undefined;
+      }
+      if (this.streamableHTTPTransport) {
+        await this.streamableHTTPTransport.close?.();
+        this.streamableHTTPTransport = undefined;
+      }
+      await this.server.close();
+      this.logger.info('MCP server closed.');
+    } catch (error) {
+      this.logger.error('Error closing MCP server:', { error });
+    }
+  }
 }