From a5a748166b3bb4db5be8a8af7440be50b1a11ba2 Mon Sep 17 00:00:00 2001
From: Ben McMorran <bemcmorr@microsoft.com>
Date: Mon, 30 Sep 2024 14:56:20 -0700
Subject: [PATCH] Add cancellation token to react to Copilot API changes
 (#12773)

* Add cancellation token to react to Copilot API changes

* Check for known LSP error code

---------

Co-authored-by: Sean McManus <seanmcm@microsoft.com>
---
 Extension/src/LanguageServer/client.ts    | 46 +++++++++++++----------
 Extension/src/LanguageServer/extension.ts | 12 +++---
 2 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/Extension/src/LanguageServer/client.ts b/Extension/src/LanguageServer/client.ts
index 481aa9499c..2b1b8125b1 100644
--- 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,29 +2206,17 @@ 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 DefaultClient.withLspCancellationHandling(
+            () => this.languageClient.sendRequest(IncludesRequest, params, token), token);
     }
 
     public async getChatContext(token: vscode.CancellationToken): Promise<ChatContextResult> {
         await withCancellation(this.ready, token);
-        let result: ChatContextResult;
-        try {
-            result = await this.languageClient.sendRequest(CppContextRequest, null, token);
-        } catch (e: any) {
-            if (e instanceof ResponseError && (e.code === RequestCancelled || e.code === ServerCancelled)) {
-                throw new vscode.CancellationError();
-            }
-
-            throw e;
-        }
-        if (token.isCancellationRequested) {
-            throw new vscode.CancellationError();
-        }
-
-        return result;
+        return DefaultClient.withLspCancellationHandling(
+            () => this.languageClient.sendRequest(CppContextRequest, null, token), token);
     }
 
     /**
@@ -2310,6 +2298,26 @@ export class DefaultClient implements Client {
         this.dispatching.resolve();
     }
 
+    private static async withLspCancellationHandling<T>(task: () => Promise<T>, token: vscode.CancellationToken): Promise<T> {
+        let result: T;
+
+        try {
+            result = await task();
+        } catch (e: any) {
+            if (e instanceof ResponseError && (e.code === RequestCancelled || e.code === ServerCancelled)) {
+                throw new vscode.CancellationError();
+            } else {
+                throw e;
+            }
+        }
+
+        if (token.isCancellationRequested) {
+            throw new vscode.CancellationError();
+        }
+
+        return result;
+    }
+
     private callTaskWithTimeout<T>(task: () => Thenable<T>, ms: number, cancelToken?: vscode.CancellationTokenSource): Promise<T> {
         let timer: NodeJS.Timeout;
 
@@ -4110,6 +4118,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
index bb30b4927a..5ed2159c79 100644
--- 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;
 }