draft QA: single-flight Pi/CUA stdio restart (for T3 #16528)

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<string, string> }) => {
+      connect(): Promise<void>;
+      callTool(name: string, args: Record<string, unknown>, signal?: AbortSignal): Promise<unknown>;
+      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();
+  });
+});

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<string, string> }) {
   let nextId = 1;
   let child: ChildProcess | undefined;
+  let starting: Promise<void> | undefined;
+  let connecting: Promise<void> | undefined;
+  let connected = false;
   const waiting = new Map<number, { resolve: (value: unknown) => 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<unknown>((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<void> => {
+    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<string, unknown>, 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."));
     },