fix(app): retain assistant image previews

packages/app/e2e/agent-tab-image-stability.spec.ts

@@ -1,6 +1,11 @@
 import type { Page } from "@playwright/test";
 import { expect, test as base } from "./fixtures";
-import { openWorkspaceWithAgents, type ArchiveTabAgent } from "./helpers/archive-tab";
+import {
+  clickSessionRow,
+  openSessions,
+  openWorkspaceWithAgents,
+  type ArchiveTabAgent,
+} from "./helpers/archive-tab";
 import { seedWorkspace, type SeedDaemonClient, type SeededWorkspace } from "./helpers/seed-client";
 
 const IMAGE_ALT = "Synthetic image";
@@ -85,3 +90,35 @@ test("switching between settled agent tabs keeps the image timeline valid", asyn
   await expect(page.getByRole("img", { name: IMAGE_ALT })).toHaveCount(1, { timeout: 30_000 });
   await expectNoObservedImagePreviewError(page);
 });
+
+test("reopening an image timeline after editing another draft keeps its image available", async ({
+  imageWorkspace: workspace,
+  page,
+}) => {
+  test.setTimeout(120_000);
+  const imageAgent = await createMockAgent(workspace.client, {
+    cwd: workspace.repoPath,
+    workspaceId: workspace.workspaceId,
+    title: `Image cleanup ${Date.now()}`,
+  });
+  const otherAgent = await createMockAgent(workspace.client, {
+    cwd: workspace.repoPath,
+    workspaceId: workspace.workspaceId,
+    title: `Draft cleanup ${Date.now()}`,
+  });
+  await workspace.client.sendAgentMessage(imageAgent.id, IMAGE_PROMPT);
+  await workspace.client.waitForFinish(imageAgent.id, 15_000);
+  await openWorkspaceWithAgents(page, [otherAgent, imageAgent]);
+  await expect(page.getByRole("img", { name: IMAGE_ALT })).toBeVisible({ timeout: 30_000 });
+
+  await page.getByRole("button", { name: otherAgent.title, exact: true }).click();
+  await page
+    .getByRole("textbox", { name: "Message agent..." })
+    .fill("Keep this draft while I check history");
+  await openSessions(page);
+  await observeImagePreviewErrors(page);
+  await clickSessionRow(page, imageAgent.title);
+
+  await expect(page.getByRole("img", { name: IMAGE_ALT })).toBeVisible({ timeout: 30_000 });
+  await expectNoObservedImagePreviewError(page);
+});

packages/app/src/assistant-image/acquisition-cache.test.ts

@@ -2,6 +2,7 @@ import { describe, expect, it } from "vitest";
 import {
   createAssistantImageAcquisitionCache,
   createAssistantImageFileAcquisitionKey,
+  createAssistantImageFilePreviewAttachmentId,
 } from "./acquisition-cache";
 
 describe("assistant image acquisition cache", () => {
@@ -86,4 +87,46 @@ describe("assistant image acquisition cache", () => {
     expect(remount).toBe(first);
     expect(laterMessage).not.toBe(first);
   });
+
+  it("scopes persisted file previews to the rendered message occurrence", () => {
+    const first = createAssistantImageFilePreviewAttachmentId({
+      occurrenceKey: "message-1:image-1",
+      mimeType: "image/png",
+      path: "/workspace/screenshot.png",
+      size: 512,
+      modifiedAt: "2026-07-27T12:00:00.000Z",
+      contentLength: 512,
+    });
+    const second = createAssistantImageFilePreviewAttachmentId({
+      occurrenceKey: "message-2:image-1",
+      mimeType: "image/png",
+      path: "/workspace/screenshot.png",
+      size: 512,
+      modifiedAt: "2026-07-27T12:00:00.000Z",
+      contentLength: 512,
+    });
+
+    expect(second).not.toBe(first);
+  });
+
+  it("retains successful values until their cache entry is evicted", async () => {
+    const retained: string[] = [];
+    const released: string[] = [];
+    const cache = createAssistantImageAcquisitionCache<string>({
+      capacity: 1,
+      onRetain(value) {
+        retained.push(value);
+        return () => released.push(value);
+      },
+    });
+
+    await cache.acquire("first", async () => "attachment-1");
+    expect({ retained, released }).toEqual({ retained: ["attachment-1"], released: [] });
+
+    await cache.acquire("second", async () => "attachment-2");
+    expect({ retained, released }).toEqual({
+      retained: ["attachment-1", "attachment-2"],
+      released: ["attachment-1"],
+    });
+  });
 });

