feat(agents): add generic provider api key rotation (#19587)

src/memory/embeddings-gemini.ts

@@ -1,13 +1,18 @@
+import type { EmbeddingProvider, EmbeddingProviderOptions } from "./embeddings.js";
+import {
+  collectProviderApiKeysForExecution,
+  executeWithApiKeyRotation,
+} from "../agents/api-key-rotation.js";
 import { requireApiKey, resolveApiKeyForProvider } from "../agents/model-auth.js";
 import { parseGeminiAuth } from "../infra/gemini-auth.js";
 import { debugEmbeddingsLog } from "./embeddings-debug.js";
-import type { EmbeddingProvider, EmbeddingProviderOptions } from "./embeddings.js";
 
 export type GeminiEmbeddingClient = {
   baseUrl: string;
   headers: Record<string, string>;
   model: string;
   modelPath: string;
+  apiKeys: string[];
 };
 
 const DEFAULT_GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta";
@@ -62,23 +67,40 @@ export async function createGeminiEmbeddingProvider(
   const embedUrl = `${baseUrl}/${client.modelPath}:embedContent`;
   const batchUrl = `${baseUrl}/${client.modelPath}:batchEmbedContents`;
 
-  const embedQuery = async (text: string): Promise<number[]> => {
-    if (!text.trim()) {
-      return [];
-    }
-    const res = await fetch(embedUrl, {
+  const fetchWithGeminiAuth = async (apiKey: string, endpoint: string, body: unknown) => {
+    const authHeaders = parseGeminiAuth(apiKey);
+    const headers = {
+      ...authHeaders.headers,
+      ...client.headers,
+    };
+    const res = await fetch(endpoint, {
       method: "POST",
-      headers: client.headers,
-      body: JSON.stringify({
-        content: { parts: [{ text }] },
-        taskType: "RETRIEVAL_QUERY",
-      }),
+      headers,
+      body: JSON.stringify(body),
     });
     if (!res.ok) {
       const payload = await res.text();
       throw new Error(`gemini embeddings failed: ${res.status} ${payload}`);
     }
-    const payload = (await res.json()) as { embedding?: { values?: number[] } };
+    return (await res.json()) as {
+      embedding?: { values?: number[] };
+      embeddings?: Array<{ values?: number[] }>;
+    };
+  };
+
+  const embedQuery = async (text: string): Promise<number[]> => {
+    if (!text.trim()) {
+      return [];
+    }
+    const payload = await executeWithApiKeyRotation({
+      provider: "google",
+      apiKeys: client.apiKeys,
+      execute: (apiKey) =>
+        fetchWithGeminiAuth(apiKey, embedUrl, {
+          content: { parts: [{ text }] },
+          taskType: "RETRIEVAL_QUERY",
+        }),
+    });
     return payload.embedding?.values ?? [];
   };
 
@@ -91,16 +113,14 @@ export async function createGeminiEmbeddingProvider(
       content: { parts: [{ text }] },
       taskType: "RETRIEVAL_DOCUMENT",
     }));
-    const res = await fetch(batchUrl, {
-      method: "POST",
-      headers: client.headers,
-      body: JSON.stringify({ requests }),
+    const payload = await executeWithApiKeyRotation({
+      provider: "google",
+      apiKeys: client.apiKeys,
+      execute: (apiKey) =>
+        fetchWithGeminiAuth(apiKey, batchUrl, {
+          requests,
+        }),
     });
-    if (!res.ok) {
-      const payload = await res.text();
-      throw new Error(`gemini embeddings failed: ${res.status} ${payload}`);
-    }
-    const payload = (await res.json()) as { embeddings?: Array<{ values?: number[] }> };
     const embeddings = Array.isArray(payload.embeddings) ? payload.embeddings : [];
     return texts.map((_, index) => embeddings[index]?.values ?? []);
   };
@@ -139,11 +159,13 @@ export async function resolveGeminiEmbeddingClient(
   const rawBaseUrl = remoteBaseUrl || providerConfig?.baseUrl?.trim() || DEFAULT_GEMINI_BASE_URL;
   const baseUrl = normalizeGeminiBaseUrl(rawBaseUrl);
   const headerOverrides = Object.assign({}, providerConfig?.headers, remote?.headers);
-  const authHeaders = parseGeminiAuth(apiKey);
   const headers: Record<string, string> = {
-    ...authHeaders.headers,
     ...headerOverrides,
   };
+  const apiKeys = collectProviderApiKeysForExecution({
+    provider: "google",
+    primaryApiKey: apiKey,
+  });
   const model = normalizeGeminiModel(options.model);
   const modelPath = buildGeminiModelPath(model);
   debugEmbeddingsLog("memory embeddings: gemini client", {
@@ -154,5 +176,5 @@ export async function resolveGeminiEmbeddingClient(
     embedEndpoint: `${baseUrl}/${modelPath}:embedContent`,
     batchEndpoint: `${baseUrl}/${modelPath}:batchEmbedContents`,
   });
-  return { baseUrl, headers, model, modelPath };
+  return { baseUrl, headers, model, modelPath, apiKeys };
 }