Negative-cache dead MCP session ids in the workers (#1627)

Author: RhysSullivan

apps/cloud/src/mcp/agent-handler.test.ts

@@ -0,0 +1,146 @@
+import { beforeEach, describe, expect, it, vi } from "@effect/vitest";
+import { Effect, Layer } from "effect";
+
+import {
+  authenticated,
+  McpAuthProvider,
+  unauthorized,
+  type Principal,
+} from "@executor-js/host-mcp";
+
+import { cloudDeadSessionCacheForTest, makeCloudMcpAgentHandler } from "./agent-handler";
+
+const principal: Principal = {
+  accountId: "acct_test",
+  organizationId: "org_test",
+  organizationName: "Test Org",
+  email: "test@example.com",
+  name: "Test",
+  avatarUrl: null,
+  roles: ["member"],
+};
+
+const AuthProviderLive = Layer.succeed(McpAuthProvider)({
+  discoveryRoutes: [],
+  resourceMetadataUrl: (request) => new URL("/.well-known/mcp", request.url).toString(),
+  authenticate: (request) =>
+    Effect.succeed(
+      request.headers.has("authorization") ? authenticated(principal) : unauthorized(),
+    ),
+});
+
+const requestFor = (method: "GET" | "POST", sessionId: string, authenticated = true): Request =>
+  new Request("https://executor.test/mcp", {
+    method,
+    headers: {
+      ...(authenticated ? { authorization: "Bearer test" } : {}),
+      "mcp-session-id": sessionId,
+      ...(method === "POST" ? { "content-type": "application/json" } : {}),
+    },
+    ...(method === "POST"
+      ? {
+          body: JSON.stringify({ jsonrpc: "2.0", id: 1, method: "tools/list", params: {} }),
+        }
+      : {}),
+  });
+
+const makeHarness = (owner: "not_found" | "terminated" = "not_found") => {
+  const ownerChecks = { count: 0 };
+  const traces = { count: 0 };
+  const stub = {
+    validateMcpSessionOwner: async () => {
+      ownerChecks.count += 1;
+      return owner;
+    },
+  };
+  const namespace = {
+    idFromString: (sessionId: string) => sessionId,
+    get: () => stub,
+  };
+  // oxlint-disable-next-line executor/no-double-cast -- test boundary: the handler only reads the MCP_SESSION namespace in these legacy dead-session cases
+  const env = { MCP_SESSION: namespace } as unknown as Env;
+  // oxlint-disable-next-line executor/no-double-cast -- test boundary: no ExecutionContext capability is used before the dead-session response
+  const ctx = {} as unknown as ExecutionContext;
+  const handler = makeCloudMcpAgentHandler({
+    authProvider: AuthProviderLive,
+    makeModernServerBuilder: () => ({ build: () => Effect.die("unused modern builder") }),
+    traceRequest: async (request, _env, _ctx, handle) => {
+      traces.count += 1;
+      return handle(request);
+    },
+  });
+  return { ctx, env, handler, ownerChecks, traces };
+};
+
+describe("cloud MCP dead-session negative cache", () => {
+  beforeEach(() => {
+    vi.useRealTimers();
+    cloudDeadSessionCacheForTest.clear();
+  });
+
+  it.each([
+    { method: "GET" as const, status: 405 },
+    { method: "POST" as const, status: 404 },
+  ])("serves a repeated $method without another DO lookup", async ({ method, status }) => {
+    const { ctx, env, handler, ownerChecks, traces } = makeHarness();
+    const sessionId = `dead-${method.toLowerCase()}`;
+
+    const first = await handler(requestFor(method, sessionId), env, ctx);
+    const second = await handler(requestFor(method, sessionId), env, ctx);
+
+    expect(first.status).toBe(status);
+    expect(second.status).toBe(status);
+    expect(ownerChecks.count).toBe(1);
+    expect(traces.count).toBe(1);
+  });
+
+  it("keeps authentication ahead of cached session existence", async () => {
+    const { ctx, env, handler, ownerChecks, traces } = makeHarness();
+    const sessionId = "dead-auth";
+
+    expect((await handler(requestFor("GET", sessionId), env, ctx)).status).toBe(405);
+    const unauthenticated = await handler(requestFor("GET", sessionId, false), env, ctx);
+
+    expect(unauthenticated.status).toBe(401);
+    expect(ownerChecks.count).toBe(1);
+    expect(traces.count).toBe(1);
+  });
+
+  it("preserves the terminated-session response on a cached hit", async () => {
+    const { ctx, env, handler, ownerChecks, traces } = makeHarness("terminated");
+    const sessionId = "dead-terminated";
+
+    const first = await handler(requestFor("POST", sessionId), env, ctx);
+    const second = await handler(requestFor("POST", sessionId), env, ctx);
+
+    expect(first.status).toBe(404);
+    expect(second.status).toBe(404);
+    expect(await second.text()).toBe(await first.text());
+    expect(ownerChecks.count).toBe(1);
+    expect(traces.count).toBe(1);
+  });
+
+  it("consults the DO again after five minutes", async () => {
+    vi.useFakeTimers();
+    vi.setSystemTime(0);
+    const { ctx, env, handler, ownerChecks } = makeHarness();
+    const sessionId = "dead-expiry";
+
+    expect((await handler(requestFor("GET", sessionId), env, ctx)).status).toBe(405);
+    vi.advanceTimersByTime(5 * 60 * 1_000 + 1);
+    expect((await handler(requestFor("GET", sessionId), env, ctx)).status).toBe(405);
+
+    expect(ownerChecks.count).toBe(2);
+  });
+
+  it("evicts the oldest entry without exceeding 4,096 sessions", () => {
+    for (let index = 0; index <= 4_096; index += 1) {
+      cloudDeadSessionCacheForTest.remember(`dead-${index}`, 0);
+    }
+
+    expect(cloudDeadSessionCacheForTest.size()).toBe(4_096);
+    expect(cloudDeadSessionCacheForTest.has("dead-0", 1)).toBe(false);
+    expect(cloudDeadSessionCacheForTest.has("dead-1", 1)).toBe(true);
+    expect(cloudDeadSessionCacheForTest.has("dead-4096", 1)).toBe(true);
+  });
+});

apps/cloud/src/mcp/agent-handler.ts

@@ -1,5 +1,5 @@
 import * as OtelTracer from "@effect/opentelemetry/Tracer";
-import { Effect, Predicate } from "effect";
+import { Effect, Layer, Predicate } from "effect";
 
 import {
   McpAuthProvider,
@@ -8,6 +8,7 @@ import {
   defaultMcpResource,
   UNAVAILABLE_RETRY_AFTER_SECONDS,
   type AuthOutcome,
+  type McpModernServerBuilder,
   type McpResource,
 } from "@executor-js/host-mcp";
 import { requestBodyFromRequest } from "@executor-js/host-mcp/tool-server";
@@ -32,9 +33,53 @@ import { createMcpSessionStub, mcpSessionStub } from "@executor-js/cloudflare/mc
 import { wrapMcpSseResponse } from "../observability/memory-metrics";
 import { WorkerTelemetryLive } from "../observability/telemetry";
 import { cloudMcpAuth } from "./auth-provider";
-import { makeCloudModernMcpServerBuilder } from "./session-durable-object";
 import { parseTraceparent } from "./traceparent";
 
+const DEAD_SESSION_CACHE_TTL_MS = 5 * 60 * 1_000;
+const DEAD_SESSION_CACHE_MAX_ENTRIES = 4_096;
+const deadSessionExpiries = new Map<string, number>();
+const timedOutSessionIds = new Set<string>();
+
+type DeadSessionReason = "not_found" | "timed_out";
+
+const isDeadSessionCached = (sessionId: string, now = Date.now()): boolean => {
+  const expiry = deadSessionExpiries.get(sessionId);
+  if (expiry === undefined) return false;
+  if (expiry > now) return true;
+  deadSessionExpiries.delete(sessionId);
+  timedOutSessionIds.delete(sessionId);
+  return false;
+};
+
+const cacheDeadSession = (sessionId: string, reason: DeadSessionReason, now = Date.now()): void => {
+  deadSessionExpiries.delete(sessionId);
+  if (deadSessionExpiries.size >= DEAD_SESSION_CACHE_MAX_ENTRIES) {
+    const oldestSessionId = deadSessionExpiries.keys().next().value;
+    if (oldestSessionId !== undefined) {
+      deadSessionExpiries.delete(oldestSessionId);
+      timedOutSessionIds.delete(oldestSessionId);
+    }
+  }
+  deadSessionExpiries.set(sessionId, now + DEAD_SESSION_CACHE_TTL_MS);
+  if (reason === "timed_out") timedOutSessionIds.add(sessionId);
+  else timedOutSessionIds.delete(sessionId);
+};
+
+const cachedDeadSessionMessage = (sessionId: string): string =>
+  timedOutSessionIds.has(sessionId) ? "Session timed out, please reconnect" : "Session not found";
+
+/** Test-only access to reset and verify the isolate-local dead-session cache. */
+export const cloudDeadSessionCacheForTest = {
+  clear: (): void => {
+    deadSessionExpiries.clear();
+    timedOutSessionIds.clear();
+  },
+  remember: (sessionId: string, now?: number): void =>
+    cacheDeadSession(sessionId, "not_found", now),
+  has: isDeadSessionCached,
+  size: (): number => deadSessionExpiries.size,
+};
+
 const jsonRpcResponse = (
   status: number,
   code: number,
@@ -97,12 +142,12 @@ const renderAuthError = (
   });
 };
 
-const authenticate = (request: Request) =>
+const authenticate = (request: Request, authProvider: Layer.Layer<McpAuthProvider>) =>
   Effect.gen(function* () {
     const auth = yield* McpAuthProvider;
     const outcome = yield* auth.authenticate(request);
     return { auth, outcome };
-  }).pipe(Effect.provide(cloudMcpAuth));
+  }).pipe(Effect.provide(authProvider));
 
 // The earlier shared envelope ran the MCP auth path inside the Effect app, whose
 // HttpMiddleware provided the OTEL tracer — that is where the `mcp.request`
