History

Handle transient MCP auth and transport hangs (#445)

Rhys Sullivan committed 6 months ago

Commit 3f6028c1 parent 2909f8f8 files · +1049 −428Browse files

Unified · Split

8 files
  1. apps/cloud/src/mcp-flow.test.ts +77 −9
  2. apps/cloud/src/mcp-miniflare.e2e.node.test.ts +496 −196
  3. apps/cloud/src/mcp-session.ts +50 −8
  4. apps/cloud/src/mcp.ts +25 −212
  5. apps/cloud/src/mcp/response-peek.ts +203 −0
  6. apps/cloud/src/mcp/responses.ts +66 −0
  7. apps/cloud/src/services/mcp-worker-transport.ts +122 −1
  8. apps/cloud/src/test-worker.ts +10 −2
apps/cloud/src/mcp-flow.test.ts+77−9
@@ -24,12 +24,7 @@ · 23 unmodified lines
2424 // MCP session coverage lives in `mcp-miniflare.e2e.node.test.ts`.
2525 // ---------------------------------------------------------------------------
2626
27−import {
28− env,
29− runDurableObjectAlarm,
30− runInDurableObject,
31− SELF,
32−} from "cloudflare:test";
27+import { env, runDurableObjectAlarm, runInDurableObject, SELF } from "cloudflare:test";
3328 import { Effect } from "effect";
3429 import { afterAll, beforeAll, describe, expect, it } from "vitest";
3530
@@ -108,6 +103,16 @@ const mcpPost = (init: McpPostInit): Promise<Response> => { · 72 unmodified lines
108103 });
109104 };
110105
106+const mcpGet = (init: { readonly bearer: string; readonly sessionId: string }): Promise<Response> =>
107+ SELF.fetch(MCP_URL, {
108+ method: "GET",
109+ headers: {
110+ accept: "text/event-stream",
111+ authorization: `Bearer ${init.bearer}`,
112+ "mcp-session-id": init.sessionId,
113+ },
114+ });
115+
111116 const seedOrg = async (id: string, name = "MCP Flow Org"): Promise<void> => {
112117 const response = await SELF.fetch(`${BASE}/__test__/seed-org`, {
113118 method: "POST",
@@ -146,9 +151,7 @@ describe("/mcp CORS preflight", () => { · 32 unmodified lines
146151
147152 expect(response.status).toBe(204);
148153 expect(response.headers.get("access-control-allow-origin")).toBe("*");
149− expect(response.headers.get("access-control-allow-methods")).toBe(
150− "GET, POST, DELETE, OPTIONS",
151− );
154+ expect(response.headers.get("access-control-allow-methods")).toBe("GET, POST, DELETE, OPTIONS");
152155 const allowedHeaders = response.headers.get("access-control-allow-headers") ?? "";
153156 expect(allowedHeaders).toContain("mcp-session-id");
154157 expect(allowedHeaders).toContain("authorization");
@@ -217,6 +220,24 @@ describe("/mcp verified token without org", () => { · 62 unmodified lines
217220 });
218221 });
219222
223+describe("/mcp transient auth failure", () => {
224+ it("returns a retryable JSON-RPC error instead of a generic 500", async () => {
225+ const response = await mcpPost({
226+ bearer: "test-system-error",
227+ body: TOOLS_LIST_REQUEST,
228+ });
229+ expect(response.status).toBe(503);
230+ expect(response.headers.get("access-control-allow-origin")).toBe("*");
231+ const body = (await response.json()) as {
232+ jsonrpc: string;
233+ error: { code: number; message: string };
234+ };
235+ expect(body.jsonrpc).toBe("2.0");
236+ expect(body.error.code).toBe(-32001);
237+ expect(body.error.message).toMatch(/temporarily unavailable/i);
238+ });
239+});
240+
220241 describe("/mcp verified token without live org access", () => {
221242 it("returns JSON-RPC -32001 before creating a session", async () => {
222243 const response = await mcpPost({
@@ -330,6 +351,53 @@ describe("/mcp session restore", () => { · 107 unmodified lines
330351 expect(body.result?.tools?.some((tool) => tool.name === "execute")).toBe(true);
331352 }, 15_000);
332353
354+ it("keeps JSON POST responses after a session is restored by a GET reconnect", async () => {
355+ const orgId = nextOrgId();
356+ const accountId = nextAccountId();
357+ const bearer = makeTestBearer(accountId, orgId);
358+ await seedOrg(orgId);
359+
360+ const initializeResponse = await mcpPost({
361+ bearer,
362+ body: INITIALIZE_REQUEST,
363+ });
364+ expect(initializeResponse.status).toBe(200);
365+ const sessionId = initializeResponse.headers.get("mcp-session-id");
366+ expect(sessionId).toBeTruthy();
367+
368+ const ns = env.MCP_SESSION;
369+ const stub = ns.get(ns.idFromString(sessionId!));
370+ await runInDurableObject(stub, async (instance) => {
371+ await Effect.runPromise(
372+ (instance as unknown as { closeRuntime: () => Effect.Effect<void> }).closeRuntime(),
373+ );
374+ });
375+
376+ const getResponse = await mcpGet({ bearer, sessionId: sessionId! });
377+ expect(getResponse.status).toBe(200);
378+ expect(getResponse.headers.get("content-type") ?? "").toContain("text/event-stream");
379+ await getResponse.body?.cancel().catch(() => undefined);
380+
381+ const response = await Promise.race([
382+ mcpPost({
383+ bearer,
384+ sessionId,
385+ body: TOOLS_LIST_REQUEST,
386+ }),
387+ new Promise<never>((_, reject) =>
388+ setTimeout(() => reject(new Error("POST did not return after GET restore")), 5_000),
389+ ),
390+ ]);
391+ expect(response.status).toBe(200);
392+ expect(response.headers.get("content-type") ?? "").toContain("application/json");
393+ const body = (await response.json()) as {
394+ readonly jsonrpc: string;
395+ readonly result?: { readonly tools?: ReadonlyArray<{ readonly name: string }> };
396+ };
397+ expect(body.jsonrpc).toBe("2.0");
398+ expect(body.result?.tools?.some((tool) => tool.name === "execute")).toBe(true);
399+ }, 15_000);
400+
333401 it("reproduces cross-account session reuse via leaked mcp-session-id", async () => {
334402 const victimOrgId = nextOrgId();
335403 const attackerOrgId = nextOrgId();
apps/cloud/src/mcp-miniflare.e2e.node.test.ts+496−196
@@ -58,20 +58,14 @@ const ApproveGroup = HttpApiGroup.make("approve").add( · 57 unmodified lines
5858 const UpstreamApi = HttpApi.make("approveApi").add(ApproveGroup);
5959
6060 const ApproveHandlers = HttpApiBuilder.group(UpstreamApi, "approve", (h) =>
61− h.handle("approveThing", () =>
62− Effect.succeed(new ApprovedResponse({ approved: true })),
63− ),
61+ h.handle("approveThing", () => Effect.succeed(new ApprovedResponse({ approved: true }))),
6462 );
6563
66−const UpstreamApiLive = HttpApiBuilder.api(UpstreamApi).pipe(
67− Layer.provide(ApproveHandlers),
68−);
64+const UpstreamApiLive = HttpApiBuilder.api(UpstreamApi).pipe(Layer.provide(ApproveHandlers));
6965
7066 const UpstreamServeLayer = HttpApiBuilder.serve().pipe(
7167 Layer.provide(UpstreamApiLive),
72− Layer.provideMerge(
73− NodeHttpServer.layer(() => createServer(), { port: 0, host: "127.0.0.1" }),
74− ),
68+ Layer.provideMerge(NodeHttpServer.layer(() => createServer(), { port: 0, host: "127.0.0.1" })),
7569 );
7670
7771 // ---------------------------------------------------------------------------
@@ -187,7 +181,9 @@ const TelemetryReceiverLive = Layer.scoped( · 109 unmodified lines
187181 return;
188182 }
189183 let body = "";
190− req.on("data", (chunk) => { body += chunk; });
184+ req.on("data", (chunk) => {
185+ body += chunk;
186+ });
191187 req.on("end", () => {
192188 try {
193189 const payload = JSON.parse(body) as OtlpPayload;
@@ -232,10 +228,7 @@ const TelemetryReceiverLive = Layer.scoped( · 38 unmodified lines
232228 Effect.map((t) => ({
233229 tracesUrl: t.tracesUrl,
234230 spans: () => [...t.store],
235− waitForSpan: async (
236− predicate: (s: CapturedSpan) => boolean,
237− timeoutMs = 5_000,
238− ) => {
231+ waitForSpan: async (predicate: (s: CapturedSpan) => boolean, timeoutMs = 5_000) => {
239232 const deadline = Date.now() + timeoutMs;
240233 for (;;) {
241234 const hit = t.store.find(predicate);
@@ -322,218 +315,525 @@ const connectClient = async ( · 80 unmodified lines
322315 return client;
323316 };
324317
318+const initializeSession = async (baseUrl: URL, bearer: string): Promise<string> => {
319+ const response = await fetch(new URL("/mcp", baseUrl), {
320+ method: "POST",
321+ headers: {
322+ accept: "application/json, text/event-stream",
323+ authorization: `Bearer ${bearer}`,
324+ "content-type": "application/json",
325+ },
326+ body: JSON.stringify({
327+ jsonrpc: "2.0",
328+ id: "init",
329+ method: "initialize",
330+ params: {
331+ protocolVersion: "2025-11-25",
332+ capabilities: {},
333+ clientInfo: { name: "mcp-miniflare-manual", version: "0" },
334+ },
335+ }),
336+ });
337+ expect(response.status).toBe(200);
338+ const sessionId = response.headers.get("mcp-session-id");
339+ expect(sessionId).toEqual(expect.any(String));
340+ await response.text();
341+
342+ const initialized = await fetch(new URL("/mcp", baseUrl), {
343+ method: "POST",
344+ headers: {
345+ accept: "application/json, text/event-stream",
346+ authorization: `Bearer ${bearer}`,
347+ "content-type": "application/json",
348+ "mcp-session-id": sessionId ?? "",
349+ },
350+ body: JSON.stringify({
351+ jsonrpc: "2.0",
352+ method: "notifications/initialized",
353+ }),
354+ });
355+ expect(initialized.status).toBe(202);
356+ await initialized.text();
357+
358+ return sessionId ?? "";
359+};
360+
325361 // ---------------------------------------------------------------------------
326362 // Tests
327363 // ---------------------------------------------------------------------------
328364
329365 layer(TestEnv, { timeout: 60_000 })("cloud MCP over real HTTP (miniflare)", (it) => {
330− it.effect("returns 401 for malformed bearer tokens through the production auth layer", () =>
331− Effect.gen(function* () {
332− const { baseUrl } = yield* Worker;
333− const response = yield* Effect.promise(() =>
334− fetch(new URL("/__test__/real-auth-mcp", baseUrl), {
335− method: "POST",
336− headers: {
337− accept: "application/json, text/event-stream",
338− authorization: "Bearer bogus",
339− "content-type": "application/json",
340− },
341− body: JSON.stringify({
342− jsonrpc: "2.0",
343− id: 1,
344− method: "initialize",
345− params: {
346− protocolVersion: "2025-06-18",
347− capabilities: {},
348− clientInfo: { name: "mcp-miniflare-invalid-bearer", version: "0" },
366+ it.effect(
367+ "returns 401 for malformed bearer tokens through the production auth layer",
368+ () =>
369+ Effect.gen(function* () {
370+ const { baseUrl } = yield* Worker;
371+ const response = yield* Effect.promise(() =>
372+ fetch(new URL("/__test__/real-auth-mcp", baseUrl), {
373+ method: "POST",
374+ headers: {
375+ accept: "application/json, text/event-stream",
376+ authorization: "Bearer bogus",
377+ "content-type": "application/json",
349378 },
379+ body: JSON.stringify({
380+ jsonrpc: "2.0",
381+ id: 1,
382+ method: "initialize",
383+ params: {
384+ protocolVersion: "2025-06-18",
385+ capabilities: {},
386+ clientInfo: { name: "mcp-miniflare-invalid-bearer", version: "0" },
387+ },
388+ }),
350389 }),
351− }),
352− );
353− expect(response.status).toBe(401);
354− const wwwAuth = response.headers.get("www-authenticate") ?? "";
355− expect(wwwAuth).toContain('Bearer error="invalid_token"');
356− expect(wwwAuth).toContain('error_description="The access token is invalid"');
357− expect(wwwAuth).toContain(
358− "https://test-resource.example.com/.well-known/oauth-protected-resource/mcp",
359− );
360− const body = yield* Effect.promise(() => response.json());
361− expect(body).toEqual({ error: "unauthorized" });
362− }), 30_000,
390+ );
391+ expect(response.status).toBe(401);
392+ const wwwAuth = response.headers.get("www-authenticate") ?? "";
393+ expect(wwwAuth).toContain('Bearer error="invalid_token"');
394+ expect(wwwAuth).toContain('error_description="The access token is invalid"');
395+ expect(wwwAuth).toContain(
396+ "https://test-resource.example.com/.well-known/oauth-protected-resource/mcp",
397+ );
398+ const body = yield* Effect.promise(() => response.json());
399+ expect(body).toEqual({ error: "unauthorized" });
400+ }),
401+ 30_000,
363402 );
364403
365− it.effect("completes the initialize handshake via SDK", () =>
366− Effect.gen(function* () {
367− const { baseUrl, seedOrg } = yield* Worker;
368− const orgId = nextOrgId();
369− yield* Effect.promise(() => seedOrg(orgId, "Miniflare Org"));
370− const client = yield* Effect.promise(() =>
371− connectClient(baseUrl, makeTestBearer(nextAccountId(), orgId)),
372− );
373− expect(client.getServerVersion()?.name).toBe("executor");
374− yield* Effect.promise(() => client.close());
375− }), 30_000,
404+ it.effect(
405+ "completes the initialize handshake via SDK",
406+ () =>
407+ Effect.gen(function* () {
408+ const { baseUrl, seedOrg } = yield* Worker;
409+ const orgId = nextOrgId();
410+ yield* Effect.promise(() => seedOrg(orgId, "Miniflare Org"));
411+ const client = yield* Effect.promise(() =>
412+ connectClient(baseUrl, makeTestBearer(nextAccountId(), orgId)),
413+ );
414+ expect(client.getServerVersion()?.name).toBe("executor");
415+ yield* Effect.promise(() => client.close());
416+ }),
417+ 30_000,
376418 );
377419
378− it.effect("lists the execute tool after handshake", () =>
379− Effect.gen(function* () {
380− const { baseUrl, seedOrg } = yield* Worker;
381− const orgId = nextOrgId();
382− yield* Effect.promise(() => seedOrg(orgId, "List Tools Org"));
383− const client = yield* Effect.promise(() =>
384− connectClient(baseUrl, makeTestBearer(nextAccountId(), orgId)),
385− );
386− const { tools } = yield* Effect.promise(() => client.listTools());
387− expect(tools.map((t) => t.name)).toContain("execute");
388− yield* Effect.promise(() => client.close());
389− }), 30_000,
420+ it.effect(
421+ "lists the execute tool after handshake",
422+ () =>
423+ Effect.gen(function* () {
424+ const { baseUrl, seedOrg } = yield* Worker;
425+ const orgId = nextOrgId();
426+ yield* Effect.promise(() => seedOrg(orgId, "List Tools Org"));
427+ const client = yield* Effect.promise(() =>
428+ connectClient(baseUrl, makeTestBearer(nextAccountId(), orgId)),
429+ );
430+ const { tools } = yield* Effect.promise(() => client.listTools());
431+ expect(tools.map((t) => t.name)).toContain("execute");
432+ yield* Effect.promise(() => client.close());
433+ }),
434+ 30_000,
390435 );
391436
392− it.effect("executes code end-to-end via tools/call", () =>
393− Effect.gen(function* () {
394− const { baseUrl, seedOrg } = yield* Worker;
395− const orgId = nextOrgId();
396− yield* Effect.promise(() => seedOrg(orgId, "Execute Org"));
397− const client = yield* Effect.promise(() =>
398− connectClient(baseUrl, makeTestBearer(nextAccountId(), orgId)),
399− );
400− const result = yield* Effect.promise(() =>
401− client.callTool({ name: "execute", arguments: { code: "return 1 + 2" } }),
402− );
403− expect(result.isError).not.toBe(true);
404− const content = (result.content ?? []) as Array<{ type: string; text?: string }>;
405− expect(content.find((c) => c.type === "text")?.text ?? "").toContain("3");
406− yield* Effect.promise(() => client.close());
407− }), 30_000,
437+ it.effect(
438+ "executes code end-to-end via tools/call",
439+ () =>
440+ Effect.gen(function* () {
441+ const { baseUrl, seedOrg } = yield* Worker;
442+ const orgId = nextOrgId();
443+ yield* Effect.promise(() => seedOrg(orgId, "Execute Org"));
444+ const client = yield* Effect.promise(() =>
445+ connectClient(baseUrl, makeTestBearer(nextAccountId(), orgId)),
446+ );
447+ const result = yield* Effect.promise(() =>
448+ client.callTool({ name: "execute", arguments: { code: "return 1 + 2" } }),
449+ );
450+ expect(result.isError).not.toBe(true);
451+ const content = (result.content ?? []) as Array<{ type: string; text?: string }>;
452+ expect(content.find((c) => c.type === "text")?.text ?? "").toContain("3");
453+ yield* Effect.promise(() => client.close());
454+ }),
455+ 30_000,
408456 );
409457
410− it.effect("round-trips approval elicitation for a POST openapi operation", () =>
411− Effect.gen(function* () {
412− const { baseUrl, seedOrg } = yield* Worker;
413− const { specJson } = yield* Upstream;
414− const orgId = nextOrgId();
415− yield* Effect.promise(() => seedOrg(orgId, "Elicit Org"));
416−
417− const client = yield* Effect.promise(() =>
418− connectClient(baseUrl, makeTestBearer(nextAccountId(), orgId), {
419− withElicitation: true,
420− }),
421− );
458+ it.effect(
459+ "replaces duplicate standalone SSE GET streams for the same session",
460+ () =>
461+ Effect.gen(function* () {
462+ const { baseUrl, seedOrg } = yield* Worker;
463+ const orgId = nextOrgId();
464+ const bearer = makeTestBearer(nextAccountId(), orgId);
465+ yield* Effect.promise(() => seedOrg(orgId, "Duplicate SSE Org"));
466+ const sessionId = yield* Effect.promise(() => initializeSession(baseUrl, bearer));
467+
468+ const getHeaders = {
469+ accept: "text/event-stream",
470+ authorization: `Bearer ${bearer}`,
471+ "mcp-protocol-version": "2025-11-25",
472+ "mcp-session-id": sessionId,
473+ };
474+ const first = yield* Effect.promise(() =>
475+ fetch(new URL("/mcp", baseUrl), { method: "GET", headers: getHeaders }),
476+ );
477+ expect(first.status).toBe(200);
422478
423− let elicitCount = 0;
424− client.setRequestHandler(ElicitRequestSchema, async () => {
425− elicitCount++;
426− return { action: "accept" as const, content: {} };
427− });
479+ const second = yield* Effect.promise(() =>
480+ fetch(new URL("/mcp", baseUrl), { method: "GET", headers: getHeaders }),
481+ );
482+ expect(second.status).toBe(200);
483+ expect(second.headers.get("content-type") ?? "").toContain("text/event-stream");
428484
429− // User code inside `execute` (1) registers the upstream as an OpenAPI
430− // source and (2) invokes its POST operation. `annotationsForOperation`
431− // marks the POST as `requiresApproval: true`, which fires
432− // `enforceApproval` in the executor; that goes through the MCP
433− // elicitation handler and lands on `client.setRequestHandler` above.
434− // Tool id is `<namespace>.<group>.<operation>` — Effect's
435− // `HttpApiGroup` name ("approve") becomes part of the sandbox path,
436− // so the invocation reads `tools.approveapi.approve.approveThing`.
437− const code = [
438− `await tools.openapi.addSource({ spec: ${JSON.stringify(specJson)}, namespace: "approveapi" });`,
439− `return await tools.approveapi.approve.approveThing({});`,
440− ].join("\n");
441− const result = yield* Effect.promise(() =>
442− client.callTool({ name: "execute", arguments: { code } }),
443− );
444− expect(result.isError).not.toBe(true);
445− expect(elicitCount).toBeGreaterThan(0);
446− const text = ((result.content ?? []) as Array<{ type: string; text?: string }>)
447− .find((c) => c.type === "text")?.text ?? "";
448− expect(text).toContain("approved");
449−
450− yield* Effect.promise(() => client.close());
451− }), 30_000,
485+ yield* Effect.promise(async () => {
486+ await first.body?.cancel().catch(() => undefined);
487+ });
488+ yield* Effect.promise(async () => {
489+ await second.body?.cancel().catch(() => undefined);
490+ });
491+ }),
492+ 30_000,
452493 );
453494
454− it.effect("reports the McpSessionDO.handleRequest span via the OTLP exporter", () =>
455− Effect.gen(function* () {
456− const { baseUrl, seedOrg } = yield* Worker;
457− const receiver = yield* TelemetryReceiver;
458− const orgId = nextOrgId();
459− yield* Effect.promise(() => seedOrg(orgId, "Telemetry Org"));
460− const client = yield* Effect.promise(() =>
461− connectClient(baseUrl, makeTestBearer(nextAccountId(), orgId)),
462− );
463− // Trigger the DO through a multi-step flow so we can assert that
464− // handleRequest spans are reported for every DO hit, not just init.
465− yield* Effect.promise(() => client.listTools());
466− yield* Effect.promise(() =>
467− client.callTool({ name: "execute", arguments: { code: "return 1 + 2" } }),
468− );
469− yield* Effect.promise(() => client.close());
470−
471− // The initialize POST carries no session-id; subsequent requests do.
472− // Assert on one of the session-id'd handleRequest spans so we verify
473− // attribute propagation beyond the degenerate init case.
474− const handleSpan = yield* Effect.promise(() =>
475− receiver.waitForSpan(
476− (s) =>
477− s.name === "McpSessionDO.handleRequest" &&
478− s.attributes["mcp.request.session_id_present"] === true,
479− ),
480− );
481− expect(handleSpan.attributes["mcp.request.method"]).toBeDefined();
482− // 200 for normal POSTs, 202 for notifications/initialized.
483− expect([200, 202]).toContain(handleSpan.attributes["mcp.response.status_code"]);
484− expect(handleSpan.attributes["mcp.response.content_type"]).toEqual(expect.any(String));
485− expect(handleSpan.attributes["mcp.transport.enable_json_response"]).toBe(true);
486−
487− // init runs once per new session and should appear on the initialize POST.
488− yield* Effect.promise(() =>
489− receiver.waitForSpan((s) => s.name === "McpSessionDO.init"),
490− );
491− }), 30_000,
492− );
495+ it.effect(
496+ "does not replace the active standalone SSE stream for an invalid GET",
497+ () =>
498+ Effect.gen(function* () {
499+ const { baseUrl, seedOrg } = yield* Worker;
500+ const orgId = nextOrgId();
501+ const bearer = makeTestBearer(nextAccountId(), orgId);
502+ yield* Effect.promise(() => seedOrg(orgId, "Invalid SSE Replacement Org"));
503+ const sessionId = yield* Effect.promise(() => initializeSession(baseUrl, bearer));
504+
505+ const getHeaders = {
506+ accept: "text/event-stream",
507+ authorization: `Bearer ${bearer}`,
508+ "mcp-protocol-version": "2025-11-25",
509+ "mcp-session-id": sessionId,
510+ };
511+ const first = yield* Effect.promise(() =>
512+ fetch(new URL("/mcp", baseUrl), { method: "GET", headers: getHeaders }),
513+ );
514+ expect(first.status).toBe(200);
515+ const firstReader = first.body?.getReader();
516+ expect(firstReader).toBeDefined();
517+
518+ const invalid = yield* Effect.promise(() =>
519+ fetch(new URL("/mcp", baseUrl), {
520+ method: "GET",
521+ headers: { ...getHeaders, "mcp-protocol-version": "1999-01-01" },
522+ }),
523+ );
524+ expect(invalid.status).toBe(400);
525+ yield* Effect.promise(() => invalid.text());
526+
527+ const firstRead = yield* Effect.promise(() =>
528+ Promise.race([
529+ firstReader!.read().then(() => "closed"),
530+ new Promise<"open">((resolve) => setTimeout(() => resolve("open"), 100)),
531+ ]),
532+ );
533+ expect(firstRead).toBe("open");
493534
494−});
535+ yield* Effect.promise(async () => {
536+ await firstReader?.cancel().catch(() => undefined);
537+ });
538+ }),
539+ 30_000,
540+ );
495541
496−layer(TestEnv, { timeout: 60_000 })("cloud MCP request-id telemetry", (it) => {
497− it.effect("exports MCP JSON-RPC request ids for request-shaped ids", () =>
498− Effect.gen(function* () {
499− const { baseUrl, seedOrg } = yield* Worker;
500− const receiver = yield* TelemetryReceiver;
501− const orgId = nextOrgId();
502− const accountId = nextAccountId();
503− const requestId = `req_${crypto.randomUUID().replace(/-/g, "")}`;
504− yield* Effect.promise(() => seedOrg(orgId, "Request Id Org"));
505−
506− const response = yield* Effect.promise(() =>
507− fetch(new URL("/mcp", baseUrl), {
542+ it.effect(
543+ "returns tools/call results while standalone SSE GET reconnects churn",
544+ () =>
545+ Effect.gen(function* () {
546+ const { baseUrl, seedOrg } = yield* Worker;
547+ const orgId = nextOrgId();
548+ const bearer = makeTestBearer(nextAccountId(), orgId);
549+ yield* Effect.promise(() => seedOrg(orgId, "SSE Reconnect Churn Org"));
550+ const sessionId = yield* Effect.promise(() => initializeSession(baseUrl, bearer));
551+
552+ const getHeaders = {
553+ accept: "text/event-stream",
554+ authorization: `Bearer ${bearer}`,
555+ "mcp-protocol-version": "2025-11-25",
556+ "mcp-session-id": sessionId,
557+ };
558+
559+ const openSse = () =>
560+ fetch(new URL("/mcp", baseUrl), { method: "GET", headers: getHeaders });
561+
562+ const postResult = fetch(new URL("/mcp", baseUrl), {
508563 method: "POST",
509564 headers: {
510565 accept: "application/json, text/event-stream",
511− authorization: `Bearer ${makeTestBearer(accountId, orgId)}`,
566+ authorization: `Bearer ${bearer}`,
512567 "content-type": "application/json",
568+ "mcp-protocol-version": "2025-11-25",
569+ "mcp-session-id": sessionId,
513570 },
514571 body: JSON.stringify({
515572 jsonrpc: "2.0",
516− id: requestId,
517− method: "initialize",
573+ id: "search",
574+ method: "tools/call",
518575 params: {
519− protocolVersion: "2025-06-18",
520− capabilities: {},
521− clientInfo: { name: "mcp-request-id-e2e", version: "0" },
576+ name: "execute",
577+ arguments: {
578+ code: [
579+ "await new Promise((resolve) => setTimeout(resolve, 1_000));",
580+ 'return await tools.search({ namespace: "vercel_api", query: "list domains", limit: 8 });',
581+ ].join("\n"),
582+ },
522583 },
523584 }),
524− }),
525− );
526− expect(response.status).toBe(200);
527− yield* Effect.promise(() => response.text());
528−
529− const annotateSpan = yield* Effect.promise(() =>
530− receiver.waitForSpan(
531− (s) =>
532− s.name === "mcp.request.annotate" &&
533− s.attributes["mcp.rpc.id"] === requestId,
534− ),
535− );
536− expect(annotateSpan.attributes["mcp.rpc.method"]).toBe("initialize");
537− }), 30_000,
585+ });
586+
587+ const reconnects = yield* Effect.promise(async () => {
588+ const responses: Array<Response> = [];
589+ for (let i = 0; i < 35; i++) {
590+ const response = await openSse();
591+ expect(response.status).toBe(200);
592+ responses.push(response);
593+ await new Promise((resolve) => setTimeout(resolve, 10));
594+ }
595+ return responses;
596+ });
597+
598+ const response = yield* Effect.promise(() =>
599+ Promise.race([
600+ postResult,
601+ new Promise<never>((_, reject) =>
602+ setTimeout(
603+ () => reject(new Error("tools/call did not return during SSE churn")),
604+ 5_000,
605+ ),
606+ ),
607+ ]),
608+ );
609+ expect(response.status).toBe(200);
610+ const body = (yield* Effect.promise(() => response.json())) as {
611+ readonly jsonrpc?: string;
612+ readonly id?: string;
613+ readonly result?: { readonly content?: ReadonlyArray<{ readonly text?: string }> };
614+ readonly error?: unknown;
615+ };
616+ expect(body).toMatchObject({ jsonrpc: "2.0", id: "search" });
617+ expect(body.error).toBeUndefined();
618+ expect(body.result).toBeDefined();
619+
620+ yield* Effect.promise(async () => {
621+ await Promise.all(
622+ reconnects.map((response) => response.body?.cancel().catch(() => undefined)),
623+ );
624+ });
625+ }),
626+ 30_000,
627+ );
628+
629+ it.effect(
630+ "returns both overlapping tools/call responses when JSON-RPC ids collide",
631+ () =>
632+ Effect.gen(function* () {
633+ const { baseUrl, seedOrg } = yield* Worker;
634+ const orgId = nextOrgId();
635+ const bearer = makeTestBearer(nextAccountId(), orgId);
636+ yield* Effect.promise(() => seedOrg(orgId, "Overlapping Request Id Org"));
637+ const sessionId = yield* Effect.promise(() => initializeSession(baseUrl, bearer));
638+
639+ const postExecute = (code: string) =>
640+ fetch(new URL("/mcp", baseUrl), {
641+ method: "POST",
642+ headers: {
643+ accept: "application/json, text/event-stream",
644+ authorization: `Bearer ${bearer}`,
645+ "content-type": "application/json",
646+ "mcp-protocol-version": "2025-11-25",
647+ "mcp-session-id": sessionId,
648+ },
649+ body: JSON.stringify({
650+ jsonrpc: "2.0",
651+ id: 1,
652+ method: "tools/call",
653+ params: { name: "execute", arguments: { code } },
654+ }),
655+ });
656+
657+ const first = postExecute(
658+ ["await new Promise((resolve) => setTimeout(resolve, 500));", 'return "first";'].join(
659+ "\n",
660+ ),
661+ );
662+ yield* Effect.promise(() => new Promise((resolve) => setTimeout(resolve, 50)));
663+ const second = postExecute('return "second";');
664+
665+ const responses = yield* Effect.promise(() =>
666+ Promise.race([
667+ Promise.all([first, second]),
668+ new Promise<never>((_, reject) =>
669+ setTimeout(
670+ () => reject(new Error("overlapping tools/call requests did not both return")),
671+ 5_000,
672+ ),
673+ ),
674+ ]),
675+ );
676+
677+ expect(responses.map((response) => response.status)).toEqual([200, 200]);
678+ const bodies = yield* Effect.promise(() =>
679+ Promise.all(
680+ responses.map(
681+ (response) =>
682+ response.json() as Promise<{
683+ readonly result?: {
684+ readonly content?: ReadonlyArray<{ readonly text?: string }>;
685+ };
686+ readonly error?: unknown;
687+ }>,
688+ ),
689+ ),
690+ );
691+ expect(
692+ bodies.some((body) => body.result?.content?.some((item) => item.text?.includes("first"))),
693+ ).toBe(true);
694+ expect(
695+ bodies.some((body) =>
696+ body.result?.content?.some((item) => item.text?.includes("second")),
697+ ),
698+ ).toBe(true);
699+ expect(bodies.every((body) => body.error === undefined)).toBe(true);
700+ }),
701+ 30_000,
702+ );
703+
704+ it.effect(
705+ "round-trips approval elicitation for a POST openapi operation",
706+ () =>
707+ Effect.gen(function* () {
708+ const { baseUrl, seedOrg } = yield* Worker;
709+ const { specJson } = yield* Upstream;
710+ const orgId = nextOrgId();
711+ yield* Effect.promise(() => seedOrg(orgId, "Elicit Org"));
712+
713+ const client = yield* Effect.promise(() =>
714+ connectClient(baseUrl, makeTestBearer(nextAccountId(), orgId), {
715+ withElicitation: true,
716+ }),
717+ );
718+
719+ let elicitCount = 0;
720+ client.setRequestHandler(ElicitRequestSchema, async () => {
721+ elicitCount++;
722+ return { action: "accept" as const, content: {} };
723+ });
724+
725+ // User code inside `execute` (1) registers the upstream as an OpenAPI
726+ // source and (2) invokes its POST operation. `annotationsForOperation`
727+ // marks the POST as `requiresApproval: true`, which fires
728+ // `enforceApproval` in the executor; that goes through the MCP
729+ // elicitation handler and lands on `client.setRequestHandler` above.
730+ // Tool id is `<namespace>.<group>.<operation>` — Effect's
731+ // `HttpApiGroup` name ("approve") becomes part of the sandbox path,
732+ // so the invocation reads `tools.approveapi.approve.approveThing`.
733+ const code = [
734+ `await tools.openapi.addSource({ spec: ${JSON.stringify(specJson)}, namespace: "approveapi" });`,
735+ `return await tools.approveapi.approve.approveThing({});`,
736+ ].join("\n");
737+ const result = yield* Effect.promise(() =>
738+ client.callTool({ name: "execute", arguments: { code } }),
739+ );
740+ expect(result.isError).not.toBe(true);
741+ expect(elicitCount).toBeGreaterThan(0);
742+ const text =
743+ ((result.content ?? []) as Array<{ type: string; text?: string }>).find(
744+ (c) => c.type === "text",
745+ )?.text ?? "";
746+ expect(text).toContain("approved");
747+
748+ yield* Effect.promise(() => client.close());
749+ }),
750+ 30_000,
751+ );
752+
753+ it.effect(
754+ "reports the McpSessionDO.handleRequest span via the OTLP exporter",
755+ () =>
756+ Effect.gen(function* () {
757+ const { baseUrl, seedOrg } = yield* Worker;
758+ const receiver = yield* TelemetryReceiver;
759+ const orgId = nextOrgId();
760+ yield* Effect.promise(() => seedOrg(orgId, "Telemetry Org"));
761+ const client = yield* Effect.promise(() =>
762+ connectClient(baseUrl, makeTestBearer(nextAccountId(), orgId)),
763+ );
764+ // Trigger the DO through a multi-step flow so we can assert that
765+ // handleRequest spans are reported for every DO hit, not just init.
766+ yield* Effect.promise(() => client.listTools());
767+ yield* Effect.promise(() =>
768+ client.callTool({ name: "execute", arguments: { code: "return 1 + 2" } }),
769+ );
770+ yield* Effect.promise(() => client.close());
771+
772+ // The initialize POST carries no session-id; subsequent requests do.
773+ // Assert on one of the session-id'd handleRequest spans so we verify
774+ // attribute propagation beyond the degenerate init case.
775+ const handleSpan = yield* Effect.promise(() =>
776+ receiver.waitForSpan(
777+ (s) =>
778+ s.name === "McpSessionDO.handleRequest" &&
779+ s.attributes["mcp.request.session_id_present"] === true,
780+ ),
781+ );
782+ expect(handleSpan.attributes["mcp.request.method"]).toBeDefined();
783+ // 200 for normal POSTs, 202 for notifications/initialized.
784+ expect([200, 202]).toContain(handleSpan.attributes["mcp.response.status_code"]);
785+ expect(handleSpan.attributes["mcp.response.content_type"]).toEqual(expect.any(String));
786+ expect(handleSpan.attributes["mcp.transport.enable_json_response"]).toBe(true);
787+
788+ // init runs once per new session and should appear on the initialize POST.
789+ yield* Effect.promise(() => receiver.waitForSpan((s) => s.name === "McpSessionDO.init"));
790+ }),
791+ 30_000,
792+ );
793+});
794+
795+layer(TestEnv, { timeout: 60_000 })("cloud MCP request-id telemetry", (it) => {
796+ it.effect(
797+ "exports MCP JSON-RPC request ids for request-shaped ids",
798+ () =>
799+ Effect.gen(function* () {
800+ const { baseUrl, seedOrg } = yield* Worker;
801+ const receiver = yield* TelemetryReceiver;
802+ const orgId = nextOrgId();
803+ const accountId = nextAccountId();
804+ const requestId = `req_${crypto.randomUUID().replace(/-/g, "")}`;
805+ yield* Effect.promise(() => seedOrg(orgId, "Request Id Org"));
806+
807+ const response = yield* Effect.promise(() =>
808+ fetch(new URL("/mcp", baseUrl), {
809+ method: "POST",
810+ headers: {
811+ accept: "application/json, text/event-stream",
812+ authorization: `Bearer ${makeTestBearer(accountId, orgId)}`,
813+ "content-type": "application/json",
814+ },
815+ body: JSON.stringify({
816+ jsonrpc: "2.0",
817+ id: requestId,
818+ method: "initialize",
819+ params: {
820+ protocolVersion: "2025-06-18",
821+ capabilities: {},
822+ clientInfo: { name: "mcp-request-id-e2e", version: "0" },
823+ },
824+ }),
825+ }),
826+ );
827+ expect(response.status).toBe(200);
828+ yield* Effect.promise(() => response.text());
829+
830+ const annotateSpan = yield* Effect.promise(() =>
831+ receiver.waitForSpan(
832+ (s) => s.name === "mcp.request.annotate" && s.attributes["mcp.rpc.id"] === requestId,
833+ ),
834+ );
835+ expect(annotateSpan.attributes["mcp.rpc.method"]).toBe("initialize");
836+ }),
837+ 30_000,
538838 );
539839 });
apps/cloud/src/mcp-session.ts+50−8
@@ -335,6 +335,23 @@ export class McpSessionDO extends DurableObject { · 334 unmodified lines
335335 }).pipe(Effect.orDie);
336336 }
337337
338+ private installRuntime(
339+ sessionMeta: SessionMeta,
340+ options: {
341+ readonly dbHandle: DbHandle;
342+ readonly enableJsonResponse: boolean;
343+ },
344+ ) {
345+ const self = this;
346+ return Effect.gen(function* () {
347+ const runtime = yield* self.createConnectedRuntime(sessionMeta, options);
348+ self.dbHandle = options.dbHandle;
349+ self.mcpServer = runtime.mcpServer;
350+ self.transport = runtime.transport;
351+ self.initialized = true;
352+ });
353+ }
354+
338355 private restoreRuntimeFromStorage(request: Request): Effect.Effect<"restored" | "missing_meta"> {
339356 const self = this;
340357 return Effect.gen(function* () {
@@ -349,14 +366,14 @@ export class McpSessionDO extends DurableObject { · 8 unmodified lines
349366 }
350367
351368 yield* self.closeRuntime();
352− self.dbHandle = makeLongLivedDb();
353− const runtime = yield* self.createConnectedRuntime(sessionMeta, {
354− dbHandle: self.dbHandle,
355− enableJsonResponse: request.method !== "GET",
369+ const dbHandle = makeLongLivedDb();
370+ yield* self.installRuntime(sessionMeta, {
371+ dbHandle,
372+ // GET always returns an SSE stream regardless of this option, but the
373+ // session-scoped transport is reused by later POSTs. Keep JSON mode on
374+ // across cold restores so a GET reconnect cannot poison future POSTs.
375+ enableJsonResponse: true,
356376 });
357− self.mcpServer = runtime.mcpServer;
358− self.transport = runtime.transport;
359− self.initialized = true;
360377 yield* Effect.promise(() => self.markActivity()).pipe(
361378 Effect.withSpan("McpSessionDO.markActivity"),
362379 );
@@ -375,6 +392,26 @@ export class McpSessionDO extends DurableObject { · 12 unmodified lines
375392 );
376393 }
377394
395+ private ensureJsonResponseTransportForPost(request: Request): Effect.Effect<void> {
396+ const self = this;
397+ return Effect.gen(function* () {
398+ if (request.method !== "POST" || self.transportJsonResponseMode === true) return;
399+
400+ const sessionMeta = yield* self.loadSessionMeta();
401+ if (!sessionMeta) return;
402+
403+ yield* self.closeRuntime();
404+ const dbHandle = makeLongLivedDb();
405+ yield* self.installRuntime(sessionMeta, {
406+ dbHandle,
407+ enableJsonResponse: true,
408+ });
409+ yield* Effect.annotateCurrentSpan({
410+ "mcp.session.transport_upgraded_json_response": true,
411+ });
412+ }).pipe(Effect.withSpan("McpSessionDO.ensureJsonResponseTransportForPost"), Effect.orDie);
413+ }
414+
378415 private validateSessionOwner(request: Request): Effect.Effect<Response | null> {
379416 return Effect.gen(this, function* () {
380417 const sessionMeta = yield* this.loadSessionMeta();
@@ -551,9 +588,14 @@ export class McpSessionDO extends DurableObject { · 170 unmodified lines
551588 });
552589 }
553590
554− const transport = this.transport;
555591 const self = this;
556592 return Effect.gen(function* () {
593+ yield* self.ensureJsonResponseTransportForPost(request);
594+ const transport = self.transport;
595+ if (!transport) {
596+ return jsonRpcError(404, -32001, "Session timed out due to inactivity — please reconnect");
597+ }
598+
557599 yield* Effect.promise(() => self.markActivity()).pipe(
558600 Effect.withSpan("McpSessionDO.markActivity"),
559601 );
apps/cloud/src/mcp.ts+25−212
@@ -30,6 +30,14 @@ import { authorizeOrganization } from "./auth/authorize-organization"; · 29 unmodified lines
3030 import { UserStoreService } from "./auth/context";
3131 import { CoreSharedServices } from "./api/core-shared-services";
3232 import { DbService } from "./services/db";
33+import { peekAndAnnotate } from "./mcp/response-peek";
34+import {
35+ authTemporarilyUnavailable,
36+ CORS_ALLOW_ORIGIN,
37+ jsonResponse,
38+ jsonRpcError,
39+ unauthorized,
40+} from "./mcp/responses";
3341
3442 // ---------------------------------------------------------------------------
3543 // Constants
@@ -45,8 +53,6 @@ const BEARER_PREFIX = "Bearer "; · 9 unmodified lines
4553 const INTERNAL_ACCOUNT_ID_HEADER = "x-executor-mcp-account-id";
4654 const INTERNAL_ORGANIZATION_ID_HEADER = "x-executor-mcp-organization-id";
4755
48−const CORS_ALLOW_ORIGIN = { "access-control-allow-origin": "*" } as const;
49−
5056 const CORS_PREFLIGHT_HEADERS = {
5157 ...CORS_ALLOW_ORIGIN,
5258 "access-control-allow-methods": "GET, POST, DELETE, OPTIONS",
@@ -89,43 +95,6 @@ export const mcpUnauthorized = ( · 36 unmodified lines
8995 description,
9096 });
9197
92−const quoteAuthParam = (value: string) =>
93− `"${value.replaceAll("\\", "\\\\").replaceAll('"', '\\"')}"`;
94−
95−const bearerChallenge = (auth: McpUnauthorizedResult) => {
96− const params =
97− auth.reason === "missing_bearer"
98− ? [`resource_metadata=${quoteAuthParam(PROTECTED_RESOURCE_METADATA_URL)}`]
99− : [
100− 'error="invalid_token"',
101− `error_description=${quoteAuthParam(
102− auth.description ?? "The access token is invalid or expired",
103− )}`,
104− `resource_metadata=${quoteAuthParam(PROTECTED_RESOURCE_METADATA_URL)}`,
105− ];
106−
107− return `Bearer ${params.join(", ")}`;
108−};
109−
110−// ---------------------------------------------------------------------------
111−// Response helpers
112−// ---------------------------------------------------------------------------
113−
114−const jsonResponse = (body: unknown, status = 200) =>
115− HttpServerResponse.unsafeJson(body, { status, headers: CORS_ALLOW_ORIGIN });
116−
117−const jsonRpcError = (status: number, code: number, message: string) =>
118− HttpServerResponse.unsafeJson({ jsonrpc: "2.0", error: { code, message }, id: null }, { status });
119−
120−const unauthorized = (auth: McpUnauthorizedResult) =>
121− HttpServerResponse.unsafeJson(
122− { error: "unauthorized" },
123− {
124− status: 401,
125− headers: { ...CORS_ALLOW_ORIGIN, "www-authenticate": bearerChallenge(auth) },
126− },
127− );
128−
12998 const corsPreflight = HttpServerResponse.empty({
13099 status: 204,
131100 headers: CORS_PREFLIGHT_HEADERS,
@@ -162,11 +131,7 @@ const verifyJwt = (token: string) => · 30 unmodified lines
162131
163132 const DbLive = DbService.Live;
164133 const UserStoreLive = UserStoreService.Live.pipe(Layer.provide(DbLive));
165−const McpOrganizationAuthServices = Layer.mergeAll(
166− DbLive,
167− UserStoreLive,
168− CoreSharedServices,
169−);
134+const McpOrganizationAuthServices = Layer.mergeAll(DbLive, UserStoreLive, CoreSharedServices);
170135
171136 export const McpOrganizationAuthLive = Layer.succeed(McpOrganizationAuth, {
172137 authorize: (accountId, organizationId) =>
@@ -193,9 +158,7 @@ export const McpAuthLive = Layer.succeed(McpAuth, { · 20 unmodified lines
193158 });
194159 return mcpUnauthorized(
195160 "invalid_token",
196− error.reason === "expired"
197− ? "The access token expired"
198− : "The access token is invalid",
161+ error.reason === "expired" ? "The access token expired" : "The access token is invalid",
199162 );
200163 });
201164 }),
@@ -443,167 +406,6 @@ const authorizationServerMetadata = Effect.promise(async () => { · 241 unmodified lines
443406 }
444407 });
445408
446−// ---------------------------------------------------------------------------
447−// Response-body peek for JSON-RPC error attrs
448−// ---------------------------------------------------------------------------
449−//
450−// The MCP protocol wraps protocol-level failures (malformed envelope, method
451−// not found, invalid params) as `{ error: { code, message } }` in the
452−// JSON-RPC response body with HTTP 200 — none of which surface at the HTTP
453−// layer or as an Effect failure. Same for tool results carrying
454−// `isError: true`. To make those visible in Axiom we peek the first
455−// JSON-RPC message out of the DO's response and stamp it onto the
456−// surrounding `mcp.request` span.
457−//
458−// Only applied to non-streaming response shapes (POST/DELETE). GET hops
459−// onto a long-lived SSE channel we don't want to consume.
460−// ---------------------------------------------------------------------------
461−
462−type SandboxOutcome = {
463− readonly status?: string;
464− readonly error?: { readonly kind?: string; readonly message?: string };
465−};
466−
467−type JsonRpcErrorBody = {
468− readonly jsonrpc?: string;
469− readonly error?: { readonly code?: number; readonly message?: string };
470− readonly result?: {
471− readonly isError?: boolean;
472− readonly structuredContent?: SandboxOutcome;
473− };
474−};
475−
476−const responseBodyShape = (body: string): string => {
477− const trimmed = body.trimStart();
478− if (!trimmed) return "empty";
479− if (trimmed.startsWith("{")) return "json-object";
480− if (trimmed.startsWith("[")) return "json-array";
481− if (trimmed.startsWith("event:") || trimmed.startsWith("data:")) return "sse";
482− if (trimmed.startsWith("<")) return "html-or-xml";
483− return "other";
484−};
485−
486−const parseFirstJsonRpc = (contentType: string, body: string): JsonRpcErrorBody | null => {
487− if (!body) return null;
488− try {
489− if (contentType.includes("text/event-stream")) {
490− // Grab the first `data:` line from the first SSE event.
491− for (const line of body.split(/\r?\n/)) {
492− if (line.startsWith("data:")) return JSON.parse(line.slice(5).trimStart());
493− }
494− return null;
495− }
496− if (contentType.includes("application/json")) {
497− return JSON.parse(body);
498− }
499− return null;
500− } catch {
501− return null;
502− }
503−};
504−
505−const rpcResponseAttrs = (payload: JsonRpcErrorBody | null): Record<string, unknown> => {
506− // Require a JSON-RPC 2.0 envelope so we don't false-positive on other
507− // JSON shapes the edge happens to return (e.g. the auth-failure body
508− // `{ "error": "unauthorized" }` — not a JSON-RPC error).
509− if (!payload || payload.jsonrpc !== "2.0") return {};
510− const attrs: Record<string, unknown> = {};
511− const err = payload.error;
512− if (err && typeof err === "object") {
513− attrs["mcp.rpc.is_error"] = true;
514− if (typeof err.code === "number") attrs["mcp.rpc.error.code"] = err.code;
515− if (typeof err.message === "string") {
516− attrs["mcp.rpc.error.message"] = err.message.slice(0, 500);
517− }
518− }
519− if (payload.result?.isError === true) {
520− attrs["mcp.tool.result.is_error"] = true;
521− }
522− const sc = payload.result?.structuredContent;
523− if (sc && typeof sc.status === "string") {
524− attrs["mcp.tool.sandbox.status"] = sc.status;
525− if (sc.error?.kind) attrs["mcp.tool.sandbox.error.kind"] = sc.error.kind;
526− if (typeof sc.error?.message === "string") {
527− attrs["mcp.tool.sandbox.error.message"] = sc.error.message.slice(0, 500);
528− }
529− }
530− return attrs;
531−};
532−
533−const peekAndAnnotate = (response: Response): Effect.Effect<Response> =>
534− Effect.gen(function* () {
535− const contentType = response.headers.get("content-type") ?? "";
536− if (response.status === 202) {
537− // MCP Streamable HTTP accepts notification-only POSTs with 202 and no body.
538− yield* Effect.annotateCurrentSpan({
539− "mcp.response.status_code": response.status,
540− "mcp.response.content_type": contentType,
541− "mcp.response.body.shape": "empty",
542− "mcp.response.body.length": 0,
543− "mcp.response.jsonrpc.detected": false,
544− });
545− const headers = new Headers(response.headers);
546− headers.delete("content-type");
547− headers.delete("content-length");
548− return new Response(null, {
549− status: response.status,
550− statusText: response.statusText,
551− headers,
552− });
553− }
554− if (!response.body) {
555− yield* Effect.annotateCurrentSpan({
556− "mcp.response.status_code": response.status,
557− "mcp.response.content_type": contentType,
558− "mcp.response.body.shape": "empty",
559− "mcp.response.body.length": 0,
560− "mcp.response.jsonrpc.detected": false,
561− });
562− return response;
563− }
564− // The DO can return a streaming SSE response for older in-memory
565− // sessions, so `response.text()` blocks on the
566− // entire downstream execution — RPC into the dynamic Worker, tool
567− // invocations back to the host, result serialisation. Carving this
568− // into its own span pins "worker drain time" on the trace so you can
569− // tell worker-side waiting apart from any pre/post work on the edge.
570− const text = yield* Effect.promise(() => response.text()).pipe(
571− Effect.withSpan("mcp.peek_response", {
572− attributes: {
573− "http.response.content_type": contentType,
574− "http.response.status_code": response.status,
575− },
576− }),
577− );
578− const payload = parseFirstJsonRpc(contentType, text);
579− yield* Effect.annotateCurrentSpan({
580− "mcp.response.status_code": response.status,
581− "mcp.response.content_type": contentType,
582− "mcp.response.body.length": text.length,
583− "mcp.response.body.shape": responseBodyShape(text),
584− "mcp.response.jsonrpc.detected": payload?.jsonrpc === "2.0",
585− });
586− const attrs = rpcResponseAttrs(payload);
587− if (Object.keys(attrs).length > 0) {
588− yield* Effect.annotateCurrentSpan(attrs);
589− }
590− // Internal-error code -32603 means our server failed handling a
591− // structurally valid request. Unlike -32601 / -32602 ("the client
592− // fucked up"), this is a real bug in our code — route to Sentry so
593− // we get alerted. Protocol-level client errors stay in Axiom.
594− if (payload?.error?.code === -32603) {
595− yield* Effect.sync(() => {
596− const msg = payload.error?.message ?? "unknown";
597− Sentry.captureException(new Error(`MCP internal error (-32603): ${msg}`));
598− });
599− }
600− return new Response(text, {
601− status: response.status,
602− statusText: response.statusText,
603− headers: response.headers,
604− });
605− });
606−
607409 // ---------------------------------------------------------------------------
608410 // DO dispatch
609411 // ---------------------------------------------------------------------------
@@ -845,18 +647,29 @@ export const mcpApp: Effect.Effect< · 235 unmodified lines
845647 if (route === "oauth-authorization-server") return yield* authorizationServerMetadata;
846648
847649 const auth = yield* McpAuth;
848− const authResult = yield* auth.verifyBearer(request);
650+ const authResult = yield* auth.verifyBearer(request).pipe(Effect.either);
651+
652+ if (authResult._tag === "Left") {
653+ yield* annotateMcpRequest(request, {
654+ token: null,
655+ parseBody: request.method === "POST",
656+ });
657+ return yield* authTemporarilyUnavailable(authResult.left);
658+ }
659+ const authValue = authResult.right;
849660
850661 // Annotate before dispatch so even 401s show up with what we know. Only
851662 // POST bodies are JSON-RPC payloads worth parsing; GET (SSE) and DELETE
852663 // don't carry one.
853664 yield* annotateMcpRequest(request, {
854− token: authResult._tag === "Authorized" ? authResult.token : null,
665+ token: authValue._tag === "Authorized" ? authValue.token : null,
855666 parseBody: request.method === "POST",
856667 });
857668
858− if (authResult._tag === "Unauthorized") return unauthorized(authResult);
859− const token = authResult.token;
669+ if (authValue._tag === "Unauthorized") {
670+ return unauthorized(authValue, PROTECTED_RESOURCE_METADATA_URL);
671+ }
672+ const token = authValue.token;
860673 switch (request.method) {
861674 case "POST":
862675 return yield* dispatchPost(request, token);
apps/cloud/src/mcp/response-peek.ts+203−0
@@ -0,0 +1,203 @@
1+import * as Sentry from "@sentry/cloudflare";
2+import { Effect } from "effect";
3+
4+import { jsonRpcWebResponse } from "./responses";
5+
6+const SSE_PEEK_TIMEOUT_MS = 10_000;
7+
8+type SandboxOutcome = {
9+ readonly status?: string;
10+ readonly error?: { readonly kind?: string; readonly message?: string };
11+};
12+
13+type JsonRpcResponseBody = {
14+ readonly jsonrpc?: string;
15+ readonly error?: { readonly code?: number; readonly message?: string };
16+ readonly result?: {
17+ readonly isError?: boolean;
18+ readonly structuredContent?: SandboxOutcome;
19+ };
20+};
21+
22+const responseBodyShape = (body: string): string => {
23+ const trimmed = body.trimStart();
24+ if (!trimmed) return "empty";
25+ if (trimmed.startsWith("{")) return "json-object";
26+ if (trimmed.startsWith("[")) return "json-array";
27+ if (trimmed.startsWith("event:") || trimmed.startsWith("data:")) return "sse";
28+ if (trimmed.startsWith("<")) return "html-or-xml";
29+ return "other";
30+};
31+
32+const parseFirstJsonRpc = (contentType: string, body: string): JsonRpcResponseBody | null => {
33+ if (!body) return null;
34+ try {
35+ if (contentType.includes("text/event-stream")) {
36+ for (const line of body.split(/\r?\n/)) {
37+ if (line.startsWith("data:")) return JSON.parse(line.slice(5).trimStart());
38+ }
39+ return null;
40+ }
41+ if (contentType.includes("application/json")) return JSON.parse(body);
42+ return null;
43+ } catch {
44+ return null;
45+ }
46+};
47+
48+const jsonRpcResponseAttrs = (payload: JsonRpcResponseBody | null): Record<string, unknown> => {
49+ if (!payload || payload.jsonrpc !== "2.0") return {};
50+ const attrs: Record<string, unknown> = {};
51+ const err = payload.error;
52+ if (err && typeof err === "object") {
53+ attrs["mcp.rpc.is_error"] = true;
54+ if (typeof err.code === "number") attrs["mcp.rpc.error.code"] = err.code;
55+ if (typeof err.message === "string") {
56+ attrs["mcp.rpc.error.message"] = err.message.slice(0, 500);
57+ }
58+ }
59+ if (payload.result?.isError === true) attrs["mcp.tool.result.is_error"] = true;
60+ const structured = payload.result?.structuredContent;
61+ if (structured && typeof structured.status === "string") {
62+ attrs["mcp.tool.sandbox.status"] = structured.status;
63+ if (structured.error?.kind) attrs["mcp.tool.sandbox.error.kind"] = structured.error.kind;
64+ if (typeof structured.error?.message === "string") {
65+ attrs["mcp.tool.sandbox.error.message"] = structured.error.message.slice(0, 500);
66+ }
67+ }
68+ return attrs;
69+};
70+
71+class ResponseBodyTimeoutError extends Error {
72+ constructor(readonly timeoutMs: number) {
73+ super(`Timed out waiting for MCP response body after ${timeoutMs}ms`);
74+ this.name = "ResponseBodyTimeoutError";
75+ }
76+}
77+
78+const readResponseText = async (response: Response, timeoutMs: number | null): Promise<string> => {
79+ if (timeoutMs === null) return await response.text();
80+
81+ const reader = response.body?.getReader();
82+ if (!reader) return "";
83+
84+ const decoder = new TextDecoder();
85+ let timeout: ReturnType<typeof setTimeout> | undefined;
86+ const timeoutPromise = new Promise<never>((_, reject) => {
87+ timeout = setTimeout(() => {
88+ void reader.cancel().catch(() => undefined);
89+ reject(new ResponseBodyTimeoutError(timeoutMs));
90+ }, timeoutMs);
91+ });
92+ const readPromise = (async () => {
93+ try {
94+ let text = "";
95+ for (;;) {
96+ const { done, value } = await reader.read();
97+ if (done) return text + decoder.decode();
98+ text += decoder.decode(value, { stream: true });
99+ }
100+ } finally {
101+ if (timeout) clearTimeout(timeout);
102+ }
103+ })();
104+
105+ return await Promise.race([readPromise, timeoutPromise]);
106+};
107+
108+const annotateEmptyResponse = (response: Response, contentType: string) =>
109+ Effect.annotateCurrentSpan({
110+ "mcp.response.status_code": response.status,
111+ "mcp.response.content_type": contentType,
112+ "mcp.response.body.shape": "empty",
113+ "mcp.response.body.length": 0,
114+ "mcp.response.jsonrpc.detected": false,
115+ });
116+
117+const withoutBodyHeaders = (response: Response) => {
118+ const headers = new Headers(response.headers);
119+ headers.delete("content-type");
120+ headers.delete("content-length");
121+ return new Response(null, {
122+ status: response.status,
123+ statusText: response.statusText,
124+ headers,
125+ });
126+};
127+
128+const responseReadFailure = (error: unknown) =>
129+ Effect.gen(function* () {
130+ const timedOut = error instanceof ResponseBodyTimeoutError;
131+ yield* Effect.annotateCurrentSpan({
132+ "mcp.response.status_code": timedOut ? 504 : 500,
133+ "mcp.response.content_type": "application/json",
134+ "mcp.response.body.shape": "json-object",
135+ "mcp.response.body.length": 0,
136+ "mcp.response.jsonrpc.detected": true,
137+ "mcp.peek_response.timed_out": timedOut,
138+ "mcp.peek_response.error": String(error),
139+ });
140+ return jsonRpcWebResponse(
141+ timedOut ? 504 : 500,
142+ -32001,
143+ timedOut
144+ ? "Timed out waiting for MCP response - please retry"
145+ : "Failed to read MCP response",
146+ );
147+ });
148+
149+const reportInternalJsonRpcError = (payload: JsonRpcResponseBody | null) =>
150+ Effect.sync(() => {
151+ if (payload?.error?.code !== -32603) return;
152+ const msg = payload.error.message ?? "unknown";
153+ Sentry.captureException(new Error(`MCP internal error (-32603): ${msg}`));
154+ });
155+
156+export const peekAndAnnotate = (response: Response): Effect.Effect<Response> =>
157+ Effect.gen(function* () {
158+ const contentType = response.headers.get("content-type") ?? "";
159+ if (response.status === 202) {
160+ yield* annotateEmptyResponse(response, contentType);
161+ return withoutBodyHeaders(response);
162+ }
163+ if (!response.body) {
164+ yield* annotateEmptyResponse(response, contentType);
165+ return response;
166+ }
167+
168+ const isSseResponse = contentType.includes("text/event-stream");
169+ const timeoutMs = isSseResponse ? SSE_PEEK_TIMEOUT_MS : null;
170+ const textResult = yield* Effect.tryPromise({
171+ try: () => readResponseText(response, timeoutMs),
172+ catch: (error) => error,
173+ }).pipe(
174+ Effect.withSpan("mcp.peek_response", {
175+ attributes: {
176+ "http.response.content_type": contentType,
177+ "http.response.status_code": response.status,
178+ "mcp.peek_response.timeout_ms": timeoutMs ?? 0,
179+ },
180+ }),
181+ Effect.either,
182+ );
183+ if (textResult._tag === "Left") return yield* responseReadFailure(textResult.left);
184+
185+ const text = textResult.right;
186+ const payload = parseFirstJsonRpc(contentType, text);
187+ yield* Effect.annotateCurrentSpan({
188+ "mcp.response.status_code": response.status,
189+ "mcp.response.content_type": contentType,
190+ "mcp.response.body.length": text.length,
191+ "mcp.response.body.shape": responseBodyShape(text),
192+ "mcp.response.jsonrpc.detected": payload?.jsonrpc === "2.0",
193+ });
194+ const attrs = jsonRpcResponseAttrs(payload);
195+ if (Object.keys(attrs).length > 0) yield* Effect.annotateCurrentSpan(attrs);
196+ yield* reportInternalJsonRpcError(payload);
197+
198+ return new Response(text, {
199+ status: response.status,
200+ statusText: response.statusText,
201+ headers: response.headers,
202+ });
203+ });
apps/cloud/src/mcp/responses.ts+66−0
@@ -0,0 +1,66 @@
1+import { HttpServerResponse } from "@effect/platform";
2+import { Effect } from "effect";
3+
4+import type { McpJwtVerificationError } from "../mcp-auth";
5+
6+export const CORS_ALLOW_ORIGIN = { "access-control-allow-origin": "*" } as const;
7+
8+type UnauthorizedAuth = {
9+ readonly reason: "missing_bearer" | "invalid_token";
10+ readonly description?: string;
11+};
12+
13+const quoteAuthParam = (value: string) =>
14+ `"${value.replaceAll("\\", "\\\\").replaceAll('"', '\\"')}"`;
15+
16+const bearerChallenge = (auth: UnauthorizedAuth, protectedResourceMetadataUrl: string) => {
17+ const params =
18+ auth.reason === "missing_bearer"
19+ ? [`resource_metadata=${quoteAuthParam(protectedResourceMetadataUrl)}`]
20+ : [
21+ 'error="invalid_token"',
22+ `error_description=${quoteAuthParam(
23+ auth.description ?? "The access token is invalid or expired",
24+ )}`,
25+ `resource_metadata=${quoteAuthParam(protectedResourceMetadataUrl)}`,
26+ ];
27+
28+ return `Bearer ${params.join(", ")}`;
29+};
30+
31+export const jsonResponse = (body: unknown, status = 200) =>
32+ HttpServerResponse.unsafeJson(body, { status, headers: CORS_ALLOW_ORIGIN });
33+
34+export const jsonRpcError = (status: number, code: number, message: string) =>
35+ HttpServerResponse.unsafeJson(
36+ { jsonrpc: "2.0", error: { code, message }, id: null },
37+ { status, headers: CORS_ALLOW_ORIGIN },
38+ );
39+
40+export const jsonRpcWebResponse = (status: number, code: number, message: string) =>
41+ new Response(JSON.stringify({ jsonrpc: "2.0", error: { code, message }, id: null }), {
42+ status,
43+ headers: { ...CORS_ALLOW_ORIGIN, "content-type": "application/json" },
44+ });
45+
46+export const unauthorized = (auth: UnauthorizedAuth, protectedResourceMetadataUrl: string) =>
47+ HttpServerResponse.unsafeJson(
48+ { error: "unauthorized" },
49+ {
50+ status: 401,
51+ headers: {
52+ ...CORS_ALLOW_ORIGIN,
53+ "www-authenticate": bearerChallenge(auth, protectedResourceMetadataUrl),
54+ },
55+ },
56+ );
57+
58+export const authTemporarilyUnavailable = (error: McpJwtVerificationError) =>
59+ Effect.gen(function* () {
60+ yield* Effect.annotateCurrentSpan({
61+ "mcp.auth.outcome": "system_error",
62+ "mcp.auth.system_error.reason": error.reason,
63+ "mcp.auth.system_error.message": String(error.cause).slice(0, 500),
64+ });
65+ return jsonRpcError(503, -32001, "Authentication temporarily unavailable - please retry");
66+ });
apps/cloud/src/services/mcp-worker-transport.ts+122−1
@@ -13,11 +13,102 @@ export type McpWorkerTransport = Readonly<{ · 12 unmodified lines
1313 close: () => Effect.Effect<void>;
1414 }>;
1515
16+type WorkerTransportInternals = {
17+ readonly standaloneSseStreamId?: string;
18+ readonly streamMapping?: Map<string, { readonly cleanup?: () => void }>;
19+};
20+
21+type JsonRpcLike = {
22+ readonly id?: unknown;
23+ readonly method?: unknown;
24+};
25+
26+type HandleRequestResult = {
27+ readonly response: Response;
28+ readonly replacedStandaloneSse: boolean;
29+};
30+
31+const closeExistingStandaloneSse = (transport: WorkerTransport): boolean => {
32+ const internals = transport as unknown as WorkerTransportInternals;
33+ const streamId = internals.standaloneSseStreamId ?? "_GET_stream";
34+ const stream = internals.streamMapping?.get(streamId);
35+ if (!stream) return false;
36+
37+ stream.cleanup?.();
38+ internals.streamMapping?.delete(streamId);
39+ return true;
40+};
41+
42+const isStandaloneSseGet = (request: Request): boolean =>
43+ request.method === "GET" && (request.headers.get("accept") ?? "").includes("text/event-stream");
44+
45+const jsonRpcRequestIdKey = (id: unknown): string | null => {
46+ switch (typeof id) {
47+ case "string":
48+ case "number":
49+ case "boolean":
50+ return `${typeof id}:${String(id)}`;
51+ default:
52+ return null;
53+ }
54+};
55+
56+const extractJsonRpcRequestIdKeys = async (request: Request): Promise<ReadonlyArray<string>> => {
57+ if (request.method !== "POST") return [];
58+ const contentType = request.headers.get("content-type") ?? "";
59+ if (!contentType.includes("application/json")) return [];
60+
61+ try {
62+ const parsed = (await request.clone().json()) as unknown;
63+ const messages = Array.isArray(parsed) ? parsed : [parsed];
64+ return messages.flatMap((message) => {
65+ if (!message || typeof message !== "object") return [];
66+ const rpc = message as JsonRpcLike;
67+ if (typeof rpc.method !== "string") return [];
68+ const key = jsonRpcRequestIdKey(rpc.id);
69+ return key ? [key] : [];
70+ });
71+ } catch {
72+ return [];
73+ }
74+};
75+
76+class JsonRpcRequestIdQueue {
77+ private readonly inFlight = new Map<string, Promise<void>>();
78+
79+ async run<A>(request: Request, run: () => Promise<A>): Promise<A> {
80+ const ids = [...new Set(await extractJsonRpcRequestIdKeys(request))];
81+ if (ids.length === 0) return await run();
82+
83+ const previous = ids.map((id) => this.inFlight.get(id)).filter((p) => p !== undefined);
84+ let release!: () => void;
85+ const current = new Promise<void>((resolve) => {
86+ release = resolve;
87+ });
88+ for (const id of ids) {
89+ this.inFlight.set(id, current);
90+ }
91+
92+ try {
93+ await Promise.all(previous.map((p) => p.catch(() => undefined)));
94+ return await run();
95+ } finally {
96+ for (const id of ids) {
97+ if (this.inFlight.get(id) === current) {
98+ this.inFlight.delete(id);
99+ }
100+ }
101+ release();
102+ }
103+ }
104+}
105+
16106 export const makeMcpWorkerTransport = (
17107 options: WorkerTransportOptions,
18108 ): Effect.Effect<McpWorkerTransport> =>
19109 Effect.sync(() => {
20110 const transport = new WorkerTransport(options);
111+ const requestIdQueue = new JsonRpcRequestIdQueue();
21112
22113 const use = <A>(name: string, fn: () => Promise<A>) =>
23114 Effect.tryPromise({
@@ -25,11 +116,41 @@ export const makeMcpWorkerTransport = ( · 1 unmodified line
25116 catch: (cause) => new McpWorkerTransportError({ cause }),
26117 }).pipe(Effect.withSpan(`mcp.worker_transport.${name}`));
27118
119+ const handleWithStandaloneSseReplacement = async (
120+ request: Request,
121+ ): Promise<HandleRequestResult> => {
122+ if (!isStandaloneSseGet(request)) {
123+ return {
124+ response: await transport.handleRequest(request),
125+ replacedStandaloneSse: false,
126+ };
127+ }
128+
129+ const initial = await transport.handleRequest(request);
130+ if (initial.status !== 409) {
131+ return { response: initial, replacedStandaloneSse: false };
132+ }
133+
134+ const replacedStandaloneSse = closeExistingStandaloneSse(transport);
135+ return {
136+ response: replacedStandaloneSse ? await transport.handleRequest(request) : initial,
137+ replacedStandaloneSse,
138+ };
139+ };
140+
28141 return {
29142 transport,
30143 connect: (server: McpServer) => use("connect", () => server.connect(transport)),
31144 handleRequest: (request: Request) =>
32− use("handle_request", () => transport.handleRequest(request)),
145+ Effect.gen(function* () {
146+ const result = yield* use("handle_request", () =>
147+ requestIdQueue.run(request, () => handleWithStandaloneSseReplacement(request)),
148+ );
149+ yield* Effect.annotateCurrentSpan({
150+ "mcp.transport.replaced_standalone_sse": result.replacedStandaloneSse,
151+ });
152+ return result.response;
153+ }),
33154 close: () =>
34155 Effect.promise(() => transport.close().catch(() => undefined)).pipe(
35156 Effect.withSpan("mcp.worker_transport.close"),
apps/cloud/src/test-worker.ts+10−2
@@ -28,6 +28,7 @@ import { · 27 unmodified lines
2828 mcpApp,
2929 mcpUnauthorized,
3030 } from "./mcp";
31+import { McpJwtVerificationError } from "./mcp-auth";
3132 import { organizations } from "./services/schema";
3233 import { parseTestBearer } from "./test-bearer";
3334 import { DoTelemetryLive } from "./services/telemetry";
@@ -36,10 +37,17 @@ export { McpSessionDO } from "./mcp-session"; · 2 unmodified lines
3637
3738 const TestMcpAuthLive = Layer.succeed(McpAuth, {
3839 verifyBearer: (request) =>
39− Effect.sync(() => {
40+ Effect.gen(function* () {
4041 const header = request.headers.get("authorization");
4142 if (!header?.startsWith("Bearer ")) return mcpUnauthorized("missing_bearer");
42− const token = parseTestBearer(header.slice("Bearer ".length));
43+ const rawToken = header.slice("Bearer ".length);
44+ if (rawToken === "test-system-error") {
45+ return yield* Effect.fail(new McpJwtVerificationError({
46+ cause: new Error("simulated jwks fetch failure"),
47+ reason: "system",
48+ }));
49+ }
50+ const token = parseTestBearer(rawToken);
4351 return token ? mcpAuthorized(token) : mcpUnauthorized("invalid_token");
4452 }),
4553 });

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t