promise resolves when abort to free space

This commit is contained in:
Andrew Pareles 2025-03-17 18:39:14 -07:00
parent ad2e5f2616
commit a320323aa1
9 changed files with 78 additions and 28 deletions

View file

@ -852,6 +852,7 @@ export class AutocompleteService extends Disposable implements IAutocompleteServ
newAutocompletion.status = 'error'
reject(message)
},
onAbort: () => { },
})
newAutocompletion.requestId = requestId

View file

@ -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

View file

@ -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<void>((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

View file

@ -111,7 +111,9 @@ export const JumpToFileButton = ({ uri }: { uri: URI | 'current' }) => {
const jumpToFileButton = uri !== 'current' && (
<IconShell1
Icon={FileSymlink}
onClick={() => { commandService.executeCommand('vscode.open', uri, { preview: true }) }}
onClick={() => {
commandService.executeCommand('vscode.open', uri, { preview: true })
}}
title="Reject changes"
/>
)

View file

@ -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 {}
},

View file

@ -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<EventLLMMessageOnTextParams>)(e => { this.llmMessageHooks.onText[e.requestId]?.(e) }))
this._register((this.channel.listen('onFinalMessage_sendLLMMessage') satisfies Event<EventLLMMessageOnFinalMessageParams>)(e => { this.llmMessageHooks.onFinalMessage[e.requestId]?.(e); this._onRequestIdDone(e.requestId) }))
this._register((this.channel.listen('onError_sendLLMMessage') satisfies Event<EventLLMMessageOnErrorParams>)(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<EventLLMMessageOnFinalMessageParams>)(e => { this.llmMessageHooks.onFinalMessage[e.requestId]?.(e); this._clearChannelHooks(e.requestId) }))
this._register((this.channel.listen('onError_sendLLMMessage') satisfies Event<EventLLMMessageOnErrorParams>)(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<EventModelListOnSuccessParams<OllamaModelResponse>>)(e => { this.listHooks.ollama.success[e.requestId]?.(e) }))
this._register((this.channel.listen('onError_list_ollama') satisfies Event<EventModelListOnErrorParams<OllamaModelResponse>>)(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<VLLMModelResponse>)
}
_onRequestIdDone(requestId: string) {
_clearChannelHooks(requestId: string) {
delete this.llmMessageHooks.onText[requestId]
delete this.llmMessageHooks.onFinalMessage[requestId]
delete this.llmMessageHooks.onError[requestId]

View file

@ -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<OnText>[0] & { requestId: s
export type EventLLMMessageOnFinalMessageParams = Parameters<OnFinalMessage>[0] & { requestId: string }
export type EventLLMMessageOnErrorParams = Parameters<OnError>[0] & { requestId: string }
// service -> main -> internal -> event (back to main)
// (browser)

View file

@ -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

View file

@ -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<OllamaModelResponse>) => {
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]
}
}