diff --git a/apps/cloud/src/mcp/session-build-semaphore.test.ts b/apps/cloud/src/mcp/session-build-semaphore.test.ts index 3d4ad7634..584b65ee0 100644 --- a/apps/cloud/src/mcp/session-build-semaphore.test.ts +++ b/apps/cloud/src/mcp/session-build-semaphore.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it, beforeEach } from "@effect/vitest"; +import { describe, expect, it, beforeEach, afterEach, vi } from "@effect/vitest"; import { acquireBuildSlot, @@ -13,6 +13,10 @@ describe("session-build-semaphore", () => { resetBuildSlotsForTest(); }); + afterEach(() => { + vi.useRealTimers(); + }); + it("grants up to the cap immediately, with no wait", async () => { const results = await Promise.all([ acquireBuildSlot().promise, @@ -214,6 +218,7 @@ describe("session-build-semaphore", () => { }); it("proceeds without a slot when the queue wait exceeds the timeout, and does not count it as active", async () => { + vi.useFakeTimers(); await Promise.all([ acquireBuildSlot().promise, acquireBuildSlot().promise, @@ -223,6 +228,10 @@ describe("session-build-semaphore", () => { expect(currentActiveBuildsForTest()).toBe(4); const timedOutHandle = acquireBuildSlot(10); + await vi.advanceTimersByTimeAsync(9); + expect(currentQueueLengthForTest()).toBe(1); + expect(currentActiveBuildsForTest()).toBe(4); + await vi.advanceTimersByTimeAsync(1); const result = await timedOutHandle.promise; expect(result).toEqual({ acquired: false, waitMs: expect.any(Number), timedOut: true }); diff --git a/e2e/cloud/mcp-session-cap-eviction.test.ts b/e2e/cloud/mcp-session-cap-eviction.test.ts index 8724edd5a..8103eef5c 100644 --- a/e2e/cloud/mcp-session-cap-eviction.test.ts +++ b/e2e/cloud/mcp-session-cap-eviction.test.ts @@ -19,8 +19,8 @@ // the whole boot (see that file for the value and its headroom story), so // this test can cross it with a bounded number of real sessions instead of // registering the production default of 32. -import { expect } from "@effect/vitest"; -import { Effect, Schedule } from "effect"; +import { expect, it } from "@effect/vitest"; +import { Effect, Option, Schedule, Schema } from "effect"; import { scenario } from "../src/scenario"; import { Mcp, Target, Telemetry } from "../src/services"; @@ -52,6 +52,68 @@ const postJson = (mcpUrl: string, bearer: string, body: unknown, sessionId?: str body: JSON.stringify(body), }); +const decodeRestartEnvelope = Schema.decodeUnknownOption( + Schema.fromJsonString( + Schema.Struct({ + jsonrpc: Schema.Literal("2.0"), + id: Schema.Null, + error: Schema.Struct({ + code: Schema.Literal(-32001), + message: Schema.Literal("MCP session is restarting, please retry"), + }), + }), + ), +); + +const isRestartResponse = (status: number, body: string): boolean => + status === 503 && Option.isSome(decodeRestartEnvelope(body)); + +it.each([ + [ + 503, + { + jsonrpc: "2.0", + id: null, + error: { code: -32001, message: "MCP session is restarting, please retry" }, + }, + true, + ], + [ + 404, + { + jsonrpc: "2.0", + id: null, + error: { code: -32001, message: "MCP session is restarting, please retry" }, + }, + false, + ], + [ + 503, + { + jsonrpc: "2.0", + id: null, + error: { code: -32603, message: "MCP session is restarting, please retry" }, + }, + false, + ], + [ + 503, + { + jsonrpc: "2.0", + id: null, + error: { code: -32001, message: "MCP session is restarting unexpectedly" }, + }, + false, + ], + [503, { error: "MCP session is restarting, please retry" }, false], +] as const)("only retries the documented restart envelope (%s, %j)", (status, body, retry) => { + expect(isRestartResponse(status, JSON.stringify(body))).toBe(retry); +}); + +it("does not retry malformed restart responses", () => { + expect(isRestartResponse(503, "MCP session is restarting, please retry")).toBe(false); +}); + /** * Opens one fresh MCP session under an already-minted bearer. `initialize` * without an existing `mcp-session-id` always mints a new session, the same @@ -73,21 +135,42 @@ const openSession = async ( label: string, recordSession: (sessionId: string) => void, ): Promise => { - const initialized = await postJson(mcpUrl, bearer, { - jsonrpc: "2.0" as const, - id: "initialize", - method: "initialize", - params: { - protocolVersion: PROTOCOL_VERSION, - capabilities: {}, - clientInfo: { name: `executor-e2e-cap-eviction-${label}`, version: "0.0.1" }, - }, - }); - const sessionId = initialized.headers.get("mcp-session-id"); - if (!sessionId) { + // The platform can reset a session Durable Object while its initialize + // is in flight, and the server answers that with the documented restart + // envelope (503, -32001, "MCP session is restarting, please retry") — the + // same contract a streamable-http client follows: same request, after the + // advertised delay. Treat it as transient here instead of failing the + // scenario on a retryable platform blip. + const RESTART_ATTEMPTS = 8; + const RESTART_DELAY_MS = 2_000; // The host advertises Retry-After: 2. + let minted: { readonly response: Response; readonly sessionId: string } | undefined; + for (let attempt = 0; attempt < RESTART_ATTEMPTS; attempt += 1) { + const response = await postJson(mcpUrl, bearer, { + jsonrpc: "2.0" as const, + id: "initialize", + method: "initialize", + params: { + protocolVersion: PROTOCOL_VERSION, + capabilities: {}, + clientInfo: { name: `executor-e2e-cap-eviction-${label}`, version: "0.0.1" }, + }, + }); + const candidate = response.headers.get("mcp-session-id"); + if (candidate !== null && candidate.length > 0) { + minted = { response, sessionId: candidate }; + break; + } + const body = await response.text(); + const isRestart = isRestartResponse(response.status, body); + if (!isRestart) break; + if (attempt === RESTART_ATTEMPTS - 1) break; + await new Promise((resolve) => setTimeout(resolve, RESTART_DELAY_MS)); + } + if (!minted) { // oxlint-disable-next-line executor/no-error-constructor -- boundary: e2e setup precondition. throw new Error(`openSession (${label}): no mcp-session-id header`); } + const { response: initialized, sessionId } = minted; // Recorded the moment the id exists — BEFORE the body read and status // assertion below, either of which can throw with the session already live // on the server. The cleanup finalizer needs the id on every one of those