From 5c126fd0b9d85258cfdd8b2d98acfe6d5f55dda6 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Tue, 28 Jul 2026 00:04:22 +0200 Subject: [PATCH] Protect preview persistence and lifecycle hydration --- packages/app/src/attachments/service.test.ts | 48 +++++++- packages/app/src/attachments/service.ts | 88 +++++++++++++-- .../timeline/session-stream-reducers.test.ts | 106 ++++++++++++++++++ packages/app/src/types/stream.ts | 73 +++++++++--- 4 files changed, 288 insertions(+), 27 deletions(-) diff --git a/packages/app/src/attachments/service.test.ts b/packages/app/src/attachments/service.test.ts index 080204ed0..05e13034f 100644 --- a/packages/app/src/attachments/service.test.ts +++ b/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 { 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((resolve) => { + reportSaveStarted = resolve; + }); + const saveGate = new Promise((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"]]); + }); }); diff --git a/packages/app/src/attachments/service.ts b/packages/app/src/attachments/service.ts index 15cff1ee2..d878b1986 100644 --- a/packages/app/src/attachments/service.ts +++ b/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>(); +const persistedDuringGarbageCollection = new Set(); +let pendingGarbageCollections = 0; +let garbageCollectionTail: Promise = Promise.resolve(); +let persistenceBarrier: Promise | null = null; +let releasePersistenceBarrier: (() => void) | null = null; + +async function waitForPersistenceBarrier(): Promise { + const barrier = persistenceBarrier; + if (!barrier) { + return; + } + await barrier; + await waitForPersistenceBarrier(); +} + +async function persistAttachment(input: SaveAttachmentInput): Promise { + 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 { - 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 { - 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 { - 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 { - 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; }): Promise { - const store = await getAttachmentStore(); - await store.garbageCollect({ referencedIds: input.referencedIds }); + pendingGarbageCollections += 1; + if (!persistenceBarrier) { + persistenceBarrier = new Promise((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?.(); + } + } } diff --git a/packages/app/src/timeline/session-stream-reducers.test.ts b/packages/app/src/timeline/session-stream-reducers.test.ts index a549571f4..dab7ec467 100644 --- a/packages/app/src/timeline/session-stream-reducers.test.ts +++ b/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, diff --git a/packages/app/src/types/stream.ts b/packages/app/src/types/stream.ts index 05f42c6fb..6586ad5e7 100644 --- a/packages/app/src/types/stream.ts +++ b/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(); 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); }