diff --git a/src/lib/code-index.test.ts b/src/lib/code-index.test.ts index 5d4c69f..04fe51f 100644 --- a/src/lib/code-index.test.ts +++ b/src/lib/code-index.test.ts @@ -6,6 +6,7 @@ // tree snapshot SHA. import { describe, it, expect, vi, beforeEach } from "vitest"; +import { VECTOR_BACKGROUND_TIMEOUT_MS } from "./vector"; import type { BattleMageConfig } from "./config"; const h = vi.hoisted(() => { @@ -50,8 +51,10 @@ vi.mock("./logger", () => ({ log: (...args: unknown[]) => h.logSpy(...args), })); -vi.mock("./vector", () => ({ - isVectorConfigured: (...args: unknown[]) => h.isVectorConfiguredMock(...args), +vi.mock("./vector", async (importOriginal) => ({ + // Real constants flow through; only the side-effectful fns are mocked. + ...(await importOriginal()), + isVectorConfigured: (...args: unknown[]) => h.isVectorConfiguredMock(...args), srcNamespace: () => "acme_backend:src", docsNamespace: (sha: string) => `acme_backend:docs:${sha}`, kbNamespace: () => "acme_backend:kb", @@ -221,6 +224,8 @@ describe("runCodeIndexTick", () => { expect(vectorUpsertSpy).toHaveBeenCalledWith( "acme_backend:src", [expect.objectContaining({ id: "src/a.ts#0", metadata: expect.objectContaining({ path: "src/a.ts", startLine: 1 }) })], + // Background pipeline uses the generous embed budget (BATTLE-MAGE-5). + { timeoutMs: VECTOR_BACKGROUND_TIMEOUT_MS }, ); expect(kvData.get("srcindex:manifest")).toEqual({ "src/a.ts": { sha: "blob-a", chunks: 1 } }); expect(kvData.get("srcindex:sha")).toBe("sha-2"); diff --git a/src/lib/code-index.ts b/src/lib/code-index.ts index 1564c95..8ab8481 100644 --- a/src/lib/code-index.ts +++ b/src/lib/code-index.ts @@ -37,6 +37,7 @@ import { log } from "./logger"; import { isVectorConfigured, srcNamespace, + VECTOR_BACKGROUND_TIMEOUT_MS, vectorUpsert, vectorDelete, } from "./vector"; @@ -324,6 +325,10 @@ export async function runCodeIndexTick( const ok = await vectorUpsert( namespace, chunks.map((c) => ({ id: c.id, text: c.text, metadata: c.metadata })), + // Background pipeline: server-side embedding of a file's chunks + // legitimately takes seconds (BATTLE-MAGE-5); the tick's + // wall-clock budget absorbs the slack. + { timeoutMs: VECTOR_BACKGROUND_TIMEOUT_MS }, ); if (!ok) { degraded = true; diff --git a/src/lib/repo-index.test.ts b/src/lib/repo-index.test.ts index eb19790..966faf3 100644 --- a/src/lib/repo-index.test.ts +++ b/src/lib/repo-index.test.ts @@ -1,4 +1,5 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; +import { VECTOR_BACKGROUND_TIMEOUT_MS } from "./vector"; // ── #127 mocks: kv / github / vector / logger for the rebuild-hook tests ── // The pure-function tests below never touch these; the mocks only feed @@ -38,8 +39,10 @@ vi.mock("./kv", () => ({ }, })); -vi.mock("./vector", () => ({ - isVectorConfigured: () => true, +vi.mock("./vector", async (importOriginal) => ({ + // Real constants flow through; only the side-effectful fns are mocked. + ...(await importOriginal()), + isVectorConfigured: () => true, docsNamespace: (sha: string) => `acme_backend:docs:${sha}`, vectorUpsert: (...a: unknown[]) => vectorUpsertSpy(...a), vectorDeleteNamespace: (...a: unknown[]) => vectorDeleteNamespaceSpy(...a), @@ -610,9 +613,12 @@ describe("getOrRebuildIndex — doc embedding hook (#127)", () => { expect(logSpy).toHaveBeenCalledWith("docs_embedded", expect.anything()), ); - expect(vectorUpsertSpy).toHaveBeenCalledWith("acme_backend:docs:new-sha", [ - expect.objectContaining({ id: "docs/setup.md#0" }), - ]); + expect(vectorUpsertSpy).toHaveBeenCalledWith( + "acme_backend:docs:new-sha", + [expect.objectContaining({ id: "docs/setup.md#0" })], + // Background pipeline uses the generous embed budget (BATTLE-MAGE-5). + { timeoutMs: VECTOR_BACKGROUND_TIMEOUT_MS }, + ); expect(kvData.get("index:vector_docs_ns")).toBe("acme_backend:docs:new-sha"); expect(vectorDeleteNamespaceSpy).toHaveBeenCalledWith( "acme_backend:docs:old-sha", diff --git a/src/lib/repo-index.ts b/src/lib/repo-index.ts index 4c77463..7b9d708 100644 --- a/src/lib/repo-index.ts +++ b/src/lib/repo-index.ts @@ -5,6 +5,7 @@ import { docsNamespace, vectorUpsert, vectorDeleteNamespace, + VECTOR_BACKGROUND_TIMEOUT_MS, } from "./vector"; import { type BattleMageConfig, @@ -521,7 +522,9 @@ async function embedDocChunks(sha: string, docPaths: string[]): Promise { } const namespace = docsNamespace(sha); - const ok = await vectorUpsert(namespace, items); + // Background pipeline: one docs corpus can be many chunks; use the + // generous embed budget, not the interactive 2s cap (BATTLE-MAGE-5). + const ok = await vectorUpsert(namespace, items, { timeoutMs: VECTOR_BACKGROUND_TIMEOUT_MS }); if (!ok) { log("docs_embed_failed", { sha, docCount, chunkCount: chunks.length }); return; diff --git a/src/lib/vector.test.ts b/src/lib/vector.test.ts index f08bd94..26b76d0 100644 --- a/src/lib/vector.test.ts +++ b/src/lib/vector.test.ts @@ -22,6 +22,8 @@ import { docsNamespace, srcNamespace, VECTOR_OP_TIMEOUT_MS, + VECTOR_BACKGROUND_TIMEOUT_MS, + UPSERT_BATCH_SIZE, __setVectorStoreFactoryForTests, type VectorStore, } from "./vector"; @@ -277,6 +279,106 @@ describe("degradation: store errors never throw", () => { }); }); +describe("upsert batching + timeout override (BATTLE-MAGE-5)", () => { + beforeEach(() => stubVectorEnv()); + + const item = (i: number) => ({ id: `c${i}`, text: `chunk ${i}` }); + + it("pins the background timeout and batch size constants", () => { + expect(VECTOR_BACKGROUND_TIMEOUT_MS).toBe(30_000); + expect(UPSERT_BATCH_SIZE).toBe(20); + expect(VECTOR_BACKGROUND_TIMEOUT_MS).toBeGreaterThan(VECTOR_OP_TIMEOUT_MS); + }); + + it("splits a large item list into sequential store calls of at most UPSERT_BATCH_SIZE", async () => { + const store = makeFakeStore(); + __setVectorStoreFactoryForTests(() => store); + const items = Array.from({ length: UPSERT_BATCH_SIZE * 2 + 5 }, (_, i) => item(i)); + const ok = await vectorUpsert("ns", items); + expect(ok).toBe(true); + expect(store.upsert).toHaveBeenCalledTimes(3); + expect(vi.mocked(store.upsert).mock.calls[0][1]).toHaveLength(UPSERT_BATCH_SIZE); + expect(vi.mocked(store.upsert).mock.calls[1][1]).toHaveLength(UPSERT_BATCH_SIZE); + expect(vi.mocked(store.upsert).mock.calls[2][1]).toHaveLength(5); + // One success log carrying the TOTAL count. + expect(logSpy).toHaveBeenCalledWith( + "vector_op", + expect.objectContaining({ op: "upsert", count: UPSERT_BATCH_SIZE * 2 + 5 }), + ); + }); + + it("stops on the first failed batch and returns false (later batches never sent)", async () => { + const store = makeFakeStore({ + upsert: vi + .fn() + .mockResolvedValueOnce(undefined) + .mockRejectedValueOnce(new Error("boom")), + }); + __setVectorStoreFactoryForTests(() => store); + const items = Array.from({ length: UPSERT_BATCH_SIZE * 3 }, (_, i) => item(i)); + const ok = await vectorUpsert("ns", items); + expect(ok).toBe(false); + expect(store.upsert).toHaveBeenCalledTimes(2); + }); + + it("a list at or under the batch size still makes exactly one store call", async () => { + const store = makeFakeStore(); + __setVectorStoreFactoryForTests(() => store); + await vectorUpsert("ns", Array.from({ length: UPSERT_BATCH_SIZE }, (_, i) => item(i))); + expect(store.upsert).toHaveBeenCalledTimes(1); + }); + + it("honors a per-call timeout override: survives past the default, fails past the override", async () => { + vi.useFakeTimers(); + __setVectorStoreFactoryForTests(() => + makeFakeStore({ upsert: () => new Promise(() => {}) }), // hangs forever + ); + const pending = vectorUpsert("ns", [item(1)], { timeoutMs: VECTOR_BACKGROUND_TIMEOUT_MS }); + await vi.advanceTimersByTimeAsync(VECTOR_OP_TIMEOUT_MS + 1); + // Default budget elapsed — the overridden call must still be alive. + let settled = false; + pending.then(() => { settled = true; }); + await Promise.resolve(); + expect(settled).toBe(false); + await vi.advanceTimersByTimeAsync(VECTOR_BACKGROUND_TIMEOUT_MS); + await expect(pending).resolves.toBe(false); + const err = logSpy.mock.calls.find((c) => c[0] === "vector_error"); + expect(err?.[1]).toMatchObject({ op: "upsert", errorClass: "VectorTimeoutError" }); + expect(String(err?.[1].errorMessage)).toContain("30000"); + }); + + it("timeoutMs is an end-to-end deadline for the whole call, not per batch", async () => { + vi.useFakeTimers(); + // Each batch takes 12s. With a 30s per-CALL budget, batch 1 (12s) and + // batch 2 (24s) fit; batch 3 would finish at 36s — the call must fail + // at the 30s deadline instead of granting each batch a fresh 30s. + __setVectorStoreFactoryForTests(() => + makeFakeStore({ + upsert: () => new Promise((resolve) => setTimeout(() => resolve(undefined), 12_000)), + }), + ); + const items = Array.from({ length: UPSERT_BATCH_SIZE * 3 }, (_, i) => ({ + id: `c${i}`, + text: `chunk ${i}`, + })); + const pending = vectorUpsert("ns", items, { timeoutMs: 30_000 }); + await vi.advanceTimersByTimeAsync(40_000); + await expect(pending).resolves.toBe(false); + const err = logSpy.mock.calls.find((c) => c[0] === "vector_error"); + expect(err?.[1]).toMatchObject({ op: "upsert", errorClass: "VectorTimeoutError" }); + }); + + it("default timeout still applies when no override is passed", async () => { + vi.useFakeTimers(); + __setVectorStoreFactoryForTests(() => + makeFakeStore({ upsert: () => new Promise(() => {}) }), + ); + const pending = vectorUpsert("ns", [item(1)]); + await vi.advanceTimersByTimeAsync(VECTOR_OP_TIMEOUT_MS + 1); + await expect(pending).resolves.toBe(false); + }); +}); + describe("degradation: timeout", () => { beforeEach(() => stubVectorEnv()); diff --git a/src/lib/vector.ts b/src/lib/vector.ts index 609f3da..856b026 100644 --- a/src/lib/vector.ts +++ b/src/lib/vector.ts @@ -172,17 +172,38 @@ function getStore(): VectorStore { /** Hard latency cap per vector op — recall must never stall a turn. */ export const VECTOR_OP_TIMEOUT_MS = 2000; +/** + * Timeout for background embed pipelines (code index, doc embedding), + * where a single upsert batch is server-side-embedded and legitimately + * needs seconds — the 2s interactive cap starved them and stalled the + * index on large files (BATTLE-MAGE-5). Cron budgets absorb the slack. + */ +export const VECTOR_BACKGROUND_TIMEOUT_MS = 30_000; + +/** + * Max items per underlying store call. One upsert request carrying a + * whole large file (or a whole docs corpus) embeds slowly and fails as + * a unit; smaller batches keep per-request latency inside the timeout + * and make retries finer-grained. Ids are deterministic, so re-sending + * a batch after a mid-list failure is harmless. + */ +export const UPSERT_BATCH_SIZE = 20; + export class VectorTimeoutError extends Error { - constructor(op: string) { - super(`vector ${op} exceeded ${VECTOR_OP_TIMEOUT_MS}ms`); + constructor(op: string, timeoutMs: number) { + super(`vector ${op} exceeded ${timeoutMs}ms`); this.name = "VectorTimeoutError"; } } -function withTimeout(op: string, promise: Promise): Promise { +function withTimeout( + op: string, + promise: Promise, + timeoutMs: number = VECTOR_OP_TIMEOUT_MS, +): Promise { let timer: ReturnType | undefined; const timeout = new Promise((_, reject) => { - timer = setTimeout(() => reject(new VectorTimeoutError(op)), VECTOR_OP_TIMEOUT_MS); + timer = setTimeout(() => reject(new VectorTimeoutError(op, timeoutMs)), timeoutMs); }); return Promise.race([promise, timeout]).finally(() => clearTimeout(timer)) as Promise; } @@ -221,20 +242,34 @@ function logError(op: VectorOp, namespace: string, startedAt: number, err: unkno export async function vectorUpsert( namespace: string, items: VectorUpsertItem[], + opts?: { timeoutMs?: number }, ): Promise { if (items.length === 0) return true; if (!isVectorConfigured()) { logUnavailable("upsert"); return false; } + const timeoutMs = opts?.timeoutMs ?? VECTOR_OP_TIMEOUT_MS; const startedAt = Date.now(); try { - await withTimeout("upsert", getStore().upsert(namespace, items)); + // Sequential batches; stop on the first failure. Completed batches + // stay written — deterministic ids make the caller's retry re-send + // them idempotently. timeoutMs is an END-TO-END deadline for the + // whole call: each batch races the remaining budget, so a + // multi-batch upsert can never exceed the caller's per-call cap. + const deadline = startedAt + timeoutMs; + for (let i = 0; i < items.length; i += UPSERT_BATCH_SIZE) { + const remaining = deadline - Date.now(); + if (remaining <= 0) throw new VectorTimeoutError("upsert", timeoutMs); + const batch = items.slice(i, i + UPSERT_BATCH_SIZE); + await withTimeout("upsert", getStore().upsert(namespace, batch), remaining); + } log("vector_op", { op: "upsert", namespace, durationMs: Date.now() - startedAt, count: items.length, + batches: Math.ceil(items.length / UPSERT_BATCH_SIZE), }); return true; } catch (err) {