Merge branch 'main' into feature/mcp

Author: bywang56

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

@@ -2,7 +2,7 @@ import { WorkspaceFolderManager } from './workspaceFolderManager'
 
 export class IdleWorkspaceManager {
     private static readonly idleThreshold = 30 * 60 * 1000 // 30 minutes
-    private static lastActivityTimestamp = 0 // treat session as idle as the start
+    private static lastActivityTimestamp = 0 // treat session as idle at the start
 
     private constructor() {}
 
@@ -30,6 +30,10 @@ export class IdleWorkspaceManager {
         }
     }
 
+    public static setSessionAsIdle(): void {
+        IdleWorkspaceManager.lastActivityTimestamp = 0
+    }
+
     public static isSessionIdle(): boolean {
         const currentTime = Date.now()
         const timeSinceLastActivity = currentTime - IdleWorkspaceManager.lastActivityTimestamp

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.isFeatureDisabled() &&
             workspaceIdentifier
         )
     }
@@ -302,7 +303,7 @@ export const WorkspaceContextServer = (): Server => features => {
                     await evaluateABTesting()
                     isWorkflowInitialized = true
 
-                    workspaceFolderManager.resetAdminOptOutStatus()
+                    workspaceFolderManager.resetAdminOptOutAndFeatureDisabledStatus()
                     if (!isUserEligibleForWorkspaceContext()) {
                         return
                     }

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

@@ -8,6 +8,7 @@ import { ArtifactManager } from './artifactManager'
 import { CodeWhispererServiceToken } from '../../shared/codeWhispererService'
 import { ListWorkspaceMetadataResponse } from '../../client/token/codewhispererbearertokenclient'
 import { IdleWorkspaceManager } from './IdleWorkspaceManager'
+import { AWSError } from 'aws-sdk'
 
 describe('WorkspaceFolderManager', () => {
     let mockServiceManager: StubbedInstance<AmazonQTokenServiceManager>
@@ -135,4 +136,93 @@ describe('WorkspaceFolderManager', () => {
             )
         })
     })
+
+    describe('isFeatureDisabled', () => {
+        it('should return true when feature is disabled', async () => {
+            // Setup
+            const workspaceFolders: WorkspaceFolder[] = [
+                {
+                    uri: 'file:///test/workspace',
+                    name: 'test-workspace',
+                },
+            ]
+
+            // Mock listWorkspaceMetadata to throw AccessDeniedException with feature not supported
+            const mockError: AWSError = {
+                name: 'AccessDeniedException',
+                message: 'Feature is not supported',
+                code: 'AccessDeniedException',
+                time: new Date(),
+                retryable: false,
+                statusCode: 403,
+            }
+
+            mockCodeWhispererService.listWorkspaceMetadata.rejects(mockError)
+
+            // Create the WorkspaceFolderManager instance
+            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 listWorkspaceMetadata which sets feature disabled state
+            await (workspaceFolderManager as any).listWorkspaceMetadata()
+
+            // Assert
+            expect(workspaceFolderManager.isFeatureDisabled()).toBe(true)
+
+            // Verify that clearAllWorkspaceResources was called
+            sinon.assert.calledOnce(clearAllWorkspaceResourcesSpy)
+        })
+
+        it('should return false when feature is not disabled', async () => {
+            // Setup
+            const workspaceFolders: WorkspaceFolder[] = [
+                {
+                    uri: 'file:///test/workspace',
+                    name: 'test-workspace',
+                },
+            ]
+
+            // Mock successful response
+            const mockResponse: ListWorkspaceMetadataResponse = {
+                workspaces: [
+                    {
+                        workspaceId: 'test-workspace-id',
+                        workspaceStatus: 'RUNNING',
+                    },
+                ],
+            }
+
+            mockCodeWhispererService.listWorkspaceMetadata.resolves(mockResponse as any)
+
+            // Create the WorkspaceFolderManager instance
+            workspaceFolderManager = WorkspaceFolderManager.createInstance(
+                mockServiceManager,
+                mockLogging,
+                mockArtifactManager,
+                mockDependencyDiscoverer,
+                workspaceFolders,
+                mockCredentialsProvider,
+                'test-workspace-identifier'
+            )
+
+            // Act - trigger listWorkspaceMetadata
+            await (workspaceFolderManager as any).listWorkspaceMetadata()
+
+            // Assert
+            expect(workspaceFolderManager.isFeatureDisabled()).toBe(false)
+        })
+    })
 })

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

