fix(mcp): route stdio bridge requests

This commit is contained in:
Danila Poyarkov 2026-05-19 15:47:38 +03:00
parent 3a85632f38
commit 02bbabaf17
2 changed files with 113 additions and 11 deletions

View file

@ -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)
})
})

View file

@ -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 {