Retry file images after reconnect

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

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" };