fix(chat): only show context on first message

Author: ctlai95

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

@@ -1177,9 +1177,6 @@ export class ChatController {
 
                 this.messenger.sendAsyncEventProgress(message.tabID, true, '')
 
-                // Save the context for the agentic loop
-                session.setContext(message.context)
-
                 await this.generateResponse(
                     {
                         message: message.message ?? '',
@@ -1480,7 +1477,12 @@ export class ChatController {
         let response: MessengerResponseType | undefined = undefined
         session.createNewTokenSource()
         try {
-            this.messenger.sendInitalStream(tabID, triggerID, triggerPayload.documentReferences)
+            if (!session.context) {
+                // Only show context for the first message in the loop
+                this.messenger.sendContextMessage(tabID, triggerID, triggerPayload.documentReferences)
+                session.setContext(triggerPayload.context)
+            }
+            this.messenger.sendInitalStream(tabID, triggerID)
             this.telemetryHelper.setConversationStreamStartTime(tabID)
             if (isSsoConnection(AuthUtil.instance.conn)) {
                 const { $metadata, generateAssistantResponseResponse } = await session.chatSso(request)

packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts

@@ -103,7 +103,30 @@ export class Messenger {
         )
     }
 
-    public sendInitalStream(
+    public sendInitalStream(tabID: string, triggerID: string) {
+        this.dispatcher.sendChatMessage(
+            new ChatMessage(
+                {
+                    message: '',
+                    messageType: 'answer-stream',
+                    followUps: undefined,
+                    followUpsHeader: undefined,
+                    relatedSuggestions: undefined,
+                    triggerID,
+                    messageID: '',
+                    userIntent: undefined,
+                    codeBlockLanguage: undefined,
+                    contextList: undefined,
+                    buttons: undefined,
+                    fileList: undefined,
+                    canBeVoted: false,
+                },
+                tabID
+            )
+        )
+    }
+
+    public sendContextMessage(
         tabID: string,
         triggerID: string,
         mergedRelevantDocuments: DocumentReference[] | undefined
@@ -112,7 +135,7 @@ export class Messenger {
             new ChatMessage(
                 {
                     message: '',
-                    messageType: 'answer-stream',
+                    messageType: 'answer',
                     followUps: undefined,
                     followUpsHeader: undefined,
                     relatedSuggestions: undefined,
@@ -125,11 +148,13 @@ export class Messenger {
                     buttons: undefined,
                     fileList: undefined,
                     canBeVoted: false,
+                    padding: false,
                 },
                 tabID
             )
         )
     }
+
     /**
      * Tries to calculate the total number of code blocks.
      * NOTES:
@@ -455,6 +480,11 @@ export class Messenger {
                             toolUse.input !== '' && { toolUses: [{ ...toolUse }] }),
                     },
                 })
+                const agenticLoopEnded = !eventCounts.has('toolUseEvent')
+                if (agenticLoopEnded) {
+                    // Reset context for the next request
+                    session.setContext(undefined)
+                }
 
                 getLogger().info(
                     `All events received. requestId=%s counts=%s`,

packages/core/src/codewhispererChat/tools/chatStream.ts

@@ -28,7 +28,7 @@ export class ChatStream extends Writable {
     ) {
         super()
         this.logger.debug(`ChatStream created for tabID: ${tabID}, triggerID: ${triggerID}`)
-        this.messenger.sendInitalStream(tabID, triggerID, undefined)
+        this.messenger.sendInitalStream(tabID, triggerID)
     }
 
     override _write(chunk: Buffer, encoding: BufferEncoding, callback: (error?: Error | null) => void): void {