diff --git a/extensions/google/embedding-batch.test.ts b/extensions/google/embedding-batch.test.ts index 6af6d9b1392..65f43ac2324 100644 --- a/extensions/google/embedding-batch.test.ts +++ b/extensions/google/embedding-batch.test.ts @@ -206,6 +206,42 @@ function makeOversizedResponse(status = 200): { } describe("Google embedding-batch bounded JSON reads", () => { + it("clamps polling to the remaining batch timeout", async () => { + vi.useFakeTimers(); + vi.setSystemTime(0); + const setTimeoutSpy = vi.spyOn(globalThis, "setTimeout"); + const fetchMock = stubBatchFetch(); + + const result = runGeminiEmbeddingBatches({ + gemini: makeGeminiClient(), + agentId: "main", + requests: singleRequest(), + wait: true, + concurrency: 1, + pollIntervalMs: 2_000, + timeoutMs: 1_000, + debug: (message) => { + if (message.includes("batches/b-0 pending")) { + vi.setSystemTime(500); + } + }, + }); + + for (let attempt = 0; attempt < 100 && setTimeoutSpy.mock.calls.length === 0; attempt++) { + await Promise.resolve(); + } + expect(setTimeoutSpy).toHaveBeenCalledTimes(1); + expect(setTimeoutSpy.mock.calls[0]?.[1]).toBe(500); + const rejection = expect(result).rejects.toThrow( + "gemini batch batches/b-0 timed out after 1000ms", + ); + await vi.runAllTimersAsync(); + await rejection; + expect( + fetchMock.mock.calls.filter(([input]) => fetchInputUrl(input).includes("/batches/")), + ).toHaveLength(0); + }); + it.each([ { stage: "upload", label: "gemini.batch-file-upload" }, { stage: "create", label: "gemini.batch-create" }, diff --git a/extensions/google/embedding-batch.ts b/extensions/google/embedding-batch.ts index b7e12c2a1c9..c283b336fec 100644 --- a/extensions/google/embedding-batch.ts +++ b/extensions/google/embedding-batch.ts @@ -15,8 +15,11 @@ import { } from "openclaw/plugin-sdk/memory-core-host-engine-embeddings"; import { assertOkOrThrowProviderError, + createProviderOperationDeadline, createProviderHttpError, readProviderJsonObjectResponse, + resolveProviderOperationTimeoutMs, + waitProviderOperationPollInterval, } from "openclaw/plugin-sdk/provider-http"; import type { GeminiEmbeddingClient, GeminiTextEmbeddingRequest } from "./embedding-provider.js"; import { parseGeminiAuth } from "./gemini-auth.js"; @@ -257,6 +260,7 @@ async function submitGeminiBatch(params: { async function fetchGeminiBatchStatus(params: { gemini: GeminiEmbeddingClient; batchName: string; + signal?: AbortSignal; }): Promise { const baseUrl = normalizeBatchBaseUrl(params.gemini); const name = params.batchName.startsWith("batches/") @@ -267,6 +271,7 @@ async function fetchGeminiBatchStatus(params: { return await withRemoteHttpResponse({ url: statusUrl, ssrfPolicy: params.gemini.ssrfPolicy, + signal: params.signal, init: { headers: buildBatchHeaders(params.gemini, { json: true }), }, @@ -350,15 +355,24 @@ async function waitForGeminiBatch(params: { debug?: (message: string, data?: Record) => void; initial?: GeminiBatchOperation; }): Promise<{ outputFileId: string }> { - const start = Date.now(); + const deadline = createProviderOperationDeadline({ + label: `gemini batch ${params.batchName}`, + timeoutMs: params.timeoutMs, + }); let current: GeminiBatchOperation | undefined = params.initial; while (true) { - const operation = - current ?? - (await fetchGeminiBatchStatus({ - gemini: params.gemini, - batchName: params.batchName, - })); + const operation = current + ? current + : await fetchGeminiBatchStatus({ + gemini: params.gemini, + batchName: params.batchName, + signal: AbortSignal.timeout( + resolveProviderOperationTimeoutMs({ + deadline, + defaultTimeoutMs: params.timeoutMs, + }), + ), + }); const state = getGeminiBatchState(operation); if (state === "succeeded") { const outputFileId = getGeminiBatchOutputFileId(operation); @@ -380,12 +394,12 @@ async function waitForGeminiBatch(params: { `gemini batch ${params.batchName} submitted; enable remote.batch.wait to await completion`, ); } - if (Date.now() - start > params.timeoutMs) { - throw new Error(`gemini batch ${params.batchName} timed out after ${params.timeoutMs}ms`); - } - params.debug?.(`gemini batch ${params.batchName} ${state}; waiting ${params.pollIntervalMs}ms`); - await new Promise((resolve) => { - setTimeout(resolve, params.pollIntervalMs); + params.debug?.( + `gemini batch ${params.batchName} ${state}; waiting up to ${params.pollIntervalMs}ms`, + ); + await waitProviderOperationPollInterval({ + deadline, + pollIntervalMs: params.pollIntervalMs, }); current = undefined; }