Unified · Split
| @@ -0,0 +1,146 @@ | |||
| 1 | import { beforeEach, describe, expect, it, vi } from "@effect/vitest"; | ||
| 2 | import { Effect, Layer } from "effect"; | ||
| 3 | |||
| 4 | import { | ||
| 5 | authenticated, | ||
| 6 | McpAuthProvider, | ||
| 7 | unauthorized, | ||
| 8 | type Principal, | ||
| 9 | } from "@executor-js/host-mcp"; | ||
| 10 | |||
| 11 | import { cloudDeadSessionCacheForTest, makeCloudMcpAgentHandler } from "./agent-handler"; | ||
| 12 | |||
| 13 | const principal: Principal = { | ||
| 14 | accountId: "acct_test", | ||
| 15 | organizationId: "org_test", | ||
| 16 | organizationName: "Test Org", | ||
| 17 | email: "test@example.com", | ||
| 18 | name: "Test", | ||
| 19 | avatarUrl: null, | ||
| 20 | roles: ["member"], | ||
| 21 | }; | ||
| 22 | |||
| 23 | const AuthProviderLive = Layer.succeed(McpAuthProvider)({ | ||
| 24 | discoveryRoutes: [], | ||
| 25 | resourceMetadataUrl: (request) => new URL("/.well-known/mcp", request.url).toString(), | ||
| 26 | authenticate: (request) => | ||
| 27 | Effect.succeed( | ||
| 28 | request.headers.has("authorization") ? authenticated(principal) : unauthorized(), | ||
| 29 | ), | ||
| 30 | }); | ||
| 31 | |||
| 32 | const requestFor = (method: "GET" | "POST", sessionId: string, authenticated = true): Request => | ||
| 33 | new Request("https://executor.test/mcp", { | ||
| 34 | method, | ||
| 35 | headers: { | ||
| 36 | ...(authenticated ? { authorization: "Bearer test" } : {}), | ||
| 37 | "mcp-session-id": sessionId, | ||
| 38 | ...(method === "POST" ? { "content-type": "application/json" } : {}), | ||
| 39 | }, | ||
| 40 | ...(method === "POST" | ||
| 41 | ? { | ||
| 42 | body: JSON.stringify({ jsonrpc: "2.0", id: 1, method: "tools/list", params: {} }), | ||
| 43 | } | ||
| 44 | : {}), | ||
| 45 | }); | ||
| 46 | |||
| 47 | const makeHarness = (owner: "not_found" | "terminated" = "not_found") => { | ||
| 48 | const ownerChecks = { count: 0 }; | ||
| 49 | const traces = { count: 0 }; | ||
| 50 | const stub = { | ||
| 51 | validateMcpSessionOwner: async () => { | ||
| 52 | ownerChecks.count += 1; | ||
| 53 | return owner; | ||
| 54 | }, | ||
| 55 | }; | ||
| 56 | const namespace = { | ||
| 57 | idFromString: (sessionId: string) => sessionId, | ||
| 58 | get: () => stub, | ||
| 59 | }; | ||
| 60 | // oxlint-disable-next-line executor/no-double-cast -- test boundary: the handler only reads the MCP_SESSION namespace in these legacy dead-session cases | ||
| 61 | const env = { MCP_SESSION: namespace } as unknown as Env; | ||
| 62 | // oxlint-disable-next-line executor/no-double-cast -- test boundary: no ExecutionContext capability is used before the dead-session response | ||
| 63 | const ctx = {} as unknown as ExecutionContext; | ||
| 64 | const handler = makeCloudMcpAgentHandler({ | ||
| 65 | authProvider: AuthProviderLive, | ||
| 66 | makeModernServerBuilder: () => ({ build: () => Effect.die("unused modern builder") }), | ||
| 67 | traceRequest: async (request, _env, _ctx, handle) => { | ||
| 68 | traces.count += 1; | ||
| 69 | return handle(request); | ||
| 70 | }, | ||
| 71 | }); | ||
| 72 | return { ctx, env, handler, ownerChecks, traces }; | ||
| 73 | }; | ||
| 74 | |||
| 75 | describe("cloud MCP dead-session negative cache", () => { | ||
| 76 | beforeEach(() => { | ||
| 77 | vi.useRealTimers(); | ||
| 78 | cloudDeadSessionCacheForTest.clear(); | ||
| 79 | }); | ||
| 80 | |||
| 81 | it.each([ | ||
| 82 | { method: "GET" as const, status: 405 }, | ||
| 83 | { method: "POST" as const, status: 404 }, | ||
| 84 | ])("serves a repeated $method without another DO lookup", async ({ method, status }) => { | ||
| 85 | const { ctx, env, handler, ownerChecks, traces } = makeHarness(); | ||
| 86 | const sessionId = `dead-${method.toLowerCase()}`; | ||
| 87 | |||
| 88 | const first = await handler(requestFor(method, sessionId), env, ctx); | ||
| 89 | const second = await handler(requestFor(method, sessionId), env, ctx); | ||
| 90 | |||
| 91 | expect(first.status).toBe(status); | ||
| 92 | expect(second.status).toBe(status); | ||
| 93 | expect(ownerChecks.count).toBe(1); | ||
| 94 | expect(traces.count).toBe(1); | ||
| 95 | }); | ||
| 96 | |||
| 97 | it("keeps authentication ahead of cached session existence", async () => { | ||
| 98 | const { ctx, env, handler, ownerChecks, traces } = makeHarness(); | ||
| 99 | const sessionId = "dead-auth"; | ||
| 100 | |||
| 101 | expect((await handler(requestFor("GET", sessionId), env, ctx)).status).toBe(405); | ||
| 102 | const unauthenticated = await handler(requestFor("GET", sessionId, false), env, ctx); | ||
| 103 | |||
| 104 | expect(unauthenticated.status).toBe(401); | ||
| 105 | expect(ownerChecks.count).toBe(1); | ||
| 106 | expect(traces.count).toBe(1); | ||
| 107 | }); | ||
| 108 | |||
| 109 | it("preserves the terminated-session response on a cached hit", async () => { | ||
| 110 | const { ctx, env, handler, ownerChecks, traces } = makeHarness("terminated"); | ||
| 111 | const sessionId = "dead-terminated"; | ||
| 112 | |||
| 113 | const first = await handler(requestFor("POST", sessionId), env, ctx); | ||
| 114 | const second = await handler(requestFor("POST", sessionId), env, ctx); | ||
| 115 | |||
| 116 | expect(first.status).toBe(404); | ||
| 117 | expect(second.status).toBe(404); | ||
| 118 | expect(await second.text()).toBe(await first.text()); | ||
| 119 | expect(ownerChecks.count).toBe(1); | ||
| 120 | expect(traces.count).toBe(1); | ||
| 121 | }); | ||
| 122 | |||
| 123 | it("consults the DO again after five minutes", async () => { | ||
| 124 | vi.useFakeTimers(); | ||
| 125 | vi.setSystemTime(0); | ||
| 126 | const { ctx, env, handler, ownerChecks } = makeHarness(); | ||
| 127 | const sessionId = "dead-expiry"; | ||
| 128 | |||
| 129 | expect((await handler(requestFor("GET", sessionId), env, ctx)).status).toBe(405); | ||
| 130 | vi.advanceTimersByTime(5 * 60 * 1_000 + 1); | ||
| 131 | expect((await handler(requestFor("GET", sessionId), env, ctx)).status).toBe(405); | ||
| 132 | |||
| 133 | expect(ownerChecks.count).toBe(2); | ||
| 134 | }); | ||
| 135 | |||
| 136 | it("evicts the oldest entry without exceeding 4,096 sessions", () => { | ||
| 137 | for (let index = 0; index <= 4_096; index += 1) { | ||
| 138 | cloudDeadSessionCacheForTest.remember(`dead-${index}`, 0); | ||
| 139 | } | ||
| 140 | |||
| 141 | expect(cloudDeadSessionCacheForTest.size()).toBe(4_096); | ||
| 142 | expect(cloudDeadSessionCacheForTest.has("dead-0", 1)).toBe(false); | ||
| 143 | expect(cloudDeadSessionCacheForTest.has("dead-1", 1)).toBe(true); | ||
| 144 | expect(cloudDeadSessionCacheForTest.has("dead-4096", 1)).toBe(true); | ||
| 145 | }); | ||
| 146 | }); | ||
| @@ -1,5 +1,5 @@ | |||
| 1 | import * as OtelTracer from "@effect/opentelemetry/Tracer"; | 1 | import * as OtelTracer from "@effect/opentelemetry/Tracer"; |
| 2 | import { Effect, Predicate } from "effect"; | 2 | import { Effect, Layer, Predicate } from "effect"; |
| 3 | 3 | ||
| 4 | import { | 4 | import { |
| 5 | McpAuthProvider, | 5 | McpAuthProvider, |
| @@ -8,6 +8,7 @@ import { · 2 unmodified lines | |||
| 8 | defaultMcpResource, | 8 | defaultMcpResource, |
| 9 | UNAVAILABLE_RETRY_AFTER_SECONDS, | 9 | UNAVAILABLE_RETRY_AFTER_SECONDS, |
| 10 | type AuthOutcome, | 10 | type AuthOutcome, |
| 11 | type McpModernServerBuilder, | ||
| 11 | type McpResource, | 12 | type McpResource, |
| 12 | } from "@executor-js/host-mcp"; | 13 | } from "@executor-js/host-mcp"; |
| 13 | import { requestBodyFromRequest } from "@executor-js/host-mcp/tool-server"; | 14 | import { requestBodyFromRequest } from "@executor-js/host-mcp/tool-server"; |
| @@ -32,9 +33,53 @@ import { createMcpSessionStub, mcpSessionStub } from "@executor-js/cloudflare/mc · 18 unmodified lines | |||
| 32 | import { wrapMcpSseResponse } from "../observability/memory-metrics"; | 33 | import { wrapMcpSseResponse } from "../observability/memory-metrics"; |
| 33 | import { WorkerTelemetryLive } from "../observability/telemetry"; | 34 | import { WorkerTelemetryLive } from "../observability/telemetry"; |
| 34 | import { cloudMcpAuth } from "./auth-provider"; | 35 | import { cloudMcpAuth } from "./auth-provider"; |
| 35 | import { makeCloudModernMcpServerBuilder } from "./session-durable-object"; | ||
| 36 | import { parseTraceparent } from "./traceparent"; | 36 | import { parseTraceparent } from "./traceparent"; |
| 37 | 37 | ||
| 38 | const DEAD_SESSION_CACHE_TTL_MS = 5 * 60 * 1_000; | ||
| 39 | const DEAD_SESSION_CACHE_MAX_ENTRIES = 4_096; | ||
| 40 | const deadSessionExpiries = new Map<string, number>(); | ||
| 41 | const timedOutSessionIds = new Set<string>(); | ||
| 42 | |||
| 43 | type DeadSessionReason = "not_found" | "timed_out"; | ||
| 44 | |||
| 45 | const isDeadSessionCached = (sessionId: string, now = Date.now()): boolean => { | ||
| 46 | const expiry = deadSessionExpiries.get(sessionId); | ||
| 47 | if (expiry === undefined) return false; | ||
| 48 | if (expiry > now) return true; | ||
| 49 | deadSessionExpiries.delete(sessionId); | ||
| 50 | timedOutSessionIds.delete(sessionId); | ||
| 51 | return false; | ||
| 52 | }; | ||
| 53 | |||
| 54 | const cacheDeadSession = (sessionId: string, reason: DeadSessionReason, now = Date.now()): void => { | ||
| 55 | deadSessionExpiries.delete(sessionId); | ||
| 56 | if (deadSessionExpiries.size >= DEAD_SESSION_CACHE_MAX_ENTRIES) { | ||
| 57 | const oldestSessionId = deadSessionExpiries.keys().next().value; | ||
| 58 | if (oldestSessionId !== undefined) { | ||
| 59 | deadSessionExpiries.delete(oldestSessionId); | ||
| 60 | timedOutSessionIds.delete(oldestSessionId); | ||
| 61 | } | ||
| 62 | } | ||
| 63 | deadSessionExpiries.set(sessionId, now + DEAD_SESSION_CACHE_TTL_MS); | ||
| 64 | if (reason === "timed_out") timedOutSessionIds.add(sessionId); | ||
| 65 | else timedOutSessionIds.delete(sessionId); | ||
| 66 | }; | ||
| 67 | |||
| 68 | const cachedDeadSessionMessage = (sessionId: string): string => | ||
| 69 | timedOutSessionIds.has(sessionId) ? "Session timed out, please reconnect" : "Session not found"; | ||
| 70 | |||
| 71 | /** Test-only access to reset and verify the isolate-local dead-session cache. */ | ||
| 72 | export const cloudDeadSessionCacheForTest = { | ||
| 73 | clear: (): void => { | ||
| 74 | deadSessionExpiries.clear(); | ||
| 75 | timedOutSessionIds.clear(); | ||
| 76 | }, | ||
| 77 | remember: (sessionId: string, now?: number): void => | ||
| 78 | cacheDeadSession(sessionId, "not_found", now), | ||
| 79 | has: isDeadSessionCached, | ||
| 80 | size: (): number => deadSessionExpiries.size, | ||
| 81 | }; | ||
| 82 | |||
| 38 | const jsonRpcResponse = ( | 83 | const jsonRpcResponse = ( |
| 39 | status: number, | 84 | status: number, |
| 40 | code: number, | 85 | code: number, |
| @@ -97,12 +142,12 @@ const renderAuthError = ( · 56 unmodified lines | |||
| 97 | }); | 142 | }); |
| 98 | }; | 143 | }; |
| 99 | 144 | ||
| 100 | const authenticate = (request: Request) => | 145 | const authenticate = (request: Request, authProvider: Layer.Layer<McpAuthProvider>) => |
| 101 | Effect.gen(function* () { | 146 | Effect.gen(function* () { |
| 102 | const auth = yield* McpAuthProvider; | 147 | const auth = yield* McpAuthProvider; |
| 103 | const outcome = yield* auth.authenticate(request); | 148 | const outcome = yield* auth.authenticate(request); |
| 104 | return { auth, outcome }; | 149 | return { auth, outcome }; |
| 105 | }).pipe(Effect.provide(cloudMcpAuth)); | 150 | }).pipe(Effect.provide(authProvider)); |
| 106 | 151 | ||
| 107 | // The earlier shared envelope ran the MCP auth path inside the Effect app, whose | 152 | // The earlier shared envelope ran the MCP auth path inside the Effect app, whose |
| 108 | // HttpMiddleware provided the OTEL tracer — that is where the `mcp.request` | 153 | // 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> = · 17 unmodified lines | |||
| 126 | ); | 171 | ); |
| 127 | }; | 172 | }; |
| 128 | 173 | ||
| 174 | type TraceCloudMcpRequest = ( | ||
| 175 | request: Request, | ||
| 176 | env: Env, | ||
| 177 | ctx: ExecutionContext, | ||
| 178 | handle: (tracedRequest: Request) => Promise<Response>, | ||
| 179 | ) => Promise<Response>; | ||
| 180 | |||
| 181 | interface CloudMcpAgentHandlerOptions { | ||
| 182 | readonly makeModernServerBuilder: ( | ||
| 183 | session: McpSessionProps["session"], | ||
| 184 | ) => McpModernServerBuilder["Service"]; | ||
| 185 | readonly authProvider?: Layer.Layer<McpAuthProvider>; | ||
| 186 | readonly traceRequest?: TraceCloudMcpRequest; | ||
| 187 | } | ||
| 188 | |||
| 129 | // The MCP resource the request targets. `server.ts` routes both the bare `/mcp` | 189 | // The MCP resource the request targets. `server.ts` routes both the bare `/mcp` |
| 130 | // and `/mcp/toolkits/<slug>` to this handler (`prepareMcpOrgScope` strips the org | 190 | // and `/mcp/toolkits/<slug>` to this handler (`prepareMcpOrgScope` strips the org |
| 131 | // selector but keeps the toolkit segment), so a session minted on a toolkit path | 191 | // selector but keeps the toolkit segment), so a session minted on a toolkit path |
| @@ -158,11 +218,14 @@ const propsForPrincipal = ( · 26 unmodified lines | |||
| 158 | }; | 218 | }; |
| 159 | }); | 219 | }); |
| 160 | 220 | ||
| 161 | export const makeCloudMcpAgentHandler = () => { | 221 | /** Build the cloud worker's authenticated legacy/modern MCP request handler. */ |
| 222 | export const makeCloudMcpAgentHandler = (options: CloudMcpAgentHandlerOptions) => { | ||
| 223 | const authProvider = options.authProvider ?? cloudMcpAuth; | ||
| 224 | const traceRequest = options.traceRequest ?? ((request, _env, _ctx, handle) => handle(request)); | ||
| 162 | const modern = makeMcpModernRequestRouter(); | 225 | const modern = makeMcpModernRequestRouter(); |
| 163 | const ALLOWED_METHODS = new Set(["GET", "POST", "DELETE", "OPTIONS"]); | 226 | const ALLOWED_METHODS = new Set(["GET", "POST", "DELETE", "OPTIONS"]); |
| 164 | 227 | ||
| 165 | return async (request: Request, env: Env, ctx: ExecutionContext): Promise<Response> => { | 228 | const handle = async (request: Request, env: Env, ctx: ExecutionContext): Promise<Response> => { |
| 166 | if (request.method === "OPTIONS") { | 229 | if (request.method === "OPTIONS") { |
| 167 | return mcpCorsPreflightResponse(request.headers.get("access-control-request-headers")); | 230 | return mcpCorsPreflightResponse(request.headers.get("access-control-request-headers")); |
| 168 | } | 231 | } |
| @@ -173,7 +236,7 @@ export const makeCloudMcpAgentHandler = () => { · 4 unmodified lines | |||
| 173 | } | 236 | } |
| 174 | const sessionId = request.headers.get("mcp-session-id"); | 237 | const sessionId = request.headers.get("mcp-session-id"); |
| 175 | 238 | ||
| 176 | const { auth, outcome } = await runTraced(request, authenticate(request)); | 239 | const { auth, outcome } = await runTraced(request, authenticate(request, authProvider)); |
| 177 | if (!Predicate.isTagged(outcome, "Authenticated")) { | 240 | if (!Predicate.isTagged(outcome, "Authenticated")) { |
| 178 | // Destroying a live session on auth grounds requires a POSITIVE | 241 | // Destroying a live session on auth grounds requires a POSITIVE |
| 179 | // determination that access is genuinely gone — only `Forbidden` carries | 242 | // determination that access is genuinely gone — only `Forbidden` carries |
| @@ -218,7 +281,7 @@ export const makeCloudMcpAgentHandler = () => { · 38 unmodified lines | |||
| 218 | resource, | 281 | resource, |
| 219 | props, | 282 | props, |
| 220 | requestStateSigningKey: requireMcpRequestStateKey(env.MCP_REQUEST_STATE_KEY), | 283 | requestStateSigningKey: requireMcpRequestStateKey(env.MCP_REQUEST_STATE_KEY), |
| 221 | builder: makeCloudModernMcpServerBuilder(props.session), | 284 | builder: options.makeModernServerBuilder(props.session), |
| 222 | sessions: env.MCP_SESSION, | 285 | sessions: env.MCP_SESSION, |
| 223 | executionOwners: mcpExecutionOwnerDirectoryFromNamespace(env.MCP_EXECUTION_OWNER), | 286 | executionOwners: mcpExecutionOwnerDirectoryFromNamespace(env.MCP_EXECUTION_OWNER), |
| 224 | }); | 287 | }); |
| @@ -238,18 +301,20 @@ export const makeCloudMcpAgentHandler = () => { · 13 unmodified lines | |||
| 238 | if (sessionId && !existingSession) { | 301 | if (sessionId && !existingSession) { |
| 239 | return deadSessionResponse(request.method, "Session not found"); | 302 | return deadSessionResponse(request.method, "Session not found"); |
| 240 | } | 303 | } |
| 241 | if (existingSession) { | 304 | if (existingSession && sessionId) { |
| 242 | const owner = await existingSession.validateMcpSessionOwner({ | 305 | const owner = await existingSession.validateMcpSessionOwner({ |
| 243 | accountId: outcome.principal.accountId, | 306 | accountId: outcome.principal.accountId, |
| 244 | organizationId: outcome.principal.organizationId, | 307 | organizationId: outcome.principal.organizationId, |
| 245 | }); | 308 | }); |
| 246 | if (owner === "not_found") { | 309 | if (owner === "not_found") { |
| 310 | cacheDeadSession(sessionId, "not_found"); | ||
| 247 | return deadSessionResponse(request.method, "Session not found"); | 311 | return deadSessionResponse(request.method, "Session not found"); |
| 248 | } | 312 | } |
| 249 | if (owner === "terminated") { | 313 | if (owner === "terminated") { |
| 250 | // DELETE-condemned but the deferred destroy alarm hasn't wiped storage | 314 | // DELETE-condemned but the deferred destroy alarm hasn't wiped storage |
| 251 | // yet. Same envelope as the post-destroy race below: the client must | 315 | // yet. Same envelope as the post-destroy race below: the client must |
| 252 | // treat the id as dead and reconnect. | 316 | // treat the id as dead and reconnect. |
| 317 | cacheDeadSession(sessionId, "timed_out"); | ||
| 253 | return deadSessionResponse(request.method, "Session timed out, please reconnect"); | 318 | return deadSessionResponse(request.method, "Session timed out, please reconnect"); |
| 254 | } | 319 | } |
| 255 | if (owner === "forbidden") { | 320 | if (owner === "forbidden") { |
| @@ -286,6 +351,7 @@ export const makeCloudMcpAgentHandler = () => { · 30 unmodified lines | |||
| 286 | // client to be told to reconnect, matching a timed-out session). | 351 | // client to be told to reconnect, matching a timed-out session). |
| 287 | // oxlint-disable-next-line executor/no-unknown-error-message -- adapter boundary: the abort reason is a plain runtime Error whose message IS the signal | 352 | // oxlint-disable-next-line executor/no-unknown-error-message -- adapter boundary: the abort reason is a plain runtime Error whose message IS the signal |
| 288 | if (Predicate.isError(error) && error.message === "destroyed") { | 353 | if (Predicate.isError(error) && error.message === "destroyed") { |
| 354 | if (sessionId) cacheDeadSession(sessionId, "timed_out"); | ||
| 289 | return deadSessionResponse(request.method, "Session timed out, please reconnect"); | 355 | return deadSessionResponse(request.method, "Session timed out, please reconnect"); |
| 290 | } | 356 | } |
| 291 | // 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 | 357 | // 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 |
| @@ -293,4 +359,20 @@ export const makeCloudMcpAgentHandler = () => { · 1 unmodified line | |||
| 293 | } | 359 | } |
| 294 | return withMcpResponseHeaders(wrapMcpSseResponse(request, env, response)); | 360 | return withMcpResponseHeaders(wrapMcpSseResponse(request, env, response)); |
| 295 | }; | 361 | }; |
| 362 | |||
| 363 | return async (request: Request, env: Env, ctx: ExecutionContext): Promise<Response> => { | ||
| 364 | const sessionId = request.headers.get("mcp-session-id"); | ||
| 365 | const cacheEligible = request.method !== "OPTIONS" && ALLOWED_METHODS.has(request.method); | ||
| 366 | if (cacheEligible && sessionId && isDeadSessionCached(sessionId)) { | ||
| 367 | return Effect.runPromise( | ||
| 368 | Effect.gen(function* () { | ||
| 369 | const { auth, outcome } = yield* authenticate(request, authProvider); | ||
| 370 | return Predicate.isTagged(outcome, "Authenticated") | ||
| 371 | ? deadSessionResponse(request.method, cachedDeadSessionMessage(sessionId)) | ||
| 372 | : renderAuthError(auth, request, outcome); | ||
| 373 | }).pipe(Effect.withTracerEnabled(false)), | ||
| 374 | ); | ||
| 375 | } | ||
| 376 | return traceRequest(request, env, ctx, (tracedRequest) => handle(tracedRequest, env, ctx)); | ||
| 377 | }; | ||
| 296 | }; | 378 | }; |
| @@ -15,7 +15,10 @@ import { isAppOwnedPath } from "./app-paths"; · 14 unmodified lines | |||
| 15 | import { makeCloudMcpAgentHandler } from "./mcp/agent-handler"; | 15 | import { makeCloudMcpAgentHandler } from "./mcp/agent-handler"; |
| 16 | import { classifyMcpPath, prepareMcpOrgScope } from "./mcp/mount"; | 16 | import { classifyMcpPath, prepareMcpOrgScope } from "./mcp/mount"; |
| 17 | import { parseTraceparent } from "./mcp/traceparent"; | 17 | import { parseTraceparent } from "./mcp/traceparent"; |
| 18 | import { McpSessionDOSqlite as McpSessionDOBase } from "./mcp/session-durable-object"; | 18 | import { |
| 19 | makeCloudModernMcpServerBuilder, | ||
| 20 | McpSessionDOSqlite as McpSessionDOBase, | ||
| 21 | } from "./mcp/session-durable-object"; | ||
| 19 | import { | 22 | import { |
| 20 | beforeSendWithOtelCorrelation, | 23 | beforeSendWithOtelCorrelation, |
| 21 | captureCause, | 24 | captureCause, |
| @@ -85,15 +88,12 @@ export { McpExecutionOwnerDirectoryDO } from "@executor-js/cloudflare/mcp/execut · 63 unmodified lines | |||
| 85 | // migration — without the OTel-SDK version-conflict that package would now | 88 | // migration — without the OTel-SDK version-conflict that package would now |
| 86 | // drag in (it pins `@opentelemetry/otlp-* ^0.200.0`, we ship ^0.214.0). | 89 | // drag in (it pins `@opentelemetry/otlp-* ^0.200.0`, we ship ^0.214.0). |
| 87 | // | 90 | // |
| 88 | // ONLY for paths the Effect app does not own. App-owned paths (/api/*, /mcp, | 91 | // App-owned paths (/api/* and /.well-known/* — see app-paths.ts) get their |
| 89 | // /.well-known/* — see app-paths.ts) get their `http.server` span from | 92 | // `http.server` span from Effect's HttpMiddleware tracer. `/mcp` is dispatched |
| 90 | // Effect's own HttpMiddleware.tracer, which parses `traceparent` itself and | 93 | // directly and uses `traceCloudMcpRequest` below so the agent handler can skip |
| 91 | // parents the workos/store/db child spans. Wrapping those here too produced | 94 | // the entire span envelope for negative-cache hits. Other paths keep this |
| 92 | // two identical sibling `http.server` spans per request (scope | 95 | // worker span. Wrapping Effect-owned paths here too produced duplicate sibling |
| 93 | // `executor-cloud-worker` next to scope `executor-cloud`) — double ingest, | 96 | // spans per request. |
| 94 | // and the waterfall showed a childless twin. The worker span remains for | ||
| 95 | // everything Effect never sees: Start SSR, the marketing proxy, /_astro | ||
| 96 | // assets. | ||
| 97 | // | 97 | // |
| 98 | // SimpleSpanProcessor exports synchronously at span end but the underlying | 98 | // SimpleSpanProcessor exports synchronously at span end but the underlying |
| 99 | // `fetch()` to Axiom is fire-and-forget; the Worker may terminate before it | 99 | // `fetch()` to Axiom is fire-and-forget; the Worker may terminate before it |
| @@ -108,7 +108,70 @@ const fetchHandler = handler.fetch as ( · 8 unmodified lines | |||
| 108 | ) => Response | Promise<Response>; | 108 | ) => Response | Promise<Response>; |
| 109 | 109 | ||
| 110 | const tracer = trace.getTracer("executor-cloud-worker"); | 110 | const tracer = trace.getTracer("executor-cloud-worker"); |
| 111 | const mcpAgentHandler = makeCloudMcpAgentHandler(); | 111 | |
| 112 | const traceCloudMcpRequest = async ( | ||
| 113 | request: Request, | ||
| 114 | _env: Env, | ||
| 115 | ctx: ExecutionContext, | ||
| 116 | handle: (tracedRequest: Request) => Promise<Response>, | ||
| 117 | ): Promise<Response> => { | ||
| 118 | if (!installTracerProvider()) return handle(request); | ||
| 119 | |||
| 120 | const url = new URL(request.url); | ||
| 121 | const inbound = parseTraceparent(request.headers.get("traceparent"), null); | ||
| 122 | const parentContext = inbound | ||
| 123 | ? trace.setSpanContext(context.active(), { | ||
| 124 | traceId: inbound.traceId, | ||
| 125 | spanId: inbound.spanId, | ||
| 126 | traceFlags: inbound.traceFlags, | ||
| 127 | isRemote: true, | ||
| 128 | }) | ||
| 129 | : context.active(); | ||
| 130 | |||
| 131 | return tracer.startActiveSpan( | ||
| 132 | `http.server ${request.method}`, | ||
| 133 | { kind: SpanKind.SERVER }, | ||
| 134 | parentContext, | ||
| 135 | async (span) => { | ||
| 136 | span.setAttribute(ATTR_HTTP_REQUEST_METHOD, request.method); | ||
| 137 | span.setAttribute(ATTR_URL_FULL, request.url); | ||
| 138 | span.setAttribute(ATTR_URL_PATH, url.pathname); | ||
| 139 | span.setAttribute(ATTR_URL_SCHEME, url.protocol.replace(/:$/, "")); | ||
| 140 | const spanContext = span.spanContext(); | ||
| 141 | const headers = new Headers(request.headers); | ||
| 142 | headers.set( | ||
| 143 | "traceparent", | ||
| 144 | `00-${spanContext.traceId}-${spanContext.spanId}-${(spanContext.traceFlags & 0xff).toString(16).padStart(2, "0")}`, | ||
| 145 | ); | ||
| 146 | // oxlint-disable-next-line executor/no-try-catch-or-throw -- adapter boundary; observe response/error for span status, keep trace export alive after the Agents bridge resolves or rejects | ||
| 147 | try { | ||
| 148 | const response = await handle(new Request(request, { headers })); | ||
| 149 | span.setAttribute(ATTR_HTTP_RESPONSE_STATUS_CODE, response.status); | ||
| 150 | if (response.status >= 500) { | ||
| 151 | span.setStatus({ code: SpanStatusCode.ERROR, message: `HTTP ${response.status}` }); | ||
| 152 | } | ||
| 153 | return response; | ||
| 154 | } catch (err) { | ||
| 155 | // oxlint-disable-next-line executor/no-instanceof-error, executor/no-unknown-error-message -- adapter boundary: Cloudflare's fetch callback throws untyped; normalized only for the OTel span record, the original error is rethrown below | ||
| 156 | const cause = err instanceof Error ? err : String(err); | ||
| 157 | span.recordException(cause); | ||
| 158 | // oxlint-disable-next-line executor/no-unknown-error-message -- adapter boundary: same normalization as the recordException line above | ||
| 159 | const message = typeof cause === "string" ? cause : cause.message; | ||
| 160 | span.setStatus({ code: SpanStatusCode.ERROR, message }); | ||
| 161 | // oxlint-disable-next-line executor/no-try-catch-or-throw -- adapter boundary; preserve original error to Cloudflare runtime | ||
| 162 | throw err; | ||
| 163 | } finally { | ||
| 164 | span.end(); | ||
| 165 | ctx.waitUntil(flushTracerProvider()); | ||
| 166 | } | ||
| 167 | }, | ||
| 168 | ); | ||
| 169 | }; | ||
| 170 | |||
| 171 | const mcpAgentHandler = makeCloudMcpAgentHandler({ | ||
| 172 | makeModernServerBuilder: makeCloudModernMcpServerBuilder, | ||
| 173 | traceRequest: traceCloudMcpRequest, | ||
| 174 | }); | ||
| 112 | 175 | ||
| 113 | const cloudflareHandler: ExportedHandler<Env> = { | 176 | const cloudflareHandler: ExportedHandler<Env> = { |
| 114 | fetch: async (request, env, ctx) => { | 177 | fetch: async (request, env, ctx) => { |
| @@ -120,10 +183,17 @@ const cloudflareHandler: ExportedHandler<Env> = { · 5 unmodified lines | |||
| 120 | // The MCP dispatch is classified up front, independent of whether | 183 | // The MCP dispatch is classified up front, independent of whether |
| 121 | // telemetry installs — an unset `AXIOM_TOKEN` (tracer not installed) must | 184 | // telemetry installs — an unset `AXIOM_TOKEN` (tracer not installed) must |
| 122 | // never take /mcp requests down with it. See `installTracerProvider`'s | 185 | // never take /mcp requests down with it. See `installTracerProvider`'s |
| 123 | // early return below: it only governs the tracing envelope for | 186 | // early return below: the handler invokes it for uncached MCP traffic, and |
| 124 | // non-MCP paths. | 187 | // this entry invokes it for non-MCP paths. |
| 125 | const url = new URL(request.url); | 188 | const url = new URL(request.url); |
| 126 | const mcpRoute = classifyMcpPath(url.pathname); | 189 | const mcpRoute = classifyMcpPath(url.pathname); |
| 190 | if (mcpRoute?.kind === "mcp") { | ||
| 191 | // The Cloudflare Agents MCP bridge needs the platform ExecutionContext | ||
| 192 | // to pass authenticated session props into the hibernatable DO. | ||
| 193 | // Discovery docs still flow through the app-level MCP envelope. | ||
| 194 | const forwarded = prepareMcpOrgScope(request); | ||
| 195 | return mcpAgentHandler(forwarded, env, ctx); | ||
| 196 | } | ||
| 127 | const tracingInstalled = installTracerProvider(); | 197 | const tracingInstalled = installTracerProvider(); |
| 128 | // Join the caller's W3C trace when the request carries one — the web UI | 198 | // Join the caller's W3C trace when the request carries one — the web UI |
| 129 | // sends traceparent on every API fetch, so the browser's spans and this | 199 | // sends traceparent on every API fetch, so the browser's spans and this |
| @@ -138,61 +208,6 @@ const cloudflareHandler: ExportedHandler<Env> = { · 8 unmodified lines | |||
| 138 | isRemote: true, | 208 | isRemote: true, |
| 139 | }) | 209 | }) |
| 140 | : context.active(); | 210 | : context.active(); |
| 141 | if (mcpRoute?.kind === "mcp") { | ||
| 142 | // The Cloudflare Agents MCP bridge needs the platform ExecutionContext | ||
| 143 | // to pass authenticated session props into the hibernatable DO. | ||
| 144 | // Discovery docs still flow through the app-level MCP envelope. | ||
| 145 | const forwarded = prepareMcpOrgScope(request); | ||
| 146 | if (!tracingInstalled) { | ||
| 147 | return mcpAgentHandler(forwarded, env, ctx); | ||
| 148 | } | ||
| 149 | // /mcp left the Effect app in the Agents-bridge migration, so no | ||
| 150 | // downstream HttpMiddleware.tracer opens the request envelope anymore — | ||
| 151 | // this worker span is now THE `http.server` span for MCP traffic. Its | ||
| 152 | // context is stamped onto the forwarded request's traceparent so the | ||
| 153 | // agent handler's Effect programs (mcp.request and children) and the | ||
| 154 | // session DO parent under it instead of exporting orphaned roots. | ||
| 155 | return tracer.startActiveSpan( | ||
| 156 | `http.server ${request.method}`, | ||
| 157 | { kind: SpanKind.SERVER }, | ||
| 158 | parentContext, | ||
| 159 | async (span) => { | ||
| 160 | span.setAttribute(ATTR_HTTP_REQUEST_METHOD, request.method); | ||
| 161 | span.setAttribute(ATTR_URL_FULL, request.url); | ||
| 162 | span.setAttribute(ATTR_URL_PATH, url.pathname); | ||
| 163 | span.setAttribute(ATTR_URL_SCHEME, url.protocol.replace(/:$/, "")); | ||
| 164 | const spanContext = span.spanContext(); | ||
| 165 | const headers = new Headers(forwarded.headers); | ||
| 166 | headers.set( | ||
| 167 | "traceparent", | ||
| 168 | `00-${spanContext.traceId}-${spanContext.spanId}-${(spanContext.traceFlags & 0xff).toString(16).padStart(2, "0")}`, | ||
| 169 | ); | ||
| 170 | // oxlint-disable-next-line executor/no-try-catch-or-throw -- adapter boundary; observe response/error for span status, keep trace export alive after the Agents bridge resolves or rejects | ||
| 171 | try { | ||
| 172 | const response = await mcpAgentHandler(new Request(forwarded, { headers }), env, ctx); | ||
| 173 | span.setAttribute(ATTR_HTTP_RESPONSE_STATUS_CODE, response.status); | ||
| 174 | if (response.status >= 500) { | ||
| 175 | span.setStatus({ code: SpanStatusCode.ERROR, message: `HTTP ${response.status}` }); | ||
| 176 | } | ||
| 177 | return response; | ||
| 178 | } catch (err) { | ||
| 179 | // Record the exception itself, not just the status bit: without it | ||
| 180 | // these spans are ERROR with zero diagnostic content. | ||
| 181 | // oxlint-disable-next-line executor/no-instanceof-error, executor/no-unknown-error-message -- adapter boundary: Cloudflare's fetch callback throws untyped; normalized only for the OTel span record, the original error is rethrown below | ||
| 182 | const cause = err instanceof Error ? err : String(err); | ||
| 183 | span.recordException(cause); | ||
| 184 | // oxlint-disable-next-line executor/no-unknown-error-message -- adapter boundary: same normalization as the recordException line above | ||
| 185 | const message = typeof cause === "string" ? cause : cause.message; | ||
| 186 | span.setStatus({ code: SpanStatusCode.ERROR, message }); | ||
| 187 | // oxlint-disable-next-line executor/no-try-catch-or-throw -- adapter boundary; preserve original error to Cloudflare runtime | ||
| 188 | throw err; | ||
| 189 | } finally { | ||
| 190 | span.end(); | ||
| 191 | ctx.waitUntil(flushTracerProvider()); | ||
| 192 | } | ||
| 193 | }, | ||
| 194 | ); | ||
| 195 | } | ||
| 196 | if (!tracingInstalled) { | 211 | if (!tracingInstalled) { |
| 197 | return fetchHandler(request, env, ctx); | 212 | return fetchHandler(request, env, ctx); |
| 198 | } | 213 | } |
| @@ -16,6 +16,7 @@ import { ErrorCaptureLive } from "./observability"; · 15 unmodified lines | |||
| 16 | import { cloudflareAccountMiddleware } from "./account/account-provider"; | 16 | import { cloudflareAccountMiddleware } from "./account/account-provider"; |
| 17 | import { makeCloudflareApprovalHandler } from "./mcp"; | 17 | import { makeCloudflareApprovalHandler } from "./mcp"; |
| 18 | import { makeCloudflareMcpAgentHandler } from "./mcp/agent-handler"; | 18 | import { makeCloudflareMcpAgentHandler } from "./mcp/agent-handler"; |
| 19 | import { makeCloudflareModernMcpServerBuilder } from "./mcp/session-durable-object"; | ||
| 19 | import { preloadQuickJs } from "./quickjs"; | 20 | import { preloadQuickJs } from "./quickjs"; |
| 20 | 21 | ||
| 21 | // =========================================================================== | 22 | // =========================================================================== |
| @@ -45,7 +46,9 @@ export const makeCloudflareApp = async (env: CloudflareEnv) => { · 23 unmodified lines | |||
| 45 | // handle the per-request scoped executor reads through the DbProvider seam. | 46 | // handle the per-request scoped executor reads through the DbProvider seam. |
| 46 | const dbHandle = await createD1ExecutorDb(env.DB, env.BLOBS); | 47 | const dbHandle = await createD1ExecutorDb(env.DB, env.BLOBS); |
| 47 | const identityLayer = cloudflareAccessIdentityLayer(config); | 48 | const identityLayer = cloudflareAccessIdentityLayer(config); |
| 48 | const mcpAgentHandler = makeCloudflareMcpAgentHandler(config); | 49 | const mcpAgentHandler = makeCloudflareMcpAgentHandler(config, { |
| 50 | makeModernServerBuilder: makeCloudflareModernMcpServerBuilder, | ||
| 51 | }); | ||
| 49 | const approvalHandler = makeCloudflareApprovalHandler(config, env); | 52 | const approvalHandler = makeCloudflareApprovalHandler(config, env); |
| 50 | 53 | ||
| 51 | const { appLayer, toWebHandler } = ExecutorApp.make({ | 54 | const { appLayer, toWebHandler } = ExecutorApp.make({ |
| @@ -0,0 +1,153 @@ | |||
| 1 | import { beforeEach, describe, expect, it, vi } from "@effect/vitest"; | ||
| 2 | import { Effect, Layer } from "effect"; | ||
| 3 | |||
| 4 | import { | ||
| 5 | authenticated, | ||
| 6 | McpAuthProvider, | ||
| 7 | unauthorized, | ||
| 8 | type Principal, | ||
| 9 | } from "@executor-js/host-mcp"; | ||
| 10 | |||
| 11 | import type { CloudflareConfig, CloudflareEnv } from "../config"; | ||
| 12 | import { cloudflareDeadSessionCacheForTest, makeCloudflareMcpAgentHandler } from "./agent-handler"; | ||
| 13 | |||
| 14 | const principal: Principal = { | ||
| 15 | accountId: "acct_test", | ||
| 16 | organizationId: "org_test", | ||
| 17 | organizationName: "Test Org", | ||
| 18 | email: "test@example.com", | ||
| 19 | name: "Test", | ||
| 20 | avatarUrl: null, | ||
| 21 | roles: ["member"], | ||
| 22 | }; | ||
| 23 | |||
| 24 | const config = { | ||
| 25 | accessTeamDomain: "test.cloudflareaccess.com", | ||
| 26 | accessAud: "aud_test", | ||
| 27 | accessNameClaim: "name", | ||
| 28 | accessGroupsClaim: "groups", | ||
| 29 | adminEmails: [], | ||
| 30 | organizationId: "org_test", | ||
| 31 | organizationName: "Test Org", | ||
| 32 | organizationSlug: "test-org", | ||
| 33 | secretKey: "test-secret-key-0123456789", | ||
| 34 | allowLocalNetwork: false, | ||
| 35 | enableDevAuth: false, | ||
| 36 | } satisfies CloudflareConfig; | ||
| 37 | |||
| 38 | const AuthProviderLive = Layer.succeed(McpAuthProvider)({ | ||
| 39 | discoveryRoutes: [], | ||
| 40 | resourceMetadataUrl: (request) => new URL("/.well-known/mcp", request.url).toString(), | ||
| 41 | authenticate: (request) => | ||
| 42 | Effect.succeed( | ||
| 43 | request.headers.has("authorization") ? authenticated(principal) : unauthorized(), | ||
| 44 | ), | ||
| 45 | }); | ||
| 46 | |||
| 47 | const requestFor = (method: "GET" | "POST", sessionId: string, authenticated = true): Request => | ||
| 48 | new Request("https://executor.test/mcp", { | ||
| 49 | method, | ||
| 50 | headers: { | ||
| 51 | ...(authenticated ? { authorization: "Bearer test" } : {}), | ||
| 52 | "mcp-session-id": sessionId, | ||
| 53 | ...(method === "POST" ? { "content-type": "application/json" } : {}), | ||
| 54 | }, | ||
| 55 | ...(method === "POST" | ||
| 56 | ? { | ||
| 57 | body: JSON.stringify({ jsonrpc: "2.0", id: 1, method: "tools/list", params: {} }), | ||
| 58 | } | ||
| 59 | : {}), | ||
| 60 | }); | ||
| 61 | |||
| 62 | const makeHarness = (owner: "not_found" | "terminated" = "not_found") => { | ||
| 63 | const ownerChecks = { count: 0 }; | ||
| 64 | const stub = { | ||
| 65 | validateMcpSessionOwner: async () => { | ||
| 66 | ownerChecks.count += 1; | ||
| 67 | return owner; | ||
| 68 | }, | ||
| 69 | }; | ||
| 70 | const namespace = { | ||
| 71 | idFromString: (sessionId: string) => sessionId, | ||
| 72 | get: () => stub, | ||
| 73 | }; | ||
| 74 | // oxlint-disable-next-line executor/no-double-cast -- test boundary: the handler only reads the MCP_SESSION namespace in these legacy dead-session cases | ||
| 75 | const env = { MCP_SESSION: namespace } as unknown as CloudflareEnv; | ||
| 76 | // oxlint-disable-next-line executor/no-double-cast -- test boundary: no ExecutionContext capability is used before the dead-session response | ||
| 77 | const ctx = {} as unknown as ExecutionContext; | ||
| 78 | const handler = makeCloudflareMcpAgentHandler(config, { | ||
| 79 | authProvider: AuthProviderLive, | ||
| 80 | makeModernServerBuilder: () => ({ build: () => Effect.die("unused modern builder") }), | ||
| 81 | }); | ||
| 82 | return { ctx, env, handler, ownerChecks }; | ||
| 83 | }; | ||
| 84 | |||
| 85 | describe("standalone Cloudflare MCP dead-session negative cache", () => { | ||
| 86 | beforeEach(() => { | ||
| 87 | vi.useRealTimers(); | ||
| 88 | cloudflareDeadSessionCacheForTest.clear(); | ||
| 89 | }); | ||
| 90 | |||
| 91 | it.each([ | ||
| 92 | { method: "GET" as const, status: 405 }, | ||
| 93 | { method: "POST" as const, status: 404 }, | ||
| 94 | ])("serves a repeated $method without another DO lookup", async ({ method, status }) => { | ||
| 95 | const { ctx, env, handler, ownerChecks } = makeHarness(); | ||
| 96 | const sessionId = `dead-${method.toLowerCase()}`; | ||
| 97 | |||
| 98 | const first = await handler(requestFor(method, sessionId), env, ctx); | ||
| 99 | const second = await handler(requestFor(method, sessionId), env, ctx); | ||
| 100 | |||
| 101 | expect(first.status).toBe(status); | ||
| 102 | expect(second.status).toBe(status); | ||
| 103 | expect(ownerChecks.count).toBe(1); | ||
| 104 | }); | ||
| 105 | |||
| 106 | it("keeps authentication ahead of cached session existence", async () => { | ||
| 107 | const { ctx, env, handler, ownerChecks } = makeHarness(); | ||
| 108 | const sessionId = "dead-auth"; | ||
| 109 | |||
| 110 | expect((await handler(requestFor("GET", sessionId), env, ctx)).status).toBe(405); | ||
| 111 | const unauthenticated = await handler(requestFor("GET", sessionId, false), env, ctx); | ||
| 112 | |||
| 113 | expect(unauthenticated.status).toBe(401); | ||
| 114 | expect(ownerChecks.count).toBe(1); | ||
| 115 | }); | ||
| 116 | |||
| 117 | it("preserves the terminated-session response on a cached hit", async () => { | ||
| 118 | const { ctx, env, handler, ownerChecks } = makeHarness("terminated"); | ||
| 119 | const sessionId = "dead-terminated"; | ||
| 120 | |||
| 121 | const first = await handler(requestFor("POST", sessionId), env, ctx); | ||
| 122 | const second = await handler(requestFor("POST", sessionId), env, ctx); | ||
| 123 | |||
| 124 | expect(first.status).toBe(404); | ||
| 125 | expect(second.status).toBe(404); | ||
| 126 | expect(await second.text()).toBe(await first.text()); | ||
| 127 | expect(ownerChecks.count).toBe(1); | ||
| 128 | }); | ||
| 129 | |||
| 130 | it("consults the DO again after five minutes", async () => { | ||
| 131 | vi.useFakeTimers(); | ||
| 132 | vi.setSystemTime(0); | ||
| 133 | const { ctx, env, handler, ownerChecks } = makeHarness(); | ||
| 134 | const sessionId = "dead-expiry"; | ||
| 135 | |||
| 136 | expect((await handler(requestFor("GET", sessionId), env, ctx)).status).toBe(405); | ||
| 137 | vi.advanceTimersByTime(5 * 60 * 1_000 + 1); | ||
| 138 | expect((await handler(requestFor("GET", sessionId), env, ctx)).status).toBe(405); | ||
| 139 | |||
| 140 | expect(ownerChecks.count).toBe(2); | ||
| 141 | }); | ||
| 142 | |||
| 143 | it("evicts the oldest entry without exceeding 4,096 sessions", () => { | ||
| 144 | for (let index = 0; index <= 4_096; index += 1) { | ||
| 145 | cloudflareDeadSessionCacheForTest.remember(`dead-${index}`, 0); | ||
| 146 | } | ||
| 147 | |||
| 148 | expect(cloudflareDeadSessionCacheForTest.size()).toBe(4_096); | ||
| 149 | expect(cloudflareDeadSessionCacheForTest.has("dead-0", 1)).toBe(false); | ||
| 150 | expect(cloudflareDeadSessionCacheForTest.has("dead-1", 1)).toBe(true); | ||
| 151 | expect(cloudflareDeadSessionCacheForTest.has("dead-4096", 1)).toBe(true); | ||
| 152 | }); | ||
| 153 | }); | ||
| @@ -1,4 +1,4 @@ | |||
| 1 | import { Effect, Predicate } from "effect"; | 1 | import { Effect, Layer, Predicate } from "effect"; |
| 2 | 2 | ||
| 3 | import { | 3 | import { |
| 4 | McpAuthProvider, | 4 | McpAuthProvider, |
| @@ -6,6 +6,7 @@ import { · 1 unmodified line | |||
| 6 | mcpModernDisabledResponse, | 6 | mcpModernDisabledResponse, |
| 7 | defaultMcpResource, | 7 | defaultMcpResource, |
| 8 | type AuthOutcome, | 8 | type AuthOutcome, |
| 9 | type McpModernServerBuilder, | ||
| 9 | type Principal, | 10 | type Principal, |
| 10 | } from "@executor-js/host-mcp"; | 11 | } from "@executor-js/host-mcp"; |
| 11 | import { requestBodyFromRequest } from "@executor-js/host-mcp/tool-server"; | 12 | import { requestBodyFromRequest } from "@executor-js/host-mcp/tool-server"; |
| @@ -29,7 +30,51 @@ import { createMcpSessionStub, mcpSessionStub } from "@executor-js/cloudflare/mc · 17 unmodified lines | |||
| 29 | 30 | ||
| 30 | import type { CloudflareConfig, CloudflareEnv } from "../config"; | 31 | import type { CloudflareConfig, CloudflareEnv } from "../config"; |
| 31 | import { cloudflareAccessMcpAuth } from "./auth"; | 32 | import { cloudflareAccessMcpAuth } from "./auth"; |
| 32 | import { makeCloudflareModernMcpServerBuilder } from "./session-durable-object"; | 33 | |
| 34 | const DEAD_SESSION_CACHE_TTL_MS = 5 * 60 * 1_000; | ||
| 35 | const DEAD_SESSION_CACHE_MAX_ENTRIES = 4_096; | ||
| 36 | const deadSessionExpiries = new Map<string, number>(); | ||
| 37 | const timedOutSessionIds = new Set<string>(); | ||
| 38 | |||
| 39 | type DeadSessionReason = "not_found" | "timed_out"; | ||
| 40 | |||
| 41 | const isDeadSessionCached = (sessionId: string, now = Date.now()): boolean => { | ||
| 42 | const expiry = deadSessionExpiries.get(sessionId); | ||
| 43 | if (expiry === undefined) return false; | ||
| 44 | if (expiry > now) return true; | ||
| 45 | deadSessionExpiries.delete(sessionId); | ||
| 46 | timedOutSessionIds.delete(sessionId); | ||
| 47 | return false; | ||
| 48 | }; | ||
| 49 | |||
| 50 | const cacheDeadSession = (sessionId: string, reason: DeadSessionReason, now = Date.now()): void => { | ||
| 51 | deadSessionExpiries.delete(sessionId); | ||
| 52 | if (deadSessionExpiries.size >= DEAD_SESSION_CACHE_MAX_ENTRIES) { | ||
| 53 | const oldestSessionId = deadSessionExpiries.keys().next().value; | ||
| 54 | if (oldestSessionId !== undefined) { | ||
| 55 | deadSessionExpiries.delete(oldestSessionId); | ||
| 56 | timedOutSessionIds.delete(oldestSessionId); | ||
| 57 | } | ||
| 58 | } | ||
| 59 | deadSessionExpiries.set(sessionId, now + DEAD_SESSION_CACHE_TTL_MS); | ||
| 60 | if (reason === "timed_out") timedOutSessionIds.add(sessionId); | ||
| 61 | else timedOutSessionIds.delete(sessionId); | ||
| 62 | }; | ||
| 63 | |||
| 64 | const cachedDeadSessionMessage = (sessionId: string): string => | ||
| 65 | timedOutSessionIds.has(sessionId) ? "Session timed out, please reconnect" : "Session not found"; | ||
| 66 | |||
| 67 | /** Test-only access to reset and verify the isolate-local dead-session cache. */ | ||
| 68 | export const cloudflareDeadSessionCacheForTest = { | ||
| 69 | clear: (): void => { | ||
| 70 | deadSessionExpiries.clear(); | ||
| 71 | timedOutSessionIds.clear(); | ||
| 72 | }, | ||
| 73 | remember: (sessionId: string, now?: number): void => | ||
| 74 | cacheDeadSession(sessionId, "not_found", now), | ||
| 75 | has: isDeadSessionCached, | ||
| 76 | size: (): number => deadSessionExpiries.size, | ||
| 77 | }; | ||
| 33 | 78 | ||
| 34 | const jsonRpcResponse = ( | 79 | const jsonRpcResponse = ( |
| 35 | status: number, | 80 | status: number, |
| @@ -79,12 +124,12 @@ const renderAuthError = ( · 43 unmodified lines | |||
| 79 | return jsonRpcResponse(503, -32001, outcome.message); | 124 | return jsonRpcResponse(503, -32001, outcome.message); |
| 80 | }; | 125 | }; |
| 81 | 126 | ||
| 82 | const authenticate = (request: Request, config: CloudflareConfig) => | 127 | const authenticate = (request: Request, authProvider: Layer.Layer<McpAuthProvider>) => |
| 83 | Effect.gen(function* () { | 128 | Effect.gen(function* () { |
| 84 | const auth = yield* McpAuthProvider; | 129 | const auth = yield* McpAuthProvider; |
| 85 | const outcome = yield* auth.authenticate(request); | 130 | const outcome = yield* auth.authenticate(request); |
| 86 | return { auth, outcome }; | 131 | return { auth, outcome }; |
| 87 | }).pipe(Effect.provide(cloudflareAccessMcpAuth(config))); | 132 | }).pipe(Effect.provide(authProvider)); |
| 88 | 133 | ||
| 89 | const propsForPrincipal = ( | 134 | const propsForPrincipal = ( |
| 90 | request: Request, | 135 | request: Request, |
| @@ -108,7 +153,21 @@ const propsForPrincipal = ( · 17 unmodified lines | |||
| 108 | }; | 153 | }; |
| 109 | }); | 154 | }); |
| 110 | 155 | ||
| 111 | export const makeCloudflareMcpAgentHandler = (config: CloudflareConfig) => { | 156 | interface CloudflareMcpAgentHandlerOptions { |
| 157 | readonly makeModernServerBuilder: ( | ||
| 158 | env: CloudflareEnv, | ||
| 159 | config: CloudflareConfig, | ||
| 160 | session: McpSessionProps["session"], | ||
| 161 | ) => McpModernServerBuilder["Service"]; | ||
| 162 | readonly authProvider?: Layer.Layer<McpAuthProvider>; | ||
| 163 | } | ||
| 164 | |||
| 165 | /** Build the standalone Cloudflare worker's authenticated MCP request handler. */ | ||
| 166 | export const makeCloudflareMcpAgentHandler = ( | ||
| 167 | config: CloudflareConfig, | ||
| 168 | options: CloudflareMcpAgentHandlerOptions, | ||
| 169 | ) => { | ||
| 170 | const authProvider = options.authProvider ?? cloudflareAccessMcpAuth(config); | ||
| 112 | const modern = makeMcpModernRequestRouter(); | 171 | const modern = makeMcpModernRequestRouter(); |
| 113 | return async (request: Request, env: CloudflareEnv, ctx: ExecutionContext): Promise<Response> => { | 172 | return async (request: Request, env: CloudflareEnv, ctx: ExecutionContext): Promise<Response> => { |
| 114 | if (request.method === "OPTIONS") { | 173 | if (request.method === "OPTIONS") { |
| @@ -116,7 +175,18 @@ export const makeCloudflareMcpAgentHandler = (config: CloudflareConfig) => { · 1 unmodified line | |||
| 116 | } | 175 | } |
| 117 | const sessionId = request.headers.get("mcp-session-id"); | 176 | const sessionId = request.headers.get("mcp-session-id"); |
| 118 | 177 | ||
| 119 | const { auth, outcome } = await Effect.runPromise(authenticate(request, config)); | 178 | if (sessionId && isDeadSessionCached(sessionId)) { |
| 179 | return Effect.runPromise( | ||
| 180 | Effect.gen(function* () { | ||
| 181 | const { auth, outcome } = yield* authenticate(request, authProvider); | ||
| 182 | return Predicate.isTagged(outcome, "Authenticated") | ||
| 183 | ? deadSessionResponse(request.method, cachedDeadSessionMessage(sessionId)) | ||
| 184 | : renderAuthError(auth, request, outcome); | ||
| 185 | }).pipe(Effect.withTracerEnabled(false)), | ||
| 186 | ); | ||
| 187 | } | ||
| 188 | |||
| 189 | const { auth, outcome } = await Effect.runPromise(authenticate(request, authProvider)); | ||
| 120 | if (!Predicate.isTagged(outcome, "Authenticated")) { | 190 | if (!Predicate.isTagged(outcome, "Authenticated")) { |
| 121 | if (Predicate.isTagged(outcome, "Forbidden") && sessionId) { | 191 | if (Predicate.isTagged(outcome, "Forbidden") && sessionId) { |
| 122 | const session = mcpSessionStub(env.MCP_SESSION, sessionId); | 192 | const session = mcpSessionStub(env.MCP_SESSION, sessionId); |
| @@ -152,7 +222,7 @@ export const makeCloudflareMcpAgentHandler = (config: CloudflareConfig) => { · 29 unmodified lines | |||
| 152 | resource: defaultMcpResource, | 222 | resource: defaultMcpResource, |
| 153 | props, | 223 | props, |
| 154 | requestStateSigningKey: requireMcpRequestStateKey(env.MCP_REQUEST_STATE_KEY), | 224 | requestStateSigningKey: requireMcpRequestStateKey(env.MCP_REQUEST_STATE_KEY), |
| 155 | builder: makeCloudflareModernMcpServerBuilder(env, config, props.session), | 225 | builder: options.makeModernServerBuilder(env, config, props.session), |
| 156 | sessions: env.MCP_SESSION, | 226 | sessions: env.MCP_SESSION, |
| 157 | executionOwners: mcpExecutionOwnerDirectoryFromNamespace(env.MCP_EXECUTION_OWNER), | 227 | executionOwners: mcpExecutionOwnerDirectoryFromNamespace(env.MCP_EXECUTION_OWNER), |
| 158 | }); | 228 | }); |
| @@ -166,17 +236,19 @@ export const makeCloudflareMcpAgentHandler = (config: CloudflareConfig) => { · 7 unmodified lines | |||
| 166 | if (sessionId && !existingSession) { | 236 | if (sessionId && !existingSession) { |
| 167 | return deadSessionResponse(request.method, "Session not found"); | 237 | return deadSessionResponse(request.method, "Session not found"); |
| 168 | } | 238 | } |
| 169 | if (existingSession) { | 239 | if (existingSession && sessionId) { |
| 170 | const owner = await existingSession.validateMcpSessionOwner({ | 240 | const owner = await existingSession.validateMcpSessionOwner({ |
| 171 | accountId: outcome.principal.accountId, | 241 | accountId: outcome.principal.accountId, |
| 172 | organizationId: outcome.principal.organizationId, | 242 | organizationId: outcome.principal.organizationId, |
| 173 | }); | 243 | }); |
| 174 | if (owner === "not_found") { | 244 | if (owner === "not_found") { |
| 245 | cacheDeadSession(sessionId, "not_found"); | ||
| 175 | return deadSessionResponse(request.method, "Session not found"); | 246 | return deadSessionResponse(request.method, "Session not found"); |
| 176 | } | 247 | } |
| 177 | if (owner === "terminated") { | 248 | if (owner === "terminated") { |
| 178 | // DELETE-condemned but the deferred destroy alarm hasn't wiped storage | 249 | // DELETE-condemned but the deferred destroy alarm hasn't wiped storage |
| 179 | // yet; the terminated id must read as dead immediately. | 250 | // yet; the terminated id must read as dead immediately. |
| 251 | cacheDeadSession(sessionId, "timed_out"); | ||
| 180 | return deadSessionResponse(request.method, "Session timed out, please reconnect"); | 252 | return deadSessionResponse(request.method, "Session timed out, please reconnect"); |
| 181 | } | 253 | } |
| 182 | if (owner === "forbidden") { | 254 | if (owner === "forbidden") { |
| @@ -0,0 +1,7 @@ | |||
| 1 | /** Node-test stand-ins for Cloudflare module exports reached during handler imports. */ | ||
| 2 | export const env: Record<string, unknown> = {}; | ||
| 3 | export class WorkerEntrypoint {} | ||
| 4 | export class DurableObject {} | ||
| 5 | export class WorkflowEntrypoint {} | ||
| 6 | export class RpcTarget {} | ||
| 7 | export const exports: Record<string, unknown> = {}; | ||
| @@ -1,6 +1,13 @@ | |||
| 1 | import { resolve } from "node:path"; | ||
| 2 | |||
| 1 | import { defineConfig } from "vitest/config"; | 3 | import { defineConfig } from "vitest/config"; |
| 2 | 4 | ||
| 3 | export default defineConfig({ | 5 | export default defineConfig({ |
| 6 | resolve: { | ||
| 7 | alias: { | ||
| 8 | "cloudflare:workers": resolve(__dirname, "./test-stubs/cloudflare-workers.ts"), | ||
| 9 | }, | ||
| 10 | }, | ||
| 4 | test: { | 11 | test: { |
| 5 | include: ["src/**/*.test.ts"], | 12 | include: ["src/**/*.test.ts"], |
| 6 | passWithNoTests: true, | 13 | passWithNoTests: true, |