From 1d161a83f4b1d73e35db29a0fb2026bcd0aacc45 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Mon, 27 Jul 2026 18:35:53 +0200 Subject: [PATCH] fix(app): keep file image previews current --- packages/app/src/agent-stream/view.tsx | 1 + .../assistant-image/acquisition-cache.test.ts | 45 ++++++++++++++++++- .../src/assistant-image/acquisition-cache.ts | 9 ++++ .../assistant-image/use-assistant-image.ts | 23 ++++++++-- packages/app/src/components/message.tsx | 15 ++++++- 5 files changed, 86 insertions(+), 7 deletions(-) diff --git a/packages/app/src/agent-stream/view.tsx b/packages/app/src/agent-stream/view.tsx index 0bc079e78c..0d6bde3c08 100644 --- a/packages/app/src/agent-stream/view.tsx +++ b/packages/app/src/agent-stream/view.tsx @@ -694,6 +694,7 @@ const AgentStreamViewComponent = forwardRef { it("evicts a rejected acquisition so the next request can retry", async () => { @@ -43,4 +46,44 @@ describe("assistant image acquisition cache", () => { size: 2, }); }); + + it("reuses an acquired image when the current locator is unavailable", async () => { + const cache = createAssistantImageAcquisitionCache({ capacity: 2 }); + let unavailableCalls = 0; + + await cache.acquire("message:image", async () => "persisted attachment"); + const cached = await cache.acquire("message:image", async () => { + unavailableCalls += 1; + throw new Error("daemon disconnected"); + }); + + expect({ cached, unavailableCalls }).toEqual({ + cached: "persisted attachment", + unavailableCalls: 0, + }); + }); + + it("scopes file acquisitions to the rendered message occurrence", () => { + const first = createAssistantImageFileAcquisitionKey({ + serverId: "server", + occurrenceKey: "message-1:image-1", + cwd: "/workspace", + path: "screenshot.png", + }); + const remount = createAssistantImageFileAcquisitionKey({ + serverId: "server", + occurrenceKey: "message-1:image-1", + cwd: "/workspace", + path: "screenshot.png", + }); + const laterMessage = createAssistantImageFileAcquisitionKey({ + serverId: "server", + occurrenceKey: "message-2:image-1", + cwd: "/workspace", + path: "screenshot.png", + }); + + expect(remount).toBe(first); + expect(laterMessage).not.toBe(first); + }); }); diff --git a/packages/app/src/assistant-image/acquisition-cache.ts b/packages/app/src/assistant-image/acquisition-cache.ts index 6882c1c578..2077ab02a9 100644 --- a/packages/app/src/assistant-image/acquisition-cache.ts +++ b/packages/app/src/assistant-image/acquisition-cache.ts @@ -3,6 +3,15 @@ export interface AssistantImageAcquisitionCache { size(): number; } +export function createAssistantImageFileAcquisitionKey(input: { + serverId?: string; + occurrenceKey: string; + cwd: string; + path: string; +}): string { + return `file:${input.serverId ?? "unknown-server"}:${input.occurrenceKey}:${input.cwd}:${input.path}`; +} + export function createAssistantImageAcquisitionCache(input: { capacity: number; }): AssistantImageAcquisitionCache { diff --git a/packages/app/src/assistant-image/use-assistant-image.ts b/packages/app/src/assistant-image/use-assistant-image.ts index c433882292..fd60222823 100644 --- a/packages/app/src/assistant-image/use-assistant-image.ts +++ b/packages/app/src/assistant-image/use-assistant-image.ts @@ -20,7 +20,10 @@ import { } from "@/utils/assistant-image-metadata"; import { resolveAssistantImageSource } from "@/utils/assistant-image-source"; import type { AssistantImageSourceResolution } from "@/utils/assistant-image-source"; -import { createAssistantImageAcquisitionCache } from "./acquisition-cache"; +import { + createAssistantImageAcquisitionCache, + createAssistantImageFileAcquisitionKey, +} from "./acquisition-cache"; import { createAssistantImageLifecycle, transitionAssistantImageLifecycle, @@ -49,6 +52,7 @@ export type AssistantImageResult = interface UseAssistantImageInput { source: string; + occurrenceKey: string; client?: DaemonClient | null; workspaceRoot?: string; serverId?: string; @@ -122,15 +126,24 @@ function createFileAcquisition(input: { client?: DaemonClient | null; resolution: AssistantImageSourceResolution | null; serverId?: string; + occurrenceKey: string; unavailableMessage: string; }): AttachmentAcquisition | null { - if (!input.client || input.resolution?.kind !== "file_rpc") { + if (input.resolution?.kind !== "file_rpc") { return null; } const { client, resolution } = input; return { - key: `file:${input.serverId ?? "unknown-server"}:${resolution.cwd}:${resolution.path}`, + key: createAssistantImageFileAcquisitionKey({ + serverId: input.serverId, + occurrenceKey: input.occurrenceKey, + cwd: resolution.cwd, + path: resolution.path, + }), locate: async () => { + if (!client) { + throw new Error(input.unavailableMessage); + } const file = await client.readFile(resolution.cwd, resolution.path); if (file.kind !== "image") { throw new Error(input.unavailableMessage); @@ -265,6 +278,7 @@ function getAcquisitionFailure(input: { export function useAssistantImage({ source, + occurrenceKey, client, workspaceRoot, serverId, @@ -281,9 +295,10 @@ export function useAssistantImage({ client, resolution, serverId, + occurrenceKey, unavailableMessage: t("message.attachments.imagePreviewUnavailable"), }), - [client, resolution, serverId, t], + [client, occurrenceKey, resolution, serverId, t], ); const dataImageAcquisition = useMemo( () => createDataImageAcquisition({ source, dataImage }), diff --git a/packages/app/src/components/message.tsx b/packages/app/src/components/message.tsx index dfb0cfb715..3c1cf35d07 100644 --- a/packages/app/src/components/message.tsx +++ b/packages/app/src/components/message.tsx @@ -723,6 +723,7 @@ export const LiveElapsed = memo(function LiveElapsed({ }); interface AssistantMessageProps { + occurrenceKey: string; message: string; timestamp: number; workspaceRoot?: string; @@ -783,6 +784,7 @@ const ASSISTANT_IMAGE_MIN_HEIGHT = 160; function AssistantMarkdownImage({ source, + occurrenceKey, alt, hasLeadingContent, client, @@ -790,6 +792,7 @@ function AssistantMarkdownImage({ serverId, }: { source: string; + occurrenceKey: string; alt?: string; hasLeadingContent: boolean; client?: DaemonClient | null; @@ -803,7 +806,13 @@ function AssistantMarkdownImage({ }), [hasLeadingContent], ); - const image = useAssistantImage({ source, client, workspaceRoot, serverId }); + const image = useAssistantImage({ + source, + occurrenceKey, + client, + workspaceRoot, + serverId, + }); const binding = image.status === "failed" ? null : image.binding; const aspectRatio = image.status === "failed" ? null : image.aspectRatio; const imageUri = binding?.uri ?? ""; @@ -1410,6 +1419,7 @@ function MarkdownListView({ baseStyle, spacing, children }: MarkdownListViewProp } export const AssistantMessage = memo(function AssistantMessage({ + occurrenceKey, message, timestamp: _timestamp, workspaceRoot, @@ -1744,6 +1754,7 @@ export const AssistantMessage = memo(function AssistantMessage({ splitMarkdownBlocks(message), [message]); const keyedBlocks = useMemo(