diff --git a/packages/app/e2e/agent-tab-image-stability.spec.ts b/packages/app/e2e/agent-tab-image-stability.spec.ts index e4e8acfa06..70d5227025 100644 --- a/packages/app/e2e/agent-tab-image-stability.spec.ts +++ b/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); +}); diff --git a/packages/app/src/assistant-image/acquisition-cache.test.ts b/packages/app/src/assistant-image/acquisition-cache.test.ts index 287eccd4d8..2b0ee9f675 100644 --- a/packages/app/src/assistant-image/acquisition-cache.test.ts +++ b/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({ + 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"], + }); + }); }); diff --git a/packages/app/src/assistant-image/acquisition-cache.ts b/packages/app/src/assistant-image/acquisition-cache.ts index 2077ab02a9..32f0586cce 100644 --- a/packages/app/src/assistant-image/acquisition-cache.ts +++ b/packages/app/src/assistant-image/acquisition-cache.ts @@ -1,8 +1,28 @@ +import { createPreviewAttachmentId } from "@/attachments/utils"; + export interface AssistantImageAcquisitionCache { acquire(key: string, locate: () => Promise): Promise; 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(input: { capacity: number; + onRetain?: (value: T) => () => void; }): AssistantImageAcquisitionCache { 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>(); + interface CacheEntry { + pending: Promise; + release: (() => void) | null; + } + const entries = new Map(); + + 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(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; diff --git a/packages/app/src/assistant-image/use-assistant-image.ts b/packages/app/src/assistant-image/use-assistant-image.ts index 32362b9cef..6125342e90 100644 --- a/packages/app/src/assistant-image/use-assistant-image.ts +++ b/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({ capacity: 500, + onRetain: (attachment) => retainAttachmentForGarbageCollection(attachment.id), }); function acquireAttachment(acquisition: AttachmentAcquisition): Promise { @@ -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, diff --git a/packages/app/src/attachments/gc-retention.test.ts b/packages/app/src/attachments/gc-retention.test.ts new file mode 100644 index 0000000000..4b39926814 --- /dev/null +++ b/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"); + }); +}); diff --git a/packages/app/src/attachments/gc-retention.ts b/packages/app/src/attachments/gc-retention.ts new file mode 100644 index 0000000000..4473350d9a --- /dev/null +++ b/packages/app/src/attachments/gc-retention.ts @@ -0,0 +1,22 @@ +const retentionCounts = new Map(); + +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 { + return new Set(retentionCounts.keys()); +} diff --git a/packages/app/src/stores/draft-store/index.ts b/packages/app/src/stores/draft-store/index.ts index c558d2fa03..5285170415 100644 --- a/packages/app/src/stores/draft-store/index.ts +++ b/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 { 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)) {