core.js

// barber core — query-conditioned pruning of retrieved/tool context.
//
// Port of barber/core.py, which is itself the validated algorithm from
// Nadir's context_selection.py. The selection logic is decision-identical to
// the benchmarked Python implementation and is verified against golden
// fixtures generated by the Python package (test/fixtures/golden.json).
// Do NOT improve the selection logic here: the defaults are benchmark-locked
// (keep=0.6, relative floor 0.35, assertive drop marker, all guards on), and
// any behavior change must land in the Python source first and re-benchmark.
//
// Cross-language parity notes (each guards a real divergence):
// - pyRound: Python round() is round-half-to-even; Math.round is half-up.
// - charLen: Python len() counts code points; String.length counts UTF-16
//   units, which differs on emoji and other astral characters.
// - Term vectors are Maps (see embedders.js).
// - Pin patterns use explicit \p{L}\p{N}_ lookarounds because Python's \b is
//   Unicode-aware and JS's \b is ASCII-only.

import { createHash } from "node:crypto";
import { lexical, tokenize } from "./embedders.js";

// ---------------------------------------------------------------------------
// Config
// ---------------------------------------------------------------------------

export const DEFAULT_CONFIG = Object.freeze({
  // Only consider a message for selection if it is at least this many chars
  // (small messages aren't retrieved-context; leave them alone). Sized for
  // Latin script: 800 CJK characters is several thousand tokens, so a Chinese
  // or Japanese workload wants this a few hundred lower or selection rarely
  // fires at all.
  minMessageChars: 800,
  // Only split/select when the message has at least this many chunks.
  minChunks: 4,
  // Keep-ratio bounds (fraction of chunks retained) after tier conditioning.
  minKeepRatio: 0.25,
  maxKeepRatio: 1.0,
  // A chunk scoring below (top_score * relativeFloor) is a drop candidate even
  // if the budget would keep it — protects against keeping pure noise.
  relativeFloor: 0.35,
  // Always keep the first & last chunk of a block (lead/tail bias).
  keepLeadTail: true,
  // Marker inserted where a run of chunks was dropped. The assertive wording
  // is benchmark-locked: it beat the neutral variant in the published A/B.
  dropMarker:
    "[… {n} passage(s) omitted as not relevant to this question — the remaining context is sufficient …]",
  // Roles whose large messages are selectable (never the latest user query).
  selectableRoles: Object.freeze(["user", "tool", "function"]),
  // Patterns whose match makes a chunk never-droppable. Defaults are English
  // (see PIN_PATTERNS); pass your own to pin another language's deontic and
  // policy vocabulary. Empty array = pin on rare query entities only.
  pinPatterns: null, // null -> PIN_PATTERNS
});

// Tier -> target keep ratio, for router integrations that know task
// complexity. Null prototype: `tier` is caller input and must not resolve
// "constructor" or "toString" to a ratio.
export const TIER_KEEP_RATIO = Object.freeze(
  Object.assign(Object.create(null), {
    simple: 0.45,
    mid: 0.55,
    medium: 0.55,
    complex: 0.7,
    reasoning: 0.85,
  })
);

// ---------------------------------------------------------------------------
// Cross-language helpers
// ---------------------------------------------------------------------------

// Python round(): round-half-to-even on the double. x is always >= 0 here.
export function pyRound(x) {
  const f = Math.floor(x);
  if (x - f === 0.5) return f % 2 === 0 ? f : f + 1;
  return Math.round(x);
}

// Python len(): code points, not UTF-16 units.
export function charLen(s) {
  let n = 0;
  for (const _ of s) n++;
  return n;
}

export function md5key(s) {
  return createHash("md5").update(s, "utf8").digest("hex").slice(0, 12);
}

// ---------------------------------------------------------------------------
// Pinning — never-drop patterns (the silent-failure guard)
// ---------------------------------------------------------------------------

