fix(google): keep embedding batch waits within timeout (#111674)

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" },

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<GeminiBatchOperation> {
   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<string, unknown>) => 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;
   }