|
| 1 | +/* -------------------------------------------------------------------------------------------- |
| 2 | + * Copyright (c) Microsoft Corporation. All Rights Reserved. |
| 3 | + * See 'LICENSE' in the project root for license information. |
| 4 | + * ------------------------------------------------------------------------------------------ */ |
| 5 | +import { CodeSnippet, ContextResolver, ResolveRequest } from '@github/copilot-language-server'; |
| 6 | +import * as vscode from 'vscode'; |
| 7 | +import { DocumentSelector } from 'vscode-languageserver-protocol'; |
| 8 | +import { getOutputChannelLogger, Logger } from '../logger'; |
| 9 | +import * as telemetry from '../telemetry'; |
| 10 | +import { CompletionContextResult } from './client'; |
| 11 | +import { CopilotCompletionContextTelemetry } from './copilotCompletionContextTelemetry'; |
| 12 | +import { getCopilotApi } from './copilotProviders'; |
| 13 | +import { clients } from './extension'; |
| 14 | + |
| 15 | +class DefaultValueFallback extends Error { |
| 16 | + static readonly DefaultValue = "DefaultValue"; |
| 17 | + constructor() { super(DefaultValueFallback.DefaultValue); } |
| 18 | +} |
| 19 | + |
| 20 | +class CancellationError extends Error { |
| 21 | + static readonly Canceled = "Canceled"; |
| 22 | + constructor() { super(CancellationError.Canceled); } |
| 23 | +} |
| 24 | + |
| 25 | +class CopilotContextProviderException extends Error { |
| 26 | +} |
| 27 | + |
| 28 | +class WellKnownErrors extends Error { |
| 29 | + static readonly ClientNotFound = "ClientNotFound"; |
| 30 | + private constructor(message: string) { super(message); } |
| 31 | + public static clientNotFound(): Error { |
| 32 | + return new WellKnownErrors(WellKnownErrors.ClientNotFound); |
| 33 | + } |
| 34 | +} |
| 35 | + |
| 36 | +// Mutually exclusive values for the kind of returned completion context. They either are: |
| 37 | +// - computed. |
| 38 | +// - obtained from the cache. |
| 39 | +// - missing since the computation took too long and no cache is present (cache miss). The value |
| 40 | +// is asynchronously computed and stored in cache. |
| 41 | +// - the token is signaled as cancelled, in which case all the operations are aborted. |
| 42 | +// - an unknown state. |
| 43 | +enum CopilotCompletionKind { |
| 44 | + Computed = 'computed', |
| 45 | + GotFromCache = 'gotFromCacheHit', |
| 46 | + MissingCacheMiss = 'missingCacheMiss', |
| 47 | + Canceled = 'canceled', |
| 48 | + Unknown = 'unknown' |
| 49 | +} |
| 50 | + |
| 51 | +export class CopilotCompletionContextProvider implements ContextResolver<CodeSnippet> { |
| 52 | + private static readonly providerId = 'cppTools'; |
| 53 | + private readonly completionContextCache: Map<string, CompletionContextResult> = new Map<string, CompletionContextResult>(); |
| 54 | + private static readonly defaultCppDocumentSelector: DocumentSelector = [{ language: 'cpp' }, { language: 'c' }, { language: 'cuda-cpp' }]; |
| 55 | + private static readonly defaultTimeBudgetFactor: number = 0.5; |
| 56 | + private completionContextCancellation = new vscode.CancellationTokenSource(); |
| 57 | + private contextProviderDisposable: vscode.Disposable | undefined; |
| 58 | + |
| 59 | + private async waitForCompletionWithTimeoutAndCancellation<T>(promise: Promise<T>, defaultValue: T | undefined, |
| 60 | + timeout: number, token: vscode.CancellationToken): Promise<[T | undefined, CopilotCompletionKind]> { |
| 61 | + const defaultValuePromise = new Promise<T>((_resolve, reject) => setTimeout(() => { |
| 62 | + if (token.isCancellationRequested) { |
| 63 | + reject(new CancellationError()); |
| 64 | + } else { |
| 65 | + reject(new DefaultValueFallback()); |
| 66 | + } |
| 67 | + }, timeout)); |
| 68 | + const cancellationPromise = new Promise<T>((_, reject) => { |
| 69 | + token.onCancellationRequested(() => { |
| 70 | + reject(new CancellationError()); |
| 71 | + }); |
| 72 | + }); |
| 73 | + let snippetsOrNothing: T | undefined; |
| 74 | + try { |
| 75 | + snippetsOrNothing = await Promise.race([promise, cancellationPromise, defaultValuePromise]); |
| 76 | + } catch (e) { |
| 77 | + if (e instanceof DefaultValueFallback) { |
| 78 | + return [defaultValue, defaultValue !== undefined ? CopilotCompletionKind.GotFromCache : CopilotCompletionKind.MissingCacheMiss]; |
| 79 | + } else if (e instanceof CancellationError) { |
| 80 | + return [undefined, CopilotCompletionKind.Canceled]; |
| 81 | + } else { |
| 82 | + throw e; |
| 83 | + } |
| 84 | + } |
| 85 | + |
| 86 | + return [snippetsOrNothing, CopilotCompletionKind.Computed]; |
| 87 | + } |
| 88 | + |
| 89 | + // Get the completion context with a timeout and a cancellation token. |
| 90 | + // The cancellationToken indicates that the value should not be returned nor cached. |
| 91 | + private async getCompletionContextWithCancellation(documentUri: string, caretOffset: number, |
| 92 | + startTime: number, out: Logger, telemetry: CopilotCompletionContextTelemetry, token: vscode.CancellationToken): |
| 93 | + Promise<CompletionContextResult | undefined> { |
| 94 | + try { |
| 95 | + const docUri = vscode.Uri.parse(documentUri); |
| 96 | + const client = clients.getClientFor(docUri); |
| 97 | + if (!client) { throw WellKnownErrors.clientNotFound(); } |
| 98 | + const getContextStartTime = performance.now(); |
| 99 | + const completionContext = await client.getCompletionContext(docUri, caretOffset, token); |
| 100 | + |
| 101 | + if (completionContext.translationUnitUri !== docUri.toString()) { |
| 102 | + out.appendLine(`Copilot: getCompletionContextWithCancellation(${docUri}:${caretOffset}): translation unit URI mismatch: ${completionContext.translationUnitUri} vs ${docUri.toString()}`); |
| 103 | + } |
| 104 | + |
| 105 | + const copilotCompletionContext = completionContext; |
| 106 | + this.completionContextCache.set(completionContext.translationUnitUri, copilotCompletionContext); |
| 107 | + const duration = CopilotCompletionContextProvider.getRoundedDuration(startTime); |
| 108 | + out.appendLine(`Copilot: getCompletionContextWithCancellation(${docUri}:${caretOffset}): from ${completionContext.translationUnitUri} cached ${completionContext.snippets.length} snippets in [ms]: ${duration}`); |
| 109 | + telemetry.addSnippetCount(completionContext.snippets.length); |
| 110 | + telemetry.addCacheComputedElapsed(duration); |
| 111 | + telemetry.addComputeContextElapsed(CopilotCompletionContextProvider.getRoundedDuration(getContextStartTime)); |
| 112 | + return copilotCompletionContext; |
| 113 | + } catch (e) { |
| 114 | + if (e instanceof CancellationError) { |
| 115 | + telemetry.addInternalCanceled(CopilotCompletionContextProvider.getRoundedDuration(startTime)); |
| 116 | + throw e; |
| 117 | + } else if (e instanceof vscode.CancellationError || (e as Error)?.message === CancellationError.Canceled) { |
| 118 | + telemetry.addCopilotCanceled(CopilotCompletionContextProvider.getRoundedDuration(startTime)); |
| 119 | + throw e; |
| 120 | + } |
| 121 | + |
| 122 | + if (e instanceof WellKnownErrors) { |
| 123 | + telemetry.addWellKnownError(e.message); |
| 124 | + } |
| 125 | + |
| 126 | + const err = e as Error; |
| 127 | + out.appendLine(`Copilot: getCompletionContextWithCancellation(${documentUri}:${caretOffset}): Error: '${err?.message}', stack '${err?.stack}`); |
| 128 | + telemetry.addError(); |
| 129 | + return undefined; |
| 130 | + } finally { |
| 131 | + telemetry.file(); |
| 132 | + } |
| 133 | + } |
| 134 | + |
| 135 | + private async fetchTimeBudgetFactor(context: ResolveRequest): Promise<number> { |
| 136 | + const budgetFactor = context.activeExperiments.get("CppToolsCopilotTimeBudget"); |
| 137 | + return (budgetFactor as number) !== undefined ? budgetFactor as number : CopilotCompletionContextProvider.defaultTimeBudgetFactor; |
| 138 | + } |
| 139 | + |
| 140 | + private static getRoundedDuration(startTime: number): number { |
| 141 | + return Math.round(performance.now() - startTime); |
| 142 | + } |
| 143 | + |
| 144 | + public static async Create() { |
| 145 | + const copilotCompletionProvider = new CopilotCompletionContextProvider(); |
| 146 | + await copilotCompletionProvider.registerCopilotContextProvider(); |
| 147 | + return copilotCompletionProvider; |
| 148 | + } |
| 149 | + |
| 150 | + public dispose(): void { |
| 151 | + this.completionContextCancellation.cancel(); |
| 152 | + this.contextProviderDisposable?.dispose(); |
| 153 | + } |
| 154 | + |
| 155 | + public removeFile(fileUri: string): void { |
| 156 | + this.completionContextCache.delete(fileUri); |
| 157 | + } |
| 158 | + |
| 159 | + public async resolve(context: ResolveRequest, copilotCancel: vscode.CancellationToken): Promise<CodeSnippet[]> { |
| 160 | + const resolveStartTime = performance.now(); |
| 161 | + const out: Logger = getOutputChannelLogger(); |
| 162 | + const timeBudgetFactor = await this.fetchTimeBudgetFactor(context); |
| 163 | + const telemetry = new CopilotCompletionContextTelemetry(); |
| 164 | + let copilotCompletionContext: CompletionContextResult | undefined; |
| 165 | + let copilotCompletionContextKind: CopilotCompletionKind = CopilotCompletionKind.Unknown; |
| 166 | + try { |
| 167 | + this.completionContextCancellation.cancel(); |
| 168 | + this.completionContextCancellation = new vscode.CancellationTokenSource(); |
| 169 | + const docUri = context.documentContext.uri; |
| 170 | + const cachedValue: CompletionContextResult | undefined = this.completionContextCache.get(docUri.toString()); |
| 171 | + const computeSnippetsPromise = this.getCompletionContextWithCancellation(docUri, |
| 172 | + context.documentContext.offset, resolveStartTime, out, telemetry.fork(), this.completionContextCancellation.token); |
| 173 | + [copilotCompletionContext, copilotCompletionContextKind] = await this.waitForCompletionWithTimeoutAndCancellation( |
| 174 | + computeSnippetsPromise, cachedValue, context.timeBudget * timeBudgetFactor, copilotCancel); |
| 175 | + if (copilotCompletionContextKind === CopilotCompletionKind.Canceled) { |
| 176 | + const duration: number = CopilotCompletionContextProvider.getRoundedDuration(resolveStartTime); |
| 177 | + out.appendLine(`Copilot: getCompletionContext(${context.documentContext.uri}:${context.documentContext.offset}): cancelled, elapsed time (ms) : ${duration}`); |
| 178 | + telemetry.addInternalCanceled(duration); |
| 179 | + throw new CancellationError(); |
| 180 | + } |
| 181 | + telemetry.addSnippetCount(copilotCompletionContext?.snippets?.length); |
| 182 | + return copilotCompletionContext?.snippets ?? []; |
| 183 | + } catch (e: any) { |
| 184 | + if (e instanceof CancellationError) { |
| 185 | + throw e; |
| 186 | + } |
| 187 | + |
| 188 | + // For any other exception's type, it is an error. |
| 189 | + telemetry.addError(); |
| 190 | + throw e; |
| 191 | + } finally { |
| 192 | + telemetry.addKind(copilotCompletionContextKind.toString()); |
| 193 | + const duration: number = CopilotCompletionContextProvider.getRoundedDuration(resolveStartTime); |
| 194 | + if (copilotCompletionContext === undefined) { |
| 195 | + out.appendLine(`Copilot: getCompletionContext(${context.documentContext.uri}:${context.documentContext.offset}): no snippets provided (${copilotCompletionContextKind.toString()}), elapsed time (ms): ${duration}`); |
| 196 | + } else { |
| 197 | + const uri = copilotCompletionContext.translationUnitUri ?? "<undefined-uri>"; |
| 198 | + out.appendLine(`Copilot: getCompletionContext(${context.documentContext.uri}:${context.documentContext.offset}): for ${uri} provided ${copilotCompletionContext.snippets?.length} snippets (${copilotCompletionContextKind.toString()}), elapsed time (ms): ${duration}`); |
| 199 | + } |
| 200 | + telemetry.addResolvedElapsed(duration); |
| 201 | + telemetry.addCacheSize(this.completionContextCache.size); |
| 202 | + telemetry.file(); |
| 203 | + } |
| 204 | + } |
| 205 | + |
| 206 | + public async registerCopilotContextProvider(): Promise<void> { |
| 207 | + try { |
| 208 | + const isCustomSnippetProviderApiEnabled = await telemetry.isExperimentEnabled("CppToolsCustomSnippetsApi"); |
| 209 | + if (isCustomSnippetProviderApiEnabled) { |
| 210 | + const copilotApi = await getCopilotApi(); |
| 211 | + if (!copilotApi) { throw new CopilotContextProviderException("getCopilotApi() returned null."); } |
| 212 | + const contextAPI = await copilotApi.getContextProviderAPI("v1"); |
| 213 | + if (!contextAPI) { throw new CopilotContextProviderException("getContextProviderAPI(v1) returned null."); } |
| 214 | + this.contextProviderDisposable = contextAPI.registerContextProvider({ |
| 215 | + id: CopilotCompletionContextProvider.providerId, |
| 216 | + selector: CopilotCompletionContextProvider.defaultCppDocumentSelector, |
| 217 | + resolver: this |
| 218 | + }); |
| 219 | + } |
| 220 | + } catch (e) { |
| 221 | + console.warn("Failed to register the Copilot Context Provider."); |
| 222 | + let msg = "Failed to register the Copilot Context Provider"; |
| 223 | + if (e instanceof CopilotContextProviderException) { |
| 224 | + msg = `${msg}: ${e.message}`; |
| 225 | + } |
| 226 | + telemetry.logCopilotEvent("registerCopilotContextProviderError", { "message": msg }); |
| 227 | + } |
| 228 | + } |
| 229 | +} |
0 commit comments