diff --git a/Extension/src/LanguageServer/client.ts b/Extension/src/LanguageServer/client.ts index 5e98f2859..a927d4761 100644 --- a/Extension/src/LanguageServer/client.ts +++ b/Extension/src/LanguageServer/client.ts @@ -926,7 +926,7 @@ export class DefaultClient implements Client { private browsePath?: string[]; private hoverProvider: HoverProvider | undefined; private copilotHoverProvider: CopilotHoverProvider | undefined; - private copilotCompletionProvider?: CopilotCompletionContextProvider; + private static copilotCompletionProvider?: CopilotCompletionContextProvider; public lastCustomBrowseConfiguration: PersistentFolderState | undefined; public lastCustomBrowseConfigurationProviderId: PersistentFolderState | undefined; @@ -1450,9 +1450,9 @@ export class DefaultClient implements Client { this.semanticTokensProviderDisposable = vscode.languages.registerDocumentSemanticTokensProvider(util.documentSelector, this.semanticTokensProvider, semanticTokensLegend); } - this.copilotCompletionProvider = CopilotCompletionContextProvider.Create(); + DefaultClient.copilotCompletionProvider = CopilotCompletionContextProvider.Create(); util.setProgress(util.getProgressCopilotSuccess()); - this.disposables.push(this.copilotCompletionProvider); + this.disposables.push(DefaultClient.copilotCompletionProvider); // Listen for messages from the language server. this.registerNotifications(); @@ -1875,6 +1875,7 @@ export class DefaultClient implements Client { public async onDidChangeSettings(_event: vscode.ConfigurationChangeEvent): Promise> { const defaultClient: Client = clients.getDefaultClient(); + DefaultClient.copilotCompletionProvider?.clear(); if (this === defaultClient) { // Only send the updated settings information once, as it includes values for all folders. void this.sendDidChangeSettings().catch(logAndReturn.undefined); @@ -2023,6 +2024,7 @@ export class DefaultClient implements Client { public onDidChangeTextDocument(textDocumentChangeEvent: vscode.TextDocumentChangeEvent): void { if (util.isCpp(textDocumentChangeEvent.document)) { + DefaultClient.copilotCompletionProvider?.clear(); // If any file has changed, we need to abort the current rename operation if (workspaceReferences !== undefined // Occurs when a document changes before cpptools starts. && workspaceReferences.renamePending) { @@ -2056,7 +2058,7 @@ export class DefaultClient implements Client { if (diagnosticsCollectionIntelliSense) { diagnosticsCollectionIntelliSense.delete(document.uri); } - this.copilotCompletionProvider?.removeFile(uri); + DefaultClient.copilotCompletionProvider?.removeFile(uri); openFileVersions.delete(uri); } diff --git a/Extension/src/LanguageServer/copilotCompletionContextCache.ts b/Extension/src/LanguageServer/copilotCompletionContextCache.ts new file mode 100644 index 000000000..33f48f37e --- /dev/null +++ b/Extension/src/LanguageServer/copilotCompletionContextCache.ts @@ -0,0 +1,82 @@ +/* -------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All Rights Reserved. + * See 'LICENSE' in the project root for license information. + * ------------------------------------------------------------------------------------------ */ + +export interface CompletionContextCachePolicy { + featureFlag: number; + maxSnippetCount: number; + maxSnippetLength: number; + doAggregateSnippets: boolean; +} + +export interface CompletionContextCacheEntry { + id: string; + result: T; + policy: CompletionContextCachePolicy; +} + +export class CompletionContextCache { + private readonly entries = new Map>(); + private generation = 0; + + public get size(): number { + return this.entries.size; + } + + public get currentGeneration(): number { + return this.generation; + } + + public set(uri: string, id: string, result: T, policy: CompletionContextCachePolicy, generation = this.generation): boolean { + if (generation !== this.generation) { + return false; + } + this.entries.set(uri, { id, result, policy }); + return true; + } + + public get(uri: string, caretOffset: number, maxCaretDistance: number, + policy: CompletionContextCachePolicy): CompletionContextCacheEntry | undefined { + const entry = this.entries.get(uri); + if (!entry || Math.abs(caretOffset - entry.result.caretOffset) > maxCaretDistance) { + return undefined; + } + const cachedPolicy = entry.policy; + return cachedPolicy.featureFlag === policy.featureFlag && + cachedPolicy.maxSnippetCount === policy.maxSnippetCount && + cachedPolicy.maxSnippetLength === policy.maxSnippetLength && + cachedPolicy.doAggregateSnippets === policy.doAggregateSnippets ? entry : undefined; + } + + public has(uri: string): boolean { + return this.entries.has(uri); + } + + public clear(): void { + this.entries.clear(); + this.generation++; + } +} + +export class DisposableStore { + private disposables: T[] = []; + private disposed = false; + + public add(disposable: T): boolean { + if (this.disposed) { + disposable.dispose(); + return false; + } + this.disposables.push(disposable); + return true; + } + + public dispose(): void { + this.disposed = true; + for (const disposable of this.disposables) { + disposable.dispose(); + } + this.disposables = []; + } +} diff --git a/Extension/src/LanguageServer/copilotCompletionContextProvider.ts b/Extension/src/LanguageServer/copilotCompletionContextProvider.ts index a0c3ee444..a7a45dba5 100644 --- a/Extension/src/LanguageServer/copilotCompletionContextProvider.ts +++ b/Extension/src/LanguageServer/copilotCompletionContextProvider.ts @@ -10,6 +10,7 @@ import { isBoolean, isNumber, isString } from '../common'; import { getOutputChannelLogger, Logger } from '../logger'; import * as telemetry from '../telemetry'; import { CopilotCompletionContextResult } from './client'; +import { CompletionContextCache, CompletionContextCachePolicy, DisposableStore } from './copilotCompletionContextCache'; import { CopilotCompletionContextTelemetry } from './copilotCompletionContextTelemetry'; import { getCopilotChatApi, getCopilotClientApi, type CopilotContextProviderAPI } from './copilotProviders'; import { clients } from './extension'; @@ -69,11 +70,9 @@ export enum CopilotCompletionKind { Unknown = 'unknown' } -type CacheEntry = [string, CopilotCompletionContextResult]; - export class CopilotCompletionContextProvider implements ContextResolver { private static readonly providerId = 'ms-vscode.cpptools'; - private readonly completionContextCache: Map = new Map(); + private readonly completionContextCache = new CompletionContextCache(); private static readonly defaultCppDocumentSelector: DocumentSelector = [{ language: 'cpp' }, { language: 'c' }, { language: 'cuda-cpp' }]; // The default time budget for providing a value from resolve(). private static readonly defaultTimeBudgetMs: number = 7; @@ -83,7 +82,7 @@ export class CopilotCompletionContextProvider implements ContextResolver(); static readonly CppContextProviderEnabledFeatures = 'enabledFeatures'; static readonly CppContextProviderTimeBudgetMs = 'timeBudgetMs'; static readonly CppContextProviderMaxSnippetCount = 'maxSnippetCount'; @@ -136,7 +135,7 @@ export class CopilotCompletionContextProvider implements ContextResolver { const documentUri = context.documentContext.uri; const caretOffset = context.documentContext.offset; @@ -161,10 +160,20 @@ export class CopilotCompletionContextProvider implements ContextResolver { if (this.computeSnippetsResolved) { this.computeSnippetsResolved = false; + const cacheGeneration = this.completionContextCache.currentGeneration; const computeSnippetsPromise = this.getCompletionContextWithCancellation(context, featureFlag, - maxSnippetCount, maxSnippetLength, doAggregateSnippets, resolveStartTime, telemetry.fork(), this.completionContextCancellation.token).finally( + maxSnippetCount, maxSnippetLength, doAggregateSnippets, resolveStartTime, telemetry.fork(), + this.completionContextCancellation.token, cacheGeneration).finally( () => this.computeSnippetsResolved = true ); const res = await this.waitForCompletionWithTimeoutAndCancellation( @@ -342,10 +357,6 @@ export class CopilotCompletionContextProvider implements ContextResolver maxCaretDistance; - } - private static createContextItems(copilotCompletionContext: CopilotCompletionContextResult | undefined): SupportedContextItem[] { return [...copilotCompletionContext?.snippets ?? [], ...copilotCompletionContext?.traits ?? []] as SupportedContextItem[]; } @@ -373,33 +384,32 @@ export class CopilotCompletionContextProvider implements ContextResolver { + it('requires a matching request policy and caret distance', () => { + const cache = new CompletionContextCache(); + cache.set('file:///source.cpp', 'entry', { caretOffset: 100, value: 'cached' }, defaultPolicy); + + strictEqual(cache.get('file:///source.cpp', 108, 8, defaultPolicy)?.result.value, 'cached'); + strictEqual(cache.get('file:///source.cpp', 109, 8, defaultPolicy), undefined); + strictEqual(cache.get('file:///source.cpp', 100, 8, { ...defaultPolicy, featureFlag: 2 }), undefined); + strictEqual(cache.get('file:///source.cpp', 100, 8, { ...defaultPolicy, maxSnippetCount: 1 }), undefined); + strictEqual(cache.get('file:///source.cpp', 100, 8, { ...defaultPolicy, maxSnippetLength: 128 }), undefined); + strictEqual(cache.get('file:///source.cpp', 100, 8, { ...defaultPolicy, doAggregateSnippets: false }), undefined); + }); + + it('clears cached results when source or configuration state changes', () => { + const cache = new CompletionContextCache(); + cache.set('file:///source.cpp', 'entry', { caretOffset: 100, value: 'cached' }, defaultPolicy); + + cache.clear(); + + strictEqual(cache.size, 0); + strictEqual(cache.get('file:///source.cpp', 100, 8, defaultPolicy), undefined); + }); + + it('does not repopulate after an in-flight computation is invalidated', () => { + const cache = new CompletionContextCache(); + const computationGeneration = cache.currentGeneration; + + cache.clear(); + + strictEqual( + cache.set('file:///source.cpp', 'stale-entry', { caretOffset: 100, value: 'stale' }, defaultPolicy, computationGeneration), + false); + strictEqual(cache.size, 0); + strictEqual( + cache.set('file:///source.cpp', 'current-entry', { caretOffset: 100, value: 'current' }, defaultPolicy), + true); + strictEqual(cache.get('file:///source.cpp', 100, 8, defaultPolicy)?.result.value, 'current'); + }); + + it('disposes registrations that complete after provider disposal', () => { + const store = new DisposableStore<{ dispose(): void }>(); + let disposed = 0; + + strictEqual(store.add({ dispose: () => disposed++ }), true); + store.dispose(); + strictEqual(disposed, 1); + strictEqual(store.add({ dispose: () => disposed++ }), false); + strictEqual(disposed, 2); + }); +});