From a320323aa1b8b2c0615a9d47abc3545d493930c0 Mon Sep 17 00:00:00 2001 From: Andrew Pareles Date: Mon, 17 Mar 2025 18:39:14 -0700 Subject: [PATCH] promise resolves when abort to free space --- .../void/browser/autocompleteService.ts | 1 + .../contrib/void/browser/chatThreadService.ts | 26 ++++++++++++--- .../contrib/void/browser/editCodeService.ts | 32 +++++++++++++++---- .../src/markdown/ApplyBlockHoverButtons.tsx | 4 ++- .../contrib/void/browser/toolsService.ts | 7 ++-- .../void/common/sendLLMMessageService.ts | 14 ++++---- .../void/common/sendLLMMessageTypes.ts | 3 +- .../llmMessage/sendLLMMessage.ts | 1 + .../electron-main/sendLLMMessageChannel.ts | 18 +++++++---- 9 files changed, 78 insertions(+), 28 deletions(-) diff --git a/src/vs/workbench/contrib/void/browser/autocompleteService.ts b/src/vs/workbench/contrib/void/browser/autocompleteService.ts index 3a621e18..894316bd 100644 --- a/src/vs/workbench/contrib/void/browser/autocompleteService.ts +++ b/src/vs/workbench/contrib/void/browser/autocompleteService.ts @@ -852,6 +852,7 @@ export class AutocompleteService extends Disposable implements IAutocompleteServ newAutocompletion.status = 'error' reject(message) }, + onAbort: () => { }, }) newAutocompletion.requestId = requestId diff --git a/src/vs/workbench/contrib/void/browser/chatThreadService.ts b/src/vs/workbench/contrib/void/browser/chatThreadService.ts index 53bf61d8..ff04587d 100644 --- a/src/vs/workbench/contrib/void/browser/chatThreadService.ts +++ b/src/vs/workbench/contrib/void/browser/chatThreadService.ts @@ -526,7 +526,8 @@ class ChatThreadService extends Disposable implements IChatThreadService { editUserMessageAndStreamResponse: IChatThreadService['editUserMessageAndStreamResponse'] = async ({ userMessage, messageIdx, threadId }) => { - const thread = this.getCurrentThread() + const thread = this.state.allThreads[threadId] + if (!thread) return // should never happen if (thread.messages?.[messageIdx]?.role !== 'user') { throw new Error(`Error: editing a message with role !=='user'`) @@ -599,8 +600,6 @@ class ChatThreadService extends Disposable implements IChatThreadService { const thread = this.state.allThreads[threadId] if (!thread) return // should never happen - const llmCancelToken = this.streamState[threadId]?.streamingToken - if (llmCancelToken !== undefined) this._llmMessageService.abort(llmCancelToken) const messageSoFar = this.streamState[threadId]?.messageSoFar ?? '' const reasoningSoFar = this.streamState[threadId]?.reasoningSoFar ?? '' @@ -613,7 +612,9 @@ class ChatThreadService extends Disposable implements IChatThreadService { // interrupt assistant message this._addMessageToThread(threadId, { role: 'assistant', content: messageSoFar, reasoning: reasoningSoFar, anthropicReasoning: null }) } - this._setStreamState(threadId, {}, 'set') + + const llmCancelToken = this.streamState[threadId]?.streamingToken + if (llmCancelToken !== undefined) this._llmMessageService.abort(llmCancelToken) } @@ -755,6 +756,7 @@ class ChatThreadService extends Disposable implements IChatThreadService { let nMessagesSent = 0 let shouldSendAnotherMessage = true let exitReason: 'end' | 'awaitingToolApproval' = 'end' as 'end' | 'awaitingToolApproval' + let aborted = false // before enter loop, call tool if (callThisToolFirst) { @@ -803,7 +805,6 @@ class ChatThreadService extends Disposable implements IChatThreadService { } resMessageIsDonePromise() - return }, onError: (error) => { const messageSoFar = this.streamState[threadId]?.messageSoFar ?? '' @@ -813,6 +814,11 @@ class ChatThreadService extends Disposable implements IChatThreadService { this._setStreamState(threadId, { error }, 'set') resMessageIsDonePromise() }, + onAbort: () => { + // stop the loop to free up the promise, but don't modify state (already handled by whatever stopped it) + resMessageIsDonePromise() + aborted = true + }, }) // should never happen, just for safety @@ -826,6 +832,10 @@ class ChatThreadService extends Disposable implements IChatThreadService { this._setStreamState(threadId, { streamingToken: llmCancelToken }, 'merge') // new stream token for the new message await messageIsDonePromise + if (aborted) { + this._metricsService.capture('Agent Loop Done', { nMessagesSent, chatMode }) + return + } } // end while // if awaiting user approval, keep isRunning true, else end isRunning @@ -846,6 +856,12 @@ class ChatThreadService extends Disposable implements IChatThreadService { const thread = this.state.allThreads[threadId] if (!thread) return // should never happen + // if the current thread is already streaming, stop it (this simply resolves the promise to free up space) + const llmCancelToken = this.streamState[threadId]?.streamingToken + if (llmCancelToken) { + if (llmCancelToken !== undefined) this._llmMessageService.abort(llmCancelToken) + } + // selections in all past chats, then in current chat (can have many duplicates here) const prevSelns: StagingSelectionItem[] = _chatSelections?.prevSelns ?? this._getAllSelections(threadId) const currSelns: StagingSelectionItem[] = _chatSelections?.currSelns ?? thread.state.stagingSelections diff --git a/src/vs/workbench/contrib/void/browser/editCodeService.ts b/src/vs/workbench/contrib/void/browser/editCodeService.ts index fe7b962f..0a111386 100644 --- a/src/vs/workbench/contrib/void/browser/editCodeService.ts +++ b/src/vs/workbench/contrib/void/browser/editCodeService.ts @@ -1311,8 +1311,10 @@ class EditCodeService extends Disposable implements IEditCodeService { if (!uri_) return uri = uri_ await this._voidModelService.initializeModel(uri) + console.log('initd2') const { model } = this._voidModelService.getModel(uri) if (!model) return + console.log('a2') currentFileStr = model.getValue(EndOfLinePreference.LF) const numLines = model.getLineCount() @@ -1333,9 +1335,10 @@ class EditCodeService extends Disposable implements IEditCodeService { const { startLine: startLine_, endLine: endLine_, _URI } = ctrlKZone uri = _URI await this._voidModelService.initializeModel(uri) + console.log('initd3') const { model } = this._voidModelService.getModel(uri) if (!model) return - + console.log('a3') currentFileStr = model.getValue(EndOfLinePreference.LF) startLine = startLine_ @@ -1456,9 +1459,6 @@ class EditCodeService extends Disposable implements IEditCodeService { const latestStreamInfoMutable: StreamLocationMutable = { line: diffZone.startLine, addedSplitYet: false, col: 1, originalCodeStartLine: 1 } - // state used in onText: - let fullTextSoFar = '' // so far (INCLUDING ignored suffix) - let prevIgnoredSuffix = '' const featureName: FeatureName = opts.from === 'ClickApply' ? 'Apply' : 'Ctrl+K' const modelSelection = this._settingsService.state.modelSelectionOfFeature[featureName] @@ -1469,6 +1469,11 @@ class EditCodeService extends Disposable implements IEditCodeService { let resMessageDonePromise: () => void = () => { } const messageDonePromise = new Promise((res_) => { resMessageDonePromise = res_ }) + // state used in onText: + let fullTextSoFar = '' // so far (INCLUDING ignored suffix) + let prevIgnoredSuffix = '' + let aborted = false + streamRequestIdRef.current = this._llmMessageService.sendLLMMessage({ messagesType: 'chatMessages', logging: { loggingName: `Edit (Writeover) - ${from}` }, @@ -1509,12 +1514,18 @@ class EditCodeService extends Disposable implements IEditCodeService { this._undoHistory(uri) resMessageDonePromise() }, + onAbort: () => { + // stop the loop to free up the promise, but don't modify state (already handled by whatever stopped it) + resMessageDonePromise() + aborted = true + }, }) // should never happen, just for safety if (streamRequestIdRef.current === null) { return } await messageDonePromise - console.log('done waiting') + if (aborted) { return } + } writeover().then(() => { @@ -1542,9 +1553,12 @@ class EditCodeService extends Disposable implements IEditCodeService { } await this._voidModelService.initializeModel(uri) + console.log('initd') + const { model } = this._voidModelService.getModel(uri) if (!model) return + console.log('a') // generate search/replace block text const originalFileCode = model.getValue(EndOfLinePreference.LF) @@ -1666,6 +1680,7 @@ class EditCodeService extends Disposable implements IEditCodeService { let shouldSendAnotherMessage = true let nMessagesSent = 0 let currStreamingBlockNum = 0 + let aborted = false while (shouldSendAnotherMessage) { shouldSendAnotherMessage = false nMessagesSent += 1 @@ -1841,13 +1856,18 @@ class EditCodeService extends Disposable implements IEditCodeService { this._undoHistory(uri) resMessageDonePromise() }, - + onAbort: () => { + // stop the loop to free up the promise, but don't modify state (already handled by whatever stopped it) + resMessageDonePromise() + aborted = true + }, }) // should never happen, just for safety if (streamRequestIdRef.current === null) { break } await messageDonePromise + if (aborted) { return } } // end while diff --git a/src/vs/workbench/contrib/void/browser/react/src/markdown/ApplyBlockHoverButtons.tsx b/src/vs/workbench/contrib/void/browser/react/src/markdown/ApplyBlockHoverButtons.tsx index 30865d49..e68b3bde 100644 --- a/src/vs/workbench/contrib/void/browser/react/src/markdown/ApplyBlockHoverButtons.tsx +++ b/src/vs/workbench/contrib/void/browser/react/src/markdown/ApplyBlockHoverButtons.tsx @@ -111,7 +111,9 @@ export const JumpToFileButton = ({ uri }: { uri: URI | 'current' }) => { const jumpToFileButton = uri !== 'current' && ( { commandService.executeCommand('vscode.open', uri, { preview: true }) }} + onClick={() => { + commandService.executeCommand('vscode.open', uri, { preview: true }) + }} title="Reject changes" /> ) diff --git a/src/vs/workbench/contrib/void/browser/toolsService.ts b/src/vs/workbench/contrib/void/browser/toolsService.ts index 6ce81ded..739e9bf6 100644 --- a/src/vs/workbench/contrib/void/browser/toolsService.ts +++ b/src/vs/workbench/contrib/void/browser/toolsService.ts @@ -11,6 +11,7 @@ import { ITerminalToolService } from './terminalToolService.js' import { ToolCallParams, ToolDirectoryItem, ToolName, ToolResultType } from '../common/toolsServiceTypes.js' import { IVoidModelService } from '../common/voidModelService.js' import { EndOfLinePreference } from '../../../../editor/common/model.js' +import { basename } from '../../../../base/common/path.js' // tool use for AI @@ -332,12 +333,14 @@ export class ToolsService implements IToolsService { }, edit: async ({ uri, changeDescription }) => { - const [_, applyDonePromise] = await editCodeService.startApplying({ + const res = await editCodeService.startApplying({ uri, applyStr: changeDescription, from: 'ClickApply', startBehavior: 'accept-conflicts', - }) ?? [] + }) + if (!res) throw new Error(`The Apply model did not start running on ${basename(uri.fsPath)}. Please try again.`) + const [_, applyDonePromise] = res await applyDonePromise return {} }, diff --git a/src/vs/workbench/contrib/void/common/sendLLMMessageService.ts b/src/vs/workbench/contrib/void/common/sendLLMMessageService.ts index 20c863fe..5b3023ce 100644 --- a/src/vs/workbench/contrib/void/common/sendLLMMessageService.ts +++ b/src/vs/workbench/contrib/void/common/sendLLMMessageService.ts @@ -38,6 +38,7 @@ export class LLMMessageService extends Disposable implements ILLMMessageService onText: {} as { [eventId: string]: ((params: EventLLMMessageOnTextParams) => void) }, onFinalMessage: {} as { [eventId: string]: ((params: EventLLMMessageOnFinalMessageParams) => void) }, onError: {} as { [eventId: string]: ((params: EventLLMMessageOnErrorParams) => void) }, + onAbort: {} as { [eventId: string]: (() => void) }, // NOT sent over the channel, result is instant when we call .abort() } // list hooks @@ -71,8 +72,8 @@ export class LLMMessageService extends Disposable implements ILLMMessageService // .listen sets up an IPC channel and takes a few ms, so we set up listeners immediately and add hooks to them instead // llm this._register((this.channel.listen('onText_sendLLMMessage') satisfies Event)(e => { this.llmMessageHooks.onText[e.requestId]?.(e) })) - this._register((this.channel.listen('onFinalMessage_sendLLMMessage') satisfies Event)(e => { this.llmMessageHooks.onFinalMessage[e.requestId]?.(e); this._onRequestIdDone(e.requestId) })) - this._register((this.channel.listen('onError_sendLLMMessage') satisfies Event)(e => { this.llmMessageHooks.onError[e.requestId]?.(e); this._onRequestIdDone(e.requestId); console.error('Error in LLMMessageService:', JSON.stringify(e)) })) + this._register((this.channel.listen('onFinalMessage_sendLLMMessage') satisfies Event)(e => { this.llmMessageHooks.onFinalMessage[e.requestId]?.(e); this._clearChannelHooks(e.requestId) })) + this._register((this.channel.listen('onError_sendLLMMessage') satisfies Event)(e => { this.llmMessageHooks.onError[e.requestId]?.(e); this._clearChannelHooks(e.requestId); console.error('Error in LLMMessageService:', JSON.stringify(e)) })) // ollama .list() this._register((this.channel.listen('onSuccess_list_ollama') satisfies Event>)(e => { this.listHooks.ollama.success[e.requestId]?.(e) })) this._register((this.channel.listen('onError_list_ollama') satisfies Event>)(e => { this.listHooks.ollama.error[e.requestId]?.(e) })) @@ -82,7 +83,7 @@ export class LLMMessageService extends Disposable implements ILLMMessageService } sendLLMMessage(params: ServiceSendLLMMessageParams) { - const { onText, onFinalMessage, onError, modelSelection, ...proxyParams } = params; + const { onText, onFinalMessage, onError, onAbort, modelSelection, ...proxyParams } = params; // throw an error if no model/provider selected (this should usually never be reached, the UI should check this first, but might happen in cases like Apply where we haven't built much UI/checks yet, good practice to have check logic on backend) if (modelSelection === null) { @@ -103,6 +104,7 @@ export class LLMMessageService extends Disposable implements ILLMMessageService this.llmMessageHooks.onText[requestId] = onText this.llmMessageHooks.onFinalMessage[requestId] = onFinalMessage this.llmMessageHooks.onError[requestId] = onError + this.llmMessageHooks.onAbort[requestId] = onAbort // used internally only const { aiInstructions } = this.voidSettingsService.state.globalSettings const { settingsOfProvider, } = this.voidSettingsService.state @@ -119,10 +121,10 @@ export class LLMMessageService extends Disposable implements ILLMMessageService return requestId } - abort(requestId: string) { + this.llmMessageHooks.onAbort[requestId]?.() // calling the abort hook here is instant (doesn't go over a channel) this.channel.call('abort', { requestId } satisfies MainLLMMessageAbortParams); - this._onRequestIdDone(requestId) + this._clearChannelHooks(requestId) } @@ -163,7 +165,7 @@ export class LLMMessageService extends Disposable implements ILLMMessageService } satisfies MainModelListParams) } - _onRequestIdDone(requestId: string) { + _clearChannelHooks(requestId: string) { delete this.llmMessageHooks.onText[requestId] delete this.llmMessageHooks.onFinalMessage[requestId] delete this.llmMessageHooks.onError[requestId] diff --git a/src/vs/workbench/contrib/void/common/sendLLMMessageTypes.ts b/src/vs/workbench/contrib/void/common/sendLLMMessageTypes.ts index 9ecb680e..f5660924 100644 --- a/src/vs/workbench/contrib/void/common/sendLLMMessageTypes.ts +++ b/src/vs/workbench/contrib/void/common/sendLLMMessageTypes.ts @@ -57,6 +57,7 @@ export type AnthropicReasoning = ({ type: 'thinking'; thinking: any; signature: export type OnText = (p: { fullText: string; fullReasoning: string }) => void export type OnFinalMessage = (p: { fullText: string; fullReasoning: string; toolCalls?: ToolCallType[]; anthropicReasoning: AnthropicReasoning[] | null }) => void // id is tool_use_id export type OnError = (p: { message: string; fullError: Error | null }) => void +export type OnAbort = () => void export type AbortRef = { current: (() => void) | null } @@ -84,6 +85,7 @@ export type ServiceSendLLMMessageParams = { logging: { loggingName: string, loggingExtras?: { [k: string]: any } }; modelSelection: ModelSelection | null; modelSelectionOptions: ModelSelectionOptions | undefined; + onAbort: OnAbort; } & SendLLMType; // params to the true sendLLMMessage function @@ -114,7 +116,6 @@ export type EventLLMMessageOnTextParams = Parameters[0] & { requestId: s export type EventLLMMessageOnFinalMessageParams = Parameters[0] & { requestId: string } export type EventLLMMessageOnErrorParams = Parameters[0] & { requestId: string } - // service -> main -> internal -> event (back to main) // (browser) diff --git a/src/vs/workbench/contrib/void/electron-main/llmMessage/sendLLMMessage.ts b/src/vs/workbench/contrib/void/electron-main/llmMessage/sendLLMMessage.ts index 25f0e19e..ad94bb72 100644 --- a/src/vs/workbench/contrib/void/electron-main/llmMessage/sendLLMMessage.ts +++ b/src/vs/workbench/contrib/void/electron-main/llmMessage/sendLLMMessage.ts @@ -85,6 +85,7 @@ export const sendLLMMessage = ({ onError_({ message: errorMessage, fullError }) } + // we should NEVER call onAbort internally, only from the outside const onAbort = () => { captureLLMEvent(`${loggingName} - Abort`, { messageLengthSoFar: _fullTextSoFar.length }) try { _aborter?.() } // aborter sometimes automatically throws an error diff --git a/src/vs/workbench/contrib/void/electron-main/sendLLMMessageChannel.ts b/src/vs/workbench/contrib/void/electron-main/sendLLMMessageChannel.ts index 12b0d984..182ce580 100644 --- a/src/vs/workbench/contrib/void/electron-main/sendLLMMessageChannel.ts +++ b/src/vs/workbench/contrib/void/electron-main/sendLLMMessageChannel.ts @@ -106,6 +106,17 @@ export class LLMMessageChannel implements IServerChannel { sendLLMMessage(mainThreadParams, this.metricsService); } + private _callAbort(params: MainLLMMessageAbortParams) { + const { requestId } = params; + if (!(requestId in this.abortRefOfRequestId)) return + this.abortRefOfRequestId[requestId].current?.() + delete this.abortRefOfRequestId[requestId] + } + + + + + _callOllamaList = (params: MainModelListParams) => { const { requestId } = params const emitters = this.listEmitters.ollama @@ -132,11 +143,4 @@ export class LLMMessageChannel implements IServerChannel { - private _callAbort(params: MainLLMMessageAbortParams) { - const { requestId } = params; - if (!(requestId in this.abortRefOfRequestId)) return - this.abortRefOfRequestId[requestId].current?.() - delete this.abortRefOfRequestId[requestId] - } - }