@@ -126,6 +171,21 @@ const runTraced = <A>(request: Request, program: Effect.Effect<A>): Promise<A> =
   );
 };
 
+type TraceCloudMcpRequest = (
+  request: Request,
+  env: Env,
+  ctx: ExecutionContext,
+  handle: (tracedRequest: Request) => Promise<Response>,
+) => Promise<Response>;
+
+interface CloudMcpAgentHandlerOptions {
+  readonly makeModernServerBuilder: (
+    session: McpSessionProps["session"],
+  ) => McpModernServerBuilder["Service"];
+  readonly authProvider?: Layer.Layer<McpAuthProvider>;
+  readonly traceRequest?: TraceCloudMcpRequest;
+}
+
 // The MCP resource the request targets. `server.ts` routes both the bare `/mcp`
 // and `/mcp/toolkits/<slug>` to this handler (`prepareMcpOrgScope` strips the org
 // selector but keeps the toolkit segment), so a session minted on a toolkit path
@@ -158,11 +218,14 @@ const propsForPrincipal = (
     };
   });
 
-export const makeCloudMcpAgentHandler = () => {
+/** Build the cloud worker's authenticated legacy/modern MCP request handler. */
+export const makeCloudMcpAgentHandler = (options: CloudMcpAgentHandlerOptions) => {
+  const authProvider = options.authProvider ?? cloudMcpAuth;
+  const traceRequest = options.traceRequest ?? ((request, _env, _ctx, handle) => handle(request));
   const modern = makeMcpModernRequestRouter();
   const ALLOWED_METHODS = new Set(["GET", "POST", "DELETE", "OPTIONS"]);
 
-  return async (request: Request, env: Env, ctx: ExecutionContext): Promise<Response> => {
+  const handle = async (request: Request, env: Env, ctx: ExecutionContext): Promise<Response> => {
     if (request.method === "OPTIONS") {
       return mcpCorsPreflightResponse(request.headers.get("access-control-request-headers"));
     }
@@ -173,7 +236,7 @@ export const makeCloudMcpAgentHandler = () => {
     }
     const sessionId = request.headers.get("mcp-session-id");
 
-    const { auth, outcome } = await runTraced(request, authenticate(request));
+    const { auth, outcome } = await runTraced(request, authenticate(request, authProvider));
     if (!Predicate.isTagged(outcome, "Authenticated")) {
       // Destroying a live session on auth grounds requires a POSITIVE
       // determination that access is genuinely gone — only `Forbidden` carries
@@ -218,7 +281,7 @@ export const makeCloudMcpAgentHandler = () => {
         resource,
         props,
         requestStateSigningKey: requireMcpRequestStateKey(env.MCP_REQUEST_STATE_KEY),
-        builder: makeCloudModernMcpServerBuilder(props.session),
+        builder: options.makeModernServerBuilder(props.session),
         sessions: env.MCP_SESSION,
         executionOwners: mcpExecutionOwnerDirectoryFromNamespace(env.MCP_EXECUTION_OWNER),
       });
@@ -238,18 +301,20 @@ export const makeCloudMcpAgentHandler = () => {
     if (sessionId && !existingSession) {
       return deadSessionResponse(request.method, "Session not found");
     }
-    if (existingSession) {
+    if (existingSession && sessionId) {
       const owner = await existingSession.validateMcpSessionOwner({
         accountId: outcome.principal.accountId,
         organizationId: outcome.principal.organizationId,
       });
       if (owner === "not_found") {
+        cacheDeadSession(sessionId, "not_found");
         return deadSessionResponse(request.method, "Session not found");
       }
       if (owner === "terminated") {
         // DELETE-condemned but the deferred destroy alarm hasn't wiped storage
         // yet. Same envelope as the post-destroy race below: the client must
         // treat the id as dead and reconnect.
+        cacheDeadSession(sessionId, "timed_out");
         return deadSessionResponse(request.method, "Session timed out, please reconnect");
       }
       if (owner === "forbidden") {
@@ -286,11 +351,28 @@ export const makeCloudMcpAgentHandler = () => {
       // client to be told to reconnect, matching a timed-out session).
       // oxlint-disable-next-line executor/no-unknown-error-message -- adapter boundary: the abort reason is a plain runtime Error whose message IS the signal
       if (Predicate.isError(error) && error.message === "destroyed") {
+        if (sessionId) cacheDeadSession(sessionId, "timed_out");
         return deadSessionResponse(request.method, "Session timed out, please reconnect");
       }
       // oxlint-disable-next-line executor/no-try-catch-or-throw -- adapter boundary: rethrow anything that isn't the condemned-DO abort to the Workers runtime unchanged
       throw error;
     }
     return withMcpResponseHeaders(wrapMcpSseResponse(request, env, response));
   };
+
+  return async (request: Request, env: Env, ctx: ExecutionContext): Promise<Response> => {
+    const sessionId = request.headers.get("mcp-session-id");
+    const cacheEligible = request.method !== "OPTIONS" && ALLOWED_METHODS.has(request.method);
+    if (cacheEligible && sessionId && isDeadSessionCached(sessionId)) {
+      return Effect.runPromise(
+        Effect.gen(function* () {
+          const { auth, outcome } = yield* authenticate(request, authProvider);
+          return Predicate.isTagged(outcome, "Authenticated")
+            ? deadSessionResponse(request.method, cachedDeadSessionMessage(sessionId))
+            : renderAuthError(auth, request, outcome);
+        }).pipe(Effect.withTracerEnabled(false)),
+      );
+    }
+    return traceRequest(request, env, ctx, (tracedRequest) => handle(tracedRequest, env, ctx));
+  };
 };