From d61b65c8e64d9c5450f1e0b47b96a49dd12fb023 Mon Sep 17 00:00:00 2001 From: Laxman Reddy Aileni Date: Mon, 5 Oct 2026 22:31:58 +0000 Subject: [PATCH] fix(amazonq): honor agentic OFF state across chat lifecycle --- .../agenticChat/agenticChatController.test.ts | 436 ++++++++++++++++++ .../agenticChat/agenticChatController.ts | 83 +++- .../agenticChat/tabBarController.test.ts | 57 ++- .../agenticChat/tabBarController.ts | 9 +- .../agenticChat/tools/chatDb/chatDb.test.ts | 166 +++++++ .../agenticChat/tools/chatDb/chatDb.ts | 133 ++++-- .../chat/chatSessionService.ts | 3 +- 7 files changed, 838 insertions(+), 49 deletions(-) diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts index 4ba552b98f..8c21402a51 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts @@ -12,6 +12,7 @@ import { ContentType, GenerateAssistantResponseCommandInput, SendMessageCommandInput, + ToolResultStatus, } from '@amzn/codewhisperer-streaming' import { QDeveloperStreaming, @@ -58,6 +59,9 @@ import { ChatDatabase } from './tools/chatDb/chatDb' import { LocalProjectContextController } from '../../shared/localProjectContextController' import { CancellationError } from '@aws/lsp-core' import { ToolApprovalException } from './tools/toolShared' +import { FsWrite } from './tools/fsWrite' +import { FsRead } from './tools/fsRead' +import { AgenticChatTriggerContext } from './context/agenticChatTriggerContext' import * as constants from './constants/constants' import { GENERIC_ERROR_MS } from './constants/constants' import { TokenLimitsCalculator } from './utils/tokenLimitsCalculator' @@ -3636,6 +3640,438 @@ ${' '.repeat(8)}} assert.strictEqual(toolInput.ruleArtifacts[1].path, '/test/rule2.json') }) }) + + describe('agentic coding (pair programming) mode regression', () => { + const PP_KEY = 'pair-programmer-mode' + let getEffectiveModeStub: sinon.SinonStub + + const getSession = (tabId: string) => chatSessionManagementService.getSession(tabId).data! + + // Make the write tool observable in the session's tool set so that OFF-mode + // exclusion (and ON-mode availability) is testable through the real #getTools. + const withFsWriteInToolset = () => { + testFeatures.agent.getTools = sinon.stub().returns( + ['mock-tool-name', 'codeReview', 'fsRead', 'fsWrite'].map(name => ({ + toolSpecification: { name, description: 'Mock tool for testing' }, + })) + ) + } + + const makeResultStream = () => + ({ + removeResultBlockAndUpdateUI: sinon.stub().resolves(), + writeResultBlock: sinon.stub().resolves(1), + overwriteResultBlock: sinon.stub().resolves(), + removeResultBlock: sinon.stub().resolves(), + getMessageBlockId: sinon.stub().returns(undefined), + hasMessage: sinon.stub().returns(false), + updateOngoingProgressResult: sinon.stub().resolves(), + getResult: sinon.stub().returns({ messageId: 'test', body: '' }), + setMessageIdToUpdateForTool: sinon.stub(), + getMessageIdToUpdateForTool: sinon.stub().returns(undefined), + addMessageOperation: sinon.stub(), + getMessageOperation: sinon.stub().returns(undefined), + }) as any + + beforeEach(() => { + // Normal initialized default: the effective mode is ON unless a test overrides it. + getEffectiveModeStub = sinon + .stub(ChatDatabase.prototype, 'getEffectiveTabPairProgrammingMode') + .returns(true) + }) + + describe('onChatPrompt mode option', () => { + it('explicit pair-programmer-mode=false corrects an ON session and excludes fsWrite', async () => { + withFsWriteInToolset() + chatController.onTabAdd({ tabId: mockTabId }) + const session = getSession(mockTabId) + assert.strictEqual(session.pairProgrammingMode, true) + + const prepareRequestSpy = sinon.spy(AgenticChatTriggerContext.prototype, 'getChatParamsFromTrigger') + await chatController.onChatPrompt( + { tabId: mockTabId, prompt: { prompt: 'Hello', options: { [PP_KEY]: 'false' } } }, + mockCancellationToken + ) + + assert.strictEqual(session.pairProgrammingMode, false) + sinon.assert.calledOnce(prepareRequestSpy) + assert.ok(!prepareRequestSpy.firstCall.args[6]?.some(tool => tool.toolSpecification.name === 'fsWrite')) + + // The write tool is now excluded from the session's tool set: dispatch is blocked. + const runToolStub = testFeatures.agent.runTool as sinon.SinonStub + runToolStub.resetHistory() + await chatController.processToolUses( + [{ toolUseId: 'tu-w', name: 'fsWrite', input: { path: '/tmp/a.ts' }, stop: true } as any], + makeResultStream(), + session, + mockTabId, + mockCancellationToken + ) + sinon.assert.notCalled(runToolStub) + }) + + it('explicit pair-programmer-mode=true turns an OFF session back ON', async () => { + getEffectiveModeStub.returns(false) + chatController.onTabAdd({ tabId: mockTabId }) + const session = getSession(mockTabId) + assert.strictEqual(session.pairProgrammingMode, false) + + await chatController.onChatPrompt( + { tabId: mockTabId, prompt: { prompt: 'Hello', options: { [PP_KEY]: 'true' } } } as any, + mockCancellationToken + ) + + assert.strictEqual(session.pairProgrammingMode, true) + }) + + it('omitted prompt options preserve the current mode', async () => { + chatController.onTabAdd({ tabId: mockTabId }) + const session = getSession(mockTabId) + assert.strictEqual(session.pairProgrammingMode, true) + + await chatController.onChatPrompt( + { tabId: mockTabId, prompt: { prompt: 'Hello' } }, + mockCancellationToken + ) + + assert.strictEqual(session.pairProgrammingMode, true) + }) + + it('any explicit non-"true" option value disables the mode', async () => { + chatController.onTabAdd({ tabId: mockTabId }) + const session = getSession(mockTabId) + assert.strictEqual(session.pairProgrammingMode, true) + + await chatController.onChatPrompt( + { tabId: mockTabId, prompt: { prompt: 'Hello', options: { [PP_KEY]: 'maybe' } } } as any, + mockCancellationToken + ) + + assert.strictEqual(session.pairProgrammingMode, false) + }) + + it('lazily creates the session honoring an effective OFF mode when no tab-add occurred', async () => { + getEffectiveModeStub.returns(false) + assert.ok(!chatSessionManagementService.hasSession(mockTabId)) + + await chatController.onChatPrompt( + { tabId: mockTabId, prompt: { prompt: 'Hello' } }, + mockCancellationToken + ) + + assert.strictEqual(getSession(mockTabId).pairProgrammingMode, false) + }) + }) + + describe('onTabAdd lazy creation and preservation', () => { + it('initializes a new session from the effective OFF mode and notifies the UI', () => { + getEffectiveModeStub.returns(false) + + chatController.onTabAdd({ tabId: mockTabId }) + + assert.strictEqual(getSession(mockTabId).pairProgrammingMode, false) + sinon.assert.calledWithMatch(testFeatures.chat.chatOptionsUpdate as sinon.SinonStub, { + tabId: mockTabId, + pairProgrammingMode: false, + }) + }) + + it('does not overwrite an OFF session on a late or repeated tab add', () => { + getEffectiveModeStub.returns(false) + chatController.onTabAdd({ tabId: mockTabId }) + assert.strictEqual(getSession(mockTabId).pairProgrammingMode, false) + + // The effective mode later reports ON, but the existing session must be preserved. + getEffectiveModeStub.returns(true) + chatController.onTabAdd({ tabId: mockTabId }) + + assert.strictEqual(getSession(mockTabId).pairProgrammingMode, false) + }) + }) + + describe('restored per-tab OFF with global ON synchronizes the session', () => { + const historyId = 'history-restore' + + // Drive the real controller restore callback through TabBarController.loadChats. + const primeRestore = () => { + // New sessions initially inherit global ON. The restored tab's + // OFF preference becomes available only after its history mapping. + let mapped = false + getEffectiveModeStub.callsFake(() => !mapped) + sinon + .stub(ChatDatabase.prototype, 'getOpenTabs') + .returns([{ historyId, conversations: [{ messages: [] }] }] as any) + sinon.stub(ChatDatabase.prototype, 'setHistoryIdMapping').callsFake(() => { + mapped = true + }) + sinon.stub(ChatDatabase.prototype, 'updateTabOpenState') + sinon.stub(ChatDatabase.prototype, 'getTabPreferences').returns({ pairProgrammingMode: false }) + sinon.stub(ChatDatabase.prototype, 'getLoadTime').returns(undefined) + sinon.stub(ChatDatabase.prototype, 'getDatabaseFileSize').returns(undefined) + testFeatures.chat.openTab = sinon.stub().resolves({ tabId: mockTabId }) as any + } + + it('restore before tab-add leaves the session OFF', async () => { + primeRestore() + + await chatController.restorePreviousChats() + + // The restore callback created and synchronized the session OFF. + assert.ok(chatSessionManagementService.hasSession(mockTabId)) + assert.strictEqual(getSession(mockTabId).pairProgrammingMode, false) + + chatController.onTabAdd({ tabId: mockTabId }) + assert.strictEqual(getSession(mockTabId).pairProgrammingMode, false) + }) + + it('tab-add before restore leaves the session OFF', async () => { + primeRestore() + + chatController.onTabAdd({ tabId: mockTabId }) + assert.strictEqual(getSession(mockTabId).pairProgrammingMode, true) + + await chatController.restorePreviousChats() + assert.strictEqual(getSession(mockTabId).pairProgrammingMode, false) + sinon.assert.calledWithMatch(testFeatures.chat.chatOptionsUpdate as sinon.SinonStub, { + tabId: mockTabId, + pairProgrammingMode: false, + }) + }) + }) + + describe('onPromptInputOptionChange', () => { + it('preserves the mode for a model-only option event', () => { + const cachedModels = [ + { id: 'model-x', name: 'X', description: 'test', tokenLimits: { maxInputTokens: 200000 } }, + ] + sinon.stub(ChatDatabase.prototype, 'getCachedModels').returns({ + models: cachedModels as any, + defaultModelId: 'model-x', + timestamp: Date.now(), + }) + chatController.onTabAdd({ tabId: mockTabId }) + const session = getSession(mockTabId) + assert.strictEqual(session.pairProgrammingMode, true) + + chatController.onPromptInputOptionChange({ + tabId: mockTabId, + optionsValues: { 'model-selection': 'model-x' }, + }) + + assert.strictEqual(session.modelId, 'model-x') + assert.strictEqual(session.pairProgrammingMode, true) + }) + + it('rejects queued restricted-tool approvals when turned OFF and never runs the tool', () => { + withFsWriteInToolset() + chatController.onTabAdd({ tabId: mockTabId }) + const session = getSession(mockTabId) + + // A restricted (write) tool and a read tool each have a pending approval. + const writeReject = sinon.spy() + const readReject = sinon.spy() + session.toolUseLookup.set('tu-w', { name: 'fsWrite', toolUseId: 'tu-w' } as any) + session.toolUseLookup.set('tu-r', { name: 'fsRead', toolUseId: 'tu-r' } as any) + session.setDeferredToolExecution('tu-w', sinon.spy(), writeReject) + session.setDeferredToolExecution('tu-r', sinon.spy(), readReject) + + chatController.onPromptInputOptionChange({ + tabId: mockTabId, + optionsValues: { [PP_KEY]: 'false' }, + }) + + assert.strictEqual(session.pairProgrammingMode, false) + // The write-tool approval is rejected and cleared; the read-tool one survives. + sinon.assert.calledOnce(writeReject) + assert.ok(writeReject.firstCall.args[0] instanceof ToolApprovalException) + assert.strictEqual(session.getDeferredToolExecution('tu-w'), undefined) + sinon.assert.notCalled(readReject) + assert.ok(session.getDeferredToolExecution('tu-r') !== undefined) + sinon.assert.notCalled(testFeatures.agent.runTool as sinon.SinonStub) + }) + }) + + describe('processToolUses mode enforcement', () => { + const writeToolUse = { toolUseId: 'tu-w', name: 'fsWrite', input: { path: '/tmp/a.ts' }, stop: true } + const readToolUse = { toolUseId: 'tu-r', name: 'fsRead', input: { paths: ['/tmp/a.ts'] }, stop: true } + + const freshSession = (mode: boolean) => { + chatController.onTabAdd({ tabId: mockTabId }) + const session = getSession(mockTabId) + session.pairProgrammingMode = mode + return session + } + + it('blocks a returned write tool when the session mode is OFF', async () => { + withFsWriteInToolset() + const session = freshSession(false) + const runToolStub = testFeatures.agent.runTool as sinon.SinonStub + runToolStub.resetHistory() + + const results = await chatController.processToolUses( + [writeToolUse as any], + makeResultStream(), + session, + mockTabId, + mockCancellationToken + ) + + sinon.assert.notCalled(runToolStub) + assert.strictEqual(results[0]?.status, ToolResultStatus.ERROR) + }) + + it('dispatches a write tool when the session mode is ON', async () => { + withFsWriteInToolset() + sinon.stub(FsWrite.prototype, 'requiresAcceptance').resolves({ requiresAcceptance: false } as any) + sinon.stub(AgenticChatTriggerContext.prototype, 'getTextDocumentFromPath').resolves(undefined) + const session = freshSession(true) + const runToolStub = testFeatures.agent.runTool as sinon.SinonStub + runToolStub.resetHistory() + runToolStub.resolves({}) + + await chatController.processToolUses( + [writeToolUse as any], + makeResultStream(), + session, + mockTabId, + mockCancellationToken + ) + + sinon.assert.called(runToolStub) + assert.strictEqual(runToolStub.firstCall.args[0], 'fsWrite') + }) + + it('blocks dispatch when the mode flips OFF while an async permission check is pending', async () => { + withFsWriteInToolset() + const session = freshSession(true) + // Simulate a concurrent mode change during the async requiresAcceptance call. + const requiresAcceptanceStub = sinon + .stub(FsWrite.prototype, 'requiresAcceptance') + .callsFake(async () => { + session.pairProgrammingMode = false + return { requiresAcceptance: false } as any + }) + sinon.stub(AgenticChatTriggerContext.prototype, 'getTextDocumentFromPath').resolves(undefined) + const runToolStub = testFeatures.agent.runTool as sinon.SinonStub + runToolStub.resetHistory() + + await chatController.processToolUses( + [writeToolUse as any], + makeResultStream(), + session, + mockTabId, + mockCancellationToken + ) + + // Passed the entry guard (mode ON) but blocked by the re-check after the await. + sinon.assert.called(requiresAcceptanceStub) + sinon.assert.notCalled(runToolStub) + }) + + it('does not wait for approval if OFF arrives while the approval card is rendered', async () => { + withFsWriteInToolset() + sinon.stub(FsWrite.prototype, 'requiresAcceptance').resolves({ requiresAcceptance: true }) + const session = freshSession(true) + const resultStream = makeResultStream() + resultStream.writeResultBlock.onFirstCall().callsFake(async () => { + chatController.onPromptInputOptionChange({ + tabId: mockTabId, + optionsValues: { [PP_KEY]: 'false' }, + }) + return 1 + }) + + const results = await chatController.processToolUses( + [writeToolUse], + resultStream, + session, + mockTabId, + mockCancellationToken + ) + + sinon.assert.notCalled(testFeatures.agent.runTool as sinon.SinonStub) + assert.strictEqual(session.getDeferredToolExecution('tu-w'), undefined) + assert.strictEqual(results[0]?.status, ToolResultStatus.ERROR) + }) + + it('rejects a real pending write approval when the mode is turned OFF', async () => { + withFsWriteInToolset() + sinon.stub(FsWrite.prototype, 'requiresAcceptance').resolves({ requiresAcceptance: true }) + const session = freshSession(true) + let approvalReady!: () => void + const ready = new Promise(resolve => { + approvalReady = resolve + }) + const setDeferred = session.setDeferredToolExecution.bind(session) + sinon.stub(session, 'setDeferredToolExecution').callsFake((id, resolve, reject) => { + setDeferred(id, resolve, reject) + approvalReady() + }) + const pending = chatController.processToolUses( + [writeToolUse], + makeResultStream(), + session, + mockTabId, + mockCancellationToken + ) + const rejected = assert.rejects(pending, ToolApprovalException) + await ready + + chatController.onPromptInputOptionChange({ + tabId: mockTabId, + optionsValues: { [PP_KEY]: 'false' }, + }) + + await rejected + sinon.assert.notCalled(testFeatures.agent.runTool as sinon.SinonStub) + assert.strictEqual(session.getDeferredToolExecution('tu-w'), undefined) + }) + + it('blocks dispatch when OFF arrives during the last document read before a write', async () => { + withFsWriteInToolset() + sinon.stub(FsWrite.prototype, 'requiresAcceptance').resolves({ requiresAcceptance: false }) + const session = freshSession(true) + sinon.stub(AgenticChatTriggerContext.prototype, 'getTextDocumentFromPath').callsFake(async () => { + chatController.onPromptInputOptionChange({ + tabId: mockTabId, + optionsValues: { [PP_KEY]: 'false' }, + }) + return undefined + }) + + const results = await chatController.processToolUses( + [writeToolUse], + makeResultStream(), + session, + mockTabId, + mockCancellationToken + ) + + sinon.assert.notCalled(testFeatures.agent.runTool as sinon.SinonStub) + assert.strictEqual(results[0]?.status, ToolResultStatus.ERROR) + }) + + it('keeps read tools usable when the session mode is OFF', async () => { + withFsWriteInToolset() + sinon.stub(FsRead.prototype, 'requiresAcceptance').resolves({ requiresAcceptance: false } as any) + const session = freshSession(false) + const runToolStub = testFeatures.agent.runTool as sinon.SinonStub + runToolStub.resetHistory() + runToolStub.resolves({}) + + await chatController.processToolUses( + [readToolUse as any], + makeResultStream(), + session, + mockTabId, + mockCancellationToken + ) + + sinon.assert.called(runToolStub) + assert.strictEqual(runToolStub.firstCall.args[0], 'fsRead') + }) + }) + }) }) // The body may include text-based progress updates from tool invocations. diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts index ac991235ed..b8a97ebde4 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts @@ -368,7 +368,13 @@ export class AgenticChatController implements ChatHandlers { 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 @@ export class AgenticChatController implements ChatHandlers { }) } - async onChatPrompt(params: ChatParams, token: CancellationToken): Promise> { + async onChatPrompt( + params: ChatParams & { prompt: { options?: Record } }, + token: CancellationToken + ): Promise> { const clientRegion = this.#features.lsp.getClientInitializeParams()?.initializationOptions?.aws?.region const maybeGovResponse = getGovCloudUnsupportedResponse(clientRegion) if (maybeGovResponse) { @@ -880,13 +889,20 @@ export class AgenticChatController implements ChatHandlers { 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(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 @@ export class AgenticChatController implements ChatHandlers { 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 @@ export class AgenticChatController implements ChatHandlers { 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 @@ export class AgenticChatController implements ChatHandlers { } } + // 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 @@ export class AgenticChatController implements ChatHandlers { 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(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,8 +4822,32 @@ export class AgenticChatController implements ChatHandlers { } } + #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) { @@ -4817,7 +4855,11 @@ export class AgenticChatController implements ChatHandlers { 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 @@ export class AgenticChatController implements ChatHandlers { } this.#chatHistoryDb.setTabModelId(params.tabId, session.modelId) - this.#chatHistoryDb.setTabPairProgrammingMode(params.tabId, session.pairProgrammingMode) } updateConfiguration = (newConfig: AmazonQWorkspaceConfig) => { @@ -4852,6 +4893,12 @@ export class AgenticChatController implements ChatHandlers { 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' }) diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tabBarController.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tabBarController.test.ts index 1b17ee35ad..a838059255 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tabBarController.test.ts +++ b/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>().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>().resolves({ tabId: 'tabOn' }) + testFeatures.chat.openTab = openTabStub + + await tabBarController.restoreTab(mockTab) + + sinon.assert.calledOnceWithExactly(restorePairProgrammingModeStub, 'tabOn', true) + }) }) describe('loadChats', () => { diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tabBarController.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tabBarController.ts index dfff7c1c49..4c39f52dd5 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tabBarController.ts +++ b/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 diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/chatDb/chatDb.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/chatDb/chatDb.test.ts index 2c0e07baff..91aa1ecdf8 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/chatDb/chatDb.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/chatDb/chatDb.test.ts @@ -708,6 +708,172 @@ describe('ChatDatabase', () => { }, 'Should not throw when no instance exists') }) }) + + describe('Pair Programming Mode Initialization', () => { + // These tests exercise the window before the LokiJS database finishes + // loading (isInitialized() === false). To make that window deterministic, + // we build a dedicated instance whose filesystem `mkdir` never resolves, + // so LokiJS autoload never completes and the instance stays uninitialized + // until we explicitly call databaseInitialize(). + let modeDb: ChatDatabase + + beforeEach(() => { + const neverResolvingMkdir = sinon.stub().returns(new Promise(() => {})) + const modeFeatures = { + ...(mockFeatures as any), + workspace: { + ...(mockFeatures.workspace as any), + fs: { + ...(mockFeatures.workspace as any).fs, + mkdir: neverResolvingMkdir, + }, + }, + } as unknown as Features + // Use `new` (not getInstance) so this controlled instance is independent + // of the singleton created by the outer beforeEach. + modeDb = new ChatDatabase(modeFeatures) + }) + + afterEach(() => { + modeDb.close() + }) + + it('pre-init reads: effective mode is false and settings are undefined with no pending choice', () => { + assert.strictEqual(modeDb.isInitialized(), false, 'Database should be uninitialized') + assert.strictEqual(modeDb.getPairProgrammingMode(), undefined, 'Global mode should be undefined') + assert.strictEqual(modeDb.getTabPairProgrammingMode('tab-1'), undefined, 'Tab mode should be undefined') + // Key regression: before init with no explicit choice, do NOT default to + // true (agentic ON). Default to false so a possibly-persisted OFF is not + // overridden before the database is ready. + assert.strictEqual( + modeDb.getEffectiveTabPairProgrammingMode('tab-1'), + false, + 'Effective mode should be false when uninitialized with no pending choice' + ) + }) + + it('early OFF: an explicit global OFF before init is surfaced on reads and survives the flush', async () => { + modeDb.setPairProgrammingMode(false) + + assert.strictEqual(modeDb.isInitialized(), false, 'Database should still be uninitialized') + assert.strictEqual(modeDb.getPairProgrammingMode(), false, 'Pending OFF should be surfaced before init') + assert.strictEqual( + modeDb.getEffectiveTabPairProgrammingMode('tab-1'), + false, + 'Effective mode should honor the pending OFF before init' + ) + + await modeDb.databaseInitialize(0) + + assert.strictEqual(modeDb.getPairProgrammingMode(), false, 'OFF should be flushed to global settings') + assert.strictEqual( + modeDb.getEffectiveTabPairProgrammingMode('tab-1'), + false, + 'Effective mode should remain false after init' + ) + }) + + it('early OFF (per-tab): an explicit tab OFF before init is surfaced on reads', () => { + modeDb.setTabPairProgrammingMode('tab-1', false) + + assert.strictEqual(modeDb.getTabPairProgrammingMode('tab-1'), false, 'Pending tab OFF should be surfaced') + // Per-tab setters also mirror into the pending global default. + assert.strictEqual(modeDb.getPairProgrammingMode(), false, 'Tab OFF should mirror into pending global') + assert.strictEqual( + modeDb.getEffectiveTabPairProgrammingMode('tab-1'), + false, + 'Effective mode should honor the pending tab OFF before init' + ) + }) + + it('mixed tab choices: distinct per-tab selections before init are preserved through the flush', async () => { + modeDb.setTabPairProgrammingMode('tab-a', true) + modeDb.setTabPairProgrammingMode('tab-b', false) + + // Pre-init reads reflect each tab's own pending selection. + assert.strictEqual(modeDb.getTabPairProgrammingMode('tab-a'), true, 'tab-a pending should be true') + assert.strictEqual(modeDb.getTabPairProgrammingMode('tab-b'), false, 'tab-b pending should be false') + assert.strictEqual(modeDb.getEffectiveTabPairProgrammingMode('tab-a'), true, 'tab-a effective pre-init') + assert.strictEqual(modeDb.getEffectiveTabPairProgrammingMode('tab-b'), false, 'tab-b effective pre-init') + + await modeDb.databaseInitialize(0) + + // After flush, per-tab selections are persisted independently. + assert.strictEqual(modeDb.getTabPairProgrammingMode('tab-a'), true, 'tab-a should persist true') + assert.strictEqual(modeDb.getTabPairProgrammingMode('tab-b'), false, 'tab-b should persist false') + assert.strictEqual(modeDb.getEffectiveTabPairProgrammingMode('tab-a'), true, 'tab-a effective post-init') + assert.strictEqual(modeDb.getEffectiveTabPairProgrammingMode('tab-b'), false, 'tab-b effective post-init') + }) + + it('final global ordering: the last global selection wins across interleaved setters after flush', async () => { + // Interleave per-tab setters (which mirror into global) with direct + // global setters. The last global write (false) must be authoritative. + modeDb.setTabPairProgrammingMode('tab-a', true) // global -> true + modeDb.setPairProgrammingMode(false) // global -> false + modeDb.setTabPairProgrammingMode('tab-b', true) // global -> true + modeDb.setPairProgrammingMode(false) // global -> false (last) + + await modeDb.databaseInitialize(0) + + assert.strictEqual( + modeDb.getPairProgrammingMode(), + false, + 'Global mode should equal the last global selection, not the last tab mirror' + ) + // Per-tab selections are still preserved independently of the global value. + assert.strictEqual(modeDb.getTabPairProgrammingMode('tab-a'), true, 'tab-a should persist true') + assert.strictEqual(modeDb.getTabPairProgrammingMode('tab-b'), true, 'tab-b should persist true') + // A brand-new tab with no selection falls back to the final global value. + assert.strictEqual( + modeDb.getEffectiveTabPairProgrammingMode('tab-c'), + false, + 'New tab should inherit the final global selection' + ) + }) + + it('loaded settings preserved: flushing pending mode merges into (does not clobber) existing settings', async () => { + modeDb.setPairProgrammingMode(false) // pending global OFF + + // Simulate settings that were loaded from disk before the flush runs. + // getSettings is what the flush reads before writing the merged record. + const loadedSettings = { + modelId: 'model-X', + pairProgrammingMode: undefined, + cachedModels: [{ id: 'm1', name: 'M1' }], + cachedDefaultModelId: 'm1', + modelCacheTimestamp: 123456, + } + const getSettingsStub = sinon.stub(modeDb, 'getSettings').returns(loadedSettings as any) + + await modeDb.databaseInitialize(0) + + getSettingsStub.restore() + + // The pending mode is applied... + assert.strictEqual(modeDb.getPairProgrammingMode(), false, 'Pending OFF should be flushed') + // ...without dropping unrelated loaded settings. + assert.strictEqual(modeDb.getModelId(), 'model-X', 'modelId should be preserved') + const cached = modeDb.getCachedModels() + assert.ok(cached, 'Cached models should be preserved') + assert.deepStrictEqual(cached.models, loadedSettings.cachedModels, 'Cached models should be intact') + assert.strictEqual(cached.defaultModelId, 'm1', 'Cached default model should be intact') + assert.strictEqual(cached.timestamp, 123456, 'Cache timestamp should be intact') + }) + + it('initialized defaults: effective mode is true when initialized with no explicit choice', async () => { + await modeDb.databaseInitialize(0) + + assert.strictEqual(modeDb.getPairProgrammingMode(), undefined, 'No global setting should exist') + assert.strictEqual(modeDb.getTabPairProgrammingMode('tab-1'), undefined, 'No tab setting should exist') + // Once initialized with genuinely no stored preference, keep the + // first-time-user default of true. + assert.strictEqual( + modeDb.getEffectiveTabPairProgrammingMode('tab-1'), + true, + 'Effective mode should default to true when initialized with no choice' + ) + }) + }) }) function uuid(): `${string}-${string}-${string}-${string}-${string}` { throw new Error('Function not implemented.') diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/chatDb/chatDb.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/chatDb/chatDb.ts index b0dd06280f..6009162d33 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/chatDb/chatDb.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/chatDb/chatDb.ts @@ -69,6 +69,21 @@ export class ChatDatabase { #loadTimeMs?: number #dbFileSize?: number #historyMaintainer: ChatHistoryMaintainer + /** + * Minimal, mode-specific pending state for pair programming mode selections + * made before the LokiJS database finishes loading (while isInitialized() is + * false). This is intentionally NOT a general settings queue: it captures only + * pair programming mode so an explicit user choice — especially turning the + * mode OFF — survives initialization and is visible on reads, without masking + * unrelated settings loaded from disk. + * + * The global entry uses an optional wrapper so "set to undefined" is + * distinguishable from "never set". The most recent write wins, which keeps + * the last global selection authoritative across interleaved global/per-tab + * setters. + */ + #pendingGlobalPairProgrammingMode?: { value: boolean | undefined } + #pendingTabPairProgrammingMode: Map = new Map() constructor(features: Features) { this.#features = features @@ -230,6 +245,42 @@ export class ChatDatabase { this.#db.addCollection(SettingsCollection) this.#initialized = true this.#loadTimeMs = Date.now() - startTime + // The database is now ready; apply any pair programming mode selections + // that were made while it was still loading. + this.flushPendingPairProgrammingMode() + } + + /** + * Applies pair programming mode selections captured before the database + * finished loading. Per-tab selections are written first; the global + * selection is applied LAST so the last in-session global selection wins, + * even though setTabPairProgrammingMode also updates the global default as a + * side effect (interleaved setters). Unrelated settings already loaded from + * disk are preserved because the underlying writes merge rather than replace. + */ + private flushPendingPairProgrammingMode(): void { + if (!this.#initialized) { + return + } + + // Snapshot and clear the pending buffers first. The setters below run in + // the initialized path and must not read or write the pending state. + const pendingTabModes = this.#pendingTabPairProgrammingMode + this.#pendingTabPairProgrammingMode = new Map() + const pendingGlobal = this.#pendingGlobalPairProgrammingMode + this.#pendingGlobalPairProgrammingMode = undefined + + // Apply per-tab selections first. setTabPairProgrammingMode also mirrors + // the value into global settings, so applying the global selection AFTER + // this loop keeps the last in-session global selection authoritative. + for (const [tabId, mode] of pendingTabModes) { + this.setTabPairProgrammingMode(tabId, mode) + } + + // Apply the final global selection last (last write wins). + if (pendingGlobal) { + this.setPairProgrammingMode(pendingGlobal.value) + } } getOpenTabs() { @@ -1020,11 +1071,23 @@ export class ChatDatabase { } getPairProgrammingMode(): boolean | undefined { + if (!this.#initialized) { + // Before the database is ready, surface the pending global selection + // (including an explicit OFF) rather than nothing. + return this.#pendingGlobalPairProgrammingMode?.value + } const settings = this.getSettings() return settings?.pairProgrammingMode } setPairProgrammingMode(pairProgrammingMode: boolean | undefined): void { + if (!this.#initialized) { + // Preserve the global selection until the database is ready. The most + // recent write wins, keeping last-global-selection semantics across + // interleaved global and per-tab setters. + this.#pendingGlobalPairProgrammingMode = { value: pairProgrammingMode } + return + } // Get existing settings to preserve other fields like modelId const settings = this.getSettings() || { modelId: undefined } this.updateSettings({ ...settings, pairProgrammingMode }) @@ -1085,13 +1148,15 @@ export class ChatDatabase { * @returns The tab's pair programming mode, or undefined if not set */ getTabPairProgrammingMode(tabId: string): boolean | undefined { - if (this.#initialized) { - const collection = this.#db.getCollection(TabCollection) - const historyId = this.#historyIdMapping.get(tabId) - if (historyId) { - const tab = collection.findOne({ historyId }) - return tab?.pairProgrammingMode - } + if (!this.#initialized) { + // Surface a pending per-tab selection made before the database is ready. + return this.#pendingTabPairProgrammingMode.get(tabId) + } + const collection = this.#db.getCollection(TabCollection) + const historyId = this.#historyIdMapping.get(tabId) + if (historyId) { + const tab = collection.findOne({ historyId }) + return tab?.pairProgrammingMode } return undefined } @@ -1102,28 +1167,34 @@ export class ChatDatabase { * @param pairProgrammingMode The pair programming mode to set */ setTabPairProgrammingMode(tabId: string, pairProgrammingMode: boolean | undefined): void { - if (this.#initialized) { - const collection = this.#db.getCollection(TabCollection) - const historyId = this.getOrCreateHistoryId(tabId) - const tab = collection.findOne({ historyId }) - - this.#features.logging.log(`Setting tab pair programming mode: tabId=${tabId}, mode=${pairProgrammingMode}`) - - if (!tab) { - this.addTabWithContext(collection, historyId, {}) - const newTab = collection.findOne({ historyId }) - if (newTab) { - newTab.pairProgrammingMode = pairProgrammingMode - collection.update(newTab) - } - } else { - tab.pairProgrammingMode = pairProgrammingMode - collection.update(tab) + if (!this.#initialized) { + // Preserve the per-tab selection and mirror it into the pending global + // selection, matching the initialized behavior where setting a tab's + // mode also updates the global default for new tabs. + this.#pendingTabPairProgrammingMode.set(tabId, pairProgrammingMode) + this.#pendingGlobalPairProgrammingMode = { value: pairProgrammingMode } + return + } + const collection = this.#db.getCollection(TabCollection) + const historyId = this.getOrCreateHistoryId(tabId) + const tab = collection.findOne({ historyId }) + + this.#features.logging.log(`Setting tab pair programming mode: tabId=${tabId}, mode=${pairProgrammingMode}`) + + if (!tab) { + this.addTabWithContext(collection, historyId, {}) + const newTab = collection.findOne({ historyId }) + if (newTab) { + newTab.pairProgrammingMode = pairProgrammingMode + collection.update(newTab) } - - // Also update global settings with the latest selection for new tab defaults - this.setPairProgrammingMode(pairProgrammingMode) + } else { + tab.pairProgrammingMode = pairProgrammingMode + collection.update(tab) } + + // Also update global settings with the latest selection for new tab defaults + this.setPairProgrammingMode(pairProgrammingMode) } /** @@ -1154,7 +1225,13 @@ export class ChatDatabase { if (globalMode !== undefined) { return globalMode } - // Default to true for first-time users + // Before initialization, with no explicit in-session choice, we cannot + // know the persisted preference yet. Default to false (mode OFF) so we + // never force agentic mode ON over a possibly-persisted OFF preference. + if (!this.#initialized) { + return false + } + // Initialized with genuinely no stored preference → first-time user default. return true } diff --git a/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts b/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts index 716bd2a7c7..5daf5f627e 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts @@ -28,7 +28,8 @@ type DeferredHandler = { reject: (err: Error) => void } export class ChatSessionService { - public pairProgrammingMode: boolean = true + // Session creation alone must not grant write access before mode initialization. + public pairProgrammingMode: boolean = false public contextListSent: boolean = false public isMemoryBankGeneration: boolean = false #modelId: string | undefined