Skip to content

Commit 76d35cd

Browse files
author
Jiatong Li
committed
fix(amazonq): skips continuous monitoring when WCS sees workspace as idle
1 parent 9949c6d commit 76d35cd

File tree

9 files changed

+278
-111
lines changed

9 files changed

+278
-111
lines changed

‎server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ import { McpManager } from './tools/mcp/mcpManager'
6767
import { AgenticChatResultStream } from './agenticChatResultStream'
6868
import { AgenticChatError } from './errors'
6969
import * as sharedUtils from '../../shared/utils'
70+
import { IdleWorkspaceManager } from '../workspaceContext/IdleWorkspaceManager'
7071

7172
describe('AgenticChatController', () => {
7273
let mcpInstanceStub: sinon.SinonStub
@@ -475,6 +476,15 @@ describe('AgenticChatController', () => {
475476
assert.strictEqual(typeof session.conversationId, 'string')
476477
})
477478

479+
it('invokes IdleWorkspaceManager recordActivityTimestamp', async () => {
480+
const recordActivityTimestampStub = sinon.stub(IdleWorkspaceManager, 'recordActivityTimestamp')
481+
482+
await chatController.onChatPrompt({ tabId: mockTabId, prompt: { prompt: 'Hello' } }, mockCancellationToken)
483+
484+
sinon.assert.calledOnce(recordActivityTimestampStub)
485+
recordActivityTimestampStub.restore()
486+
})
487+
478488
it('includes chat history from the database in the request input', async () => {
479489
// Mock chat history
480490
const mockHistory = [

‎server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,7 @@ import { ActiveUserTracker } from '../../shared/activeUserTracker'
230230
import { UserContext } from '../../client/token/codewhispererbearertokenclient'
231231
import { CodeWhispererServiceToken } from '../../shared/codeWhispererService'
232232
import { DisplayFindings } from './tools/qCodeAnalysis/displayFindings'
233+
import { IdleWorkspaceManager } from '../workspaceContext/IdleWorkspaceManager'
233234

234235
type ChatHandlers = Omit<
235236
LspHandlers<Chat>,
@@ -722,6 +723,8 @@ export class AgenticChatController implements ChatHandlers {
722723
// Phase 1: Initial Setup - This happens only once
723724
params.prompt.prompt = sanitizeInput(params.prompt.prompt || '')
724725

726+
await IdleWorkspaceManager.recordActivityTimestamp()
727+
725728
const maybeDefaultResponse = !params.prompt.command && getDefaultChatResponse(params.prompt.prompt)
726729
if (maybeDefaultResponse) {
727730
return maybeDefaultResponse
@@ -3266,6 +3269,9 @@ export class AgenticChatController implements ChatHandlers {
32663269
const metric = new Metric<AddMessageEvent>({
32673270
cwsprChatConversationType: 'Chat',
32683271
})
3272+
3273+
await IdleWorkspaceManager.recordActivityTimestamp()
3274+
32693275
const triggerContext = await this.#getInlineChatTriggerContext(params)
32703276

32713277
let response: ChatCommandOutput

‎server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.test.ts‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ import { INVALID_TOKEN } from '../../shared/constants'
6161
import { AmazonQError } from '../../shared/amazonQServiceManager/errors'
6262
import * as path from 'path'
6363
import { CONTEXT_CHARACTERS_LIMIT } from './constants'
64+
import { IdleWorkspaceManager } from '../workspaceContext/IdleWorkspaceManager'
6465

6566
const updateConfiguration = async (
6667
features: TestFeatures,
@@ -770,6 +771,22 @@ describe('CodeWhisperer Server', () => {
770771
assert.rejects(promise, ResponseError)
771772
})
772773

774+
it('invokes IdleWorkspaceManager recordActivityTimestamp', async () => {
775+
const recordActivityTimestampStub = sinon.stub(IdleWorkspaceManager, 'recordActivityTimestamp')
776+
777+
await features.doInlineCompletionWithReferences(
778+
{
779+
textDocument: { uri: SOME_FILE.uri },
780+
position: { line: 0, character: 0 },
781+
context: { triggerKind: InlineCompletionTriggerKind.Invoked },
782+
},
783+
CancellationToken.None
784+
)
785+
786+
sinon.assert.calledOnce(recordActivityTimestampStub)
787+
recordActivityTimestampStub.restore()
788+
})
789+
773790
describe('Supplemental Context', () => {
774791
it('should send supplemental context when using token authentication', async () => {
775792
const test_service = sinon.createStubInstance(

‎server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ import {
5656
import { DocumentChangedListener } from './documentChangedListener'
5757
import { EditCompletionHandler } from './editCompletionHandler'
5858
import { EMPTY_RESULT } from './constants'
59+
import { IdleWorkspaceManager } from '../workspaceContext/IdleWorkspaceManager'
5960

6061
const mergeSuggestionsWithRightContext = (
6162
rightFileContext: string,
@@ -141,6 +142,8 @@ export const CodewhispererServerFactory =
141142
// 2. it is not designed to handle concurrent changes to these state variables.
142143
// when one handler is at the API call stage, it has not yet update the session state
143144
// but another request can start, causing the state to be incorrect.
145+
await IdleWorkspaceManager.recordActivityTimestamp()
146+
144147
if (isOnInlineCompletionHandlerInProgress) {
145148
logging.log(`Skip concurrent inline completion`)
146149
return EMPTY_RESULT

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

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
import { IdleWorkspaceManager } from './IdleWorkspaceManager'
2+
import { WorkspaceFolderManager } from './workspaceFolderManager'
3+
import sinon, { stubInterface, StubbedInstance } from 'ts-sinon'
4+
5+
describe('IdleWorkspaceManager', () => {
6+
let clock: sinon.SinonFakeTimers
7+
let mockWorkspaceFolderManager: StubbedInstance<WorkspaceFolderManager>
8+
9+
beforeEach(() => {
10+
clock = sinon.useFakeTimers()
11+
mockWorkspaceFolderManager = stubInterface<WorkspaceFolderManager>()
12+
sinon.stub(WorkspaceFolderManager, 'getInstance').returns(mockWorkspaceFolderManager)
13+
sinon.stub(console, 'error')
14+
})
15+
16+
afterEach(() => {
17+
clock.restore()
18+
sinon.restore()
19+
})
20+
21+
describe('isSessionIdle', () => {
22+
it('should return false when session is not idle', () => {
23+
IdleWorkspaceManager.recordActivityTimestamp()
24+
25+
const result = IdleWorkspaceManager.isSessionIdle()
26+
27+
expect(result).toBe(false)
28+
})
29+
30+
it('should return true when session exceeds idle threshold', () => {
31+
IdleWorkspaceManager.recordActivityTimestamp()
32+
clock.tick(31 * 60 * 1000) // 31 minutes
33+
34+
const result = IdleWorkspaceManager.isSessionIdle()
35+
36+
expect(result).toBe(true)
37+
})
38+
})
39+
40+
describe('recordActivityTimestamp', () => {
41+
it('should update activity timestamp', async () => {
42+
const initialTime = Date.now()
43+
clock.tick(5000)
44+
45+
await IdleWorkspaceManager.recordActivityTimestamp()
46+
47+
expect(IdleWorkspaceManager.isSessionIdle()).toBe(false)
48+
})
49+
50+
it('should not trigger workspace check when session was not idle', async () => {
51+
mockWorkspaceFolderManager.isContinuousMonitoringStopped.returns(false)
52+
53+
await IdleWorkspaceManager.recordActivityTimestamp()
54+
55+
sinon.assert.notCalled(mockWorkspaceFolderManager.checkRemoteWorkspaceStatusAndReact)
56+
})
57+
58+
it('should trigger workspace check when session was idle and monitoring is active', async () => {
59+
// Make session idle first
60+
clock.tick(31 * 60 * 1000)
61+
mockWorkspaceFolderManager.isContinuousMonitoringStopped.returns(false)
62+
mockWorkspaceFolderManager.checkRemoteWorkspaceStatusAndReact.resolves()
63+
64+
await IdleWorkspaceManager.recordActivityTimestamp()
65+
66+
sinon.assert.calledOnce(mockWorkspaceFolderManager.checkRemoteWorkspaceStatusAndReact)
67+
})
68+
69+
it('should not trigger workspace check when session was idle but monitoring is stopped', async () => {
70+
clock.tick(31 * 60 * 1000)
71+
mockWorkspaceFolderManager.isContinuousMonitoringStopped.returns(true)
72+
73+
await IdleWorkspaceManager.recordActivityTimestamp()
74+
75+
sinon.assert.notCalled(mockWorkspaceFolderManager.checkRemoteWorkspaceStatusAndReact)
76+
})
77+
})
78+
})

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
import { WorkspaceFolderManager } from './workspaceFolderManager'
2+
3+
export class IdleWorkspaceManager {
4+
private static readonly idleThreshold = 30 * 60 * 1000 // 30 minutes
5+
private static lastActivityTimestamp = Date.now()
6+
7+
private constructor() {}
8+
9+
/**
10+
* Records activity timestamp and triggers workspace status check if session was idle.
11+
*
12+
* When transitioning from idle to active, proactively checks remote workspace status
13+
* (if continuous monitoring is enabled) without blocking the current operation.
14+
*/
15+
public static async recordActivityTimestamp(): Promise<void> {
16+
const wasSessionIdle = IdleWorkspaceManager.isSessionIdle()
17+
IdleWorkspaceManager.lastActivityTimestamp = Date.now()
18+
19+
const workspaceFolderManager = WorkspaceFolderManager.getInstance()
20+
if (workspaceFolderManager && wasSessionIdle && !workspaceFolderManager.isContinuousMonitoringStopped()) {
21+
// Proactively check the remote workspace status instead of waiting for the next scheduled check
22+
// Fire and forget - don't await to avoid blocking
23+
workspaceFolderManager.checkRemoteWorkspaceStatusAndReact().catch(err => {
24+
console.error('Failed to check remote workspace status', err)
25+
})
26+
}
27+
}
28+
29+
public static isSessionIdle(): boolean {
30+
const currentTime = Date.now()
31+
const timeSinceLastActivity = currentTime - IdleWorkspaceManager.lastActivityTimestamp
32+
return timeSinceLastActivity > IdleWorkspaceManager.idleThreshold
33+
}
34+
}

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,6 @@ export const WorkspaceContextServer = (): Server => features => {
221221
isLoggedInUsingBearerToken(credentialsProvider) &&
222222
abTestingEnabled &&
223223
!workspaceFolderManager.getOptOutStatus() &&
224-
!workspaceFolderManager.getServiceQuotaExceededStatus() &&
225224
workspaceIdentifier
226225
)
227226
}
@@ -303,7 +302,7 @@ export const WorkspaceContextServer = (): Server => features => {
303302
await evaluateABTesting()
304303
isWorkflowInitialized = true
305304

306-
workspaceFolderManager.resetAdminOptOutAndServiceQuotaStatus()
305+
workspaceFolderManager.resetAdminOptOutStatus()
307306
if (!isUserEligibleForWorkspaceContext()) {
308307
return
309308
}

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

Lines changed: 49 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@ import { DependencyDiscoverer } from './dependency/dependencyDiscoverer'
66
import { WorkspaceFolder } from 'vscode-languageserver-protocol'
77
import { ArtifactManager } from './artifactManager'
88
import { CodeWhispererServiceToken } from '../../shared/codeWhispererService'
9-
import { CreateWorkspaceResponse } from '../../client/token/codewhispererbearertokenclient'
10-
import { AWSError } from 'aws-sdk'
9+
import { ListWorkspaceMetadataResponse } from '../../client/token/codewhispererbearertokenclient'
10+
import { IdleWorkspaceManager } from './IdleWorkspaceManager'
1111

1212
describe('WorkspaceFolderManager', () => {
1313
let mockServiceManager: StubbedInstance<AmazonQTokenServiceManager>
@@ -33,8 +33,8 @@ describe('WorkspaceFolderManager', () => {
3333
sinon.restore()
3434
})
3535

36-
describe('getServiceQuotaExceededStatus', () => {
37-
it('should return true when service quota is exceeded', async () => {
36+
describe('checkRemoteWorkspaceStatusAndReact', () => {
37+
it('should check and react when IDE session is not idle', async () => {
3838
// Setup
3939
const workspaceFolders: WorkspaceFolder[] = [
4040
{
@@ -43,17 +43,20 @@ describe('WorkspaceFolderManager', () => {
4343
},
4444
]
4545

46-
// Mock the createWorkspace method to throw a ServiceQuotaExceededException
47-
const mockError: AWSError = {
48-
name: 'ServiceQuotaExceededException',
49-
message: 'You have too many active running workspaces.',
50-
code: 'ServiceQuotaExceededException',
51-
time: new Date(),
52-
retryable: false,
53-
statusCode: 400,
46+
// Mock IdleSessionManager to return false (not idle)
47+
sinon.stub(IdleWorkspaceManager, 'isSessionIdle').returns(false)
48+
49+
// Mock successful response
50+
const mockResponse: ListWorkspaceMetadataResponse = {
51+
workspaces: [
52+
{
53+
workspaceId: 'test-workspace-id',
54+
workspaceStatus: 'CREATED',
55+
},
56+
],
5457
}
5558

56-
mockCodeWhispererService.createWorkspace.rejects(mockError)
59+
mockCodeWhispererService.listWorkspaceMetadata.resolves(mockResponse as any)
5760

5861
// Create the WorkspaceFolderManager instance using the static createInstance method
5962
workspaceFolderManager = WorkspaceFolderManager.createInstance(
@@ -66,23 +69,24 @@ describe('WorkspaceFolderManager', () => {
6669
'test-workspace-identifier'
6770
)
6871

69-
// Spy on clearAllWorkspaceResources and related methods
70-
const clearAllWorkspaceResourcesSpy = sinon.stub(
72+
// Spy on resetWebSocketClient
73+
const resetWebSocketClientSpy = sinon.stub(workspaceFolderManager as any, 'resetWebSocketClient')
74+
75+
// Spy on handleWorkspaceCreatedState
76+
const handleWorkspaceCreatedStateSpy = sinon.stub(
7177
workspaceFolderManager as any,
72-
'clearAllWorkspaceResources'
78+
'handleWorkspaceCreatedState'
7379
)
7480

75-
// Act - trigger the createNewWorkspace method which sets isServiceQuotaExceeded
76-
await (workspaceFolderManager as any).createNewWorkspace()
81+
// Act - trigger the checkRemoteWorkspaceStatusAndReact method
82+
await workspaceFolderManager.checkRemoteWorkspaceStatusAndReact()
7783

78-
// Assert
79-
expect(workspaceFolderManager.getServiceQuotaExceededStatus()).toBe(true)
80-
81-
// Verify that clearAllWorkspaceResources was called
82-
sinon.assert.calledOnce(clearAllWorkspaceResourcesSpy)
84+
// Verify that resetWebSocketClient was called once
85+
sinon.assert.notCalled(resetWebSocketClientSpy)
86+
sinon.assert.calledOnce(handleWorkspaceCreatedStateSpy)
8387
})
8488

85-
it('should return false when service quota is not exceeded', async () => {
89+
it('should skip checking and reacting when IDE session is idle', async () => {
8690
// Setup
8791
const workspaceFolders: WorkspaceFolder[] = [
8892
{
@@ -91,15 +95,20 @@ describe('WorkspaceFolderManager', () => {
9195
},
9296
]
9397

98+
// Mock IdleSessionManager to return true (idle)
99+
sinon.stub(IdleWorkspaceManager, 'isSessionIdle').returns(true)
100+
94101
// Mock successful response
95-
const mockResponse: CreateWorkspaceResponse = {
96-
workspace: {
97-
workspaceId: 'test-workspace-id',
98-
workspaceStatus: 'RUNNING',
99-
},
102+
const mockResponse: ListWorkspaceMetadataResponse = {
103+
workspaces: [
104+
{
105+
workspaceId: 'test-workspace-id',
106+
workspaceStatus: 'CREATED',
107+
},
108+
],
100109
}
101110

102-
mockCodeWhispererService.createWorkspace.resolves(mockResponse as any)
111+
mockCodeWhispererService.listWorkspaceMetadata.resolves(mockResponse as any)
103112

104113
// Create the WorkspaceFolderManager instance using the static createInstance method
105114
workspaceFolderManager = WorkspaceFolderManager.createInstance(
@@ -112,20 +121,18 @@ describe('WorkspaceFolderManager', () => {
112121
'test-workspace-identifier'
113122
)
114123

115-
// Spy on clearAllWorkspaceResources
116-
const clearAllWorkspaceResourcesSpy = sinon.stub(
117-
workspaceFolderManager as any,
118-
'clearAllWorkspaceResources'
119-
)
124+
// Spy on resetWebSocketClient
125+
const resetWebSocketClientSpy = sinon.stub(workspaceFolderManager as any, 'resetWebSocketClient')
120126

121-
// Act - trigger the createNewWorkspace method
122-
await (workspaceFolderManager as any).createNewWorkspace()
127+
// Act - trigger the checkRemoteWorkspaceStatusAndReact method
128+
await workspaceFolderManager.checkRemoteWorkspaceStatusAndReact()
123129

124-
// Assert
125-
expect(workspaceFolderManager.getServiceQuotaExceededStatus()).toBe(false)
126-
127-
// Verify that clearAllWorkspaceResources was not called
128-
sinon.assert.notCalled(clearAllWorkspaceResourcesSpy)
130+
// Verify that resetWebSocketClient was called once
131+
sinon.assert.calledOnce(resetWebSocketClientSpy)
132+
sinon.assert.calledWith(
133+
mockLogging.log,
134+
sinon.match(/Session is idle, skipping remote workspace status check/)
135+
)
129136
})
130137
})
131138
})

0 commit comments

Comments
 (0)