diff --git a/CHANGELOG.md b/CHANGELOG.md index 9cb5ea2f5..cf950b4d8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,76 @@ All notable changes to GBrain will be documented in this file. +## [0.34.4.0] - 2026-05-14 + +**`gbrain embed --stale` now actually works on large brains. `--source` flag stops silently dropping. The retry loop stops storming itself.** + +The v0.33.3 cherry-pick gave `embed --stale` cursor pagination so it would stop dying on Supabase's 2-minute pooler timeout. This release hardens the hot path so it survives real production load: a 30-minute wall-clock budget threads `AbortSignal` into the retry sleep, the worker claim loop, AND the gateway HTTP call (so a worker stuck mid-fetch on a 30-second OpenAI timeout cancels within seconds instead of running past budget). The retry-after delay now jitters ±30% so 20 concurrent workers don't relock on the next 429 wave. The 429 detector unwraps the gateway's `AITransientError` cause chain (the prior bare `e.status === 429` check silently never fired against normalized errors). The wrapper passes `maxRetries: 0` through to the AI SDK so its default 2-retry stack stops multiplying this wrapper's 5 attempts into 15 cycles per call. + +The `gbrain embed --stale --source X` flag was silently broken on multi-source brains — `embedAll(sourceId)` dropped the `sourceId` argument when calling `embedAllStale`, so operators got a full-brain scan even when they asked for one source. Threading `sourceId?` through `countStaleChunks` + `listStaleChunks` + the engine interfaces (Postgres + PGLite) fixes that. + +Migration v66 (the partial index `idx_chunks_embedding_null`) now uses `CREATE INDEX CONCURRENTLY` on Postgres with a handler-based engine branch and pre-drop of any invalid remnant (mirrors the v14 pattern). Plain `CREATE INDEX IF NOT EXISTS` would have taken a ShareLock on the 373K-row `content_chunks` table for the entire build, blocking every concurrent sync/embed/autopilot write. PGLite stays on plain `CREATE INDEX` since it has no concurrent writers. + +Twenty-six new test cases pin every new code path: 8 unit cases for `embedBatchWithBackoff` (parse ms/s retry-after, fallback, non-rate-limit rethrow, jitter range, budget-abort-mid-fetch, cause-chain unwrap, maxRetries:0 passthrough), 3 CLI flag-wiring tests, 1 end-to-end budget test, 6 Postgres E2E tests against a fresh database (static/failed-page/page-split/source-scope/dup-slug), 4 migration v66 source-shape tests, 4 PGLite engine parity tests, plus a fix to every pre-existing stale-row mock that was silently passing TypeScript's structural typing while violating `StaleChunkRow`'s newly-required `source_id` + `page_id` fields. + +### What changes for users + +| Behavior | Before | After v0.34.4.0 | +|---|---|---| +| `embed --stale` budget on millions-of-chunks brain | unbounded — could run hours past Minion timeout | bounded to `GBRAIN_EMBED_TIME_BUDGET_MS` (default 30m); AbortSignal cancels mid-fetch | +| 20 concurrent workers hitting 429 | sleep in lockstep, slam API together on wake | ±30% jitter decorrelates the herd | +| Gateway-wrapped 429 detection | `e.status === 429` silently false against `AITransientError` | unwraps cause chain (depth ≤5); message-match as fallback | +| AI SDK retries per embed call | 3 SDK × 5 wrapper = up to 15 cycles per call | `maxRetries: 0` passthrough; wrapper is single source of truth | +| `embed --stale --source X` | flag silently dropped; full-brain scan | actually scopes to that source | +| Migration v66 on 373K-row table | plain `CREATE INDEX` takes ShareLock for the whole build | `CREATE INDEX CONCURRENTLY` on Postgres; invalid-remnant DROP first | +| Cursor pagination test coverage | 4 mocks ignored opts, never exercised the cursor | 6 real-Postgres E2E + PGLite parity unit tests | + +### Itemized changes + +#### Added +- `GBRAIN_EMBED_TIME_BUDGET_MS` env knob (default 30 min) bounds `embed --stale` wall-clock runtime; partial index makes "next run picks up" automatic. +- `EmbedOpts` in `src/core/ai/gateway.ts` (and `EmbedBatchOptions` in `src/core/embedding.ts`) gain `abortSignal` + `maxRetries` passthrough so callers can cancel mid-fetch and disable the AI SDK's retry stack. +- `detect429FromCause`, `parseRetryDelayMs`, `abortableSleep` exported as `@internal` from `src/commands/embed.ts` so tests can exercise the retry helpers directly. +- `test/e2e/embed-stale-pagination.test.ts` — 6 Postgres E2E tests covering the cursor's static, failure-skip, page-split, source-scoped, and duplicate-slug-across-sources behaviors. + +#### Changed +- Migration v66 (`embed_stale_partial_index`) rewritten to handler-based engine branching; Postgres uses `CREATE INDEX CONCURRENTLY` with a pre-drop of any invalid remnant via `pg_index.indisvalid`. Renumbered v58→v59→v60→v66 across four merge waves as master's v0.33.3, v0.34.0, v0.34.1, and v0.34.2 each claimed slots ahead of this branch. +- `BrainEngine.countStaleChunks(opts?: {sourceId?})` and `listStaleChunks({sourceId?, ...})` now scope to a single source when requested. Both Postgres and PGLite engines updated. +- `embedBatchWithBackoff` now: jitters parsed retry delays ±30%, detects 429 via cause-chain unwrap, passes `maxRetries: 0` through the gateway, and cancels mid-sleep when an `AbortSignal` fires. +- `embedAllStale` accepts `sourceId` and creates a budget-bound `AbortController` threaded into the retry sleep, the per-key worker claim loop, and the gateway embed call. + +#### Fixed +- `gbrain embed --stale --source X` silently dropped the `sourceId` argument and scanned the entire brain. Now correctly scoped end-to-end. +- Existing test mocks omitted `source_id` and `page_id` on stale-row literals; TypeScript's structural typing accepted them but the runtime cursor needs both. Every mock fixed. +- The wave-1 cherry-pick's `embedBatchWithBackoff` had three latent bugs (thundering herd, no wall-clock cap, SDK-retry stacking) and one silently-broken structural check (`e.status === 429` against wrapped errors). All addressed. + +### For contributors + +- 100% AI-assessed coverage of new code paths. Tests: 504 → 505 files (+1 new E2E file) and ~26 new test cases inside expanded files. `bun run test` passes 6385/0/0; full E2E suite passes 90 files / 603 tests against a fresh Postgres. +- Plan reviewed via `/plan-eng-review` (9 decisions made interactively), `/codex` consult (12 findings folded in), and a re-review pass (D8 AbortSignal threading into gateway). Plan file: `~/.claude/plans/system-instruction-you-are-working-iterative-torvalds.md`. +- The cherry-pick lineage: PR #987 (garrytan-agents) → cherry-picked as `ecae761a` onto this branch → opened as PR #991 → all 9 plan decisions + 12 codex findings + 1 re-review finding implemented + tested. + +## To take advantage of v0.34.4.0 + +`gbrain upgrade` does it automatically. The migration v66 runs via `gbrain apply-migrations` and uses `CREATE INDEX CONCURRENTLY` on Postgres so it does not block concurrent writes. + +1. **Run the orchestrator (only if `gbrain doctor` warns about a partial migration):** + ```bash + gbrain apply-migrations --yes + ``` +2. **Verify the outcome:** + ```bash + gbrain doctor --json | grep -o 'Version [0-9]*' # expect Version 66 + ``` +3. **Operator knob** for the new wall-clock budget — only set if 30 min is wrong for your brain: + ```bash + export GBRAIN_EMBED_TIME_BUDGET_MS=1800000 # 30 min default + ``` +4. **If `gbrain embed --stale --source X` was returning unexpected counts pre-upgrade,** the bug was that `--source` was being silently dropped. After upgrade it scopes correctly — re-run with the same flags and you should see the actual per-source NULL count. + +5. **If any step fails or `gbrain embed --stale` still hangs,** please file an issue: + https://github.com/garrytan/gbrain/issues with `gbrain doctor` output and `~/.gbrain/upgrade-errors.jsonl` if it exists. + ## [0.34.3.0] - 2026-05-14 **Supervisor stops crashing every 24h on large brains. Watchdog stops false-firing on git mmap.** @@ -339,6 +409,7 @@ initSchema runs the v60-v65 chain in ~786ms total.** Contributed by @Hansen1018 (#870), @ding-modding (#909), @DukeDawg (#864), @toilalesondev (#861 + #876), @yoelgal (#875). + ## [0.34.0.0] - 2026-05-14 **Recursive code intelligence ships. Plan-mode subagents get one-call blast and flow.** diff --git a/TODOS.md b/TODOS.md index 786ccf5d4..77dfaebab 100644 --- a/TODOS.md +++ b/TODOS.md @@ -1,5 +1,12 @@ # TODOS + +## embed --stale follow-ups (v0.34.4.0) + +- [ ] **v0.35.x: Concurrent NULL→non-NULL upsert race in `embed.ts:429-443` + `postgres-engine.ts:1231`'s `COALESCE(EXCLUDED.embedding, content_chunks.embedding)`.** Two `embed --stale` workers (or `embed --stale` racing with a sync that re-embeds the same chunk) can have the slower writer overwrite the faster one's fresher embedding. Window is small (20 workers, all from the same `listStaleChunks` snapshot) but exists. Tractable fix: a `WHERE content_chunks.embedded_at < EXCLUDED.embedded_at OR content_chunks.embedding IS NULL` predicate on the upsert. Out of scope for v0.34.4.0 because the upsert is not in the diff; pre-existing bug. Filed during v0.34.4.0 codex outside-voice review. + +- [ ] **v0.35.x: New stale rows inserted behind the keyset cursor.** A sync or `gbrain put_page` mid-`embed --stale` creates chunks with `embedding IS NULL` at `(page_id, chunk_index)` already passed by the cursor. Picked up on next run via the partial index; documented limitation. Possible fix: a second pass at end-of-run that does a fresh `countStaleChunks()` and re-enters the loop while count > 0 and budget allows. Filed during v0.34.4.0 codex outside-voice review. + ## MCP fix wave follow-ups (v0.34.1) - [ ] **v0.34.x: Source-scope `takes_*` ops (pre-existing leak surfaced during v0.34.1 adversarial review).** `takes_list`, `takes_search`, `takes_scorecard`, `takes_calibration` in `src/core/operations.ts:1248-1335` thread `ctx.takesHoldersAllowList` but never `ctx.sourceId`. An auth'd OAuth client scoped to `source_id='canon-a'` can call `takes_list --page_slug=foo` (slug in `canon-b`) and read takes attached to foreign-source pages. Pre-existing, not introduced by v0.34.1, but the wave was framed as "P0 source-isolation seal on the read path" and `takes_*` surfaces were missed. Fix: extend `TakesListOpts` in `src/core/engine.ts:186` with `sourceId?: string` + `sourceIds?: string[]`; thread `sourceScopeOpts(ctx)` at each op handler; engine `listTakes`/`searchTakes` filter via the `pages` JOIN. @@ -16,6 +23,7 @@ - [ ] **v0.34.x: `hybrid.ts:223` explicit-pick refactor.** The SearchOpts rebuild manually picks fields from HybridSearchOpts. This is the bug shape that caused the original v0.34.1 P0 leak — a new SearchOpts field is silently dropped if not manually added here. The wave added `sourceId` + `sourceIds` to the pick; future fields will keep hitting this footgun. Fix: refactor to spread + TypeScript `Pick<>` helper that narrows HybridSearchOpts → SearchOpts type-safely. + ## functional-area-resolver follow-ups (v0.32.3.0) - [ ] **v0.33.x: Dogfood `functional-area-resolver` on gbrain's own `skills/RESOLVER.md`** when it crosses ~12KB (currently 8KB). Apply the pattern to the Operational section first (largest). Filed during v0.32.3.0 CEO review. diff --git a/VERSION b/VERSION index 9501a6f3d..97745b5df 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -0.34.3.0 \ No newline at end of file +0.34.4.0 diff --git a/package.json b/package.json index b21ab81e4..e9254752b 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "gbrain", - "version": "0.34.3.0", + "version": "0.34.4.0", "description": "Postgres-native personal knowledge brain with hybrid RAG search", "type": "module", "main": "src/core/index.ts", diff --git a/src/commands/embed.ts b/src/commands/embed.ts index 541cd08c1..302144fc1 100644 --- a/src/commands/embed.ts +++ b/src/commands/embed.ts @@ -246,7 +246,8 @@ async function embedAll( // chunks that already have embeddings. // ───────────────────────────────────────────────────────────── if (staleOnly) { - return await embedAllStale(engine, dryRun, result, onProgress); + // D7: thread sourceId so `gbrain embed --stale --source X` actually scopes. + return await embedAllStale(engine, sourceId, dryRun, result, onProgress); } // v0.31.12: when sourceId is set, scope listPages to that source. @@ -359,13 +360,17 @@ async function embedAll( */ async function embedAllStale( engine: BrainEngine, + sourceId: string | undefined, dryRun: boolean, result: EmbedResult, onProgress?: (done: number, total: number, embedded: number) => void, ) { + // D7: thread sourceId so source-scoped runs only count + visit + // that source's NULL embeddings. + const sourceOpt = sourceId ? { sourceId } : undefined; + // Pre-flight: 0 stale chunks → nothing to do, no further DB reads. - // Cheapest possible exit on the autopilot common case. - const staleCount = await engine.countStaleChunks(); + const staleCount = await engine.countStaleChunks(sourceOpt); if (staleCount === 0) { if (dryRun) { console.log('[dry-run] Would embed 0 chunks (0 stale found)'); @@ -375,91 +380,260 @@ async function embedAllStale( return; } - // Pull only the stale chunks (no embedding column). - const staleRows = await engine.listStaleChunks(); - // v0.31.12: group by composite key (source_id::slug) so same-slug pages - // in different sources are embedded independently with correct source_id. - const byKey = new Map(); - for (const row of staleRows) { - const key = `${row.source_id}::${row.slug}`; - const list = byKey.get(key); - if (list) list.push(row); - else byKey.set(key, [row]); - } - - const keys = Array.from(byKey.keys()); - const totalStaleChunks = staleRows.length; - result.total_chunks += totalStaleChunks; - // skipped is "chunks we considered and skipped due to having an embedding". - // We never considered the non-stale chunks here, so leave skipped at 0. - // Callers reading EmbedResult who care about coverage should call - // engine.getStats() / engine.getHealth() afterward. - if (dryRun) { - result.would_embed += totalStaleChunks; - result.pages_processed += keys.length; - if (onProgress) { - // Emit a single tick to satisfy the contract (CLI progress reporters - // expect at least one start/finish pair). - onProgress(keys.length, keys.length, 0); - } - console.log(`[dry-run] Would embed ${totalStaleChunks} chunks across ${keys.length} pages`); + result.would_embed += staleCount; + result.total_chunks += staleCount; + if (onProgress) onProgress(1, 1, 0); + console.log(`[dry-run] Would embed ${staleCount} stale chunks`); return; } + // v0.33.3: cursor-paginated stale loading. Instead of pulling all 48K+ + // rows in one query (which times out on Supabase's 2-min pooler timeout), + // we page through 2000 rows at a time via keyset pagination on + // (page_id, chunk_index). Each query finishes in <1s. + const PAGE_SIZE = 2000; const CONCURRENCY = parseInt(process.env.GBRAIN_EMBED_CONCURRENCY || '20', 10); - let processed = 0; - async function embedOneKey(key: string) { - const stale = byKey.get(key)!; - // v0.31.12: extract source_id + slug from the stale row so getChunks - // and upsertChunks target the correct (source_id, slug) row. Pre-fix, - // both calls defaulted to source_id='default', silently discarding - // embeddings for every non-default source (e.g. media-corpus). - const sourceId = stale[0]?.source_id ?? 'default'; - const slug = stale[0].slug; - try { - const embeddings = await embedBatch(stale.map(c => c.chunk_text)); - // CRITICAL: passing ONLY the stale indices to upsertChunks would - // delete every non-stale chunk on the same page (the != ALL filter - // wipes any chunk_index NOT in the input). To preserve them, we - // re-fetch existing chunks for this page and merge. Bounded by the - // stale slug count, not by total slugs — autopilot common case - // is 0 stale (pre-flight short-circuit, never reaches this path). - const existing = await engine.getChunks(slug, { sourceId }); - const staleIdxToEmbedding = new Map(); - for (let j = 0; j < stale.length; j++) { - staleIdxToEmbedding.set(stale[j].chunk_index, embeddings[j]); + // D3 + D3a + D8: wall-clock budget. 30 min default; env override. + // The old single-shot `LIMIT 100000` query implicitly capped runtime + // by failing on timeout; pagination removed that cap. AbortController + // threads cancellation into (a) the retry sleep below, (b) the worker + // claim loop, and (c) the gateway embed call so an in-flight HTTP + // request also unwinds. + const BUDGET_MS = parseInt(process.env.GBRAIN_EMBED_TIME_BUDGET_MS || `${30 * 60 * 1000}`, 10); + const budgetController = new AbortController(); + const budgetTimer = setTimeout(() => budgetController.abort(), BUDGET_MS); + const budgetSignal = budgetController.signal; + + let totalProcessedPages = 0; + let afterPageId = 0; + let afterChunkIndex = -1; + let totalChunksLoaded = 0; + let budgetExitNotified = false; + + try { + // eslint-disable-next-line no-constant-condition + while (true) { + if (budgetSignal.aborted) { + if (!budgetExitNotified) { + console.error(`\n [embed] wall-clock budget (${BUDGET_MS}ms) exceeded; exiting cleanly. Re-run picks up via partial index.`); + budgetExitNotified = true; + } + break; } - const merged: ChunkInput[] = existing.map(c => ({ - chunk_index: c.chunk_index, - chunk_text: c.chunk_text, - chunk_source: c.chunk_source, - // For stale chunks: pass the new embedding. - // For non-stale chunks: pass undefined → COALESCE preserves existing embedding. - embedding: staleIdxToEmbedding.get(c.chunk_index) ?? undefined, - token_count: c.token_count || Math.ceil(c.chunk_text.length / 4), - })); - await engine.upsertChunks(slug, merged, { sourceId }); - result.embedded += stale.length; - } catch (e: unknown) { - console.error(`\n Error embedding ${slug}: ${e instanceof Error ? e.message : e}`); + + const batch = await engine.listStaleChunks({ + batchSize: PAGE_SIZE, + afterPageId, + afterChunkIndex, + ...(sourceId && { sourceId }), + }); + if (batch.length === 0) break; + totalChunksLoaded += batch.length; + + // Advance cursor to last row in this batch. + const last = batch[batch.length - 1]; + afterPageId = last.page_id; + afterChunkIndex = last.chunk_index; + + // Group by composite key (source_id::slug). + const byKey = new Map(); + for (const row of batch) { + const key = `${row.source_id}::${row.slug}`; + const list = byKey.get(key); + if (list) list.push(row); + else byKey.set(key, [row]); + } + + const keys = Array.from(byKey.keys()); + result.total_chunks += batch.length; + + let nextIdx = 0; + async function embedOneKey(key: string) { + const stale = byKey.get(key)!; + const keySourceId = stale[0]?.source_id ?? 'default'; + const slug = stale[0].slug; + try { + const embeddings = await embedBatchWithBackoff(stale.map(c => c.chunk_text), { abortSignal: budgetSignal }); + // Re-fetch existing chunks and merge to avoid deleting non-stale chunks. + const existing = await engine.getChunks(slug, { sourceId: keySourceId }); + const staleIdxToEmbedding = new Map(); + for (let j = 0; j < stale.length; j++) { + staleIdxToEmbedding.set(stale[j].chunk_index, embeddings[j]); + } + const merged: ChunkInput[] = existing.map(c => ({ + chunk_index: c.chunk_index, + chunk_text: c.chunk_text, + chunk_source: c.chunk_source, + embedding: staleIdxToEmbedding.get(c.chunk_index) ?? undefined, + token_count: c.token_count || Math.ceil(c.chunk_text.length / 4), + })); + await engine.upsertChunks(slug, merged, { sourceId: keySourceId }); + result.embedded += stale.length; + } catch (e: unknown) { + // Budget-fired aborts are expected on the way out; don't spam + // per-page "Error embedding" lines when we're shutting down. + if (budgetSignal.aborted) return; + console.error(`\n Error embedding ${slug}: ${e instanceof Error ? e.message : e}`); + } + totalProcessedPages++; + result.pages_processed++; + // Use staleCount as the estimated total for progress (not exact after + // pagination starts, but directionally correct). + onProgress?.(totalProcessedPages, Math.ceil(staleCount / PAGE_SIZE) * keys.length, result.embedded); + } + + async function worker() { + // D3a: workers check the budget before claiming the next key. + // A stuck mid-fetch worker also has the abortSignal threaded into + // its embedBatch call, so the in-flight HTTP cancels too. + while (nextIdx < keys.length && !budgetSignal.aborted) { + const idx = nextIdx++; + await embedOneKey(keys[idx]); + } + } + + const numWorkers = Math.min(CONCURRENCY, keys.length); + await Promise.all(Array.from({ length: numWorkers }, () => worker())); + + // If we got fewer rows than PAGE_SIZE, we've reached the end. + if (batch.length < PAGE_SIZE) break; } - processed++; - result.pages_processed++; - onProgress?.(processed, keys.length, result.embedded); + } finally { + clearTimeout(budgetTimer); } - let nextIdx = 0; - async function worker() { - while (nextIdx < keys.length) { - const idx = nextIdx++; - await embedOneKey(keys[idx]); + console.log(`Embedded ${result.embedded} chunks across ${totalProcessedPages} pages`); +} + +/** + * v0.33.3: rate-limit-aware embedBatch wrapper. + * + * The OpenAI SDK has built-in retry with exponential backoff, but its + * backoff window (max ~4s) is too short for TPM (tokens-per-minute) + * rate limits on large pages (~90K tokens). This wrapper catches + * 429-shaped errors, parses the retry delay from the error message + * (e.g. "Please try again in 248ms"), and sleeps before retrying. + * + * v0.33.4 hardening (codex + re-review findings): + * - D4: detect 429 via the wrapped error's `cause.status` (the gateway's + * normalizeAIError stores the original error there). Bare `e.status` + * never fires against an `AITransientError` wrap. Message-match stays + * as a fallback. + * - D4a: pass `maxRetries: 0` through `embedBatch` so the AI SDK's + * default 2-retry stack doesn't multiply this wrapper's 5 attempts. + * - D2: jitter the parsed delay ±30% so 20 concurrent workers don't + * resynchronize on the next 429 wave. + * - D3a/D8: when an external AbortSignal fires (wall-clock budget), the + * sleep wakes up early AND the abortSignal is threaded into the gateway + * embed call so an in-flight HTTP request cancels too. + * + * Up to MAX_RATE_LIMIT_RETRIES attempts with the parsed (jittered) delay + * (or a 60s fallback when the message can't be parsed). + * + * @internal Exported for unit tests; not part of the public surface. + */ +export const MAX_RATE_LIMIT_RETRIES = 5; +export const RATE_LIMIT_FALLBACK_MS = 60_000; +export const RATE_LIMIT_PAD_MS = 500; +export const RATE_LIMIT_JITTER = 0.3; + +export interface EmbedBatchWithBackoffOpts { + abortSignal?: AbortSignal; +} + +/** + * Walk the cause chain looking for a 429 status. The current + * `normalizeAIError` wraps once into `AITransientError` with `cause = original`, + * so one level is sufficient — but iterate to handle future wrap layers + * defensively (max 5 levels to bound a malformed cyclic chain). + * + * @internal exported for unit tests. + */ +export function detect429FromCause(e: unknown): boolean { + let cur: unknown = e; + for (let depth = 0; depth < 5 && cur !== undefined && cur !== null; depth++) { + const obj = cur as { status?: unknown; statusCode?: unknown; cause?: unknown }; + if (obj.status === 429 || obj.statusCode === 429) return true; + cur = obj.cause; + } + return false; +} + +/** + * Parse a Retry-After hint out of an OpenAI-style 429 message. Falls back + * to `RATE_LIMIT_FALLBACK_MS` when the message can't be parsed. Adds + * `RATE_LIMIT_PAD_MS` padding and `RATE_LIMIT_JITTER` randomization so + * concurrent workers don't resynchronize. + * + * @internal exported for unit tests. + */ +export function parseRetryDelayMs(msg: string, rng: () => number = Math.random): number { + let delayMs = RATE_LIMIT_FALLBACK_MS; + const msMatch = msg.match(/try again in (\d+)ms/i); + const secMatch = msg.match(/try again in ([\d.]+)s/i); + if (msMatch) delayMs = parseInt(msMatch[1], 10) + RATE_LIMIT_PAD_MS; + else if (secMatch) delayMs = Math.ceil(parseFloat(secMatch[1]) * 1000) + RATE_LIMIT_PAD_MS; + // D2: ±30% jitter to decorrelate the herd of 20 workers. + const jitterFactor = 1 + (rng() * 2 - 1) * RATE_LIMIT_JITTER; + return Math.max(1, Math.floor(delayMs * jitterFactor)); +} + +/** + * Sleep for `ms` milliseconds. Resolves early (not rejects) when `signal` + * fires, so the retry loop's caller can re-check `signal.aborted` and + * exit cleanly without an unhandled rejection. + * + * @internal exported for unit tests. + */ +export function abortableSleep(ms: number, signal?: AbortSignal): Promise { + return new Promise((resolve) => { + if (signal?.aborted) { + resolve(); + return; } - } + const timer = setTimeout(() => { + signal?.removeEventListener('abort', onAbort); + resolve(); + }, ms); + const onAbort = () => { + clearTimeout(timer); + signal?.removeEventListener('abort', onAbort); + resolve(); + }; + signal?.addEventListener('abort', onAbort, { once: true }); + }); +} - const numWorkers = Math.min(CONCURRENCY, keys.length); - await Promise.all(Array.from({ length: numWorkers }, () => worker())); +export async function embedBatchWithBackoff( + texts: string[], + opts: EmbedBatchWithBackoffOpts = {}, +): Promise { + const signal = opts.abortSignal; + for (let attempt = 0; attempt <= MAX_RATE_LIMIT_RETRIES; attempt++) { + if (signal?.aborted) throw new Error('embed budget aborted'); + try { + // D4a + D8: maxRetries:0 disables the SDK's stacked retries (so this + // wrapper is the single source of truth) and abortSignal threads + // through to the gateway so an in-flight HTTP request cancels mid-fetch. + return await embedBatch(texts, { maxRetries: 0, ...(signal && { abortSignal: signal }) }); + } catch (e: unknown) { + // If the budget fired we may have been aborted mid-fetch; bubble out. + if (signal?.aborted) throw e; + const msg = e instanceof Error ? e.message : String(e); + // D4: structured detection first (handles gateway-wrapped errors via + // cause chain); message-match as fallback for providers whose wrappers + // strip `cause.status`. + const isRateLimit = detect429FromCause(e) + || /rate.?limit|429/i.test(msg); + if (!isRateLimit || attempt === MAX_RATE_LIMIT_RETRIES) throw e; - console.log(`Embedded ${result.embedded} chunks across ${keys.length} pages`); + const delayMs = parseRetryDelayMs(msg); + console.error(` [rate-limit] attempt ${attempt + 1}/${MAX_RATE_LIMIT_RETRIES}, waiting ${delayMs}ms...`); + await abortableSleep(delayMs, signal); + } + } + // Unreachable, but TypeScript needs it. + return embedBatch(texts); } diff --git a/src/core/ai/gateway.ts b/src/core/ai/gateway.ts index eef0ee2d6..406497d9f 100644 --- a/src/core/ai/gateway.ts +++ b/src/core/ai/gateway.ts @@ -785,7 +785,34 @@ const MIN_SUB_BATCH = 1; * declared `safety_factor` so a transient miss doesn't permanently cap * throughput. */ -export async function embed(texts: string[]): Promise { +/** + * Per-call passthroughs for `embed()`. v0.33.4 added both fields so the + * `embed --stale` retry wrapper can (a) cancel mid-fetch when the wall-clock + * budget fires, and (b) suppress the AI SDK's default 2-retry stack so the + * wrapper's retry-after-aware loop is the single source of truth. + * + * Both are optional; production callers that don't pass them get unchanged + * pre-v0.33.4 behavior. + */ +export interface EmbedOpts { + /** + * Propagated to Vercel AI SDK's `embedMany({abortSignal})`. When the + * caller's wall-clock budget fires, an in-flight HTTP request is + * cancelled within seconds instead of waiting out the provider's HTTP + * timeout (~30s on OpenAI). + */ + abortSignal?: AbortSignal; + /** + * Propagated to Vercel AI SDK's `embedMany({maxRetries})`. Default in + * the SDK is 2 (so up to 3 attempts per call). Pass `0` to disable + * SDK retries when a higher-level wrapper owns the retry policy — + * otherwise SDK and wrapper retries stack and amplify rate-limit + * pressure (3 × N wrapper attempts). + */ + maxRetries?: number; +} + +export async function embed(texts: string[], opts?: EmbedOpts): Promise { if (!texts || texts.length === 0) return []; const cfg = requireConfig(); @@ -807,7 +834,7 @@ export async function embed(texts: string[]): Promise { const allEmbeddings: Float32Array[] = []; for (const batch of batches) { - const result = await embedSubBatch(batch, model, providerOpts, expected, recipe, modelId); + const result = await embedSubBatch(batch, model, providerOpts, expected, recipe, modelId, opts); allEmbeddings.push(...result); } @@ -925,12 +952,18 @@ async function embedSubBatch( expectedDims: number, recipe: Recipe, modelId: string, + opts?: EmbedOpts, ): Promise { try { const result = await _embedTransport({ model, values: texts, providerOptions: providerOpts, + // v0.33.4: caller-supplied abortSignal + maxRetries passthrough. + // Undefined fields are ignored by the AI SDK so the call shape stays + // identical for production callers that don't opt in. + ...(opts?.abortSignal !== undefined && { abortSignal: opts.abortSignal }), + ...(opts?.maxRetries !== undefined && { maxRetries: opts.maxRetries }), }); const first = result.embeddings?.[0]; @@ -950,8 +983,8 @@ async function embedSubBatch( if (isTokenLimitError(err) && texts.length > MIN_SUB_BATCH) { shrinkOnMiss(recipe); const mid = Math.ceil(texts.length / 2); - const left = await embedSubBatch(texts.slice(0, mid), model, providerOpts, expectedDims, recipe, modelId); - const right = await embedSubBatch(texts.slice(mid), model, providerOpts, expectedDims, recipe, modelId); + const left = await embedSubBatch(texts.slice(0, mid), model, providerOpts, expectedDims, recipe, modelId, opts); + const right = await embedSubBatch(texts.slice(mid), model, providerOpts, expectedDims, recipe, modelId, opts); return [...left, ...right]; } throw normalizeAIError(err, `embed(${recipe.id}:${modelId})`); diff --git a/src/core/embedding.ts b/src/core/embedding.ts index 40b52104a..553aa1916 100644 --- a/src/core/embedding.ts +++ b/src/core/embedding.ts @@ -29,6 +29,21 @@ export interface EmbedBatchOptions { * tick a reporter; Minion handlers can call job.updateProgress here. */ onBatchComplete?: (done: number, total: number) => void; + /** + * v0.33.4 (D8): propagate the caller's `AbortSignal` into Vercel AI SDK's + * `embedMany({abortSignal})` so a wall-clock budget can cancel mid-fetch. + * Without this, a worker stuck mid-HTTP on a ~30s OpenAI timeout ignores + * the budget until the fetch resolves. + */ + abortSignal?: AbortSignal; + /** + * v0.33.4 (D4a): cap on AI SDK's per-call retries. Default in `embedMany` + * is 2 (so up to 3 attempts). Pass `0` from higher-level wrappers that + * own their own retry policy, otherwise wrapper × SDK retries stack + * (e.g. 3 SDK attempts × 5 wrapper attempts = 15 cycles per embedBatch) + * and amplify rate-limit pressure. + */ + maxRetries?: number; } /** @@ -43,14 +58,20 @@ export async function embedBatch( options: EmbedBatchOptions = {}, ): Promise { if (!texts || texts.length === 0) return []; + // Build the gateway-call passthrough once; undefined fields stay undefined + // so non-opt-in callers see unchanged pre-v0.33.4 behavior. + const gwOpts = { + ...(options.abortSignal !== undefined && { abortSignal: options.abortSignal }), + ...(options.maxRetries !== undefined && { maxRetries: options.maxRetries }), + }; // Fast path: small batch, no progress callback — single gateway call. if (texts.length <= BATCH_SIZE && !options.onBatchComplete) { - return gatewayEmbed(texts); + return gatewayEmbed(texts, gwOpts); } const results: Float32Array[] = []; for (let i = 0; i < texts.length; i += BATCH_SIZE) { const slice = texts.slice(i, i + BATCH_SIZE); - const out = await gatewayEmbed(slice); + const out = await gatewayEmbed(slice, gwOpts); results.push(...out); options.onBatchComplete?.(results.length, texts.length); } diff --git a/src/core/engine.ts b/src/core/engine.ts index 24b4703c9..101db1445 100644 --- a/src/core/engine.ts +++ b/src/core/engine.ts @@ -533,20 +533,38 @@ export interface BrainEngine { */ getChunks(slug: string, opts?: { sourceId?: string }): Promise; /** - * Count chunks across the entire brain where embedded_at IS NULL. + * Count chunks across the brain where embedding IS NULL. * Pre-flight short-circuit for `embed --stale` so a 100%-embedded brain * does no further work after a single SELECT count(*) (~50 bytes wire). + * + * `opts.sourceId` scopes the count to a single source. When omitted, + * counts across every source in the brain. Operators running + * `gbrain embed --stale --source media-corpus` expect only that + * source's NULLs touched; the caller threads `sourceId` here. */ - countStaleChunks(): Promise; + countStaleChunks(opts?: { sourceId?: string }): Promise; /** - * Return every chunk where embedded_at IS NULL, with the metadata needed + * Return every chunk where embedding IS NULL, with the metadata needed * to call embedBatch + upsertChunks. The `embedding` column is omitted * by design — stale rows have NULL embeddings, so shipping them wastes * wire bytes for no gain. Caller groups by slug, embeds, and re-upserts. * - * Bounded by an internal LIMIT of 100000 to mirror listPages. + * v0.33.3: cursor-paginated — yields up to `batchSize` rows per call + * (default 2000) to stay within Supabase's statement_timeout. Pass the + * last row's `(page_id, chunk_index)` as `afterPageId`/`afterChunkIndex` + * to fetch the next page. When fewer than `batchSize` rows come back, + * the caller has reached the end. + * + * `opts.sourceId` scopes the scan to a single source (matches the + * countStaleChunks contract). Paired with embedAllStale's --source + * support. */ - listStaleChunks(): Promise; + listStaleChunks(opts?: { + batchSize?: number; + afterPageId?: number; + afterChunkIndex?: number; + sourceId?: string; + }): Promise; /** * Delete every chunk for a page. Internal page-id lookup is sourceId-scoped * when `opts.sourceId` is given; otherwise the bare-slug subquery returns diff --git a/src/core/migrate.ts b/src/core/migrate.ts index 907723793..3c074aa4e 100644 --- a/src/core/migrate.ts +++ b/src/core/migrate.ts @@ -3150,6 +3150,66 @@ export const MIGRATIONS: Migration[] = [ ON oauth_clients USING GIN (federated_read); `, }, + { + version: 66, + name: 'embed_stale_partial_index', + // Renumbered v58→v59→v60→v66 across merge waves: + // - v58 was taken by master's v0.33.3 edges_backfilled_at. + // - v59 was taken by master's v0.34.0 code_traversal_cache. + // - v60-v65 were taken by master's v0.34.1 oauth_clients source-isolation cluster. + // All landed before this branch could ship. + // + // Partial index for `embedding IS NULL` on content_chunks. + // + // The `embed --stale` command scans for chunks missing embeddings. + // Without this index, the query does a full table scan of 300K+ rows + // to find the ~48K NULLs, taking >2 min and hitting Supabase's + // statement_timeout. With the partial index, the scan is instant. + // + // Also used by countStaleChunks() for the pre-flight check. + // + // Engine-aware via handler (mirrors v14): Postgres uses + // CREATE INDEX CONCURRENTLY to avoid the ShareLock on `content_chunks` + // that a plain CREATE INDEX takes for the duration of the build. + // On a 373K-row table this lock blocks every concurrent write (sync, + // embed, autopilot). CONCURRENTLY refuses to run inside a transaction + // AND postgres.js's multi-statement `.unsafe()` wraps in an implicit + // transaction, so each statement runs as a separate call. A failed + // CONCURRENTLY leaves an invalid index with the target name; the + // handler pre-drops any invalid remnant via pg_index.indisvalid. + // PGLite has no concurrent writers, so plain CREATE is safe. + idempotent: true, + sql: '', + handler: async (engine) => { + if (engine.kind === 'postgres') { + await engine.runMigration( + 66, + `DO $$ BEGIN + IF EXISTS ( + SELECT 1 FROM pg_index i + JOIN pg_class c ON c.oid = i.indexrelid + WHERE c.relname = 'idx_chunks_embedding_null' AND NOT i.indisvalid + ) THEN + EXECUTE 'DROP INDEX CONCURRENTLY IF EXISTS idx_chunks_embedding_null'; + END IF; + END $$;` + ); + await engine.runMigration( + 66, + `CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_chunks_embedding_null + ON content_chunks (page_id, chunk_index) + WHERE embedding IS NULL;` + ); + } else { + await engine.runMigration( + 66, + `CREATE INDEX IF NOT EXISTS idx_chunks_embedding_null + ON content_chunks (page_id, chunk_index) + WHERE embedding IS NULL;` + ); + } + }, + }, ]; export const LATEST_VERSION = MIGRATIONS.length > 0 diff --git a/src/core/pglite-engine.ts b/src/core/pglite-engine.ts index 797358190..7299ba9bb 100644 --- a/src/core/pglite-engine.ts +++ b/src/core/pglite-engine.ts @@ -1309,25 +1309,65 @@ export class PGLiteEngine implements BrainEngine { return (rows as Record[]).map(r => rowToChunk(r)); } - async countStaleChunks(): Promise { + async countStaleChunks(opts?: { sourceId?: string }): Promise { + // D7: source-scoped count for `gbrain embed --stale --source X`. + if (opts?.sourceId === undefined) { + const { rows } = await this.db.query( + `SELECT count(*)::int AS count + FROM content_chunks + WHERE embedding IS NULL`, + ); + const count = (rows[0] as { count: number } | undefined)?.count ?? 0; + return Number(count); + } const { rows } = await this.db.query( `SELECT count(*)::int AS count - FROM content_chunks - WHERE embedding IS NULL`, + FROM content_chunks cc + JOIN pages p ON p.id = cc.page_id + WHERE cc.embedding IS NULL + AND p.source_id = $1`, + [opts.sourceId], ); const count = (rows[0] as { count: number } | undefined)?.count ?? 0; return Number(count); } - async listStaleChunks(): Promise { + async listStaleChunks(opts?: { + batchSize?: number; + afterPageId?: number; + afterChunkIndex?: number; + sourceId?: string; + }): Promise { + const limit = opts?.batchSize ?? 2000; + const afterPid = opts?.afterPageId ?? 0; + const afterIdx = opts?.afterChunkIndex ?? -1; + // D7: optional source-scoped cursor scan. PGLite mirrors postgres-engine + // so the engine-parity E2E catches drift. + if (opts?.sourceId === undefined) { + const { rows } = await this.db.query( + `SELECT p.slug, cc.chunk_index, cc.chunk_text, cc.chunk_source, + cc.model, cc.token_count, p.source_id, cc.page_id + FROM content_chunks cc + JOIN pages p ON p.id = cc.page_id + WHERE cc.embedding IS NULL + AND (cc.page_id, cc.chunk_index) > ($1, $2) + ORDER BY cc.page_id, cc.chunk_index + LIMIT $3`, + [afterPid, afterIdx, limit], + ); + return rows as unknown as StaleChunkRow[]; + } const { rows } = await this.db.query( `SELECT p.slug, cc.chunk_index, cc.chunk_text, cc.chunk_source, - cc.model, cc.token_count, p.source_id + cc.model, cc.token_count, p.source_id, cc.page_id FROM content_chunks cc JOIN pages p ON p.id = cc.page_id WHERE cc.embedding IS NULL - ORDER BY p.id, cc.chunk_index - LIMIT 100000`, + AND p.source_id = $1 + AND (cc.page_id, cc.chunk_index) > ($2, $3) + ORDER BY cc.page_id, cc.chunk_index + LIMIT $4`, + [opts.sourceId, afterPid, afterIdx, limit], ); return rows as unknown as StaleChunkRow[]; } diff --git a/src/core/postgres-engine.ts b/src/core/postgres-engine.ts index e57aab13c..cdc12c29c 100644 --- a/src/core/postgres-engine.ts +++ b/src/core/postgres-engine.ts @@ -1308,26 +1308,70 @@ export class PostgresEngine implements BrainEngine { return rows.map((r) => rowToChunk(r as Record)); } - async countStaleChunks(): Promise { + async countStaleChunks(opts?: { sourceId?: string }): Promise { const sql = this.sql; + // Fast path: no source filter → bare count query, no join. + // Slow path: source-scoped count → join pages. + // D7: closes the bug where `gbrain embed --stale --source X` silently + // dropped X and counted across every source. + if (opts?.sourceId === undefined) { + const [row] = await sql` + SELECT count(*)::int AS count + FROM content_chunks + WHERE embedding IS NULL + `; + return Number((row as { count?: number } | undefined)?.count ?? 0); + } const [row] = await sql` SELECT count(*)::int AS count - FROM content_chunks - WHERE embedding IS NULL + FROM content_chunks cc + JOIN pages p ON p.id = cc.page_id + WHERE cc.embedding IS NULL + AND p.source_id = ${opts.sourceId} `; return Number((row as { count?: number } | undefined)?.count ?? 0); } - async listStaleChunks(): Promise { + async listStaleChunks(opts?: { + batchSize?: number; + afterPageId?: number; + afterChunkIndex?: number; + sourceId?: string; + }): Promise { const sql = this.sql; + const limit = opts?.batchSize ?? 2000; + const afterPid = opts?.afterPageId ?? 0; + const afterIdx = opts?.afterChunkIndex ?? -1; + // Cursor-paginated: keyset pagination on (page_id, chunk_index). + // The partial index idx_chunks_embedding_null makes the WHERE fast; + // LIMIT keeps each round-trip well within statement_timeout. + // + // D7: optional source_id filter. NULL/undefined = scan all sources + // (pre-existing behavior); a value scopes to that source so + // `gbrain embed --stale --source X` actually does what it says. + if (opts?.sourceId === undefined) { + const rows = await sql` + SELECT p.slug, cc.chunk_index, cc.chunk_text, cc.chunk_source, + cc.model, cc.token_count, p.source_id, cc.page_id + FROM content_chunks cc + JOIN pages p ON p.id = cc.page_id + WHERE cc.embedding IS NULL + AND (cc.page_id, cc.chunk_index) > (${afterPid}, ${afterIdx}) + ORDER BY cc.page_id, cc.chunk_index + LIMIT ${limit} + `; + return rows as unknown as StaleChunkRow[]; + } const rows = await sql` SELECT p.slug, cc.chunk_index, cc.chunk_text, cc.chunk_source, - cc.model, cc.token_count, p.source_id + cc.model, cc.token_count, p.source_id, cc.page_id FROM content_chunks cc JOIN pages p ON p.id = cc.page_id WHERE cc.embedding IS NULL - ORDER BY p.id, cc.chunk_index - LIMIT 100000 + AND p.source_id = ${opts.sourceId} + AND (cc.page_id, cc.chunk_index) > (${afterPid}, ${afterIdx}) + ORDER BY cc.page_id, cc.chunk_index + LIMIT ${limit} `; return rows as unknown as StaleChunkRow[]; } diff --git a/src/core/types.ts b/src/core/types.ts index 468ef135e..381e9b3bc 100644 --- a/src/core/types.ts +++ b/src/core/types.ts @@ -349,6 +349,8 @@ export interface StaleChunkRow { token_count: number | null; /** v0.31.12: source_id so embed --stale can thread it through getChunks/upsertChunks. */ source_id: string; + /** v0.33.3: page_id for cursor pagination in listStaleChunks. */ + page_id: number; } export interface ChunkInput { diff --git a/test/e2e/embed-stale-pagination.test.ts b/test/e2e/embed-stale-pagination.test.ts new file mode 100644 index 000000000..d33d2ae70 --- /dev/null +++ b/test/e2e/embed-stale-pagination.test.ts @@ -0,0 +1,290 @@ +/** + * E2E regression: cursor-paginated `embed --stale` (D7 + IRON RULE from the + * v0.33.4 plan-eng-review). Seeds >PAGE_SIZE chunks with `embedding IS NULL` + * and walks the production cursor to verify: + * + * - Static case: every chunk visited exactly once across multiple batches. + * - Cursor monotonically advances on `(page_id, chunk_index)`. + * - Migration v66's partial index `idx_chunks_embedding_null` exists. + * - D7: source-scoped scan returns ONLY that source's NULLs even when + * same-slug pages exist across sources. + * - Failed-page semantics: a failed upsert keeps `embedding IS NULL` and + * the next run picks it up (covers the "cursor advances past failures + * in the same run" intended behavior). + * + * Requires DATABASE_URL. Skips gracefully otherwise. + * + * Run: DATABASE_URL=... bun test test/e2e/embed-stale-pagination.test.ts + */ + +import { describe, test, expect, beforeAll, afterAll } from 'bun:test'; +import { hasDatabase, setupDB, teardownDB, getEngine } from './helpers.ts'; +import { getConnection } from '../../src/core/db.ts'; + +const getConn = getConnection; + +const skip = !hasDatabase(); +const describeE2E = skip ? describe.skip : describe; + +if (skip) { + console.log('Skipping E2E embed-stale-pagination tests (DATABASE_URL not set)'); +} + +const PAGE_SIZE = 2000; // matches embed.ts production constant + +/** + * Seed N chunks across M pages, all with `embedding IS NULL`. + * Uses raw SQL because the public engine API doesn't expose page+chunk + * insertion at this granularity without going through importFromContent + * (which would compute embeddings if a model is configured). + */ +async function seedNullChunks(opts: { + sourceId: string; + pageCount: number; + chunksPerPage: number; + slugPrefix: string; +}) { + const sql = getConn(); + // Bulk-insert all pages in one round-trip via unnest, then bulk-insert + // all chunks. Avoids 3000+ sequential round-trips on the "static case" + // seed which would otherwise blow the test's 5s default timeout. + const slugs = Array.from({ length: opts.pageCount }, (_, p) => + `${opts.slugPrefix}-${p.toString().padStart(4, '0')}`, + ); + const titles = slugs; + const sourceIds = slugs.map(() => opts.sourceId); + const pageRows = await sql` + INSERT INTO pages (slug, title, type, compiled_truth, timeline, frontmatter, source_id) + SELECT s, t, 'page', '', '', '{}'::jsonb, src + FROM unnest(${sql.array(slugs)}::text[], ${sql.array(titles)}::text[], ${sql.array(sourceIds)}::text[]) AS u(s, t, src) + RETURNING id, slug + `; + const slugToId = new Map(); + for (const r of pageRows as unknown as Array<{ id: number; slug: string }>) { + slugToId.set(r.slug, r.id); + } + // Build the flat (page_id, chunk_index, chunk_text) arrays. + const pageIds: number[] = []; + const indices: number[] = []; + const texts: string[] = []; + for (let p = 0; p < opts.pageCount; p++) { + const slug = slugs[p]; + const pid = slugToId.get(slug); + if (pid === undefined) throw new Error(`seed: missing pageId for ${slug}`); + for (let c = 0; c < opts.chunksPerPage; c++) { + pageIds.push(pid); + indices.push(c); + texts.push(`chunk-${p}-${c}`); + } + } + // `model` column is NOT NULL with a default ('text-embedding-3-large') + // — omit it from the INSERT so the default applies. `embedding` is + // nullable; we intentionally leave it NULL so the chunk is "stale". + await sql` + INSERT INTO content_chunks (page_id, chunk_index, chunk_text, chunk_source, token_count, embedding, embedded_at) + SELECT pid, idx, txt, 'compiled_truth', 5, NULL, NULL + FROM unnest(${sql.array(pageIds)}::int[], ${sql.array(indices)}::int[], ${sql.array(texts)}::text[]) AS u(pid, idx, txt) + `; +} + +/** Ensure the migration v66 partial index exists in the seeded DB. */ +async function indexExists(name: string): Promise { + const sql = getConn(); + const rows = await sql` + SELECT 1 FROM pg_indexes WHERE indexname = ${name} + `; + return rows.length > 0; +} + +describeE2E('embed --stale cursor pagination (D7 + REGRESSION)', () => { + beforeAll(async () => { + await setupDB(); + // setupDB() seeds one default source; D7 cases add a second source + // via raw SQL on demand. + }); + + afterAll(async () => { + // setupDB's ALL_TABLES truncate list does NOT include `sources` + // (the `default` row is treated as fixed). This test adds an + // `other-source` row for the D7 cases; clean it up so later tests + // running on the same DB don't see a never-synced source (which + // mechanical.test.ts:`gbrain doctor` would correctly fail on). + try { + await getConn()`DELETE FROM sources WHERE id <> 'default'`; + } catch { + // If the connection is already torn down, that's fine. + } + await teardownDB(); + }); + + test('migration v66 created partial index idx_chunks_embedding_null', async () => { + const exists = await indexExists('idx_chunks_embedding_null'); + expect(exists).toBe(true); + }); + + test('static case: every chunk visited exactly once across multiple batches', async () => { + const engine = getEngine(); + // Truncate so other tests don't bleed in. + await getConn()`TRUNCATE content_chunks, pages CASCADE`; + // Seed PAGE_SIZE + 500 chunks to force at least 2 cursor pages. + const TOTAL_PAGES = Math.ceil((PAGE_SIZE + 500) / 5); // 5 chunks/page + const CHUNKS_PER_PAGE = 5; + await seedNullChunks({ + sourceId: 'default', + pageCount: TOTAL_PAGES, + chunksPerPage: CHUNKS_PER_PAGE, + slugPrefix: 'static-case', + }); + const expectedTotal = TOTAL_PAGES * CHUNKS_PER_PAGE; + + const count = await engine.countStaleChunks(); + expect(count).toBe(expectedTotal); + + // Walk the cursor manually (production code in embedAllStale does + // the same; we re-implement the loop here so the test asserts on + // the engine contract, not the caller's wrapper). + const visited = new Set(); + let lastPageId = -1; + let lastChunkIndex = -1; + let afterPageId = 0; + let afterChunkIndex = -1; + let cursorMonotonic = true; + let batchCount = 0; + + // eslint-disable-next-line no-constant-condition + while (true) { + const batch = await engine.listStaleChunks({ + batchSize: PAGE_SIZE, + afterPageId, + afterChunkIndex, + }); + if (batch.length === 0) break; + batchCount++; + for (const row of batch) { + const key = `${row.page_id}::${row.chunk_index}`; + // No duplicate visits. + expect(visited.has(key)).toBe(false); + visited.add(key); + // Monotonic advance on (page_id, chunk_index). + const adv = row.page_id > lastPageId + || (row.page_id === lastPageId && row.chunk_index > lastChunkIndex); + if (!adv) cursorMonotonic = false; + lastPageId = row.page_id; + lastChunkIndex = row.chunk_index; + } + const tail = batch[batch.length - 1]; + afterPageId = tail.page_id; + afterChunkIndex = tail.chunk_index; + if (batch.length < PAGE_SIZE) break; + } + + expect(visited.size).toBe(expectedTotal); + expect(cursorMonotonic).toBe(true); + expect(batchCount).toBeGreaterThanOrEqual(2); // forced split + }); + + test('source-scoped scan only returns the target source (D7)', async () => { + const engine = getEngine(); + await getConn()`TRUNCATE content_chunks, pages CASCADE`; + // Insert the second source row (default already exists from setupDB). + await getConn()` + INSERT INTO sources (id, name, local_path) + VALUES ('other-source', 'other-source', '/tmp/other') + ON CONFLICT (id) DO NOTHING + `; + // Same slug-prefix in both sources to verify the filter is true source + // separation, not slug separation. + await seedNullChunks({ sourceId: 'default', pageCount: 3, chunksPerPage: 2, slugPrefix: 'shared' }); + await seedNullChunks({ sourceId: 'other-source', pageCount: 5, chunksPerPage: 2, slugPrefix: 'shared' }); + + // Global count: 8 pages × 2 chunks = 16. + expect(await engine.countStaleChunks()).toBe(16); + // Default-only: 3 × 2 = 6. + expect(await engine.countStaleChunks({ sourceId: 'default' })).toBe(6); + // Other-only: 5 × 2 = 10. + expect(await engine.countStaleChunks({ sourceId: 'other-source' })).toBe(10); + + // listStaleChunks should return ONLY the requested source. + const defaultRows = await engine.listStaleChunks({ sourceId: 'default', batchSize: 100 }); + expect(defaultRows).toHaveLength(6); + for (const row of defaultRows) expect(row.source_id).toBe('default'); + + const otherRows = await engine.listStaleChunks({ sourceId: 'other-source', batchSize: 100 }); + expect(otherRows).toHaveLength(10); + for (const row of otherRows) expect(row.source_id).toBe('other-source'); + }); + + test('duplicate slug across sources: cursor on (page_id, chunk_index) keeps them separate', async () => { + const engine = getEngine(); + await getConn()`TRUNCATE content_chunks, pages CASCADE`; + await getConn()` + INSERT INTO sources (id, name, local_path) + VALUES ('other-source', 'other-source', '/tmp/other') + ON CONFLICT (id) DO NOTHING + `; + // Same slug, different sources. + await seedNullChunks({ sourceId: 'default', pageCount: 1, chunksPerPage: 3, slugPrefix: 'collide' }); + await seedNullChunks({ sourceId: 'other-source', pageCount: 1, chunksPerPage: 3, slugPrefix: 'collide' }); + + const allRows = await engine.listStaleChunks({ batchSize: 100 }); + expect(allRows).toHaveLength(6); + const seen = new Set(); + for (const row of allRows) { + // Each (source_id, slug, chunk_index) is unique. + const key = `${row.source_id}::${row.slug}::${row.chunk_index}`; + expect(seen.has(key)).toBe(false); + seen.add(key); + } + // Two distinct source_ids. + const sourceIds = new Set(allRows.map(r => r.source_id)); + expect(sourceIds.size).toBe(2); + }); + + test('failed embedding stays NULL; next run picks it up', async () => { + const engine = getEngine(); + await getConn()`TRUNCATE content_chunks, pages CASCADE`; + await seedNullChunks({ sourceId: 'default', pageCount: 1, chunksPerPage: 3, slugPrefix: 'fail-case' }); + + // First walk — simulate that we read all 3 chunks but did not upsert + // anything (i.e. the embedding step "failed" for every page). Cursor + // advance is observed; the chunks should still be NULL. + const firstRun = await engine.listStaleChunks({ batchSize: 100 }); + expect(firstRun).toHaveLength(3); + + // Verify the NULL rows are still there. + expect(await engine.countStaleChunks()).toBe(3); + + // Second walk picks them up because the partial index still finds them. + const secondRun = await engine.listStaleChunks({ batchSize: 100 }); + expect(secondRun).toHaveLength(3); + // Same (page_id, chunk_index) set. + const firstKeys = new Set(firstRun.map(r => `${r.page_id}::${r.chunk_index}`)); + const secondKeys = new Set(secondRun.map(r => `${r.page_id}::${r.chunk_index}`)); + expect(firstKeys).toEqual(secondKeys); + }); + + test('page split across batches: a multi-chunk page can land in two cursor pages', async () => { + const engine = getEngine(); + await getConn()`TRUNCATE content_chunks, pages CASCADE`; + // One page with 5 chunks; use batchSize=2 to force splits. + await seedNullChunks({ sourceId: 'default', pageCount: 1, chunksPerPage: 5, slugPrefix: 'split' }); + const visited: number[] = []; + let after_pid = 0; + let after_idx = -1; + // eslint-disable-next-line no-constant-condition + while (true) { + const batch = await engine.listStaleChunks({ + batchSize: 2, + afterPageId: after_pid, + afterChunkIndex: after_idx, + }); + if (batch.length === 0) break; + for (const r of batch) visited.push(r.chunk_index); + const tail = batch[batch.length - 1]; + after_pid = tail.page_id; + after_idx = tail.chunk_index; + if (batch.length < 2) break; + } + expect(visited).toEqual([0, 1, 2, 3, 4]); + }); +}); diff --git a/test/embed.serial.test.ts b/test/embed.serial.test.ts index 05f800a8f..96e2f20bf 100644 --- a/test/embed.serial.test.ts +++ b/test/embed.serial.test.ts @@ -7,18 +7,30 @@ import type { BrainEngine } from '../src/core/engine.ts'; let activeEmbedCalls = 0; let maxConcurrentEmbedCalls = 0; let totalEmbedCalls = 0; +// D5: capture per-call opts so tests can assert maxRetries / abortSignal +// passthrough into the gateway path. +let lastEmbedBatchOpts: unknown = undefined; +// D5: pluggable behavior for tests that need to simulate 429s or aborts. +let embedBatchBehavior: ((texts: string[], opts?: unknown) => Promise) | null = null; mock.module('../src/core/embedding.ts', () => ({ - embedBatch: async (texts: string[]) => { + embedBatch: async (texts: string[], opts?: unknown) => { activeEmbedCalls++; totalEmbedCalls++; + lastEmbedBatchOpts = opts; if (activeEmbedCalls > maxConcurrentEmbedCalls) { maxConcurrentEmbedCalls = activeEmbedCalls; } - // Simulate API latency so concurrent workers actually overlap. - await new Promise(r => setTimeout(r, 30)); - activeEmbedCalls--; - return texts.map(() => new Float32Array(1536)); + try { + if (embedBatchBehavior) { + return await embedBatchBehavior(texts, opts); + } + // Default: simulate API latency so concurrent workers actually overlap. + await new Promise(r => setTimeout(r, 30)); + return texts.map(() => new Float32Array(1536)); + } finally { + activeEmbedCalls--; + } }, })); @@ -47,10 +59,13 @@ beforeEach(() => { activeEmbedCalls = 0; maxConcurrentEmbedCalls = 0; totalEmbedCalls = 0; + lastEmbedBatchOpts = undefined; + embedBatchBehavior = null; }); afterEach(() => { delete process.env.GBRAIN_EMBED_CONCURRENCY; + delete process.env.GBRAIN_EMBED_TIME_BUDGET_MS; }); describe('runEmbed --all (parallel)', () => { @@ -111,8 +126,9 @@ describe('runEmbed --all (parallel)', () => { ['stale', [{ chunk_index: 0, chunk_text: 'hi', chunk_source: 'compiled_truth', embedded_at: null, token_count: 1 }]], ]); // Stale path uses countStaleChunks + listStaleChunks (SQL-side filter), not listPages. + // D5a: source_id + page_id required on StaleChunkRow as of v0.33.3 cursor pagination. const stale = [ - { slug: 'stale', chunk_index: 0, chunk_text: 'hi', chunk_source: 'compiled_truth', model: null, token_count: 1 }, + { slug: 'stale', chunk_index: 0, chunk_text: 'hi', chunk_source: 'compiled_truth' as const, model: null, token_count: 1, source_id: 'default', page_id: 1 }, ]; const engine = mockEngine({ @@ -150,9 +166,10 @@ describe('runEmbedCore --dry-run never calls the embedding model', () => { ]), ); // SQL-side stale path: 6 stale rows across 3 pages. - const stale = pages.flatMap(p => [ - { slug: p.slug, chunk_index: 0, chunk_text: 'a', chunk_source: 'compiled_truth', model: null, token_count: 1 }, - { slug: p.slug, chunk_index: 1, chunk_text: 'b', chunk_source: 'compiled_truth', model: null, token_count: 1 }, + // D5a: source_id + page_id required on StaleChunkRow. + const stale = pages.flatMap((p, pi) => [ + { slug: p.slug, chunk_index: 0, chunk_text: 'a', chunk_source: 'compiled_truth' as const, model: null, token_count: 1, source_id: 'default', page_id: pi + 1 }, + { slug: p.slug, chunk_index: 1, chunk_text: 'b', chunk_source: 'compiled_truth' as const, model: null, token_count: 1, source_id: 'default', page_id: pi + 1 }, ]); const upserts: string[] = []; @@ -177,17 +194,21 @@ describe('runEmbedCore --dry-run never calls the embedding model', () => { // skipped is 0 in the new SQL-side path: we never considered non-stale chunks. expect(result.skipped).toBe(0); expect(result.total_chunks).toBe(6); // only stale chunks counted in SQL-side path - expect(result.pages_processed).toBe(3); + // v0.33.3 cherry-pick: dry-run skips the cursor walk and only does a + // countStaleChunks call. pages_processed is 0 because we don't enumerate + // pages in dry-run (cheaper pre-flight). + expect(result.pages_processed).toBe(0); }); test('dry-run --stale correctly identifies stale chunks (SQL-side path)', async () => { const { runEmbedCore } = await import('../src/commands/embed.ts'); // SQL-side stale: only the 3 chunks where embedding IS NULL come back, // grouped by slug. 'fresh' page has no stale rows so it's not in the result. + // D5a: source_id + page_id required on StaleChunkRow. const stale = [ - { slug: 'partial', chunk_index: 1, chunk_text: 'b', chunk_source: 'compiled_truth', model: null, token_count: 1 }, - { slug: 'all-stale', chunk_index: 0, chunk_text: 'a', chunk_source: 'compiled_truth', model: null, token_count: 1 }, - { slug: 'all-stale', chunk_index: 1, chunk_text: 'b', chunk_source: 'compiled_truth', model: null, token_count: 1 }, + { slug: 'partial', chunk_index: 1, chunk_text: 'b', chunk_source: 'compiled_truth' as const, model: null, token_count: 1, source_id: 'default', page_id: 1 }, + { slug: 'all-stale', chunk_index: 0, chunk_text: 'a', chunk_source: 'compiled_truth' as const, model: null, token_count: 1, source_id: 'default', page_id: 2 }, + { slug: 'all-stale', chunk_index: 1, chunk_text: 'b', chunk_source: 'compiled_truth' as const, model: null, token_count: 1, source_id: 'default', page_id: 2 }, ]; const engine = mockEngine({ @@ -205,7 +226,9 @@ describe('runEmbedCore --dry-run never calls the embedding model', () => { // Callers wanting full coverage should call engine.getStats()/getHealth() afterward. expect(result.skipped).toBe(0); expect(result.total_chunks).toBe(3); - expect(result.pages_processed).toBe(2); // 'partial' + 'all-stale' + // v0.33.3 cherry-pick: pages_processed=0 in dry-run because we skip + // the cursor walk (countStaleChunks-only pre-flight). + expect(result.pages_processed).toBe(0); }); test('dry-run --slugs on a single page counts stale chunks, no API calls', async () => { @@ -241,10 +264,11 @@ describe('runEmbedCore --dry-run never calls the embedding model', () => { { chunk_index: 1, chunk_text: 'y', chunk_source: 'compiled_truth', embedded_at: null, token_count: 1 }, ]], ]); + // D5a: source_id + page_id required on StaleChunkRow. const stale = [ - { slug: 'a', chunk_index: 0, chunk_text: 'a', chunk_source: 'compiled_truth', model: null, token_count: 1 }, - { slug: 'b', chunk_index: 0, chunk_text: 'x', chunk_source: 'compiled_truth', model: null, token_count: 1 }, - { slug: 'b', chunk_index: 1, chunk_text: 'y', chunk_source: 'compiled_truth', model: null, token_count: 1 }, + { slug: 'a', chunk_index: 0, chunk_text: 'a', chunk_source: 'compiled_truth' as const, model: null, token_count: 1, source_id: 'default', page_id: 1 }, + { slug: 'b', chunk_index: 0, chunk_text: 'x', chunk_source: 'compiled_truth' as const, model: null, token_count: 1, source_id: 'default', page_id: 2 }, + { slug: 'b', chunk_index: 1, chunk_text: 'y', chunk_source: 'compiled_truth' as const, model: null, token_count: 1, source_id: 'default', page_id: 2 }, ]; const engine = mockEngine({ @@ -300,10 +324,11 @@ describe('runEmbedCore --stale egress fix (SQL-side filter)', () => { const { runEmbedCore } = await import('../src/commands/embed.ts'); let listPagesCalled = false; + // D5a: source_id + page_id required on StaleChunkRow. const stale = [ - { slug: 'page-a', chunk_index: 0, chunk_text: 'x', chunk_source: 'compiled_truth' as const, model: null, token_count: null }, - { slug: 'page-b', chunk_index: 1, chunk_text: 'y', chunk_source: 'compiled_truth' as const, model: null, token_count: null }, - { slug: 'page-b', chunk_index: 2, chunk_text: 'z', chunk_source: 'compiled_truth' as const, model: null, token_count: null }, + { slug: 'page-a', chunk_index: 0, chunk_text: 'x', chunk_source: 'compiled_truth' as const, model: null, token_count: null, source_id: 'default', page_id: 1 }, + { slug: 'page-b', chunk_index: 1, chunk_text: 'y', chunk_source: 'compiled_truth' as const, model: null, token_count: null, source_id: 'default', page_id: 2 }, + { slug: 'page-b', chunk_index: 2, chunk_text: 'z', chunk_source: 'compiled_truth' as const, model: null, token_count: null, source_id: 'default', page_id: 2 }, ]; // page-b has a FRESH chunk at index 0 that must be preserved through the upsert. const fullChunks: Record = { @@ -348,16 +373,17 @@ describe('runEmbedCore --stale egress fix (SQL-side filter)', () => { expect(staleChunkInUpsert.embedding).toBeInstanceOf(Float32Array); }); - test('--stale dry-run: counts stale via countStaleChunks, reports via listStaleChunks, no embedBatch or upsertChunks', async () => { + test('--stale dry-run: counts stale via countStaleChunks (no listStaleChunks call), no embedBatch or upsertChunks', async () => { const { runEmbedCore } = await import('../src/commands/embed.ts'); - const stale = [ - { slug: 'page-a', chunk_index: 0, chunk_text: 'x', chunk_source: 'compiled_truth' as const, model: null, token_count: null }, - { slug: 'page-b', chunk_index: 0, chunk_text: 'y', chunk_source: 'compiled_truth' as const, model: null, token_count: null }, - ]; + // v0.33.3 cherry-pick contract: dry-run path uses countStaleChunks + // ONLY — it does not call listStaleChunks. The pre-flight count is + // what gets reported; pages_processed stays at 0 because we + // intentionally skip the cursor walk in dry-run. + let listStaleCalled = false; const upserts: string[] = []; const engine = mockEngine({ countStaleChunks: async () => 2, - listStaleChunks: async () => stale, + listStaleChunks: async () => { listStaleCalled = true; return []; }, upsertChunks: async (slug: string) => { upserts.push(slug); }, }); @@ -366,7 +392,9 @@ describe('runEmbedCore --stale egress fix (SQL-side filter)', () => { expect(totalEmbedCalls).toBe(0); expect(upserts).toEqual([]); expect(result.would_embed).toBe(2); - expect(result.pages_processed).toBe(2); + // Cheaper dry-run: skips the cursor walk entirely. + expect(listStaleCalled).toBe(false); + expect(result.pages_processed).toBe(0); expect(result.dryRun).toBe(true); }); @@ -400,3 +428,298 @@ describe('runEmbedCore --stale egress fix (SQL-side filter)', () => { expect(result.embedded).toBe(2); }); }); + +// ──────────────────────────────────────────────────────────────── +// D5: embedBatchWithBackoff retry wrapper — 8 cases per plan +// (D2 jitter, D4 cause-unwrap, D4a maxRetries:0 passthrough, +// D8 abortSignal threading, plus the pure helpers). +// ──────────────────────────────────────────────────────────────── + +describe('embedBatchWithBackoff (D2/D4/D4a/D8)', () => { + test('case 1: parses "try again in 248ms" form and retries', async () => { + const { embedBatchWithBackoff } = await import('../src/commands/embed.ts'); + let calls = 0; + embedBatchBehavior = async () => { + calls++; + if (calls === 1) { + const err = new Error('Rate limit reached. Please try again in 50ms.'); + (err as any).cause = { status: 429 }; + throw err; + } + return [new Float32Array(1536)]; + }; + const result = await embedBatchWithBackoff(['x']); + expect(calls).toBe(2); + expect(result).toHaveLength(1); + }); + + test('case 2: parses "try again in 1.5s" form and retries', async () => { + const { embedBatchWithBackoff, parseRetryDelayMs, RATE_LIMIT_JITTER, RATE_LIMIT_PAD_MS } = await import('../src/commands/embed.ts'); + let calls = 0; + embedBatchBehavior = async () => { + calls++; + if (calls === 1) { + const err = new Error('429 — please try again in 0.05s'); + (err as any).cause = { status: 429 }; + throw err; + } + return [new Float32Array(1536)]; + }; + const result = await embedBatchWithBackoff(['x']); + expect(calls).toBe(2); + expect(result).toHaveLength(1); + // Pure-helper sanity check on the "s" form path while we're here. + const delay = parseRetryDelayMs('try again in 1.5s', () => 0.5); + // 1.5s = 1500ms + 500ms pad = 2000ms; jitter at rng=0.5 → 1.0 multiplier. + const expected = (1500 + RATE_LIMIT_PAD_MS) * (1 + (0.5 * 2 - 1) * RATE_LIMIT_JITTER); + expect(delay).toBe(Math.floor(expected)); + }); + + test('case 3: unparseable rate-limit message uses RATE_LIMIT_FALLBACK_MS', async () => { + const { parseRetryDelayMs, RATE_LIMIT_FALLBACK_MS, RATE_LIMIT_JITTER } = await import('../src/commands/embed.ts'); + // Min delay = fallback × (1 - jitter); max = fallback × (1 + jitter). + const minExpected = Math.floor(RATE_LIMIT_FALLBACK_MS * (1 - RATE_LIMIT_JITTER)); + const maxExpected = Math.floor(RATE_LIMIT_FALLBACK_MS * (1 + RATE_LIMIT_JITTER)); + for (let i = 0; i < 20; i++) { + const d = parseRetryDelayMs('429 too many requests'); + expect(d).toBeGreaterThanOrEqual(minExpected); + expect(d).toBeLessThanOrEqual(maxExpected); + } + }); + + test('case 4: non-rate-limit error rethrows immediately without retry', async () => { + const { embedBatchWithBackoff } = await import('../src/commands/embed.ts'); + let calls = 0; + embedBatchBehavior = async () => { + calls++; + throw new Error('500 internal server error'); + }; + await expect(embedBatchWithBackoff(['x'])).rejects.toThrow('500 internal server error'); + // Single attempt — no retries on non-429. + expect(calls).toBe(1); + }); + + test('case 5: jitter range — same parsed delay produces non-identical sleeps across runs', async () => { + const { parseRetryDelayMs } = await import('../src/commands/embed.ts'); + const samples = new Set(); + for (let i = 0; i < 50; i++) { + samples.add(parseRetryDelayMs('try again in 100ms')); + } + // 50 random samples with ±30% jitter should yield many distinct values. + expect(samples.size).toBeGreaterThan(5); + }); + + test('case 6: wall-clock budget mid-batch wakes the retry sleep and cancels mid-fetch', async () => { + const { embedBatchWithBackoff } = await import('../src/commands/embed.ts'); + const controller = new AbortController(); + let calls = 0; + embedBatchBehavior = async (_texts, opts) => { + calls++; + // The wrapper MUST pass the abortSignal into the gateway opts. + expect((opts as { abortSignal?: AbortSignal } | undefined)?.abortSignal).toBe(controller.signal); + if (calls === 1) { + const err = new Error('Rate limit reached. Please try again in 5000ms.'); + (err as any).cause = { status: 429 }; + throw err; + } + return [new Float32Array(1536)]; + }; + // Fire the budget abort during the retry sleep — abortableSleep should + // wake up early instead of waiting the full 5000ms. + setTimeout(() => controller.abort(), 50); + const t0 = Date.now(); + await expect(embedBatchWithBackoff(['x'], { abortSignal: controller.signal })).rejects.toThrow(); + const elapsed = Date.now() - t0; + // Should exit within ~200ms, not the 5000ms+ the retry-after would suggest. + expect(elapsed).toBeLessThan(500); + }); + + test('case 7: AITransientError-shaped wrap with 429 cause triggers retry; 500 cause does not', async () => { + const { embedBatchWithBackoff, detect429FromCause } = await import('../src/commands/embed.ts'); + + // Pure helper checks first. + expect(detect429FromCause({ cause: { status: 429 } })).toBe(true); + expect(detect429FromCause({ cause: { statusCode: 429 } })).toBe(true); + expect(detect429FromCause({ cause: { status: 500 } })).toBe(false); + expect(detect429FromCause({ status: 500 })).toBe(false); + expect(detect429FromCause(undefined)).toBe(false); + expect(detect429FromCause(null)).toBe(false); + // Deep wrap (defensive — current normalizeAIError wraps once). + expect(detect429FromCause({ cause: { cause: { status: 429 } } })).toBe(true); + + // End-to-end: 429 wrapped as AITransientError-like shape → retry. + // Use a small retry-after in the wrapper message so the parsed delay + // is fast (keeps the test under the 5s timeout). The fallback delay + // of 60s would otherwise dominate. + let calls = 0; + embedBatchBehavior = async () => { + calls++; + if (calls === 1) { + // Simulate normalizeAIError wrap: message has a parseable retry-after, + // status only on cause (the structural detection path under test). + const wrapper = new Error('try again in 10ms'); + (wrapper as any).cause = { status: 429 }; + throw wrapper; + } + return [new Float32Array(1536)]; + }; + const result = await embedBatchWithBackoff(['x']); + expect(calls).toBe(2); + expect(result).toHaveLength(1); + + // 500 wrapped → no retry, rethrow immediately. + embedBatchBehavior = async () => { + const wrapper = new Error('AI transient error'); + (wrapper as any).cause = { status: 500 }; + throw wrapper; + }; + await expect(embedBatchWithBackoff(['x'])).rejects.toThrow('AI transient error'); + }); + + test('case 8: wrapper passes maxRetries:0 through to embedBatch (no SDK retry stack)', async () => { + const { embedBatchWithBackoff } = await import('../src/commands/embed.ts'); + embedBatchBehavior = async () => [new Float32Array(1536)]; + await embedBatchWithBackoff(['x']); + expect(lastEmbedBatchOpts).toBeDefined(); + expect((lastEmbedBatchOpts as { maxRetries?: number }).maxRetries).toBe(0); + }); +}); + +// ──────────────────────────────────────────────────────────────── +// D5/D7: embedAllStale sourceId threading — invariant tests +// ──────────────────────────────────────────────────────────────── + +// ──────────────────────────────────────────────────────────────── +// Gap scan: CLI flag wiring + end-to-end budget firing (beyond plan) +// ──────────────────────────────────────────────────────────────── + +describe('runEmbed CLI flag wiring (--stale --source)', () => { + test('--source on CLI threads sourceId into countStaleChunks', async () => { + let receivedOpts: unknown; + const engine = mockEngine({ + countStaleChunks: async (opts: unknown) => { + receivedOpts = opts; + return 0; // short-circuit so we don't hit listStaleChunks + }, + }); + await runEmbed(engine, ['--stale', '--source', 'media-corpus']); + expect(receivedOpts).toEqual({ sourceId: 'media-corpus' }); + }); + + test('--stale without --source passes undefined opts (back-compat fast path)', async () => { + let receivedOpts: unknown; + const engine = mockEngine({ + countStaleChunks: async (opts: unknown) => { + receivedOpts = opts; + return 0; + }, + }); + await runEmbed(engine, ['--stale']); + expect(receivedOpts).toBeUndefined(); + }); +}); + +describe('embedAllStale wall-clock budget end-to-end (D3 + D3a)', () => { + test('GBRAIN_EMBED_TIME_BUDGET_MS=N cuts the outer loop short on stuck workers', async () => { + const { runEmbedCore } = await import('../src/commands/embed.ts'); + // Tiny budget: 100ms. Each embed call sleeps 50ms; with budget + multiple + // small batches, the second listStaleChunks call should see the abort + // signal AND the worker loop should not claim further keys. + process.env.GBRAIN_EMBED_TIME_BUDGET_MS = '100'; + process.env.GBRAIN_EMBED_CONCURRENCY = '1'; + + let listCallCount = 0; + let totalRowsReturned = 0; + // Return rows in chunks of 1 so the outer while-loop ticks frequently. + // 10 rows total across 10 "batches"; the budget should kill the loop + // partway through. + const allRows = Array.from({ length: 10 }, (_, i) => ({ + slug: `b-${i}`, + chunk_index: 0, + chunk_text: `t${i}`, + chunk_source: 'compiled_truth' as const, + model: null, + token_count: 1, + source_id: 'default', + page_id: i + 1, + })); + + const engine = mockEngine({ + countStaleChunks: async () => allRows.length, + listStaleChunks: async (opts: { afterPageId?: number } = {}) => { + listCallCount++; + const startIdx = (opts.afterPageId ?? 0); // 0 means start + const idx = allRows.findIndex(r => r.page_id > startIdx); + if (idx === -1) return []; + const row = allRows[idx]; + totalRowsReturned++; + return [row]; + }, + getChunks: async () => [], + upsertChunks: async () => {}, + }); + + // embedBatch takes 80ms per call — budget exhausts after ~1 page. + embedBatchBehavior = async (texts) => { + await new Promise(r => setTimeout(r, 80)); + return texts.map(() => new Float32Array(1536)); + }; + + const t0 = Date.now(); + const result = await runEmbedCore(engine, { stale: true }); + const elapsed = Date.now() - t0; + + // Should not have visited all 10 pages. + expect(result.pages_processed).toBeLessThan(10); + // Total wall-clock should be roughly the budget + the time for in-flight + // workers to drain (1 worker × 80ms latency). Generous upper bound: 1500ms. + expect(elapsed).toBeLessThan(1500); + }); +}); + +describe('embedAllStale --source threading (D7)', () => { + test('countStaleChunks receives the sourceId opt', async () => { + const { runEmbedCore } = await import('../src/commands/embed.ts'); + let receivedOpts: unknown; + const engine = mockEngine({ + countStaleChunks: async (opts: unknown) => { + receivedOpts = opts; + return 0; // short-circuit + }, + }); + await runEmbedCore(engine, { stale: true, sourceId: 'media-corpus' }); + expect(receivedOpts).toEqual({ sourceId: 'media-corpus' }); + }); + + test('countStaleChunks receives undefined opts when --source omitted (back-compat)', async () => { + const { runEmbedCore } = await import('../src/commands/embed.ts'); + let receivedOpts: unknown; + const engine = mockEngine({ + countStaleChunks: async (opts: unknown) => { + receivedOpts = opts; + return 0; + }, + }); + await runEmbedCore(engine, { stale: true }); + expect(receivedOpts).toBeUndefined(); + }); + + test('listStaleChunks receives the sourceId in opts when running source-scoped', async () => { + const { runEmbedCore } = await import('../src/commands/embed.ts'); + let firstCallOpts: unknown; + const stale = [ + { slug: 'p', chunk_index: 0, chunk_text: 'x', chunk_source: 'compiled_truth' as const, model: null, token_count: 1, source_id: 'media-corpus', page_id: 1 }, + ]; + const engine = mockEngine({ + countStaleChunks: async () => 1, + listStaleChunks: async (opts: unknown) => { + if (firstCallOpts === undefined) firstCallOpts = opts; + return stale; + }, + getChunks: async () => stale.map(s => ({ chunk_index: s.chunk_index, chunk_text: s.chunk_text, chunk_source: s.chunk_source, embedded_at: null, token_count: 1 })), + upsertChunks: async () => {}, + }); + await runEmbedCore(engine, { stale: true, sourceId: 'media-corpus' }); + expect((firstCallOpts as { sourceId?: string }).sourceId).toBe('media-corpus'); + }); +}); diff --git a/test/migrate.test.ts b/test/migrate.test.ts index f6a39de4d..38a7bf44f 100644 --- a/test/migrate.test.ts +++ b/test/migrate.test.ts @@ -719,6 +719,71 @@ describe('migrate — runner behavioral (v14 handler + v15 backfill)', () => { }); }); +// ──────────────────────────────────────────────────────────────────────── +// v0.33.4 D6 — migration v66 (embed_stale_partial_index) +// Mirrors v14's CONCURRENTLY + invalid-remnant pattern for the +// content_chunks partial index. Verifies both the source shape and the +// actual schema state post-migration on PGLite. +// ──────────────────────────────────────────────────────────────────────── + +describe('migrate v66 — embed_stale_partial_index (D6)', () => { + const v66 = MIGRATIONS.find(m => m.version === 66); + + test('v66 exists and uses a handler (engine-aware branching, mirrors v14)', () => { + expect(v66).toBeDefined(); + expect(v66!.name).toBe('embed_stale_partial_index'); + expect(typeof v66!.handler).toBe('function'); + expect(v66!.sql).toBe(''); + }); + + test('v66 handler source: CONCURRENTLY + invalid-index cleanup on Postgres branch', async () => { + const { readFileSync } = await import('fs'); + const src = readFileSync('src/core/migrate.ts', 'utf-8'); + const v66Start = src.indexOf("name: 'embed_stale_partial_index'"); + expect(v66Start).toBeGreaterThan(-1); + const v66Block = src.slice(v66Start, v66Start + 3000); + expect(v66Block).toContain('pg_index'); + expect(v66Block).toContain('indisvalid'); + expect(v66Block).toContain('DROP INDEX CONCURRENTLY IF EXISTS idx_chunks_embedding_null'); + expect(v66Block).toContain('CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_chunks_embedding_null'); + // Partial index predicate must match the production query in + // postgres-engine.ts / pglite-engine.ts: `WHERE embedding IS NULL`. + expect(v66Block).toContain('WHERE embedding IS NULL'); + // DROP IF EXISTS must precede CREATE IF NOT EXISTS so a failed prior + // CONCURRENTLY build is cleaned before re-create. + const dropIdx = v66Block.indexOf('DROP INDEX CONCURRENTLY IF EXISTS'); + const createIdx = v66Block.indexOf('CREATE INDEX CONCURRENTLY IF NOT EXISTS'); + expect(dropIdx).toBeLessThan(createIdx); + // Branches on engine.kind (handler-pattern from v14). + expect(v66Block).toContain('engine.kind'); + }); + + test('v66 idempotent flag is true (re-run safety)', () => { + expect(v66!.idempotent).toBe(true); + }); +}); + +describe('migrate runner v66 — partial index materialized on PGLite', () => { + let engine: PGLiteEngine; + + beforeAll(async () => { + engine = new PGLiteEngine(); + await engine.connect({}); + await engine.initSchema(); + }); + + afterAll(async () => { + await engine.disconnect(); + }); + + test('v66 created idx_chunks_embedding_null on PGLite via handler branch', async () => { + const rows = await (engine as any).db.query( + `SELECT indexname FROM pg_indexes WHERE indexname = 'idx_chunks_embedding_null'` + ); + expect(rows.rows.length).toBe(1); + }); +}); + describe('migrate: v8 (links_dedup) regression — must be fast on 1K duplicate rows', () => { let engine: PGLiteEngine; diff --git a/test/pglite-engine.test.ts b/test/pglite-engine.test.ts index 0663fd8be..19d87146b 100644 --- a/test/pglite-engine.test.ts +++ b/test/pglite-engine.test.ts @@ -392,6 +392,135 @@ describe('PGLiteEngine: Chunks', () => { }); }); +// ──────────────────────────────────────────────────────────────────────── +// v0.33.4 D7 + IRON RULE — countStaleChunks + listStaleChunks contract +// PGLite parity for the Postgres E2E in test/e2e/embed-stale-pagination. +// Pins the tuple-compare `(cc.page_id, cc.chunk_index) > ($1, $2)` against +// the WASM build (Postgres 17.5 in WASM has had quirks here historically). +// ──────────────────────────────────────────────────────────────────────── + +describe('PGLiteEngine: stale chunk pagination (D7 + REGRESSION)', () => { + beforeEach(truncateAll); + + test('countStaleChunks: zero-state baseline', async () => { + expect(await engine.countStaleChunks()).toBe(0); + }); + + test('countStaleChunks counts chunks with NULL embedding only', async () => { + await engine.putPage('test/stale-a', testPage); + await engine.upsertChunks('test/stale-a', [ + { chunk_index: 0, chunk_text: 'no embed', chunk_source: 'compiled_truth' }, + { chunk_index: 1, chunk_text: 'has embed', chunk_source: 'compiled_truth', embedding: new Float32Array(1536).fill(0.1) }, + ]); + expect(await engine.countStaleChunks()).toBe(1); + }); + + test('listStaleChunks: cursor pagination across page boundaries', async () => { + // Seed 3 pages × 3 chunks = 9 stale rows; walk with batchSize=2. + for (const slug of ['c-a', 'c-b', 'c-c']) { + await engine.putPage(`test/${slug}`, testPage); + await engine.upsertChunks(`test/${slug}`, [ + { chunk_index: 0, chunk_text: 'a', chunk_source: 'compiled_truth' }, + { chunk_index: 1, chunk_text: 'b', chunk_source: 'compiled_truth' }, + { chunk_index: 2, chunk_text: 'c', chunk_source: 'compiled_truth' }, + ]); + } + const visited = new Set(); + let after_pid = 0; + let after_idx = -1; + let lastPid = -1; + let lastIdx = -1; + let cursorMonotonic = true; + // eslint-disable-next-line no-constant-condition + while (true) { + const batch = await engine.listStaleChunks({ + batchSize: 2, + afterPageId: after_pid, + afterChunkIndex: after_idx, + }); + if (batch.length === 0) break; + for (const row of batch) { + const key = `${row.page_id}::${row.chunk_index}`; + expect(visited.has(key)).toBe(false); + visited.add(key); + const advance = row.page_id > lastPid + || (row.page_id === lastPid && row.chunk_index > lastIdx); + if (!advance) cursorMonotonic = false; + lastPid = row.page_id; + lastIdx = row.chunk_index; + } + const tail = batch[batch.length - 1]; + after_pid = tail.page_id; + after_idx = tail.chunk_index; + if (batch.length < 2) break; + } + expect(visited.size).toBe(9); + expect(cursorMonotonic).toBe(true); + }); + + test('listStaleChunks: page split across batches (1 page, 5 chunks, batchSize=2)', async () => { + await engine.putPage('test/split', testPage); + await engine.upsertChunks('test/split', [ + { chunk_index: 0, chunk_text: 'a', chunk_source: 'compiled_truth' }, + { chunk_index: 1, chunk_text: 'b', chunk_source: 'compiled_truth' }, + { chunk_index: 2, chunk_text: 'c', chunk_source: 'compiled_truth' }, + { chunk_index: 3, chunk_text: 'd', chunk_source: 'compiled_truth' }, + { chunk_index: 4, chunk_text: 'e', chunk_source: 'compiled_truth' }, + ]); + const collected: number[] = []; + let after_pid = 0; + let after_idx = -1; + // eslint-disable-next-line no-constant-condition + while (true) { + const batch = await engine.listStaleChunks({ + batchSize: 2, + afterPageId: after_pid, + afterChunkIndex: after_idx, + }); + if (batch.length === 0) break; + for (const r of batch) collected.push(r.chunk_index); + const tail = batch[batch.length - 1]; + after_pid = tail.page_id; + after_idx = tail.chunk_index; + if (batch.length < 2) break; + } + expect(collected).toEqual([0, 1, 2, 3, 4]); + }); + + test('countStaleChunks + listStaleChunks honor sourceId filter (D7)', async () => { + // PGLite default seed has 'default' source. Add 'other-source' and + // seed identical slugs in both. + await engine.executeRaw( + `INSERT INTO sources (id, name, local_path) VALUES ('other', 'other', '/tmp/other') ON CONFLICT (id) DO NOTHING`, + ); + // Default source page via the engine API. + await engine.putPage('test/shared', testPage, { sourceId: 'default' }); + await engine.upsertChunks('test/shared', [ + { chunk_index: 0, chunk_text: 'default-0', chunk_source: 'compiled_truth' }, + { chunk_index: 1, chunk_text: 'default-1', chunk_source: 'compiled_truth' }, + ], { sourceId: 'default' }); + // Same slug, different source. + await engine.putPage('test/shared', testPage, { sourceId: 'other' }); + await engine.upsertChunks('test/shared', [ + { chunk_index: 0, chunk_text: 'other-0', chunk_source: 'compiled_truth' }, + { chunk_index: 1, chunk_text: 'other-1', chunk_source: 'compiled_truth' }, + { chunk_index: 2, chunk_text: 'other-2', chunk_source: 'compiled_truth' }, + ], { sourceId: 'other' }); + + expect(await engine.countStaleChunks()).toBe(5); + expect(await engine.countStaleChunks({ sourceId: 'default' })).toBe(2); + expect(await engine.countStaleChunks({ sourceId: 'other' })).toBe(3); + + const defaultRows = await engine.listStaleChunks({ sourceId: 'default', batchSize: 100 }); + expect(defaultRows).toHaveLength(2); + for (const r of defaultRows) expect(r.source_id).toBe('default'); + + const otherRows = await engine.listStaleChunks({ sourceId: 'other', batchSize: 100 }); + expect(otherRows).toHaveLength(3); + for (const r of otherRows) expect(r.source_id).toBe('other'); + }); +}); + // ───────────────────────────────────────────────────────────────── // Links + Graph // ─────────────────────────────────────────────────────────────────