From be5e8ce978b3ff2f9240ce1e4e3a7d52d54fb701 Mon Sep 17 00:00:00 2001 From: NIO Date: Mon, 29 Jun 2026 05:20:28 +0800 Subject: [PATCH] fix(google): bound embedding-batch JSON response reads (#97535) * fix(google): bound embedding-batch JSON response reads * fix(google): repair embedding-batch test types and lint --------- Co-authored-by: NIO (cherry picked from commit 91297bf420822bbfd8883af157fe9d1c77774dde) --- extensions/google/embedding-batch.test.ts | 246 ++++++++++++++++++++++ extensions/google/embedding-batch.ts | 14 +- 2 files changed, 256 insertions(+), 4 deletions(-) create mode 100644 extensions/google/embedding-batch.test.ts diff --git a/extensions/google/embedding-batch.test.ts b/extensions/google/embedding-batch.test.ts new file mode 100644 index 000000000000..336e9e125dd2 --- /dev/null +++ b/extensions/google/embedding-batch.test.ts @@ -0,0 +1,246 @@ +// Google tests cover embedding batch bounded JSON response reads. +import { afterEach, describe, expect, it, vi } from "vitest"; +import { runGeminiEmbeddingBatches } from "./embedding-batch.js"; +import type { GeminiEmbeddingClient } from "./embedding-provider.js"; + +// Pass-through so onResponse receives real Response objects (required by +// readProviderJsonResponse which needs a real .body ReadableStream). +vi.mock("openclaw/plugin-sdk/memory-core-host-engine-embeddings", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + withRemoteHttpResponse: async (params: { + url: string; + ssrfPolicy?: unknown; + init?: RequestInit; + onResponse: (response: Response) => Promise; + }): Promise => { + const response = await fetch(params.url, params.init); + return await params.onResponse(response); + }, + }; +}); + +afterEach(() => { + vi.restoreAllMocks(); + vi.unstubAllGlobals(); +}); + +function fetchInputUrl(input: RequestInfo | URL): string { + if (typeof input === "string") { + return input; + } + if (input instanceof URL) { + return input.href; + } + return input.url; +} + +function jsonResponse(body: unknown, status = 200): Response { + return new Response(JSON.stringify(body), { + status, + headers: { "Content-Type": "application/json" }, + }); +} + +function makeGeminiClient(): GeminiEmbeddingClient { + return { + baseUrl: "https://gemini-compatible.example/v1beta", + model: "text-embedding-004", + modelPath: "models/text-embedding-004", + headers: { "x-goog-api-key": "test-key" }, + apiKeys: ["test-key"], + ssrfPolicy: undefined, + }; +} + +type GeminiBatchRequest = Parameters[0]["requests"][number]; + +function singleRequest(): GeminiBatchRequest[] { + return [ + { + custom_id: "r0", + request: { + model: "models/text-embedding-004", + content: { parts: [{ text: "hello" }] }, + taskType: "RETRIEVAL_DOCUMENT", + }, + }, + ]; +} + +function makeOversizedResponse(): { + response: Response; + getReadCount: () => number; + wasCanceled: () => boolean; +} { + const chunkSize = 1024 * 1024; + const chunkCount = 20; // 20 MiB — over 16 MiB cap + let readCount = 0; + let canceled = false; + return { + response: new Response( + new ReadableStream({ + pull(controller) { + if (readCount >= chunkCount) { + controller.close(); + return; + } + readCount += 1; + controller.enqueue(new Uint8Array(chunkSize)); + }, + cancel() { + canceled = true; + }, + }), + { status: 200, headers: { "Content-Type": "application/json" } }, + ), + getReadCount: () => readCount, + wasCanceled: () => canceled, + }; +} + +describe("Google embedding-batch bounded JSON reads", () => { + it("bounds oversized file-upload JSON response and cancels the stream", async () => { + const streamed = makeOversizedResponse(); + vi.stubGlobal( + "fetch", + vi.fn(async (input: RequestInfo | URL) => { + if (fetchInputUrl(input).includes("/upload/")) { + return streamed.response; + } + return new Response("unexpected", { status: 500 }); + }), + ); + + await expect( + runGeminiEmbeddingBatches({ + gemini: makeGeminiClient(), + agentId: "main", + requests: singleRequest(), + wait: true, + concurrency: 1, + pollIntervalMs: 50, + timeoutMs: 5_000, + }), + ).rejects.toThrow(/gemini\.batch-file-upload/); + + expect(streamed.wasCanceled()).toBe(true); + expect(streamed.getReadCount()).toBeLessThan(20); + }); + + it("bounds oversized batch-create JSON response and cancels the stream", async () => { + const streamed = makeOversizedResponse(); + vi.stubGlobal( + "fetch", + vi.fn(async (input: RequestInfo | URL) => { + const url = fetchInputUrl(input); + if (url.includes("/upload/")) { + return jsonResponse({ name: "files/f-ok" }); + } + if (url.includes(":asyncBatchEmbedContent")) { + return streamed.response; + } + return new Response("unexpected", { status: 500 }); + }), + ); + + await expect( + runGeminiEmbeddingBatches({ + gemini: makeGeminiClient(), + agentId: "main", + requests: singleRequest(), + wait: true, + concurrency: 1, + pollIntervalMs: 50, + timeoutMs: 5_000, + }), + ).rejects.toThrow(/gemini\.batch-create/); + + expect(streamed.wasCanceled()).toBe(true); + expect(streamed.getReadCount()).toBeLessThan(20); + }); + + it("bounds oversized batch-status poll JSON response and cancels the stream", async () => { + const streamed = makeOversizedResponse(); + let statusCalled = false; + vi.stubGlobal( + "fetch", + vi.fn(async (input: RequestInfo | URL) => { + const url = fetchInputUrl(input); + if (url.includes("/upload/")) { + return jsonResponse({ name: "files/f-ok" }); + } + if (url.includes(":asyncBatchEmbedContent")) { + return jsonResponse({ name: "batches/b-0", state: "PENDING" }); + } + if (url.includes("/batches/") && !statusCalled) { + statusCalled = true; + return streamed.response; + } + return new Response("unexpected", { status: 500 }); + }), + ); + + await expect( + runGeminiEmbeddingBatches({ + gemini: makeGeminiClient(), + agentId: "main", + requests: singleRequest(), + wait: true, + concurrency: 1, + pollIntervalMs: 50, + timeoutMs: 5_000, + }), + ).rejects.toThrow(/gemini\.batch-status/); + + expect(streamed.wasCanceled()).toBe(true); + expect(streamed.getReadCount()).toBeLessThan(20); + }); + + it("parses small responses on all three JSON paths correctly", async () => { + // Use a unit-length vector so sanitizeAndNormalizeEmbedding preserves values. + const outputLine = JSON.stringify({ + key: "r0", + embedding: { values: [1, 0, 0] }, + }); + let statusCalled = false; + vi.stubGlobal( + "fetch", + vi.fn(async (input: RequestInfo | URL) => { + const url = fetchInputUrl(input); + if (url.includes("/upload/")) { + return jsonResponse({ name: "files/f-ok" }); + } + if (url.includes(":asyncBatchEmbedContent")) { + return jsonResponse({ name: "batches/b-0", state: "PENDING" }); + } + if (url.includes("/batches/") && !statusCalled) { + statusCalled = true; + return jsonResponse({ + name: "batches/b-0", + state: "SUCCEEDED", + outputConfig: { file: "files/out-0" }, + }); + } + if (url.includes(":download")) { + return new Response(outputLine, { status: 200 }); + } + return new Response("unexpected", { status: 500 }); + }), + ); + + const result = await runGeminiEmbeddingBatches({ + gemini: makeGeminiClient(), + agentId: "main", + requests: singleRequest(), + wait: true, + concurrency: 1, + pollIntervalMs: 50, + timeoutMs: 5_000, + }); + + expect(result.get("r0")).toEqual([1, 0, 0]); + }); +}); diff --git a/extensions/google/embedding-batch.ts b/extensions/google/embedding-batch.ts index ab4c89228e1e..290f6ae6980e 100644 --- a/extensions/google/embedding-batch.ts +++ b/extensions/google/embedding-batch.ts @@ -9,7 +9,10 @@ import { sanitizeAndNormalizeEmbedding, withRemoteHttpResponse, } from "openclaw/plugin-sdk/memory-core-host-engine-embeddings"; -import { createProviderHttpError } from "openclaw/plugin-sdk/provider-http"; +import { + createProviderHttpError, + readProviderJsonResponse, +} from "openclaw/plugin-sdk/provider-http"; import { normalizeStringEntries } from "openclaw/plugin-sdk/string-coerce-runtime"; import type { GeminiEmbeddingClient, GeminiTextEmbeddingRequest } from "./embedding-provider.js"; @@ -126,7 +129,10 @@ async function submitGeminiBatch(params: { const text = await fileRes.text(); throw new Error(`gemini batch file upload failed: ${fileRes.status} ${text}`); } - return (await fileRes.json()) as { name?: string; file?: { name?: string } }; + return readProviderJsonResponse<{ name?: string; file?: { name?: string } }>( + fileRes, + "gemini.batch-file-upload", + ); }, }); const fileId = filePayload.name ?? filePayload.file?.name; @@ -158,7 +164,7 @@ async function submitGeminiBatch(params: { }, onResponse: async (batchRes) => { if (batchRes.ok) { - return (await batchRes.json()) as GeminiBatchStatus; + return readProviderJsonResponse(batchRes, "gemini.batch-create"); } const text = await batchRes.text(); if (batchRes.status === 404) { @@ -191,7 +197,7 @@ async function fetchGeminiBatchStatus(params: { if (!res.ok) { throw await createProviderHttpError(res, "gemini batch status failed"); } - return (await res.json()) as GeminiBatchStatus; + return readProviderJsonResponse(res, "gemini.batch-status"); }, }); }