packages/app/src/assistant-image/acquisition-cache.ts

@@ -1,8 +1,28 @@
+import { createPreviewAttachmentId } from "@/attachments/utils";
+
 export interface AssistantImageAcquisitionCache<T> {
   acquire(key: string, locate: () => Promise<T>): Promise<T>;
   size(): number;
 }
 
+export function createAssistantImageFilePreviewAttachmentId(input: {
+  occurrenceKey: string;
+  mimeType: string;
+  path: string;
+  size: number;
+  modifiedAt?: string | null;
+  contentLength: number;
+}): string {
+  return createPreviewAttachmentId({
+    mimeType: input.mimeType,
+    path: input.path,
+    size: input.size,
+    modifiedAt: input.modifiedAt,
+    contentLength: input.contentLength,
+    contentKey: input.occurrenceKey,
+  });
+}
+
 export function createAssistantImageFileAcquisitionKey(input: {
   serverId?: string;
   occurrenceKey: string;
@@ -14,11 +34,24 @@ export function createAssistantImageFileAcquisitionKey(input: {
 
 export function createAssistantImageAcquisitionCache<T>(input: {
   capacity: number;
+  onRetain?: (value: T) => () => void;
 }): AssistantImageAcquisitionCache<T> {
   if (!Number.isInteger(input.capacity) || input.capacity < 1) {
     throw new Error("Assistant image acquisition cache capacity must be a positive integer.");
   }
-  const entries = new Map<string, Promise<T>>();
+  interface CacheEntry {
+    pending: Promise<T>;
+    release: (() => void) | null;
+  }
+  const entries = new Map<string, CacheEntry>();
+
+  const evict = (key: string, entry: CacheEntry) => {
+    if (entries.get(key) === entry) {
+      entries.delete(key);
+    }
+    entry.release?.();
+    entry.release = null;
+  };
 
   return {
     acquire(key, locate) {
@@ -26,23 +59,28 @@ export function createAssistantImageAcquisitionCache<T>(input: {
       if (cached) {
         entries.delete(key);
         entries.set(key, cached);
-        return cached;
+        return cached.pending;
       }
       const pending = locate();
-      entries.set(key, pending);
+      const entry: CacheEntry = { pending, release: null };
+      entries.set(key, entry);
       if (entries.size > input.capacity) {
         const leastRecentlyUsedKey = entries.keys().next().value;
         if (leastRecentlyUsedKey !== undefined) {
-          entries.delete(leastRecentlyUsedKey);
+          const leastRecentlyUsedEntry = entries.get(leastRecentlyUsedKey);
+          if (leastRecentlyUsedEntry) {
+            evict(leastRecentlyUsedKey, leastRecentlyUsedEntry);
+          }
         }
       }
       void (async () => {
         try {
-          await pending;
-        } catch {
-          if (entries.get(key) === pending) {
-            entries.delete(key);
+          const value = await pending;
+          if (entries.get(key) === entry) {
+            entry.release = input.onRetain?.(value) ?? null;
           }
+        } catch {
+          evict(key, entry);
         }
       })();
       return pending;

packages/app/src/assistant-image/use-assistant-image.ts

@@ -3,6 +3,7 @@ import { Image } from "react-native";
 import { useTranslation } from "react-i18next";
 import type { DaemonClient } from "@getpaseo/client/internal/daemon-client";
 import type { AttachmentMetadata } from "@/attachments/types";
+import { retainAttachmentForGarbageCollection } from "@/attachments/gc-retention";
 import {
   persistAttachmentFromBytes,
   persistAttachmentFromDataUrl,
@@ -23,6 +24,7 @@ import type { AssistantImageSourceResolution } from "@/utils/assistant-image-sou
 import {
   createAssistantImageAcquisitionCache,
   createAssistantImageFileAcquisitionKey,
+  createAssistantImageFilePreviewAttachmentId,
 } from "./acquisition-cache";
 import {
   createAssistantImageLifecycle,
@@ -84,6 +86,7 @@ interface DataImage {
 
 const attachmentAcquisitionCache = createAssistantImageAcquisitionCache<AttachmentMetadata>({
   capacity: 500,
+  onRetain: (attachment) => retainAttachmentForGarbageCollection(attachment.id),
 });
 
 function acquireAttachment(acquisition: AttachmentAcquisition): Promise<AttachmentMetadata> {
@@ -153,7 +156,8 @@ function createFileAcquisition(input: {
         throw new Error(input.unavailableMessage);
       }
       return await persistAttachmentFromBytes({
-        id: createPreviewAttachmentId({
+        id: createAssistantImageFilePreviewAttachmentId({
+          occurrenceKey: input.occurrenceKey,
           mimeType: file.mime,
           path: file.path || resolution.path,
           size: file.size,

packages/app/src/attachments/gc-retention.test.ts

@@ -0,0 +1,15 @@
+import { describe, expect, it } from "vitest";
+import { collectRetainedAttachmentIds, retainAttachmentForGarbageCollection } from "./gc-retention";
+
+describe("attachment garbage-collection retention", () => {
+  it("keeps an attachment referenced until its final owner releases it", () => {
+    const releaseFirst = retainAttachmentForGarbageCollection("preview-1");
+    const releaseSecond = retainAttachmentForGarbageCollection("preview-1");
+
+    expect(collectRetainedAttachmentIds()).toContain("preview-1");
+    releaseFirst();
+    expect(collectRetainedAttachmentIds()).toContain("preview-1");
+    releaseSecond();
+    expect(collectRetainedAttachmentIds()).not.toContain("preview-1");
+  });
+});

packages/app/src/attachments/gc-retention.ts

@@ -0,0 +1,22 @@
+const retentionCounts = new Map<string, number>();
+
+export function retainAttachmentForGarbageCollection(attachmentId: string): () => void {
+  retentionCounts.set(attachmentId, (retentionCounts.get(attachmentId) ?? 0) + 1);
+  let released = false;
+  return () => {
+    if (released) {
+      return;
+    }
+    released = true;
+    const nextCount = (retentionCounts.get(attachmentId) ?? 1) - 1;
+    if (nextCount <= 0) {
+      retentionCounts.delete(attachmentId);
+      return;
+    }
+    retentionCounts.set(attachmentId, nextCount);
+  };
+}
+
+export function collectRetainedAttachmentIds(): ReadonlySet<string> {
+  return new Set(retentionCounts.keys());
+}

packages/app/src/stores/draft-store/index.ts

@@ -8,6 +8,7 @@ import {
   persistAttachmentFromDataUrl,
   persistAttachmentFromFileUri,
 } from "@/attachments/service";
+import { collectRetainedAttachmentIds } from "@/attachments/gc-retention";
 import { useCreateFlowStore } from "@/stores/create-flow-store";
 import { useSessionStore, type SessionState } from "@/stores/session-store";
 import { useWorkspaceAttachmentsStore } from "@/attachments/workspace-attachments-store";
@@ -138,6 +139,9 @@ async function runAttachmentGc(): Promise<void> {
   for (const id of useDraftStore.getState().collectActiveAttachmentIds()) {
     referencedIds.add(id);
   }
+  for (const id of collectRetainedAttachmentIds()) {
+    referencedIds.add(id);
+  }
 
   const pendingByDraftId = useCreateFlowStore.getState().pendingByDraftId;
   for (const pendingCreate of Object.values(pendingByDraftId)) {