2 files changed, 51 insertions, 22 deletions

diff --git a/packages/server/src/core/client.ts b/packages/server/src/core/client.ts
index 69b815ab..9006c675 100644
--- a/packages/server/src/core/client.ts
+++ b/packages/server/src/core/client.ts
@@ -23,6 +23,7 @@ import { getResponseText } from '../utils/generateContentResponseUtilities.js';
 import { checkNextSpeaker } from '../utils/nextSpeakerChecker.js';
 import { reportError } from '../utils/errorReporting.js';
 import { GeminiChat } from './geminiChat.js';
+import { retryWithBackoff } from '../utils/retry.js';
 
 export class GeminiClient {
   private client: GoogleGenAI;
@@ -194,16 +195,20 @@ export class GeminiClient {
         ...config,
       };
 
-      const result = await this.client.models.generateContent({
-        model,
-        config: {
-          ...requestConfig,
-          systemInstruction,
-          responseSchema: schema,
-          responseMimeType: 'application/json',
-        },
-        contents,
-      });
+      const apiCall = () =>
+        this.client.models.generateContent({
+          model,
+          config: {
+            ...requestConfig,
+            systemInstruction,
+            responseSchema: schema,
+            responseMimeType: 'application/json',
+          },
+          contents,
+        });
+
+      const result = await retryWithBackoff(apiCall);
+
       const text = getResponseText(result);
       if (!text) {
         const error = new Error(
diff --git a/packages/server/src/core/geminiChat.ts b/packages/server/src/core/geminiChat.ts
index 877d0825..b34b6f35 100644
--- a/packages/server/src/core/geminiChat.ts
+++ b/packages/server/src/core/geminiChat.ts
@@ -16,6 +16,7 @@ import {
   GoogleGenAI,
   createUserContent,
 } from '@google/genai';
+import { retryWithBackoff } from '../utils/retry.js';
 import { isFunctionResponse } from '../utils/messageInspectors.js';
 
 /**
@@ -152,11 +153,16 @@ export class GeminiChat {
   ): Promise<GenerateContentResponse> {
     await this.sendPromise;
     const userContent = createUserContent(params.message);
-    const responsePromise = this.modelsModule.generateContent({
-      model: this.model,
-      contents: this.getHistory(true).concat(userContent),
-      config: { ...this.config, ...params.config },
-    });
+
+    const apiCall = () =>
+      this.modelsModule.generateContent({
+        model: this.model,
+        contents: this.getHistory(true).concat(userContent),
+        config: { ...this.config, ...params.config },
+      });
+
+    const responsePromise = retryWithBackoff(apiCall);
+
     this.sendPromise = (async () => {
       const response = await responsePromise;
       const outputContent = response.candidates?.[0]?.content;
@@ -216,19 +222,37 @@ export class GeminiChat {
   ): Promise<AsyncGenerator<GenerateContentResponse>> {
     await this.sendPromise;
     const userContent = createUserContent(params.message);
-    const streamResponse = this.modelsModule.generateContentStream({
-      model: this.model,
-      contents: this.getHistory(true).concat(userContent),
-      config: { ...this.config, ...params.config },
+
+    const apiCall = () =>
+      this.modelsModule.generateContentStream({
+        model: this.model,
+        contents: this.getHistory(true).concat(userContent),
+        config: { ...this.config, ...params.config },
+      });
+
+    // Note: Retrying streams can be complex. If generateContentStream itself doesn't handle retries
+    // for transient issues internally before yielding the async generator, this retry will re-initiate
+    // the stream. For simple 429/500 errors on initial call, this is fine.
+    // If errors occur mid-stream, this setup won't resume the stream; it will restart it.
+    const streamResponse = await retryWithBackoff(apiCall, {
+      shouldRetry: (error: Error) => {
+        // Check error messages for status codes, or specific error names if known
+        if (error && error.message) {
+          if (error.message.includes('429')) return true;
+          if (error.message.match(/5\d{2}/)) return true;
+        }
+        return false; // Don't retry other errors by default
+      },
     });
+
     // Resolve the internal tracking of send completion promise - `sendPromise`
     // for both success and failure response. The actual failure is still
     // propagated by the `await streamResponse`.
-    this.sendPromise = streamResponse
+    this.sendPromise = Promise.resolve(streamResponse)
       .then(() => undefined)
       .catch(() => undefined);
-    const response = await streamResponse;
-    const result = this.processStreamResponse(response, userContent);
+
+    const result = this.processStreamResponse(streamResponse, userContent);
     return result;
   }