fix(amazonqFeatureDev): race conditions on session creation aws#4724

Problem:
When a customer starts a /dev chat with a message, e.g. /dev <do something>, the
connector logic calls two event handlers in amazonqFeatureDev logic parallely
that try to create a session, this creates a race condition where the customer
could get to code generation with the wrong session, i.e. one with no
conversationId.

Solution:
Lock the getSession function to mitigate concurrent calls.

This is a similar approach to the one taken in JetBrain's toolkit:
https://github.com/aws/aws-toolkit-jetbrains/blob/ec168fc6e9306f22f07945fd516d6acdc7fe396e/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/storage/ChatSessionStorage.kt#L18

Author: sannicm

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