diff --git a/paseo-omp/server/provider/mcp-transport.ts b/paseo-omp/server/provider/mcp-transport.ts index efd73407..74b6930d 100644 --- a/paseo-omp/server/provider/mcp-transport.ts +++ b/paseo-omp/server/provider/mcp-transport.ts @@ -7,7 +7,8 @@ import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/ import { deserializeMessage, serializeMessage } from "@modelcontextprotocol/sdk/shared/stdio.js"; import type { FetchLike, Transport } from "@modelcontextprotocol/sdk/shared/transport.js"; import type { JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js"; -import { terminateSpawnedProcessTree } from "./omp-rpc-process"; +import { terminateSpawnedProcessTree, PROCESS_STOP_TIMEOUT_MS } from "./omp-rpc-process"; +import type { ProcessTreeCleanup } from "./omp-rpc-process"; import { OmpCleanupFailure } from "./security"; const MAX_MCP_TRANSPORT_FRAME_BYTES = 1024 * 1024; @@ -224,21 +225,54 @@ export class SupervisedStdioClientTransport implements Transport { this.notifyClose(); return; } - const treeCleanup = this.startTreeCleanup(); - if (!this.exited) { + + if (this.exited) return; + + class TimeOutError extends Error {} + + const endStdinPromise = new Promise((resolve, reject) => { + if (this.exited) { + resolve("verified"); + return; + } + + let timeout: TimerHandle; + this.child.once("close", () => { + clearTimeout(timeout); + resolve("verified"); + }); + this.child.once("error", (error: Error) => { + clearTimeout(timeout); + reject(error); + }); + + timeout = setTimeout(() => reject(new TimeOutError("MCP process did not close in time")), PROCESS_STOP_TIMEOUT_MS); + try { - child.stdin.end(); - } catch { - // Process-tree cleanup remains authoritative when stdin is already closed. + this.child.stdin.end(); + } catch (error) { + reject(error); } + }); + + try { + await endStdinPromise; + } catch (error) { + if (error instanceof TimeOutError) { + console.warn(`MCP stdio process did not close in time (${PROCESS_STOP_TIMEOUT_MS}ms), starting tree cleanup`); + let terminated = await this.startTreeCleanup(); + + if (!terminated) throw new Error("MCP stdio process tree cleanup failed"); + } + } + + if ( + !this.spawnFailedWithoutProcess && + !this.exited && + !(await this.waitForExit(PROCESS_STOP_TIMEOUT_MS)) + ) { + throw new Error("MCP stdio process did not close after tree cleanup"); } - const terminated = await treeCleanup; - const exited = - this.spawnFailedWithoutProcess || - this.exited || - (await waitForExit(this.exit.promise, PROCESS_EXIT_TIMEOUT_MS)); - this.notifyClose(); - if (!terminated || !exited) throw new Error("MCP stdio process tree cleanup failed"); } } diff --git a/paseo-omp/server/provider/omp-rpc-transport.ts b/paseo-omp/server/provider/omp-rpc-transport.ts index 34d9c120..313bf82c 100644 --- a/paseo-omp/server/provider/omp-rpc-transport.ts +++ b/paseo-omp/server/provider/omp-rpc-transport.ts @@ -473,16 +473,48 @@ export class OmpRpcProcess { this.flushProtocolViolations(); this.clearChunk(); this.failPending(new Error("OMP RPC process was closed")); - const cleanupPromise = this.startTreeCleanup(); - if (!this.exited) { + + if (this.exited) return; + + class TimeOutError extends Error {} + + const endStdinPromise = new Promise((resolve, reject) => { + if (this.exited) { + resolve("verified"); + return; + } + + let timeout: TimerHandle; + this.child.once("close", () => { + clearTimeout(timeout); + resolve("verified"); + }); + this.child.once("error", (error: Error) => { + clearTimeout(timeout); + reject(error); + }); + + timeout = setTimeout(() => reject(new TimeOutError("OMP RPC process did not close in time")), PROCESS_STOP_TIMEOUT_MS); + try { this.child.stdin.end(); - } catch { - // Continue waiting for process-tree cleanup when the input channel is already closed. + } catch (error) { + reject(error); + } + }); + + try { + await endStdinPromise; + } catch (error) { + if (error instanceof TimeOutError) { + console.warn(`OMP RPC process did not close in time (${PROCESS_STOP_TIMEOUT_MS}ms), starting tree cleanup`); + let cleanup = await this.startTreeCleanup(); + + if (cleanup === "uncertain" && this.exited) cleanup = "verified"; + if (cleanup !== "verified") throw new Error("OMP RPC process tree cleanup failed"); } } - const cleanup = await cleanupPromise; - if (cleanup !== "verified") throw new Error("OMP RPC process tree cleanup failed"); + if ( !this.spawnFailedWithoutProcess && !this.exited &&