@hicaru/pi-rlm 0.3.20 → 0.3.22

package/src/ui/tree/tree-model.ts

@@ -5,18 +5,21 @@
  * the result and only rebuilds when the underlying store reports a change.
  *
  * Nothing is ever hidden: every sub-call renders as its own row (parity with
- * pi, which shows each concurrent tool call individually) — except runs of
- * CONSECUTIVE IDENTICAL sibling leaves (same label+model+status), which
- * collapse into one expandable "label ×N" group row so a 20-item llm_batch is
- * one line, not 20 and a wholesale batch failure is one `✗ label ×N` line.
- * Errors group exactly like successes; distinct failures keep their own rows
- * and reasons, and singletons render as plain rows. Interleaved siblings (✗ ✓
- * ✗ with different keys between) stay in encounter order — position is never
- * rewritten. Collapsed subtrees are skipped at the user's explicit request
- * (chevron flips). Token rows are own-spend only — a row never blends models.
+ * pi, which shows each concurrent tool call individually) — except identical
+ * sibling llm leaves (same label+model+status), which consolidate into ONE
+ * expandable "label ×N" group row placed at the first member's position: a
+ * 16-item batch failure is a single `✗ label ×N · reason` line however many
+ * other rows interleave the run. rlm nodes and nodes with children never
+ * group; grouped llm leaves move up to the group head, every other row keeps
+ * its encounter order. Errors group exactly like successes; diverging reasons
+ * collapse to "N failure reasons" (per-item reasons stay in the detail
+ * modal). Singletons render as plain rows. Collapsed subtrees are skipped at
+ * the user's explicit request (chevron flips). Token rows are own-spend only
+ * — a row never blends models.
  */
 
 import type { RlmSubcall, RlmRunStatus, SubcallPhase, SubcallStatus } from "../../tool/rlm-details.ts";
