brane 0.1.1 → 0.3.1

package/dist/memory/embedder.js

@@ -0,0 +1,145 @@
+/**
+ * Feature extraction on `onnxruntime-node` directly — the slice of
+ * `@huggingface/transformers` brane actually used.
+ *
+ * The omnibus package brought 138 MB of code for four steps: tokenize, run one
+ * ONNX session, mean-pool over the attention mask, L2-normalize. This file is
+ * those four steps. It is written to produce byte-identical vectors to the
+ * pipeline it replaces, because `memory_vec` is already full of the old ones —
+ * see `src/test/embedder.test.ts`, which pins that against real chunks.
+ *
+ * The on-disk cache layout is deliberately the one transformers.js used
+ * (`<cacheDir>/<repo>/<file>`), so an existing `~/.brane/models` keeps working
+ * and nobody re-downloads 133 MB on upgrade.
+ */
+import { createWriteStream, existsSync, mkdirSync, renameSync } from "node:fs";
+import { availableParallelism, homedir } from "node:os";
+import { dirname, join } from "node:path";
+import { Readable } from "node:stream";
+import { pipeline as streamPipeline } from "node:stream/promises";
+import * as ort from "onnxruntime-node";
+import { WordPieceTokenizer } from "./wordpiece.js";
+/** The model brane embeds with. */
+export const MODEL_REPO = "Xenova/bge-small-en-v1.5";
+/** Its hidden size — the width of every vector in `memory_vec`. */
+export const VECTOR_DIM = 384;
+const HF_ENDPOINT = process.env.HF_ENDPOINT || "https://huggingface.co";
+function cacheDir() {
+    return process.env.BRANE_MODEL_CACHE || join(homedir(), ".brane", "models");
+}
+/**
+ * Download one file to the cache if it isn't already there.
+ *
+ * Writes to a temporary name and renames into place: a model interrupted
+ * halfway is 133 MB of plausible-looking ONNX that fails at session load with
+ * nothing to suggest a truncated download. The rename makes presence mean
+ * complete. (Never removes the partial file — the golden rule applies to
+ * brane's own caches too; a stale `.partial` costs disk, and disk is cheap.)
+ */
+export async function ensureFile(repo, file) {
+    const target = join(cacheDir(), repo, file);
+    if (existsSync(target))
+        return target;
+    mkdirSync(dirname(target), { recursive: true });
+    const url = `${HF_ENDPOINT}/${repo}/resolve/main/${file}`;
+    const res = await fetch(url);
+    if (!res.ok || !res.body) {
+        throw new Error(`Failed to fetch ${url}: ${res.status} ${res.statusText}`);
+    }
+    const partial = `${target}.partial-${process.pid}`;
+    await streamPipeline(Readable.fromWeb(res.body), createWriteStream(partial));
+    renameSync(partial, target);
+    return target;
+}
+/**
+ * How many cores one embedding session may use.
+ *
+ * This number was measured, not chosen. Capping ORT at a single intra-op
+ * thread — to be a good neighbour to the daemon's event loop and the 10 fps
+ * tmux capture — made embedding **6.5x slower** than the transformers.js
+ * pipeline it replaces (308 ms/text against 47), which would have turned a
+ * dependency cleanup into a visible regression in indexing. Scaling on a
+ * 16-core machine: 1→308, 2→152, 4→80, 8→47.6, unbounded→49.7 ms/text. It
+ * saturates at 8; letting ORT take everything is marginally *worse* than
+ * taking eight.
+ *
+ * So: leave two cores for everything else, and never exceed the point where
+ * more threads stop buying anything.
+ */
+const MAX_INTRA_OP_THREADS = 8;
+function intraOpThreads() {
+    return Math.max(1, Math.min(MAX_INTRA_OP_THREADS, availableParallelism() - 2));
+}
+/** Load the model (downloading on a cold cache) and return an embedder. */
+export async function createEmbedder() {
+    const tokenizerPath = await ensureFile(MODEL_REPO, "tokenizer.json");
+    const modelPath = await ensureFile(MODEL_REPO, "onnx/model.onnx");
+    const tokenizer = WordPieceTokenizer.fromFile(tokenizerPath);
+    const session = await ort.InferenceSession.create(modelPath, {
+        executionProviders: ["cpu"],
+        graphOptimizationLevel: "all",
+        intraOpNumThreads: intraOpThreads(),
+    });
+    return {
+        async embed(texts) {
+            if (texts.length === 0)
+                return [];
+            const { input_ids, attention_mask, token_type_ids } = tokenizer.encodeBatch(texts);
+            const batch = input_ids.length;
+            const seqLen = input_ids[0].length;
+            const dims = [batch, seqLen];
+            const toTensor = (rows) => new ort.Tensor("int64", BigInt64Array.from(rows.flat().map(BigInt)), dims);
+            const output = await session.run({
+                input_ids: toTensor(input_ids),
+                attention_mask: toTensor(attention_mask),
+                token_type_ids: toTensor(token_type_ids),
+            });
+            const hidden = output.last_hidden_state;
+            const data = hidden.data;
+            const embedDim = hidden.dims[2];
+            const vectors = [];
+            for (let b = 0; b < batch; b++) {
+                const vec = meanPool(data, b * seqLen * embedDim, seqLen, embedDim, attention_mask[b]);
+                l2Normalize(vec);
+                vectors.push(vec);
+            }
+            return vectors;
+        },
+        async dispose() {
+            await session.release();
+        },
+    };
+}
+/**
+ * Masked mean pool one sequence out of a `[batch, seq, dim]` hidden state.
+ *
+ * Padding positions carry mask 0 and so contribute nothing to either the sum
+ * or the divisor — which is what lets a text embed identically whether it was
+ * batched with longer neighbours or sent alone. Divides by the token count,
+ * with no epsilon, matching transformers.js's `mean_pooling` exactly; a
+ * sequence always contains at least [CLS] and [SEP], so the count is never 0.
+ */
+export function meanPool(hidden, base, seqLen, embedDim, mask) {
+    const vec = new Float32Array(embedDim);
+    let count = 0;
+    for (let j = 0; j < seqLen; j++)
+        count += mask[j];
+    for (let k = 0; k < embedDim; k++) {
+        let sum = 0;
+        for (let j = 0; j < seqLen; j++) {
+            sum += hidden[base + j * embedDim + k] * mask[j];
+        }
+        vec[k] = sum / count;
+    }
+    return vec;
+}
+/** L2-normalize in place, so cosine distance in `memory_vec` is a dot product. */
+export function l2Normalize(vec) {
+    let norm = 0;
+    for (let k = 0; k < vec.length; k++)
+        norm += vec[k] * vec[k];
+    norm = Math.sqrt(norm);
+    for (let k = 0; k < vec.length; k++)
+        vec[k] /= norm;
+}
+//# sourceMappingURL=embedder.js.map

