Skip to content

Commit 7632d5c

Browse files
wuzixuan06EtienneLescot
wuzixuan06
authored and
EtienneLescot
committed
fix: preserve one undo step per agent turn
1 parent 6d1f6bc commit 7632d5c

5 files changed

Lines changed: 133 additions & 205 deletions

File tree

‎electron/ai-edition/chat-service.ts‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -523,9 +523,8 @@ export async function runChat(
523523
success: true,
524524
assistantMessage,
525525
cancelled: result.cancelled,
526-
// Non-streaming callers still receive the completed document. Both this
527-
// fallback and the document event above enforce the edit setting; the
528-
// panel never reapplies this aggregate after individual streamed edits.
526+
// Return the accumulated edits even on Stop. The panel applies this once
527+
// at the end of the turn, preserving one save and one undo step per turn.
529528
document: result.mutated && editsAllowed ? result.document : undefined,
530529
toolCalls: appliedToolCalls.length ? appliedToolCalls : undefined,
531530
userMessageCheckpointId: userMessage.checkpointId ?? undefined,

‎src/components/ai-edition/LeftPanel.stop.test.tsx‎

Lines changed: 90 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import { TooltipProvider } from "@/components/ui/tooltip";
66
import { EditorDialogsProvider } from "@/contexts/EditorDialogsContext";
77
import { createEmptyDocument } from "@/lib/ai-edition/schema";
88
import { useProjectStore } from "@/lib/ai-edition/store/projectStore";
9-
import { clearHistory, undo } from "@/lib/ai-edition/store/undo";
9+
import { clearHistory, redo, undo } from "@/lib/ai-edition/store/undo";
1010
import { past } from "@/lib/ai-edition/store/undoStack";
1111
import type { AiEditionChatEvent, AiEditionChatResult } from "@/native/contracts";
1212
import { ChatStripPanel } from "./LeftPanel";
@@ -19,6 +19,7 @@ const api = vi.hoisted(() => ({
1919
select: vi.fn(),
2020
create: vi.fn(),
2121
error: vi.fn(),
22+
warning: vi.fn(),
2223
}));
2324
vi.mock("@/native/client", () => ({
2425
nativeBridgeClient: {
@@ -41,7 +42,7 @@ vi.mock("@/contexts/I18nContext", () => ({
4142
useI18n: () => ({ locale: "en" }),
4243
useScopedT: () => (key: string) => key,
4344
}));
44-
vi.mock("sonner", () => ({ toast: { error: api.error, warning: vi.fn() } }));
45+
vi.mock("sonner", () => ({ toast: { error: api.error, warning: api.warning } }));
4546
let onEvent: (event: AiEditionChatEvent) => void;
4647
let replies: ((value: AiEditionChatResult) => void)[];
4748
const doc = (title: string, projectId = "p") => createEmptyDocument({ projectId, title });
@@ -125,27 +126,39 @@ function emit(
125126
}
126127

127128
describe("composer Stop", () => {
128-
it("keeps edits and undo steps, retains partial reply, ignores old events and starts the next turn", async () => {
129+
it("commits stopped edits once, retains the live reply, ignores old events and starts the next turn", async () => {
129130
panel();
130131
const id = await start();
131132
expect(screen.getByRole("button", { name: "chat.stop" })).toBeEnabled();
132133
emit(id, { kind: "document", document: doc("First") });
133134
emit(id, { kind: "document", document: doc("Second") });
134135
emit(id, { kind: "text", delta: "Partial" });
135-
await waitFor(() => expect(past).toHaveLength(2));
136+
await act(async () => {
137+
// Flush event work before checking that no intermediate edits were saved.
138+
});
139+
expect(screen.getByText("Partial")).toBeVisible();
140+
expect(api.save).not.toHaveBeenCalled();
141+
expect(past).toHaveLength(0);
142+
expect(useProjectStore.getState().document?.project.title).toBe("Before");
136143
fireEvent.click(screen.getByRole("button", { name: "chat.stop" }));
137144
fireEvent.click(screen.getByRole("button", { name: "chat.stop" }));
138145
expect(api.cancel).toHaveBeenCalledOnce();
139146
expect(api.cancel).toHaveBeenCalledWith("p", "s", id);
140-
await act(async () => replies[0](result("Partial", true)));
147+
await act(async () => replies[0]({ ...result("Partial", true), document: doc("Second") }));
141148
expect(screen.getByText("chat.stopped")).toBeVisible();
142149
expect(screen.getByText("Partial")).toBeVisible();
143150
expect(screen.getByRole("button", { name: "chat.send" })).toBeInTheDocument();
144-
expect(api.save).toHaveBeenCalledTimes(2);
151+
expect(api.save).toHaveBeenCalledOnce();
152+
expect(past).toHaveLength(1);
153+
expect(useProjectStore.getState().document?.project.title).toBe("Second");
145154
act(() => {
146155
undo();
147156
});
148-
expect(useProjectStore.getState().document?.project.title).toBe("First");
157+
expect(useProjectStore.getState().document?.project.title).toBe("Before");
158+
act(() => {
159+
redo();
160+
});
161+
expect(useProjectStore.getState().document?.project.title).toBe("Second");
149162
const nextId = await start();
150163
expect(nextId).not.toBe(id);
151164
emit(id, { kind: "text", delta: "STALE" });
@@ -157,6 +170,70 @@ describe("composer Stop", () => {
157170
expect(api.error).not.toHaveBeenCalled();
158171
});
159172

173+
it.each([false, true])("keeps a long turn in one undo step (stopped: %s)", async (stopped) => {
174+
useProjectStore.getState().setDocument(doc("Pre-turn"), { history: true });
175+
panel();
176+
const id = await start();
177+
for (let i = 0; i < 60; i++) {
178+
emit(id, { kind: "document", document: doc(`Tool ${i}`) });
179+
}
180+
await act(async () => {
181+
// Let any queued per-tool writes run so the regression cannot hide in microtasks.
182+
});
183+
expect(api.save).not.toHaveBeenCalled();
184+
expect(past).toHaveLength(1);
185+
if (stopped) fireEvent.click(screen.getByRole("button", { name: "chat.stop" }));
186+
await act(async () =>
187+
replies[0]({ ...result("Done", stopped), document: doc("Final aggregate") }),
188+
);
189+
expect(api.save).toHaveBeenCalledOnce();
190+
expect(past).toHaveLength(2);
191+
expect(useProjectStore.getState().document?.project.title).toBe("Final aggregate");
192+
act(() => {
193+
undo();
194+
});
195+
expect(useProjectStore.getState().document?.project.title).toBe("Pre-turn");
196+
act(() => {
197+
undo();
198+
});
199+
expect(useProjectStore.getState().document?.project.title).toBe("Before");
200+
});
201+
202+
it("applies a stopped aggregate when no document event arrived", async () => {
203+
panel();
204+
await start();
205+
fireEvent.click(screen.getByRole("button", { name: "chat.stop" }));
206+
await act(async () =>
207+
replies[0]({ ...result("Partial", true), document: doc("Completed edit") }),
208+
);
209+
expect(api.save).toHaveBeenCalledOnce();
210+
expect(past).toHaveLength(1);
211+
expect(useProjectStore.getState().document?.project.title).toBe("Completed edit");
212+
act(() => {
213+
undo();
214+
});
215+
expect(useProjectStore.getState().document?.project.title).toBe("Before");
216+
});
217+
218+
it("keeps manual edits during a turn until Apply anyway commits the aggregate once", async () => {
219+
panel();
220+
const id = await start();
221+
emit(id, { kind: "document", document: doc("First") });
222+
act(() => useProjectStore.getState().setDocument(doc("Manual"), { history: true }));
223+
emit(id, { kind: "document", document: doc("Second") });
224+
await act(async () => replies[0]({ ...result("Done"), document: doc("Second") }));
225+
expect(api.save).not.toHaveBeenCalled();
226+
expect(useProjectStore.getState().document?.project.title).toBe("Manual");
227+
expect(api.warning).toHaveBeenCalledWith("chat.agentEditConflict", expect.any(Object));
228+
await act(async () => api.warning.mock.calls[0][1].action.onClick());
229+
await waitFor(() => expect(api.save).toHaveBeenCalledOnce());
230+
expect(useProjectStore.getState().document?.project.title).toBe("Second");
231+
act(() => {
232+
undo();
233+
});
234+
expect(useProjectStore.getState().document?.project.title).toBe("Manual");
235+
});
236+
160237
it("stops while creating the first session without sending a model request", async () => {
161238
api.list.mockResolvedValue([]);
162239
let created!: (value: ReturnType<typeof session>) => void;
@@ -184,11 +261,12 @@ describe("composer Stop", () => {
184261
act(() => useProjectStore.setState({ projectId: "other", document: doc("Other", "other") }));
185262
await waitFor(() => expect(api.cancel).toHaveBeenCalledWith("p", "s", id));
186263
await start();
187-
await act(async () => replies[0](result("Old reply", true)));
264+
await act(async () => replies[0]({ ...result("Old reply", true), document: doc("Old") }));
188265
emit(id, { kind: "document", document: doc("Old") });
189266
expect(screen.queryByText("Old reply")).toBeNull();
190267
expect(screen.getByRole("button", { name: "chat.stop" })).toBeEnabled();
191268
expect(useProjectStore.getState().document?.project.title).toBe("Other");
269+
expect(api.save).not.toHaveBeenCalled();
192270
await act(async () => replies[1](result("New reply")));
193271
});
194272

@@ -197,14 +275,16 @@ describe("composer Stop", () => {
197275
const id = await start();
198276
fireEvent.click(screen.getByRole("button", { name: "chat.newConversation" }));
199277
await waitFor(() => expect(api.cancel).toHaveBeenCalledWith("p", "s", id));
200-
await act(async () => replies[0](result("Old", true)));
278+
await act(async () => replies[0]({ ...result("Old", true), document: doc("Old") }));
201279
expect(screen.queryByText("Old")).toBeNull();
280+
expect(api.save).not.toHaveBeenCalled();
202281
fireEvent.change(screen.getByRole("textbox"), { target: { value: "Next" } });
203282
fireEvent.keyDown(screen.getByRole("textbox"), { key: "Enter" });
204283
await waitFor(() => expect(api.run).toHaveBeenCalledTimes(2));
205284
const nextId = api.run.mock.calls[1][4];
206285
view.unmount();
207286
expect(api.cancel).toHaveBeenLastCalledWith("p", "new", nextId);
208-
await act(async () => replies[1](result("Disposed", true)));
287+
await act(async () => replies[1]({ ...result("Disposed", true), document: doc("Disposed") }));
288+
expect(api.save).not.toHaveBeenCalled();
209289
});
210290
});

‎src/components/ai-edition/LeftPanel.tsx‎

Lines changed: 41 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,10 @@ import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover
66
import { Tooltip } from "@/components/ui/tooltip";
77
import { useEditorDialogActions, useEditorDialogSection } from "@/contexts/EditorDialogsContext";
88
import { useScopedT } from "@/contexts/I18nContext";
9-
import { applyAgentDocumentIfCurrent } from "@/lib/ai-edition/store/agentDocumentApply";
10-
import { AgentDocumentStream } from "@/lib/ai-edition/store/agentDocumentStream";
9+
import {
10+
applyAgentDocumentIfCurrent,
11+
runAgentTurn,
12+
} from "@/lib/ai-edition/store/agentDocumentApply";
1113
import { useProjectStore } from "@/lib/ai-edition/store/projectStore";
1214
import { useChatPromptBus } from "@/lib/ai-edition/store/useChatPromptBus";
1315
import { nativeBridgeClient } from "@/native/client";
@@ -446,7 +448,6 @@ export function ChatStripPanel() {
446448
sessionId: string | null;
447449
stopping: boolean;
448450
thinking: string;
449-
stream: AgentDocumentStream;
450451
} | null>(null);
451452
const viewRef = useRef({ projectId, sessionId: null as string | null });
452453
const [llmConfig, setLlmConfig] = useState<AiEditionLlmConfig | null>(null);
@@ -535,7 +536,6 @@ export function ChatStripPanel() {
535536
const run = runRef.current;
536537
if (!run) return;
537538
runRef.current = null;
538-
run.stream.close();
539539
if (run.sessionId) {
540540
void nativeBridgeClient.aiEdition
541541
.chatCancel(run.projectId, run.sessionId, run.id)
@@ -558,8 +558,7 @@ export function ChatStripPanel() {
558558
viewRef.current.sessionId !== run.sessionId
559559
)
560560
return;
561-
if (event.kind === "document") run.stream.push(event.document);
562-
else if (event.kind === "thinking") {
561+
if (event.kind === "thinking") {
563562
run.thinking += event.delta;
564563
setThinkingText(run.thinking);
565564
} else if (event.kind === "text") setLiveText((prev) => prev + event.delta);
@@ -673,7 +672,6 @@ export function ChatStripPanel() {
673672
sessionId: activeSessionId,
674673
stopping: false,
675674
thinking: "",
676-
stream: new AgentDocumentStream(),
677675
};
678676
runRef.current = run;
679677
latestRunIdRef.current = run.id;
@@ -726,41 +724,45 @@ export function ChatStripPanel() {
726724
}
727725
if (!isCurrent() || run.stopping) return;
728726
run.sessionId = sessionId;
729-
run.stream = new AgentDocumentStream();
730-
const documentSnapshot = useProjectStore.getState().document ?? undefined;
731-
const result = await nativeBridgeClient.aiEdition.chatRun(
732-
projectId,
733-
sessionId,
734-
text,
735-
documentSnapshot,
736-
run.id,
727+
// Capture the document and its revision together before awaiting the turn.
728+
// Stop returns the completed tools' aggregate too: commit it once so the
729+
// entire turn has one save and one undo step, regardless of tool count.
730+
const { result, applyDocument } = await runAgentTurn((documentSnapshot) =>
731+
nativeBridgeClient.aiEdition.chatRun(projectId, sessionId, text, documentSnapshot, run.id),
737732
);
738733
if (!isCurrent()) return;
739-
// Non-streaming callers keep their existing result path. Completed tool
740-
// edits already in history are never saved again as one aggregate.
741-
if (!run.stream.seen && result.document && !result.cancelled)
742-
run.stream.push(result.document);
743-
await run.stream.finish();
744-
if (!isCurrent()) return;
745-
if (run.stream.failure === "conflict")
746-
toast.warning(t("chat.agentEditConflict"), {
747-
duration: Number.POSITIVE_INFINITY,
748-
action: {
749-
label: t("chat.applyAnyway"),
750-
onClick: () => {
751-
if (
752-
latestRunIdRef.current !== run.id ||
753-
viewRef.current.projectId !== projectId ||
754-
viewRef.current.sessionId !== sessionId
755-
)
756-
return;
757-
void run.stream.applyRemaining().then(() => {
758-
if (run.stream.failure === "save-failed") toast.error(t("chat.applyEditsFailed"));
759-
});
734+
if (result.document) {
735+
const applyEdits = async (options?: { ignoreConflict?: boolean }) => {
736+
try {
737+
const applied = await applyDocument(options);
738+
if (applied === "save-failed") toast.error(t("chat.applyEditsFailed"));
739+
return applied;
740+
} catch (error) {
741+
toast.error(t("chat.applyEditsFailed"), {
742+
description: error instanceof Error ? error.message : String(error),
743+
});
744+
return "malformed" as const;
745+
}
746+
};
747+
const applied = await applyEdits();
748+
if (!isCurrent()) return;
749+
if (applied === "conflict")
750+
toast.warning(t("chat.agentEditConflict"), {
751+
duration: Number.POSITIVE_INFINITY,
752+
action: {
753+
label: t("chat.applyAnyway"),
754+
onClick: () => {
755+
if (
756+
latestRunIdRef.current !== run.id ||
757+
viewRef.current.projectId !== projectId ||
758+
viewRef.current.sessionId !== sessionId
759+
)
760+
return;
761+
void applyEdits({ ignoreConflict: true });
762+
},
760763
},
761-
},
762-
});
763-
else if (run.stream.failure === "save-failed") toast.error(t("chat.applyEditsFailed"));
764+
});
765+
}
764766
setMessages((prev) => adoptStoredUserMessage(prev, optimisticUserId, result));
765767
// A rewind confirmation opened during the turn names the message too. Stored without a
766768
// checkpoint, the message has nothing to rewind to and the confirmation closes.
@@ -801,7 +803,6 @@ export function ChatStripPanel() {
801803
description: err instanceof Error ? err.message : String(err),
802804
});
803805
} finally {
804-
run.stream.close();
805806
if (isCurrent()) {
806807
runRef.current = null;
807808
setBusy(false);

0 commit comments

Comments
 (0)