Instantiate a new chat client on every request (#351)

aws · Jun 26, 2024 · a090744 · a090744
1 parent d34ef83
commit a090744
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 17 deletions.
diff --git a/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.test.ts b/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.test.ts
@@ -11,7 +11,6 @@ import { ChatSessionService } from './chatSessionService'
 describe('Chat Session Service', () => {
     let generateAssistantResponseStub: sinon.SinonStub<any, any>
     let abortStub: sinon.SinonStub<any, any>
-    let destroyClientStub: sinon.SinonStub<any, any>
     let chatSessionService: ChatSessionService
     const mockCredentialsProvider: CredentialsProvider = {
         hasCredentials: sinon.stub().returns(true),
@@ -45,14 +44,11 @@ describe('Chat Session Service', () => {
             .stub(CodeWhispererStreaming.prototype, 'generateAssistantResponse')
             .callsFake(() => Promise.resolve(mockRequestResponse))
 
-        destroyClientStub = sinon.stub(CodeWhispererStreaming.prototype, 'destroy')
-
         chatSessionService = new ChatSessionService(mockCredentialsProvider)
     })
 
     afterEach(() => {
         generateAssistantResponseStub.restore()
-        destroyClientStub.restore()
         abortStub.restore()
     })
 
@@ -94,13 +90,12 @@ describe('Chat Session Service', () => {
         sinon.assert.calledOnce(abortStub)
     })
 
-    it('dispose() calls client.destroy and aborts outgoing requests', async () => {
+    it('dispose() calls aborts outgoing requests', async () => {
         await chatSessionService.generateAssistantResponse(mockRequestParams)
 
         chatSessionService.dispose()
 
         sinon.assert.calledOnce(abortStub)
-        sinon.assert.calledOnce(destroyClientStub)
     })
 
     it('clear() reset session id and aborts outgoing request', async () => {

diff --git a/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts b/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.ts
@@ -13,8 +13,9 @@ export class ChatSessionService {
     public shareCodeWhispererContentWithAWS = false
     readonly #codeWhispererRegion = 'us-east-1'
     readonly #codeWhispererEndpoint = 'https://codewhisperer.us-east-1.amazonaws.com/'
-    #client: CodeWhispererStreaming
     #abortController?: AbortController
+    #credentialsProvider: CredentialsProvider
+    #config?: CodeWhispererStreamingClientConfig
     #sessionId?: string
 
     public get sessionId(): string | undefined {
@@ -26,14 +27,8 @@ export class ChatSessionService {
     }
 
     constructor(credentialsProvider: CredentialsProvider, config?: CodeWhispererStreamingClientConfig) {
-        this.#client = new CodeWhispererStreaming({
-            region: this.#codeWhispererRegion,
-            endpoint: this.#codeWhispererEndpoint,
-            token: () => Promise.resolve({ token: getBearerTokenFromProvider(credentialsProvider) }),
-            retryStrategy: new ConfiguredRetryStrategy(0, (attempt: number) => 500 + attempt ** 10),
-            customUserAgent: '%Amazon-Q-For-LanguageServers%',
-            ...config,
-        })
+        this.#credentialsProvider = credentialsProvider
+        this.#config = config
     }
 
     public async generateAssistantResponse(
@@ -45,7 +40,16 @@ export class ChatSessionService {
             request.conversationState.conversationId = this.#sessionId
         }
 
-        const response = await this.#client.generateAssistantResponse(request, {
+        const client = new CodeWhispererStreaming({
+            region: this.#codeWhispererRegion,
+            endpoint: this.#codeWhispererEndpoint,
+            token: () => Promise.resolve({ token: getBearerTokenFromProvider(this.#credentialsProvider) }),
+            retryStrategy: new ConfiguredRetryStrategy(0, (attempt: number) => 500 + attempt ** 10),
+            customUserAgent: '%Amazon-Q-For-LanguageServers%',
+            ...this.#config,
+        })
+
+        const response = await client.generateAssistantResponse(request, {
             abortSignal: this.#abortController?.signal,
         })
 
@@ -61,7 +65,6 @@ export class ChatSessionService {
 
     public dispose(): void {
         this.#abortController?.abort()
-        this.#client.destroy()
     }
 
     public abortRequest(): void {