/** * #3390 regression: page straddling a stale-batch BOUNDARY (adversarial review * blocker 2). * * The embed loop stamps `pages.embedding_signature` only when * `stale.length === existing.length` — i.e. when every chunk of the page landed * in the SAME batch. `listStaleChunks` is a plain keyset LIMIT with no page * alignment, so on any corpus bigger than one batch the page split across the * boundary is embedded correctly but never stamped. * * Pre-fix consequence: the completion probe counted that page stale, the * command printed "Migration incomplete" and exited 1 on a perfectly-migrated * brain, and the re-run RE-INVALIDATED and RE-PAID for those pages — * contradicting the "already-migrated chunks are never re-embedded" contract. * * Shape here: 3 pages x 2 chunks = 6 chunks with --batch-size 3, so the * boundary falls mid-page-2. Asserts exit 0 on the first run, every page * stamped, and ZERO embed work on the second run. * * Named `.serial.test.ts`: holds a temp GBRAIN_HOME + an installed fake embed * transport for its whole beforeAll→afterAll lifecycle, which withEnv() can't * wrap. */ import { describe, test, expect, beforeAll, afterAll } from 'bun:test'; import { mkdtempSync, mkdirSync, writeFileSync, rmSync } from 'fs'; import { join } from 'path'; import { tmpdir } from 'os'; import { PGLiteEngine } from '../src/core/pglite-engine.ts'; import { configureGateway, resetGateway, __setEmbedTransportForTests, } from '../src/core/ai/gateway.ts'; import { runEmbedCore } from '../src/commands/embed.ts'; import { runMigrateEmbeddings } from '../src/commands/migrate-embeddings.ts'; import { MIGRATION_STATE_KEY, MIGRATION_COMPLETED_KEY } from '../src/core/embedding-migration.ts'; const FROM_DIMS = 1280; const TO_DIMS = 1536; const PAGES = ['b-1', 'b-2', 'b-3']; const PROBE_TEXT = 'gbrain embedding migration probe'; let engine: PGLiteEngine; let tmpHome: string; const savedEnv: Record = {}; let currentDims = FROM_DIMS; let embeddedTexts: string[] = []; class ExitError extends Error { constructor(public code: number) { super(`exit ${code}`); } } const exitSeam = (code: number): never => { throw new ExitError(code); }; async function runMigrate(args: string[]): Promise { try { await runMigrateEmbeddings(engine, args, { exit: exitSeam }); throw new Error('runMigrateEmbeddings returned without exiting'); } catch (e) { if (e instanceof ExitError) return e.code; throw e; } } beforeAll(async () => { for (const k of ['GBRAIN_HOME', 'GBRAIN_EMBEDDING_MODEL', 'GBRAIN_EMBEDDING_DIMENSIONS', 'OPENAI_API_KEY', 'ZEROENTROPY_API_KEY', 'DATABASE_URL']) { savedEnv[k] = process.env[k]; delete process.env[k]; } tmpHome = mkdtempSync(join(tmpdir(), 'gbrain-migrate-boundary-')); process.env.GBRAIN_HOME = tmpHome; mkdirSync(join(tmpHome, '.gbrain'), { recursive: true }); writeFileSync(join(tmpHome, '.gbrain', 'config.json'), JSON.stringify({ engine: 'pglite', embedding_model: 'zeroentropyai:zembed-1', embedding_dimensions: FROM_DIMS, zeroentropy_api_key: 'ze-test-fake', openai_api_key: 'sk-test-fake', }, null, 2)); resetGateway(); configureGateway({ embedding_model: 'zeroentropyai:zembed-1', embedding_dimensions: FROM_DIMS, env: { ZEROENTROPY_API_KEY: 'ze-test-fake', OPENAI_API_KEY: 'sk-test-fake' }, }); __setEmbedTransportForTests(async ({ values }: { values: string[] }) => { for (const v of values) if (v !== PROBE_TEXT) embeddedTexts.push(v); return { embeddings: values.map(() => new Array(currentDims).fill(0).map((_, i) => Math.sin(i) * 0.01 + 0.003)), usage: { tokens: values.length * 4 }, } as never; }); engine = new PGLiteEngine(); await engine.connect({ embedding_dimensions: FROM_DIMS } as never); await engine.initSchema(); // 3 pages x 2 chunks each = 6 chunks. for (const slug of PAGES) { await engine.putPage(slug, { type: 'note', title: slug, compiled_truth: `# ${slug}` }); await engine.upsertChunks(slug, [ { chunk_index: 0, chunk_text: `${slug} chunk zero`, chunk_source: 'compiled_truth', token_count: 4 }, { chunk_index: 1, chunk_text: `${slug} chunk one`, chunk_source: 'compiled_truth', token_count: 4 }, ]); } await runEmbedCore(engine, { stale: true, quiet: true }); }, 60000); afterAll(async () => { __setEmbedTransportForTests(null); resetGateway(); await engine.disconnect(); rmSync(tmpHome, { recursive: true, force: true }); for (const [k, v] of Object.entries(savedEnv)) { if (v === undefined) delete process.env[k]; else process.env[k] = v; } }); describe('migration across a stale-batch boundary', () => { test('seed is fully embedded at the source width', async () => { expect(await engine.countStaleChunks()).toBe(0); const n = await engine.executeRaw<{ n: number }>( `SELECT count(*)::int AS n FROM content_chunks WHERE embedding IS NOT NULL`, ); expect(Number(n[0]?.n)).toBe(6); }); test('batch-size 3 splits a page across the boundary yet still exits 0 with every page stamped', async () => { currentDims = TO_DIMS; embeddedTexts = []; const code = await runMigrate([ '--to', 'openai:text-embedding-3-small', '--yes', '--batch-size', '3', ]); // Pre-fix this was 1 ("Migration incomplete") even though all 6 chunks // were correctly embedded — the boundary page was never stamped. expect(code).toBe(0); expect(embeddedTexts.length).toBe(6); // all six chunks re-embedded once expect(await engine.countStaleChunks()).toBe(0); // Every page carries the TARGET signature, including the boundary page. const sig = `openai:text-embedding-3-small:${TO_DIMS}`; const stamped = await engine.executeRaw<{ n: number }>( `SELECT count(*)::int AS n FROM pages WHERE embedding_signature = $1`, [sig], ); expect(Number(stamped[0]?.n)).toBe(PAGES.length); // Completion bookkeeping ran (it only runs when the backlog drained). expect(await engine.getConfig(MIGRATION_STATE_KEY)).toBeFalsy(); expect(await engine.getConfig(MIGRATION_COMPLETED_KEY)).toBeTruthy(); }, 60000); test('second run does ZERO work — the "never re-embedded twice" contract holds across a boundary', async () => { embeddedTexts = []; const code = await runMigrate([ '--to', 'openai:text-embedding-3-small', '--yes', '--batch-size', '3', ]); expect(code).toBe(0); // Pre-fix the unstamped boundary page was re-invalidated and PAID FOR again. expect(embeddedTexts.length).toBe(0); }, 60000); });