fix(agents): recover images when regeneration fails

src/features/agents/messageEditing.test.tsx

@@ -25,7 +25,8 @@ vi.mock("@tauri-apps/api/webview", () => ({
 }));
 
 import { SessionEditor } from "./SessionEditor";
-import { updateComposer, useAgents } from "./store";
+import { composerDraft, sendMessage, updateComposer, useAgents } from "./store";
+import { replyActions } from "./replyActions";
 
 const task: AgentTask = {
   id: "editing",
@@ -64,6 +65,46 @@ async function edit() {
   return { user, box, send: screen.getAllByRole("button", { name: "Send" })[0]! };
 }
 
+it.each([false, true])("recovers images after failed regeneration (new draft: %s)", async (draft) => {
+  const image = { mimeType: "image/png", data: "AAA", name: "shot.png" };
+  const added = { mimeType: "image/png", data: "BBB", name: "added.png" };
+  useAgents.setState({
+    entries: new Map([
+      [
+        task.id,
+        [
+          { type: "user", text: "Original message", attachments: [image.name], images: [image], steered: false },
+          entries[1]!,
+        ],
+      ],
+    ]),
+  });
+  ipc.rollback.mockImplementation(async () => {
+    useAgents.setState({ entries: new Map([[task.id, []]]) });
+    return { status: "ok", data: null };
+  });
+  let finish!: (value: unknown) => void;
+  ipc.prompt.mockReturnValueOnce(new Promise((resolve) => (finish = resolve)));
+  const running = replyActions
+    .find((a) => a.id === "agents.regenerateReply")!
+    .run({ ctx: {}, arg: { task: task.id, turn: 0 } });
+  await waitFor(() => expect(ipc.prompt).toHaveBeenCalledTimes(1));
+  const files = draft ? [{ path: "/project/a.ts", rel: "a.ts" }] : [];
+  if (draft) updateComposer(task.id, { text: "Another draft @a.ts", files, images: [added] });
+  finish({ status: "error", error: { kind: "internal", message: "Send failed" } });
+  await running;
+  const expected = {
+    text: draft ? "Original message\nAnother draft @a.ts" : "Original message",
+    files,
+    images: draft ? [image, added] : [image],
+  };
+  expect(composerDraft(task.id)).toEqual(expected);
+  expect(useAgents.getState().focus?.key).toBe(task.id);
+  await sendMessage(task.id, false);
+  expect(ipc.prompt).toHaveBeenLastCalledWith(task.id, { ...expected, files: files.map((f) => f.path) }, false);
+  expect(ipc.rollback).toHaveBeenCalledTimes(1);
+});
+
 describe("a sent message with images", () => {
   afterEach(() => {
     entries[0] = { type: "user", text: "Original message", attachments: [], sentMs: 1, doneMs: 2, steered: false };

src/features/agents/store.ts

@@ -555,9 +555,12 @@ export async function rerunFrom(task: string, turn: number, text: string, images
   try {
     await call(commands.agentPrompt(task, { text, files: [], images }, false));
   } catch (e) {
-    // Rollback removes the inline editor; recover its text in the session composer.
+    // Rollback removes the original turn; recover its prompt in the session composer.
     const current = composerDraft(task);
-    updateComposer(task, { text: current.text ? `${text}\n${current.text}` : text });
+    updateComposer(task, {
+      text: current.text ? `${text}\n${current.text}` : text,
+      images: [...images, ...current.images],
+    });
     focusComposer(task);
     throw e;
   }