fix(amazonq): stop continuous monitor when WCS sees ServiceQuotaExceeded (#1957)

Co-authored-by: Jiatong Li <[email protected]>

Author: LiGaCu

server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.ts

@@ -221,6 +221,7 @@ export const WorkspaceContextServer = (): Server => features => {
             isLoggedInUsingBearerToken(credentialsProvider) &&
             abTestingEnabled &&
             !workspaceFolderManager.getOptOutStatus() &&
+            !workspaceFolderManager.getServiceQuotaExceededStatus() &&
             workspaceIdentifier
         )
     }
@@ -302,6 +303,7 @@ export const WorkspaceContextServer = (): Server => features => {
                     await evaluateABTesting()
                     isWorkflowInitialized = true
 
+                    workspaceFolderManager.resetAdminOptOutAndServiceQuotaStatus()
                     if (!isUserEligibleForWorkspaceContext()) {
                         return
                     }

server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.test.ts

@@ -0,0 +1,131 @@
+import { WorkspaceFolderManager } from './workspaceFolderManager'
+import sinon, { stubInterface, StubbedInstance } from 'ts-sinon'
+import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager'
+import { CredentialsProvider, Logging } from '@aws/language-server-runtimes/server-interface'
+import { DependencyDiscoverer } from './dependency/dependencyDiscoverer'
+import { WorkspaceFolder } from 'vscode-languageserver-protocol'
+import { ArtifactManager } from './artifactManager'
+import { CodeWhispererServiceToken } from '../../shared/codeWhispererService'
+import { CreateWorkspaceResponse } from '../../client/token/codewhispererbearertokenclient'
+import { AWSError } from 'aws-sdk'
+
+describe('WorkspaceFolderManager', () => {
+    let mockServiceManager: StubbedInstance<AmazonQTokenServiceManager>
+    let mockLogging: StubbedInstance<Logging>
+    let mockCredentialsProvider: StubbedInstance<CredentialsProvider>
+    let mockDependencyDiscoverer: StubbedInstance<DependencyDiscoverer>
+    let mockArtifactManager: StubbedInstance<ArtifactManager>
+    let mockCodeWhispererService: StubbedInstance<CodeWhispererServiceToken>
+    let workspaceFolderManager: WorkspaceFolderManager
+
+    beforeEach(() => {
+        mockServiceManager = stubInterface<AmazonQTokenServiceManager>()
+        mockLogging = stubInterface<Logging>()
+        mockCredentialsProvider = stubInterface<CredentialsProvider>()
+        mockDependencyDiscoverer = stubInterface<DependencyDiscoverer>()
+        mockArtifactManager = stubInterface<ArtifactManager>()
+        mockCodeWhispererService = stubInterface<CodeWhispererServiceToken>()
+
+        mockServiceManager.getCodewhispererService.returns(mockCodeWhispererService)
+    })
+
+    afterEach(() => {
+        sinon.restore()
+    })
+
+    describe('getServiceQuotaExceededStatus', () => {
+        it('should return true when service quota is exceeded', async () => {
+            // Setup
+            const workspaceFolders: WorkspaceFolder[] = [
+                {
+                    uri: 'file:///test/workspace',
+                    name: 'test-workspace',
+                },
+            ]
+
+            // Mock the createWorkspace method to throw a ServiceQuotaExceededException
+            const mockError: AWSError = {
+                name: 'ServiceQuotaExceededException',
+                message: 'You have too many active running workspaces.',
+                code: 'ServiceQuotaExceededException',
+                time: new Date(),
+                retryable: false,
+                statusCode: 400,
+            }
+
+            mockCodeWhispererService.createWorkspace.rejects(mockError)
+
+            // Create the WorkspaceFolderManager instance using the static createInstance method
+            workspaceFolderManager = WorkspaceFolderManager.createInstance(
+                mockServiceManager,
+                mockLogging,
+                mockArtifactManager,
+                mockDependencyDiscoverer,
+                workspaceFolders,
+                mockCredentialsProvider,
+                'test-workspace-identifier'
+            )
+
+            // Spy on clearAllWorkspaceResources and related methods
+            const clearAllWorkspaceResourcesSpy = sinon.stub(
+                workspaceFolderManager as any,
+                'clearAllWorkspaceResources'
+            )
+
+            // Act - trigger the createNewWorkspace method which sets isServiceQuotaExceeded
+            await (workspaceFolderManager as any).createNewWorkspace()
+
+            // Assert
+            expect(workspaceFolderManager.getServiceQuotaExceededStatus()).toBe(true)
+
+            // Verify that clearAllWorkspaceResources was called
+            sinon.assert.calledOnce(clearAllWorkspaceResourcesSpy)
+        })
+
+        it('should return false when service quota is not exceeded', async () => {
+            // Setup
+            const workspaceFolders: WorkspaceFolder[] = [
+                {
+                    uri: 'file:///test/workspace',
+                    name: 'test-workspace',
+                },
+            ]
+
+            // Mock successful response
+            const mockResponse: CreateWorkspaceResponse = {
+                workspace: {
+                    workspaceId: 'test-workspace-id',
+                    workspaceStatus: 'RUNNING',
+                },
+            }
+
+            mockCodeWhispererService.createWorkspace.resolves(mockResponse as any)
+
+            // Create the WorkspaceFolderManager instance using the static createInstance method
+            workspaceFolderManager = WorkspaceFolderManager.createInstance(
+                mockServiceManager,
+                mockLogging,
+                mockArtifactManager,
+                mockDependencyDiscoverer,
+                workspaceFolders,
+                mockCredentialsProvider,
+                'test-workspace-identifier'
+            )
+
+            // Spy on clearAllWorkspaceResources
+            const clearAllWorkspaceResourcesSpy = sinon.stub(
+                workspaceFolderManager as any,
+                'clearAllWorkspaceResources'
+            )
+
+            // Act - trigger the createNewWorkspace method
+            await (workspaceFolderManager as any).createNewWorkspace()
+
+            // Assert
+            expect(workspaceFolderManager.getServiceQuotaExceededStatus()).toBe(false)
+
+            // Verify that clearAllWorkspaceResources was not called
+            sinon.assert.notCalled(clearAllWorkspaceResourcesSpy)
+        })
+    })
+})

server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.ts

@@ -19,6 +19,7 @@ import { DependencyDiscoverer } from './dependency/dependencyDiscoverer'
 import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager'
 import { URI } from 'vscode-uri'
 import path = require('path')
+import { isAwsError } from '../../shared/utils'
 
 interface WorkspaceState {
     remoteWorkspaceState: WorkspaceStatus
@@ -47,12 +48,14 @@ export class WorkspaceFolderManager {
     private credentialsProvider: CredentialsProvider
     private readonly INITIAL_CHECK_INTERVAL = 40 * 1000 // 40 seconds
     private readonly INITIAL_CONNECTION_TIMEOUT = 2 * 60 * 1000 // 2 minutes
-    private readonly CONTINUOUS_MONITOR_INTERVAL = 5 * 60 * 1000 // 30 minutes
+    private readonly CONTINUOUS_MONITOR_INTERVAL = 30 * 60 * 1000 // 30 minutes
     private readonly MESSAGE_PUBLISH_INTERVAL: number = 100 // 100 milliseconds
     private continuousMonitorInterval: NodeJS.Timeout | undefined
     private optOutMonitorInterval: NodeJS.Timeout | undefined
     private messageQueueConsumerInterval: NodeJS.Timeout | undefined
     private isOptedOut: boolean = false
+    // Tracks if the user has reached their maximum allowed remote workspaces quota
+    private isServiceQuotaExceeded: boolean = false
 
     static createInstance(
         serviceManager: AmazonQTokenServiceManager,
@@ -135,6 +138,15 @@ export class WorkspaceFolderManager {
         return this.isOptedOut
     }
 
+    getServiceQuotaExceededStatus(): boolean {
+        return this.isServiceQuotaExceeded
+    }
+
+    resetAdminOptOutAndServiceQuotaStatus(): void {
+        this.isOptedOut = false
+        this.isServiceQuotaExceeded = false
+    }
+
     getWorkspaceState(): WorkspaceState {
         return this.workspaceState
     }
@@ -343,7 +355,7 @@ export class WorkspaceFolderManager {
         await this.checkRemoteWorkspaceStatusAndReact(true)
 
         // Set up continuous monitoring which periodically invokes checkRemoteWorkspaceStatusAndReact
-        if (!this.isOptedOut && this.continuousMonitorInterval === undefined) {
+        if (!this.isOptedOut && !this.isServiceQuotaExceeded && this.continuousMonitorInterval === undefined) {
             this.logging.log(`Starting continuous monitor for workspace [${this.workspaceIdentifier}]`)
             this.continuousMonitorInterval = setInterval(async () => {
                 try {
@@ -592,6 +604,13 @@ export class WorkspaceFolderManager {
 
     private async createNewWorkspace() {
         const createWorkspaceResult = await this.createWorkspace(this.workspaceIdentifier)
+
+        this.isServiceQuotaExceeded = createWorkspaceResult.isServiceQuotaExceeded
+        if (this.isServiceQuotaExceeded) {
+            // Stop continuous monitor and all actions
+            this.clearAllWorkspaceResources()
+        }
+
         const workspaceDetails = createWorkspaceResult.response
         if (!workspaceDetails) {
             this.logging.warn(`Failed to create remote workspace for [${this.workspaceIdentifier}]`)
@@ -712,23 +731,31 @@ export class WorkspaceFolderManager {
         return { metadata, optOut, error }
     }
 
-    private async createWorkspace(workspaceRoot: WorkspaceRoot) {
+    private async createWorkspace(workspaceRoot: WorkspaceRoot): Promise<{
+        response: CreateWorkspaceResponse | undefined | null
+        isServiceQuotaExceeded: boolean
+        error: any
+    }> {
         let response: CreateWorkspaceResponse | undefined | null
+        let isServiceQuotaExceeded = false
+        let error: any
         try {
             response = await this.serviceManager.getCodewhispererService().createWorkspace({
                 workspaceRoot: workspaceRoot,
             })
-            return { response, error: null }
         } catch (e: any) {
             this.logging.warn(
                 `Error while creating workspace (${workspaceRoot}): ${e.message}. Error is ${e.retryable ? '' : 'not'} retryable}`
             )
-            const error = {
+            if (isAwsError(e) && e.code === 'ServiceQuotaExceededException') {
+                isServiceQuotaExceeded = true
+            }
+            error = {
                 message: e.message,
                 retryable: e.retryable ?? false,
                 originalError: e,
             }
-            return { response: null, error }
         }
+        return { response, isServiceQuotaExceeded, error }
     }
 }