From a515f7f8fc2e444d1b60238b828d4db74e7a3188 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Thu, 27 Aug 2026 01:02:56 -0700 Subject: [PATCH] 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 | 249 +++++++++++++++------- extensions/google/embedding-batch.ts | 65 +++--- extensions/google/embedding-provider.ts | 32 ++- 3 files changed, 220 insertions(+), 126 deletions(-) diff --git a/extensions/google/embedding-batch.test.ts b/extensions/google/embedding-batch.test.ts index bfcdae4b87f3..157d9254bb0b 100644 --- a/extensions/google/embedding-batch.test.ts +++ b/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 = []; - 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((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 = []; + const tenantHeaders: Array = []; + const realSdk = await vi.importActual( + "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) => diff --git a/extensions/google/embedding-batch.ts b/extensions/google/embedding-batch.ts index c420d60deb78..ff880ba05896 100644 --- a/extensions/google/embedding-batch.ts +++ b/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 { - 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 { - 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; }): Promise { - 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, diff --git a/extensions/google/embedding-provider.ts b/extensions/google/embedding-provider.ts index 587a4ac28711..e792be37313b 100644 --- a/extensions/google/embedding-provider.ts +++ b/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 }> {