22 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
33 * SPDX-License-Identifier: Apache-2.0
44 */
5-
65import {
76 InlineCompletionListWithReferences ,
87 InlineCompletionWithReferencesParams ,
@@ -80,7 +79,7 @@ export class RecommendationService {
8079 nextToken : request . partialResultToken ,
8180 } ,
8281 } )
83- let result : InlineCompletionListWithReferences = await languageClient . sendRequest (
82+ const result : InlineCompletionListWithReferences = await languageClient . sendRequest (
8483 inlineCompletionWithReferencesRequestType . method ,
8584 request ,
8685 token
@@ -120,18 +119,10 @@ export class RecommendationService {
120119 getLogger ( ) . info (
121120 'Suggestion type is COMPLETIONS. Start fetching for more items if partialResultToken exists.'
122121 )
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+ } )
135126 }
136127 } else {
137128 // 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 {
140131 getLogger ( ) . info ( 'Suggestion type is EDITS. Skip fetching for more items.' )
141132 this . sessionManager . updateActiveEditsStreakToken ( result . partialResultToken )
142133 }
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 ( )
148134 } catch ( error : any ) {
149135 getLogger ( ) . error ( 'Error getting recommendations: %O' , error )
150136 // bearer token expired
@@ -167,4 +153,31 @@ export class RecommendationService {
167153 }
168154 }
169155 }
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+ }
170183}
0 commit comments