feat(amazonq): support JupyterLab conversation history on refresh (#2325)

Allow multiple loadChats() calls in JupyterLab environment.

Co-authored-by: ruotingx <[email protected]>

Author: ruotingx

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)
+        })
     })
 })
 

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