Skip to content

Commit 519403e

Browse files
RomneyDaRomneyDa
authored
Merge pull request #7278 from continuedev/dallin/fix-yolo-mode-apply
fix: handle apply errors in tools (+ fix auto accept diffs flake)
2 parents 99f49b4 + fa9eb17 commit 519403e

File tree

9 files changed

+298
-186
lines changed

9 files changed

+298
-186
lines changed

‎extensions/vscode/src/apply/ApplyManager.ts‎

Lines changed: 21 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -50,45 +50,36 @@ export class ApplyManager {
5050
toolCallId,
5151
});
5252

53-
try {
54-
const hasExistingDocument = !!activeTextEditor.document.getText().trim();
55-
56-
if (hasExistingDocument) {
57-
// Currently `isSearchAndReplace` will always provide a full file rewrite
58-
// as the contents of `text`, so we can just instantly apply
59-
if (isSearchAndReplace) {
60-
const diffLinesGenerator = generateLines(
61-
myersDiff(activeTextEditor.document.getText(), text),
62-
);
63-
64-
await this.verticalDiffManager.streamDiffLines(
65-
diffLinesGenerator,
66-
true,
67-
streamId,
68-
toolCallId,
69-
);
70-
} else {
71-
await this.handleExistingDocument(
72-
activeTextEditor,
73-
text,
74-
streamId,
75-
toolCallId,
76-
);
77-
}
53+
const hasExistingDocument = !!activeTextEditor.document.getText().trim();
54+
if (hasExistingDocument) {
55+
// Currently `isSearchAndReplace` will always provide a full file rewrite
56+
// as the contents of `text`, so we can just instantly apply
57+
if (isSearchAndReplace) {
58+
const diffLinesGenerator = generateLines(
59+
myersDiff(activeTextEditor.document.getText(), text),
60+
);
61+
62+
await this.verticalDiffManager.streamDiffLines(
63+
diffLinesGenerator,
64+
true,
65+
streamId,
66+
toolCallId,
67+
);
7868
} else {
79-
await this.handleEmptyDocument(
69+
await this.handleExistingDocument(
8070
activeTextEditor,
8171
text,
8272
streamId,
8373
toolCallId,
8474
);
8575
}
86-
} finally {
87-
await this.webviewProtocol.request("updateApplyState", {
76+
} else {
77+
await this.handleEmptyDocument(
78+
activeTextEditor,
79+
text,
8880
streamId,
89-
status: "done",
9081
toolCallId,
91-
});
82+
);
9283
}
9384
}
9485

‎gui/src/hooks/ParallelListeners.tsx‎

Lines changed: 5 additions & 91 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,10 @@
11
import { useCallback, useContext, useEffect, useRef } from "react";
22
import { IdeMessengerContext } from "../context/IdeMessenger";
33

4-
import { EDIT_MODE_STREAM_ID } from "core/edit/constants";
54
import { FromCoreProtocol } from "core/protocol";
65
import { useAppDispatch, useAppSelector } from "../redux/hooks";
76
import { setConfigLoading, setConfigResult } from "../redux/slices/configSlice";
8-
import {
9-
setLastNonEditSessionEmpty,
10-
updateEditStateApplyState,
11-
} from "../redux/slices/editState";
7+
import { setLastNonEditSessionEmpty } from "../redux/slices/editState";
128
import { updateIndexingStatus } from "../redux/slices/indexingSlice";
139
import {
1410
initializeProfilePreferences,
@@ -17,27 +13,21 @@ import {
1713
setSelectedProfile,
1814
} from "../redux/slices/profilesSlice";
1915
import {
20-
acceptToolCall,
2116
addContextItemsAtIndex,
2217
setHasReasoningEnabled,
2318
setIsSessionMetadataLoading,
24-
updateApplyState,
2519
} from "../redux/slices/sessionSlice";
2620
import { setTTSActive } from "../redux/slices/uiSlice";
27-
import { exitEdit } from "../redux/thunks/edit";
28-
import { streamResponseAfterToolCall } from "../redux/thunks/streamResponseAfterToolCall";
2921

30-
import { store } from "../redux/store";
3122
import { cancelStream } from "../redux/thunks/cancelStream";
23+
import { handleApplyStateUpdate } from "../redux/thunks/handleApplyStateUpdate";
3224
import { refreshSessionMetadata } from "../redux/thunks/session";
3325
import { updateFileSymbolsFromHistory } from "../redux/thunks/updateFileSymbols";
34-
import { findToolCallById, logToolUsage } from "../redux/util";
3526
import {
3627
setDocumentStylesFromLocalStorage,
3728
setDocumentStylesFromTheme,
3829
} from "../styles/theme";
3930
import { isJetBrains } from "../util";
40-
import { logAgentModeEditOutcome } from "../util/editOutcomeLogger";
4131
import { setLocalStorage } from "../util/localStorage";
4232
import { migrateLocalStorage } from "../util/migrateLocalStorage";
4333
import { useWebviewListener } from "./useWebviewListener";
@@ -51,9 +41,7 @@ function ParallelListeners() {
5141
(store) => store.profiles.selectedProfileId,
5242
);
5343
const hasDoneInitialConfigLoad = useRef(false);
54-
const autoAcceptEditToolDiffs = useAppSelector(
55-
(store) => store.config.config.ui?.autoAcceptEditToolDiffs,
56-
);
44+
5745
// Load symbols for chat on any session change
5846
const sessionId = useAppSelector((state) => state.session.id);
5947

@@ -235,83 +223,9 @@ function ParallelListeners() {
235223
useWebviewListener(
236224
"updateApplyState",
237225
async (state) => {
238-
if (state.streamId === EDIT_MODE_STREAM_ID) {
239-
dispatch(updateEditStateApplyState(state));
240-
241-
if (state.status === "closed") {
242-
const toolCallState = findToolCallById(
243-
store.getState().session.history,
244-
state.toolCallId!,
245-
);
246-
if (toolCallState) {
247-
logToolUsage(toolCallState, true, true, ideMessenger);
248-
}
249-
void dispatch(exitEdit({}));
250-
}
251-
} else {
252-
// chat or agent
253-
dispatch(updateApplyState(state));
254-
255-
// Handle apply status updates - use toolCallId from event payload
256-
if (state.toolCallId) {
257-
if (state.status === "done" && autoAcceptEditToolDiffs) {
258-
ideMessenger.post("acceptDiff", {
259-
streamId: state.streamId,
260-
filepath: state.filepath,
261-
});
262-
}
263-
if (state.status === "closed") {
264-
// Find the tool call to check if it was canceled
265-
const toolCallState = findToolCallById(
266-
store.getState().session.history,
267-
state.toolCallId,
268-
);
269-
270-
if (toolCallState) {
271-
const accepted = toolCallState.status !== "canceled";
272-
273-
logToolUsage(toolCallState, accepted, true, ideMessenger);
274-
275-
// Log edit outcome for Agent Mode
276-
const applyState = store
277-
.getState()
278-
.session.codeBlockApplyStates.states.find(
279-
(s) => s.streamId === state.streamId,
280-
);
281-
282-
if (applyState) {
283-
void logAgentModeEditOutcome(
284-
toolCallState,
285-
applyState,
286-
accepted,
287-
ideMessenger,
288-
);
289-
}
290-
291-
if (accepted) {
292-
dispatch(
293-
acceptToolCall({
294-
toolCallId: state.toolCallId,
295-
}),
296-
);
297-
void dispatch(
298-
streamResponseAfterToolCall({
299-
toolCallId: state.toolCallId,
300-
}),
301-
);
302-
}
303-
}
304-
// const output: ContextItem = {
305-
// name: "Edit tool output",
306-
// content: "Completed edit",
307-
// description: "",
308-
// };
309-
// dispatch(setToolCallOutput([]));
310-
}
311-
}
312-
}
226+
void dispatch(handleApplyStateUpdate(state));
313227
},
314-
[autoAcceptEditToolDiffs, ideMessenger],
228+
[],
315229
);
316230

317231
useEffect(() => {

‎gui/src/pages/gui/chat-tests/EditToolScenarios.test.tsx‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,16 @@
11
import { BuiltInToolNames } from "core/tools/builtIn";
2+
import { generateToolCallButtonTestId } from "../../../components/mainInput/Lump/LumpToolbar/PendingToolCallToolbar";
23
import {
34
addAndSelectMockLlm,
45
triggerConfigUpdate,
56
} from "../../../util/test/config";
67
import { renderWithProviders } from "../../../util/test/render";
78
import { Chat } from "../Chat";
8-
import { generateToolCallButtonTestId } from "../../../components/mainInput/Lump/LumpToolbar/PendingToolCallToolbar";
99

1010
import { waitFor } from "@testing-library/dom";
1111
import { act } from "@testing-library/react";
1212
import { ChatMessage } from "core";
13-
import {
14-
setToolPolicy,
15-
toggleToolSetting,
16-
} from "../../../redux/slices/uiSlice";
13+
import { setToolPolicy } from "../../../redux/slices/uiSlice";
1714
import {
1815
getElementByTestId,
1916
getElementByText,
@@ -63,6 +60,7 @@ test(
6360

6461
ideMessenger.responses["getWorkspaceDirs"] = [EDIT_WORKSPACE_DIR];
6562
const messengerPostSpy = vi.spyOn(ideMessenger, "post");
63+
const messengerRequestSpy = vi.spyOn(ideMessenger, "request");
6664

6765
addAndSelectMockLlm(store, ideMessenger);
6866

@@ -92,7 +90,7 @@ test(
9290

9391
// Tool call, check that applyToFile was called for edit
9492
await waitFor(() => {
95-
expect(messengerPostSpy).toHaveBeenCalledWith("applyToFile", {
93+
expect(messengerRequestSpy).toHaveBeenCalledWith("applyToFile", {
9694
streamId: expect.any(String),
9795
filepath: EDIT_FILE_URI,
9896
text: EDIT_CHANGES,
@@ -101,7 +99,7 @@ test(
10199
});
102100

103101
// Extract stream ID and initiate mock streaming
104-
const streamId = messengerPostSpy.mock.calls.find(
102+
const streamId = messengerRequestSpy.mock.calls.find(
105103
(call) => call[0] === "applyToFile",
106104
)?.[1]?.streamId;
107105
expect(streamId).toBeDefined();
@@ -172,6 +170,7 @@ test("Edit run with no policy and yolo mode", { timeout: 15000 }, async () => {
172170

173171
ideMessenger.responses["getWorkspaceDirs"] = [EDIT_WORKSPACE_DIR];
174172
const messengerPostSpy = vi.spyOn(ideMessenger, "post");
173+
const messengerRequestSpy = vi.spyOn(ideMessenger, "request");
175174

176175
addAndSelectMockLlm(store, ideMessenger);
177176

@@ -219,7 +218,7 @@ test("Edit run with no policy and yolo mode", { timeout: 15000 }, async () => {
219218
);
220219
// Tool call, check that applyToFile was called for edit
221220
await waitFor(() => {
222-
expect(messengerPostSpy).toHaveBeenCalledWith("applyToFile", {
221+
expect(messengerRequestSpy).toHaveBeenCalledWith("applyToFile", {
223222
streamId: expect.any(String),
224223
filepath: EDIT_FILE_URI,
225224
text: EDIT_CHANGES,
@@ -228,7 +227,7 @@ test("Edit run with no policy and yolo mode", { timeout: 15000 }, async () => {
228227
});
229228

230229
// Extract stream ID and initiate mock streaming
231-
const streamId = messengerPostSpy.mock.calls.find(
230+
const streamId = messengerRequestSpy.mock.calls.find(
232231
(call) => call[0] === "applyToFile",
233232
)?.[1]?.streamId;
234233
expect(streamId).toBeDefined();

0 commit comments

Comments
 (0)