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

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

@@ -375,6 +375,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",
@@ -384,7 +385,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>",
@@ -396,6 +397,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,61 @@ 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 {
-    constructor() {
-        super(i18n('AWS.amazonq.doc.error.readmeTooLarge'), {
-            code: ReadmeTooLargeError.name,
-        })
+export class ReadmeTooLargeError extends DocServiceError {
+    constructor(remainingIterations: number) {
+        super(i18n('AWS.amazonq.doc.error.readmeTooLarge'), ReadmeTooLargeError.name, remainingIterations)
     }
 }
 
-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 {
-    constructor() {
-        super(i18n('AWS.amazonq.doc.error.workspaceEmpty'), {
-            code: WorkspaceEmptyError.name,
-        })
+export class WorkspaceEmptyError extends DocServiceError {
+    constructor(remainingIterations: number) {
+        super(i18n('AWS.amazonq.doc.error.workspaceEmpty'), WorkspaceEmptyError.name, remainingIterations)
     }
 }
 
-export class NoChangeRequiredException extends ToolkitError {
-    constructor() {
-        super(i18n('AWS.amazonq.doc.error.noChangeRequiredException'), {
-            code: NoChangeRequiredException.name,
-        })
+export class NoChangeRequiredException extends DocServiceError {
+    constructor(remainingIterations: number) {
+        super(
+            i18n('AWS.amazonq.doc.error.noChangeRequiredException'),
+            NoChangeRequiredException.name,
+            remainingIterations
+        )
     }
 }
 
-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 {
-    constructor() {
-        super(i18n('AWS.amazonq.doc.error.contentLengthError'), { code: ContentLengthError.name })
+export class ContentLengthError extends DocServiceError {
+    constructor(remainingIterations: number) {
+        super(i18n('AWS.amazonq.doc.error.contentLengthError'), ContentLengthError.name, remainingIterations)
     }
 }
 
-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)
     }
 }