fix: Support updating output schema (#1048)

Co-authored-by: Konstantin Konstantinov <konstantin@mach5technology.com>

src/server/mcp.test.ts

@@ -480,6 +480,95 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
             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
          */

src/server/mcp.ts

@@ -906,6 +906,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 = objectFromShape(updates.paramsSchema);
+                if (typeof updates.outputSchema !== 'undefined') registeredTool.outputSchema = objectFromShape(updates.outputSchema);
                 if (typeof updates.callback !== 'undefined') registeredTool.handler = updates.callback;
                 if (typeof updates.annotations !== 'undefined') registeredTool.annotations = updates.annotations;
                 if (typeof updates._meta !== 'undefined') registeredTool._meta = updates._meta;