@@ -12,10 +12,16 @@ import {
1212import { CancellationToken , InlineCompletionContext , Position , TextDocument } from 'vscode'
1313import { LanguageClient } from 'vscode-languageclient'
1414import { SessionManager } from './sessionManager'
15- import { AuthUtil , CodeWhispererStatusBarManager , vsCodeState } from 'aws-core-vscode/codewhisperer'
15+ import {
16+ AuthUtil ,
17+ CodeWhispererConstants ,
18+ CodeWhispererStatusBarManager ,
19+ vsCodeState ,
20+ } from 'aws-core-vscode/codewhisperer'
1621import { TelemetryHelper } from './telemetryHelper'
1722import { ICursorUpdateRecorder } from './cursorUpdateManager'
1823import { getLogger } from 'aws-core-vscode/shared'
24+ import { asyncCallWithTimeout } from '../../util/timeoutUtil'
1925
2026export interface GetAllRecommendationsOptions {
2127 emitTelemetry ?: boolean
@@ -35,6 +41,23 @@ export class RecommendationService {
3541 this . cursorUpdateRecorder = recorder
3642 }
3743
44+ async getRecommendationsWithTimeout (
45+ languageClient : LanguageClient ,
46+ request : InlineCompletionWithReferencesParams ,
47+ token : CancellationToken
48+ ) {
49+ const resultPromise : Promise < InlineCompletionListWithReferences > = languageClient . sendRequest (
50+ inlineCompletionWithReferencesRequestType . method ,
51+ request ,
52+ token
53+ )
54+ return await asyncCallWithTimeout < InlineCompletionListWithReferences > (
55+ resultPromise ,
56+ `${ inlineCompletionWithReferencesRequestType . method } time out` ,
57+ CodeWhispererConstants . promiseTimeoutLimit * 1000
58+ )
59+ }
60+
3861 async getAllRecommendations (
3962 languageClient : LanguageClient ,
4063 document : TextDocument ,
@@ -93,11 +116,9 @@ export class RecommendationService {
93116 } ,
94117 } )
95118 const t0 = performance . now ( )
96- const result : InlineCompletionListWithReferences = await languageClient . sendRequest (
97- inlineCompletionWithReferencesRequestType . method ,
98- request ,
99- token
100- )
119+
120+ const result = await this . getRecommendationsWithTimeout ( languageClient , request , token )
121+
101122 getLogger ( ) . info ( 'Received inline completion response from LSP: %O' , {
102123 sessionId : result . sessionId ,
103124 latency : performance . now ( ) - t0 ,
@@ -181,11 +202,7 @@ export class RecommendationService {
181202 while ( nextToken ) {
182203 const request = { ...initialRequest , partialResultToken : nextToken }
183204
184- const result : InlineCompletionListWithReferences = await languageClient . sendRequest (
185- inlineCompletionWithReferencesRequestType . method ,
186- request ,
187- token
188- )
205+ const result = await this . getRecommendationsWithTimeout ( languageClient , request , token )
189206 // when pagination is in progress, but user has already accepted or rejected an inline completion
190207 // then stop pagination
191208 if ( this . sessionManager . getActiveSession ( ) === undefined || vsCodeState . isCodeWhispererEditing ) {
0 commit comments