// ENGLISH ONLY, and unavoidably so: "must" and "ne doit pas" are not one
// regex. These are the DEFAULT — override `config.pinPatterns` with your own
// language's vocabulary. Nothing else here is language-locked (the tokenizer
// is Unicode-aware), so this list is all a non-English caller has to supply.
export const PIN_PATTERNS = [
  // deontic/safety
  /(?<![\p{L}\p{N}_])(?:must|never|do not|don't|shall not|prohibited|required|only if)(?![\p{L}\p{N}_])/iu,
  /(?<![\p{L}\p{N}_])(?:PII|HIPAA|PCI|SSN|password|secret|api[_ ]?key)(?![\p{L}\p{N}_])/iu,
];
// NOTE: deliberately no pinning on bare digits — most chunks contain a number,
// so digit-pinning neuters selection. Numeric chunks that matter are kept by
// relevance scoring; rare query entities protect the answer-bearing chunk.

function isPinned(chunk, rareQueryEntities, patterns = PIN_PATTERNS) {
  for (const p of patterns) {
    if (p.test(chunk)) return true;
  }
  const low = chunk.toLowerCase();
  for (const e of rareQueryEntities) {
    if (low.includes(e)) return true;
  }
  return false;
}

// ---------------------------------------------------------------------------
// Chunking & scoring
// ---------------------------------------------------------------------------

// Line-oriented shapes an agent harness produces. Prose has blank lines
// between paragraphs; a file read, a grep, a diff, or an ls does not, so the
// paragraph splitter returns one chunk and selection declines the block.
const LINE_NO = /^\s*\d+\t/; // `cat -n` / Read line prefix
const HUNK = /^(?:diff --git |@@ )/; // unified diff

function group(lines, isBoundary, drop) {
  const out = [];
  let cur = [];
  for (const l of lines) {
    if (isBoundary(l)) {
      if (drop) {
        if (cur.length) {
          out.push(cur.join("\n"));
          cur = [];
        }
        continue;
      }
      if (cur.length) {
        out.push(cur.join("\n"));
        cur = [];
      }
    }
    cur.push(l);
  }
  if (cur.length) out.push(cur.join("\n"));
  return out;
}

function splitLines(text, minChunks) {
  // A JSON body is one value: dropping lines out of the middle yields
  // something that no longer parses. Leave it whole.
  const first = text.replace(/^\s+/, "").charAt(0);
  if (first === "{" || first === "[") return [];
  const lines = text.split("\n");
  if (lines.length < minChunks) return [];

  // Line-numbered file read: strip the prefix to find the blank lines of the
  // underlying file, then chunk on those.
  const head = lines.slice(0, 40);
  if (head.filter((l) => LINE_NO.test(l)).length >= head.length * 0.8) {
    const groups = group(lines, (l) => !l.replace(LINE_NO, "").trim(), true);
    if (groups.length >= minChunks) return groups;
  }

  // Unified diff: one chunk per hunk.
  if (lines.slice(0, 5).some((l) => HUNK.test(l))) {
    const groups = group(lines, (l) => HUNK.test(l), false);
    if (groups.length >= minChunks) return groups;
  }

  // Anything else line-oriented (grep hits, file lists, logs): fixed windows.
  const step = Math.max(1, Math.floor(lines.length / 12));
  const groups = [];
  for (let i = 0; i < lines.length; i += step) {
    groups.push(lines.slice(i, i + step).join("\n"));
  }
  return groups.length >= minChunks ? groups : [];
}

// Split a block into selectable chunks: prefer blank-line / delimiter
// boundaries (RAG concatenations, tool rows), fall back to sentences, then to
// line structure for the line-oriented output an agent harness produces.
// The group is CAPTURING so split hands the delimiter back instead of eating
// it: some of these are separators, but some are part of the chunk they
// introduce. Mirrors _DELIM in core.py.
const DELIM = /(\n\s*\n|\n-{3,}\n|\n#{1,6}\s|\[\d+\]\s)/;

// The part of a delimiter that BELONGS TO the chunk it introduces. A blank line
// separates two chunks and belongs to neither; a heading marker (`## `) and a
// citation marker (`[1] `) are the first characters of the passage that
// follows, and dropping them silently rewrites content barber promised to keep
// verbatim. Mirrors _lead_of in core.py.
// The lead is the delimiter's own bytes, never a reconstruction: rebuilding it
// as `stripped + " "` fabricated a space whenever \s matched a newline, tab, or
// NBSP, so `##\nheading` came back as `## heading` — a string not in the source.
function leadOf(delim) {
  if (!delim.trim()) return "";                 // blank line: pure separator
  return delim.replace(/^\n+/, "");             // drop only the anchoring newline
}

export function splitChunks(text, minChunks = DEFAULT_CONFIG.minChunks) {
  // String.split with one capturing group yields [text, delim, text, delim, ...]
  const pieces = text.split(DELIM);
  const parts = [];
  let bodies = 0;                 // parts backed by real content, not a bare marker
  const first = pieces[0].trim();
  if (first) {
    parts.push(first);
    bodies++;
  }
  for (let i = 1; i < pieces.length; i += 2) {
    const raw = i + 1 < pieces.length ? pieces[i + 1] : "";
    const lead = leadOf(pieces[i]);
    // With a lead attached, keep the body's LEADING whitespace so lead + body is
    // a contiguous slice of the source and a chunk is verifiably a substring of
    // the input. Without a lead the delimiter was a pure separator and both ends
    // trim, which is the behaviour every fixture was generated against.
    const body = raw.trim() ? (lead ? raw.replace(/\s+$/, "") : raw.trim()) : "";
    if (!body) {
      // Keep a bare marker's bytes, but do NOT count it as a chunk: counting it
      // let a single trailing `---` reach the >= 2 test below and return a
      // 2-chunk block, which _select_block then declined for being under
      // minChunks — silently disabling trimming for a block the sentence
      // splitter used to handle.
      if (lead.trim()) parts.push(lead.trim());
      continue;
    }
    parts.push(lead + body);
    bodies++;
  }
  // Delimiters found: this is a chunked block, and how many chunks it has is
  // the answer. Too few to bother cutting is a real signal, not a failure to
  // look harder — do not fall through.
  if (bodies >= 2) return parts;
  const sents = text
    .trim()
    .split(/(?<=[.!?])\s+/)
    .filter(Boolean);
  if (sents.length >= minChunks) return sents;
  // Nothing found any structure, so selection was about to decline the block
  // whole. Line layout is the last place to look.
  const byLine = splitLines(text, minChunks);
  return byLine.length ? byLine : sents;
}

function cosBag(a, b) {
  if (!a.size || !b.size) return 0.0;
  let num = 0;
  for (const [k, va] of a) {
    const vb = b.get(k);
    if (vb !== undefined) num += va * vb;
  }
  let na = 0;
  for (const v of a.values()) na += v * v;
  let nb = 0;
  for (const v of b.values()) nb += v * v;
  na = Math.sqrt(na);
  nb = Math.sqrt(nb);
  return na && nb ? num / (na * nb) : 0.0;
}

function cosVec(a, b) {
  const n = Math.min(a.length, b.length);
  let num = 0,
    na = 0,
    nb = 0;
  for (let i = 0; i < n; i++) num += a[i] * b[i];
  for (let i = 0; i < a.length; i++) na += a[i] * a[i];
  for (let i = 0; i < b.length; i++) nb += b[i] * b[i];
  na = Math.sqrt(na);
  nb = Math.sqrt(nb);
  return na && nb ? num / (na * nb) : 0.0;
}

const isDense = (v) => Array.isArray(v) || ArrayBuffer.isView(v);

// ---------------------------------------------------------------------------
// Core selection over a single block
// ---------------------------------------------------------------------------

export function newStats() {
  return { blocksProcessed: 0, chunksIn: 0, chunksKept: 0, charsIn: 0, charsOut: 0 };
}

const fmtMarker = (template, run) => template.replaceAll("{n}", String(run));

function selectBlock(text, query, embedFn, keepRatio, cfg, _queryEntities, stats) {
  const chunks = splitChunks(text, cfg.minChunks);
  if (chunks.length < cfg.minChunks) return [text, false];

  // rare query entities = query tokens appearing in <=2 chunks of THIS block
  const qEnts = new Set(tokenize(query).filter((w) => w.length >= 4));
  const lowChunks = chunks.map((c) => c.toLowerCase());
  const rareQueryEntities = new Set();
  for (const e of qEnts) {
    let d = 0;
    for (const lc of lowChunks) if (lc.includes(e)) d++;
    if (d >= 1 && d <= 2) rareQueryEntities.add(e);
  }

  const vecs = embedFn([...chunks, query]);
  if (typeof vecs?.then === "function") {
    // ponytail: sync-only pipeline, mirror of the sync Python API; add an
    // async trim() variant if a real async encoder integration asks for it.
    const err = new TypeError(
      "barber: async embedders are not supported yet — embed(texts) must return vectors synchronously"
    );
    err.barberWiring = true;   // wiring bug, not a runtime failure: do not fail open
    throw err;
  }
  const qv = vecs[vecs.length - 1];
  const cvs = vecs.slice(0, -1);
  const sim = isDense(qv) ? cosVec : cosBag;
  const scores = cvs.map((cv) => sim(cv, qv));
  // A *relative* floor needs something to be relative to. When no chunk shares
  // a token with the query -- a non-Latin-script query on the lexical embedder,
  // a block on a wholly different subject -- every score is 0 and the ranking
  // carries no signal at all. The old `top || 1e-9` turned that into a floor of
  // 3.5e-10 that every zero score failed, collapsing the block to lead+tail
  // even at keep=1.0, where the caller asked to drop nothing. With no signal to
  // discriminate on, the budget alone decides.
  let top = -Infinity;
  for (const s of scores) if (s > top) top = s;

  const n = chunks.length;
  const keepN = Math.max(1, Math.min(n, pyRound(n * keepRatio)));

  // rank chunks; always keep pinned + (optionally) lead/tail.
  // Comparator ties resolve to original index order, same as Python's stable
  // sorted(..., reverse=True).
  const pinned = chunks.map((c) => isPinned(c, rareQueryEntities, cfg.pinPatterns ?? PIN_PATTERNS));
  const order = [...scores.keys()].sort((a, b) => scores[b] - scores[a]);
  const keep = new Set(order.slice(0, keepN));
  if (cfg.keepLeadTail) {
    keep.add(0);
    keep.add(n - 1);
  }
  for (let i = 0; i < n; i++) if (pinned[i]) keep.add(i);
  // drop anything far below the top even if budget kept it
  const floor = top > 0 ? top * cfg.relativeFloor : 0;
  const kept = new Set();
  for (const i of keep) {
    if (scores[i] >= floor || pinned[i] || i === 0 || i === n - 1) kept.add(i);
  }

  if (kept.size >= n) return [text, false];

  // rebuild preserving order, collapsing dropped runs into one marker
  const out = [];
  let run = 0;
  for (let i = 0; i < n; i++) {
    if (kept.has(i)) {
      if (run) {
        out.push(fmtMarker(cfg.dropMarker, run));
        run = 0;
      }
      out.push(chunks[i]);
    } else {
      run++;
    }
  }
  if (run) out.push(fmtMarker(cfg.dropMarker, run));

  stats.blocksProcessed += 1;
  stats.chunksIn += n;
  stats.chunksKept += kept.size;
  const newText = out.join("\n\n");
  stats.charsIn += charLen(text);
  stats.charsOut += charLen(newText);
  return [newText, true];
}

// ---------------------------------------------------------------------------
// Build the pipeline transform
// ---------------------------------------------------------------------------

// `toolResults` also walks tool-result bodies — right for token accounting,
// wrong for finding the user's question (a tool result is not something the
// human asked).
export function textOf(content, toolResults = false) {
  if (typeof content === "string") return content;
  if (Array.isArray(content)) {
    const out = [];
    for (const p of content) {
      if (!p || typeof p !== "object") continue;
      if (p.type === "text") out.push(p.text ?? "");
      else if (toolResults && p.type === "tool_result") out.push(textOf(p.content, true));
    }
    return out.join(" ");
  }
  return "";
}

// Apply `selectText` to the trimmable string bodies inside a content list.
// Trimmable means text parts and tool-result bodies; everything else (images,
// tool_use inputs, anything unrecognised) is copied through untouched.
function selectParts(parts, selectText) {
  const out = [];
  let changed = false;
  for (const p of parts) {
    if (!p || typeof p !== "object") {
      out.push(p);
      continue;
    }
    if (p.type === "text" && typeof p.text === "string") {
      const [next, ch] = selectText(p.text);
      out.push(ch ? { ...p, text: next } : p);
      changed = changed || ch;
    } else if (p.type === "tool_result" && typeof p.content === "string") {
      const [next, ch] = selectText(p.content);
      out.push(ch ? { ...p, content: next } : p);
      changed = changed || ch;
    } else if (p.type === "tool_result" && Array.isArray(p.content)) {
      const [next, ch] = selectParts(p.content, selectText);
      out.push(ch ? { ...p, content: next } : p);
      changed = changed || ch;
    } else {
      out.push(p);
    }
  }
  return [out, changed];
}

// Returns ["context_selection", fn] where fn(messages) -> [messages, changed],
// matching the message-transform hook of a Nadir-style optimizer pipeline.
export function makeSelectionTransform({
  embedFn = null,
  tier = "mid",
  complexityScore = null,
  cfg = null,
  decisionCache = null,
} = {}) {
  const config = { ...DEFAULT_CONFIG, ...cfg };
  const embed = embedFn ?? lexical();
  const cache = decisionCache ?? new Map();

  let keepRatio = TIER_KEEP_RATIO[tier] ?? 0.55;
  if (complexityScore != null) {
    // optionally blend the router's continuous score
    keepRatio = Math.max(
      config.minKeepRatio,
      Math.min(config.maxKeepRatio, 0.3 + 0.6 * Number(complexityScore))
    );
  }
  keepRatio = Math.max(config.minKeepRatio, Math.min(config.maxKeepRatio, keepRatio));

  const queryOf = (messages) => {
    for (let i = messages.length - 1; i >= 0; i--) {
      if (messages[i].role === "user") {
        const c = messages[i].content;
        const q = typeof c === "string" ? c : textOf(c);
        // In an agent loop the last user message is usually a tool result
        // carrying no text of its own. That is not the question; keep walking
        // back to the last thing the human actually asked.
        if (q.trim()) return q;
      }
    }
    return "";
  };

  const selectMessages = (messages) => {
    const query = queryOf(messages);
    if (!query) return [messages, false];
    // "rare" query tokens = load-bearing entities to pin
    const queryEntities = new Set(tokenize(query).filter((w) => w.length >= 4));

    const stats = newStats();
    let changedAny = false;
    const outMsgs = [];
    // the latest user message is the QUERY — never prune it
    let lastUserIdx = -1;
    for (let i = 0; i < messages.length; i++) {
      if (messages[i].role === "user") lastUserIdx = i;
    }

    // Run (or replay) the decision for one string block.
    const selectText = (text) => {
      if (charLen(text) < config.minMessageChars) return [text, false];
      // Key on the block ONLY — not the query, and not keepRatio.
      // Freeze-on-first-sight: the first turn to see a block decides it,
      // every later turn replays that decision verbatim, so the stable
      // prefix never mutates and the provider prompt cache stays warm.
      const key = md5key(text);
      if (!cache.has(key)) {
        cache.set(
          key,
          selectBlock(text, query, embed, keepRatio, config, queryEntities, stats)
        );
      }
      return cache.get(key);
    };

    for (let idx = 0; idx < messages.length; idx++) {
      const m = messages[idx];
      const content = m.content;
      if (idx === lastUserIdx || !config.selectableRoles.includes(m.role)) {
        outMsgs.push(m);
        continue;
      }

      // Content parts, not a plain string: this is the shape every agent
      // harness actually sends, because a tool result is a content block. The
      // big blocks in an agent session live in here, so skipping lists meant
      // selection never fired on an agent transcript at all.
      if (Array.isArray(content)) {
        const [newParts, chAny] = selectParts(content, selectText);
        if (chAny) {
          outMsgs.push({ ...m, content: newParts });
          changedAny = true;
        } else {
          outMsgs.push(m);
        }
        continue;
      }

      if (typeof content !== "string" || charLen(content) < config.minMessageChars) {
        outMsgs.push(m);
        continue;
      }

      const [newContent, ch] = selectText(content);
      if (ch) {
        outMsgs.push({ ...m, content: newContent });
        changedAny = true;
      } else {
        outMsgs.push(m);
      }
    }

    transform.lastStats = stats;
    return [outMsgs, changedAny];
  };

  // Fail-open wrapper, matching Python's `transform` in core.py. A context
  // optimizer must never be the reason a request dies: an injected embedder
  // that throws should cost the trim, not the conversation. Logged, not
  // swallowed silently — an encoder that is permanently broken and quietly
  // saves nothing is its own kind of outage.
  //
  // `sweep()` deliberately does NOT fail open, in either language: it is an
  // explicit call at a compaction point, not a pipeline hook, and a caller who
  // asked to rewrite history should hear that it did not happen.
  //
  // A wiring bug (an async embedder) still throws. It is deterministic and
  // fires on the first call, so swallowing it would hide a broken integration
  // behind numbers that say "saved nothing" forever.
  const transform = (messages) => {
    try {
      return selectMessages(messages);
    } catch (e) {
      if (e?.barberWiring) throw e;
      console.warn("barber: selection failed, context passed through untrimmed", e);
      return [messages, false];
    }
  };

  transform.lastStats = newStats();
  return ["context_selection", transform];
}