fix(core): apply retry logic to CodeAssistServer for all users (#20507)

This commit is contained in:
Bryan Morgan
2026-02-27 09:26:53 -05:00
committed by GitHub
parent 82336de7a4
commit 522e95439c
2 changed files with 58 additions and 26 deletions

View File

@@ -73,19 +73,26 @@ describe('CodeAssistServer', () => {
LlmRole.MAIN,
);
expect(mockRequest).toHaveBeenCalledWith(
expect.objectContaining({
url: expect.stringContaining(':generateContent'),
method: 'POST',
headers: {
'Content-Type': 'application/json',
'x-custom-header': 'test-value',
},
responseType: 'json',
body: expect.any(String),
signal: undefined,
}),
);
expect(mockRequest).toHaveBeenCalledWith({
url: expect.stringContaining(':generateContent'),
method: 'POST',
headers: {
'Content-Type': 'application/json',
'x-custom-header': 'test-value',
},
responseType: 'json',
body: expect.any(String),
signal: undefined,
retryConfig: {
retry: 3,
noResponseRetries: 3,
statusCodesToRetry: [
[429, 429],
[499, 499],
[500, 599],
],
},
});
const requestBody = JSON.parse(mockRequest.mock.calls[0][0].body);
expect(requestBody.user_prompt_id).toBe('user-prompt-id');
@@ -393,19 +400,26 @@ describe('CodeAssistServer', () => {
results.push(res);
}
expect(mockRequest).toHaveBeenCalledWith(
expect.objectContaining({
url: expect.stringContaining(':streamGenerateContent'),
method: 'POST',
params: { alt: 'sse' },
responseType: 'stream',
body: expect.any(String),
headers: {
'Content-Type': 'application/json',
},
signal: undefined,
}),
);
expect(mockRequest).toHaveBeenCalledWith({
url: expect.stringContaining(':streamGenerateContent'),
method: 'POST',
params: { alt: 'sse' },
responseType: 'stream',
body: expect.any(String),
headers: {
'Content-Type': 'application/json',
},
signal: undefined,
retryConfig: {
retry: 3,
noResponseRetries: 3,
statusCodesToRetry: [
[429, 429],
[499, 499],
[500, 599],
],
},
});
expect(results).toHaveLength(2);
expect(results[0].candidates?.[0].content?.parts?.[0].text).toBe('Hello');

View File

@@ -305,6 +305,15 @@ export class CodeAssistServer implements ContentGenerator {
responseType: 'json',
body: JSON.stringify(req),
signal,
retryConfig: {
retry: 3,
noResponseRetries: 3,
statusCodesToRetry: [
[429, 429],
[499, 499],
[500, 599],
],
},
});
return res.data;
}
@@ -352,6 +361,15 @@ export class CodeAssistServer implements ContentGenerator {
responseType: 'stream',
body: JSON.stringify(req),
signal,
retryConfig: {
retry: 3,
noResponseRetries: 3,
statusCodesToRetry: [
[429, 429],
[499, 499],
[500, 599],
],
},
});
return (async function* (): AsyncGenerator<T> {