* Cap resident MCP session runtimes per isolate with LRU eviction Isolates were letting resident session runtimes pile up unbounded past the point idle disposal alone could keep up, so memory pressure kept climbing until the platform reset the whole isolate. init now evicts the least-recently-active evictable session once the isolate hits a soft cap, reusing the same disposal path and eligibility signals the idle alarm already used. A cap that finds nothing evictable never blocks or fails init; it only marks the init span. * Evict resident MCP sessions through a real cross-DO request, not in-context Cap eviction now asks the candidate's own Durable Object stub to tear itself down instead of running its teardown inside the evicting session's request context, where its I/O objects aren't valid. Duplicate or failed requests are safe no-ops; a recently-requested candidate is skipped by the next pick. * Fix cap-admission race, mid-disposal race, and e2e session leak Reserve an in-flight cold-build slot at admission so concurrent cold inits at the cap evict instead of all passing the check before any of their builds finish incrementing residency. Track in-progress runtime disposal and flip initialized before the first async close, so a request landing mid-teardown awaits the same disposal and rebuilds instead of running against a half-closed runtime; a second concurrent disposal trigger waits on the same teardown instead of running it twice. Close every MCP session opened by the cap-eviction e2e scenario via Effect.ensuring. * Fix interrupt-path leaks in cold-build reservation, disposal, and e2e cleanup - release the cold-build reservation with one Effect.ensuring around all of init, so an interrupt between cap admission and the build can't leak it - make closeRuntime's teardown uninterruptible so disposingRuntime never resolves against half-released resources - record e2e session ids for cleanup as soon as they're minted, not after the handshake completes, so a failed notification can't orphan a session * Absorb rejected closes in teardown and record e2e sessions before the body read
Unified · Split
| @@ -118,6 +118,9 @@ declare global { · 117 unmodified lines | |||
| 118 | MCP_RESOURCE_ORIGIN?: string; | 118 | MCP_RESOURCE_ORIGIN?: string; |
| 119 | MCP_SESSION_TIMEOUT_MS?: string; | 119 | MCP_SESSION_TIMEOUT_MS?: string; |
| 120 | MCP_PAUSED_SESSION_IDLE_TIMEOUT_MS?: string; | 120 | MCP_PAUSED_SESSION_IDLE_TIMEOUT_MS?: string; |
| 121 | /** Test-only override for the isolate-wide resident-runtime soft cap | ||
| 122 | * (see `RESIDENT_RUNTIME_SOFT_CAP`). Unset in production. */ | ||
| 123 | MCP_RESIDENT_RUNTIME_SOFT_CAP?: string; | ||
| 121 | NODE_ENV?: string; | 124 | NODE_ENV?: string; |
| 122 | 125 | ||
| 123 | // Shared with frontend | 126 | // Shared with frontend |
| @@ -103,6 +103,10 @@ const positiveMilliseconds = (raw: string | undefined): number | undefined => { · 102 unmodified lines | |||
| 103 | return Math.floor(parsed); | 103 | return Math.floor(parsed); |
| 104 | }; | 104 | }; |
| 105 | 105 | ||
| 106 | /** Same shape as `positiveMilliseconds`, for a plain count rather than a | ||
| 107 | * duration — used only by the resident-runtime soft-cap override below. */ | ||
| 108 | const positiveInteger = positiveMilliseconds; | ||
| 109 | |||
| 106 | type CloudSessionDbHandle = DbServiceShape & { | 110 | type CloudSessionDbHandle = DbServiceShape & { |
| 107 | readonly sql: Sql; | 111 | readonly sql: Sql; |
| 108 | readonly end: () => Promise<void>; | 112 | readonly end: () => Promise<void>; |
| @@ -196,6 +200,25 @@ export class McpSessionDOSqlite extends McpAgentSessionDOBase<Env, CloudSessionD · 87 unmodified lines | |||
| 196 | return mcpExecutionOwnerDirectoryFromNamespace(env.MCP_EXECUTION_OWNER); | 200 | return mcpExecutionOwnerDirectoryFromNamespace(env.MCP_EXECUTION_OWNER); |
| 197 | } | 201 | } |
| 198 | 202 | ||
| 203 | // Test-only override so e2e can exercise a REAL cross-DO eviction request in | ||
| 204 | // workerd without registering 32 sessions. Unset in production, where this | ||
| 205 | // falls through to the base class's `RESIDENT_RUNTIME_SOFT_CAP`. | ||
| 206 | protected override residentRuntimeSoftCap(): number { | ||
| 207 | return positiveInteger(env.MCP_RESIDENT_RUNTIME_SOFT_CAP) ?? super.residentRuntimeSoftCap(); | ||
| 208 | } | ||
| 209 | |||
| 210 | protected override supportsCapEviction(): boolean { | ||
| 211 | return true; | ||
| 212 | } | ||
| 213 | |||
| 214 | protected override requestSelfEviction(): Promise<void> { | ||
| 215 | // Routed through this session's OWN stub (never a direct in-process call) | ||
| 216 | // so `requestCapEviction`'s teardown runs under an IoContext scoped to | ||
| 217 | // this request, not whatever request happened to trigger the eviction | ||
| 218 | // check — see the base class's `requestSelfEviction` doc comment. | ||
| 219 | return mcpSessionStub(env.MCP_SESSION, this.sessionId).requestCapEviction(); | ||
| 220 | } | ||
| 221 | |||
| 199 | protected override forwardModelResumeToOwner( | 222 | protected override forwardModelResumeToOwner( |
| 200 | owner: McpExecutionOwnerRoute, | 223 | owner: McpExecutionOwnerRoute, |
| 201 | identity: McpApprovalOwner, | 224 | identity: McpApprovalOwner, |
| @@ -80,6 +80,18 @@ export class McpSessionDO extends McpAgentSessionDOBase<CloudflareEnv, CfSession · 79 unmodified lines | |||
| 80 | return mcpExecutionOwnerDirectoryFromNamespace(this.cfEnv.MCP_EXECUTION_OWNER); | 80 | return mcpExecutionOwnerDirectoryFromNamespace(this.cfEnv.MCP_EXECUTION_OWNER); |
| 81 | } | 81 | } |
| 82 | 82 | ||
| 83 | protected override supportsCapEviction(): boolean { | ||
| 84 | return true; | ||
| 85 | } | ||
| 86 | |||
| 87 | protected override requestSelfEviction(): Promise<void> { | ||
| 88 | // Routed through this session's OWN stub (never a direct in-process call) | ||
| 89 | // so `requestCapEviction`'s teardown runs under an IoContext scoped to | ||
| 90 | // this request, not whatever request happened to trigger the eviction | ||
| 91 | // check — see the base class's `requestSelfEviction` doc comment. | ||
| 92 | return mcpSessionStub(this.cfEnv.MCP_SESSION, this.sessionId).requestCapEviction(); | ||
| 93 | } | ||
| 94 | |||
| 83 | protected override forwardModelResumeToOwner( | 95 | protected override forwardModelResumeToOwner( |
| 84 | owner: McpExecutionOwnerRoute, | 96 | owner: McpExecutionOwnerRoute, |
| 85 | identity: McpApprovalOwner, | 97 | identity: McpApprovalOwner, |
| @@ -0,0 +1,256 @@ | |||
| 1 | // Cloud: crossing the isolate's resident-runtime soft cap evicts an idle | ||
| 2 | // session's runtime through a REAL cross-Durable-Object request, not a | ||
| 3 | // same-context call. | ||
| 4 | // | ||
| 5 | // The defect this pins: the original design ran the evicted (candidate) | ||
| 6 | // session's teardown — closing its postgres.js socket, storage writes, span | ||
| 7 | // flush — directly inside the EVICTING session's own request/IoContext. In | ||
| 8 | // production workerd, I/O objects are bound to the IoContext that created | ||
| 9 | // them, so a cross-context call like that throws "Cannot perform I/O on | ||
| 10 | // behalf of a different request" or silently soft-fails. That failure mode | ||
| 11 | // cannot reproduce against an in-process unit-test double (same JS object, | ||
| 12 | // same context either way) — it only shows up against a real Durable Object | ||
| 13 | // stub. The fix routes eviction through the candidate's OWN stub | ||
| 14 | // (`requestCapEviction`, an RPC method mirroring `forwardModelResumeToOwner`), | ||
| 15 | // so the candidate's teardown runs in the candidate's own context, and this | ||
| 16 | // scenario is what actually exercises that stub in workerd. | ||
| 17 | // | ||
| 18 | // e2e/setup/resident-runtime-cap.ts lowers MCP_RESIDENT_RUNTIME_SOFT_CAP for | ||
| 19 | // the whole boot (see that file for the value and its headroom story), so | ||
| 20 | // this test can cross it with a bounded number of real sessions instead of | ||
| 21 | // registering the production default of 32. | ||
| 22 | import { expect } from "@effect/vitest"; | ||
| 23 | import { Effect, Schedule } from "effect"; | ||
| 24 | |||
| 25 | import { scenario } from "../src/scenario"; | ||
| 26 | import { Mcp, Target, Telemetry } from "../src/services"; | ||
| 27 | import type { Identity } from "../src/target"; | ||
| 28 | import { E2E_MCP_RESIDENT_RUNTIME_SOFT_CAP } from "../setup/resident-runtime-cap"; | ||
| 29 | |||
| 30 | const PROTOCOL_VERSION = "2025-03-26"; | ||
| 31 | const JSON_AND_SSE = "application/json, text/event-stream"; | ||
| 32 | |||
| 33 | // Comfortably past the cap: even if a handful of other scenarios' sessions | ||
| 34 | // are still incidentally resident when this file runs, enough of THESE | ||
| 35 | // sessions cross it that at least one eviction targets a session opened here. | ||
| 36 | const SESSIONS_TO_OPEN = E2E_MCP_RESIDENT_RUNTIME_SOFT_CAP + 10; | ||
| 37 | |||
| 38 | const emailOf = (identity: Identity): string => identity.credentials?.email ?? identity.label; | ||
| 39 | |||
| 40 | const mcpHeaders = (bearer: string, sessionId?: string) => ({ | ||
| 41 | accept: JSON_AND_SSE, | ||
| 42 | authorization: `Bearer ${bearer}`, | ||
| 43 | "content-type": "application/json", | ||
| 44 | "mcp-protocol-version": PROTOCOL_VERSION, | ||
| 45 | ...(sessionId ? { "mcp-session-id": sessionId } : {}), | ||
| 46 | }); | ||
| 47 | |||
| 48 | const postJson = (mcpUrl: string, bearer: string, body: unknown, sessionId?: string) => | ||
| 49 | fetch(mcpUrl, { | ||
| 50 | method: "POST", | ||
| 51 | headers: mcpHeaders(bearer, sessionId), | ||
| 52 | body: JSON.stringify(body), | ||
| 53 | }); | ||
| 54 | |||
| 55 | /** | ||
| 56 | * Opens one fresh MCP session under an already-minted bearer. `initialize` | ||
| 57 | * without an existing `mcp-session-id` always mints a new session, the same | ||
| 58 | * way separate browser tabs sharing one login would — so many of these under | ||
| 59 | * one identity is a cheap way to grow the isolate's resident-runtime count | ||
| 60 | * without a full OAuth round trip per session. | ||
| 61 | * | ||
| 62 | * `recordSession` is called the moment the session id is known — before the | ||
| 63 | * `notifications/initialized` round trip below, not after this function | ||
| 64 | * returns. A session is live on the server as soon as `initialize` responds | ||
| 65 | * with an `mcp-session-id`, regardless of whether the handshake ever | ||
| 66 | * completes; recording it only on a full return left a failed notification | ||
| 67 | * (or an interrupt landing between the two requests) with no cleanup entry, | ||
| 68 | * orphaning a real session on the target isolate. | ||
| 69 | */ | ||
| 70 | const openSession = async ( | ||
| 71 | mcpUrl: string, | ||
| 72 | bearer: string, | ||
| 73 | label: string, | ||
| 74 | recordSession: (sessionId: string) => void, | ||
| 75 | ): Promise<string> => { | ||
| 76 | const initialized = await postJson(mcpUrl, bearer, { | ||
| 77 | jsonrpc: "2.0" as const, | ||
| 78 | id: "initialize", | ||
| 79 | method: "initialize", | ||
| 80 | params: { | ||
| 81 | protocolVersion: PROTOCOL_VERSION, | ||
| 82 | capabilities: {}, | ||
| 83 | clientInfo: { name: `executor-e2e-cap-eviction-${label}`, version: "0.0.1" }, | ||
| 84 | }, | ||
| 85 | }); | ||
| 86 | const sessionId = initialized.headers.get("mcp-session-id"); | ||
| 87 | if (!sessionId) { | ||
| 88 | // oxlint-disable-next-line executor/no-error-constructor -- boundary: e2e setup precondition. | ||
| 89 | throw new Error(`openSession (${label}): no mcp-session-id header`); | ||
| 90 | } | ||
| 91 | // Recorded the moment the id exists — BEFORE the body read and status | ||
| 92 | // assertion below, either of which can throw with the session already live | ||
| 93 | // on the server. The cleanup finalizer needs the id on every one of those | ||
| 94 | // paths, not just a fully successful return. | ||
| 95 | recordSession(sessionId); | ||
| 96 | await initialized.text(); | ||
| 97 | expect(initialized.status, `initialize (${label}) opens a session`).toBe(200); | ||
| 98 | const notification = await postJson( | ||
| 99 | mcpUrl, | ||
| 100 | bearer, | ||
| 101 | { jsonrpc: "2.0" as const, method: "notifications/initialized" }, | ||
| 102 | sessionId, | ||
| 103 | ); | ||
| 104 | await notification.text(); | ||
| 105 | expect(notification.status, `(${label}) completes the handshake`).toBe(202); | ||
| 106 | return sessionId; | ||
| 107 | }; | ||
| 108 | |||
| 109 | const executeBody = (id: string, code: string) => ({ | ||
| 110 | jsonrpc: "2.0" as const, | ||
| 111 | id, | ||
| 112 | method: "tools/call", | ||
| 113 | params: { name: "execute", arguments: { code } }, | ||
| 114 | }); | ||
| 115 | |||
| 116 | /** Run `execute` and return the response text once the call has fully settled. */ | ||
| 117 | const execute = async ( | ||
| 118 | mcpUrl: string, | ||
| 119 | bearer: string, | ||
| 120 | sessionId: string, | ||
| 121 | id: string, | ||
| 122 | code: string, | ||
| 123 | ): Promise<string> => { | ||
| 124 | const response = await postJson(mcpUrl, bearer, executeBody(id, code), sessionId); | ||
| 125 | const body = await response.text(); | ||
| 126 | expect(response.status, `execute ${id} is served`).toBe(200); | ||
| 127 | return body; | ||
| 128 | }; | ||
| 129 | |||
| 130 | scenario( | ||
| 131 | "MCP session · crossing the resident-runtime cap evicts a session through a real cross-DO request", | ||
| 132 | { timeout: 180_000 }, | ||
| 133 | Effect.gen(function* () { | ||
| 134 | const target = yield* Target; | ||
| 135 | const mcp = yield* Mcp; | ||
| 136 | const telemetry = yield* Telemetry; | ||
| 137 | |||
| 138 | const identity = yield* target.newIdentity(); | ||
| 139 | const bearer = yield* mcp.mintBearer(emailOf(identity)); | ||
| 140 | |||
| 141 | // Opened sessions are recorded here as each one succeeds, so the cleanup | ||
| 142 | // below can close exactly what was actually opened even if the scenario | ||
| 143 | // fails partway through. Cap eviction already tears most of these down as | ||
| 144 | // a side effect of the scenario itself, but termination is idempotent | ||
| 145 | // (see mcp-destroyed-session-envelope.test.ts) — closing an already-torn- | ||
| 146 | // -down session is a harmless no-op, not a double-free. | ||
| 147 | const openedSessionIds: string[] = []; | ||
| 148 | |||
| 149 | const scenarioBody = Effect.gen(function* () { | ||
| 150 | // Open more sessions than the cap allows, at limited concurrency. None | ||
| 151 | // of them run any work, so every one is immediately eviction-eligible — | ||
| 152 | // crossing the cap must pick at least one and tear it down through its | ||
| 153 | // own stub. | ||
| 154 | const sessionIds = yield* Effect.forEach( | ||
| 155 | Array.from({ length: SESSIONS_TO_OPEN }, (_, index) => index), | ||
| 156 | (index) => | ||
| 157 | Effect.promise(() => | ||
| 158 | openSession(target.mcpUrl, bearer, `session-${index}`, (sessionId) => { | ||
| 159 | openedSessionIds.push(sessionId); | ||
| 160 | }), | ||
| 161 | ), | ||
| 162 | { concurrency: 8 }, | ||
| 163 | ); | ||
| 164 | |||
| 165 | expect(sessionIds.length, "every session opened").toBe(SESSIONS_TO_OPEN); | ||
| 166 | expect(new Set(sessionIds).size, "every session got a distinct id").toBe(SESSIONS_TO_OPEN); | ||
| 167 | |||
| 168 | // ---- a real cap eviction fired, against a session opened here ------- | ||
| 169 | // Same span the idle path emits (`mcp.session.idle_runtime_dispose`); | ||
| 170 | // `mcp.session.dispose_reason` is what disambiguates the trigger. | ||
| 171 | const capDisposals = yield* telemetry | ||
| 172 | .searchSpans({ operation: "mcp.session.idle_runtime_dispose" }) | ||
| 173 | .pipe( | ||
| 174 | Effect.map((spans) => | ||
| 175 | spans.filter( | ||
| 176 | (span) => | ||
| 177 | span.span.tags["mcp.session.dispose_reason"] === "cap" && | ||
| 178 | sessionIds.some((id) => (span.span.tags["mcp.session.id"] ?? "").includes(id)), | ||
| 179 | ), | ||
| 180 | ), | ||
| 181 | Effect.filterOrFail( | ||
| 182 | (spans) => spans.length > 0, | ||
| 183 | () => "no cap-triggered idle_runtime_dispose span exported for any session opened here", | ||
| 184 | ), | ||
| 185 | // The eviction request is fire-and-forget (`ctx.waitUntil`) from the | ||
| 186 | // evictor's `init`, and its own span flush is off that same | ||
| 187 | // background path — same polling grace the idle-disposal scenario | ||
| 188 | // uses for its alarm-driven flush. | ||
| 189 | Effect.retry(Schedule.both(Schedule.spaced("500 millis"), Schedule.recurs(40))), | ||
| 190 | ); | ||
| 191 | |||
| 192 | expect( | ||
| 193 | capDisposals.length, | ||
| 194 | "crossing the resident-runtime cap evicted at least one session opened here", | ||
| 195 | ).toBeGreaterThan(0); | ||
| 196 | |||
| 197 | const disposal = capDisposals[0]!; | ||
| 198 | expect( | ||
| 199 | disposal.span.tags["mcp.isolate.resident_runtimes"], | ||
| 200 | "the cap disposal records the isolate's resident-runtime gauge, same as the idle path", | ||
| 201 | ).toBeDefined(); | ||
| 202 | |||
| 203 | // ---- the evicted session still works — restore is transparent ------- | ||
| 204 | const evictedSessionId = sessionIds.find((id) => | ||
| 205 | (disposal.span.tags["mcp.session.id"] ?? "").includes(id), | ||
| 206 | ); | ||
| 207 | expect( | ||
| 208 | evictedSessionId, | ||
| 209 | "the disposed span's session id matches a session opened here", | ||
| 210 | ).toBeDefined(); | ||
| 211 | |||
| 212 | const marker = `after-cap-evict-${evictedSessionId}`; | ||
| 213 | const restored = yield* Effect.promise(() => | ||
| 214 | execute( | ||
| 215 | target.mcpUrl, | ||
| 216 | bearer, | ||
| 217 | evictedSessionId!, | ||
| 218 | "execute-after-cap-eviction", | ||
| 219 | `return ${JSON.stringify(marker)};`, | ||
| 220 | ), | ||
| 221 | ); | ||
| 222 | expect( | ||
| 223 | restored, | ||
| 224 | "the evicted session serves the next call correctly after restoring underneath the client", | ||
| 225 | ).toContain(marker); | ||
| 226 | }); | ||
| 227 | |||
| 228 | yield* scenarioBody.pipe( | ||
| 229 | // `Effect.ensuring`, not a trailing statement: a failure partway through | ||
| 230 | // (an assertion above, a timed-out span search) must not leak the | ||
| 231 | // sessions already opened. Read `openedSessionIds` at cleanup time, not | ||
| 232 | // capture time — `Effect.suspend` so the array is read when the | ||
| 233 | // finalizer actually runs, after the scenario body has finished pushing | ||
| 234 | // to it, rather than snapshotted empty at construction. | ||
| 235 | Effect.ensuring( | ||
| 236 | Effect.suspend(() => | ||
| 237 | Effect.forEach( | ||
| 238 | openedSessionIds, | ||
| 239 | (sessionId) => | ||
| 240 | Effect.tryPromise(async () => { | ||
| 241 | const closed = await fetch(target.mcpUrl, { | ||
| 242 | method: "DELETE", | ||
| 243 | headers: { | ||
| 244 | authorization: `Bearer ${bearer}`, | ||
| 245 | "mcp-session-id": sessionId, | ||
| 246 | }, | ||
| 247 | }); | ||
| 248 | await closed.text(); | ||
| 249 | }).pipe(Effect.ignore), | ||
| 250 | { concurrency: 8, discard: true }, | ||
| 251 | ), | ||
| 252 | ), | ||
| 253 | ), | ||
| 254 | ); | ||
| 255 | }), | ||
| 256 | ); | ||
| @@ -16,6 +16,7 @@ import { · 15 unmodified lines | |||
| 16 | E2E_EXECUTION_RATE_LIMIT, | 16 | E2E_EXECUTION_RATE_LIMIT, |
| 17 | E2E_EXECUTION_RATE_LIMIT_CHECK_TIMEOUT_MS, | 17 | E2E_EXECUTION_RATE_LIMIT_CHECK_TIMEOUT_MS, |
| 18 | } from "./execution-limits"; | 18 | } from "./execution-limits"; |
| 19 | import { E2E_MCP_RESIDENT_RUNTIME_SOFT_CAP } from "./resident-runtime-cap"; | ||
| 19 | 20 | ||
| 20 | export const cloudDir = fileURLToPath(new URL("../../apps/cloud/", import.meta.url)); | 21 | export const cloudDir = fileURLToPath(new URL("../../apps/cloud/", import.meta.url)); |
| 21 | 22 | ||
| @@ -107,6 +108,10 @@ export const bootCloud = async (options: CloudBootOptions): Promise<CloudBooted> · 85 unmodified lines | |||
| 107 | MCP_RESOURCE_ORIGIN: options.publicUrl, | 108 | MCP_RESOURCE_ORIGIN: options.publicUrl, |
| 108 | MCP_SESSION_TIMEOUT_MS: process.env.MCP_SESSION_TIMEOUT_MS, | 109 | MCP_SESSION_TIMEOUT_MS: process.env.MCP_SESSION_TIMEOUT_MS, |
| 109 | MCP_PAUSED_SESSION_IDLE_TIMEOUT_MS: process.env.MCP_PAUSED_SESSION_IDLE_TIMEOUT_MS, | 110 | MCP_PAUSED_SESSION_IDLE_TIMEOUT_MS: process.env.MCP_PAUSED_SESSION_IDLE_TIMEOUT_MS, |
| 111 | // See resident-runtime-cap.ts for why this value, and why it is safe to | ||
| 112 | // set unconditionally for the whole boot (same treatment as the execution | ||
| 113 | // rate limit below). | ||
| 114 | MCP_RESIDENT_RUNTIME_SOFT_CAP: String(E2E_MCP_RESIDENT_RUNTIME_SOFT_CAP), | ||
| 110 | ALLOW_LOCAL_NETWORK: "true", | 115 | ALLOW_LOCAL_NETWORK: "true", |
| 111 | // A first-party GitHub app for the first-party-oauth scenario: proves the | 116 | // A first-party GitHub app for the first-party-oauth scenario: proves the |
| 112 | // env → HostConfig → executor plumbing end to end. The scenario asserts the | 117 | // env → HostConfig → executor plumbing end to end. The scenario asserts the |
| @@ -0,0 +1,21 @@ | |||
| 1 | // The e2e worker's isolate-wide resident-runtime soft cap | ||
| 2 | // (MCP_RESIDENT_RUNTIME_SOFT_CAP). One constant with two consumers, the boot | ||
| 3 | // recipe env (cloud.boot.ts) and the cap-eviction scenario | ||
| 4 | // (cloud/mcp-session-cap-eviction.test.ts), so they cannot drift apart. | ||
| 5 | // | ||
| 6 | // Same squeeze as EXECUTION_RATE_LIMIT_PER_HOUR (execution-limits.ts). It must | ||
| 7 | // be LOW enough that the cap-eviction scenario can cross it with a bounded | ||
| 8 | // number of real sessions opened under one identity (prod's 32 is reachable | ||
| 9 | // but wasteful to open on every run), and HIGH enough that no OTHER cloud | ||
| 10 | // scenario's incidental concurrent MCP-session count trips it: this is an | ||
| 11 | // isolate-wide counter, not per-org, and the dev server is shared across the | ||
| 12 | // whole cloud e2e run (`fileParallelism: false` keeps files serial, but a | ||
| 13 | // single busy file can still open a double-digit number of sessions before | ||
| 14 | // any of them idle out). The busiest current file (mcp-protocol.test.ts) opens | ||
| 15 | // on the order of a dozen sessions across its scenarios; this stays well above | ||
| 16 | // that with headroom. If a scenario ever incidentally trips a cap eviction (a | ||
| 17 | // span with `mcp.session.dispose_reason: "cap"` for a session it didn't | ||
| 18 | // expect), that is not a correctness bug — eviction is designed to be | ||
| 19 | // transparent, restoring on the next call — but it means this constant needs | ||
| 20 | // to grow: raise it here, never by special-casing a scenario against it. | ||
| 21 | export const E2E_MCP_RESIDENT_RUNTIME_SOFT_CAP = 24; | ||
| @@ -1,6 +1,7 @@ | |||
| 1 | // 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 | 1 | // 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 |
| 2 | import { afterEach, describe, expect, it } from "@effect/vitest"; | 2 | import { afterEach, beforeEach, describe, expect, it } from "@effect/vitest"; |
| 3 | import { Cause, Effect, Exit, Schema } from "effect"; | 3 | import { Cause, Effect, Exit, Schema } from "effect"; |
| 4 | import type * as Tracer from "effect/Tracer"; | ||
| 4 | import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; | 5 | import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; |
| 5 | import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js"; | 6 | import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js"; |
| 6 | import type { JSONRPCMessage, MessageExtraInfo } from "@modelcontextprotocol/sdk/types.js"; | 7 | import type { JSONRPCMessage, MessageExtraInfo } from "@modelcontextprotocol/sdk/types.js"; |
| @@ -14,6 +15,20 @@ import { · 7 unmodified lines | |||
| 14 | type McpSessionModelResumeResult, | 15 | type McpSessionModelResumeResult, |
| 15 | type SessionMeta, | 16 | type SessionMeta, |
| 16 | } from "./agent-session-durable-object"; | 17 | } from "./agent-session-durable-object"; |
| 18 | import { | ||
| 19 | currentInFlightColdBuildCount, | ||
| 20 | currentResidentRuntimeCount, | ||
| 21 | EVICTION_REQUEST_GRACE_MS, | ||
| 22 | markEvictionRequested, | ||
| 23 | pickEvictionCandidate, | ||
| 24 | registerResidentSession, | ||
| 25 | releaseResidentSession, | ||
| 26 | resetInFlightColdBuildCountForTest, | ||
| 27 | resetResidentRuntimeCountForTest, | ||
| 28 | resetResidentSessionRegistryForTest, | ||
| 29 | residentSessionIdsForTest, | ||
| 30 | touchResidentSession, | ||
| 31 | } from "./session-runtime-residency"; | ||
| 17 | 32 | ||
| 18 | class MemoryStorage { | 33 | class MemoryStorage { |
| 19 | private readonly data = new Map<string, unknown>(); | 34 | private readonly data = new Map<string, unknown>(); |
| @@ -85,11 +100,40 @@ class MemoryStorage { · 65 unmodified lines | |||
| 85 | return this; | 100 | return this; |
| 86 | } | 101 | } |
| 87 | 102 | ||
| 103 | /** Give this storage's Durable Object a distinct name, so multiple harness | ||
| 104 | * sessions in the same test resolve to distinct `sessionIdForTelemetry()` | ||
| 105 | * values instead of all colliding on the default. */ | ||
| 106 | withIdName(name: string): this { | ||
| 107 | this.idName = name; | ||
| 108 | return this; | ||
| 109 | } | ||
| 110 | |||
| 88 | get storage(): MemoryStorage { | 111 | get storage(): MemoryStorage { |
| 89 | return this; | 112 | return this; |
| 90 | } | 113 | } |
| 91 | 114 | ||
| 92 | waitUntil(_promise: Promise<unknown>): void {} | 115 | private readonly waitUntilPromises: Promise<unknown>[] = []; |
| 116 | |||
| 117 | /** | ||
| 118 | * `ctx.waitUntil` extends work past the response instead of blocking it — | ||
| 119 | * cap eviction requests it. A real DO holds the runtime open until every | ||
| 120 | * queued promise settles; this fake instead just remembers them so a test | ||
| 121 | * can explicitly wait for the background work it cares about via | ||
| 122 | * `drainWaitUntil`, rather than the two racing unpredictably. | ||
| 123 | */ | ||
| 124 | waitUntil(promise: Promise<unknown>): void { | ||
| 125 | this.waitUntilPromises.push(promise); | ||
| 126 | } | ||
| 127 | |||
| 128 | /** Test-only: await every `waitUntil`-queued promise queued so far, | ||
| 129 | * including ones queued BY those promises while draining (an eviction | ||
| 130 | * request settling can itself queue more background work). */ | ||
| 131 | async drainWaitUntil(): Promise<void> { | ||
| 132 | while (this.waitUntilPromises.length > 0) { | ||
| 133 | const pending = this.waitUntilPromises.splice(0, this.waitUntilPromises.length); | ||
| 134 | await Promise.allSettled(pending); | ||
| 135 | } | ||
| 136 | } | ||
| 93 | } | 137 | } |
| 94 | 138 | ||
| 95 | type HarnessSession = { | 139 | type HarnessSession = { |
| @@ -920,3 +964,720 @@ describe("McpAgentSessionDOBase alarm name resolution", () => { · 824 unmodified lines | |||
| 920 | ); | 964 | ); |
| 921 | }); | 965 | }); |
| 922 | }); | 966 | }); |
| 967 | |||
| 968 | // The isolate-wide residency gauge used to be purely observational. These pin | ||
| 969 | // the enforcing half: past the soft cap, `init` evicts the least-recently | ||
| 970 | // -active EVICTABLE session to make room, at most one per init, and never | ||
| 971 | // blocks or fails an init over it even when nothing can be evicted. | ||
| 972 | // | ||
| 973 | // `residentRuntimeCount`/the resident-session registry are isolate (module) | ||
| 974 | // scope, exactly like production — so every test here resets them, both | ||
| 975 | // before (in case an earlier describe block in this file leaked residents by | ||
| 976 | // calling `init()` without ever disposing) and after (so it does not leak | ||
| 977 | // into whatever runs next). | ||
| 978 | describe("McpAgentSessionDOBase residency cap eviction", () => { | ||
| 979 | type ResidencySession = { | ||
| 980 | ctx: MemoryStorage; | ||
| 981 | captureCause: (cause: Cause.Cause<unknown>) => void; | ||
| 982 | dbHandle: { readonly end: () => void } | null; | ||
| 983 | engine: ExecutionEngine<Cause.YieldableError> | null; | ||
| 984 | getConnections: () => Iterable<unknown>; | ||
| 985 | getSessionId: () => string; | ||
| 986 | init: () => Promise<void>; | ||
| 987 | // Re-exposed for the harness only, same idiom as `evictResidentRuntimeForCap` | ||
| 988 | // below: the `AbortController` backing the CURRENTLY in-flight `init` | ||
| 989 | // call's root fiber, so a test can interrupt it deterministically at a | ||
| 990 | // chosen suspension point instead of only exercising the failure/defect | ||
| 991 | // path. | ||
| 992 | initAbortController: AbortController | null; | ||
| 993 | // Same idiom, for the disposal path — see `disposeAbortController`'s doc | ||
| 994 | // comment on the class. | ||
| 995 | disposeAbortController: AbortController | null; | ||
| 996 | initialized: boolean; | ||
| 997 | lastActivityMs: number; | ||
| 998 | pendingApprovalLeases: Map<string, never>; | ||
| 999 | props: Record<string, unknown>; | ||
| 1000 | residentRuntimeSoftCap: () => number; | ||
| 1001 | server?: McpServer; | ||
| 1002 | sessionIdForTelemetry: () => string; | ||
| 1003 | sessionTimeoutMs: () => number; | ||
| 1004 | withTelemetry: <A, E>(effect: Effect.Effect<A, E>) => Effect.Effect<A, E>; | ||
| 1005 | buildMcpServer: () => Effect.Effect<{ | ||
| 1006 | mcpServer: McpServer; | ||
| 1007 | engine: ExecutionEngine<Cause.YieldableError>; | ||
| 1008 | }>; | ||
| 1009 | openSessionDb: () => { readonly end: () => void } | Promise<{ readonly end: () => void }>; | ||
| 1010 | resolveSessionMeta: () => Effect.Effect<SessionMeta>; | ||
| 1011 | supportsCapEviction: () => boolean; | ||
| 1012 | requestSelfEviction: () => Promise<void>; | ||
| 1013 | // Re-exposed for the harness only (private on the real class, same idiom | ||
| 1014 | // as `resolveAndStoreSessionMeta` above): the candidate side of eviction, | ||
| 1015 | // wired below to stand in for a landed self-addressed stub call. | ||
| 1016 | evictResidentRuntimeForCap: () => Promise<void>; | ||
| 1017 | }; | ||
| 1018 | |||
| 1019 | const residencySessionMeta = (organizationId: string): SessionMeta => ({ | ||
| 1020 | organizationId, | ||
| 1021 | organizationName: "Org 1", | ||
| 1022 | userId: "user-1", | ||
| 1023 | resource: defaultMcpResource, | ||
| 1024 | }); | ||
| 1025 | |||
| 1026 | /** | ||
| 1027 | * A session harness built to actually run `init()`, the entry point that | ||
| 1028 | * builds (and, on cold restore, rebuilds) a resident runtime. | ||
| 1029 | * | ||
| 1030 | * `requestSelfEviction` stands in for what, in production, is a real | ||
| 1031 | * self-addressed Durable Object stub call — `mcpSessionStub(...).requestCapEviction()` | ||
| 1032 | * — landing back on this same instance and invoking `requestCapEviction`, | ||
| 1033 | * which runs `evictResidentRuntimeForCap` in the (correctly-scoped) IoContext | ||
| 1034 | * that call created. This harness has no Durable Object stub machinery to | ||
| 1035 | * exercise that routing with, so it wires straight to this session's OWN | ||
| 1036 | * `evictResidentRuntimeForCap` instead: it is the candidate-side handler | ||
| 1037 | * that matters for these tests (does the re-check + teardown behave | ||
| 1038 | * correctly), not the RPC transport that gets a request there — that | ||
| 1039 | * transport is covered by the e2e idle-disposal run against real workerd, | ||
| 1040 | * not here. `supportsCapEviction` defaults to `true` so a harness session is | ||
| 1041 | * an eviction candidate unless a test deliberately opts out. | ||
| 1042 | */ | ||
| 1043 | const makeResidencySession = (input: { | ||
| 1044 | readonly id: string; | ||
| 1045 | readonly cap?: number; | ||
| 1046 | readonly hasActiveStream?: boolean; | ||
| 1047 | readonly pausedExecutionCount?: number; | ||
| 1048 | readonly supportsCapEviction?: boolean; | ||
| 1049 | }): { readonly session: ResidencySession; readonly storage: MemoryStorage } => { | ||
| 1050 | const storage = new MemoryStorage().withIdName(`streamable-http:${input.id}`); | ||
| 1051 | const session = Object.create(McpAgentSessionDOBase.prototype) as ResidencySession; | ||
| 1052 | session.ctx = storage; | ||
| 1053 | session.captureCause = () => undefined; | ||
| 1054 | session.dbHandle = null; | ||
| 1055 | session.engine = null; | ||
| 1056 | session.getConnections = () => (input.hasActiveStream ? [{ close: () => undefined }] : []); | ||
| 1057 | session.getSessionId = () => input.id; | ||
| 1058 | session.initialized = false; | ||
| 1059 | session.lastActivityMs = 0; | ||
| 1060 | session.pendingApprovalLeases = new Map<string, never>(); | ||
| 1061 | session.props = { session: { organizationId: input.id, userId: "user-1" } }; | ||
| 1062 | session.sessionTimeoutMs = () => 60_000; | ||
| 1063 | session.residentRuntimeSoftCap = () => input.cap ?? 32; | ||
| 1064 | session.resolveSessionMeta = () => Effect.succeed(residencySessionMeta(input.id)); | ||
| 1065 | session.openSessionDb = () => ({ end: () => undefined }); | ||
| 1066 | session.buildMcpServer = () => | ||
| 1067 | Effect.succeed({ | ||
| 1068 | mcpServer: makeServer(), | ||
| 1069 | engine: { | ||
| 1070 | ...makeEngine().engine, | ||
| 1071 | pausedExecutionCount: () => Effect.succeed(input.pausedExecutionCount ?? 0), | ||
| 1072 | }, | ||
| 1073 | }); | ||
| 1074 | session.supportsCapEviction = () => input.supportsCapEviction ?? true; | ||
| 1075 | session.requestSelfEviction = () => session.evictResidentRuntimeForCap(); | ||
| 1076 | return { session, storage }; | ||
| 1077 | }; | ||
| 1078 | |||
| 1079 | /** A tracer that records every span + its attributes, to assert on the | ||
| 1080 | * `mcp.isolate.cap_overflow` attribute the `McpSessionDO.init` span | ||
| 1081 | * carries when nothing was evictable. */ | ||
| 1082 | const makeRecordingTracer = (): { | ||
| 1083 | readonly tracer: Tracer.Tracer; | ||
| 1084 | readonly spans: ReadonlyArray<{ | ||
| 1085 | readonly name: string; | ||
| 1086 | readonly attributes: ReadonlyMap<string, unknown>; | ||
| 1087 | }>; | ||
| 1088 | } => { | ||
| 1089 | const spans: Array<{ name: string; attributes: Map<string, unknown> }> = []; | ||
| 1090 | const tracer: Tracer.Tracer = { | ||
| 1091 | span: (options) => { | ||
| 1092 | const attributes = new Map<string, unknown>(); | ||
| 1093 | spans.push({ name: options.name, attributes }); | ||
| 1094 | let status: Tracer.SpanStatus = { _tag: "Started", startTime: options.startTime }; | ||
| 1095 | return { | ||
| 1096 | _tag: "Span", | ||
| 1097 | name: options.name, | ||
| 1098 | spanId: `span-${spans.length}`, | ||
| 1099 | traceId: "trace-1", | ||
| 1100 | parent: options.parent, | ||
| 1101 | annotations: options.annotations, | ||
| 1102 | get status() { | ||
| 1103 | return status; | ||
| 1104 | }, | ||
| 1105 | attributes, | ||
| 1106 | links: options.links, | ||
| 1107 | sampled: options.sampled, | ||
| 1108 | kind: options.kind, | ||
| 1109 | end: (endTime, exit) => { | ||
| 1110 | status = { _tag: "Ended", startTime: options.startTime, endTime, exit }; | ||
| 1111 | }, | ||
| 1112 | attribute: (key, value) => { | ||
| 1113 | attributes.set(key, value); | ||
| 1114 | }, | ||
| 1115 | event: () => undefined, | ||
| 1116 | addLinks: () => undefined, | ||
| 1117 | }; | ||
| 1118 | }, | ||
| 1119 | }; | ||
| 1120 | return { tracer, spans }; | ||
| 1121 | }; | ||
| 1122 | |||
| 1123 | beforeEach(() => { | ||
| 1124 | // Earlier describe blocks in this file build sessions via `init()` too, | ||
| 1125 | // without ever disposing them (that is out of scope for what they test) — | ||
| 1126 | // so without this, residency here would start from whatever they left | ||
| 1127 | // behind rather than zero. | ||
| 1128 | resetResidentRuntimeCountForTest(); | ||
| 1129 | resetResidentSessionRegistryForTest(); | ||
| 1130 | resetInFlightColdBuildCountForTest(); | ||
| 1131 | }); | ||
| 1132 | |||
| 1133 | afterEach(() => { | ||
| 1134 | resetResidentRuntimeCountForTest(); | ||
| 1135 | resetResidentSessionRegistryForTest(); | ||
| 1136 | resetInFlightColdBuildCountForTest(); | ||
| 1137 | }); | ||
| 1138 | |||
| 1139 | it("evicts the least-recently-active evictable session once the cap is reached, and the newer session survives", async () => { | ||
| 1140 | const { session: sessionA } = makeResidencySession({ id: "session-cap-a", cap: 2 }); | ||
| 1141 | const { session: sessionB } = makeResidencySession({ id: "session-cap-b", cap: 2 }); | ||
| 1142 | const { session: sessionC, storage: storageC } = makeResidencySession({ | ||
| 1143 | id: "session-cap-c", | ||
| 1144 | cap: 2, | ||
| 1145 | }); | ||
| 1146 | |||
| 1147 | await sessionA.init(); | ||
| 1148 | await sessionB.init(); | ||
| 1149 | expect(currentResidentRuntimeCount(), "two sessions resident, right at the cap").toBe(2); | ||
| 1150 | |||
| 1151 | // Force a deterministic LRU order: real-clock timestamps from two `init` | ||
| 1152 | // calls a tick apart are not reliable enough to assert on. | ||
| 1153 | touchResidentSession("session-cap-a", Date.now() - 10_000); | ||
| 1154 | touchResidentSession("session-cap-b", Date.now()); | ||
| 1155 | |||
| 1156 | await sessionC.init(); | ||
| 1157 | // `evictForCapIfNeeded` fires the eviction REQUEST via `ctx.waitUntil` and | ||
| 1158 | // returns without waiting for it, so `sessionC.init()` resolving proves | ||
| 1159 | // nothing about whether session A has actually been torn down yet — only | ||
| 1160 | // draining session C's queued background work does. | ||
| 1161 | await storageC.drainWaitUntil(); | ||
| 1162 | |||
| 1163 | expect(sessionA.initialized, "the least-recently-active session was evicted").toBe(false); | ||
| 1164 | expect(sessionA.engine, "its engine was released").toBeNull(); | ||
| 1165 | expect(sessionB.initialized, "the newer session survives untouched").toBe(true); | ||
| 1166 | expect(sessionC.initialized, "the session that triggered eviction still built").toBe(true); | ||
| 1167 | expect(currentResidentRuntimeCount(), "one evicted, two resident").toBe(2); | ||
| 1168 | }); | ||
| 1169 | |||
| 1170 | it("proceeds with init and records overflow when nothing resident is currently evictable", async () => { | ||
| 1171 | const { tracer, spans } = makeRecordingTracer(); | ||
| 1172 | const { session: sessionBusy } = makeResidencySession({ | ||
| 1173 | id: "session-busy", | ||
| 1174 | cap: 1, | ||
| 1175 | hasActiveStream: true, | ||
| 1176 | }); | ||
| 1177 | const { session: sessionNew } = makeResidencySession({ id: "session-new", cap: 1 }); | ||
| 1178 | sessionNew.withTelemetry = (effect) => Effect.withTracer(effect, tracer); | ||
| 1179 | |||
| 1180 | await sessionBusy.init(); | ||
| 1181 | expect(currentResidentRuntimeCount()).toBe(1); | ||
| 1182 | |||
| 1183 | await sessionNew.init(); | ||
| 1184 | |||
| 1185 | expect(sessionBusy.initialized, "a streaming session is never evicted").toBe(true); | ||
| 1186 | expect(sessionNew.initialized, "init proceeds despite the cap").toBe(true); | ||
| 1187 | expect( | ||
| 1188 | currentResidentRuntimeCount(), | ||
| 1189 | "the soft cap is exceeded rather than blocking init", | ||
| 1190 | ).toBe(2); | ||
| 1191 | |||
| 1192 | const initSpan = spans.find((span) => span.name === "McpSessionDO.init"); | ||
| 1193 | expect(initSpan?.attributes.get("mcp.isolate.cap_overflow")).toBe(true); | ||
| 1194 | }); | ||
| 1195 | |||
| 1196 | it("cold-restores on the next request after being evicted for the cap", async () => { | ||
| 1197 | const { session: sessionA } = makeResidencySession({ id: "session-restore-a", cap: 1 }); | ||
| 1198 | const { session: sessionB, storage: storageB } = makeResidencySession({ | ||
| 1199 | id: "session-restore-b", | ||
| 1200 | cap: 1, | ||
| 1201 | }); | ||
| 1202 | |||
| 1203 | await sessionA.init(); | ||
| 1204 | await sessionB.init(); | ||
| 1205 | await storageB.drainWaitUntil(); | ||
| 1206 | expect(sessionA.initialized, "evicted to make room under the cap").toBe(false); | ||
| 1207 | |||
| 1208 | await sessionA.init(); | ||
| 1209 | |||
| 1210 | expect(sessionA.initialized, "a later request rebuilds the evicted session").toBe(true); | ||
| 1211 | expect(sessionA.engine, "a fresh engine is installed").not.toBeNull(); | ||
| 1212 | expect(sessionA.server, "a fresh server is installed").toBeDefined(); | ||
| 1213 | }); | ||
| 1214 | |||
| 1215 | // Two different sessions' `init`s can both pick the SAME LRU candidate | ||
| 1216 | // before either eviction request lands — nothing serializes the picks. The | ||
| 1217 | // candidate must survive that safely: its own re-check plus `closeRuntime`'s | ||
| 1218 | // idempotency is what the base class's doc comments claim makes a repeat | ||
| 1219 | // request a no-op rather than a double-teardown. Requests here run | ||
| 1220 | // sequentially rather than raced, deliberately: the guarantee under test is | ||
| 1221 | // that a SECOND request against an already-evicted (or being-evicted) | ||
| 1222 | // candidate is a safe no-op, not a claim about ordering within a genuine | ||
| 1223 | // race — the same style already used by the "releasing twice" test below. | ||
| 1224 | it("treats a repeat eviction request against the same candidate as a safe no-op", async () => { | ||
| 1225 | const { session: candidate } = makeResidencySession({ id: "session-double-evict" }); | ||
| 1226 | await candidate.init(); | ||
| 1227 | expect(candidate.initialized).toBe(true); | ||
| 1228 | |||
| 1229 | await candidate.requestSelfEviction(); | ||
| 1230 | |||
| 1231 | expect(candidate.initialized, "the first request tears the candidate down").toBe(false); | ||
| 1232 | expect(candidate.engine).toBeNull(); | ||
| 1233 | expect(currentResidentRuntimeCount()).toBe(0); | ||
| 1234 | |||
| 1235 | // A repeat request — the shape two overlapping evictor `init`s produce, or | ||
| 1236 | // one that was merely slow to land after the candidate already tore | ||
| 1237 | // itself down some other way — must also be a no-op, not a crash or a | ||
| 1238 | // second decrement of the (already-zero) resident count. | ||
| 1239 | await expect( | ||
| 1240 | candidate.requestSelfEviction(), | ||
| 1241 | "a repeat request against an already-evicted candidate does not throw", | ||
| 1242 | ).resolves.toBeUndefined(); | ||
| 1243 | expect(candidate.initialized).toBe(false); | ||
| 1244 | expect( | ||
| 1245 | currentResidentRuntimeCount(), | ||
| 1246 | "the resident count is not decremented a second time", | ||
| 1247 | ).toBe(0); | ||
| 1248 | }); | ||
| 1249 | |||
| 1250 | // `residentRuntimeCount` only moves once a build actually FINISHES, so N | ||
| 1251 | // overlapping cold inits admitted at the same moment each read it, see | ||
| 1252 | // themselves still under the cap, and none of them evicts — residency then | ||
| 1253 | // blows past the cap by N once every build lands. `evictForCapIfNeeded` | ||
| 1254 | // closes that gap with an in-flight reservation counter: this pins that a | ||
| 1255 | // second concurrent admission sees the first's still-building reservation | ||
| 1256 | // and evicts, instead of also passing the count-only check for free. | ||
| 1257 | 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 () => { | ||
| 1258 | const { session: sessionA } = makeResidencySession({ id: "session-inflight-a", cap: 2 }); | ||
| 1259 | await sessionA.init(); | ||
| 1260 | expect(currentResidentRuntimeCount(), "one session resident, one below the cap").toBe(1); | ||
| 1261 | |||
| 1262 | const buildEntered = makeDeferred(); | ||
| 1263 | const buildGate = makeDeferred(); | ||
| 1264 | const { session: sessionB } = makeResidencySession({ id: "session-inflight-b", cap: 2 }); | ||
| 1265 | sessionB.buildMcpServer = () => | ||
| 1266 | Effect.gen(function* () { | ||
| 1267 | buildEntered.resolve(); | ||
| 1268 | yield* Effect.promise(() => buildGate.promise); | ||
| 1269 | return { mcpServer: makeServer(), engine: makeEngine().engine }; | ||
| 1270 | }); | ||
| 1271 | |||
| 1272 | const initB = sessionB.init(); | ||
| 1273 | // Deterministic: wait for sessionB's admission to actually be inside its | ||
| 1274 | // (gated) build — holding its in-flight reservation — instead of guessing | ||
| 1275 | // a number of microtask ticks. | ||
| 1276 | await buildEntered.promise; | ||
| 1277 | expect(currentInFlightColdBuildCount(), "sessionB's admission reserved a cold-build slot").toBe( | ||
| 1278 | 1, | ||
| 1279 | ); | ||
| 1280 | |||
| 1281 | const { session: sessionC, storage: storageC } = makeResidencySession({ | ||
| 1282 | id: "session-inflight-c", | ||
| 1283 | cap: 2, | ||
| 1284 | }); | ||
| 1285 | await sessionC.init(); | ||
| 1286 | await storageC.drainWaitUntil(); | ||
| 1287 | |||
| 1288 | // sessionB is not yet counted-as-resident — its build is still gated — so | ||
| 1289 | // sessionA is the only session actually registered as resident right now | ||
| 1290 | // and the only thing `pickEvictionCandidate` can choose. Its eviction is | ||
| 1291 | // what proves sessionC's admission saw sessionB's in-flight reservation | ||
| 1292 | // and treated the isolate as already at the cap. | ||
| 1293 | expect( | ||
| 1294 | sessionA.initialized, | ||
| 1295 | "evicted because the in-flight reservation counted against the cap", | ||
| 1296 | ).toBe(false); | ||
| 1297 | expect(sessionC.initialized, "the admission that triggered eviction still built").toBe(true); | ||
| 1298 | |||
| 1299 | buildGate.resolve(); | ||
| 1300 | await initB; | ||
| 1301 | expect(sessionB.initialized, "sessionB's gated build eventually completes").toBe(true); | ||
| 1302 | expect( | ||
| 1303 | currentInFlightColdBuildCount(), | ||
| 1304 | "the reservation was released once the build landed", | ||
| 1305 | ).toBe(0); | ||
| 1306 | }); | ||
| 1307 | |||
| 1308 | it("releases the in-flight cold-build reservation when the build fails, instead of leaking residual cap pressure", async () => { | ||
| 1309 | const { session: sessionA } = makeResidencySession({ id: "session-inflight-fail-a", cap: 2 }); | ||
| 1310 | await sessionA.init(); | ||
| 1311 | expect(currentResidentRuntimeCount()).toBe(1); | ||
| 1312 | |||
| 1313 | const { session: sessionFailing } = makeResidencySession({ | ||
| 1314 | id: "session-inflight-fail-b", | ||
| 1315 | cap: 2, | ||
| 1316 | }); | ||
| 1317 | sessionFailing.buildMcpServer = () => Effect.die(new Error("cold build boundary failure")); | ||
| 1318 | |||
| 1319 | await expect(sessionFailing.init()).rejects.toThrow(/cold build boundary failure/); | ||
| 1320 | expect( | ||
| 1321 | currentInFlightColdBuildCount(), | ||
| 1322 | "the reservation was released on the build's failure path, not leaked", | ||
| 1323 | ).toBe(0); | ||
| 1324 | expect(sessionFailing.initialized, "the failed build never became resident").toBe(false); | ||
| 1325 | |||
| 1326 | const { session: sessionC, storage: storageC } = makeResidencySession({ | ||
| 1327 | id: "session-inflight-fail-c", | ||
| 1328 | cap: 2, | ||
| 1329 | }); | ||
| 1330 | await sessionC.init(); | ||
| 1331 | await storageC.drainWaitUntil(); | ||
| 1332 | |||
| 1333 | expect( | ||
| 1334 | sessionA.initialized, | ||
| 1335 | "no residual in-flight pressure from the failed build, so sessionC's admission stays under the cap", | ||
| 1336 | ).toBe(true); | ||
| 1337 | expect( | ||
| 1338 | currentResidentRuntimeCount(), | ||
| 1339 | "sessionA and sessionC both resident, under the cap", | ||
| 1340 | ).toBe(2); | ||
| 1341 | }); | ||
| 1342 | |||
| 1343 | // The reservation used to be released only by an `Effect.ensuring` scoped to | ||
| 1344 | // the build block that starts right after `evictForCapIfNeeded` admits — | ||
| 1345 | // leaving a gap between the admission itself and that narrower wrap's own | ||
| 1346 | // coverage beginning. An interrupt landing in that gap leaked the | ||
| 1347 | // reservation permanently: nothing ever decremented it, so the isolate's | ||
| 1348 | // in-flight counter drifted up forever and every later admission saw false | ||
| 1349 | // cap pressure. The fix moves the release to a single `Effect.ensuring` | ||
| 1350 | // around init's entire program, so there is no window between acquiring the | ||
| 1351 | // reservation and being covered by its release. This interrupts `init` | ||
| 1352 | // itself (via the `AbortController` `initAbortController` exposes for | ||
| 1353 | // exactly this) while it is genuinely suspended opening the session db | ||
| 1354 | // handle — the first async step after admission — rather than merely | ||
| 1355 | // failing the build, which the pre-existing test above already covers. | ||
| 1356 | it("releases the in-flight cold-build reservation when init is interrupted after cap admission, before the build finishes", async () => { | ||
| 1357 | const { session: sessionA } = makeResidencySession({ id: "session-interrupt-a", cap: 2 }); | ||
| 1358 | await sessionA.init(); | ||
| 1359 | expect(currentResidentRuntimeCount(), "one session resident, one below the cap").toBe(1); | ||
| 1360 | |||
| 1361 | const dbHandleEntered = makeDeferred(); | ||
| 1362 | const dbHandleGate = makeDeferred(); | ||
| 1363 | const { session: sessionB } = makeResidencySession({ id: "session-interrupt-b", cap: 2 }); | ||
| 1364 | sessionB.openSessionDb = () => { | ||
| 1365 | dbHandleEntered.resolve(); | ||
| 1366 | // Never resolves — the interrupt below is what ends this suspension, | ||
| 1367 | // not the gate. `dbHandleGate.promise` only carries a `void` payload; | ||
| 1368 | // `.then` produces the right shape without a cast, and without ever | ||
| 1369 | // actually resolving — the callback here never runs. | ||
| 1370 | return dbHandleGate.promise.then(() => ({ end: () => undefined })); | ||
| 1371 | }; | ||
| 1372 | |||
| 1373 | const initPromise = sessionB.init(); | ||
| 1374 | // Deterministic: wait for sessionB's admission to have actually reserved | ||
| 1375 | // a slot and for its build to be genuinely suspended opening the db | ||
| 1376 | // handle, instead of guessing a number of microtask ticks. | ||
| 1377 | await dbHandleEntered.promise; | ||
| 1378 | expect( | ||
| 1379 | currentInFlightColdBuildCount(), | ||
| 1380 | "admission reserved a slot before the build's own db-handle open resolves", | ||
| 1381 | ).toBe(1); | ||
| 1382 | |||
| 1383 | sessionB.initAbortController?.abort(); | ||
| 1384 | |||
| 1385 | await expect( | ||
| 1386 | initPromise, | ||
| 1387 | "an interrupted init rejects rather than silently resolving", | ||
| 1388 | ).rejects.toThrow(); | ||
| 1389 | expect( | ||
| 1390 | currentInFlightColdBuildCount(), | ||
| 1391 | "the reservation was released by init's outer ensuring instead of leaking", | ||
| 1392 | ).toBe(0); | ||
| 1393 | expect(sessionB.initialized, "the interrupted build never became resident").toBe(false); | ||
| 1394 | |||
| 1395 | // The reservation not leaking is what lets a later admission still see | ||
| 1396 | // the isolate as under the cap. | ||
| 1397 | const { session: sessionC, storage: storageC } = makeResidencySession({ | ||
| 1398 | id: "session-interrupt-c", | ||
| 1399 | cap: 2, | ||
| 1400 | }); | ||
| 1401 | await sessionC.init(); | ||
| 1402 | await storageC.drainWaitUntil(); | ||
| 1403 | |||
| 1404 | expect( | ||
| 1405 | sessionA.initialized, | ||
| 1406 | "no residual in-flight pressure from the interrupted init, so sessionC's admission stays under the cap", | ||
| 1407 | ).toBe(true); | ||
| 1408 | expect( | ||
| 1409 | currentResidentRuntimeCount(), | ||
| 1410 | "sessionA and sessionC both resident, under the cap", | ||
| 1411 | ).toBe(2); | ||
| 1412 | }); | ||
| 1413 | |||
| 1414 | // `initialized` used to stay `true` across the async closes inside | ||
| 1415 | // `closeRuntime` (`server.close()`, `dbHandle.end()`), so a request landing | ||
| 1416 | // mid-teardown took `init`'s early-return path and ran against a | ||
| 1417 | // deleted server / null engine. `disposingRuntime` closes that: `init` | ||
| 1418 | // awaits any in-progress disposal before deciding whether to rebuild. | ||
| 1419 | it("a request landing mid-disposal awaits the in-progress teardown, then rebuilds and serves — never racing the closes", async () => { | ||
| 1420 | const { session } = makeResidencySession({ id: "session-mid-disposal" }); | ||
| 1421 | await session.init(); | ||
| 1422 | expect(session.initialized).toBe(true); | ||
| 1423 | const originalEngine = session.engine; | ||
| 1424 | |||
| 1425 | const closeEntered = makeDeferred(); | ||
| 1426 | const closeGate = makeDeferred(); | ||
| 1427 | let closeCalls = 0; | ||
| 1428 | const server = makeServer(); | ||
| 1429 | server.close = () => { | ||
| 1430 | closeCalls += 1; | ||
| 1431 | closeEntered.resolve(); | ||
| 1432 | return closeGate.promise; | ||
| 1433 | }; | ||
| 1434 | session.server = server; | ||
| 1435 | |||
| 1436 | const disposal = session.evictResidentRuntimeForCap(); | ||
| 1437 | // Deterministic: wait for disposal to actually be mid-teardown (inside | ||
| 1438 | // `server.close()`), instead of guessing a number of microtask ticks. | ||
| 1439 | await closeEntered.promise; | ||
| 1440 | |||
| 1441 | const rebuild = session.init(); | ||
| 1442 | let rebuildSettled = false; | ||
| 1443 | void rebuild.then(() => { | ||
| 1444 | rebuildSettled = true; | ||
| 1445 | }); | ||
| 1446 | await Promise.resolve(); | ||
| 1447 | await Promise.resolve(); | ||
| 1448 | expect( | ||
| 1449 | rebuildSettled, | ||
| 1450 | "init awaits the in-progress disposal instead of racing its still-pending close", | ||
| 1451 | ).toBe(false); | ||
| 1452 | |||
| 1453 | closeGate.resolve(); | ||
| 1454 | await disposal; | ||
| 1455 | await rebuild; | ||
| 1456 | |||
| 1457 | expect(closeCalls, "the runtime was closed exactly once").toBe(1); | ||
| 1458 | expect(session.initialized, "the request that landed mid-disposal rebuilt and is serving").toBe( | ||
| 1459 | true, | ||
| 1460 | ); | ||
| 1461 | expect( | ||
| 1462 | session.engine, | ||
| 1463 | "a fresh engine was installed, not the one that was mid-teardown", | ||
| 1464 | ).not.toBe(originalEngine); | ||
| 1465 | }); | ||
| 1466 | |||
| 1467 | // The `Effect.ensuring` that resolves `disposingRuntime` used to sit outside | ||
| 1468 | // any interruptibility guard, so an interrupt landing while `closeRuntime` | ||
| 1469 | // was genuinely suspended (mid `server.close()`, same seam as the test | ||
| 1470 | // above) still ran that `ensuring` and resolved `disposingRuntime` | ||
| 1471 | // immediately — before `dbHandle.end()`, the engine clear, or the residency | ||
| 1472 | // release ever ran. A concurrently waiting `init` would see the gate open | ||
| 1473 | // and rebuild against that half-released state. Wrapping the teardown body | ||
| 1474 | // in `Effect.uninterruptible` makes the interrupt request wait until the | ||
| 1475 | // body — including the still-gated `server.close()` — actually finishes, | ||
| 1476 | // so `disposingRuntime` cannot resolve early. This interrupts the disposal | ||
| 1477 | // itself (via `disposeAbortController`, exposed for exactly this) instead | ||
| 1478 | // of just failing a step, which the double-close test above doesn't cover. | ||
| 1479 | it("keeps a waiting init blocked when the disposal is interrupted mid-teardown, until the uninterruptible teardown actually finishes", async () => { | ||
| 1480 | const { session } = makeResidencySession({ id: "session-interrupted-disposal" }); | ||
| 1481 | await session.init(); | ||
| 1482 | expect(session.initialized).toBe(true); | ||
| 1483 | const originalEngine = session.engine; | ||
| 1484 | |||
| 1485 | const closeEntered = makeDeferred(); | ||
| 1486 | const closeGate = makeDeferred(); | ||
| 1487 | let closeCalls = 0; | ||
| 1488 | const server = makeServer(); | ||
| 1489 | server.close = () => { | ||
| 1490 | closeCalls += 1; | ||
| 1491 | closeEntered.resolve(); | ||
| 1492 | return closeGate.promise; | ||
| 1493 | }; | ||
| 1494 | session.server = server; | ||
| 1495 | |||
| 1496 | const disposal = session.evictResidentRuntimeForCap(); | ||
| 1497 | // Deterministic: wait for disposal to actually be mid-teardown (inside | ||
| 1498 | // `server.close()`), instead of guessing a number of microtask ticks. | ||
| 1499 | await closeEntered.promise; | ||
| 1500 | |||
| 1501 | // A concurrent init is the observable proof: it only proceeds once | ||
| 1502 | // `disposingRuntime` resolves, so it stays pending for exactly as long as | ||
| 1503 | // the teardown is genuinely still running. | ||
| 1504 | const rebuild = session.init(); | ||
| 1505 | let rebuildSettled = false; | ||
| 1506 | void rebuild.then(() => { | ||
| 1507 | rebuildSettled = true; | ||
| 1508 | }); | ||
| 1509 | |||
| 1510 | // Interrupt the disposal fiber while it is still suspended inside the | ||
| 1511 | // gated `server.close()`. On the old (interruptible) code this would let | ||
| 1512 | // the outer `Effect.ensuring` resolve `disposingRuntime` right here, | ||
| 1513 | // before `closeGate` ever resolves — letting `rebuild` proceed against a | ||
| 1514 | // still-open server / un-cleared engine. | ||
| 1515 | session.disposeAbortController?.abort(); | ||
| 1516 | // Interrupt delivery hops through more than a couple of microtasks (it is | ||
| 1517 | // not itself under test here), so give it a real macrotask tick before | ||
| 1518 | // asserting anything stayed blocked — a few `Promise.resolve()`s alone | ||
| 1519 | // are not enough ticks for the abort to even be delivered, uninterruptible | ||
| 1520 | // or not, and would make this assertion true vacuously either way. | ||
| 1521 | await new Promise((resolve) => setTimeout(resolve, 20)); | ||
| 1522 | expect( | ||
| 1523 | rebuildSettled, | ||
| 1524 | "the interrupt must not resolve disposingRuntime early — the teardown is still genuinely suspended inside server.close()", | ||
| 1525 | ).toBe(false); | ||
| 1526 | expect( | ||
| 1527 | session.engine, | ||
| 1528 | "the interrupt must not let the engine be cleared before server.close() actually returns", | ||
| 1529 | ).toBe(originalEngine); | ||
| 1530 | |||
| 1531 | closeGate.resolve(); | ||
| 1532 | // The interrupt may still surface once the now-uninterruptible body has | ||
| 1533 | // actually finished; only completeness of the teardown itself — not the | ||
| 1534 | // outer promise's resolve/reject outcome — is under test here. | ||
| 1535 | try { | ||
| 1536 | await disposal; | ||
| 1537 | } catch { | ||
| 1538 | // Expected: the interrupt requested above can still land once the | ||
| 1539 | // uninterruptible teardown finishes. | ||
| 1540 | } | ||
| 1541 | await rebuild; | ||
| 1542 | |||
| 1543 | expect(closeCalls, "server.close was actually called, not skipped").toBe(1); | ||
| 1544 | expect( | ||
| 1545 | session.initialized, | ||
| 1546 | "the request that was waiting on the interrupted disposal rebuilt and is serving", | ||
| 1547 | ).toBe(true); | ||
| 1548 | expect( | ||
| 1549 | session.engine, | ||
| 1550 | "a fresh engine was installed once the teardown genuinely finished, not the one that was mid-teardown", | ||
| 1551 | ).not.toBe(originalEngine); | ||
| 1552 | }); | ||
| 1553 | |||
| 1554 | it("does not double-close when two disposal triggers land concurrently (idle alarm and a cap eviction landing together)", async () => { | ||
| 1555 | const { session } = makeResidencySession({ id: "session-concurrent-dispose" }); | ||
| 1556 | await session.init(); | ||
| 1557 | |||
| 1558 | const closeGate = makeDeferred(); | ||
| 1559 | let closeCalls = 0; | ||
| 1560 | const server = makeServer(); | ||
| 1561 | server.close = () => { | ||
| 1562 | closeCalls += 1; | ||
| 1563 | return closeGate.promise; | ||
| 1564 | }; | ||
| 1565 | session.server = server; | ||
| 1566 | |||
| 1567 | const first = session.evictResidentRuntimeForCap(); | ||
| 1568 | const second = session.evictResidentRuntimeForCap(); | ||
| 1569 | |||
| 1570 | closeGate.resolve(); | ||
| 1571 | await Promise.all([first, second]); | ||
| 1572 | |||
| 1573 | expect( | ||
| 1574 | closeCalls, | ||
| 1575 | "the runtime was closed exactly once even though two disposals overlapped", | ||
| 1576 | ).toBe(1); | ||
| 1577 | expect(session.initialized).toBe(false); | ||
| 1578 | expect(currentResidentRuntimeCount(), "released exactly once").toBe(0); | ||
| 1579 | }); | ||
| 1580 | |||
| 1581 | describe("resident session registry release", () => { | ||
| 1582 | it("is idempotent", () => { | ||
| 1583 | registerResidentSession({ | ||
| 1584 | sessionId: "idempotent-release", | ||
| 1585 | lastActivityMs: Date.now(), | ||
| 1586 | canEvict: () => true, | ||
| 1587 | dispose: async () => undefined, | ||
| 1588 | }); | ||
| 1589 | |||
| 1590 | expect(() => releaseResidentSession("idempotent-release")).not.toThrow(); | ||
| 1591 | expect( | ||
| 1592 | () => releaseResidentSession("idempotent-release"), | ||
| 1593 | "releasing twice is a no-op", | ||
| 1594 | ).not.toThrow(); | ||
| 1595 | expect(residentSessionIdsForTest()).not.toContain("idempotent-release"); | ||
| 1596 | expect(pickEvictionCandidate()).toBeUndefined(); | ||
| 1597 | }); | ||
| 1598 | |||
| 1599 | // The eviction REQUEST is fire-and-forget (`ctx.waitUntil`), so the | ||
| 1600 | // evictor's `init` has no synchronous failure to react to any more — it | ||
| 1601 | // cannot tell a request that failed from one still in flight. The new | ||
| 1602 | // failure story: optimistically LEAVE the registry entry (only the | ||
| 1603 | // candidate's own successful teardown removes it, since as far as the | ||
| 1604 | // evictor knows the candidate might still be actually resident), and mark | ||
| 1605 | // it recently-requested so a following pick skips it rather than | ||
| 1606 | // re-requesting the same stuck candidate forever. | ||
| 1607 | it("keeps the registry slot when the candidate's own eviction request fails, without touching the resident count", async () => { | ||
| 1608 | const { session: sessionFailing } = makeResidencySession({ | ||
| 1609 | id: "session-failing", | ||
| 1610 | cap: 1, | ||
| 1611 | }); | ||
| 1612 | // Simulate the REQUEST itself failing — the self-addressed stub call | ||
| 1613 | // never lands (the production analogue: `mcpSessionStub(...).requestCapEviction()` | ||
| 1614 | // rejects), never even reaching the candidate's own re-check. | ||
| 1615 | // oxlint-disable-next-line executor/no-promise-reject -- test double: simulates a rejected self-eviction stub call, not application error modeling. | ||
| 1616 | sessionFailing.requestSelfEviction = () => Promise.reject(new Error("stub unreachable")); | ||
| 1617 | const { session: sessionB, storage: storageB } = makeResidencySession({ | ||
| 1618 | id: "session-b-after-failure", | ||
| 1619 | cap: 1, | ||
| 1620 | }); | ||
| 1621 | |||
| 1622 | await sessionFailing.init(); | ||
| 1623 | await sessionB.init(); | ||
| 1624 | await storageB.drainWaitUntil(); | ||
| 1625 | |||
| 1626 | expect( | ||
| 1627 | sessionFailing.initialized, | ||
| 1628 | "a failed eviction REQUEST never runs the candidate's own teardown", | ||
| 1629 | ).toBe(true); | ||
| 1630 | expect(sessionB.initialized, "the triggering init is never blocked by the failure").toBe( | ||
| 1631 | true, | ||
| 1632 | ); | ||
| 1633 | expect( | ||
| 1634 | currentResidentRuntimeCount(), | ||
| 1635 | "the failing candidate's own runtime is still actually resident", | ||
| 1636 | ).toBe(2); | ||
| 1637 | expect( | ||
| 1638 | residentSessionIdsForTest(), | ||
| 1639 | "the registry slot is kept — only the candidate's own successful teardown removes it", | ||
| 1640 | ).toContain("session-failing"); | ||
| 1641 | expect( | ||
| 1642 | pickEvictionCandidate()?.sessionId, | ||
| 1643 | "marked recently-requested, so the next pick does not re-target the same stuck candidate", | ||
| 1644 | ).not.toBe("session-failing"); | ||
| 1645 | }); | ||
| 1646 | |||
| 1647 | it("skips an entry with a recent pending eviction request, and picks the next evictable one instead", () => { | ||
| 1648 | const now = Date.now(); | ||
| 1649 | registerResidentSession({ | ||
| 1650 | sessionId: "recently-requested", | ||
| 1651 | lastActivityMs: now - 10_000, | ||
| 1652 | canEvict: () => true, | ||
| 1653 | dispose: async () => undefined, | ||
| 1654 | }); | ||
| 1655 | registerResidentSession({ | ||
| 1656 | sessionId: "next-candidate", | ||
| 1657 | lastActivityMs: now - 5_000, | ||
| 1658 | canEvict: () => true, | ||
| 1659 | dispose: async () => undefined, | ||
| 1660 | }); | ||
| 1661 | markEvictionRequested("recently-requested", now); | ||
| 1662 | |||
| 1663 | expect(pickEvictionCandidate(now)?.sessionId).toBe("next-candidate"); | ||
| 1664 | }); | ||
| 1665 | |||
| 1666 | it("picks a previously-requested entry again once the grace period elapses", () => { | ||
| 1667 | const now = Date.now(); | ||
| 1668 | registerResidentSession({ | ||
| 1669 | sessionId: "stale-request", | ||
| 1670 | lastActivityMs: now - 10_000, | ||
| 1671 | canEvict: () => true, | ||
| 1672 | dispose: async () => undefined, | ||
| 1673 | }); | ||
| 1674 | markEvictionRequested("stale-request", now - EVICTION_REQUEST_GRACE_MS - 1); | ||
| 1675 | |||
| 1676 | expect(pickEvictionCandidate(now)?.sessionId).toBe("stale-request"); | ||
| 1677 | }); | ||
| 1678 | |||
| 1679 | it("marking an eviction request on an entry that no longer exists is a no-op", () => { | ||
| 1680 | expect(() => markEvictionRequested("never-registered")).not.toThrow(); | ||
| 1681 | }); | ||
| 1682 | }); | ||
| 1683 | }); | ||
| @@ -36,8 +36,19 @@ import { · 35 unmodified lines | |||
| 36 | } from "./session-alarm-policy"; | 36 | } from "./session-alarm-policy"; |
| 37 | import { | 37 | import { |
| 38 | acquireResidentRuntime, | 38 | acquireResidentRuntime, |
| 39 | currentInFlightColdBuildCount, | ||
| 40 | currentResidentRuntimeCount, | ||
| 41 | markEvictionRequested, | ||
| 42 | pickEvictionCandidate, | ||
| 43 | registerResidentSession, | ||
| 44 | releaseColdBuildSlot, | ||
| 39 | releaseResidentRuntime, | 45 | releaseResidentRuntime, |
| 46 | releaseResidentSession, | ||
| 47 | reserveColdBuildSlot, | ||
| 48 | type ResidentSessionEntry, | ||
| 40 | residencyAttributes, | 49 | residencyAttributes, |
| 50 | RESIDENT_RUNTIME_SOFT_CAP, | ||
| 51 | touchResidentSession, | ||
| 41 | } from "./session-runtime-residency"; | 52 | } from "./session-runtime-residency"; |
| 42 | 53 | ||
| 43 | export type IncomingTraceHeaders = IncomingPropagationHeaders; | 54 | export type IncomingTraceHeaders = IncomingPropagationHeaders; |
| @@ -261,6 +272,43 @@ export abstract class McpAgentSessionDOBase< · 217 unmodified lines | |||
| 261 | * gauge. Tracked separately from `engine` because `closeRuntime` runs on | 272 | * gauge. Tracked separately from `engine` because `closeRuntime` runs on |
| 262 | * paths where nothing was ever built, and it must not decrement then. */ | 273 | * paths where nothing was ever built, and it must not decrement then. */ |
| 263 | private countedAsResident = false; | 274 | private countedAsResident = false; |
| 275 | /** Whether THIS init's cold-build admission is currently holding a reserved | ||
| 276 | * slot in the isolate-wide in-flight counter. Mirrors `countedAsResident`'s | ||
| 277 | * guard discipline: gates {@link releaseColdBuildSlotIfReserved} so it | ||
| 278 | * releases exactly once per reservation, from whichever of its two callers | ||
| 279 | * (success or failure/interrupt) gets there first. */ | ||
| 280 | private reservedColdBuildSlot = false; | ||
| 281 | /** | ||
| 282 | * The `AbortController` backing the currently in-flight `init` call's | ||
| 283 | * `Effect.runPromise`, so its root fiber can be interrupted from outside | ||
| 284 | * the Promise it returns. Nothing in production ever calls `.abort()` on | ||
| 285 | * it today — it exists so a unit test can interrupt `init` deterministically | ||
| 286 | * at a chosen suspension point (e.g. mid `openSessionDbHandle`) and assert | ||
| 287 | * the cold-build reservation is still released, rather than only ever being | ||
| 288 | * exercised through failure/defect paths. | ||
| 289 | */ | ||
| 290 | private initAbortController: AbortController | null = null; | ||
| 291 | /** Same purpose as {@link initAbortController}, for `disposeIdleRuntime`'s | ||
| 292 | * root effect (`closeRuntime` plus its alarm/activity bookkeeping) — lets | ||
| 293 | * a unit test interrupt a disposal deterministically to confirm the | ||
| 294 | * uninterruptible teardown in `closeRuntime` actually finishes rather than | ||
| 295 | * being cut short. */ | ||
| 296 | private disposeAbortController: AbortController | null = null; | ||
| 297 | /** | ||
| 298 | * The in-progress runtime disposal, if one is running right now — from | ||
| 299 | * whichever of `closeRuntime`'s callers (the idle alarm, a cap eviction | ||
| 300 | * request, `cleanup`) got there first. Set at the very start of | ||
| 301 | * `closeRuntime`'s body, before its first async close, and cleared in a | ||
| 302 | * finally once that disposal settles. | ||
| 303 | * | ||
| 304 | * Exists to close two races at once: `init` awaits this before deciding | ||
| 305 | * whether to rebuild, so a request that lands mid-teardown never races the | ||
| 306 | * closes it is waiting on; and a second `closeRuntime` call that lands | ||
| 307 | * while this is set (the idle alarm and a cap eviction request landing | ||
| 308 | * together) waits on the SAME promise instead of running the teardown a | ||
| 309 | * second time. | ||
| 310 | */ | ||
| 311 | private disposingRuntime: Promise<void> | null = null; | ||
| 264 | private onStartPromise: Promise<void> | null = null; | 312 | private onStartPromise: Promise<void> | null = null; |
| 265 | private lastActivityMs = 0; | 313 | private lastActivityMs = 0; |
| 266 | private resolvedSessionName: string | undefined = undefined; | 314 | private resolvedSessionName: string | undefined = undefined; |
| @@ -459,6 +507,49 @@ export abstract class McpAgentSessionDOBase< · 192 unmodified lines | |||
| 459 | }); | 507 | }); |
| 460 | } | 508 | } |
| 461 | 509 | ||
| 510 | /** | ||
| 511 | * Whether this host can route an eviction REQUEST to this session's own | ||
| 512 | * Durable Object instance rather than tearing it down directly in some | ||
| 513 | * OTHER session's request context. Hosts that override `requestSelfEviction` | ||
| 514 | * with a real self-addressed stub call return `true` here too. A host that | ||
| 515 | * returns `false` (the default) is never registered as an eviction | ||
| 516 | * candidate at all — see `init` — so it degrades to purely observational: | ||
| 517 | * the residency gauge and cap-overflow attribute still work, cap eviction | ||
| 518 | * just never picks it. | ||
| 519 | */ | ||
| 520 | protected supportsCapEviction(): boolean { | ||
| 521 | return false; | ||
| 522 | } | ||
| 523 | |||
| 524 | /** | ||
| 525 | * Ask THIS session's own Durable Object instance — running in ITS OWN | ||
| 526 | * request/IoContext — to tear down its resident runtime because the | ||
| 527 | * isolate is over its cap. | ||
| 528 | * | ||
| 529 | * This must never be `evictResidentRuntimeForCap()` called directly on | ||
| 530 | * `this` from within another session's request. Even though that would be | ||
| 531 | * the exact same JS object in the exact same isolate (Durable Objects with | ||
| 532 | * the same id ARE the same instance), a plain method call does not create a | ||
| 533 | * new IoContext — it runs inside whatever IoContext is already current, | ||
| 534 | * which belongs to the CALLING session's request. workerd binds I/O objects | ||
| 535 | * (a postgres.js socket, a storage transaction, a span flush) to the | ||
| 536 | * IoContext that created them, so tearing this session down that way throws | ||
| 537 | * "Cannot perform I/O on behalf of a different request" or silently | ||
| 538 | * miscredits the I/O to the wrong request. Routing through this session's | ||
| 539 | * own Durable Object STUB (`mcpSessionStub(...).requestCapEviction()`) goes | ||
| 540 | * through the Workers RPC/fetch machinery instead, which gives the call a | ||
| 541 | * freshly-created IoContext bound to itself — so the teardown it triggers | ||
| 542 | * runs correctly scoped, no matter which session's `init` sent the request. | ||
| 543 | * | ||
| 544 | * Overridden per host, because only a concrete host knows its own | ||
| 545 | * self-addressed namespace binding. The base default is a no-op so a host | ||
| 546 | * that never overrides it (and therefore never overrides | ||
| 547 | * `supportsCapEviction` to `true`) is simply never asked. | ||
| 548 | */ | ||
| 549 | protected requestSelfEviction(): Promise<void> { | ||
| 550 | return Promise.resolve(); | ||
| 551 | } | ||
| 552 | |||
| 462 | protected readonly browserApprovalStore: BrowserApprovalStore = { | 553 | protected readonly browserApprovalStore: BrowserApprovalStore = { |
| 463 | takeResponse: (executionId) => this.takeApprovalResponse(executionId), | 554 | takeResponse: (executionId) => this.takeApprovalResponse(executionId), |
| 464 | waitForResponse: (executionId) => this.waitForApprovalResponse(executionId), | 555 | waitForResponse: (executionId) => this.waitForApprovalResponse(executionId), |
| @@ -540,6 +631,11 @@ export abstract class McpAgentSessionDOBase< · 75 unmodified lines | |||
| 540 | 631 | ||
| 541 | private async markActivity(now = Date.now()): Promise<void> { | 632 | private async markActivity(now = Date.now()): Promise<void> { |
| 542 | this.lastActivityMs = now; | 633 | this.lastActivityMs = now; |
| 634 | // Keeps the isolate-wide eviction registry's LRU order current. A no-op | ||
| 635 | // when this session has no registry entry yet (nothing resident) or none | ||
| 636 | // any more (already disposed) — `touchResidentSession` is a lookup-then-set | ||
| 637 | // that quietly does nothing on a miss. | ||
| 638 | touchResidentSession(this.sessionIdForTelemetry(), now); | ||
| 543 | await Promise.all([ | 639 | await Promise.all([ |
| 544 | this.ctx.storage.put(LAST_ACTIVITY_KEY, now), | 640 | this.ctx.storage.put(LAST_ACTIVITY_KEY, now), |
| 545 | this.ctx.storage.setAlarm(now + this.sessionTimeoutMs()), | 641 | this.ctx.storage.setAlarm(now + this.sessionTimeoutMs()), |
| @@ -662,15 +758,26 @@ export abstract class McpAgentSessionDOBase< · 116 unmodified lines | |||
| 662 | } | 758 | } |
| 663 | 759 | ||
| 664 | /** | 760 | /** |
| 665 | * Drop this session's execution runtime because it has gone idle, returning | 761 | * Drop this session's execution runtime, returning its memory to the |
| 666 | * its memory to the isolate. Nothing durable is discarded, so the next | 762 | * isolate. Nothing durable is discarded, so the next request restores the |
| 667 | * request restores the session and the client sees only restore latency. | 763 | * session and the client sees only restore latency. |
| 764 | * | ||
| 765 | * `reason` disambiguates WHY on the shared span and log line without | ||
| 766 | * splitting them: `"idle"` is the alarm-driven path (this session itself | ||
| 767 | * went quiet), `"cap"` is another session's `init` evicting this one because | ||
| 768 | * the isolate was over its resident-runtime ceiling. The mechanism — | ||
| 769 | * `closeRuntime` plus dropping the durable alarm/activity bookkeeping — is | ||
| 770 | * identical either way; only the trigger differs, and callers are expected | ||
| 771 | * to have already established (via their own eligibility check) that | ||
| 772 | * disposing this session right now is safe. | ||
| 668 | */ | 773 | */ |
| 669 | private async disposeIdleRuntime(input: { | 774 | private async disposeIdleRuntime(input: { |
| 670 | readonly idleMs: number; | 775 | readonly idleMs: number; |
| 671 | readonly pausedExecutionCount: number; | 776 | readonly pausedExecutionCount: number; |
| 672 | readonly activeStreamCount: number; | 777 | readonly activeStreamCount: number; |
| 778 | readonly reason?: "idle" | "cap"; | ||
| 673 | }): Promise<void> { | 779 | }): Promise<void> { |
| 780 | const reason = input.reason ?? "idle"; | ||
| 674 | console.info( | 781 | console.info( |
| 675 | JSON.stringify({ | 782 | JSON.stringify({ |
| 676 | event: "mcp_session_idle_runtime_dispose", | 783 | event: "mcp_session_idle_runtime_dispose", |
| @@ -678,6 +785,7 @@ export abstract class McpAgentSessionDOBase< · 1 unmodified line | |||
| 678 | idleMs: input.idleMs, | 785 | idleMs: input.idleMs, |
| 679 | pausedExecutionCount: input.pausedExecutionCount, | 786 | pausedExecutionCount: input.pausedExecutionCount, |
| 680 | activeStreamCount: input.activeStreamCount, | 787 | activeStreamCount: input.activeStreamCount, |
| 788 | reason, | ||
| 681 | }), | 789 | }), |
| 682 | ); | 790 | ); |
| 683 | const self = this; | 791 | const self = this; |
| @@ -691,12 +799,15 @@ export abstract class McpAgentSessionDOBase< · 7 unmodified lines | |||
| 691 | // isolate is actually holding now rather than what it held a moment ago. | 799 | // isolate is actually holding now rather than what it held a moment ago. |
| 692 | yield* Effect.annotateCurrentSpan(residencyAttributes()); | 800 | yield* Effect.annotateCurrentSpan(residencyAttributes()); |
| 693 | }).pipe( | 801 | }).pipe( |
| 802 | // Span name kept stable for dashboard continuity across both triggers; | ||
| 803 | // `mcp.session.dispose_reason` is what disambiguates them. | ||
| 694 | Effect.withSpan("mcp.session.idle_runtime_dispose", { | 804 | Effect.withSpan("mcp.session.idle_runtime_dispose", { |
| 695 | attributes: { | 805 | attributes: { |
| 696 | "mcp.session.id": self.sessionId, | 806 | "mcp.session.id": self.sessionId, |
| 697 | "mcp.session.idle_ms": input.idleMs, | 807 | "mcp.session.idle_ms": input.idleMs, |
| 698 | "mcp.session.paused_execution_count": input.pausedExecutionCount, | 808 | "mcp.session.paused_execution_count": input.pausedExecutionCount, |
| 699 | "mcp.session.active_stream_count": input.activeStreamCount, | 809 | "mcp.session.active_stream_count": input.activeStreamCount, |
| 810 | "mcp.session.dispose_reason": reason, | ||
| 700 | }, | 811 | }, |
| 701 | }), | 812 | }), |
| 702 | ); | 813 | ); |
| @@ -704,7 +815,194 @@ export abstract class McpAgentSessionDOBase< · 1 unmodified line | |||
| 704 | // trace. It still has to be flushed explicitly — the alarm is not on any | 815 | // trace. It still has to be flushed explicitly — the alarm is not on any |
| 705 | // request's response path, and without the flush the span dies with the | 816 | // request's response path, and without the flush the span dies with the |
| 706 | // isolate and the mechanism stays unobservable in production. | 817 | // isolate and the mechanism stays unobservable in production. |
| 707 | await Effect.runPromise(this.withSpanFlush(this.withTelemetry(program))); | 818 | // See `disposeAbortController`'s doc comment: nothing in production aborts |
| 819 | // this signal today; it exists so a unit test can interrupt this exact | ||
| 820 | // fiber and confirm `closeRuntime`'s now-uninterruptible teardown still | ||
| 821 | // runs to completion instead of being cut short. | ||
| 822 | const abortController = new AbortController(); | ||
| 823 | this.disposeAbortController = abortController; | ||
| 824 | await Effect.runPromise(this.withSpanFlush(this.withTelemetry(program)), { | ||
| 825 | signal: abortController.signal, | ||
| 826 | }).finally(() => { | ||
| 827 | // Only clear the field if it is still THIS call's controller — see the | ||
| 828 | // matching comment in `init` for why an overlapping later call's | ||
| 829 | // controller must not be clobbered. | ||
| 830 | if (this.disposeAbortController === abortController) this.disposeAbortController = null; | ||
| 831 | }); | ||
| 832 | } | ||
| 833 | |||
| 834 | /** | ||
| 835 | * Isolate is over its resident-runtime soft cap and this session was the | ||
| 836 | * LRU-eligible pick (see `pickEvictionCandidate`). Re-checks every | ||
| 837 | * disqualifying signal `decideSessionAlarm` treats as active work — LIVE, | ||
| 838 | * not the snapshot `canEvict` used to be picked — because a session can | ||
| 839 | * start a request in the gap between being picked and being disposed here. | ||
| 840 | * `canEvict`'s snapshot is a cheap, synchronous, necessarily-optimistic | ||
| 841 | * filter for CHOOSING among candidates; this is the authoritative gate that | ||
| 842 | * actually decides whether disposing this session right now is safe, and it | ||
| 843 | * includes the one signal `canEvict` cannot see synchronously — undelivered | ||
| 844 | * stream responses still in storage. A session that is no longer eligible is | ||
| 845 | * left alone: eviction becomes a no-op instead of a wrongful teardown. | ||
| 846 | */ | ||
| 847 | private async evictResidentRuntimeForCap(): Promise<void> { | ||
| 848 | const [pausedExecutionCount, runningExecutionCount] = await Promise.all([ | ||
| 849 | this.pausedExecutionCount(), | ||
| 850 | this.runningExecutionCount(), | ||
| 851 | ]); | ||
| 852 | const activeStreamCount = this.activeStreamCount(); | ||
| 853 | if (pausedExecutionCount > 0 || runningExecutionCount > 0 || activeStreamCount > 0) return; | ||
| 854 | const idleMs = this.lastActivityMs > 0 ? Date.now() - this.lastActivityMs : 0; | ||
| 855 | await this.disposeIdleRuntime({ | ||
| 856 | idleMs, | ||
| 857 | pausedExecutionCount, | ||
| 858 | activeStreamCount, | ||
| 859 | reason: "cap", | ||
| 860 | }); | ||
| 861 | } | ||
| 862 | |||
| 863 | /** | ||
| 864 | * Cheap, synchronous eligibility filter consulted by `pickEvictionCandidate` | ||
| 865 | * to choose AMONG resident sessions. Deliberately conservative rather than | ||
| 866 | * exhaustive: it mirrors the same "no active work" signals | ||
| 867 | * `decideSessionAlarm` treats as disqualifying, restricted to what can be | ||
| 868 | * answered without an async storage read (undelivered stream responses, | ||
| 869 | * `runningExecutionCount`, requires one). `evictResidentRuntimeForCap` | ||
| 870 | * re-checks the full set — including that signal — right before actually | ||
| 871 | * disposing, so a false "evictable" here can only ever produce a safe | ||
| 872 | * no-op, never a wrongful eviction. | ||
| 873 | */ | ||
| 874 | private canEvictResidentRuntime(): boolean { | ||
| 875 | if (this.activeStreamCount() > 0) return false; | ||
| 876 | if (!this.engine) return true; | ||
| 877 | // 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. | ||
| 878 | try { | ||
| 879 | return Effect.runSync(this.engine.pausedExecutionCount()) === 0; | ||
| 880 | } catch { | ||
| 881 | return false; | ||
| 882 | } | ||
| 883 | } | ||
| 884 | |||
| 885 | /** | ||
| 886 | * The soft cap this instance enforces. A protected method rather than a bare | ||
| 887 | * reference to `RESIDENT_RUNTIME_SOFT_CAP`, matching `sessionTimeoutMs` and | ||
| 888 | * `maxPausedSessionIdleMs` elsewhere in this class, so tests can install a | ||
| 889 | * small cap and exercise real eviction without registering 32 sessions. | ||
| 890 | */ | ||
| 891 | protected residentRuntimeSoftCap(): number { | ||
| 892 | return RESIDENT_RUNTIME_SOFT_CAP; | ||
| 893 | } | ||
| 894 | |||
| 895 | /** | ||
| 896 | * Before this session's own runtime is built, make room if the isolate is | ||
| 897 | * already at its resident-runtime cap. Picks AT MOST ONE other session — | ||
| 898 | * never loops — and never blocks or fails THIS init on the outcome: the | ||
| 899 | * eviction REQUEST is fired via `ctx.waitUntil` and this init proceeds | ||
| 900 | * immediately, because a memory-pressure mechanism must never itself become | ||
| 901 | * the reason a session fails or is delayed starting. If nothing is | ||
| 902 | * currently evictable (every resident is streaming, paused, or already has | ||
| 903 | * a request outstanding), this session still builds and the overflow is | ||
| 904 | * only recorded on the init span. | ||
| 905 | * | ||
| 906 | * `candidate.dispose` does not run the candidate's teardown here — it SENDS | ||
| 907 | * the candidate a request that ITS OWN Durable Object instance executes in | ||
| 908 | * its own context (see `requestSelfEviction`'s doc comment). Because the | ||
| 909 | * request is fire-and-forget, this init has no way to react to it failing: | ||
| 910 | * the registry entry is left in place either way (never released here on | ||
| 911 | * failure), since whatever the candidate holds is, as far as this init | ||
| 912 | * knows, still actually resident — only the candidate's own successful | ||
| 913 | * teardown removes its entry. `markEvictionRequested` timestamps the entry | ||
| 914 | * up front so a candidate whose request is stuck or failing does not squat | ||
| 915 | * the LRU pick forever: the next init's `pickEvictionCandidate` skips it | ||
| 916 | * (see the grace period there) and picks the NEXT candidate instead. | ||
| 917 | * | ||
| 918 | * Safe to fire at the same candidate twice — two different sessions' inits | ||
| 919 | * both picking the same LRU entry before either's request lands — because | ||
| 920 | * the candidate's own handler (`evictResidentRuntimeForCap`) re-checks | ||
| 921 | * liveness and its teardown (`closeRuntime`) is idempotent: a second | ||
| 922 | * request either finds the runtime already gone (a no-op) or finds it | ||
| 923 | * newly busy again and leaves it alone. | ||
| 924 | * | ||
| 925 | * The check reads `currentResidentRuntimeCount() + currentInFlightColdBuildCount()`, | ||
| 926 | * not `currentResidentRuntimeCount()` alone. Residency only moves once a | ||
| 927 | * cold build actually finishes, so without the in-flight term, N | ||
| 928 | * overlapping cold inits arriving at the cap would each read the same | ||
| 929 | * still-under-cap count before any of them finishes, none would evict, and | ||
| 930 | * residency would land N over the cap once every build completed. Admitting | ||
| 931 | * (reserving a slot) unconditionally below — whether or not eviction fires | ||
| 932 | * for THIS init — is what lets the NEXT concurrent init see this one | ||
| 933 | * reflected in the sum. | ||
| 934 | */ | ||
| 935 | private evictForCapIfNeeded(): Effect.Effect<void> { | ||
| 936 | const self = this; | ||
| 937 | return Effect.gen(function* () { | ||
| 938 | const overCap = | ||
| 939 | currentResidentRuntimeCount() + currentInFlightColdBuildCount() >= | ||
| 940 | self.residentRuntimeSoftCap(); | ||
| 941 | // Admission: this init is about to start a cold build, so it counts as | ||
| 942 | // in-flight from here whether or not the check above finds anything to | ||
| 943 | // evict. Released exactly once `evictForCapIfNeeded`'s caller (`init`) | ||
| 944 | // either finishes building or fails/is interrupted — see | ||
| 945 | // `releaseColdBuildSlotIfReserved`. | ||
| 946 | self.reservedColdBuildSlot = true; | ||
| 947 | reserveColdBuildSlot(); | ||
| 948 | if (!overCap) return; | ||
| 949 | const candidate = pickEvictionCandidate(); | ||
| 950 | if (!candidate) { | ||
| 951 | yield* Effect.annotateCurrentSpan({ "mcp.isolate.cap_overflow": true }); | ||
| 952 | return; | ||
| 953 | } | ||
| 954 | yield* Effect.sync(() => self.queueCapEvictionRequest(candidate)); | ||
| 955 | }); | ||
| 956 | } | ||
| 957 | |||
| 958 | /** | ||
| 959 | * Release this init's in-flight cold-build reservation (see | ||
| 960 | * `evictForCapIfNeeded`), if it is still holding one. Idempotent, mirroring | ||
| 961 | * `countedAsResident`'s guard discipline: called from both of its possible | ||
| 962 | * finishing points — the build's success path in `init`, and an | ||
| 963 | * `Effect.ensuring` finalizer covering the build's failure/interrupt path — | ||
| 964 | * so whichever happens releases the slot exactly once, and the other | ||
| 965 | * becomes a safe no-op. | ||
| 966 | */ | ||
| 967 | private releaseColdBuildSlotIfReserved(): Effect.Effect<void> { | ||
| 968 | const self = this; | ||
| 969 | return Effect.sync(() => { | ||
| 970 | if (!self.reservedColdBuildSlot) return; | ||
| 971 | self.reservedColdBuildSlot = false; | ||
| 972 | releaseColdBuildSlot(); | ||
| 973 | }); | ||
| 974 | } | ||
| 975 | |||
| 976 | /** | ||
| 977 | * Fires the eviction REQUEST at `candidate`'s own stub and returns | ||
| 978 | * immediately — same fire-and-forget shape as | ||
| 979 | * `queuePendingApprovalLeaseStart`/`queuePendingApprovalLeaseExpiration` | ||
| 980 | * below. `markEvictionRequested` is stamped up front, before the request | ||
| 981 | * even lands, so a stuck or slow candidate cannot be re-picked by the next | ||
| 982 | * init in the meantime (see `pickEvictionCandidate`'s grace period). | ||
| 983 | */ | ||
| 984 | private queueCapEvictionRequest(candidate: ResidentSessionEntry): void { | ||
| 985 | markEvictionRequested(candidate.sessionId); | ||
| 986 | this.ctx.waitUntil( | ||
| 987 | Effect.runPromise( | ||
| 988 | Effect.tryPromise({ | ||
| 989 | try: () => candidate.dispose("cap"), | ||
| 990 | catch: (cause: unknown) => cause, | ||
| 991 | }).pipe( | ||
| 992 | Effect.catch((cause: unknown) => | ||
| 993 | Effect.sync(() => { | ||
| 994 | console.warn( | ||
| 995 | JSON.stringify({ | ||
| 996 | event: "mcp_session_cap_eviction_request_failed", | ||
| 997 | sessionId: candidate.sessionId, | ||
| 998 | }), | ||
| 999 | ); | ||
| 1000 | console.error("[mcp-session] cap eviction request failed:", cause); | ||
| 1001 | }), | ||
| 1002 | ), | ||
| 1003 | ), | ||
| 1004 | ), | ||
| 1005 | ); | ||
| 708 | } | 1006 | } |
| 709 | 1007 | ||
| 710 | private resolveAndStoreSessionMeta(token: McpSessionInit) { | 1008 | private resolveAndStoreSessionMeta(token: McpSessionInit) { |
| @@ -897,7 +1195,32 @@ export abstract class McpAgentSessionDOBase< · 186 unmodified lines | |||
| 897 | 1195 | ||
| 898 | private closeRuntime(options: { readonly closeStreams?: boolean } = {}): Effect.Effect<void> { | 1196 | private closeRuntime(options: { readonly closeStreams?: boolean } = {}): Effect.Effect<void> { |
| 899 | const self = this; | 1197 | const self = this; |
| 1198 | // A disposal is already tearing this runtime down — the idle alarm and a | ||
| 1199 | // cap eviction request landing together, or a plain repeat call on any of | ||
| 1200 | // `closeRuntime`'s existing callers. `server.close()`/`dbHandle.end()` | ||
| 1201 | // are not safe to run twice concurrently on the same resources, and a | ||
| 1202 | // second pass through the body below would double-release the residency | ||
| 1203 | // counters. Wait for the SAME in-progress disposal instead of starting a | ||
| 1204 | // second one; do not run the teardown body again. | ||
| 1205 | if (self.disposingRuntime) { | ||
| 1206 | const inProgress = self.disposingRuntime; | ||
| 1207 | return Effect.promise(() => inProgress); | ||
| 1208 | } | ||
| 1209 | let resolveDisposal!: () => void; | ||
| 1210 | const disposal = new Promise<void>((resolve) => { | ||
| 1211 | resolveDisposal = resolve; | ||
| 1212 | }); | ||
| 1213 | self.disposingRuntime = disposal; | ||
| 900 | return Effect.gen(function* () { | 1214 | return Effect.gen(function* () { |
| 1215 | // Flip this BEFORE the first async close below (`server.close()`), not | ||
| 1216 | // after. `init` awaits `disposingRuntime` (set above) before deciding | ||
| 1217 | // whether to rebuild, and it only gets the right answer because | ||
| 1218 | // `initialized` is already `false` by the time that await resolves — | ||
| 1219 | // otherwise a request that interleaved during the closes below would | ||
| 1220 | // still see `initialized === true`, take `init`'s early-return path, | ||
| 1221 | // and run against a server/engine that are mid-teardown or already | ||
| 1222 | // gone. | ||
| 1223 | self.initialized = false; | ||
| 901 | yield* self.releaseAllPendingApprovalLeases(); | 1224 | yield* self.releaseAllPendingApprovalLeases(); |
| 902 | if (options.closeStreams ?? true) { | 1225 | if (options.closeStreams ?? true) { |
| 903 | yield* Effect.sync(() => self.closeActiveStreams()); | 1226 | yield* Effect.sync(() => self.closeActiveStreams()); |
| @@ -905,21 +1228,58 @@ export abstract class McpAgentSessionDOBase< · 1 unmodified line | |||
| 905 | if (self.server) { | 1228 | if (self.server) { |
| 906 | const server = self.server; | 1229 | const server = self.server; |
| 907 | delete (self as { server?: McpServer }).server; | 1230 | delete (self as { server?: McpServer }).server; |
| 908 | yield* Effect.promise(() => server.close()).pipe(Effect.ignore); | 1231 | // `tryPromise`, not `promise`: a rejected close must land in the error |
| 1232 | // channel where `ignore` absorbs it. With `Effect.promise` a rejection | ||
| 1233 | // becomes a defect, which `ignore` does NOT absorb — the teardown | ||
| 1234 | // would stop here while the `ensuring` below still resolved | ||
| 1235 | // `disposingRuntime`, telling a waiting `init` the resources were | ||
| 1236 | // released when the steps after this one never ran. | ||
| 1237 | yield* Effect.tryPromise({ | ||
| 1238 | try: () => server.close(), | ||
| 1239 | catch: (cause: unknown) => cause, | ||
| 1240 | }).pipe(Effect.ignore); | ||
| 909 | } | 1241 | } |
| 910 | Reflect.set(self, "_transport", undefined); | 1242 | Reflect.set(self, "_transport", undefined); |
| 911 | self.engine = null; | 1243 | self.engine = null; |
| 912 | if (self.dbHandle) { | 1244 | if (self.dbHandle) { |
| 913 | const dbHandle = self.dbHandle; | 1245 | const dbHandle = self.dbHandle; |
| 914 | self.dbHandle = null; | 1246 | self.dbHandle = null; |
| 915 | yield* Effect.promise(() => Promise.resolve(dbHandle.end())).pipe(Effect.ignore); | 1247 | // Same `tryPromise` reasoning as the server close above. |
| 1248 | yield* Effect.tryPromise({ | ||
| 1249 | try: () => Promise.resolve(dbHandle.end()), | ||
| 1250 | catch: (cause: unknown) => cause, | ||
| 1251 | }).pipe(Effect.ignore); | ||
| 916 | } | 1252 | } |
| 917 | self.initialized = false; | ||
| 918 | if (self.countedAsResident) { | 1253 | if (self.countedAsResident) { |
| 919 | self.countedAsResident = false; | 1254 | self.countedAsResident = false; |
| 920 | releaseResidentRuntime(); | 1255 | releaseResidentRuntime(); |
| 1256 | // Pairs with `registerResidentSession` in `init`. Idempotent, and | ||
| 1257 | // gated on the same flag that guards the counter release, so a | ||
| 1258 | // `closeRuntime` that runs on a path where nothing was ever built | ||
| 1259 | // never removes an entry it did not add. | ||
| 1260 | releaseResidentSession(self.sessionIdForTelemetry()); | ||
| 921 | } | 1261 | } |
| 922 | }); | 1262 | }).pipe( |
| 1263 | // Uninterruptible once teardown starts, so it always runs to | ||
| 1264 | // completion. The `Effect.ensuring` below resolves `disposingRuntime` | ||
| 1265 | // unconditionally — a waiting `init` treats that resolution as "the | ||
| 1266 | // resources are actually released" and proceeds to rebuild — which is | ||
| 1267 | // only correct if nothing here can be interrupted or half-run partway | ||
| 1268 | // through. Every step above is already bounded and safe to run | ||
| 1269 | // uninterruptibly: the two closes route rejections into the error | ||
| 1270 | // channel via `tryPromise` and `ignore` them (a bare `Effect.promise` | ||
| 1271 | // would turn a rejection into a defect `ignore` cannot absorb), and | ||
| 1272 | // `releaseAllPendingApprovalLeases` ignores its own failures | ||
| 1273 | // (`deleteExecutionOwnerEntry` is `Effect.ignore`d too), so nothing | ||
| 1274 | // here can defect either. | ||
| 1275 | Effect.uninterruptible, | ||
| 1276 | Effect.ensuring( | ||
| 1277 | Effect.sync(() => { | ||
| 1278 | if (self.disposingRuntime === disposal) self.disposingRuntime = null; | ||
| 1279 | resolveDisposal(); | ||
| 1280 | }), | ||
| 1281 | ), | ||
| 1282 | ); | ||
| 923 | } | 1283 | } |
| 924 | 1284 | ||
| 925 | private ensureRuntimeForApproval(): Effect.Effect<boolean> { | 1285 | private ensureRuntimeForApproval(): Effect.Effect<boolean> { |
| @@ -978,6 +1338,16 @@ export abstract class McpAgentSessionDOBase< · 52 unmodified lines | |||
| 978 | } | 1338 | } |
| 979 | 1339 | ||
| 980 | async init(): Promise<void> { | 1340 | async init(): Promise<void> { |
| 1341 | if (this.disposingRuntime) { | ||
| 1342 | // A disposal (idle alarm, cap eviction) is mid-teardown for this | ||
| 1343 | // session's own Durable Object instance. `closeRuntime` flips | ||
| 1344 | // `initialized` to `false` before its first async close specifically so | ||
| 1345 | // this does not race it: wait for the disposal to actually finish — | ||
| 1346 | // server closed, engine and db handle released, residency counters | ||
| 1347 | // updated — before deciding whether a rebuild is even needed, instead | ||
| 1348 | // of starting one against half-torn-down state. | ||
| 1349 | await this.disposingRuntime; | ||
| 1350 | } | ||
| 981 | if (this.initialized) return; | 1351 | if (this.initialized) return; |
| 982 | const props = isSessionProps(this.props) ? this.props : null; | 1352 | const props = isSessionProps(this.props) ? this.props : null; |
| 983 | if (!props) { | 1353 | if (!props) { |
| @@ -988,8 +1358,26 @@ export abstract class McpAgentSessionDOBase< · 4 unmodified lines | |||
| 988 | const program = Effect.gen(function* () { | 1358 | const program = Effect.gen(function* () { |
| 989 | yield* self.prepareErrorCaptureScope(); | 1359 | yield* self.prepareErrorCaptureScope(); |
| 990 | const sessionMeta = yield* self.resolveAndStoreSessionMeta(props.session); | 1360 | const sessionMeta = yield* self.resolveAndStoreSessionMeta(props.session); |
| 991 | const dbHandle = yield* self.openSessionDbHandle(); | 1361 | // Before building anything that will itself occupy isolate memory |
| 992 | const { mcpServer, engine } = yield* self.buildRuntime(sessionMeta, dbHandle); | 1362 | // (a live db handle, the engine, the built tool catalog), make room if |
| 1363 | // this isolate is already at its resident-runtime cap. Deliberately | ||
| 1364 | // BEFORE `openSessionDbHandle`, not just before `buildRuntime`: the db | ||
| 1365 | // handle is one of the three things a resident runtime holds. | ||
| 1366 | yield* self.evictForCapIfNeeded(); | ||
| 1367 | // The in-flight cold-build reservation `evictForCapIfNeeded` just took | ||
| 1368 | // is released exactly once — success, failure, or interrupt — by the | ||
| 1369 | // single `Effect.ensuring` wrapped around this ENTIRE program below, | ||
| 1370 | // not a narrower one scoped to just this build block. A narrower wrap | ||
| 1371 | // still leaves a gap between the reservation being taken above and the | ||
| 1372 | // wrap starting here: an interrupt landing in exactly that gap would | ||
| 1373 | // leak the reservation forever (the module counter only drifts up), | ||
| 1374 | // which is why the release is anchored to the same scope as the | ||
| 1375 | // acquisition instead. | ||
| 1376 | const { dbHandle, mcpServer, engine } = yield* Effect.gen(function* () { | ||
| 1377 | const dbHandle = yield* self.openSessionDbHandle(); | ||
| 1378 | const { mcpServer, engine } = yield* self.buildRuntime(sessionMeta, dbHandle); | ||
| 1379 | return { dbHandle, mcpServer, engine }; | ||
| 1380 | }); | ||
| 993 | self.dbHandle = dbHandle; | 1381 | self.dbHandle = dbHandle; |
| 994 | self.server = mcpServer; | 1382 | self.server = mcpServer; |
| 995 | self.engine = engine; | 1383 | self.engine = engine; |
| @@ -997,6 +1385,26 @@ export abstract class McpAgentSessionDOBase< · 1 unmodified line | |||
| 997 | if (!self.countedAsResident) { | 1385 | if (!self.countedAsResident) { |
| 998 | self.countedAsResident = true; | 1386 | self.countedAsResident = true; |
| 999 | acquireResidentRuntime(); | 1387 | acquireResidentRuntime(); |
| 1388 | // Only a host that can route an eviction request back to THIS | ||
| 1389 | // session's own Durable Object instance (see `requestSelfEviction`) | ||
| 1390 | // registers as a candidate at all. A host that cannot is still | ||
| 1391 | // counted in the gauge above — cap-overflow tracking stays accurate — | ||
| 1392 | // it is just never picked, degrading to observational rather than | ||
| 1393 | // running teardown in the wrong context. | ||
| 1394 | if (self.supportsCapEviction()) { | ||
| 1395 | // Paired with `releaseResidentSession` in `closeRuntime`. Registered | ||
| 1396 | // as soon as this session counts as resident, so a cap check running | ||
| 1397 | // in another session's `init` moments later already sees it as a | ||
| 1398 | // candidate. `markActivity` below immediately corrects the initial | ||
| 1399 | // timestamp via `touchResidentSession`, so `Date.now()` here only | ||
| 1400 | // needs to be a safe placeholder, not the true last-activity time. | ||
| 1401 | registerResidentSession({ | ||
| 1402 | sessionId: self.sessionIdForTelemetry(), | ||
| 1403 | lastActivityMs: Date.now(), | ||
| 1404 | canEvict: () => self.canEvictResidentRuntime(), | ||
| 1405 | dispose: () => self.requestSelfEviction(), | ||
| 1406 | }); | ||
| 1407 | } | ||
| 1000 | } | 1408 | } |
| 1001 | // The gauge on the way up. Paired with the same attributes on | 1409 | // The gauge on the way up. Paired with the same attributes on |
| 1002 | // `mcp.session.idle_runtime_dispose`, this is what shows whether idle | 1410 | // `mcp.session.idle_runtime_dispose`, this is what shows whether idle |
| @@ -1010,6 +1418,13 @@ export abstract class McpAgentSessionDOBase< · 7 unmodified lines | |||
| 1010 | .bestEffortBookkeeping("init.mark_activity", () => self.markActivity()) | 1418 | .bestEffortBookkeeping("init.mark_activity", () => self.markActivity()) |
| 1011 | .pipe(Effect.withSpan("McpSessionDO.markActivity")); | 1419 | .pipe(Effect.withSpan("McpSessionDO.markActivity")); |
| 1012 | }).pipe( | 1420 | }).pipe( |
| 1421 | // Covers the ENTIRE program above, not just the build block: anything | ||
| 1422 | // from `evictForCapIfNeeded`'s admission onward that ends this effect — | ||
| 1423 | // success, failure, or interrupt — releases the cold-build reservation | ||
| 1424 | // exactly once (`releaseColdBuildSlotIfReserved` is idempotent). Wrapping | ||
| 1425 | // only the build block would leave the gap between admission and the | ||
| 1426 | // build starting uncovered. | ||
| 1427 | Effect.ensuring(self.releaseColdBuildSlotIfReserved()), | ||
| 1013 | // ONE capture owner for an init defect. `init` can only reject its | 1428 | // ONE capture owner for an init defect. `init` can only reject its |
| 1014 | // Promise, and the host's DO-level error instrumentation captures that | 1429 | // Promise, and the host's DO-level error instrumentation captures that |
| 1015 | // rejection too — so the DO claims the cause below and the host drops its | 1430 | // rejection too — so the DO claims the cause below and the host drops its |
| @@ -1047,12 +1462,53 @@ export abstract class McpAgentSessionDOBase< · 31 unmodified lines | |||
| 1047 | }), | 1462 | }), |
| 1048 | ); | 1463 | ); |
| 1049 | const traced = this.withTelemetry(program, props?.propagation); | 1464 | const traced = this.withTelemetry(program, props?.propagation); |
| 1465 | // See `initAbortController`'s doc comment: nothing in production aborts | ||
| 1466 | // this signal today, but wiring it through `runPromise` gives a unit test | ||
| 1467 | // a real handle to interrupt this exact fiber deterministically instead | ||
| 1468 | // of only ever exercising the failure/defect path. | ||
| 1469 | const abortController = new AbortController(); | ||
| 1470 | self.initAbortController = abortController; | ||
| 1050 | return Effect.runPromise( | 1471 | return Effect.runPromise( |
| 1051 | traced.pipe( | 1472 | traced.pipe( |
| 1052 | // oxlint-disable-next-line executor/no-effect-escape-hatch -- boundary: Durable Object init method can only reject its Promise | 1473 | // oxlint-disable-next-line executor/no-effect-escape-hatch -- boundary: Durable Object init method can only reject its Promise |
| 1053 | Effect.orDie, | 1474 | Effect.orDie, |
| 1054 | (effect) => self.withSpanFlush(effect), | 1475 | (effect) => self.withSpanFlush(effect), |
| 1055 | ), | 1476 | ), |
| 1477 | { signal: abortController.signal }, | ||
| 1478 | ).finally(() => { | ||
| 1479 | // Only clear the field if it is still THIS call's controller — an | ||
| 1480 | // overlapping later `init` call may already have installed its own by | ||
| 1481 | // the time this one settles, and clobbering that with `null` would | ||
| 1482 | // leave a unit test unable to reach it. | ||
| 1483 | if (self.initAbortController === abortController) self.initAbortController = null; | ||
| 1484 | }); | ||
| 1485 | } | ||
| 1486 | |||
| 1487 | /** | ||
| 1488 | * The candidate side of cap eviction. Called only through THIS session's own | ||
| 1489 | * Durable Object stub — see `requestSelfEviction` — never invoked directly | ||
| 1490 | * on an in-process reference by another session, which is the whole point: | ||
| 1491 | * routing the call through the stub gives it a correctly-scoped IoContext | ||
| 1492 | * for `evictResidentRuntimeForCap`'s teardown to run in. | ||
| 1493 | * | ||
| 1494 | * Safe to call more than once, including two overlapping calls (two | ||
| 1495 | * different sessions' `init`s both having picked this one as their LRU | ||
| 1496 | * candidate before either request lands): `evictResidentRuntimeForCap` | ||
| 1497 | * re-checks liveness every time, and `closeRuntime` is idempotent, so a | ||
| 1498 | * repeat call either finds the runtime already gone (a no-op) or finds it | ||
| 1499 | * newly busy again and leaves it alone. | ||
| 1500 | */ | ||
| 1501 | async requestCapEviction(): Promise<void> { | ||
| 1502 | const self = this; | ||
| 1503 | return Effect.runPromise( | ||
| 1504 | Effect.gen(function* () { | ||
| 1505 | yield* self.prepareErrorCaptureScope(); | ||
| 1506 | yield* Effect.promise(() => self.evictResidentRuntimeForCap()); | ||
| 1507 | }).pipe( | ||
| 1508 | Effect.withSpan("McpSessionDO.requestCapEviction"), | ||
| 1509 | // oxlint-disable-next-line executor/no-effect-escape-hatch -- boundary: DO RPC exposes Promise results | ||
| 1510 | Effect.orDie, | ||
| 1511 | ), | ||
| 1056 | ); | 1512 | ); |
| 1057 | } | 1513 | } |
| 1058 | 1514 | ||
| @@ -1299,7 +1755,12 @@ export abstract class McpAgentSessionDOBase< · 240 unmodified lines | |||
| 1299 | return; | 1755 | return; |
| 1300 | } | 1756 | } |
| 1301 | 1757 | ||
| 1302 | await this.disposeIdleRuntime({ idleMs, pausedExecutionCount, activeStreamCount }); | 1758 | await this.disposeIdleRuntime({ |
| 1759 | idleMs, | ||
| 1760 | pausedExecutionCount, | ||
| 1761 | activeStreamCount, | ||
| 1762 | reason: "idle", | ||
| 1763 | }); | ||
| 1303 | } | 1764 | } |
| 1304 | 1765 | ||
| 1305 | private validateApprovalIdentity( | 1766 | private validateApprovalIdentity( |
| @@ -48,6 +48,189 @@ export const resetResidentRuntimeCountForTest = (): void => { · 47 unmodified lines | |||
| 48 | peakResidentRuntimeCount = 0; | 48 | peakResidentRuntimeCount = 0; |
| 49 | }; | 49 | }; |
| 50 | 50 | ||
| 51 | /** | ||
| 52 | * How many cold builds are currently in flight in this isolate — reserved at | ||
| 53 | * admission, the same moment `evictForCapIfNeeded` runs, and released exactly | ||
| 54 | * once the reserving session either becomes counted-as-resident | ||
| 55 | * (`acquireResidentRuntime`) or its build fails or is interrupted. | ||
| 56 | * | ||
| 57 | * `residentRuntimeCount` only moves once a build actually finishes, so N | ||
| 58 | * overlapping cold inits at the cap each read it, see themselves still under | ||
| 59 | * the cap, and none of them evicts anything — residency then blows past the | ||
| 60 | * cap by N once every build lands. This counter closes that gap: the cap | ||
| 61 | * check becomes `currentResidentRuntimeCount() + currentInFlightColdBuildCount()`, | ||
| 62 | * so a second (or third, ...) concurrent admission sees the first's | ||
| 63 | * reservation and triggers eviction instead of also passing the check for | ||
| 64 | * free. | ||
| 65 | */ | ||
| 66 | let inFlightColdBuildCount = 0; | ||
| 67 | |||
| 68 | /** Reserve a cold-build slot at admission. Paired with exactly one of | ||
| 69 | * {@link releaseColdBuildSlot} per reservation — see the doc comment above. */ | ||
| 70 | export const reserveColdBuildSlot = (): number => { | ||
| 71 | inFlightColdBuildCount += 1; | ||
| 72 | return inFlightColdBuildCount; | ||
| 73 | }; | ||
| 74 | |||
| 75 | /** Release a previously-reserved cold-build slot. Floors at zero so a stray | ||
| 76 | * extra release (there should never be one; callers gate this on their own | ||
| 77 | * per-init flag) cannot drive the counter negative. */ | ||
| 78 | export const releaseColdBuildSlot = (): number => { | ||
| 79 | inFlightColdBuildCount = Math.max(0, inFlightColdBuildCount - 1); | ||
| 80 | return inFlightColdBuildCount; | ||
| 81 | }; | ||
| 82 | |||
| 83 | export const currentInFlightColdBuildCount = (): number => inFlightColdBuildCount; | ||
| 84 | |||
| 85 | /** Test-only: isolate-scoped module state outlives a single test case. */ | ||
| 86 | export const resetInFlightColdBuildCountForTest = (): void => { | ||
| 87 | inFlightColdBuildCount = 0; | ||
| 88 | }; | ||
| 89 | |||
| 90 | /** | ||
| 91 | * The isolate-wide ceiling on resident session runtimes, past which `init` | ||
| 92 | * evicts the least-recently-active evictable session to make room instead of | ||
| 93 | * letting residency climb unbounded. | ||
| 94 | * | ||
| 95 | * Production isolates have been observed OOM-killing every co-resident | ||
| 96 | * session somewhere between roughly 10 and 64 residents, and isolate memory | ||
| 97 | * is already sitting at the platform's 128MB limit at the fleet-wide median | ||
| 98 | * before this mechanism runs at all — there is no headroom left to discover | ||
| 99 | * the ceiling empirically per-isolate. 32 sits inside that failure band, so it | ||
| 100 | * bounds the COUNT-driven failure mode (an isolate cannot pile up unbounded | ||
| 101 | * residents no matter how many quiet-but-connected sessions land on it) while | ||
| 102 | * per-session footprint reduction, a separate effort, addresses the rest. | ||
| 103 | * It is a soft cap: an isolate that cannot find anything evictable is still | ||
| 104 | * allowed past it (see `pickEvictionCandidate`) rather than failing an init. | ||
| 105 | */ | ||
| 106 | export const RESIDENT_RUNTIME_SOFT_CAP = 32; | ||
| 107 | |||
| 108 | /** | ||
| 109 | * One resident session runtime, as far as the isolate-wide eviction registry | ||
| 110 | * needs to know about it: when it was last active, whether it is safe to tear | ||
| 111 | * down right now, and how to ask it to tear itself down. | ||
| 112 | * | ||
| 113 | * `canEvict` is consulted by `pickEvictionCandidate` to choose AMONG entries | ||
| 114 | * and is expected to be cheap and synchronous — it does not have to be the | ||
| 115 | * final word. `dispose` sends this session an eviction REQUEST — it does not | ||
| 116 | * run this session's own teardown itself. The owning Durable Object instance | ||
| 117 | * is the only thing allowed to tear down its own runtime (see | ||
| 118 | * `requestSelfEviction` on the base class for why: teardown does I/O bound to | ||
| 119 | * this session's own request context, and a caller running it directly would | ||
| 120 | * be running that I/O under ITS OWN context instead). `dispose` is expected to | ||
| 121 | * re-check liveness itself before actually releasing anything, using whatever | ||
| 122 | * signals it has (including ones `canEvict` could not afford to read), so a | ||
| 123 | * candidate that became active between the pick and the dispose call is left | ||
| 124 | * alone rather than torn down. It resolving does not guarantee the candidate | ||
| 125 | * was actually torn down — only that the request was sent — so callers that | ||
| 126 | * fire it in the background (see `evictForCapIfNeeded`) call | ||
| 127 | * `markEvictionRequested` rather than assuming success. | ||
| 128 | */ | ||
| 129 | export type ResidentSessionEntry = { | ||
| 130 | readonly sessionId: string; | ||
| 131 | lastActivityMs: number; | ||
| 132 | evictionRequestedAt?: number; | ||
| 133 | readonly canEvict: () => boolean; | ||
| 134 | readonly dispose: (reason: "cap") => Promise<void>; | ||
| 135 | }; | ||
| 136 | |||
| 137 | /** | ||
| 138 | * Every session runtime currently resident in THIS isolate, keyed by session | ||
| 139 | * id. Dependency-free by design: this module is a leaf that both the session | ||
| 140 | * Durable Object and its tests import directly, and it must never need to | ||
| 141 | * know what an Effect, a Durable Object, or a storage API is. | ||
| 142 | * | ||
| 143 | * An entry lives here for exactly as long as its runtime is resident. It is | ||
| 144 | * added once, at the same moment `acquireResidentRuntime` counts it, and | ||
| 145 | * removed once, at the same moment `releaseResidentRuntime` releases it — the | ||
| 146 | * two are meant to move together, though this module does not enforce that; | ||
| 147 | * see `closeRuntime` on the Durable Object for where they are actually paired. | ||
| 148 | */ | ||
| 149 | const residentSessions = new Map<string, ResidentSessionEntry>(); | ||
| 150 | |||
| 151 | /** Start tracking a newly-resident session runtime for isolate-wide eviction. */ | ||
| 152 | export const registerResidentSession = (entry: ResidentSessionEntry): void => { | ||
| 153 | residentSessions.set(entry.sessionId, entry); | ||
| 154 | }; | ||
| 155 | |||
| 156 | /** Record that a resident session just did something, so it sorts last for eviction. */ | ||
| 157 | export const touchResidentSession = (sessionId: string, lastActivityMs = Date.now()): void => { | ||
| 158 | const entry = residentSessions.get(sessionId); | ||
| 159 | if (entry) entry.lastActivityMs = lastActivityMs; | ||
| 160 | }; | ||
| 161 | |||
| 162 | /** | ||
| 163 | * Stop tracking a session's runtime because it is no longer resident. | ||
| 164 | * | ||
| 165 | * Idempotent on purpose: `Map.delete` on an absent key is already a no-op, so | ||
| 166 | * this can be called from every path that might end a session's residency — | ||
| 167 | * a clean disposal, a failed one, a repeat call — without any caller having to | ||
| 168 | * first ask whether the entry is still there. | ||
| 169 | */ | ||
| 170 | export const releaseResidentSession = (sessionId: string): void => { | ||
| 171 | residentSessions.delete(sessionId); | ||
| 172 | }; | ||
| 173 | |||
| 174 | /** | ||
| 175 | * How long a resident session stays skipped by `pickEvictionCandidate` after | ||
| 176 | * an eviction request was sent to it. The request is fire-and-forget (see | ||
| 177 | * `evictForCapIfNeeded` on the base class) — the sender never learns whether | ||
| 178 | * it succeeded, failed, or is still in flight — so this grace window is the | ||
| 179 | * only thing standing between one stuck or repeatedly-failing candidate and | ||
| 180 | * it squatting the LRU pick forever, re-requested by every subsequent init. | ||
| 181 | * Long enough that a healthy eviction (self-request, own teardown, registry | ||
| 182 | * removal) always finishes well inside it; short enough that a genuinely | ||
| 183 | * stuck candidate stops blocking the LRU pick soon after. | ||
| 184 | */ | ||
| 185 | export const EVICTION_REQUEST_GRACE_MS = 15_000; | ||
| 186 | |||
| 187 | /** | ||
| 188 | * Record that an eviction request was just sent to this entry, so | ||
| 189 | * `pickEvictionCandidate` skips it until the grace period elapses. A no-op | ||
| 190 | * when the entry is already gone (it was evicted, or removed some other way) | ||
| 191 | * — nothing left to mark, and nothing wrong with that. | ||
| 192 | */ | ||
| 193 | export const markEvictionRequested = (sessionId: string, requestedAtMs = Date.now()): void => { | ||
| 194 | const entry = residentSessions.get(sessionId); | ||
| 195 | if (entry) entry.evictionRequestedAt = requestedAtMs; | ||
| 196 | }; | ||
| 197 | |||
| 198 | /** | ||
| 199 | * The least-recently-active resident session that is currently safe to evict, | ||
| 200 | * or `undefined` when every resident is streaming, paused, recently asked to | ||
| 201 | * evict itself already, or otherwise ineligible right now. | ||
| 202 | * | ||
| 203 | * `undefined` is a legitimate, expected answer — it means the isolate is over | ||
| 204 | * its soft cap but everything resident is doing real work (or already has an | ||
| 205 | * eviction request outstanding), and the caller must let the new session | ||
| 206 | * build anyway rather than block or fail on it. | ||
| 207 | */ | ||
| 208 | export const pickEvictionCandidate = (nowMs = Date.now()): ResidentSessionEntry | undefined => { | ||
| 209 | let candidate: ResidentSessionEntry | undefined; | ||
| 210 | for (const entry of residentSessions.values()) { | ||
| 211 | if (!entry.canEvict()) continue; | ||
| 212 | if ( | ||
| 213 | entry.evictionRequestedAt !== undefined && | ||
| 214 | nowMs - entry.evictionRequestedAt < EVICTION_REQUEST_GRACE_MS | ||
| 215 | ) { | ||
| 216 | continue; | ||
| 217 | } | ||
| 218 | if (!candidate || entry.lastActivityMs < candidate.lastActivityMs) { | ||
| 219 | candidate = entry; | ||
| 220 | } | ||
| 221 | } | ||
| 222 | return candidate; | ||
| 223 | }; | ||
| 224 | |||
| 225 | /** Test-only: isolate-scoped module state outlives a single test case. */ | ||
| 226 | export const resetResidentSessionRegistryForTest = (): void => { | ||
| 227 | residentSessions.clear(); | ||
| 228 | }; | ||
| 229 | |||
| 230 | /** Test-only: read without mutating, to assert on registry membership directly. */ | ||
| 231 | export const residentSessionIdsForTest = (): ReadonlyArray<string> => | ||
| 232 | Array.from(residentSessions.keys()); | ||
| 233 | |||
| 51 | type MemoryCapablePerformance = { | 234 | type MemoryCapablePerformance = { |
| 52 | readonly memory?: { | 235 | readonly memory?: { |
| 53 | readonly usedJSHeapSize?: unknown; | 236 | readonly usedJSHeapSize?: unknown; |
| @@ -36,6 +36,14 @@ export interface McpSessionStub { · 35 unmodified lines | |||
| 36 | response: ResumeResponse, | 36 | response: ResumeResponse, |
| 37 | incoming?: IncomingTraceHeaders, | 37 | incoming?: IncomingTraceHeaders, |
| 38 | ) => Promise<McpSessionModelResumeResult>; | 38 | ) => Promise<McpSessionModelResumeResult>; |
| 39 | /** | ||
| 40 | * Ask this session's OWN Durable Object instance to tear down its resident | ||
| 41 | * runtime because the isolate is over its cap. Routed through the stub | ||
| 42 | * rather than called on an in-process reference so the teardown runs in | ||
| 43 | * this session's own IoContext — see `requestSelfEviction` on the base | ||
| 44 | * class for why that boundary matters. | ||
| 45 | */ | ||
| 46 | readonly requestCapEviction: () => Promise<void>; | ||
| 39 | } | 47 | } |
| 40 | 48 | ||
| 41 | export const mcpSessionStub = <Id>( | 49 | export const mcpSessionStub = <Id>( |