fix(sessions): keep attachment references a rewritten message still points at

.changeset/sessions-attachment-refs.md

@@ -0,0 +1,7 @@
+---
+"agents": patch
+---
+
+Sessions: keep the attachment references of a message written back with pointers in it.
+
+A message whose media was already offloaded carries `attachment:sha256:…` pointers. Writing that stored form back (an `updateMessage`, or a copy through `appendMessage` or `importMessage`) derived references only from media extracted on that write, so it recorded none, and the payload could be collected while the row still pointed at it. References now cover every pointer the stored message contains. Rows already orphaned are not repaired.

packages/agents/src/sessions/attachment-ingest.ts

@@ -143,6 +143,13 @@ export interface ExtractionResult {
   message: SessionMessage;
   /** Payloads to write, deduplicated by address, in encounter order. */
   attachments: PendingAttachment[];
+  /**
+   * Every address a read of the stored message restores: the payloads
+   * extracted on this pass and any pointer the message already carried. A
+   * message can be written back with its pointers in place, and its
+   * references must follow what the row says, not what this pass extracted.
+   */
+  references: string[];
 }
 
 /**
@@ -203,10 +210,41 @@ export function extractAttachments(message: SessionMessage): ExtractionResult {
     return changed ? next : value;
   };
 
-  if (message.parts.length === 0) return { message, attachments };
+  if (message.parts.length === 0) {
+    return { message, attachments, references: [] };
+  }
   const parts = walk(message.parts, 0) as SessionMessagePart[];
-  if (parts === message.parts) return { message, attachments };
-  return { message: { ...message, parts }, attachments };
+  const references = [...pointersOf(parts)];
+  if (parts === message.parts) return { message, attachments, references };
+  return { message: { ...message, parts }, attachments, references };
+}
+
+/**
+ * The attachment addresses a read of this stored form restores. It walks
+ * exactly as {@link resolveAttachments} does, so a message holds a reference
+ * for every pointer a read would put back and for nothing else.
+ */
+function pointersOf(
+  value: unknown,
+  into = new Set<string>(),
+  depth = 0
+): Set<string> {
+  if (depth > MAX_WALK_DEPTH || value === null || typeof value !== "object") {
+    return into;
+  }
+  if (Array.isArray(value)) {
+    for (const entry of value) pointersOf(entry, into, depth + 1);
+    return into;
+  }
+  const record = value as Record<string, unknown>;
+  const hash =
+    parseAttachmentUrl(record.url) ?? parseAttachmentUrl(record.data);
+  if (hash) {
+    into.add(hash);
+    return into;
+  }
+  for (const entry of Object.values(record)) pointersOf(entry, into, depth + 1);
+  return into;
 }
 
 /**

packages/agents/src/sessions/core.ts

@@ -940,7 +940,11 @@ export class SessionsCore {
     // Inline media leaves the message before it is serialized, so the row
     // holds a pointer and never the payload. Addresses are computed here, out
     // of the transaction; the transaction only writes.
-    const { message: staged, attachments } = extractAttachments(message);
+    const {
+      message: staged,
+      attachments,
+      references
+    } = extractAttachments(message);
     const slices = splitContent(JSON.stringify(staged));
     const seq = tail.nextSeq;
     this.io.transaction(() => {
@@ -964,11 +968,7 @@ export class SessionsCore {
         ]
       );
       this.#writeContinuations(sessionId, message.id, slices);
-      this.#attachments.addRefs(
-        sessionId,
-        message.id,
-        attachments.map((attachment) => attachment.hash)
-      );
+      this.#attachments.addRefs(sessionId, message.id, references);
       this.#indexFts(sessionId, staged, false);
     });
 
@@ -1024,7 +1024,11 @@ export class SessionsCore {
           (this.#continuations(sessionId, [message.id]).get(message.id) ?? "");
     // Compare in stored form: a re-sent identical image extracts to the same
     // address, so an unchanged update still writes nothing.
-    const { message: staged, attachments } = extractAttachments(message);
+    const {
+      message: staged,
+      attachments,
+      references
+    } = extractAttachments(message);
     const json = JSON.stringify(staged);
     if (oldContent === json) return "unchanged";
 
@@ -1058,11 +1062,7 @@ export class SessionsCore {
       this.#writeContinuations(sessionId, message.id, slices);
       // Payloads are stored before references move, so a hash this message
       // still uses is never momentarily unreferenced and collected.
-      this.#attachments.replaceRefs(
-        sessionId,
-        message.id,
-        attachments.map((attachment) => attachment.hash)
-      );
+      this.#attachments.replaceRefs(sessionId, message.id, references);
       this.#indexFts(sessionId, staged, true);
     });
     const memo = this.#pathTokens.get(sessionId);
@@ -1293,7 +1293,11 @@ export class SessionsCore {
     message: SessionMessage,
     options: { parentId: string | null; createdAt: number }
   ): boolean {
-    const { message: staged, attachments } = extractAttachments(message);
+    const {
+      message: staged,
+      attachments,
+      references
+    } = extractAttachments(message);
     const slices = splitContent(JSON.stringify(staged));
     const tail = this.#tail(sessionId);
     let inserted = 0;
@@ -1319,11 +1323,7 @@ export class SessionsCore {
       );
       if (inserted === 0) return;
       this.#writeContinuations(sessionId, message.id, slices);
-      this.#attachments.addRefs(
-        sessionId,
-        message.id,
-        attachments.map((attachment) => attachment.hash)
-      );
+      this.#attachments.addRefs(sessionId, message.id, references);
       this.#indexFts(sessionId, staged, false);
     });
     if (inserted === 0) return false;

packages/agents/src/tests/capabilities/sessions.ts

@@ -138,6 +138,20 @@ export class SessionHarnessObject extends DurableObject<Cloudflare.Env> {
     return rows.length > 0 ? Number(rows[0].content_chunks) : null;
   }
 
+  /** One message row exactly as stored: pointers in place of payloads. */
+  storedMessage(sessionId: string, messageId: string): SessionMessage {
+    return JSON.parse(
+      this.ctx.storage.sql
+        .exec<{ content: string }>(
+          `SELECT content FROM cf_agents_session_messages
+           WHERE session_id = ? AND id = ?`,
+          sessionId,
+          messageId
+        )
+        .one().content
+    ) as SessionMessage;
+  }
+
   /** Stored size of one message row, excluding anything it points at. */
   messageRowBytes(sessionId: string, messageId: string): number {
     return Number(

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

@@ -220,6 +220,77 @@ describe("Sessions attachments", () => {
     });
   });
 
