Skip to content

Harden Copilot completion context caching #14747

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Sean McManus (sean-mcmanus) wants to merge 2 commits into
base: main
Choose a base branch
from
+211 −50
Draft

Harden Copilot completion context caching #14747

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d337205
Harden Copilot completion context caching
sean-mcmanus Sep 6, 2026
291afc8
Restore Copilot context diagnostics
sean-mcmanus Sep 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions Extension/src/LanguageServer/client.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<WorkspaceBrowseConfiguration | undefined> | undefined;
public lastCustomBrowseConfigurationProviderId: PersistentFolderState<string | undefined> | undefined;
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -1875,6 +1875,7 @@ export class DefaultClient implements Client {

public async onDidChangeSettings(_event: vscode.ConfigurationChangeEvent): Promise<Record<string, string>> {
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);
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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);
}

Expand Down
82 changes: 82 additions & 0 deletions Extension/src/LanguageServer/copilotCompletionContextCache.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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<T> {
id: string;
result: T;
policy: CompletionContextCachePolicy;
}

export class CompletionContextCache<T extends { caretOffset: number }> {
private readonly entries = new Map<string, CompletionContextCacheEntry<T>>();
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<T> | 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<T extends { dispose(): unknown }> {
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 = [];
}
}
98 changes: 52 additions & 46 deletions Extension/src/LanguageServer/copilotCompletionContextProvider.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -69,11 +70,9 @@ export enum CopilotCompletionKind {
Unknown = 'unknown'
}

type CacheEntry = [string, CopilotCompletionContextResult];

export class CopilotCompletionContextProvider implements ContextResolver<SupportedContextItem> {
private static readonly providerId = 'ms-vscode.cpptools';
private readonly completionContextCache: Map<string, CacheEntry> = new Map();
private readonly completionContextCache = new CompletionContextCache<CopilotCompletionContextResult>();
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;
Expand All @@ -83,7 +82,7 @@ export class CopilotCompletionContextProvider implements ContextResolver<Support
private static readonly defaultMaxSnippetLength = 3 * 1024;
private static readonly defaultDoAggregateSnippets = true;
private completionContextCancellation = new vscode.CancellationTokenSource();
private contextProviderDisposables: vscode.Disposable[] | undefined;
private readonly contextProviderDisposables = new DisposableStore<vscode.Disposable>();
static readonly CppContextProviderEnabledFeatures = 'enabledFeatures';
static readonly CppContextProviderTimeBudgetMs = 'timeBudgetMs';
static readonly CppContextProviderMaxSnippetCount = 'maxSnippetCount';
Expand Down Expand Up @@ -136,7 +135,7 @@ export class CopilotCompletionContextProvider implements ContextResolver<Support
// The cancellationToken indicates that the value should not be returned nor cached.
private async getCompletionContextWithCancellation(context: ResolveRequest, featureFlag: CopilotCompletionContextFeatures,
maxSnippetCount: number, maxSnippetLength: number, doAggregateSnippets: boolean, startTime: number, telemetry: CopilotCompletionContextTelemetry,
internalToken: vscode.CancellationToken):
internalToken: vscode.CancellationToken, cacheGeneration: number):
Promise<CopilotCompletionContextResult | undefined> {
const documentUri = context.documentContext.uri;
const caretOffset = context.documentContext.offset;
Expand All @@ -161,10 +160,20 @@ export class CopilotCompletionContextProvider implements ContextResolver<Support
if (resultMismatch) { logMessage += `(mismatch TU vs result)`; }
}
const cacheEntryId = randomUUID().toString();
this.completionContextCache.set(copilotCompletionContext.sourceFileUri, [cacheEntryId, copilotCompletionContext]);
const cached = this.completionContextCache.set(
copilotCompletionContext.sourceFileUri,
cacheEntryId,
copilotCompletionContext,
{ featureFlag: snippetsFeatureFlag, maxSnippetCount, maxSnippetLength, doAggregateSnippets },
cacheGeneration);
const duration = CopilotCompletionContextProvider.getRoundedDuration(startTime);
telemetry.addCacheComputedData(duration, cacheEntryId);
logMessage += ` cached in ${duration}ms ${copilotCompletionContext.traits.length} trait(s)`;
if (cached) {
telemetry.addCacheComputedData(duration, cacheEntryId);
logMessage += ` cached in ${duration}ms`;
} else {
logMessage += ` invalidated before caching in ${duration}ms`;
}
logMessage += ` ${copilotCompletionContext.traits.length} trait(s)`;
if (copilotCompletionContext.areSnippetsMissing) { logMessage += `(missing code snippets)`; }
else {
logMessage += ` and ${copilotCompletionContext.snippets.length} snippet(s)`;
Expand All @@ -176,7 +185,7 @@ export class CopilotCompletionContextProvider implements ContextResolver<Support
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));
Expand Down Expand Up @@ -312,16 +321,20 @@ export class CopilotCompletionContextProvider implements ContextResolver<Support

public dispose(): void {
this.completionContextCancellation.cancel();
if (this.contextProviderDisposables) {
for (const disposable of this.contextProviderDisposables) {
disposable.dispose();
}
this.contextProviderDisposables = undefined;
}
this.contextProviderDisposables.dispose();
}

public removeFile(fileUri: string): void {
this.completionContextCache.delete(fileUri);
void fileUri;
this.clear();
}

public clear(): void {
this.completionContextCache.clear();
CopilotCompletionContextProvider.paramsCacheCreated = false;
for (const key of Object.keys(CopilotCompletionContextProvider.paramsCache)) {
delete CopilotCompletionContextProvider.paramsCache[key];
}
}

