diff --git a/packages/app/src/assistant-image/use-assistant-image.test.tsx b/packages/app/src/assistant-image/use-assistant-image.test.tsx index 0ea29f9070..03a6759a68 100644 --- a/packages/app/src/assistant-image/use-assistant-image.test.tsx +++ b/packages/app/src/assistant-image/use-assistant-image.test.tsx @@ -4,11 +4,17 @@ import { act, cleanup, renderHook, waitFor } from "@testing-library/react"; import { afterEach, describe, expect, it } from "vitest"; import type { ImageLoadEvent } from "react-native"; +import type { DaemonClient } from "@getpaseo/client/internal/daemon-client"; +import type { AttachmentStore } from "@/attachments/types"; +import { __setAttachmentStoreForTests } from "@/attachments/store"; import { setAssistantImageMetadata } from "@/utils/assistant-image-metadata"; import { useAssistantImage } from "./use-assistant-image"; describe("useAssistantImage", () => { - afterEach(cleanup); + afterEach(() => { + cleanup(); + __setAttachmentStoreForTests(null); + }); it("keeps a successfully loaded image when cached dimensions are available", async () => { const source = "https://example.com/expiring-image.png"; @@ -101,4 +107,67 @@ describe("useAssistantImage", () => { expect(remounted.result.current).toMatchObject({ status: "loaded", aspectRatio: 2 }); }); + + it("restarts a failed file acquisition when the daemon client reconnects", async () => { + const attachmentStore: AttachmentStore = { + storageType: "web-indexeddb", + async save(input) { + return { + id: input.id ?? "assistant-image", + mimeType: input.mimeType ?? "image/png", + storageType: "web-indexeddb", + storageKey: input.id ?? "assistant-image", + fileName: input.fileName, + byteSize: 4, + createdAt: 1, + }; + }, + async encodeBase64() { + return ""; + }, + async resolvePreviewUrl({ attachment }) { + return `blob:${attachment.id}`; + }, + async delete() {}, + async garbageCollect() {}, + }; + __setAttachmentStoreForTests(attachmentStore); + + let readCount = 0; + const connectedClient = { + async readFile(_cwd: string, path: string) { + readCount += 1; + return { + kind: "image" as const, + path, + mime: "image/png", + size: 4, + modifiedAt: "1", + bytes: new Uint8Array([1, 2, 3, 4]), + }; + }, + } as unknown as DaemonClient; + const initialProps: { client: DaemonClient | null } = { client: null }; + const { result, rerender } = renderHook( + ({ client }: { client: DaemonClient | null }) => + useAssistantImage({ + source: "/workspace/reconnect.png", + occurrenceKey: "agent:message:reconnect-image", + client, + workspaceRoot: "/workspace", + serverId: "server", + }), + { initialProps }, + ); + + await waitFor(() => expect(result.current.status).toBe("failed"), { timeout: 3_000 }); + + rerender({ client: connectedClient }); + + await waitFor(() => { + expect(readCount).toBe(1); + expect(result.current.status).toBe("loading"); + expect(result.current.status === "loading" && result.current.binding).not.toBeNull(); + }); + }); }); diff --git a/packages/app/src/assistant-image/use-assistant-image.ts b/packages/app/src/assistant-image/use-assistant-image.ts index 10e2432ea6..ecf3d2a8b7 100644 --- a/packages/app/src/assistant-image/use-assistant-image.ts +++ b/packages/app/src/assistant-image/use-assistant-image.ts @@ -179,16 +179,10 @@ function useAttachmentAcquisition( state: cached ? { status: "loaded", attachment: cached } : { status: "waiting" }, }; }); - const acquireCurrent = useStableEvent(async () => { - if (!acquisition) { - return null; - } - return await acquireAttachment(acquisition); - }); useEffect(() => { let disposed = false; - if (!acquisitionKey) { + if (!acquisition || !acquisitionKey) { setEntry({ key: null, state: { status: "waiting" } }); return; } @@ -203,13 +197,7 @@ function useAttachmentAcquisition( void (async () => { try { const attachment = await runAssistantImageOperationWithRetry({ - operation: async () => { - const current = await acquireCurrent(); - if (!current) { - throw new Error("Assistant image acquisition is unavailable."); - } - return current; - }, + operation: async () => await acquireAttachment(acquisition), shouldStop: () => disposed, }); if (!disposed) { @@ -224,7 +212,7 @@ function useAttachmentAcquisition( return () => { disposed = true; }; - }, [acquireCurrent, acquisitionKey]); + }, [acquisition, acquisitionKey]); if (!acquisitionKey) { return { status: "waiting" };