import { DurableObject } from "cloudflare:workers"; import { describe, expect, it } from "@effect/vitest"; import { Effect } from "effect"; import { McpExecutionOwnerDirectoryDO, mcpExecutionOwnerDirectoryFromNamespace, type McpExecutionOwnerDirectoryNamespace, type McpExecutionOwnerRecord, } from "./execution-owner-directory"; class FakeStorage implements DurableObjectStorage { private readonly values = new Map(); readonly sql = {} as DurableObjectStorage["sql"]; readonly kv = {} as DurableObjectStorage["kv"]; alarmAt: number | null = null; async get(key: string): Promise; async get(keys: string[]): Promise>; async get(keyOrKeys: string | string[]): Promise> { if (Array.isArray(keyOrKeys)) { return new Map(keyOrKeys.map((key) => [key, this.values.get(key) as T])); } return this.values.get(keyOrKeys) as T | undefined; } async put(key: string, value: T): Promise; async put(entries: Record | Map): Promise; async put( keyOrEntries: string | Record | Map, value?: T, ): Promise { if (typeof keyOrEntries === "string") { this.values.set(keyOrEntries, value); return; } const entries = keyOrEntries instanceof Map ? keyOrEntries.entries() : Object.entries(keyOrEntries); for (const [key, entry] of entries) this.values.set(key, entry); } async delete(key: string): Promise; async delete(keys: string[]): Promise; async delete(keyOrKeys: string | string[]): Promise { if (!Array.isArray(keyOrKeys)) return this.values.delete(keyOrKeys); let deleted = 0; for (const key of keyOrKeys) { if (this.values.delete(key)) deleted += 1; } return deleted; } async list(): Promise> { return new Map([...this.values.entries()].map(([key, value]) => [key, value as T])); } async deleteAll(): Promise { this.values.clear(); this.alarmAt = null; } transaction(_closure: (txn: DurableObjectTransaction) => Promise): Promise { return Promise.resolve(undefined as T); } transactionSync(_closure: () => T): T { return undefined as T; } async sync(): Promise {} async getAlarm(): Promise { return this.alarmAt; } async getCurrentBookmark(): Promise { return "test-bookmark"; } async getBookmarkForTime(_timestamp: number | Date): Promise { return "test-bookmark"; } onNextSessionRestoreBookmark(_bookmark: string): Promise { return Promise.resolve("test-bookmark"); } async setAlarm(scheduledTime: number | Date): Promise { this.alarmAt = scheduledTime instanceof Date ? scheduledTime.getTime() : scheduledTime; } async deleteAlarm(): Promise { this.alarmAt = null; } } class FakeDurableObjectState implements DurableObjectState { readonly id: DurableObjectId; readonly props: unknown = undefined; readonly facets = {} as DurableObjectState["facets"]; readonly ctx = this; private readonly waitUntilPromises: Promise[] = []; constructor( name: string, readonly storage: FakeStorage, ) { const id: Pick = { equals: (other: DurableObjectId) => other.toString() === name, name, toString: () => name, }; this.id = id as DurableObjectId; } waitUntil(promise: Promise): void { this.waitUntilPromises.push(promise); } blockConcurrencyWhile(callback: () => Promise): Promise { return callback(); } acceptWebSocket(_ws: WebSocket, _tags?: string[]): void {} getWebSockets(_tag?: string): WebSocket[] { return []; } getTags(_ws: WebSocket): string[] { return []; } setWebSocketAutoResponse(_pair?: WebSocketRequestResponsePair): void {} getWebSocketAutoResponse(): WebSocketRequestResponsePair | null { return null; } getWebSocketAutoResponseTimestamp(_ws: WebSocket): Date | null { return null; } setHibernatableWebSocketEventTimeout(_timeoutMs?: number): void {} getHibernatableWebSocketEventTimeout(): number | null { return null; } abort(_reason?: string): void {} } const makeDirectory = () => { const storage = new FakeStorage(); const directory = new McpExecutionOwnerDirectoryDO( new FakeDurableObjectState("execution-owner", storage), {}, ); return { directory, storage }; }; class FakeDirectoryNamespace implements McpExecutionOwnerDirectoryNamespace { private readonly directories = new Map(); idFromName(name: string): string { return name; } get(id: string): unknown { let directory = this.directories.get(id); if (!directory) { directory = makeDirectory().directory; this.directories.set(id, directory); } if (!(directory instanceof DurableObject)) return {}; return { put: (entry: McpExecutionOwnerRecord) => directory.put(entry), get: (executionId: string) => directory.get(executionId), delete: (executionId: string) => directory.delete(executionId), }; } } const record = (input?: Partial): McpExecutionOwnerRecord => ({ executionId: "exec_1", owner: { sessionId: "session_a" }, accountId: "acct_1", organizationId: "org_1", expiresAt: new Date(Date.now() + 60_000).toISOString(), ttlMs: 60_000, ...input, }); describe("McpExecutionOwnerDirectoryDO", () => { it("exports a DurableObject subclass for RPC method dispatch", () => { expect(Object.getPrototypeOf(McpExecutionOwnerDirectoryDO)).toBe(DurableObject); expect(makeDirectory().directory).toBeInstanceOf(DurableObject); }); it("stores and reads an unexpired execution owner record", async () => { const { directory, storage } = makeDirectory(); const entry = record(); await directory.put(entry); expect(await directory.get("exec_1")).toEqual(entry); expect(storage.alarmAt).toBe(Date.parse(entry.expiresAt)); }); it("treats expired records as absent and deletes them", async () => { const { directory, storage } = makeDirectory(); const entry = record({ expiresAt: new Date(Date.now() - 1_000).toISOString(), }); await directory.put(entry); expect(await directory.get("exec_1")).toBeNull(); expect(await storage.get("owner")).toBeUndefined(); expect(storage.alarmAt).toBeNull(); }); it("alarm deletes the owner record", async () => { const { directory, storage } = makeDirectory(); await directory.put(record()); await directory.alarm(); expect(await storage.get("owner")).toBeUndefined(); expect(storage.alarmAt).toBeNull(); }); it("stores, reads, and deletes records through the namespace RPC path", async () => { const directory = mcpExecutionOwnerDirectoryFromNamespace(new FakeDirectoryNamespace()); const entry = record({ executionId: "exec_rpc" }); expect(directory).not.toBeNull(); await Effect.runPromise(directory!.put(entry)); expect(await Effect.runPromise(directory!.get("exec_rpc"))).toEqual(entry); await Effect.runPromise(directory!.delete("exec_rpc")); expect(await Effect.runPromise(directory!.get("exec_rpc"))).toBeNull(); }); });