featureDev: Handle invalidStateEvent during planning

Author: jsalles

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

@@ -18,6 +18,7 @@ import { CodeReference } from '../../amazonq/webview/ui/connector'
 import { ApiError, ContentLengthError, UnknownApiError } from '../errors'
 import { ToolkitError, isAwsError, isCodeWhispererStreamingServiceException } from '../../shared/errors'
 import { getCodewhispererConfig } from '../../codewhisperer/client/codewhisperer'
+import { LLMResponseType } from '../types'
 
 // Create a client for featureDev proxy client based off of aws sdk v2
 export async function createFeatureDevProxyClient(): Promise<FeatureDevProxyClient> {
@@ -136,7 +137,13 @@ export class FeatureDevClient {
             throw new UnknownApiError(e instanceof Error ? e.message : 'Unknown error', 'CreateUploadUrl')
         }
     }
-    public async generatePlan(conversationId: string, uploadId: string, userMessage: string) {
+    public async generatePlan(
+        conversationId: string,
+        uploadId: string,
+        userMessage: string
+    ): Promise<
+        { responseType: 'EMPTY'; approach?: string } | { responseType: 'VALID' | 'INVALID_STATE'; approach: string }
+    > {
         try {
             const streamingClient = await this.getStreamingClient()
             const params = {
@@ -155,8 +162,9 @@ export class FeatureDevClient {
             getLogger().debug(`${featureName}: Generated plan: %O`, {
                 requestId: response.$metadata.requestId,
             })
+            let responseType: LLMResponseType = 'EMPTY'
             if (!response.planningResponseStream) {
-                return undefined
+                return { responseType }
             }
 
             const assistantResponse: string[] = []
@@ -167,12 +175,14 @@ export class FeatureDevClient {
                     getLogger().debug('Received Invalid State Event: %O', responseItem.invalidStateEvent)
                     assistantResponse.splice(0)
                     assistantResponse.push(responseItem.invalidStateEvent.message ?? '')
+                    responseType = 'INVALID_STATE'
                     break
                 } else if (responseItem.assistantResponseEvent !== undefined) {
+                    responseType = 'VALID'
                     assistantResponse.push(responseItem.assistantResponseEvent.content ?? '')
                 }
             }
-            return assistantResponse.join('')
+            return { responseType, approach: assistantResponse.join('') }
         } catch (e) {
             if (isCodeWhispererStreamingServiceException(e)) {
                 getLogger().error(

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

@@ -262,18 +262,20 @@ export class FeatureDevController {
             canBeVoted: true,
         })
 
-        this.messenger.sendAnswer({
-            type: 'answer',
-            tabID,
-            message: `Would you like to generate a suggestion for this? You’ll review a file diff before inserting into your project.`,
-        })
+        if (interactions.responseType === 'VALID') {
+            this.messenger.sendAnswer({
+                type: 'answer',
+                tabID,
+                message: `Would you like to generate a suggestion for this? You’ll review a file diff before inserting into your project.`,
+            })
 
-        // Follow up with action items and complete the request stream
-        this.messenger.sendAnswer({
-            type: 'system-prompt', // show the followups on the right side
-            followUps: this.getFollowUpOptions(session.state.phase),
-            tabID: tabID,
-        })
+            // Follow up with action items and complete the request stream
+            this.messenger.sendAnswer({
+                type: 'system-prompt', // show the followups on the right side
+                followUps: this.getFollowUpOptions(session.state.phase),
+                tabID: tabID,
+            })
+        }
 
         // Unlock the prompt again so that users can iterate
         this.messenger.sendAsyncEventProgress(tabID, false, undefined)

packages/core/src/amazonqFeatureDev/session/sessionState.ts

@@ -115,7 +115,7 @@ export class RefinementState implements SessionState {
                     throw new UserMessageNotFoundError()
                 }
 
-                const approach = await this.config.proxyClient.generatePlan(
+                const { responseType, approach } = await this.config.proxyClient.generatePlan(
                     this.config.conversationId,
                     this.config.uploadId,
                     action.msg
@@ -126,7 +126,7 @@ export class RefinementState implements SessionState {
                         'There has been a problem generating an approach. Please open a conversation in a new tab'
                 )
 
-                action.telemetry.recordUserApproachTelemetry(span, this.conversationId)
+                action.telemetry.recordUserApproachTelemetry(span, this.conversationId, responseType)
                 return {
                     nextState: new RefinementState(
                         {
@@ -139,6 +139,7 @@ export class RefinementState implements SessionState {
                     ),
                     interaction: {
                         content: `${this.approach}\n`,
+                        responseType,
                     },
                 }
             } catch (e) {

packages/core/src/amazonqFeatureDev/types.ts

@@ -15,6 +15,7 @@ import { CodeReference } from '../amazonq/webview/ui/connector'
 export type Interaction = {
     // content to be sent back to the chat UI
     content?: string
+    responseType?: LLMResponseType
 }
 
 export interface SessionStateInteraction {
@@ -95,3 +96,5 @@ export function createUri(filePath: string, tabID?: string) {
         ...(tabID ? { query: `tabID=${tabID}` } : {}),
     })
 }
+
+export type LLMResponseType = 'EMPTY' | 'INVALID_STATE' | 'VALID'

packages/core/src/amazonqFeatureDev/util/telemetryHelper.ts

@@ -5,6 +5,7 @@
 
 import { getLogger } from '../../shared/logger/logger'
 import { AmazonqApproachInvoke, AmazonqCodeGenerationInvoke, Metric } from '../../shared/telemetry/telemetry'
+import { LLMResponseType } from '../types'
 
 const performance = globalThis.performance ?? require('perf_hooks').performance
 
@@ -31,11 +32,16 @@ export class TelemetryHelper {
         this.sessionStartTime = performance.now()
     }
 
-    public recordUserApproachTelemetry(span: Metric<AmazonqApproachInvoke>, amazonqConversationId: string) {
+    public recordUserApproachTelemetry(
+        span: Metric<AmazonqApproachInvoke>,
+        amazonqConversationId: string,
+        responseType: LLMResponseType
+    ) {
         const event = {
             amazonqConversationId,
             amazonqGenerateApproachIteration: this.generateApproachIteration,
             amazonqGenerateApproachLatency: performance.now() - this.generateApproachLastInvocationTime,
+            amazonqGenerateApproachResponseType: responseType,
         }
         getLogger().debug(`recordUserApproachTelemetry: %O`, event)
         span.record(event)

packages/core/src/test/amazonqFeatureDev/session/sessionState.test.ts

@@ -118,14 +118,15 @@ describe('sessionState', () => {
         it('transitions to RefinementState and returns an approach', async () => {
             sinon.stub(performance, 'now').returns(0)
 
-            mockGeneratePlan = sinon.stub().resolves(testApproach)
+            mockGeneratePlan = sinon.stub().resolves({ responseType: 'TEST_RESPONSE_TYPE', approach: testApproach })
             const state = new RefinementState(testConfig, testApproach, tabId, 0)
             const result = await state.interact(testAction)
 
             assert.deepStrictEqual(result, {
                 nextState: new RefinementState(testConfig, testApproach, tabId, 1),
                 interaction: {
                     content: `${testApproach}\n`,
+                    responseType: 'TEST_RESPONSE_TYPE',
                 },
             })
 
@@ -138,7 +139,7 @@ describe('sessionState', () => {
         })
 
         it('transitions to RefinementState but does not return an approach', async () => {
-            mockGeneratePlan = sinon.stub().resolves(undefined)
+            mockGeneratePlan = sinon.stub().resolves({ responseType: 'TEST_RESPONSE_TYPE', approach: undefined })
             const state = new RefinementState(testConfig, testApproach, tabId, 0)
             const result = await state.interact(testAction)
             const invokeFailureApproach =
@@ -148,14 +149,17 @@ describe('sessionState', () => {
                 nextState: new RefinementState(testConfig, invokeFailureApproach, tabId, 1),
                 interaction: {
                     content: `${invokeFailureApproach}\n`,
+                    responseType: 'TEST_RESPONSE_TYPE',
                 },
             })
         })
 
         it('invalid html gets sanitized', async () => {
             const invalidHTMLApproach =
                 '<head><script src="https://foo"></script></head><body><h1>hello world</h1></body>'
-            mockGeneratePlan = sinon.stub().resolves(invalidHTMLApproach)
+            mockGeneratePlan = sinon
+                .stub()
+                .resolves({ responseType: 'TEST_RESPONSE_TYPE', approach: invalidHTMLApproach })
             const state = new RefinementState(testConfig, invalidHTMLApproach, tabId, 0)
             const result = await state.interact(testAction)
 
@@ -165,6 +169,7 @@ describe('sessionState', () => {
                 nextState: new RefinementState(testConfig, expectedApproach, tabId, 1),
                 interaction: {
                     content: `${expectedApproach}\n`,
+                    responseType: 'TEST_RESPONSE_TYPE',
                 },
             })
         })