package/dist/memory/embedding-worker.js

@@ -0,0 +1,71 @@
+import { parentPort } from "node:worker_threads";
+import { createEmbedder, VECTOR_DIM } from "./embedder.js";
+if (!parentPort) {
+    throw new Error("embedding-worker.ts must be loaded as a worker thread");
+}
+let embedderPromise = null;
+async function loadExtractor() {
+    console.log("Loading embedding model...");
+    // The first load may fetch the model from HuggingFace, which can be
+    // transiently rate-limited (HTTP 429) or otherwise flaky. Retry with
+    // exponential backoff so a cold cache survives a hiccup instead of failing
+    // the whole process (notably CI). Once the model is on disk this loop
+    // succeeds on the first attempt and never touches the network.
+    const MAX_ATTEMPTS = 4;
+    let lastErr;
+    for (let attempt = 1; attempt <= MAX_ATTEMPTS; attempt++) {
+        try {
+            const embedder = await createEmbedder();
+            console.log("Embedding model ready.");
+            return embedder;
+        }
+        catch (err) {
+            lastErr = err;
+            if (attempt === MAX_ATTEMPTS)
+                break;
+            const delayMs = 1000 * 2 ** (attempt - 1); // 1s, 2s, 4s
+            console.warn(`Embedding model load failed (attempt ${attempt}/${MAX_ATTEMPTS}), ` +
+                `retrying in ${delayMs}ms:`, err instanceof Error ? err.message : err);
+            await new Promise((resolve) => setTimeout(resolve, delayMs));
+        }
+    }
+    throw lastErr;
+}
+function getExtractor() {
+    if (!embedderPromise) {
+        embedderPromise = loadExtractor().catch((err) => {
+            embedderPromise = null;
+            throw err;
+        });
+    }
+    return embedderPromise;
+}
+parentPort.on("message", async (msg) => {
+    if (msg.kind === "preload") {
+        getExtractor().catch((err) => {
+            console.error("Failed to preload embedding model:", err);
+        });
+        return;
+    }
+    if (msg.kind === "embed") {
+        try {
+            const embedder = await getExtractor();
+            const vectors = await embedder.embed(msg.texts);
+            // Concatenate into one buffer so the whole batch crosses the worker
+            // boundary as a single transfer rather than N copies; `embedBatch`
+            // slices it back apart on the other side.
+            const flat = new Float32Array(vectors.length * VECTOR_DIM);
+            for (let i = 0; i < vectors.length; i++)
+                flat.set(vectors[i], i * VECTOR_DIM);
+            parentPort.postMessage({ kind: "embed-result", id: msg.id, flat }, [flat.buffer]);
+        }
+        catch (err) {
+            parentPort.postMessage({
+                kind: "embed-error",
+                id: msg.id,
+                message: err instanceof Error ? err.message : String(err),
+            });
+        }
+    }
+});
+//# sourceMappingURL=embedding-worker.js.map

