History

Serialize MCP session DO startup (#1313)

* Add e2e coverage for concurrent SSE GET + POST reconnect after idle dispose * Serialize MCP session startup * Drop repro notes doc * Neutral wording for the session startup regression test

Rhys Sullivan committed 4 months ago

Commit fb808c81 parent f275d5f3 files · +377 −46Browse files

Unified · Split

3 files
  1. e2e/cloud/repro-transport-brick.test.ts +215 −0
  2. packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts +134 −6
  3. packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts +28 −40
e2e/cloud/repro-transport-brick.test.ts+215−0
@@ -0,0 +1,215 @@
1+// Cloud: regression test for a session-DO startup race.
2+//
3+// 1. Open an MCP session (real SDK client init, capture mcp-session-id),
4+// make a call, close the client.
5+// 2. Idle past the session timeout so the DO's alarm runs disposeIdleRuntime
6+// (runtime torn down, `initialized = false`, transport cleared).
7+// 3. Fire an SSE GET (listen stream) and a POST (tools/list) concurrently on
8+// the same session id.
9+// 4. The session must survive: the concurrent pair and follow-up calls
10+// should succeed.
11+//
12+// Mechanism guarded against: a request carrying a session id runs
13+// validateMcpSessionOwner, which restores via onStart when the runtime is
14+// disposed. The SDK's serve() streaming handler separately wakes the DO
15+// through agent.fetch, and the Agent base also drives onStart. Without
16+// serialization across both entry paths, two onStart calls interleave and the
17+// second server.connect throws "Already connected to a transport", after
18+// which every request on the session fails the same way. The concurrent pair
19+// itself can answer 200/200 (one restore wins the race), so the assertions
20+// cover the follow-up calls, which is where the failure persists.
21+//
22+import { writeFileSync } from "node:fs";
23+import { join } from "node:path";
24+
25+import { expect } from "@effect/vitest";
26+import { Effect } from "effect";
27+import { Client } from "@modelcontextprotocol/sdk/client/index.js";
28+import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
29+
30+import { scenario } from "../src/scenario";
31+import { Mcp, RunDir, Target } from "../src/services";
32+import type { Identity } from "../src/target";
33+import { configuredMcpSessionTimeoutMs } from "../setup/mcp-session-timeouts";
34+
35+const emailOf = (identity: Identity): string => identity.credentials?.email ?? identity.label;
36+
37+// Give the DO's idle alarm room to fire disposeIdleRuntime after the timeout.
38+const IDLE_DISPOSE_BUFFER_MS = 2_000;
39+const IDLE_DISPOSE_GAP_MS = configuredMcpSessionTimeoutMs() + IDLE_DISPOSE_BUFFER_MS;
40+const SCENARIO_TIMEOUT_MS = IDLE_DISPOSE_GAP_MS + 120_000;
41+
42+interface Connected {
43+ readonly client: Client;
44+ readonly transport: StreamableHTTPClientTransport;
45+}
46+
47+const connectClient = async (mcpUrl: string, bearer: string): Promise<Connected> => {
48+ const client = new Client({ name: "executor-e2e-brick", version: "0.0.1" }, { capabilities: {} });
49+ const transport = new StreamableHTTPClientTransport(new URL(mcpUrl), {
50+ requestInit: { headers: { authorization: `Bearer ${bearer}` } },
51+ });
52+ await client.connect(transport);
53+ return { client, transport };
54+};
55+
56+interface RawResult {
57+ readonly status: number;
58+ readonly body: string;
59+}
60+
61+/** A raw POST of a JSON-RPC request onto an existing session id. */
62+const rawPost = async (
63+ mcpUrl: string,
64+ bearer: string,
65+ sessionId: string,
66+ message: unknown,
67+): Promise<RawResult> => {
68+ const res = await fetch(mcpUrl, {
69+ method: "POST",
70+ headers: {
71+ authorization: `Bearer ${bearer}`,
72+ "content-type": "application/json",
73+ // The streamable-http transport requires the client accept both.
74+ accept: "application/json, text/event-stream",
75+ "mcp-session-id": sessionId,
76+ "mcp-protocol-version": "2025-06-18",
77+ },
78+ body: JSON.stringify(message),
79+ });
80+ const body = await res.text().catch(() => "");
81+ return { status: res.status, body };
82+};
83+
84+/** A raw SSE GET (the "listen" stream) on an existing session id. We only need
85+ * the response head — status + start of body — to see whether it 500s. */
86+const rawSseGet = async (mcpUrl: string, bearer: string, sessionId: string): Promise<RawResult> => {
87+ const res = await fetch(mcpUrl, {
88+ method: "GET",
89+ headers: {
90+ authorization: `Bearer ${bearer}`,
91+ accept: "text/event-stream",
92+ "mcp-session-id": sessionId,
93+ "mcp-protocol-version": "2025-06-18",
94+ },
95+ });
96+ // A 200 opens a long-lived stream; don't hang reading it. A 500 has a short
97+ // (usually empty) body we can read fully.
98+ let body = "";
99+ if (res.status !== 200) {
100+ body = await res.text().catch(() => "");
101+ } else {
102+ // Cancel the stream immediately; we only wanted the status.
103+ await res.body?.cancel().catch(() => undefined);
104+ }
105+ return { status: res.status, body };
106+};
107+
108+const toolsListMessage = { jsonrpc: "2.0" as const, id: "brick-tools-list", method: "tools/list" };
109+
110+const isBrick = (r: RawResult): boolean =>
111+ r.status >= 500 || /"code"\s*:\s*-32603/.test(r.body) || /Already connected/i.test(r.body);
112+
113+scenario(
114+ "REGRESSION · concurrent SSE GET + POST after idle-dispose keeps the session alive",
115+ { timeout: SCENARIO_TIMEOUT_MS },
116+ Effect.gen(function* () {
117+ const target = yield* Target;
118+ const mcp = yield* Mcp;
119+ const runDir = yield* RunDir;
120+ const identity = yield* target.newIdentity();
121+ const bearer = yield* mcp.mintBearer(emailOf(identity));
122+ const diagnostics: unknown[] = [];
123+ const record = (entry: Record<string, unknown>): void => {
124+ diagnostics.push(entry);
125+ console.log(JSON.stringify(entry));
126+ // Vitest swallows worker stdout on passing runs; the artifact dir is the
127+ // durable copy of what each attempt observed.
128+ writeFileSync(join(runDir, "repro-diagnostics.json"), JSON.stringify(diagnostics, null, 2));
129+ };
130+
131+ // 1. Open a real session (SDK init handshake persists the session in the
132+ // DO), do a call, close the client.
133+ const first = yield* Effect.promise(() => connectClient(target.mcpUrl, bearer));
134+ const sessionId = first.transport.sessionId;
135+ expect(sessionId, "the client got a session id").toEqual(expect.any(String));
136+ if (sessionId === undefined) return yield* Effect.die("missing session id");
137+
138+ yield* Effect.promise(() =>
139+ first.client.callTool({ name: "execute", arguments: { code: "return 1;" } }),
140+ );
141+ yield* Effect.promise(() => first.client.close().catch(() => undefined));
142+
143+ // 2. Idle past the timeout so the DO's alarm runs disposeIdleRuntime.
144+ yield* Effect.sleep(IDLE_DISPOSE_GAP_MS);
145+
146+ // 3. Fire the SSE GET and the POST CONCURRENTLY on the same session id.
147+ // Raw fetch gives precise concurrency control (both start before either
148+ // resolves) that the SDK client's request queue would serialize away.
149+ const [getResult, postResult] = yield* Effect.promise(() =>
150+ Promise.all([
151+ rawSseGet(target.mcpUrl, bearer, sessionId),
152+ rawPost(target.mcpUrl, bearer, sessionId, toolsListMessage),
153+ ]),
154+ );
155+
156+ record({
157+ event: "repro_concurrent_pair",
158+ sessionId,
159+ get: getResult,
160+ post: postResult,
161+ });
162+
163+ const concurrentBricked = isBrick(getResult) || isBrick(postResult);
164+
165+ // 4. Prove the session survives: subsequent SEQUENTIAL POSTs also succeed.
166+ // This is the same control path as the idle-restore scenario from #1302.
167+ //
168+ // Observed locally: the concurrent pair frequently returns 200/200 (one
169+ // of the two restores wins the race and answers), but the race leaves
170+ // the DO's McpServer in the "Already connected to a transport" state, so
171+ // the very NEXT request 500s — and every request after that. The
172+ // regression gate checks the same two follow-up calls now stay healthy.
173+ const followUp = yield* Effect.promise(() =>
174+ rawPost(target.mcpUrl, bearer, sessionId, {
175+ ...toolsListMessage,
176+ id: "brick-follow-up",
177+ }),
178+ );
179+ const followUp2 = yield* Effect.promise(() =>
180+ rawPost(target.mcpUrl, bearer, sessionId, {
181+ ...toolsListMessage,
182+ id: "brick-follow-up-2",
183+ }),
184+ );
185+ record({ event: "repro_follow_up", sessionId, followUp, followUp2 });
186+
187+ const permanentlyBricked = isBrick(followUp) && isBrick(followUp2);
188+
189+ expect(getResult.status, "the concurrent SSE GET opens successfully").toBe(200);
190+ expect(postResult.status, "the concurrent POST answers successfully").toBe(200);
191+ expect(isBrick(getResult), "the concurrent SSE GET does not expose the transport brick").toBe(
192+ false,
193+ );
194+ expect(isBrick(postResult), "the concurrent POST does not expose the transport brick").toBe(
195+ false,
196+ );
197+ expect(followUp.status, "the first follow-up POST answers successfully").toBe(200);
198+ expect(followUp2.status, "the second follow-up POST answers successfully").toBe(200);
199+ expect(isBrick(followUp), "the first follow-up POST does not expose the transport brick").toBe(
200+ false,
201+ );
202+ expect(
203+ isBrick(followUp2),
204+ "the second follow-up POST does not expose the transport brick",
205+ ).toBe(false);
206+ expect(
207+ permanentlyBricked,
208+ "after the concurrent SSE-GET + POST post-idle, the session remains usable",
209+ ).toBe(false);
210+ // Diagnostic (not a gate): whether the collision surfaced on the concurrent
211+ // pair itself vs. only on the follow-up depends on which restore won the
212+ // race. Record it so the artifact shows the timing that occurred this run.
213+ record({ event: "repro_variant", concurrentBricked, permanentlyBricked });
214+ }),
215+);
packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts+134−6
@@ -1,12 +1,18 @@
11 import { describe, expect, it } from "@effect/vitest";
2−import { Effect } from "effect";
2+import { Cause, Effect } from "effect";
33 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
44 import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
55 import type { JSONRPCMessage, MessageExtraInfo } from "@modelcontextprotocol/sdk/types.js";
66
77 import { defaultMcpResource } from "@executor-js/host-mcp";
8+import type { ExecutionEngine, ExecutionResult, ResumeResponse } from "@executor-js/execution";
89
9−import { McpAgentSessionDOBase, type SessionMeta } from "./agent-session-durable-object";
10+import {
11+ McpAgentSessionDOBase,
12+ type McpApprovalOwner,
13+ type McpSessionModelResumeResult,
14+ type SessionMeta,
15+} from "./agent-session-durable-object";
1016
1117 class MemoryStorage {
1218 private readonly data = new Map<string, unknown>();
@@ -77,7 +83,7 @@ type HarnessSession = { · 64 unmodified lines
7783 alarm: () => Promise<void>;
7884 ctx: MemoryStorage;
7985 dbHandle: { readonly end: () => void } | null;
80− engine: { readonly pausedExecutionCount: () => Effect.Effect<number> } | null;
86+ engine: ExecutionEngine<Cause.YieldableError> | null;
8187 getSessionId: () => string;
8288 initialized: boolean;
8389 lastActivityMs: number;
@@ -85,9 +91,15 @@ type HarnessSession = { · 1 unmodified line
8591 onStart: () => Promise<void>;
8692 pendingApprovalLeases: Map<string, never>;
8793 props: Record<string, unknown>;
94+ runMcpAgentOnStart: () => Promise<void>;
8895 server?: McpServer;
8996 sessionMeta: SessionMeta;
9097 sessionTimeoutMs: () => number;
98+ resumeExecutionForModel: (
99+ executionId: string,
100+ identity: McpApprovalOwner,
101+ response: ResumeResponse,
102+ ) => Promise<McpSessionModelResumeResult>;
91103 validateMcpSessionOwner: (identity: {
92104 readonly accountId: string;
93105 readonly organizationId: string;
@@ -130,6 +142,44 @@ const makeDeferred = (): { readonly promise: Promise<void>; readonly resolve: () · 36 unmodified lines
130142 return { promise, resolve };
131143 };
132144
145+type ResumeCall = {
146+ readonly executionId: string;
147+ readonly response: ResumeResponse;
148+};
149+
150+const completed = (result: unknown): ExecutionResult => ({
151+ status: "completed",
152+ result: { result },
153+});
154+
155+const makeEngine = (
156+ resultForResume: (executionId: string, response: ResumeResponse) => ExecutionResult | null = () =>
157+ completed("resume-result"),
158+): { readonly calls: ResumeCall[]; readonly engine: ExecutionEngine<Cause.YieldableError> } => {
159+ const calls: ResumeCall[] = [];
160+ return {
161+ calls,
162+ engine: {
163+ execute: () => Effect.succeed({ result: "execute-result" }),
164+ executeWithPause: () => Effect.succeed(completed("execute-result")),
165+ resume: (executionId, response) =>
166+ Effect.sync(() => {
167+ calls.push({ executionId, response });
168+ return resultForResume(executionId, response);
169+ }),
170+ getPausedExecution: () => Effect.succeed(null),
171+ pausedExecutionCount: () => Effect.succeed(0),
172+ hasPausedExecutions: () => Effect.succeed(false),
173+ getDescription: Effect.succeed("test engine"),
174+ },
175+ };
176+};
177+
178+const approval = {
179+ action: "accept",
180+ content: { approved: true },
181+} satisfies ResumeResponse;
182+
133183 const makeHarnessSession = async (): Promise<HarnessSession> => {
134184 const sessionId = "session-reconnect";
135185 const sessionMeta: SessionMeta = {
@@ -145,7 +195,7 @@ const makeHarnessSession = async (): Promise<HarnessSession> => { · 9 unmodified lines
145195 const session = Object.create(McpAgentSessionDOBase.prototype) as HarnessSession;
146196 session.ctx = storage;
147197 session.dbHandle = { end: () => undefined };
148− session.engine = { pausedExecutionCount: () => Effect.succeed(0) };
198+ session.engine = makeEngine().engine;
149199 session.getSessionId = () => sessionId;
150200 session.initialized = true;
151201 session.lastActivityMs = Date.now() - 10;
@@ -155,10 +205,11 @@ const makeHarnessSession = async (): Promise<HarnessSession> => { · 3 unmodified lines
155205 session.server = server;
156206 session.sessionMeta = sessionMeta;
157207 session.sessionTimeoutMs = () => 1;
158− session.onStart = async () => {
208+ session.runMcpAgentOnStart = async () => {
159209 const restored = session.server ?? makeServer();
160210 session.server = restored;
161211 await restored.connect(new RestoredTransport());
212+ session.engine = makeEngine().engine;
162213 session.initialized = true;
163214 };
164215
@@ -183,7 +234,7 @@ describe("McpAgentSessionDOBase transport restore", () => { · 18 unmodified lines
183234 let onStartCalls = 0;
184235 let restoredServer: McpServer | undefined;
185236
186− session.onStart = async () => {
237+ session.runMcpAgentOnStart = async () => {
187238 onStartCalls += 1;
188239 const restored = session.server ?? makeServer();
189240 restoredServer ??= restored;
@@ -213,4 +264,81 @@ describe("McpAgentSessionDOBase transport restore", () => { · 23 unmodified lines
213264 expect(onStartCalls).toBe(1);
214265 expect(session.server).toBe(restoredServer);
215266 });
267+
268+ it("single-flights SDK onStart callers with same-session restore", async () => {
269+ const session = await makeHarnessSession();
270+ const firstStartEntered = makeDeferred();
271+ const finishStart = makeDeferred();
272+ let onStartCalls = 0;
273+
274+ session.runMcpAgentOnStart = async () => {
275+ onStartCalls += 1;
276+ const restored = session.server ?? makeServer();
277+ session.server = restored;
278+ firstStartEntered.resolve();
279+ await finishStart.promise;
280+ await restored.connect(new RestoredTransport());
281+ session.initialized = true;
282+ };
283+
284+ await session.alarm();
285+
286+ const restore = session.validateMcpSessionOwner({
287+ accountId: "user-1",
288+ organizationId: "org-1",
289+ });
290+ const sdkStart = session.onStart();
291+
292+ await firstStartEntered.promise;
293+ await Promise.resolve();
294+ finishStart.resolve();
295+
296+ await expect(Promise.all([restore, sdkStart])).resolves.toEqual(["ok", undefined]);
297+ expect(onStartCalls).toBe(1);
298+ });
299+
300+ it("single-flights model resume restore with SDK onStart", async () => {
301+ const session = await makeHarnessSession();
302+ const firstStartEntered = makeDeferred();
303+ const finishStart = makeDeferred();
304+ const restoredEngine = makeEngine(() => completed("model-result"));
305+ let onStartCalls = 0;
306+
307+ session.runMcpAgentOnStart = async () => {
308+ onStartCalls += 1;
309+ const restored = session.server ?? makeServer();
310+ session.server = restored;
311+ firstStartEntered.resolve();
312+ await finishStart.promise;
313+ await restored.connect(new RestoredTransport());
314+ session.engine = restoredEngine.engine;
315+ session.initialized = true;
316+ };
317+
318+ await session.alarm();
319+
320+ const resume = session.resumeExecutionForModel(
321+ "exec-model",
322+ { accountId: "user-1", organizationId: "org-1" },
323+ approval,
324+ );
325+ const sdkStart = session.onStart();
326+
327+ await firstStartEntered.promise;
328+ await Promise.resolve();
329+ finishStart.resolve();
330+
331+ const [resumeResult] = await Promise.all([resume, sdkStart]);
332+ expect(resumeResult).toMatchObject({
333+ status: "result",
334+ result: {
335+ structuredContent: {
336+ status: "completed",
337+ result: "model-result",
338+ },
339+ },
340+ });
341+ expect(onStartCalls).toBe(1);
342+ expect(restoredEngine.calls).toEqual([{ executionId: "exec-model", response: approval }]);
343+ });
216344 });
packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts+28−40
@@ -212,7 +212,7 @@ export abstract class McpAgentSessionDOBase< · 211 unmodified lines
212212 private dbHandle: TDbHandle | null = null;
213213 private sessionMeta: SessionMeta | null = null;
214214 private initialized = false;
215− private restoreTransportRuntimePromise: Promise<void> | null = null;
215+ private onStartPromise: Promise<void> | null = null;
216216 private lastActivityMs = 0;
217217 private approvalResponses = new Map<string, ResumeResponse>();
218218 private approvalWaiters = new Map<string, Deferred.Deferred<ResumeResponse>>();
@@ -556,55 +556,43 @@ export abstract class McpAgentSessionDOBase< · 337 unmodified lines
556556 const sessionMeta = yield* self.loadSessionMeta();
557557 if (!sessionMeta) return false;
558558
559− yield* self.closeRuntime();
560− const dbHandle = yield* self.openSessionDbHandle();
561− const { mcpServer, engine } = yield* self.buildRuntime(sessionMeta, dbHandle);
562− self.dbHandle = dbHandle;
563− self.server = mcpServer;
564− self.engine = engine;
565− self.initialized = true;
566− yield* Effect.promise(() => self.markActivity()).pipe(
567− Effect.withSpan("McpSessionDO.markActivity"),
559+ yield* Effect.promise(() => self.onStart()).pipe(
560+ Effect.withSpan("McpSessionDO.restore_runtime_for_approval"),
568561 );
569− return true;
562+ return self.initialized && !!self.engine;
570563 }).pipe(Effect.withSpan("McpSessionDO.ensure_runtime_for_approval"));
571564 }
572565
573− private restoreTransportRuntimeOnce(): Effect.Effect<void> {
566+ private startRuntimeFromOnStart(props?: McpSessionProps): Effect.Effect<void> {
574567 const self = this;
575568 return Effect.gen(function* () {
576569 yield* self.closeRuntime();
577− const restored = yield* Effect.exit(Effect.promise(() => self.onStart()));
578− if (Exit.isFailure(restored)) {
570+ const started = yield* Effect.exit(Effect.promise(() => self.runMcpAgentOnStart(props)));
571+ if (Exit.isFailure(started)) {
579572 yield* self.closeRuntime();
580− return yield* Effect.failCause(restored.cause);
573+ return yield* Effect.failCause(started.cause);
581574 }
582575 });
583576 }
584577
585− private restoreTransportRuntime(): Effect.Effect<void> {
586− const self = this;
587− return Effect.promise(() => {
588− if (self.restoreTransportRuntimePromise) return self.restoreTransportRuntimePromise;
578+ protected runMcpAgentOnStart(props?: McpSessionProps): Promise<void> {
579+ return super.onStart(props);
580+ }
589581
590− const restoring = Promise.resolve().then(() =>
591− Effect.runPromise(self.restoreTransportRuntimeOnce()),
592− );
593− self.restoreTransportRuntimePromise = restoring;
594− restoring.then(
595− () => {
596− if (self.restoreTransportRuntimePromise === restoring) {
597− self.restoreTransportRuntimePromise = null;
598− }
599− },
600− () => {
601− if (self.restoreTransportRuntimePromise === restoring) {
602− self.restoreTransportRuntimePromise = null;
603− }
604− },
605− );
606− return restoring;
607− });
582+ override async onStart(props?: McpSessionProps): Promise<void> {
583+ if (this.onStartPromise) return this.onStartPromise;
584+
585+ const starting = Effect.runPromise(this.startRuntimeFromOnStart(props));
586+ this.onStartPromise = starting;
587+ starting.then(
588+ () => {
589+ if (this.onStartPromise === starting) this.onStartPromise = null;
590+ },
591+ () => {
592+ if (this.onStartPromise === starting) this.onStartPromise = null;
593+ },
594+ );
595+ return starting;
608596 }
609597
610598 async init(): Promise<void> {
@@ -669,9 +657,9 @@ export abstract class McpAgentSessionDOBase< · 58 unmodified lines
669657 Effect.withSpan("McpSessionDO.markActivity"),
670658 );
671659 } else {
672− yield* self
673− .restoreTransportRuntime()
674− .pipe(Effect.withSpan("McpSessionDO.restore_transport_runtime"));
660+ yield* Effect.promise(() => self.onStart()).pipe(
661+ Effect.withSpan("McpSessionDO.restore_transport_runtime"),
662+ );
675663 }
676664 return identity.accountId === sessionMeta.userId &&
677665 identity.organizationId === sessionMeta.organizationId