From 4c60dbe11654b255f7ad95b114a316b61a4d2816 Mon Sep 17 00:00:00 2001
From: Lei Gao <leigaol@amazon.com>
Date: Mon, 28 Jul 2025 16:17:27 -0700
Subject: [PATCH] add timeout to local LSP call inline completion

---
 .../src/app/inline/recommendationService.ts   | 39 +++++++++++++------
 packages/amazonq/src/util/timeoutUtil.ts      | 15 +++++++
 2 files changed, 43 insertions(+), 11 deletions(-)
 create mode 100644 packages/amazonq/src/util/timeoutUtil.ts

diff --git a/packages/amazonq/src/app/inline/recommendationService.ts b/packages/amazonq/src/app/inline/recommendationService.ts
index 794d6c46183..1ff060c5a70 100644
--- a/packages/amazonq/src/app/inline/recommendationService.ts
+++ b/packages/amazonq/src/app/inline/recommendationService.ts
@@ -12,10 +12,16 @@ import {
 import { CancellationToken, InlineCompletionContext, Position, TextDocument } from 'vscode'
 import { LanguageClient } from 'vscode-languageclient'
 import { SessionManager } from './sessionManager'
-import { AuthUtil, CodeWhispererStatusBarManager, vsCodeState } from 'aws-core-vscode/codewhisperer'
+import {
+    AuthUtil,
+    CodeWhispererConstants,
+    CodeWhispererStatusBarManager,
+    vsCodeState,
+} from 'aws-core-vscode/codewhisperer'
 import { TelemetryHelper } from './telemetryHelper'
 import { ICursorUpdateRecorder } from './cursorUpdateManager'
 import { getLogger } from 'aws-core-vscode/shared'
+import { asyncCallWithTimeout } from '../../util/timeoutUtil'
 
 export interface GetAllRecommendationsOptions {
     emitTelemetry?: boolean
@@ -35,6 +41,23 @@ export class RecommendationService {
         this.cursorUpdateRecorder = recorder
     }
 
+    async getRecommendationsWithTimeout(
+        languageClient: LanguageClient,
+        request: InlineCompletionWithReferencesParams,
+        token: CancellationToken
+    ) {
+        const resultPromise: Promise<InlineCompletionListWithReferences> = languageClient.sendRequest(
+            inlineCompletionWithReferencesRequestType.method,
+            request,
+            token
+        )
+        return await asyncCallWithTimeout<InlineCompletionListWithReferences>(
+            resultPromise,
+            `${inlineCompletionWithReferencesRequestType.method} time out`,
+            CodeWhispererConstants.promiseTimeoutLimit * 1000
+        )
+    }
+
     async getAllRecommendations(
         languageClient: LanguageClient,
         document: TextDocument,
@@ -93,11 +116,9 @@ export class RecommendationService {
                 },
             })
             const t0 = performance.now()
-            const result: InlineCompletionListWithReferences = await languageClient.sendRequest(
-                inlineCompletionWithReferencesRequestType.method,
-                request,
-                token
-            )
+
+            const result = await this.getRecommendationsWithTimeout(languageClient, request, token)
+
             getLogger().info('Received inline completion response from LSP: %O', {
                 sessionId: result.sessionId,
                 latency: performance.now() - t0,
@@ -181,11 +202,7 @@ export class RecommendationService {
         while (nextToken) {
             const request = { ...initialRequest, partialResultToken: nextToken }
 
-            const result: InlineCompletionListWithReferences = await languageClient.sendRequest(
-                inlineCompletionWithReferencesRequestType.method,
-                request,
-                token
-            )
+            const result = await this.getRecommendationsWithTimeout(languageClient, request, token)
             // when pagination is in progress, but user has already accepted or rejected an inline completion
             // then stop pagination
             if (this.sessionManager.getActiveSession() === undefined || vsCodeState.isCodeWhispererEditing) {
diff --git a/packages/amazonq/src/util/timeoutUtil.ts b/packages/amazonq/src/util/timeoutUtil.ts
new file mode 100644
index 00000000000..c42d1e3be01
--- /dev/null
+++ b/packages/amazonq/src/util/timeoutUtil.ts
@@ -0,0 +1,15 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export function asyncCallWithTimeout<T>(asyncPromise: Promise<T>, message: string, timeLimit: number): Promise<T> {
+    let timeoutHandle: NodeJS.Timeout
+    const timeoutPromise = new Promise((_resolve, reject) => {
+        timeoutHandle = setTimeout(() => reject(new Error(message)), timeLimit)
+    })
+    return Promise.race([asyncPromise, timeoutPromise]).then((result) => {
+        clearTimeout(timeoutHandle)
+        return result as T
+    })
+}