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 fecbdebe13..5aa341d232 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 @@ -50,6 +50,7 @@ describe('TabBarController', () => { afterEach(() => { sinon.restore() clock.restore() + delete process.env.JUPYTER_LAB // Clean up JupyterLab environment variables testFeatures.dispose() }) @@ -540,7 +541,7 @@ describe('TabBarController', () => { }) }) - it('should only load chats once', async () => { + it('should only load chats once in non-JupyterLab environments', async () => { const mockTabs = [{ historyId: 'history1', conversations: [{ messages: [] }] }] as unknown as Tab[] ;(chatHistoryDb.getOpenTabs as sinon.SinonStub).returns(mockTabs) @@ -559,6 +560,22 @@ describe('TabBarController', () => { result: 'Succeeded', }) }) + + it('should allow multiple loads in JupyterLab environment', async () => { + // Set JupyterLab environment + process.env.JUPYTER_LAB = 'true' + + const mockTabs = [{ historyId: 'history1', conversations: [{ messages: [] }] }] as unknown as Tab[] + ;(chatHistoryDb.getOpenTabs as sinon.SinonStub).returns(mockTabs) + + const restoreTabStub = sinon.stub(tabBarController, 'restoreTab') + + await tabBarController.loadChats() + await tabBarController.loadChats() // Second call should NOT be ignored in JupyterLab + + sinon.assert.calledTwice(restoreTabStub) + sinon.assert.calledTwice(telemetryService.emitLoadHistory as sinon.SinonStub) + }) }) }) 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 62f5b2ab0b..4531a9c589 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tabBarController.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tabBarController.ts @@ -313,10 +313,17 @@ export class TabBarController { * When IDE is opened, restore chats that were previously open in IDE for the current workspace. */ async loadChats() { - if (this.#loadedChats) { + const isJupyterLab = this.isJupyterLabEnvironment() + + // For non-JupyterLab environments, prevent multiple loads + if (!isJupyterLab && this.#loadedChats) { return } - this.#loadedChats = true + + if (!isJupyterLab) { + this.#loadedChats = true + } + const openConversations = this.#chatHistoryDb.getOpenTabs() if (openConversations) { for (const conversation of openConversations) { @@ -332,6 +339,18 @@ export class TabBarController { } } + /** + * Determines if the environment is JupyterLab. + */ + private isJupyterLabEnvironment(): boolean { + try { + return process.env.JUPYTER_LAB === 'true' + } catch (error) { + this.#features.logging.error(`Failed to read JUPYTER_LAB environment variable: ${error}`) + return false + } + } + public static enableChatExport(params?: InitializeParams) { if (params?.initializationOptions?.aws?.awsClientCapabilities?.window?.showSaveFileDialog) { // Export Chat UX flow relies on show Save File dialog protocol supported by client