package/dist/memory/embeddings.js

@@ -0,0 +1,91 @@
+import { Worker } from "node:worker_threads";
+// Re-exported, not redeclared: the dimensionality belongs to the model, and
+// `storage.ts` bakes it into the vec0 schema. Two copies could disagree.
+export { VECTOR_DIM } from "./embedder.js";
+import { VECTOR_DIM } from "./embedder.js";
+let worker = null;
+let nextId = 0;
+const pending = new Map();
+function failAllPending(err) {
+    for (const p of pending.values())
+        p.reject(err);
+    pending.clear();
+}
+function getWorker() {
+    if (worker)
+        return worker;
+    const w = new Worker(new URL("./embedding-worker.js", import.meta.url));
+    w.on("message", (msg) => {
+        const p = pending.get(msg.id);
+        if (!p)
+            return;
+        pending.delete(msg.id);
+        if (msg.kind === "embed-result")
+            p.resolve(msg.flat);
+        else
+            p.reject(new Error(msg.message));
+    });
+    w.on("error", (err) => {
+        failAllPending(err);
+        worker = null;
+    });
+    w.on("exit", () => {
+        failAllPending(new Error("embedding worker exited"));
+        worker = null;
+    });
+    worker = w;
+    return w;
+}
+/** Preload the embedding model in the background. Safe to call multiple times. */
+export function preloadEmbeddingModel() {
+    try {
+        getWorker().postMessage({ kind: "preload" });
+    }
+    catch (err) {
+        console.error("Failed to preload embedding model:", err);
+    }
+}
+/**
+ * Tear down the embedding worker. Calling `worker.terminate()` kills the
+ * worker thread and releases its ONNX session — crucially without invoking
+ * the main process's `__cxa_finalize`, which is where the
+ * `pthread_mutex_lock: Invalid argument` abort used to come from.
+ */
+export async function disposeEmbeddingModel() {
+    if (!worker)
+        return;
+    const w = worker;
+    worker = null;
+    try {
+        await w.terminate();
+    }
+    catch {
+        // Worker already gone; nothing to do.
+    }
+    failAllPending(new Error("embedding worker terminated"));
+}
+function sendEmbed(texts) {
+    const w = getWorker();
+    const id = nextId++;
+    return new Promise((resolve, reject) => {
+        pending.set(id, { resolve, reject });
+        w.postMessage({ kind: "embed", id, texts });
+    });
+}
+/** Embed a single text string, returns normalized 384-dim vector. */
+export async function embedText(text) {
+    return sendEmbed([text]);
+}
+/** Embed multiple texts in one call. Returns array of 384-dim vectors. */
+export async function embedBatch(texts) {
+    if (texts.length === 0)
+        return [];
+    const flat = await sendEmbed(texts);
+    const vectors = [];
+    for (let i = 0; i < texts.length; i++) {
+        const start = i * VECTOR_DIM;
+        vectors.push(flat.slice(start, start + VECTOR_DIM));
+    }
+    return vectors;
+}
+//# sourceMappingURL=embeddings.js.map

