diff --git a/packages/amazonq/.changes/next-release/Feature-3f18ffa1-f90f-4ceb-8af9-a913a56cd5ce.json b/packages/amazonq/.changes/next-release/Feature-3f18ffa1-f90f-4ceb-8af9-a913a56cd5ce.json new file mode 100644 index 00000000000..8a226747b15 --- /dev/null +++ b/packages/amazonq/.changes/next-release/Feature-3f18ffa1-f90f-4ceb-8af9-a913a56cd5ce.json @@ -0,0 +1,4 @@ +{ + "type": "Feature", + "description": "Amazon Q /dev: Add stop generation action" +} diff --git a/packages/core/package.nls.json b/packages/core/package.nls.json index d83d4b1ff02..d108766cd63 100644 --- a/packages/core/package.nls.json +++ b/packages/core/package.nls.json @@ -299,6 +299,7 @@ "AWS.amazonq.featureDev.pillText.generatingCode": "Generating code...", "AWS.amazonq.featureDev.pillText.requestingChanges": "Requesting changes ...", "AWS.amazonq.featureDev.pillText.insertCode": "Accept code", + "AWS.amazonq.featureDev.pillText.stoppingCodeGeneration": "Stopping code generation...", "AWS.amazonq.featureDev.pillText.sendFeedback": "Send feedback", "AWS.amazonq.featureDev.pillText.selectFiles": "Select files for context", "AWS.amazonq.featureDev.pillText.retry": "Retry", @@ -315,7 +316,7 @@ "AWS.amazonq.featureDev.answer.sessionClosed": "Okay, I've ended this chat session. You can open a new tab to chat or start another workflow.", "AWS.amazonq.featureDev.answer.newTaskChanges": "What new task would you like to work on?", "AWS.amazonq.featureDev.placeholder.chatInputDisabled": "Chat input is disabled", - "AWS.amazonq.featureDev.placeholder.additionalImprovements": "Choose an option to proceed", + "AWS.amazonq.featureDev.placeholder.additionalImprovements": "Describe your task or issue in detail", "AWS.amazonq.featureDev.placeholder.feedback": "Provide feedback or comments", "AWS.amazonq.featureDev.placeholder.describe": "Describe your task or issue in detail", "AWS.amazonq.featureDev.placeholder.sessionClosed": "Open a new chat tab to continue" diff --git a/packages/core/src/amazonq/webview/ui/apps/featureDevChatConnector.ts b/packages/core/src/amazonq/webview/ui/apps/featureDevChatConnector.ts index 3e013566c91..6ee1828801a 100644 --- a/packages/core/src/amazonq/webview/ui/apps/featureDevChatConnector.ts +++ b/packages/core/src/amazonq/webview/ui/apps/featureDevChatConnector.ts @@ -280,6 +280,7 @@ export class Connector { this.sendMessageToExtension({ tabID: tabID, command: 'stop-response', + tabType: 'featuredev', }) } diff --git a/packages/core/src/amazonq/webview/ui/connector.ts b/packages/core/src/amazonq/webview/ui/connector.ts index 3e105b31c1b..1d832c8e023 100644 --- a/packages/core/src/amazonq/webview/ui/connector.ts +++ b/packages/core/src/amazonq/webview/ui/connector.ts @@ -176,6 +176,17 @@ export class Connector { this.gumbyChatConnector.transform(tabID) } + onStopChatResponse = (tabID: string): void => { + switch (this.tabsStorage.getTab(tabID)?.type) { + case 'featuredev': + this.featureDevChatConnector.onStopChatResponse(tabID) + break + case 'cwc': + this.cwChatConnector.onStopChatResponse(tabID) + break + } + } + handleMessageReceive = async (message: MessageEvent): Promise => { if (message.data === undefined) { return @@ -457,17 +468,6 @@ export class Connector { } } - onStopChatResponse = (tabID: string): void => { - switch (this.tabsStorage.getTab(tabID)?.type) { - case 'featuredev': - this.featureDevChatConnector.onStopChatResponse(tabID) - break - case 'cwc': - this.cwChatConnector.onStopChatResponse(tabID) - break - } - } - sendFeedback = (tabId: string, feedbackPayload: FeedbackPayload): void | undefined => { switch (this.tabsStorage.getTab(tabId)?.type) { case 'featuredev': diff --git a/packages/core/src/amazonq/webview/ui/followUps/handler.ts b/packages/core/src/amazonq/webview/ui/followUps/handler.ts index 32b7600ea94..032ae47d50b 100644 --- a/packages/core/src/amazonq/webview/ui/followUps/handler.ts +++ b/packages/core/src/amazonq/webview/ui/followUps/handler.ts @@ -46,6 +46,7 @@ export class FollowUpInteractionHandler { if (followUp.prompt !== undefined) { this.mynahUI.updateStore(tabID, { loadingChat: true, + cancelButtonWhenLoading: false, promptInputDisabledState: true, }) this.mynahUI.addChatItem(tabID, { diff --git a/packages/core/src/amazonq/webview/ui/main.ts b/packages/core/src/amazonq/webview/ui/main.ts index 0c12c528c97..84e76273a59 100644 --- a/packages/core/src/amazonq/webview/ui/main.ts +++ b/packages/core/src/amazonq/webview/ui/main.ts @@ -197,6 +197,7 @@ export const createMynahUI = ( mynahUI.updateStore(tabID, { loadingChat: true, promptInputDisabledState: true, + cancelButtonWhenLoading: true, }) if (message && messageId) { @@ -307,6 +308,7 @@ export const createMynahUI = ( ) { mynahUI.updateStore(tabID, { loadingChat: true, + cancelButtonWhenLoading: false, promptInputDisabledState: true, }) @@ -462,6 +464,13 @@ export const createMynahUI = ( onTabRemove: connector.onTabRemove, onTabChange: connector.onTabChange, // TODO: update mynah-ui this type doesn't seem correct https://github.com/aws/mynah-ui/blob/3777a39eb534a91fd6b99d6cf421ce78ee5c7526/src/main.ts#L372 + onStopChatResponse: (tabID: string) => { + mynahUI.updateStore(tabID, { + loadingChat: false, + promptInputDisabledState: false, + }) + connector.onStopChatResponse(tabID) + }, onChatPrompt: (tabID: string, prompt: ChatPrompt, eventId: string | undefined) => { if ((prompt.prompt ?? '') === '' && (prompt.command ?? '') === '') { return diff --git a/packages/core/src/amazonq/webview/ui/messages/controller.ts b/packages/core/src/amazonq/webview/ui/messages/controller.ts index 26665a4f7d1..6565423bf56 100644 --- a/packages/core/src/amazonq/webview/ui/messages/controller.ts +++ b/packages/core/src/amazonq/webview/ui/messages/controller.ts @@ -76,6 +76,7 @@ export class MessageController { this.mynahUI.updateStore(selectedTab.id, { loadingChat: true, + cancelButtonWhenLoading: false, promptInputDisabledState: true, }) this.mynahUI.addChatItem(selectedTab.id, message) @@ -107,6 +108,7 @@ export class MessageController { this.mynahUI.updateStore(newTabID, { loadingChat: true, + cancelButtonWhenLoading: false, promptInputDisabledState: true, }) diff --git a/packages/core/src/amazonq/webview/ui/messages/handler.ts b/packages/core/src/amazonq/webview/ui/messages/handler.ts index e77006cea9e..d85774d23f6 100644 --- a/packages/core/src/amazonq/webview/ui/messages/handler.ts +++ b/packages/core/src/amazonq/webview/ui/messages/handler.ts @@ -36,6 +36,7 @@ export class TextMessageHandler { this.mynahUI.updateStore(tabID, { loadingChat: true, + cancelButtonWhenLoading: false, promptInputDisabledState: true, }) diff --git a/packages/core/src/amazonq/webview/ui/quickActions/handler.ts b/packages/core/src/amazonq/webview/ui/quickActions/handler.ts index 9202949c6ff..a219f71d286 100644 --- a/packages/core/src/amazonq/webview/ui/quickActions/handler.ts +++ b/packages/core/src/amazonq/webview/ui/quickActions/handler.ts @@ -79,6 +79,7 @@ export class QuickActionHandler { if (this.tabsStorage.getTab(affectedTabId)?.type !== 'unknown') { affectedTabId = this.mynahUI.updateStore('', { loadingChat: true, + cancelButtonWhenLoading: false, }) } @@ -103,6 +104,7 @@ export class QuickActionHandler { // disable chat prompt this.mynahUI.updateStore(affectedTabId, { loadingChat: true, + cancelButtonWhenLoading: false, }) this.connector.transform(affectedTabId) @@ -161,6 +163,7 @@ export class QuickActionHandler { this.mynahUI.updateStore(affectedTabId, { loadingChat: true, promptInputDisabledState: true, + cancelButtonWhenLoading: false, }) void this.connector.requestGenerativeAIAnswer(affectedTabId, '', { diff --git a/packages/core/src/amazonq/webview/ui/texts/constants.ts b/packages/core/src/amazonq/webview/ui/texts/constants.ts index d0fb74469ee..5ac1f34fda9 100644 --- a/packages/core/src/amazonq/webview/ui/texts/constants.ts +++ b/packages/core/src/amazonq/webview/ui/texts/constants.ts @@ -19,7 +19,7 @@ export const uiComponentsTexts = { save: 'Save', cancel: 'Cancel', submit: 'Submit', - stopGenerating: 'Stop generating', + stopGenerating: 'Stop', copyToClipboard: 'Copied to clipboard', noMoreTabsTooltip: 'You can only open ten conversation tabs at a time.', codeSuggestionWithReferenceTitle: 'Some suggestions contain code with references.', diff --git a/packages/core/src/amazonqFeatureDev/client/codewhispererruntime-2022-11-11.json b/packages/core/src/amazonqFeatureDev/client/codewhispererruntime-2022-11-11.json index c8690b720ee..79b9a0f66fd 100644 --- a/packages/core/src/amazonqFeatureDev/client/codewhispererruntime-2022-11-11.json +++ b/packages/core/src/amazonqFeatureDev/client/codewhispererruntime-2022-11-11.json @@ -353,8 +353,8 @@ }, "AppStudioStatePropertyValueString": { "type": "string", - "max": 1024, - "min": 1, + "max": 10240, + "min": 0, "sensitive": true }, "ArtifactId": { @@ -439,7 +439,8 @@ "acceptedCharacterCount": { "shape": "Integer" }, "acceptedLineCount": { "shape": "Integer" }, "acceptedSnippetHasReference": { "shape": "Boolean" }, - "hasProjectLevelContext": { "shape": "Boolean" } + "hasProjectLevelContext": { "shape": "Boolean" }, + "userIntent": { "shape": "UserIntent" } } }, "ChatInteractWithMessageEventInteractionTargetString": { @@ -613,6 +614,17 @@ }, "exception": true }, + "ConsoleState": { + "type": "structure", + "members": { + "region": { "shape": "String" }, + "consoleUrl": { "shape": "SensitiveString" }, + "serviceId": { "shape": "String" }, + "serviceConsolePage": { "shape": "String" }, + "serviceSubconsolePage": { "shape": "String" }, + "taskName": { "shape": "SensitiveString" } + } + }, "ContentChecksumType": { "type": "string", "enum": ["SHA_256"] @@ -657,7 +669,8 @@ "contentLength": { "shape": "CreateUploadUrlRequestContentLengthLong" }, "artifactType": { "shape": "ArtifactType" }, "uploadIntent": { "shape": "UploadIntent" }, - "uploadContext": { "shape": "UploadContext" } + "uploadContext": { "shape": "UploadContext" }, + "uploadId": { "shape": "UploadId" } } }, "CreateUploadUrlRequestContentChecksumString": { @@ -811,7 +824,8 @@ "members": { "document": { "shape": "TextDocument" }, "cursorState": { "shape": "CursorState" }, - "relevantDocuments": { "shape": "RelevantDocumentList" } + "relevantDocuments": { "shape": "RelevantDocumentList" }, + "useRelevantDocuments": { "shape": "Boolean" } } }, "EnvState": { @@ -1093,6 +1107,21 @@ "max": 10, "min": 0 }, + "InlineChatEvent": { + "type": "structure", + "members": { + "inputLength": { "shape": "PrimitiveInteger" }, + "numSelectedLines": { "shape": "PrimitiveInteger" }, + "codeIntent": { "shape": "Boolean" }, + "userDecision": { "shape": "InlineChatUserDecision" }, + "responseStartLatency": { "shape": "Double" }, + "responseEndLatency": { "shape": "Double" } + } + }, + "InlineChatUserDecision": { + "type": "string", + "enum": ["ACCEPT", "REJECT", "DISMISS"] + }, "Integer": { "type": "integer", "box": true @@ -1535,7 +1564,10 @@ "required": ["conversationState", "workspaceState"], "members": { "conversationState": { "shape": "ConversationState" }, - "workspaceState": { "shape": "WorkspaceState" } + "workspaceState": { "shape": "WorkspaceState" }, + "taskAssistPlan": { "shape": "TaskAssistPlan" }, + "codeGenerationId": { "shape": "CodeGenerationId" }, + "currentCodeGenerationId": { "shape": "CodeGenerationId" } } }, "StartTaskAssistCodeGenerationResponse": { @@ -1648,6 +1680,46 @@ "type": "string", "enum": ["DECLARATION", "USAGE"] }, + "TaskAssistPlan": { + "type": "list", + "member": { "shape": "TaskAssistPlanStep" }, + "min": 0 + }, + "TaskAssistPlanStep": { + "type": "structure", + "required": ["filePath", "description"], + "members": { + "filePath": { "shape": "TaskAssistPlanStepFilePathString" }, + "description": { "shape": "TaskAssistPlanStepDescriptionString" }, + "startLine": { "shape": "TaskAssistPlanStepStartLineInteger" }, + "endLine": { "shape": "TaskAssistPlanStepEndLineInteger" }, + "action": { "shape": "TaskAssistPlanStepAction" } + } + }, + "TaskAssistPlanStepAction": { + "type": "string", + "enum": ["MODIFY", "CREATE", "DELETE", "UNKNOWN"] + }, + "TaskAssistPlanStepDescriptionString": { + "type": "string", + "max": 1024, + "min": 1 + }, + "TaskAssistPlanStepEndLineInteger": { + "type": "integer", + "box": true, + "min": 0 + }, + "TaskAssistPlanStepFilePathString": { + "type": "string", + "max": 1024, + "min": 1 + }, + "TaskAssistPlanStepStartLineInteger": { + "type": "integer", + "box": true, + "min": 0 + }, "TaskAssistPlanningUploadContext": { "type": "structure", "required": ["conversationId"], @@ -1668,7 +1740,8 @@ "chatInteractWithMessageEvent": { "shape": "ChatInteractWithMessageEvent" }, "chatUserModificationEvent": { "shape": "ChatUserModificationEvent" }, "terminalUserInteractionEvent": { "shape": "TerminalUserInteractionEvent" }, - "featureDevEvent": { "shape": "FeatureDevEvent" } + "featureDevEvent": { "shape": "FeatureDevEvent" }, + "inlineChatEvent": { "shape": "InlineChatEvent" } }, "union": true }, @@ -1777,7 +1850,7 @@ }, "TransformationDownloadArtifactType": { "type": "string", - "enum": ["ClientInstructions", "Logs"] + "enum": ["ClientInstructions", "Logs", "GeneratedCode"] }, "TransformationDownloadArtifacts": { "type": "list", @@ -1808,7 +1881,15 @@ }, "TransformationLanguage": { "type": "string", - "enum": ["JAVA_8", "JAVA_11", "JAVA_17", "C_SHARP"] + "enum": ["JAVA_8", "JAVA_11", "JAVA_17", "C_SHARP", "COBOL", "PL_I", "JCL"] + }, + "TransformationLanguages": { + "type": "list", + "member": { "shape": "TransformationLanguage" } + }, + "TransformationMainframeRuntimeEnv": { + "type": "string", + "enum": ["MAINFRAME"] }, "TransformationOperatingSystemFamily": { "type": "string", @@ -1841,24 +1922,40 @@ }, "TransformationProgressUpdateStatus": { "type": "string", - "enum": ["IN_PROGRESS", "COMPLETED", "FAILED", "PAUSED"] + "enum": ["IN_PROGRESS", "COMPLETED", "FAILED", "PAUSED", "AWAITING_CLIENT_ACTION"] + }, + "TransformationProjectArtifactDescriptor": { + "type": "structure", + "members": { + "sourceCodeArtifact": { "shape": "TransformationSourceCodeArtifactDescriptor" } + }, + "union": true }, "TransformationProjectState": { "type": "structure", "members": { "language": { "shape": "TransformationLanguage" }, "runtimeEnv": { "shape": "TransformationRuntimeEnv" }, - "platformConfig": { "shape": "TransformationPlatformConfig" } + "platformConfig": { "shape": "TransformationPlatformConfig" }, + "projectArtifact": { "shape": "TransformationProjectArtifactDescriptor" } } }, "TransformationRuntimeEnv": { "type": "structure", "members": { "java": { "shape": "TransformationJavaRuntimeEnv" }, - "dotNet": { "shape": "TransformationDotNetRuntimeEnv" } + "dotNet": { "shape": "TransformationDotNetRuntimeEnv" }, + "mainframe": { "shape": "TransformationMainframeRuntimeEnv" } }, "union": true }, + "TransformationSourceCodeArtifactDescriptor": { + "type": "structure", + "members": { + "languages": { "shape": "TransformationLanguages" }, + "runtimeEnv": { "shape": "TransformationRuntimeEnv" } + } + }, "TransformationSpec": { "type": "structure", "members": { @@ -1912,11 +2009,11 @@ }, "TransformationType": { "type": "string", - "enum": ["LANGUAGE_UPGRADE"] + "enum": ["LANGUAGE_UPGRADE", "DOCUMENT_GENERATION"] }, "TransformationUploadArtifactType": { "type": "string", - "enum": ["Dependencies"] + "enum": ["Dependencies", "ClientBuildResult"] }, "TransformationUploadContext": { "type": "structure", @@ -1998,7 +2095,9 @@ "gitState": { "shape": "GitState" }, "envState": { "shape": "EnvState" }, "appStudioContext": { "shape": "AppStudioState" }, - "diagnostic": { "shape": "Diagnostic" } + "diagnostic": { "shape": "Diagnostic" }, + "consoleState": { "shape": "ConsoleState" }, + "userSettings": { "shape": "UserSettings" } } }, "UserIntent": { @@ -2011,7 +2110,8 @@ "CITE_SOURCES", "EXPLAIN_LINE_BY_LINE", "EXPLAIN_CODE_SELECTION", - "GENERATE_CLOUDFORMATION_TEMPLATE" + "GENERATE_CLOUDFORMATION_TEMPLATE", + "GENERATE_UNIT_TESTS" ] }, "UserModificationEvent": { @@ -2026,6 +2126,12 @@ "timestamp": { "shape": "Timestamp" } } }, + "UserSettings": { + "type": "structure", + "members": { + "hasConsentedToCrossRegionCalls": { "shape": "Boolean" } + } + }, "UserTriggerDecisionEvent": { "type": "structure", "required": [ diff --git a/packages/core/src/amazonqFeatureDev/client/featureDev.ts b/packages/core/src/amazonqFeatureDev/client/featureDev.ts index 847db903ab1..8cd5ee97f90 100644 --- a/packages/core/src/amazonqFeatureDev/client/featureDev.ts +++ b/packages/core/src/amazonqFeatureDev/client/featureDev.ts @@ -10,7 +10,6 @@ import { ServiceOptions } from '../../shared/awsClientBuilder' import globals from '../../shared/extensionGlobals' import { getLogger } from '../../shared/logger' import * as FeatureDevProxyClient from './featuredevproxyclient' -import apiConfig = require('./codewhispererruntime-2022-11-11.json') import { featureName } from '../constants' import { CodeReference } from '../../amazonq/webview/ui/connector' import { @@ -25,6 +24,7 @@ import { getCodewhispererConfig } from '../../codewhisperer/client/codewhisperer import { createCodeWhispererChatStreamingClient } from '../../shared/clients/codewhispererChatClient' import { getClientId, getOptOutPreference, getOperatingSystem } from '../../shared/telemetry/util' import { extensionVersion } from '../../shared/vscode/env' +import apiConfig = require('./codewhispererruntime-2022-11-11.json') // Re-enable once BE is able to handle retries. const writeAPIRetryOptions = { @@ -87,7 +87,12 @@ export class FeatureDevClient { } } - public async createUploadUrl(conversationId: string, contentChecksumSha256: string, contentLength: number) { + public async createUploadUrl( + conversationId: string, + contentChecksumSha256: string, + contentLength: number, + uploadId: string + ) { try { const client = await this.getClient(writeAPIRetryOptions) const params = { @@ -96,6 +101,7 @@ export class FeatureDevClient { conversationId, }, }, + uploadId, contentChecksum: contentChecksumSha256, contentChecksumType: 'SHA_256', artifactType: 'SourceCode', @@ -105,7 +111,7 @@ export class FeatureDevClient { getLogger().debug(`Executing createUploadUrl with %O`, omit(params, 'contentChecksum')) const response = await client.createUploadUrl(params).promise() getLogger().debug(`${featureName}: Created upload url: %O`, { - uploadId: response.uploadId, + uploadId: uploadId, requestId: response.$response.requestId, }) return response @@ -124,10 +130,17 @@ export class FeatureDevClient { } } - public async startCodeGeneration(conversationId: string, uploadId: string, message: string) { + public async startCodeGeneration( + conversationId: string, + uploadId: string, + message: string, + codeGenerationId: string, + currentCodeGenerationId?: string + ) { try { const client = await this.getClient(writeAPIRetryOptions) const params = { + codeGenerationId, conversationState: { conversationId, currentMessage: { @@ -139,6 +152,9 @@ export class FeatureDevClient { uploadId, programmingLanguage: { languageName: 'javascript' }, }, + } as FeatureDevProxyClient.Types.StartTaskAssistCodeGenerationRequest + if (currentCodeGenerationId) { + params.currentCodeGenerationId = currentCodeGenerationId } getLogger().debug(`Executing startTaskAssistCodeGeneration with %O`, params) const response = await client.startTaskAssistCodeGeneration(params).promise() diff --git a/packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts b/packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts index 831d66fb2f8..aa300b9ddba 100644 --- a/packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts +++ b/packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts @@ -47,6 +47,8 @@ import { openDeletedDiff, openDiff } from '../../../amazonq/commons/diff' import { i18n } from '../../../shared/i18n-helper' import globals from '../../../shared/extensionGlobals' +export const TotalSteps = 3 + export interface ChatControllerEventEmitters { readonly processHumanChatMessage: EventEmitter readonly followUpClicked: EventEmitter @@ -311,6 +313,25 @@ export class FeatureDevController { } } + /** + * + * This function dispose cancellation token to free resources and provide a new token. + * Since user can abort a call in the same session, when the processing ends, we need provide a new one + * to start with the new prompt and allow the ability to stop again. + * + * @param session + */ + + private disposeToken(session: Session | undefined) { + if (session?.state?.tokenSource?.token.isCancellationRequested) { + session?.state.tokenSource?.dispose() + if (session?.state?.tokenSource) { + session.state.tokenSource = new vscode.CancellationTokenSource() + } + getLogger().debug('Request cancelled, skipping further processing') + } + } + // TODO add type private async processUserChatMessage(message: any) { if (message.message === undefined) { @@ -346,6 +367,7 @@ export class FeatureDevController { await this.onCodeGeneration(session, message.message, message.tabID) } } catch (err: any) { + this.disposeToken(session) await this.processErrorChatMessage(err, message, session) // Lock the chat input until they explicitly click one of the follow ups this.messenger.sendChatInputEnabled(message.tabID, false) @@ -376,6 +398,11 @@ export class FeatureDevController { await session.send(message) const filePaths = session.state.filePaths ?? [] const deletedFiles = session.state.deletedFiles ?? [] + // Only add the follow up accept/deny buttons when the tab hasn't been closed/request hasn't been cancelled + if (session?.state?.tokenSource?.token.isCancellationRequested) { + return + } + if (filePaths.length === 0 && deletedFiles.length === 0) { this.messenger.sendAnswer({ message: i18n('AWS.amazonq.featureDev.pillText.unableGenerateChanges'), @@ -402,11 +429,6 @@ export class FeatureDevController { return } - // Only add the follow up accept/deny buttons when the tab hasn't been closed/request hasn't been cancelled - if (session?.state.tokenSource.token.isCancellationRequested) { - return - } - this.messenger.sendCodeResult( filePaths, deletedFiles, @@ -439,25 +461,82 @@ export class FeatureDevController { this.messenger.sendUpdatePlaceholder(tabID, i18n('AWS.amazonq.featureDev.pillText.selectOption')) } finally { // Finish processing the event - this.messenger.sendAsyncEventProgress(tabID, false, undefined) - - // Lock the chat input until they explicitly click one of the follow ups - this.messenger.sendChatInputEnabled(tabID, false) - if (!this.isAmazonQVisible) { - const open = 'Open chat' - const resp = await vscode.window.showInformationMessage( - i18n('AWS.amazonq.featureDev.answer.qGeneratedCode'), - open + if (session?.state?.tokenSource?.token.isCancellationRequested) { + this.workOnNewTask( + session, + session.state.codeGenerationRemainingIterationCount || + TotalSteps - (session.state?.currentIteration || 0), + session.state.codeGenerationTotalIterationCount || TotalSteps, + session?.state?.tokenSource?.token.isCancellationRequested ) - if (resp === open) { - await vscode.commands.executeCommand('aws.AmazonQChatView.focus') - // TODO add focusing on the specific tab once that's implemented + this.disposeToken(session) + } else { + this.messenger.sendAsyncEventProgress(tabID, false, undefined) + + // Lock the chat input until they explicitly click one of the follow ups + this.messenger.sendChatInputEnabled(tabID, false) + + if (!this.isAmazonQVisible) { + const open = 'Open chat' + const resp = await vscode.window.showInformationMessage( + i18n('AWS.amazonq.featureDev.answer.qGeneratedCode'), + open + ) + if (resp === open) { + await vscode.commands.executeCommand('aws.AmazonQChatView.focus') + // TODO add focusing on the specific tab once that's implemented + } } } } } + private workOnNewTask( + message: any, + remainingIterations: number = 0, + totalIterations?: number, + isStoppedGeneration: boolean = false + ) { + if (isStoppedGeneration) { + this.messenger.sendAnswer({ + message: + (remainingIterations ?? 0) <= 0 + ? "I stopped generating your code. You don't have more iterations left, however, you can start a new session." + : `I stopped generating your code. If you want to continue working on this task, provide another description. You have ${remainingIterations} out of ${totalIterations} code generations left.`, + type: 'answer-part', + tabID: message.tabID, + }) + } + + if ((remainingIterations <= 0 && isStoppedGeneration) || !isStoppedGeneration) { + this.messenger.sendAnswer({ + type: 'system-prompt', + tabID: message.tabID, + followUps: [ + { + pillText: i18n('AWS.amazonq.featureDev.pillText.newTask'), + type: FollowUpTypes.NewTask, + status: 'info', + }, + { + pillText: i18n('AWS.amazonq.featureDev.pillText.closeSession'), + type: FollowUpTypes.CloseSession, + status: 'info', + }, + ], + }) + this.messenger.sendChatInputEnabled(message.tabID, false) + this.messenger.sendUpdatePlaceholder(message.tabID, i18n('AWS.amazonq.featureDev.pillText.selectOption')) + return + } + // Ensure that chat input is enabled so that they can provide additional iterations if they choose + this.messenger.sendChatInputEnabled(message.tabID, true) + this.messenger.sendUpdatePlaceholder( + message.tabID, + i18n('AWS.amazonq.featureDev.placeholder.additionalImprovements') + ) + } // TODO add type private async insertCode(message: any) { let session @@ -477,7 +556,6 @@ export class FeatureDevController { result: 'Succeeded', }) await session.insertChanges() - this.messenger.sendAnswer({ type: 'answer', tabID: message.tabID, @@ -485,26 +563,10 @@ export class FeatureDevController { canBeVoted: true, }) - this.messenger.sendAnswer({ - type: 'system-prompt', - tabID: message.tabID, - followUps: [ - { - pillText: i18n('AWS.amazonq.featureDev.pillText.newTask'), - type: FollowUpTypes.NewTask, - status: 'info', - }, - { - pillText: i18n('AWS.amazonq.featureDev.pillText.closeSession'), - type: FollowUpTypes.CloseSession, - status: 'info', - }, - ], - }) - - this.messenger.sendUpdatePlaceholder( - message.tabID, - i18n('AWS.amazonq.featureDev.placeholder.additionalImprovements') + this.workOnNewTask( + message, + session.state.codeGenerationRemainingIterationCount, + session.state.codeGenerationTotalIterationCount ) } catch (err: any) { this.messenger.sendErrorMessage( @@ -726,8 +788,22 @@ export class FeatureDevController { } private async stopResponse(message: any) { + telemetry.ui_click.emit({ elementId: 'amazonq_stopCodeGeneration' }) + this.messenger.sendAnswer({ + message: i18n('AWS.amazonq.featureDev.pillText.stoppingCodeGeneration'), + type: 'answer-part', + tabID: message.tabID, + }) + this.messenger.sendUpdatePlaceholder( + message.tabID, + i18n('AWS.amazonq.featureDev.pillText.stoppingCodeGeneration') + ) + this.messenger.sendChatInputEnabled(message.tabID, false) + const session = await this.sessionStorage.getSession(message.tabID) - session.state.tokenSource.cancel() + if (session.state?.tokenSource) { + session.state?.tokenSource?.cancel() + } } private async tabOpened(message: any) { diff --git a/packages/core/src/amazonqFeatureDev/session/session.ts b/packages/core/src/amazonqFeatureDev/session/session.ts index 27f26b18d15..204e974eee0 100644 --- a/packages/core/src/amazonqFeatureDev/session/session.ts +++ b/packages/core/src/amazonqFeatureDev/session/session.ts @@ -26,7 +26,6 @@ import { ReferenceLogViewProvider } from '../../codewhisperer/service/referenceL import { AuthUtil } from '../../codewhisperer/util/authUtil' import { getLogger } from '../../shared' import { logWithConversationId } from '../userFacingText' - export class Session { private _state?: SessionState | Omit private task: string = '' @@ -89,6 +88,7 @@ export class Session { ...this.getSessionStateConfig(), conversationId: this.conversationId, uploadId: '', + currentCodeGenerationId: undefined, }, [], [], @@ -130,13 +130,14 @@ export class Session { fs: this.config.fs, messenger: this.messenger, telemetry: this.telemetry, + tokenSource: this.state.tokenSource, uploadHistory: this.state.uploadHistory, }) if (resp.nextState) { - // Cancel the request before moving to a new state - this.state.tokenSource.cancel() - + if (!this.state?.tokenSource?.token.isCancellationRequested) { + this.state?.tokenSource?.cancel() + } // Move to the next state this._state = resp.nextState } @@ -182,6 +183,10 @@ export class Session { return this._state } + get currentCodeGenerationId() { + return this.state.currentCodeGenerationId + } + get uploadId() { if (!('uploadId' in this.state)) { throw new Error("UploadId has to be initialized before it's read") diff --git a/packages/core/src/amazonqFeatureDev/session/sessionState.ts b/packages/core/src/amazonqFeatureDev/session/sessionState.ts index f23074f5009..939234b5947 100644 --- a/packages/core/src/amazonqFeatureDev/session/sessionState.ts +++ b/packages/core/src/amazonqFeatureDev/session/sessionState.ts @@ -44,6 +44,8 @@ import { collectFiles, getWorkspaceFoldersByPrefixes } from '../../shared/utilit import { i18n } from '../../shared/i18n-helper' import { Messenger } from '../controllers/chat/messenger/messenger' +const EmptyCodeGenID = 'EMPTY_CURRENT_CODE_GENERATION_ID' + export class ConversationNotStartedState implements Omit { public tokenSource: vscode.CancellationTokenSource public readonly phase = DevPhase.INIT @@ -133,10 +135,12 @@ function getDeletedFileInfos(deletedFiles: string[], workspaceFolders: CurrentWs abstract class CodeGenBase { private pollCount = 360 private requestDelay = 5000 - readonly tokenSource: vscode.CancellationTokenSource + public tokenSource: vscode.CancellationTokenSource public phase: SessionStatePhase = DevPhase.CODEGEN public readonly conversationId: string public readonly uploadId: string + public currentCodeGenerationId?: string + public isCancellationRequested?: boolean constructor( protected config: SessionStateConfig, @@ -145,6 +149,7 @@ abstract class CodeGenBase { this.tokenSource = new vscode.CancellationTokenSource() this.conversationId = config.conversationId this.uploadId = config.uploadId + this.currentCodeGenerationId = config.currentCodeGenerationId || EmptyCodeGenID } async generateCode({ @@ -168,7 +173,7 @@ abstract class CodeGenBase { }> { for ( let pollingIteration = 0; - pollingIteration < this.pollCount && !this.tokenSource.token.isCancellationRequested; + pollingIteration < this.pollCount && !this.isCancellationRequested; ++pollingIteration ) { const codegenResult = await this.config.proxyClient.getCodeGeneration(this.conversationId, codeGenerationId) @@ -253,7 +258,7 @@ abstract class CodeGenBase { } } } - if (!this.tokenSource.token.isCancellationRequested) { + if (!this.isCancellationRequested) { // still in progress const errorMessage = i18n('AWS.amazonq.featureDev.error.codeGen.timeout') throw new ToolkitError(errorMessage, { code: 'CodeGenTimeout' }) @@ -273,7 +278,7 @@ export class CodeGenState extends CodeGenBase implements SessionState { public deletedFiles: DeletedFileInfo[], public references: CodeReference[], tabID: string, - private currentIteration: number, + public currentIteration: number, public uploadHistory: UploadHistory, public codeGenerationRemainingIterationCount?: number, public codeGenerationTotalIterationCount?: number @@ -284,6 +289,12 @@ export class CodeGenState extends CodeGenBase implements SessionState { async interact(action: SessionStateAction): Promise { return telemetry.amazonq_codeGenerationInvoke.run(async (span) => { try { + action.tokenSource?.token.onCancellationRequested(() => { + this.isCancellationRequested = true + if (action.tokenSource) { + this.tokenSource = action.tokenSource + } + }) span.record({ amazonqConversationId: this.config.conversationId, credentialStartUrl: AuthUtil.instance.startUrl, @@ -291,22 +302,26 @@ export class CodeGenState extends CodeGenBase implements SessionState { action.telemetry.setGenerateCodeIteration(this.currentIteration) action.telemetry.setGenerateCodeLastInvocationTime() - - const { codeGenerationId } = await this.config.proxyClient.startCodeGeneration( + const codeGenerationId = randomUUID() + await this.config.proxyClient.startCodeGeneration( this.config.conversationId, this.config.uploadId, - action.msg + action.msg, + codeGenerationId, + this.currentCodeGenerationId ) - action.messenger.sendAnswer({ - message: i18n('AWS.amazonq.featureDev.pillText.generatingCode'), - type: 'answer-part', - tabID: this.tabID, - }) - action.messenger.sendUpdatePlaceholder( - this.tabID, - i18n('AWS.amazonq.featureDev.pillText.generatingCode') - ) + if (!this.isCancellationRequested) { + action.messenger.sendAnswer({ + message: i18n('AWS.amazonq.featureDev.pillText.generatingCode'), + type: 'answer-part', + tabID: this.tabID, + }) + action.messenger.sendUpdatePlaceholder( + this.tabID, + i18n('AWS.amazonq.featureDev.pillText.generatingCode') + ) + } const codeGeneration = await this.generateCode({ messenger: action.messenger, @@ -316,6 +331,11 @@ export class CodeGenState extends CodeGenBase implements SessionState { workspaceFolders: this.config.workspaceFolders, }) + if (codeGeneration && !action.tokenSource?.token.isCancellationRequested) { + this.config.currentCodeGenerationId = codeGenerationId + this.currentCodeGenerationId = codeGenerationId + } + this.filePaths = codeGeneration.newFiles this.deletedFiles = codeGeneration.deletedFiles this.references = codeGeneration.references @@ -344,6 +364,8 @@ export class CodeGenState extends CodeGenBase implements SessionState { this.codeGenerationRemainingIterationCount, this.codeGenerationTotalIterationCount, action.uploadHistory, + this.tokenSource, + this.currentCodeGenerationId, codeGenerationId ) return { @@ -453,24 +475,27 @@ export class MockCodeGenState implements SessionState { } export class PrepareCodeGenState implements SessionState { - public tokenSource: vscode.CancellationTokenSource public readonly phase = DevPhase.CODEGEN public uploadId: string public conversationId: string + public tokenSource: vscode.CancellationTokenSource constructor( private config: SessionStateConfig, public filePaths: NewFileInfo[], public deletedFiles: DeletedFileInfo[], public references: CodeReference[], public tabID: string, - private currentIteration: number, + public currentIteration: number, public codeGenerationRemainingIterationCount?: number, public codeGenerationTotalIterationCount?: number, public uploadHistory: UploadHistory = {}, + public superTokenSource: vscode.CancellationTokenSource = new vscode.CancellationTokenSource(), + public currentCodeGenerationId?: string, public codeGenerationId?: string ) { - this.tokenSource = new vscode.CancellationTokenSource() + this.tokenSource = superTokenSource || new vscode.CancellationTokenSource() this.uploadId = config.uploadId + this.currentCodeGenerationId = currentCodeGenerationId this.conversationId = config.conversationId this.uploadHistory = uploadHistory this.codeGenerationId = codeGenerationId @@ -488,7 +513,6 @@ export class PrepareCodeGenState implements SessionState { }) action.messenger.sendUpdatePlaceholder(this.tabID, i18n('AWS.amazonq.featureDev.pillText.uploadingCode')) - const uploadId = await telemetry.amazonq_createUpload.run(async (span) => { span.record({ amazonqConversationId: this.config.conversationId, @@ -500,30 +524,34 @@ export class PrepareCodeGenState implements SessionState { action.telemetry, span ) - - const { uploadUrl, uploadId, kmsKeyArn } = await this.config.proxyClient.createUploadUrl( + const uploadId = randomUUID() + const { uploadUrl, kmsKeyArn } = await this.config.proxyClient.createUploadUrl( this.config.conversationId, zipFileChecksum, - zipFileBuffer.length + zipFileBuffer.length, + uploadId ) await uploadCode(uploadUrl, zipFileBuffer, zipFileChecksum, kmsKeyArn) - action.messenger.sendAnswer({ - message: i18n('AWS.amazonq.featureDev.pillText.contextGatheringCompleted'), - type: 'answer-part', - tabID: this.tabID, - }) + if (!action.tokenSource?.token.isCancellationRequested) { + action.messenger.sendAnswer({ + message: i18n('AWS.amazonq.featureDev.pillText.contextGatheringCompleted'), + type: 'answer-part', + tabID: this.tabID, + }) - action.messenger.sendUpdatePlaceholder( - this.tabID, - i18n('AWS.amazonq.featureDev.pillText.contextGatheringCompleted') - ) + action.messenger.sendUpdatePlaceholder( + this.tabID, + i18n('AWS.amazonq.featureDev.pillText.contextGatheringCompleted') + ) + } return uploadId }) this.uploadId = uploadId + const nextState = new CodeGenState( - { ...this.config, uploadId }, + { ...this.config, uploadId: this.uploadId, currentCodeGenerationId: this.currentCodeGenerationId }, this.filePaths, this.deletedFiles, this.references, diff --git a/packages/core/src/amazonqFeatureDev/types.ts b/packages/core/src/amazonqFeatureDev/types.ts index 75564350882..cf046743425 100644 --- a/packages/core/src/amazonqFeatureDev/types.ts +++ b/packages/core/src/amazonqFeatureDev/types.ts @@ -21,6 +21,7 @@ export type Interaction = { export interface SessionStateInteraction { nextState: SessionState | Omit | undefined interaction: Interaction + currentCodeGenerationId?: string } export enum DevPhase { @@ -60,7 +61,9 @@ export interface SessionState { readonly references?: CodeReference[] readonly phase?: SessionStatePhase readonly uploadId: string - readonly tokenSource: CancellationTokenSource + readonly currentIteration?: number + currentCodeGenerationId?: string + tokenSource?: CancellationTokenSource readonly codeGenerationId?: string readonly tabID: string interact(action: SessionStateAction): Promise @@ -76,6 +79,7 @@ export interface SessionStateConfig { conversationId: string proxyClient: FeatureDevClient uploadId: string + currentCodeGenerationId?: string } export interface SessionStateAction { @@ -85,6 +89,7 @@ export interface SessionStateAction { fs: VirtualFileSystem telemetry: TelemetryHelper uploadHistory?: UploadHistory + tokenSource?: CancellationTokenSource } export type NewFileZipContents = { zipFilePath: string; fileContent: string } diff --git a/packages/core/src/test/amazonqFeatureDev/controllers/chat/controller.test.ts b/packages/core/src/test/amazonqFeatureDev/controllers/chat/controller.test.ts index 2cba1816ffd..e7e8ecdb6fc 100644 --- a/packages/core/src/test/amazonqFeatureDev/controllers/chat/controller.test.ts +++ b/packages/core/src/test/amazonqFeatureDev/controllers/chat/controller.test.ts @@ -508,4 +508,15 @@ describe('Controller', () => { }) }) }) + + describe('stopResponse', () => { + it('should emit ui_click telemetry with elementId amazonq_stopCodeGeneration', async () => { + const getSessionStub = sinon.stub(controllerSetup.sessionStorage, 'getSession').resolves(session) + controllerSetup.emitters.stopResponse.fire({ tabID, conversationID }) + await waitUntil(() => { + return Promise.resolve(getSessionStub.callCount > 0) + }, {}) + assertTelemetry('ui_click', { elementId: 'amazonq_stopCodeGeneration' }) + }) + }) }) diff --git a/packages/core/src/test/amazonqFeatureDev/session/sessionState.test.ts b/packages/core/src/test/amazonqFeatureDev/session/sessionState.test.ts index c62266364f8..ec9cf8ea396 100644 --- a/packages/core/src/test/amazonqFeatureDev/session/sessionState.test.ts +++ b/packages/core/src/test/amazonqFeatureDev/session/sessionState.test.ts @@ -38,10 +38,12 @@ const mockSessionStateConfig = ({ conversationId, uploadId, workspaceFolder, + currentCodeGenerationId, }: { conversationId: string uploadId: string workspaceFolder: vscode.WorkspaceFolder + currentCodeGenerationId?: string }): SessionStateConfig => ({ workspaceRoots: ['fake-source'], workspaceFolders: [workspaceFolder], @@ -54,12 +56,14 @@ const mockSessionStateConfig = ({ exportResultArchive: () => mockExportResultArchive(), } as unknown as FeatureDevClient, uploadId, + currentCodeGenerationId, }) describe('sessionState', () => { const conversationId = 'conversation-id' const uploadId = 'upload-id' const tabId = 'tab-id' + const currentCodeGenerationId = '' let testConfig: SessionStateConfig beforeEach(async () => { @@ -67,6 +71,7 @@ describe('sessionState', () => { conversationId, uploadId, workspaceFolder: await createTestWorkspaceFolder('fake-root'), + currentCodeGenerationId, }) }) @@ -107,18 +112,18 @@ describe('sessionState', () => { codeGenerationRemainingIterationCount: 2, codeGenerationTotalIterationCount: 3, }) + mockExportResultArchive = sinon.stub().resolves({ newFileContents: [], deletedFiles: [], references: [] }) const testAction = mockSessionStateAction() const state = new CodeGenState(testConfig, [], [], [], tabId, 0, {}, 2, 3) const result = await state.interact(testAction) - const nextState = new PrepareCodeGenState(testConfig, [], [], [], tabId, 1, 2, 3) + const nextState = new PrepareCodeGenState(testConfig, [], [], [], tabId, 1, 2, 3, undefined) - assert.deepStrictEqual(result, { - nextState, - interaction: {}, - }) + assert.deepStrictEqual(result.nextState?.deletedFiles, nextState.deletedFiles) + assert.deepStrictEqual(result.nextState?.filePaths, result.nextState?.filePaths) + assert.deepStrictEqual(result.nextState?.references, result.nextState?.references) }) it('fails when codeGenerationStatus failed ', async () => {