feat(featureDev): handle iteration limit errors during planning / code generation #4559

Problem
We want to handle iteration limit errors for plan generation and code generation for featureDev; currently the user experience is impacted as proper prompts in chat are not surfaced, instead Retry is prompted incorrectly.

Solution
Added handling of iteration limit errors during both planning phase and code generation phase, to show specific prompts in planning and code generation phase. This will unblock customers in case of limit errors.

Author: kumsmrit

.changes/next-release/Bug Fix-1777c060-e176-4a85-9527-3833250136e3.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Amazon Q Feature Dev: Fix followups after hitting iteration limit"
+}

packages/core/src/amazonqFeatureDev/client/featureDev.ts

@@ -13,7 +13,13 @@ 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 { ApiError, ContentLengthError, UnknownApiError } from '../errors'
+import {
+    ApiError,
+    CodeIterationLimitError,
+    ContentLengthError,
+    PlanIterationLimitError,
+    UnknownApiError,
+} from '../errors'
 import { ToolkitError, isAwsError, isCodeWhispererStreamingServiceException } from '../../shared/errors'
 import { getCodewhispererConfig } from '../../codewhisperer/client/codewhisperer'
 import { LLMResponseType } from '../types'
@@ -168,6 +174,13 @@ export class FeatureDevClient {
                         e.$metadata.requestId ?? 'unknown'
                     }`
                 )
+                if (
+                    (e.name === 'ThrottlingException' &&
+                        e.message.includes('limit for number of iterations on an implementation plan')) ||
+                    e.name === 'ServiceQuotaExceededException'
+                ) {
+                    throw new PlanIterationLimitError()
+                }
                 throw new ApiError(
                     e.message,
                     'GeneratePlan',
@@ -206,6 +219,14 @@ export class FeatureDevClient {
                     (e as any).requestId
                 }`
             )
+            if (
+                isAwsError(e) &&
+                ((e.code === 'ThrottlingException' &&
+                    e.message.includes('limit for number of iterations on a code generation')) ||
+                    e.code === 'ServiceQuotaExceededException')
+            ) {
+                throw new CodeIterationLimitError()
+            }
             throw new ToolkitError((e as Error).message, { code: 'StartCodeGenerationFailed' })
         }
     }

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

@@ -11,7 +11,13 @@ import { EventEmitter } from 'vscode'
 import { telemetry } from '../../../shared/telemetry/telemetry'
 import { createSingleFileDialog } from '../../../shared/ui/common/openDialog'
 import { featureDevScheme } from '../../constants'
-import { ContentLengthError, SelectedFolderNotInWorkspaceFolderError, createUserFacingErrorMessage } from '../../errors'
+import {
+    CodeIterationLimitError,
+    ContentLengthError,
+    PlanIterationLimitError,
+    SelectedFolderNotInWorkspaceFolderError,
+    createUserFacingErrorMessage,
+} from '../../errors'
 import { defaultRetryLimit } from '../../limits'
 import { Session } from '../../session/session'
 import { featureName } from '../../constants'
@@ -236,6 +242,38 @@ export class FeatureDevController {
                         },
                     ],
                 })
+            } else if (err instanceof PlanIterationLimitError) {
+                this.messenger.sendErrorMessage(err.message, message.tabID, this.retriesRemaining(session))
+                this.messenger.sendAnswer({
+                    type: 'system-prompt',
+                    tabID: message.tabID,
+                    followUps: [
+                        {
+                            pillText: 'Discuss a new plan',
+                            type: FollowUpTypes.NewTask,
+                            status: 'info',
+                        },
+                        {
+                            pillText: 'Generate code',
+                            type: FollowUpTypes.GenerateCode,
+                            status: 'info',
+                        },
+                    ],
+                })
+            } else if (err instanceof CodeIterationLimitError) {
+                this.messenger.sendErrorMessage(err.message, message.tabID, this.retriesRemaining(session))
+                this.messenger.sendAnswer({
+                    type: 'system-prompt',
+                    tabID: message.tabID,
+                    followUps: [
+                        {
+                            pillText: 'Insert code',
+                            type: FollowUpTypes.InsertCode,
+                            icon: 'ok' as MynahIcons,
+                            status: 'success',
+                        },
+                    ],
+                })
             } else {
                 const errorMessage = createUserFacingErrorMessage(
                     `${featureName} request failed: ${err.cause?.message ?? err.message}`

packages/core/src/amazonqFeatureDev/errors.ts

@@ -76,6 +76,24 @@ export class ContentLengthError extends ToolkitError {
     }
 }
 
+export class PlanIterationLimitError extends ToolkitError {
+    constructor() {
+        super(
+            'You have reached the free tier limit for number of iterations on an implementation plan. Please proceed to generating code or start to discuss a new plan.',
+            { code: 'PlanIterationLimitError' }
+        )
+    }
+}
+
+export class CodeIterationLimitError extends ToolkitError {
+    constructor() {
+        super(
+            'You have reached the free tier limit for number of iterations on a code generation. Please proceed to accept the code or start a new conversation.',
+            { code: 'CodeIterationLimitError' }
+        )
+    }
+}
+
 export class UnknownApiError extends ToolkitError {
     constructor(message: string, api: string) {
         super(message, { code: `${api}-Unknown` })