From 02bbabaf1784a0111e9feb4d5c94b28f1f674b03 Mon Sep 17 00:00:00 2001 From: Danila Poyarkov Date: Tue, 19 May 2026 15:47:38 +0300 Subject: [PATCH] fix(mcp): route stdio bridge requests --- packages/mcp/src/server.ts | 66 +++++++++++++++++++++++++++------ tests/engine/mcp/server.test.ts | 58 +++++++++++++++++++++++++++++ 2 files changed, 113 insertions(+), 11 deletions(-) diff --git a/packages/mcp/src/server.ts b/packages/mcp/src/server.ts index 838894114..ade91ecfa 100644 --- a/packages/mcp/src/server.ts +++ b/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() + + 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) { + 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 + try { + msg = JSON.parse(data) as Record + } 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) }) }) diff --git a/tests/engine/mcp/server.test.ts b/tests/engine/mcp/server.test.ts index 62d79fc1a..69950c426 100644 --- a/tests/engine/mcp/server.test.ts +++ b/tests/engine/mcp/server.test.ts @@ -93,6 +93,19 @@ function connectMockBrowser(port: number, graph: SceneGraph): Promise(ws: WebSocket): Promise { + 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): Promise { 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((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 {