Keep MCP tool callback updates inside host executors

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<string | undefined> {
+  const result = await client.callTool({ name, arguments: {} })
+  const content = result.content[0]
+  return content?.type === 'text' ? content.text : undefined
+}

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<typeof callback>) =>
-      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<string, unknown>,
+  ) => void
+  registration.update = ((updates: Record<string, unknown>) => {
+    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
 }