fix(amazonq): enable input after prompt is rejected, prevent iteration limit exceeded (aws#6395)

## Problem
When a user encounters an error while updating a README, they see a
"Retry" button, which isn't correct because the previous prompt is not
retried. Instead, the user should be able to input a new prompt.

When a user runs out of iterations in a session, the chat allows them to
continue making README update requests, which results in an "Iteration
limit exceeded error".

## Solution
If a user gets an error while updating a README, the chat input is
enabled and the user can try another prompt.

Once the user hits their iteration limit for the session, they are not
able to suggest more changes and have to accept/reject the existing set
of changes. This will reduce the "iteration limit exceeded" errors
encountered by users.


A E2E test is added to verify that prompt errors during README updates
are handled correctly with the relevant error message, and validates
that users can enter a new prompt.


Screenshots:
_Prompt enabled after error updating README_
<img width="520" alt="image"
src="https://github.com/user-attachments/assets/407109d2-3f10-4e8c-9312-ecab57e625b4"
/>

_'Make changes' button hidden once iteration limit is reached_
<img width="506" alt="image"
src="https://github.com/user-attachments/assets/acf960c0-1c43-475c-a478-c95336ab857b"
/>

---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).
- License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

Author: avi-alpert

packages/amazonq/.changes/next-release/Bug Fix-11625969-6700-4d10-805f-b6eaea2287a2.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Amazon Q /doc: Prevent users from requesting changes if no iterations remain"
+}

packages/amazonq/.changes/next-release/Bug Fix-fdf7f991-1da6-405d-8a8a-3ce1af1db29d.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Amazon Q /doc: Ask for user prompt if error occurs while updating documentation"
+}

packages/amazonq/test/e2e/amazonq/doc.test.ts

@@ -146,5 +146,33 @@ describe('Amazon Q Doc', async function () {
                 FollowUpTypes.RejectChanges,
             ])
         })
+
+        it('Handle unrelated prompt error', async () => {
+            await tab.waitForButtons([FollowUpTypes.UpdateDocumentation])
+
+            tab.clickButton(FollowUpTypes.UpdateDocumentation)
+
+            await tab.waitForButtons([FollowUpTypes.SynchronizeDocumentation, FollowUpTypes.EditDocumentation])
+
+            tab.clickButton(FollowUpTypes.EditDocumentation)
+
+            await tab.waitForButtons([FollowUpTypes.ProceedFolderSelection])
+
+            tab.clickButton(FollowUpTypes.ProceedFolderSelection)
+
+            tab.addChatMessage({ prompt: 'tell me about the weather' })
+
+            await tab.waitForEvent(() =>
+                tab.getChatItems().some(({ body }) => body?.startsWith(i18n('AWS.amazonq.doc.error.promptUnrelated')))
+            )
+
+            await tab.waitForEvent(() => {
+                const store = tab.getStore()
+                return (
+                    !store.promptInputDisabledState &&
+                    store.promptInputPlaceholder === i18n('AWS.amazonq.doc.placeholder.editReadme')
+                )
+            })
+        })
     })
 })

packages/core/package.nls.json

@@ -376,6 +376,7 @@
     "AWS.amazonq.doc.answer.readmeCreated": "I've created a README for your code.",
     "AWS.amazonq.doc.answer.readmeUpdated": "I've updated your README.",
     "AWS.amazonq.doc.answer.codeResult": "You can accept the changes to your files, or describe any additional changes you'd like me to make.",
+    "AWS.amazonq.doc.answer.acceptOrReject": "You can accept or reject the changes to your files.",
     "AWS.amazonq.doc.answer.scanning": "Scanning source files",
     "AWS.amazonq.doc.answer.summarizing": "Summarizing source files",
     "AWS.amazonq.doc.answer.generating": "Generating documentation",
