fix(memory): optimize voyage batch memory usage with streaming and deduplicate code

Cherry-picked from PR #2519. Fixed lint error: changed this.runWithConcurrency
to use imported runWithConcurrency function after extraction to internal.ts

src/memory/batch-voyage.test.ts

@@ -0,0 +1,170 @@
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { ReadableStream } from "node:stream/web";
+import type { VoyageBatchOutputLine, VoyageBatchRequest } from "./batch-voyage.js";
+import type { VoyageEmbeddingClient } from "./embeddings-voyage.js";
+
+// Mock internal.js if needed, but runWithConcurrency is simple enough to keep real.
+// We DO need to mock retryAsync to avoid actual delays/retries logic complicating tests
+vi.mock("../infra/retry.js", () => ({
+  retryAsync: async <T>(fn: () => Promise<T>) => fn(),
+}));
+
+describe("runVoyageEmbeddingBatches", () => {
+  afterEach(() => {
+    vi.resetAllMocks();
+    vi.unstubAllGlobals();
+  });
+
+  const mockClient: VoyageEmbeddingClient = {
+    baseUrl: "https://api.voyageai.com/v1",
+    headers: { Authorization: "Bearer test-key" },
+    model: "voyage-4-large",
+  };
+
+  const mockRequests: VoyageBatchRequest[] = [
+    { custom_id: "req-1", body: { input: "text1" } },
+    { custom_id: "req-2", body: { input: "text2" } },
+  ];
+
+  it("successfully submits batch, waits, and streams results", async () => {
+    const fetchMock = vi.fn();
+    vi.stubGlobal("fetch", fetchMock);
+
+    // Sequence of fetch calls:
+    // 1. Upload file
+    fetchMock.mockResolvedValueOnce({
+      ok: true,
+      json: async () => ({ id: "file-123" }),
+    });
+
+    // 2. Create batch
+    fetchMock.mockResolvedValueOnce({
+      ok: true,
+      json: async () => ({ id: "batch-abc", status: "pending" }),
+    });
+
+    // 3. Poll status (pending) - Optional depending on wait loop, let's say it finishes immediately for this test
+    // Actually the code does: initial check (if completed) -> wait loop.
+    // If create returns "pending", it enters waitForVoyageBatch.
+    // waitForVoyageBatch fetches status.
+
+    // 3. Poll status (completed)
+    fetchMock.mockResolvedValueOnce({
+      ok: true,
+      json: async () => ({
+        id: "batch-abc",
+        status: "completed",
+        output_file_id: "file-out-999",
+      }),
+    });
+
+    // 4. Download content (Streaming)
+    const outputLines: VoyageBatchOutputLine[] = [
+      {
+        custom_id: "req-1",
+        response: { status_code: 200, body: { data: [{ embedding: [0.1, 0.1] }] } },
+      },
+      {
+        custom_id: "req-2",
+        response: { status_code: 200, body: { data: [{ embedding: [0.2, 0.2] }] } },
+      },
+    ];
+
+    // Create a stream that emits the NDJSON lines
+    const stream = new ReadableStream({
+      start(controller) {
+        const text = outputLines.map((l) => JSON.stringify(l)).join("\n");
+        controller.enqueue(new TextEncoder().encode(text));
+        controller.close();
+      },
+    });
+
+    fetchMock.mockResolvedValueOnce({
+      ok: true,
+      body: stream,
+    });
+
+    const { runVoyageEmbeddingBatches } = await import("./batch-voyage.js");
+
+    const results = await runVoyageEmbeddingBatches({
+      client: mockClient,
+      agentId: "agent-1",
+      requests: mockRequests,
+      wait: true,
+      pollIntervalMs: 1, // fast poll
+      timeoutMs: 1000,
+      concurrency: 1,
+    });
+
+    expect(results.size).toBe(2);
+    expect(results.get("req-1")).toEqual([0.1, 0.1]);
+    expect(results.get("req-2")).toEqual([0.2, 0.2]);
+
+    // Verify calls
+    expect(fetchMock).toHaveBeenCalledTimes(4);
+
+    // Verify File Upload
+    expect(fetchMock.mock.calls[0][0]).toContain("/files");
+    const uploadBody = fetchMock.mock.calls[0][1].body as FormData;
+    expect(uploadBody).toBeInstanceOf(FormData);
+    expect(uploadBody.get("purpose")).toBe("batch");
+
+    // Verify Batch Create
+    expect(fetchMock.mock.calls[1][0]).toContain("/batches");
+    const createBody = JSON.parse(fetchMock.mock.calls[1][1].body);
+    expect(createBody.input_file_id).toBe("file-123");
+    expect(createBody.completion_window).toBe("12h");
+
+    // Verify Content Fetch
+    expect(fetchMock.mock.calls[3][0]).toContain("/files/file-out-999/content");
+  });
+
+  it("handles empty lines and stream chunks correctly", async () => {
+    const fetchMock = vi.fn();
+    vi.stubGlobal("fetch", fetchMock);
+
+    // 1. Upload
+    fetchMock.mockResolvedValueOnce({ ok: true, json: async () => ({ id: "f1" }) });
+    // 2. Create (completed immediately)
+    fetchMock.mockResolvedValueOnce({
+      ok: true,
+      json: async () => ({ id: "b1", status: "completed", output_file_id: "out1" }),
+    });
+    // 3. Download Content (Streaming with chunks and newlines)
+    const stream = new ReadableStream({
+      start(controller) {
+        const line1 = JSON.stringify({
+          custom_id: "req-1",
+          response: { body: { data: [{ embedding: [1] }] } },
+        });
+        const line2 = JSON.stringify({
+          custom_id: "req-2",
+          response: { body: { data: [{ embedding: [2] }] } },
+        });
+
+        // Split across chunks
+        controller.enqueue(new TextEncoder().encode(line1 + "\n"));
+        controller.enqueue(new TextEncoder().encode("\n")); // empty line
+        controller.enqueue(new TextEncoder().encode(line2)); // no newline at EOF
+        controller.close();
+      },
+    });
+
+    fetchMock.mockResolvedValueOnce({ ok: true, body: stream });
+
+    const { runVoyageEmbeddingBatches } = await import("./batch-voyage.js");
+
+    const results = await runVoyageEmbeddingBatches({
+      client: mockClient,
+      agentId: "a1",
+      requests: mockRequests,
+      wait: true,
+      pollIntervalMs: 1,
+      timeoutMs: 1000,
+      concurrency: 1,
+    });
+
+    expect(results.get("req-1")).toEqual([1]);
+    expect(results.get("req-2")).toEqual([2]);
+  });
+});

