refactor(agents): keep ThinkHarness records in the shared harness store

.changeset/think-harness-shared-store.md

@@ -0,0 +1,5 @@
+---
+"agents": patch
+---
+
+`ThinkHarness` now keeps its sessions and operations in the shared `agents/harness/store`, and moves records from its old tables on first start so work in flight still recovers. See [ThinkHarness: Recovery](https://github.com/cloudflare/agents/blob/main/docs/agents/harnesses/think.md#recovery).

packages/agents/src/harness/container/tests/store.test.ts

@@ -146,6 +146,54 @@ describe("HarnessStore", () => {
       expect(store.read("from", "engine")).toHaveLength(2);
     }));
 
+  it("keeps a session's creation time when one is given, for imported records", () =>
+    withStore((store) => {
+      store.createSession({ id: "new" });
+      store.createSession({ id: "imported", createdAt: 1000 });
+      expect(store.session("imported")).toMatchObject({
+        createdAt: 1000,
+        updatedAt: 1000
+      });
+      expect(store.sessions().map((s) => s.id)).toEqual(["imported", "new"]);
+    }));
+
+  it("forgets a session's settled operations and keeps its open ones", () =>
+    withStore((store) => {
+      for (const id of ["done", "gone", "running", "queued"]) {
+        store.enqueue({ session: "s", id, input: null });
+      }
+      store.enqueue({ session: "other", id: "done", input: null });
+      store.settle("s", "done", { status: "done", result: null });
+      store.settle("s", "gone", { status: "unanswered", reason: "withdrawn" });
+      store.settle("other", "done", { status: "done", result: null });
+      store.start("s", "running");
+      expect(store.deleteSettled("s")).toBe(2);
+      expect(store.operations({ session: "s" }).map((o) => o.id)).toEqual([
+        "running",
+        "queued"
+      ]);
+      expect(store.operation("other", "done")?.status).toBe("done");
+      // The id is free again.
+      expect(
+        store.enqueue({ session: "s", id: "done", input: 1 }).accepted
+      ).toBe(true);
+    }));
+
+  it("deletes one operation, so its id may be enqueued again", () =>
+    withStore((store) => {
+      store.enqueue({ session: "s", id: "op", input: "old" });
+      store.enqueue({ session: "other", id: "op", input: "kept" });
+      expect(store.deleteOperation("s", "op")).toBe(true);
+      expect(store.deleteOperation("s", "op")).toBe(false);
+      expect(store.operation("other", "op")?.input).toBe("kept");
+      expect(
+        store.enqueue({ session: "s", id: "op", input: "new" })
+      ).toMatchObject({
+        accepted: true,
+        record: { input: "new", status: "queued" }
+      });
+    }));
+
   it("deletes a session with its operations and logs", () =>
     withStore((store) => {
       store.createSession({ id: "s" });

packages/agents/src/harness/store/store.ts

@@ -267,14 +267,17 @@ export class HarnessStore {
    * Create a session, or return the existing one with this id unchanged.
    *
    * @param input - The id, an optional parent, and the initial state.
+   *   `createdAt` keeps the creation time of a session imported from
+   *   elsewhere, so the session list stays in creation order; default now.
    * @returns The stored session.
    */
   createSession(input: {
     readonly id: string;
     readonly parent?: string;
     readonly state?: JsonValue;
+    readonly createdAt?: number;
   }): SessionRecord {
-    const now = Date.now();
+    const now = input.createdAt ?? Date.now();
     this.#storage.sql.exec(
       `INSERT INTO ${this.#sessions} (id, parent, state, created_at, updated_at)
        VALUES (?, ?, ?, ?, ?) ON CONFLICT (id) DO NOTHING`,
@@ -503,6 +506,41 @@ export class HarnessStore {
     );
   }
 
+  /**
+   * Delete one operation, whatever its status. Its id may be enqueued
+   * again. For replacing a record imported from elsewhere; a harness that
+   * deletes an open operation must stop running it itself.
+   *
+   * @param session - The session id.
+   * @param id - The operation id.
+   * @returns False when there was no such operation.
+   */
+  deleteOperation(session: string, id: string): boolean {
+    return (
+      this.#storage.sql.exec(
+        `DELETE FROM ${this.#operations} WHERE session = ? AND id = ?`,
+        session,
+        id
+      ).rowsWritten > 0
+    );
+  }
+
+  /**
+   * Forget a session's settled operations, as when the session starts
+   * over. Open operations are kept. A forgotten operation's id may be
+   * enqueued again.
+   *
+   * @param session - The session id.
+   * @returns How many operations were deleted.
+   */
+  deleteSettled(session: string): number {
+    return this.#storage.sql.exec(
+      `DELETE FROM ${this.#operations}
+       WHERE session = ? AND status IN ('done', 'unanswered')`,
+      session
+    ).rowsWritten;
+  }
+
   /**
    * Replace an operation's harness-owned bookkeeping.
    *

packages/agents/src/harness/think/harness.ts

@@ -577,7 +577,7 @@ export class ThinkHarness<TOOLS extends ToolSet = ToolSet>
       if (live?.operationId === op.operationId) {
         live.abort.abort();
       } else if (op.status === "queued") {
-        store.update(session, op.operationId, {
+        store.settle(session, op.operationId, {
           status: "unanswered",
           reason: "withdrawn"
         });
@@ -945,7 +945,7 @@ export class ThinkHarness<TOOLS extends ToolSet = ToolSet>
         case "await-input":
         case "end":
           await this.#commit(session, undefined, () => {
-            this.#tables().update(session, operationId, {
+            this.#tables().settle(session, operationId, {
               status: "done",
               text: textOf(message)
             });
@@ -974,8 +974,7 @@ export class ThinkHarness<TOOLS extends ToolSet = ToolSet>
       readonly parentId?: string;
       readonly messageId: string;
     }) =>
-      store.update(session, operationId, {
-        status: "running",
+      store.start(session, operationId, {
         parentId: change.parentId ?? null,
         messageId: change.messageId,
         pendingModel: true
@@ -1052,7 +1051,7 @@ export class ThinkHarness<TOOLS extends ToolSet = ToolSet>
           placed.push(answered);
           if (continues) running({ messageId: answered.id });
           else {
-            store.update(session, operationId, {
+            store.settle(session, operationId, {
               status: "done",
               text: textOf(answered)
             });
@@ -1094,7 +1093,7 @@ export class ThinkHarness<TOOLS extends ToolSet = ToolSet>
       this.#events.emit(session, { type: "message", message, operationId });
     }
     if (unanswered !== undefined) {
-      store.update(session, operationId, {
+      store.settle(session, operationId, {
         status: "unanswered",
         reason: unanswered
       });
@@ -1267,11 +1266,12 @@ export class ThinkHarness<TOOLS extends ToolSet = ToolSet>
     if (live.abort.signal.aborted) {
       await this.#commit(session, streamId, (put) => {
         if (keep) put(keep, { parentId: op.parentId });
-        this.#tables().update(session, operationId, {
-          status: "unanswered",
-          reason: "aborted",
-          streamId: null
-        });
+        this.#tables().settle(
+          session,
+          operationId,
+          { status: "unanswered", reason: "aborted" },
+          { streamId: null }
+        );
       });
       await this.#settled(session, operationId, keep);
       return {};
@@ -1311,11 +1311,12 @@ export class ThinkHarness<TOOLS extends ToolSet = ToolSet>
       }
       await this.#commit(session, streamId, (put) => {
         if (keep) put(keep, { parentId: op.parentId });
-        this.#tables().update(session, operationId, {
-          status: "unanswered",
-          reason: errorText(failure),
-          streamId: null
-        });
+        this.#tables().settle(
+          session,
+          operationId,
+          { status: "unanswered", reason: errorText(failure) },
+          { streamId: null }
+        );
       });
       await this.#settled(session, operationId, keep);
       return {};
@@ -1389,14 +1390,17 @@ export class ThinkHarness<TOOLS extends ToolSet = ToolSet>
     }
     await this.#commit(session, streamId, (put) => {
       if (partial) put(partial, { parentId: op.parentId });
-      this.#tables().update(session, operationId, {
-        streamId: null,
-        interruptions: attempt,
-        ...(decision === "abandon" && {
-          status: "unanswered",
-          reason: "interrupted"
-        })
-      });
+      const progress = { streamId: null, interruptions: attempt };
+      if (decision === "abandon") {
+        this.#tables().settle(
+          session,
+          operationId,
+          { status: "unanswered", reason: "interrupted" },
+          progress
+        );
+      } else {
+        this.#tables().update(session, operationId, progress);
+      }
     });
     if (partial) {
       this.#events.emit(session, {
@@ -1624,11 +1628,12 @@ export class ThinkHarness<TOOLS extends ToolSet = ToolSet>
         : undefined;
     await this.#commit(session, op.streamId, (put) => {
       if (keep) put(keep, { parentId: op.parentId });
-      this.#tables().update(session, operationId, {
-        status: "unanswered",
-        reason,
-        streamId: null
-      });
+      this.#tables().settle(
+        session,
+        operationId,
+        { status: "unanswered", reason },
+        { streamId: null }
+      );
     });
     if (keep) {
       this.#events.emit(session, {
@@ -1698,8 +1703,7 @@ export class ThinkHarness<TOOLS extends ToolSet = ToolSet>
   // ── Storage ──────────────────────────────────────────────────────────────
 
   #tables(): OperationRecords {
-    this.#records ??= new OperationRecords(this.lifecycle.storage.sql);
-    this.#records.ensureTables();
+    this.#records ??= new OperationRecords(this.lifecycle.storage);
     return this.#records;
   }