fix(chat): preserve context in agentic loop

Author: ctlai95

packages/core/src/codewhispererChat/clients/chat/v0/chat.ts

@@ -14,16 +14,19 @@ import { ToolkitError } from '../../../../shared/errors'
 import { createCodeWhispererChatStreamingClient } from '../../../../shared/clients/codewhispererChatClient'
 import { createQDeveloperStreamingClient } from '../../../../shared/clients/qDeveloperChatClient'
 import { UserWrittenCodeTracker } from '../../../../codewhisperer/tracker/userWrittenCodeTracker'
+import { PromptMessage } from '../../../controllers/chat/model'
 
 export class ChatSession {
     private sessionId?: string
     /**
      * _readFiles = list of files read from the project to gather context before generating response.
      * _showDiffOnFileWrite = Controls whether to show diff view (true) or file context view (false) to the user
+     * _context = Additional context to be passed to the LLM for generating the response
      */
     private _readFiles: string[] = []
     private _toolUse: ToolUse | undefined
     private _showDiffOnFileWrite: boolean = false
+    private _context: PromptMessage['context']
 
     contexts: Map<string, { first: number; second: number }[]> = new Map()
     // TODO: doesn't handle the edge case when two files share the same relativePath string but from different root
@@ -41,6 +44,14 @@ export class ChatSession {
         this._toolUse = toolUse
     }
 
+    public get context(): PromptMessage['context'] {
+        return this._context
+    }
+
+    public setContext(context: PromptMessage['context']) {
+        this._context = context
+    }
+
     public tokenSource!: vscode.CancellationTokenSource
 
     constructor() {

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

@@ -709,7 +709,7 @@ export class ChatController {
                         toolResults: toolResults,
                         origin: Origin.IDE,
                         chatHistory: this.chatHistoryManager.getHistory(),
-                        context: [],
+                        context: session.context ?? [],
                         relevantTextDocuments: [],
                         additionalContents: [],
                         documentReferences: [],
@@ -1081,6 +1081,10 @@ export class ChatController {
                     type: 'chat_message',
                     context,
                 })
+
+                // Save the context for the agentic loop
+                session.setContext(message.context)
+
                 await this.generateResponse(
                     {
                         message: message.message ?? '',