diff --git a/packages/app/src/assistant-image/acquisition-cache.test.ts b/packages/app/src/assistant-image/acquisition-cache.test.ts index 9a245f91f..ce1961b63 100644 --- a/packages/app/src/assistant-image/acquisition-cache.test.ts +++ b/packages/app/src/assistant-image/acquisition-cache.test.ts @@ -1,4 +1,8 @@ import { describe, expect, it } from "vitest"; +import { + collectRetainedAttachmentIds, + retainAttachmentForGarbageCollection, +} from "@/attachments/gc-retention"; import { createAssistantImageAcquisitionCache, createAssistantImageFileAcquisitionKey, @@ -162,4 +166,28 @@ describe("assistant image acquisition cache", () => { }); second.release(); }); + + it("protects every actively consumed attachment from garbage collection past capacity", async () => { + const cache = createAssistantImageAcquisitionCache<{ id: string }>({ + capacity: 1, + onRetain: (attachment) => retainAttachmentForGarbageCollection(attachment.id), + }); + + const first = cache.acquireRetained("first", async () => ({ id: "mounted-image-1" })); + await first.promise; + const second = cache.acquireRetained("second", async () => ({ id: "mounted-image-2" })); + await second.promise; + + expect(collectRetainedAttachmentIds()).toEqual(new Set(["mounted-image-1", "mounted-image-2"])); + + first.release(); + expect(collectRetainedAttachmentIds()).toEqual(new Set(["mounted-image-2"])); + second.release(); + await expect( + cache.acquire("cleanup", async () => { + throw new Error("cleanup"); + }), + ).rejects.toThrow("cleanup"); + expect(collectRetainedAttachmentIds()).toEqual(new Set()); + }); }); diff --git a/packages/app/src/assistant-image/use-assistant-image.ts b/packages/app/src/assistant-image/use-assistant-image.ts index 38090c45f..0c0818e8f 100644 --- a/packages/app/src/assistant-image/use-assistant-image.ts +++ b/packages/app/src/assistant-image/use-assistant-image.ts @@ -146,10 +146,6 @@ function rememberLoadedImage(uri: string, aspectRatio: number): void { } } -function acquireAttachment(acquisition: AssistantImageAcquisition): Promise { - return attachmentAcquisitionCache.acquire(acquisition.key, acquisition.locate); -} - function useAttachmentAcquisition( acquisition: AssistantImageAcquisition | null, ): AttachmentAcquisitionState { @@ -167,22 +163,40 @@ function useAttachmentAcquisition( useEffect(() => { let disposed = false; + let releaseCurrent: (() => void) | null = null; if (!acquisition || !acquisitionKey) { setEntry({ key: null, state: { status: "waiting" } }); return; } - const cached = attachmentAcquisitionCache.peek(acquisitionKey); - if (cached) { - setEntry({ key: acquisitionKey, state: { status: "loaded", attachment: cached } }); - return; + const acquireCurrent = () => { + releaseCurrent?.(); + const retained = attachmentAcquisitionCache.acquireRetained( + acquisitionKey, + acquisition.locate, + ); + releaseCurrent = retained.release; + return retained; + }; + const initial = acquireCurrent(); + if (initial.value) { + setEntry({ key: acquisitionKey, state: { status: "loaded", attachment: initial.value } }); + return () => { + disposed = true; + releaseCurrent?.(); + }; } setEntry({ key: acquisitionKey, state: { status: "loading" } }); void (async () => { + let firstAttempt: ReturnType | null = initial; try { const attachment = await runAssistantImageOperationWithRetry({ - operation: async () => await acquireAttachment(acquisition), + operation: async () => { + const retained = firstAttempt ?? acquireCurrent(); + firstAttempt = null; + return await retained.promise; + }, shouldStop: () => disposed, }); if (!disposed) { @@ -196,6 +210,7 @@ function useAttachmentAcquisition( })(); return () => { disposed = true; + releaseCurrent?.(); }; }, [acquisition, acquisitionKey]);