@@ -385,7 +386,7 @@
     "AWS.amazonq.doc.error.noFolderSelected": "It looks like you didn't choose a folder. Choose a folder to continue.",
     "AWS.amazonq.doc.error.contentLengthError": "Your workspace is too large for me to review. Your workspace must be within the quota, even if you choose a smaller folder. For more information on quotas, see the <a href=\"https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/doc-generation.html#quotas\" target=\"_blank\">Amazon Q Developer documentation.</a>",
     "AWS.amazonq.doc.error.readmeTooLarge": "The README in your folder is too large for me to review. Try reducing the size of your README, or choose a folder with a smaller README. For more information on quotas, see the <a href=\"https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/doc-generation.html#quotas\" target=\"_blank\">Amazon Q Developer documentation.</a>",
-    "AWS.amazonq.doc.error.readmeUpdateTooLarge": "The updated README is too large. Try reducing the size of your README, or asking for a smaller update. For more information on quotas, see the <a href=\"https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/doc-generation.html#quotas\" target=\"_blank\">Amazon Q Developer documentation.</a>",
+    "AWS.amazonq.doc.error.readmeUpdateTooLarge": "The updated README exceeds document size limits. Try reducing the size of your current README or working on a smaller task that won't produce as much content. For more information on quotas, see the <a href=\"https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/doc-generation.html#quotas\" target=\"_blank\">Amazon Q Developer documentation.</a>",
     "AWS.amazonq.doc.error.workspaceEmpty": "The folder you chose did not contain any source files in a supported language. Choose another folder and try again. For more information on supported languages, see the <a href=\"https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/doc-generation.html\" target=\"_blank\">Amazon Q Developer documentation.</a>",
     "AWS.amazonq.doc.error.promptTooVague": "I need more information to make changes to your README. Try providing some of the following details:\n- Which sections you want to modify\n- The content you want to add or remove\n- Specific issues that need correcting\n\nFor more information on prompt best practices, see the <a href=\"https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/doc-generation.html\" target=\"_blank\">Amazon Q Developer documentation.</a>",
     "AWS.amazonq.doc.error.promptUnrelated": "These changes don't seem related to documentation. Try describing your changes again, using the following best practices:\n- Changes should relate to how project functionality is reflected in the README\n- Content you refer to should be available in your codebase\n\n For more information on prompt best practices, see the <a href=\"https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/doc-generation.html\" target=\"_blank\">Amazon Q Developer documentation.</a>",
@@ -397,6 +398,9 @@
     "AWS.amazonq.doc.pillText.newTask": "Start a new documentation task",
     "AWS.amazonq.doc.pillText.update": "Update README to reflect code",
     "AWS.amazonq.doc.pillText.makeChange": "Make a specific change",
+    "AWS.amazonq.doc.pillText.accept": "Accept",
+    "AWS.amazonq.doc.pillText.reject": "Reject",
+    "AWS.amazonq.doc.pillText.makeChanges": "Make changes",
     "AWS.amazonq.inline.invokeChat": "Inline chat",
     "AWS.toolkit.lambda.walkthrough.quickpickTitle": "Application Builder Walkthrough",
     "AWS.toolkit.lambda.walkthrough.title": "Get started building your application",

packages/core/src/amazonq/commons/connector/baseMessenger.ts

@@ -85,6 +85,7 @@ export class Messenger {
             type: 'answer',
             tabID: tabID,
             message: i18n('AWS.amazonq.featureDev.error.monthlyLimitReached'),
+            disableChatInput: true,
         })
         this.sendUpdatePlaceholder(tabID, i18n('AWS.amazonq.featureDev.placeholder.chatInputDisabled'))
     }

packages/core/src/amazonqDoc/constants.ts

@@ -92,6 +92,43 @@ export const FolderSelectorFollowUps = [
     },
 ]
 
