test(voyage): move bounded reads to public runner (#121957)

extensions/voyage/embedding-batch.test.ts

@@ -2,9 +2,6 @@
 import { describe, expect, it, vi } from "vitest";
 import { runVoyageEmbeddingBatches } from "./embedding-batch.js";
 import type { VoyageEmbeddingClient } from "./embedding-provider.js";
-import { voyageEmbeddingBatchTesting as testing } from "./test-support.js";
-
-const { fetchVoyageBatchStatus, readVoyageBatchError, VOYAGE_BATCH_RESPONSE_MAX_BYTES } = testing;
 
 function jsonResponse(body: unknown): Response {
   return new Response(JSON.stringify(body), {
@@ -20,25 +17,6 @@ function buildClient(): VoyageEmbeddingClient {
   };
 }
 
-/**
- * Build deps whose withRemoteHttpResponse drives the real onResponse against a
- * caller-provided Response, so the bounded readers run exactly as in production.
- */
-function buildDeps(response: Response): Parameters<typeof fetchVoyageBatchStatus>[0]["deps"] {
-  return {
-    now: () => 0,
-    sleep: async () => {},
-    postJsonWithRetry: async () => {
-      throw new Error("postJsonWithRetry should not be called in these tests");
-    },
-    uploadBatchJsonlFile: (async () => {
-      throw new Error("uploadBatchJsonlFile should not be called in these tests");
-    }) as never,
-    withRemoteHttpResponse: (async (params: { onResponse: (res: Response) => Promise<unknown> }) =>
-      await params.onResponse(response)) as never,
-  };
-}
-
 /**
  * A streaming JSON-ish body that proves an oversized response stops being read
  * before the whole advertised payload is buffered into memory. getReadCount
@@ -51,7 +29,6 @@ function streamingResponse(params: { chunkCount: number; chunkSize: number; stat
 } {
   let reads = 0;
   let canceled = false;
-  const encoder = new TextEncoder();
   const stream = new ReadableStream<Uint8Array>({
     pull(controller) {
       if (reads >= params.chunkCount) {
@@ -59,7 +36,7 @@ function streamingResponse(params: { chunkCount: number; chunkSize: number; stat
         return;
       }
       reads += 1;
-      controller.enqueue(encoder.encode("a".repeat(params.chunkSize)));
+      controller.enqueue(new Uint8Array(params.chunkSize));
     },
     cancel() {
       canceled = true;
@@ -134,143 +111,103 @@ describe("voyage batch bounded reads", () => {
     expect(statusFetch).not.toHaveBeenCalled();
   });
 
-  it("uses a 16 MiB cap for successful status/error-file responses", () => {
-    expect(VOYAGE_BATCH_RESPONSE_MAX_BYTES).toBe(16 * 1024 * 1024);
-  });
-
-  it("parses a well-formed batch status response under the byte cap", async () => {
-    const response = new Response(JSON.stringify({ id: "batch_1", status: "completed" }), {
-      status: 200,
-      headers: { "content-type": "application/json" },
-    });
-
-    const status = await fetchVoyageBatchStatus({
-      client: buildClient(),
-      batchId: "batch_1",
-      deps: buildDeps(response),
-    });
-
-    expect(status).toEqual({ id: "batch_1", status: "completed" });
-  });
-
-  it("caps an oversized batch status stream instead of buffering the whole body", async () => {
-    const streamed = streamingResponse({ chunkCount: 64, chunkSize: 1024 });
+  it("caps an oversized batch status stream through the public runner", async () => {
+    const streamed = streamingResponse({ chunkCount: 20, chunkSize: 1024 * 1024 });
 
     await expect(
-      fetchVoyageBatchStatus({
+      runVoyageEmbeddingBatches({
         client: buildClient(),
-        batchId: "batch_1",
-        deps: buildDeps(streamed.response),
-        maxResponseBytes: 4096,
+        agentId: "main",
+        requests: [{ custom_id: "req-0", body: { input: "hello" } }],
+        wait: true,
+        pollIntervalMs: 1,
+        timeoutMs: 60_000,
+        concurrency: 1,
+        deps: {
+          now: () => 0,
+          sleep: async () => {},
+          uploadBatchJsonlFile: async () => "input-0",
+          postJsonWithRetry: async () => ({ id: "batch-0", status: "in_progress" }),
+          withRemoteHttpResponse: (async (params: {
+            onResponse: (response: Response) => Promise<unknown>;
+          }) => await params.onResponse(streamed.response)) as never,
+        },
       }),
-    ).rejects.toThrow(/voyage-batch-status: JSON response exceeds 4096 bytes/);
+    ).rejects.toThrow(/voyage-batch-status: JSON response exceeds \d+ bytes/);
 
-    // Stream was cancelled mid-flight: fewer chunks read than the full payload.
-    expect(streamed.getReadCount()).toBeLessThan(64);
+    expect(streamed.getReadCount()).toBeLessThan(20);
     expect(streamed.wasCanceled()).toBe(true);
   });
 
-  it("preserves the full NDJSON parse chain for an under-cap error file", async () => {
-    // Multi-line NDJSON with a blank line proves the bounded read does not
-    // disturb the original trim/split("\n")/JSON.parse/extractBatchErrorMessage
-    // pipeline: the first useful error message is still extracted byte-for-byte
-    // identically to the pre-change `await res.text()` path.
-    const body = [
-      JSON.stringify({ custom_id: "req-0", response: { status_code: 200 } }),
-      "",
-      JSON.stringify({ custom_id: "req-1", error: { message: "voyage upstream rejected" } }),
-      JSON.stringify({ custom_id: "req-2", error: { message: "second error ignored" } }),
-      "",
-    ].join("\n");
-    const response = new Response(body, {
-      status: 200,
-      headers: { "content-type": "application/x-ndjson" },
-    });
+  it("fail-softs an oversized error file through the public runner", async () => {
+    const streamed = streamingResponse({ chunkCount: 20, chunkSize: 1024 * 1024 });
 
-    const message = await readVoyageBatchError({
-      client: buildClient(),
-      errorFileId: "file_1",
-      deps: buildDeps(response),
-    });
-
-    // extractBatchErrorMessage returns the first line carrying a message, so the
-    // success line is skipped and the second error is not surfaced.
-    expect(message).toBe("voyage upstream rejected");
-  });
-
-  it("returns undefined for an empty error file via the original empty-body branch", async () => {
-    // Whitespace-only body must still hit the `!text.trim()` short-circuit after
-    // decoding the bounded buffer, returning undefined exactly as before.
-    const response = new Response("   \n", {
-      status: 200,
-      headers: { "content-type": "application/x-ndjson" },
-    });
-
-    const message = await readVoyageBatchError({
-      client: buildClient(),
-      errorFileId: "file_1",
-      deps: buildDeps(response),
-    });
-
-    expect(message).toBeUndefined();
-  });
-
-  it("fail-softs an oversized error file into formatUnavailableBatchError by design", async () => {
-    const streamed = streamingResponse({ chunkCount: 64, chunkSize: 1024 });
-
-    // Intended behavior: an over-cap error file must NOT throw out of
-    // readVoyageBatchError. An unbounded error body would otherwise OOM the
-    // worker, so the bounded overflow error is caught and degraded into a
-    // diagnostic string via formatUnavailableBatchError. We accept the lost
-    // detail; the overflow message names the cap so the truncation is visible.
-    const readError = async () =>
-      await readVoyageBatchError({
+    await expect(
+      runVoyageEmbeddingBatches({
         client: buildClient(),
-        errorFileId: "file_1",
-        deps: buildDeps(streamed.response),
-        maxResponseBytes: 4096,
-      });
-
-    await expect(readError()).resolves.toMatch(
-      /error file unavailable: voyage batch error file content exceeds 4096 bytes/,
+        agentId: "main",
+        requests: [{ custom_id: "req-0", body: { input: "hello" } }],
+        wait: true,
+        pollIntervalMs: 1,
+        timeoutMs: 60_000,
+        concurrency: 1,
+        deps: {
+          uploadBatchJsonlFile: async () => "input-0",
+          postJsonWithRetry: async () => ({
+            id: "batch-0",
+            status: "completed",
+            output_file_id: "output-0",
+            error_file_id: "error-0",
+          }),
+          withRemoteHttpResponse: (async (params: {
+            url: string;
+            onResponse: (response: Response) => Promise<unknown>;
+          }) => {
+            expect(params.url).toContain("/files/error-0/content");
+            return await params.onResponse(streamed.response);
+          }) as never,
+        },
+      }),
+    ).rejects.toThrow(
+      /voyage batch batch-0 completed: error file unavailable: voyage batch error file content exceeds \d+ bytes/,
     );
 
-    // The bounded reader still cancels the stream mid-flight rather than
-    // buffering the whole advertised payload before failing soft.
-    expect(streamed.getReadCount()).toBeLessThan(64);
+    expect(streamed.getReadCount()).toBeLessThan(20);
     expect(streamed.wasCanceled()).toBe(true);
   });
 
-  it("normalizes and bounds a non-OK diagnostic body", async () => {
-    const streamed = streamingResponse({ chunkCount: 64, chunkSize: 1024, status: 500 });
+  it("normalizes and bounds a non-OK status diagnostic through the public runner", async () => {
+    const streamed = streamingResponse({
+      chunkCount: 20,
+      chunkSize: 1024 * 1024,
+      status: 500,
+    });
 
     await expect(
-      fetchVoyageBatchStatus({
+      runVoyageEmbeddingBatches({
         client: buildClient(),
-        batchId: "batch_1",
-        deps: buildDeps(streamed.response),
+        agentId: "main",
+        requests: [{ custom_id: "req-0", body: { input: "hello" } }],
+        wait: true,
+        pollIntervalMs: 1,
+        timeoutMs: 60_000,
+        concurrency: 1,
+        deps: {
+          now: () => 0,
+          sleep: async () => {},
+          uploadBatchJsonlFile: async () => "input-0",
+          postJsonWithRetry: async () => ({ id: "batch-0", status: "in_progress" }),
+          withRemoteHttpResponse: (async (params: {
+            onResponse: (response: Response) => Promise<unknown>;
+          }) => await params.onResponse(streamed.response)) as never,
+        },
       }),
     ).rejects.toMatchObject({ name: "ProviderHttpError", status: 500, statusCode: 500 });
 
-    expect(streamed.getReadCount()).toBeLessThan(64);
+    expect(streamed.getReadCount()).toBeLessThan(20);
     expect(streamed.wasCanceled()).toBe(true);
   });
 
-  it("preserves a small non-OK diagnostic", async () => {
-    const response = new Response("voyage upstream is down", {
-      status: 503,
-      headers: { "content-type": "text/plain" },
-    });
-
-    await expect(
-      fetchVoyageBatchStatus({
-        client: buildClient(),
-        batchId: "batch_1",
-        deps: buildDeps(response),
-      }),
-    ).rejects.toThrow("voyage.batch-status (503): voyage upstream is down");
-  });
-
   it("uses the shared output reader and stops after the expected result", async () => {
     let canceled = false;
     const encoder = new TextEncoder();

extensions/voyage/embedding-batch.ts

@@ -131,10 +131,8 @@ async function fetchVoyageBatchStatus(params: {
   client: VoyageEmbeddingClient;
   batchId: string;
   deps: VoyageBatchDeps;
-  maxResponseBytes?: number;
   signal?: AbortSignal;
 }): Promise<VoyageBatchStatus> {
-  const maxBytes = params.maxResponseBytes ?? VOYAGE_BATCH_RESPONSE_MAX_BYTES;
   return await params.deps.withRemoteHttpResponse(
     buildVoyageBatchRequest({
       client: params.client,
@@ -143,7 +141,7 @@ async function fetchVoyageBatchStatus(params: {
       onResponse: async (res) => {
         await assertOkOrThrowProviderError(res, "voyage.batch-status");
         return await readProviderJsonResponse<VoyageBatchStatus>(res, "voyage-batch-status", {
-          maxBytes,
+          maxBytes: VOYAGE_BATCH_RESPONSE_MAX_BYTES,
         });
       },
     }),
@@ -154,9 +152,7 @@ async function readVoyageBatchError(params: {
   client: VoyageEmbeddingClient;
   errorFileId: string;
   deps: VoyageBatchDeps;
-  maxResponseBytes?: number;
 }): Promise<string | undefined> {
-  const maxBytes = params.maxResponseBytes ?? VOYAGE_BATCH_RESPONSE_MAX_BYTES;
   try {
     return await params.deps.withRemoteHttpResponse(
       buildVoyageBatchRequest({
@@ -164,7 +160,7 @@ async function readVoyageBatchError(params: {
         path: `files/${params.errorFileId}/content`,
         onResponse: async (res) => {
           await assertOkOrThrowProviderError(res, "voyage.batch-error-file-content");
-          const bytes = await readResponseWithLimit(res, maxBytes, {
+          const bytes = await readResponseWithLimit(res, VOYAGE_BATCH_RESPONSE_MAX_BYTES, {
             onOverflow: ({ maxBytes: maxBytesLocal }) =>
               new Error(`voyage batch error file content exceeds ${maxBytesLocal} bytes`),
           });
@@ -363,13 +359,3 @@ export async function runVoyageEmbeddingBatches(
     },
   });
 }
-
-const testing = {
-  fetchVoyageBatchStatus,
-  readVoyageBatchError,
-  VOYAGE_BATCH_RESPONSE_MAX_BYTES,
-} as const;
-
-if (process.env.VITEST === "true") {
-  Reflect.set(globalThis, Symbol.for("openclaw.voyageEmbeddingBatchTestApi"), testing);
-}

extensions/voyage/test-support.ts

@@ -1,22 +0,0 @@
-import type { VoyageEmbeddingClient } from "./embedding-provider.js";
-
-type VoyageBatchTestParams = {
-  client: VoyageEmbeddingClient;
-  deps: Record<string, unknown>;
-  maxResponseBytes?: number;
-};
-
-type VoyageEmbeddingBatchTestApi = {
-  fetchVoyageBatchStatus: (params: VoyageBatchTestParams & { batchId: string }) => Promise<unknown>;
-  readVoyageBatchError: (
-    params: VoyageBatchTestParams & { errorFileId: string },
-  ) => Promise<string | undefined>;
-  VOYAGE_BATCH_RESPONSE_MAX_BYTES: number;
-};
-
-const api = Reflect.get(globalThis, Symbol.for("openclaw.voyageEmbeddingBatchTestApi"));
-if (!api) {
-  throw new Error("Voyage embedding batch test API is unavailable");
-}
-
-export const voyageEmbeddingBatchTesting = api as VoyageEmbeddingBatchTestApi;