History

Cold-init e2e: prove the session works and count the org reads deterministically

Rhys Sullivan committed 2 months ago

Commit c8490051 parent e6753341 file · +105 −14Browse files

Unified · Split

e2e/cloud/mcp-session-cold-init.test.ts+105−14
@@ -10,10 +10,27 @@ · 9 unmodified lines
1010 // client-visible failure on a healthy database, because the only unhealthy
1111 // thing was a socket nothing needed to open.
1212 //
13−// The contract asserted here is deliberately about the DATA PATH, not the
14−// failure: `McpSessionDOSqlite.resolveSessionMeta` reports where the org
15−// identity came from, and the whole request performs exactly ONE org read (the
16−// worker's own authorization check) instead of two.
13+// Two contracts, in the order a user meets them:
14+//
15+// 1. the session opens and WORKS — initialize mints a session id and the same
16+// id then serves `tools/list`, all off the identity the props carried;
17+// 2. the whole request performs exactly ONE organization read (the worker's
18+// own authorization check) and the DO opens no database connection of its
19+// own to name the org.
20+//
21+// (2) is asserted on the EXPORTED spans, and the two planes — worker and
22+// Durable Object — export independently, so the count is only taken once a
23+// worker-plane span for this same request has landed. Without that wait a
24+// still-pending worker batch would make a two-read request look like a one-read
25+// request, and the assertion would pass for the wrong reason.
26+//
27+// The failure half of the fix (an unreachable directory becomes a bounded retry
28+// and a retryable 503 rather than a bare 500) is not reachable from here: the
29+// harness runs one single-process PGlite shared by the worker and the DO, and
30+// the worker's own authorization reads that same row on the same request — so
31+// freezing the database fails the request before init identically on both sides
32+// of the fix. Those branches are pinned in
33+// `apps/cloud/src/mcp/session-meta.node.test.ts`.
1734 import { randomBytes } from "node:crypto";
1835
1936 import { expect } from "@effect/vitest";
@@ -36,6 +53,18 @@ const INITIALIZE_REQUEST = { · 16 unmodified lines
3653 },
3754 };
3855
56+const INITIALIZED_NOTIFICATION = {
57+ jsonrpc: "2.0" as const,
58+ method: "notifications/initialized",
59+};
60+
61+const TOOLS_LIST_REQUEST = {
62+ jsonrpc: "2.0" as const,
63+ id: 2,
64+ method: "tools/list",
65+ params: {},
66+};
67+
3968 const emailOf = (identity: Identity): string => identity.credentials?.email ?? identity.label;
4069
4170 /** A client-supplied W3C trace context, so every span this one request produces
@@ -46,6 +75,27 @@ const newTraceContext = (): { readonly traceId: string; readonly traceparent: st · 4 unmodified lines
4675 return { traceId, traceparent: `00-${traceId}-${spanId}-01` };
4776 };
4877
78+const mcpPost = (
79+ url: string,
80+ init: {
81+ readonly bearer: string;
82+ readonly sessionId?: string;
83+ readonly traceparent?: string;
84+ readonly body: unknown;
85+ },
86+): Promise<Response> =>
87+ fetch(url, {
88+ method: "POST",
89+ headers: {
90+ accept: JSON_AND_SSE,
91+ "content-type": "application/json",
92+ authorization: `Bearer ${init.bearer}`,
93+ ...(init.sessionId ? { "mcp-session-id": init.sessionId } : {}),
94+ ...(init.traceparent ? { traceparent: init.traceparent } : {}),
95+ },
96+ body: JSON.stringify(init.body),
97+ });
98+
4999 scenario(
50100 "MCP session cold init · the org identity rides in the session props instead of a second Postgres read",
51101 { timeout: 120_000 },
@@ -58,28 +108,56 @@ scenario( · 6 unmodified lines
58108 const bearer = yield* mcp.mintBearer(emailOf(identity));
59109 const trace = newTraceContext();
60110
111+ // ---- 1. the session opens, and it works -------------------------------
61112 const response = yield* Effect.promise(() =>
62− fetch(target.mcpUrl, {
63− method: "POST",
64− headers: {
65− accept: JSON_AND_SSE,
66− "content-type": "application/json",
67− authorization: `Bearer ${bearer}`,
68− traceparent: trace.traceparent,
69− },
70− body: JSON.stringify(INITIALIZE_REQUEST),
113+ mcpPost(target.mcpUrl, {
114+ bearer,
115+ traceparent: trace.traceparent,
116+ body: INITIALIZE_REQUEST,
71117 }),
72118 );
73119 yield* Effect.promise(() => response.text());
74120 expect(response.status, "initialize opens a session").toBe(200);
75− expect(response.headers.get("mcp-session-id"), "the session id is minted").toBeTruthy();
121+ const sessionId = response.headers.get("mcp-session-id");
122+ expect(sessionId, "the session id is minted").toBeTruthy();
76123
124+ const initialized = yield* Effect.promise(() =>
125+ mcpPost(target.mcpUrl, {
126+ bearer,
127+ sessionId: sessionId ?? "",
128+ body: INITIALIZED_NOTIFICATION,
129+ }),
130+ );
131+ yield* Effect.promise(() => initialized.text());
132+ expect(initialized.status, "the client completes the handshake").toBe(202);
133+
134+ // The session built from the props-carried identity actually serves work —
135+ // the production symptom was an `initialize` that never got this far.
136+ const tools = yield* Effect.promise(() =>
137+ mcpPost(target.mcpUrl, {
138+ bearer,
139+ sessionId: sessionId ?? "",
140+ body: TOOLS_LIST_REQUEST,
141+ }),
142+ );
143+ const toolsBody = yield* Effect.promise(() => tools.text());
144+ expect(tools.status, "the session serves requests once open").toBe(200);
145+ expect(toolsBody, "the session advertises the execute tool").toContain("execute");
146+
147+ // ---- 2. one organization read for the whole init request --------------
77148 // The DO really did resolve meta on this request (a cold init), and it
78149 // resolved it from the props the worker handed over.
79150 const resolveSpan = yield* telemetry.expectSpan({
80151 traceId: trace.traceId,
81152 operation: "McpSessionDOSqlite.resolveSessionMeta",
82153 });
154+ expect(resolveSpan.span.status, "the cold init resolved its meta without failing").toBe("ok");
155+
156+ // The worker plane exports on its own batch, independently of the DO's.
157+ // Wait for a worker-plane span from this same request before counting, or
158+ // an unflushed worker batch would hide the very read being counted.
159+ yield* telemetry.expectSpan({ traceId: trace.traceId, operation: "mcp.request" });
160+
83161 // One org read for the whole request: the worker's own authorization
84162 // lookup. A second one means the DO reopened a connection to re-read a row
85163 // the request already had.
@@ -92,6 +170,19 @@ scenario( · 6 unmodified lines
92170 "only the worker's authorization check reads the organization row",
93171 ).toBe(1);
94172
173+ // …and the DO's own database step never ran at all. `resolveSessionMeta`
174+ // has already landed and this span is its child, so its absence here is
175+ // absence, not lag.
176+ const doDatabaseReads = yield* telemetry.searchSpans({
177+ traceId: trace.traceId,
178+ operation: "mcp.session.resolve_organization",
179+ });
180+ expect(
181+ doDatabaseReads.length,
182+ "the session DO opens no connection of its own to name the organization",
183+ ).toBe(0);
184+
185+ // …because the identity it used is the one the worker handed it.
95186 expect(
96187 resolveSpan.span.tags["mcp.session.meta_source"],
97188 "the session meta comes from the props the worker already resolved",