+export const CodeChangeFollowUps = [
+    {
+        pillText: i18n('AWS.amazonq.doc.pillText.accept'),
+        prompt: i18n('AWS.amazonq.doc.pillText.accept'),
+        type: FollowUpTypes.AcceptChanges,
+        icon: 'ok' as MynahIcons,
+        status: 'success' as Status,
+    },
+    {
+        pillText: i18n('AWS.amazonq.doc.pillText.makeChanges'),
+        prompt: i18n('AWS.amazonq.doc.pillText.makeChanges'),
+        type: FollowUpTypes.MakeChanges,
+        icon: 'refresh' as MynahIcons,
+        status: 'info' as Status,
+    },
+    {
+        pillText: i18n('AWS.amazonq.doc.pillText.reject'),
+        prompt: i18n('AWS.amazonq.doc.pillText.reject'),
+        type: FollowUpTypes.RejectChanges,
+        icon: 'cancel' as MynahIcons,
+        status: 'error' as Status,
+    },
+]
+
+export const NewSessionFollowUps = [
+    {
+        pillText: i18n('AWS.amazonq.doc.pillText.newTask'),
+        type: FollowUpTypes.NewTask,
+        status: 'info' as Status,
+    },
+    {
+        pillText: i18n('AWS.amazonq.doc.pillText.closeSession'),
+        type: FollowUpTypes.CloseSession,
+        status: 'info' as Status,
+    },
+]
+
 export const SynchronizeDocumentation = {
     pillText: i18n('AWS.amazonq.doc.pillText.update'),
     prompt: i18n('AWS.amazonq.doc.pillText.update'),

packages/core/src/amazonqDoc/controllers/chat/controller.ts

@@ -10,7 +10,9 @@ import {
     EditDocumentation,
     FolderSelectorFollowUps,
     Mode,
+    NewSessionFollowUps,
     SynchronizeDocumentation,
+    CodeChangeFollowUps,
     docScheme,
     featureName,
     findReadmePath,
@@ -22,7 +24,6 @@ import { Session } from '../../session/session'
 import { i18n } from '../../../shared/i18n-helper'
 import path from 'path'
 import { createSingleFileDialog } from '../../../shared/ui/common/openDialog'
-import { MynahIcons } from '@aws/mynah-ui'
 
 import {
     MonthlyConversationLimitError,
@@ -298,18 +299,7 @@ export class DocController {
                         tabID: data?.tabID,
                         disableChatInput: true,
                         message: 'Your changes have been discarded.',
-                        followUps: [
-                            {
-                                pillText: i18n('AWS.amazonq.doc.pillText.newTask'),
-                                type: FollowUpTypes.NewTask,
-                                status: 'info',
-                            },
-                            {
-                                pillText: i18n('AWS.amazonq.doc.pillText.closeSession'),
-                                type: FollowUpTypes.CloseSession,
-                                status: 'info',
-                            },
-                        ],
+                        followUps: NewSessionFollowUps,
                     })
                     break
                 case FollowUpTypes.ProceedFolderSelection:
@@ -412,13 +402,19 @@ export class DocController {
         const errorMessage = createUserFacingErrorMessage(`${err.cause?.message ?? err.message}`)
         // eslint-disable-next-line unicorn/no-null
         this.messenger.sendUpdatePromptProgress(message.tabID, null)
+        if (err.constructor.name === MonthlyConversationLimitError.name) {
+            this.messenger.sendMonthlyLimitError(message.tabID)
+        } else {
+            const enableUserInput = this.mode === Mode.EDIT && err.remainingIterations > 0
 
-        switch (err.constructor.name) {
-            case MonthlyConversationLimitError.name:
-                this.messenger.sendMonthlyLimitError(message.tabID)
-                break
-            default:
-                this.messenger.sendErrorMessage(errorMessage, message.tabID, 0, session?.conversationIdUnsafe, false)
+            this.messenger.sendErrorMessage(
+                errorMessage,
+                message.tabID,
+                0,
+                session?.conversationIdUnsafe,
+                false,
+                enableUserInput
+            )
         }
     }
 
@@ -427,8 +423,6 @@ export class DocController {
             await this.onDocsGeneration(session, message.message, message.tabID)
         } catch (err: any) {
             this.processErrorChatMessage(err, message, session)
-            // Lock the chat input until they explicitly click one of the follow ups
-            this.messenger.sendChatInputEnabled(message.tabID, false)
         }
     }
 
@@ -461,12 +455,8 @@ export class DocController {
             }
 
             await this.generateDocumentation({ message, session })
-            this.messenger.sendChatInputEnabled(message?.tabID, false)
-            this.messenger.sendUpdatePlaceholder(message.tabID, i18n('AWS.amazonq.doc.pillText.selectOption'))
         } catch (err: any) {
             this.processErrorChatMessage(err, message, session)
-            // Lock the chat input until they explicitly click one of the follow ups
-            this.messenger.sendChatInputEnabled(message.tabID, false)
         }
     }
 