+import { ERROR_PREFIX, isErrorText } from "../../util/errors.ts";
 
 /** Immutable per-run view the model consumes (built by RunRegistry from a live store). */
 export interface RunSnapshot {
@@ -77,6 +80,8 @@ export interface GroupRow {
   readonly icon: SubcallStatus | "queued";
   readonly expandable: boolean;
   readonly expanded: boolean;
+  /** First-line failure reason shared by every member — error groups only, else "N failure reasons". */
+  readonly reason?: string;
 }
 
 /** Internal build-time entry: a real node or an accumulating group. */
@@ -94,21 +99,65 @@ const groupable = (sc: RlmSubcall, byParent: ReadonlyMap<string | undefined, Rlm
 
 const groupKey = (sc: RlmSubcall): string => `${sc.label}|${sc.model ?? ""}|${sc.status}`;
 
-/** Merge consecutive identical sibling leaves into group entries; keep order. */
+/** Longest reason shown inline in a group header — full text lives in the modal. */
+const REASON_MAX_CHARS = 48;
+
+/** First line of an error text, "Error: " prefix stripped — undefined when there is nothing usable. */
+function errorLineOf(sc: RlmSubcall): string | undefined {
+  const raw = sc.detail ?? sc.resultPreview;
+  if (raw === undefined || raw === "") return undefined;
+  const body = isErrorText(raw) ? raw.slice(ERROR_PREFIX.length + 1) : raw;
+  const nl = body.indexOf("\n");
+  const line = (nl === -1 ? body : body.slice(0, nl)).trim();
+  return line === "" ? undefined : line;
+}
+
+/** The reason all members share, or "N failure reasons" when their errors diverge. */
+function groupReason(members: readonly RlmSubcall[]): string | undefined {
+  let first: string | undefined;
+  const distinct = new Set<string>();
+  for (const m of members) {
+    const line = errorLineOf(m);
+    if (line === undefined) continue;
+    if (first === undefined) first = line;
+    distinct.add(line);
+  }
+  if (first === undefined || distinct.size === 0) return undefined;
+  const reason = distinct.size === 1 ? first : `${String(distinct.size)} failure reasons`;
+  return reason.length > REASON_MAX_CHARS ? `${reason.slice(0, REASON_MAX_CHARS - 1)}…` : reason;
+}
+
+/**
+ * Consolidate every identical sibling leaf (same label+model+status) into ONE
+ * group entry at its first member's position — a 16-item batch failure stays
+ * a single `✗ label ×16` line however many other rows interleave the run.
+ * Non-grouping siblings keep their encounter order; members keep start order
+ * for the expanded view.
+ */
 function partition(children: readonly RlmSubcall[], byParent: ReadonlyMap<string | undefined, RlmSubcall[]>): readonly Entry[] {
+  const groups = new Map<string, Extract<Entry, { type: "group" }>>();
   const out: Entry[] = [];
   for (const sc of children) {
-    if (groupable(sc, byParent)) {
-      const key = groupKey(sc);
-      const last = out[out.length - 1];
-      if (last !== undefined && last.type === "group" && last.key === key) {
-        last.members.push(sc);
-        continue;
-      }
-      out.push({ type: "group", key, label: sc.label, model: sc.model, status: sc.status, members: [sc] });
-    } else {
+    if (!groupable(sc, byParent)) {
       out.push({ type: "node", sc });
+      continue;
+    }
+    const key = groupKey(sc);
+    const existing = groups.get(key);
+    if (existing !== undefined) {
+      existing.members.push(sc);
+      continue;
     }
+    const group: Extract<Entry, { type: "group" }> = {
+      type: "group",
+      key,
+      label: sc.label,
+      model: sc.model,
+      status: sc.status,
+      members: [sc],
+    };
+    groups.set(key, group);
+    out.push(group);
   }
   return out;
 }
@@ -218,6 +267,7 @@ export function buildRows(
       icon: iconOf(entry.status, entry.members.some((m) => m.phase === "queued") ? "queued" : undefined),
       expandable: true,
       expanded,
+      reason: entry.status === "error" ? groupReason(entry.members) : undefined,
     });
     if (!expanded) return;
     for (let i = 0; i < entry.members.length; i++) {

package/src/ui/tree/tree-rows.ts

@@ -55,7 +55,8 @@ function formatGroup(row: GroupRow, selected: boolean, width: number, theme: The
   const chevron = row.expanded ? GLYPHS.expanded : GLYPHS.collapsed;
   const cursor = selected ? theme.fg("accent", "❯") : " ";
   const icon = row.icon === "done" ? theme.fg("success", GLYPHS.done) : row.icon === "error" ? theme.fg("error", GLYPHS.error) : theme.fg("warning", spinnerFrame());
-  const left = `${cursor} ${row.prefix}${chevron} ${icon} ${row.label} ×${row.count}`;
+  const reason = row.reason === undefined ? "" : theme.fg("warning", ` · ${row.reason}`);
+  const left = `${cursor} ${row.prefix}${chevron} ${icon} ${row.label} ×${row.count}${reason}`;
   return assembleLine(left, row.tokens, row.tokensIn, row.tokensOut, row.model, selected, width, theme);
 }
 

package/src/util/abort.ts

@@ -0,0 +1,34 @@
+/**
+ * abort — combine abort signals without leaking listeners.
+ *
+ * pi hands every tool execution a per-call signal (esc) while the session adds
+ * longer-lived controllers (/rlm-stop rotation) — in-flight work must react to both.
+ * The longer-lived signal outlives any single call, so its relay listener MUST detach
+ * when the awaited work settles, or every call accumulates one listener forever.
+ */
+
+export interface RacedSignal {
+  /** Fires when either input aborts; undefined when no input signal was given. */
+  readonly signal: AbortSignal | undefined;
+  /** Detach relay listeners — call when the awaited work settles. */
+  dispose(): void;
+}
+
+export function raceAbort(a: AbortSignal | undefined, b: AbortSignal | undefined): RacedSignal {
+  if (a === undefined || b === undefined) return { signal: a ?? b, dispose: () => {} };
+  const combined = new AbortController();
+  const relay = (src: AbortSignal): (() => void) => () => combined.abort(src.reason);
+  const onA = relay(a);
+  const onB = relay(b);
+  if (a.aborted) combined.abort(a.reason);
+  if (b.aborted) combined.abort(b.reason);
+  a.addEventListener("abort", onA);
+  b.addEventListener("abort", onB);
+  return {
+    signal: combined.signal,
+    dispose: () => {
+      a.removeEventListener("abort", onA);
+      b.removeEventListener("abort", onB);
+    },
+  };
+}

package/src/util/bm25.ts

@@ -3,8 +3,10 @@
  *
  * BM25 DUALITY, documented not accidental: this file and `sandbox/py/retrieval.py`
  * (`_Bm25Index`) implement the SAME scoring for two runtimes — identical constants, identical
- * tokenizer, identical idf/norm formulas — so a note ranked here lands in the same order the
- * sandbox-side `search` would put it. Keep the two files in lockstep (AGENTS.md convention).
+ * tokenizer + stemmer, identical idf/norm formulas, identical PRF/phrase-bonus pipeline — so a
+ * note ranked here lands in the same order the sandbox-side `search` would put it. Keep the two
+ * files in lockstep (AGENTS.md convention). Python-only extras (window overlap, glob
+ * pre-filter, adjacent-window merge) are window mechanics, not scoring, and stop at that file.
  */
 
 const BM25_K1 = 1.2; // mirrors retrieval.py:_BM25_K1
@@ -14,18 +16,69 @@ const BM25_B = 0.75; // mirrors retrieval.py:_BM25_B
 const TOKEN_SPLIT = /[^0-9A-Za-z]+/;
 const CAMEL_SPLIT = /(?<=[a-z0-9])(?=[A-Z])/;
 
+// BM25 V2 constants (twin: retrieval.py — identical values there).
+const PRF_FEEDBACK_DOCS = 3; // top first-pass docs harvested for expansion terms
+const PRF_EXPANSION_TERMS = 8; // max terms added by pseudo-relevance feedback
+const PRF_EXPANSION_WEIGHT = 0.4; // Rocchio beta: expansion terms contribute at this weight
+const PRF_MIN_DOCS = 12; // below this the corpus is too small to harvest from
+const PHRASE_BONUS_WEIGHT = 0.25; // per adjacent query-bigram occurrence in a doc
+const RERANK_POOL_MULT = 3; // phrase-bonus pool = top (k * mult) docs, capped
+const RERANK_POOL_CAP = 60;
+
+/**
+ * Light deterministic suffix stripper (TWIN: retrieval.py `_stem` — identical rules).
+ * A matching aid, not linguistics: different surface forms converge (files/file → fil,
+ * running/run → run, studies/study → studi); identical forms always map to themselves.
+ */
+function stem(t: string): string {
+  if (t.length <= 3) return t;
+  let r: string;
+  if (t.endsWith("ies")) {
+    r = t.slice(0, -3) + "i"; // studies → studi
+  } else if (t.endsWith("sses")) {
+    r = t.slice(0, -2); // classes → class
+  } else if (t.endsWith("es")) {
+    const stem2 = t.slice(0, -2);
+    r = /(x|ch|sh)$/.test(stem2) ? stem2 : t.slice(0, -1); // boxes → box, files → file
+  } else if (t.endsWith("s") && !/(ss|us|is)$/.test(t)) {
+    r = t.length > 4 ? t.slice(0, -1) : t; // cats → cat; keeps "was"/"its"
+  } else {
+    r = t;
+  }
+  if (r.endsWith("ing") && r.length >= 6) {
+    let base = r.slice(0, -3); // running → runn
+    if (base.length >= 4) {
+      if (base.length >= 2 && base[base.length - 1] === base[base.length - 2]) {
+        base = base.slice(0, -1); // runn → run
+      }
+      r = base; // "string" stays whole
+    }
+  } else if (r.endsWith("ed") && r.length >= 5) {
+    let base = r.slice(0, -2); // mapped → mapp
+    if (base.length >= 4) {
+      if (base.length >= 2 && base[base.length - 1] === base[base.length - 2]) {
+        base = base.slice(0, -1); // mapp → map
+      }
+      r = base;
+    }
+  }
+  if (r.endsWith("y") && r.length > 3) r = r.slice(0, -1) + "i"; // study → studi (meets studies)
+  if (r.endsWith("e") && r.length > 3) r = r.slice(0, -1); // file → fil (meets files)
+  return r.length >= 3 ? r : t;
+}
+
 export function bm25Tokenize(text: string): readonly string[] {
   const out: string[] = [];
   for (const raw of text.split(TOKEN_SPLIT)) {
     if (raw === "") continue;
     const lowered = raw.toLowerCase();
-    out.push(lowered);
+    out.push(stem(lowered));
     if (raw.length > 3) {
       const parts = raw.split(CAMEL_SPLIT);
       if (parts.length > 1) {
         for (const part of parts) {
           const piece = part.toLowerCase();
-          if (piece !== "" && piece !== lowered) out.push(piece);
+          if (piece !== "" && piece !== lowered) out.push(stem(piece));
         }
       }
     }
@@ -43,14 +96,42 @@ export interface Bm25Hit<T> {
   readonly score: number;
 }
 
+/** Ranking knobs (frozen option bags; defaults enable the full V2 pipeline). */
+export interface Bm25RankOptions {
+  /** Pseudo-relevance-feedback query expansion (auto-off below PRF_MIN_DOCS docs). */
+  readonly prf?: boolean;
+  /** Adjacent-bigram phrase bonus over a re-rank pool. */
+  readonly phraseBonus?: boolean;
+}
+export const BM25_RANK_DEFAULTS: Readonly<Bm25RankOptions> = Object.freeze({
+  prf: true,
+  phraseBonus: true,
+});
+
+/**
+ * Positive-scoring docs as (idx, score) pairs, best first. The Python twin's score dict only
+ * ever holds docs with ≥1 matching term; a dense TS array would otherwise let zero-score docs
+ * into PRF feedback and the phrase pool — a divergence the parity suite catches.
+ */
+function positivePairs(scores: readonly number[]): readonly (readonly [number, number])[] {
+  const pairs: (readonly [number, number])[] = [];
+  for (let i = 0; i < scores.length; i++) {
+    if (scores[i] > 0) pairs.push([i, scores[i]]);
+  }
+  pairs.sort((a, b) => b[1] - a[1] || a[0] - b[0]);
+  return pairs;
+}
+
 /**
  * Rank entries against a query, best first, top-k, score > 0 only.
- * Pre-allocated score/index arrays; no growth in the scoring loops.
+ * Pipeline mirrors retrieval.py:search — weighted scoring → PRF expansion → phrase bonus
+ * over a re-rank pool → top-k. Pre-allocated arrays; no growth in the scoring loops.
  */
 export function bm25Rank<T>(
   query: string,
   entries: readonly Bm25Entry<T>[],
   k: number,
+  options: Bm25RankOptions = BM25_RANK_DEFAULTS,
 ): readonly Bm25Hit<T>[] {
   const n = entries.length;
   if (n === 0 || k <= 0) return [];
@@ -70,28 +151,96 @@ export function bm25Rank<T>(
   let totalLen = 0;
   for (let i = 0; i < n; i++) totalLen += docLens[i];
   const avgLen = totalLen > 0 ? totalLen / n : 1.0;
+  const idfOf = (term: string): number => {
+    const df = postings.get(term)?.length ?? 0;
+    return Math.log(1.0 + (n - df + 0.5) / (df + 0.5));
+  };
+  const score = (weights: ReadonlyMap<string, number>): number[] => {
+    const scores = new Array<number>(n).fill(0);
+    for (const [term, weight] of weights) {
+      const posting = postings.get(term);
+      if (posting === undefined) continue;
+      const idf = idfOf(term);
+      for (const [idx, tf] of posting) {
+        const norm = BM25_K1 * (1.0 - BM25_B + BM25_B * (docLens[idx] / avgLen));
+        scores[idx] += weight * idf * (tf * (BM25_K1 + 1.0)) / (tf + norm);
+      }
+    }
+    return scores;
+  };
 
-  const scores = new Array<number>(n).fill(0);
-  const seen = new Set<string>();
+  const weights = new Map<string, number>();
   for (const term of bm25Tokenize(query)) {
-    if (seen.has(term)) continue; // Python twin scores set(terms)
-    seen.add(term);
-    const posting = postings.get(term);
-    if (posting === undefined) continue;
-    const df = posting.length;
-    const idf = Math.log(1.0 + (n - df + 0.5) / (df + 0.5));
-    for (const [idx, tf] of posting) {
-      const norm = BM25_K1 * (1.0 - BM25_B + BM25_B * (docLens[idx] / avgLen));
-      scores[idx] += (idf * (tf * (BM25_K1 + 1.0))) / (tf + norm);
+    if (!weights.has(term)) weights.set(term, 1.0); // Python twin: first occurrence wins
+  }
+  let scores = score(weights);
+  const basePositive = positivePairs(scores);
+  if (options.prf !== false && n >= PRF_MIN_DOCS && basePositive.length > 0) {
+    const feedback = basePositive.slice(0, PRF_FEEDBACK_DOCS).map(([idx]) => idx);
+    const tf = new Map<string, number>();
+    for (const idx of feedback) {
+      for (const term of bm25Tokenize(entries[idx].text)) {
+        tf.set(term, (tf.get(term) ?? 0) + 1);
+      }
+    }
+    // Exclude query terms BEFORE slicing — the twin (retrieval.py:_expansion_terms) excludes
+    // first, then takes the top-8; slicing first would let query terms eat expansion slots.
+    const ranked: readonly (readonly [number, string])[] = Array.from(tf, ([term, f]) => [f * idfOf(term), term] as const)
+      .filter(([, term]) => !weights.has(term))
+      .sort((a, b) => b[0] - a[0] || (a[1] < b[1] ? -1 : a[1] > b[1] ? 1 : 0))
+      .slice(0, PRF_EXPANSION_TERMS);
+    let added = false;
+    for (const [, term] of ranked) {
+      if (!weights.has(term)) {
+        weights.set(term, PRF_EXPANSION_WEIGHT);
+        added = true;
+      }
     }
+    if (added) scores = score(weights);
+  }
+  if (scores.every((s) => s <= 0)) return [];
+
+  // Pool of (idx, score) pairs — the phrase bonus MUTATES the score (twin parity), so the
+  // reported value must come from the pool pairs, not the pre-bonus score array.
+  const positive = positivePairs(scores);
+  const poolN = Math.min(positive.length, RERANK_POOL_CAP, Math.max(k * RERANK_POOL_MULT, PRF_FEEDBACK_DOCS));
+  let pool = positive.slice(0, poolN);
+  const seq = bm25Tokenize(query); // raw sequence — repeated terms make self-bigrams, like Python
+  const bigrams: (readonly [string, string])[] = [];
+  if (options.phraseBonus !== false) {
+    for (let i = 0; i < seq.length - 1; i++) {
+      if (seq[i] !== seq[i + 1] && !bigrams.some(([a, b]) => a === seq[i] && b === seq[i + 1])) {
+        bigrams.push([seq[i], seq[i + 1]]);
+      }
+    }
+  }
+  if (bigrams.length > 0) {
+    const boosted = pool.map(([idx, s]) => {
+      const toks = bm25Tokenize(entries[idx].text);
+      const pos = new Map<string, number[]>();
+      for (let i = 0; i < toks.length; i++) {
+        const list = pos.get(toks[i]);
+        if (list !== undefined) list.push(i);
+        else pos.set(toks[i], [i]);
+      }
+      let bonus = 0;
+      for (const [a, b] of bigrams) {
+        const pa = pos.get(a);
+        const pb = pos.get(b);
+        if (pa !== undefined && pb !== undefined && pa.some((x) => pb.some((y) => y === x + 1))) {
+          bonus += PHRASE_BONUS_WEIGHT * Math.max(idfOf(a), idfOf(b));
+        }
+      }
+      return [idx, s + bonus] as const;
+    });
+    boosted.sort((x, y) => y[1] - x[1] || x[0] - y[0]);
+    pool = boosted;
   }
 
-  const order = Array.from({ length: n }, (_, i) => i);
-  order.sort((a, b) => scores[b] - scores[a] || a - b); // deterministic tie-break
   const out: Bm25Hit<T>[] = [];
-  for (let i = 0; i < n && out.length < k; i++) {
-    const idx = order[i];
-    if (scores[idx] > 0) out.push({ item: entries[idx].item, score: scores[idx] });
+  for (const [idx, s] of pool) {
+    if (out.length >= k) break;
+    if (s > 0) out.push({ item: entries[idx].item, score: s });
   }
   return out;
 }

package/src/util/errors.ts

@@ -12,7 +12,7 @@ export function err<T = never, E = string>(error: E): Result<T, E> {
   return { ok: false, error };
 }
 
-const ERROR_PREFIX = "Error:";
+export const ERROR_PREFIX = "Error:";
 
 export function formatError(message: string): string {
   return `${ERROR_PREFIX} ${message}`;