use session in chatHistory

Author: chungjac

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

@@ -785,10 +785,8 @@ export class ChatController {
         const promptTypeValue = message.optionsValues['prompt-type']
         // TODO: display message: You turned off pair programmer mode. Q will not include code diffs or run commands in the chat.
         if (promptTypeValue === 'pair-programming-on') {
-            this.chatHistoryStorage.setTools(message.tabID, tools)
             session.setPairProgrammingModeOn(true)
         } else {
-            this.chatHistoryStorage.setTools(message.tabID, noWriteTools)
             session.setPairProgrammingModeOn(false)
         }
     }

packages/core/src/codewhispererChat/storages/chatHistory.ts

@@ -4,15 +4,15 @@
  */
 import {
     ChatMessage,
-    Tool,
     ToolResult,
     ToolResultStatus,
     UserInputMessage,
     UserInputMessageContext,
 } from '@amzn/codewhisperer-streaming'
 import { randomUUID } from '../../shared/crypto'
 import { getLogger } from '../../shared/logger/logger'
-import { tools } from '../constants'
+import { tools, noWriteTools } from '../constants'
+import { ChatSessionStorage } from './chatSession'
 
 // Maximum number of messages to keep in history
 const MaxConversationHistoryLength = 100
@@ -27,12 +27,12 @@ export class ChatHistoryManager {
     private history: ChatMessage[] = []
     private logger = getLogger()
     private lastUserMessage?: ChatMessage
-    private tools: Tool[] = []
+    private sessionStorage: ChatSessionStorage
 
     constructor(tabId?: string) {
         this.conversationId = randomUUID()
         this.tabId = tabId ?? randomUUID()
-        this.tools = tools
+        this.sessionStorage = new ChatSessionStorage()
     }
 
     /**
@@ -46,10 +46,6 @@ export class ChatHistoryManager {
         this.conversationId = conversationId
     }
 
-    public setTools(tools: Tool[]) {
-        this.tools = tools
-    }
-
     /**
      * Get the tab ID
      */
@@ -182,11 +178,14 @@ export class ChatHistoryManager {
                             status: ToolResultStatus.ERROR,
                         }))
 
+                        const session = this.sessionStorage.getSession(this.tabId)
+                        const availableTools = session.pairProgrammingModeOn ? tools : noWriteTools
+
                         newUserMessage.userInputMessage.userInputMessageContext = {
                             shellState: undefined,
                             envState: undefined,
                             toolResults: toolResults,
-                            tools: this.tools.length === 0 ? undefined : [...this.tools],
+                            tools: availableTools.length === 0 ? undefined : [...availableTools],
                         }
 
                         return newUserMessage
@@ -203,11 +202,14 @@ export class ChatHistoryManager {
      * Adds tool results to the conversation.
      */
     addToolResults(toolResults: ToolResult[]): void {
+        const session = this.sessionStorage.getSession(this.tabId)
+        const availableTools = session.pairProgrammingModeOn ? tools : noWriteTools
+
         const userInputMessageContext: UserInputMessageContext = {
             shellState: undefined,
             envState: undefined,
             toolResults: toolResults,
-            tools: this.tools.length === 0 ? undefined : [...this.tools],
+            tools: availableTools.length === 0 ? undefined : [...availableTools],
         }
 
         const msg: UserInputMessage = {
@@ -252,12 +254,15 @@ export class ChatHistoryManager {
                     status: ToolResultStatus.ERROR,
                 }))
 
+                const session = this.sessionStorage.getSession(this.tabId)
+                const availableTools = session.pairProgrammingModeOn ? tools : noWriteTools
+
                 // Create a new user message with cancelled tool results
                 const userInputMessageContext: UserInputMessageContext = {
                     shellState: undefined,
                     envState: undefined,
                     toolResults: toolResults,
-                    tools: this.tools.length === 0 ? undefined : [...this.tools],
+                    tools: availableTools.length === 0 ? undefined : [...availableTools],
                 }
 
                 const userMessage: ChatMessage = {

packages/core/src/codewhispererChat/storages/chatHistoryStorage.ts

@@ -3,7 +3,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { Tool } from '@amzn/codewhisperer-streaming'
 import { ChatHistoryManager } from './chatHistory'
 
 /**
@@ -41,8 +40,4 @@ export class ChatHistoryStorage {
     public deleteHistory(tabId: string) {
         this.histories.delete(tabId)
     }
-
-    public setTools(tabId: string, tools: Tool[]) {
-        this.histories.get(tabId)?.setTools(tools)
-    }
 }