Keep mounted image attachments retained

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());
+  });
 });

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<AttachmentMetadata> {
-  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<typeof acquireCurrent> | 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]);