From c3ec8601a986d1fb168c84f39091dd4c571f3df6 Mon Sep 17 00:00:00 2001 From: Kevin Rajan <7121943+kvnloo@users.noreply.github.com> Date: Mon, 12 Oct 2026 03:56:24 -0500 Subject: [PATCH 1/2] fix(pi): single-flight Cua Driver stdio restart and handshake The async child_process import permits concurrent Pi tool requests to spawn multiple clients after the driver exits. Share the pending start and initialization, reset handshake status on exit/close, and keep reconnect explicitly session-scoped. Fork-only CHECK candidate for Julius Marminge's existing #16528; source finding PER-1778. GPT-6 / ChatGPT under Kevin's direction; no upstream PR. --- .../Adapters/piT3McpExtensionSource.ts | 48 +++++++++++++++---- 1 file changed, 39 insertions(+), 9 deletions(-) diff --git a/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.ts b/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.ts index 63fba58ccdf6..091212d037d9 100644 --- a/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.ts +++ b/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.ts @@ -234,6 +234,9 @@ function mcpImages(result: unknown): ToolContent[] { function createStdioMcpClient(spec: { command: string; args: string[]; env: Record }) { let nextId = 1; let child: ChildProcess | undefined; + let starting: Promise | undefined; + let connecting: Promise | undefined; + let connected = false; const waiting = new Map void; reject: (error: Error) => void }>(); const failAll = (error: Error) => { @@ -274,13 +277,27 @@ function createStdioMcpClient(spec: { command: string; args: string[]; env: Reco const stopped = () => { if (child !== current) return; child = undefined; + connected = false; failAll(new Error("Cua Driver stopped. Call the tool again to start a new session.")); }; current.on("exit", stopped); current.on("error", stopped); }; + // Importing node:child_process is async; share the pending start so concurrent + // Pi tool calls cannot spawn two clients and orphan the first child process. + const ensureStarted = async () => { + if (child !== undefined) return; + if (starting === undefined) { + const attempt = start(); + starting = attempt; + void attempt.finally(() => { + if (starting === attempt) starting = undefined; + }).catch(() => undefined); + } + await starting; + }; const request = async (method: string, params?: unknown, signal?: AbortSignal) => { - if (child === undefined) await start(); + await ensureStarted(); return new Promise((resolve, reject) => { const id = nextId++; waiting.set(id, { resolve, reject }); @@ -294,13 +311,25 @@ function createStdioMcpClient(spec: { command: string; args: string[]; env: Reco child?.stdin?.write(JSON.stringify({ jsonrpc: "2.0", method, params }) + "\\n"); }; - const connect = async (signal?: AbortSignal) => { - await request( - "initialize", - { protocolVersion: PROTOCOL, capabilities: {}, clientInfo: { name: "t3-pi-mcp", version: "1.0.0" } }, - signal, - ); - notify("notifications/initialized", {}); + const connect = (signal?: AbortSignal): Promise => { + if (child !== undefined && connected) return Promise.resolve(); + if (connecting !== undefined) return connecting; + // The replacement driver needs one handshake, even if concurrent tool + // calls hit the reconnect boundary on the same event loop tick. + const attempt = (async () => { + await request( + "initialize", + { protocolVersion: PROTOCOL, capabilities: {}, clientInfo: { name: "t3-pi-mcp", version: "1.0.0" } }, + signal, + ); + notify("notifications/initialized", {}); + connected = true; + })(); + connecting = attempt; + void attempt.finally(() => { + if (connecting === attempt) connecting = undefined; + }).catch(() => undefined); + return attempt; }; return { @@ -311,12 +340,13 @@ function createStdioMcpClient(spec: { command: string; args: string[]; env: Reco }, async callTool(name: string, args: Record, signal?: AbortSignal) { // A driver that exited (the session ended) starts again and needs a handshake. - if (child === undefined) await connect(signal); + if (child === undefined || !connected) await connect(signal); return request("tools/call", { name, arguments: args }, signal); }, close() { const current = child; child = undefined; + connected = false; current?.kill(); failAll(new Error("Cua Driver stopped. Call the tool again to start a new session.")); }, From e0f34e9b45d88925d44a7b39b40c225bdb84c7bc Mon Sep 17 00:00:00 2001 From: Kevin Rajan <7121943+kvnloo@users.noreply.github.com> Date: Mon, 12 Oct 2026 03:58:04 -0500 Subject: [PATCH 2/2] test(pi): exercise concurrent Cua Driver restart through emitted extension client Extract the generated stdio client, stub child_process with deterministic in-memory JSON-RPC pipes, and assert two simultaneous tool calls produce one restarted child and one MCP handshake. No live driver, desktop, or user data; fork-only CHECK pilot, source finding PER-1778. --- .../Adapters/piT3McpExtensionSource.test.ts | 101 ++++++++++++++++++ 1 file changed, 101 insertions(+) diff --git a/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.test.ts b/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.test.ts index 9dc9fa4526e6..530d69bf4f0a 100644 --- a/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.test.ts +++ b/apps/server/src/orchestration-v2/Adapters/piT3McpExtensionSource.test.ts @@ -57,3 +57,104 @@ describe("Pi upstream output-budget workaround", () => { ); }); }); + + +describe("Cua stdio restart single flight", () => { + it("reconnects two simultaneous Pi tool calls with one child and one handshake", async () => { + const children: Array<{ emitExit: () => void }> = []; + let initializations = 0; + + // Isolate the actual emitted Pi stdio client, replacing only child_process + // with a deterministic in-memory pipe. The await preserves the start race. + const beginning = PI_T3_MCP_EXTENSION_SOURCE.indexOf("function createStdioMcpClient("); + const end = PI_T3_MCP_EXTENSION_SOURCE.indexOf("\nfunction cuaSpec(", beginning); + assert.isAtLeast(beginning, 0); + assert.isAbove(end, beginning); + const clientSource = PI_T3_MCP_EXTENSION_SOURCE.slice(beginning, end); + const spawnStatement = 'const { spawn } = await import("node:child_process");'; + assert.include(clientSource, spawnStatement); + + const spawn = () => { + const stdoutListeners: Array<(chunk: string) => void> = []; + const exitListeners: Array<() => void> = []; + const child = { + stdout: { + setEncoding: (_encoding: string) => undefined, + on: (event: string, callback: (chunk: string) => void) => { + if (event === "data") stdoutListeners.push(callback); + }, + }, + stdin: { + write: (line: string) => { + const request = JSON.parse(line) as { + readonly id?: number; + readonly method: string; + readonly params?: { readonly name?: string }; + }; + if (request.method === "initialize") initializations += 1; + if (request.id === undefined) return; + const result = + request.method === "tools/call" + ? { content: [{ type: "text", text: request.params?.name ?? "" }] } + : {}; + queueMicrotask(() => { + const reply = JSON.stringify({ jsonrpc: "2.0", id: request.id, result }) + "\n"; + for (const listener of stdoutListeners) listener(reply); + }); + }, + }, + on: (event: string, callback: () => void) => { + if (event === "exit") exitListeners.push(callback); + }, + kill: () => { + for (const callback of exitListeners) callback(); + }, + emitExit: () => { + for (const callback of exitListeners) callback(); + }, + }; + children.push(child); + return child; + }; + + const source = NodeModule.stripTypeScriptTypes( + clientSource.replace( + spawnStatement, + "const { spawn } = await Promise.resolve(__fakeChildProcess);", + ), + ); + const makeClient = NodeVM.runInNewContext(source + "\ncreateStdioMcpClient", { + __fakeChildProcess: { spawn }, + process: { env: {} }, + PROTOCOL: "2025-06-18", + AbortSignal, + Promise, + }) as (spec: { command: string; args: string[]; env: Record }) => { + connect(): Promise; + callTool(name: string, args: Record, signal?: AbortSignal): Promise; + close(): void; + }; + + const client = makeClient({ command: "fake-driver", args: [], env: {} }); + await client.connect(); + assert.equal(children.length, 1); + assert.equal(initializations, 1); + + children[0]!.emitExit(); + const [one, two] = await Promise.all([ + client.callTool("list_apps", {}, AbortSignal.timeout(1_000)), + client.callTool("get_window_state", {}, AbortSignal.timeout(1_000)), + ]); + assert.equal(children.length, 2, "concurrent reconnect must not spawn two children"); + assert.equal(initializations, 2, "reconnect must initialize exactly once"); + assert.equal( + (one as { content: Array<{ text: string }> }).content[0]?.text, + "list_apps", + ); + assert.equal( + (two as { content: Array<{ text: string }> }).content[0]?.text, + "get_window_state", + ); + client.close(); + }); +});