diff --git a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts b/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts index 75b93ff47bf..f8f08f71148 100644 --- a/packages/core/src/codewhispererChat/clients/chat/v0/chat.ts +++ b/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 = 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() { diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts index a56bf5c811e..4f58b3ab85a 100644 --- a/packages/core/src/codewhispererChat/controllers/chat/controller.ts +++ b/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 ?? '',