From d42b25bee060401235a4a134ee03283c20acad86 Mon Sep 17 00:00:00 2001 From: Garry Tan Date: Sat, 8 Aug 2026 19:51:51 -0700 Subject: [PATCH] =?UTF-8?q?fix(review):=20wave-4=20batch=20=E2=80=94=20tak?= =?UTF-8?q?es=20join=20the=20dim-transition=20set,=20dual-plane=20think-ta?= =?UTF-8?q?ke=20append,=20source-scoped=20embed=20lane,=20claim=20bounds?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two criticals: 1. Dim-transition hole: `takes` was missing from TEXT_EMBEDDING_DIM_PINNED_TABLES, so migrating embeddings to a different dim left takes.embedding at the old width — the embed takes lane then failed every write forever while still paying the gateway each run. takes now transitions with query_cache + facts (partial HNSW index recreated, #1734 dim-cap gate honored), and invalidateStaleSignatureEmbeddings re-stales takes on drifted pages in BOTH engines (same page-signature predicate + source scope as the chunks sweep) so a same-dim provider swap can't score old-space take vectors against new-space queries. 2. Fence-plane clobber: persistThinkTake wrote DB-only rows while `takes add` allocates row numbers from the markdown fence — the next fence-allocated row at the same number deterministically OVERWROTE the think-take via ON CONFLICT (page_id, row_num) DO UPDATE. Both writers now share one dual-plane helper (src/core/takes-append.ts): withPageLock → scoped page lookup → fence allocates the row number (file-first) → DB mirror. Headless/DB-only brains fall back to MAX(row_num)+1 with the TAKE_FILE_PLANE_UNAVAILABLE warning. Also in this batch: countStaleTakes/listStaleTakes accept a sourceId scope (page-join filter) and the embed takes lane threads it so `embed --stale --source X` touches only X's takes; think-take claims are flattened to one line and bounded at THINK_TAKE_CLAIM_MAX_CHARS (2000) with TAKE_CLAIM_TRUNCATED; updateTakeEmbeddingsBatch routes through batchRetry with its own BATCH_AUDIT_SITES entry (both engines); the think op's remote gate uses the fail-closed `ctx.remote !== false` pattern; the takes-lane failure line reports written vs skipped (row retired mid-flight) honestly; and `gbrain think --take` resolves the caller's ambient source (env/dotfile/local_path/brain-default tiers) so duplicate anchor slugs land on the caller's-context page. Includes the 5 wave-4 coverage-audit test files, new pins for the dim transition + takes invalidation + lane scoping + claim bound, and the "Think-ops follow-ups (Wave 4, 2026-08)" TODO section. Co-Authored-By: Claude Fable 5 --- TODOS.md | 54 ++++++ src/commands/embed.ts | 30 ++- src/commands/takes.ts | 74 ++++---- src/commands/think.ts | 22 ++- src/core/engine.ts | 26 ++- src/core/operations.ts | 2 +- src/core/pglite-engine.ts | 47 ++++- src/core/postgres-engine.ts | 39 +++- src/core/retrieval-upgrade-planner.ts | 29 ++- src/core/retry.ts | 1 + src/core/takes-append.ts | 175 ++++++++++++++++++ src/core/think/index.ts | 85 +++++---- test/core/retry.test.ts | 2 +- test/e2e/migrate-embeddings-postgres.test.ts | 10 +- test/embed-takes-lane.serial.test.ts | 170 +++++++++++++++++ test/embedding-migration.test.ts | 48 +++++ test/migrations-takes-embedding-rerun.test.ts | 82 ++++++++ test/retrieval-upgrade-planner.test.ts | 65 +++++++ test/takes-command-source-scope.test.ts | 10 + test/think-pipeline.serial.test.ts | 32 +++- test/think-take-cli.serial.test.ts | 81 ++++++++ test/think-take-concurrency.serial.test.ts | 83 +++++++++ test/think-take-op-federated.test.ts | 129 +++++++++++++ 23 files changed, 1186 insertions(+), 110 deletions(-) create mode 100644 src/core/takes-append.ts create mode 100644 test/embed-takes-lane.serial.test.ts create mode 100644 test/migrations-takes-embedding-rerun.test.ts create mode 100644 test/think-take-cli.serial.test.ts create mode 100644 test/think-take-concurrency.serial.test.ts create mode 100644 test/think-take-op-federated.test.ts diff --git a/TODOS.md b/TODOS.md index 34ff1362b..c890f6352 100644 --- a/TODOS.md +++ b/TODOS.md @@ -4886,3 +4886,57 @@ respective shapes. Small, mechanical; pinned by `test/init-embed-check.test.ts` + the models-doctor tests. **Depends on:** nothing. + +## Think-ops follow-ups (Wave 4, 2026-08) + +### Next-take-row allocation helper (unify 4 hand-rolled copies) +**Priority:** P3 + +**What:** the `MAX(row_num)+1` next-row allocation for takes is hand-rolled in +four places with inconsistent locking: `src/core/takes-append.ts` (DB-only +fallback, under withPageLock), both engines' `supersedeTake` (in-SQL), and +`src/core/consolidate.ts`. Extract one shared helper with a consistent locking +story so a fifth writer can't invent a fifth allocation. + +**Why:** every divergent copy is a future duplicate-row or clobber bug of the +exact class Wave 4 fixed (think-take rows overwritten by fence-allocated rows). + +**How to start:** put the allocator next to `appendTake` in +`src/core/takes-append.ts`; migrate call sites one at a time, pinning each with +the existing takes-engine / consolidate tests. + +### Takes-lane terminal-state warning parity + failure-path tests +**Priority:** P3 + +**What:** the chunks lane warns loudly when a catch-up run finishes with +persistently-failing chunks (embed.ts's post-loop `countStaleChunks` probe); +the takes lane has no equivalent — the same takes can fail every run silently +(deferral note sits at the takes-lane batch loop in `src/commands/embed.ts`). +Add the mirror warning plus failure-path tests: batch-throw (embedPageTexts +rejects) and probe-throw (countStaleTakes rejects) paths. + +### CLI take success-path output test +**Priority:** P3 + +**What:** `test/think-take-cli.serial.test.ts` covers the two exit(1) branches +but not the success path (`Take: appended row #N ...` line + exit 0). The +review provided a chat-transport seam stub (see +`test/think-take-op-federated.test.ts`'s `__setChatTransportForTests` canned +envelope) — drive `runThinkCli` with it and assert the success output. + +### v126 HNSW dim-cap branch test (3072d) + real-PG run +**Priority:** P3 + +**What:** migration v126's `hnswIndexExpected` gate (skip HNSW above the #1734 +cap) has no test at a >cap dim (e.g. 3072d: column rebuilt, index deliberately +absent, exact scan still works). Add the PGLite test plus a DATABASE_URL-gated +real-PG run — PGLite cannot surface real pgvector index-build failures. + +### Federated think --take write-policy note +**Priority:** P4 + +**What:** a take synthesized under a federated grant [A,B] persists into the +ANCHOR page's source — documented behavior (the anchor lookup is +grant-confined; the write inherits the anchor's source). Revisit if per-source +take provenance ever lands; a grant-scoped caller writing brain-holder takes +into a mounted source may want an explicit provenance column instead. diff --git a/src/commands/embed.ts b/src/commands/embed.ts index fb5d6a20e..39b7312a6 100644 --- a/src/commands/embed.ts +++ b/src/commands/embed.ts @@ -440,14 +440,18 @@ export async function runEmbedCore(engine: BrainEngine, opts: EmbedOpts): Promis // budget hit (or caller abort) inside the pages lane skips takes // gracefully — the next run picks them up via the NULL-embedding // predicate. Inside this try so the single-flight locks are still held - // (a takes backfill is subject to the same per-source mutual exclusion) - // and the pacer telemetry in the finally covers takes DB writes too. + // — and the lane threads opts.sourceId into its enumerators, so a + // `--source X` run touches ONLY X's takes: the work is confined to the + // same source(s) whose locks this run actually holds (per-source + // mutual exclusion holds for the takes lane too). Pacer telemetry in + // the finally covers takes DB writes as well. if (!laneAborted && !isAborted(opts.signal)) { await embedStaleTakes(engine, result, { dryRun: !!opts.dryRun, pacer, signal: opts.signal, quiet: opts.quiet, + sourceId: opts.sourceId, }); } } finally { @@ -1403,13 +1407,16 @@ export const TAKES_EMBED_BATCH_SIZE = 64; async function embedStaleTakes( engine: BrainEngine, result: EmbedResult, - opts: { dryRun: boolean; pacer: DbPacer; signal?: AbortSignal; quiet?: boolean }, + opts: { dryRun: boolean; pacer: DbPacer; signal?: AbortSignal; quiet?: boolean; sourceId?: string }, ): Promise { + // Source scope: thread the run's --source through both enumerators so a + // scoped run never enumerates (or pays to embed) another source's takes. + const scope = opts.sourceId !== undefined ? { sourceId: opts.sourceId } : undefined; let staleCount: number; try { // Defensive Number(): int8 counts can arrive as BigInt/string from a // driver; a null/undefined (partial engine double in tests) reads as 0. - staleCount = Number(await engine.countStaleTakes()) || 0; + staleCount = Number(await engine.countStaleTakes(scope)) || 0; } catch (e: unknown) { // Pre-v37 brains (or a wedged probe) must not fail the whole embed run // over the optional takes lane. @@ -1424,7 +1431,7 @@ async function embedStaleTakes( return; } - const staleTakes = (await engine.listStaleTakes()) ?? []; + const staleTakes = (await engine.listStaleTakes(scope)) ?? []; if (staleTakes.length === 0) return; result.takes_embedded ??= 0; @@ -1451,8 +1458,19 @@ async function embedStaleTakes( result.takes_embedded += updated; if (failed > 0) { recordFailure(result, failed, `takes:${batch[0]?.page_slug ?? '?'}`, firstError); - serr(` [takes] ${failed} claim(s) failed to embed in batch; embedded the other ${updated}`); + // `updated` is the RETURNING-based count from the writer; rows the + // writer skipped (superseded mid-flight, `AND active` guard) are + // called out separately so the arithmetic is honest. + const skipped = writes.length - updated; + serr( + ` [takes] ${failed} claim(s) failed to embed in batch; ${updated} written` + + (skipped > 0 ? `, ${skipped} skipped (row retired mid-flight)` : ''), + ); } + // NOTE: unlike the chunks lane's catch-up terminal-state warning + // (see the countStaleChunks probe after the pages loop), persistently + // failing takes are not yet flagged as stuck — deferred, see + // TODOS.md "Think-ops follow-ups (Wave 4, 2026-08)". } catch (e: unknown) { if (isAborted(opts.signal)) break; // shutdown, not a failure recordFailure(result, batch.length, `takes:${batch[0]?.page_slug ?? '?'}`, e); diff --git a/src/commands/takes.ts b/src/commands/takes.ts index 8d7eca1cc..13e346d2d 100644 --- a/src/commands/takes.ts +++ b/src/commands/takes.ts @@ -19,16 +19,22 @@ * 6. releases the lock (auto via withPageLock) */ -import { readFileSync, writeFileSync, existsSync, mkdirSync } from 'node:fs'; -import { join, dirname } from 'node:path'; +import { existsSync } from 'node:fs'; import type { BrainEngine, TakeKind } from '../core/engine.ts'; import { parseTakesFence, - upsertTakeRow, supersedeRow, type ParsedTake, } from '../core/takes-fence.ts'; import { withPageLock } from '../core/page-lock.ts'; +import { + appendTake, + TakePageNotFoundError, + pageFilePath, + readBodyOrEmpty, + writeBody, + resolveBrainDirOrNull, +} from '../core/takes-append.ts'; import { resolveSourceId } from '../core/source-resolver.ts'; import { resolveOwnerHolder } from '../core/owner-holder.ts'; @@ -44,26 +50,20 @@ function flagPresent(args: string[], name: string): boolean { return args.includes(name); } +// CLI-strict wrapper around the shared resolveBrainDirOrNull (takes-append.ts): +// unresolvable → loud error + exit(1). Core callers (persistThinkTake) use the +// null-returning form and fall back to DB-only allocation instead. async function resolveBrainDir(engine: BrainEngine | null, explicitDir: string | null): Promise { - if (explicitDir) { - if (!existsSync(explicitDir)) { - console.error(`--dir path does not exist: ${explicitDir}`); - process.exit(1); - } - return explicitDir; - } - if (engine) { - const configured = await engine.getConfig('sync.repo_path'); - if (configured && existsSync(configured)) return configured; + if (explicitDir && !existsSync(explicitDir)) { + console.error(`--dir path does not exist: ${explicitDir}`); + process.exit(1); } + const dir = engine ? await resolveBrainDirOrNull(engine, explicitDir) : (explicitDir ?? null); + if (dir) return dir; console.error('No brain directory configured. Pass --dir or run `gbrain init` first.'); process.exit(1); } -function pageFilePath(brainDir: string, slug: string): string { - return join(brainDir, `${slug}.md`); -} - function ensureKind(raw: string | undefined): TakeKind { if (!raw) { console.error('Missing --kind. Expected one of: fact, take, bet, hunch.'); @@ -117,15 +117,8 @@ async function resolveTakesSourceId(engine: BrainEngine): Promise { return resolveSourceId(engine, null); } -function readBodyOrEmpty(path: string): string { - if (!existsSync(path)) return ''; - return readFileSync(path, 'utf-8'); -} - -function writeBody(path: string, body: string): void { - mkdirSync(dirname(path), { recursive: true }); - writeFileSync(path, body, 'utf-8'); -} +// readBodyOrEmpty / writeBody / pageFilePath now live in +// src/core/takes-append.ts (shared with the dual-plane appendTake helper). // --- Subcommands --- @@ -209,22 +202,23 @@ async function cmdAdd(engine: BrainEngine, args: string[], sourceId?: string): P const dirArg = flagValue(args, '--dir'); const brainDir = await resolveBrainDir(engine, dirArg ?? null); - await withPageLock(slug, async () => { - const path = pageFilePath(brainDir, slug); - const body = readBodyOrEmpty(path); - const { body: nextBody, rowNum } = upsertTakeRow(body, { - claim, kind, holder, weight, source, sinceDate: since, active: true, + // Dual-plane append via the shared core helper (takes-append.ts): page + // lock → scoped page lookup → fence allocates the row number (file-first) + // → DB mirror. Same path persistThinkTake uses, so the planes can't drift. + try { + const { rowNum } = await appendTake(engine, { + slug, sourceId, claim, kind, holder, weight, + source, sinceDate: since, brainDir, }); - writeBody(path, nextBody); - - // Mirror to DB. Page may not be in DB yet if not synced — caller must run sync first. - const pageId = await getPageId(engine, slug, sourceId); - await engine.addTakesBatch([{ - page_id: pageId, row_num: rowNum, claim, kind, holder, weight, - since_date: since, source, active: true, superseded_by: null, - }]); console.log(`Added take #${rowNum} to ${slug}.`); - }); + } catch (e) { + if (e instanceof TakePageNotFoundError) { + // Page may not be in DB yet if not synced — caller must run sync first. + console.error(`Page not found in brain: ${slug}${sourceId ? ` (source=${sourceId})` : ''}. Run \`gbrain sync\` first.`); + process.exit(1); + } + throw e; + } } async function cmdUpdate(engine: BrainEngine, args: string[], sourceId?: string): Promise { diff --git a/src/commands/think.ts b/src/commands/think.ts index 906176806..1576c62b1 100644 --- a/src/commands/think.ts +++ b/src/commands/think.ts @@ -10,6 +10,7 @@ import { runThink, persistSynthesis, persistThinkTake, stripGapsSection } from ' import { loadConfig, isThinClient } from '../core/config.ts'; import { callRemoteTool, unpackToolResult } from '../core/mcp-client.ts'; import { canonicalLookup } from '../core/model-pricing.ts'; +import { resolveSourceWithTier, ALL_SOURCES } from '../core/source-resolver.ts'; function flagValue(args: string[], name: string): string | undefined { const i = args.indexOf(name); @@ -166,7 +167,26 @@ prints what would have been the input (exit 0). // nothing consumed it). Same honesty contract as --save: persist or // exit non-zero — an explicit --take that writes nothing must be loud. if (take) { - const persistedTake = await persistThinkTake(engine, result, { anchor }); + // Wave-4: resolve the caller's ambient source context (GBRAIN_SOURCE / + // .gbrain-source dotfile / local_path / brain default — the same + // resolveSourceId chain every other CLI surface uses) so a duplicate + // anchor slug across sources lands on the CALLER's-context page, not + // a planner-chosen one. The seed_default tier means NOTHING was + // configured anywhere — keep the historical unscoped posture there + // (undefined → unscoped getPage), and __all__ has span-everything + // semantics, which for a single-page lookup is also unscoped. + // Resolution errors (e.g. GBRAIN_SOURCE naming an unregistered + // source) propagate to the catch below — fail-closed, matching the + // `gbrain takes` posture. + const resolvedSource = await resolveSourceWithTier(engine, undefined); + const takeSourceId = + resolvedSource.tier === 'seed_default' || resolvedSource.source_id === ALL_SOURCES + ? undefined + : resolvedSource.source_id; + const persistedTake = await persistThinkTake(engine, result, { + anchor, + ...(takeSourceId !== undefined ? { sourceId: takeSourceId } : {}), + }); takeRow = persistedTake.rowNum; takeInserted = persistedTake.inserted; for (const w of persistedTake.warnings) result.warnings.push(w); diff --git a/src/core/engine.ts b/src/core/engine.ts index f83b98355..f5efe5fb6 100644 --- a/src/core/engine.ts +++ b/src/core/engine.ts @@ -1050,6 +1050,11 @@ export interface BrainEngine { * default grandfather clause would silently keep them mixed into the new * index. `gbrain migrate embeddings` and `embed --stale * --include-null-signature` set this. + * + * ALSO re-stales `takes` rows (embedding + embedded_at → NULL) on the same + * pages, under the SAME signature predicate + source scope — takes carry + * text-embedding-space vectors too, and a same-dim provider swap has no + * schema transition to drop them. The returned count remains chunks-only. */ invalidateStaleSignatureEmbeddings(opts: { signature: string; sourceId?: string; includeNullSignature?: boolean }): Promise; /** @@ -1552,11 +1557,18 @@ export interface BrainEngine { /** Look up embeddings by take id (mirrors getEmbeddingsByChunkIds). */ getTakeEmbeddings(ids: number[]): Promise>; - /** Pre-flight count for `gbrain embed --stale`. WHERE active AND embedding IS NULL. */ - countStaleTakes(): Promise; + /** + * Pre-flight count for `gbrain embed --stale`. WHERE active AND embedding + * IS NULL. `opts.sourceId` scopes via the page join (p.source_id) so a + * `--source X` run only counts X's takes. + */ + countStaleTakes(opts?: { sourceId?: string }): Promise; - /** List stale takes (no embedding column in payload — same pattern as listStaleChunks). */ - listStaleTakes(): Promise; + /** + * List stale takes (no embedding column in payload — same pattern as + * listStaleChunks). `opts.sourceId` scopes via the page join. + */ + listStaleTakes(opts?: { sourceId?: string }): Promise; /** * #2089: batch-write take embeddings — the writer half of the takes @@ -1566,8 +1578,12 @@ export interface BrainEngine { * row via supersedeTake, and the retired row must not receive the * vector. Returns the number of rows actually updated (superseded / * deleted ids are silently skipped). Empty input → 0 without a query. + * + * Wrapped in `batchRetry` like the other batch primitives (addTakesBatch + * precedent), so `opts` (auditSite, AbortSignal) is honored; same + * no-double-wrap contract as `BatchOpts`. */ - updateTakeEmbeddingsBatch(rows: Array<{ take_id: number; embedding: Float32Array }>): Promise; + updateTakeEmbeddingsBatch(rows: Array<{ take_id: number; embedding: Float32Array }>, opts?: BatchOpts): Promise; /** * Update a take's mutable fields. May NOT change claim/kind/holder per the diff --git a/src/core/operations.ts b/src/core/operations.ts index f68f64386..f3679a79f 100644 --- a/src/core/operations.ts +++ b/src/core/operations.ts @@ -2287,7 +2287,7 @@ const think: Operation = { }, mutating: true, handler: async (ctx, p) => { - const remote = ctx.remote ?? true; + const remote = ctx.remote !== false; // fail-closed: anything not strictly false is untrusted (CLAUDE.md invariant) // Codex P1 #7 + privacy: remote callers cannot persist via MCP. const safeSave = remote ? false : Boolean(p.save); const safeTake = remote ? false : Boolean(p.take); diff --git a/src/core/pglite-engine.ts b/src/core/pglite-engine.ts index 087889f6d..82e9fd76f 100644 --- a/src/core/pglite-engine.ts +++ b/src/core/pglite-engine.ts @@ -2767,6 +2767,23 @@ export class PGLiteEngine implements BrainEngine { RETURNING cc.page_id`, params, ); + // Takes carry TEXT-embedding-space vectors keyed off the same pages, so a + // same-dim provider swap (no schema transition to drop the column) must + // re-stale them too — old-space take vectors must never score against + // new-space queries. Scoped by the SAME page-signature predicate as the + // chunks UPDATE above: routine `embed --stale` runs (where nothing + // drifted) touch zero takes, so there is no re-embed churn. The `embed + // --stale` takes lane picks the NULLed rows up via `embedding IS NULL`. + // Return value stays the CHUNK count (existing contract). + await this.db.query( + `UPDATE takes t + SET embedding = NULL, embedded_at = NULL + FROM pages p + WHERE t.page_id = p.id + AND t.embedding IS NOT NULL + AND ${sigClause}${srcClause}`, + params, + ); return (rows as unknown[]).length; } @@ -5225,27 +5242,45 @@ export class PGLiteEngine implements BrainEngine { return out; } - async countStaleTakes(): Promise { + async countStaleTakes(opts?: { sourceId?: string }): Promise { + // Source scope goes through the page join (takes has no source_id of its + // own) so `embed --stale --source X` only counts X's takes. + const srcClause = opts?.sourceId !== undefined ? ` AND p.source_id = $1` : ''; + const params = opts?.sourceId !== undefined ? [opts.sourceId] : []; const { rows } = await this.db.query( - `SELECT count(*)::int AS count FROM takes WHERE active AND embedding IS NULL` + `SELECT count(*)::int AS count + FROM takes t + JOIN pages p ON p.id = t.page_id + WHERE t.active AND t.embedding IS NULL${srcClause}`, + params, ); return Number((rows[0] as { count?: number } | undefined)?.count ?? 0); } - async listStaleTakes(): Promise { + async listStaleTakes(opts?: { sourceId?: string }): Promise { + const srcClause = opts?.sourceId !== undefined ? ` AND p.source_id = $1` : ''; + const params = opts?.sourceId !== undefined ? [opts.sourceId] : []; const { rows } = await this.db.query( `SELECT t.id AS take_id, p.slug AS page_slug, t.row_num, t.claim FROM takes t JOIN pages p ON p.id = t.page_id - WHERE t.active AND t.embedding IS NULL + WHERE t.active AND t.embedding IS NULL${srcClause} ORDER BY t.id - LIMIT 100000` + LIMIT 100000`, + params, ); return rows as unknown as StaleTakeRow[]; } - async updateTakeEmbeddingsBatch(rows: Array<{ take_id: number; embedding: Float32Array }>): Promise { + async updateTakeEmbeddingsBatch(rows: Array<{ take_id: number; embedding: Float32Array }>, opts?: BatchOpts): Promise { if (rows.length === 0) return 0; + // Wave-4: batch primitive → batchRetry, matching the addTakesBatch + // precedent above so a connection blip mid-backfill retries + audits + // instead of failing the whole takes lane batch. + return this.batchRetry(opts?.auditSite ?? 'updateTakeEmbeddingsBatch', opts?.signal, () => this._updateTakeEmbeddingsBatchOnce(rows), rows.length); + } + + private async _updateTakeEmbeddingsBatchOnce(rows: Array<{ take_id: number; embedding: Float32Array }>): Promise { let updated = 0; for (const r of rows) { // pgvector text literal (searchTakesVector precedent). `AND active` diff --git a/src/core/postgres-engine.ts b/src/core/postgres-engine.ts index 90558db17..b09a8f213 100644 --- a/src/core/postgres-engine.ts +++ b/src/core/postgres-engine.ts @@ -2719,6 +2719,23 @@ export class PostgresEngine implements BrainEngine { RETURNING cc.page_id`, params as Parameters[1], ); + // Takes carry TEXT-embedding-space vectors keyed off the same pages, so a + // same-dim provider swap (no schema transition to drop the column) must + // re-stale them too — old-space take vectors must never score against + // new-space queries. Scoped by the SAME page-signature predicate as the + // chunks UPDATE above: routine `embed --stale` runs (where nothing + // drifted) touch zero takes, so there is no re-embed churn. The `embed + // --stale` takes lane picks the NULLed rows up via `embedding IS NULL`. + // Return value stays the CHUNK count (existing contract). + await this.sql.unsafe( + `UPDATE takes t + SET embedding = NULL, embedded_at = NULL + FROM pages p + WHERE t.page_id = p.id + AND t.embedding IS NOT NULL + AND ${sigClause}${srcClause}`, + params as Parameters[1], + ); return (rows as unknown[]).length; } @@ -5155,29 +5172,43 @@ export class PostgresEngine implements BrainEngine { return out; } - async countStaleTakes(): Promise { + async countStaleTakes(opts?: { sourceId?: string }): Promise { + // Source scope goes through the page join (takes has no source_id of its + // own) so `embed --stale --source X` only counts X's takes. const sql = this.sql; const [row] = await sql` - SELECT count(*)::int AS count FROM takes WHERE active AND embedding IS NULL + SELECT count(*)::int AS count + FROM takes t + JOIN pages p ON p.id = t.page_id + WHERE t.active AND t.embedding IS NULL + ${opts?.sourceId !== undefined ? sql`AND p.source_id = ${opts.sourceId}` : sql``} `; return Number((row as { count?: number } | undefined)?.count ?? 0); } - async listStaleTakes(): Promise { + async listStaleTakes(opts?: { sourceId?: string }): Promise { const sql = this.sql; const rows = await sql` SELECT t.id AS take_id, p.slug AS page_slug, t.row_num, t.claim FROM takes t JOIN pages p ON p.id = t.page_id WHERE t.active AND t.embedding IS NULL + ${opts?.sourceId !== undefined ? sql`AND p.source_id = ${opts.sourceId}` : sql``} ORDER BY t.id LIMIT 100000 `; return rows as unknown as StaleTakeRow[]; } - async updateTakeEmbeddingsBatch(rows: Array<{ take_id: number; embedding: Float32Array }>): Promise { + async updateTakeEmbeddingsBatch(rows: Array<{ take_id: number; embedding: Float32Array }>, opts?: BatchOpts): Promise { if (rows.length === 0) return 0; + // Wave-4: batch primitive → batchRetry, matching the addTakesBatch + // precedent so a connection blip mid-backfill retries + audits + // instead of failing the whole takes lane batch. + return this.batchRetry(opts?.auditSite ?? 'updateTakeEmbeddingsBatch', opts?.signal, () => this._updateTakeEmbeddingsBatchOnce(rows), rows.length); + } + + private async _updateTakeEmbeddingsBatchOnce(rows: Array<{ take_id: number; embedding: Float32Array }>): Promise { const sql = this.sql; let updated = 0; for (const r of rows) { diff --git a/src/core/retrieval-upgrade-planner.ts b/src/core/retrieval-upgrade-planner.ts index 1edce6522..ef0c96af6 100644 --- a/src/core/retrieval-upgrade-planner.ts +++ b/src/core/retrieval-upgrade-planner.ts @@ -601,17 +601,22 @@ export async function runSchemaTransition(engine: BrainEngine, targetDim: number ); } - // #3390: the OTHER two dim-pinned columns that carry TEXT-embedding-space - // vectors. Both are created at brain-birth width (migrate.ts v55 for - // query_cache, v42 for facts) and NO migration ever ALTERs them, so before - // this fix a dimension change left them at the old width: + // #3390: the OTHER dim-pinned columns that carry TEXT-embedding-space + // vectors. All are created at brain-birth width (migrate.ts v55 for + // query_cache, v42 for facts, v37/v126 for takes) and no LATER migration + // ever ALTERs them, so before this fix a dimension change left them at + // the old width: // - query_cache.embedding stayed narrow → every store() AND lookup() // silently swallowed the width error (by design, so the cache can // never break search), i.e. a PERMANENT 0% hit rate. // - facts.embedding stayed narrow → every per-fact embed write failed // ($N::vector into the old width), and the doctor check that would // warn is skipped on PGLite (the DEFAULT engine). - // Both are text-embedding-space columns, so they MUST move with + // - takes.embedding stayed narrow → the #2089 embed takes lane failed + // every write forever while still paying the gateway each run + // (migration v126 only fixes the dim at migration-run time; a LATER + // `gbrain migrate embeddings --to ` has to move it here). + // All are text-embedding-space columns, so they MUST move with // content_chunks.embedding. The image/multimodal columns above are the // deliberate exception (separate models, independent dims). for (const t of TEXT_EMBEDDING_DIM_PINNED_TABLES) { @@ -647,6 +652,17 @@ const TEXT_EMBEDDING_DIM_PINNED_TABLES: ReadonlyArray<{ ON facts USING hnsw (embedding ${opclass}) WHERE embedding IS NOT NULL AND expired_at IS NULL`, }, + { + // Partial predicate mirrors migrate.ts v37 (and the v126 rebuild) + // verbatim. hnswIndexExpected in transitionDimPinnedColumn carries the + // #1734 dim-cap gate, matching v126's recreateIndexSql behavior. + table: 'takes', + index: 'idx_takes_embedding_hnsw', + indexSql: (opclass) => + `CREATE INDEX IF NOT EXISTS idx_takes_embedding_hnsw + ON takes USING hnsw (embedding ${opclass}) + WHERE active AND embedding IS NOT NULL`, + }, ]; /** @@ -658,7 +674,8 @@ const TEXT_EMBEDDING_DIM_PINNED_TABLES: ReadonlyArray<{ * Dropping the column discards the stored vectors, which is correct: they are * in the OLD embedding space and unusable after the swap. query_cache is a * cache (refills on the next query); facts re-embed on their next write / - * `gbrain extract` pass. + * `gbrain extract` pass; takes re-embed via the `gbrain embed --stale` takes + * lane (staleness predicate is `embedding IS NULL`, which the drop satisfies). */ async function transitionDimPinnedColumn( tx: { executeRaw: (sql: string, params?: unknown[]) => Promise }, diff --git a/src/core/retry.ts b/src/core/retry.ts index 5defca720..8980802db 100644 --- a/src/core/retry.ts +++ b/src/core/retry.ts @@ -77,6 +77,7 @@ export const BATCH_AUDIT_SITES = [ 'addLinksBatch', 'addTimelineEntriesBatch', 'addTakesBatch', + 'updateTakeEmbeddingsBatch', 'upsertChunks', // extract.ts per-site labels. 'extract.links_inc', diff --git a/src/core/takes-append.ts b/src/core/takes-append.ts new file mode 100644 index 000000000..e4216858b --- /dev/null +++ b/src/core/takes-append.ts @@ -0,0 +1,175 @@ +/** + * Shared dual-plane take append (Wave-4 review fix). + * + * Markdown is the source of truth for takes: `gbrain takes add` allocates + * row numbers from the FILE fence (upsertTakeRow → max fence rowNum + 1), + * then mirrors to the DB via addTakesBatch's ON CONFLICT (page_id, row_num) + * DO UPDATE. Any writer that appends DB-only rows (the original + * persistThinkTake did) allocates from DB MAX(row_num) instead — and the + * next fence-allocated row at the same number deterministically OVERWRITES + * it. This module is the single append path both surfaces call, so the two + * planes cannot drift: + * + * withPageLock(slug) + * → resolve the page (scoped: sourceIds > sourceId > unscoped) + * → resolve the brain dir (explicit override > sync.repo_path config) + * → file plane: upsertTakeRow allocates the row number (file-first) + * → writeBody + * → DB plane: addTakesBatch mirror at the fence-allocated row number + * + * Fallback posture (headless / DB-only brains): when no brain dir resolves, + * the row is allocated DB-only (MAX(row_num)+1) and the + * 'TAKE_FILE_PLANE_UNAVAILABLE' warning is returned — documented posture: + * DB-only rows can be renumbered/overwritten by a later fence reconcile. + */ + +import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'node:fs'; +import { join, dirname } from 'node:path'; +import type { BrainEngine, TakeKind } from './engine.ts'; +import { upsertTakeRow } from './takes-fence.ts'; +import { withPageLock } from './page-lock.ts'; + +/** Warning code returned when the append fell back to DB-only allocation. */ +export const TAKE_FILE_PLANE_UNAVAILABLE = 'TAKE_FILE_PLANE_UNAVAILABLE'; + +/** Thrown when the (scoped) page lookup finds nothing. Callers map this to + * their own surface signal (CLI error+exit, TAKE_ANCHOR_NOT_FOUND, ...). */ +export class TakePageNotFoundError extends Error { + readonly slug: string; + readonly sourceId?: string; + constructor(slug: string, sourceId?: string) { + super(`Page not found in brain: ${slug}${sourceId ? ` (source=${sourceId})` : ''}`); + this.name = 'TakePageNotFoundError'; + this.slug = slug; + this.sourceId = sourceId; + } +} + +export function pageFilePath(brainDir: string, slug: string): string { + return join(brainDir, `${slug}.md`); +} + +export function readBodyOrEmpty(path: string): string { + if (!existsSync(path)) return ''; + return readFileSync(path, 'utf-8'); +} + +export function writeBody(path: string, body: string): void { + mkdirSync(dirname(path), { recursive: true }); + writeFileSync(path, body, 'utf-8'); +} + +/** + * Resolve the on-disk brain dir for the file plane, or null when none is + * available (headless install, unset/dangling sync.repo_path). An explicit + * dir that does not exist ALSO returns null — the CLI wraps this with its + * own loud error+exit; core callers fall back to DB-only allocation. + */ +export async function resolveBrainDirOrNull( + engine: BrainEngine, + explicitDir?: string | null, +): Promise { + if (explicitDir) return existsSync(explicitDir) ? explicitDir : null; + try { + const configured = await engine.getConfig('sync.repo_path'); + if (configured && existsSync(configured)) return configured; + } catch { + // Config read failure → treat as no file plane (fallback posture). + } + return null; +} + +export interface AppendTakeInput { + slug: string; + /** Scalar source scope (sourceIds wins when both set — sourceScopeOpts precedence). */ + sourceId?: string; + /** Federated source scope for the page lookup. */ + sourceIds?: string[]; + claim: string; + kind: TakeKind; + holder: string; + weight?: number; + source?: string; + sinceDate?: string; + /** + * Pre-resolved brain dir (CLI --dir, or a caller that already validated + * sync.repo_path). When omitted, resolves via resolveBrainDirOrNull; when + * that yields nothing, the append falls back to DB-only allocation with + * the TAKE_FILE_PLANE_UNAVAILABLE warning. + */ + brainDir?: string; +} + +export interface AppendTakeResult { + rowNum: number; + inserted: number; + warnings: string[]; +} + +/** + * Append one take row to a page on BOTH planes (file fence + DB), serialized + * under withPageLock(slug) so concurrent appends can't race the allocation. + * + * The page lookup runs FIRST (inside the lock, before any file write) so a + * missing page throws {@link TakePageNotFoundError} without leaving a + * half-written fence behind. + */ +export async function appendTake( + engine: BrainEngine, + input: AppendTakeInput, +): Promise { + const warnings: string[] = []; + return withPageLock(input.slug, async () => { + // Scoped page lookup: federated array > scalar > unscoped. + const pageOpts: { sourceId?: string; sourceIds?: string[] } = {}; + if (input.sourceIds !== undefined) pageOpts.sourceIds = input.sourceIds; + else if (input.sourceId !== undefined) pageOpts.sourceId = input.sourceId; + const page = await engine.getPage(input.slug, pageOpts); + if (!page) throw new TakePageNotFoundError(input.slug, input.sourceId); + + const brainDir = input.brainDir ?? await resolveBrainDirOrNull(engine); + + let rowNum: number; + if (brainDir) { + // File plane first: the fence allocates the row number. + const path = pageFilePath(brainDir, input.slug); + const body = readBodyOrEmpty(path); + const up = upsertTakeRow(body, { + claim: input.claim, + kind: input.kind, + holder: input.holder, + weight: input.weight ?? 0.5, + source: input.source, + sinceDate: input.sinceDate, + active: true, + }); + writeBody(path, up.body); + rowNum = up.rowNum; + } else { + // DB-only fallback (no file plane). MAX(row_num)+1 under the same page + // lock. A later fence reconcile may renumber/overwrite these rows — + // documented posture for headless brains. + warnings.push(TAKE_FILE_PLANE_UNAVAILABLE); + const rows = await engine.executeRaw<{ next: number | string }>( + `SELECT (COALESCE(MAX(row_num), 0) + 1)::int AS next FROM takes WHERE page_id = $1`, + [page.id], + ); + rowNum = Number(rows[0]?.next ?? 1); + } + + const inserted = await engine.addTakesBatch([{ + page_id: page.id, + row_num: rowNum, + claim: input.claim, + kind: input.kind, + holder: input.holder, + weight: input.weight ?? 0.5, + since_date: input.sinceDate, + source: input.source, + active: true, + superseded_by: null, + }]); + + return { rowNum, inserted, warnings }; + }); +} diff --git a/src/core/think/index.ts b/src/core/think/index.ts index 43208f84c..d50e464d4 100644 --- a/src/core/think/index.ts +++ b/src/core/think/index.ts @@ -669,30 +669,49 @@ export interface PersistThinkTakeResult { warnings: string[]; } +/** + * Cap on the persisted think-take claim. Rationale: the prompt renderer + * (sanitize.ts:sanitizeTakeForPrompt) truncates claims at 500 chars anyway, + * and fence cells are single-line — a multi-page synthesis answer stored + * verbatim as one claim bloats the fence + DB for content nothing ever + * reads past. 2000 keeps ample headroom over the 500-char read while + * bounding the cell. + */ +export const THINK_TAKE_CLAIM_MAX_CHARS = 2000; + /** * Persist `gbrain think --take` as the next append-only take row on the * anchor page (#2556 — the flag was declared but nothing consumed it, so * `think --take` silently wrote nothing on both the CLI and MCP paths). * * The synthesis answer is the claim; holder='brain' because this is - * gbrain's own analysis; kind='take' (KIND_VALUES member). DB-plane write - * via addTakesBatch — the same path extract-takes-from-pages uses; the - * engines maintain the takes fence at that write site. The row-number - * computation + insert are serialized under withPageLock(anchor) so a - * concurrent `takes add` / second think --take can't race MAX(row_num)+1 - * into a duplicate row. + * gbrain's own analysis; kind='take' (KIND_VALUES member). Dual-plane + * write via the shared appendTake helper (takes-append.ts) — the SAME path + * `gbrain takes add` uses: withPageLock(anchor slug) → fence allocates the + * row number from the FILE (markdown is the source of truth) → writeBody → + * addTakesBatch mirror. When no brain dir resolves (headless / DB-only + * brains), appendTake falls back to DB-only MAX(row_num)+1 allocation and + * surfaces TAKE_FILE_PLANE_UNAVAILABLE (documented posture: DB-only rows + * can be renumbered/overwritten by a later fence reconcile). + * + * The claim is flattened to a single line (fence cells are single-line; + * newlines → ' ') and truncated at THINK_TAKE_CLAIM_MAX_CHARS with the + * TAKE_CLAIM_TRUNCATED warning. * * Signals (never throws for expected shapes, mirroring persistSynthesis): - * - TAKE_REQUIRES_ANCHOR — no anchor given. - * - TAKE_EMPTY_NOT_PERSISTED — synthesis failed or empty answer. - * - TAKE_ANCHOR_NOT_FOUND: — anchor page absent in the caller's scope + * - TAKE_REQUIRES_ANCHOR — no anchor given. + * - TAKE_EMPTY_NOT_PERSISTED — synthesis failed or empty answer. + * - TAKE_ANCHOR_NOT_FOUND: — anchor page absent in the caller's scope * (scope resolved via the sourceIds > sourceId precedence, matching * sourceScopeOpts). Remote callers never reach here — the op handler's * fail-closed gate zeroes `take` for them. + * - TAKE_CLAIM_TRUNCATED — claim exceeded the cap (row still written). + * - TAKE_FILE_PLANE_UNAVAILABLE — DB-only fallback fired (row still written). * * Design provenance: community PR #2618 (its successor #3164 was declined * for bundled feature scope; the maintainer confirmed this half is a - * legitimate fix). Re-implemented with the page-lock serialization added. + * legitimate fix). Re-implemented with the page-lock serialization added; + * Wave-4 review moved the write onto the dual-plane appendTake path. */ export async function persistThinkTake( engine: BrainEngine, @@ -707,36 +726,34 @@ export async function persistThinkTake( return { rowNum: null, inserted: 0, warnings: ['TAKE_EMPTY_NOT_PERSISTED'] }; } - const pageOpts: { sourceId?: string; sourceIds?: string[] } = {}; - if (opts.sourceIds !== undefined) pageOpts.sourceIds = opts.sourceIds; - else if (opts.sourceId !== undefined) pageOpts.sourceId = opts.sourceId; - const page = await engine.getPage(anchor, pageOpts); - if (!page) { - return { rowNum: null, inserted: 0, warnings: [`TAKE_ANCHOR_NOT_FOUND: ${anchor}`] }; + const warnings: string[] = []; + // Fence cells are single-line: flatten first, then bound the length. + let claim = result.answer.trim().replace(/\s*\r?\n\s*/g, ' '); + if (claim.length > THINK_TAKE_CLAIM_MAX_CHARS) { + claim = claim.slice(0, THINK_TAKE_CLAIM_MAX_CHARS); + warnings.push('TAKE_CLAIM_TRUNCATED'); } - const { withPageLock } = await import('../page-lock.ts'); - let rowNum = 0; - let inserted = 0; - await withPageLock(anchor, async () => { - const rows = await engine.executeRaw<{ next: number | string }>( - `SELECT (COALESCE(MAX(row_num), 0) + 1)::int AS next FROM takes WHERE page_id = $1`, - [page.id], - ); - rowNum = Number(rows[0]?.next ?? 1); - inserted = await engine.addTakesBatch([{ - page_id: page.id, - row_num: rowNum, - claim: result.answer.trim(), + const { appendTake, TakePageNotFoundError } = await import('../takes-append.ts'); + try { + const appended = await appendTake(engine, { + slug: anchor, + ...(opts.sourceIds !== undefined ? { sourceIds: opts.sourceIds } : {}), + ...(opts.sourceId !== undefined ? { sourceId: opts.sourceId } : {}), + claim, kind: 'take', holder: 'brain', weight: 0.5, source: 'gbrain think', - active: true, - }]); - }); - - return { rowNum, inserted, warnings: [] }; + }); + warnings.push(...appended.warnings); + return { rowNum: appended.rowNum, inserted: appended.inserted, warnings }; + } catch (e) { + if (e instanceof TakePageNotFoundError) { + return { rowNum: null, inserted: 0, warnings: [`TAKE_ANCHOR_NOT_FOUND: ${anchor}`] }; + } + throw e; + } } // ───────────────────────────────────────────────────────────────── diff --git a/test/core/retry.test.ts b/test/core/retry.test.ts index 77746da9d..13eacaa4f 100644 --- a/test/core/retry.test.ts +++ b/test/core/retry.test.ts @@ -339,7 +339,7 @@ describe('BATCH_AUDIT_SITES typed enum + isBatchAuditSite guard (D10c codex)', ( // Pin the set so a future "cleanup" PR can't silently drop a site and // break audit-attribution for the corresponding caller. const expected = new Set([ - 'addLinksBatch', 'addTimelineEntriesBatch', 'addTakesBatch', 'upsertChunks', + 'addLinksBatch', 'addTimelineEntriesBatch', 'addTakesBatch', 'updateTakeEmbeddingsBatch', 'upsertChunks', 'extract.links_inc', 'extract.timeline_inc', 'extract.links_fs', 'extract.timeline_fs', 'extract.links_db', 'extract.timeline_db', diff --git a/test/e2e/migrate-embeddings-postgres.test.ts b/test/e2e/migrate-embeddings-postgres.test.ts index f33e67dc2..ef165c9ff 100644 --- a/test/e2e/migrate-embeddings-postgres.test.ts +++ b/test/e2e/migrate-embeddings-postgres.test.ts @@ -174,12 +174,14 @@ d('embedding migration (live Postgres + pgvector)', () => { expect(persisted).toEqual([[toModel, targetDims]]); expect(await columnDims()).toBe(targetDims); - // All THREE dim-pinned text-embedding-space columns move together on real - // pgvector (query_cache + facts are created at brain-birth width and no - // migration ever ALTERs them — before the fix they stayed narrow, which - // silently killed the query cache and every per-fact embed write). + // ALL dim-pinned text-embedding-space columns move together on real + // pgvector (query_cache + facts + takes are created at brain-birth width + // and no later migration ALTERs them — before the fix they stayed + // narrow, which silently killed the query cache, every per-fact embed + // write, and every takes-lane embed write). expect(await embeddingColWidth('query_cache')).toBe(targetDims); expect(await embeddingColWidth('facts')).toBe(targetDims); + expect(await embeddingColWidth('takes')).toBe(targetDims); // And the rebuilt columns actually ACCEPT a vector at the new width. const nv = `[${new Array(targetDims).fill(0.01).join(',')}]`; await engine.executeRaw( diff --git a/test/embed-takes-lane.serial.test.ts b/test/embed-takes-lane.serial.test.ts new file mode 100644 index 000000000..e65359565 --- /dev/null +++ b/test/embed-takes-lane.serial.test.ts @@ -0,0 +1,170 @@ +/** + * #2089 — the takes lane of `gbrain embed --stale` (embedStaleTakes via + * runEmbedCore). The wave shipped the lane with ZERO tests driving it: the + * engine writer (updateTakeEmbeddingsBatch) is covered, but the lane's + * control flow — dry-run counting, id coercion into the writer, and the + * budget/abort cutShort skip — was untested. + * + * Hermetic: mock.module on core/embedding.ts (same pattern as + * embed.serial.test.ts) + the gateway embed-transport seam for the creds + * preflight. Serial: mock.module + env mutation are process-global. + */ +import { test, expect, mock, beforeEach, afterEach } from 'bun:test'; +import type { BrainEngine } from '../src/core/engine.ts'; + +let totalEmbedCalls = 0; + +mock.module('../src/core/embedding.ts', () => ({ + embedBatch: async (texts: string[]) => { + totalEmbedCalls++; + await new Promise(r => setTimeout(r, 5)); + return texts.map(() => new Float32Array(1536)); + }, + currentEmbeddingSignature: () => 'test:model:1536', +})); + +// Import AFTER mocking. +const { runEmbedCore } = await import('../src/commands/embed.ts'); +const { __setEmbedTransportForTests } = await import('../src/core/ai/gateway.ts'); +__setEmbedTransportForTests(async () => ({ embeddings: [], usage: { tokens: 0 } } as any)); + +/** Proxy mock engine (embed.serial.test.ts pattern) with call tracking. */ +function mockEngine(overrides: Partial> = {}): BrainEngine { + const calls: { method: string; args: any[] }[] = []; + const track = (method: string) => (...args: any[]) => { + calls.push({ method, args }); + if (overrides[method]) return overrides[method](...args); + return Promise.resolve(null); + }; + const engine = new Proxy({} as any, { + get(_, prop: string) { + if (prop === '_calls') return calls; + if (overrides[prop]) { + return (...args: any[]) => { + calls.push({ method: prop, args }); + return overrides[prop](...args); + }; + } + return track(prop); + }, + }); + return engine; +} + +function callCount(engine: BrainEngine, method: string): number { + return (engine as any)._calls.filter((c: any) => c.method === method).length; +} + +beforeEach(() => { + totalEmbedCalls = 0; +}); + +afterEach(() => { + delete process.env.GBRAIN_EMBED_TIME_BUDGET_MS; +}); + +test('#2089 --stale takes lane: dry-run counts only; real run embeds and batch-writes with coerced numeric ids', async () => { + // Driver-shaped ids on purpose: string + bigint must arrive at the writer + // as plain numbers (the lane's Number() coercion). + const staleTakes = [ + { take_id: '11' as any, page_slug: 'p/a', row_num: 1, claim: 'claim one' }, + { take_id: 12n as any, page_slug: 'p/a', row_num: 2, claim: 'claim two' }, + { take_id: 13, page_slug: 'p/b', row_num: 1, claim: 'claim three' }, + ]; + let written: Array<{ take_id: number; embedding: Float32Array }> = []; + const makeEngine = () => mockEngine({ + countStaleChunks: async () => 0, // pages lane: nothing stale + invalidateStaleSignatureEmbeddings: async () => 0, + countStaleTakes: async () => staleTakes.length, + listStaleTakes: async () => staleTakes, + updateTakeEmbeddingsBatch: async (rows: any[]) => { written = rows; return rows.length; }, + }); + + // Dry-run: count-only — no gateway call, no listStaleTakes, no writes. + const dryEngine = makeEngine(); + const dry = await runEmbedCore(dryEngine, { stale: true, dryRun: true, quiet: true }); + expect(dry.takes_would_embed).toBe(3); + expect(dry.takes_embedded ?? 0).toBe(0); + expect(totalEmbedCalls).toBe(0); + expect(callCount(dryEngine, 'listStaleTakes')).toBe(0); + expect(callCount(dryEngine, 'updateTakeEmbeddingsBatch')).toBe(0); + expect(written).toHaveLength(0); + + // Real run: 3 claims fit one 64-cap batch → one gateway call, 3 writes. + const res = await runEmbedCore(makeEngine(), { stale: true, quiet: true }); + expect(res.takes_embedded).toBe(3); + expect(res.failures).toBe(0); + expect(written).toHaveLength(3); + expect(written.map(w => w.take_id)).toEqual([11, 12, 13]); + for (const w of written) { + expect(typeof w.take_id).toBe('number'); + expect(w.embedding).toBeInstanceOf(Float32Array); + expect(w.embedding.length).toBe(1536); + } +}, 20_000); + +test('Wave-4 source scoping: `--stale --source X` threads sourceId into BOTH takes enumerators', async () => { + const enumeratorArgs: Array<{ method: string; opts: unknown }> = []; + let written: Array<{ take_id: number; embedding: Float32Array }> = []; + const engine = mockEngine({ + countStaleChunks: async () => 0, // pages lane: nothing stale + invalidateStaleSignatureEmbeddings: async () => 0, + countStaleTakes: async (opts: unknown) => { + enumeratorArgs.push({ method: 'countStaleTakes', opts }); + return 1; + }, + listStaleTakes: async (opts: unknown) => { + enumeratorArgs.push({ method: 'listStaleTakes', opts }); + return [{ take_id: 21, page_slug: 'p/scoped', row_num: 1, claim: 'scoped claim' }]; + }, + updateTakeEmbeddingsBatch: async (rows: any[]) => { written = rows; return rows.length; }, + }); + + const res = await runEmbedCore(engine, { stale: true, quiet: true, sourceId: 'tenant-a' }); + + // Both enumerators saw the run's source scope — a `--source X` run must + // never enumerate (or pay to embed) another source's takes. + expect(enumeratorArgs).toEqual([ + { method: 'countStaleTakes', opts: { sourceId: 'tenant-a' } }, + { method: 'listStaleTakes', opts: { sourceId: 'tenant-a' } }, + ]); + expect(res.takes_embedded).toBe(1); + expect(written.map(w => w.take_id)).toEqual([21]); + + // Unscoped run: enumerators get NO scope object (all-source semantics). + enumeratorArgs.length = 0; + const engine2 = mockEngine({ + countStaleChunks: async () => 0, + invalidateStaleSignatureEmbeddings: async () => 0, + countStaleTakes: async (opts: unknown) => { + enumeratorArgs.push({ method: 'countStaleTakes', opts }); + return 0; + }, + }); + await runEmbedCore(engine2, { stale: true, quiet: true }); + expect(enumeratorArgs).toEqual([{ method: 'countStaleTakes', opts: undefined }]); +}, 20_000); + +test('#2089 cutShort: a wall-clock-budget abort in the pages lane SKIPS the takes lane (next run picks them up)', async () => { + // 1ms budget; the pages lane's first page-load sleeps 30ms so the budget + // signal is guaranteed to have fired by the time the lane reports back. + process.env.GBRAIN_EMBED_TIME_BUDGET_MS = '1'; + const engine = mockEngine({ + countStaleChunks: async () => 5, // pass the 0-stale early return + invalidateStaleSignatureEmbeddings: async () => 0, + listStaleChunks: async () => { + await new Promise(r => setTimeout(r, 30)); + return []; + }, + countStaleTakes: async () => 99, // must never be reached + }); + + const res = await runEmbedCore(engine, { stale: true, quiet: true }); + + // The takes lane never even probed the stale count — a timed-out run must + // not start NEW work (#2089 contract: graceful skip, not extra load). + expect(callCount(engine, 'countStaleTakes')).toBe(0); + expect(res.takes_embedded).toBeUndefined(); + expect(res.takes_would_embed).toBeUndefined(); + expect(totalEmbedCalls).toBe(0); +}, 20_000); diff --git a/test/embedding-migration.test.ts b/test/embedding-migration.test.ts index 90e7ddd84..dbee42c87 100644 --- a/test/embedding-migration.test.ts +++ b/test/embedding-migration.test.ts @@ -165,6 +165,54 @@ describe('#3391 includeNullSignature widening', () => { expect(await engine.invalidateStaleSignatureEmbeddings({ signature: 'new:model:1' })).toBe(0); expect(await engine.countStaleChunks({ signature: 'new:model:1' })).toBe(0); }); + + // Wave-4 review: invalidation also re-stales TAKES on the same pages under + // the SAME signature predicate — a same-dim provider swap has no schema + // transition to drop takes.embedding, and old-space take vectors must not + // score against new-space queries. Routine runs (nothing drifted) touch + // zero takes, so there is no re-embed churn. + test('invalidation re-stales takes on drifted pages; fresh kept; grandfather lifts with the flag', async () => { + await seedEmbedded('legacy', 'abcde', null); + await seedEmbedded('drifted', 'fghij', 'old:model:1'); + await seedEmbedded('fresh', 'klmno', 'new:model:1'); + + // One EMBEDDED take per page (text-embedding space, same as the chunks). + for (const slug of ['legacy', 'drifted', 'fresh']) { + const page = await engine.getPage(slug); + await engine.addTakesBatch([{ + page_id: page!.id, row_num: 1, claim: `take on ${slug}`, + kind: 'fact', holder: 'world', weight: 0.9, + }]); + } + await engine.executeRaw( + `UPDATE takes + SET embedding = ('[' || array_to_string(array_fill(0.0::real, ARRAY[$1::int]), ',') || ']')::vector, + embedded_at = now()`, + [colDim], + ); + expect(await engine.countStaleTakes()).toBe(0); + + // Default (grandfather): only the drifted page's take re-stales; the + // returned count remains CHUNKS-only (existing contract). + expect(await engine.invalidateStaleSignatureEmbeddings({ signature: 'new:model:1' })).toBe(1); + let staleTakes = await engine.listStaleTakes(); + expect(staleTakes.map(t => t.page_slug)).toEqual(['drifted']); + const cleared = await engine.executeRaw<{ n: number }>( + `SELECT count(*)::int AS n FROM takes WHERE embedding IS NULL AND embedded_at IS NULL`, + ); + expect(Number(cleared[0]?.n)).toBe(1); // embedded_at cleared alongside + + // Widened (#3391): the legacy (NULL-sig) take re-stales too; fresh kept. + await engine.invalidateStaleSignatureEmbeddings({ signature: 'new:model:1', includeNullSignature: true }); + staleTakes = await engine.listStaleTakes(); + expect(staleTakes.map(t => t.page_slug).sort()).toEqual(['drifted', 'legacy']); + const keptTake = await engine.executeRaw<{ n: number }>( + `SELECT count(*)::int AS n FROM takes t + JOIN pages p ON p.id = t.page_id + WHERE p.slug = 'fresh' AND t.embedding IS NOT NULL`, + ); + expect(Number(keptTake[0]?.n)).toBe(1); + }); }); describe('planEmbeddingMigration', () => { diff --git a/test/migrations-takes-embedding-rerun.test.ts b/test/migrations-takes-embedding-rerun.test.ts new file mode 100644 index 000000000..172ba62f1 --- /dev/null +++ b/test/migrations-takes-embedding-rerun.test.ts @@ -0,0 +1,82 @@ +/** + * Migration v126 (#2089) — the DEFAULT-dims no-op path, and the guard that + * protects backfilled vectors. + * + * The shipped serial suite (migrations-takes-embedding-dims.serial.test.ts) + * pins the 1024 REBUILD path only, and its idempotency re-run happens before + * any embedding exists — so the most consequential property of the + * `currentDim === embeddingDim` early return was unpinned: a forced re-run + * of v126 on a correct column must NOT wipe already-backfilled take + * embeddings (the rebuild branch NULLs + drops the column). This test runs + * on a fresh default-dims (1536) brain — the overwhelmingly common install — + * asserts v126 no-opped at init, backfills a vector via the new writer, then + * forces a re-run and asserts the vector SURVIVES. + */ +import { test, expect, beforeAll, afterAll } from 'bun:test'; +import { PGLiteEngine } from '../src/core/pglite-engine.ts'; +import { runMigrations } from '../src/core/migrate.ts'; + +const DEFAULT_DIMS = 1536; + +let engine: PGLiteEngine; + +beforeAll(async () => { + engine = new PGLiteEngine(); + await engine.connect({}); + await engine.initSchema(); +}, 60_000); + +afterAll(async () => { + await engine.disconnect(); +}); + +async function takesEmbeddingDim(): Promise { + const rows = await engine.executeRaw<{ atttypmod: number }>( + `SELECT atttypmod FROM pg_attribute + WHERE attrelid = 'takes'::regclass AND attname = 'embedding'`, + ); + expect(rows.length).toBe(1); + return Number(rows[0].atttypmod); +} + +test('v126 no-ops on a default-dims brain, and a forced re-run PRESERVES backfilled take embeddings', async () => { + // Fresh default brain: column already VECTOR(1536) from v37 → v126 no-op. + expect(await takesEmbeddingDim()).toBe(DEFAULT_DIMS); + const idx = await engine.executeRaw<{ indexdef: string }>( + `SELECT indexdef FROM pg_indexes + WHERE tablename = 'takes' AND indexname = 'idx_takes_embedding_hnsw'`, + ); + expect(idx.length).toBe(1); // v37's HNSW index intact + + // Backfill one claim through the #2089 writer. + const page = await engine.putPage('companies/rerun-example', { + title: 'Rerun Example', type: 'company' as const, + compiled_truth: '## Takes\n\nA placeholder company page.\n', + }); + await engine.addTakesBatch([ + { page_id: page.id, row_num: 1, claim: 'rerun-preservation claim', kind: 'fact', holder: 'world', weight: 0.9 }, + ]); + const stale = (await engine.listStaleTakes()).filter(t => t.claim === 'rerun-preservation claim'); + expect(stale).toHaveLength(1); + const emb = new Float32Array(DEFAULT_DIMS); + emb[3] = 1; + expect(await engine.updateTakeEmbeddingsBatch([ + { take_id: Number(stale[0].take_id), embedding: emb }, + ])).toBe(1); + + // Force v126 to run again (version rollback → re-migrate). + await engine.setConfig('version', '125'); + const { applied } = await runMigrations(engine); + expect(applied).toBeGreaterThanOrEqual(1); + + // The no-op guard held: same dims, index intact, and — the load-bearing + // bit — the backfilled embedding was NOT nulled/dropped by a re-rebuild. + expect(await takesEmbeddingDim()).toBe(DEFAULT_DIMS); + const kept = await engine.executeRaw<{ has_embedding: boolean }>( + `SELECT (embedding IS NOT NULL) AS has_embedding FROM takes WHERE id = $1`, + [Number(stale[0].take_id)], + ); + expect(kept[0]?.has_embedding).toBe(true); + const hits = await engine.searchTakesVector(emb, { limit: 5 }); + expect(hits.map(h => h.claim)).toContain('rerun-preservation claim'); +}, 30_000); diff --git a/test/retrieval-upgrade-planner.test.ts b/test/retrieval-upgrade-planner.test.ts index 025a95b2d..389d89196 100644 --- a/test/retrieval-upgrade-planner.test.ts +++ b/test/retrieval-upgrade-planner.test.ts @@ -288,6 +288,71 @@ describe('applyRetrievalUpgrade — state machine + atomicity (D12, D18)', () => expect(byName.get('embedding_multimodal')).toBe('vector(1024)'); }); + // Wave-4 review: takes joined TEXT_EMBEDDING_DIM_PINNED_TABLES. Before the + // fix, a dim-change transition left takes.embedding at the old width and + // the #2089 embed takes lane failed every write forever (while still + // paying the gateway each run). + test('runSchemaTransition rebuilds takes.embedding at the target dim + recreates the partial HNSW index', async () => { + await setLegacyDefaultConfig(); + await seedPages(150); + + // Seed one take and give it a vector at the OLD width (1536). + const page = await engine.putPage('takes/transition-example', { + title: 'Takes transition example', type: 'note', + compiled_truth: 'A safe placeholder page for the takes dim transition.', + }); + await engine.addTakesBatch([{ + page_id: page.id, row_num: 1, claim: 'pre-transition claim', + kind: 'fact', holder: 'world', weight: 0.9, + }]); + const stale = (await engine.listStaleTakes()).filter(t => t.claim === 'pre-transition claim'); + expect(stale).toHaveLength(1); + const takeId = Number(stale[0].take_id); + // Probe the CURRENT column width for the pre-transition write — the + // suite shares one engine, and an earlier apply test may already have + // transitioned the takes column (schema persists across resetPgliteState). + const cur = await engine.executeRaw<{ atttypmod: number }>( + `SELECT atttypmod FROM pg_attribute + WHERE attrelid = 'takes'::regclass AND attname = 'embedding' + AND attnum > 0 AND NOT attisdropped`, + ); + const curDim = Number(cur[0]?.atttypmod); + expect(curDim).toBeGreaterThan(0); + const oldVec = new Float32Array(curDim); + oldVec[0] = 1; + expect(await engine.updateTakeEmbeddingsBatch([{ take_id: takeId, embedding: oldVec }])).toBe(1); + + const plan = await planRetrievalUpgrade(engine); + await applyRetrievalUpgrade(engine, plan); + + // Column rebuilt at the target dim (atttypmod IS the pgvector dimension). + const dim = await engine.executeRaw<{ atttypmod: number }>( + `SELECT atttypmod FROM pg_attribute + WHERE attrelid = 'takes'::regclass AND attname = 'embedding' + AND attnum > 0 AND NOT attisdropped`, + ); + expect(Number(dim[0]?.atttypmod)).toBe(ZE_TARGET_EMBEDDING_DIM); + + // Partial HNSW index recreated (dim under the #1734 cap), v37 predicate. + const idx = await engine.executeRaw<{ indexdef: string }>( + `SELECT indexdef FROM pg_indexes + WHERE tablename = 'takes' AND indexname = 'idx_takes_embedding_hnsw'`, + ); + expect(idx).toHaveLength(1); + expect(idx[0].indexdef).toContain('hnsw'); + expect(idx[0].indexdef).toMatch(/WHERE\s+\(?active/i); + + // Old-space vector discarded — the take is back in the stale pool. + const restaled = (await engine.listStaleTakes()).filter(t => Number(t.take_id) === takeId); + expect(restaled).toHaveLength(1); + + // And the rebuilt column ACCEPTS a vector at the new width end-to-end. + const newVec = new Float32Array(ZE_TARGET_EMBEDDING_DIM); + newVec[3] = 1; + expect(await engine.updateTakeEmbeddingsBatch([{ take_id: takeId, embedding: newVec }])).toBe(1); + expect((await engine.getTakeEmbeddings([takeId])).get(takeId)?.length).toBe(ZE_TARGET_EMBEDDING_DIM); + }); + test('runSchemaTransition restores partial WHERE on idx_chunks_embedding_image', async () => { await setLegacyDefaultConfig(); await seedPages(150); diff --git a/test/takes-command-source-scope.test.ts b/test/takes-command-source-scope.test.ts index 85c714bcf..4b38dfd09 100644 --- a/test/takes-command-source-scope.test.ts +++ b/test/takes-command-source-scope.test.ts @@ -19,6 +19,16 @@ function makeEngine(opts: { knownSources?: string[] } = {}) { const pageLookups: unknown[][] = []; const engine = { getConfig: async () => null, + // Wave-4: cmdAdd routes through the shared appendTake helper + // (takes-append.ts), whose scoped page lookup is engine.getPage — + // the same seam persistThinkTake uses. Record [slug, scope] so the + // scoping assertions below stay shape-identical to the old raw-SQL pin. + getPage: async (slug: string, pageOpts?: { sourceId?: string; sourceIds?: string[] }) => { + pageLookups.push([slug, pageOpts?.sourceIds ?? pageOpts?.sourceId]); + if (slug === 'shared/page' && pageOpts?.sourceId === 'dept') return { id: 22, slug }; + if (slug === 'shared/page' && (pageOpts?.sourceId === 'default' || pageOpts?.sourceId === undefined)) return { id: 11, slug }; + return null; + }, executeRaw: async (sql: string, params: unknown[] = []) => { if (sql.includes('FROM sources WHERE id = $1')) { // Default (no `knownSources` override): every id "exists", matching diff --git a/test/think-pipeline.serial.test.ts b/test/think-pipeline.serial.test.ts index 203b3d049..9041e9bba 100644 --- a/test/think-pipeline.serial.test.ts +++ b/test/think-pipeline.serial.test.ts @@ -1,7 +1,7 @@ import { describe, test, expect, beforeAll, afterAll } from 'bun:test'; import { PGLiteEngine } from '../src/core/pglite-engine.ts'; import { operationsByName } from '../src/core/operations.ts'; -import { runThink, persistSynthesis, persistThinkTake, type ThinkLLMClient } from '../src/core/think/index.ts'; +import { runThink, persistSynthesis, persistThinkTake, THINK_TAKE_CLAIM_MAX_CHARS, type ThinkLLMClient } from '../src/core/think/index.ts'; import { sanitizeTakeForPrompt, renderTakesBlock } from '../src/core/think/sanitize.ts'; import { resolveCitations, parseInlineCitations, normalizeStructuredCitations } from '../src/core/think/cite-render.ts'; import { runGather } from '../src/core/think/gather.ts'; @@ -490,7 +490,11 @@ describe('persistThinkTake — #2556 think --take actually persists', () => { engine, synthResult('This page should remember the synthesized placeholder insight.'), { anchor: 'notes/think-take-target-example' }, ); - expect(persisted).toEqual({ rowNum: 1, inserted: 1, warnings: [] }); + // Wave-4 dual-plane contract: this hermetic engine has no brain repo + // (sync.repo_path unset), so appendTake falls back to DB-only allocation + // and surfaces TAKE_FILE_PLANE_UNAVAILABLE. The row still lands. + expect(persisted).toMatchObject({ rowNum: 1, inserted: 1 }); + expect(persisted.warnings).toEqual(['TAKE_FILE_PLANE_UNAVAILABLE']); const takes = await engine.listTakes({ page_id: target.id }); expect(takes).toHaveLength(1); expect(takes[0]).toMatchObject({ @@ -506,6 +510,30 @@ describe('persistThinkTake — #2556 think --take actually persists', () => { expect(second.rowNum).toBe(2); }); + test('bounds the claim: 2500-char multi-line answer → single-line 2000-char claim + TAKE_CLAIM_TRUNCATED', async () => { + const target = await engine.putPage('notes/think-take-truncate-example', { + title: 'Think take truncate target', type: 'note', + compiled_truth: 'A safe placeholder page for claim-bound persistence.', + }); + // Multi-line answer well past the cap: fence cells are single-line, and + // sanitizeTakeForPrompt only ever reads the first 500 chars anyway. + const line = 'a'.repeat(99); + const answer = Array.from({ length: 25 }, () => line).join('\n'); // 25*99 + 24 = 2499 chars + expect(answer.length).toBeGreaterThan(THINK_TAKE_CLAIM_MAX_CHARS); + + const persisted = await persistThinkTake(engine, synthResult(answer), { + anchor: 'notes/think-take-truncate-example', + }); + expect(persisted.rowNum).toBe(1); + expect(persisted.inserted).toBe(1); + expect(persisted.warnings).toContain('TAKE_CLAIM_TRUNCATED'); + + const takes = await engine.listTakes({ page_id: target.id }); + expect(takes).toHaveLength(1); + expect(takes[0].claim.length).toBe(THINK_TAKE_CLAIM_MAX_CHARS); + expect(takes[0].claim).not.toContain('\n'); // flattened before bounding + }); + test('refuses empty/no-LLM synthesis instead of writing a blank take', async () => { const target = await engine.putPage('notes/think-take-empty-example', { title: 'Think take empty target', type: 'note', diff --git a/test/think-take-cli.serial.test.ts b/test/think-take-cli.serial.test.ts new file mode 100644 index 000000000..ba196e6de --- /dev/null +++ b/test/think-take-cli.serial.test.ts @@ -0,0 +1,81 @@ +/** + * #2556 — `gbrain think --take` CLI exit-honesty pins. + * + * The shipped tests cover persistThinkTake (the core) and the MCP op + * handler, but nothing drives src/commands/think.ts's take path — the two + * exit(1) branches ("--take requires --anchor" and "--take requested but no + * take row was written") were untested. An explicit --take that writes + * nothing MUST be loud (same honesty contract as --save / #1698 F2). + * + * Serial: spies on process.exit + console.error (process-global), and the + * no-LLM path mutates ANTHROPIC_API_KEY/GBRAIN_HOME via withoutAnthropicKey. + */ +import { test, expect, beforeAll, afterAll, spyOn } from 'bun:test'; +import { PGLiteEngine } from '../src/core/pglite-engine.ts'; +import { runThinkCli } from '../src/commands/think.ts'; +import { withoutAnthropicKey } from './helpers/no-anthropic-key.ts'; + +let engine: PGLiteEngine; +let anchorPageId: number; + +beforeAll(async () => { + engine = new PGLiteEngine(); + await engine.connect({}); + await engine.initSchema(); + const p = await engine.putPage('notes/cli-take-anchor-example', { + title: 'CLI take anchor', type: 'note', + compiled_truth: 'A safe placeholder page for the CLI take exit tests.', + }); + anchorPageId = p.id; +}, 60_000); + +afterAll(async () => { + await engine.disconnect(); +}); + +/** + * Run runThinkCli with process.exit + console.error captured. The exit spy + * throws a sentinel so control flow actually stops at the exit site (a + * non-throwing mock would let the command keep running past exit(1)). + * Note: an exit(1) inside the runThink try-block is caught by the #1698 + * catch, which console.errors the sentinel message and exits again — the + * FIRST captured code is the one asserted. + */ +async function runExpectingExit(args: string[]): Promise<{ code: number | null; stderr: string }> { + const errors: string[] = []; + let code: number | null = null; + const errSpy = spyOn(console, 'error').mockImplementation((...a: unknown[]) => { + errors.push(a.map(String).join(' ')); + }); + const exitSpy = spyOn(process, 'exit').mockImplementation(((c?: number) => { + if (code === null) code = c ?? 0; + throw new Error(`EXIT:${c}`); + }) as never); + try { + await runThinkCli(engine, args); + } catch (e) { + if (!(e instanceof Error) || !e.message.startsWith('EXIT:')) throw e; + } finally { + exitSpy.mockRestore(); + errSpy.mockRestore(); + } + return { code, stderr: errors.join('\n') }; +} + +test('--take without --anchor exits 1 with the actionable message (parse-time guard)', async () => { + const { code, stderr } = await runExpectingExit(['what do we know', '--take']); + expect(code).toBe(1); + expect(stderr).toContain('--take requires --anchor'); +}); + +test('--take with no LLM (empty synthesis) exits 1 loudly and persists NOTHING', async () => { + const { code, stderr } = await withoutAnthropicKey(() => + runExpectingExit(['what do we know about this page', '--take', '--anchor', 'notes/cli-take-anchor-example']), + ); + // #2556 honesty contract: explicit --take that wrote nothing → exit 1 + // with the take-specific message, not a silent 0. + expect(code).toBe(1); + expect(stderr).toContain('--take requested but no take row was written'); + // And the DB really is untouched — no blank/stub take row. + expect(await engine.listTakes({ page_id: anchorPageId })).toHaveLength(0); +}, 30_000); diff --git a/test/think-take-concurrency.serial.test.ts b/test/think-take-concurrency.serial.test.ts new file mode 100644 index 000000000..85b4de3c9 --- /dev/null +++ b/test/think-take-concurrency.serial.test.ts @@ -0,0 +1,83 @@ +/** + * #2556 — persistThinkTake concurrency pin. + * + * The function's doc-comment CLAIMS the MAX(row_num)+1 computation + insert + * are serialized under withPageLock(anchor) so two concurrent `think --take` + * calls (or a takes add racing a think --take) can't produce duplicate + * row_nums. The shipped tests only exercise SEQUENTIAL calls (row 1 then + * row 2), which passes even with no lock at all. This test drives the two + * calls CONCURRENTLY via Promise.all and asserts they land on distinct + * consecutive rows — the behavioral proof the lock is actually keyed and + * actually serializes. + * + * Serial: the page lock is a real file under $GBRAIN_HOME/.gbrain/page-locks + * (process-global fs + env), and the loser of the race polls at 200ms. + */ +import { test, expect, beforeAll, afterAll } from 'bun:test'; +import { PGLiteEngine } from '../src/core/pglite-engine.ts'; +import { persistThinkTake } from '../src/core/think/index.ts'; +import { withEnv, emptyHome } from './helpers/with-env.ts'; + +let engine: PGLiteEngine; + +beforeAll(async () => { + engine = new PGLiteEngine(); + await engine.connect({}); + await engine.initSchema(); +}, 60_000); + +afterAll(async () => { + await engine.disconnect(); +}); + +function synthResult(answer: string): any { + return { + question: 'what should this page remember?', + answer, citations: [], gaps: [], pagesGathered: 0, takesGathered: 0, + graphHits: 0, modelUsed: 'stub', rounds: 1, warnings: [], synthesisOk: true, + diagnostics: { pagesFromHybrid: 0, takesFromKeyword: 0, takesFromVector: 0, graphHits: 0 }, + }; +} + +test('two CONCURRENT persistThinkTake calls on the same anchor serialize to rows 1 and 2 (no duplicate row_num)', async () => { + const page = await engine.putPage('notes/think-take-race-example', { + title: 'Think take race target', type: 'note', + compiled_truth: 'A safe placeholder page for concurrent think take persistence.', + }); + + // GBRAIN_HOME → fresh temp dir so the page-lock files land in a hermetic + // location (gbrainPath honors GBRAIN_HOME) instead of the user's ~/.gbrain. + await withEnv({ GBRAIN_HOME: emptyHome() }, async () => { + const [a, b] = await Promise.all([ + persistThinkTake(engine, synthResult('first concurrent synthesized insight'), { + anchor: 'notes/think-take-race-example', + }), + persistThinkTake(engine, synthResult('second concurrent synthesized insight'), { + anchor: 'notes/think-take-race-example', + }), + ]); + + // Both persisted and the row numbers are EXACTLY {1, 2} — a lost lock + // would yield {1, 1} (duplicate) or a failed insert. Wave-4 dual-plane + // contract: no brain repo is configured on this hermetic engine, so both + // appends fall back to DB-only allocation and surface + // TAKE_FILE_PLANE_UNAVAILABLE (still serialized under the page lock). + expect(a.inserted).toBe(1); + expect(b.inserted).toBe(1); + expect(a.warnings).toEqual(['TAKE_FILE_PLANE_UNAVAILABLE']); + expect(b.warnings).toEqual(['TAKE_FILE_PLANE_UNAVAILABLE']); + expect([a.rowNum, b.rowNum].sort()).toEqual([1, 2]); + }); + + // DB ground truth: two rows, distinct row_nums, both holder=brain. + const takes = await engine.listTakes({ page_id: page.id }); + expect(takes).toHaveLength(2); + expect(new Set(takes.map(t => t.row_num)).size).toBe(2); + const dupCheck = await engine.executeRaw<{ n: number }>( + `SELECT COUNT(*)::int AS n FROM takes + WHERE page_id = $1 + GROUP BY row_num HAVING COUNT(*) > 1`, + [page.id], + ); + expect(dupCheck).toHaveLength(0); +}, 15_000); diff --git a/test/think-take-op-federated.test.ts b/test/think-take-op-federated.test.ts new file mode 100644 index 000000000..a203e997a --- /dev/null +++ b/test/think-take-op-federated.test.ts @@ -0,0 +1,129 @@ +/** + * #2556 — think op handler: the scope-dialect mapping under a FEDERATED ctx. + * + * The op handler maps thinkSourceScopeOpts's `allowedSources` (runThink's + * dialect) onto persistThinkTake's `sourceIds` (the engine's getPage + * dialect). The inline comment warns that spreading thinkScope would + * "silently drop the federated array and unscope the anchor lookup" — but + * no shipped test constructs a ctx with `auth.allowedSources`, so that exact + * regression (a cross-source take write) would land green. + * + * These tests drive the FULL op handler to a successful synthesis using the + * gateway chat-transport test seam (no real LLM call), with the anchor page + * living in tenant-a: + * - grant includes tenant-a → take row lands (mapping threads the array) + * - grant is tenant-b only → TAKE_ANCHOR_NOT_FOUND, nothing written + * (fail-closed; an unscoped getPage would wrongly find the page) + */ +import { describe, test, expect, beforeAll, afterAll } from 'bun:test'; +import { PGLiteEngine } from '../src/core/pglite-engine.ts'; +import { operationsByName } from '../src/core/operations.ts'; +import { + __setChatTransportForTests, + __setEmbedTransportForTests, +} from '../src/core/ai/gateway.ts'; +import { withEnv, emptyHome } from './helpers/with-env.ts'; + +let engine: PGLiteEngine; +let anchorPageId: number; + +const ANSWER = 'Federated synthesis insight for the anchor page.'; + +beforeAll(async () => { + engine = new PGLiteEngine(); + await engine.connect({}); + await engine.initSchema(); + + for (const id of ['tenant-a', 'tenant-b']) { + await engine.executeRaw( + `INSERT INTO sources (id, name, local_path, config, created_at) + VALUES ($1, $1, NULL, '{}'::jsonb, NOW()) ON CONFLICT (id) DO NOTHING`, + [id], + ); + } + const p = await engine.putPage('people/fed-anchor-example', { + title: 'Federated anchor', type: 'person', + compiled_truth: 'A safe placeholder page living in tenant-a.', + }, { sourceId: 'tenant-a' }); + anchorPageId = p.id; + + // Canned successful synthesis — chat() calls the transport directly, + // skipping provider resolution/SDK; think parses the JSON envelope. + __setChatTransportForTests(async () => ({ + text: JSON.stringify({ answer: ANSWER, citations: [], gaps: [] }), + blocks: [], + stopReason: 'end' as const, + usage: { input_tokens: 10, output_tokens: 5, cache_read_tokens: 0, cache_creation_tokens: 0 }, + model: 'anthropic:claude-opus-4-7', + providerId: 'anthropic', + }) as any); + // Neutralize any query-embedding attempt during gather (fake key below + // must never reach the network). + __setEmbedTransportForTests((async (args: any) => ({ + embeddings: ((args?.values ?? []) as string[]).map(() => { + const v = new Array(1536).fill(0); + v[0] = 1; + return v; + }), + usage: { tokens: 1 }, + })) as any); +}, 60_000); + +afterAll(async () => { + __setChatTransportForTests(null); + __setEmbedTransportForTests(null); + await engine.disconnect(); +}); + +function fedCtx(allowedSources: string[]): any { + return { + engine, + config: {} as any, + dryRun: false, + remote: false, // trusted local — safeTake honored; scope still MUST confine + auth: { allowedSources }, + logger: { info() {}, warn() {}, error() {}, debug() {} } as any, + }; +} + +async function runThinkOpTake(allowedSources: string[]): Promise { + const op = operationsByName['think']; + // Fake key so tryBuildGatewayClient builds a client (transport is stubbed); + // empty GBRAIN_HOME so the real user config can't leak into model routing + // and the page-lock files land in a hermetic temp dir. + return withEnv({ ANTHROPIC_API_KEY: 'sk-test-fake', GBRAIN_HOME: emptyHome() }, () => + op.handler(fedCtx(allowedSources), { + question: 'what do we know about the federated anchor?', + anchor: 'people/fed-anchor-example', + take: true, + }), + ); +} + +describe('think op — federated allowedSources → sourceIds mapping (#2556)', () => { + test('in-grant federated ctx: take row lands on the anchor in the granted source', async () => { + const res = await runThinkOpTake(['tenant-a']); + expect(res.take_row).toBe(1); + expect(res.take_inserted).toBe(1); + expect(res.remote_persisted_blocked).toBe(false); + + const takes = await engine.listTakes({ page_id: anchorPageId }); + expect(takes).toHaveLength(1); + expect(takes[0]).toMatchObject({ + row_num: 1, claim: ANSWER, kind: 'take', holder: 'brain', source: 'gbrain think', + }); + }, 30_000); + + test('out-of-grant federated ctx: anchor lookup fail-closes — no cross-source take write', async () => { + const res = await runThinkOpTake(['tenant-b']); + // If the handler spread thinkScope (dropping the federated array), the + // getPage lookup would be UNSCOPED, find the tenant-a page, and write a + // cross-source row — this pins the mapping instead. + expect(res.take_row).toBeNull(); + expect(res.take_inserted).toBe(0); + expect((res.warnings as string[]).some(w => w.startsWith('TAKE_ANCHOR_NOT_FOUND'))).toBe(true); + + // Still exactly the one row from the in-grant test above. + expect(await engine.listTakes({ page_id: anchorPageId })).toHaveLength(1); + }, 30_000); +});