diff --git a/packages/amazonq/src/app/inline/EditRendering/displayImage.ts b/packages/amazonq/src/app/inline/EditRendering/displayImage.ts index 2c5986afa51..28612d825b5 100644 --- a/packages/amazonq/src/app/inline/EditRendering/displayImage.ts +++ b/packages/amazonq/src/app/inline/EditRendering/displayImage.ts @@ -326,7 +326,7 @@ export async function displaySvgDecoration( selectedCompletionInfo: undefined, }, new vscode.CancellationTokenSource().token, - { emitTelemetry: false, showUi: false } + { emitTelemetry: false, showUi: false, editsStreakToken: session.editsStreakPartialResultToken } ) } }, diff --git a/packages/amazonq/src/app/inline/recommendationService.ts b/packages/amazonq/src/app/inline/recommendationService.ts index 1b121da9047..ad16a46dd33 100644 --- a/packages/amazonq/src/app/inline/recommendationService.ts +++ b/packages/amazonq/src/app/inline/recommendationService.ts @@ -20,6 +20,7 @@ import { globals, getLogger } from 'aws-core-vscode/shared' export interface GetAllRecommendationsOptions { emitTelemetry?: boolean showUi?: boolean + editsStreakToken?: number | string } export class RecommendationService { @@ -47,13 +48,16 @@ export class RecommendationService { // Record that a regular request is being made this.cursorUpdateRecorder?.recordCompletionRequest() - const request: InlineCompletionWithReferencesParams = { + let request: InlineCompletionWithReferencesParams = { textDocument: { uri: document.uri.toString(), }, position, context, } + if (options.editsStreakToken) { + request = { ...request, partialResultToken: options.editsStreakToken } + } const requestStartTime = globals.clock.Date.now() const statusBar = CodeWhispererStatusBarManager.instance @@ -112,19 +116,31 @@ export class RecommendationService { firstCompletionDisplayLatency ) - // If there are more results to fetch, handle them in the background - try { - while (result.partialResultToken) { - const paginatedRequest = { ...request, partialResultToken: result.partialResultToken } - result = await languageClient.sendRequest( - inlineCompletionWithReferencesRequestType.method, - paginatedRequest, - token - ) - this.sessionManager.updateSessionSuggestions(result.items) + const isInlineEdit = result.items.some((item) => item.isInlineEdit) + if (!isInlineEdit) { + // If the suggestion is COMPLETIONS and there are more results to fetch, handle them in the background + getLogger().info( + 'Suggestion type is COMPLETIONS. Start fetching for more items if partialResultToken exists.' + ) + try { + while (result.partialResultToken) { + const paginatedRequest = { ...request, partialResultToken: result.partialResultToken } + result = await languageClient.sendRequest( + inlineCompletionWithReferencesRequestType.method, + paginatedRequest, + token + ) + this.sessionManager.updateSessionSuggestions(result.items) + } + } catch (error) { + languageClient.warn(`Error when getting suggestions: ${error}`) } - } catch (error) { - languageClient.warn(`Error when getting suggestions: ${error}`) + } else { + // Skip fetching for more items if the suggesion is EDITS. If it is EDITS suggestion, only fetching for more + // suggestions when the user start to accept a suggesion. + // Save editsStreakPartialResultToken for the next EDITS suggestion trigger if user accepts. + getLogger().info('Suggestion type is EDITS. Skip fetching for more items.') + this.sessionManager.updateActiveEditsStreakToken(result.partialResultToken) } // Close session and finalize telemetry regardless of pagination path diff --git a/packages/amazonq/src/app/inline/sessionManager.ts b/packages/amazonq/src/app/inline/sessionManager.ts index 7b3971ae2c1..b660a5d94a7 100644 --- a/packages/amazonq/src/app/inline/sessionManager.ts +++ b/packages/amazonq/src/app/inline/sessionManager.ts @@ -14,6 +14,8 @@ export interface CodeWhispererSession { requestStartTime: number firstCompletionDisplayLatency?: number startPosition: vscode.Position + // partialResultToken for the next trigger if user accepts an EDITS suggestion + editsStreakPartialResultToken?: number | string } export class SessionManager { @@ -69,6 +71,13 @@ export class SessionManager { this._acceptedSuggestionCount += 1 } + public updateActiveEditsStreakToken(partialResultToken?: number | string) { + if (!this.activeSession || !partialResultToken) { + return + } + this.activeSession.editsStreakPartialResultToken = partialResultToken + } + public clear() { this.activeSession = undefined }