Protect preview persistence and lifecycle hydration

packages/app/src/attachments/service.test.ts

@@ -1,7 +1,11 @@
 import { afterEach, describe, expect, it } from "vitest";
 import type { AttachmentMetadata, AttachmentStore, SaveAttachmentInput } from "@/attachments/types";
 import { __setAttachmentStoreForTests } from "./store";
-import { encodeAttachmentsForSend, persistAttachmentFromBytes } from "./service";
+import {
+  encodeAttachmentsForSend,
+  garbageCollectAttachments,
+  persistAttachmentFromBytes,
+} from "./service";
 
 function createAttachment(input: Partial<AttachmentMetadata> = {}): AttachmentMetadata {
   return {
@@ -94,4 +98,46 @@ describe("attachment service", () => {
       { data: "att_send:base64", mimeType: "image/jpeg" },
     ]);
   });
+
+  it("does not collect an attachment persisted while garbage collection is starting", async () => {
+    let releaseSave: () => void = () => undefined;
+    let reportSaveStarted: () => void = () => undefined;
+    const saveStarted = new Promise<void>((resolve) => {
+      reportSaveStarted = resolve;
+    });
+    const saveGate = new Promise<void>((resolve) => {
+      releaseSave = resolve;
+    });
+    const garbageCollections: string[][] = [];
+    const store: AttachmentStore = {
+      ...createRecordingStore(),
+      async save(input) {
+        reportSaveStarted();
+        await saveGate;
+        return createAttachment({ id: input.id });
+      },
+      async garbageCollect({ referencedIds }) {
+        garbageCollections.push([...referencedIds]);
+      },
+    };
+    __setAttachmentStoreForTests(store);
+
+    const persist = persistAttachmentFromBytes({
+      id: "assistant-preview",
+      bytes: new Uint8Array([1, 2, 3]),
+      mimeType: "image/png",
+    });
+    await saveStarted;
+    const collect = garbageCollectAttachments({ referencedIds: new Set() });
+
+    try {
+      await Promise.resolve();
+      expect(garbageCollections).toEqual([]);
+    } finally {
+      releaseSave();
+      await Promise.all([persist, collect]);
+    }
+
+    expect(garbageCollections).toEqual([["assistant-preview"]]);
+  });
 });

packages/app/src/attachments/service.ts

@@ -1,5 +1,40 @@
-import type { AttachmentMetadata } from "@/attachments/types";
+import { collectRetainedAttachmentIds } from "@/attachments/gc-retention";
 import { getAttachmentStore } from "@/attachments/store";
