2
2
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3
3
* SPDX-License-Identifier: Apache-2.0
4
4
*/
5
-
6
5
import {
7
6
InlineCompletionListWithReferences ,
8
7
InlineCompletionWithReferencesParams ,
@@ -80,7 +79,7 @@ export class RecommendationService {
80
79
nextToken : request . partialResultToken ,
81
80
} ,
82
81
} )
83
- let result : InlineCompletionListWithReferences = await languageClient . sendRequest (
82
+ const result : InlineCompletionListWithReferences = await languageClient . sendRequest (
84
83
inlineCompletionWithReferencesRequestType . method ,
85
84
request ,
86
85
token
@@ -120,18 +119,10 @@ export class RecommendationService {
120
119
getLogger ( ) . info (
121
120
'Suggestion type is COMPLETIONS. Start fetching for more items if partialResultToken exists.'
122
121
)
123
- try {
124
- while ( result . partialResultToken ) {
125
- const paginatedRequest = { ...request , partialResultToken : result . partialResultToken }
126
- result = await languageClient . sendRequest (
127
- inlineCompletionWithReferencesRequestType . method ,
128
- paginatedRequest ,
129
- token
130
- )
131
- this . sessionManager . updateSessionSuggestions ( result . items )
132
- }
133
- } catch ( error ) {
134
- languageClient . warn ( `Error when getting suggestions: ${ error } ` )
122
+ if ( result . partialResultToken ) {
123
+ this . processRemainingRequests ( languageClient , request , result , token ) . catch ( ( error ) => {
124
+ languageClient . warn ( `Error when getting suggestions: ${ error } ` )
125
+ } )
135
126
}
136
127
} else {
137
128
// Skip fetching for more items if the suggesion is EDITS. If it is EDITS suggestion, only fetching for more
@@ -140,11 +131,6 @@ export class RecommendationService {
140
131
getLogger ( ) . info ( 'Suggestion type is EDITS. Skip fetching for more items.' )
141
132
this . sessionManager . updateActiveEditsStreakToken ( result . partialResultToken )
142
133
}
143
-
144
- // Close session and finalize telemetry regardless of pagination path
145
- this . sessionManager . closeSession ( )
146
- TelemetryHelper . instance . setAllPaginationEndTime ( )
147
- options . emitTelemetry && TelemetryHelper . instance . tryRecordClientComponentLatency ( )
148
134
} catch ( error : any ) {
149
135
getLogger ( ) . error ( 'Error getting recommendations: %O' , error )
150
136
// bearer token expired
@@ -167,4 +153,31 @@ export class RecommendationService {
167
153
}
168
154
}
169
155
}
156
+
157
+ private async processRemainingRequests (
158
+ languageClient : LanguageClient ,
159
+ initialRequest : InlineCompletionWithReferencesParams ,
160
+ firstResult : InlineCompletionListWithReferences ,
161
+ token : CancellationToken
162
+ ) : Promise < void > {
163
+ let nextToken = firstResult . partialResultToken
164
+ while ( nextToken ) {
165
+ const request = { ...initialRequest , partialResultToken : nextToken }
166
+
167
+ const result : InlineCompletionListWithReferences = await languageClient . sendRequest (
168
+ inlineCompletionWithReferencesRequestType . method ,
169
+ request ,
170
+ token
171
+ )
172
+ this . sessionManager . updateSessionSuggestions ( result . items )
173
+ nextToken = result . partialResultToken
174
+ }
175
+
176
+ this . sessionManager . closeSession ( )
177
+
178
+ // refresh inline completion items to render paginated responses
179
+ // All pagination requests completed
180
+ TelemetryHelper . instance . setAllPaginationEndTime ( )
181
+ TelemetryHelper . instance . tryRecordClientComponentLatency ( )
182
+ }
170
183
}
0 commit comments