@@ -590,40 +580,21 @@ export class DocController {
                 this.messenger.sendAnswer({
                     type: 'answer',
                     tabID: tabID,
-                    message: `${this.mode === Mode.CREATE ? i18n('AWS.amazonq.doc.answer.readmeCreated') : i18n('AWS.amazonq.doc.answer.readmeUpdated')} ${i18n('AWS.amazonq.doc.answer.codeResult')}`,
+                    message: `${this.mode === Mode.CREATE ? i18n('AWS.amazonq.doc.answer.readmeCreated') : i18n('AWS.amazonq.doc.answer.readmeUpdated')} ${remainingIterations > 0 ? i18n('AWS.amazonq.doc.answer.codeResult') : i18n('AWS.amazonq.doc.answer.acceptOrReject')}`,
                     disableChatInput: true,
                 })
-            }
 
-            this.messenger.sendAnswer({
-                message: undefined,
-                type: 'system-prompt',
-                disableChatInput: true,
-                followUps: [
-                    {
-                        pillText: 'Accept',
-                        prompt: 'Accept',
-                        type: FollowUpTypes.AcceptChanges,
-                        icon: 'ok' as MynahIcons,
-                        status: 'success',
-                    },
-                    {
-                        pillText: 'Make changes',
-                        prompt: 'Make changes',
-                        type: FollowUpTypes.MakeChanges,
-                        icon: 'refresh' as MynahIcons,
-                        status: 'info',
-                    },
-                    {
-                        pillText: 'Reject',
-                        prompt: 'Reject',
-                        type: FollowUpTypes.RejectChanges,
-                        icon: 'cancel' as MynahIcons,
-                        status: 'error',
-                    },
-                ],
-                tabID: tabID,
-            })
+                this.messenger.sendAnswer({
+                    message: undefined,
+                    type: 'system-prompt',
+                    disableChatInput: true,
+                    followUps:
+                        remainingIterations > 0
+                            ? CodeChangeFollowUps
+                            : CodeChangeFollowUps.filter((followUp) => followUp.type !== FollowUpTypes.MakeChanges),
+                    tabID: tabID,
+                })
+            }
         } finally {
             if (session?.state?.tokenSource?.token.isCancellationRequested) {
                 await this.newTask({ tabID })
@@ -642,10 +613,8 @@ export class DocController {
             type: 'answer',
             tabID: message.tabID,
             message: 'Follow instructions to re-authenticate ...',
+            disableChatInput: true,
         })
-
-        // Explicitly ensure the user goes through the re-authenticate flow
-        this.messenger.sendChatInputEnabled(message.tabID, false)
     }
 
     private tabClosed(message: any) {
@@ -670,18 +639,7 @@ export class DocController {
                 type: 'answer',
                 disableChatInput: true,
                 tabID: message.tabID,
-                followUps: [
-                    {
-                        pillText: i18n('AWS.amazonq.doc.pillText.newTask'),
-                        type: FollowUpTypes.NewTask,
-                        status: 'info',
-                    },
-                    {
-                        pillText: i18n('AWS.amazonq.doc.pillText.closeSession'),
-                        type: FollowUpTypes.CloseSession,
-                        status: 'info',
-                    },
-                ],
+                followUps: NewSessionFollowUps,
             })
 
             this.messenger.sendUpdatePlaceholder(message.tabID, i18n('AWS.amazonq.doc.pillText.selectOption'))

