From f0486a6826dc487c516b96e6ca2e2f88485b2e66 Mon Sep 17 00:00:00 2001
From: l0ng-ai <24760907+l0ng-ai@users.noreply.github.com>
Date: Thu, 29 Oct 2026 21:26:57 +0800
Subject: [PATCH] fix(agents): recover edited draft when sending after rollback
 fails

---
 src/features/agents/messageEditing.test.tsx | 29 ++++++++++++++++++++-
 src/features/agents/store.ts                | 10 ++++++-
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/src/features/agents/messageEditing.test.tsx b/src/features/agents/messageEditing.test.tsx
index 6d98ffad..919d5a72 100644
--- a/src/features/agents/messageEditing.test.tsx
+++ b/src/features/agents/messageEditing.test.tsx
@@ -25,7 +25,7 @@ vi.mock("@tauri-apps/api/webview", () => ({
 }));
 
 import { SessionEditor } from "./SessionEditor";
-import { useAgents } from "./store";
+import { updateComposer, useAgents } from "./store";
 
 const task: AgentTask = {
   id: "editing",
@@ -110,6 +110,33 @@ describe("editing a sent agent message", () => {
     );
   });
 
+  it.each(["", "Another draft"])("recovers the edit after rollback if sending fails (draft: %j)", async (draft) => {
+    ipc.rollback.mockImplementation(async () => {
+      // The backend emits Truncated before replying, removing the inline editor.
+      useAgents.setState({ entries: new Map([[task.id, []]]) });
+      return { status: "ok", data: null };
+    });
+    let finish!: (value: unknown) => void;
+    ipc.prompt.mockReturnValue(
+      new Promise((resolve) => {
+        finish = resolve;
+      }),
+    );
+    const { user, box, send } = await edit();
+    act(() => updateComposer(task.id, { text: draft }));
+    await user.click(send);
+    expect(box).not.toBeInTheDocument();
+    await act(async () => finish({ status: "error", error: { kind: "internal", message: "Send failed" } }));
+    expect(screen.getByRole("textbox", { name: "Message" })).toHaveValue(
+      draft ? `Edited message\n${draft}` : "Edited message",
+    );
+    expect(screen.getByRole("textbox", { name: "Message" })).toHaveFocus();
+    ipc.prompt.mockResolvedValue({ status: "ok", data: null });
+    await user.click(screen.getByRole("button", { name: "Send" }));
+    expect(ipc.prompt).toHaveBeenCalledTimes(2);
+    expect(ipc.rollback).toHaveBeenCalledTimes(1);
+  });
+
   it.each([{ metaKey: true }, { ctrlKey: true }])("still sends with modified Enter (%j)", async (modifier) => {
     const { box } = await edit();
     fireEvent.keyDown(box, { key: "Enter", ...modifier });
diff --git a/src/features/agents/store.ts b/src/features/agents/store.ts
index d284121f..b0d04ed2 100644
--- a/src/features/agents/store.ts
+++ b/src/features/agents/store.ts
@@ -415,7 +415,15 @@ export async function rerunFrom(task: string, turn: number, text: string): Promi
     await settled(task);
   }
   await call(commands.agentRollback(task, turn));
-  await call(commands.agentPrompt(task, { text, files: [], images: [] }, false));
+  try {
+    await call(commands.agentPrompt(task, { text, files: [], images: [] }, false));
+  } catch (e) {
+    // Rollback removes the inline editor; recover its text in the session composer.
+    const current = composerDraft(task);
+    updateComposer(task, { text: current.text ? `${text}\n${current.text}` : text });
+    focusComposer(task);
+    throw e;
+  }
 }
 
 /** Resolves once the task's turn has ended (or after a while). */