src/memory/batch-voyage.ts

@@ -1,6 +1,9 @@
+import { createInterface } from "node:readline";
+import { Readable } from "node:stream";
+
 import { retryAsync } from "../infra/retry.js";
 import type { VoyageEmbeddingClient } from "./embeddings-voyage.js";
-import { hashText } from "./internal.js";
+import { hashText, runWithConcurrency } from "./internal.js";
 
 /**
  * Voyage Batch API Input Line format.
@@ -153,40 +156,26 @@ async function fetchVoyageBatchStatus(params: {
   return (await res.json()) as VoyageBatchStatus;
 }
 
-async function fetchVoyageFileContent(params: {
-  client: VoyageEmbeddingClient;
-  fileId: string;
-}): Promise<string> {
-  const baseUrl = getVoyageBaseUrl(params.client);
-  const res = await fetch(`${baseUrl}/files/${params.fileId}/content`, {
-    headers: getVoyageHeaders(params.client, { json: true }),
-  });
-  if (!res.ok) {
-    const text = await res.text();
-    throw new Error(`voyage batch file content failed: ${res.status} ${text}`);
-  }
-  return await res.text();
-}
-
-function parseVoyageBatchOutput(text: string): VoyageBatchOutputLine[] {
-  if (!text.trim()) return [];
-  return text
-    .split("\n")
-    .map((line) => line.trim())
-    .filter(Boolean)
-    .map((line) => JSON.parse(line) as VoyageBatchOutputLine);
-}
-
 async function readVoyageBatchError(params: {
   client: VoyageEmbeddingClient;
   errorFileId: string;
 }): Promise<string | undefined> {
   try {
-    const content = await fetchVoyageFileContent({
+    const baseUrl = getVoyageBaseUrl(params.client);
-      client: params.client,
+    const res = await fetch(`${baseUrl}/files/${params.errorFileId}/content`, {
-      fileId: params.errorFileId,
+      headers: getVoyageHeaders(params.client, { json: true }),
     });
-    const lines = parseVoyageBatchOutput(content);
+    if (!res.ok) {
+      const text = await res.text();
+      throw new Error(`voyage batch error file content failed: ${res.status} ${text}`);
+    }
+    const text = await res.text();
+    if (!text.trim()) return undefined;
+    const lines = text
+      .split("\n")
+      .map((line) => line.trim())
+      .filter(Boolean)
+      .map((line) => JSON.parse(line) as VoyageBatchOutputLine);
     const first = lines.find((line) => line.error?.message || line.response?.body?.error);
     const message =
       first?.error?.message ??
@@ -247,33 +236,6 @@ async function waitForVoyageBatch(params: {
   }
 }
 
-async function runWithConcurrency<T>(tasks: Array<() => Promise<T>>, limit: number): Promise<T[]> {
-  if (tasks.length === 0) return [];
-  const resolvedLimit = Math.max(1, Math.min(limit, tasks.length));
-  const results: T[] = Array.from({ length: tasks.length });
-  let next = 0;
-  let firstError: unknown = null;
-
-  const workers = Array.from({ length: resolvedLimit }, async () => {
-    while (true) {
-      if (firstError) return;
-      const index = next;
-      next += 1;
-      if (index >= tasks.length) return;
-      try {
-        results[index] = await tasks[index]();
-      } catch (err) {
-        firstError = err;
-        return;
-      }
-    }
-  });
-
-  await Promise.allSettled(workers);
-  if (firstError) throw firstError;
-  return results;
-}
-
 export async function runVoyageEmbeddingBatches(params: {
   client: VoyageEmbeddingClient;
   agentId: string;
@@ -331,39 +293,52 @@ export async function runVoyageEmbeddingBatches(params: {
       throw new Error(`voyage batch ${batchInfo.id} completed without output file`);
     }
 
-    const content = await fetchVoyageFileContent({
+    const baseUrl = getVoyageBaseUrl(params.client);
-      client: params.client,
+    const contentRes = await fetch(`${baseUrl}/files/${completed.outputFileId}/content`, {
-      fileId: completed.outputFileId,
+      headers: getVoyageHeaders(params.client, { json: true }),
     });
-    const outputLines = parseVoyageBatchOutput(content);
+    if (!contentRes.ok) {
+      const text = await contentRes.text();
+      throw new Error(`voyage batch file content failed: ${contentRes.status} ${text}`);
+    }
+
     const errors: string[] = [];
     const remaining = new Set(group.map((request) => request.custom_id));
 
-    for (const line of outputLines) {
+    if (contentRes.body) {
-      const customId = line.custom_id;
+      const reader = createInterface({
-      if (!customId) continue;
+        input: Readable.fromWeb(contentRes.body as any),
-      remaining.delete(customId);
+        terminal: false,
-      if (line.error?.message) {
+      });
-        errors.push(`${customId}: ${line.error.message}`);
+
-        continue;
+      for await (const rawLine of reader) {
+        if (!rawLine.trim()) continue;
+        const line = JSON.parse(rawLine) as VoyageBatchOutputLine;
+        const customId = line.custom_id;
+        if (!customId) continue;
+        remaining.delete(customId);
+        if (line.error?.message) {
+          errors.push(`${customId}: ${line.error.message}`);
+          continue;
+        }
+        const response = line.response;
+        const statusCode = response?.status_code ?? 0;
+        if (statusCode >= 400) {
+          const message =
+            response?.body?.error?.message ??
+            (typeof response?.body === "string" ? response.body : undefined) ??
+            "unknown error";
+          errors.push(`${customId}: ${message}`);
+          continue;
+        }
+        const data = response?.body?.data ?? [];
+        const embedding = data[0]?.embedding ?? [];
+        if (embedding.length === 0) {
+          errors.push(`${customId}: empty embedding`);
+          continue;
+        }
+        byCustomId.set(customId, embedding);
       }
-      const response = line.response;
-      const statusCode = response?.status_code ?? 0;
-      if (statusCode >= 400) {
-        const message =
-          response?.body?.error?.message ??
-          (typeof response?.body === "string" ? response.body : undefined) ??
-          "unknown error";
-        errors.push(`${customId}: ${message}`);
-        continue;
-      }
-      const data = response?.body?.data ?? [];
-      const embedding = data[0]?.embedding ?? [];
-      if (embedding.length === 0) {
-        errors.push(`${customId}: empty embedding`);
-        continue;
-      }
-      byCustomId.set(customId, embedding);
     }
 
     if (errors.length > 0) {

src/memory/embeddings-voyage.test.ts

@@ -0,0 +1,100 @@
+import { afterEach, describe, expect, it, vi } from "vitest";
+
+vi.mock("../agents/model-auth.js", () => ({
+  resolveApiKeyForProvider: vi.fn(),
+  requireApiKey: (auth: { apiKey?: string; mode?: string }, provider: string) => {
+    if (auth?.apiKey) return auth.apiKey;
+    throw new Error(`No API key resolved for provider "${provider}" (auth mode: ${auth?.mode}).`);
+  },
+}));
+
+const createFetchMock = () =>
+  vi.fn(async () => ({
+    ok: true,
+    status: 200,
+    json: async () => ({ data: [{ embedding: [0.1, 0.2, 0.3] }] }),
+  })) as unknown as typeof fetch;
+
+describe("voyage embedding provider", () => {
+  afterEach(() => {
+    vi.resetAllMocks();
+    vi.resetModules();
+    vi.unstubAllGlobals();
+  });
+
+  it("configures client with correct defaults and headers", async () => {
+    const fetchMock = createFetchMock();
+    vi.stubGlobal("fetch", fetchMock);
+
+    const { createVoyageEmbeddingProvider } = await import("./embeddings-voyage.js");
+    const authModule = await import("../agents/model-auth.js");
+
+    vi.mocked(authModule.resolveApiKeyForProvider).mockResolvedValue({
+      apiKey: "voyage-key-123",
+      mode: "api-key",
+      source: "test",
+    });
+
+    const result = await createVoyageEmbeddingProvider({
+      config: {} as never,
+      provider: "voyage",
+      model: "voyage-4-large",
+      fallback: "none",
+    });
+
+    await result.provider.embedQuery("test query");
+
+    expect(authModule.resolveApiKeyForProvider).toHaveBeenCalledWith(
+      expect.objectContaining({ provider: "voyage" }),
+    );
+
+    const [url, init] = fetchMock.mock.calls[0] ?? [];
+    expect(url).toBe("https://api.voyageai.com/v1/embeddings");
+
+    const headers = (init?.headers ?? {}) as Record<string, string>;
+    expect(headers.Authorization).toBe("Bearer voyage-key-123");
+    expect(headers["Content-Type"]).toBe("application/json");
+
+    const body = JSON.parse(init?.body as string);
+    expect(body).toEqual({
+      model: "voyage-4-large",
+      input: ["test query"],
+    });
+  });
+
+  it("respects remote overrides for baseUrl and apiKey", async () => {
+    const fetchMock = createFetchMock();
+    vi.stubGlobal("fetch", fetchMock);
+
+    const { createVoyageEmbeddingProvider } = await import("./embeddings-voyage.js");
+
+    const result = await createVoyageEmbeddingProvider({
+      config: {} as never,
+      provider: "voyage",
+      model: "voyage-4-lite",
+      fallback: "none",
+      remote: {
+        baseUrl: "https://proxy.example.com",
+        apiKey: "remote-override-key",
+        headers: { "X-Custom": "123" },
+      },
+    });
+
+    await result.provider.embedQuery("test");
+
+    const [url, init] = fetchMock.mock.calls[0] ?? [];
+    expect(url).toBe("https://proxy.example.com/embeddings");
+
+    const headers = (init?.headers ?? {}) as Record<string, string>;
+    expect(headers.Authorization).toBe("Bearer remote-override-key");
+    expect(headers["X-Custom"]).toBe("123");
+  });
+
+  it("normalizes model names", async () => {
+    const { normalizeVoyageModel } = await import("./embeddings-voyage.js");
+    expect(normalizeVoyageModel("voyage/voyage-large-2")).toBe("voyage-large-2");
+    expect(normalizeVoyageModel("voyage-4-large")).toBe("voyage-4-large");
+    expect(normalizeVoyageModel("  voyage-lite  ")).toBe("voyage-lite");
+    expect(normalizeVoyageModel("")).toBe("voyage-4-large"); // Default
+  });
+});

src/memory/internal.ts

@@ -203,3 +203,33 @@ export function cosineSimilarity(a: number[], b: number[]): number {
   if (normA === 0 || normB === 0) return 0;
   return dot / (Math.sqrt(normA) * Math.sqrt(normB));
 }
+
+export async function runWithConcurrency<T>(
+  tasks: Array<() => Promise<T>>,
+  limit: number,
+): Promise<T[]> {
+  if (tasks.length === 0) return [];
+  const resolvedLimit = Math.max(1, Math.min(limit, tasks.length));
+  const results: T[] = Array.from({ length: tasks.length });
+  let next = 0;
+  let firstError: unknown = null;
+
+  const workers = Array.from({ length: resolvedLimit }, async () => {
+    while (true) {
+      if (firstError) return;
+      const index = next;
+      next += 1;
+      if (index >= tasks.length) return;
+      try {
+        results[index] = await tasks[index]();
+      } catch (err) {
+        firstError = err;
+        return;
+      }
+    }
+  });
+
+  await Promise.allSettled(workers);
+  if (firstError) throw firstError;
+  return results;
+}

src/memory/manager.ts

@@ -42,6 +42,7 @@ import {
   type MemoryFileEntry,
   normalizeRelPath,
   parseEmbedding,
+  runWithConcurrency,
 } from "./internal.js";
 import { bm25RankToScore, buildFtsQuery, mergeHybridResults } from "./hybrid.js";
 import { searchKeyword, searchVector } from "./manager-search.js";
@@ -1023,7 +1024,7 @@ export class MemoryIndexManager {
         });
       }
     });
-    await this.runWithConcurrency(tasks, this.getIndexConcurrency());
+    await runWithConcurrency(tasks, this.getIndexConcurrency());
 
     const staleRows = this.db
       .prepare(`SELECT path FROM files WHERE source = ?`)
@@ -1118,7 +1119,7 @@ export class MemoryIndexManager {
         });
       }
     });
-    await this.runWithConcurrency(tasks, this.getIndexConcurrency());
+    await runWithConcurrency(tasks, this.getIndexConcurrency());
 
     const staleRows = this.db
       .prepare(`SELECT path FROM files WHERE source = ?`)
@@ -2021,33 +2022,6 @@ export class MemoryIndexManager {
     }
   }
 
-  private async runWithConcurrency<T>(tasks: Array<() => Promise<T>>, limit: number): Promise<T[]> {
-    if (tasks.length === 0) return [];
-    const resolvedLimit = Math.max(1, Math.min(limit, tasks.length));
-    const results: T[] = Array.from({ length: tasks.length });
-    let next = 0;
-    let firstError: unknown = null;
-
-    const workers = Array.from({ length: resolvedLimit }, async () => {
-      while (true) {
-        if (firstError) return;
-        const index = next;
-        next += 1;
-        if (index >= tasks.length) return;
-        try {
-          results[index] = await tasks[index]();
-        } catch (err) {
-          firstError = err;
-          return;
-        }
-      }
-    });
-
-    await Promise.allSettled(workers);
-    if (firstError) throw firstError;
-    return results;
-  }
-
   private async withBatchFailureLock<T>(fn: () => Promise<T>): Promise<T> {
     let release: () => void;
     const wait = this.batchFailureLock;