From 8dc96eff6ac1f03ddedde49767d10b737db0951b Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Wed, 24 Jun 2026 06:50:27 +0200 Subject: [PATCH] fix(vscode): scope notebook autocomplete aborts --- .../AutocompleteInlineCompletionProvider.ts | 20 ++-- .../unit/autocomplete-abort-scope.test.ts | 112 ++++++++++++++++++ 2 files changed, 125 insertions(+), 7 deletions(-) create mode 100644 packages/kilo-vscode/tests/unit/autocomplete-abort-scope.test.ts diff --git a/packages/kilo-vscode/src/services/autocomplete/classic-auto-complete/AutocompleteInlineCompletionProvider.ts b/packages/kilo-vscode/src/services/autocomplete/classic-auto-complete/AutocompleteInlineCompletionProvider.ts index b3fd6a1121..f1e8e86ffd 100644 --- a/packages/kilo-vscode/src/services/autocomplete/classic-auto-complete/AutocompleteInlineCompletionProvider.ts +++ b/packages/kilo-vscode/src/services/autocomplete/classic-auto-complete/AutocompleteInlineCompletionProvider.ts @@ -131,8 +131,8 @@ export class AutocompleteInlineCompletionProvider implements vscode.InlineComple private debouncedPendingRequest: PendingRequest | null = null private isFirstCall: boolean = true public readonly ignoreController: Promise - /** Abort controller for the current in-flight FIM request */ - private fimAbortController: AbortController | null = null + /** Abort controllers for in-flight FIM requests, scoped by file/notebook context. */ + private fimAbortControllers = new Map() private acceptedCommand: vscode.Disposable | null = null private contextService: ContextRetrievalService | null = null private debounceDelayMs: number = INITIAL_DEBOUNCE_DELAY_MS @@ -323,8 +323,10 @@ export class AutocompleteInlineCompletionProvider implements vscode.InlineComple } this.settleDebouncedPendingRequest() this.pendingRequests.length = 0 - this.fimAbortController?.abort() - this.fimAbortController = null + for (const controller of this.fimAbortControllers.values()) { + controller.abort() + } + this.fimAbortControllers.clear() this.telemetry?.dispose() this.contextService?.dispose() this.contextService = null @@ -580,10 +582,10 @@ export class AutocompleteInlineCompletionProvider implements vscode.InlineComple suffix: string, languageId: string, ): Promise { - // Abort any previous in-flight FIM request before starting a new one - this.fimAbortController?.abort() + // Abort only the request superseded within this file/notebook scope. + this.fimAbortControllers.get(scope)?.abort() const controller = new AbortController() - this.fimAbortController = controller + this.fimAbortControllers.set(scope, controller) const startTime = performance.now() @@ -657,6 +659,10 @@ export class AutocompleteInlineCompletionProvider implements vscode.InlineComple this.fatalNotified = true this.onFatalError?.(this.backoff.getFatalStatus()) } + } finally { + if (this.fimAbortControllers.get(scope) === controller) { + this.fimAbortControllers.delete(scope) + } } } diff --git a/packages/kilo-vscode/tests/unit/autocomplete-abort-scope.test.ts b/packages/kilo-vscode/tests/unit/autocomplete-abort-scope.test.ts new file mode 100644 index 0000000000..43ec508d4b --- /dev/null +++ b/packages/kilo-vscode/tests/unit/autocomplete-abort-scope.test.ts @@ -0,0 +1,112 @@ +import { describe, expect, it } from "bun:test" +import * as vscode from "vscode" +import { AutocompleteInlineCompletionProvider } from "../../src/services/autocomplete/classic-auto-complete/AutocompleteInlineCompletionProvider" + +function createProvider() { + ;(vscode.window as any).onDidChangeActiveTextEditor = () => ({ dispose: () => {} }) + ;(vscode.window as any).onDidChangeTextEditorSelection = () => ({ dispose: () => {} }) + return new AutocompleteInlineCompletionProvider( + {} as any, + "kilo/mistralai/codestral-2508", + { getConnectionState: () => "connected" } as any, + () => {}, + () => ({ enableAutoTrigger: true }), + "/repo", + ) +} + +function prompt() { + return { + prefix: "", + suffix: "", + modelName: "codestral", + completionOptions: {}, + selectedCompletionInfo: undefined, + } as any +} + +async function tick() { + await new Promise((resolve) => setTimeout(resolve, 0)) +} + +describe("autocomplete FIM abort scope", () => { + it("does not abort in-flight requests from another scope", async () => { + const provider = createProvider() + const signals: AbortSignal[] = [] + const resolvers: Array<() => void> = [] + ;(provider as any).fimPromptBuilder = { + getFromFIM: async ( + _connection: unknown, + _model: string, + _prompt: unknown, + process: (text: string) => unknown, + signal: AbortSignal, + ) => { + signals.push(signal) + return new Promise((resolve) => { + resolvers.push(() => { + resolve({ + suggestion: process("value"), + cost: 0, + inputTokens: 0, + outputTokens: 0, + }) + }) + }) + }, + } + + const first = provider.fetchAndCacheSuggestion("scope-a", prompt(), "", "", "python") + await tick() + + const second = provider.fetchAndCacheSuggestion("scope-b", prompt(), "", "", "python") + await tick() + + expect(signals[0]?.aborted).toBe(false) + expect(signals[1]?.aborted).toBe(false) + + resolvers.forEach((resolve) => resolve()) + await Promise.all([first, second]) + provider.dispose() + }) + + it("aborts older in-flight requests in the same scope", async () => { + const provider = createProvider() + const signals: AbortSignal[] = [] + const resolvers: Array<() => void> = [] + ;(provider as any).fimPromptBuilder = { + getFromFIM: async ( + _connection: unknown, + _model: string, + _prompt: unknown, + process: (text: string) => unknown, + signal: AbortSignal, + ) => { + signals.push(signal) + return new Promise((resolve) => { + resolvers.push(() => { + resolve({ + suggestion: process("value"), + cost: 0, + inputTokens: 0, + outputTokens: 0, + }) + }) + }) + }, + } + + const first = provider.fetchAndCacheSuggestion("scope-a", prompt(), "", "", "python") + await tick() + + const second = provider.fetchAndCacheSuggestion("scope-a", prompt(), "", "", "python") + await tick() + + expect(signals[0]?.aborted).toBe(true) + expect(signals[1]?.aborted).toBe(false) + + resolvers.forEach((resolve) => resolve()) + await Promise.all([first, second]) + provider.dispose() + }) +})