From d337205fdeafcca9ac7ac24df4526a47b97797ad Mon Sep 17 00:00:00 2001 From: Sean McManus Date: Sun, 6 Sep 2026 05:00:00 -0700 Subject: [PATCH 1/2] Harden Copilot completion context caching --- Extension/src/LanguageServer/client.ts | 10 +- .../copilotCompletionContextCache.ts | 86 +++++++++++++ .../copilotCompletionContextProvider.ts | 114 ++++++++++-------- .../copilotCompletionContextTelemetry.ts | 2 +- .../copilotCompletionContextCache.test.ts | 80 ++++++++++++ 5 files changed, 234 insertions(+), 58 deletions(-) create mode 100644 Extension/src/LanguageServer/copilotCompletionContextCache.ts create mode 100644 Extension/test/unit/copilotCompletionContextCache.test.ts 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..86f99ad5b --- /dev/null +++ b/Extension/src/LanguageServer/copilotCompletionContextCache.ts @@ -0,0 +1,86 @@ +/* -------------------------------------------------------------------------------------------- + * 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 function formatCompletionContextLocation(requestUri: string, sourceFileUri: string, caretOffset: number): string { + return `(response.uriMatchesRequest:${sourceFileUri === requestUri})(response.caretOffset:${caretOffset})`; +} + +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..0cf199662 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, formatCompletionContextLocation } 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,14 +135,14 @@ export class CopilotCompletionContextProvider implements ContextResolver { const documentUri = context.documentContext.uri; const caretOffset = context.documentContext.offset; - let logMessage = `Copilot: getCompletionContext(${documentUri}:${caretOffset}):`; + let logMessage = `Copilot: getCompletionContext:`; try { const snippetsFeatureFlag = CopilotCompletionContextProvider.normalizeFeatureFlag(featureFlag); - telemetry.addRequestMetadata(documentUri, caretOffset, context.completionId, + telemetry.addRequestMetadata(caretOffset, context.completionId, context.documentContext.languageId, { featureFlag: snippetsFeatureFlag }); const docUri = vscode.Uri.parse(documentUri); const getClientForTime = performance.now(); @@ -161,22 +160,33 @@ export class CopilotCompletionContextProvider implements ContextResolver"}:${copilotCompletionContext.caretOffset})`; + logMessage += formatCompletionContextLocation( + documentUri, copilotCompletionContext.sourceFileUri, copilotCompletionContext.caretOffset); } telemetry.addResponseMetadata(copilotCompletionContext.areSnippetsMissing, copilotCompletionContext.snippets.length, copilotCompletionContext.traits.length, copilotCompletionContext.caretOffset, copilotCompletionContext.featureFlag); telemetry.addComputeContextElapsed(CopilotCompletionContextProvider.getRoundedDuration(getCompletionContextStartTime)); - return copilotCompletionContext; + return cached ? copilotCompletionContext : undefined; } catch (e: any) { if (e instanceof vscode.CancellationError || e.message === CancellationError.Canceled) { telemetry.addInternalCanceled(CopilotCompletionContextProvider.getRoundedDuration(startTime)); @@ -189,7 +199,8 @@ 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 +359,6 @@ export class CopilotCompletionContextProvider implements ContextResolver maxCaretDistance; - } - private static createContextItems(copilotCompletionContext: CopilotCompletionContextResult | undefined): SupportedContextItem[] { return [...copilotCompletionContext?.snippets ?? [], ...copilotCompletionContext?.traits ?? []] as SupportedContextItem[]; } @@ -353,7 +366,7 @@ export class CopilotCompletionContextProvider implements ContextResolver { const proposedEdits = context.documentContext.proposedEdits; const resolveStartTime = performance.now(); - let logMessage = `Copilot: resolve(${context.documentContext.uri}:${context.documentContext.offset}):`; + let logMessage = `Copilot: resolve:`; const cppTimeBudgetMs = await this.fetchTimeBudgetMs(context); const maxCaretDistance = await this.fetchMaxDistanceToCaret(context); const maxSnippetCount = await this.fetchMaxSnippetCount(context); @@ -367,39 +380,38 @@ 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('redacts request and result URIs from diagnostics', () => { + const requestCanary = 'file:///private/request-canary.cpp'; + const resultCanary = 'file:///private/result-canary.cpp'; + const diagnostic = formatCompletionContextLocation(requestCanary, resultCanary, 42); + + strictEqual(diagnostic, '(response.uriMatchesRequest:false)(response.caretOffset:42)'); + doesNotMatch(diagnostic, /request-canary|result-canary/); + }); + + 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); + }); +}); From 291afc8a9560e584748ef9ede90fe09b3a67cc59 Mon Sep 17 00:00:00 2001 From: Sean McManus Date: Mon, 7 Sep 2026 09:39:43 -0700 Subject: [PATCH 2/2] Restore Copilot context diagnostics --- .../copilotCompletionContextCache.ts | 4 ---- .../copilotCompletionContextProvider.ts | 18 ++++++++---------- .../copilotCompletionContextTelemetry.ts | 2 +- .../unit/copilotCompletionContextCache.test.ts | 13 ++----------- 4 files changed, 11 insertions(+), 26 deletions(-) diff --git a/Extension/src/LanguageServer/copilotCompletionContextCache.ts b/Extension/src/LanguageServer/copilotCompletionContextCache.ts index 86f99ad5b..33f48f37e 100644 --- a/Extension/src/LanguageServer/copilotCompletionContextCache.ts +++ b/Extension/src/LanguageServer/copilotCompletionContextCache.ts @@ -59,10 +59,6 @@ export class CompletionContextCache { } } -export function formatCompletionContextLocation(requestUri: string, sourceFileUri: string, caretOffset: number): string { - return `(response.uriMatchesRequest:${sourceFileUri === requestUri})(response.caretOffset:${caretOffset})`; -} - export class DisposableStore { private disposables: T[] = []; private disposed = false; diff --git a/Extension/src/LanguageServer/copilotCompletionContextProvider.ts b/Extension/src/LanguageServer/copilotCompletionContextProvider.ts index 0cf199662..a7a45dba5 100644 --- a/Extension/src/LanguageServer/copilotCompletionContextProvider.ts +++ b/Extension/src/LanguageServer/copilotCompletionContextProvider.ts @@ -10,7 +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, formatCompletionContextLocation } from './copilotCompletionContextCache'; +import { CompletionContextCache, CompletionContextCachePolicy, DisposableStore } from './copilotCompletionContextCache'; import { CopilotCompletionContextTelemetry } from './copilotCompletionContextTelemetry'; import { getCopilotChatApi, getCopilotClientApi, type CopilotContextProviderAPI } from './copilotProviders'; import { clients } from './extension'; @@ -139,10 +139,10 @@ export class CopilotCompletionContextProvider implements ContextResolver { const documentUri = context.documentContext.uri; const caretOffset = context.documentContext.offset; - let logMessage = `Copilot: getCompletionContext:`; + let logMessage = `Copilot: getCompletionContext(${documentUri}:${caretOffset}):`; try { const snippetsFeatureFlag = CopilotCompletionContextProvider.normalizeFeatureFlag(featureFlag); - telemetry.addRequestMetadata(caretOffset, context.completionId, + telemetry.addRequestMetadata(documentUri, caretOffset, context.completionId, context.documentContext.languageId, { featureFlag: snippetsFeatureFlag }); const docUri = vscode.Uri.parse(documentUri); const getClientForTime = performance.now(); @@ -178,8 +178,7 @@ export class CopilotCompletionContextProvider implements ContextResolver"}:${copilotCompletionContext.caretOffset})`; } telemetry.addResponseMetadata(copilotCompletionContext.areSnippetsMissing, copilotCompletionContext.snippets.length, @@ -199,8 +198,7 @@ export class CopilotCompletionContextProvider implements ContextResolver { const proposedEdits = context.documentContext.proposedEdits; const resolveStartTime = performance.now(); - let logMessage = `Copilot: resolve:`; + let logMessage = `Copilot: resolve(${context.documentContext.uri}:${context.documentContext.offset}):`; const cppTimeBudgetMs = await this.fetchTimeBudgetMs(context); const maxCaretDistance = await this.fetchMaxDistanceToCaret(context); const maxSnippetCount = await this.fetchMaxSnippetCount(context); @@ -380,7 +378,7 @@ export class CopilotCompletionContextProvider implements ContextResolver { strictEqual(cache.get('file:///source.cpp', 100, 8, defaultPolicy)?.result.value, 'current'); }); - it('redacts request and result URIs from diagnostics', () => { - const requestCanary = 'file:///private/request-canary.cpp'; - const resultCanary = 'file:///private/result-canary.cpp'; - const diagnostic = formatCompletionContextLocation(requestCanary, resultCanary, 42); - - strictEqual(diagnostic, '(response.uriMatchesRequest:false)(response.caretOffset:42)'); - doesNotMatch(diagnostic, /request-canary|result-canary/); - }); - it('disposes registrations that complete after provider disposal', () => { const store = new DisposableStore<{ dispose(): void }>(); let disposed = 0;