@@ -12,10 +12,16 @@ import {
12
12
import { CancellationToken , InlineCompletionContext , Position , TextDocument } from 'vscode'
13
13
import { LanguageClient } from 'vscode-languageclient'
14
14
import { 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'
16
21
import { TelemetryHelper } from './telemetryHelper'
17
22
import { ICursorUpdateRecorder } from './cursorUpdateManager'
18
23
import { getLogger } from 'aws-core-vscode/shared'
24
+ import { asyncCallWithTimeout } from '../../util/timeoutUtil'
19
25
20
26
export interface GetAllRecommendationsOptions {
21
27
emitTelemetry ?: boolean
@@ -35,6 +41,23 @@ export class RecommendationService {
35
41
this . cursorUpdateRecorder = recorder
36
42
}
37
43
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
+
38
61
async getAllRecommendations (
39
62
languageClient : LanguageClient ,
40
63
document : TextDocument ,
@@ -93,11 +116,9 @@ export class RecommendationService {
93
116
} ,
94
117
} )
95
118
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
+
101
122
getLogger ( ) . info ( 'Received inline completion response from LSP: %O' , {
102
123
sessionId : result . sessionId ,
103
124
latency : performance . now ( ) - t0 ,
@@ -181,11 +202,7 @@ export class RecommendationService {
181
202
while ( nextToken ) {
182
203
const request = { ...initialRequest , partialResultToken : nextToken }
183
204
184
- const result : InlineCompletionListWithReferences = await languageClient . sendRequest (
185
- inlineCompletionWithReferencesRequestType . method ,
186
- request ,
187
- token
188
- )
205
+ const result = await this . getRecommendationsWithTimeout ( languageClient , request , token )
189
206
// when pagination is in progress, but user has already accepted or rejected an inline completion
190
207
// then stop pagination
191
208
if ( this . sessionManager . getActiveSession ( ) === undefined || vsCodeState . isCodeWhispererEditing ) {
0 commit comments