From 7165245e3c8b30271bf366c820a083494eeaf062 Mon Sep 17 00:00:00 2001
From: Tai Lai <tlaiamz@amazon.com>
Date: Mon, 31 Mar 2025 17:04:02 -0700
Subject: [PATCH] fix(chat): preserve context in agentic loop

---
 .../src/codewhispererChat/clients/chat/v0/chat.ts     | 11 +++++++++++
 .../codewhispererChat/controllers/chat/controller.ts  |  6 +++++-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts b/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts
index 05a23403dda..c7e309d2350 100644
--- a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts
+++ b/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts
@@ -14,6 +14,7 @@ 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
@@ -22,12 +23,14 @@ export class ChatSession {
      * _filePath = The path helps the system locate exactly where to make the necessary changes in the project structure
      * _tempFilePath = Used to show the code diff view in the editor including LLM changes.
      * _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 _filePath: string | undefined
     private _tempFilePath: string | undefined
     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
@@ -45,6 +48,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() {
diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts
index ee21f252ffb..ec05d2d492a 100644
--- a/packages/core/src/codewhispererChat/controllers/chat/controller.ts
+++ b/packages/core/src/codewhispererChat/controllers/chat/controller.ts
@@ -1010,7 +1010,7 @@ export class ChatController {
                         toolResults: toolResults,
                         origin: Origin.IDE,
                         chatHistory: this.chatHistoryManager.getHistory(),
-                        context: [],
+                        context: session.context ?? [],
                         relevantTextDocuments: [],
                         additionalContents: [],
                         documentReferences: [],
@@ -1042,6 +1042,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 ?? '',