fix(google): preserve gateway queries in embedding batches (#130768)

* fix(google): preserve gateway queries in embedding batches

Normalize embedding destination paths without changing query values, and reuse the canonical endpoint resolver throughout the batch lifecycle.

* test(google): validate provider availability in batch fixture

extensions/google/embedding-batch.test.ts

@@ -1,8 +1,10 @@
 // Google tests cover embedding batch bounded JSON response reads.
 import { createServer } from "node:http";
+import * as embeddingSdk from "openclaw/plugin-sdk/memory-core-host-engine-embeddings";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import { runGeminiEmbeddingBatches } from "./embedding-batch.js";
 import type { GeminiEmbeddingClient } from "./embedding-provider.js";
+import { geminiMemoryEmbeddingProviderAdapter } from "./memory-embedding-adapter.js";
 
 // Pass-through so onResponse receives real Response objects (required by
 // readProviderJsonResponse which needs a real .body ReadableStream).
@@ -321,15 +323,25 @@ describe("Google embedding-batch bounded JSON reads", () => {
     expect(response.bodyUsed).toBe(true);
   });
 
-  it("normalizes raw Google Operations and uses the canonical download route", async () => {
+  it.each([
+    { baseUrl: "https://generativelanguage.googleapis.com/v1beta", version: "v1beta", query: "" },
+    {
+      baseUrl: "https://generativelanguage.googleapis.com/v1alpha/?tenant=remote",
+      version: "v1alpha",
+      query: "tenant=remote&",
+    },
+  ])("uses canonical Google file routes for $baseUrl", async ({ baseUrl, version, query }) => {
     const fetchMock = stubBatchFetch();
 
-    const result = await runBatch();
+    const result = await runBatch(singleRequest(), makeGeminiClient(baseUrl));
 
     expect(result.get("r0")).toEqual([1, 0, 0]);
-    expect(fetchMock.mock.calls.map(([input]) => fetchInputUrl(input))).toContain(
-      "https://generativelanguage.googleapis.com/download/v1beta/files/out-0:download?alt=media",
-    );
+    expect(fetchMock.mock.calls.map(([input]) => fetchInputUrl(input))).toEqual([
+      `https://generativelanguage.googleapis.com/upload/${version}/files?${query}uploadType=multipart`,
+      `https://generativelanguage.googleapis.com/${version}/models/gemini-embedding-001:asyncBatchEmbedContent${query ? `?${query.slice(0, -1)}` : ""}`,
+      `https://generativelanguage.googleapis.com/${version}/batches/b-0${query ? `?${query.slice(0, -1)}` : ""}`,
+      `https://generativelanguage.googleapis.com/download/${version}/files/out-0:download?${query}alt=media`,
+    ]);
     for (const [, init] of fetchMock.mock.calls) {
       expect(new Headers(init?.headers).get("x-goog-api-key")).toBe("test-key");
     }
@@ -351,79 +363,164 @@ describe("Google embedding-batch bounded JSON reads", () => {
     );
   });
 
-  it("runs the complete batch lifecycle over loopback HTTP", async () => {
-    let createBody: unknown;
-    const authHeaders: Array<string | undefined> = [];
-    const server = createServer((request, response) => {
-      void (async () => {
-        const url = new URL(request.url ?? "/", "http://127.0.0.1");
-        const apiKey = request.headers["x-goog-api-key"];
-        authHeaders.push(Array.isArray(apiKey) ? apiKey.join(", ") : apiKey);
-        const respondJson = (body: unknown) => {
-          response.writeHead(200, { "content-type": "application/json" });
-          response.end(JSON.stringify(body));
-        };
-        if (url.pathname === "/upload/v1beta/files") {
-          request.resume();
-          await new Promise<void>((resolve) => {
-            request.once("end", () => resolve());
-          });
-          respondJson({ file: { name: "files/input-0" } });
-          return;
-        }
-        if (url.pathname.endsWith(":asyncBatchEmbedContent")) {
-          let body = "";
-          request.setEncoding("utf8");
-          for await (const chunk of request) {
-            body += chunk;
-          }
-          createBody = JSON.parse(body) as unknown;
-          respondJson({
-            name: "batches/b-0",
-            done: false,
-            metadata: { state: "BATCH_STATE_PENDING" },
-          });
-          return;
-        }
-        if (url.pathname === "/v1beta/batches/b-0") {
-          respondJson({
-            name: "batches/b-0",
-            done: true,
-            metadata: { state: "BATCH_STATE_SUCCEEDED" },
-            response: { responsesFile: "files/output-0" },
-          });
-          return;
-        }
-        if (url.pathname === "/v1beta/files/output-0:download") {
-          response.writeHead(200, { "content-type": "application/jsonl" });
-          const line = JSON.stringify({
-            key: "r0",
-            response: { embedding: { values: [1, 0, 0] } },
-          });
-          response.write(line.slice(0, 17));
-          response.end(line.slice(17));
-          return;
-        }
-        response.writeHead(404).end();
-      })().catch((error: unknown) => {
-        response.writeHead(500).end(error instanceof Error ? error.message : String(error));
-      });
-    });
-    const port = await listenLoopbackServer(server);
-
-    try {
-      const result = await runBatch(
-        singleRequest(),
-        makeGeminiClient(`http://127.0.0.1:${port}/v1beta`),
+  it.each([
+    { basePath: "/v1beta", prefix: "", query: "" },
+    { basePath: "/gateway/v1beta/", prefix: "/gateway", query: "?tenant=remote" },
+    { basePath: "/gateway/v1beta/", prefix: "/gateway", query: "?tenant=remote&route=a/" },
+    { basePath: "/gateway/v1beta", prefix: "/gateway", query: "?tenant=/openai/team/" },
+    { basePath: "/gateway/v1beta/openai", prefix: "/gateway", query: "?tenant=remote" },
+  ])(
+    "runs the public adapter over HTTP for $basePath with query $query",
+    async ({ basePath, prefix, query }) => {
+      let createBody: unknown;
+      let uploadBody = "";
+      const observedUrls: string[] = [];
+      const authHeaders: Array<string | undefined> = [];
+      const tenantHeaders: Array<string | undefined> = [];
+      const realSdk = await vi.importActual<typeof embeddingSdk>(
+        "openclaw/plugin-sdk/memory-core-host-engine-embeddings",
       );
+      const remoteHttp = vi
+        .spyOn(embeddingSdk, "withRemoteHttpResponse")
+        .mockImplementation(realSdk.withRemoteHttpResponse);
+      const server = createServer((request, response) => {
+        void (async () => {
+          const url = new URL(request.url ?? "/", "http://127.0.0.1");
+          observedUrls.push(`${request.method} ${url.pathname}${url.search}`);
+          const apiKey = request.headers["x-goog-api-key"];
+          authHeaders.push(Array.isArray(apiKey) ? apiKey.join(", ") : apiKey);
+          const tenant = request.headers["x-proof-tenant"];
+          tenantHeaders.push(Array.isArray(tenant) ? tenant.join(", ") : tenant);
+          const expectedQuery = new URLSearchParams(query);
+          const configuredQuerySurvives = [...expectedQuery].every(
+            ([name, value]) => url.searchParams.get(name) === value,
+          );
+          if (!configuredQuerySurvives) {
+            response.writeHead(400).end(`configured query changed: ${url.pathname}${url.search}`);
+            request.resume();
+            return;
+          }
+          const respondJson = (body: unknown) => {
+            response.writeHead(200, { "content-type": "application/json" });
+            response.end(JSON.stringify(body));
+          };
+          if (url.pathname === `${prefix}/v1beta/models/gemini-embedding-001:embedContent`) {
+            request.resume();
+            respondJson({ embedding: { values: [1, 0, 0] } });
+            return;
+          }
+          if (
+            url.pathname === `${prefix}/upload/v1beta/files` &&
+            url.searchParams.get("uploadType") === "multipart"
+          ) {
+            request.setEncoding("utf8");
+            for await (const chunk of request) {
+              uploadBody += chunk;
+            }
+            respondJson({ file: { name: "files/input-0" } });
+            return;
+          }
+          if (
+            url.pathname === `${prefix}/v1beta/models/gemini-embedding-001:asyncBatchEmbedContent`
+          ) {
+            let body = "";
+            request.setEncoding("utf8");
+            for await (const chunk of request) {
+              body += chunk;
+            }
+            createBody = JSON.parse(body) as unknown;
+            respondJson({
+              name: "batches/b-0",
+              done: false,
+              metadata: { state: "BATCH_STATE_PENDING" },
+            });
+            return;
+          }
+          if (url.pathname === `${prefix}/v1beta/batches/b-0`) {
+            respondJson({
+              name: "batches/b-0",
+              done: true,
+              metadata: { state: "BATCH_STATE_SUCCEEDED" },
+              response: { responsesFile: "files/output-0" },
+            });
+            return;
+          }
+          if (
+            url.pathname === `${prefix}/v1beta/files/output-0:download` &&
+            url.searchParams.get("alt") === "media"
+          ) {
+            response.writeHead(200, { "content-type": "application/jsonl" });
+            const line = JSON.stringify({
+              key: "0",
+              response: { embedding: { values: [1, 0, 0] } },
+            });
+            response.write(line.slice(0, 17));
+            response.end(line.slice(17));
+            return;
+          }
+          response.writeHead(404).end();
+        })().catch((error: unknown) => {
+          response.writeHead(500).end(error instanceof Error ? error.message : String(error));
+        });
+      });
+      const port = await listenLoopbackServer(server);
 
-      expect(result).toEqual(new Map([["r0", [1, 0, 0]]]));
-      expect(createBody).toMatchObject({ batch: { inputConfig: { file_name: "files/input-0" } } });
-      expect(authHeaders).toEqual(["test-key", "test-key", "test-key", "test-key"]);
-    } finally {
-      await closeServer(server);
-    }
-  });
+      try {
+        const adapter = await geminiMemoryEmbeddingProviderAdapter.create({
+          config: {},
+          provider: "gemini",
+          model: "gemini-embedding-001",
+          fallback: "none",
+          remote: {
+            baseUrl: `http://127.0.0.1:${port}${basePath}${query}`,
+            apiKey: "test-key",
+            headers: { "X-Proof-Tenant": "remote" },
+          },
+        });
+        if (!adapter.provider) {
+          throw new Error("Expected a Gemini embedding provider");
+        }
+        await expect(adapter.provider.embed("hello", { inputType: "query" })).resolves.toEqual([
+          1, 0, 0,
+        ]);
+        const result = await adapter.runtime?.batchEmbed?.({
+          agentId: "main",
+          chunks: [{ text: "hello" }],
+          wait: true,
+          concurrency: 1,
+          pollIntervalMs: 1,
+          timeoutMs: 5_000,
+          debug: () => {},
+        });
+
+        expect(result).toEqual([[1, 0, 0]]);
+        const uploadedRequest = uploadBody.split("\r\n\r\n")[2]?.split("\r\n")[0];
+        expect(JSON.parse(uploadedRequest ?? "null")).toEqual({
+          key: "0",
+          request: {
+            content: { parts: [{ text: "hello" }] },
+            taskType: "RETRIEVAL_DOCUMENT",
+            model: "models/gemini-embedding-001",
+          },
+        });
+        expect(createBody).toMatchObject({
+          batch: { inputConfig: { file_name: "files/input-0" } },
+        });
+        expect(authHeaders).toEqual(Array(5).fill("test-key"));
+        expect(tenantHeaders).toEqual(Array(5).fill("remote"));
+        expect(observedUrls.map((value) => value.split("?")[0])).toEqual([
+          `POST ${prefix}/v1beta/models/gemini-embedding-001:embedContent`,
+          `POST ${prefix}/upload/v1beta/files`,
+          `POST ${prefix}/v1beta/models/gemini-embedding-001:asyncBatchEmbedContent`,
+          `GET ${prefix}/v1beta/batches/b-0`,
+          `GET ${prefix}/v1beta/files/output-0:download`,
+        ]);
+      } finally {
+        remoteHttp.mockRestore();
+        await closeServer(server);
+      }
+    },
+  );
 
   it("honors terminal LRO fields when metadata is stale", async () => {
     stubBatchFetch((stage) =>

extensions/google/embedding-batch.ts

@@ -7,8 +7,8 @@ import {
   debugEmbeddingsLog,
   EmbeddingBatchUnavailableError,
   formatBatchErrorDetail,
-  normalizeBatchBaseUrl,
   readEmbeddingBatchJsonl,
+  resolveEmbeddingEndpointUrl,
   withRemoteHttpResponse,
   type EmbeddingBatchExecutionParams,
 } from "openclaw/plugin-sdk/memory-core-host-engine-embeddings";
@@ -82,30 +82,32 @@ function hashText(text: string): string {
   return crypto.createHash("sha256").update(text).digest("hex");
 }
 
-function getGeminiVersionedRouteBase(baseUrl: string, route: "upload" | "download"): string | null {
-  const trimmed = baseUrl.replace(/\/$/, "");
-  const match = trimmed.match(/^(.*)\/(v\d+(?:alpha|beta)?)$/);
-  return match ? `${match[1]}/${route}/${match[2]}` : null;
-}
-
-function getGeminiUploadUrl(baseUrl: string): string {
-  return getGeminiVersionedRouteBase(baseUrl, "upload") ?? `${baseUrl.replace(/\/$/, "")}/upload`;
-}
-
-function getGeminiDownloadUrl(baseUrl: string, fileId: string): string {
-  const file = fileId.startsWith("files/") ? fileId : `files/${fileId}`;
-  const trimmed = baseUrl.replace(/\/$/, "");
-  let officialGoogleOrigin = false;
-  try {
-    officialGoogleOrigin =
-      new URL(trimmed).origin.toLowerCase() === "https://generativelanguage.googleapis.com";
-  } catch {
-    // Custom base URLs are preserved below.
+function getGeminiBatchFileUrl(
+  baseUrl: string,
+  route: "upload" | "download",
+  fileId: string,
+): string {
+  const base = new URL(baseUrl);
+  const pathname = base.pathname.replace(/\/+$/, "");
+  // Google file routes precede the API version; custom download gateways own their prefix.
+  if (route === "upload" || base.origin === "https://generativelanguage.googleapis.com") {
+    const version = pathname.match(/^(.*)\/(v\d+(?:alpha|beta)?)$/);
+    base.pathname = version
+      ? `${version[1]}/${route}/${version[2]}`
+      : route === "upload"
+        ? `${pathname}/upload`
+        : pathname;
   }
-  const downloadBase = officialGoogleOrigin
-    ? (getGeminiVersionedRouteBase(trimmed, "download") ?? trimmed)
-    : trimmed;
-  return `${downloadBase}/${file}:download?alt=media`;
+  const endpoint =
+    route === "upload"
+      ? fileId
+      : `${fileId.startsWith("files/") ? fileId : `files/${fileId}`}:download`;
+  const url = new URL(resolveEmbeddingEndpointUrl(base.href, endpoint));
+  url.searchParams.set(
+    route === "upload" ? "uploadType" : "alt",
+    route === "upload" ? "multipart" : "media",
+  );
+  return url.href;
 }
 
 function getGeminiBatchState(operation: GeminiBatchOperation): GeminiBatchState {
@@ -180,7 +182,7 @@ async function submitGeminiBatch(params: {
   requests: GeminiBatchRequest[];
   agentId: string;
 }): Promise<GeminiBatchOperation> {
-  const baseUrl = normalizeBatchBaseUrl(params.gemini);
+  const baseUrl = params.gemini.baseUrl;
   const jsonl = params.requests
     .map((request) =>
       JSON.stringify({
@@ -192,7 +194,7 @@ async function submitGeminiBatch(params: {
   const displayName = `memory-embeddings-${hashText(String(Date.now()))}`;
   const uploadPayload = buildGeminiUploadBody({ jsonl, displayName });
 
-  const uploadUrl = `${getGeminiUploadUrl(baseUrl)}/files?uploadType=multipart`;
+  const uploadUrl = getGeminiBatchFileUrl(baseUrl, "upload", "files");
   debugEmbeddingsLog("memory embeddings: gemini batch upload", {
     uploadUrl,
     baseUrl,
@@ -230,7 +232,10 @@ async function submitGeminiBatch(params: {
     },
   };
 
-  const batchEndpoint = `${baseUrl}/${params.gemini.modelPath}:asyncBatchEmbedContent`;
+  const batchEndpoint = resolveEmbeddingEndpointUrl(
+    baseUrl,
+    `${params.gemini.modelPath}:asyncBatchEmbedContent`,
+  );
   debugEmbeddingsLog("memory embeddings: gemini batch create", {
     batchEndpoint,
     fileId,
@@ -265,11 +270,10 @@ async function fetchGeminiBatchStatus(params: {
   batchName: string;
   signal?: AbortSignal;
 }): Promise<GeminiBatchOperation> {
-  const baseUrl = normalizeBatchBaseUrl(params.gemini);
   const name = params.batchName.startsWith("batches/")
     ? params.batchName
     : `batches/${params.batchName}`;
-  const statusUrl = `${baseUrl}/${name}`;
+  const statusUrl = resolveEmbeddingEndpointUrl(params.gemini.baseUrl, name);
   debugEmbeddingsLog("memory embeddings: gemini batch status", { statusUrl });
   return await withRemoteHttpResponse({
     url: statusUrl,
@@ -323,8 +327,7 @@ async function fetchGeminiBatchOutput(params: {
   errors: string[];
   byCustomId: Map<string, number[]>;
 }): Promise<void> {
-  const baseUrl = normalizeBatchBaseUrl(params.gemini);
-  const downloadUrl = getGeminiDownloadUrl(baseUrl, params.fileId);
+  const downloadUrl = getGeminiBatchFileUrl(params.gemini.baseUrl, "download", params.fileId);
   debugEmbeddingsLog("memory embeddings: gemini batch download", { downloadUrl });
   await withRemoteHttpResponse({
     url: downloadUrl,

extensions/google/embedding-provider.ts

@@ -251,41 +251,35 @@ async function fetchGeminiEmbeddingPayload(params: {
 }
 
 function normalizeGeminiBaseUrl(raw: string): string {
-  const trimmed = raw.replace(/\/+$/, "");
-  const openAiIndex = trimmed.indexOf("/openai");
-  if (openAiIndex > -1) {
-    const queryIndex = trimmed.indexOf("?", openAiIndex);
-    return normalizeGoogleApiBaseUrl(
-      `${trimmed.slice(0, openAiIndex)}${queryIndex < 0 ? "" : trimmed.slice(queryIndex)}`,
-    );
-  }
-  return normalizeGoogleApiBaseUrl(trimmed);
-}
-
-function buildGeminiModelPath(model: string): string {
-  return model.startsWith("models/") ? model : `models/${model}`;
-}
-
-function normalizeGoogleApiBaseUrl(baseUrl: string): string {
-  const trimmed = baseUrl.trim().replace(/\/+$/, "");
+  const trimmed = raw.trim();
   if (!trimmed) {
     return DEFAULT_GOOGLE_API_BASE_URL;
   }
   try {
     const url = new URL(trimmed);
     url.hash = "";
+    // OpenAI endpoint aliases and trailing slashes belong to the path, not tenant query values.
+    const openAiIndex = url.pathname.indexOf("/openai");
+    url.pathname = (openAiIndex < 0 ? url.pathname : url.pathname.slice(0, openAiIndex)).replace(
+      /\/+$/,
+      "",
+    );
     if (
       url.origin.toLowerCase() === "https://generativelanguage.googleapis.com" &&
-      url.pathname.replace(/\/+$/, "") === ""
+      url.pathname === "/"
     ) {
       url.pathname = "/v1beta";
     }
-    return url.toString().replace(/\/+$/, "");
+    return url.search ? url.href : url.href.replace(/\/$/, "");
   } catch {
     return trimmed;
   }
 }
 
+function buildGeminiModelPath(model: string): string {
+  return model.startsWith("models/") ? model : `models/${model}`;
+}
+
 export async function createGeminiEmbeddingProvider(
   options: MemoryEmbeddingProviderCreateOptions,
 ): Promise<{ provider: MemoryEmbeddingProvider; client: GeminiEmbeddingClient }> {