packages/core/src/amazonqDoc/errors.ts

@@ -7,69 +7,57 @@ import { ToolkitError } from '../shared/errors'
 import { i18n } from '../shared/i18n-helper'
 
 export class DocServiceError extends ToolkitError {
-    constructor(message: string, code: string) {
+    remainingIterations?: number
+    constructor(message: string, code: string, remainingIterations?: number) {
         super(message, { code })
+        this.remainingIterations = remainingIterations
     }
 }
 
-export class ReadmeTooLargeError extends ToolkitError {
+export class ReadmeTooLargeError extends DocServiceError {
     constructor() {
-        super(i18n('AWS.amazonq.doc.error.readmeTooLarge'), {
-            code: ReadmeTooLargeError.name,
-        })
+        super(i18n('AWS.amazonq.doc.error.readmeTooLarge'), ReadmeTooLargeError.name)
     }
 }
 
-export class ReadmeUpdateTooLargeError extends ToolkitError {
-    constructor() {
-        super(i18n('AWS.amazonq.doc.error.readmeUpdateTooLarge'), {
-            code: ReadmeUpdateTooLargeError.name,
-        })
+export class ReadmeUpdateTooLargeError extends DocServiceError {
+    constructor(remainingIterations: number) {
+        super(i18n('AWS.amazonq.doc.error.readmeUpdateTooLarge'), ReadmeUpdateTooLargeError.name, remainingIterations)
     }
 }
 
-export class WorkspaceEmptyError extends ToolkitError {
+export class WorkspaceEmptyError extends DocServiceError {
     constructor() {
-        super(i18n('AWS.amazonq.doc.error.workspaceEmpty'), {
-            code: WorkspaceEmptyError.name,
-        })
+        super(i18n('AWS.amazonq.doc.error.workspaceEmpty'), WorkspaceEmptyError.name)
     }
 }
 
-export class NoChangeRequiredException extends ToolkitError {
+export class NoChangeRequiredException extends DocServiceError {
     constructor() {
-        super(i18n('AWS.amazonq.doc.error.noChangeRequiredException'), {
-            code: NoChangeRequiredException.name,
-        })
+        super(i18n('AWS.amazonq.doc.error.noChangeRequiredException'), NoChangeRequiredException.name)
     }
 }
 
-export class PromptRefusalException extends ToolkitError {
-    constructor() {
-        super(i18n('AWS.amazonq.doc.error.promptRefusal'), {
-            code: PromptRefusalException.name,
-        })
+export class PromptRefusalException extends DocServiceError {
+    constructor(remainingIterations: number) {
+        super(i18n('AWS.amazonq.doc.error.promptRefusal'), PromptRefusalException.name, remainingIterations)
     }
 }
 
-export class ContentLengthError extends ToolkitError {
+export class ContentLengthError extends DocServiceError {
     constructor() {
-        super(i18n('AWS.amazonq.doc.error.contentLengthError'), { code: ContentLengthError.name })
+        super(i18n('AWS.amazonq.doc.error.contentLengthError'), ContentLengthError.name)
     }
 }
 
-export class PromptTooVagueError extends ToolkitError {
-    constructor() {
-        super(i18n('AWS.amazonq.doc.error.promptTooVague'), {
-            code: PromptTooVagueError.name,
-        })
+export class PromptTooVagueError extends DocServiceError {
+    constructor(remainingIterations: number) {
+        super(i18n('AWS.amazonq.doc.error.promptTooVague'), PromptTooVagueError.name, remainingIterations)
     }
 }
 
-export class PromptUnrelatedError extends ToolkitError {
-    constructor() {
-        super(i18n('AWS.amazonq.doc.error.promptUnrelated'), {
-            code: PromptUnrelatedError.name,
-        })
+export class PromptUnrelatedError extends DocServiceError {
+    constructor(remainingIterations: number) {
+        super(i18n('AWS.amazonq.doc.error.promptUnrelated'), PromptUnrelatedError.name, remainingIterations)
     }
 }