From b727c6b97565b4be9e94d31eae79b658e0feb730 Mon Sep 17 00:00:00 2001 From: Rhys Sullivan <39114868+RhysSullivan@users.noreply.github.com> Date: Thu, 27 Aug 2026 13:24:30 -0700 Subject: [PATCH 1/3] MCP session cold init: carry the org identity in the session props --- apps/cloud/src/auth/context.ts | 10 +- apps/cloud/src/auth/errors.ts | 97 +++++++- .../src/auth/user-store-error.node.test.ts | 128 +++++++++++ apps/cloud/src/mcp/agent-handler.ts | 27 +++ apps/cloud/src/mcp/auth-provider.test.ts | 6 +- apps/cloud/src/mcp/auth-provider.ts | 31 ++- apps/cloud/src/mcp/auth.ts | 21 +- apps/cloud/src/mcp/session-durable-object.ts | 41 ++-- apps/cloud/src/mcp/session-meta.node.test.ts | 167 ++++++++++++++ apps/cloud/src/mcp/session-meta.ts | 209 ++++++++++++++++++ apps/cloud/src/observability/index.ts | 47 ++++ .../src/mcp/session-durable-object.ts | 9 +- e2e/cloud/mcp-session-cold-init.test.ts | 100 +++++++++ .../mcp/agent-session-durable-object.test.ts | 115 +++++++++- .../src/mcp/agent-session-durable-object.ts | 46 +++- 15 files changed, 1000 insertions(+), 54 deletions(-) create mode 100644 apps/cloud/src/auth/user-store-error.node.test.ts create mode 100644 apps/cloud/src/mcp/session-meta.node.test.ts create mode 100644 apps/cloud/src/mcp/session-meta.ts create mode 100644 e2e/cloud/mcp-session-cold-init.test.ts diff --git a/apps/cloud/src/auth/context.ts b/apps/cloud/src/auth/context.ts index dd713b189..ce8aa1804 100644 --- a/apps/cloud/src/auth/context.ts +++ b/apps/cloud/src/auth/context.ts @@ -1,7 +1,7 @@ import { Context, Effect, Layer } from "effect"; import { makeUserStore } from "../auth/user-store"; import { DbService } from "../db/db"; -import { UserStoreError, tryPromiseService, withServiceLogging } from "./errors"; +import { tryPromiseService, userStoreErrorFromFailure, withServiceLogging } from "./errors"; // --------------------------------------------------------------------------- // UserStoreService — wraps the Drizzle-backed user store with Effect @@ -10,13 +10,15 @@ import { UserStoreError, tryPromiseService, withServiceLogging } from "./errors" type RawStore = ReturnType; // `op` names the store call so every span reads `user_store.` -// instead of one undifferentiated "user_store" bucket, and failures log which -// query actually failed. +// instead of one undifferentiated "user_store" bucket, failures log which +// query actually failed, and — because the same `op` is threaded onto the +// public error alongside the classified driver reason — an error report is +// diagnosable without the trace. const makeService = (store: RawStore) => ({ use: (op: string, fn: (s: RawStore) => Promise) => withServiceLogging( `user_store.${op}`, - () => new UserStoreError(), + (failure) => userStoreErrorFromFailure(op, failure), tryPromiseService(() => fn(store)), ), }); diff --git a/apps/cloud/src/auth/errors.ts b/apps/cloud/src/auth/errors.ts index 927ae59f6..debf0b063 100644 --- a/apps/cloud/src/auth/errors.ts +++ b/apps/cloud/src/auth/errors.ts @@ -1,10 +1,103 @@ import { Data, Effect, Option, Predicate, Schema } from "effect"; +// How a user-store call failed, classified from the driver cause. Safe to put +// on the wire and on a Sentry tag: it names a failure MODE, never a query, a +// value, or a customer. +export const USER_STORE_FAILURE_REASONS = [ + "connect_timeout", + "connection_closed", + "query", + "unknown", +] as const; + +export type UserStoreFailureReason = (typeof USER_STORE_FAILURE_REASONS)[number]; + +/** + * The public failure of every cloud user-store call. + * + * It carries the two fields that make an issue diagnosable from the error + * alone: which store call failed, and how. Before those existed the error had + * an empty field set, so Sentry showed a titleless, messageless issue and the + * only cause detail (the pretty-printed Effect cause in a Sentry `extra`) is + * scrubbed server-side — the failing operation and the driver reason existed + * only in the trace store. Same shape as `WorkOSError.status`: a small, safe + * classification field threaded at the service boundary. + */ export class UserStoreError extends Schema.TaggedErrorClass()( "UserStoreError", - {}, + { + /** The store call that failed, e.g. `getOrganization`. */ + operation: Schema.String, + /** How it failed, classified from the driver cause chain. */ + reason: Schema.Literals(USER_STORE_FAILURE_REASONS), + }, { httpApiStatus: 500 }, -) {} +) { + override get message(): string { + return `user store ${this.operation} failed: ${this.reason}`; + } +} + +/** Reasons a retry can plausibly clear: the query never reached a healthy + * server. A `query` failure is deterministic and must not be retried. */ +export const isTransientUserStoreReason = (reason: UserStoreFailureReason): boolean => + reason === "connect_timeout" || reason === "connection_closed"; + +// postgres.js tags its connection failures with a string `code` +// (`CONNECT_TIMEOUT`, `CONNECTION_CLOSED`, …) and its query failures with the +// SQLSTATE. Drizzle re-throws both wrapped in its own "Failed query" error with +// the driver error in `.cause`, and the service adapter wraps that again, so +// the classification walks the chain rather than inspecting one level. +const MAX_CAUSE_DEPTH = 8; + +const REASON_BY_DRIVER_CODE: Readonly> = { + CONNECT_TIMEOUT: "connect_timeout", + ETIMEDOUT: "connect_timeout", + CONNECTION_CLOSED: "connection_closed", + CONNECTION_ENDED: "connection_closed", + CONNECTION_DESTROYED: "connection_closed", + ECONNREFUSED: "connection_closed", + ECONNRESET: "connection_closed", +}; + +const stringCodeOf = (value: unknown): string | undefined => { + if (typeof value !== "object" || value === null) return undefined; + const code = (value as { readonly code?: unknown }).code; + return typeof code === "string" ? code : undefined; +}; + +const driverCodesOf = (failure: unknown): readonly string[] => { + const codes: string[] = []; + let current: unknown = isServiceAdapterError(failure) ? failure.cause : failure; + for ( + let depth = 0; + depth < MAX_CAUSE_DEPTH && current !== undefined && current !== null; + depth++ + ) { + const code = stringCodeOf(current); + if (code !== undefined) codes.push(code); + current = + typeof current === "object" ? (current as { readonly cause?: unknown }).cause : undefined; + } + return codes; +}; + +/** Classify a raw store failure. A recognised connection code wins; any other + * driver code (a SQLSTATE) is a deterministic query failure; nothing at all is + * `unknown`. */ +export const userStoreReasonFromCause = (failure: unknown): UserStoreFailureReason => { + const codes = driverCodesOf(failure); + for (const code of codes) { + const reason = REASON_BY_DRIVER_CODE[code]; + if (reason !== undefined) return reason; + } + return codes.length > 0 ? "query" : "unknown"; +}; + +/** Build the public `UserStoreError` for a store-adapter failure, naming the + * operation the call site already knows and classifying the driver cause. */ +export const userStoreErrorFromFailure = (operation: string, failure: unknown): UserStoreError => + new UserStoreError({ operation, reason: userStoreReasonFromCause(failure) }); export class WorkOSError extends Schema.TaggedErrorClass()( "WorkOSError", diff --git a/apps/cloud/src/auth/user-store-error.node.test.ts b/apps/cloud/src/auth/user-store-error.node.test.ts new file mode 100644 index 000000000..880c6fa37 --- /dev/null +++ b/apps/cloud/src/auth/user-store-error.node.test.ts @@ -0,0 +1,128 @@ +// `UserStoreError` is the public failure of every cloud user-store call, and it +// used to carry NOTHING: no operation, no reason, no message. Sentry grouped +// every store failure — a connect timeout, a missing table, a constraint +// violation — into one titleless issue, and the only cause detail (the +// pretty-printed Effect cause stuffed into a Sentry `extra`) is scrubbed +// server-side, so the issue was undiagnosable from Sentry alone. +// +// This pins the two safe classification fields it must carry instead: +// `operation` (already in hand at the call site) and `reason` (classified from +// the driver cause the way `statusFromWorkOSCause` classifies WorkOS causes). +import { createServer, type Server, type Socket } from "node:net"; + +import { describe, expect, it } from "@effect/vitest"; +import { Effect, Layer, Result } from "effect"; +import postgres from "postgres"; + +import { UserStoreService } from "./context"; +import { ServiceAdapterError, userStoreReasonFromCause, type UserStoreError } from "./errors"; +import { DbService } from "../db/db"; + +// A socket that completes the TCP handshake and then says nothing — exactly +// what a wedged Hyperdrive/Postgres endpoint looks like to postgres.js, and the +// only way to obtain the driver's REAL connect-timeout error object rather than +// a hand-written imitation of it. +const blackHolePort = async (): Promise<{ + readonly port: number; + readonly close: () => Promise; +}> => { + const sockets = new Set(); + const server: Server = createServer((socket) => { + sockets.add(socket); + socket.on("close", () => sockets.delete(socket)); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const address = server.address(); + // oxlint-disable-next-line executor/no-error-constructor, executor/no-try-catch-or-throw -- boundary: the fixture cannot run without a bound port + if (address === null || typeof address === "string") throw new Error("no port"); + return { + port: address.port, + close: () => + new Promise((resolve) => { + // The timed-out client leaves its half-open socket attached; without + // dropping it first, `close` waits for a peer that will never speak. + for (const socket of sockets) socket.destroy(); + server.close(() => resolve()); + }), + }; +}; + +const realConnectTimeoutError = async (): Promise => { + const hole = await blackHolePort(); + const sql = postgres(`postgresql://postgres:postgres@127.0.0.1:${hole.port}/postgres`, { + max: 1, + connect_timeout: 1, + fetch_types: false, + onnotice: () => undefined, + }); + // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: capturing the driver's own thrown error IS the fixture + try { + await sql`select 1`; + // oxlint-disable-next-line executor/no-error-constructor, executor/no-try-catch-or-throw -- boundary: the fixture is unusable if the socket answered + throw new Error("expected the connect to time out"); + } catch (error) { + return error; + } finally { + // oxlint-disable-next-line executor/no-promise-catch -- boundary: best-effort teardown of a connection that never opened + await sql.end({ timeout: 0 }).catch(() => undefined); + await hole.close(); + } +}; + +// Drizzle re-throws driver failures wrapped in its own error with the failing +// SQL in the message and the driver error in `.cause` — the shape production +// actually reports (`Failed query: select … -> write CONNECT_TIMEOUT …`). +const wrappedLikeDrizzle = (cause: unknown): Error => + // oxlint-disable-next-line executor/no-error-constructor -- boundary: reproducing the driver wrapper shape the store really fails with + Object.assign(new Error('Failed query: select "id" from "organizations" where "id" = $1'), { + cause, + }); + +const stubDb = Layer.succeed(DbService)({ db: {} as never }); + +const failingStoreCall = ( + operation: string, + failure: unknown, +): Effect.Effect> => + Effect.gen(function* () { + const users = yield* UserStoreService; + // oxlint-disable-next-line executor/no-promise-reject -- boundary: the store adapter lifts a REJECTING promise; rejecting is the fixture + return yield* users.use(operation, () => Promise.reject(failure)); + }).pipe( + Effect.provide(UserStoreService.Live.pipe(Layer.provide(stubDb))), + Effect.result, + ) as Effect.Effect>; + +describe("UserStoreError classification", () => { + it("classifies the driver cause chain", async () => { + const driverError = await realConnectTimeoutError(); + + expect(userStoreReasonFromCause(wrappedLikeDrizzle(driverError))).toBe("connect_timeout"); + expect( + userStoreReasonFromCause(new ServiceAdapterError({ cause: wrappedLikeDrizzle(driverError) })), + ).toBe("connect_timeout"); + expect( + userStoreReasonFromCause( + // oxlint-disable-next-line executor/no-error-constructor -- boundary: a SQLSTATE failure shape + Object.assign(new Error('relation "organizations" does not exist'), { code: "42P01" }), + ), + ).toBe("query"); + // oxlint-disable-next-line executor/no-error-constructor -- boundary: a failure with nothing to classify + expect(userStoreReasonFromCause(new Error("something else"))).toBe("unknown"); + }, 20_000); + + it("carries operation, reason and a stable message on the public error", async () => { + const driverError = await realConnectTimeoutError(); + + const result = await Effect.runPromise( + failingStoreCall("getOrganization", wrappedLikeDrizzle(driverError)), + ); + + expect(Result.isFailure(result)).toBe(true); + if (!Result.isFailure(result)) return; + expect(result.failure.operation).toBe("getOrganization"); + expect(result.failure.reason).toBe("connect_timeout"); + expect(result.failure.message).toContain("getOrganization"); + expect(result.failure.message).toContain("connect_timeout"); + }, 20_000); +}); diff --git a/apps/cloud/src/mcp/agent-handler.ts b/apps/cloud/src/mcp/agent-handler.ts index 6d46bbd80..0ec697c91 100644 --- a/apps/cloud/src/mcp/agent-handler.ts +++ b/apps/cloud/src/mcp/agent-handler.ts @@ -27,9 +27,12 @@ import { mcpSessionStub } from "@executor-js/cloudflare/mcp/session-stub"; import { wrapMcpSseResponse } from "../observability/memory-metrics"; import { WorkerTelemetryLive } from "../observability/telemetry"; import { cloudMcpAuth } from "./auth-provider"; +import { isMcpSessionMetaUnavailable } from "./session-meta"; import { McpSessionDOSqlite } from "./session-durable-object"; import { parseTraceparent } from "./traceparent"; +const MCP_SESSION_UNAVAILABLE_MESSAGE = "Session storage temporarily unavailable - please retry"; + const corsPreflightResponse = (): Response => new Response(null, { status: 204, @@ -171,6 +174,13 @@ const propsForPrincipal = ( return { session: { organizationId: principal.organizationId, + // The org record the live membership check resolved microseconds ago, + // handed to the session DO so it never opens a connection of its own to + // re-read it. An unnamed org (no auth plane could resolve one) is + // omitted rather than sent empty, so the DO can tell "not carried" from + // "carried, and blank". + ...(principal.organizationName ? { organizationName: principal.organizationName } : {}), + ...(principal.organizationSlug ? { organizationSlug: principal.organizationSlug } : {}), userId: principal.accountId, elicitationMode: readElicitationMode(request), artifactsEnabled: readArtifactsEnabled(request), @@ -306,6 +316,23 @@ export const makeCloudMcpAgentHandler = () => { // vocabulary — a deploy, a storage timeout, a cancelled // blockConcurrencyWhile — which reaches here through the agents SDK's own // `getServerByName` retry and used to 500 identically. + // The session DO could not reach the organization directory to name the + // org (after its own bounded retry). Transient by construction, so it + // gets the same retryable envelope a WorkOS blip gets on the auth path — + // not an unclassified 500 the agents SDK then retries the whole DO + // operation over, which is what turned a 10s connect timeout into a + // half-minute client hang. + // + // Checked BEFORE the platform classifier: this is an application failure + // that merely escapes through the same seam, and it names its own cause. + // The classifier only recognizes the runtime's own reset vocabulary, so + // the two never contend — the order just keeps it that way if either + // vocabulary grows. + if (isMcpSessionMetaUnavailable(error)) { + return jsonRpcErrorBody(503, -32001, MCP_SESSION_UNAVAILABLE_MESSAGE, { + retryAfterSeconds: UNAVAILABLE_RETRY_AFTER_SECONDS, + }); + } const failure = classifyDurableObjectError(error); if (!failure) { // oxlint-disable-next-line executor/no-try-catch-or-throw -- adapter boundary: rethrow anything that isn't a recognized platform failure to the Workers runtime unchanged diff --git a/apps/cloud/src/mcp/auth-provider.test.ts b/apps/cloud/src/mcp/auth-provider.test.ts index 454573ade..bcd6009ef 100644 --- a/apps/cloud/src/mcp/auth-provider.test.ts +++ b/apps/cloud/src/mcp/auth-provider.test.ts @@ -70,9 +70,11 @@ const stubOrgAuthNoMembership = Layer.succeed(McpOrganizationAuth)({ authorize: () => Effect.succeed(null), }); -// `authorize` SUCCEEDS with an org id — active membership. +// `authorize` SUCCEEDS with the resolved org record — active membership. The +// record, not just the id: the session props carry the org's name and slug so +// the session DO never re-reads the row. const stubOrgAuthActive = Layer.succeed(McpOrganizationAuth)({ - authorize: () => Effect.succeed(ORG_ID), + authorize: () => Effect.succeed({ id: ORG_ID, name: "Stub Org", slug: "stub-org" }), }); // A failure that is not a WorkOSError at all (e.g. the per-request DB layer diff --git a/apps/cloud/src/mcp/auth-provider.ts b/apps/cloud/src/mcp/auth-provider.ts index 8a5944f05..92384da16 100644 --- a/apps/cloud/src/mcp/auth-provider.ts +++ b/apps/cloud/src/mcp/auth-provider.ts @@ -55,6 +55,7 @@ import { McpAuthLive, McpOrganizationAuth, McpOrganizationAuthLive, + type AuthorizedMcpOrganization, type McpAuthResult, type VerifiedToken, } from "./auth"; @@ -88,16 +89,24 @@ const ORGANIZATION_AUTHORIZE_UNAVAILABLE = /** * Enrich a cloud {@link VerifiedToken} (which carries only accountId + - * organizationId) into the full {@link Principal} the seam validates. The - * envelope only uses `accountId` + `organizationId` for ownership; cloud - * resolves org name/email inside the DO, so the cosmetic identity fields carry - * empty placeholders. `organizationId` is guaranteed non-null here because the - * Forbidden branch already rejected the no-org case before Authenticated. + * organizationId) into the full {@link Principal} the seam validates. + * + * The org name and slug come from the record the live membership check just + * resolved — this is the whole point of `authorize` returning the record rather + * than an id. They used to be dropped here (`organizationName: ""`), which left + * the session Durable Object to re-read the same row over a fresh database + * connection on every cold init. `email` stays a placeholder: the envelope only + * uses `accountId` + `organizationId` for ownership, and nothing downstream + * reads it. */ -const principalFromToken = (token: VerifiedToken, organizationId: string): Principal => ({ +const principalFromToken = ( + token: VerifiedToken, + organization: AuthorizedMcpOrganization, +): Principal => ({ accountId: token.accountId, - organizationId, - organizationName: "", + organizationId: organization.id, + organizationName: organization.name, + ...(organization.slug === undefined ? {} : { organizationSlug: organization.slug }), email: "", name: null, avatarUrl: null, @@ -223,9 +232,9 @@ export const cloudMcpAuthProviderLayer: Layer.Layer< // caller genuinely holds no active membership (revoked / never a member) // — a real Forbidden, which the handler may act on by condemning the // session. - const organizationId = authorizeResult.success; - if (!organizationId) return forbidden(NO_ORGANIZATION_MESSAGE, -32001); - return authenticated(principalFromToken(token, organizationId)); + const organization = authorizeResult.success; + if (!organization) return forbidden(NO_ORGANIZATION_MESSAGE, -32001); + return authenticated(principalFromToken(token, organization)); }); const toOutcome = (request: Request, result: McpAuthResult): Effect.Effect => { diff --git a/apps/cloud/src/mcp/auth.ts b/apps/cloud/src/mcp/auth.ts index 347a79faf..ee6bb8ca1 100644 --- a/apps/cloud/src/mcp/auth.ts +++ b/apps/cloud/src/mcp/auth.ts @@ -160,19 +160,32 @@ export class McpAuth extends Context.Service< } >()("@executor-js/cloud/McpAuth") {} +/** + * The organization an MCP request was authorized against. The full record, not + * just its id: the same request later needs the org's display name and slug to + * open a session, and re-reading the row for them (from the session Durable + * Object, on a fresh database connection) is a redundant failure point on a + * request that already has the answer. + */ +export type AuthorizedMcpOrganization = { + readonly id: string; + readonly name: string; + readonly slug?: string; +}; + export class McpOrganizationAuth extends Context.Service< McpOrganizationAuth, { /** * Authorize `accountId` against an org SELECTOR — a WorkOS org id * (`org_…`, from the token or a legacy URL) or the org's URL slug (the - * form the install card prints). Returns the resolved org id when the + * form the install card prints). Returns the resolved organization when the * caller holds an active membership, `null` otherwise. */ readonly authorize: ( accountId: string, organizationSelector: string, - ) => Effect.Effect; + ) => Effect.Effect; } >()("@executor-js/cloud/McpOrganizationAuth") {} @@ -216,7 +229,9 @@ export const McpOrganizationAuthLive = Layer.succeed(McpOrganizationAuth)({ Effect.flatMap((organizationId) => organizationId ? authorizeOrganization(accountId, organizationId).pipe( - Effect.map((org) => (org ? org.id : null)), + Effect.map((org) => + org ? ({ id: org.id, name: org.name, slug: org.slug } as const) : null, + ), ) : Effect.succeed(null), ), diff --git a/apps/cloud/src/mcp/session-durable-object.ts b/apps/cloud/src/mcp/session-durable-object.ts index 50354810f..5ff5ee2f0 100644 --- a/apps/cloud/src/mcp/session-durable-object.ts +++ b/apps/cloud/src/mcp/session-durable-object.ts @@ -58,7 +58,7 @@ import { buildExecuteDescription, type ResumeResponse } from "@executor-js/execu // `SessionAuthLive` instead.) import { CoreSharedServices } from "../auth/workos"; import { UserStoreService } from "../auth/context"; -import { resolveOrganization } from "../auth/organization"; +import { resolveSessionMetaForToken } from "./session-meta"; import { DbService, combinedSchema, @@ -107,10 +107,6 @@ type CloudSessionDbHandle = DbServiceShape & { readonly end: () => Promise; }; -class OrganizationNotFoundError extends Data.TaggedError("OrganizationNotFoundError")<{ - readonly organizationId: string; -}> {} - class McpModelResumeForwardError extends Data.TaggedError("McpModelResumeForwardError")<{ readonly cause: unknown; }> {} @@ -209,28 +205,27 @@ export class McpSessionDOSqlite extends McpAgentSessionDOBase { + protected override resolveSessionMeta( + token: McpSessionInit, + storedMeta: SessionMeta | null, + ): Effect.Effect { + // The database handle is opened LAZILY: on the props and stored paths — the + // overwhelming majority of inits — nothing here touches Postgres at all, + // which is the whole point. postgres.js only dials on first query, so + // building the handle costs nothing; `ensuring` still closes it. const dbHandle = makeEphemeralDb(); - return Effect.gen(function* () { - const org = yield* resolveOrganization(token.organizationId); - if (!org) { - return yield* new OrganizationNotFoundError({ organizationId: token.organizationId }); - } - return { - organizationId: org.id, - organizationName: org.name, - organizationSlug: org.slug, - userId: token.userId, - resource: token.resource, - elicitationMode: token.elicitationMode, - artifactsEnabled: token.artifactsEnabled, - searchToolsEnabled: token.searchToolsEnabled, - } satisfies SessionMeta; - }).pipe( + return resolveSessionMetaForToken(token, storedMeta).pipe( Effect.withSpan("McpSessionDOSqlite.resolveSessionMeta"), Effect.provide(makeSessionServices(dbHandle)), Effect.ensuring(Effect.promise(() => dbHandle.end())), - // oxlint-disable-next-line executor/no-effect-escape-hatch -- boundary: a vanished org is a defect; the worker already verified the bearer + // The base's `resolveSessionMeta` seam has no error channel, and a + // Durable Object's `init` can only reject its Promise — so a failure has + // to leave as a defect. What changed is WHAT leaves: an unreachable + // organization directory is now a bounded, classified + // `McpSessionMetaUnavailableError` whose message the worker recognises + // and renders as a retryable 503 (see `agent-handler.ts`), instead of an + // unclassified Postgres cause that produced a 500 and a client hang. + // oxlint-disable-next-line executor/no-effect-escape-hatch -- boundary: the DO init seam is Promise-only; the failure is classified before it dies Effect.orDie, ); } diff --git a/apps/cloud/src/mcp/session-meta.node.test.ts b/apps/cloud/src/mcp/session-meta.node.test.ts new file mode 100644 index 000000000..d1dcdac23 --- /dev/null +++ b/apps/cloud/src/mcp/session-meta.node.test.ts @@ -0,0 +1,167 @@ +// Where an MCP session's organization identity comes from, and what happens +// when the only remaining source — the database — is unreachable. +// +// The production defect: every cold session-DO init read the `organizations` +// row over a brand-new Postgres connection, even though the worker had resolved +// that exact row microseconds earlier on the same request, and even when the DO +// already held the answer in its own storage. A connect timeout on that +// unnecessary connection became an unclassified defect and killed `initialize`. +import { describe, expect, it } from "@effect/vitest"; +import { Effect, Layer, Predicate, Result } from "effect"; + +import { defaultMcpResource } from "@executor-js/host-mcp"; +import type { McpSessionInit, SessionMeta } from "@executor-js/cloudflare/mcp/agent-durable-object"; + +import { UserStoreService } from "../auth/context"; +import { UserStoreError } from "../auth/errors"; +import { WorkOSClient, type WorkOSClientService } from "../auth/workos"; +import { + isMcpSessionMetaUnavailable, + resolveSessionMetaForToken, + SESSION_META_DB_RETRIES, +} from "./session-meta"; + +const TOKEN: McpSessionInit = { + organizationId: "org_test", + userId: "user_test", + elicitationMode: "model", + resource: defaultMcpResource, + artifactsEnabled: true, +}; + +const STORED: SessionMeta = { + organizationId: "org_test", + organizationName: "Stored Org", + organizationSlug: "stored-org", + userId: "user_test", + resource: defaultMcpResource, +}; + +/** A user store that always fails the way a wedged Hyperdrive endpoint does, + * counting how many times it was asked. */ +const countingConnectTimeoutStore = (): { + readonly layer: Layer.Layer; + readonly calls: () => number; +} => { + let calls = 0; + return { + calls: () => calls, + layer: Layer.succeed(UserStoreService)({ + use: (operation: string) => + Effect.suspend(() => { + calls += 1; + return Effect.fail(new UserStoreError({ operation, reason: "connect_timeout" })); + }), + } as UserStoreService["Service"]), + }; +}; + +const namingStore = (): { + readonly layer: Layer.Layer; + readonly calls: () => number; +} => { + let calls = 0; + return { + calls: () => calls, + layer: Layer.succeed(UserStoreService)({ + use: (_operation: string, fn: (store: never) => Promise) => + Effect.suspend(() => { + calls += 1; + return Effect.promise(() => + fn({ + getOrganization: async (id: string) => ({ + id, + name: "Database Org", + slug: "database-org", + }), + } as never), + ); + }), + } as UserStoreService["Service"]), + }; +}; + +const unusedWorkOS = Layer.succeed( + WorkOSClient, + new Proxy({} as WorkOSClientService, { + get: (_target, prop) => () => Effect.die(`unexpected WorkOSClient.${String(prop)} call`), + }), +); + +// The props source — the one the overwhelming majority of inits take — is +// covered black-box by `e2e/cloud/mcp-session-cold-init.test.ts`. What stays +// here is what that scenario cannot reach: the sources it falls back to, and +// what an unreachable database does to them. +describe("resolveSessionMetaForToken", () => { + it("falls back to the meta this session already stored, without touching the store", async () => { + const store = countingConnectTimeoutStore(); + + const meta = await Effect.runPromise( + resolveSessionMetaForToken(TOKEN, STORED).pipe( + Effect.provide(Layer.mergeAll(store.layer, unusedWorkOS)), + ), + ); + + expect(meta.organizationName).toBe("Stored Org"); + expect(meta.organizationSlug).toBe("stored-org"); + expect(store.calls(), "a restore reuses what it persisted").toBe(0); + }); + + it("reads the database only when nothing else names the org", async () => { + const store = namingStore(); + + const meta = await Effect.runPromise( + resolveSessionMetaForToken(TOKEN, null).pipe( + Effect.provide(Layer.mergeAll(store.layer, unusedWorkOS)), + ), + ); + + expect(meta.organizationName).toBe("Database Org"); + expect(store.calls()).toBe(1); + }); + + it("retries a connect timeout a bounded number of times, then fails retryably", async () => { + const store = countingConnectTimeoutStore(); + + const result = await Effect.runPromise( + resolveSessionMetaForToken(TOKEN, null).pipe( + Effect.provide(Layer.mergeAll(store.layer, unusedWorkOS)), + Effect.result, + ), + ); + + expect(Result.isFailure(result), "an unreachable directory is a failure, not a defect").toBe( + true, + ); + if (!Result.isFailure(result)) return; + expect(Predicate.isTagged(result.failure, "McpSessionMetaUnavailableError")).toBe(true); + expect( + isMcpSessionMetaUnavailable(result.failure), + "the worker can recognise it across the Durable Object boundary", + ).toBe(true); + expect(store.calls(), "bounded: the first attempt plus its retries").toBe( + SESSION_META_DB_RETRIES + 1, + ); + }); + + it("does not retry a deterministic query failure", async () => { + let calls = 0; + const store = Layer.succeed(UserStoreService)({ + use: (operation: string) => + Effect.suspend(() => { + calls += 1; + return Effect.fail(new UserStoreError({ operation, reason: "query" })); + }), + } as UserStoreService["Service"]); + + const result = await Effect.runPromise( + resolveSessionMetaForToken(TOKEN, null).pipe( + Effect.provide(Layer.mergeAll(store, unusedWorkOS)), + Effect.result, + ), + ); + + expect(Result.isFailure(result)).toBe(true); + expect(calls, "a query the server answered will answer the same way again").toBe(1); + }); +}); diff --git a/apps/cloud/src/mcp/session-meta.ts b/apps/cloud/src/mcp/session-meta.ts new file mode 100644 index 000000000..a0bd1df81 --- /dev/null +++ b/apps/cloud/src/mcp/session-meta.ts @@ -0,0 +1,209 @@ +// --------------------------------------------------------------------------- +// Where an MCP session's organization identity comes from. +// +// Three sources, in the order they are preferred: +// +// props — the org record the worker resolved while authorizing THIS +// request, carried in the session props. Free, and always the +// freshest answer available. +// stored — the meta this session Durable Object already persisted for the +// same organization on an earlier init. Free, and correct for a +// session that already exists. +// database — an actual read of the `organizations` row. +// +// Only the third can fail, and it is the one that used to run unconditionally. +// Every cold DO init opened a brand-new Postgres connection through Hyperdrive +// purely to re-read a row the worker had loaded microseconds earlier on the +// same request and then discarded; when that connection could not be +// established the read hung for the full connect budget and the failure was +// turned into a defect, killing `initialize` on a database that was, at that +// same moment, answering the worker's own queries in milliseconds. +// +// So: prefer what the request already knows, and when the database genuinely is +// the only source, give it a bounded retry and a CLASSIFIED failure — the same +// transient-vs-definitive split the WorkOS membership check already uses in +// `auth-provider.ts` — instead of an unclassified defect. +// --------------------------------------------------------------------------- + +import { Data, Effect, Predicate, Result, Schedule } from "effect"; + +import type { McpSessionInit, SessionMeta } from "@executor-js/cloudflare/mcp/agent-durable-object"; + +import { UserStoreService } from "../auth/context"; +import { WorkOSClient } from "../auth/workos"; +import { + isDefinitiveWorkOSDenial, + isTransientUserStoreReason, + type UserStoreError, +} from "../auth/errors"; +import { resolveOrganization } from "../auth/organization"; + +export type SessionMetaSource = "props" | "stored" | "database"; + +export const SESSION_META_SOURCE_ATTRIBUTE = "mcp.session.meta_source"; + +/** + * The wire marker for "the organization directory was unreachable, try again". + * + * A Durable Object's `init` can only reject its Promise, so this travels to the + * worker as an ordinary error message across the DO boundary — the same + * mechanism the condemned-session `"destroyed"` abort already uses. The worker + * matches on it to answer a retryable 503 instead of letting an unclassified + * 500 (and the agents SDK's own DO-operation retry on top of it) turn a + * transient blip into a half-minute client hang. + */ +export const MCP_SESSION_META_UNAVAILABLE = "mcp_session_meta_unavailable"; + +export class McpSessionMetaUnavailableError extends Data.TaggedError( + "McpSessionMetaUnavailableError", +)<{ + readonly reason: string; + readonly attempts: number; +}> { + override get message(): string { + return `${MCP_SESSION_META_UNAVAILABLE}: organization directory unavailable (${this.reason}) after ${this.attempts} attempts`; + } +} + +export class OrganizationNotFoundError extends Data.TaggedError("OrganizationNotFoundError")<{ + readonly organizationId: string; +}> {} + +/** Does this failure, seen at the worker, mean "the session DO could not reach + * the organization directory"? Matches on the message because that is what + * survives the Durable Object RPC boundary. */ +export const isMcpSessionMetaUnavailable = (error: unknown): boolean => + // oxlint-disable-next-line executor/no-unknown-error-message -- adapter boundary: a Durable Object rejection reaches the worker as a plain Error whose message IS the signal (same mechanism as the "destroyed" abort) + Predicate.isError(error) && error.message.includes(MCP_SESSION_META_UNAVAILABLE); + +/** + * Retries for the database path. Deliberately small: the point is to ride out a + * single bad connection attempt, not to sit on a client's `initialize` while a + * database stays down. Exhausting them answers the client quickly and + * retryably, which is strictly better than hanging. + */ +export const SESSION_META_DB_RETRIES = 2; + +const RETRY_SCHEDULE = Schedule.both( + Schedule.exponential("200 millis"), + Schedule.recurs(SESSION_META_DB_RETRIES), +); + +const isUserStoreError = Predicate.isTagged("UserStoreError") as ( + error: unknown, +) => error is UserStoreError; + +/** + * Is this org-lookup failure worth another attempt? A connection that never + * opened is; a query the server answered with an error, or a WorkOS denial, is + * not — retrying those only burns the client's init budget. + */ +export const isRetryableOrganizationLookupFailure = (failure: unknown): boolean => { + if (isUserStoreError(failure)) return isTransientUserStoreReason(failure.reason); + if (isDefinitiveWorkOSDenial(failure)) return false; + // A WorkOS blip (429/5xx/timeout/network) — the same class the MCP auth path + // already treats as retryable. + return Predicate.isTagged(failure, "WorkOSError"); +}; + +const failureReason = (failure: unknown): string => + isUserStoreError(failure) ? failure.reason : "upstream"; + +const metaFromIdentity = ( + token: McpSessionInit, + organization: { readonly name: string; readonly slug?: string }, +): SessionMeta => ({ + organizationId: token.organizationId, + organizationName: organization.name, + ...(organization.slug === undefined ? {} : { organizationSlug: organization.slug }), + userId: token.userId, + resource: token.resource, + elicitationMode: token.elicitationMode, + artifactsEnabled: token.artifactsEnabled, + searchToolsEnabled: token.searchToolsEnabled, +}); + +/** + * Read the organization row, retrying only failures a retry can clear, and + * surfacing an exhausted retry as a typed, retryable failure rather than a + * defect. + */ +const organizationFromDatabase = ( + organizationId: string, +): Effect.Effect< + { readonly id: string; readonly name: string; readonly slug?: string }, + McpSessionMetaUnavailableError | OrganizationNotFoundError, + UserStoreService | WorkOSClient +> => + Effect.gen(function* () { + let attempts = 0; + // `Effect.retry` retries every failure it sees, so definitive failures are + // lifted OUT of the error channel before the schedule ever runs and only + // the retryable ones are left in it. + const attempt = Effect.suspend(() => { + attempts += 1; + return resolveOrganization(organizationId).pipe( + Effect.result, + Effect.flatMap((outcome) => + Result.isFailure(outcome) && isRetryableOrganizationLookupFailure(outcome.failure) + ? Effect.fail(outcome.failure) + : Effect.succeed(outcome), + ), + ); + }); + + // Two nested Results: the outer one is "the retries ran out", the inner one + // is "the lookup failed definitively on the first look". Both mean the same + // thing to the caller. + const retried = yield* attempt.pipe(Effect.retry(RETRY_SCHEDULE), Effect.result); + const outcome = Result.isFailure(retried) ? Result.fail(retried.failure) : retried.success; + + if (Result.isFailure(outcome)) { + return yield* new McpSessionMetaUnavailableError({ + reason: failureReason(outcome.failure), + attempts, + }); + } + const organization = outcome.success; + if (!organization) return yield* new OrganizationNotFoundError({ organizationId }); + return organization; + }).pipe( + Effect.withSpan("mcp.session.resolve_organization", { + attributes: { "mcp.auth.organization_id": organizationId }, + }), + ); + +/** + * Build the session meta for an init, preferring the org identity the request + * already carries over any read of the organization directory. + * + * `storedMeta` is this DO's own persisted meta for the SAME organization (the + * base Durable Object only offers a matching one), or `null`. + */ +export const resolveSessionMetaForToken = ( + token: McpSessionInit, + storedMeta: SessionMeta | null, +): Effect.Effect< + SessionMeta, + McpSessionMetaUnavailableError | OrganizationNotFoundError, + UserStoreService | WorkOSClient +> => + Effect.gen(function* () { + const fromProps = token.organizationName; + if (fromProps) { + yield* Effect.annotateCurrentSpan(SESSION_META_SOURCE_ATTRIBUTE, "props"); + return metaFromIdentity(token, { name: fromProps, slug: token.organizationSlug }); + } + + if (storedMeta?.organizationName) { + yield* Effect.annotateCurrentSpan(SESSION_META_SOURCE_ATTRIBUTE, "stored"); + return metaFromIdentity(token, { + name: storedMeta.organizationName, + slug: storedMeta.organizationSlug, + }); + } + + yield* Effect.annotateCurrentSpan(SESSION_META_SOURCE_ATTRIBUTE, "database"); + const organization = yield* organizationFromDatabase(token.organizationId); + return metaFromIdentity(token, organization); + }); diff --git a/apps/cloud/src/observability/index.ts b/apps/cloud/src/observability/index.ts index a56118a6c..296516ae9 100644 --- a/apps/cloud/src/observability/index.ts +++ b/apps/cloud/src/observability/index.ts @@ -174,14 +174,61 @@ export const sentryPayloadForCause = ( return { primary: input, pretty: null }; }; +// Safe classification fields our tagged errors carry (`UserStoreError.operation` +// / `.reason`, `WorkOSError.status`). They are promoted to Sentry TAGS because +// the pretty cause is only an `extra`, and Sentry's server-side scrubber +// replaces that extra with "[Filtered]" — leaving an issue with no failing +// operation and no reason in it at all. Tags survive, group, and are +// searchable. Values are failure modes and operation names; never a query, a +// value, or anything customer-derived. +const CLASSIFICATION_TAG_FIELDS = ["operation", "reason", "status"] as const; + +const MAX_CLASSIFICATION_TAG_CHARS = 120; + +const MAX_CAUSE_NESTING = 3; + +/** Every error value a cause carries, failures and defects alike. A defect can + * itself be a `Cause` (an inner `runPromise` rejecting with its own squashed + * cause), so the walk unwraps a few levels. */ +const errorValuesOf = (input: unknown, depth = 0): readonly unknown[] => { + if (depth >= MAX_CAUSE_NESTING) return []; + if (!Cause.isCause(input)) return [input]; + const values: unknown[] = []; + for (const reason of input.reasons) { + if (Cause.isFailReason(reason)) values.push(reason.error); + else if (Cause.isDieReason(reason)) values.push(...errorValuesOf(reason.defect, depth + 1)); + } + return values; +}; + +/** Read the classification fields off the tagged errors inside a cause. First + * writer wins, so the innermost reported error names the issue. */ +const classificationTagsOf = (input: unknown): Readonly> => { + const tags: Record = {}; + for (const candidate of errorValuesOf(input)) { + if (typeof candidate !== "object" || candidate === null) continue; + const tagged = candidate as Record; + for (const field of CLASSIFICATION_TAG_FIELDS) { + const value = tagged[field]; + if (tags[field] !== undefined) continue; + if (typeof value === "string" || typeof value === "number") { + tags[field] = String(value).slice(0, MAX_CLASSIFICATION_TAG_CHARS); + } + } + } + return tags; +}; + export const captureCause = ( input: unknown, context: OtelCorrelationContext | null = null, ): string | undefined => { const { primary, pretty } = sentryPayloadForCause(input); + const classification = classificationTagsOf(input); tagCurrentSentryScopeWithOtelContext(context); return Sentry.captureException(primary, (scope) => { tagSentryScopeWithOtelContext(scope, context); + for (const [key, value] of Object.entries(classification)) scope.setTag(key, value); if (pretty !== null) scope.setExtra("cause", pretty); return scope; }); diff --git a/apps/host-cloudflare/src/mcp/session-durable-object.ts b/apps/host-cloudflare/src/mcp/session-durable-object.ts index 5e373b0ca..4fccd6af9 100644 --- a/apps/host-cloudflare/src/mcp/session-durable-object.ts +++ b/apps/host-cloudflare/src/mcp/session-durable-object.ts @@ -102,9 +102,14 @@ export class McpSessionDO extends McpAgentSessionDOBase handle.close() }; } - protected override resolveSessionMeta(token: McpSessionInit): Effect.Effect { + protected override resolveSessionMeta( + token: McpSessionInit, + _storedMeta: SessionMeta | null, + ): Effect.Effect { // Single-tenant: every Access principal belongs to the one configured org, - // so there is nothing to resolve — stamp the configured org name. + // so there is nothing to resolve — stamp the configured org name. Nothing + // to reuse from the stored meta either; config is already the cheapest and + // freshest source there is. return Effect.succeed({ organizationId: token.organizationId, organizationName: this.cfConfig.organizationName, diff --git a/e2e/cloud/mcp-session-cold-init.test.ts b/e2e/cloud/mcp-session-cold-init.test.ts new file mode 100644 index 000000000..9219f10f2 --- /dev/null +++ b/e2e/cloud/mcp-session-cold-init.test.ts @@ -0,0 +1,100 @@ +// Cloud: the org identity a session needs at init travels in the session props +// the worker already resolved — it is NOT re-read from Postgres inside the +// session Durable Object. +// +// The defect this pins: on EVERY cold DO init the DO opened a brand-new +// Postgres connection purely to re-read the `organizations` row the worker had +// just read microseconds earlier on the same request (the worker threw it away: +// the auth principal hardcoded an empty organization name). When that fresh +// connection could not be established the whole `initialize` died — a hard, +// client-visible failure on a healthy database, because the only unhealthy +// thing was a socket nothing needed to open. +// +// The contract asserted here is deliberately about the DATA PATH, not the +// failure: `McpSessionDOSqlite.resolveSessionMeta` reports where the org +// identity came from, and the whole request performs exactly ONE org read (the +// worker's own authorization check) instead of two. +import { randomBytes } from "node:crypto"; + +import { expect } from "@effect/vitest"; +import { Effect } from "effect"; + +import { scenario } from "../src/scenario"; +import { Mcp, Target, Telemetry } from "../src/services"; +import type { Identity } from "../src/target"; + +const JSON_AND_SSE = "application/json, text/event-stream"; + +const INITIALIZE_REQUEST = { + jsonrpc: "2.0" as const, + id: 1, + method: "initialize", + params: { + protocolVersion: "2025-03-26", + capabilities: {}, + clientInfo: { name: "executor-e2e-session-cold-init", version: "0.0.1" }, + }, +}; + +const emailOf = (identity: Identity): string => identity.credentials?.email ?? identity.label; + +/** A client-supplied W3C trace context, so every span this one request produces + * — worker plane and Durable Object alike — is addressable by one trace id. */ +const newTraceContext = (): { readonly traceId: string; readonly traceparent: string } => { + const traceId = randomBytes(16).toString("hex"); + const spanId = randomBytes(8).toString("hex"); + return { traceId, traceparent: `00-${traceId}-${spanId}-01` }; +}; + +scenario( + "MCP session cold init · the org identity rides in the session props instead of a second Postgres read", + { timeout: 120_000 }, + Effect.gen(function* () { + const target = yield* Target; + const mcp = yield* Mcp; + const telemetry = yield* Telemetry; + + const identity = yield* target.newIdentity(); + const bearer = yield* mcp.mintBearer(emailOf(identity)); + const trace = newTraceContext(); + + const response = yield* Effect.promise(() => + fetch(target.mcpUrl, { + method: "POST", + headers: { + accept: JSON_AND_SSE, + "content-type": "application/json", + authorization: `Bearer ${bearer}`, + traceparent: trace.traceparent, + }, + body: JSON.stringify(INITIALIZE_REQUEST), + }), + ); + yield* Effect.promise(() => response.text()); + expect(response.status, "initialize opens a session").toBe(200); + expect(response.headers.get("mcp-session-id"), "the session id is minted").toBeTruthy(); + + // The DO really did resolve meta on this request (a cold init), and it + // resolved it from the props the worker handed over. + const resolveSpan = yield* telemetry.expectSpan({ + traceId: trace.traceId, + operation: "McpSessionDOSqlite.resolveSessionMeta", + }); + // One org read for the whole request: the worker's own authorization + // lookup. A second one means the DO reopened a connection to re-read a row + // the request already had. + const orgReads = yield* telemetry.searchSpans({ + traceId: trace.traceId, + operation: "user_store.getOrganization", + }); + expect( + orgReads.length, + "only the worker's authorization check reads the organization row", + ).toBe(1); + + expect( + resolveSpan.span.tags["mcp.session.meta_source"], + "the session meta comes from the props the worker already resolved", + ).toBe("props"); + }), +); diff --git a/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts b/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts index 125b3990d..05c257cb2 100644 --- a/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts +++ b/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts @@ -1,6 +1,6 @@ // oxlint-disable executor/no-error-constructor, executor/no-try-catch-or-throw -- boundary: the storage fake reproduces the plain Errors the Cloudflare runtime throws, and rejecting is the only way a DurableObjectStorage reports them import { describe, expect, it } from "@effect/vitest"; -import { Cause, Effect } from "effect"; +import { Cause, Effect, Exit } from "effect"; import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js"; import type { JSONRPCMessage, MessageExtraInfo } from "@modelcontextprotocol/sdk/types.js"; @@ -308,6 +308,119 @@ describe("McpAgentSessionDOBase apps capability persistence", () => { }); }); +// A cold restore used to re-resolve the org identity through the host's backing +// store (on cloud: a brand-new Postgres connection) BEFORE it ever looked at +// the meta this DO had already persisted for the very session it is restoring. +// A transient failure of that lookup killed `init` and the restore with it — +// for a row the DO was already holding. The DO's own storage is the +// authoritative copy of the org identity of a session it already minted, so it +// is offered to the host first; the host still rebuilds everything the CONNECT +// carries (resource, elicitation mode, capability flags) from the token. +describe("McpAgentSessionDOBase cold-restore meta reuse", () => { + type RestoreSession = { + ctx: MemoryStorage; + getSessionId: () => string; + loadSessionMeta: () => Effect.Effect; + resolveSessionMeta: ( + token: unknown, + storedMeta: SessionMeta | null, + ) => Effect.Effect; + resolveAndStoreSessionMeta: (token: unknown) => Effect.Effect; + }; + + const storedMeta: SessionMeta = { + organizationId: "org-1", + organizationName: "Org One", + organizationSlug: "org-one", + userId: "user-1", + resource: defaultMcpResource, + }; + + const token = { + organizationId: "org-1", + userId: "user-1", + elicitationMode: "model" as const, + resource: defaultMcpResource, + }; + + const makeRestoreSession = async ( + stored: SessionMeta | null, + ): Promise<{ session: RestoreSession; storage: MemoryStorage }> => { + const storage = new MemoryStorage(); + if (stored) await storage.put("session-meta", stored); + const session = Object.create(McpAgentSessionDOBase.prototype) as RestoreSession; + session.ctx = storage; + session.getSessionId = () => "session-restore"; + return { session, storage }; + }; + + // The host stands in for cloud with an unreachable database: it can only + // answer when the DO hands it what it already knows. + const hostWithUnreachableStore = + (seen: { storedMeta: SessionMeta | null; calls: number }) => + (tokenIn: unknown, stored: SessionMeta | null): Effect.Effect => { + seen.calls += 1; + seen.storedMeta = stored; + if (!stored) return Effect.die("organization lookup: CONNECT_TIMEOUT"); + const t = tokenIn as { readonly userId: string; readonly organizationId: string }; + return Effect.succeed({ + organizationId: t.organizationId, + organizationName: stored.organizationName, + organizationSlug: stored.organizationSlug, + userId: t.userId, + resource: defaultMcpResource, + } satisfies SessionMeta); + }; + + it("restores from its own stored meta when the backing store is unreachable", async () => { + const { session } = await makeRestoreSession(storedMeta); + const seen = { storedMeta: null as SessionMeta | null, calls: 0 }; + session.resolveSessionMeta = hostWithUnreachableStore(seen); + + const resolved = await Effect.runPromise(session.resolveAndStoreSessionMeta(token)); + + expect(seen.calls).toBe(1); + expect(seen.storedMeta).toMatchObject({ organizationId: "org-1", organizationName: "Org One" }); + expect(resolved.organizationName).toBe("Org One"); + expect(resolved.organizationSlug).toBe("org-one"); + }); + + // Stored meta is only a shortcut for the SAME organization. A session id + // reused across orgs must never inherit the previous org's identity. + it("offers nothing when the stored meta belongs to another organization", async () => { + const { session } = await makeRestoreSession({ ...storedMeta, organizationId: "org-other" }); + const seen = { storedMeta: null as SessionMeta | null, calls: 0 }; + session.resolveSessionMeta = hostWithUnreachableStore(seen); + + const exit = await Effect.runPromiseExit(session.resolveAndStoreSessionMeta(token)); + + expect(Exit.isFailure(exit)).toBe(true); + expect(seen.storedMeta).toBeNull(); + }); + + // A brand-new session has nothing stored; the host must resolve from scratch. + it("offers nothing on a first init", async () => { + const { session } = await makeRestoreSession(null); + const seen = { storedMeta: null as SessionMeta | null, calls: 0 }; + session.resolveSessionMeta = (tokenIn, stored) => { + seen.calls += 1; + seen.storedMeta = stored; + const t = tokenIn as { readonly userId: string; readonly organizationId: string }; + return Effect.succeed({ + organizationId: t.organizationId, + organizationName: "Freshly Resolved", + userId: t.userId, + resource: defaultMcpResource, + } satisfies SessionMeta); + }; + + const resolved = await Effect.runPromise(session.resolveAndStoreSessionMeta(token)); + + expect(seen.storedMeta).toBeNull(); + expect(resolved.organizationName).toBe("Freshly Resolved"); + }); +}); + describe("McpAgentSessionDOBase transport restore", () => { it("preserves hibernated response streams when a cold isolate starts", async () => { const session = await makeHarnessSession(); diff --git a/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts b/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts index a1dd27a3f..d7dfcb0e8 100644 --- a/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts +++ b/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts @@ -39,6 +39,13 @@ export type IncomingTraceHeaders = IncomingPropagationHeaders; export interface McpSessionInit { readonly organizationId: string; + /** The organization's display name, as the worker resolved it while + * authorizing this very request. Carried so the session DO never has to + * re-read a row the request already loaded. Absent when the auth plane could + * not name the org, in which case the host resolves it itself. */ + readonly organizationName?: string; + /** The organization's URL slug, from the same resolved record. */ + readonly organizationSlug?: string; readonly userId: string; readonly elicitationMode: McpElicitationMode; /** Whether this session serves artifacts, read off `?artifacts=` at connect @@ -247,7 +254,20 @@ export abstract class McpAgentSessionDOBase< protected abstract openSessionDb(): TDbHandle | Promise; - protected abstract resolveSessionMeta(token: McpSessionInit): Effect.Effect; + /** + * Build the session's {@link SessionMeta} for this init. + * + * `storedMeta` is what this DO already persisted for the SAME organization on + * an earlier init, or `null`. It is offered first so a host never has to + * re-resolve an org identity it is already holding — on cloud that resolution + * is a fresh Postgres connection, and a cold restore used to die on it. Every + * field the CONNECT carries (resource, elicitation mode, capability flags) + * still comes from `token`; only the org identity may be reused. + */ + protected abstract resolveSessionMeta( + token: McpSessionInit, + storedMeta: SessionMeta | null, + ): Effect.Effect; protected abstract buildMcpServer( sessionMeta: SessionMeta, @@ -524,12 +544,21 @@ export abstract class McpAgentSessionDOBase< private resolveAndStoreSessionMeta(token: McpSessionInit) { const self = this; return Effect.gen(function* () { - const resolved = yield* self.resolveSessionMeta(token); - // `init` runs again on every cold restore, and `resolveSessionMeta` - // rebuilds meta from the bearer token — which carries no negotiated - // capabilities. Carry the stored value forward, or restoring the session - // would erase the very bit that survives the restore. + // Read what this DO already knows BEFORE asking the host to resolve + // anything. `init` runs again on every cold restore, and the stored meta + // is this session's own durable record of the organization it was minted + // for — re-deriving that identity from the host's backing store is the + // single most failure-prone step of a restore (on cloud, a brand-new + // Postgres connection) and it is redundant for a session that already + // exists. It is offered only for the SAME organization; a token naming a + // different org resolves from scratch. const stored = yield* self.loadSessionMeta(); + const reusable = stored && stored.organizationId === token.organizationId ? stored : null; + // The stored meta also carries the capabilities negotiated at + // `initialize`, which the bearer token knows nothing about. Carry them + // forward, or restoring the session would erase the very bit that + // survives the restore. + const resolved = yield* self.resolveSessionMeta(token, reusable); const sessionMeta: SessionMeta = { ...resolved, ...(token.webOrigin ? { webOrigin: token.webOrigin } : {}), @@ -803,6 +832,11 @@ export abstract class McpAgentSessionDOBase< .bestEffortBookkeeping("init.mark_activity", () => self.markActivity()) .pipe(Effect.withSpan("McpSessionDO.markActivity")); }).pipe( + // ONE capture owner for an init defect. `init` can only reject its + // Promise, and the host's DO-level error instrumentation captures that + // rejection too — so the DO claims the cause below and the host drops its + // own echo, rather than both filing the same failure as two issues with + // the same trace id and span id. Effect.tapCause((cause) => Effect.gen(function* () { // A Cloudflare platform reset of an in-flight init is not a defect — From e675334a9392098475faefafcdc315030e6ba7f2 Mon Sep 17 00:00:00 2001 From: Rhys Sullivan <39114868+RhysSullivan@users.noreply.github.com> Date: Thu, 27 Aug 2026 13:32:31 -0700 Subject: [PATCH 2/3] Promote classification tags only from the errors that define them --- apps/cloud/src/observability/index.ts | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/apps/cloud/src/observability/index.ts b/apps/cloud/src/observability/index.ts index 296516ae9..d42f7f5ae 100644 --- a/apps/cloud/src/observability/index.ts +++ b/apps/cloud/src/observability/index.ts @@ -12,7 +12,7 @@ import * as Sentry from "@sentry/cloudflare"; import type { ErrorEvent, Scope } from "@sentry/cloudflare"; -import { Cause, Effect, Layer } from "effect"; +import { Cause, Effect, Layer, Predicate } from "effect"; import type * as Tracer from "effect/Tracer"; import { ErrorCapture } from "@executor-js/api"; @@ -183,6 +183,14 @@ export const sentryPayloadForCause = ( // value, or anything customer-derived. const CLASSIFICATION_TAG_FIELDS = ["operation", "reason", "status"] as const; +/** The errors those fields are read from. An allowlist, because the fields are + * only known to be safe on the errors this app defines. */ +const CLASSIFIED_ERROR_TAGS = [ + "UserStoreError", + "WorkOSError", + "McpSessionMetaUnavailableError", +] as const; + const MAX_CLASSIFICATION_TAG_CHARS = 120; const MAX_CAUSE_NESTING = 3; @@ -207,6 +215,12 @@ const classificationTagsOf = (input: unknown): Readonly> const tags: Record = {}; for (const candidate of errorValuesOf(input)) { if (typeof candidate !== "object" || candidate === null) continue; + // Only the errors whose fields we know are safe classifications. Without + // this, any foreign object in the cause that happens to carry an + // `operation` / `reason` / `status` field would have that value promoted + // onto a tag, and a foreign field is not known to be free of a query, a + // value, or anything customer-derived. + if (!CLASSIFIED_ERROR_TAGS.some((tag) => Predicate.isTagged(candidate, tag))) continue; const tagged = candidate as Record; for (const field of CLASSIFICATION_TAG_FIELDS) { const value = tagged[field]; From c849005608eb33f6748d4e3e1b30fb14945b7570 Mon Sep 17 00:00:00 2001 From: Rhys Sullivan <39114868+RhysSullivan@users.noreply.github.com> Date: Thu, 27 Aug 2026 14:46:24 -0700 Subject: [PATCH 3/3] Cold-init e2e: prove the session works and count the org reads deterministically --- e2e/cloud/mcp-session-cold-init.test.ts | 119 +++++++++++++++++++++--- 1 file changed, 105 insertions(+), 14 deletions(-) diff --git a/e2e/cloud/mcp-session-cold-init.test.ts b/e2e/cloud/mcp-session-cold-init.test.ts index 9219f10f2..c24d4368b 100644 --- a/e2e/cloud/mcp-session-cold-init.test.ts +++ b/e2e/cloud/mcp-session-cold-init.test.ts @@ -10,10 +10,27 @@ // client-visible failure on a healthy database, because the only unhealthy // thing was a socket nothing needed to open. // -// The contract asserted here is deliberately about the DATA PATH, not the -// failure: `McpSessionDOSqlite.resolveSessionMeta` reports where the org -// identity came from, and the whole request performs exactly ONE org read (the -// worker's own authorization check) instead of two. +// Two contracts, in the order a user meets them: +// +// 1. the session opens and WORKS — initialize mints a session id and the same +// id then serves `tools/list`, all off the identity the props carried; +// 2. the whole request performs exactly ONE organization read (the worker's +// own authorization check) and the DO opens no database connection of its +// own to name the org. +// +// (2) is asserted on the EXPORTED spans, and the two planes — worker and +// Durable Object — export independently, so the count is only taken once a +// worker-plane span for this same request has landed. Without that wait a +// still-pending worker batch would make a two-read request look like a one-read +// request, and the assertion would pass for the wrong reason. +// +// The failure half of the fix (an unreachable directory becomes a bounded retry +// and a retryable 503 rather than a bare 500) is not reachable from here: the +// harness runs one single-process PGlite shared by the worker and the DO, and +// the worker's own authorization reads that same row on the same request — so +// freezing the database fails the request before init identically on both sides +// of the fix. Those branches are pinned in +// `apps/cloud/src/mcp/session-meta.node.test.ts`. import { randomBytes } from "node:crypto"; import { expect } from "@effect/vitest"; @@ -36,6 +53,18 @@ const INITIALIZE_REQUEST = { }, }; +const INITIALIZED_NOTIFICATION = { + jsonrpc: "2.0" as const, + method: "notifications/initialized", +}; + +const TOOLS_LIST_REQUEST = { + jsonrpc: "2.0" as const, + id: 2, + method: "tools/list", + params: {}, +}; + const emailOf = (identity: Identity): string => identity.credentials?.email ?? identity.label; /** A client-supplied W3C trace context, so every span this one request produces @@ -46,6 +75,27 @@ const newTraceContext = (): { readonly traceId: string; readonly traceparent: st return { traceId, traceparent: `00-${traceId}-${spanId}-01` }; }; +const mcpPost = ( + url: string, + init: { + readonly bearer: string; + readonly sessionId?: string; + readonly traceparent?: string; + readonly body: unknown; + }, +): Promise => + fetch(url, { + method: "POST", + headers: { + accept: JSON_AND_SSE, + "content-type": "application/json", + authorization: `Bearer ${init.bearer}`, + ...(init.sessionId ? { "mcp-session-id": init.sessionId } : {}), + ...(init.traceparent ? { traceparent: init.traceparent } : {}), + }, + body: JSON.stringify(init.body), + }); + scenario( "MCP session cold init · the org identity rides in the session props instead of a second Postgres read", { timeout: 120_000 }, @@ -58,28 +108,56 @@ scenario( const bearer = yield* mcp.mintBearer(emailOf(identity)); const trace = newTraceContext(); + // ---- 1. the session opens, and it works ------------------------------- const response = yield* Effect.promise(() => - fetch(target.mcpUrl, { - method: "POST", - headers: { - accept: JSON_AND_SSE, - "content-type": "application/json", - authorization: `Bearer ${bearer}`, - traceparent: trace.traceparent, - }, - body: JSON.stringify(INITIALIZE_REQUEST), + mcpPost(target.mcpUrl, { + bearer, + traceparent: trace.traceparent, + body: INITIALIZE_REQUEST, }), ); yield* Effect.promise(() => response.text()); expect(response.status, "initialize opens a session").toBe(200); - expect(response.headers.get("mcp-session-id"), "the session id is minted").toBeTruthy(); + const sessionId = response.headers.get("mcp-session-id"); + expect(sessionId, "the session id is minted").toBeTruthy(); + const initialized = yield* Effect.promise(() => + mcpPost(target.mcpUrl, { + bearer, + sessionId: sessionId ?? "", + body: INITIALIZED_NOTIFICATION, + }), + ); + yield* Effect.promise(() => initialized.text()); + expect(initialized.status, "the client completes the handshake").toBe(202); + + // The session built from the props-carried identity actually serves work — + // the production symptom was an `initialize` that never got this far. + const tools = yield* Effect.promise(() => + mcpPost(target.mcpUrl, { + bearer, + sessionId: sessionId ?? "", + body: TOOLS_LIST_REQUEST, + }), + ); + const toolsBody = yield* Effect.promise(() => tools.text()); + expect(tools.status, "the session serves requests once open").toBe(200); + expect(toolsBody, "the session advertises the execute tool").toContain("execute"); + + // ---- 2. one organization read for the whole init request -------------- // The DO really did resolve meta on this request (a cold init), and it // resolved it from the props the worker handed over. const resolveSpan = yield* telemetry.expectSpan({ traceId: trace.traceId, operation: "McpSessionDOSqlite.resolveSessionMeta", }); + expect(resolveSpan.span.status, "the cold init resolved its meta without failing").toBe("ok"); + + // The worker plane exports on its own batch, independently of the DO's. + // Wait for a worker-plane span from this same request before counting, or + // an unflushed worker batch would hide the very read being counted. + yield* telemetry.expectSpan({ traceId: trace.traceId, operation: "mcp.request" }); + // One org read for the whole request: the worker's own authorization // lookup. A second one means the DO reopened a connection to re-read a row // the request already had. @@ -92,6 +170,19 @@ scenario( "only the worker's authorization check reads the organization row", ).toBe(1); + // …and the DO's own database step never ran at all. `resolveSessionMeta` + // has already landed and this span is its child, so its absence here is + // absence, not lag. + const doDatabaseReads = yield* telemetry.searchSpans({ + traceId: trace.traceId, + operation: "mcp.session.resolve_organization", + }); + expect( + doDatabaseReads.length, + "the session DO opens no connection of its own to name the organization", + ).toBe(0); + + // …because the identity it used is the one the worker handed it. expect( resolveSpan.span.tags["mcp.session.meta_source"], "the session meta comes from the props the worker already resolved",