vultr/openclaw
1
0
Fork 0
mirror of https://github.com/openclaw/openclaw.git synced 2026-07-20 22:11:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(memory): harden embedding batch workflows (#103472)

Browse Source 
* fix(memory): harden embedding batch workflows

Co-authored-by: 徐闻涵0668001344 <xu.wenhan1@xydigit.com>
Co-authored-by: ben.li <li.yang6@xydigit.com>

* test(memory): align batch redaction expectation

* test(memory): tighten batch dependency types

* refactor(google): clarify batch state control flow

* test(google): make batch stage fallback lint-safe

* build(plugin-sdk): refresh memory batch surface budget

---------

Co-authored-by: Peter Steinberger <steipete@gmail.com>
Co-authored-by: ben.li <li.yang6@xydigit.com>
This commit is contained in:
ZOOWH
2026-07-12 07:43:58 +08:00
committed by GitHub
parent 847460600f
commit fb702e5cae
19 changed files with 1298 additions and 764 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
docs/.generated/plugin-sdk-api-baseline.sha256
View File

@@ -1,2 +1,2 @@
d8383ea9819598d92e1d1052eeb9adb61f79cd5b16e0991fc60e17a99612848d plugin-sdk-api-baseline.json
b9c7eed49500d9f5d371a25c0bc7b968d55e6272081a6e9be2113f81c2415689 plugin-sdk-api-baseline.jsonl
e54bcec7d07d05d00010ca038776c16d33346b2b69c5c920f331b805584553a0 plugin-sdk-api-baseline.json
af38688ddabaeeef60aa19917c2d96ab507afc2f05df49dc3a009f8f75bddb12 plugin-sdk-api-baseline.jsonl

561
extensions/google/embedding-batch.test.ts
View File

