fix(amazonq): honor agentic OFF state across chat lifecycle

server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

@@ -3,9 +3,9 @@
  * Will be deleted or merged.
  */
 
 import * as crypto from 'crypto'
 import * as path from 'path'
 import * as os from 'os'
 import {
     ChatTriggerType,
     Origin,
@@ -368,7 +368,13 @@
             features,
             this.#chatHistoryDb,
             telemetryService,
-            (tabId: string) => this.sendPinnedContext(tabId)
+            (tabId: string) => this.sendPinnedContext(tabId),
+            (tabId, enabled) => {
+                const { data: session, success } = this.#getOrCreateSession(tabId)
+                if (success) {
+                    this.#setPairProgrammingMode(session, enabled)
+                }
+            }
         )
 
         // Inject McpManager.getResources as a callback to avoid importing McpManager directly
@@ -868,7 +874,10 @@
         })
     }
 
-    async onChatPrompt(params: ChatParams, token: CancellationToken): Promise<ChatResult | ResponseError<ChatResult>> {
+    async onChatPrompt(
+        params: ChatParams & { prompt: { options?: Record<string, unknown> } },
+        token: CancellationToken
+    ): Promise<ChatResult | ResponseError<ChatResult>> {
         const clientRegion = this.#features.lsp.getClientInitializeParams()?.initializationOptions?.aws?.region
         const maybeGovResponse = getGovCloudUnsupportedResponse(clientRegion)
         if (maybeGovResponse) {
@@ -880,13 +889,20 @@
 
         IdleWorkspaceManager.recordActivityTimestamp()
 
-        const sessionResult = this.#chatSessionManagementService.getSession(params.tabId)
+        const sessionResult = this.#getOrCreateSession(params.tabId)
         const { data: session, success } = sessionResult
 
         if (!success) {
             return new ResponseError<ChatResult>(ErrorCodes.InternalError, sessionResult.error)
         }
 
+        // Mynah sends the displayed mode with each prompt, including restored tabs
+        // that did not emit an option-change event. Older clients may omit it.
+        const requestedMode = params.prompt.options?.['pair-programmer-mode']
+        if (requestedMode !== undefined) {
+            this.#setPairProgrammingMode(session, requestedMode === 'true')
+        }
+
         // Memory Bank Creation Flow - Delegate to MemoryBankController
         if (this.#memoryBankController.isMemoryBankCreationRequest(params.prompt.prompt)) {
             this.#features.logging.info(`Memory Bank creation request detected for tabId: ${params.tabId}`)
@@ -1965,6 +1981,9 @@
         session: ChatSessionService,
         toolName: string
     ) {
+        // The mode can change while the approval card is being rendered. Avoid
+        // registering an approval that the OFF event already tried to cancel.
+        this.#assertToolAvailable(session, toolUse.name)
         const deferred = this.#createDeferred()
         session.setDeferredToolExecution(toolUse.toolUseId!, deferred.resolve, deferred.reject)
         this.#log(`Prompting for tool approval for tool: ${toolName ?? toolUse.name}`)
@@ -2002,10 +2021,7 @@
 
             try {
                 // TODO: Can we move this check in the event parser before the stream completes?
-                const availableToolNames = this.#getTools(session).map(tool => tool.toolSpecification.name)
-                if (!availableToolNames.includes(toolUse.name)) {
-                    throw new Error(`Tool ${toolUse.name} is not available in the current mode`)
-                }
+                this.#assertToolAvailable(session, toolUse.name)
 
                 this.recordChunk(`tool_execution_start - ${toolUse.name}`)
                 this.#toolStartTime = Date.now()
@@ -2225,6 +2241,10 @@
                     }
                 }
 
+                // Mode may change while permission checks, approvals, or document
+                // reads are pending. Do not dispatch a tool using the earlier decision.
+                this.#assertToolAvailable(session, toolUse.name)
+
                 // After approval, add the path to the approved paths in the session
                 const inputPath = (toolUse.input as any)?.path || (toolUse.input as any)?.cwd
                 if (inputPath) {
@@ -4050,19 +4070,13 @@
             this.sendPinnedContext(params.tabId)
         }
 
-        const sessionResult = this.#chatSessionManagementService.createSession(params.tabId)
+        const sessionResult = this.#getOrCreateSession(params.tabId)
         const { data: session, success } = sessionResult
         if (!success) {
             return new ResponseError<ChatResult>(ErrorCodes.InternalError, sessionResult.error)
         }
 
-        // Get the saved pair programming mode from the database or default to true if not found
-        const savedPairProgrammingMode = this.#chatHistoryDb.getPairProgrammingMode()
-        session.pairProgrammingMode = savedPairProgrammingMode !== undefined ? savedPairProgrammingMode : true
-        if (session) {
-            // Set the logging object on the session
-            session.setLogging(this.#features.logging)
-        }
+        session.setLogging(this.#features.logging)
 
         // Update the client with the initial pair programming mode
         this.#features.chat.chatOptionsUpdate({
@@ -4808,16 +4822,44 @@
         }
     }
 
+    #getOrCreateSession(tabId: string) {
+        const existed = this.#chatSessionManagementService.hasSession(tabId)
+        const result = this.#chatSessionManagementService.getSession(tabId)
+        if (result.success && !existed) {
+            result.data.pairProgrammingMode = this.#chatHistoryDb.getEffectiveTabPairProgrammingMode(tabId)
+        }
+        return result
+    }
+
+    #setPairProgrammingMode(session: ChatSessionService, enabled: boolean) {
+        session.pairProgrammingMode = enabled
+        if (!enabled) {
+            const allowedTools = new Set(this.#getTools(session).map(tool => tool.toolSpecification.name))
+            for (const [toolUseId, toolUse] of session.toolUseLookup) {
+                if (toolUse.name && !allowedTools.has(toolUse.name)) {
+                    session
+                        .getDeferredToolExecution(toolUseId)
+                        ?.reject(new ToolApprovalException('Tool canceled: agentic coding is off', true))
+                    session.removeDeferredToolExecution(toolUseId)
+                }
+            }
+        }
+    }
+
     onPromptInputOptionChange(params: PromptInputOptionChangeParams) {
-        const sessionResult = this.#chatSessionManagementService.getSession(params.tabId)
+        const sessionResult = this.#getOrCreateSession(params.tabId)
         const { data: session, success } = sessionResult
 
         if (!success) {
             this.#log('onPromptInputOptionChange: on valid session found')
             return
         }
 
-        session.pairProgrammingMode = params.optionsValues['pair-programmer-mode'] === 'true'
+        const requestedMode = params.optionsValues['pair-programmer-mode']
+        if (requestedMode !== undefined) {
+            this.#setPairProgrammingMode(session, requestedMode === 'true')
+            this.#chatHistoryDb.setTabPairProgrammingMode(params.tabId, session.pairProgrammingMode)
+        }
         const newModelId = params.optionsValues['model-selection']
 
         // Set model (automatically recalculates token limits)
@@ -4828,7 +4870,6 @@
         }
 
         this.#chatHistoryDb.setTabModelId(params.tabId, session.modelId)
-        this.#chatHistoryDb.setTabPairProgrammingMode(params.tabId, session.pairProgrammingMode)
     }
 
     updateConfiguration = (newConfig: AmazonQWorkspaceConfig) => {
@@ -4852,6 +4893,12 @@
         this.#subscriptionStatusPromise = undefined
     }
 
+    #assertToolAvailable(session: ChatSessionService, toolName: string | undefined) {
+        if (!this.#getTools(session).some(tool => tool.toolSpecification.name === toolName)) {
+            throw new Error(`Tool ${toolName} is not available in the current mode`)
+        }
+    }
+
     #getTools(session: ChatSessionService) {
         const builtInWriteTools = new Set(this.#features.agent.getBuiltInWriteToolNames())
         const allTools = this.#features.agent.getTools({ format: 'bedrock' })

server/aws-lsp-codewhisperer/src/language-server/agenticChat/tabBarController.test.ts

@@ -23,6 +23,9 @@ describe('TabBarController', () => {
     let tabBarController: TabBarController
     let clock: sinon.SinonFakeTimers
     let telemetryService: TelemetryService
+    // Fifth TabBarController arg: the controller-provided callback that restores
+    // the per-tab agentic coding (pair programming) mode onto the live session.
+    let restorePairProgrammingModeStub: sinon.SinonStub
 
     beforeEach(() => {
         testFeatures = new TestFeatures()
@@ -38,6 +41,8 @@ describe('TabBarController', () => {
             getDatabaseFileSize: sinon.stub(),
             getLoadTime: sinon.stub(),
             getTabPreferences: sinon.stub().returns({}),
+            // Default effective mode ON; restore tests override per tab as needed.
+            getEffectiveTabPairProgrammingMode: sinon.stub().returns(true),
         } as unknown as ChatDatabase
 
         telemetryService = {
@@ -46,7 +51,14 @@ describe('TabBarController', () => {
             emitLoadHistory: sinon.stub(),
         } as any
 
-        tabBarController = new TabBarController(testFeatures, chatHistoryDb, telemetryService, sinon.stub())
+        restorePairProgrammingModeStub = sinon.stub()
+        tabBarController = new TabBarController(
+            testFeatures,
+            chatHistoryDb,
+            telemetryService,
+            sinon.stub(),
+            restorePairProgrammingModeStub
+        )
         clock = sinon.useFakeTimers()
     })
 
@@ -517,6 +529,49 @@ describe('TabBarController', () => {
             // Verify only the last 250 messages were passed
             assert.strictEqual(passedMessages.length, 250)
         })
+
+        it('restores per-tab mode via the callback using the new tabId, after mapping and before the UI update', async () => {
+            const historyId = 'history-ppm'
+            const mockTab = { historyId, conversations: [{ messages: [] }] } as unknown as Tab
+
+            // Restored tab had agentic coding turned OFF.
+            ;(chatHistoryDb.getEffectiveTabPairProgrammingMode as sinon.SinonStub).returns(false)
+            // Preferences carry the mode so the UI chatOptionsUpdate is emitted.
+            ;(chatHistoryDb.getTabPreferences as sinon.SinonStub).returns({ pairProgrammingMode: false })
+
+            const openTabStub = sinon.stub<[OpenTabParams], Promise<OpenTabResult>>().resolves({ tabId: 'newTabId' })
+            testFeatures.chat.openTab = openTabStub
+
+            await tabBarController.restoreTab(mockTab)
+
+            // Execution state is restored onto the new tabId (not the historyId).
+            sinon.assert.calledOnceWithExactly(restorePairProgrammingModeStub, 'newTabId', false)
+            sinon.assert.calledWith(chatHistoryDb.getEffectiveTabPairProgrammingMode as sinon.SinonStub, 'newTabId')
+
+            // Mapping is established first; the session is updated before the UI is told.
+            sinon.assert.callOrder(
+                chatHistoryDb.setHistoryIdMapping as sinon.SinonStub,
+                chatHistoryDb.getEffectiveTabPairProgrammingMode as sinon.SinonStub,
+                restorePairProgrammingModeStub,
+                testFeatures.chat.chatOptionsUpdate as sinon.SinonStub
+            )
+            sinon.assert.calledWithMatch(testFeatures.chat.chatOptionsUpdate as sinon.SinonStub, {
+                tabId: 'newTabId',
+                pairProgrammingMode: false,
+            })
+        })
+
+        it('restores an ON per-tab mode through the callback', async () => {
+            const mockTab = { historyId: 'history-on', conversations: [{ messages: [] }] } as unknown as Tab
+            ;(chatHistoryDb.getEffectiveTabPairProgrammingMode as sinon.SinonStub).returns(true)
+
+            const openTabStub = sinon.stub<[OpenTabParams], Promise<OpenTabResult>>().resolves({ tabId: 'tabOn' })
+            testFeatures.chat.openTab = openTabStub
+
+            await tabBarController.restoreTab(mockTab)
+
+            sinon.assert.calledOnceWithExactly(restorePairProgrammingModeStub, 'tabOn', true)
+        })
     })
 
     describe('loadChats', () => {

server/aws-lsp-codewhisperer/src/language-server/agenticChat/tabBarController.ts

@@ -43,17 +43,20 @@ export class TabBarController {
     #chatHistoryDb: ChatDatabase
     #telemetryService: TelemetryService
     #sendPinnedContext: (tabId: string) => void
+    #restorePairProgrammingMode: (tabId: string, enabled: boolean) => void
 
     constructor(
         features: Features,
         chatHistoryDb: ChatDatabase,
         telemetryService: TelemetryService,
-        sendPinnedContext: (tabId: string) => void
+        sendPinnedContext: (tabId: string) => void,
+        restorePairProgrammingMode: (tabId: string, enabled: boolean) => void
     ) {
         this.#features = features
         this.#chatHistoryDb = chatHistoryDb
         this.#telemetryService = telemetryService
         this.#sendPinnedContext = sendPinnedContext
+        this.#restorePairProgrammingMode = restorePairProgrammingMode
     }
 
     /**
@@ -311,6 +314,10 @@ export class TabBarController {
 
             // Restore per-tab preferences (model selection and agentic coding mode)
             const preferences = this.#chatHistoryDb.getTabPreferences(selectedTab.historyId)
+            // Update the execution state before displaying the restored mode. The
+            // history mapping is not available when the client first adds this tab.
+            const pairProgrammingMode = this.#chatHistoryDb.getEffectiveTabPairProgrammingMode(tabId)
+            this.#restorePairProgrammingMode(tabId, pairProgrammingMode)
             if (preferences.modelId !== undefined || preferences.pairProgrammingMode !== undefined) {
                 // Validate modelId against current available models
                 let validModelId = preferences.modelId