Fix execution owner directory Durable Object

Author: RhysSullivan

packages/hosts/cloudflare/src/mcp/agent-session-model-resume.test.ts

@@ -223,9 +223,7 @@ class InMemoryExecutionOwnerDirectoryNamespace implements McpExecutionOwnerDirec
   get(id: string): McpExecutionOwnerDirectoryDO {
     let directory = this.directories.get(id);
     if (!directory) {
-      directory = new McpExecutionOwnerDirectoryDO({
-        storage: new FakeStorage(),
-      });
+      directory = new McpExecutionOwnerDirectoryDO(new FakeDurableObjectState(id), {});
       this.directories.set(id, directory);
     }
     return directory;

packages/hosts/cloudflare/src/mcp/execution-owner-directory.test.ts

@@ -1,24 +1,88 @@
+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 {
+class FakeStorage implements DurableObjectStorage {
   private readonly values = new Map<string, unknown>();
+  readonly sql = {} as DurableObjectStorage["sql"];
+  readonly kv = {} as DurableObjectStorage["kv"];
   alarmAt: number | null = null;
 
-  async get<T>(key: string): Promise<T | undefined> {
-    return this.values.get(key) as T | undefined;
+  async get<T>(key: string): Promise<T | undefined>;
+  async get<T>(keys: string[]): Promise<Map<string, T>>;
+  async get<T>(keyOrKeys: string | string[]): Promise<T | undefined | Map<string, T>> {
+    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<T>(key: string, value: T): Promise<void> {
-    this.values.set(key, value);
+  async put<T>(key: string, value: T): Promise<void>;
+  async put<T>(entries: Record<string, T> | Map<string, T>): Promise<void>;
+  async put<T>(
+    keyOrEntries: string | Record<string, T> | Map<string, T>,
+    value?: T,
+  ): Promise<void> {
+    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<boolean> {
-    return this.values.delete(key);
+  async delete(key: string): Promise<boolean>;
+  async delete(keys: string[]): Promise<number>;
+  async delete(keyOrKeys: string | string[]): Promise<boolean | number> {
+    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<T = unknown>(): Promise<Map<string, T>> {
+    return new Map([...this.values.entries()].map(([key, value]) => [key, value as T]));
+  }
+
+  async deleteAll(): Promise<void> {
+    this.values.clear();
+    this.alarmAt = null;
+  }
+
+  transaction<T>(_closure: (txn: DurableObjectTransaction) => Promise<T>): Promise<T> {
+    return Promise.resolve(undefined as T);
+  }
+
+  transactionSync<T>(_closure: () => T): T {
+    return undefined as T;
+  }
+
+  async sync(): Promise<void> {}
+
+  async getAlarm(): Promise<number | null> {
+    return this.alarmAt;
+  }
+
+  async getCurrentBookmark(): Promise<string> {
+    return "test-bookmark";
+  }
+
+  async getBookmarkForTime(_timestamp: number | Date): Promise<string> {
+    return "test-bookmark";
+  }
+
+  onNextSessionRestoreBookmark(_bookmark: string): Promise<string> {
+    return Promise.resolve("test-bookmark");
   }
 
   async setAlarm(scheduledTime: number | Date): Promise<void> {
@@ -30,14 +94,93 @@ class FakeStorage {
   }
 }
 
+class FakeDurableObjectState implements DurableObjectState {
+  readonly id: DurableObjectId;
+  readonly props: unknown = undefined;
+  readonly facets = {} as DurableObjectState["facets"];
+  readonly ctx = this;
+  private readonly waitUntilPromises: Promise<unknown>[] = [];
+
+  constructor(
+    name: string,
+    readonly storage: FakeStorage,
+  ) {
+    const id: Pick<DurableObjectId, "equals" | "name" | "toString"> = {
+      equals: (other: DurableObjectId) => other.toString() === name,
+      name,
+      toString: () => name,
+    };
+    this.id = id as DurableObjectId;
+  }
+
+  waitUntil(promise: Promise<unknown>): void {
+    this.waitUntilPromises.push(promise);
+  }
+
+  blockConcurrencyWhile<T>(callback: () => Promise<T>): Promise<T> {
+    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({
-    storage,
-  });
+  const directory = new McpExecutionOwnerDirectoryDO(
+    new FakeDurableObjectState("execution-owner", storage),
+    {},
+  );
   return { directory, storage };
 };
 
+class FakeDirectoryNamespace implements McpExecutionOwnerDirectoryNamespace<string> {
+  private readonly directories = new Map<string, McpExecutionOwnerDirectoryDO>();
+
+  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>): McpExecutionOwnerRecord => ({
   executionId: "exec_1",
   owner: { sessionId: "session_a" },
@@ -49,6 +192,11 @@ const record = (input?: Partial<McpExecutionOwnerRecord>): McpExecutionOwnerReco
 });
 
 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();
@@ -81,4 +229,18 @@ describe("McpExecutionOwnerDirectoryDO", () => {
     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();
+  });
 });

packages/hosts/cloudflare/src/mcp/execution-owner-directory.ts

@@ -1,3 +1,4 @@
+import { DurableObject } from "cloudflare:workers";
 import { Data, Effect } from "effect";
 
 export type McpExecutionOwnerRoute = {
@@ -30,17 +31,7 @@ export interface McpExecutionOwnerDirectoryNamespace<Id> {
   readonly get: (id: Id) => unknown;
 }
 
-interface McpExecutionOwnerDirectoryStorage {
-  readonly get: <T>(key: string) => Promise<T | undefined>;
-  readonly put: <T>(key: string, value: T) => Promise<void>;
-  readonly delete: (key: string) => Promise<boolean>;
-  readonly setAlarm: (scheduledTime: number | Date) => Promise<void>;
-  readonly deleteAlarm: () => Promise<void>;
-}
-
-interface McpExecutionOwnerDirectoryState {
-  readonly storage: McpExecutionOwnerDirectoryStorage;
-}
+type McpExecutionOwnerDirectoryStorage = DurableObjectState["storage"];
 
 const toMcpExecutionOwnerDirectoryStub = (stub: unknown): McpExecutionOwnerDirectoryStub =>
   stub as McpExecutionOwnerDirectoryStub;
@@ -73,10 +64,11 @@ const isOwnerRecord = (value: unknown): value is McpExecutionOwnerRecord =>
 
 const expiryMs = (record: McpExecutionOwnerRecord): number => Date.parse(record.expiresAt);
 
-export class McpExecutionOwnerDirectoryDO {
+export class McpExecutionOwnerDirectoryDO extends DurableObject<unknown> {
   private readonly storage: McpExecutionOwnerDirectoryStorage;
 
-  constructor(ctx: McpExecutionOwnerDirectoryState) {
+  constructor(ctx: DurableObjectState, env: unknown) {
+    super(ctx, env);
     this.storage = ctx.storage;
   }
 
@@ -104,7 +96,7 @@ export class McpExecutionOwnerDirectoryDO {
     await Promise.all([this.storage.delete(RECORD_KEY), this.storage.deleteAlarm()]);
   }
 
-  async alarm(): Promise<void> {
+  override async alarm(): Promise<void> {
     await Promise.all([this.storage.delete(RECORD_KEY), this.storage.deleteAlarm()]);
   }
 }