Merge master into feature/cw-proactive-scan

Author: aws-toolkit-automation

packages/core/src/amazonqFeatureDev/storages/chatSession.ts

@@ -3,11 +3,14 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import AsyncLock from 'async-lock'
 import { Messenger } from '../controllers/chat/messenger/messenger'
 import { Session } from '../session/session'
 import { createSessionConfig } from '../session/sessionConfigFactory'
 
 export class ChatSessionStorage {
+    private lock = new AsyncLock()
+
     private sessions: Map<string, Session> = new Map()
 
     constructor(private readonly messenger: Messenger) {}
@@ -20,12 +23,19 @@ export class ChatSessionStorage {
     }
 
     public async getSession(tabID: string): Promise<Session> {
-        const sessionFromStorage = this.sessions.get(tabID)
-        if (sessionFromStorage === undefined) {
-            // If a session doesn't already exist just create it
-            return this.createSession(tabID)
-        }
-        return sessionFromStorage
+        /**
+         * The lock here is added in order to mitigate amazon Q's eventing fire & forget design when integrating with mynah-ui that creates a race condition here.
+         * The race condition happens when handleDevFeatureCommand in src/amazonq/webview/ui/quickActions/handler.ts is firing two events after each other to amazonqFeatureDev controller
+         * This eventually may make code generation fail as at the moment of that event it may get from the storage a session that has not been properly updated.
+         */
+        return this.lock.acquire(tabID, async () => {
+            const sessionFromStorage = this.sessions.get(tabID)
+            if (sessionFromStorage === undefined) {
+                // If a session doesn't already exist just create it
+                return this.createSession(tabID)
+            }
+            return sessionFromStorage
+        })
     }
 
     // Find all sessions that are currently waiting to be authenticated

packages/core/src/test/amazonqFeatureDev/session/chatSessionStorage.test.ts

@@ -0,0 +1,26 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import * as assert from 'assert'
+
+import { Messenger } from '../../../amazonqFeatureDev/controllers/chat/messenger/messenger'
+import { ChatSessionStorage } from '../../../amazonqFeatureDev/storages/chatSession'
+import { createMessenger } from '../utils'
+
+describe('chatSession', () => {
+    const tabID = '1234'
+    let chatStorage: ChatSessionStorage
+    let messenger: Messenger
+
+    beforeEach(() => {
+        messenger = createMessenger()
+        chatStorage = new ChatSessionStorage(messenger)
+    })
+
+    it('locks getSession', async () => {
+        const results = await Promise.allSettled([chatStorage.getSession(tabID), chatStorage.getSession(tabID)])
+        assert.equal(results.length, 2)
+        assert.deepStrictEqual(results[0], results[1])
+    })
+})