diff --git a/packages/mcp/src/server.test.ts b/packages/mcp/src/server.test.ts index 232970725..2da6137a8 100644 --- a/packages/mcp/src/server.test.ts +++ b/packages/mcp/src/server.test.ts @@ -113,4 +113,98 @@ describe('Pascal MCP tool execution', () => { await server.close() } }) + + test('wraps registerTool callback updates and fails closed on renames', async () => { + const bridge = new SceneBridge() + bridge.loadDefault() + const executed: string[] = [] + const server = createPascalMcpServer({ + bridge, + executeTool: async ({ name, execute }) => { + executed.push(name) + return execute() + }, + }) + const registration = server.registerTool('update_probe', { inputSchema: {} }, async () => ({ + content: [{ type: 'text', text: 'initial' }], + })) + const [serverTransport, clientTransport] = InMemoryTransport.createLinkedPair() + const client = new Client({ name: 'registered-tool-update-test', version: '0.0.0' }) + await Promise.all([server.connect(serverTransport), client.connect(clientTransport)]) + + try { + expect(await toolText(client, 'update_probe')).toBe('initial') + registration.update({ + callback: async () => ({ content: [{ type: 'text', text: 'replacement' }] }), + }) + expect(await toolText(client, 'update_probe')).toBe('replacement') + expect(() => registration.update({ name: 'renamed_probe' })).toThrow( + 'MCP tool renaming is unsupported', + ) + expect(() => registration.update({ name: 'renamed_again_probe' })).toThrow( + 'MCP tool renaming is unsupported', + ) + expect(await toolText(client, 'update_probe')).toBe('replacement') + registration.remove() + expect((await client.listTools()).tools.map((tool) => tool.name)).not.toContain( + 'update_probe', + ) + expect(executed).toEqual(['update_probe', 'update_probe', 'update_probe']) + } finally { + await client.close() + await server.close() + } + }) + + test('wraps deprecated tool callback updates and fails closed on renames', async () => { + const bridge = new SceneBridge() + bridge.loadDefault() + const executed: string[] = [] + const server = createPascalMcpServer({ + bridge, + executeTool: async ({ name, execute }) => { + executed.push(name) + return execute() + }, + }) + const registration = server.tool('legacy_update_probe', async () => ({ + content: [{ type: 'text', text: 'initial' }], + })) + const [serverTransport, clientTransport] = InMemoryTransport.createLinkedPair() + const client = new Client({ name: 'legacy-tool-update-test', version: '0.0.0' }) + await Promise.all([server.connect(serverTransport), client.connect(clientTransport)]) + + try { + expect(await toolText(client, 'legacy_update_probe')).toBe('initial') + registration.update({ + callback: async () => ({ content: [{ type: 'text', text: 'replacement' }] }), + }) + expect(await toolText(client, 'legacy_update_probe')).toBe('replacement') + expect(() => registration.update({ name: 'legacy_renamed_probe' })).toThrow( + 'MCP tool renaming is unsupported', + ) + expect(() => registration.update({ name: 'legacy_renamed_again_probe' })).toThrow( + 'MCP tool renaming is unsupported', + ) + expect(await toolText(client, 'legacy_update_probe')).toBe('replacement') + registration.remove() + expect((await client.listTools()).tools.map((tool) => tool.name)).not.toContain( + 'legacy_update_probe', + ) + expect(executed).toEqual([ + 'legacy_update_probe', + 'legacy_update_probe', + 'legacy_update_probe', + ]) + } finally { + await client.close() + await server.close() + } + }) }) + +async function toolText(client: Client, name: string): Promise { + const result = await client.callTool({ name, arguments: {} }) + const content = result.content[0] + return content?.type === 'text' ? content.text : undefined +} diff --git a/packages/mcp/src/server.ts b/packages/mcp/src/server.ts index b03c530bf..0a3aba19b 100644 --- a/packages/mcp/src/server.ts +++ b/packages/mcp/src/server.ts @@ -1,4 +1,4 @@ -import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js' +import { McpServer, type RegisteredTool } from '@modelcontextprotocol/sdk/server/mcp.js' import type { SceneBridge } from './bridge/scene-bridge' import { createSceneOperations, type SceneOperations } from './operations' import { registerPrompts } from './prompts' @@ -21,7 +21,11 @@ export type CreatePascalMcpServerOptions = { store?: SceneStore name?: string version?: string - /** Wrap every tool handler, for example to serialize access to a stateful bridge. */ + /** + * Wrap every regular tool handler, including callback updates. + * Tool renames fail closed because the SDK registration lifecycle cannot safely rename twice. + * Experimental task-based tool registrations are outside this hook. + */ executeTool?: PascalMcpToolExecutor } @@ -42,13 +46,15 @@ export function createPascalMcpServer(opts: CreatePascalMcpServerOptions): McpSe function installToolExecutor(server: McpServer, executeTool: PascalMcpToolExecutor): void { const registerTool = server.registerTool.bind(server) - const wrappedRegisterTool: McpServer['registerTool'] = (name, config, callback) => - registerTool(name, config, ((...args: Parameters) => - executeTool({ - name, - signal: toolRequestSignal(args), - execute: () => Promise.resolve(Reflect.apply(callback, undefined, args)), - })) as typeof callback) + const wrappedRegisterTool: McpServer['registerTool'] = (name, config, callback) => { + const runtimeCallback = callback as unknown as RuntimeToolCallback + const registration = registerTool( + name, + config, + wrapToolCallback(name, runtimeCallback, executeTool) as typeof callback, + ) + return wrapRegisteredTool(registration, name, runtimeCallback, executeTool) + } server.registerTool = wrappedRegisterTool const tool = server.tool.bind(server) @@ -57,16 +63,56 @@ function installToolExecutor(server: McpServer, executeTool: PascalMcpToolExecut if (typeof callback !== 'function') { return Reflect.apply(tool, undefined, [name, ...args]) } - args[args.length - 1] = (...callbackArgs: unknown[]) => - executeTool({ - name, - signal: toolRequestSignal(callbackArgs), - execute: () => Promise.resolve(Reflect.apply(callback, undefined, callbackArgs)), - }) - return Reflect.apply(tool, undefined, [name, ...args]) + const runtimeCallback = callback as RuntimeToolCallback + args[args.length - 1] = wrapToolCallback(name, runtimeCallback, executeTool) + const registration = Reflect.apply(tool, undefined, [name, ...args]) as RegisteredTool + return wrapRegisteredTool(registration, name, runtimeCallback, executeTool) }) as McpServer['tool'] } +type RuntimeToolCallback = (...args: unknown[]) => unknown + +function wrapToolCallback( + name: string, + callback: RuntimeToolCallback, + executeTool: PascalMcpToolExecutor, +): RuntimeToolCallback { + return (...args) => + executeTool({ + name, + signal: toolRequestSignal(args), + execute: () => Promise.resolve(Reflect.apply(callback, undefined, args)), + }) +} + +function wrapRegisteredTool( + registration: RegisteredTool, + initialName: string, + initialCallback: RuntimeToolCallback, + executeTool: PascalMcpToolExecutor, +): RegisteredTool { + let currentCallback = initialCallback + const update = registration.update.bind(registration) as ( + updates: Record, + ) => void + registration.update = ((updates: Record) => { + if (typeof updates.name === 'string') { + throw new Error('MCP tool renaming is unsupported when executeTool is configured') + } + const callbackUpdate = updates.callback + if (typeof callbackUpdate === 'function') { + currentCallback = callbackUpdate as RuntimeToolCallback + } + update({ + ...updates, + ...(typeof callbackUpdate === 'function' + ? { callback: wrapToolCallback(initialName, currentCallback, executeTool) } + : {}), + }) + }) as RegisteredTool['update'] + return registration +} + function toolRequestSignal(args: readonly unknown[]): AbortSignal { return (args.at(-1) as { signal: AbortSignal }).signal }