tryharness 0.7.18 → 0.7.23

package/lib/atlas.mjs

@@ -0,0 +1,304 @@
+import fs from "node:fs";
+import path from "node:path";
+import { createHash } from "node:crypto";
+import { spawnSync } from "node:child_process";
+import sqliteWasm from "node-sqlite3-wasm";
+import { fromBlob, toBlob } from "./vectors.mjs";
+
+// The relay's local index of the code its agents work in: every text file in the folders the swarm uses, cut into
+// pieces at top-level declarations and headings, each piece embedded once by the search model and found again by words
+// (SQLite's full-text index) and by meaning. Pieces are keyed by their repository path and text, not by commit, so a
+// new worktree or a small edit costs only the pieces that changed. Files that look like keys or secrets are never read.
+
+const { Database } = sqliteWasm;
+export const ATLAS_VERSION = "embeddinggemma-2:q8:768:v1";
+const MAX_FILE = 300_000;
+const MIN_PIECE = 300;
+const MAX_PIECE = 2_400;
+const TEXT_EXT = /\.(?:[cm]?[jt]sx?|json|md|mdx|txt|ya?ml|toml|sql|sol|rs|go|py|rb|java|kt|swift|c|cc|cpp|h|hpp|cs|php|sh|zsh|css|scss|html|vue|svelte|astro|graphql|proto|ini|cfg|conf|env\.example|lua|ex|exs|dart|scala|zig)$/i;
+const SKIP_PATH = /(?:^|\/)(?:node_modules|dist|build|out|\.next|coverage|vendor|target|\.git)\/|(?:^|\/)(?:package-lock\.json|pnpm-lock\.yaml|yarn\.lock|Cargo\.lock|go\.sum)$|\.min\.[jc]ss?$|\.map$/;
+const SECRET_PATH = /(?:^|\/)(?:\.env(?:\.[\w-]+)?|id_[rd]sa\w*|[\w.-]+\.(?:pem|key|p12|pfx|keystore|jks)|\.npmrc|\.pypirc|credentials(?:\.json)?)$/i;
+const SECRET_TEXT = /-----BEGIN (?:[A-Z]+ )?PRIVATE KEY-----/;
+const DECLARATION = /^(?:export\s|import\s|async\s|function\s|class\s|const\s|let\s|var\s|type\s|interface\s|enum\s|def\s|fn\s|pub\s|impl\s|struct\s|mod\s|func\s|package\s|contract\s|library\s|module\s|describe\(|it\(|test\(|#{1,6}\s|\[[\w.-]+\]$|create\s|alter\s)/i;
+const STOP = new Set("the a an and or of to in on for with is are be it this that from by at as was were can should would could will do does did not no yes we you i me my our your they them their its into out up so if then else than but about how what why when where which who whom also just like make made use used using need needs want".split(" "));
+
+/** Whether a repository path may be indexed: a text file, not generated, not a secret. */
+export const indexable = (rel, size = 0) => TEXT_EXT.test(rel) && !SKIP_PATH.test(rel) && !SECRET_PATH.test(rel) && size <= MAX_FILE;
+
+/** A file's pieces: { text, line } cut at top-level declarations and headings, small ones merged, big ones split at blank lines. */
+export function pieces(text) {
+  const lines = text.split("\n");
+  const segments = [];
+  let start = 0;
+  for (let i = 1; i <= lines.length; i += 1) {
+    const line = lines[i];
+    const boundary = i === lines.length || (line && !/^\s/.test(line) && (lines[i - 1].trim() === "" || DECLARATION.test(line)) && !/^[)}\]]/.test(line));
+    if (boundary) { segments.push({ line: start + 1, text: lines.slice(start, i).join("\n") }); start = i; }
+  }
+  const merged = [];
+  for (const segment of segments) {
+    const last = merged.at(-1);
+    if (last && last.text.length < MIN_PIECE) { last.text += `\n${segment.text}`; continue; }
+    merged.push({ ...segment });
+  }
+  const out = [];
+  for (const segment of merged) {
+    if (!segment.text.trim()) continue;
+    if (segment.text.length <= MAX_PIECE) { out.push(segment); continue; }
+    let buffer = "";
+    let line = segment.line;
+    let at = segment.line;
+    for (const block of segment.text.split(/\n(?=\s*\n)/)) {
+      if (buffer && buffer.length + block.length > MAX_PIECE) { out.push({ line, text: buffer }); buffer = ""; line = at; }
+      for (let i = 0; i < block.length; i += MAX_PIECE) {
+        const part = block.slice(i, i + MAX_PIECE);
+        if (buffer && buffer.length + part.length > MAX_PIECE) { out.push({ line, text: buffer }); buffer = ""; line = at; }
+        buffer += (buffer ? "\n" : "") + part;
+      }
+      at += block.split("\n").length;
+    }
+    if (buffer.trim()) out.push({ line, text: buffer });
+  }
+  return out;
+}
+
+export const pieceHash = (rel, text) => createHash("sha1").update(`${rel}\0${text}`).digest("hex");
+
+/** Words and identifiers worth searching for in a request, most specific first (identifiers, paths, then words). */
+export function queryTerms(text, max = 24) {
+  const raw = String(text).match(/[\p{L}\p{N}_$./-]{3,}/gu) || [];
+  const terms = [];
+  for (const token of raw) {
+    const clean = token.replace(/^[./-]+|[./-]+$/g, "");
+    if (clean.length < 3 || STOP.has(clean.toLowerCase())) continue;
+    for (const part of clean.split(/[./-]/)) if (part.length >= 3 && !STOP.has(part.toLowerCase())) terms.push(part);
+  }
+  const specific = (term) => (/[A-Z_$]/.test(term.slice(1)) || /\d/.test(term) ? 0 : 1);
+  return [...new Set(terms)].sort((a, b) => specific(a) - specific(b)).slice(0, max);
+}
+
+export function openAtlas(file) {
+  const db = new Database(file);
+  // The atlas is a cache the relay can always rebuild from the repositories and session files, so writes skip the disk sync.
+  db.exec("pragma journal_mode = MEMORY; pragma synchronous = OFF; pragma busy_timeout = 5000;");
+  db.exec(`
+    create table if not exists pieces (hash text primary key, path text not null, line integer not null, text text not null);
+    create table if not exists piece_vectors (hash text not null, version text not null, vector blob not null, primary key(hash, version));
+    create table if not exists blobs (blob text primary key, path text not null, pieces text not null);
+    create table if not exists files (root text not null, path text not null, blob text not null, indexed_at text not null, primary key(root, path));
+    create virtual table if not exists pieces_fts using fts5(hash unindexed, path, text, tokenize='unicode61 tokenchars ''_$''');
+    create table if not exists runs (agent text not null, session_id text not null, root text not null, at integer not null, brief text not null,
+      edited text not null, read text not null, before text not null, source_mtime real not null, primary key(agent, session_id));
+    create index if not exists runs_root_at on runs(root, at);
+    create table if not exists run_vectors (agent text not null, session_id text not null, version text not null, vector blob not null, primary key(agent, session_id, version));
+    create table if not exists meta (key text primary key, value text not null);
+    create table if not exists media (hash text primary key, path text not null, kind text not null, start_ms integer not null, end_ms integer not null, mtime real not null);
+    create index if not exists media_path on media(path);
+  `);
+  // A blob's pieces never change, so lookups are kept (a few MB for a large repository and its recent versions).
+  const blobPieces = new Map();
+  const run = (sql, ...values) => db.run(sql, values.length ? values : undefined);
+  const all = (sql, ...values) => db.all(sql, values.length ? values : undefined);
+  const get = (sql, ...values) => db.get(sql, values.length ? values : undefined);
+  return {
+    close: () => db.close(),
+    blob(blob) {
+      if (blobPieces.has(blob)) return blobPieces.get(blob);
+      const row = get("select pieces from blobs where blob = ?", blob);
+      if (!row) return null;
+      if (blobPieces.size >= 200_000) blobPieces.clear();
+      blobPieces.set(blob, JSON.parse(row.pieces));
+      return blobPieces.get(blob);
+    },
+    /** Runs fn inside one transaction. */
+    batch(fn) {
+      db.exec("begin");
+      try { const out = fn(); db.exec("commit"); return out; } catch (error) { db.exec("rollback"); throw error; }
+    },
+    putBlob(blob, rel, parts) {
+      const hashes = [];
+      for (const part of parts) {
+        const hash = pieceHash(rel, part.text);
+        hashes.push(hash);
+        if (!get("select 1 as one from pieces where hash = ?", hash)) {
+          run("insert into pieces(hash, path, line, text) values(?, ?, ?, ?)", hash, rel, part.line, part.text);
+          run("insert into pieces_fts(hash, path, text) values(?, ?, ?)", hash, rel.replace(/[/._-]/g, " "), part.text);
+        }
+      }
+      run("insert or replace into blobs(blob, path, pieces) values(?, ?, ?)", blob, rel, JSON.stringify(hashes));
+      return hashes;
+    },
+    setFiles(root, rows) {
+      db.exec("begin");
+      try {
+        run("delete from files where root = ?", root);
+        const now = new Date().toISOString();
+        for (const row of rows) run("insert into files(root, path, blob, indexed_at) values(?, ?, ?, ?)", root, row.path, row.blob, now);
+        db.exec("commit");
+      } catch (error) { db.exec("rollback"); throw error; }
+    },
+    files: (root) => all("select path, blob from files where root = ?", root),
+    roots: () => all("select root, count(*) as files, max(indexed_at) as indexed_at from files group by root").map((row) => ({ ...row, files: Number(row.files) })),
+    toEmbed: (version, limit = 64) => all("select p.hash, p.path, p.text from pieces p left join piece_vectors v on v.hash = p.hash and v.version = ? where v.hash is null limit ?", version, limit),
+    missing: (version) => Number(get("select count(*) as n from pieces p left join piece_vectors v on v.hash = p.hash and v.version = ? where v.hash is null", version).n),
+    putVector: (hash, version, vector) => run("insert or replace into piece_vectors(hash, version, vector) values(?, ?, ?)", hash, version, toBlob(vector)),
+    vectors(hashes, version) {
+      const out = new Map();
+      const list = [...hashes];
+      for (let i = 0; i < list.length; i += 500) {
+        const slice = list.slice(i, i + 500);
+        for (const row of all(`select hash, vector from piece_vectors where version = ? and hash in (${slice.map(() => "?").join(",")})`, version, ...slice)) out.set(row.hash, fromBlob(row.vector));
+      }
+      return out;
+    },
+    piece: (hash) => get("select hash, path, line, text from pieces where hash = ?", hash),
+    meta(key, value) {
+      if (value === undefined) return get("select value from meta where key = ?", key)?.value ?? null;
+      run("insert or replace into meta(key, value) values(?, ?)", key, String(value));
+      return value;
+    },
+    runSeen: (agent, sessionId) => get("select source_mtime from runs where agent = ? and session_id = ?", agent, sessionId)?.source_mtime ?? null,
+    putRun(row) {
+      run(`insert or replace into runs(agent, session_id, root, at, brief, edited, read, before, source_mtime) values(?, ?, ?, ?, ?, ?, ?, ?, ?)`,
+        row.agent, row.sessionId, row.root, row.at, row.brief.slice(0, 4_000), JSON.stringify(row.edited), JSON.stringify(row.read), JSON.stringify(row.before), row.sourceMtime);
+    },
+    runs: (root, since = 0) => all("select agent, session_id, root, at, brief, edited, read, before from runs where root = ? and at >= ? order by at", root, since)
+      .map((row) => ({ ...row, at: Number(row.at), edited: JSON.parse(row.edited), read: JSON.parse(row.read), before: JSON.parse(row.before) })),
+    runsToEmbed: (version, limit = 32) => all("select r.agent, r.session_id, r.brief from runs r left join run_vectors v on v.agent = r.agent and v.session_id = r.session_id and v.version = ? where v.agent is null limit ?", version, limit),
+    putRunVector: (agent, sessionId, version, vector) => run("insert or replace into run_vectors(agent, session_id, version, vector) values(?, ?, ?, ?)", agent, sessionId, version, toBlob(vector)),
+    runVectors: (root, version) => all("select r.agent, r.session_id, r.at, r.brief, r.edited, v.vector from runs r join run_vectors v on v.agent = r.agent and v.session_id = r.session_id and v.version = ? where r.root = ?", version, root)
+      .map((row) => ({ agent: row.agent, sessionId: row.session_id, at: Number(row.at), brief: row.brief, edited: JSON.parse(row.edited), vector: fromBlob(row.vector) })),
+    /** Recordings' windows and images (media.mjs): one row each, its vector beside the pieces'. */
+    mediaSeen: (file) => get("select max(mtime) as mtime from media where path = ?", file)?.mtime ?? null,
+    mediaPaths: () => all("select distinct path from media").map((row) => row.path),
+    dropMedia: (file) => run("delete from media where path = ?", file),
+    putMedia: (row) => run("insert or replace into media(hash, path, kind, start_ms, end_ms, mtime) values(?, ?, ?, ?, ?, ?)", row.hash, row.path, row.kind, row.startMs, row.endMs, row.mtime),
+    mediaVectors: (version) => all("select m.hash, m.path, m.kind, m.start_ms, m.end_ms, v.vector from media m join piece_vectors v on v.hash = m.hash and v.version = ?", version)
+      .map((row) => ({ hash: row.hash, path: row.path, kind: row.kind, startMs: Number(row.start_ms), endMs: Number(row.end_ms), vector: fromBlob(row.vector) })),
+    /** Full-text matches: { hash, rank } best first (bm25, lower is better), over every indexed piece. */
+    words(terms, limit = 400) {
+      if (!terms.length) return [];
+      const match = terms.map((term) => `"${term.replaceAll('"', '""')}"`).join(" OR ");
+      return all("select hash, bm25(pieces_fts, 0, 2.0, 1.0) as rank from pieces_fts where pieces_fts match ? order by rank limit ?", match, limit);
+    },
+  };
+}
+
+const git = (cwd, args, input) => spawnSync("git", args, { cwd, input, encoding: "utf8", maxBuffer: 256 * 2 ** 20 });
+
+/** The repository's files at a commit (or the working tree when commit is null): [{ path, blob, size }]. */
+export function treeFiles(root, commit = null) {
+  if (commit) {
+    const listed = git(root, ["ls-tree", "-r", "-l", commit]);
+    if (listed.status !== 0) return [];
+    return listed.stdout.split("\n").filter(Boolean).map((line) => {
+      const [meta, rel] = line.split("\t");
+      const [, type, blob, size] = meta.split(/\s+/);
+      return type === "blob" ? { path: rel, blob, size: Number(size) || 0 } : null;
+    }).filter((row) => row && indexable(row.path, row.size));
+  }
+  // The working tree: git's index already holds the blob of every unchanged tracked file; only files changed on disk
+  // and new files git does not ignore are hashed, so a pass over a large repository costs about what `git status` does.
+  const staged = git(root, ["ls-files", "-s", "-z"]);
+  if (staged.status !== 0) return [];
+  const blobs = new Map();
+  for (const entry of staged.stdout.split("\0")) {
+    const tab = entry.indexOf("\t");
+    if (tab < 0) continue;
+    const [, blob, stage] = entry.slice(0, tab).split(" ");
+    if (stage === "0") blobs.set(entry.slice(tab + 1), blob);
+  }
+  const changed = git(root, ["diff", "--name-only", "-z"]).stdout.split("\0").filter(Boolean);
+  const untracked = git(root, ["ls-files", "--others", "--exclude-standard", "-z"]).stdout.split("\0").filter(Boolean);
+  const rows = [];
+  const rehash = [];
+  for (const rel of new Set([...blobs.keys(), ...untracked])) {
+    if (!indexable(rel, 0)) continue;
+    let stat;
+    try { stat = fs.statSync(path.join(root, rel)); } catch { continue; }
+    if (!stat.isFile() || stat.size > MAX_FILE) continue;
+    const row = { path: rel, size: stat.size, blob: blobs.get(rel) };
+    rows.push(row);
+    if (!row.blob || changed.includes(rel)) rehash.push(row);
+  }
+  if (rehash.length) {
+    const hashed = git(root, ["hash-object", "--stdin-paths"], rehash.map((row) => path.join(root, row.path)).join("\n"));
+    const out = hashed.status === 0 ? hashed.stdout.split("\n") : [];
+    rehash.forEach((row, i) => { row.blob = out[i]; });
+  }
+  return rows.filter((row) => /^[0-9a-f]{40}$/.test(row.blob || ""));
+}
+
+/** File contents for blobs, read through one `git cat-file --batch`. */
+export function readBlobs(root, blobs) {
+  const out = new Map();
+  if (!blobs.length) return out;
+  const result = spawnSync("git", ["cat-file", "--batch"], { cwd: root, input: blobs.join("\n") + "\n", maxBuffer: 2 * 2 ** 30 });
+  const buffer = result.stdout;
+  let at = 0;
+  while (at < buffer.length) {
+    const end = buffer.indexOf(10, at);
+    if (end < 0) break;
+    const [blob, type, size] = buffer.subarray(at, end).toString("utf8").split(" ");
+    if (type === "missing" || !size) { at = end + 1; continue; }
+    const length = Number(size);
+    const body = buffer.subarray(end + 1, end + 1 + length);
+    at = end + 1 + length + 1;
+    if (!body.includes(0) && !SECRET_TEXT.test(body.toString("latin1"))) out.set(blob, body.toString("utf8"));
+  }
+  return out;
+}
+
+/** Adds the files' pieces to the atlas (new blobs only); returns how many blobs were new. */
+export function addFiles(atlas, root, rows) {
+  const fresh = rows.filter((row) => !atlas.blob(row.blob));
+  const texts = readBlobs(root, [...new Set(fresh.map((row) => row.blob))]);
+  for (let i = 0; i < fresh.length; i += 500) atlas.batch(() => {
+    for (const row of fresh.slice(i, i + 500)) {
+      const text = texts.get(row.blob);
+      atlas.putBlob(row.blob, row.path, text === undefined ? [] : pieces(text));
+    }
+  });
+  return fresh.length;
+}
+
+/** Embeds every piece that has no vector yet, in batches, on the given lane. */
+export async function embedMissing(atlas, inference, { lane = "background", version = ATLAS_VERSION, batch = 32, onProgress } = {}) {
+  let done = 0;
+  for (;;) {
+    const rows = atlas.toEmbed(version, batch);
+    if (!rows.length) return done;
+    const vectors = await inference.embed(rows.map((row) => ({ kind: "document", title: row.path, text: row.text })), { lane });
+    atlas.batch(() => rows.forEach((row, i) => atlas.putVector(row.hash, version, vectors[i])));
+    done += rows.length;
+    onProgress?.(done);
+  }
+}
+
+/**
+ * Files in scope ranked for a request, with the signals each one scored on. scope: [{ path, blob }]. Returns
+ * [{ path, blob, dense, lexical, named, best: { hash, line } }] sorted by the given scorer (default: dense+lexical+named).
+ */
+export async function searchFiles(atlas, inference, query, scope, { version = ATLAS_VERSION, lane = "interactive", queryVector = null, vectors = null, wordLimit = 400 } = {}) {
+  const byHash = new Map();
+  const files = scope.map((row) => ({ path: row.path, blob: row.blob, hashes: atlas.blob(row.blob) ?? [] }));
+  for (const file of files) for (const hash of file.hashes) { if (!byHash.has(hash)) byHash.set(hash, []); byHash.get(hash).push(file); }
+  const qv = queryVector ?? (await inference.embed([{ kind: "code-query", text: String(query).slice(0, 6_000) }], { lane }))[0];
+  const pieceVectors = vectors ?? atlas.vectors(byHash.keys(), version);
+  const terms = queryTerms(query);
+  const lexical = new Map();
+  for (const row of atlas.words(terms, wordLimit)) if (byHash.has(row.hash)) lexical.set(row.hash, -row.rank);
+  const lowered = String(query).toLowerCase();
+  return files.map((file) => {
+    let dense = -1, best = null, words = 0;
+    for (const hash of file.hashes) {
+      const vector = pieceVectors.get(hash);
+      if (vector) { let dot = 0; for (let i = 0; i < qv.length; i += 1) dot += qv[i] * vector[i]; if (dot > dense) { dense = dot; best = hash; } }
+      if (lexical.has(hash)) words = Math.max(words, lexical.get(hash));
+    }
+    const base = path.basename(file.path).toLowerCase();
+    const stem = base.replace(/\.[^.]+$/, "");
+    const named = lowered.includes(file.path.toLowerCase()) ? 2 : lowered.includes(base) ? 1.5 : stem.length > 3 && new RegExp(`\\b${stem.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}\\b`).test(lowered) ? 1 : 0;
+    return { path: file.path, blob: file.blob, dense, lexical: words, named, best };
+  });
+}