diff --git a/gitnexus/src/core/embeddings/staged-embedding-recovery-child.ts b/gitnexus/src/core/embeddings/staged-embedding-recovery-child.ts new file mode 100644 index 000000000..708a8d511 --- /dev/null +++ b/gitnexus/src/core/embeddings/staged-embedding-recovery-child.ts @@ -0,0 +1,118 @@ +/** Isolated strict native reader. Never import this entrypoint into analyze. */ +import fs from 'node:fs'; +import path from 'node:path'; +import lbug from '@ladybugdb/core'; +import { createLbugDatabase, toNativeSafePath } from '../lbug/lbug-config.js'; +import { + abortCachedEmbeddingsBuilder, + createCachedEmbeddingsBuilder, + finalizeCachedEmbeddingsSnapshot, + ingestCachedEmbeddingRow, +} from './embedding-restore-spill.js'; +import type { StagedEmbeddingExport } from './staged-embedding-recovery.js'; + +async function extract(): Promise { + const [dbPath, exportDir, dimensionsArg] = process.argv.slice(2); + const dimensions = Number(dimensionsArg); + if (!dbPath || !exportDir || !Number.isInteger(dimensions) || dimensions <= 0) { + throw new Error('invalid staged embedding extraction arguments'); + } + const stat = fs.lstatSync(dbPath); + if (!stat.isFile() || stat.isSymbolicLink()) + throw new Error('staged embedding DB is not a regular file'); + const builder = createCachedEmbeddingsBuilder({ inMemoryRowLimit: 0, spillDir: exportDir }); + const rejectedNodeIds = new Set(); + // Avoid openLbugConnection's test-fixture lock sweep: a recovery source must + // never have its WAL removed, even when an external slot resembles a fixture. + const db = createLbugDatabase(lbug, toNativeSafePath(dbPath), { throwOnWalReplayFailure: true }); + const handle = { db, conn: new lbug.Connection(db) }; + try { + const queried = await handle.conn.query( + 'MATCH (e:CodeEmbedding) RETURN e.nodeId AS nodeId, e.chunkIndex AS chunkIndex, e.startLine AS startLine, e.endLine AS endLine, e.embedding AS embedding, e.contentHash AS contentHash', + ); + const results = Array.isArray(queried) ? queried : [queried]; + try { + if (results.length !== 1) throw new Error('unexpected staged embedding query result'); + const result = results[0]; + while (await result.hasNext()) { + const raw = await result.getNext(); + const rec = raw as Record & unknown[]; + const nodeId = rec.nodeId ?? rec[0]; + if (typeof nodeId !== 'string' || !nodeId) + throw new Error('invalid staged embedding node id'); + const chunkIndex = rec.chunkIndex ?? rec[1]; + const startLine = rec.startLine ?? rec[2]; + const endLine = rec.endLine ?? rec[3]; + const embedding = rec.embedding ?? rec[4]; + const contentHash = rec.contentHash ?? rec[5]; + const vector = + Array.isArray(embedding) || + (ArrayBuffer.isView(embedding) && !(embedding instanceof DataView)) + ? Array.from(embedding as ArrayLike) + : undefined; + if ( + !Number.isInteger(chunkIndex) || + Number(chunkIndex) < 0 || + !Number.isInteger(startLine) || + Number(startLine) < 0 || + !Number.isInteger(endLine) || + Number(endLine) < Number(startLine) || + typeof contentHash !== 'string' || + !contentHash || + !vector || + vector.length !== dimensions || + vector.some( + (value) => + typeof value !== 'number' || + !Number.isFinite(value) || + !Number.isFinite(Math.fround(value)), + ) + ) { + rejectedNodeIds.add(nodeId); + continue; + } + ingestCachedEmbeddingRow( + builder, + { nodeId, chunkIndex, startLine, endLine, embedding: vector, contentHash }, + true, + ); + } + } finally { + for (const result of results) await result.close(); + } + // Both closes must succeed. Suppressed native teardown errors are unsafe. + await handle.conn.close(); + await handle.db.close(); + const snapshot = finalizeCachedEmbeddingsSnapshot(builder); + if (snapshot.spill) fs.renameSync(snapshot.spill.path, path.join(exportDir, 'vectors.bin')); + const manifest: StagedEmbeddingExport = { + version: 1, + dimensions, + rows: snapshot.rows, + rejectedNodeIds: [...rejectedNodeIds], + }; + fs.writeFileSync(path.join(exportDir, 'manifest.json'), JSON.stringify(manifest), { + flag: 'wx', + mode: 0o600, + }); + } catch (err) { + abortCachedEmbeddingsBuilder(builder); + // Cleanup is best effort on a rejected source, never used to approve output. + try { + await handle.conn.close(); + } catch { + /* rejected */ + } + try { + await handle.db.close(); + } catch { + /* rejected */ + } + throw err; + } +} + +extract().catch((err: unknown) => { + process.stderr.write(`${err instanceof Error ? err.message : String(err)}\n`); + process.exitCode = 1; +}); diff --git a/gitnexus/src/core/embeddings/staged-embedding-recovery.ts b/gitnexus/src/core/embeddings/staged-embedding-recovery.ts new file mode 100644 index 000000000..b295f4beb --- /dev/null +++ b/gitnexus/src/core/embeddings/staged-embedding-recovery.ts @@ -0,0 +1,244 @@ +/** Recover paid embedding rows without opening an interrupted native DB in analyze. */ +import { spawn } from 'node:child_process'; +import fs from 'node:fs'; +import os from 'node:os'; +import path from 'node:path'; +import { fileURLToPath } from 'node:url'; +import { + abortCachedEmbeddingsBuilder, + createCachedEmbeddingsBuilder, + EmbeddingSpillReader, + emptyCachedEmbeddingsSnapshot, + finalizeCachedEmbeddingsSnapshot, + ingestCachedEmbeddingRow, + materializeCachedEmbeddings, + type CachedEmbeddingMeta, + type CachedEmbeddingsSnapshot, +} from './embedding-restore-spill.js'; + +export interface StagedEmbeddingRecoveryOptions { + dimensions: number; + /** Active checkpoint window and any incomplete inherited restore groups. */ + excludedNodeIds?: Iterable; + timeoutMs?: number; +} + +export interface StagedEmbeddingExport { + version: 1; + dimensions: number; + rows: CachedEmbeddingMeta[]; + rejectedNodeIds: string[]; +} + +const RESTORE_BATCH_SIZE = 200; +const DEFAULT_EXTRACTION_TIMEOUT_MS = 120_000; + +/** A contiguous prefix is safe only when its node is outside every unsafe window. */ +export function validateRecoveredNodeGroups( + rows: readonly CachedEmbeddingMeta[], + excludedNodeIds: ReadonlySet = new Set(), +): Set { + const groups = new Map; invalid: boolean }>(); + for (const row of rows) { + let group = groups.get(row.nodeId); + if (!group) { + group = { hash: row.contentHash, ordinals: new Set(), invalid: false }; + groups.set(row.nodeId, group); + } + if ( + typeof row.nodeId !== 'string' || + !row.nodeId || + typeof row.contentHash !== 'string' || + !row.contentHash || + row.contentHash !== group.hash || + !Number.isInteger(row.chunkIndex) || + row.chunkIndex < 0 || + group.ordinals.has(row.chunkIndex) || + !Number.isInteger(row.startLine) || + row.startLine < 0 || + !Number.isInteger(row.endLine) || + row.endLine < row.startLine + ) + group.invalid = true; + group.ordinals.add(row.chunkIndex); + } + const accepted = new Set(); + for (const [nodeId, group] of groups) { + if (group.invalid || excludedNodeIds.has(nodeId)) continue; + // Unique ordinals with max n-1 and zero present have no holes. + if (!group.ordinals.has(0)) continue; + if ([...group.ordinals].some((ordinal) => ordinal >= group.ordinals.size)) continue; + accepted.add(nodeId); + } + return accepted; +} + +/** Whole recovered nodes replace whole published groups; vectors stay in bounded batches. */ +export function mergeRecoveredEmbeddings( + live: CachedEmbeddingsSnapshot, + recovered: CachedEmbeddingsSnapshot, +): CachedEmbeddingsSnapshot { + const builder = createCachedEmbeddingsBuilder({ inMemoryRowLimit: 0 }); + try { + appendSnapshotRows( + live, + live.rows.filter((row) => !recovered.embeddingNodeIds.has(row.nodeId)), + builder, + ); + appendSnapshotRows(recovered, recovered.rows, builder); + return finalizeCachedEmbeddingsSnapshot(builder); + } catch (err) { + abortCachedEmbeddingsBuilder(builder); + throw err; + } +} + +function appendSnapshotRows( + snapshot: CachedEmbeddingsSnapshot, + rows: readonly CachedEmbeddingMeta[], + builder: ReturnType, +): void { + const reader = snapshot.spill ? new EmbeddingSpillReader(snapshot.spill) : undefined; + try { + for (let i = 0; i < rows.length; i += RESTORE_BATCH_SIZE) { + const batch = materializeCachedEmbeddings( + snapshot, + rows.slice(i, i + RESTORE_BATCH_SIZE), + reader, + ); + for (const row of batch) + ingestCachedEmbeddingRow(builder, row as unknown as Record, true); + } + } finally { + reader?.close(); + } +} + +/** + * Caller must validate checkpoint identity, schema, exact generation, and hold the + * index lock. A subprocess contains native WAL replay/query/destructor failures. + * A failed strict open is never retried with WAL removed or validation disabled. + */ +export async function recoverStagedEmbeddings( + dbPath: string, + options: StagedEmbeddingRecoveryOptions, +): Promise { + if (!Number.isInteger(options.dimensions) || options.dimensions <= 0) { + throw new Error('invalid staged embedding dimensions'); + } + const dbStat = fs.lstatSync(dbPath); + if (!dbStat.isFile() || dbStat.isSymbolicLink()) + throw new Error('staged embedding DB is not a regular file'); + const exportDir = fs.mkdtempSync(path.join(os.tmpdir(), 'gitnexus-stage-export-')); + try { + await runExtractionChild(dbPath, exportDir, options); + const manifestPath = path.join(exportDir, 'manifest.json'); + const manifest = JSON.parse(fs.readFileSync(manifestPath, 'utf8')) as StagedEmbeddingExport; + if ( + manifest.version !== 1 || + manifest.dimensions !== options.dimensions || + !Array.isArray(manifest.rows) || + !Array.isArray(manifest.rejectedNodeIds) || + manifest.rejectedNodeIds.some((id) => typeof id !== 'string') || + manifest.rows.some((row, i) => !row || row.vectorIndex !== i) + ) + throw new Error('invalid staged embedding export manifest'); + if (manifest.rows.length === 0) return emptyCachedEmbeddingsSnapshot(); + const spill = { + path: path.join(exportDir, 'vectors.bin'), + dims: options.dimensions, + rowCount: manifest.rows.length, + }; + const vectorStat = fs.lstatSync(spill.path); + if ( + !vectorStat.isFile() || + vectorStat.isSymbolicLink() || + vectorStat.size !== 12 + spill.rowCount * spill.dims * 4 + ) { + throw new Error('invalid staged embedding export size'); + } + const excluded = new Set(options.excludedNodeIds ?? []); + for (const nodeId of manifest.rejectedNodeIds) excluded.add(nodeId); + const accepted = validateRecoveredNodeGroups(manifest.rows, excluded); + const exported: CachedEmbeddingsSnapshot = { + rows: manifest.rows, + embeddings: [], + embeddingNodeIds: accepted, + spill, + }; + const builder = createCachedEmbeddingsBuilder({ inMemoryRowLimit: 0 }); + const reader = new EmbeddingSpillReader(spill); + try { + const rows = manifest.rows.filter((row) => accepted.has(row.nodeId)); + for (let i = 0; i < rows.length; i += RESTORE_BATCH_SIZE) { + for (const row of materializeCachedEmbeddings( + exported, + rows.slice(i, i + RESTORE_BATCH_SIZE), + reader, + )) { + if ( + row.embedding.length !== options.dimensions || + row.embedding.some((value) => !Number.isFinite(value)) + ) { + throw new Error('invalid staged embedding vector'); + } + ingestCachedEmbeddingRow(builder, row as unknown as Record, true); + } + } + return finalizeCachedEmbeddingsSnapshot(builder); + } catch (err) { + abortCachedEmbeddingsBuilder(builder); + throw err; + } finally { + reader.close(); + } + } finally { + fs.rmSync(exportDir, { recursive: true, force: true }); + } +} + +function runExtractionChild( + dbPath: string, + exportDir: string, + options: StagedEmbeddingRecoveryOptions, +): Promise { + const compiledPath = fileURLToPath( + new URL('./staged-embedding-recovery-child.js', import.meta.url), + ); + const sourcePath = compiledPath.replace(/\.js$/, '.ts'); + const childPath = fs.existsSync(compiledPath) ? compiledPath : sourcePath; + const args = childPath.endsWith('.ts') + ? ['--import', import.meta.resolve('tsx'), childPath] + : [childPath]; + args.push(dbPath, exportDir, String(options.dimensions)); + return new Promise((resolve, reject) => { + const child = spawn(process.execPath, args, { + stdio: ['ignore', 'ignore', 'pipe'], + windowsHide: true, + }); + let stderr = ''; + let timedOut = false; + child.stderr.on('data', (chunk: Buffer) => { + if (stderr.length < 4096) stderr += chunk.toString().slice(0, 4096 - stderr.length); + }); + const timer = setTimeout(() => { + timedOut = true; + // A process boundary is safe to kill even while native code is executing. + child.kill('SIGKILL'); + }, options.timeoutMs ?? DEFAULT_EXTRACTION_TIMEOUT_MS); + child.once('error', (err) => { + clearTimeout(timer); + reject(err); + }); + child.once('close', (code, signal) => { + clearTimeout(timer); + if (code === 0 && !signal && !timedOut) resolve(); + else + reject( + new Error( + `staged embedding extraction failed${timedOut ? ' (timeout)' : ` (${signal ?? code})`}${stderr ? `: ${stderr.trim()}` : ''}`, + ), + ); + }); + }); +}