Add cancellation token to react to Copilot API changes

microsoft · Sep 26, 2024 · 5b193a6 · 5b193a6
1 parent 2431425
commit 5b193a6
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 9 deletions.
diff --git a/Extension/src/LanguageServer/client.ts b/Extension/src/LanguageServer/client.ts
@@ -790,7 +790,7 @@ export interface Client {
     getShowConfigureIntelliSenseButton(): boolean;
     setShowConfigureIntelliSenseButton(show: boolean): void;
     addTrustedCompiler(path: string): Promise<void>;
-    getIncludes(maxDepth: number): Promise<GetIncludesResult>;
+    getIncludes(maxDepth: number, token: vscode.CancellationToken): Promise<GetIncludesResult>;
     getChatContext(token: vscode.CancellationToken): Promise<ChatContextResult>;
 }
 
@@ -2206,10 +2206,10 @@ export class DefaultClient implements Client {
         await this.languageClient.sendNotification(DidOpenNotification, params);
     }
 
-    public async getIncludes(maxDepth: number): Promise<GetIncludesResult> {
+    public async getIncludes(maxDepth: number, token: vscode.CancellationToken): Promise<GetIncludesResult> {
         const params: GetIncludesParams = { maxDepth: maxDepth };
         await this.ready;
-        return this.languageClient.sendRequest(IncludesRequest, params);
+        return this.languageClient.sendRequest(IncludesRequest, params, token);
     }
 
     public async getChatContext(token: vscode.CancellationToken): Promise<ChatContextResult> {
@@ -4110,6 +4110,6 @@ class NullClient implements Client {
     getShowConfigureIntelliSenseButton(): boolean { return false; }
     setShowConfigureIntelliSenseButton(show: boolean): void { }
     addTrustedCompiler(path: string): Promise<void> { return Promise.resolve(); }
-    getIncludes(): Promise<GetIncludesResult> { return Promise.resolve({} as GetIncludesResult); }
+    getIncludes(maxDepth: number, token: vscode.CancellationToken): Promise<GetIncludesResult> { return Promise.resolve({} as GetIncludesResult); }
     getChatContext(token: vscode.CancellationToken): Promise<ChatContextResult> { return Promise.resolve({} as ChatContextResult); }
 }
diff --git a/Extension/src/LanguageServer/extension.ts b/Extension/src/LanguageServer/extension.ts
@@ -36,8 +36,8 @@ import { makeLspRange, rangeEquals, showInstallCompilerWalkthrough } from './uti
 interface CopilotApi {
     registerRelatedFilesProvider(
         providerId: { extensionId: string; languageId: string },
-        callback: (uri: vscode.Uri) => Promise<{ entries: vscode.Uri[]; traits?: { name: string; value: string }[] }>
-    ): void;
+        callback: (uri: vscode.Uri, token: vscode.CancellationToken) => Promise<{ entries: vscode.Uri[]; traits?: { name: string; value: string }[] }>
+    ): vscode.Disposable;
 }
 
 nls.config({ messageFormat: nls.MessageFormat.bundle, bundleFormat: nls.BundleFormat.standalone })();
@@ -270,8 +270,8 @@ export async function activate(): Promise<void> {
                 for (const languageId of ['c', 'cpp', 'cuda-cpp']) {
                     api.registerRelatedFilesProvider(
                         { extensionId: util.extensionContext.extension.id, languageId },
-                        async (_uri: vscode.Uri) =>
-                            ({ entries: (await clients.ActiveClient.getIncludes(1))?.includedFiles.map(file => vscode.Uri.file(file)) ?? [] })
+                        async (_uri: vscode.Uri, token: vscode.CancellationToken) =>
+                            ({ entries: (await clients.ActiveClient.getIncludes(1, token))?.includedFiles.map(file => vscode.Uri.file(file)) ?? [] })
                     );
                 }
             } catch {
@@ -1403,7 +1403,9 @@ export async function preReleaseCheck(): Promise<void> {
 }
 
 export async function getIncludes(maxDepth: number): Promise<any> {
-    const includes = await clients.ActiveClient.getIncludes(maxDepth);
+    const tokenSource = new vscode.CancellationTokenSource();
+    const includes = await clients.ActiveClient.getIncludes(maxDepth, tokenSource.token);
+    tokenSource.dispose();
     return includes;
 }