From 522e95439c1a42f50680935e269ca284a0784a53 Mon Sep 17 00:00:00 2001
From: Bryan Morgan <bryanmorgan@google.com>
Date: Fri, 27 Feb 2026 09:26:53 -0500
Subject: [PATCH] fix(core): apply retry logic to CodeAssistServer for all
 users (#20507)

---
 packages/core/src/code_assist/server.test.ts | 66 ++++++++++++--------
 packages/core/src/code_assist/server.ts      | 18 ++++++
 2 files changed, 58 insertions(+), 26 deletions(-)

diff --git a/packages/core/src/code_assist/server.test.ts b/packages/core/src/code_assist/server.test.ts
index d79526d1c3..aa5cbae7fe 100644
--- a/packages/core/src/code_assist/server.test.ts
+++ b/packages/core/src/code_assist/server.test.ts
@@ -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');
diff --git a/packages/core/src/code_assist/server.ts b/packages/core/src/code_assist/server.ts
index ff5fb76e07..dfcc693ca7 100644
--- a/packages/core/src/code_assist/server.ts
+++ b/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> {