diff --git a/src/server/mcp.test.ts b/src/server/mcp.test.ts index 537bdc3ae..6d6421cf1 100644 --- a/src/server/mcp.test.ts +++ b/src/server/mcp.test.ts @@ -453,6 +453,95 @@ describe('tool()', () => { expect(notifications).toHaveLength(0); }); + /*** + * Test: Updating Tool with outputSchema + */ + test('should update tool with outputSchema', async () => { + const mcpServer = new McpServer({ + name: 'test server', + version: '1.0' + }); + const notifications: Notification[] = []; + const client = new Client({ + name: 'test client', + version: '1.0' + }); + client.fallbackNotificationHandler = async notification => { + notifications.push(notification); + }; + + // Register initial tool + const tool = mcpServer.registerTool( + 'test', + { + outputSchema: { + result: z.number() + } + }, + async () => ({ + content: [{ type: 'text', text: '' }], + structuredContent: { + result: 42 + } + }) + ); + + // Update the tool with a different outputSchema + tool.update({ + outputSchema: { + result: z.number(), + sum: z.number() + }, + callback: async () => ({ + content: [{ type: 'text', text: '' }], + structuredContent: { + result: 42, + sum: 100 + } + }) + }); + + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + + await Promise.all([client.connect(clientTransport), mcpServer.connect(serverTransport)]); + + // Verify the outputSchema was updated + const listResult = await client.request( + { + method: 'tools/list' + }, + ListToolsResultSchema + ); + + expect(listResult.tools[0].outputSchema).toMatchObject({ + type: 'object', + properties: { + result: { type: 'number' }, + sum: { type: 'number' } + } + }); + + // Call the tool to verify it works with the updated outputSchema + const callResult = await client.request( + { + method: 'tools/call', + params: { + name: 'test', + arguments: {} + } + }, + CallToolResultSchema + ); + + expect(callResult.structuredContent).toEqual({ + result: 42, + sum: 100 + }); + + // Update happened before transport was connected, so no notifications should be expected + expect(notifications).toHaveLength(0); + }); + /*** * Test: Tool List Changed Notifications */ diff --git a/src/server/mcp.ts b/src/server/mcp.ts index 765ba864f..7255cf8e1 100644 --- a/src/server/mcp.ts +++ b/src/server/mcp.ts @@ -672,6 +672,7 @@ export class McpServer { if (typeof updates.title !== 'undefined') registeredTool.title = updates.title; if (typeof updates.description !== 'undefined') registeredTool.description = updates.description; if (typeof updates.paramsSchema !== 'undefined') registeredTool.inputSchema = z.object(updates.paramsSchema); + if (typeof updates.outputSchema !== 'undefined') registeredTool.outputSchema = z.object(updates.outputSchema); if (typeof updates.callback !== 'undefined') registeredTool.callback = updates.callback; if (typeof updates.annotations !== 'undefined') registeredTool.annotations = updates.annotations; if (typeof updates._meta !== 'undefined') registeredTool._meta = updates._meta;