diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 6fd8a6dfb..7fa728da8 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -37,6 +37,10 @@ jobs: working-directory: ./client run: npm test + - name: Run server tests + working-directory: ./server + run: npm test + - run: npm run build publish: diff --git a/server/package.json b/server/package.json index 5b398d64c..4b5d298e9 100644 --- a/server/package.json +++ b/server/package.json @@ -22,6 +22,7 @@ ], "scripts": { "build": "tsc && shx cp -R static build", + "test": "npm run build && node --test test/*.test.mjs", "start": "node build/index.js", "dev": "tsx watch --clear-screen=false src/index.ts", "dev:windows": "tsx watch --clear-screen=false src/index.ts < NUL" diff --git a/server/src/index.ts b/server/src/index.ts index bdfe49019..1da4fb6d0 100644 --- a/server/src/index.ts +++ b/server/src/index.ts @@ -222,6 +222,12 @@ const webAppTransports: Map = new Map(); / const serverTransports: Map = new Map(); // Server Transports by web app sessionId const sessionHeaderHolders: Map = new Map(); // For dynamic header updates +const cleanupTransportSession = (sessionId: string) => { + webAppTransports.delete(sessionId); + serverTransports.delete(sessionId); + sessionHeaderHolders.delete(sessionId); +}; + // Use provided token from environment or generate a new one const sessionToken = process.env.MCP_PROXY_AUTH_TOKEN || randomBytes(32).toString("hex"); @@ -660,34 +666,40 @@ app.get( const endpoint = `${prefix}/message`; const webAppTransport = new SSEServerTransport(endpoint, res); - webAppTransports.set(webAppTransport.sessionId, webAppTransport); + const sessionId = webAppTransport.sessionId; + webAppTransports.set(sessionId, webAppTransport); console.log("Created client transport"); - serverTransports.set(webAppTransport.sessionId, serverTransport); + serverTransports.set(sessionId, serverTransport); console.log("Created server transport"); await webAppTransport.start(); - (serverTransport as StdioClientTransport).stderr!.on("data", (chunk) => { + const stderr = (serverTransport as StdioClientTransport).stderr!; + const handleStderr = (chunk: Buffer) => { if (chunk.toString().includes("MODULE_NOT_FOUND")) { // Server command not found, remove transports const message = "Command not found, transports removed"; - webAppTransport.send({ - jsonrpc: "2.0", - method: "notifications/message", - params: { - level: "emergency", - logger: "proxy", - data: { - message, + void webAppTransport + .send({ + jsonrpc: "2.0", + method: "notifications/message", + params: { + level: "emergency", + logger: "proxy", + data: { + message, + }, }, - }, - }); - webAppTransport.close(); - serverTransport.close(); - webAppTransports.delete(webAppTransport.sessionId); - serverTransports.delete(webAppTransport.sessionId); - sessionHeaderHolders.delete(webAppTransport.sessionId); + }) + .catch((error) => { + console.error("Error sending command failure to client:", error); + }); + void Promise.allSettled([ + webAppTransport.close(), + serverTransport.close(), + ]); + cleanupTransportSession(sessionId); console.error(message); } else { // Inspect message and attempt to assign a RFC 5424 Syslog Protocol level @@ -722,23 +734,35 @@ app.get( } else { level = "info"; } - webAppTransport.send({ - jsonrpc: "2.0", - method: "notifications/message", - params: { - level, - logger: "stdio", - data: { - message, + void webAppTransport + .send({ + jsonrpc: "2.0", + method: "notifications/message", + params: { + level, + logger: "stdio", + data: { + message, + }, }, - }, - }); + }) + .catch((error) => { + if (webAppTransports.has(sessionId)) { + console.error("Error forwarding STDERR to client:", error); + } + }); } - }); + }; + stderr.on("data", handleStderr); mcpProxy({ transportToClient: webAppTransport, transportToServer: serverTransport, + onCleanup: () => { + stderr.off("data", handleStderr); + cleanupTransportSession(sessionId); + console.log(`Transports removed for sessionId ${sessionId}`); + }, }); } catch (error) { if (is401Error(error)) { @@ -773,12 +797,13 @@ app.get( const endpoint = `${prefix}/message`; const webAppTransport = new SSEServerTransport(endpoint, res); - webAppTransports.set(webAppTransport.sessionId, webAppTransport); + const sessionId = webAppTransport.sessionId; + webAppTransports.set(sessionId, webAppTransport); console.log("Created client transport"); - serverTransports.set(webAppTransport.sessionId, serverTransport!); // eslint-disable-line @typescript-eslint/no-non-null-assertion + serverTransports.set(sessionId, serverTransport!); // eslint-disable-line @typescript-eslint/no-non-null-assertion if (headerHolder) { - sessionHeaderHolders.set(webAppTransport.sessionId, headerHolder); + sessionHeaderHolders.set(sessionId, headerHolder); } console.log("Created server transport"); @@ -788,6 +813,10 @@ app.get( transportToClient: webAppTransport, transportToServer: serverTransport, headerHolder, + onCleanup: () => { + cleanupTransportSession(sessionId); + console.log(`Transports removed for sessionId ${sessionId}`); + }, }); } catch (error) { if (is401Error(error)) { diff --git a/server/src/mcpProxy.ts b/server/src/mcpProxy.ts index 3195553e1..c3d921dd5 100644 --- a/server/src/mcpProxy.ts +++ b/server/src/mcpProxy.ts @@ -81,13 +81,25 @@ export default function mcpProxy({ transportToClient, transportToServer, headerHolder, + onCleanup, }: { transportToClient: Transport; transportToServer: Transport; headerHolder?: ProxyHeaderHolder; + onCleanup?: () => void; }) { let transportToClientClosed = false; let transportToServerClosed = false; + let cleanupComplete = false; + + const cleanup = () => { + if (cleanupComplete) { + return; + } + + cleanupComplete = true; + onCleanup?.(); + }; let reportedServerSession = false; @@ -132,19 +144,24 @@ export default function mcpProxy({ }; transportToClient.onclose = () => { + transportToClientClosed = true; + cleanup(); + if (transportToServerClosed) { return; } - transportToClientClosed = true; transportToServer.close().catch(onServerError); }; transportToServer.onclose = () => { + transportToServerClosed = true; + cleanup(); + if (transportToClientClosed) { return; } - transportToServerClosed = true; + transportToClient.close().catch(onClientError); }; diff --git a/server/test/fixtures/stdio-shutdown-server.mjs b/server/test/fixtures/stdio-shutdown-server.mjs new file mode 100644 index 000000000..d8716a317 --- /dev/null +++ b/server/test/fixtures/stdio-shutdown-server.mjs @@ -0,0 +1,16 @@ +let shuttingDown = false; + +const shutdown = () => { + if (shuttingDown) { + return; + } + + shuttingDown = true; + console.error("STDIO fixture shutting down"); + setTimeout(() => process.exit(0), 25); +}; + +process.on("SIGINT", shutdown); +process.on("SIGTERM", shutdown); +process.stdin.on("end", shutdown); +process.stdin.resume(); diff --git a/server/test/mcpProxy.test.mjs b/server/test/mcpProxy.test.mjs new file mode 100644 index 000000000..2af2a27cf --- /dev/null +++ b/server/test/mcpProxy.test.mjs @@ -0,0 +1,57 @@ +import assert from "node:assert/strict"; +import test from "node:test"; + +import mcpProxy from "../build/mcpProxy.js"; + +class TestTransport { + closeCalls = 0; + + async start() {} + + async send() {} + + async close() { + this.closeCalls += 1; + this.onclose?.(); + } +} + +test("client disconnect cleans up once and closes the server transport", async () => { + const clientTransport = new TestTransport(); + const serverTransport = new TestTransport(); + let cleanupCalls = 0; + + mcpProxy({ + transportToClient: clientTransport, + transportToServer: serverTransport, + onCleanup: () => { + cleanupCalls += 1; + }, + }); + + clientTransport.onclose(); + await Promise.resolve(); + + assert.equal(cleanupCalls, 1); + assert.equal(serverTransport.closeCalls, 1); +}); + +test("server disconnect cleans up once and closes the client transport", async () => { + const clientTransport = new TestTransport(); + const serverTransport = new TestTransport(); + let cleanupCalls = 0; + + mcpProxy({ + transportToClient: clientTransport, + transportToServer: serverTransport, + onCleanup: () => { + cleanupCalls += 1; + }, + }); + + serverTransport.onclose(); + await Promise.resolve(); + + assert.equal(cleanupCalls, 1); + assert.equal(clientTransport.closeCalls, 1); +}); diff --git a/server/test/stdioReconnect.test.mjs b/server/test/stdioReconnect.test.mjs new file mode 100644 index 000000000..23826cdf2 --- /dev/null +++ b/server/test/stdioReconnect.test.mjs @@ -0,0 +1,95 @@ +import assert from "node:assert/strict"; +import { spawn } from "node:child_process"; +import { once } from "node:events"; +import { fileURLToPath } from "node:url"; +import test from "node:test"; + +const TEST_PORT = 16322; +const TEST_TOKEN = "stdio-reconnect-test-token"; +const BASE_URL = `http://127.0.0.1:${TEST_PORT}`; +const SERVER_PATH = fileURLToPath( + new URL("../build/index.js", import.meta.url), +); +const FIXTURE_PATH = fileURLToPath( + new URL("./fixtures/stdio-shutdown-server.mjs", import.meta.url), +); + +const delay = (milliseconds) => + new Promise((resolve) => setTimeout(resolve, milliseconds)); + +async function waitForHealth(proxy, logs) { + const deadline = Date.now() + 5_000; + + while (Date.now() < deadline) { + if (proxy.exitCode !== null) { + throw new Error(`Proxy exited before becoming ready:\n${logs.join("")}`); + } + + try { + const response = await fetch(`${BASE_URL}/health`); + if (response.ok) { + return; + } + } catch { + // The proxy may still be binding its listen socket. + } + + await delay(50); + } + + throw new Error(`Proxy did not become ready:\n${logs.join("")}`); +} + +async function connectStdioSession() { + const controller = new AbortController(); + const url = new URL(`${BASE_URL}/stdio`); + url.searchParams.set("transportType", "stdio"); + url.searchParams.set("command", process.execPath); + url.searchParams.set("args", FIXTURE_PATH); + + const response = await fetch(url, { + headers: { "X-MCP-Proxy-Auth": `Bearer ${TEST_TOKEN}` }, + signal: controller.signal, + }); + + assert.equal(response.status, 200); + return controller; +} + +test("STDIO session can reconnect after client disconnect", async (t) => { + const logs = []; + const proxy = spawn(process.execPath, [SERVER_PATH], { + env: { + ...process.env, + HOST: "127.0.0.1", + SERVER_PORT: String(TEST_PORT), + MCP_PROXY_AUTH_TOKEN: TEST_TOKEN, + }, + stdio: ["ignore", "pipe", "pipe"], + }); + + proxy.stdout.on("data", (chunk) => logs.push(chunk.toString())); + proxy.stderr.on("data", (chunk) => logs.push(chunk.toString())); + + t.after(async () => { + if (proxy.exitCode === null) { + proxy.kill(); + await once(proxy, "exit"); + } + }); + + await waitForHealth(proxy, logs); + + const firstSession = await connectStdioSession(); + firstSession.abort(); + await delay(200); + + assert.equal(proxy.exitCode, null, logs.join("")); + assert.equal((await fetch(`${BASE_URL}/health`)).status, 200); + + const secondSession = await connectStdioSession(); + secondSession.abort(); + await delay(200); + + assert.equal(proxy.exitCode, null, logs.join("")); +});