Skip to content

Commit 15d1b1f

Browse files
ashishrp-awsashishrp-aws
authored
feat(amazonq): adding classification based retry strategy for chat (#2234) (#2409)
1 parent a908195 commit 15d1b1f

22 files changed

+1832
-212
lines changed

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1013,6 +1013,14 @@ export class AgenticChatController implements ChatHandlers {
10131013
})
10141014
session.setConversationType('AgenticChat')
10151015

1016+
// Set up delay notification callback to show retry progress to users
1017+
session.setDelayNotificationCallback(notification => {
1018+
if (notification.thresholdExceeded) {
1019+
this.#log(`Updating progress message: ${notification.message}`)
1020+
void chatResultStream.updateProgressMessage(notification.message)
1021+
}
1022+
})
1023+
10161024
const additionalContext = await this.#additionalContextProvider.getAdditionalContext(
10171025
triggerContext,
10181026
params.tabId,

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

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,28 @@ export class AgenticChatResultStream {
194194
}
195195
}
196196

197+
async updateProgressMessage(message: string) {
198+
for (const block of this.#state.chatResultBlocks) {
199+
if (block.messageId?.startsWith(progressPrefix) && block.header?.status?.icon === 'progress') {
200+
const blockId = this.getMessageBlockId(block.messageId)
201+
if (blockId !== undefined) {
202+
const updatedBlock = {
203+
...block,
204+
header: {
205+
...block.header,
206+
status: {
207+
...block.header.status,
208+
text: message,
209+
},
210+
},
211+
}
212+
await this.overwriteResultBlock(updatedBlock, blockId)
213+
}
214+
break
215+
}
216+
}
217+
}
218+
197219
hasMessage(messageId: string): boolean {
198220
return this.#state.chatResultBlocks.some(block => block.messageId === messageId)
199221
}

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,41 @@ The summary should have following main sections:
8383
</example_output>
8484
`
8585

86+
// Retry Strategy Constants
87+
export const RETRY_BASE_DELAY_MS = 1000
88+
export const RETRY_MAX_DELAY_MS = 10000
89+
export const RETRY_JITTER_MIN = 0.5
90+
export const RETRY_JITTER_MAX = 1.0
91+
export const RETRY_DELAY_NOTIFICATION_THRESHOLD_MS = 2000
92+
export const RETRY_BACKOFF_MULTIPLIER = 2
93+
94+
// HTTP Status Codes
95+
export const HTTP_STATUS_TOO_MANY_REQUESTS = 429
96+
export const HTTP_STATUS_INTERNAL_SERVER_ERROR = 500
97+
98+
// Error Messages
99+
export const MONTHLY_LIMIT_ERROR_MARKER = 'MONTHLY_REQUEST_COUNT'
100+
export const CONTENT_LENGTH_EXCEEDS_THRESHOLD = 'CONTENT_LENGTH_EXCEEDS_THRESHOLD'
101+
export const HIGH_LOAD_ERROR_MESSAGE =
102+
'Encountered unexpectedly high load when processing the request, please try again.'
103+
export const SERVICE_UNAVAILABLE_EXCEPTION = 'ServiceUnavailableException'
104+
export const INSUFFICIENT_MODEL_CAPACITY = 'INSUFFICIENT_MODEL_CAPACITY'
105+
export const INVALID_MODEL_ID = 'INVALID_MODEL_ID'
106+
export const SERVICE_QUOTA_EXCEPTION = 'ServiceQuotaExceededException'
107+
export const MAXIMUM_CHAT_CONTENT_MESSAGE = 'Exceeded max chat context length.'
108+
109+
// Delay tracking constants
110+
export const MINOR_DELAY_THRESHOLD_MS = 2000 // 2 seconds
111+
export const MAJOR_DELAY_THRESHOLD_MS = 5000 // 5 seconds
112+
export const MAX_RETRY_DELAY_MS = 10000 // 10 seconds
113+
114+
// Stalled stream protection constants
115+
export const STALLED_STREAM_GRACE_PERIOD_MS = 300000 // 5 minutes
116+
export const STALLED_STREAM_CHECK_INTERVAL_MS = 1000 // 1 second
117+
118+
// Request attempt tracking
119+
export const MAX_REQUEST_ATTEMPTS = 3
120+
86121
// FsRead limits
87122
export const FSREAD_MAX_PER_FILE = 200_000
88123
export const FSREAD_MAX_TOTAL = 400_000

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

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,182 @@
1+
import { QDelayTrackingInterceptor, DelayNotification } from './delayInterceptor'
2+
import { expect } from 'chai'
3+
import * as sinon from 'sinon'
4+
5+
describe('QDelayTrackingInterceptor', () => {
6+
let interceptor: QDelayTrackingInterceptor
7+
let mockLogging: any
8+
let mockCallback: sinon.SinonSpy
9+
10+
beforeEach(() => {
11+
mockLogging = {
12+
log: sinon.spy(),
13+
debug: sinon.spy(),
14+
}
15+
mockCallback = sinon.spy()
16+
interceptor = new QDelayTrackingInterceptor(mockLogging)
17+
})
18+
19+
describe('setDelayNotificationCallback', () => {
20+
it('should set callback and log debug message', () => {
21+
interceptor.setDelayNotificationCallback(mockCallback)
22+
23+
expect(mockLogging.debug.calledWith('QDelayTrackingInterceptor: setDelayNotificationCallback called')).to.be
24+
.true
25+
})
26+
})
27+
28+
describe('beforeAttempt', () => {
29+
it('should log first attempt without delay calculation', () => {
30+
interceptor.beforeAttempt(1)
31+
32+
expect(mockLogging.debug.calledWith('QDelayTrackingInterceptor: Attempt 1')).to.be.true
33+
expect(
34+
mockLogging.debug.calledWith(
35+
'QDelayTrackingInterceptor: First attempt or no lastAttemptTime, skipping delay calculation'
36+
)
37+
).to.be.true
38+
})
39+
40+
it('should calculate delay for subsequent attempts', () => {
41+
const clock = sinon.useFakeTimers()
42+
43+
// First attempt
44+
interceptor.beforeAttempt(1)
45+
46+
// Wait a bit and make second attempt
47+
clock.tick(3000) // 3 seconds
48+
interceptor.beforeAttempt(2)
49+
50+
expect(mockLogging.debug.args.some((args: any) => args[0].includes('Delay'))).to.be.true
51+
52+
clock.restore()
53+
})
54+
55+
it('should send major delay notification for long delays', () => {
56+
interceptor.setDelayNotificationCallback(mockCallback)
57+
58+
const clock = sinon.useFakeTimers(1000)
59+
60+
// First attempt
61+
interceptor.beforeAttempt(1)
62+
63+
// Simulate 6 second delay (major threshold)
64+
clock.tick(6000)
65+
interceptor.beforeAttempt(2)
66+
67+
expect(mockCallback.calledOnce).to.be.true
68+
const call = mockCallback.getCall(0)
69+
expect(call.args[0].message).to.include('retrying within 10s')
70+
expect(call.args[0].attemptNumber).to.equal(2)
71+
expect(call.args[0].delay).to.equal(6)
72+
expect(call.args[0].thresholdExceeded).to.be.true
73+
74+
clock.restore()
75+
})
76+
77+
it('should send minor delay notification for medium delays', () => {
78+
interceptor.setDelayNotificationCallback(mockCallback)
79+
80+
const clock = sinon.useFakeTimers(1000)
81+
82+
// First attempt
83+
interceptor.beforeAttempt(1)
84+
85+
// Simulate 3 second delay (minor threshold)
86+
clock.tick(3000)
87+
interceptor.beforeAttempt(2)
88+
89+
expect(mockCallback.calledOnce).to.be.true
90+
const call = mockCallback.getCall(0)
91+
expect(call.args[0].message).to.include('retrying within 5s')
92+
expect(call.args[0].attemptNumber).to.equal(2)
93+
expect(call.args[0].delay).to.equal(3)
94+
expect(call.args[0].thresholdExceeded).to.be.true
95+
96+
clock.restore()
97+
})
98+
99+
it('should not notify for short delays', () => {
100+
interceptor.setDelayNotificationCallback(mockCallback)
101+
102+
const clock = sinon.useFakeTimers(1000)
103+
104+
// First attempt
105+
interceptor.beforeAttempt(1)
106+
107+
// Simulate 1 second delay (below threshold)
108+
clock.tick(1000)
109+
interceptor.beforeAttempt(2)
110+
111+
expect(mockCallback.called).to.be.false
112+
expect(
113+
mockLogging.debug.calledWith('QDelayTrackingInterceptor: Delay 1000ms below threshold, no notification')
114+
).to.be.true
115+
116+
clock.restore()
117+
})
118+
119+
it('should cap delay at maximum retry delay', () => {
120+
interceptor.setDelayNotificationCallback(mockCallback)
121+
122+
const clock = sinon.useFakeTimers(1000)
123+
124+
// First attempt
125+
interceptor.beforeAttempt(1)
126+
127+
// Simulate very long delay (15 seconds)
128+
clock.tick(15000)
129+
interceptor.beforeAttempt(2)
130+
131+
expect(mockCallback.calledOnce).to.be.true
132+
const call = mockCallback.getCall(0)
133+
expect(call.args[0].message).to.include('retrying within 10s')
134+
expect(call.args[0].attemptNumber).to.equal(2)
135+
expect(call.args[0].delay).to.equal(10) // Capped at 10 seconds
136+
expect(call.args[0].thresholdExceeded).to.be.true
137+
138+
clock.restore()
139+
})
140+
141+
it('should log when no callback is set', () => {
142+
const clock = sinon.useFakeTimers(1000)
143+
144+
// First attempt
145+
interceptor.beforeAttempt(1)
146+
147+
// Simulate delay above threshold
148+
clock.tick(3000)
149+
interceptor.beforeAttempt(2)
150+
151+
expect(mockLogging.debug.calledWith('QDelayTrackingInterceptor: No delay notification callback set')).to.be
152+
.true
153+
154+
clock.restore()
155+
})
156+
})
157+
158+
describe('reset', () => {
159+
it('should reset lastAttemptTime', () => {
160+
// Make an attempt to set lastAttemptTime
161+
interceptor.beforeAttempt(1)
162+
163+
// Reset
164+
interceptor.reset()
165+
166+
// Next attempt should be treated as first
167+
interceptor.beforeAttempt(1)
168+
169+
expect(
170+
mockLogging.debug.calledWith(
171+
'QDelayTrackingInterceptor: First attempt or no lastAttemptTime, skipping delay calculation'
172+
)
173+
).to.be.true
174+
})
175+
})
176+
177+
describe('name', () => {
178+
it('should return correct name', () => {
179+
expect(interceptor.name()).to.equal('Q Language Server Delay Tracking Interceptor')
180+
})
181+
})
182+
})

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

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
import { Logging } from '@aws/language-server-runtimes/server-interface'
2+
import { MINOR_DELAY_THRESHOLD_MS, MAJOR_DELAY_THRESHOLD_MS, MAX_RETRY_DELAY_MS } from '../constants/constants'
3+
4+
export interface DelayNotification {
5+
message: string
6+
attemptNumber: number
7+
delay: number
8+
thresholdExceeded: boolean
9+
}
10+
11+
/**
12+
* Delay tracking interceptor that matches CLI's DelayTrackingInterceptor behavior.
13+
* Tracks retry delays and provides user notifications.
14+
*/
15+
export class QDelayTrackingInterceptor {
16+
private logging?: Logging
17+
private minorDelayThreshold: number = MINOR_DELAY_THRESHOLD_MS
18+
private majorDelayThreshold: number = MAJOR_DELAY_THRESHOLD_MS
19+
private maxRetryDelay: number = MAX_RETRY_DELAY_MS
20+
private lastAttemptTime?: number
21+
private onDelayNotification?: (notification: DelayNotification) => void
22+
23+
constructor(logging?: Logging) {
24+
this.logging = logging
25+
}
26+
27+
/**
28+
* Sets the delay notification callback for UI integration
29+
*/
30+
public setDelayNotificationCallback(callback: (notification: DelayNotification) => void): void {
31+
this.logging?.debug(`QDelayTrackingInterceptor: setDelayNotificationCallback called`)
32+
this.onDelayNotification = callback
33+
}
34+
35+
/**
36+
* Called before each request attempt to track delays and notify users
37+
*/
38+
public beforeAttempt(attemptNumber: number): void {
39+
this.logging?.debug(`QDelayTrackingInterceptor: Attempt ${attemptNumber}`)
40+
const now = Date.now()
41+
42+
if (this.lastAttemptTime && attemptNumber > 1) {
43+
const delay = Math.min(now - this.lastAttemptTime, this.maxRetryDelay)
44+
this.logging?.debug(
45+
`QDelayTrackingInterceptor: Delay ${delay}ms, thresholds: minor=${this.minorDelayThreshold}ms, major=${this.majorDelayThreshold}ms`
46+
)
47+
48+
let message: string
49+
if (delay >= this.majorDelayThreshold) {
50+
message = `Retry #${attemptNumber}, retrying within ${Math.ceil(this.maxRetryDelay / 1000)}s..`
51+
} else if (delay >= this.minorDelayThreshold) {
52+
message = `Retry #${attemptNumber}, retrying within 5s..`
53+
} else {
54+
// No notification for short delays
55+
this.logging?.debug(`QDelayTrackingInterceptor: Delay ${delay}ms below threshold, no notification`)
56+
this.lastAttemptTime = now
57+
return
58+
}
59+
60+
this.logging?.debug(`QDelayTrackingInterceptor: Delay message: ${message}`)
61+
62+
// Notify UI about the delay
63+
if (this.onDelayNotification) {
64+
this.logging?.debug(`QDelayTrackingInterceptor: Sending delay notification`)
65+
this.onDelayNotification({
66+
message,
67+
attemptNumber,
68+
delay: Math.ceil(delay / 1000),
69+
thresholdExceeded: delay >= this.minorDelayThreshold,
70+
})
71+
} else {
72+
this.logging?.debug(`QDelayTrackingInterceptor: No delay notification callback set`)
73+
}
74+
} else {
75+
this.logging?.debug(
76+
`QDelayTrackingInterceptor: First attempt or no lastAttemptTime, skipping delay calculation`
77+
)
78+
}
79+
80+
this.lastAttemptTime = now
81+
}
82+
83+
/**
84+
* Reset tracking state
85+
*/
86+
public reset(): void {
87+
this.lastAttemptTime = undefined
88+
}
89+
90+
public name(): string {
91+
return 'Q Language Server Delay Tracking Interceptor'
92+
}
93+
}

0 commit comments

Comments
 (0)