fix(mcp): route stdio bridge requests

packages/mcp/src/server.ts

@@ -1,17 +1,16 @@
-
 import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
-import { resolveCommand } from 'package-manager-detector/commands'
-import { detect, getUserAgent } from 'package-manager-detector/detect'
 import { Hono } from 'hono'
 import { cors } from 'hono/cors'
-import { WebSocketServer } from 'ws'
+import { resolveCommand } from 'package-manager-detector/commands'
+import { detect, getUserAgent } from 'package-manager-detector/detect'
+import { WebSocketServer, type WebSocket } from 'ws'
 
+import packageJson from '../package.json'
 import { bearerToken, isAuthorized, mcpRequestToken } from './auth'
 import { createBrowserRpcBridge } from './browser-rpc'
 import { MCP_CORS_HEADERS, MCP_CORS_METHODS, MCP_EXPOSED_HEADERS } from './http-options'
 import { preprocessRpc } from './jsx-preprocess'
 import { createMcpSessionManager } from './mcp-sessions'
-import packageJson from '../package.json'
 import { registerTools } from './tool/registration'
 
 export const MCP_VERSION: string = packageJson.version
@@ -73,19 +72,64 @@ export function startServer(options: ServerOptions = {}) {
   // --- WebSocket: browser connects here ---
 
   const wss = new WebSocketServer({ port: wsPort, host: '127.0.0.1' })
+  const wsClients = new Set<WebSocket>()
+
+  function sendRegisterToken(ws: WebSocket) {
+    const token = browserRpc.currentRpcToken()
+    if (token && ws.readyState === ws.OPEN) {
+      ws.send(JSON.stringify({ type: 'register', token }))
+    }
+  }
+
+  function broadcastRegisterToken() {
+    for (const client of wsClients) sendRegisterToken(client)
+  }
+
+  async function handleClientRequest(ws: WebSocket, msg: Record<string, unknown>) {
+    const id = typeof msg.id === 'string' ? msg.id : null
+    if (!id) return
+    const { type: _type, id: _id, ...body } = msg
+    try {
+      const result = await sendToBrowser(body)
+      if (ws.readyState === ws.OPEN) {
+        ws.send(JSON.stringify({ type: 'response', id, ok: true, ...(result as object) }))
+      }
+    } catch (e) {
+      if (ws.readyState === ws.OPEN) {
+        ws.send(
+          JSON.stringify({
+            type: 'response',
+            id,
+            ok: false,
+            error: e instanceof Error ? e.message : String(e)
+          })
+        )
+      }
+    }
+  }
 
   wss.on('connection', (ws) => {
-    const token = browserRpc.currentRpcToken()
-    if (token) ws.send(JSON.stringify({ type: 'register', token }))
+    wsClients.add(ws)
+    sendRegisterToken(ws)
 
     ws.on('message', (raw) => {
-      browserRpc.handleMessage(
-        typeof raw === 'string' ? raw : Buffer.from(raw as Buffer).toString('utf-8'),
-        ws
-      )
+      const data = typeof raw === 'string' ? raw : Buffer.from(raw as Buffer).toString('utf-8')
+      let msg: Record<string, unknown>
+      try {
+        msg = JSON.parse(data) as Record<string, unknown>
+      } catch {
+        return
+      }
+      if (msg.type === 'request') {
+        void handleClientRequest(ws, msg)
+        return
+      }
+      browserRpc.handleMessage(data, ws)
+      if (msg.type === 'register') broadcastRegisterToken()
     })
 
     ws.on('close', () => {
+      wsClients.delete(ws)
       browserRpc.handleClose(ws)
     })
   })

tests/engine/mcp/server.test.ts

@@ -93,6 +93,19 @@ function connectMockBrowser(port: number, graph: SceneGraph): Promise<MockBrowse
   })
 }
 
+function readWsJson<T>(ws: WebSocket): Promise<T> {
+  return new Promise((resolve, reject) => {
+    ws.once('message', (raw) => {
+      try {
+        resolve(JSON.parse(raw.toString()) as T)
+      } catch (error) {
+        reject(error)
+      }
+    })
+    ws.once('error', reject)
+  })
+}
+
 function waitForWsListening(wss: InstanceType<typeof WebSocket.Server>): Promise<number> {
   return new Promise((resolve) => {
     if (wss.address()) {
@@ -250,6 +263,51 @@ describe('MCP server', () => {
   })
 })
 
+describe('MCP WebSocket stdio bridge routing', () => {
+  test('forwards stdio client requests to the registered desktop app', async () => {
+    const { wss, close: closeServer } = startServer({ httpPort: 0, wsPort: 0 })
+    const wsPort = await waitForWsListening(wss)
+    const graph = new SceneGraph()
+    const browser = await connectMockBrowser(wsPort, graph)
+    const clientWs = new WebSocket(`ws://127.0.0.1:${wsPort}`)
+
+    try {
+      await new Promise<void>((resolve, reject) => {
+        clientWs.once('open', () => resolve())
+        clientWs.once('error', reject)
+      })
+      const register = await readWsJson<{ type: string; token?: string }>(clientWs)
+      expect(register.type).toBe('register')
+      expect(register.token).toBeTruthy()
+
+      clientWs.send(
+        JSON.stringify({
+          type: 'request',
+          id: 'stdio-1',
+          command: 'tool',
+          args: { name: 'get_current_page', args: {} }
+        })
+      )
+      const response = await readWsJson<{
+        type: string
+        id: string
+        ok?: boolean
+        result?: { name: string }
+      }>(clientWs)
+
+      expect(response.type).toBe('response')
+      expect(response.id).toBe('stdio-1')
+      expect(response.ok).toBe(true)
+      expect(response.result?.name).toBe('Page 1')
+      expect(browser.requests.at(-1)?.command).toBe('tool')
+    } finally {
+      clientWs.close()
+      browser.close()
+      closeServer()
+    }
+  })
+})
+
 describe('MCP server with mcpRoot', () => {
   test('registers open_file and new_document tools when mcpRoot is set', async () => {
     const {