// 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 themimport { afterEach, beforeEach, describe, expect, it } from "@effect/vitest";import { Cause, Effect, Exit, Schema } from "effect";import type * as Tracer from "effect/Tracer";import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";import type { JSONRPCMessage, MessageExtraInfo } from "@modelcontextprotocol/sdk/types.js"; import { defaultMcpResource } from "@executor-js/host-mcp";import type { ExecutionEngine, ExecutionResult, ResumeResponse } from "@executor-js/execution"; import { McpAgentSessionDOBase, type McpApprovalOwner, type McpSessionModelResumeResult, type SessionMeta,} from "./agent-session-durable-object";import { currentInFlightColdBuildCount, currentResidentRuntimeCount, EVICTION_REQUEST_GRACE_MS, markEvictionRequested, pickEvictionCandidate, registerResidentSession, releaseResidentSession, resetInFlightColdBuildCountForTest, resetResidentRuntimeCountForTest, resetResidentSessionRegistryForTest, residentSessionIdsForTest, touchResidentSession,} from "./session-runtime-residency"; class MemoryStorage { private readonly data = new Map<string, unknown>(); alarm: number | undefined; private idName: string | undefined = "streamable-http:session-reconnect"; readonly sql = { exec: () => [], }; async get<T>(key: string): Promise<T | undefined> { return this.data.get(key) as T | undefined; } async put(key: string, value: unknown): Promise<void> { this.data.set(key, value); } async setAlarm(time: number | Date): Promise<void> { this.alarm = typeof time === "number" ? time : time.getTime(); } async deleteAlarm(): Promise<void> { this.alarm = undefined; } async delete(key: string | readonly string[]): Promise<void> { if (typeof key === "string") { this.data.delete(key); return; } for (const entry of key) { this.data.delete(entry); } } async deleteAll(): Promise<void> { this.data.clear(); } async list<T>( options: { readonly prefix?: string; readonly limit?: number } = {}, ): Promise<Map<string, T>> { const rows = new Map<string, T>(); for (const [key, value] of this.data) { if (options.prefix && !key.startsWith(options.prefix)) continue; rows.set(key, value as T); if (options.limit && rows.size >= options.limit) break; } return rows; } async blockConcurrencyWhile<T>(callback: () => T | Promise<T>): Promise<T> { return callback(); } get id(): { readonly name: string | undefined } { return { name: this.idName }; } /** * Model a Durable Object invocation the runtime does not give a * `ctx.id.name` — the shape an alarm fires in when it is running against an * alarm record that never carried one. */ withoutIdName(): this { this.idName = undefined; return this; } /** Give this storage's Durable Object a distinct name, so multiple harness * sessions in the same test resolve to distinct `sessionIdForTelemetry()` * values instead of all colliding on the default. */ withIdName(name: string): this { this.idName = name; return this; } get storage(): MemoryStorage { return this; } private readonly waitUntilPromises: Promise<unknown>[] = []; /** * `ctx.waitUntil` extends work past the response instead of blocking it — * cap eviction requests it. A real DO holds the runtime open until every * queued promise settles; this fake instead just remembers them so a test * can explicitly wait for the background work it cares about via * `drainWaitUntil`, rather than the two racing unpredictably. */ waitUntil(promise: Promise<unknown>): void { this.waitUntilPromises.push(promise); } /** Test-only: await every `waitUntil`-queued promise queued so far, * including ones queued BY those promises while draining (an eviction * request settling can itself queue more background work). */ async drainWaitUntil(): Promise<void> { while (this.waitUntilPromises.length > 0) { const pending = this.waitUntilPromises.splice(0, this.waitUntilPromises.length); await Promise.allSettled(pending); } }} type HarnessSession = { alarm: () => Promise<void>; ctx: MemoryStorage; dbHandle: { readonly end: () => void } | null; engine: ExecutionEngine<Cause.YieldableError> | null; getConnections?: () => Iterable<unknown>; getSessionId: () => string; initialized: boolean; lastActivityMs: number; maxPausedSessionIdleMs: () => number; onStart: () => Promise<void>; pendingApprovalLeases: Map<string, never>; props: Record<string, unknown>; runMcpAgentOnStart: () => Promise<void>; server?: McpServer; sessionMeta: SessionMeta; sessionTimeoutMs: () => number; resumeExecutionForModel: ( executionId: string, identity: McpApprovalOwner, response: ResumeResponse, ) => Promise<McpSessionModelResumeResult>; validateMcpSessionOwner: (identity: { readonly accountId: string; readonly organizationId: string; }) => Promise<"ok" | "not_found" | "forbidden" | "terminated">;}; class StaleCloseTransport implements Transport { onclose?: () => void; onerror?: (error: Error) => void; onmessage?: (message: JSONRPCMessage, extra?: MessageExtraInfo) => void; async start(): Promise<void> {} async close(): Promise<void> {} async send(_message: JSONRPCMessage): Promise<void> {}} class RestoredTransport implements Transport { onclose?: () => void; onerror?: (error: Error) => void; onmessage?: (message: JSONRPCMessage, extra?: MessageExtraInfo) => void; async start(): Promise<void> {} async close(): Promise<void> { this.onclose?.(); } async send(_message: JSONRPCMessage): Promise<void> {}} const makeServer = () => new McpServer({ name: "executor-test", version: "1.0.0" }); /** The DO's structured logs are JSON lines; assertions read them back. */const decodeLogLine = Schema.decodeUnknownSync(Schema.fromJsonString(Schema.Unknown)); const makeDeferred = (): { readonly promise: Promise<void>; readonly resolve: () => void } => { let resolve: () => void = () => undefined; const promise = new Promise<void>((settle) => { resolve = settle; }); return { promise, resolve };}; type ResumeCall = { readonly executionId: string; readonly response: ResumeResponse;}; const completed = (result: unknown): ExecutionResult => ({ status: "completed", result: { result },}); const makeEngine = ( resultForResume: (executionId: string, response: ResumeResponse) => ExecutionResult | null = () => completed("resume-result"),): { readonly calls: ResumeCall[]; readonly engine: ExecutionEngine<Cause.YieldableError> } => { const calls: ResumeCall[] = []; return { calls, engine: { execute: () => Effect.succeed({ result: "execute-result" }), executeWithPause: () => Effect.succeed(completed("execute-result")), resume: (executionId, response) => Effect.sync(() => { calls.push({ executionId, response }); return resultForResume(executionId, response); }), getPausedExecution: () => Effect.succeed(null), pausedExecutionCount: () => Effect.succeed(0), hasPausedExecutions: () => Effect.succeed(false), getDescription: Effect.succeed("test engine"), // The fake forks nothing, so there is no sandbox fiber to end. shutdown: Effect.void, }, };}; const approval = { action: "accept", content: { approved: true },} satisfies ResumeResponse; const makeHarnessSession = async (): Promise<HarnessSession> => { const sessionId = "session-reconnect"; const sessionMeta: SessionMeta = { organizationId: "org-1", organizationName: "Org 1", userId: "user-1", resource: defaultMcpResource, }; const storage = new MemoryStorage(); const server = makeServer(); await server.connect(new StaleCloseTransport()); const session = Object.create(McpAgentSessionDOBase.prototype) as HarnessSession; session.ctx = storage; session.dbHandle = { end: () => undefined }; session.engine = makeEngine().engine; session.getSessionId = () => sessionId; session.initialized = true; session.lastActivityMs = Date.now() - 10; session.maxPausedSessionIdleMs = () => 1_000; session.pendingApprovalLeases = new Map<string, never>(); session.props = {}; session.server = server; session.sessionMeta = sessionMeta; session.sessionTimeoutMs = () => 1; session.runMcpAgentOnStart = async () => { const restored = session.server ?? makeServer(); session.server = restored; await restored.connect(new RestoredTransport()); session.engine = makeEngine().engine; session.initialized = true; }; return session;}; // The negotiated MCP-Apps capability arrives once, at `initialize`, and lives// in the rebuilt server's memory. These pin the storage round-trip that lets a// cold-restored session rebuild with it instead of silently downgrading every// artifact to a deep link.describe("McpAgentSessionDOBase apps capability persistence", () => { type CapabilitySession = HarnessSession & { persistAppsEnabled: (appsEnabled: boolean) => Effect.Effect<void>; loadSessionMeta: () => Effect.Effect<SessionMeta | null>; resolveSessionMeta: (token: unknown) => Effect.Effect<SessionMeta>; resolveAndStoreSessionMeta: (token: unknown) => Effect.Effect<SessionMeta>; }; const baseMeta: SessionMeta = { organizationId: "org-1", organizationName: "Org 1", userId: "user-1", resource: defaultMcpResource, }; const makeCapabilitySession = async ( stored: SessionMeta = baseMeta, ): Promise<{ session: CapabilitySession; storage: MemoryStorage }> => { const storage = new MemoryStorage(); await storage.put("session-meta", stored); const session = Object.create(McpAgentSessionDOBase.prototype) as CapabilitySession; session.ctx = storage; session.getSessionId = () => "session-caps"; return { session, storage }; }; it("persists the negotiated capability so a later restore can read it back", async () => { const { session, storage } = await makeCapabilitySession(); await Effect.runPromise(session.persistAppsEnabled(true)); expect(await storage.get<SessionMeta>("session-meta")).toMatchObject({ organizationId: "org-1", appsEnabled: true, }); }); it("records a client that loses apps support just as durably", async () => { const { session, storage } = await makeCapabilitySession({ ...baseMeta, appsEnabled: true }); await Effect.runPromise(session.persistAppsEnabled(false)); expect(await storage.get<SessionMeta>("session-meta")).toMatchObject({ appsEnabled: false }); }); // `init` runs again on every cold restore and rebuilds meta from the bearer // token, which carries no capabilities. If that overwrite won, restoring the // session would erase the very bit meant to survive it. it("carries the stored capability through the re-resolve on cold restore", async () => { const { session, storage } = await makeCapabilitySession({ ...baseMeta, appsEnabled: true }); // What the token resolves to: no `appsEnabled` anywhere in sight. session.resolveSessionMeta = () => Effect.succeed(baseMeta); const resolved = await Effect.runPromise( session.resolveAndStoreSessionMeta({ organizationId: "org-1", userId: "user-1" }), ); expect(resolved.appsEnabled).toBe(true); expect(await storage.get<SessionMeta>("session-meta")).toMatchObject({ appsEnabled: true }); }); it("leaves a session with no negotiated capability untouched", async () => { const { session, storage } = await makeCapabilitySession(); session.resolveSessionMeta = () => Effect.succeed(baseMeta); const resolved = await Effect.runPromise( session.resolveAndStoreSessionMeta({ organizationId: "org-1", userId: "user-1" }), ); expect(resolved.appsEnabled).toBeUndefined(); expect(await storage.get<SessionMeta>("session-meta")).not.toHaveProperty("appsEnabled"); }); // Persistence is best-effort observation of a capability, never a reason to // fail the session that was merely trying to render something. it("stays silent when there is no stored meta to merge into", async () => { const storage = new MemoryStorage(); const session = Object.create(McpAgentSessionDOBase.prototype) as CapabilitySession; session.ctx = storage; session.getSessionId = () => "session-caps"; await expect(Effect.runPromise(session.persistAppsEnabled(true))).resolves.toBeUndefined(); expect(await storage.get<SessionMeta>("session-meta")).toBeUndefined(); });}); // 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<SessionMeta | null>; resolveSessionMeta: ( token: unknown, storedMeta: SessionMeta | null, ) => Effect.Effect<SessionMeta>; resolveAndStoreSessionMeta: (token: unknown) => Effect.Effect<SessionMeta>; }; 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<SessionMeta> => { 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(); let closeCalls = 0; session.initialized = false; session.engine = null; session.dbHandle = null; delete session.server; session.getConnections = () => [ { close: () => { closeCalls += 1; }, }, ];Lines 1–500 of 1683. Load more
// oxlint-disable executor/no-error-constructor, executor/no-try-catch-or-throw -- boundary: the storage fake reproduces the plain Errors the Cloudflare runtime throws, and rejecting is the only way a DurableObjectStorage reports them import { afterEach, beforeEach, describe, expect, it } from "@effect/vitest"; import { Cause, Effect, Exit, Schema } from "effect"; import type * as Tracer from "effect/Tracer"; import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js"; import type { JSONRPCMessage, MessageExtraInfo } from "@modelcontextprotocol/sdk/types.js"; import { defaultMcpResource } from "@executor-js/host-mcp"; import type { ExecutionEngine, ExecutionResult, ResumeResponse } from "@executor-js/execution"; import { McpAgentSessionDOBase, type McpApprovalOwner, type McpSessionModelResumeResult, type SessionMeta, } from "./agent-session-durable-object"; import { currentInFlightColdBuildCount, currentResidentRuntimeCount, EVICTION_REQUEST_GRACE_MS, markEvictionRequested, pickEvictionCandidate, registerResidentSession, releaseResidentSession, resetInFlightColdBuildCountForTest, resetResidentRuntimeCountForTest, resetResidentSessionRegistryForTest, residentSessionIdsForTest, touchResidentSession, } from "./session-runtime-residency"; class MemoryStorage { private readonly data = new Map<string, unknown>(); alarm: number | undefined; private idName: string | undefined = "streamable-http:session-reconnect"; readonly sql = { exec: () => [], }; async get<T>(key: string): Promise<T | undefined> { return this.data.get(key) as T | undefined; } async put(key: string, value: unknown): Promise<void> { this.data.set(key, value); } async setAlarm(time: number | Date): Promise<void> { this.alarm = typeof time === "number" ? time : time.getTime(); } async deleteAlarm(): Promise<void> { this.alarm = undefined; } async delete(key: string | readonly string[]): Promise<void> { if (typeof key === "string") { this.data.delete(key); return; } for (const entry of key) { this.data.delete(entry); } } async deleteAll(): Promise<void> { this.data.clear(); } async list<T>( options: { readonly prefix?: string; readonly limit?: number } = {}, ): Promise<Map<string, T>> { const rows = new Map<string, T>(); for (const [key, value] of this.data) { if (options.prefix && !key.startsWith(options.prefix)) continue; rows.set(key, value as T); if (options.limit && rows.size >= options.limit) break; } return rows; } async blockConcurrencyWhile<T>(callback: () => T | Promise<T>): Promise<T> { return callback(); } get id(): { readonly name: string | undefined } { return { name: this.idName }; } /** * Model a Durable Object invocation the runtime does not give a * `ctx.id.name` — the shape an alarm fires in when it is running against an * alarm record that never carried one. */ withoutIdName(): this { this.idName = undefined; return this; } /** Give this storage's Durable Object a distinct name, so multiple harness * sessions in the same test resolve to distinct `sessionIdForTelemetry()` * values instead of all colliding on the default. */ withIdName(name: string): this { this.idName = name; return this; } get storage(): MemoryStorage { return this; } private readonly waitUntilPromises: Promise<unknown>[] = []; /** * `ctx.waitUntil` extends work past the response instead of blocking it — * cap eviction requests it. A real DO holds the runtime open until every * queued promise settles; this fake instead just remembers them so a test * can explicitly wait for the background work it cares about via * `drainWaitUntil`, rather than the two racing unpredictably. */ waitUntil(promise: Promise<unknown>): void { this.waitUntilPromises.push(promise); } /** Test-only: await every `waitUntil`-queued promise queued so far, * including ones queued BY those promises while draining (an eviction * request settling can itself queue more background work). */ async drainWaitUntil(): Promise<void> { while (this.waitUntilPromises.length > 0) { const pending = this.waitUntilPromises.splice(0, this.waitUntilPromises.length); await Promise.allSettled(pending); } } } type HarnessSession = { alarm: () => Promise<void>; ctx: MemoryStorage; dbHandle: { readonly end: () => void } | null; engine: ExecutionEngine<Cause.YieldableError> | null; getConnections?: () => Iterable<unknown>; getSessionId: () => string; initialized: boolean; lastActivityMs: number; maxPausedSessionIdleMs: () => number; onStart: () => Promise<void>; pendingApprovalLeases: Map<string, never>; props: Record<string, unknown>; runMcpAgentOnStart: () => Promise<void>; server?: McpServer; sessionMeta: SessionMeta; sessionTimeoutMs: () => number; resumeExecutionForModel: ( executionId: string, identity: McpApprovalOwner, response: ResumeResponse, ) => Promise<McpSessionModelResumeResult>; validateMcpSessionOwner: (identity: { readonly accountId: string; readonly organizationId: string; }) => Promise<"ok" | "not_found" | "forbidden" | "terminated">; }; class StaleCloseTransport implements Transport { onclose?: () => void; onerror?: (error: Error) => void; onmessage?: (message: JSONRPCMessage, extra?: MessageExtraInfo) => void; async start(): Promise<void> {} async close(): Promise<void> {} async send(_message: JSONRPCMessage): Promise<void> {} } class RestoredTransport implements Transport { onclose?: () => void; onerror?: (error: Error) => void; onmessage?: (message: JSONRPCMessage, extra?: MessageExtraInfo) => void; async start(): Promise<void> {} async close(): Promise<void> { this.onclose?.(); } async send(_message: JSONRPCMessage): Promise<void> {} } const makeServer = () => new McpServer({ name: "executor-test", version: "1.0.0" }); /** The DO's structured logs are JSON lines; assertions read them back. */ const decodeLogLine = Schema.decodeUnknownSync(Schema.fromJsonString(Schema.Unknown)); const makeDeferred = (): { readonly promise: Promise<void>; readonly resolve: () => void } => { let resolve: () => void = () => undefined; const promise = new Promise<void>((settle) => { resolve = settle; }); return { promise, resolve }; }; type ResumeCall = { readonly executionId: string; readonly response: ResumeResponse; }; const completed = (result: unknown): ExecutionResult => ({ status: "completed", result: { result }, }); const makeEngine = ( resultForResume: (executionId: string, response: ResumeResponse) => ExecutionResult | null = () => completed("resume-result"), ): { readonly calls: ResumeCall[]; readonly engine: ExecutionEngine<Cause.YieldableError> } => { const calls: ResumeCall[] = []; return { calls, engine: { execute: () => Effect.succeed({ result: "execute-result" }), executeWithPause: () => Effect.succeed(completed("execute-result")), resume: (executionId, response) => Effect.sync(() => { calls.push({ executionId, response }); return resultForResume(executionId, response); }), getPausedExecution: () => Effect.succeed(null), pausedExecutionCount: () => Effect.succeed(0), hasPausedExecutions: () => Effect.succeed(false), getDescription: Effect.succeed("test engine"), // The fake forks nothing, so there is no sandbox fiber to end. shutdown: Effect.void, }, }; }; const approval = { action: "accept", content: { approved: true }, } satisfies ResumeResponse; const makeHarnessSession = async (): Promise<HarnessSession> => { const sessionId = "session-reconnect"; const sessionMeta: SessionMeta = { organizationId: "org-1", organizationName: "Org 1", userId: "user-1", resource: defaultMcpResource, }; const storage = new MemoryStorage(); const server = makeServer(); await server.connect(new StaleCloseTransport()); const session = Object.create(McpAgentSessionDOBase.prototype) as HarnessSession; session.ctx = storage; session.dbHandle = { end: () => undefined }; session.engine = makeEngine().engine; session.getSessionId = () => sessionId; session.initialized = true; session.lastActivityMs = Date.now() - 10; session.maxPausedSessionIdleMs = () => 1_000; session.pendingApprovalLeases = new Map<string, never>(); session.props = {}; session.server = server; session.sessionMeta = sessionMeta; session.sessionTimeoutMs = () => 1; session.runMcpAgentOnStart = async () => { const restored = session.server ?? makeServer(); session.server = restored; await restored.connect(new RestoredTransport()); session.engine = makeEngine().engine; session.initialized = true; }; return session; }; // The negotiated MCP-Apps capability arrives once, at `initialize`, and lives // in the rebuilt server's memory. These pin the storage round-trip that lets a // cold-restored session rebuild with it instead of silently downgrading every // artifact to a deep link. describe("McpAgentSessionDOBase apps capability persistence", () => { type CapabilitySession = HarnessSession & { persistAppsEnabled: (appsEnabled: boolean) => Effect.Effect<void>; loadSessionMeta: () => Effect.Effect<SessionMeta | null>; resolveSessionMeta: (token: unknown) => Effect.Effect<SessionMeta>; resolveAndStoreSessionMeta: (token: unknown) => Effect.Effect<SessionMeta>; }; const baseMeta: SessionMeta = { organizationId: "org-1", organizationName: "Org 1", userId: "user-1", resource: defaultMcpResource, }; const makeCapabilitySession = async ( stored: SessionMeta = baseMeta, ): Promise<{ session: CapabilitySession; storage: MemoryStorage }> => { const storage = new MemoryStorage(); await storage.put("session-meta", stored); const session = Object.create(McpAgentSessionDOBase.prototype) as CapabilitySession; session.ctx = storage; session.getSessionId = () => "session-caps"; return { session, storage }; }; it("persists the negotiated capability so a later restore can read it back", async () => { const { session, storage } = await makeCapabilitySession(); await Effect.runPromise(session.persistAppsEnabled(true)); expect(await storage.get<SessionMeta>("session-meta")).toMatchObject({ organizationId: "org-1", appsEnabled: true, }); }); it("records a client that loses apps support just as durably", async () => { const { session, storage } = await makeCapabilitySession({ ...baseMeta, appsEnabled: true }); await Effect.runPromise(session.persistAppsEnabled(false)); expect(await storage.get<SessionMeta>("session-meta")).toMatchObject({ appsEnabled: false }); }); // `init` runs again on every cold restore and rebuilds meta from the bearer // token, which carries no capabilities. If that overwrite won, restoring the // session would erase the very bit meant to survive it. it("carries the stored capability through the re-resolve on cold restore", async () => { const { session, storage } = await makeCapabilitySession({ ...baseMeta, appsEnabled: true }); // What the token resolves to: no `appsEnabled` anywhere in sight. session.resolveSessionMeta = () => Effect.succeed(baseMeta); const resolved = await Effect.runPromise( session.resolveAndStoreSessionMeta({ organizationId: "org-1", userId: "user-1" }), ); expect(resolved.appsEnabled).toBe(true); expect(await storage.get<SessionMeta>("session-meta")).toMatchObject({ appsEnabled: true }); }); it("leaves a session with no negotiated capability untouched", async () => { const { session, storage } = await makeCapabilitySession(); session.resolveSessionMeta = () => Effect.succeed(baseMeta); const resolved = await Effect.runPromise( session.resolveAndStoreSessionMeta({ organizationId: "org-1", userId: "user-1" }), ); expect(resolved.appsEnabled).toBeUndefined(); expect(await storage.get<SessionMeta>("session-meta")).not.toHaveProperty("appsEnabled"); }); // Persistence is best-effort observation of a capability, never a reason to // fail the session that was merely trying to render something. it("stays silent when there is no stored meta to merge into", async () => { const storage = new MemoryStorage(); const session = Object.create(McpAgentSessionDOBase.prototype) as CapabilitySession; session.ctx = storage; session.getSessionId = () => "session-caps"; await expect(Effect.runPromise(session.persistAppsEnabled(true))).resolves.toBeUndefined(); expect(await storage.get<SessionMeta>("session-meta")).toBeUndefined(); }); }); // 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<SessionMeta | null>; resolveSessionMeta: ( token: unknown, storedMeta: SessionMeta | null, ) => Effect.Effect<SessionMeta>; resolveAndStoreSessionMeta: (token: unknown) => Effect.Effect<SessionMeta>; }; 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<SessionMeta> => { 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(); let closeCalls = 0; session.initialized = false; session.engine = null; session.dbHandle = null; delete session.server; session.getConnections = () => [ { close: () => { closeCalls += 1; }, }, ];