private computeSnippetsResolved: boolean = true;
Expand All @@ -332,8 +345,10 @@ export class CopilotCompletionContextProvider implements ContextResolver<Support
copilotCancel: vscode.CancellationToken): Promise<[CopilotCompletionContextResult | undefined, CopilotCompletionKind]> {
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(
Expand All @@ -342,10 +357,6 @@ export class CopilotCompletionContextProvider implements ContextResolver<Support
} else { return [defaultValue, defaultValue ? CopilotCompletionKind.GotFromCache : CopilotCompletionKind.MissingCacheMiss]; }
}

private static isStaleCacheHit(caretOffset: number, cacheCaretOffset: number, maxCaretDistance: number): boolean {
return Math.abs(caretOffset - caretOffset) > maxCaretDistance;
}

private static createContextItems(copilotCompletionContext: CopilotCompletionContextResult | undefined): SupportedContextItem[] {
return [...copilotCompletionContext?.snippets ?? [], ...copilotCompletionContext?.traits ?? []] as SupportedContextItem[];
}
Expand Down Expand Up @@ -373,33 +384,32 @@ export class CopilotCompletionContextProvider implements ContextResolver<Support
maxSnippetCount, maxSnippetLength, doAggregateSnippets
});
if (featureFlag === undefined) { return []; }
const cacheEntry: CacheEntry | undefined = this.completionContextCache.get(docUri.toString());
const cachePolicy: CompletionContextCachePolicy = {
featureFlag: CopilotCompletionContextProvider.normalizeFeatureFlag(featureFlag),
maxSnippetCount,
maxSnippetLength,
doAggregateSnippets
};
const hadCacheEntry = this.completionContextCache.has(docUri.toString());
const cacheEntry = this.completionContextCache.get(
docUri.toString(), docOffset, maxCaretDistance, cachePolicy);
if (proposedEdits) {
const defaultValue = cacheEntry?.[1];
const isStaleCache = defaultValue !== undefined ? CopilotCompletionContextProvider.isStaleCacheHit(docOffset, defaultValue.caretOffset, maxCaretDistance) : true;
const contextItems = isStaleCache ? [] : CopilotCompletionContextProvider.createContextItems(defaultValue);
copilotCompletionContext = isStaleCache ? undefined : defaultValue;
copilotCompletionContextKind = isStaleCache ? CopilotCompletionKind.StaleCacheHit : CopilotCompletionKind.GotFromCache;
copilotCompletionContext = cacheEntry?.result;
copilotCompletionContextKind = cacheEntry ? CopilotCompletionKind.GotFromCache :
hadCacheEntry ? CopilotCompletionKind.StaleCacheHit : CopilotCompletionKind.MissingCacheMiss;
telemetry.addSpeculativeRequestMetadata(proposedEdits.length);
if (cacheEntry?.[0]) {
telemetry.addCacheHitEntryGuid(cacheEntry[0]);
if (cacheEntry) {
telemetry.addCacheHitEntryGuid(cacheEntry.id);
}
return contextItems;
return CopilotCompletionContextProvider.createContextItems(copilotCompletionContext);
}
const [resultContext, resultKind] = await this.resolveResultAndKind(context, featureFlag,
telemetry.fork(), cacheEntry?.[1], resolveStartTime, cppTimeBudgetMs, maxSnippetCount, maxSnippetLength, doAggregateSnippets, copilotCancel);
telemetry.fork(), cacheEntry?.result, resolveStartTime, cppTimeBudgetMs, maxSnippetCount, maxSnippetLength, doAggregateSnippets, copilotCancel);
copilotCompletionContext = resultContext;
copilotCompletionContextKind = resultKind;
logMessage += `(id: ${copilotCompletionContext?.requestId})`;
// Fix up copilotCompletionContextKind accounting for stale-cache-hits.
if (copilotCompletionContextKind === CopilotCompletionKind.GotFromCache &&
copilotCompletionContext && cacheEntry) {
telemetry.addCacheHitEntryGuid(cacheEntry[0]);
const cachedData = cacheEntry[1];
if (CopilotCompletionContextProvider.isStaleCacheHit(docOffset, cachedData.caretOffset, maxCaretDistance)) {
copilotCompletionContextKind = CopilotCompletionKind.StaleCacheHit;
copilotCompletionContext.snippets = [];
}
if (copilotCompletionContextKind === CopilotCompletionKind.GotFromCache && cacheEntry) {
telemetry.addCacheHitEntryGuid(cacheEntry.id);
}
// Handle cancellation.
if (copilotCompletionContextKind === CopilotCompletionKind.Canceled) {
Expand Down Expand Up @@ -459,9 +469,7 @@ ${copilotCompletionContext?.areSnippetsMissing ? "(missing code snippets)" : ""}
}
const disposable = await this.installContextProvider(api, contextProvider);
if (disposable) {
this.contextProviderDisposables = this.contextProviderDisposables ?? [];
this.contextProviderDisposables.push(disposable);
return true;
return this.contextProviderDisposables.add(disposable);
} else {
throw new CopilotContextProviderException("getContextProviderAPI() is not available in Copilot Chat.");
}
Expand Down Expand Up @@ -492,9 +500,7 @@ ${copilotCompletionContext?.areSnippetsMissing ? "(missing code snippets)" : ""}
}
const disposable = await this.installContextProvider(api, contextProvider);
if (disposable) {
this.contextProviderDisposables = this.contextProviderDisposables ?? [];
this.contextProviderDisposables.push(disposable);
return true;
return this.contextProviderDisposables.add(disposable);
} else {
throw new CopilotContextProviderException("getContextProviderAPI() is not available in Copilot client.");
}
Expand Down
Loading
Loading