mirror of
https://github.com/garrytan/gbrain.git
synced 2026-07-27 22:15:33 +00:00
v0.34.4.0 fix(embed): cursor-paginated --stale hardening wave (D2/D3/D4/D6/D7/D8 + regression test) (#991)
* perf(embed): cursor-paginated stale loading + rate-limit backoff + partial index Three fixes for embed --stale on large brains (300K+ chunks): ## 1. Cursor-paginated listStaleChunks (embed timeout fix) The previous implementation pulled ALL stale rows (up to 100K) in one query. On a 373K-row content_chunks table with 48K stale rows, this query took >2 min and hit Supabase's 2-min statement_timeout, causing embed --stale to silently fail with zero progress. Fix: keyset pagination on (page_id, chunk_index) with a default batch size of 2000 rows. Each query finishes in <1s. The embedAllStale loop pages through batches, embeds each batch, then advances the cursor. ## 2. Rate-limit-aware retry (429 backoff) The OpenAI SDK's built-in retry has a ~4s max backoff window, which is too short for TPM (tokens-per-minute) limits on large pages (~90K tokens). The embed loop would fail after 3 SDK retries and skip the page entirely. Fix: embedBatchWithBackoff wrapper parses the retry delay from the 429 error message (e.g. 'try again in 248ms') and sleeps for that duration + 500ms padding. Up to 5 retries with parsed delays (60s fallback when unparseable). ## 3. Migration v58: partial index for NULL embeddings `CREATE INDEX idx_chunks_embedding_null ON content_chunks (page_id, chunk_index) WHERE embedding IS NULL` — makes countStaleChunks() and the paginated listStaleChunks() instant instead of full-table-scanning 373K rows. ## Testing Verified on a 99K-page / 373K-chunk brain with 48K stale chunks. Before: embed --stale hung for 2+ min then timed out (0 progress). After: loads 2K rows in <1s, embeds concurrently, pages through all stale chunks without timeout. * fix(embed): wave of hardening + tests on cursor-paginated --stale path Lands the 9 decisions + regression test set from /plan-eng-review on PR #991's embed-perf cherry-pick. Implements the codex outside-voice findings folded in during plan review. Architecture / correctness: - D2 jitter on the parsed retry-after delay (±30%) so 20 concurrent workers don't relock on the next 429 wave (thundering herd fix). - D3 + D3a + D8 wall-clock budget (GBRAIN_EMBED_TIME_BUDGET_MS, default 30 min) threaded as an AbortSignal into THREE places: the retry sleep (abortableSleep), the per-key worker claim loop, and the gateway embed call itself (so a worker mid-fetch on a ~30s OpenAI HTTP timeout cancels within seconds instead of waiting it out). - D4 structured 429 detection that unwraps the gateway's AITransientError wrap via cause chain (depth-limited to 5). Naive `e.status === 429` was silently false against normalized errors; message-match stays as fallback. detect429FromCause exported as @internal helper. - D4a `maxRetries: 0` passthrough through embedBatch → gateway → embedMany so the AI SDK's default 2-retry stack doesn't multiply this wrapper's 5 attempts (was up to 15 total cycles per call). - D6 migration v59 (embed_stale_partial_index) rewritten to use CREATE INDEX CONCURRENTLY + handler-based engine-branching (mirrors v14 invalid-remnant pattern). Plain CREATE INDEX would have taken ShareLock on the 373K-row content_chunks table for the duration of the build. - D7 sourceId threaded through countStaleChunks + listStaleChunks + embedAllStale. `gbrain embed --stale --source X` was silently dropping the flag pre-fix and counting/embedding across every source. Both Postgres and PGLite engines updated. Tests added: - D5 8 unit cases for embedBatchWithBackoff in test/embed.serial.test.ts: ms / s retry-after parse, fallback, non-rate-limit rethrow, jitter variance, budget abort during sleep+fetch, normalized-error cause unwrap, maxRetries:0 passthrough verification. - D5a fixed every pre-existing stale-row mock to include source_id + page_id (required on StaleChunkRow as of v0.33.3 cursor pagination — TypeScript's structural typing was hiding these). - D7 unit cases asserting CLI `--source X` parses + threads sourceId. - Gap scan: end-to-end wall-clock budget firing in the outer pagination loop via runEmbedCore. - D6 migration v59 test cases in test/migrate.test.ts: source-shape assertion (CONCURRENTLY + invalid-remnant DROP-before-CREATE ordering), PGLite handler-branch idempotency, partial-index materialization. - REGRESSION: new test/e2e/embed-stale-pagination.test.ts covering static (every chunk visited exactly once), failed-page (cursor advances past failures, next run picks up), page-split-across-batches, source-scoped scan, duplicate-slug-across-sources. - PGLite parity cases for cursor pagination, page split, source filter in test/pglite-engine.test.ts (pins tuple-compare against WASM build). Gate: - bun run test: 6305 pass / 0 fail / 0 skip across all 8 shards + serial. - DATABASE_URL=... bun run test:e2e: 90 files, 603 tests, 0 failures. Plan: ~/.claude/plans/system-instruction-you-are-working-iterative-torvalds.md Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * chore: bump version and changelog (v0.34.3.0) Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> --------- Co-authored-by: garrytan-agents <garrytan-agents@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Garry Tan
co-authored by
garrytan-agents
Claude Opus 4.7
garrytan-agents
Claude Opus 4.7
parent
668254b05e
commit
24881f60fc
@@ -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.**
|
||||
|
||||
@@ -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.
|
||||
|
||||
+1
-1
@@ -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",
|
||||
|
||||
+251
-77
@@ -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<string, typeof staleRows>();
|
||||
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<number, Float32Array>();
|
||||
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<string, typeof batch>();
|
||||
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<number, Float32Array>();
|
||||
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<void> {
|
||||
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<Float32Array[]> {
|
||||
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);
|
||||
}
|
||||
|
||||
+37
-4
@@ -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<Float32Array[]> {
|
||||
/**
|
||||
* 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<Float32Array[]> {
|
||||
if (!texts || texts.length === 0) return [];
|
||||
|
||||
const cfg = requireConfig();
|
||||
@@ -807,7 +834,7 @@ export async function embed(texts: string[]): Promise<Float32Array[]> {
|
||||
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<Float32Array[]> {
|
||||
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})`);
|
||||
|
||||
+23
-2
@@ -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<Float32Array[]> {
|
||||
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);
|
||||
}
|
||||
|
||||
+23
-5
@@ -533,20 +533,38 @@ export interface BrainEngine {
|
||||
*/
|
||||
getChunks(slug: string, opts?: { sourceId?: string }): Promise<Chunk[]>;
|
||||
/**
|
||||
* 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<number>;
|
||||
countStaleChunks(opts?: { sourceId?: string }): Promise<number>;
|
||||
/**
|
||||
* 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<StaleChunkRow[]>;
|
||||
listStaleChunks(opts?: {
|
||||
batchSize?: number;
|
||||
afterPageId?: number;
|
||||
afterChunkIndex?: number;
|
||||
sourceId?: string;
|
||||
}): Promise<StaleChunkRow[]>;
|
||||
/**
|
||||
* Delete every chunk for a page. Internal page-id lookup is sourceId-scoped
|
||||
* when `opts.sourceId` is given; otherwise the bare-slug subquery returns
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -1309,25 +1309,65 @@ export class PGLiteEngine implements BrainEngine {
|
||||
return (rows as Record<string, unknown>[]).map(r => rowToChunk(r));
|
||||
}
|
||||
|
||||
async countStaleChunks(): Promise<number> {
|
||||
async countStaleChunks(opts?: { sourceId?: string }): Promise<number> {
|
||||
// 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<StaleChunkRow[]> {
|
||||
async listStaleChunks(opts?: {
|
||||
batchSize?: number;
|
||||
afterPageId?: number;
|
||||
afterChunkIndex?: number;
|
||||
sourceId?: string;
|
||||
}): Promise<StaleChunkRow[]> {
|
||||
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[];
|
||||
}
|
||||
|
||||
@@ -1308,26 +1308,70 @@ export class PostgresEngine implements BrainEngine {
|
||||
return rows.map((r) => rowToChunk(r as Record<string, unknown>));
|
||||
}
|
||||
|
||||
async countStaleChunks(): Promise<number> {
|
||||
async countStaleChunks(opts?: { sourceId?: string }): Promise<number> {
|
||||
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<StaleChunkRow[]> {
|
||||
async listStaleChunks(opts?: {
|
||||
batchSize?: number;
|
||||
afterPageId?: number;
|
||||
afterChunkIndex?: number;
|
||||
sourceId?: string;
|
||||
}): Promise<StaleChunkRow[]> {
|
||||
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[];
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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<string, number>();
|
||||
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<boolean> {
|
||||
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<string>();
|
||||
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<string>();
|
||||
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]);
|
||||
});
|
||||
});
|
||||
+350
-27
@@ -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<Float32Array[]>) | 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<string, any[]> = {
|
||||
@@ -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<number>();
|
||||
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 <id> 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');
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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<string>();
|
||||
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
|
||||
// ─────────────────────────────────────────────────────────────────
|
||||
|
||||
Reference in New Issue
Block a user