fix(sessions): share the change-feed mirror between hosts (#2205)

.changeset/session-mirror.md

@@ -0,0 +1,7 @@
+---
+"agents": minor
+"@cloudflare/ai-chat": patch
+"@cloudflare/think": patch
+---
+
+Add `session.mirror()`, which keeps an in-memory transcript in step with one session's change feed. `AIChatAgent` and `Think` now share it instead of each carrying its own copy of the reduction; `Think` keeps its branch, compaction, and prompt-refresh handling through the `intercept` and `onApplied` hooks.

docs/agents/sessions.md

@@ -361,6 +361,25 @@ const unsubscribe = sessions.subscribe(async (event) => {
 
 `import` fires once per `importMessage()` and carries the row; it is deliberately not an `append`, so a projection does not patch itself per imported row during a migration — it marks itself stale and re-derives once. `compaction` fires when an overlay is stored directly through `addCompaction()`; `compact()` reports its own overlay as `compact`.
 
+A host that keeps an in-memory transcript does not need to write that reduction itself. `session.mirror()` subscribes for one session and patches an array in place: an append or update replaces the entry with the same id (or pushes a new row), a delete filters, and a clear empties it:
+
+```ts
+const stop = sessions.session("").mirror<UIMessage>({
+  get: () => this.messages,
+  set: (messages) => {
+    this.messages = messages;
+  },
+  transform: (message) => toCachedMessage(message),
+  intercept: async (event) => {
+    if (event.type !== "compact") return false;
+    this.messages = await reloadFromStorage();
+    return true;
+  }
+});
+```
+
+`get()` is read on every event, so reassigning the array is safe. `intercept` handles an event instead of the default reduction. Use it for changes an in-place patch cannot express, such as a branch append or a compaction. `onApplied` runs after an append or update was written, with the entry it replaced.
+
 This is a local cache-coherence feed, not a cross-object event log. Capability diagnostics are also emitted through Lifecycle observability.
 
 ## Errors

packages/agents/src/sessions/handle.ts

@@ -11,6 +11,7 @@
 import type { CompactResult } from "./compaction-helpers";
 import { COMPACTION_PREFIX } from "./compaction-helpers";
 import type { SessionsCore } from "./core";
+import { mirrorSessionChanges, type SessionMirrorOptions } from "./mirror";
 import { byteLength, sanitizeMessage } from "./sanitize";
 import type {
   AppendOptions,
@@ -67,6 +68,22 @@ export class Session {
     return this;
   }
 
+  // ── Change feed ──────────────────────────────────────────────────────────
+
+  /**
+   * Keep an in-memory transcript in step with this session's change feed,
+   * so write sites never hand-patch it. Returns the unsubscribe function.
+   */
+  mirror<M extends { id: string }>(
+    options: SessionMirrorOptions<M>
+  ): () => void {
+    return mirrorSessionChanges(
+      (listener) => this.#core.subscribe(listener),
+      this.sessionId,
+      options
+    );
+  }
+
   // ── Reads ────────────────────────────────────────────────────────────────
 
   /**

packages/agents/src/sessions/index.ts

@@ -7,6 +7,7 @@
 
 export { Sessions } from "./sessions";
 export { Session, type CompactionFunction } from "./handle";
+export type { SessionMirrorOptions } from "./mirror";
 export {
   COMPACTION_PREFIX,
   createCompactFunction,

packages/agents/src/sessions/mirror.ts

@@ -0,0 +1,85 @@
+import type { SessionChangeEvent, SessionMessage } from "./types";
+
+/**
+ * How a host keeps an in-memory transcript in step with one session's change
+ * feed. The mirror applies the reduction every such cache needs; the hooks
+ * carry the decisions only the host can make.
+ */
+export interface SessionMirrorOptions<M extends { id: string }> {
+  /**
+   * The cached transcript. Read on every event, so a host that reassigns its
+   * array (a public `messages` field, say) is followed rather than
+   * shadowed.
+   */
+  get(): M[];
+  /** Install a new array: after a delete, or emptied by a clear. */
+  set(messages: M[]): void;
+  /** The cached form of a stored message. Defaults to the stored message. */
+  transform?(message: SessionMessage): M;
+  /**
+   * Handle an event instead of the default reduction, returning `true` when
+   * it did. A host whose cache cannot express a change in place — a branch
+   * append, a compaction — re-reads storage here.
+   */
+  intercept?(event: SessionChangeEvent): boolean | Promise<boolean>;
+  /**
+   * After the default reduction wrote `message` for an append or update.
+   * `previous` is the cached entry it replaced, if there was one.
+   */
+  onApplied?(
+    event: Extract<SessionChangeEvent, { type: "append" | "update" }>,
+    message: M,
+    previous: M | undefined
+  ): void;
+}
+
+/**
+ * The default reduction, per event:
+ *
+ * - `append` of a new row: replace the cached entry with its id, or push it.
+ *   A duplicate append (`inserted: false`) changed nothing.
+ * - `update`: replace the cached entry with its id; a row the cache does not
+ *   hold stays out.
+ * - `delete`: drop the removed ids. `clear`: empty the cache.
+ * - `import`, `compact`, `compaction`: nothing, since no in-place patch
+ *   expresses them. A host that must react intercepts them.
+ */
+export function mirrorSessionChanges<M extends { id: string }>(
+  subscribe: (
+    listener: (event: SessionChangeEvent) => Promise<void>
+  ) => () => void,
+  sessionId: string,
+  options: SessionMirrorOptions<M>
+): () => void {
+  const transform =
+    options.transform ?? ((message: SessionMessage) => message as unknown as M);
+  return subscribe(async (event) => {
+    if (event.sessionId !== sessionId) return;
+    if (await options.intercept?.(event)) return;
+    switch (event.type) {
+      case "append":
+      case "update": {
+        if (event.type === "append" && !event.inserted) return;
+        const cache = options.get();
+        const index = cache.findIndex((m) => m.id === event.message.id);
+        if (index === -1 && event.type === "update") return;
+        const message = transform(event.message);
+        const previous = index === -1 ? undefined : cache[index];
+        if (index === -1) cache.push(message);
+        else cache[index] = message;
+        options.onApplied?.(event, message, previous);
+        return;
+      }
+      case "delete": {
+        const removed = new Set(event.messageIds);
+        options.set(options.get().filter((m) => !removed.has(m.id)));
+        return;
+      }
+      case "clear":
+        options.set([]);
+        return;
+      default:
+        return;
+    }
+  });
+}

packages/agents/src/tests/sessions/capability.test.ts

@@ -505,6 +505,102 @@ describe("Sessions capability", () => {
     });
   });
 
+  describe("mirror", () => {
+    type Cached = SessionMessage & { cached: true };
+
+    it("reduces appends, updates, deletes, and clears onto a host array", async () => {
+      const stub = env.SessionHarnessObject.getByName(crypto.randomUUID());
+      await runInDurableObject(stub, async (instance: SessionHarnessObject) => {
+        const session = instance.sessions.session();
+        const host = { messages: [] as Cached[] };
+        const applied: Array<[string, string, boolean]> = [];
+        session.mirror<Cached>({
+          get: () => host.messages,
+          set: (messages) => {
+            host.messages = messages;
+          },
+          transform: (message) => ({ ...message, cached: true }),
+          onApplied: (event, message, previous) => {
+            applied.push([event.type, message.id, previous !== undefined]);
+          }
+        });
+
+        await session.appendMessage(text("a", "one"));
+        await session.appendMessage(text("b", "two"));
+        await session.appendMessage(text("a", "duplicate"));
+        expect(host.messages.map((m) => [m.id, m.cached])).toEqual([
+          ["a", true],
+          ["b", true]
+        ]);
+
+        await session.updateMessage(text("b", "two, edited"));
+        expect(host.messages[1].parts).toEqual([
+          { type: "text", text: "two, edited" }
+        ]);
+
+        // A host that reassigns its array is followed, not shadowed.
+        host.messages = host.messages.filter((m) => m.id !== "a");
+        await session.updateMessage(text("a", "not cached"));
+        expect(host.messages.map((m) => m.id)).toEqual(["b"]);
+
+        await session.appendMessage(text("c", "three"));
+        await session.deleteMessages(["b"]);
+        expect(host.messages.map((m) => m.id)).toEqual(["c"]);
+
+        await session.clearMessages();
+        expect(host.messages).toEqual([]);
+
+        expect(applied).toEqual([
+          ["append", "a", false],
+          ["append", "b", false],
+          ["update", "b", true],
+          ["append", "c", false]
+        ]);
+      });
+    });
+
+    it("lets the host intercept events and ignores other sessions", async () => {
+      const stub = env.SessionHarnessObject.getByName(crypto.randomUUID());
+      await runInDurableObject(stub, async (instance: SessionHarnessObject) => {
+        const session = instance.sessions.session("mine");
+        let cache: SessionMessage[] = [];
+        const intercepted: string[] = [];
+        const unsubscribe = session.mirror({
+          get: () => cache,
+          set: (messages) => {
+            cache = messages;
+          },
+          intercept: (event) => {
+            if (event.type === "append" && event.parentId !== undefined) {
+              intercepted.push(`branch:${event.message.id}`);
+              return true;
+            }
+            if (event.type === "import") {
+              intercepted.push(`import:${event.message.id}`);
+            }
+            return false;
+          }
+        });
+
+        await session.appendMessage(text("root", "hi"));
+        await session.appendMessage(text("alt", "branch"), {
+          parentId: "root"
+        });
+        await session.importMessage(text("moved", "in"), {
+          parentId: "root",
+          createdAt: 1
+        });
+        await instance.sessions.session("theirs").appendMessage(text("x", "!"));
+        expect(cache.map((m) => m.id)).toEqual(["root"]);
+        expect(intercepted).toEqual(["branch:alt", "import:moved"]);
+
+        unsubscribe();
+        await session.appendMessage(text("after", "gone"));
+        expect(cache.map((m) => m.id)).toEqual(["root"]);
+      });
+    });
+  });
+
   it("budgets recent history with a floor and honest truncation", async () => {
     const stub = env.SessionHarnessObject.getByName(crypto.randomUUID());
     await runInDurableObject(stub, async (instance: SessionHarnessObject) => {

packages/ai-chat/src/index.ts

@@ -2564,37 +2564,12 @@ export class AIChatAgent<
    * cache.
    */
   #subscribeToSessionChanges(): void {
-    this.sessions.subscribe(async (event) => {
-      if (event.sessionId !== this.#session.sessionId) return;
-      switch (event.type) {
-        case "append": {
-          if (!event.inserted) return;
-          const message = this.#messageForCache(event.message);
-          const index = this.messages.findIndex((m) => m.id === message.id);
-          if (index === -1) this.messages.push(message);
-          else this.messages[index] = message;
-          return;
-        }
-        case "update": {
-          const index = this.messages.findIndex(
-            (m) => m.id === event.message.id
-          );
-          if (index === -1) return;
-          this.messages[index] = this.#messageForCache(event.message);
-          return;
-        }
-        case "delete": {
-          const removed = new Set(event.messageIds);
-          this.messages = this.messages.filter((m) => !removed.has(m.id));
-          return;
-        }
-        case "clear": {
-          this.messages = [];
-          return;
-        }
-        default:
-          return;
-      }
+    this.#session.mirror<UIMessage>({
+      get: () => this.messages,
+      set: (messages) => {
+        this.messages = messages;
+      },
+      transform: (message) => this.#messageForCache(message)
     });
   }