feat: Add HTTP retry support to the SDK

PiperOrigin-RevId: 867654877

Author: Annhiluc

api-report/genai-node.api.md

@@ -1854,6 +1854,7 @@ export interface HttpOptions {
     baseUrlResourceScope?: ResourceScope;
     extraBody?: Record<string, unknown>;
     headers?: Record<string, string>;
+    retryOptions?: HttpRetryOptions;
     timeout?: number;
 }
 
@@ -1866,6 +1867,11 @@ export class HttpResponse {
     responseInternal: Response;
 }
 
+// @public
+export interface HttpRetryOptions {
+    attempts?: number;
+}
+
 // @public
 interface Image_2 {
     gcsUri?: string;

api-report/genai-web.api.md

@@ -1854,6 +1854,7 @@ export interface HttpOptions {
     baseUrlResourceScope?: ResourceScope;
     extraBody?: Record<string, unknown>;
     headers?: Record<string, string>;
+    retryOptions?: HttpRetryOptions;
     timeout?: number;
 }
 
@@ -1866,6 +1867,11 @@ export class HttpResponse {
     responseInternal: Response;
 }
 
+// @public
+export interface HttpRetryOptions {
+    attempts?: number;
+}
+
 // @public
 interface Image_2 {
     gcsUri?: string;

api-report/genai.api.md

@@ -1854,6 +1854,7 @@ export interface HttpOptions {
     baseUrlResourceScope?: ResourceScope;
     extraBody?: Record<string, unknown>;
     headers?: Record<string, string>;
+    retryOptions?: HttpRetryOptions;
     timeout?: number;
 }
 
@@ -1866,6 +1867,11 @@ export class HttpResponse {
     responseInternal: Response;
 }
 
+// @public
+export interface HttpRetryOptions {
+    attempts?: number;
+}
+
 // @public
 interface Image_2 {
     gcsUri?: string;

package-lock.json

@@ -105,6 +105,7 @@
     "web/package.json"
   ],
   "devDependencies": {
+    "@cfworker/json-schema": "^4.1.1",
     "@eslint/js": "9.20.0",
     "@microsoft/api-extractor": "^7.52.9",
     "@modelcontextprotocol/sdk": "^1.25.2",
@@ -114,7 +115,6 @@
     "@types/node-fetch": "^2.6.13",
     "@types/unist": "^3.0.3",
     "@types/ws": "^8.5.14",
-    "@cfworker/json-schema": "^4.1.1",
     "c8": "^10.1.3",
     "eslint": "8.57.0",
     "gts": "^5.2.0",
@@ -142,6 +142,7 @@
   },
   "dependencies": {
     "google-auth-library": "^10.3.0",
+    "p-retry": "^7.1.1",
     "protobufjs": "^7.5.4",
     "ws": "^8.18.0"
   },

package.json

@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import pRetry, {AbortError} from 'p-retry';
 import {Auth} from './_auth.js';
 import * as common from './_common.js';
 import {Downloader} from './_downloader.js';
@@ -22,6 +23,20 @@ const LIBRARY_LABEL = `google-genai-sdk/${SDK_VERSION}`;
 const VERTEX_AI_API_DEFAULT_VERSION = 'v1beta1';
 const GOOGLE_AI_API_DEFAULT_VERSION = 'v1beta';
 
+// Default retry options.
+// The config is based on https://cloud.google.com/storage/docs/retry-strategy.
+const DEFAULT_RETRY_ATTEMPTS = 5; // Including the initial call
+// LINT.IfChange
+const DEFAULT_RETRY_HTTP_STATUS_CODES = [
+  408, // Request timeout
+  429, // Too many requests
+  500, // Internal server error
+  502, // Bad gateway
+  503, // Service unavailable
+  504, // Gateway timeout
+];
+// LINT.ThenChange(//depot/google3/third_party/py/google/genai/_api_client.py)
+
 /**
  * Options for initializing the ApiClient. The ApiClient uses the parameters
  * for authentication purposes as well as to infer if SDK should send the
@@ -659,8 +674,33 @@ export class ApiClient implements GeminiNextGenAPIClientAdapter {
     url: string,
     requestInit: RequestInit,
   ): Promise<Response> {
-    return fetch(url, requestInit).catch((e) => {
-      throw new Error(`exception ${e} sending request`);
+    if (
+      !this.clientOptions.httpOptions ||
+      !this.clientOptions.httpOptions.retryOptions
+    ) {
+      return fetch(url, requestInit);
+    }
+
+    const retryOptions = this.clientOptions.httpOptions.retryOptions;
+    const runFetch = async () => {
+      const response = await fetch(url, requestInit);
+
+      if (response.ok) {
+        return response;
+      }
+
+      if (DEFAULT_RETRY_HTTP_STATUS_CODES.includes(response.status)) {
+        throw new Error(`Retryable HTTP Error: ${response.statusText}`);
+      }
+
+      throw new AbortError(
+        `Non-retryable exception ${response.statusText} sending request`,
+      );
+    };
+
+    return pRetry(runFetch, {
+      // Retry attempts is one less than the number of total attempts.
+      retries: (retryOptions.attempts ?? DEFAULT_RETRY_ATTEMPTS) - 1,
     });
   }
 

src/_api_client.ts

@@ -1614,6 +1614,13 @@ export function createModelContent(
     parts: _toParts(partOrString),
   };
 }
+/** HTTP retry options to be used in each of the requests. */
+export declare interface HttpRetryOptions {
+  /** Maximum number of attempts, including the original request.
+      If 0 or 1, it means no retries. If not specified, default to 5. */
+  attempts?: number;
+}
+
 /** HTTP options to be used in each of the requests. */
 export declare interface HttpOptions {
   /** The base URL for the AI platform service endpoint. */
@@ -1631,6 +1638,8 @@ export declare interface HttpOptions {
       - VertexAI backend API docs: https://cloud.google.com/vertex-ai/docs/reference/rest
       - GeminiAPI backend API docs: https://ai.google.dev/api/rest */
   extraBody?: Record<string, unknown>;
+  /** HTTP retry options for the request. */
+  retryOptions?: HttpRetryOptions;
 }
 
 /** Schema is used to define the format of input/output data.

src/types.ts

@@ -385,6 +385,8 @@ describe('processStreamResponse', () => {
 
 describe('ApiClient', () => {
   describe('constructor', () => {
+    jasmine.DEFAULT_TIMEOUT_INTERVAL = 60000; // 60 seconds.
+
     it('should initialize with provided values', () => {
       const client = new ApiClient({
         auth: new FakeAuth(),
@@ -721,6 +723,94 @@ describe('ApiClient', () => {
       expect(headers['x-goog-api-client']).toContain('google-genai-sdk/');
       expect(client.getApiVersion()).toBe('v1beta1');
     });
+
+    it('should retry requests if retry options are set', async () => {
+      const client = new ApiClient({
+        auth: new FakeAuth(),
+        project: 'vertex-project',
+        location: 'vertex-location',
+        vertexai: true,
+        apiVersion: 'v1beta1',
+        httpOptions: {
+          retryOptions: {
+            attempts: 2,
+          },
+        },
+        uploader: new CrossUploader(),
+        downloader: new CrossDownloader(),
+      });
+      const fetchSpy = spyOn(global, 'fetch').and.returnValue(
+        Promise.resolve(
+          new Response(
+            JSON.stringify({'error': 'Internal Server Error'}),
+            fetch500Options,
+          ),
+        ),
+      );
+      await client
+        .request({path: 'test-path', httpMethod: 'POST'})
+        .catch((e) => {
+          console.log(e);
+        });
+      expect(fetchSpy).toHaveBeenCalledTimes(2);
+    });
+
+    it('should not retry requests if retry options are not set', async () => {
+      const client = new ApiClient({
+        auth: new FakeAuth(),
+        project: 'vertex-project',
+        location: 'vertex-location',
+        vertexai: true,
+        apiVersion: 'v1beta1',
+        uploader: new CrossUploader(),
+        downloader: new CrossDownloader(),
+      });
+      const fetchSpy = spyOn(global, 'fetch').and.returnValue(
+        Promise.resolve(
+          new Response(
+            JSON.stringify({'error': 'Internal Server Error'}),
+            fetch500Options,
+          ),
+        ),
+      );
+      await client
+        .request({path: 'test-path', httpMethod: 'POST'})
+        .catch((e) => {
+          expect(e.name).toEqual('ApiError');
+          expect(e.message).toContain('Internal Server Error');
+          expect(e.status).toEqual(500);
+        });
+      expect(fetchSpy).toHaveBeenCalledTimes(1);
+    });
+
+    it('should retry requests with default retry options if retry options are not set', async () => {
+      const client = new ApiClient({
+        auth: new FakeAuth(),
+        project: 'vertex-project',
+        location: 'vertex-location',
+        vertexai: true,
+        apiVersion: 'v1beta1',
+        httpOptions: {
+          retryOptions: {},
+        },
+        uploader: new CrossUploader(),
+        downloader: new CrossDownloader(),
+      });
+      const fetchSpy = spyOn(global, 'fetch').and.returnValue(
+        Promise.resolve(
+          new Response(
+            JSON.stringify({'error': 'Internal Server Error'}),
+            fetch500Options,
+          ),
+        ),
+      );
+      await client
+        .request({path: 'test-path', httpMethod: 'POST'})
+        .catch((e) => {
+          console.log(e);
+        });
+      expect(fetchSpy).toHaveBeenCalledTimes(5); // Default retry attempts is 5.
+    });
   });
 
   describe('post/get methods', () => {