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

packages/core/src/code_assist/server.test.ts

@@ -73,8 +73,7 @@ describe('CodeAssistServer', () => {
       LlmRole.MAIN,
     );
 
-    expect(mockRequest).toHaveBeenCalledWith(
+    expect(mockRequest).toHaveBeenCalledWith({
-      expect.objectContaining({
       url: expect.stringContaining(':generateContent'),
       method: 'POST',
       headers: {
@@ -84,8 +83,16 @@ describe('CodeAssistServer', () => {
       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,8 +400,7 @@ describe('CodeAssistServer', () => {
       results.push(res);
     }
 
-    expect(mockRequest).toHaveBeenCalledWith(
+    expect(mockRequest).toHaveBeenCalledWith({
-      expect.objectContaining({
       url: expect.stringContaining(':streamGenerateContent'),
       method: 'POST',
       params: { alt: 'sse' },
@@ -404,8 +410,16 @@ describe('CodeAssistServer', () => {
         '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');

packages/core/src/code_assist/server.ts

@@ -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> {