@@ -55,6 +55,7 @@ export class WorkspaceFolderManager {
     private optOutMonitorInterval: NodeJS.Timeout | undefined
     private messageQueueConsumerInterval: NodeJS.Timeout | undefined
     private isOptedOut: boolean = false
+    private featureDisabled: boolean = false // Serve as a server-side control. If true, stop WCS features
     private isCheckingRemoteWorkspaceStatus: boolean = false
     private isArtifactUploadedToRemoteWorkspace: boolean = false
 
@@ -139,8 +140,13 @@ export class WorkspaceFolderManager {
         return this.isOptedOut
     }
 
-    resetAdminOptOutStatus(): void {
+    resetAdminOptOutAndFeatureDisabledStatus(): void {
         this.isOptedOut = false
+        this.featureDisabled = false
+    }
+
+    isFeatureDisabled(): boolean {
+        return this.featureDisabled
     }
 
     getWorkspaceState(): WorkspaceState {
@@ -326,6 +332,7 @@ export class WorkspaceFolderManager {
         // Reset workspace ID to force operations to wait for new remote workspace information
         this.resetRemoteWorkspaceId()
 
+        IdleWorkspaceManager.setSessionAsIdle()
         this.isArtifactUploadedToRemoteWorkspace = false
 
         // Set up message queue consumer
@@ -371,7 +378,9 @@ export class WorkspaceFolderManager {
                         return resolve(false)
                     }
 
-                    const { metadata, optOut } = await this.listWorkspaceMetadata(this.workspaceIdentifier)
+                    const { metadata, optOut, featureDisabled } = await this.listWorkspaceMetadata(
+                        this.workspaceIdentifier
+                    )
 
                     if (optOut) {
                         this.logging.log(`User opted out during initial connection`)
@@ -381,6 +390,13 @@ export class WorkspaceFolderManager {
                         return resolve(false)
                     }
 
+                    if (featureDisabled) {
+                        this.logging.log(`Feature disabled during initial connection`)
+                        this.featureDisabled = true
+                        this.clearAllWorkspaceResources()
+                        return resolve(false)
+                    }
+
                     if (!metadata) {
                         // Continue polling by exiting only this iteration
                         return
@@ -437,7 +453,9 @@ export class WorkspaceFolderManager {
             }
 
             this.logging.log(`Checking remote workspace status for workspace [${this.workspaceIdentifier}]`)
-            const { metadata, optOut, error } = await this.listWorkspaceMetadata(this.workspaceIdentifier)
+            const { metadata, optOut, featureDisabled, error } = await this.listWorkspaceMetadata(
+                this.workspaceIdentifier
+            )
 
             if (optOut) {
                 this.logging.log('User opted out, clearing all resources and starting opt-out monitor')
@@ -447,6 +465,13 @@ export class WorkspaceFolderManager {
                 return
             }
 
+            if (featureDisabled) {
+                this.logging.log('Feature disabled, clearing all resources and stoping server-side indexing features')
+                this.featureDisabled = true
+                this.clearAllWorkspaceResources()
+                return
+            }
+
             if (error) {
                 // Do not do anything if we received an exception but not caused by optOut
                 return
@@ -528,7 +553,14 @@ export class WorkspaceFolderManager {
         if (this.optOutMonitorInterval === undefined) {
             const intervalId = setInterval(async () => {
                 try {
-                    const { optOut } = await this.listWorkspaceMetadata()
+                    const { optOut, featureDisabled } = await this.listWorkspaceMetadata()
+
+                    if (featureDisabled) {
+                        // Stop opt-out monitor when WCS feature is disabled from server-side
+                        this.featureDisabled = true
+                        clearInterval(intervalId)
+                        this.optOutMonitorInterval = undefined
+                    }
 
                     if (!optOut) {
                         this.isOptedOut = false
@@ -735,10 +767,12 @@ export class WorkspaceFolderManager {
     private async listWorkspaceMetadata(workspaceRoot?: WorkspaceRoot): Promise<{
         metadata: WorkspaceMetadata | undefined | null
         optOut: boolean
+        featureDisabled: boolean
         error: any
     }> {
         let metadata: WorkspaceMetadata | undefined | null
         let optOut = false
+        let featureDisabled = false
         let error: any
         try {
             const params = workspaceRoot ? { workspaceRoot } : {}
@@ -754,8 +788,11 @@ export class WorkspaceFolderManager {
                 this.logging.log(`User's administrator opted out server-side workspace context`)
                 optOut = true
             }
+            if (isAwsError(e) && e.code === 'AccessDeniedException' && e.message.includes('Feature is not supported')) {
+                featureDisabled = true
+            }
         }
-        return { metadata, optOut, error }
+        return { metadata, optOut, featureDisabled, error }
     }
 
     private async createWorkspace(workspaceRoot: WorkspaceRoot): Promise<{