+  it("keeps a pointer's reference when the message is written back", async () => {
+    const stub = env.SessionHarnessObject.getByName(crypto.randomUUID());
+    await runInDurableObject(stub, async (instance: SessionHarnessObject) => {
+      const session = instance.sessions.session();
+      const url = dataUrl("image/png", 80_000);
+      await session.appendMessage(fileMessage("m1", url, "image/png"));
+
+      // The stored form: the same message with a pointer where the bytes
+      // were. A host can hold this form — a read whose payload did not
+      // resolve, or code that works on the row — and write it back.
+      const pointer = instance.storedMessage("", "m1");
+      expect((pointer.parts[1] as { url: string }).url).toMatch(
+        /^attachment:sha256:[0-9a-f]{64}$/
+      );
+
+      // An update that changes anything else must keep the reference. It
+      // used to be replaced with what the write extracted — nothing — and
+      // the payload was collected under a live pointer.
+      await session.updateMessage({
+        ...pointer,
+        parts: [
+          { type: "text", text: "see attached (edited)" },
+          pointer.parts[1]
+        ]
+      });
+      expect(instance.attachmentRecords()).toHaveLength(1);
+      expect(instance.attachmentRefCount()).toBe(1);
+      const [full] = await session.getHistory();
+      expect((full.parts[1] as { url: string }).url).toBe(url);
+
+      // A copy under a new id takes its own reference, so deleting the
+      // original does not take the bytes with it.
+      await session.appendMessage({ ...pointer, id: "m2" });
+      expect(instance.attachmentRefCount()).toBe(2);
+      await session.deleteMessages(["m1"]);
+      expect(instance.attachmentRecords()).toHaveLength(1);
+      const copy = await session.getMessage("m2");
+      expect(copy).not.toBeNull();
+      expect((copy!.parts[1] as { url: string }).url).toBe(url);
+    });
+  });
+
+  it("references a pointer nested inside another part", async () => {
+    const stub = env.SessionHarnessObject.getByName(crypto.randomUUID());
+    await runInDurableObject(stub, async (instance: SessionHarnessObject) => {
+      const session = instance.sessions.session();
+      const url = dataUrl("image/png", 40_000);
+      await session.appendMessage(fileMessage("m1", url, "image/png"));
+      const pointerUrl = (
+        instance.storedMessage("", "m1").parts[1] as { url: string }
+      ).url;
+
+      // A copy of the pointer a few levels down, where a read still restores
+      // it, holds its own reference.
+      let node: unknown = { url: pointerUrl };
+      for (let i = 0; i < 3; i++) node = { type: "nest", child: node };
+      await session.appendMessage({
+        id: "m2",
+        role: "user",
+        parts: [node as SessionMessage["parts"][number]]
+      });
+      expect(instance.attachmentRefCount()).toBe(2);
+
+      // So dropping the message that carried the bytes keeps them for it.
+      await session.deleteMessages(["m1"]);
+      expect(instance.attachmentRefCount()).toBe(1);
+      expect(instance.attachmentRecords()).toHaveLength(1);
+      const [copy] = await session.getHistory();
+      expect(JSON.stringify(copy)).toContain(url);
+    });
+  });
   it("returns null for an update whose target is gone and stores nothing", async () => {
     const stub = env.SessionHarnessObject.getByName(crypto.randomUUID());
     await runInDurableObject(stub, async (instance: SessionHarnessObject) => {