package/dist/memory/index.js

@@ -0,0 +1,232 @@
+import { embedBatch } from "./embeddings.js";
+import { initMemorySchema, getLastIndexedSeq, setLastIndexedSeq, upsertChunks, } from "./storage.js";
+import { getAgentAdapter, TRACE_RESOLVABLE_AGENTS } from "../agents.js";
+import { fromPortablePath } from "../db/schema.js";
+/** Max events to embed per batch (memory bound). */
+const BATCH_SIZE = 64;
+/**
+ * Cooperative-scheduling bounds for `resolvePendingTraces`. The resolver runs
+ * in-process on the daemon's main thread and each resolution does a blocking
+ * transcript read, so a large backfill could otherwise starve narrate/listen/
+ * terminal. We yield to the event loop between reads and cap reads per pass;
+ * resolvable events beyond the cap roll to the next index tick (resolved rows
+ * drop out of the gate, so progress stays monotonic).
+ */
+const RESOLVE_YIELD_EVERY_READS = 16;
+const RESOLVE_YIELD_EVERY_ROWS = 2048;
+const MAX_RESOLVES_PER_PASS = 2000;
+/** Yield control to the event loop so the daemon can service requests. */
+const yieldToEventLoop = () => new Promise((r) => setImmediate(r));
+/**
+ * Resolve per-message provenance (meta.session.offset + meta.model) for
+ * narration/agent_signal events that don't yet have it. Runs *before*
+ * embedding and is gated purely on the *absence* of meta.model — never on
+ * the embedding cursor — so it both maintains new events and self-heals
+ * (backfills) any historical event still missing the attribution, with no
+ * events stranded behind a cursor. Idempotent: a resolved event is skipped
+ * forever, and agent session logs are immutable so a value never goes stale.
+ *
+ * meta.session is *author-relative*: each event's author is its `source`
+ * narrative, so we resolve against the author's agent/fold/session — which
+ * makes foreign narrations attribute to the authoring agent, not the host.
+ */
+export async function resolvePendingTraces(db, signal, opts = {}) {
+    // No trace-capable agent → nothing is ever resolvable; skip the scan.
+    if (TRACE_RESOLVABLE_AGENTS.length === 0)
+        return 0;
+    // Push every precondition into SQL so we only materialize rows the resolver
+    // could actually act on — instead of pulling the whole unresolved backlog
+    // into JS to skip most of it. A row qualifies only when:
+    //   - its author agent (events.source → narratives) has a resolveTrace
+    //     adapter (foreign narrations attribute to the *author*, not the host),
+    //   - meta.model is still absent (gate is absence-of-model, never a cursor,
+    //     so the resolver self-heals), and
+    //   - a session id is available — either the live-stamped meta.session.id or
+    //     the author narrative's agent_session_id fallback. Without one there is
+    //     no transcript to read, so such rows are pure waste to materialize.
+    // The CASE is ordered so json_extract only runs on valid JSON: it *throws*
+    // "malformed JSON" on a corrupt blob, which would otherwise abort the pass.
+    // appendEvent only writes valid JSON or NULL, so corrupt rows are reachable
+    // only via a direct DB edit — they're skipped, not allowed to take down the
+    // whole pass.
+    const agentPlaceholders = TRACE_RESOLVABLE_AGENTS.map(() => "?").join(", ");
+    const rows = db
+        .prepare(`SELECT e.narrative_id AS narrative_id, e.seq AS seq, e.content AS content,
+              e.source AS source, e.meta AS meta,
+              n.agent_type AS agent_type, n.fold_path AS fold_path,
+              n.agent_session_id AS agent_session_id
+       FROM events e
+       JOIN narratives n ON n.id = e.source
+       WHERE e.type IN ('narration', 'agent_signal')
+         AND e.content IS NOT NULL AND e.content != ''
+         AND n.agent_type IN (${agentPlaceholders})
+         AND n.fold_path IS NOT NULL
+         AND CASE
+               WHEN e.meta IS NOT NULL AND json_valid(e.meta) = 0 THEN 0
+               WHEN e.meta IS NOT NULL AND json_extract(e.meta, '$.model') IS NOT NULL THEN 0
+               WHEN e.meta IS NOT NULL AND json_extract(e.meta, '$.session.id') IS NOT NULL THEN 1
+               WHEN n.agent_session_id IS NOT NULL THEN 1
+               ELSE 0
+             END = 1`)
+        .all(...TRACE_RESOLVABLE_AGENTS);
+    // Memoize each author's whole-session trace index so a transcript is read and
+    // parsed at most once per pass, however many of its events are pending. A
+    // `null` (missing/unreadable session file) is cached too, so a dead session
+    // isn't re-read for every one of its events.
+    const indexCache = new Map();
+    // Phase 1: compute all updates with NO DB write lock held (building a trace
+    // index does file IO of the agent's own session log, which can be slow).
+    // Yields to the event loop around transcript parses and caps parses per pass
+    // so this pass can't starve the daemon's main thread (see bounds above).
+    const updates = [];
+    let reads = 0;
+    for (let i = 0; i < rows.length; i++) {
+        if (signal?.aborted)
+            break;
+        // Coarse yield so even a long run of cheap (cache-hit) rows can't block.
+        if (i > 0 && i % RESOLVE_YIELD_EVERY_ROWS === 0)
+            await yieldToEventLoop();
+        const row = rows[i];
+        if (!row.agent_type || !row.fold_path)
+            continue; // SQL-guarded; defensive
+        const adapter = getAgentAdapter(row.agent_type);
+        if (!adapter.buildTraceIndex)
+            continue; // unsupported agent — leave unset
+        const meta = row.meta ? JSON.parse(row.meta) : {};
+        // Prefer the session id stamped live at write time; fall back to the
+        // author narrative's current session (best-effort for pre-capture events).
+        const sessionId = meta.session?.id ?? row.agent_session_id ?? undefined;
+        if (!sessionId)
+            continue; // SQL-guarded; defensive
+        const foldPath = fromPortablePath(row.fold_path);
+        // JSON-stringify the pair for an unambiguous key - no separator can collide.
+        const cacheKey = JSON.stringify([sessionId, foldPath]);
+        let index = indexCache.get(cacheKey);
+        if (index === undefined) {
+            // Cache miss → parse this transcript once (the expensive op). Bound the
+            // parses per pass and yield around them; the rest roll to the next tick.
+            // Cheap cache-hit lookups below never consume this budget.
+            if (reads >= MAX_RESOLVES_PER_PASS)
+                break;
+            index = adapter.buildTraceIndex({ sessionId, foldPath });
+            indexCache.set(cacheKey, index);
+            reads++;
+            if (reads % RESOLVE_YIELD_EVERY_READS === 0)
+                await yieldToEventLoop();
+        }
+        if (!index)
+            continue; // session file gone/unreadable
+        const trace = index.get(row.content);
+        if (!trace || (!trace.model && !trace.offset))
+            continue;
+        const offset = trace.offset ?? meta.session?.offset;
+        meta.session = offset ? { id: sessionId, offset } : { id: sessionId };
+        if (trace.model)
+            meta.model = trace.model;
+        updates.push({ narrative_id: row.narrative_id, seq: row.seq, meta: JSON.stringify(meta) });
+    }
+    if (opts.dryRun || updates.length === 0)
+        return updates.length;
+    // Phase 2: apply every update in a single transaction (atomic, minimal WAL
+    // churn, short lock window). Purely additive — never touches content.
+    const stmt = db.prepare("UPDATE events SET meta = ? WHERE narrative_id = ? AND seq = ?");
+    const apply = db.transaction((items) => {
+        for (const u of items)
+            stmt.run(u.meta, u.narrative_id, u.seq);
+    });
+    apply(updates);
+    return updates.length;
+}
+/**
+ * Run incremental indexing across all narratives in brane.db.
+ *
+ * For each narrative, compares max(seq) against last_indexed_seq cursor.
+ * Only embeds events with seq > cursor. Idempotent — safe to run repeatedly.
+ */
+export async function runIndex(db, rebuild, signal) {
+    initMemorySchema(db);
+    // Resolve per-message model/session provenance before embedding. Gated on
+    // absence of meta.model, independent of the embedding cursor (see fn docs).
+    const tracesResolved = await resolvePendingTraces(db, signal);
+    if (rebuild) {
+        db.exec(`
+      DELETE FROM memory_chunks;
+      DELETE FROM memory_vec;
+      DELETE FROM memory_content_df;
+      DELETE FROM memory_stats;
+      DELETE FROM memory_cursors;
+      INSERT INTO memory_fts(memory_fts) VALUES('rebuild');
+    `);
+    }
+    // Find all narratives with events
+    const narratives = db
+        .prepare(`SELECT narrative_id, MAX(seq) as max_seq, COUNT(*) as event_count
+       FROM events
+       GROUP BY narrative_id`)
+        .all();
+    let narrativesIndexed = 0;
+    let totalChunks = 0;
+    for (const nar of narratives) {
+        if (signal?.aborted)
+            break;
+        const lastSeq = getLastIndexedSeq(db, nar.narrative_id);
+        if (nar.max_seq <= lastSeq)
+            continue;
+        // Fetch unindexed events with textual content
+        const events = db
+            .prepare(`SELECT narrative_id, seq, time, type, content
+         FROM events
+         WHERE narrative_id = ? AND seq > ? AND content IS NOT NULL AND content != ''
+         ORDER BY seq ASC`)
+            .all(nar.narrative_id, lastSeq);
+        if (events.length === 0) {
+            // No textual events, but advance cursor to avoid re-scanning
+            setLastIndexedSeq(db, nar.narrative_id, nar.max_seq);
+            continue;
+        }
+        console.log(`[${narrativesIndexed + 1}] Embedding ${nar.narrative_id} (${events.length} events)...`);
+        // Process in batches
+        for (let i = 0; i < events.length; i += BATCH_SIZE) {
+            if (signal?.aborted)
+                break;
+            const batch = events.slice(i, i + BATCH_SIZE);
+            const texts = batch.map((e) => e.content);
+            const vectors = await embedBatch(texts);
+            const chunks = batch.map((e, j) => ({
+                id: `${e.narrative_id}-${e.seq}`,
+                narrative_id: e.narrative_id,
+                seq: e.seq,
+                content: e.content,
+                time: e.time,
+                vector: vectors[j],
+            }));
+            upsertChunks(db, chunks);
+            totalChunks += chunks.length;
+        }
+        setLastIndexedSeq(db, nar.narrative_id, nar.max_seq);
+        narrativesIndexed++;
+    }
+    return {
+        narratives_scanned: narratives.length,
+        narratives_indexed: narrativesIndexed,
+        chunks_created: totalChunks,
+        traces_resolved: tracesResolved,
+    };
+}
+/**
+ * Check if any narrative has unindexed events.
+ * Cheap check — no embedding, just seq comparisons.
+ */
+export function hasUnindexedEvents(db) {
+    initMemorySchema(db);
+    const rows = db
+        .prepare(`SELECT e.narrative_id, MAX(e.seq) as max_seq
+       FROM events e
+       LEFT JOIN memory_cursors mc ON mc.narrative_id = e.narrative_id
+       WHERE e.seq > COALESCE(mc.last_seq, 0)
+       GROUP BY e.narrative_id
+       LIMIT 1`)
+        .all();
+    return rows.length > 0;
+}
+//# sourceMappingURL=index.js.map