+import type { AttachmentMetadata, SaveAttachmentInput } from "@/attachments/types";
+
+const activePersistence = new Set<Promise<AttachmentMetadata>>();
+const persistedDuringGarbageCollection = new Set<string>();
+let pendingGarbageCollections = 0;
+let garbageCollectionTail: Promise<void> = Promise.resolve();
+let persistenceBarrier: Promise<void> | null = null;
+let releasePersistenceBarrier: (() => void) | null = null;
+
+async function waitForPersistenceBarrier(): Promise<void> {
+  const barrier = persistenceBarrier;
+  if (!barrier) {
+    return;
+  }
+  await barrier;
+  await waitForPersistenceBarrier();
+}
+
+async function persistAttachment(input: SaveAttachmentInput): Promise<AttachmentMetadata> {
+  await waitForPersistenceBarrier();
+  const pending = (async () => {
+    const store = await getAttachmentStore();
+    const attachment = await store.save(input);
+    if (pendingGarbageCollections > 0) {
+      persistedDuringGarbageCollection.add(attachment.id);
+    }
+    return attachment;
+  })();
+  activePersistence.add(pending);
+  try {
+    return await pending;
+  } finally {
+    activePersistence.delete(pending);
+  }
+}
 
 export async function persistAttachmentFromBlob(input: {
   blob: Blob;
@@ -7,8 +42,7 @@ export async function persistAttachmentFromBlob(input: {
   fileName?: string | null;
   id?: string;
 }): Promise<AttachmentMetadata> {
-  const store = await getAttachmentStore();
-  return await store.save({
+  return await persistAttachment({
     id: input.id,
     mimeType: input.mimeType,
     fileName: input.fileName,
@@ -22,8 +56,7 @@ export async function persistAttachmentFromDataUrl(input: {
   fileName?: string | null;
   id?: string;
 }): Promise<AttachmentMetadata> {
-  const store = await getAttachmentStore();
-  return await store.save({
+  return await persistAttachment({
     id: input.id,
     mimeType: input.mimeType,
     fileName: input.fileName,
@@ -37,8 +70,7 @@ export async function persistAttachmentFromBytes(input: {
   fileName?: string | null;
   id?: string;
 }): Promise<AttachmentMetadata> {
-  const store = await getAttachmentStore();
-  return await store.save({
+  return await persistAttachment({
     id: input.id,
     mimeType: input.mimeType,
     fileName: input.fileName,
@@ -52,8 +84,7 @@ export async function persistAttachmentFromFileUri(input: {
   fileName?: string | null;
   id?: string;
 }): Promise<AttachmentMetadata> {
-  const store = await getAttachmentStore();
-  return await store.save({
+  return await persistAttachment({
     id: input.id,
     mimeType: input.mimeType,
     fileName: input.fileName,
@@ -133,6 +164,41 @@ export async function deleteAttachments(
 export async function garbageCollectAttachments(input: {
   referencedIds: ReadonlySet<string>;
 }): Promise<void> {
-  const store = await getAttachmentStore();
-  await store.garbageCollect({ referencedIds: input.referencedIds });
+  pendingGarbageCollections += 1;
+  if (!persistenceBarrier) {
+    persistenceBarrier = new Promise<void>((resolve) => {
+      releasePersistenceBarrier = resolve;
+    });
+  }
+
+  const previousGarbageCollection = garbageCollectionTail;
+  const currentGarbageCollection = (async () => {
+    await previousGarbageCollection;
+    while (activePersistence.size > 0) {
+      await Promise.allSettled(activePersistence);
+    }
+    const referencedIds = new Set(input.referencedIds);
+    for (const id of collectRetainedAttachmentIds()) {
+      referencedIds.add(id);
+    }
+    for (const id of persistedDuringGarbageCollection) {
+      referencedIds.add(id);
+    }
+    const store = await getAttachmentStore();
+    await store.garbageCollect({ referencedIds });
+  })();
+  garbageCollectionTail = currentGarbageCollection.catch(() => undefined);
+
+  try {
+    await currentGarbageCollection;
+  } finally {
+    pendingGarbageCollections -= 1;
+    if (pendingGarbageCollections === 0) {
+      persistedDuringGarbageCollection.clear();
+      const release = releasePersistenceBarrier;
+      releasePersistenceBarrier = null;
+      persistenceBarrier = null;
+      release?.();
+    }
+  }
 }

packages/app/src/timeline/session-stream-reducers.test.ts

@@ -416,6 +416,112 @@ describe("processTimelineResponse", () => {
     ]);
   });
 
+  it("keeps one newer todo state when bootstrap contains its older state", () => {
+    const liveTodo = hydrateStreamState(
+      [
+        {
+          event: {
+            type: "timeline",
+            provider: "codex",
+            item: {
+              type: "todo",
+              items: [{ text: "Verify hydration", completed: true }],
+            },
+          } as AgentStreamEventPayload,
+          timestamp: new Date(2000),
+          timelineCursor: { epoch: "epoch-1", seq: 2 },
+        },
+      ],
+      { source: "canonical" },
+    );
+
+    const result = processTimelineResponse({
+      ...baseTimelineInput,
+      currentHead: liveTodo,
+      hasAuthoritativeBaseline: false,
+      isInitializing: true,
+      hasActiveInitDeferred: true,
+      payload: {
+        ...baseTimelineInput.payload,
+        direction: "tail",
+        reset: false,
+        epoch: "epoch-1",
+        startCursor: { seq: 1 },
+        endCursor: { seq: 1 },
+        entries: [
+          {
+            seqStart: 1,
+            seqEnd: 1,
+            provider: "codex",
+            item: {
+              type: "todo",
+              items: [{ text: "Verify hydration", completed: false }],
+            },
+            timestamp: new Date(1000).toISOString(),
+          },
+        ],
+      },
+    });
+
+    expect([...result.tail, ...result.head].filter((item) => item.kind === "todo_list")).toEqual([
+      expect.objectContaining({
+        timelineCursor: { epoch: "epoch-1", seq: 2 },
+        items: [{ text: "Verify hydration", completed: true }],
+      }),
+    ]);
+  });
+
+  it("keeps one completed compaction when bootstrap contains its loading state", () => {
+    const liveCompaction = hydrateStreamState(
+      [
+        {
+          event: {
+            type: "timeline",
+            provider: "codex",
+            item: { type: "compaction", status: "completed", trigger: "auto", preTokens: 1200 },
+          } as AgentStreamEventPayload,
+          timestamp: new Date(2000),
+          timelineCursor: { epoch: "epoch-1", seq: 2 },
+        },
+      ],
+      { source: "canonical" },
+    );
+
+    const result = processTimelineResponse({
+      ...baseTimelineInput,
+      currentHead: liveCompaction,
+      hasAuthoritativeBaseline: false,
+      isInitializing: true,
+      hasActiveInitDeferred: true,
+      payload: {
+        ...baseTimelineInput.payload,
+        direction: "tail",
+        reset: false,
+        epoch: "epoch-1",
+        startCursor: { seq: 1 },
+        endCursor: { seq: 1 },
+        entries: [
+          {
+            seqStart: 1,
+            seqEnd: 1,
+            provider: "codex",
+            item: { type: "compaction", status: "loading", trigger: "auto" },
+            timestamp: new Date(1000).toISOString(),
+          },
+        ],
+      },
+    });
+
+    expect([...result.tail, ...result.head].filter((item) => item.kind === "compaction")).toEqual([
+      expect.objectContaining({
+        timelineCursor: { epoch: "epoch-1", seq: 2 },
+        status: "completed",
+        trigger: "auto",
+        preTokens: 1200,
+      }),
+    ]);
+  });
+
   it("replaces a painted replica with an empty authoritative timeline", () => {
     const result = processTimelineResponse({
       ...baseTimelineInput,

packages/app/src/types/stream.ts

@@ -410,6 +410,61 @@ function removeUserMessageAt(items: UserMessageItem[], index: number): UserMessa
   return [...items.slice(0, index), ...items.slice(index + 1)];
 }
 
+function mergeRetainedLifecycleItem(tail: StreamItem[], retained: StreamItem): StreamItem[] | null {
+  if (!retained.timelineCursor) {
+    return null;
+  }
+  if (isAgentToolCallItem(retained)) {
+    const tailIndex = findExistingAgentToolCallIndex(tail, retained.payload.data.callId);
+    const existing = tail[tailIndex];
+    if (tailIndex < 0 || !existing || !isAgentToolCallItem(existing)) {
+      return null;
+    }
+    const next = [...tail];
+    next[tailIndex] = mergeAgentToolCallItem(
+      existing,
+      retained.payload.data,
+      retained.timestamp,
+      retained.timelineCursor,
+    );
+    return next;
+  }
+  if (retained.kind === "todo_list") {
+    const tailIndex = tail.length - 1;
+    const existing = tail[tailIndex];
+    if (!existing || existing.kind !== "todo_list" || existing.provider !== retained.provider) {
+      return null;
+    }
+    const next = [...tail];
+    next[tailIndex] = {
+      ...existing,
+      timelineCursor: retained.timelineCursor,
+      timestamp: retained.timestamp,
+      items: retained.items,
+    };
+    return next;
+  }
+  if (retained.kind === "compaction" && retained.status === "completed") {
+    const tailIndex = tail.findIndex(
+      (item) => item.kind === "compaction" && item.status === "loading",
+    );
+    const existing = tail[tailIndex];
+    if (tailIndex < 0 || !existing || existing.kind !== "compaction") {
+      return null;
+    }
+    const next = [...tail];
+    next[tailIndex] = {
+      ...existing,
+      timelineCursor: retained.timelineCursor,
+      status: "completed",
+      trigger: retained.trigger ?? existing.trigger,
+      preTokens: retained.preTokens ?? existing.preTokens,
+    };
+    return next;
+  }
+  return null;
+}
+
 function reconcileReplacementHeadAgainstTail(
   tail: StreamItem[],
   retainedHead: StreamItem[],
@@ -417,23 +472,11 @@ function reconcileReplacementHeadAgainstTail(
   let reconciledTail = tail;
   const reconciledHeadIndexes = new Set<number>();
   for (const [headIndex, item] of retainedHead.entries()) {
-    if (!isAgentToolCallItem(item) || !item.timelineCursor) {
+    const nextTail = mergeRetainedLifecycleItem(reconciledTail, item);
+    if (!nextTail) {
       continue;
     }
-    const tailIndex = findExistingAgentToolCallIndex(reconciledTail, item.payload.data.callId);
-    const existing = reconciledTail[tailIndex];
-    if (tailIndex < 0 || !existing || !isAgentToolCallItem(existing)) {
-      continue;
-    }
-    if (reconciledTail === tail) {
-      reconciledTail = [...tail];
-    }
-    reconciledTail[tailIndex] = mergeAgentToolCallItem(
-      existing,
-      item.payload.data,
-      item.timestamp,
-      item.timelineCursor,
-    );
+    reconciledTail = nextTail;
     reconciledHeadIndexes.add(headIndex);
   }