diff --git a/apps/cloud/src/env-augment.d.ts b/apps/cloud/src/env-augment.d.ts index 8105ce9c4..773ca4d46 100644 --- a/apps/cloud/src/env-augment.d.ts +++ b/apps/cloud/src/env-augment.d.ts @@ -118,6 +118,9 @@ declare global { MCP_RESOURCE_ORIGIN?: string; MCP_SESSION_TIMEOUT_MS?: string; MCP_PAUSED_SESSION_IDLE_TIMEOUT_MS?: string; + /** Test-only override for the isolate-wide resident-runtime soft cap + * (see `RESIDENT_RUNTIME_SOFT_CAP`). Unset in production. */ + MCP_RESIDENT_RUNTIME_SOFT_CAP?: string; NODE_ENV?: string; // Shared with frontend diff --git a/apps/cloud/src/mcp/session-durable-object.ts b/apps/cloud/src/mcp/session-durable-object.ts index 5ff5ee2f0..39c99d05f 100644 --- a/apps/cloud/src/mcp/session-durable-object.ts +++ b/apps/cloud/src/mcp/session-durable-object.ts @@ -102,6 +102,10 @@ const positiveMilliseconds = (raw: string | undefined): number | undefined => { return Math.floor(parsed); }; +/** Same shape as `positiveMilliseconds`, for a plain count rather than a + * duration — used only by the resident-runtime soft-cap override below. */ +const positiveInteger = positiveMilliseconds; + type CloudSessionDbHandle = DbServiceShape & { readonly sql: Sql; readonly end: () => Promise; @@ -181,6 +185,25 @@ export class McpSessionDOSqlite extends McpAgentSessionDOBase { + // Routed through this session's OWN stub (never a direct in-process call) + // so `requestCapEviction`'s teardown runs under an IoContext scoped to + // this request, not whatever request happened to trigger the eviction + // check — see the base class's `requestSelfEviction` doc comment. + return mcpSessionStub(env.MCP_SESSION, this.sessionId).requestCapEviction(); + } + protected override forwardModelResumeToOwner( owner: McpExecutionOwnerRoute, identity: McpApprovalOwner, diff --git a/apps/host-cloudflare/src/mcp/session-durable-object.ts b/apps/host-cloudflare/src/mcp/session-durable-object.ts index 4fccd6af9..ee3695491 100644 --- a/apps/host-cloudflare/src/mcp/session-durable-object.ts +++ b/apps/host-cloudflare/src/mcp/session-durable-object.ts @@ -80,6 +80,18 @@ export class McpSessionDO extends McpAgentSessionDOBase { + // Routed through this session's OWN stub (never a direct in-process call) + // so `requestCapEviction`'s teardown runs under an IoContext scoped to + // this request, not whatever request happened to trigger the eviction + // check — see the base class's `requestSelfEviction` doc comment. + return mcpSessionStub(this.cfEnv.MCP_SESSION, this.sessionId).requestCapEviction(); + } + protected override forwardModelResumeToOwner( owner: McpExecutionOwnerRoute, identity: McpApprovalOwner, diff --git a/e2e/cloud/mcp-session-cap-eviction.test.ts b/e2e/cloud/mcp-session-cap-eviction.test.ts new file mode 100644 index 000000000..6cbddde6a --- /dev/null +++ b/e2e/cloud/mcp-session-cap-eviction.test.ts @@ -0,0 +1,256 @@ +// Cloud: crossing the isolate's resident-runtime soft cap evicts an idle +// session's runtime through a REAL cross-Durable-Object request, not a +// same-context call. +// +// The defect this pins: the original design ran the evicted (candidate) +// session's teardown — closing its postgres.js socket, storage writes, span +// flush — directly inside the EVICTING session's own request/IoContext. In +// production workerd, I/O objects are bound to the IoContext that created +// them, so a cross-context call like that throws "Cannot perform I/O on +// behalf of a different request" or silently soft-fails. That failure mode +// cannot reproduce against an in-process unit-test double (same JS object, +// same context either way) — it only shows up against a real Durable Object +// stub. The fix routes eviction through the candidate's OWN stub +// (`requestCapEviction`, an RPC method mirroring `forwardModelResumeToOwner`), +// so the candidate's teardown runs in the candidate's own context, and this +// scenario is what actually exercises that stub in workerd. +// +// e2e/setup/resident-runtime-cap.ts lowers MCP_RESIDENT_RUNTIME_SOFT_CAP for +// 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 { scenario } from "../src/scenario"; +import { Mcp, Target, Telemetry } from "../src/services"; +import type { Identity } from "../src/target"; +import { E2E_MCP_RESIDENT_RUNTIME_SOFT_CAP } from "../setup/resident-runtime-cap"; + +const PROTOCOL_VERSION = "2025-03-26"; +const JSON_AND_SSE = "application/json, text/event-stream"; + +// Comfortably past the cap: even if a handful of other scenarios' sessions +// are still incidentally resident when this file runs, enough of THESE +// sessions cross it that at least one eviction targets a session opened here. +const SESSIONS_TO_OPEN = E2E_MCP_RESIDENT_RUNTIME_SOFT_CAP + 10; + +const emailOf = (identity: Identity): string => identity.credentials?.email ?? identity.label; + +const mcpHeaders = (bearer: string, sessionId?: string) => ({ + accept: JSON_AND_SSE, + authorization: `Bearer ${bearer}`, + "content-type": "application/json", + "mcp-protocol-version": PROTOCOL_VERSION, + ...(sessionId ? { "mcp-session-id": sessionId } : {}), +}); + +const postJson = (mcpUrl: string, bearer: string, body: unknown, sessionId?: string) => + fetch(mcpUrl, { + method: "POST", + headers: mcpHeaders(bearer, sessionId), + body: JSON.stringify(body), + }); + +/** + * Opens one fresh MCP session under an already-minted bearer. `initialize` + * without an existing `mcp-session-id` always mints a new session, the same + * way separate browser tabs sharing one login would — so many of these under + * one identity is a cheap way to grow the isolate's resident-runtime count + * without a full OAuth round trip per session. + * + * `recordSession` is called the moment the session id is known — before the + * `notifications/initialized` round trip below, not after this function + * returns. A session is live on the server as soon as `initialize` responds + * with an `mcp-session-id`, regardless of whether the handshake ever + * completes; recording it only on a full return left a failed notification + * (or an interrupt landing between the two requests) with no cleanup entry, + * orphaning a real session on the target isolate. + */ +const openSession = async ( + mcpUrl: string, + bearer: string, + 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) { + // oxlint-disable-next-line executor/no-error-constructor -- boundary: e2e setup precondition. + throw new Error(`openSession (${label}): no mcp-session-id header`); + } + // 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 + // paths, not just a fully successful return. + recordSession(sessionId); + await initialized.text(); + expect(initialized.status, `initialize (${label}) opens a session`).toBe(200); + const notification = await postJson( + mcpUrl, + bearer, + { jsonrpc: "2.0" as const, method: "notifications/initialized" }, + sessionId, + ); + await notification.text(); + expect(notification.status, `(${label}) completes the handshake`).toBe(202); + return sessionId; +}; + +const executeBody = (id: string, code: string) => ({ + jsonrpc: "2.0" as const, + id, + method: "tools/call", + params: { name: "execute", arguments: { code } }, +}); + +/** Run `execute` and return the response text once the call has fully settled. */ +const execute = async ( + mcpUrl: string, + bearer: string, + sessionId: string, + id: string, + code: string, +): Promise => { + const response = await postJson(mcpUrl, bearer, executeBody(id, code), sessionId); + const body = await response.text(); + expect(response.status, `execute ${id} is served`).toBe(200); + return body; +}; + +scenario( + "MCP session · crossing the resident-runtime cap evicts a session through a real cross-DO request", + { timeout: 180_000 }, + Effect.gen(function* () { + const target = yield* Target; + const mcp = yield* Mcp; + const telemetry = yield* Telemetry; + + const identity = yield* target.newIdentity(); + const bearer = yield* mcp.mintBearer(emailOf(identity)); + + // Opened sessions are recorded here as each one succeeds, so the cleanup + // below can close exactly what was actually opened even if the scenario + // fails partway through. Cap eviction already tears most of these down as + // a side effect of the scenario itself, but termination is idempotent + // (see mcp-destroyed-session-envelope.test.ts) — closing an already-torn- + // -down session is a harmless no-op, not a double-free. + const openedSessionIds: string[] = []; + + const scenarioBody = Effect.gen(function* () { + // Open more sessions than the cap allows, at limited concurrency. None + // of them run any work, so every one is immediately eviction-eligible — + // crossing the cap must pick at least one and tear it down through its + // own stub. + const sessionIds = yield* Effect.forEach( + Array.from({ length: SESSIONS_TO_OPEN }, (_, index) => index), + (index) => + Effect.promise(() => + openSession(target.mcpUrl, bearer, `session-${index}`, (sessionId) => { + openedSessionIds.push(sessionId); + }), + ), + { concurrency: 8 }, + ); + + expect(sessionIds.length, "every session opened").toBe(SESSIONS_TO_OPEN); + expect(new Set(sessionIds).size, "every session got a distinct id").toBe(SESSIONS_TO_OPEN); + + // ---- a real cap eviction fired, against a session opened here ------- + // Same span the idle path emits (`mcp.session.idle_runtime_dispose`); + // `mcp.session.dispose_reason` is what disambiguates the trigger. + const capDisposals = yield* telemetry + .searchSpans({ operation: "mcp.session.idle_runtime_dispose" }) + .pipe( + Effect.map((spans) => + spans.filter( + (span) => + span.span.tags["mcp.session.dispose_reason"] === "cap" && + sessionIds.some((id) => (span.span.tags["mcp.session.id"] ?? "").includes(id)), + ), + ), + Effect.filterOrFail( + (spans) => spans.length > 0, + () => "no cap-triggered idle_runtime_dispose span exported for any session opened here", + ), + // The eviction request is fire-and-forget (`ctx.waitUntil`) from the + // evictor's `init`, and its own span flush is off that same + // background path — same polling grace the idle-disposal scenario + // uses for its alarm-driven flush. + Effect.retry(Schedule.both(Schedule.spaced("500 millis"), Schedule.recurs(40))), + ); + + expect( + capDisposals.length, + "crossing the resident-runtime cap evicted at least one session opened here", + ).toBeGreaterThan(0); + + const disposal = capDisposals[0]!; + expect( + disposal.span.tags["mcp.isolate.resident_runtimes"], + "the cap disposal records the isolate's resident-runtime gauge, same as the idle path", + ).toBeDefined(); + + // ---- the evicted session still works — restore is transparent ------- + const evictedSessionId = sessionIds.find((id) => + (disposal.span.tags["mcp.session.id"] ?? "").includes(id), + ); + expect( + evictedSessionId, + "the disposed span's session id matches a session opened here", + ).toBeDefined(); + + const marker = `after-cap-evict-${evictedSessionId}`; + const restored = yield* Effect.promise(() => + execute( + target.mcpUrl, + bearer, + evictedSessionId!, + "execute-after-cap-eviction", + `return ${JSON.stringify(marker)};`, + ), + ); + expect( + restored, + "the evicted session serves the next call correctly after restoring underneath the client", + ).toContain(marker); + }); + + yield* scenarioBody.pipe( + // `Effect.ensuring`, not a trailing statement: a failure partway through + // (an assertion above, a timed-out span search) must not leak the + // sessions already opened. Read `openedSessionIds` at cleanup time, not + // capture time — `Effect.suspend` so the array is read when the + // finalizer actually runs, after the scenario body has finished pushing + // to it, rather than snapshotted empty at construction. + Effect.ensuring( + Effect.suspend(() => + Effect.forEach( + openedSessionIds, + (sessionId) => + Effect.tryPromise(async () => { + const closed = await fetch(target.mcpUrl, { + method: "DELETE", + headers: { + authorization: `Bearer ${bearer}`, + "mcp-session-id": sessionId, + }, + }); + await closed.text(); + }).pipe(Effect.ignore), + { concurrency: 8, discard: true }, + ), + ), + ), + ); + }), +); diff --git a/e2e/setup/cloud.boot.ts b/e2e/setup/cloud.boot.ts index 4c24cd3ed..8314e7bc0 100644 --- a/e2e/setup/cloud.boot.ts +++ b/e2e/setup/cloud.boot.ts @@ -16,6 +16,7 @@ import { E2E_EXECUTION_RATE_LIMIT, E2E_EXECUTION_RATE_LIMIT_CHECK_TIMEOUT_MS, } from "./execution-limits"; +import { E2E_MCP_RESIDENT_RUNTIME_SOFT_CAP } from "./resident-runtime-cap"; export const cloudDir = fileURLToPath(new URL("../../apps/cloud/", import.meta.url)); @@ -107,6 +108,10 @@ export const bootCloud = async (options: CloudBootOptions): Promise MCP_RESOURCE_ORIGIN: options.publicUrl, MCP_SESSION_TIMEOUT_MS: process.env.MCP_SESSION_TIMEOUT_MS, MCP_PAUSED_SESSION_IDLE_TIMEOUT_MS: process.env.MCP_PAUSED_SESSION_IDLE_TIMEOUT_MS, + // See resident-runtime-cap.ts for why this value, and why it is safe to + // set unconditionally for the whole boot (same treatment as the execution + // rate limit below). + MCP_RESIDENT_RUNTIME_SOFT_CAP: String(E2E_MCP_RESIDENT_RUNTIME_SOFT_CAP), ALLOW_LOCAL_NETWORK: "true", // A first-party GitHub app for the first-party-oauth scenario: proves the // env → HostConfig → executor plumbing end to end. The scenario asserts the diff --git a/e2e/setup/resident-runtime-cap.ts b/e2e/setup/resident-runtime-cap.ts new file mode 100644 index 000000000..6e3d8223f --- /dev/null +++ b/e2e/setup/resident-runtime-cap.ts @@ -0,0 +1,21 @@ +// The e2e worker's isolate-wide resident-runtime soft cap +// (MCP_RESIDENT_RUNTIME_SOFT_CAP). One constant with two consumers, the boot +// recipe env (cloud.boot.ts) and the cap-eviction scenario +// (cloud/mcp-session-cap-eviction.test.ts), so they cannot drift apart. +// +// Same squeeze as EXECUTION_RATE_LIMIT_PER_HOUR (execution-limits.ts). It must +// be LOW enough that the cap-eviction scenario can cross it with a bounded +// number of real sessions opened under one identity (prod's 32 is reachable +// but wasteful to open on every run), and HIGH enough that no OTHER cloud +// scenario's incidental concurrent MCP-session count trips it: this is an +// isolate-wide counter, not per-org, and the dev server is shared across the +// whole cloud e2e run (`fileParallelism: false` keeps files serial, but a +// single busy file can still open a double-digit number of sessions before +// any of them idle out). The busiest current file (mcp-protocol.test.ts) opens +// on the order of a dozen sessions across its scenarios; this stays well above +// that with headroom. If a scenario ever incidentally trips a cap eviction (a +// span with `mcp.session.dispose_reason: "cap"` for a session it didn't +// expect), that is not a correctness bug — eviction is designed to be +// transparent, restoring on the next call — but it means this constant needs +// to grow: raise it here, never by special-casing a scenario against it. +export const E2E_MCP_RESIDENT_RUNTIME_SOFT_CAP = 24; diff --git a/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts b/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts index 26c2eb863..37c6b0510 100644 --- a/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts +++ b/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts @@ -1,6 +1,7 @@ // oxlint-disable executor/no-error-constructor, executor/no-try-catch-or-throw -- boundary: the storage fake reproduces the plain Errors the Cloudflare runtime throws, and rejecting is the only way a DurableObjectStorage reports them -import { afterEach, describe, expect, it } from "@effect/vitest"; +import { afterEach, beforeEach, describe, expect, it } from "@effect/vitest"; import { Cause, Effect, Exit, Schema } from "effect"; +import type * as Tracer from "effect/Tracer"; import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js"; import type { JSONRPCMessage, MessageExtraInfo } from "@modelcontextprotocol/sdk/types.js"; @@ -14,6 +15,20 @@ import { type McpSessionModelResumeResult, type SessionMeta, } from "./agent-session-durable-object"; +import { + currentInFlightColdBuildCount, + currentResidentRuntimeCount, + EVICTION_REQUEST_GRACE_MS, + markEvictionRequested, + pickEvictionCandidate, + registerResidentSession, + releaseResidentSession, + resetInFlightColdBuildCountForTest, + resetResidentRuntimeCountForTest, + resetResidentSessionRegistryForTest, + residentSessionIdsForTest, + touchResidentSession, +} from "./session-runtime-residency"; class MemoryStorage { private readonly data = new Map(); @@ -85,11 +100,40 @@ class MemoryStorage { return this; } + /** Give this storage's Durable Object a distinct name, so multiple harness + * sessions in the same test resolve to distinct `sessionIdForTelemetry()` + * values instead of all colliding on the default. */ + withIdName(name: string): this { + this.idName = name; + return this; + } + get storage(): MemoryStorage { return this; } - waitUntil(_promise: Promise): void {} + private readonly waitUntilPromises: Promise[] = []; + + /** + * `ctx.waitUntil` extends work past the response instead of blocking it — + * cap eviction requests it. A real DO holds the runtime open until every + * queued promise settles; this fake instead just remembers them so a test + * can explicitly wait for the background work it cares about via + * `drainWaitUntil`, rather than the two racing unpredictably. + */ + waitUntil(promise: Promise): void { + this.waitUntilPromises.push(promise); + } + + /** Test-only: await every `waitUntil`-queued promise queued so far, + * including ones queued BY those promises while draining (an eviction + * request settling can itself queue more background work). */ + async drainWaitUntil(): Promise { + while (this.waitUntilPromises.length > 0) { + const pending = this.waitUntilPromises.splice(0, this.waitUntilPromises.length); + await Promise.allSettled(pending); + } + } } type HarnessSession = { @@ -920,3 +964,720 @@ describe("McpAgentSessionDOBase alarm name resolution", () => { ); }); }); + +// The isolate-wide residency gauge used to be purely observational. These pin +// the enforcing half: past the soft cap, `init` evicts the least-recently +// -active EVICTABLE session to make room, at most one per init, and never +// blocks or fails an init over it even when nothing can be evicted. +// +// `residentRuntimeCount`/the resident-session registry are isolate (module) +// scope, exactly like production — so every test here resets them, both +// before (in case an earlier describe block in this file leaked residents by +// calling `init()` without ever disposing) and after (so it does not leak +// into whatever runs next). +describe("McpAgentSessionDOBase residency cap eviction", () => { + type ResidencySession = { + ctx: MemoryStorage; + captureCause: (cause: Cause.Cause) => void; + dbHandle: { readonly end: () => void } | null; + engine: ExecutionEngine | null; + getConnections: () => Iterable; + getSessionId: () => string; + init: () => Promise; + // Re-exposed for the harness only, same idiom as `evictResidentRuntimeForCap` + // below: the `AbortController` backing the CURRENTLY in-flight `init` + // call's root fiber, so a test can interrupt it deterministically at a + // chosen suspension point instead of only exercising the failure/defect + // path. + initAbortController: AbortController | null; + // Same idiom, for the disposal path — see `disposeAbortController`'s doc + // comment on the class. + disposeAbortController: AbortController | null; + initialized: boolean; + lastActivityMs: number; + pendingApprovalLeases: Map; + props: Record; + residentRuntimeSoftCap: () => number; + server?: McpServer; + sessionIdForTelemetry: () => string; + sessionTimeoutMs: () => number; + withTelemetry: (effect: Effect.Effect) => Effect.Effect; + buildMcpServer: () => Effect.Effect<{ + mcpServer: McpServer; + engine: ExecutionEngine; + }>; + openSessionDb: () => { readonly end: () => void } | Promise<{ readonly end: () => void }>; + resolveSessionMeta: () => Effect.Effect; + supportsCapEviction: () => boolean; + requestSelfEviction: () => Promise; + // Re-exposed for the harness only (private on the real class, same idiom + // as `resolveAndStoreSessionMeta` above): the candidate side of eviction, + // wired below to stand in for a landed self-addressed stub call. + evictResidentRuntimeForCap: () => Promise; + }; + + const residencySessionMeta = (organizationId: string): SessionMeta => ({ + organizationId, + organizationName: "Org 1", + userId: "user-1", + resource: defaultMcpResource, + }); + + /** + * A session harness built to actually run `init()`, the entry point that + * builds (and, on cold restore, rebuilds) a resident runtime. + * + * `requestSelfEviction` stands in for what, in production, is a real + * self-addressed Durable Object stub call — `mcpSessionStub(...).requestCapEviction()` + * — landing back on this same instance and invoking `requestCapEviction`, + * which runs `evictResidentRuntimeForCap` in the (correctly-scoped) IoContext + * that call created. This harness has no Durable Object stub machinery to + * exercise that routing with, so it wires straight to this session's OWN + * `evictResidentRuntimeForCap` instead: it is the candidate-side handler + * that matters for these tests (does the re-check + teardown behave + * correctly), not the RPC transport that gets a request there — that + * transport is covered by the e2e idle-disposal run against real workerd, + * not here. `supportsCapEviction` defaults to `true` so a harness session is + * an eviction candidate unless a test deliberately opts out. + */ + const makeResidencySession = (input: { + readonly id: string; + readonly cap?: number; + readonly hasActiveStream?: boolean; + readonly pausedExecutionCount?: number; + readonly supportsCapEviction?: boolean; + }): { readonly session: ResidencySession; readonly storage: MemoryStorage } => { + const storage = new MemoryStorage().withIdName(`streamable-http:${input.id}`); + const session = Object.create(McpAgentSessionDOBase.prototype) as ResidencySession; + session.ctx = storage; + session.captureCause = () => undefined; + session.dbHandle = null; + session.engine = null; + session.getConnections = () => (input.hasActiveStream ? [{ close: () => undefined }] : []); + session.getSessionId = () => input.id; + session.initialized = false; + session.lastActivityMs = 0; + session.pendingApprovalLeases = new Map(); + session.props = { session: { organizationId: input.id, userId: "user-1" } }; + session.sessionTimeoutMs = () => 60_000; + session.residentRuntimeSoftCap = () => input.cap ?? 32; + session.resolveSessionMeta = () => Effect.succeed(residencySessionMeta(input.id)); + session.openSessionDb = () => ({ end: () => undefined }); + session.buildMcpServer = () => + Effect.succeed({ + mcpServer: makeServer(), + engine: { + ...makeEngine().engine, + pausedExecutionCount: () => Effect.succeed(input.pausedExecutionCount ?? 0), + }, + }); + session.supportsCapEviction = () => input.supportsCapEviction ?? true; + session.requestSelfEviction = () => session.evictResidentRuntimeForCap(); + return { session, storage }; + }; + + /** A tracer that records every span + its attributes, to assert on the + * `mcp.isolate.cap_overflow` attribute the `McpSessionDO.init` span + * carries when nothing was evictable. */ + const makeRecordingTracer = (): { + readonly tracer: Tracer.Tracer; + readonly spans: ReadonlyArray<{ + readonly name: string; + readonly attributes: ReadonlyMap; + }>; + } => { + const spans: Array<{ name: string; attributes: Map }> = []; + const tracer: Tracer.Tracer = { + span: (options) => { + const attributes = new Map(); + spans.push({ name: options.name, attributes }); + let status: Tracer.SpanStatus = { _tag: "Started", startTime: options.startTime }; + return { + _tag: "Span", + name: options.name, + spanId: `span-${spans.length}`, + traceId: "trace-1", + parent: options.parent, + annotations: options.annotations, + get status() { + return status; + }, + attributes, + links: options.links, + sampled: options.sampled, + kind: options.kind, + end: (endTime, exit) => { + status = { _tag: "Ended", startTime: options.startTime, endTime, exit }; + }, + attribute: (key, value) => { + attributes.set(key, value); + }, + event: () => undefined, + addLinks: () => undefined, + }; + }, + }; + return { tracer, spans }; + }; + + beforeEach(() => { + // Earlier describe blocks in this file build sessions via `init()` too, + // without ever disposing them (that is out of scope for what they test) — + // so without this, residency here would start from whatever they left + // behind rather than zero. + resetResidentRuntimeCountForTest(); + resetResidentSessionRegistryForTest(); + resetInFlightColdBuildCountForTest(); + }); + + afterEach(() => { + resetResidentRuntimeCountForTest(); + resetResidentSessionRegistryForTest(); + resetInFlightColdBuildCountForTest(); + }); + + it("evicts the least-recently-active evictable session once the cap is reached, and the newer session survives", async () => { + const { session: sessionA } = makeResidencySession({ id: "session-cap-a", cap: 2 }); + const { session: sessionB } = makeResidencySession({ id: "session-cap-b", cap: 2 }); + const { session: sessionC, storage: storageC } = makeResidencySession({ + id: "session-cap-c", + cap: 2, + }); + + await sessionA.init(); + await sessionB.init(); + expect(currentResidentRuntimeCount(), "two sessions resident, right at the cap").toBe(2); + + // Force a deterministic LRU order: real-clock timestamps from two `init` + // calls a tick apart are not reliable enough to assert on. + touchResidentSession("session-cap-a", Date.now() - 10_000); + touchResidentSession("session-cap-b", Date.now()); + + await sessionC.init(); + // `evictForCapIfNeeded` fires the eviction REQUEST via `ctx.waitUntil` and + // returns without waiting for it, so `sessionC.init()` resolving proves + // nothing about whether session A has actually been torn down yet — only + // draining session C's queued background work does. + await storageC.drainWaitUntil(); + + expect(sessionA.initialized, "the least-recently-active session was evicted").toBe(false); + expect(sessionA.engine, "its engine was released").toBeNull(); + expect(sessionB.initialized, "the newer session survives untouched").toBe(true); + expect(sessionC.initialized, "the session that triggered eviction still built").toBe(true); + expect(currentResidentRuntimeCount(), "one evicted, two resident").toBe(2); + }); + + it("proceeds with init and records overflow when nothing resident is currently evictable", async () => { + const { tracer, spans } = makeRecordingTracer(); + const { session: sessionBusy } = makeResidencySession({ + id: "session-busy", + cap: 1, + hasActiveStream: true, + }); + const { session: sessionNew } = makeResidencySession({ id: "session-new", cap: 1 }); + sessionNew.withTelemetry = (effect) => Effect.withTracer(effect, tracer); + + await sessionBusy.init(); + expect(currentResidentRuntimeCount()).toBe(1); + + await sessionNew.init(); + + expect(sessionBusy.initialized, "a streaming session is never evicted").toBe(true); + expect(sessionNew.initialized, "init proceeds despite the cap").toBe(true); + expect( + currentResidentRuntimeCount(), + "the soft cap is exceeded rather than blocking init", + ).toBe(2); + + const initSpan = spans.find((span) => span.name === "McpSessionDO.init"); + expect(initSpan?.attributes.get("mcp.isolate.cap_overflow")).toBe(true); + }); + + it("cold-restores on the next request after being evicted for the cap", async () => { + const { session: sessionA } = makeResidencySession({ id: "session-restore-a", cap: 1 }); + const { session: sessionB, storage: storageB } = makeResidencySession({ + id: "session-restore-b", + cap: 1, + }); + + await sessionA.init(); + await sessionB.init(); + await storageB.drainWaitUntil(); + expect(sessionA.initialized, "evicted to make room under the cap").toBe(false); + + await sessionA.init(); + + expect(sessionA.initialized, "a later request rebuilds the evicted session").toBe(true); + expect(sessionA.engine, "a fresh engine is installed").not.toBeNull(); + expect(sessionA.server, "a fresh server is installed").toBeDefined(); + }); + + // Two different sessions' `init`s can both pick the SAME LRU candidate + // before either eviction request lands — nothing serializes the picks. The + // candidate must survive that safely: its own re-check plus `closeRuntime`'s + // idempotency is what the base class's doc comments claim makes a repeat + // request a no-op rather than a double-teardown. Requests here run + // sequentially rather than raced, deliberately: the guarantee under test is + // that a SECOND request against an already-evicted (or being-evicted) + // candidate is a safe no-op, not a claim about ordering within a genuine + // race — the same style already used by the "releasing twice" test below. + it("treats a repeat eviction request against the same candidate as a safe no-op", async () => { + const { session: candidate } = makeResidencySession({ id: "session-double-evict" }); + await candidate.init(); + expect(candidate.initialized).toBe(true); + + await candidate.requestSelfEviction(); + + expect(candidate.initialized, "the first request tears the candidate down").toBe(false); + expect(candidate.engine).toBeNull(); + expect(currentResidentRuntimeCount()).toBe(0); + + // A repeat request — the shape two overlapping evictor `init`s produce, or + // one that was merely slow to land after the candidate already tore + // itself down some other way — must also be a no-op, not a crash or a + // second decrement of the (already-zero) resident count. + await expect( + candidate.requestSelfEviction(), + "a repeat request against an already-evicted candidate does not throw", + ).resolves.toBeUndefined(); + expect(candidate.initialized).toBe(false); + expect( + currentResidentRuntimeCount(), + "the resident count is not decremented a second time", + ).toBe(0); + }); + + // `residentRuntimeCount` only moves once a build actually FINISHES, so N + // overlapping cold inits admitted at the same moment each read it, see + // themselves still under the cap, and none of them evicts — residency then + // blows past the cap by N once every build lands. `evictForCapIfNeeded` + // closes that gap with an in-flight reservation counter: this pins that a + // second concurrent admission sees the first's still-building reservation + // and evicts, instead of also passing the count-only check for free. + it("reserves an in-flight cold-build slot at admission, so a second concurrent admission at the cap evicts instead of passing the check for free", async () => { + const { session: sessionA } = makeResidencySession({ id: "session-inflight-a", cap: 2 }); + await sessionA.init(); + expect(currentResidentRuntimeCount(), "one session resident, one below the cap").toBe(1); + + const buildEntered = makeDeferred(); + const buildGate = makeDeferred(); + const { session: sessionB } = makeResidencySession({ id: "session-inflight-b", cap: 2 }); + sessionB.buildMcpServer = () => + Effect.gen(function* () { + buildEntered.resolve(); + yield* Effect.promise(() => buildGate.promise); + return { mcpServer: makeServer(), engine: makeEngine().engine }; + }); + + const initB = sessionB.init(); + // Deterministic: wait for sessionB's admission to actually be inside its + // (gated) build — holding its in-flight reservation — instead of guessing + // a number of microtask ticks. + await buildEntered.promise; + expect(currentInFlightColdBuildCount(), "sessionB's admission reserved a cold-build slot").toBe( + 1, + ); + + const { session: sessionC, storage: storageC } = makeResidencySession({ + id: "session-inflight-c", + cap: 2, + }); + await sessionC.init(); + await storageC.drainWaitUntil(); + + // sessionB is not yet counted-as-resident — its build is still gated — so + // sessionA is the only session actually registered as resident right now + // and the only thing `pickEvictionCandidate` can choose. Its eviction is + // what proves sessionC's admission saw sessionB's in-flight reservation + // and treated the isolate as already at the cap. + expect( + sessionA.initialized, + "evicted because the in-flight reservation counted against the cap", + ).toBe(false); + expect(sessionC.initialized, "the admission that triggered eviction still built").toBe(true); + + buildGate.resolve(); + await initB; + expect(sessionB.initialized, "sessionB's gated build eventually completes").toBe(true); + expect( + currentInFlightColdBuildCount(), + "the reservation was released once the build landed", + ).toBe(0); + }); + + it("releases the in-flight cold-build reservation when the build fails, instead of leaking residual cap pressure", async () => { + const { session: sessionA } = makeResidencySession({ id: "session-inflight-fail-a", cap: 2 }); + await sessionA.init(); + expect(currentResidentRuntimeCount()).toBe(1); + + const { session: sessionFailing } = makeResidencySession({ + id: "session-inflight-fail-b", + cap: 2, + }); + sessionFailing.buildMcpServer = () => Effect.die(new Error("cold build boundary failure")); + + await expect(sessionFailing.init()).rejects.toThrow(/cold build boundary failure/); + expect( + currentInFlightColdBuildCount(), + "the reservation was released on the build's failure path, not leaked", + ).toBe(0); + expect(sessionFailing.initialized, "the failed build never became resident").toBe(false); + + const { session: sessionC, storage: storageC } = makeResidencySession({ + id: "session-inflight-fail-c", + cap: 2, + }); + await sessionC.init(); + await storageC.drainWaitUntil(); + + expect( + sessionA.initialized, + "no residual in-flight pressure from the failed build, so sessionC's admission stays under the cap", + ).toBe(true); + expect( + currentResidentRuntimeCount(), + "sessionA and sessionC both resident, under the cap", + ).toBe(2); + }); + + // The reservation used to be released only by an `Effect.ensuring` scoped to + // the build block that starts right after `evictForCapIfNeeded` admits — + // leaving a gap between the admission itself and that narrower wrap's own + // coverage beginning. An interrupt landing in that gap leaked the + // reservation permanently: nothing ever decremented it, so the isolate's + // in-flight counter drifted up forever and every later admission saw false + // cap pressure. The fix moves the release to a single `Effect.ensuring` + // around init's entire program, so there is no window between acquiring the + // reservation and being covered by its release. This interrupts `init` + // itself (via the `AbortController` `initAbortController` exposes for + // exactly this) while it is genuinely suspended opening the session db + // handle — the first async step after admission — rather than merely + // failing the build, which the pre-existing test above already covers. + it("releases the in-flight cold-build reservation when init is interrupted after cap admission, before the build finishes", async () => { + const { session: sessionA } = makeResidencySession({ id: "session-interrupt-a", cap: 2 }); + await sessionA.init(); + expect(currentResidentRuntimeCount(), "one session resident, one below the cap").toBe(1); + + const dbHandleEntered = makeDeferred(); + const dbHandleGate = makeDeferred(); + const { session: sessionB } = makeResidencySession({ id: "session-interrupt-b", cap: 2 }); + sessionB.openSessionDb = () => { + dbHandleEntered.resolve(); + // Never resolves — the interrupt below is what ends this suspension, + // not the gate. `dbHandleGate.promise` only carries a `void` payload; + // `.then` produces the right shape without a cast, and without ever + // actually resolving — the callback here never runs. + return dbHandleGate.promise.then(() => ({ end: () => undefined })); + }; + + const initPromise = sessionB.init(); + // Deterministic: wait for sessionB's admission to have actually reserved + // a slot and for its build to be genuinely suspended opening the db + // handle, instead of guessing a number of microtask ticks. + await dbHandleEntered.promise; + expect( + currentInFlightColdBuildCount(), + "admission reserved a slot before the build's own db-handle open resolves", + ).toBe(1); + + sessionB.initAbortController?.abort(); + + await expect( + initPromise, + "an interrupted init rejects rather than silently resolving", + ).rejects.toThrow(); + expect( + currentInFlightColdBuildCount(), + "the reservation was released by init's outer ensuring instead of leaking", + ).toBe(0); + expect(sessionB.initialized, "the interrupted build never became resident").toBe(false); + + // The reservation not leaking is what lets a later admission still see + // the isolate as under the cap. + const { session: sessionC, storage: storageC } = makeResidencySession({ + id: "session-interrupt-c", + cap: 2, + }); + await sessionC.init(); + await storageC.drainWaitUntil(); + + expect( + sessionA.initialized, + "no residual in-flight pressure from the interrupted init, so sessionC's admission stays under the cap", + ).toBe(true); + expect( + currentResidentRuntimeCount(), + "sessionA and sessionC both resident, under the cap", + ).toBe(2); + }); + + // `initialized` used to stay `true` across the async closes inside + // `closeRuntime` (`server.close()`, `dbHandle.end()`), so a request landing + // mid-teardown took `init`'s early-return path and ran against a + // deleted server / null engine. `disposingRuntime` closes that: `init` + // awaits any in-progress disposal before deciding whether to rebuild. + it("a request landing mid-disposal awaits the in-progress teardown, then rebuilds and serves — never racing the closes", async () => { + const { session } = makeResidencySession({ id: "session-mid-disposal" }); + await session.init(); + expect(session.initialized).toBe(true); + const originalEngine = session.engine; + + const closeEntered = makeDeferred(); + const closeGate = makeDeferred(); + let closeCalls = 0; + const server = makeServer(); + server.close = () => { + closeCalls += 1; + closeEntered.resolve(); + return closeGate.promise; + }; + session.server = server; + + const disposal = session.evictResidentRuntimeForCap(); + // Deterministic: wait for disposal to actually be mid-teardown (inside + // `server.close()`), instead of guessing a number of microtask ticks. + await closeEntered.promise; + + const rebuild = session.init(); + let rebuildSettled = false; + void rebuild.then(() => { + rebuildSettled = true; + }); + await Promise.resolve(); + await Promise.resolve(); + expect( + rebuildSettled, + "init awaits the in-progress disposal instead of racing its still-pending close", + ).toBe(false); + + closeGate.resolve(); + await disposal; + await rebuild; + + expect(closeCalls, "the runtime was closed exactly once").toBe(1); + expect(session.initialized, "the request that landed mid-disposal rebuilt and is serving").toBe( + true, + ); + expect( + session.engine, + "a fresh engine was installed, not the one that was mid-teardown", + ).not.toBe(originalEngine); + }); + + // The `Effect.ensuring` that resolves `disposingRuntime` used to sit outside + // any interruptibility guard, so an interrupt landing while `closeRuntime` + // was genuinely suspended (mid `server.close()`, same seam as the test + // above) still ran that `ensuring` and resolved `disposingRuntime` + // immediately — before `dbHandle.end()`, the engine clear, or the residency + // release ever ran. A concurrently waiting `init` would see the gate open + // and rebuild against that half-released state. Wrapping the teardown body + // in `Effect.uninterruptible` makes the interrupt request wait until the + // body — including the still-gated `server.close()` — actually finishes, + // so `disposingRuntime` cannot resolve early. This interrupts the disposal + // itself (via `disposeAbortController`, exposed for exactly this) instead + // of just failing a step, which the double-close test above doesn't cover. + it("keeps a waiting init blocked when the disposal is interrupted mid-teardown, until the uninterruptible teardown actually finishes", async () => { + const { session } = makeResidencySession({ id: "session-interrupted-disposal" }); + await session.init(); + expect(session.initialized).toBe(true); + const originalEngine = session.engine; + + const closeEntered = makeDeferred(); + const closeGate = makeDeferred(); + let closeCalls = 0; + const server = makeServer(); + server.close = () => { + closeCalls += 1; + closeEntered.resolve(); + return closeGate.promise; + }; + session.server = server; + + const disposal = session.evictResidentRuntimeForCap(); + // Deterministic: wait for disposal to actually be mid-teardown (inside + // `server.close()`), instead of guessing a number of microtask ticks. + await closeEntered.promise; + + // A concurrent init is the observable proof: it only proceeds once + // `disposingRuntime` resolves, so it stays pending for exactly as long as + // the teardown is genuinely still running. + const rebuild = session.init(); + let rebuildSettled = false; + void rebuild.then(() => { + rebuildSettled = true; + }); + + // Interrupt the disposal fiber while it is still suspended inside the + // gated `server.close()`. On the old (interruptible) code this would let + // the outer `Effect.ensuring` resolve `disposingRuntime` right here, + // before `closeGate` ever resolves — letting `rebuild` proceed against a + // still-open server / un-cleared engine. + session.disposeAbortController?.abort(); + // Interrupt delivery hops through more than a couple of microtasks (it is + // not itself under test here), so give it a real macrotask tick before + // asserting anything stayed blocked — a few `Promise.resolve()`s alone + // are not enough ticks for the abort to even be delivered, uninterruptible + // or not, and would make this assertion true vacuously either way. + await new Promise((resolve) => setTimeout(resolve, 20)); + expect( + rebuildSettled, + "the interrupt must not resolve disposingRuntime early — the teardown is still genuinely suspended inside server.close()", + ).toBe(false); + expect( + session.engine, + "the interrupt must not let the engine be cleared before server.close() actually returns", + ).toBe(originalEngine); + + closeGate.resolve(); + // The interrupt may still surface once the now-uninterruptible body has + // actually finished; only completeness of the teardown itself — not the + // outer promise's resolve/reject outcome — is under test here. + try { + await disposal; + } catch { + // Expected: the interrupt requested above can still land once the + // uninterruptible teardown finishes. + } + await rebuild; + + expect(closeCalls, "server.close was actually called, not skipped").toBe(1); + expect( + session.initialized, + "the request that was waiting on the interrupted disposal rebuilt and is serving", + ).toBe(true); + expect( + session.engine, + "a fresh engine was installed once the teardown genuinely finished, not the one that was mid-teardown", + ).not.toBe(originalEngine); + }); + + it("does not double-close when two disposal triggers land concurrently (idle alarm and a cap eviction landing together)", async () => { + const { session } = makeResidencySession({ id: "session-concurrent-dispose" }); + await session.init(); + + const closeGate = makeDeferred(); + let closeCalls = 0; + const server = makeServer(); + server.close = () => { + closeCalls += 1; + return closeGate.promise; + }; + session.server = server; + + const first = session.evictResidentRuntimeForCap(); + const second = session.evictResidentRuntimeForCap(); + + closeGate.resolve(); + await Promise.all([first, second]); + + expect( + closeCalls, + "the runtime was closed exactly once even though two disposals overlapped", + ).toBe(1); + expect(session.initialized).toBe(false); + expect(currentResidentRuntimeCount(), "released exactly once").toBe(0); + }); + + describe("resident session registry release", () => { + it("is idempotent", () => { + registerResidentSession({ + sessionId: "idempotent-release", + lastActivityMs: Date.now(), + canEvict: () => true, + dispose: async () => undefined, + }); + + expect(() => releaseResidentSession("idempotent-release")).not.toThrow(); + expect( + () => releaseResidentSession("idempotent-release"), + "releasing twice is a no-op", + ).not.toThrow(); + expect(residentSessionIdsForTest()).not.toContain("idempotent-release"); + expect(pickEvictionCandidate()).toBeUndefined(); + }); + + // The eviction REQUEST is fire-and-forget (`ctx.waitUntil`), so the + // evictor's `init` has no synchronous failure to react to any more — it + // cannot tell a request that failed from one still in flight. The new + // failure story: optimistically LEAVE the registry entry (only the + // candidate's own successful teardown removes it, since as far as the + // evictor knows the candidate might still be actually resident), and mark + // it recently-requested so a following pick skips it rather than + // re-requesting the same stuck candidate forever. + it("keeps the registry slot when the candidate's own eviction request fails, without touching the resident count", async () => { + const { session: sessionFailing } = makeResidencySession({ + id: "session-failing", + cap: 1, + }); + // Simulate the REQUEST itself failing — the self-addressed stub call + // never lands (the production analogue: `mcpSessionStub(...).requestCapEviction()` + // rejects), never even reaching the candidate's own re-check. + // oxlint-disable-next-line executor/no-promise-reject -- test double: simulates a rejected self-eviction stub call, not application error modeling. + sessionFailing.requestSelfEviction = () => Promise.reject(new Error("stub unreachable")); + const { session: sessionB, storage: storageB } = makeResidencySession({ + id: "session-b-after-failure", + cap: 1, + }); + + await sessionFailing.init(); + await sessionB.init(); + await storageB.drainWaitUntil(); + + expect( + sessionFailing.initialized, + "a failed eviction REQUEST never runs the candidate's own teardown", + ).toBe(true); + expect(sessionB.initialized, "the triggering init is never blocked by the failure").toBe( + true, + ); + expect( + currentResidentRuntimeCount(), + "the failing candidate's own runtime is still actually resident", + ).toBe(2); + expect( + residentSessionIdsForTest(), + "the registry slot is kept — only the candidate's own successful teardown removes it", + ).toContain("session-failing"); + expect( + pickEvictionCandidate()?.sessionId, + "marked recently-requested, so the next pick does not re-target the same stuck candidate", + ).not.toBe("session-failing"); + }); + + it("skips an entry with a recent pending eviction request, and picks the next evictable one instead", () => { + const now = Date.now(); + registerResidentSession({ + sessionId: "recently-requested", + lastActivityMs: now - 10_000, + canEvict: () => true, + dispose: async () => undefined, + }); + registerResidentSession({ + sessionId: "next-candidate", + lastActivityMs: now - 5_000, + canEvict: () => true, + dispose: async () => undefined, + }); + markEvictionRequested("recently-requested", now); + + expect(pickEvictionCandidate(now)?.sessionId).toBe("next-candidate"); + }); + + it("picks a previously-requested entry again once the grace period elapses", () => { + const now = Date.now(); + registerResidentSession({ + sessionId: "stale-request", + lastActivityMs: now - 10_000, + canEvict: () => true, + dispose: async () => undefined, + }); + markEvictionRequested("stale-request", now - EVICTION_REQUEST_GRACE_MS - 1); + + expect(pickEvictionCandidate(now)?.sessionId).toBe("stale-request"); + }); + + it("marking an eviction request on an entry that no longer exists is a no-op", () => { + expect(() => markEvictionRequested("never-registered")).not.toThrow(); + }); + }); +}); diff --git a/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts b/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts index 4afc54651..dfe108719 100644 --- a/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts +++ b/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts @@ -36,8 +36,19 @@ import { } from "./session-alarm-policy"; import { acquireResidentRuntime, + currentInFlightColdBuildCount, + currentResidentRuntimeCount, + markEvictionRequested, + pickEvictionCandidate, + registerResidentSession, + releaseColdBuildSlot, releaseResidentRuntime, + releaseResidentSession, + reserveColdBuildSlot, + type ResidentSessionEntry, residencyAttributes, + RESIDENT_RUNTIME_SOFT_CAP, + touchResidentSession, } from "./session-runtime-residency"; export type IncomingTraceHeaders = IncomingPropagationHeaders; @@ -261,6 +272,43 @@ export abstract class McpAgentSessionDOBase< * gauge. Tracked separately from `engine` because `closeRuntime` runs on * paths where nothing was ever built, and it must not decrement then. */ private countedAsResident = false; + /** Whether THIS init's cold-build admission is currently holding a reserved + * slot in the isolate-wide in-flight counter. Mirrors `countedAsResident`'s + * guard discipline: gates {@link releaseColdBuildSlotIfReserved} so it + * releases exactly once per reservation, from whichever of its two callers + * (success or failure/interrupt) gets there first. */ + private reservedColdBuildSlot = false; + /** + * The `AbortController` backing the currently in-flight `init` call's + * `Effect.runPromise`, so its root fiber can be interrupted from outside + * the Promise it returns. Nothing in production ever calls `.abort()` on + * it today — it exists so a unit test can interrupt `init` deterministically + * at a chosen suspension point (e.g. mid `openSessionDbHandle`) and assert + * the cold-build reservation is still released, rather than only ever being + * exercised through failure/defect paths. + */ + private initAbortController: AbortController | null = null; + /** Same purpose as {@link initAbortController}, for `disposeIdleRuntime`'s + * root effect (`closeRuntime` plus its alarm/activity bookkeeping) — lets + * a unit test interrupt a disposal deterministically to confirm the + * uninterruptible teardown in `closeRuntime` actually finishes rather than + * being cut short. */ + private disposeAbortController: AbortController | null = null; + /** + * The in-progress runtime disposal, if one is running right now — from + * whichever of `closeRuntime`'s callers (the idle alarm, a cap eviction + * request, `cleanup`) got there first. Set at the very start of + * `closeRuntime`'s body, before its first async close, and cleared in a + * finally once that disposal settles. + * + * Exists to close two races at once: `init` awaits this before deciding + * whether to rebuild, so a request that lands mid-teardown never races the + * closes it is waiting on; and a second `closeRuntime` call that lands + * while this is set (the idle alarm and a cap eviction request landing + * together) waits on the SAME promise instead of running the teardown a + * second time. + */ + private disposingRuntime: Promise | null = null; private onStartPromise: Promise | null = null; private lastActivityMs = 0; private resolvedSessionName: string | undefined = undefined; @@ -459,6 +507,49 @@ export abstract class McpAgentSessionDOBase< }); } + /** + * Whether this host can route an eviction REQUEST to this session's own + * Durable Object instance rather than tearing it down directly in some + * OTHER session's request context. Hosts that override `requestSelfEviction` + * with a real self-addressed stub call return `true` here too. A host that + * returns `false` (the default) is never registered as an eviction + * candidate at all — see `init` — so it degrades to purely observational: + * the residency gauge and cap-overflow attribute still work, cap eviction + * just never picks it. + */ + protected supportsCapEviction(): boolean { + return false; + } + + /** + * Ask THIS session's own Durable Object instance — running in ITS OWN + * request/IoContext — to tear down its resident runtime because the + * isolate is over its cap. + * + * This must never be `evictResidentRuntimeForCap()` called directly on + * `this` from within another session's request. Even though that would be + * the exact same JS object in the exact same isolate (Durable Objects with + * the same id ARE the same instance), a plain method call does not create a + * new IoContext — it runs inside whatever IoContext is already current, + * which belongs to the CALLING session's request. workerd binds I/O objects + * (a postgres.js socket, a storage transaction, a span flush) to the + * IoContext that created them, so tearing this session down that way throws + * "Cannot perform I/O on behalf of a different request" or silently + * miscredits the I/O to the wrong request. Routing through this session's + * own Durable Object STUB (`mcpSessionStub(...).requestCapEviction()`) goes + * through the Workers RPC/fetch machinery instead, which gives the call a + * freshly-created IoContext bound to itself — so the teardown it triggers + * runs correctly scoped, no matter which session's `init` sent the request. + * + * Overridden per host, because only a concrete host knows its own + * self-addressed namespace binding. The base default is a no-op so a host + * that never overrides it (and therefore never overrides + * `supportsCapEviction` to `true`) is simply never asked. + */ + protected requestSelfEviction(): Promise { + return Promise.resolve(); + } + protected readonly browserApprovalStore: BrowserApprovalStore = { takeResponse: (executionId) => this.takeApprovalResponse(executionId), waitForResponse: (executionId) => this.waitForApprovalResponse(executionId), @@ -540,6 +631,11 @@ export abstract class McpAgentSessionDOBase< private async markActivity(now = Date.now()): Promise { this.lastActivityMs = now; + // Keeps the isolate-wide eviction registry's LRU order current. A no-op + // when this session has no registry entry yet (nothing resident) or none + // any more (already disposed) — `touchResidentSession` is a lookup-then-set + // that quietly does nothing on a miss. + touchResidentSession(this.sessionIdForTelemetry(), now); await Promise.all([ this.ctx.storage.put(LAST_ACTIVITY_KEY, now), this.ctx.storage.setAlarm(now + this.sessionTimeoutMs()), @@ -662,15 +758,26 @@ export abstract class McpAgentSessionDOBase< } /** - * Drop this session's execution runtime because it has gone idle, returning - * its memory to the isolate. Nothing durable is discarded, so the next - * request restores the session and the client sees only restore latency. + * Drop this session's execution runtime, returning its memory to the + * isolate. Nothing durable is discarded, so the next request restores the + * session and the client sees only restore latency. + * + * `reason` disambiguates WHY on the shared span and log line without + * splitting them: `"idle"` is the alarm-driven path (this session itself + * went quiet), `"cap"` is another session's `init` evicting this one because + * the isolate was over its resident-runtime ceiling. The mechanism — + * `closeRuntime` plus dropping the durable alarm/activity bookkeeping — is + * identical either way; only the trigger differs, and callers are expected + * to have already established (via their own eligibility check) that + * disposing this session right now is safe. */ private async disposeIdleRuntime(input: { readonly idleMs: number; readonly pausedExecutionCount: number; readonly activeStreamCount: number; + readonly reason?: "idle" | "cap"; }): Promise { + const reason = input.reason ?? "idle"; console.info( JSON.stringify({ event: "mcp_session_idle_runtime_dispose", @@ -678,6 +785,7 @@ export abstract class McpAgentSessionDOBase< idleMs: input.idleMs, pausedExecutionCount: input.pausedExecutionCount, activeStreamCount: input.activeStreamCount, + reason, }), ); const self = this; @@ -691,12 +799,15 @@ export abstract class McpAgentSessionDOBase< // isolate is actually holding now rather than what it held a moment ago. yield* Effect.annotateCurrentSpan(residencyAttributes()); }).pipe( + // Span name kept stable for dashboard continuity across both triggers; + // `mcp.session.dispose_reason` is what disambiguates them. Effect.withSpan("mcp.session.idle_runtime_dispose", { attributes: { "mcp.session.id": self.sessionId, "mcp.session.idle_ms": input.idleMs, "mcp.session.paused_execution_count": input.pausedExecutionCount, "mcp.session.active_stream_count": input.activeStreamCount, + "mcp.session.dispose_reason": reason, }, }), ); @@ -704,7 +815,194 @@ export abstract class McpAgentSessionDOBase< // trace. It still has to be flushed explicitly — the alarm is not on any // request's response path, and without the flush the span dies with the // isolate and the mechanism stays unobservable in production. - await Effect.runPromise(this.withSpanFlush(this.withTelemetry(program))); + // See `disposeAbortController`'s doc comment: nothing in production aborts + // this signal today; it exists so a unit test can interrupt this exact + // fiber and confirm `closeRuntime`'s now-uninterruptible teardown still + // runs to completion instead of being cut short. + const abortController = new AbortController(); + this.disposeAbortController = abortController; + await Effect.runPromise(this.withSpanFlush(this.withTelemetry(program)), { + signal: abortController.signal, + }).finally(() => { + // Only clear the field if it is still THIS call's controller — see the + // matching comment in `init` for why an overlapping later call's + // controller must not be clobbered. + if (this.disposeAbortController === abortController) this.disposeAbortController = null; + }); + } + + /** + * Isolate is over its resident-runtime soft cap and this session was the + * LRU-eligible pick (see `pickEvictionCandidate`). Re-checks every + * disqualifying signal `decideSessionAlarm` treats as active work — LIVE, + * not the snapshot `canEvict` used to be picked — because a session can + * start a request in the gap between being picked and being disposed here. + * `canEvict`'s snapshot is a cheap, synchronous, necessarily-optimistic + * filter for CHOOSING among candidates; this is the authoritative gate that + * actually decides whether disposing this session right now is safe, and it + * includes the one signal `canEvict` cannot see synchronously — undelivered + * stream responses still in storage. A session that is no longer eligible is + * left alone: eviction becomes a no-op instead of a wrongful teardown. + */ + private async evictResidentRuntimeForCap(): Promise { + const [pausedExecutionCount, runningExecutionCount] = await Promise.all([ + this.pausedExecutionCount(), + this.runningExecutionCount(), + ]); + const activeStreamCount = this.activeStreamCount(); + if (pausedExecutionCount > 0 || runningExecutionCount > 0 || activeStreamCount > 0) return; + const idleMs = this.lastActivityMs > 0 ? Date.now() - this.lastActivityMs : 0; + await this.disposeIdleRuntime({ + idleMs, + pausedExecutionCount, + activeStreamCount, + reason: "cap", + }); + } + + /** + * Cheap, synchronous eligibility filter consulted by `pickEvictionCandidate` + * to choose AMONG resident sessions. Deliberately conservative rather than + * exhaustive: it mirrors the same "no active work" signals + * `decideSessionAlarm` treats as disqualifying, restricted to what can be + * answered without an async storage read (undelivered stream responses, + * `runningExecutionCount`, requires one). `evictResidentRuntimeForCap` + * re-checks the full set — including that signal — right before actually + * disposing, so a false "evictable" here can only ever produce a safe + * no-op, never a wrongful eviction. + */ + private canEvictResidentRuntime(): boolean { + if (this.activeStreamCount() > 0) return false; + if (!this.engine) return true; + // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: this is a best-effort eviction-selection filter; a broken engine read must fail toward "not evictable", never toward crashing whichever OTHER session's init is picking a candidate. + try { + return Effect.runSync(this.engine.pausedExecutionCount()) === 0; + } catch { + return false; + } + } + + /** + * The soft cap this instance enforces. A protected method rather than a bare + * reference to `RESIDENT_RUNTIME_SOFT_CAP`, matching `sessionTimeoutMs` and + * `maxPausedSessionIdleMs` elsewhere in this class, so tests can install a + * small cap and exercise real eviction without registering 32 sessions. + */ + protected residentRuntimeSoftCap(): number { + return RESIDENT_RUNTIME_SOFT_CAP; + } + + /** + * Before this session's own runtime is built, make room if the isolate is + * already at its resident-runtime cap. Picks AT MOST ONE other session — + * never loops — and never blocks or fails THIS init on the outcome: the + * eviction REQUEST is fired via `ctx.waitUntil` and this init proceeds + * immediately, because a memory-pressure mechanism must never itself become + * the reason a session fails or is delayed starting. If nothing is + * currently evictable (every resident is streaming, paused, or already has + * a request outstanding), this session still builds and the overflow is + * only recorded on the init span. + * + * `candidate.dispose` does not run the candidate's teardown here — it SENDS + * the candidate a request that ITS OWN Durable Object instance executes in + * its own context (see `requestSelfEviction`'s doc comment). Because the + * request is fire-and-forget, this init has no way to react to it failing: + * the registry entry is left in place either way (never released here on + * failure), since whatever the candidate holds is, as far as this init + * knows, still actually resident — only the candidate's own successful + * teardown removes its entry. `markEvictionRequested` timestamps the entry + * up front so a candidate whose request is stuck or failing does not squat + * the LRU pick forever: the next init's `pickEvictionCandidate` skips it + * (see the grace period there) and picks the NEXT candidate instead. + * + * Safe to fire at the same candidate twice — two different sessions' inits + * both picking the same LRU entry before either's request lands — because + * the candidate's own handler (`evictResidentRuntimeForCap`) re-checks + * liveness and its teardown (`closeRuntime`) is idempotent: a second + * request either finds the runtime already gone (a no-op) or finds it + * newly busy again and leaves it alone. + * + * The check reads `currentResidentRuntimeCount() + currentInFlightColdBuildCount()`, + * not `currentResidentRuntimeCount()` alone. Residency only moves once a + * cold build actually finishes, so without the in-flight term, N + * overlapping cold inits arriving at the cap would each read the same + * still-under-cap count before any of them finishes, none would evict, and + * residency would land N over the cap once every build completed. Admitting + * (reserving a slot) unconditionally below — whether or not eviction fires + * for THIS init — is what lets the NEXT concurrent init see this one + * reflected in the sum. + */ + private evictForCapIfNeeded(): Effect.Effect { + const self = this; + return Effect.gen(function* () { + const overCap = + currentResidentRuntimeCount() + currentInFlightColdBuildCount() >= + self.residentRuntimeSoftCap(); + // Admission: this init is about to start a cold build, so it counts as + // in-flight from here whether or not the check above finds anything to + // evict. Released exactly once `evictForCapIfNeeded`'s caller (`init`) + // either finishes building or fails/is interrupted — see + // `releaseColdBuildSlotIfReserved`. + self.reservedColdBuildSlot = true; + reserveColdBuildSlot(); + if (!overCap) return; + const candidate = pickEvictionCandidate(); + if (!candidate) { + yield* Effect.annotateCurrentSpan({ "mcp.isolate.cap_overflow": true }); + return; + } + yield* Effect.sync(() => self.queueCapEvictionRequest(candidate)); + }); + } + + /** + * Release this init's in-flight cold-build reservation (see + * `evictForCapIfNeeded`), if it is still holding one. Idempotent, mirroring + * `countedAsResident`'s guard discipline: called from both of its possible + * finishing points — the build's success path in `init`, and an + * `Effect.ensuring` finalizer covering the build's failure/interrupt path — + * so whichever happens releases the slot exactly once, and the other + * becomes a safe no-op. + */ + private releaseColdBuildSlotIfReserved(): Effect.Effect { + const self = this; + return Effect.sync(() => { + if (!self.reservedColdBuildSlot) return; + self.reservedColdBuildSlot = false; + releaseColdBuildSlot(); + }); + } + + /** + * Fires the eviction REQUEST at `candidate`'s own stub and returns + * immediately — same fire-and-forget shape as + * `queuePendingApprovalLeaseStart`/`queuePendingApprovalLeaseExpiration` + * below. `markEvictionRequested` is stamped up front, before the request + * even lands, so a stuck or slow candidate cannot be re-picked by the next + * init in the meantime (see `pickEvictionCandidate`'s grace period). + */ + private queueCapEvictionRequest(candidate: ResidentSessionEntry): void { + markEvictionRequested(candidate.sessionId); + this.ctx.waitUntil( + Effect.runPromise( + Effect.tryPromise({ + try: () => candidate.dispose("cap"), + catch: (cause: unknown) => cause, + }).pipe( + Effect.catch((cause: unknown) => + Effect.sync(() => { + console.warn( + JSON.stringify({ + event: "mcp_session_cap_eviction_request_failed", + sessionId: candidate.sessionId, + }), + ); + console.error("[mcp-session] cap eviction request failed:", cause); + }), + ), + ), + ), + ); } private resolveAndStoreSessionMeta(token: McpSessionInit) { @@ -897,7 +1195,32 @@ export abstract class McpAgentSessionDOBase< private closeRuntime(options: { readonly closeStreams?: boolean } = {}): Effect.Effect { const self = this; + // A disposal is already tearing this runtime down — the idle alarm and a + // cap eviction request landing together, or a plain repeat call on any of + // `closeRuntime`'s existing callers. `server.close()`/`dbHandle.end()` + // are not safe to run twice concurrently on the same resources, and a + // second pass through the body below would double-release the residency + // counters. Wait for the SAME in-progress disposal instead of starting a + // second one; do not run the teardown body again. + if (self.disposingRuntime) { + const inProgress = self.disposingRuntime; + return Effect.promise(() => inProgress); + } + let resolveDisposal!: () => void; + const disposal = new Promise((resolve) => { + resolveDisposal = resolve; + }); + self.disposingRuntime = disposal; return Effect.gen(function* () { + // Flip this BEFORE the first async close below (`server.close()`), not + // after. `init` awaits `disposingRuntime` (set above) before deciding + // whether to rebuild, and it only gets the right answer because + // `initialized` is already `false` by the time that await resolves — + // otherwise a request that interleaved during the closes below would + // still see `initialized === true`, take `init`'s early-return path, + // and run against a server/engine that are mid-teardown or already + // gone. + self.initialized = false; yield* self.releaseAllPendingApprovalLeases(); if (options.closeStreams ?? true) { yield* Effect.sync(() => self.closeActiveStreams()); @@ -905,21 +1228,58 @@ export abstract class McpAgentSessionDOBase< if (self.server) { const server = self.server; delete (self as { server?: McpServer }).server; - yield* Effect.promise(() => server.close()).pipe(Effect.ignore); + // `tryPromise`, not `promise`: a rejected close must land in the error + // channel where `ignore` absorbs it. With `Effect.promise` a rejection + // becomes a defect, which `ignore` does NOT absorb — the teardown + // would stop here while the `ensuring` below still resolved + // `disposingRuntime`, telling a waiting `init` the resources were + // released when the steps after this one never ran. + yield* Effect.tryPromise({ + try: () => server.close(), + catch: (cause: unknown) => cause, + }).pipe(Effect.ignore); } Reflect.set(self, "_transport", undefined); self.engine = null; if (self.dbHandle) { const dbHandle = self.dbHandle; self.dbHandle = null; - yield* Effect.promise(() => Promise.resolve(dbHandle.end())).pipe(Effect.ignore); + // Same `tryPromise` reasoning as the server close above. + yield* Effect.tryPromise({ + try: () => Promise.resolve(dbHandle.end()), + catch: (cause: unknown) => cause, + }).pipe(Effect.ignore); } - self.initialized = false; if (self.countedAsResident) { self.countedAsResident = false; releaseResidentRuntime(); + // Pairs with `registerResidentSession` in `init`. Idempotent, and + // gated on the same flag that guards the counter release, so a + // `closeRuntime` that runs on a path where nothing was ever built + // never removes an entry it did not add. + releaseResidentSession(self.sessionIdForTelemetry()); } - }); + }).pipe( + // Uninterruptible once teardown starts, so it always runs to + // completion. The `Effect.ensuring` below resolves `disposingRuntime` + // unconditionally — a waiting `init` treats that resolution as "the + // resources are actually released" and proceeds to rebuild — which is + // only correct if nothing here can be interrupted or half-run partway + // through. Every step above is already bounded and safe to run + // uninterruptibly: the two closes route rejections into the error + // channel via `tryPromise` and `ignore` them (a bare `Effect.promise` + // would turn a rejection into a defect `ignore` cannot absorb), and + // `releaseAllPendingApprovalLeases` ignores its own failures + // (`deleteExecutionOwnerEntry` is `Effect.ignore`d too), so nothing + // here can defect either. + Effect.uninterruptible, + Effect.ensuring( + Effect.sync(() => { + if (self.disposingRuntime === disposal) self.disposingRuntime = null; + resolveDisposal(); + }), + ), + ); } private ensureRuntimeForApproval(): Effect.Effect { @@ -978,6 +1338,16 @@ export abstract class McpAgentSessionDOBase< } async init(): Promise { + if (this.disposingRuntime) { + // A disposal (idle alarm, cap eviction) is mid-teardown for this + // session's own Durable Object instance. `closeRuntime` flips + // `initialized` to `false` before its first async close specifically so + // this does not race it: wait for the disposal to actually finish — + // server closed, engine and db handle released, residency counters + // updated — before deciding whether a rebuild is even needed, instead + // of starting one against half-torn-down state. + await this.disposingRuntime; + } if (this.initialized) return; const props = isSessionProps(this.props) ? this.props : null; if (!props) { @@ -988,8 +1358,26 @@ export abstract class McpAgentSessionDOBase< const program = Effect.gen(function* () { yield* self.prepareErrorCaptureScope(); const sessionMeta = yield* self.resolveAndStoreSessionMeta(props.session); - const dbHandle = yield* self.openSessionDbHandle(); - const { mcpServer, engine } = yield* self.buildRuntime(sessionMeta, dbHandle); + // Before building anything that will itself occupy isolate memory + // (a live db handle, the engine, the built tool catalog), make room if + // this isolate is already at its resident-runtime cap. Deliberately + // BEFORE `openSessionDbHandle`, not just before `buildRuntime`: the db + // handle is one of the three things a resident runtime holds. + yield* self.evictForCapIfNeeded(); + // The in-flight cold-build reservation `evictForCapIfNeeded` just took + // is released exactly once — success, failure, or interrupt — by the + // single `Effect.ensuring` wrapped around this ENTIRE program below, + // not a narrower one scoped to just this build block. A narrower wrap + // still leaves a gap between the reservation being taken above and the + // wrap starting here: an interrupt landing in exactly that gap would + // leak the reservation forever (the module counter only drifts up), + // which is why the release is anchored to the same scope as the + // acquisition instead. + const { dbHandle, mcpServer, engine } = yield* Effect.gen(function* () { + const dbHandle = yield* self.openSessionDbHandle(); + const { mcpServer, engine } = yield* self.buildRuntime(sessionMeta, dbHandle); + return { dbHandle, mcpServer, engine }; + }); self.dbHandle = dbHandle; self.server = mcpServer; self.engine = engine; @@ -997,6 +1385,26 @@ export abstract class McpAgentSessionDOBase< if (!self.countedAsResident) { self.countedAsResident = true; acquireResidentRuntime(); + // Only a host that can route an eviction request back to THIS + // session's own Durable Object instance (see `requestSelfEviction`) + // registers as a candidate at all. A host that cannot is still + // counted in the gauge above — cap-overflow tracking stays accurate — + // it is just never picked, degrading to observational rather than + // running teardown in the wrong context. + if (self.supportsCapEviction()) { + // Paired with `releaseResidentSession` in `closeRuntime`. Registered + // as soon as this session counts as resident, so a cap check running + // in another session's `init` moments later already sees it as a + // candidate. `markActivity` below immediately corrects the initial + // timestamp via `touchResidentSession`, so `Date.now()` here only + // needs to be a safe placeholder, not the true last-activity time. + registerResidentSession({ + sessionId: self.sessionIdForTelemetry(), + lastActivityMs: Date.now(), + canEvict: () => self.canEvictResidentRuntime(), + dispose: () => self.requestSelfEviction(), + }); + } } // The gauge on the way up. Paired with the same attributes on // `mcp.session.idle_runtime_dispose`, this is what shows whether idle @@ -1010,6 +1418,13 @@ export abstract class McpAgentSessionDOBase< .bestEffortBookkeeping("init.mark_activity", () => self.markActivity()) .pipe(Effect.withSpan("McpSessionDO.markActivity")); }).pipe( + // Covers the ENTIRE program above, not just the build block: anything + // from `evictForCapIfNeeded`'s admission onward that ends this effect — + // success, failure, or interrupt — releases the cold-build reservation + // exactly once (`releaseColdBuildSlotIfReserved` is idempotent). Wrapping + // only the build block would leave the gap between admission and the + // build starting uncovered. + Effect.ensuring(self.releaseColdBuildSlotIfReserved()), // ONE capture owner for an init defect. `init` can only reject its // Promise, and the host's DO-level error instrumentation captures that // rejection too — so the DO claims the cause below and the host drops its @@ -1047,12 +1462,53 @@ export abstract class McpAgentSessionDOBase< }), ); const traced = this.withTelemetry(program, props?.propagation); + // See `initAbortController`'s doc comment: nothing in production aborts + // this signal today, but wiring it through `runPromise` gives a unit test + // a real handle to interrupt this exact fiber deterministically instead + // of only ever exercising the failure/defect path. + const abortController = new AbortController(); + self.initAbortController = abortController; return Effect.runPromise( traced.pipe( // oxlint-disable-next-line executor/no-effect-escape-hatch -- boundary: Durable Object init method can only reject its Promise Effect.orDie, (effect) => self.withSpanFlush(effect), ), + { signal: abortController.signal }, + ).finally(() => { + // Only clear the field if it is still THIS call's controller — an + // overlapping later `init` call may already have installed its own by + // the time this one settles, and clobbering that with `null` would + // leave a unit test unable to reach it. + if (self.initAbortController === abortController) self.initAbortController = null; + }); + } + + /** + * The candidate side of cap eviction. Called only through THIS session's own + * Durable Object stub — see `requestSelfEviction` — never invoked directly + * on an in-process reference by another session, which is the whole point: + * routing the call through the stub gives it a correctly-scoped IoContext + * for `evictResidentRuntimeForCap`'s teardown to run in. + * + * Safe to call more than once, including two overlapping calls (two + * different sessions' `init`s both having picked this one as their LRU + * candidate before either request lands): `evictResidentRuntimeForCap` + * re-checks liveness every time, and `closeRuntime` is idempotent, so a + * repeat call either finds the runtime already gone (a no-op) or finds it + * newly busy again and leaves it alone. + */ + async requestCapEviction(): Promise { + const self = this; + return Effect.runPromise( + Effect.gen(function* () { + yield* self.prepareErrorCaptureScope(); + yield* Effect.promise(() => self.evictResidentRuntimeForCap()); + }).pipe( + Effect.withSpan("McpSessionDO.requestCapEviction"), + // oxlint-disable-next-line executor/no-effect-escape-hatch -- boundary: DO RPC exposes Promise results + Effect.orDie, + ), ); } @@ -1299,7 +1755,12 @@ export abstract class McpAgentSessionDOBase< return; } - await this.disposeIdleRuntime({ idleMs, pausedExecutionCount, activeStreamCount }); + await this.disposeIdleRuntime({ + idleMs, + pausedExecutionCount, + activeStreamCount, + reason: "idle", + }); } private validateApprovalIdentity( diff --git a/packages/hosts/cloudflare/src/mcp/session-runtime-residency.ts b/packages/hosts/cloudflare/src/mcp/session-runtime-residency.ts index e1b61d166..06ff5c223 100644 --- a/packages/hosts/cloudflare/src/mcp/session-runtime-residency.ts +++ b/packages/hosts/cloudflare/src/mcp/session-runtime-residency.ts @@ -48,6 +48,189 @@ export const resetResidentRuntimeCountForTest = (): void => { peakResidentRuntimeCount = 0; }; +/** + * How many cold builds are currently in flight in this isolate — reserved at + * admission, the same moment `evictForCapIfNeeded` runs, and released exactly + * once the reserving session either becomes counted-as-resident + * (`acquireResidentRuntime`) or its build fails or is interrupted. + * + * `residentRuntimeCount` only moves once a build actually finishes, so N + * overlapping cold inits at the cap each read it, see themselves still under + * the cap, and none of them evicts anything — residency then blows past the + * cap by N once every build lands. This counter closes that gap: the cap + * check becomes `currentResidentRuntimeCount() + currentInFlightColdBuildCount()`, + * so a second (or third, ...) concurrent admission sees the first's + * reservation and triggers eviction instead of also passing the check for + * free. + */ +let inFlightColdBuildCount = 0; + +/** Reserve a cold-build slot at admission. Paired with exactly one of + * {@link releaseColdBuildSlot} per reservation — see the doc comment above. */ +export const reserveColdBuildSlot = (): number => { + inFlightColdBuildCount += 1; + return inFlightColdBuildCount; +}; + +/** Release a previously-reserved cold-build slot. Floors at zero so a stray + * extra release (there should never be one; callers gate this on their own + * per-init flag) cannot drive the counter negative. */ +export const releaseColdBuildSlot = (): number => { + inFlightColdBuildCount = Math.max(0, inFlightColdBuildCount - 1); + return inFlightColdBuildCount; +}; + +export const currentInFlightColdBuildCount = (): number => inFlightColdBuildCount; + +/** Test-only: isolate-scoped module state outlives a single test case. */ +export const resetInFlightColdBuildCountForTest = (): void => { + inFlightColdBuildCount = 0; +}; + +/** + * The isolate-wide ceiling on resident session runtimes, past which `init` + * evicts the least-recently-active evictable session to make room instead of + * letting residency climb unbounded. + * + * Production isolates have been observed OOM-killing every co-resident + * session somewhere between roughly 10 and 64 residents, and isolate memory + * is already sitting at the platform's 128MB limit at the fleet-wide median + * before this mechanism runs at all — there is no headroom left to discover + * the ceiling empirically per-isolate. 32 sits inside that failure band, so it + * bounds the COUNT-driven failure mode (an isolate cannot pile up unbounded + * residents no matter how many quiet-but-connected sessions land on it) while + * per-session footprint reduction, a separate effort, addresses the rest. + * It is a soft cap: an isolate that cannot find anything evictable is still + * allowed past it (see `pickEvictionCandidate`) rather than failing an init. + */ +export const RESIDENT_RUNTIME_SOFT_CAP = 32; + +/** + * One resident session runtime, as far as the isolate-wide eviction registry + * needs to know about it: when it was last active, whether it is safe to tear + * down right now, and how to ask it to tear itself down. + * + * `canEvict` is consulted by `pickEvictionCandidate` to choose AMONG entries + * and is expected to be cheap and synchronous — it does not have to be the + * final word. `dispose` sends this session an eviction REQUEST — it does not + * run this session's own teardown itself. The owning Durable Object instance + * is the only thing allowed to tear down its own runtime (see + * `requestSelfEviction` on the base class for why: teardown does I/O bound to + * this session's own request context, and a caller running it directly would + * be running that I/O under ITS OWN context instead). `dispose` is expected to + * re-check liveness itself before actually releasing anything, using whatever + * signals it has (including ones `canEvict` could not afford to read), so a + * candidate that became active between the pick and the dispose call is left + * alone rather than torn down. It resolving does not guarantee the candidate + * was actually torn down — only that the request was sent — so callers that + * fire it in the background (see `evictForCapIfNeeded`) call + * `markEvictionRequested` rather than assuming success. + */ +export type ResidentSessionEntry = { + readonly sessionId: string; + lastActivityMs: number; + evictionRequestedAt?: number; + readonly canEvict: () => boolean; + readonly dispose: (reason: "cap") => Promise; +}; + +/** + * Every session runtime currently resident in THIS isolate, keyed by session + * id. Dependency-free by design: this module is a leaf that both the session + * Durable Object and its tests import directly, and it must never need to + * know what an Effect, a Durable Object, or a storage API is. + * + * An entry lives here for exactly as long as its runtime is resident. It is + * added once, at the same moment `acquireResidentRuntime` counts it, and + * removed once, at the same moment `releaseResidentRuntime` releases it — the + * two are meant to move together, though this module does not enforce that; + * see `closeRuntime` on the Durable Object for where they are actually paired. + */ +const residentSessions = new Map(); + +/** Start tracking a newly-resident session runtime for isolate-wide eviction. */ +export const registerResidentSession = (entry: ResidentSessionEntry): void => { + residentSessions.set(entry.sessionId, entry); +}; + +/** Record that a resident session just did something, so it sorts last for eviction. */ +export const touchResidentSession = (sessionId: string, lastActivityMs = Date.now()): void => { + const entry = residentSessions.get(sessionId); + if (entry) entry.lastActivityMs = lastActivityMs; +}; + +/** + * Stop tracking a session's runtime because it is no longer resident. + * + * Idempotent on purpose: `Map.delete` on an absent key is already a no-op, so + * this can be called from every path that might end a session's residency — + * a clean disposal, a failed one, a repeat call — without any caller having to + * first ask whether the entry is still there. + */ +export const releaseResidentSession = (sessionId: string): void => { + residentSessions.delete(sessionId); +}; + +/** + * How long a resident session stays skipped by `pickEvictionCandidate` after + * an eviction request was sent to it. The request is fire-and-forget (see + * `evictForCapIfNeeded` on the base class) — the sender never learns whether + * it succeeded, failed, or is still in flight — so this grace window is the + * only thing standing between one stuck or repeatedly-failing candidate and + * it squatting the LRU pick forever, re-requested by every subsequent init. + * Long enough that a healthy eviction (self-request, own teardown, registry + * removal) always finishes well inside it; short enough that a genuinely + * stuck candidate stops blocking the LRU pick soon after. + */ +export const EVICTION_REQUEST_GRACE_MS = 15_000; + +/** + * Record that an eviction request was just sent to this entry, so + * `pickEvictionCandidate` skips it until the grace period elapses. A no-op + * when the entry is already gone (it was evicted, or removed some other way) + * — nothing left to mark, and nothing wrong with that. + */ +export const markEvictionRequested = (sessionId: string, requestedAtMs = Date.now()): void => { + const entry = residentSessions.get(sessionId); + if (entry) entry.evictionRequestedAt = requestedAtMs; +}; + +/** + * The least-recently-active resident session that is currently safe to evict, + * or `undefined` when every resident is streaming, paused, recently asked to + * evict itself already, or otherwise ineligible right now. + * + * `undefined` is a legitimate, expected answer — it means the isolate is over + * its soft cap but everything resident is doing real work (or already has an + * eviction request outstanding), and the caller must let the new session + * build anyway rather than block or fail on it. + */ +export const pickEvictionCandidate = (nowMs = Date.now()): ResidentSessionEntry | undefined => { + let candidate: ResidentSessionEntry | undefined; + for (const entry of residentSessions.values()) { + if (!entry.canEvict()) continue; + if ( + entry.evictionRequestedAt !== undefined && + nowMs - entry.evictionRequestedAt < EVICTION_REQUEST_GRACE_MS + ) { + continue; + } + if (!candidate || entry.lastActivityMs < candidate.lastActivityMs) { + candidate = entry; + } + } + return candidate; +}; + +/** Test-only: isolate-scoped module state outlives a single test case. */ +export const resetResidentSessionRegistryForTest = (): void => { + residentSessions.clear(); +}; + +/** Test-only: read without mutating, to assert on registry membership directly. */ +export const residentSessionIdsForTest = (): ReadonlyArray => + Array.from(residentSessions.keys()); + type MemoryCapablePerformance = { readonly memory?: { readonly usedJSHeapSize?: unknown; diff --git a/packages/hosts/cloudflare/src/mcp/session-stub.ts b/packages/hosts/cloudflare/src/mcp/session-stub.ts index 3a003ff0c..2e17f3a41 100644 --- a/packages/hosts/cloudflare/src/mcp/session-stub.ts +++ b/packages/hosts/cloudflare/src/mcp/session-stub.ts @@ -36,6 +36,14 @@ export interface McpSessionStub { response: ResumeResponse, incoming?: IncomingTraceHeaders, ) => Promise; + /** + * Ask this session's OWN Durable Object instance to tear down its resident + * runtime because the isolate is over its cap. Routed through the stub + * rather than called on an in-process reference so the teardown runs in + * this session's own IoContext — see `requestSelfEviction` on the base + * class for why that boundary matters. + */ + readonly requestCapEviction: () => Promise; } export const mcpSessionStub = (