From 3c87a8eea72fa60e54dc8dd4ab7f065bfa65169a Mon Sep 17 00:00:00 2001 From: l0ng-ai <24760907+l0ng-ai@users.noreply.github.com> Date: Thu, 29 Oct 2026 21:09:19 +0800 Subject: [PATCH] fix(agents): restore Codex message editing and cancel repaint --- crates/ide-agent/src/native/codex/mod.rs | 47 ++++++- crates/ide-agent/src/native/codex/tests.rs | 98 +++++++++++++ src/features/agents/Session.module.css | 4 +- src/features/agents/SessionEditor.tsx | 48 +++++-- src/features/agents/messageEditing.test.tsx | 144 ++++++++++++++++++++ 5 files changed, 327 insertions(+), 14 deletions(-) create mode 100644 src/features/agents/messageEditing.test.tsx diff --git a/crates/ide-agent/src/native/codex/mod.rs b/crates/ide-agent/src/native/codex/mod.rs index a83fd459..7f040e4a 100644 --- a/crates/ide-agent/src/native/codex/mod.rs +++ b/crates/ide-agent/src/native/codex/mod.rs @@ -199,7 +199,12 @@ fn invalid(message: impl Into) -> 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 { 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 { let session = str_param(params, "sessionId")?; diff --git a/crates/ide-agent/src/native/codex/tests.rs b/crates/ide-agent/src/native/codex/tests.rs index f98f6921..0cc94cf9 100644 --- a/crates/ide-agent/src/native/codex/tests.rs +++ b/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, Vec) { + 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; diff --git a/src/features/agents/Session.module.css b/src/features/agents/Session.module.css index 420f835b..512462dd 100644 --- a/src/features/agents/Session.module.css +++ b/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; } diff --git a/src/features/agents/SessionEditor.tsx b/src/features/agents/SessionEditor.tsx index e60faa8f..18ea0ab9 100644 --- a/src/features/agents/SessionEditor.tsx +++ b/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; }) { 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(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 (
@@ -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 && ( +
+ + + Couldn't send the message: {error} + +
+ )}
{later && ( @@ -804,11 +831,12 @@ function MessageEditor({ )} - -
diff --git a/src/features/agents/messageEditing.test.tsx b/src/features/agents/messageEditing.test.tsx new file mode 100644 index 00000000..6d98ffad --- /dev/null +++ b/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()), + 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(); + 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"); + }, + ); +});