Skip to content

Commit 2e59395

Browse files
ashishrp-awsashishrp-aws
committed
feat(amazonq): adding classification based retry strategy for chat requests
1 parent 013aa59 commit 2e59395

17 files changed

+1936
-120
lines changed

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -905,6 +905,13 @@ export class AgenticChatController implements ChatHandlers {
905905
})
906906
session.setConversationType('AgenticChat')
907907

908+
// Set up delay notification callback to show retry progress to users
909+
session.setDelayNotificationCallback(notification => {
910+
if (notification.thresholdExceeded) {
911+
void chatResultStream.updateProgressMessage(notification.message)
912+
}
913+
})
914+
908915
const additionalContext = await this.#additionalContextProvider.getAdditionalContext(
909916
triggerContext,
910917
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/retry/delayTracker.test.ts‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
import { DelayTracker, DelayNotification } from './delayTracker'
2+
import { Logging } from '@aws/language-server-runtimes/server-interface'
3+
import { expect } from 'chai'
4+
import * as sinon from 'sinon'
5+
6+
describe('DelayTracker', () => {
7+
let tracker: DelayTracker
8+
let mockLogging: any
9+
let mockOnNotification: sinon.SinonStub
10+
11+
beforeEach(() => {
12+
mockLogging = {
13+
info: sinon.stub(),
14+
debug: sinon.stub(),
15+
warn: sinon.stub(),
16+
error: sinon.stub(),
17+
}
18+
19+
mockOnNotification = sinon.stub()
20+
tracker = new DelayTracker(mockLogging, mockOnNotification)
21+
})
22+
23+
describe('trackInitialDelay', () => {
24+
it('should track delay above minor threshold', () => {
25+
tracker.trackInitialDelay(3000)
26+
27+
sinon.assert.calledWith(
28+
mockLogging.info,
29+
sinon.match(/Initial request delayed: ⏳ Request delayed by 3.0s/)
30+
)
31+
sinon.assert.calledWith(
32+
mockOnNotification,
33+
sinon.match({
34+
type: 'initial_delay',
35+
delayMs: 3000,
36+
attemptNumber: 1,
37+
message: sinon.match(/⏳ Request delayed by 3.0s/),
38+
thresholdExceeded: true,
39+
})
40+
)
41+
})
42+
43+
it('should not track delay below minor threshold', () => {
44+
tracker.trackInitialDelay(1000)
45+
46+
sinon.assert.notCalled(mockLogging.info)
47+
sinon.assert.notCalled(mockOnNotification)
48+
})
49+
50+
it('should generate appropriate context messages', () => {
51+
tracker.trackInitialDelay(15000)
52+
53+
sinon.assert.calledWith(
54+
mockOnNotification,
55+
sinon.match({
56+
message: sinon.match(/Service is under heavy load/),
57+
})
58+
)
59+
})
60+
})
61+
62+
describe('trackRetryDelay', () => {
63+
it('should track delay above major threshold', () => {
64+
tracker.trackRetryDelay(6000, 2)
65+
66+
sinon.assert.calledWith(mockLogging.info, sinon.match(/Retry attempt delayed: ⚠️ Retry #2 delayed by 6.0s/))
67+
sinon.assert.calledWith(mockOnNotification, {
68+
type: 'retry_delay',
69+
delayMs: 6000,
70+
attemptNumber: 2,
71+
message: sinon.match(/⚠️ Retry #2 delayed by 6.0s/),
72+
thresholdExceeded: true,
73+
})
74+
})
75+
76+
it('should track delay when doubled from previous', () => {
77+
tracker.trackRetryDelay(4000, 2, 1000)
78+
79+
sinon.assert.calledWith(mockOnNotification, {
80+
type: 'retry_delay',
81+
delayMs: 4000,
82+
attemptNumber: 2,
83+
message: sinon.match(/increased from 1.0s/),
84+
thresholdExceeded: true,
85+
})
86+
})
87+
88+
it('should not track delay below thresholds', () => {
89+
tracker.trackRetryDelay(3000, 2, 2000)
90+
91+
sinon.assert.notCalled(mockLogging.info)
92+
sinon.assert.notCalled(mockOnNotification)
93+
})
94+
})
95+
96+
describe('summarizeDelayImpact', () => {
97+
it('should summarize delays above major threshold', () => {
98+
tracker.summarizeDelayImpact(8000, 3, 2000)
99+
100+
sinon.assert.calledWith(
101+
mockLogging.info,
102+
sinon.match(
103+
/Delay summary: 📊 Request completed with delays: total delay: 8.0s, 3 attempts, final retry delay: 2.0s/
104+
)
105+
)
106+
sinon.assert.calledWith(mockOnNotification, {
107+
type: 'summary',
108+
delayMs: 8000,
109+
attemptNumber: 3,
110+
message: sinon.match(/📊 Request completed with delays/),
111+
thresholdExceeded: true,
112+
})
113+
})
114+
115+
it('should summarize without final retry delay', () => {
116+
tracker.summarizeDelayImpact(6000, 2)
117+
118+
sinon.assert.calledWith(
119+
mockOnNotification,
120+
sinon.match({
121+
message: sinon.match(str => !str.includes('final retry delay')),
122+
})
123+
)
124+
})
125+
126+
it('should not summarize delays below major threshold', () => {
127+
tracker.summarizeDelayImpact(3000, 2)
128+
129+
sinon.assert.notCalled(mockLogging.info)
130+
sinon.assert.notCalled(mockOnNotification)
131+
})
132+
})
133+
134+
describe('without callbacks', () => {
135+
it('should work without logging or notification callbacks', () => {
136+
const simpleTracker = new DelayTracker()
137+
138+
expect(() => {
139+
simpleTracker.trackInitialDelay(3000)
140+
simpleTracker.trackRetryDelay(6000, 2)
141+
simpleTracker.summarizeDelayImpact(8000, 3)
142+
}).not.to.throw()
143+
})
144+
})
145+
})

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

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
import { Logging } from '@aws/language-server-runtimes/server-interface'
2+
3+
export interface DelayNotification {
4+
type: 'initial_delay' | 'retry_delay' | 'summary'
5+
delayMs: number
6+
attemptNumber: number
7+
message: string
8+
thresholdExceeded: boolean
9+
}
10+
11+
export class DelayTracker {
12+
private minorDelayThreshold: number = 2000 // 2 seconds
13+
private majorDelayThreshold: number = 5000 // 5 seconds
14+
private logging?: Logging
15+
private onNotification?: (notification: DelayNotification) => void
16+
17+
constructor(logging?: Logging, onNotification?: (notification: DelayNotification) => void) {
18+
this.logging = logging
19+
this.onNotification = onNotification
20+
}
21+
22+
private generateContextMessage(delayMs: number, isRetry: boolean): string {
23+
const delaySeconds = delayMs / 1000
24+
25+
if (delaySeconds <= 2) {
26+
return 'Service is responding normally.'
27+
} else if (delaySeconds <= 5) {
28+
return 'Service may be experiencing light load.'
29+
} else if (delaySeconds <= 10) {
30+
return isRetry
31+
? 'Service is throttling requests. Consider switching models.'
32+
: 'Service is applying rate limiting.'
33+
} else if (delaySeconds <= 20) {
34+
return 'Service is under heavy load. Delays may continue. Consider switching models.'
35+
} else {
36+
return 'Service is experiencing significant load. Consider trying again later.'
37+
}
38+
}
39+
40+
trackInitialDelay(delayMs: number): void {
41+
if (delayMs >= this.minorDelayThreshold) {
42+
const message = `⏳ Request delayed by ${(delayMs / 1000).toFixed(1)}s due to rate limiting. ${this.generateContextMessage(delayMs, false)}`
43+
44+
const notification: DelayNotification = {
45+
type: 'initial_delay',
46+
delayMs,
47+
attemptNumber: 1,
48+
message,
49+
thresholdExceeded: delayMs >= this.minorDelayThreshold,
50+
}
51+
52+
this.logging?.info(`Initial request delayed: ${message}`)
53+
this.onNotification?.(notification)
54+
}
55+
}
56+
57+
trackRetryDelay(delayMs: number, attemptNumber: number, previousDelayMs?: number): void {
58+
const shouldNotify = delayMs >= this.majorDelayThreshold || (previousDelayMs && delayMs >= previousDelayMs * 2)
59+
60+
if (shouldNotify) {
61+
let message = `⚠️ Retry #${attemptNumber} delayed by ${(delayMs / 1000).toFixed(1)}s`
62+
63+
if (previousDelayMs) {
64+
const change = delayMs > previousDelayMs ? 'increased' : 'decreased'
65+
message += ` (${change} from ${(previousDelayMs / 1000).toFixed(1)}s)`
66+
}
67+
68+
message += `. ${this.generateContextMessage(delayMs, true)}`
69+
70+
const notification: DelayNotification = {
71+
type: 'retry_delay',
72+
delayMs,
73+
attemptNumber,
74+
message,
75+
thresholdExceeded: true,
76+
}
77+
78+
this.logging?.info(`Retry attempt delayed: ${message}`)
79+
this.onNotification?.(notification)
80+
}
81+
}
82+
83+
summarizeDelayImpact(totalDelayMs: number, attemptCount: number, finalRetryDelayMs?: number): void {
84+
if (totalDelayMs >= this.majorDelayThreshold) {
85+
let message = `📊 Request completed with delays: total delay: ${(totalDelayMs / 1000).toFixed(1)}s, ${attemptCount} attempts`
86+
87+
if (finalRetryDelayMs) {
88+
message += `, final retry delay: ${(finalRetryDelayMs / 1000).toFixed(1)}s`
89+
}
90+
91+
const notification: DelayNotification = {
92+
type: 'summary',
93+
delayMs: totalDelayMs,
94+
attemptNumber: attemptCount,
95+
message,
96+
thresholdExceeded: true,
97+
}
98+
99+
this.logging?.info(`Delay summary: ${message}`)
100+
this.onNotification?.(notification)
101+
}
102+
}
103+
}

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
// Export all retry-related classes and types
7+
export { QRetryClassifier, RetryAction } from './retryClassifier'
8+
export { DelayTracker, DelayNotification } from './delayTracker'
9+
export { AdaptiveRetryConfig, RetryConfig } from './retryConfig'
10+
export { ClientSideRateLimiter, TokenCost } from './rateLimiter'
11+
export { RetryErrorHandler } from './retryErrorHandler'
12+
export { RetryTelemetryController } from './retryTelemetry'

0 commit comments

Comments
 (0)