fix(agents): restore Codex message editing and cancel repaint

crates/ide-agent/src/native/codex/mod.rs

@@ -199,7 +199,12 @@ fn invalid(message: impl Into<String>) -> ResponseError {
 fn from_codex(e: RpcError) -> ResponseError {
     match e {
         RpcError::Server { code, message, data } => {
-            let code = if message.to_lowercase().contains("auth") { AUTH_REQUIRED } else { code };
+            // Protocol errors can list method names containing "auth"; they are not sign-in failures.
+            let protocol = matches!(
+                code,
+                codes::PARSE_ERROR | codes::INVALID_REQUEST | codes::METHOD_NOT_FOUND | codes::INVALID_PARAMS
+            );
+            let code = if !protocol && message.to_lowercase().contains("auth") { AUTH_REQUIRED } else { code };
             ResponseError { code, message, data }
         }
         RpcError::Closed => error(codes::INTERNAL_ERROR, "Codex exited"),
@@ -426,10 +431,48 @@ impl CodexShim {
     async fn rollback(&self, params: &Value) -> Result<Value, ResponseError> {
         let session = str_param(params, "sessionId")?;
         let turns = params.get("turns").and_then(Value::as_u64).filter(|n| *n > 0).ok_or_else(|| invalid("turns"))?;
-        self.call("thread/rollback", json!({ "threadId": session, "numTurns": turns })).await?;
+        match self.call("thread/rollback", json!({ "threadId": session, "numTurns": turns })).await {
+            Ok(_) => {}
+            Err(e)
+                if e.code == codes::METHOD_NOT_FOUND
+                    || (e.code == codes::INVALID_REQUEST
+                        && e.message.contains("unknown variant `thread/rollback`")) =>
+            {
+                self.revert(session, turns).await?;
+            }
+            Err(e) => return Err(e),
+        }
         Ok(json!({}))
     }
 
+    /// New Codex versions replace rollback counts with an exclusive turn id.
+    async fn revert(&self, session: &str, mut remaining: u64) -> Result<(), ResponseError> {
+        let mut cursor = Value::Null;
+        let mut seen = HashSet::new();
+        loop {
+            let page = self
+                .call(
+                    "thread/turns/list",
+                    json!({ "threadId": session, "sortDirection": "desc", "limit": remaining.min(100), "cursor": cursor }),
+                )
+                .await?;
+            let turns = page.get("data").and_then(Value::as_array).ok_or_else(|| invalid("Codex returned no turns"))?;
+            if remaining <= turns.len() as u64 {
+                let before = str_param(&turns[remaining as usize - 1], "id")?;
+                self.call("thread/revert", json!({ "threadId": session, "beforeTurnId": before })).await?;
+                return Ok(());
+            }
+            remaining -= turns.len() as u64;
+            let next = page.get("nextCursor").and_then(Value::as_str);
+            let Some(next) = next.filter(|next| seen.insert((*next).to_owned())) else {
+                return Err(invalid(
+                    "Codex history has fewer turns than the conversation; reload the task before editing",
+                ));
+            };
+            cursor = json!(next);
+        }
+    }
+
     /// Stops following the session's thread (codex unloads it once nothing follows it).
     async fn close(&self, params: &Value) -> Result<Value, ResponseError> {
         let session = str_param(params, "sessionId")?;

crates/ide-agent/src/native/codex/tests.rs

@@ -593,6 +593,104 @@ async fn resume_replays_history_then_forks_and_rolls_back() {
     assert_eq!(sent(&frames, "thread/unsubscribe")[0]["params"], json!({"threadId": thread}));
 }
 
+async fn rollback_replies(turns: u32, replies: Vec<(&str, Value)>) -> (Result<Value, RpcError>, Vec<Value>) {
+    let frames = replies
+        .into_iter()
+        .enumerate()
+        .flat_map(|(i, (method, mut reply))| {
+            let id = i + 1;
+            reply["id"] = json!(id);
+            [(true, json!({ "id": id, "method": method })), (false, reply)]
+        })
+        .collect();
+    let transcript = Transcript { meta: Value::Null, frames };
+    let (ours, theirs) = tokio::io::duplex(1 << 16);
+    let acp = attach(Arc::new(Client::default()), |events| {
+        let (read, write) = tokio::io::split(ours);
+        Ok(Peer::connect_framed(read, write, events, Framing::Lines))
+    })
+    .unwrap();
+    let (read, write) = tokio::io::split(theirs);
+    let player = tokio::spawn(async move { play(&transcript, read, write).await });
+    let result = request(&acp, "_openidea/rollback", json!({ "sessionId": "s", "turns": turns })).await;
+    let (sent, _read, _write) = player.await.unwrap();
+    acp.kill();
+    (result, sent)
+}
+
+fn removed_rollback() -> Value {
+    json!({ "error": { "code": -32600, "message": "Invalid request: unknown variant `thread/rollback`, expected one of `thread/revert`, `account/gatewayOAuth/read`, `getAuthStatus`" } })
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn removed_rollback_reverts_before_the_target_turn_across_pages() {
+    let (result, frames) = rollback_replies(
+        3,
+        vec![
+            ("thread/rollback", removed_rollback()),
+            (
+                "thread/turns/list",
+                json!({ "result": { "data": [{ "id": "last" }, { "id": "middle" }], "nextCursor": "older" } }),
+            ),
+            (
+                "thread/turns/list",
+                json!({ "result": { "data": [{ "id": "edited" }, { "id": "retained" }], "nextCursor": null } }),
+            ),
+            ("thread/revert", json!({ "result": { "thread": { "id": "s", "turns": [] } } })),
+        ],
+    )
+    .await;
+    assert_eq!(result.unwrap(), json!({}));
+    let pages = sent(&frames, "thread/turns/list");
+    assert_eq!(pages[0]["params"], json!({ "threadId": "s", "sortDirection": "desc", "limit": 3, "cursor": null }));
+    assert_eq!(pages[1]["params"]["cursor"], "older");
+    assert_eq!(pages[1]["params"]["limit"], 1);
+    assert_eq!(sent(&frames, "thread/revert")[0]["params"], json!({ "threadId": "s", "beforeTurnId": "edited" }));
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn method_not_found_reverts_the_latest_turn() {
+    let (result, frames) = rollback_replies(
+        1,
+        vec![
+            ("thread/rollback", json!({ "error": { "code": -32601, "message": "Method not found" } })),
+            ("thread/turns/list", json!({ "result": { "data": [{ "id": "last" }], "nextCursor": null } })),
+            ("thread/revert", json!({ "result": {} })),
+        ],
+    )
+    .await;
+    assert_eq!(result.unwrap(), json!({}));
+    assert_eq!(sent(&frames, "thread/revert")[0]["params"]["beforeTurnId"], "last");
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn rollback_refusals_are_not_retried_as_revert() {
+    let (result, frames) = rollback_replies(
+        1,
+        vec![("thread/rollback", json!({ "error": { "code": -32600, "message": "a turn is still running" } }))],
+    )
+    .await;
+    assert!(result.unwrap_err().to_string().contains("a turn is still running"));
+    assert_eq!(frames.len(), 1);
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn incomplete_or_repeated_history_pages_do_not_revert_the_wrong_turn() {
+    for next in [Value::Null, json!("older")] {
+        let (result, frames) = rollback_replies(
+            3,
+            vec![
+                ("thread/rollback", removed_rollback()),
+                ("thread/turns/list", json!({ "result": { "data": [{ "id": "last" }], "nextCursor": "older" } })),
+                ("thread/turns/list", json!({ "result": { "data": [], "nextCursor": next } })),
+            ],
+        )
+        .await;
+        assert!(result.unwrap_err().to_string().contains("reload the task"));
+        assert!(sent(&frames, "thread/revert").is_empty());
+    }
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn steering_adds_to_the_running_turn() {
     let (results, _client, frames, _) = replay(Scenario::Steer).await;

src/features/agents/Session.module.css

@@ -117,8 +117,8 @@
 .stream:focus {
   outline: none;
 }
-/* Off-screen entries skip layout, so resizing the pane doesn't reflow a long conversation. */
-.stream > * {
+/* Skip off-screen output layout; editable messages must repaint immediately when focus/content changes. */
+.stream > :not(.message) {
   content-visibility: auto;
   contain-intrinsic-block-size: auto 80px;
 }

src/features/agents/SessionEditor.tsx

@@ -23,6 +23,7 @@ import {
   showContextMenu,
 } from "@/ui";
 import { useTx, tx } from "@/platform/i18n";
+import { isImeKey } from "@/platform/ime";
 import { TaskStatusIcon } from "./AgentsView";
 import { Composer } from "./Composer";
 import { canSend, emptyDraft, retryParts, usageLabel, workTime } from "./composerModel";
@@ -697,9 +698,9 @@ function UserMessage({
         initial={entry.text}
         later={turn < turnCount - 1}
         onCancel={() => setEditing(false)}
-        onSend={(text) => {
+        onSend={async (text) => {
+          await rerunFrom(task.id, turn, text);
           setEditing(false);
-          void rerunFrom(task.id, turn, text).catch(fail);
         }}
       />
     );
@@ -770,12 +771,27 @@ function MessageEditor({
   /** Turns follow it (they are dropped on send). */
   later: boolean;
   onCancel: () => void;
-  onSend: (text: string) => void;
+  onSend: (text: string) => Promise<void>;
 }) {
   const tx = useTx();
   const [text, setText] = useState(initial);
-  const send = () => {
-    if (text.trim()) onSend(text.trim());
+  const [sending, setSending] = useState(false);
+  const [error, setError] = useState<string | null>(null);
+  const pending = useRef(false);
+  const send = async () => {
+    if (!text.trim() || pending.current) return;
+    pending.current = true;
+    setSending(true);
+    setError(null);
+    try {
+      await onSend(text.trim());
+    } catch (e) {
+      setError(e instanceof Error ? e.message : String(e));
+      notifications().notify({ level: "error", title: "Couldn't send the message", body: String(e) });
+    } finally {
+      pending.current = false;
+      setSending(false);
+    }
   };
   return (
     <div className={styles.message}>
@@ -784,19 +800,30 @@ function MessageEditor({
         value={text}
         rows={Math.min(Math.max(text.split("\n").length, 2), 12)}
         autoFocus
+        readOnly={sending}
         aria-label={tx("Message")}
         onChange={(e) => setText(e.target.value)}
         onFocus={(e) => e.currentTarget.setSelectionRange(text.length, text.length)}
         onKeyDown={(e) => {
+          if (isImeKey(e.nativeEvent)) return;
           if (e.key === "Escape") {
             e.preventDefault();
-            onCancel();
+            e.stopPropagation();
+            if (!pending.current) onCancel();
           } else if (e.key === "Enter" && (e.metaKey || e.ctrlKey)) {
             e.preventDefault();
-            send();
+            void send();
           }
         }}
       />
+      {error && (
+        <div className={styles.error} role="alert">
+          <Icon name="error" size={14} />
+          <span>
+            <T>Couldn't send the message</T>: {error}
+          </span>
+        </div>
+      )}
       <div className={styles.editBar}>
         {later && (
           <span className={styles.muted}>
@@ -804,11 +831,12 @@ function MessageEditor({
           </span>
         )}
         <span className={styles.editButtons}>
-          <Button size="small" onClick={onCancel}>
+          <Button size="small" disabled={sending} onClick={onCancel}>
             Cancel
           </Button>
-          <Button size="small" variant="primary" disabled={!text.trim()} onClick={send}>
-            Send
+          <Button size="small" variant="primary" disabled={!text.trim() || sending} onClick={() => void send()}>
+            {sending && <Spinner size={12} />}
+            <T>Send</T>
           </Button>
         </span>
       </div>

src/features/agents/messageEditing.test.tsx

@@ -0,0 +1,144 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { act, fireEvent, render, screen, waitFor } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import type { AgentEntry, AgentTask } from "@/bindings";
+
+const ipc = vi.hoisted(() => ({ rollback: vi.fn(), prompt: vi.fn(), cancel: vi.fn() }));
+const entries: AgentEntry[] = [
+  { type: "user", text: "Original message", attachments: [], sentMs: 1, doneMs: 2, steered: false },
+  { type: "agent", text: "Original answer" },
+];
+
+vi.mock("@/platform/ipc", async (original) => ({
+  ...(await original<typeof import("@/platform/ipc")>()),
+  commands: {
+    agentEntries: async () => ({ status: "ok", data: entries }),
+    agentChanges: async () => ({ status: "ok", data: [] }),
+    vcsBranchInfo: async () => ({ status: "ok", data: { head: "main" } }),
+    agentRollback: ipc.rollback,
+    agentPrompt: ipc.prompt,
+    agentCancel: ipc.cancel,
+  },
+}));
+vi.mock("@tauri-apps/api/webview", () => ({
+  getCurrentWebview: () => ({ onDragDropEvent: () => Promise.resolve(() => {}) }),
+}));
+
+import { SessionEditor } from "./SessionEditor";
+import { useAgents } from "./store";
+
+const task: AgentTask = {
+  id: "editing",
+  title: "Test task",
+  provider: "codex",
+  status: "ready",
+  isolation: "local",
+  model: null,
+  cwd: "/project",
+  branch: null,
+  base: null,
+  baseCommit: null,
+  error: null,
+  createdMs: 0,
+  updatedMs: 0,
+  config: [],
+  commands: [],
+  usage: null,
+  queue: [],
+  capabilities: { fork: true, rollback: true, steer: true, images: true },
+};
+
+beforeEach(() => {
+  ipc.rollback.mockReset().mockResolvedValue({ status: "ok", data: null });
+  ipc.prompt.mockReset().mockResolvedValue({ status: "ok", data: null });
+  ipc.cancel.mockReset().mockResolvedValue({ status: "ok", data: null });
+  useAgents.setState({ tasks: new Map([[task.id, task]]), entries: new Map(), composers: new Map(), focus: null });
+});
+
+async function edit() {
+  render(<SessionEditor input={{ kind: "agent-session", key: task.id, data: { task: task.id } }} />);
+  const user = userEvent.setup();
+  await user.click(await screen.findByRole("button", { name: "Edit" }));
+  const box = screen.getAllByRole("textbox", { name: "Message" })[0]!;
+  fireEvent.change(box, { target: { value: "Edited message" } });
+  return { user, box, send: screen.getAllByRole("button", { name: "Send" })[0]! };
+}
+
+describe("editing a sent agent message", () => {
+  it("keeps Enter, Shift+Enter and IME confirmation for editing", async () => {
+    const { box } = await edit();
+    fireEvent.keyDown(box, { key: "Enter" });
+    fireEvent.keyDown(box, { key: "Enter", shiftKey: true });
+    fireEvent.keyDown(box, { key: "Enter", metaKey: true, isComposing: true });
+    fireEvent.keyDown(box, { key: "Enter", metaKey: true, keyCode: 229 });
+    expect(ipc.rollback).not.toHaveBeenCalled();
+    expect(box).toHaveValue("Edited message");
+  });
+
+  it("keeps the draft and shows a rollback failure beside it", async () => {
+    ipc.rollback.mockResolvedValue({ status: "error", error: { kind: "internal", message: "Rollback failed" } });
+    const { user, box, send } = await edit();
+    await user.click(send);
+    expect(await screen.findByRole("alert")).toHaveTextContent("Rollback failed");
+    expect(box).toHaveValue("Edited message");
+    expect(box).toBeInTheDocument();
+    expect(send).toBeEnabled();
+    expect(ipc.prompt).not.toHaveBeenCalled();
+    ipc.rollback.mockResolvedValue({ status: "ok", data: null });
+    await user.click(send);
+    await waitFor(() => expect(ipc.prompt).toHaveBeenCalledTimes(1));
+    expect(screen.queryByRole("alert")).not.toBeInTheDocument();
+  });
+
+  it("rolls back before sending the edited text when Send is clicked", async () => {
+    let finish!: (value: unknown) => void;
+    ipc.rollback.mockReturnValue(
+      new Promise((resolve) => {
+        finish = resolve;
+      }),
+    );
+    const { user, box, send } = await edit();
+    await user.click(send);
+    expect(ipc.rollback).toHaveBeenCalledWith(task.id, 0);
+    expect(ipc.prompt).not.toHaveBeenCalled();
+    expect(send).toBeDisabled();
+    fireEvent.keyDown(box, { key: "Enter", metaKey: true });
+    expect(ipc.rollback).toHaveBeenCalledTimes(1);
+    await act(async () => finish({ status: "ok", data: null }));
+    await waitFor(() =>
+      expect(ipc.prompt).toHaveBeenCalledWith(task.id, { text: "Edited message", files: [], images: [] }, false),
+    );
+  });
+
+  it.each([{ metaKey: true }, { ctrlKey: true }])("still sends with modified Enter (%j)", async (modifier) => {
+    const { box } = await edit();
+    fireEvent.keyDown(box, { key: "Enter", ...modifier });
+    await waitFor(() => expect(ipc.prompt).toHaveBeenCalledTimes(1));
+  });
+
+  it("waits for a running turn to stop before rolling it back", async () => {
+    useAgents.setState({ tasks: new Map([[task.id, { ...task, status: "running" }]]) });
+    const { user, send } = await edit();
+    await user.click(send);
+    expect(ipc.cancel).toHaveBeenCalledWith(task.id, null);
+    expect(ipc.rollback).not.toHaveBeenCalled();
+    act(() => useAgents.setState({ tasks: new Map([[task.id, task]]) }));
+    await waitFor(() => expect(ipc.prompt).toHaveBeenCalledTimes(1));
+  });
+
+  it.each(["Cancel", "Escape"])(
+    "%s restores the message and discards edits without touching the conversation",
+    async (method) => {
+      const { user, box } = await edit();
+      if (method === "Cancel") await user.click(screen.getByRole("button", { name: "Cancel" }));
+      else fireEvent.keyDown(box, { key: "Escape" });
+      expect(box).not.toBeInTheDocument();
+      expect(screen.getByText("Original message")).toBeInTheDocument();
+      expect(screen.getByText("Original answer")).toBeInTheDocument();
+      expect(ipc.rollback).not.toHaveBeenCalled();
+      expect(ipc.prompt).not.toHaveBeenCalled();
+      await user.click(screen.getByRole("button", { name: "Edit" }));
+      expect(screen.getAllByRole("textbox", { name: "Message" })[0]).toHaveValue("Original message");
+    },
+  );
+});