@@ -1,4 +1,5 @@
// Google tests cover embedding batch bounded JSON response reads.
import { createServer } from "node:http";
import { afterEach, describe, expect, it, vi } from "vitest";
import { runGeminiEmbeddingBatches } from "./embedding-batch.js";
import type { GeminiEmbeddingClient } from "./embedding-provider.js";
@@ -44,12 +45,35 @@ function jsonResponse(body: unknown, status = 200): Response {
});
}
function makeGeminiClient(): GeminiEmbeddingClient {
async function listenLoopbackServer(server: ReturnType<typeof createServer>): Promise<number> {
return await new Promise((resolve, reject) => {
server.once("error", reject);
server.listen(0, "127.0.0.1", () => {
server.off("error", reject);
const address = server.address();
if (!address || typeof address === "string") {
reject(new Error("expected loopback TCP address"));
return;
}
resolve(address.port);
});
});
}
async function closeServer(server: ReturnType<typeof createServer>): Promise<void> {
await new Promise<void>((resolve, reject) => {
server.close((error) => (error ? reject(error) : resolve()));
});
}
function makeGeminiClient(
baseUrl = "https://generativelanguage.googleapis.com/v1beta",
): GeminiEmbeddingClient {
return {
baseUrl: "https://gemini-compatible.example/v1beta",
baseUrl,
model: "text-embedding-004",
modelPath: "models/text-embedding-004",
headers: { "x-goog-api-key": "test-key" },
headers: { "x-goog-api-client": "test-client" },
apiKeys: ["test-key"],
ssrfPolicy: undefined,
};
@@ -57,20 +81,100 @@ function makeGeminiClient(): GeminiEmbeddingClient {
type GeminiBatchRequest = Parameters<typeof runGeminiEmbeddingBatches>[0]["requests"][number];
function singleRequest(): GeminiBatchRequest[] {
return [
{
custom_id: "r0",
request: {
model: "models/text-embedding-004",
content: { parts: [{ text: "hello" }] },
taskType: "RETRIEVAL_DOCUMENT",
},
function batchRequest(customId: string, text: string): GeminiBatchRequest {
return {
custom_id: customId,
request: {
model: "models/text-embedding-004",
content: { parts: [{ text }] },
taskType: "RETRIEVAL_DOCUMENT",
},
];
};
}
function makeOversizedResponse(): {
function singleRequest(): GeminiBatchRequest[] {
return [batchRequest("r0", "hello")];
}
type BatchStage = "upload" | "create" | "status" | "download";
function batchStageForUrl(url: string): BatchStage {
if (url.includes("/upload/")) {
return "upload";
}
if (url.includes(":asyncBatchEmbedContent")) {
return "create";
}
if (url.includes("/batches/")) {
return "status";
}
if (url.includes(":download")) {
return "download";
}
throw new Error(`unexpected Gemini batch URL: ${url}`);
}
function defaultBatchResponse(stage: BatchStage): Response {
switch (stage) {
case "upload":
return jsonResponse({ file: { name: "files/f-ok" } });
case "create":
return jsonResponse({
name: "batches/b-0",
done: false,
metadata: { state: "BATCH_STATE_PENDING" },
});
case "status":
return jsonResponse({
name: "batches/b-0",
done: true,
metadata: { state: "BATCH_STATE_SUCCEEDED" },
response: { responsesFile: "files/out-0" },
});
case "download":
return new Response(
JSON.stringify({ key: "r0", response: { embedding: { values: [1, 0, 0] } } }),
{ status: 200 },
);
}
throw new Error("unexpected Gemini batch stage");
}
function stubBatchFetch(
override?: (stage: BatchStage, url: string, init?: RequestInit) => Response | undefined,
): ReturnType<typeof vi.fn> {
const fetchMock = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => {
const url = fetchInputUrl(input);
const stage = batchStageForUrl(url);
return override?.(stage, url, init) ?? defaultBatchResponse(stage);
});
vi.stubGlobal("fetch", fetchMock);
return fetchMock;
}
function runBatch(
requests = singleRequest(),
gemini = makeGeminiClient(),
): Promise<Map<string, number[]>> {
return runGeminiEmbeddingBatches({
gemini,
agentId: "main",
requests,
wait: true,
concurrency: 1,
pollIntervalMs: 1,
timeoutMs: 5_000,
});
}
async function captureRejection(promise: Promise<unknown>): Promise<unknown> {
return await promise.then(
() => undefined,
(error: unknown) => error,
);
}
function makeOversizedResponse(status = 200): {
response: Response;
getReadCount: () => number;
wasCanceled: () => boolean;
@@ -94,7 +198,7 @@ function makeOversizedResponse(): {
canceled = true;
},
}),
{ status: 200, headers: { "Content-Type": "application/json" } },
{ status, headers: { "Content-Type": "application/json" } },
),
getReadCount: () => readCount,
wasCanceled: () => canceled,
@@ -102,225 +206,268 @@ function makeOversizedResponse(): {
}
describe("Google embedding-batch bounded JSON reads", () => {
it("bounds oversized file-upload JSON response and cancels the stream", async () => {
it.each([
{ stage: "upload", label: "gemini.batch-file-upload" },
{ stage: "create", label: "gemini.batch-create" },
{ stage: "status", label: "gemini.batch-status" },
] as const)("bounds oversized successful $stage JSON", async ({ stage, label }) => {
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 });
}),
);
stubBatchFetch((candidate) => (candidate === stage ? streamed.response : undefined));
await expect(
runGeminiEmbeddingBatches({
gemini: makeGeminiClient(),
agentId: "main",
requests: singleRequest(),
wait: true,
concurrency: 1,
pollIntervalMs: 50,
timeoutMs: 5_000,
}),
).rejects.toThrow(/gemini\.batch-file-upload/);
const error = await captureRejection(runBatch());
expect(error).toBeInstanceOf(Error);
expect((error as Error).message).toContain(label);
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 });
}),
);
it.each([
{ stage: "upload", label: "gemini.batch-file-upload" },
{ stage: "create", label: "gemini.batch-create" },
{ stage: "status", label: "gemini.batch-status" },
{ stage: "download", label: "gemini.batch-file-content" },
] as const)("bounds oversized $stage errors", async ({ stage, label }) => {
const streamed = makeOversizedResponse(503);
stubBatchFetch((candidate) => (candidate === stage ? streamed.response : undefined));
await expect(
runGeminiEmbeddingBatches({
gemini: makeGeminiClient(),
agentId: "main",
requests: singleRequest(),
wait: true,
concurrency: 1,
pollIntervalMs: 50,
timeoutMs: 5_000,
}),
).rejects.toThrow(/gemini\.batch-create/);
const error = await captureRejection(runBatch());
expect(error).toMatchObject({ name: "ProviderHttpError", status: 503, statusCode: 503 });
expect((error as Error).message).toContain(label);
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 });
}),
it("marks create 404 as unavailable while preserving the structured cause", async () => {
const response = jsonResponse(
{ error: { code: 404, message: "Input file was not found", status: "NOT_FOUND" } },
404,
);
stubBatchFetch((stage) => (stage === "create" ? response : undefined));
await expect(
runGeminiEmbeddingBatches({
gemini: makeGeminiClient(),
agentId: "main",
requests: singleRequest(),
wait: true,
concurrency: 1,
pollIntervalMs: 50,
timeoutMs: 5_000,
}),
).rejects.toThrow(/gemini\.batch-status/);
const error = await captureRejection(runBatch());
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] },
expect(error).toMatchObject({
name: "EmbeddingBatchUnavailableError",
code: "embedding_batch_unavailable",
});
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((error as Error).message).toContain("asyncBatchEmbedContent not available");
expect((error as Error).cause).toMatchObject({
name: "ProviderHttpError",
status: 404,
code: "NOT_FOUND",
});
expect(((error as Error).cause as Error).message).toContain("Input file was not found");
expect((error as Error).message).not.toContain("switch providers");
expect(response.bodyUsed).toBe(true);
});
it("normalizes raw Google Operations and uses the canonical download route", async () => {
const fetchMock = stubBatchFetch();
const result = await runBatch();
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",
);
for (const [, init] of fetchMock.mock.calls) {
expect(new Headers(init?.headers).get("x-goog-api-key")).toBe("test-key");
}
const createCall = fetchMock.mock.calls.find(([input]) =>
fetchInputUrl(input).includes(":asyncBatchEmbedContent"),
);
expect(JSON.parse(String(createCall?.[1]?.body))).toMatchObject({
batch: { inputConfig: { file_name: "files/f-ok" } },
});
});
it("streams multi-line JSONL batch output across delayed chunks", async () => {
// Genuinely chunked delivery: body is emitted as overlapping mid-line
// chunks with an await between writes, so the readline path cannot rely on
// a single buffered string from `new Response(fullText)`.
const requests: Parameters<typeof runGeminiEmbeddingBatches>[0]["requests"] = [
{
custom_id: "r0",
request: {
model: "models/text-embedding-004",
content: { parts: [{ text: "hello" }] },
taskType: "RETRIEVAL_DOCUMENT",
},
},
{
custom_id: "r1",
request: {
model: "models/text-embedding-004",
content: { parts: [{ text: "world" }] },
taskType: "RETRIEVAL_DOCUMENT",
},
},
];
const line0 = JSON.stringify({ key: "r0", embedding: { values: [1, 0, 0] } });
const line1 = JSON.stringify({ custom_id: "r1", embedding: { values: [0, 1, 0] } });
const jsonlBody = `${line0}\n${line1}\n`;
const mid = Math.floor(line0.length / 2);
const chunks = [
jsonlBody.slice(0, mid),
jsonlBody.slice(mid, line0.length + 1 + Math.floor(line1.length / 3)),
jsonlBody.slice(line0.length + 1 + Math.floor(line1.length / 3)),
];
expect(chunks.join("")).toBe(jsonlBody);
it("preserves a configured gateway prefix for output downloads", async () => {
const fetchMock = stubBatchFetch();
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: "SUCCEEDED",
outputConfig: { file: "files/out-0" },
await runBatch(singleRequest(), makeGeminiClient("https://gateway.example/gemini/v1beta"));
expect(fetchMock.mock.calls.map(([input]) => fetchInputUrl(input))).toContain(
"https://gateway.example/gemini/v1beta/files/out-0:download?alt=media",
);
});
it("runs the complete batch lifecycle over loopback HTTP", async () => {
let createBody: unknown;
const authHeaders: Array<string | undefined> = [];
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<void>((resolve) => {
request.once("end", () => resolve());
});
respondJson({ file: { name: "files/input-0" } });
return;
}
if (url.includes(":download")) {
const stream = new ReadableStream<Uint8Array>({
async start(controller) {
const encoder = new TextEncoder();
for (const chunk of chunks) {
controller.enqueue(encoder.encode(chunk));
await new Promise((resolve) => {
setTimeout(resolve, 5);
});
}
controller.close();
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`),
);
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);
}
});
it("honors terminal LRO fields when metadata is stale", async () => {
stubBatchFetch((stage) =>
stage === "create"
? jsonResponse({
name: "batches/b-0",
done: true,
metadata: { state: "BATCH_STATE_RUNNING" },
response: { responsesFile: "files/out-0" },
})
: undefined,
);
await expect(runBatch()).resolves.toEqual(new Map([["r0", [1, 0, 0]]]));
});
it("keeps a terminal Operation error ahead of stale success metadata", async () => {
stubBatchFetch((stage) =>
stage === "create"
? jsonResponse({
name: "batches/b-0",
done: true,
metadata: { state: "BATCH_STATE_SUCCEEDED" },
response: { responsesFile: "files/out-0" },
error: { code: 13, message: "provider job failed" },
})
: undefined,
);
await expect(runBatch()).rejects.toThrow("gemini batch batches/b-0 failed");
});
it("keeps shipped compatible-endpoint output aliases", async () => {
const requests = [batchRequest("r0", "hello"), batchRequest("r1", "world")];
stubBatchFetch((stage) =>
stage === "download"
? new Response(
[
JSON.stringify({ custom_id: "r0", embedding: { values: [1, 0] } }),
JSON.stringify({ request_id: "r1", embedding: { values: [0, 1] } }),
].join("\n"),
)
: undefined,
);
await expect(runBatch(requests)).resolves.toEqual(
new Map([
["r0", [1, 0]],
["r1", [0, 1]],
]),
);
});
it("falls back from an empty top-level output error", async () => {
stubBatchFetch((stage) =>
stage === "download"
? new Response(
JSON.stringify({
key: "r0",
error: { message: "" },
response: { error: { message: "nested output error" } },
}),
)
: undefined,
);
await expect(runBatch()).rejects.toThrow("nested output error");
});
it.each([
{ state: "BATCH_STATE_FAILED", normalized: "failed" },
{ state: "JOB_STATE_CANCELLED", normalized: "cancelled" },
{ state: "BATCH_STATE_EXPIRED", normalized: "expired" },
])("surfaces $state Operation failures", async ({ state, normalized }) => {
stubBatchFetch((stage) =>
stage === "create"
? jsonResponse({
name: "batches/b-0",
done: true,
metadata: { state },
})
: undefined,
);
await expect(runBatch()).rejects.toThrow(`gemini batch batches/b-0 ${normalized}`);
});
it("rejects conflicting output files in one Operation", async () => {
stubBatchFetch((stage) =>
stage === "create"
? jsonResponse({
name: "batches/b-0",
done: true,
metadata: {
state: "BATCH_STATE_SUCCEEDED",
output: { responsesFile: "files/metadata-output" },
},
});
return new Response(stream, { status: 200 });
}
return new Response("unexpected", { status: 500 });
}),
response: { responsesFile: "files/response-output" },
})
: undefined,
);
const result = await runGeminiEmbeddingBatches({
gemini: makeGeminiClient(),
agentId: "main",
requests,
wait: true,
concurrency: 1,
pollIntervalMs: 50,
timeoutMs: 5_000,
});
expect(result.get("r0")).toEqual([1, 0, 0]);
expect(result.get("r1")).toEqual([0, 1, 0]);
await expect(runBatch()).rejects.toThrow("conflicting output files");
});
});

330
extensions/google/embedding-batch.ts
View File

@@ -1,48 +1,48 @@
// Google plugin module implements embedding batch behavior.
import crypto from "node:crypto";
import { createInterface } from "node:readline";
import { Readable } from "node:stream";
import {
buildEmbeddingBatchGroupOptions,
runEmbeddingBatchGroups,
buildBatchHeaders,
debugEmbeddingsLog,
EmbeddingBatchUnavailableError,
formatBatchErrorDetail,
normalizeBatchBaseUrl,
readEmbeddingBatchJsonl,
sanitizeAndNormalizeEmbedding,
withRemoteHttpResponse,
type EmbeddingBatchExecutionParams,
} from "openclaw/plugin-sdk/memory-core-host-engine-embeddings";
import {
assertOkOrThrowProviderError,
createProviderHttpError,
readProviderJsonResponse,
readProviderJsonObjectResponse,
} from "openclaw/plugin-sdk/provider-http";
import type { GeminiEmbeddingClient, GeminiTextEmbeddingRequest } from "./embedding-provider.js";
type EmbeddingBatchExecutionParams = {
wait: boolean;
pollIntervalMs: number;
timeoutMs: number;
concurrency: number;
debug?: (message: string, data?: Record<string, unknown>) => void;
};
import { parseGeminiAuth } from "./gemini-auth.js";
type GeminiBatchRequest = {
custom_id: string;
request: GeminiTextEmbeddingRequest;
};
type GeminiBatchStatus = {
type GeminiBatchOperation = {
name?: string;
state?: string;
outputConfig?: { file?: string; fileId?: string };
done?: boolean;
metadata?: {
state?: string;
output?: {
responsesFile?: string;
};
};
error?: { message?: string };
response?: { responsesFile?: string };
error?: { code?: number; message?: string };
};
type GeminiBatchState = "pending" | "succeeded" | "failed" | "cancelled" | "expired" | "unknown";
type GeminiBatchOutputLine = {
// Alternate ids and direct embeddings are shipped compatible-endpoint shapes.
key?: string;
custom_id?: string;
request_id?: string;
@@ -55,15 +55,93 @@ type GeminiBatchOutputLine = {
};
const GEMINI_BATCH_MAX_REQUESTS = 50000;
function bindGeminiBatchAuth(client: GeminiEmbeddingClient): GeminiEmbeddingClient {
const apiKey = client.apiKeys[0];
if (!apiKey) {
throw new Error("gemini batch requires an API key");
}
// Files and batch operations are credential-scoped. Keep one selected
// credential for upload, creation, polling, and output download.
return {
...client,
headers: {
...parseGeminiAuth(apiKey).headers,
...client.headers,
},
};
}
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 {
if (baseUrl.includes("/v1beta")) {
return baseUrl.replace(/\/v1beta\/?$/, "/upload/v1beta");
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.
}
return `${baseUrl.replace(/\/$/, "")}/upload`;
const downloadBase = officialGoogleOrigin
? (getGeminiVersionedRouteBase(trimmed, "download") ?? trimmed)
: trimmed;
return `${downloadBase}/${file}:download?alt=media`;
}
function getGeminiBatchState(operation: GeminiBatchOperation): GeminiBatchState {
// REST discovery uses BATCH_STATE_* while the public guide and SDK expose
// JOB_STATE_* for the same operation metadata.
const rawState = operation.metadata?.state?.replace(/^(?:BATCH|JOB)_STATE_/, "");
if (rawState === "FAILED") {
return "failed";
}
if (rawState === "CANCELLED" || rawState === "CANCELED") {
return "cancelled";
}
if (rawState === "EXPIRED") {
return "expired";
}
if (operation.error) {
return "failed";
}
if (operation.done === false) {
return "pending";
}
if (operation.done === true) {
return "succeeded";
}
if (rawState === "SUCCEEDED") {
return "succeeded";
}
if (rawState === "PENDING" || rawState === "RUNNING") {
return "pending";
}
return "unknown";
}
function getGeminiBatchOutputFileId(operation: GeminiBatchOperation): string | undefined {
// Google currently documents response.responsesFile while the official SDK
// consumes metadata.output.responsesFile. Accept both raw Operation shapes.
const responseFile = operation.response?.responsesFile;
const metadataFile = operation.metadata?.output?.responsesFile;
if (responseFile && metadataFile && responseFile !== metadataFile) {
throw new Error("gemini batch operation returned conflicting output files");
}
return responseFile ?? metadataFile;
}
function buildGeminiUploadBody(params: { jsonl: string; displayName: string }): {
@@ -95,7 +173,7 @@ async function submitGeminiBatch(params: {
gemini: GeminiEmbeddingClient;
requests: GeminiBatchRequest[];
agentId: string;
}): Promise<GeminiBatchStatus> {
}): Promise<GeminiBatchOperation> {
const baseUrl = normalizeBatchBaseUrl(params.gemini);
const jsonl = params.requests
.map((request) =>
@@ -126,17 +204,13 @@ async function submitGeminiBatch(params: {
body: uploadPayload.body,
},
onResponse: async (fileRes) => {
if (!fileRes.ok) {
const text = await fileRes.text();
throw new Error(`gemini batch file upload failed: ${fileRes.status} ${text}`);
}
return readProviderJsonResponse<{ name?: string; file?: { name?: string } }>(
fileRes,
"gemini.batch-file-upload",
);
await assertOkOrThrowProviderError(fileRes, "gemini.batch-file-upload");
return (await readProviderJsonObjectResponse(fileRes, "gemini.batch-file-upload")) as {
file?: { name?: string };
};
},
});
const fileId = filePayload.name ?? filePayload.file?.name;
const fileId = filePayload.file?.name;
if (!fileId) {
throw new Error("gemini batch file upload failed: missing file id");
}
@@ -164,16 +238,18 @@ async function submitGeminiBatch(params: {
body: JSON.stringify(batchBody),
},
onResponse: async (batchRes) => {
if (batchRes.ok) {
return readProviderJsonResponse<GeminiBatchStatus>(batchRes, "gemini.batch-create");
}
const text = await batchRes.text();
if (batchRes.status === 404) {
throw new Error(
"gemini batch create failed: 404 (asyncBatchEmbedContent not available for this model/baseUrl). Disable remote.batch.enabled or switch providers.",
const cause = await createProviderHttpError(batchRes, "gemini.batch-create");
throw new EmbeddingBatchUnavailableError(
"gemini asyncBatchEmbedContent not available for this request",
{ cause },
);
}
throw new Error(`gemini batch create failed: ${batchRes.status} ${text}`);
await assertOkOrThrowProviderError(batchRes, "gemini.batch-create");
return (await readProviderJsonObjectResponse(
batchRes,
"gemini.batch-create",
)) as GeminiBatchOperation;
},
});
}
@@ -181,7 +257,7 @@ async function submitGeminiBatch(params: {
async function fetchGeminiBatchStatus(params: {
gemini: GeminiEmbeddingClient;
batchName: string;
}): Promise<GeminiBatchStatus> {
}): Promise<GeminiBatchOperation> {
const baseUrl = normalizeBatchBaseUrl(params.gemini);
const name = params.batchName.startsWith("batches/")
? params.batchName
@@ -195,68 +271,39 @@ async function fetchGeminiBatchStatus(params: {
headers: buildBatchHeaders(params.gemini, { json: true }),
},
onResponse: async (res) => {
if (!res.ok) {
throw await createProviderHttpError(res, "gemini batch status failed");
}
return readProviderJsonResponse<GeminiBatchStatus>(res, "gemini.batch-status");
await assertOkOrThrowProviderError(res, "gemini.batch-status");
return (await readProviderJsonObjectResponse(
res,
"gemini.batch-status",
)) as GeminiBatchOperation;
},
});
}
/**
* Streams the JSONL batch output body line by line, parsing each embedding
* directly into the result maps. Avoids holding the full raw text and all
* parsed objects in memory simultaneously — batch outputs can reach hundreds
* of MB for large embedding jobs.
*
* Mirrors the pattern in extensions/voyage/embedding-batch.ts.
*/
async function readGeminiBatchOutputContent(
contentRes: Response,
remaining: Set<string>,
errors: string[],
byCustomId: Map<string, number[]>,
): Promise<void> {
if (!contentRes.body) {
function applyGeminiBatchOutputLine(params: {
line: GeminiBatchOutputLine;
remaining: Set<string>;
errors: string[];
byCustomId: Map<string, number[]>;
}): void {
const customId = params.line.key ?? params.line.custom_id ?? params.line.request_id;
// Only the first response for a submitted id may mutate results.
if (!customId || !params.remaining.delete(customId)) {
return;
}
const inputStream = Readable.fromWeb(
contentRes.body as unknown as import("stream/web").ReadableStream,
);
const reader = createInterface({ input: inputStream, terminal: false });
try {
for await (const rawLine of reader) {
const trimmed = rawLine.trim();
if (!trimmed) {
continue;
}
const line = JSON.parse(trimmed) as GeminiBatchOutputLine;
const customId = line.key ?? line.custom_id ?? line.request_id;
if (!customId) {
continue;
}
remaining.delete(customId);
if (line.error?.message) {
errors.push(`${customId}: ${line.error.message}`);
continue;
}
if (line.response?.error?.message) {
errors.push(`${customId}: ${line.response.error.message}`);
continue;
}
const embedding = sanitizeAndNormalizeEmbedding(
line.embedding?.values ?? line.response?.embedding?.values ?? [],
);
if (embedding.length === 0) {
errors.push(`${customId}: empty embedding`);
continue;
}
byCustomId.set(customId, embedding);
}
} finally {
reader.close();
inputStream.destroy();
const error = params.line.error?.message || params.line.response?.error?.message;
if (error) {
params.errors.push(`${customId}: ${error}`);
return;
}
const embedding = sanitizeAndNormalizeEmbedding(
params.line.embedding?.values ?? params.line.response?.embedding?.values ?? [],
);
if (embedding.length === 0) {
params.errors.push(`${customId}: empty embedding`);
return;
}
params.byCustomId.set(customId, embedding);
}
async function fetchGeminiBatchOutput(params: {
@@ -267,8 +314,7 @@ async function fetchGeminiBatchOutput(params: {
byCustomId: Map<string, number[]>;
}): Promise<void> {
const baseUrl = normalizeBatchBaseUrl(params.gemini);
const file = params.fileId.startsWith("files/") ? params.fileId : `files/${params.fileId}`;
const downloadUrl = `${baseUrl}/${file}:download`;
const downloadUrl = getGeminiDownloadUrl(baseUrl, params.fileId);
debugEmbeddingsLog("memory embeddings: gemini batch download", { downloadUrl });
await withRemoteHttpResponse({
url: downloadUrl,
@@ -277,10 +323,20 @@ async function fetchGeminiBatchOutput(params: {
headers: buildBatchHeaders(params.gemini, { json: true }),
},
onResponse: async (res) => {
if (!res.ok) {
throw await createProviderHttpError(res, "gemini batch file content failed");
}
await readGeminiBatchOutputContent(res, params.remaining, params.errors, params.byCustomId);
await assertOkOrThrowProviderError(res, "gemini.batch-file-content");
await readEmbeddingBatchJsonl<GeminiBatchOutputLine>(res, {
label: "gemini.batch-file-content",
maxRecords: params.remaining.size,
onRecord: (line) => {
applyGeminiBatchOutputLine({
line,
remaining: params.remaining,
errors: params.errors,
byCustomId: params.byCustomId,
});
return params.errors.length === 0 && params.remaining.size > 0;
},
});
},
});
}
@@ -292,34 +348,37 @@ async function waitForGeminiBatch(params: {
pollIntervalMs: number;
timeoutMs: number;
debug?: (message: string, data?: Record<string, unknown>) => void;
initial?: GeminiBatchStatus;
initial?: GeminiBatchOperation;
}): Promise<{ outputFileId: string }> {
const start = Date.now();
let current: GeminiBatchStatus | undefined = params.initial;
let current: GeminiBatchOperation | undefined = params.initial;
while (true) {
const status =
const operation =
current ??
(await fetchGeminiBatchStatus({
gemini: params.gemini,
batchName: params.batchName,
}));
const state = status.state ?? "UNKNOWN";
if (["SUCCEEDED", "COMPLETED", "DONE"].includes(state)) {
const outputFileId =
status.outputConfig?.file ??
status.outputConfig?.fileId ??
status.metadata?.output?.responsesFile;
const state = getGeminiBatchState(operation);
if (state === "succeeded") {
const outputFileId = getGeminiBatchOutputFileId(operation);
if (!outputFileId) {
throw new Error(`gemini batch ${params.batchName} completed without output file`);
}
return { outputFileId };
}
if (["FAILED", "CANCELLED", "CANCELED", "EXPIRED"].includes(state)) {
const message = status.error?.message ?? "unknown error";
throw new Error(`gemini batch ${params.batchName} ${state}: ${message}`);
if (state === "failed" || state === "cancelled" || state === "expired") {
const rawMessage =
operation.error?.message ??
(operation.error?.code === undefined ? "unknown error" : `code ${operation.error.code}`);
throw new Error(
`gemini batch ${params.batchName} ${state}: ${formatBatchErrorDetail(rawMessage) ?? "unknown error"}`,
);
}
if (!params.wait) {
throw new Error(`gemini batch ${params.batchName} still ${state}; wait disabled`);
throw new Error(
`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`);
@@ -339,6 +398,7 @@ export async function runGeminiEmbeddingBatches(
requests: GeminiBatchRequest[];
} & EmbeddingBatchExecutionParams,
): Promise<Map<string, number[]>> {
const gemini = bindGeminiBatchAuth(params.gemini);
return await runEmbeddingBatchGroups({
...buildEmbeddingBatchGroupOptions(params, {
maxRequests: GEMINI_BATCH_MAX_REQUESTS,
@@ -346,7 +406,7 @@ export async function runGeminiEmbeddingBatches(
}),
runGroup: async ({ group, groupIndex, groups, byCustomId, pollIntervalMs, timeoutMs }) => {
const batchInfo = await submitGeminiBatch({
gemini: params.gemini,
gemini,
requests: group,
agentId: params.agentId,
});
@@ -357,48 +417,26 @@ export async function runGeminiEmbeddingBatches(
params.debug?.("memory embeddings: gemini batch created", {
batchName,
state: batchInfo.state,
state: getGeminiBatchState(batchInfo),
group: groupIndex + 1,
groups,
requests: group.length,
});
if (
!params.wait &&
batchInfo.state &&
!["SUCCEEDED", "COMPLETED", "DONE"].includes(batchInfo.state)
) {
throw new Error(
`gemini batch ${batchName} submitted; enable remote.batch.wait to await completion`,
);
}
const completed =
batchInfo.state && ["SUCCEEDED", "COMPLETED", "DONE"].includes(batchInfo.state)
? {
outputFileId:
batchInfo.outputConfig?.file ??
batchInfo.outputConfig?.fileId ??
batchInfo.metadata?.output?.responsesFile ??
"",
}
: await waitForGeminiBatch({
gemini: params.gemini,
batchName,
wait: params.wait,
pollIntervalMs,
timeoutMs,
debug: params.debug,
initial: batchInfo,
});
if (!completed.outputFileId) {
throw new Error(`gemini batch ${batchName} completed without output file`);
}
const completed = await waitForGeminiBatch({
gemini,
batchName,
wait: params.wait,
pollIntervalMs,
timeoutMs,
debug: params.debug,
initial: batchInfo,
});
const errors: string[] = [];
const remaining = new Set(group.map((request) => request.custom_id));
await fetchGeminiBatchOutput({
gemini: params.gemini,
gemini,
fileId: completed.outputFileId,
remaining,
errors,
@@ -406,7 +444,9 @@ export async function runGeminiEmbeddingBatches(
});
if (errors.length > 0) {
throw new Error(`gemini batch ${batchName} failed: ${errors.join("; ")}`);
throw new Error(
`gemini batch ${batchName} failed: ${formatBatchErrorDetail(errors[0]) ?? "unknown error"}`,
);
}
if (remaining.size > 0) {
throw new Error(`gemini batch ${batchName} missing ${remaining.size} embedding responses`);

26
extensions/google/embedding-provider.ts
View File

@@ -25,6 +25,7 @@ import {
asOptionalRecord as asRecord,
normalizeOptionalString,
} from "openclaw/plugin-sdk/string-coerce-runtime";
import { parseGeminiAuth } from "./gemini-auth.js";
import { resolveGoogleApiClientHeaders } from "./google-api-client-header.js";
export type GeminiEmbeddingClient = {
@@ -45,31 +46,6 @@ const GEMINI_MAX_INPUT_TOKENS: Record<string, number> = {
"gemini-embedding-2-preview": 8192,
};
function parseGeminiAuth(apiKey: string): { headers: Record<string, string> } {
if (apiKey.startsWith("{")) {
try {
const parsed = JSON.parse(apiKey) as { token?: string };
if (typeof parsed.token === "string" && parsed.token) {
return {
headers: {
Authorization: `Bearer ${parsed.token}`,
"Content-Type": "application/json",
},
};
}
} catch {
// Fall back to API-key auth below.
}
}
return {
headers: {
"x-goog-api-key": apiKey,
"Content-Type": "application/json",
},
};
}
type GeminiTaskType = NonNullable<MemoryEmbeddingProviderCreateOptions["taskType"]>;
// --- gemini-embedding-2-preview support ---

24
extensions/memory-core/src/memory/index.test.ts
View File

@@ -760,6 +760,30 @@ describe("memory index", () => {
}
});
it("disables batch immediately when the provider reports it unavailable", async () => {
providerRuntimeBatchErrors = [
Object.assign(new Error("provider batch unavailable"), {
code: "embedding_batch_unavailable",
}),
];
const manager = await getFreshManager(
createCfg({ provider: "batch-wide-test", batchEnabled: true }),
);
try {
await manager.sync({ reason: "test" });
expect(providerRuntimeBatchCalls).toHaveLength(1);
expect(embedBatchCalls).toBe(1);
expect(manager.status().batch).toMatchObject({
enabled: false,
failures: 2,
lastError: "provider batch unavailable",
});
} finally {
await manager.close?.();
}
});
it.each([
["frozen errors", Object.freeze(new Error("provider runtime retry failed"))],
["primitive rejections", "provider runtime retry failed"],

3
extensions/memory-core/src/memory/manager-embedding-ops.ts
View File

@@ -4,6 +4,7 @@ import { formatErrorMessage } from "openclaw/plugin-sdk/error-runtime";
import {
enforceEmbeddingMaxInputTokens,
hasNonTextEmbeddingParts,
isEmbeddingBatchUnavailableError,
type EmbeddingInput,
type MemoryEmbeddingProviderRuntime,
} from "openclaw/plugin-sdk/memory-core-host-engine-embeddings";
@@ -676,7 +677,7 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps {
}
const message = formatErrorMessage(result.error);
const forceDisable = /asyncBatchEmbedContent not available/i.test(message);
const forceDisable = isEmbeddingBatchUnavailableError(result.error);
const failure = await this.recordBatchFailure({
provider: params.provider,
message,

115
extensions/openai/embedding-batch.test.ts
View File

@@ -89,6 +89,62 @@ describe("OpenAI embedding batch output", () => {
);
});
it("reads a completed error file before downloading successful output", async () => {
let outputFetched = false;
const fetchImpl = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => {
const url = fetchInputUrl(input);
if (url.endsWith("/files") && init?.method === "POST") {
return jsonResponse({ id: "input-0" });
}
if (url.endsWith("/batches") && init?.method === "POST") {
return jsonResponse({
id: "batch-0",
status: "completed",
output_file_id: "output-0",
error_file_id: "error-0",
});
}
if (url.endsWith("/files/error-0/content")) {
return new Response(
JSON.stringify({
custom_id: "0",
response: { status_code: 500, message: "provider rejected request" },
error: null,
}),
);
}
if (url.endsWith("/files/output-0/content")) {
outputFetched = true;
}
return new Response("unexpected request", { status: 500 });
});
await expect(
runOpenAiEmbeddingBatches({
openAi: {
baseUrl: "https://openai-compatible.example/v1",
headers: { Authorization: "Bearer test" },
model: "text-embedding-3-small",
fetchImpl,
},
agentId: "main",
requests: [
{
custom_id: "0",
method: "POST",
url: "/v1/embeddings",
body: { model: "text-embedding-3-small", input: "payload" },
},
],
wait: true,
concurrency: 1,
pollIntervalMs: 1,
timeoutMs: 1_000,
}),
).rejects.toThrow("openai batch batch-0 completed: provider rejected request");
expect(outputFetched).toBe(false);
});
it("splits provider uploads by serialized JSONL byte cap", async () => {
const requests: Parameters<typeof runOpenAiEmbeddingBatches>[0]["requests"] = Array.from(
{ length: 3 },
@@ -586,6 +642,63 @@ describe("OpenAI embedding batch output", () => {
expect(outputChunksSent).toBeLessThan(outputChunkCount);
});
it("retries normalized transient batch-status errors", async () => {
let statusCalls = 0;
const fetchImpl = vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => {
const url = fetchInputUrl(input);
if (url.endsWith("/files") && init?.method === "POST") {
return jsonResponse({ id: "file-0" });
}
if (url.endsWith("/batches") && init?.method === "POST") {
return jsonResponse({ id: "batch-0", status: "in_progress" });
}
if (url.endsWith("/batches/batch-0")) {
statusCalls += 1;
return statusCalls === 1
? jsonResponse({ error: { message: "retry status" } }, 503)
: jsonResponse({
id: "batch-0",
status: "completed",
output_file_id: "output-0",
});
}
if (url.endsWith("/files/output-0/content")) {
return new Response(
JSON.stringify({
custom_id: "0",
response: { status_code: 200, body: { data: [{ embedding: [1] }] } },
}),
);
}
return new Response("unexpected request", { status: 500 });
});
const result = await runOpenAiEmbeddingBatches({
openAi: {
baseUrl: "https://openai-compatible.example/v1",
headers: { Authorization: "Bearer test" },
model: "text-embedding-3-small",
fetchImpl,
},
agentId: "main",
requests: [
{
custom_id: "0",
method: "POST",
url: "/v1/embeddings",
body: { model: "text-embedding-3-small", input: "payload" },
},
],
wait: true,
concurrency: 1,
pollIntervalMs: 1,
timeoutMs: 1_000,
});
expect(result).toEqual(new Map([["0", [1]]]));
expect(statusCalls).toBe(2);
});
it("bounds batch resource error bodies without using response.text()", async () => {
const tracked = cancelTrackedResponse(`${"batch status unavailable ".repeat(1024)}tail`, {
status: 400,
@@ -633,7 +746,7 @@ describe("OpenAI embedding batch output", () => {
pollIntervalMs: 1,
timeoutMs: 60_000,
}),
).rejects.toThrow(/openai batch status failed: 400 batch status unavailable/);
).rejects.toMatchObject({ name: "ProviderHttpError", status: 400, statusCode: 400 });
expect(tracked.wasCanceled()).toBe(true);
expect(textSpy).not.toHaveBeenCalled();
});

190
extensions/openai/embedding-batch.ts
View File

@@ -5,13 +5,17 @@ import {
buildEmbeddingBatchGroupOptions,
EMBEDDING_BATCH_ENDPOINT,
extractBatchErrorMessage,
formatBatchErrorDetail,
formatUnavailableBatchError,
normalizeBatchBaseUrl,
postJsonWithRetry,
readEmbeddingBatchJsonl,
resolveBatchCompletionFromStatus,
resolveCompletedBatchResult,
runEmbeddingBatchGroups,
throwIfBatchCompletionError,
throwIfBatchTerminalFailure,
type EmbeddingBatchExecutionParams,
type EmbeddingBatchStatus,
type BatchCompletionResult,
type ProviderBatchOutputLine,
@@ -19,21 +23,13 @@ import {
withRemoteHttpResponse,
} from "openclaw/plugin-sdk/memory-core-host-engine-embeddings";
import {
assertOkOrThrowProviderError,
readProviderJsonResponse,
readProviderTextResponse,
readResponseTextLimited,
} from "openclaw/plugin-sdk/provider-http";
import { normalizeStringEntries } from "openclaw/plugin-sdk/string-coerce-runtime";
import type { OpenAiEmbeddingClient } from "./embedding-provider.js";
type EmbeddingBatchExecutionParams = {
wait: boolean;
pollIntervalMs: number;
timeoutMs: number;
concurrency: number;
debug?: (message: string, data?: Record<string, unknown>) => void;
};
type OpenAiBatchRequest = {
custom_id: string;
method: "POST";
@@ -60,8 +56,6 @@ const OPENAI_BATCH_MAX_REQUESTS = 50000;
// splitter avoids boundary-size uploads while preserving source-wide batching.
const OPENAI_BATCH_MAX_JSONL_BYTES = 190 * 1024 * 1024;
const OPENAI_BATCH_MAX_POLL_BACKOFF_MS = 5 * 60_000;
const OPENAI_BATCH_ERROR_BODY_LIMIT_BYTES = 8 * 1024;
const OPENAI_BATCH_OUTPUT_LINE_MAX_BYTES = 4 * 1024 * 1024;
async function submitOpenAiBatch(params: {
openAi: OpenAiEmbeddingClient;
@@ -100,7 +94,7 @@ async function fetchOpenAiBatchStatus(params: {
return await fetchOpenAiBatchResource({
openAi: params.openAi,
path: `/batches/${params.batchId}`,
errorPrefix: "openai batch status",
label: "openai.batch-status",
parse: async (res) => readProviderJsonResponse<OpenAiBatchStatus>(res, "openai.batch-status"),
});
}
@@ -112,111 +106,11 @@ async function fetchOpenAiFileContent(params: {
return await fetchOpenAiBatchResource({
openAi: params.openAi,
path: `/files/${params.fileId}/content`,
errorPrefix: "openai batch file content",
label: "openai.batch-file-content",
parse: async (res) => await readProviderTextResponse(res, "openai.batch-file-content"),
});
}
async function readOpenAiBatchOutputLines(
response: Response,
params: {
maxLines: number;
onLine: (line: OpenAiBatchOutputLine) => boolean;
},
): Promise<void> {
let lineCount = 0;
const emitOutputLine = (line: OpenAiBatchOutputLine): boolean => {
lineCount += 1;
if (lineCount > params.maxLines) {
throw new Error(`openai.batch-file-content: JSONL output exceeds ${params.maxLines} records`);
}
return params.onLine(line);
};
const emitParsedLine = (line: string): boolean =>
emitOutputLine(parseOpenAiBatchOutputLine(line));
const reader = response.body?.getReader();
if (!reader) {
const text = await readProviderTextResponse(response, "openai.batch-file-content", {
maxBytes: OPENAI_BATCH_OUTPUT_LINE_MAX_BYTES,
});
for (const line of parseOpenAiBatchOutput(text)) {
if (!emitOutputLine(line)) {
break;
}
}
return;
}
const decoder = new TextDecoder();
const encoder = new TextEncoder();
let line = "";
let lineBytes = 0;
const appendSegment = (segment: string) => {
if (!segment) {
return;
}
lineBytes += encoder.encode(segment).byteLength;
if (lineBytes > OPENAI_BATCH_OUTPUT_LINE_MAX_BYTES) {
throw new Error(
`openai.batch-file-content: JSONL line exceeds ${OPENAI_BATCH_OUTPUT_LINE_MAX_BYTES} bytes`,
);
}
line += segment;
};
const emitLine = (): boolean => {
const trimmed = line.trim();
line = "";
lineBytes = 0;
if (trimmed) {
return emitParsedLine(trimmed);
}
return true;
};
const consumeText = (text: string): boolean => {
let offset = 0;
while (true) {
const newline = text.indexOf("\n", offset);
if (newline === -1) {
appendSegment(text.slice(offset));
return true;
}
appendSegment(text.slice(offset, newline));
if (!emitLine()) {
return false;
}
offset = newline + 1;
}
};
try {
while (true) {
const { done, value } = await reader.read();
if (done) {
break;
}
if (value && value.byteLength > 0) {
if (!consumeText(decoder.decode(value, { stream: true }))) {
await reader.cancel().catch(() => {});
return;
}
}
}
if (!consumeText(decoder.decode())) {
return;
}
if (line.trim()) {
emitLine();
}
} catch (error) {
await reader.cancel().catch(() => {});
throw error;
} finally {
reader.releaseLock();
}
}
async function readOpenAiBatchOutputFile(params: {
openAi: OpenAiEmbeddingClient;
fileId: string;
@@ -226,11 +120,12 @@ async function readOpenAiBatchOutputFile(params: {
return await fetchOpenAiBatchResource({
openAi: params.openAi,
path: `/files/${params.fileId}/content`,
errorPrefix: "openai batch file content",
label: "openai.batch-file-content",
parse: async (res) =>
await readOpenAiBatchOutputLines(res, {
maxLines: params.maxLines,
onLine: params.onLine,
await readEmbeddingBatchJsonl<OpenAiBatchOutputLine>(res, {
label: "openai.batch-file-content",
maxRecords: params.maxLines,
onRecord: params.onLine,
}),
});
}
@@ -238,7 +133,7 @@ async function readOpenAiBatchOutputFile(params: {
async function fetchOpenAiBatchResource<T>(params: {
openAi: OpenAiEmbeddingClient;
path: string;
errorPrefix: string;
label: string;
parse: (res: Response) => Promise<T>;
}): Promise<T> {
const baseUrl = normalizeBatchBaseUrl(params.openAi);
@@ -250,10 +145,7 @@ async function fetchOpenAiBatchResource<T>(params: {
headers: buildBatchHeaders(params.openAi, { json: true }),
},
onResponse: async (res) => {
if (!res.ok) {
const text = await readResponseTextLimited(res, OPENAI_BATCH_ERROR_BODY_LIMIT_BYTES);
throw new Error(`${params.errorPrefix} failed: ${res.status} ${text}`);
}
await assertOkOrThrowProviderError(res, params.label);
return await params.parse(res);
},
});
@@ -263,6 +155,10 @@ function formatOpenAiBatchError(error: unknown): string {
return error instanceof Error ? error.message : String(error);
}
function formatOpenAiBatchDiagnostic(error: unknown): string {
return formatBatchErrorDetail(formatOpenAiBatchError(error)) ?? "unknown error";
}
function isOpenAiBatchUploadTooLargeError(error: unknown): boolean {
const message = formatOpenAiBatchError(error);
if (!/openai batch file upload failed/i.test(message)) {
@@ -303,7 +199,7 @@ async function readOpenAiBatchError(params: {
fileId: params.errorFileId,
});
const lines = parseOpenAiBatchOutput(content);
return extractBatchErrorMessage(lines);
return formatBatchErrorDetail(extractBatchErrorMessage(lines));
} catch (err) {
return formatUnavailableBatchError(err);
}
@@ -336,14 +232,17 @@ function formatOpenAiBatchProgress(status: OpenAiBatchStatus): string {
return `; progress ${completed}/${counts.total} failed=${failed}`;
}
function formatOpenAiBatchPollError(error: unknown): string {
return error instanceof Error ? error.message : String(error);
}
function isRetryableOpenAiBatchPollError(error: unknown): boolean {
const message = formatOpenAiBatchPollError(error);
const message = formatOpenAiBatchError(error);
const status =
error && typeof error === "object" ? (error as { status?: unknown }).status : undefined;
return (
/openai batch status failed: (408|409|425|429|5\d\d)\b/i.test(message) ||
(typeof status === "number" &&
(status === 408 ||
status === 409 ||
status === 425 ||
status === 429 ||
(status >= 500 && status <= 599))) ||
/\b(ECONNRESET|ECONNREFUSED|ETIMEDOUT|EAI_AGAIN)\b|fetch failed|network error/i.test(message)
);
}
@@ -380,9 +279,7 @@ async function waitForOpenAiBatch(params: {
}
const delayMs = pollBackoff.nextDelayMs();
params.debug?.(
`openai batch ${params.batchId} status check failed: ${formatOpenAiBatchPollError(
error,
)}; waiting ${delayMs}ms`,
`openai batch ${params.batchId} status check failed: ${formatOpenAiBatchDiagnostic(error)}; waiting ${delayMs}ms`,
);
await new Promise((resolve) => {
setTimeout(resolve, delayMs);
@@ -391,6 +288,15 @@ async function waitForOpenAiBatch(params: {
continue;
}
const state = status.status ?? "unknown";
await throwIfBatchCompletionError({
provider: "openai",
status: { ...status, id: params.batchId },
readError: async (errorFileId) =>
await readOpenAiBatchError({
openAi: params.openAi,
errorFileId,
}),
});
if (state === "completed") {
return resolveBatchCompletionFromStatus({
provider: "openai",
@@ -446,7 +352,7 @@ export async function runOpenAiEmbeddingBatches(
requests: group.length,
parts: parts.map((part) => part.length),
depth,
error: formatOpenAiBatchError(error),
error: formatOpenAiBatchDiagnostic(error),
});
},
runGroup: async ({ group, groupIndex, groups, byCustomId, pollIntervalMs, timeoutMs }) => {
@@ -468,6 +374,13 @@ export async function runOpenAiEmbeddingBatches(
requests: group.length,
});
await throwIfBatchCompletionError({
provider: "openai",
status: batchInfo,
readError: async (errorFileId) =>
await readOpenAiBatchError({ openAi: params.openAi, errorFileId }),
});
const completed = await resolveCompletedBatchResult({
provider: "openai",
status: batchInfo,
@@ -492,13 +405,18 @@ export async function runOpenAiEmbeddingBatches(
fileId: completed.outputFileId,
maxLines: group.length,
onLine: (line) => {
applyEmbeddingBatchOutputLine({ line, remaining, errors, byCustomId });
return remaining.size > 0;
// Only the first response for a submitted id may mutate results.
if (line.custom_id && remaining.has(line.custom_id)) {
applyEmbeddingBatchOutputLine({ line, remaining, errors, byCustomId });
}
return errors.length === 0 && remaining.size > 0;
},
});
if (errors.length > 0) {
throw new Error(`openai batch ${batchInfo.id} failed: ${errors.join("; ")}`);
throw new Error(
`openai batch ${batchInfo.id} failed: ${formatBatchErrorDetail(errors[0]) ?? "unknown error"}`,
);
}
if (remaining.size > 0) {
throw new Error(

201
extensions/voyage/embedding-batch.test.ts
View File

@@ -1,14 +1,15 @@
// Voyage batch tests cover bounded status/error response reads.
import { describe, expect, it } from "vitest";
import { testing } from "./embedding-batch.js";
import { runVoyageEmbeddingBatches, testing } from "./embedding-batch.js";
import type { VoyageEmbeddingClient } from "./embedding-provider.js";
const {
fetchVoyageBatchStatus,
readVoyageBatchError,
readBatchOutputContent,
VOYAGE_BATCH_RESPONSE_MAX_BYTES,
} = testing;
const { fetchVoyageBatchStatus, readVoyageBatchError, VOYAGE_BATCH_RESPONSE_MAX_BYTES } = testing;
function jsonResponse(body: unknown): Response {
return new Response(JSON.stringify(body), {
headers: { "content-type": "application/json" },
});
}
function buildClient(): VoyageEmbeddingClient {
return {
@@ -26,9 +27,9 @@ function buildDeps(response: Response): Parameters<typeof fetchVoyageBatchStatus
return {
now: () => 0,
sleep: async () => {},
postJsonWithRetry: (async () => {
postJsonWithRetry: async () => {
throw new Error("postJsonWithRetry should not be called in these tests");
}) as never,
},
uploadBatchJsonlFile: (async () => {
throw new Error("uploadBatchJsonlFile should not be called in these tests");
}) as never,
@@ -74,7 +75,7 @@ function streamingResponse(params: { chunkCount: number; chunkSize: number; stat
}
describe("voyage batch bounded reads", () => {
it("uses a 16 MiB cap for batch status/error responses", () => {
it("uses a 16 MiB cap for successful status/error-file responses", () => {
expect(VOYAGE_BATCH_RESPONSE_MAX_BYTES).toBe(16 * 1024 * 1024);
});
@@ -181,12 +182,7 @@ describe("voyage batch bounded reads", () => {
expect(streamed.wasCanceled()).toBe(true);
});
it("caps an oversized non-OK (error) diagnostic body instead of buffering it whole", async () => {
// Regression for the non-OK gap: `assertVoyageResponseOk` previously read the
// 4xx/5xx diagnostic body with an unbounded `await res.text()`. A hostile
// endpoint can return a 500 with a never-ending body, so that read must be
// bounded too. Drive a streaming 500 through the real status path and assert
// the bounded overflow error fires and the stream is cancelled mid-flight.
it("normalizes and bounds a non-OK diagnostic body", async () => {
const streamed = streamingResponse({ chunkCount: 64, chunkSize: 1024, status: 500 });
await expect(
@@ -194,18 +190,14 @@ describe("voyage batch bounded reads", () => {
client: buildClient(),
batchId: "batch_1",
deps: buildDeps(streamed.response),
maxResponseBytes: 4096,
}),
).rejects.toThrow(/voyage batch status failed: 500 \(error body exceeds 4096 bytes\)/);
).rejects.toMatchObject({ name: "ProviderHttpError", status: 500, statusCode: 500 });
// Stream was cancelled mid-flight rather than draining the whole body.
expect(streamed.getReadCount()).toBeLessThan(64);
expect(streamed.wasCanceled()).toBe(true);
});
it("preserves the diagnostic shape for a small non-OK (error) body", async () => {
// Under-cap non-OK body must still surface the original
// `${context}: ${status} ${text}` diagnostic byte-for-byte.
it("preserves a small non-OK diagnostic", async () => {
const response = new Response("voyage upstream is down", {
status: 503,
headers: { "content-type": "text/plain" },
@@ -217,94 +209,103 @@ describe("voyage batch bounded reads", () => {
batchId: "batch_1",
deps: buildDeps(response),
}),
).rejects.toThrow(/voyage batch status failed: 503 voyage upstream is down/);
});
});
describe("readBatchOutputContent stream cleanup", () => {
it("returns without error on a null body", async () => {
const response = new Response(null, { status: 200 });
const remaining = new Set<string>();
const errors: string[] = [];
const byCustomId = new Map<string, number[]>();
await expect(
readBatchOutputContent(response, remaining, errors, byCustomId),
).resolves.toBeUndefined();
).rejects.toThrow("voyage.batch-status (503): voyage upstream is down");
});
it("applies every valid JSONL line and tracks remaining and errors", async () => {
const body = [
JSON.stringify({
custom_id: "req-0",
response: { status_code: 200, body: { data: [{ embedding: [0.1, 0.2] }] } },
}),
JSON.stringify({
custom_id: "req-1",
response: { status_code: 200, body: { data: [{ embedding: [0.3, 0.4] }] } },
}),
JSON.stringify({ custom_id: "req-2", error: { message: "voyage failed" } }),
"",
].join("\n");
const response = new Response(body, {
status: 200,
headers: { "content-type": "application/x-ndjson" },
});
const remaining = new Set(["req-0", "req-1", "req-2", "req-3"]);
const errors: string[] = [];
const byCustomId = new Map<string, number[]>();
await readBatchOutputContent(response, remaining, errors, byCustomId);
// valid lines should be removed from remaining
expect(remaining.has("req-0")).toBe(false);
expect(remaining.has("req-1")).toBe(false);
// Error line is tracked; applyEmbeddingBatchOutputLine also removes
// error custom_ids from remaining (both consumed and errored lines)
expect(remaining.has("req-2")).toBe(false);
expect(remaining.has("req-3")).toBe(true);
// error line should be tracked in errors array
expect(errors).toHaveLength(1);
expect(errors[0]).toContain("req-2");
});
it.runIf(process.platform === "linux")(
"cancels the response stream when JSON.parse encounters malformed JSONL",
async () => {
let wasCanceled = false;
const encoder = new TextEncoder();
const stream = new ReadableStream<Uint8Array>({
start(controller) {
// First line: valid
it("uses the shared output reader and stops after the expected result", async () => {
let canceled = false;
const encoder = new TextEncoder();
const output = new Response(
new ReadableStream<Uint8Array>({
pull(controller) {
controller.enqueue(
encoder.encode(
JSON.stringify({
`${JSON.stringify({
custom_id: "req-0",
response: { status_code: 200, body: { data: [{ embedding: [0.1] }] } },
}) + "\n",
response: { status_code: 200, body: { data: [{ embedding: [1, 2] }] } },
})}\n`,
),
);
// Second line: malformed JSON — this should trigger the stream cancel
controller.enqueue(encoder.encode("{not valid json}\n"));
},
cancel() {
wasCanceled = true;
canceled = true;
},
});
const response = new Response(stream, {
status: 200,
headers: { "content-type": "application/x-ndjson" },
});
const remaining = new Set(["req-0"]);
const errors: string[] = [];
const byCustomId = new Map<string, number[]>();
}),
);
await expect(
readBatchOutputContent(response, remaining, errors, byCustomId),
).rejects.toThrow();
const result = await runVoyageEmbeddingBatches({
client: buildClient(),
agentId: "main",
requests: [{ custom_id: "req-0", body: { input: "hello" } }],
wait: true,
pollIntervalMs: 1,
timeoutMs: 1_000,
concurrency: 1,
deps: {
uploadBatchJsonlFile: async () => "input-0",
postJsonWithRetry: async () => ({
id: "batch-0",
status: "completed",
output_file_id: "output-0",
}),
withRemoteHttpResponse: (async (params: {
onResponse: (response: Response) => Promise<unknown>;
}) => await params.onResponse(output)) as never,
},
});
// The stream should have been destroyed/canceled, not left dangling
expect(wasCanceled).toBe(true);
},
);
expect(result).toEqual(new Map([["req-0", [1, 2]]]));
expect(canceled).toBe(true);
});
it("reads a completed error file before downloading successful output", async () => {
let outputFetched = false;
await expect(
runVoyageEmbeddingBatches({
client: buildClient(),
agentId: "main",
requests: [{ custom_id: "req-0", body: { input: "hello" } }],
wait: true,
pollIntervalMs: 1,
timeoutMs: 1_000,
concurrency: 1,
deps: {
uploadBatchJsonlFile: async () => "input-0",
postJsonWithRetry: async () => ({
id: "batch-0",
status: "in_progress",
}),
withRemoteHttpResponse: (async (params: {
url: string;
onResponse: (response: Response) => Promise<unknown>;
}) => {
if (params.url.endsWith("/batches/batch-0")) {
return await params.onResponse(
jsonResponse({
id: "batch-0",
status: "completed",
output_file_id: "output-0",
error_file_id: "error-0",
}),
);
}
if (params.url.endsWith("/files/output-0/content")) {
outputFetched = true;
}
return await params.onResponse(
new Response(
JSON.stringify({
custom_id: "req-0",
response: { status_code: 500, message: "provider rejected request" },
error: null,
}),
),
);
}) as never,
},
}),
).rejects.toThrow("voyage batch batch-0 completed: provider rejected request");
expect(outputFetched).toBe(false);
});
});

123
extensions/voyage/embedding-batch.ts
View File

@@ -1,18 +1,19 @@
// Voyage plugin module implements embedding batch behavior.
import { createInterface } from "node:readline";
import { Readable } from "node:stream";
import {
applyEmbeddingBatchOutputLine,
buildBatchHeaders,
buildEmbeddingBatchGroupOptions,
EMBEDDING_BATCH_ENDPOINT,
extractBatchErrorMessage,
formatBatchErrorDetail,
formatUnavailableBatchError,
normalizeBatchBaseUrl,
postJsonWithRetry,
readEmbeddingBatchJsonl,
resolveBatchCompletionFromStatus,
resolveCompletedBatchResult,
runEmbeddingBatchGroups,
throwIfBatchCompletionError,
throwIfBatchTerminalFailure,
type EmbeddingBatchExecutionParams,
type EmbeddingBatchStatus,
@@ -21,7 +22,10 @@ import {
uploadBatchJsonlFile,
withRemoteHttpResponse,
} from "openclaw/plugin-sdk/memory-core-host-engine-embeddings";
import { readProviderJsonResponse } from "openclaw/plugin-sdk/provider-http";
import {
assertOkOrThrowProviderError,
readProviderJsonResponse,
} from "openclaw/plugin-sdk/provider-http";
import { readResponseWithLimit } from "openclaw/plugin-sdk/response-limit-runtime";
import { normalizeStringEntries } from "openclaw/plugin-sdk/string-coerce-runtime";
import type { VoyageEmbeddingClient } from "./embedding-provider.js";
@@ -43,15 +47,14 @@ type VoyageBatchOutputLine = ProviderBatchOutputLine;
const VOYAGE_BATCH_ENDPOINT = EMBEDDING_BATCH_ENDPOINT;
const VOYAGE_BATCH_COMPLETION_WINDOW = "12h";
const VOYAGE_BATCH_MAX_REQUESTS = 50000;
// Voyage batch status/error responses are untrusted external bodies. Cap them
// the same way other bundled providers do (16 MiB) so a misbehaving or hostile
// endpoint cannot stream an unbounded body into memory before we parse it.
// Successful status/error-file responses are untrusted external bodies. Cap
// them at 16 MiB; non-OK diagnostics use the shared bounded provider prefix.
const VOYAGE_BATCH_RESPONSE_MAX_BYTES = 16 * 1024 * 1024;
type VoyageBatchDeps = {
now: () => number;
sleep: (ms: number) => Promise<void>;
postJsonWithRetry: typeof postJsonWithRetry;
postJsonWithRetry: typeof postJsonWithRetry<VoyageBatchStatus>;
uploadBatchJsonlFile: typeof uploadBatchJsonlFile;
withRemoteHttpResponse: typeof withRemoteHttpResponse;
};
@@ -71,27 +74,6 @@ function resolveVoyageBatchDeps(overrides: Partial<VoyageBatchDeps> | undefined)
};
}
async function assertVoyageResponseOk(
res: Response,
context: string,
maxBytes: number = VOYAGE_BATCH_RESPONSE_MAX_BYTES,
): Promise<void> {
if (!res.ok) {
// The non-OK diagnostic body is just as untrusted as the success body: a
// misbehaving or hostile endpoint can return a 4xx/5xx with an unbounded
// body, and the old `await res.text()` buffered it whole before we threw.
// Read it through the same bounded reader (16 MiB cap, stream cancelled on
// overflow) while preserving the original `${context}: ${status} ${text}`
// diagnostic shape for backward compatibility.
const bytes = await readResponseWithLimit(res, maxBytes, {
onOverflow: ({ maxBytes: maxBytesLocal }) =>
new Error(`${context}: ${res.status} (error body exceeds ${maxBytesLocal} bytes)`),
});
const text = new TextDecoder().decode(bytes);
throw new Error(`${context}: ${res.status} ${text}`);
}
}
function buildVoyageBatchRequest<T>(params: {
client: VoyageEmbeddingClient;
path: string;
@@ -122,7 +104,7 @@ async function submitVoyageBatch(params: {
});
// 2. Create batch job using Voyage Batches API
return await params.deps.postJsonWithRetry<VoyageBatchStatus>({
return await params.deps.postJsonWithRetry({
url: `${baseUrl}/batches`,
headers: buildBatchHeaders(params.client, { json: true }),
ssrfPolicy: params.client.ssrfPolicy,
@@ -155,7 +137,7 @@ async function fetchVoyageBatchStatus(params: {
client: params.client,
path: `batches/${params.batchId}`,
onResponse: async (res) => {
await assertVoyageResponseOk(res, "voyage batch status failed", maxBytes);
await assertOkOrThrowProviderError(res, "voyage.batch-status");
return await readProviderJsonResponse<VoyageBatchStatus>(res, "voyage-batch-status", {
maxBytes,
});
@@ -177,7 +159,7 @@ async function readVoyageBatchError(params: {
client: params.client,
path: `files/${params.errorFileId}/content`,
onResponse: async (res) => {
await assertVoyageResponseOk(res, "voyage batch error file content failed", maxBytes);
await assertOkOrThrowProviderError(res, "voyage.batch-error-file-content");
const bytes = await readResponseWithLimit(res, maxBytes, {
onOverflow: ({ maxBytes: maxBytesLocal }) =>
new Error(`voyage batch error file content exceeds ${maxBytesLocal} bytes`),
@@ -189,7 +171,7 @@ async function readVoyageBatchError(params: {
const lines = normalizeStringEntries(text.split("\n")).map(
(line) => JSON.parse(line) as VoyageBatchOutputLine,
);
return extractBatchErrorMessage(lines);
return formatBatchErrorDetail(extractBatchErrorMessage(lines));
},
}),
);
@@ -219,6 +201,16 @@ async function waitForVoyageBatch(params: {
deps: params.deps,
}));
const state = status.status ?? "unknown";
await throwIfBatchCompletionError({
provider: "voyage",
status: { ...status, id: params.batchId },
readError: async (errorFileId) =>
await readVoyageBatchError({
client: params.client,
errorFileId,
deps: params.deps,
}),
});
if (state === "completed") {
return resolveBatchCompletionFromStatus({
provider: "voyage",
@@ -282,6 +274,13 @@ export async function runVoyageEmbeddingBatches(
requests: group.length,
});
await throwIfBatchCompletionError({
provider: "voyage",
status: batchInfo,
readError: async (errorFileId) =>
await readVoyageBatchError({ client: params.client, errorFileId, deps }),
});
const completed = await resolveCompletedBatchResult({
provider: "voyage",
status: batchInfo,
@@ -310,19 +309,26 @@ export async function runVoyageEmbeddingBatches(
headers: buildBatchHeaders(params.client, { json: true }),
},
onResponse: async (contentRes) => {
// Same bounded non-OK diagnostic read as the status/error-file paths:
// the failure body is untrusted, so cap it instead of `await text()`.
await assertVoyageResponseOk(contentRes, "voyage batch file content failed");
await assertOkOrThrowProviderError(contentRes, "voyage.batch-file-content");
if (!contentRes.body) {
return;
}
await readBatchOutputContent(contentRes, remaining, errors, byCustomId);
await readEmbeddingBatchJsonl<VoyageBatchOutputLine>(contentRes, {
label: "voyage.batch-file-content",
maxRecords: group.length,
onRecord: (line) => {
// Only the first response for a submitted id may mutate results.
if (line.custom_id && remaining.has(line.custom_id)) {
applyEmbeddingBatchOutputLine({ line, remaining, errors, byCustomId });
}
return errors.length === 0 && remaining.size > 0;
},
});
},
});
if (errors.length > 0) {
throw new Error(`voyage batch ${batchInfo.id} failed: ${errors.join("; ")}`);
throw new Error(
`voyage batch ${batchInfo.id} failed: ${formatBatchErrorDetail(errors[0]) ?? "unknown error"}`,
);
}
if (remaining.size > 0) {
throw new Error(
@@ -333,45 +339,8 @@ export async function runVoyageEmbeddingBatches(
});
}
/**
* Stream-read a Voyage batch output file response body line by line, applying
* each parsed output line through `applyEmbeddingBatchOutputLine`.
*
* The response body is an untrusted external stream. This helper wraps the
* readline iteration in a try-finally so the readline interface and underlying
* Node.js Readable are always closed / destroyed, even when JSON.parse or
* applyEmbeddingBatchOutputLine throws.
*/
async function readBatchOutputContent(
contentRes: Response,
remaining: Set<string>,
errors: string[],
byCustomId: Map<string, number[]>,
): Promise<void> {
if (!contentRes.body) {
return;
}
const inputStream = Readable.fromWeb(
contentRes.body as unknown as import("stream/web").ReadableStream,
);
const reader = createInterface({ input: inputStream, terminal: false });
try {
for await (const rawLine of reader) {
if (!rawLine.trim()) {
continue;
}
const line = JSON.parse(rawLine) as VoyageBatchOutputLine;
applyEmbeddingBatchOutputLine({ line, remaining, errors, byCustomId });
}
} finally {
reader.close();
inputStream.destroy();
}
}
export const testing = {
fetchVoyageBatchStatus,
readVoyageBatchError,
readBatchOutputContent,
VOYAGE_BATCH_RESPONSE_MAX_BYTES,
} as const;

11
packages/memory-host-sdk/src/engine-embeddings.ts
View File

@@ -17,9 +17,15 @@ export type {
MemoryEmbeddingProviderRuntime,
} from "./host/openclaw-runtime-memory.js";
export { createLocalEmbeddingProvider, DEFAULT_LOCAL_MODEL } from "./host/embeddings.js";
export { extractBatchErrorMessage, formatUnavailableBatchError } from "./host/batch-error-utils.js";
export {
EmbeddingBatchUnavailableError,
extractBatchErrorMessage,
formatBatchErrorDetail,
formatUnavailableBatchError,
isEmbeddingBatchUnavailableError,
} from "./host/batch-error-utils.js";
export { postJsonWithRetry } from "./host/batch-http.js";
export { applyEmbeddingBatchOutputLine } from "./host/batch-output.js";
export { applyEmbeddingBatchOutputLine, readEmbeddingBatchJsonl } from "./host/batch-output.js";
export {
EMBEDDING_BATCH_ENDPOINT,
type EmbeddingBatchStatus,
@@ -33,6 +39,7 @@ export {
export {
resolveBatchCompletionFromStatus,
resolveCompletedBatchResult,
throwIfBatchCompletionError,
throwIfBatchTerminalFailure,
type BatchCompletionResult,
} from "./host/batch-status.js";

82
packages/memory-host-sdk/src/host/batch-error-utils.test.ts
View File

@@ -1,6 +1,12 @@
// Memory Host SDK tests cover batch error utils behavior.
import { describe, expect, it } from "vitest";
import { extractBatchErrorMessage, formatUnavailableBatchError } from "../engine-embeddings.js";
import {
EmbeddingBatchUnavailableError,
extractBatchErrorMessage,
formatBatchErrorDetail,
formatUnavailableBatchError,
isEmbeddingBatchUnavailableError,
} from "../engine-embeddings.js";
describe("extractBatchErrorMessage", () => {
it("returns the first top-level error message", () => {
@@ -22,6 +28,54 @@ describe("extractBatchErrorMessage", () => {
expect(extractBatchErrorMessage([{ response: { body: "provider plain-text error" } }])).toBe(
"provider plain-text error",
);
expect(
extractBatchErrorMessage([{ response: { body: "", message: "response fallback" } }]),
).toBe("response fallback");
});
it("accepts Voyage response messages", () => {
expect(
extractBatchErrorMessage([
{ response: { status_code: 500, message: "Internal Server Error" }, error: null },
]),
).toBe("Internal Server Error");
expect(
extractBatchErrorMessage([
{ response: { status_code: 500, message: "nested fallback" }, error: { message: "" } },
]),
).toBe("nested fallback");
});
});
describe("EmbeddingBatchUnavailableError", () => {
it("survives duplicate module instances through its stable code", () => {
const error = new EmbeddingBatchUnavailableError("not available", {
cause: new Error("provider detail"),
});
expect(error).toMatchObject({
name: "EmbeddingBatchUnavailableError",
code: "embedding_batch_unavailable",
message: "not available",
});
expect(error.cause).toBeInstanceOf(Error);
expect(isEmbeddingBatchUnavailableError(error)).toBe(true);
expect(isEmbeddingBatchUnavailableError({ code: "embedding_batch_unavailable" })).toBe(true);
expect(isEmbeddingBatchUnavailableError(new Error("other"))).toBe(false);
});
});
describe("formatBatchErrorDetail", () => {
it("preserves short details and redacts and bounds long provider text", () => {
expect(formatBatchErrorDetail("short detail")).toBe("short detail");
const secret = `sk-${"a".repeat(24)}`;
const formatted = formatBatchErrorDetail(`API_TOKEN=${secret} ${"😀".repeat(400)}`);
expect(formatted).toMatch(/\.\.\. \[truncated\]$/);
expect(formatted?.length).toBeLessThanOrEqual(500);
expect(formatted).not.toContain(secret);
expect(formatted).not.toMatch(/[\uD800-\uDFFF]/u);
});
});
@@ -31,34 +85,14 @@ describe("formatUnavailableBatchError", () => {
expect(formatUnavailableBatchError("unreachable")).toBe("error file unavailable: unreachable");
});
it.each([
{
boundary: "leading",
secret: `abcde😀${"x".repeat(9)}wxyz`,
expected: "abcde...wxyz",
},
{
boundary: "trailing",
secret: `abcdef${"x".repeat(9)}😀abc`,
expected: "abcdef...abc",
},
{
boundary: "intact pairs",
secret: `abcd😀${"x".repeat(9)}😀ab`,
expected: "abcd😀...😀ab",
},
{
boundary: "ASCII",
secret: "abcdef1234567890ghij",
expected: "abcdef...ghij",
},
])("keeps exported $boundary token hints UTF-16 safe", ({ secret, expected }) => {
it("redacts exported tokens without malformed UTF-16", () => {
const secret = `abcd😀${"x".repeat(18)}😀ab`;
const serialized = JSON.stringify({
error: formatUnavailableBatchError(new Error(`API_TOKEN=${secret}`)),
});
const parsed = JSON.parse(serialized) as { error: string };
expect(parsed.error).toBe(`error file unavailable: API_TOKEN=${expected}`);
expect(parsed.error).toBe("error file unavailable: API_TOKEN=***");
expect(parsed.error).not.toMatch(/[\uD800-\uDFFF]/u);
expect(serialized).not.toContain(secret);
});

64
packages/memory-host-sdk/src/host/batch-error-utils.ts
View File

@@ -1,40 +1,68 @@
// Memory Host SDK helper module supports batch error utils behavior.
import { truncateUtf16Safe } from "@openclaw/normalization-core/utf16-slice";
import type { EmbeddingBatchOutputLine } from "./batch-output.js";
import { formatErrorMessage } from "./error-utils.js";
// Extracts provider batch error text from output and unavailable error files.
/** Minimal batch output line shape that can carry provider error messages. */
type BatchOutputErrorLike = {
error?: { message?: string };
response?: {
body?:
| string
| {
error?: { message?: string };
};
};
};
const BATCH_ERROR_DETAIL_MAX_CHARS = 500;
const BATCH_ERROR_DETAIL_TRUNCATED_SUFFIX = "... [truncated]";
const EMBEDDING_BATCH_UNAVAILABLE_CODE = "embedding_batch_unavailable";
/** Signals that a provider cannot run the configured embedding batch operation. */
export class EmbeddingBatchUnavailableError extends Error {
readonly code = EMBEDDING_BATCH_UNAVAILABLE_CODE;
constructor(message: string, options?: ErrorOptions) {
super(message, options);
this.name = "EmbeddingBatchUnavailableError";
}
}
export function isEmbeddingBatchUnavailableError(error: unknown): boolean {
if (!error || typeof error !== "object") {
return false;
}
try {
return (error as { code?: unknown }).code === EMBEDDING_BATCH_UNAVAILABLE_CODE;
} catch {
return false;
}
}
/** Pull a nested response error message without assuming a fixed provider body shape. */
function getResponseErrorMessage(line: BatchOutputErrorLike | undefined): string | undefined {
function getResponseErrorMessage(line: EmbeddingBatchOutputLine | undefined): string | undefined {
const body = line?.response?.body;
if (typeof body === "string") {
return body || undefined;
return body || line?.response?.message || undefined;
}
if (!body || typeof body !== "object") {
return undefined;
return line?.response?.message || undefined;
}
return typeof body.error?.message === "string" ? body.error.message : undefined;
return body.error?.message || line?.response?.message || undefined;
}
/** Return the first useful error message from batch output lines. */
export function extractBatchErrorMessage(lines: BatchOutputErrorLike[]): string | undefined {
export function extractBatchErrorMessage(lines: EmbeddingBatchOutputLine[]): string | undefined {
const first = lines.find((line) => line.error?.message || getResponseErrorMessage(line));
return first?.error?.message ?? getResponseErrorMessage(first);
return first?.error?.message || getResponseErrorMessage(first);
}
/** Redact and bound provider-controlled batch diagnostics before logging them. */
export function formatBatchErrorDetail(detail: string | undefined): string | undefined {
if (!detail) {
return undefined;
}
const formatted = formatErrorMessage(detail);
if (formatted.length <= BATCH_ERROR_DETAIL_MAX_CHARS) {
return formatted;
}
const prefixLength = BATCH_ERROR_DETAIL_MAX_CHARS - BATCH_ERROR_DETAIL_TRUNCATED_SUFFIX.length;
return `${truncateUtf16Safe(formatted, prefixLength)}${BATCH_ERROR_DETAIL_TRUNCATED_SUFFIX}`;
}
/** Format a failed error-file read without hiding the underlying read problem. */
export function formatUnavailableBatchError(err: unknown): string | undefined {
const message = formatErrorMessage(err);
const message = formatBatchErrorDetail(formatErrorMessage(err));
return message ? `error file unavailable: ${message}` : undefined;
}

143
packages/memory-host-sdk/src/host/batch-output.test.ts
View File

@@ -1,6 +1,129 @@
// Memory Host SDK tests cover batch output behavior.
import { describe, expect, it } from "vitest";
import { applyEmbeddingBatchOutputLine } from "./batch-output.js";
import { applyEmbeddingBatchOutputLine, readEmbeddingBatchJsonl } from "./batch-output.js";
function streamingResponse(chunks: Uint8Array[]): {
response: Response;
wasCanceled: () => boolean;
} {
let index = 0;
let canceled = false;
return {
response: new Response(
new ReadableStream<Uint8Array>({
pull(controller) {
const chunk = chunks[index];
index += 1;
if (chunk) {
controller.enqueue(chunk);
return;
}
controller.close();
},
cancel() {
canceled = true;
},
}),
),
wasCanceled: () => canceled,
};
}
describe("readEmbeddingBatchJsonl", () => {
it("frames split UTF-8, CRLF, and an unterminated final record", async () => {
const bytes = new TextEncoder().encode('{"value":"😀"}\r\n{"value":2}');
const streamed = streamingResponse(Array.from(bytes, (byte) => Uint8Array.of(byte)));
const records: Array<{ value: string | number }> = [];
await readEmbeddingBatchJsonl<{ value: string | number }>(streamed.response, {
label: "test.batch-output",
maxRecords: 2,
onRecord: (record) => {
records.push(record);
return true;
},
});
expect(records).toEqual([{ value: "😀" }, { value: 2 }]);
expect(streamed.wasCanceled()).toBe(false);
});
it("bounds a newline-free record and cancels the stream", async () => {
const encoder = new TextEncoder();
const streamed = streamingResponse([
encoder.encode('{"value":"'),
encoder.encode("x".repeat(32)),
encoder.encode('"}\n'),
]);
await expect(
readEmbeddingBatchJsonl(streamed.response, {
label: "test.batch-output",
maxRecords: 1,
maxRecordBytes: 16,
onRecord: () => true,
}),
).rejects.toThrow("test.batch-output: JSONL record exceeds 16 bytes");
expect(streamed.wasCanceled()).toBe(true);
});
it("counts blank physical records against the output budget", async () => {
const streamed = streamingResponse([new TextEncoder().encode("{}\n\n{}\n")]);
await expect(
readEmbeddingBatchJsonl(streamed.response, {
label: "test.batch-output",
maxRecords: 2,
onRecord: () => true,
}),
).rejects.toThrow("test.batch-output: JSONL output exceeds 2 records");
expect(streamed.wasCanceled()).toBe(true);
});
it("cancels immediately when the consumer has enough records", async () => {
const encoder = new TextEncoder();
const streamed = streamingResponse([encoder.encode('{"value":1}\n'), encoder.encode("{")]);
const records: unknown[] = [];
await readEmbeddingBatchJsonl(streamed.response, {
label: "test.batch-output",
maxRecords: 2,
onRecord: (record) => {
records.push(record);
return false;
},
});
expect(records).toEqual([{ value: 1 }]);
expect(streamed.wasCanceled()).toBe(true);
});
it.each(["{bad}\n", "null\n", "[]\n"])(
"cancels after an invalid JSONL record %#",
async (input) => {
const malformed = streamingResponse([new TextEncoder().encode(input)]);
await expect(
readEmbeddingBatchJsonl(malformed.response, {
label: "test.batch-output",
maxRecords: 1,
onRecord: () => true,
}),
).rejects.toThrow("test.batch-output: malformed JSONL record");
expect(malformed.wasCanceled()).toBe(true);
},
);
it("accepts a null response body", async () => {
await expect(
readEmbeddingBatchJsonl(new Response(null), {
label: "test.batch-output",
maxRecords: 0,
onRecord: () => true,
}),
).resolves.toBeUndefined();
});
});
describe("applyEmbeddingBatchOutputLine", () => {
it("stores embedding for successful response", () => {
@@ -47,7 +170,7 @@ describe("applyEmbeddingBatchOutputLine", () => {
});
it("records non-2xx response errors and empty embedding errors", () => {
const remaining = new Set(["req-3", "req-4"]);
const remaining = new Set(["req-3", "req-4", "req-5"]);
const errors: string[] = [];
const byCustomId = new Map<string, number[]>();
@@ -77,7 +200,21 @@ describe("applyEmbeddingBatchOutputLine", () => {
byCustomId,
});
expect(errors).toEqual(["req-3: internal", "req-4: empty embedding"]);
applyEmbeddingBatchOutputLine({
line: {
custom_id: "req-5",
response: { status_code: 500, message: "provider response failed", body: "" },
},
remaining,
errors,
byCustomId,
});
expect(errors).toEqual([
"req-3: internal",
"req-4: empty embedding",
"req-5: provider response failed",
]);
expect(byCustomId.size).toBe(0);
});
});

121
packages/memory-host-sdk/src/host/batch-output.ts
View File

@@ -1,11 +1,126 @@
// Parses provider batch output lines into the custom-id embedding map.
const DEFAULT_BATCH_OUTPUT_RECORD_MAX_BYTES = 4 * 1024 * 1024;
const INITIAL_BATCH_OUTPUT_RECORD_BYTES = 64 * 1024;
type ReadEmbeddingBatchJsonlOptions<T> = {
label: string;
maxRecords: number;
maxRecordBytes?: number;
onRecord: (record: T) => boolean;
};
/** Stream bounded JSONL records without buffering the provider output file. */
export async function readEmbeddingBatchJsonl<T>(
response: Response,
options: ReadEmbeddingBatchJsonlOptions<T>,
): Promise<void> {
const reader = response.body?.getReader();
if (!reader) {
return;
}
const maxRecordBytes = options.maxRecordBytes ?? DEFAULT_BATCH_OUTPUT_RECORD_MAX_BYTES;
const decoder = new TextDecoder();
let recordCount = 0;
let recordBytes = 0;
let recordBuffer: Uint8Array | undefined;
const appendRecordPart = (part: Uint8Array) => {
if (part.byteLength === 0) {
return;
}
const nextRecordBytes = recordBytes + part.byteLength;
if (nextRecordBytes > maxRecordBytes) {
throw new Error(`${options.label}: JSONL record exceeds ${maxRecordBytes} bytes`);
}
if (!recordBuffer || recordBuffer.byteLength < nextRecordBytes) {
const nextCapacity = Math.min(
maxRecordBytes,
Math.max(
nextRecordBytes,
recordBuffer
? recordBuffer.byteLength * 2
: Math.min(INITIAL_BATCH_OUTPUT_RECORD_BYTES, maxRecordBytes),
),
);
const nextBuffer = new Uint8Array(nextCapacity);
if (recordBuffer) {
nextBuffer.set(recordBuffer.subarray(0, recordBytes));
}
recordBuffer = nextBuffer;
}
recordBuffer.set(part, recordBytes);
recordBytes = nextRecordBytes;
};
const emitRecord = (): boolean => {
// Count physical rows before trimming so blank rows cannot bypass the
// provider's one-output-row-per-input budget.
recordCount += 1;
if (recordCount > options.maxRecords) {
throw new Error(`${options.label}: JSONL output exceeds ${options.maxRecords} records`);
}
const text = decoder.decode(recordBuffer?.subarray(0, recordBytes)).trim();
recordBytes = 0;
if (!text) {
return true;
}
let parsed: unknown;
try {
parsed = JSON.parse(text) as unknown;
} catch {
throw new Error(`${options.label}: malformed JSONL record`);
}
if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
throw new Error(`${options.label}: malformed JSONL record`);
}
return options.onRecord(parsed as T);
};
const cancel = async () => {
await reader.cancel().catch(() => {});
};
try {
while (true) {
const { done, value } = await reader.read();
if (done) {
break;
}
// Frame raw bytes so split UTF-8 sequences cannot evade the byte cap.
let offset = 0;
for (let index = 0; index < value.byteLength; index += 1) {
if (value[index] !== 0x0a) {
continue;
}
appendRecordPart(value.subarray(offset, index));
if (!emitRecord()) {
await cancel();
return;
}
offset = index + 1;
}
appendRecordPart(value.subarray(offset));
}
if (recordBytes > 0) {
emitRecord();
}
} catch (error) {
await cancel();
throw error;
} finally {
reader.releaseLock();
}
}
/** Minimal OpenAI-compatible embedding batch output line. */
export type EmbeddingBatchOutputLine = {
custom_id?: string;
error?: { message?: string };
error?: { message?: string } | null;
response?: {
status_code?: number;
message?: string;
body?:
| {
data?: Array<{
@@ -44,7 +159,9 @@ export function applyEmbeddingBatchOutputLine(params: {
? (response.body as { error?: { message?: string } }).error?.message
: undefined;
const messageFromString = typeof response?.body === "string" ? response.body : undefined;
params.errors.push(`${customId}: ${messageFromObject ?? messageFromString ?? "unknown error"}`);
params.errors.push(
`${customId}: ${messageFromObject || messageFromString || response?.message || "unknown error"}`,
);
return;
}

16
packages/memory-host-sdk/src/host/batch-status.test.ts
View File

@@ -3,6 +3,7 @@ import { describe, expect, it } from "vitest";
import {
resolveBatchCompletionFromStatus,
resolveCompletedBatchResult,
throwIfBatchCompletionError,
throwIfBatchTerminalFailure,
} from "./batch-status.js";
@@ -33,6 +34,21 @@ describe("batch-status helpers", () => {
).rejects.toThrow("voyage batch b2 failed: bad input");
});
it("reads a completed error file before requiring successful output", async () => {
await expect(
throwIfBatchCompletionError({
provider: "openai",
status: {
id: "b3",
status: "completed",
output_file_id: null,
error_file_id: "err-file",
},
readError: async () => "all requests failed",
}),
).rejects.toThrow("openai batch b3 completed: all requests failed");
});
it("returns completed result directly without waiting", async () => {
const waitForBatch = async () => ({ outputFileId: "out-2" });
const result = await resolveCompletedBatchResult({

30
packages/memory-host-sdk/src/host/batch-status.ts
View File

@@ -1,15 +1,8 @@
// Batch status helpers shared by remote embedding providers.
import type { EmbeddingBatchStatus } from "./batch-provider-common.js";
const TERMINAL_FAILURE_STATES = new Set(["failed", "expired", "cancelled", "canceled"]);
/** Minimal provider batch status used for completion and terminal-failure checks. */
type BatchStatusLike = {
id?: string;
status?: string;
output_file_id?: string | null;
error_file_id?: string | null;
};
/** File ids returned once a batch has completed. */
export type BatchCompletionResult = {
outputFileId: string;
@@ -20,7 +13,7 @@ export type BatchCompletionResult = {
export function resolveBatchCompletionFromStatus(params: {
provider: string;
batchId: string;
status: BatchStatusLike;
status: EmbeddingBatchStatus;
}): BatchCompletionResult {
if (!params.status.output_file_id) {
throw new Error(`${params.provider} batch ${params.batchId} completed without output file`);
@@ -31,10 +24,25 @@ export function resolveBatchCompletionFromStatus(params: {
};
}
/** Fail a completed partial/all-error batch before requiring its success file. */
export async function throwIfBatchCompletionError(params: {
provider: string;
status: EmbeddingBatchStatus;
readError: (errorFileId: string) => Promise<string | undefined>;
}): Promise<void> {
if (params.status.status !== "completed" || !params.status.error_file_id) {
return;
}
const detail = await params.readError(params.status.error_file_id);
throw new Error(
`${params.provider} batch ${params.status.id ?? "<unknown>"} completed: ${detail ?? "provider error file present"}`,
);
}
/** Throw when a provider reports a terminal failure, including error-file detail if available. */
export async function throwIfBatchTerminalFailure(params: {
provider: string;
status: BatchStatusLike;
status: EmbeddingBatchStatus;
readError: (errorFileId: string) => Promise<string | undefined>;
}): Promise<void> {
const state = params.status.status ?? "unknown";
@@ -51,7 +59,7 @@ export async function throwIfBatchTerminalFailure(params: {
/** Resolve the completed batch files, optionally waiting according to caller policy. */
export async function resolveCompletedBatchResult(params: {
provider: string;
status: BatchStatusLike;
status: EmbeddingBatchStatus;
wait: boolean;
waitForBatch: () => Promise<BatchCompletionResult>;
}): Promise<BatchCompletionResult> {

4
scripts/plugin-sdk-surface-report.mjs
View File

@@ -195,12 +195,12 @@ export function readPluginSdkSurfaceBudgets(env = process.env) {
),
publicExports: readPluginSdkSurfaceBudgetEnv(
"OPENCLAW_PLUGIN_SDK_MAX_PUBLIC_EXPORTS",
10557,
10562,
env,
),
publicFunctionExports: readPluginSdkSurfaceBudgetEnv(
"OPENCLAW_PLUGIN_SDK_MAX_PUBLIC_FUNCTION_EXPORTS",
5251,
5255,
env,
),
publicDeprecatedExports: readPluginSdkSurfaceBudgetEnv(

14
src/plugin-sdk/memory-core-host-engine-embeddings.ts
View File

@@ -10,20 +10,24 @@ export {
createRemoteEmbeddingProvider,
debugEmbeddingsLog,
DEFAULT_LOCAL_MODEL,
EmbeddingBatchUnavailableError,
EMBEDDING_BATCH_ENDPOINT,
enforceEmbeddingMaxInputTokens,
estimateStructuredEmbeddingInputBytes,
estimateUtf8Bytes,
extractBatchErrorMessage,
fetchRemoteEmbeddingVectors,
formatBatchErrorDetail,
formatUnavailableBatchError,
getMemoryMultimodalExtensions,
hasNonTextEmbeddingParts,
isEmbeddingBatchUnavailableError,
isMissingEmbeddingApiKeyError,
mapBatchEmbeddingsByIndex,
normalizeBatchBaseUrl,
normalizeEmbeddingModelWithPrefixes,
postJsonWithRetry,
readEmbeddingBatchJsonl,
resolveBatchCompletionFromStatus,
resolveCompletedBatchResult,
resolveRemoteEmbeddingBearerClient,
@@ -31,23 +35,17 @@ export {
runEmbeddingBatchGroups,
sanitizeAndNormalizeEmbedding,
sanitizeEmbeddingCacheHeaders,
throwIfBatchCompletionError,
throwIfBatchTerminalFailure,
uploadBatchJsonlFile,
withRemoteHttpResponse,
} from "../../packages/memory-host-sdk/src/engine-embeddings.js";
/** Provider batch status payload shared by memory embedding batch helpers. */
export type EmbeddingBatchStatus = {
id?: string;
status?: string;
output_file_id?: string | null;
error_file_id?: string | null;
};
export type {
BatchCompletionResult,
BatchHttpClientConfig,
EmbeddingBatchExecutionParams,
EmbeddingBatchStatus,
EmbeddingInput,
ProviderBatchOutputLine,
RemoteEmbeddingClient,