@hicaru/pi-rlm 0.3.20 → 0.3.21

package/src/sandbox/py/retrieval.py

@@ -3,6 +3,11 @@
 `search` / `grep_context` / `outline` are what the prompt tells the model to reach for BEFORE
 delegating anything, so they carry no worker state: each takes the already-materialised
 `entries` list and returns pointers, never file bodies.
+
+BM25 V2 (twin: util/bm25.ts — keep the two files in lockstep, AGENTS.md duality convention):
+light suffix stemming, bigram phrase bonus, and pseudo-relevance-feedback query expansion
+(Rocchio without embeddings — the LLM Agent Memory Survey's query-reformulation line, fully
+deterministic). Windows overlap by half so evidence straddling a boundary stops diluting.
 """
 
 from __future__ import annotations
@@ -70,6 +75,7 @@ def _snippet_window(text: str, terms: set[str]) -> str:
 
 
 _INDEX_WINDOW_LINES = 40       # a window is the retrieval unit: big enough to carry meaning
+_INDEX_WINDOW_STRIDE = 20      # windows overlap by half so boundary-straddled evidence survives
 _INDEX_MAX_WINDOWS = 20_000    # ceiling so a huge add_context() cannot exhaust worker memory
 _SNIPPET_CHARS = 400
 _SNIPPET_LEAD = 100            # chars of lead-in kept before the earliest matched term
@@ -77,6 +83,16 @@ _GREP_HARD_CAP = 200           # absolute ceiling on returned grep hits, whateve
 _BM25_K1 = 1.2
 _BM25_B = 0.75
 
+# BM25 V2 constants (twin: util/bm25.ts — identical values there).
+_PRF_FEEDBACK_DOCS = 3         # top first-pass windows harvested for expansion terms
+_PRF_EXPANSION_TERMS = 8       # max terms added by pseudo-relevance feedback
+_PRF_EXPANSION_WEIGHT = 0.4    # Rocchio beta: expansion terms contribute at this weight
+_PRF_MIN_DOCS = 12             # below this the corpus is too small to harvest from
+_PHRASE_BONUS_WEIGHT = 0.25    # per adjacent query-bigram occurrence in a window
+_RERANK_POOL_MULT = 3          # phrase-bonus pool = top (k * mult) windows, capped
+_RERANK_POOL_CAP = 60
+_STEM_MIN_LEN = 4              # tokens shorter than this are never stemmed
+
 _TOKEN_SPLIT = re.compile(r"[^0-9A-Za-z]+")     # also splits snake_case and paths
 _CAMEL_SPLIT = re.compile(r"(?<=[a-z0-9])(?=[A-Z])")
 
@@ -94,21 +110,63 @@ _OUTLINE_LINE = re.compile(
 )
 
 
+def _stem(t: str) -> str:
+    """Light deterministic suffix stripper (TWIN: util/bm25.ts `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.
+    Final `y`→`i` and trailing-`e` deletion are what make the folds meet.
+    """
+    if len(t) <= 3:
+        return t
+    if t.endswith("ies"):
+        r = t[:-3] + "i"                                   # studies → studi
+    elif t.endswith("sses"):
+        r = t[:-2]                                         # classes → class
+    elif t.endswith("es"):
+        stem2 = t[:-2]
+        if stem2.endswith(("x", "ch", "sh")):
+            r = stem2                                      # boxes → box, matches → match
+        else:
+            r = t[:-1]                                     # files → file
+    elif t.endswith("s") and not t.endswith(("ss", "us", "is")):
+        r = t[:-1] if len(t) > 4 else t                    # cats → cat; keeps "was"/"its"
+    else:
+        r = t
+    if r.endswith("ing") and len(r) >= 6:
+        base = r[:-3]                                      # running → runn
+        if len(base) >= 4:
+            if len(base) >= 2 and base[-1] == base[-2]:
+                base = base[:-1]                           # runn → run
+            r = base                                       # "string" stays whole
+    elif r.endswith("ed") and len(r) >= 5:
+        base = r[:-2]                                      # mapped → mapp
+        if len(base) >= 4:
+            if len(base) >= 2 and base[-1] == base[-2]:
+                base = base[:-1]                           # mapp → map
+            r = base
+    if r.endswith("y") and len(r) > 3:
+        r = r[:-1] + "i"                                   # study → studi (meets studies)
+    if r.endswith("e") and len(r) > 3:
+        r = r[:-1]                                         # file → fil (meets files)
+    return r if len(r) >= 3 else t
+
+
 def _tokenize(text: str) -> list[str]:
-    """Lowercased alphanumeric runs, plus camelCase parts so `resolveModelId` matches `model id`."""
+    """Lowered alphanumeric runs + camelCase parts, all stemmed (`resolveModelId` ↔ model ids)."""
     out: list[str] = []
     for raw in _TOKEN_SPLIT.split(text):
         if not raw:
             continue
         lowered = raw.lower()
-        out.append(lowered)
+        out.append(_stem(lowered))
         if len(raw) > 3:
             parts = _CAMEL_SPLIT.split(raw)
             if len(parts) > 1:
                 for part in parts:
                     piece = part.lower()
                     if piece and piece != lowered:
-                        out.append(piece)
+                        out.append(_stem(piece))
     return out
 
 
@@ -132,13 +190,14 @@ def _context_entries(context: Any) -> list[tuple[str, str]]:
 
 
 class _Bm25Index:
-    """Okapi BM25 over fixed-line windows of `context`. Built lazily, discarded on change."""
+    """Okapi BM25 over overlapping fixed-line windows of `context`. Built lazily, rebuilt on change."""
 
-    __slots__ = ("paths", "starts", "texts", "postings", "doc_len", "avg_len", "truncated")
+    __slots__ = ("paths", "starts", "ends", "texts", "postings", "doc_len", "avg_len", "truncated")
 
     def __init__(self, entries: list[tuple[str, str]]) -> None:
         self.paths: list[str] = []
         self.starts: list[int] = []
+        self.ends: list[int] = []
         self.texts: list[str] = []
         self.postings: dict[str, list[tuple[int, int]]] = {}
         self.doc_len: list[int] = []
@@ -148,14 +207,19 @@ class _Bm25Index:
             if not content:
                 continue
             lines = content.split("\n")
-            for start in range(0, len(lines), _INDEX_WINDOW_LINES):
+            for start in range(0, len(lines), _INDEX_WINDOW_STRIDE):
+                window_lines = lines[start:start + _INDEX_WINDOW_LINES]
+                # The tail window fully covered by the previous one adds nothing — skip it.
+                if start > 0 and len(window_lines) <= _INDEX_WINDOW_STRIDE:
+                    break
                 if len(self.texts) >= _INDEX_MAX_WINDOWS:
                     self.truncated = True
                     break
-                window = "\n".join(lines[start:start + _INDEX_WINDOW_LINES])
+                window = "\n".join(window_lines)
                 idx = len(self.texts)
                 self.paths.append(path)
                 self.starts.append(start + 1)
+                self.ends.append(start + len(window_lines))
                 self.texts.append(window)
                 terms = _tokenize(window)
                 self.doc_len.append(len(terms))
@@ -170,46 +234,98 @@ class _Bm25Index:
         total = len(self.doc_len)
         self.avg_len = (sum(self.doc_len) / total) if total else 1.0
 
-    def query(self, terms: list[str], k: int, path_glob: str | None) -> list[dict[str, Any]]:
+    def idf(self, term: str) -> float:
+        total = len(self.doc_len)
+        if total == 0:
+            return 0.0
+        df = len(self.postings.get(term, ()))
+        return math.log(1.0 + (total - df + 0.5) / (df + 0.5))
+
+    def glob_allow(self, path_glob: str | None) -> set[int] | None:
+        """Pre-filter: window indices whose path matches the glob (None = everything).
+
+        Pre-filtering before scoring (not after) so glob-matched windows outside the global
+        top-k still surface — a post-scoring filter could return fewer than k.
+        """
+        if not path_glob:
+            return None
+        return {i for i, p in enumerate(self.paths) if fnmatch.fnmatch(p, path_glob)}
+
+    def score(self, weights: dict[str, float], allowed: set[int] | None) -> dict[int, float]:
+        """Weighted Okapi accumulation: scores[idx] += weight * idf * tf-sat (set-ordered input)."""
         total = len(self.texts)
         if total == 0:
-            return []
+            return {}
         scores: dict[int, float] = {}
-        for term in set(terms):
+        for term, weight in weights.items():
             posting = self.postings.get(term)
             if not posting:
                 continue
-            df = len(posting)
-            idf = math.log(1.0 + (total - df + 0.5) / (df + 0.5))
+            idf = self.idf(term)
             for idx, tf in posting:
+                if allowed is not None and idx not in allowed:
+                    continue
                 norm = _BM25_K1 * (1.0 - _BM25_B + _BM25_B * self.doc_len[idx] / self.avg_len)
-                scores[idx] = scores.get(idx, 0.0) + idf * (tf * (_BM25_K1 + 1.0)) / (tf + norm)
-        if path_glob:
-            scores = {i: s for i, s in scores.items() if fnmatch.fnmatch(self.paths[i], path_glob)}
-        if not scores:
-            return []
-        top = heapq.nlargest(k, scores.items(), key=lambda kv: kv[1])
-        out: list[dict[str, Any]] = [None] * len(top)  # type: ignore[list-item]
-        term_set = set(terms)
-        for i, (idx, score) in enumerate(top):
-            text = self.texts[idx]
-            snip = _snippet_window(text, term_set)
-            # Both `snippet` and `text` so agents never KeyError mixing search vs grep shapes.
-            out[i] = {
-                "path": self.paths[idx],
-                "line": self.starts[idx],
-                "score": round(score, 3),
-                "snippet": snip,
-                "text": snip,
-            }
-        return out
+                scores[idx] = scores.get(idx, 0.0) + weight * idf * (tf * (_BM25_K1 + 1.0)) / (tf + norm)
+        return scores
+
+
+def _top(scores: dict[int, float], n: int) -> list[tuple[int, float]]:
+    """Top-n by (score desc, idx asc) — the deterministic tie-break the TS twin mirrors."""
+    return sorted(scores.items(), key=lambda kv: (-kv[1], kv[0]))[:n]
+
+
+def _expansion_terms(index: _Bm25Index, feedback: list[int], exclude: set[str]) -> list[str]:
+    """Pseudo-relevance-feedback terms: tf-in-feedback × global idf, best first (twin-mirrored).
+
+    Harvested from the top first-pass windows; deterministic ties break by term codepoints.
+    """
+    tf: dict[str, int] = {}
+    for idx in feedback:
+        for term in _tokenize(index.texts[idx]):
+            tf[term] = tf.get(term, 0) + 1
+    scored: list[tuple[float, str]] = []
+    for term, f in tf.items():
+        if term in exclude:
+            continue
+        scored.append((f * index.idf(term), term))
+    scored.sort(key=lambda st: (-st[0], st[1]))
+    return [term for _, term in scored[:_PRF_EXPANSION_TERMS]]
+
+
+def _phrase_bigrams(terms: list[str]) -> list[tuple[str, str]]:
+    seen: set[tuple[str, str]] = set()
+    out: list[tuple[str, str]] = []
+    for a, b in zip(terms, terms[1:]):
+        if a != b and (a, b) not in seen:
+            seen.add((a, b))
+            out.append((a, b))
+    return out
+
+
+def _adjacent(index: _Bm25Index, idx: int, bigrams: list[tuple[str, str]]) -> float:
+    """Phrase bonus for one window: weight × max-idf per query bigram occurring adjacently."""
+    if not bigrams:
+        return 0.0
+    toks = _tokenize(index.texts[idx])
+    pos: dict[str, list[int]] = {}
+    for i, t in enumerate(toks):
+        pos.setdefault(t, []).append(i)
+    bonus = 0.0
+    for a, b in bigrams:
+        pa, pb = pos.get(a), pos.get(b)
+        if pa and pb and any(x + 1 == y for x in pa for y in pb):
+            bonus += _PHRASE_BONUS_WEIGHT * max(index.idf(a), index.idf(b))
+    return bonus
 
 
 def search(entries: list[tuple[str, str]], index: _Bm25Index, query: str, k: int = 10, path_glob: str | None = None) -> list[dict[str, Any]]:
-    """Rank `context` windows against a natural-language query (BM25).
+    """Rank `context` windows against a natural-language query (BM25 V2 pipeline).
 
-    Returns [{path, line, score, snippet, text}] — pointers, not bodies.
-    `text` is an alias of `snippet` (same as grep_context hits).
+    Pipeline: weighted scoring → PRF expansion (corpora ≥ _PRF_MIN_DOCS windows) → phrase
+    bonus over a re-rank pool → top-k → adjacent-window merge. Returns [{path, line, score,
+    snippet, text}] — pointers, not bodies; merged spans add `end`, an over-cap index adds
+    `index_truncated`. `text` aliases `snippet` (same as grep_context hits).
     """
     terms = _tokenize(str(query))
     if not terms:
@@ -218,7 +334,57 @@ def search(entries: list[tuple[str, str]], index: _Bm25Index, query: str, k: int
         limit = max(1, min(int(k), 100))
     except (TypeError, ValueError):
         limit = 10
-    return index.query(terms, limit, path_glob)
+    allowed = index.glob_allow(path_glob)
+    weights: dict[str, float] = {}
+    for t in terms:
+        weights.setdefault(t, 1.0)
+    scores = index.score(weights, allowed)
+    if len(index.doc_len) >= _PRF_MIN_DOCS and scores:
+        feedback = [idx for idx, _ in _top(scores, _PRF_FEEDBACK_DOCS)]
+        expansion = _expansion_terms(index, feedback, set(weights))
+        if expansion:
+            for t in expansion:
+                weights[t] = _PRF_EXPANSION_WEIGHT
+            scores = index.score(weights, allowed)
+    if not scores:
+        return []
+    pool_n = min(len(scores), _RERANK_POOL_CAP, max(limit * _RERANK_POOL_MULT, _PRF_FEEDBACK_DOCS))
+    pool = _top(scores, pool_n)
+    bigrams = _phrase_bigrams(terms)
+    if bigrams:
+        boosted = [(idx, score + _adjacent(index, idx, bigrams)) for idx, score in pool]
+        boosted.sort(key=lambda kv: (-kv[1], kv[0]))
+        pool = boosted
+    top = pool[:limit]
+
+    # Adjacent same-path windows merge into one wider hit (diversity: one file cannot flood k).
+    merged: list[dict[str, Any]] = []
+    for idx, score in top:
+        path = index.paths[idx]
+        start = index.starts[idx]
+        end = index.ends[idx]
+        if merged and merged[-1]["path"] == path and start <= int(merged[-1]["end"]) + 1:
+            merged[-1]["end"] = max(int(merged[-1]["end"]), end)
+            if score > float(merged[-1]["score"]):
+                snip = _snippet_window(index.texts[idx], set(terms))
+                merged[-1]["snippet"] = snip
+                merged[-1]["text"] = snip
+            merged[-1]["score"] = round(max(float(merged[-1]["score"]), score), 3)
+        else:
+            snip = _snippet_window(index.texts[idx], set(terms))
+            merged.append({
+                "path": path,
+                "line": start,
+                "end": end,
+                "score": round(score, 3),
+                "snippet": snip,
+                "text": snip,
+            })
+    if index.truncated:
+        for hit in merged:
+            hit["index_truncated"] = True
+    return merged
+
 
 def grep_context(
     entries: list[tuple[str, str]],

package/src/sandbox/py/scaffold.py

@@ -261,15 +261,30 @@ class WorkerScaffold:
     def _entries(self) -> list[tuple[str, str]]:
         return _context_entries(self.ns.get("context"))
 
+    @staticmethod
+    def _fingerprint(entries: list[tuple[str, str]]) -> int:
+        """FNV-1a over per-entry (path, length, head-256, tail-256) — catches in-place edits.
+
+        Sampling keeps this O(entries) even for 64MB payloads; a mid-file same-length edit that
+        dodges both sampled ends is accepted residual risk (the old id+length stamp missed ALL
+        same-length edits, not just unsampled ones).
+        """
+        h = 0x811C9DC5
+        for path, content in entries:
+            sample = path + "\x00" + content[:256] + "\x00" + (content[-256:] if len(content) > 256 else "")
+            for byte in sample.encode("utf-8", "replace"):
+                h = ((h ^ byte) * 0x01000193) & 0xFFFFFFFF
+            h = ((h ^ (len(content) & 0xFFFFFFFF)) * 0x01000193) & 0xFFFFFFFF
+        return h
+
     def _get_index(self) -> _Bm25Index:
-        """Build the BM25 index on first use; rebuild when `context` was replaced or resized.
+        """Build the BM25 index on first use; rebuild when `context` changed in any detectable way.
 
-        Identity+length is a cheap stamp that catches the two ways context actually changes:
-        add_context() extending the list, and the model re-binding the name. In-place edits
-        that preserve length are not detected — documented, and rare in practice.
+        Identity + content fingerprint: catches add_context() growth, model re-binds, AND
+        same-length in-place edits at the sampled ends (the old stamp caught only the first two).
         """
         ctx = self.ns.get("context")
-        stamp = (id(ctx), len(ctx) if isinstance(ctx, (list, str)) else 0)
+        stamp = (id(ctx), self._fingerprint(self._entries()))
         if self._index is None or self._index_stamp != stamp:
             self._index = _Bm25Index(self._entries())
             self._index_stamp = stamp

package/src/sandbox/py/worker.py

@@ -119,7 +119,7 @@ class Worker(WorkerScaffold):
         self._context_payload: Any = []  # empty list — the only starting value that needs no bootstrap branch
         self._nudged: set[str] = set()
         self._index: _Bm25Index | None = None
-        self._index_stamp: tuple[int, int] | None = None  # (id(context), len(context))
+        self._index_stamp: tuple[int, int] | None = None  # (id(context), content fingerprint)
         self._restore_scaffold()
 
     def _capture_answer(self, content: Any) -> None:

package/src/sandbox/sandbox-manager.ts

@@ -75,6 +75,25 @@ export class SandboxManager {
     }
   }
 
+  /**
+   * Recall W1: materialize one SessionArchive segment into the sandbox under
+   * `ctx/session-log/`. Same seam as refreshFileFromDisk (unique segment path ⇒ the upsert
+   * appends), so the worker's free search/grep_context BM25 covers elided history. Fail-soft:
+   * returns false when the worker is absent/dead (the host snapshot still replays on recreate).
+   */
+  async upsertArchiveSegment(segmentPath: string, text: string, cwd: string): Promise<boolean> {
+    this.contextPayload = upsertContextFile(this.contextPayload, segmentPath, text, cwd);
+    if (!this.sandbox || this.disposed) return false;
+    const code = patchContextExecCode(segmentPath, text, cwd);
+    try {
+      await this.execQueued(code);
+      return true;
+    } catch {
+      this.contextLoaded = false;
+      return false;
+    }
+  }
+
   /**
    * Lazy get-or-create the sandbox. On first call, spawns PythonSandbox with the
    * given handlers. Subsequent calls return the existing sandbox immediately.

package/src/text/parsing.ts

@@ -55,6 +55,31 @@ const STATE_FENCE = /(`{3,})[ \t]*state[ \t]*\r?\n([\s\S]*?)\1/g;
  *  example JSON never trips the scanner. */
 const BARE_PATCH = /\{"state_patch"\s*:/g;
 
+/** Recall W2 fence-echo fix: ```json fences are the EXAMPLE/quote channel — the system
+ *  prompt itself shows `{"state_patch": …}` in a json-fenced example, and models echo it.
+ *  A bare-object candidate inside a ```json fence is a quotation, never a commit; only
+ *  fence-free (mangled) payloads or ```state bodies are authoring. One span scanner serves
+ *  both the parse path (findStatePatches) and the answer-scrub path (stripStateFences). */
+const JSON_FENCE = /(`{3,})[ \t]*json\b[^\r\n]*\r?\n[\s\S]*?\1/g;
+
+function jsonFenceSpans(text: string): readonly { readonly start: number; readonly end: number }[] {
+  const spans: { start: number; end: number }[] = [];
+  let m: RegExpExecArray | null;
+  JSON_FENCE.lastIndex = 0;
+  while ((m = JSON_FENCE.exec(text)) !== null) {
+    spans.push({ start: m.index, end: m.index + m[0].length });
+  }
+  return spans;
+}
+
+function insideSpans(spans: readonly { readonly start: number; readonly end: number }[], index: number): boolean {
+  for (let i = 0; i < spans.length; i++) {
+    const span = spans[i];
+    if (span !== undefined && index >= span.start && index < span.end) return true;
+  }
+  return false;
+}
+
 /** String-aware balanced-brace scan from `start` (an index of `{`). Honors string literals and
  *  backslash escapes so braces inside JSON strings cannot unbalance the count. Returns the
  *  complete object slice, or undefined when braces never balance before EOF. */
@@ -113,9 +138,13 @@ export function findStatePatches(text: string): readonly StateFenceResult[] {
     }
   }
   // Tolerant harvest over fence-free remainder (well-formed payloads already taken above).
+  // Candidates inside ```json fences are quotes/examples (recall W2) — skipped, so the
+  // contract example a model echoes cannot burn session-wide patch retries.
   const rest = sansFences(text);
+  const jsonSpans = jsonFenceSpans(rest);
   BARE_PATCH.lastIndex = 0;
   while ((m = BARE_PATCH.exec(rest)) !== null) {
+    if (insideSpans(jsonSpans, m.index)) continue;
     const obj = balancedJsonObject(rest, m.index);
     if (obj === undefined) continue;
     try {
@@ -133,13 +162,16 @@ export function findStatePatches(text: string): readonly StateFenceResult[] {
  *  bookkeeping). Deterministic scrub on the answer path; parse semantics stay in findStatePatches.
  *  Also removes bare {"state_patch"…} objects (mangled-fence leaks), a `state` token glued to
  *  preceding prose, and orphan ``` lines left behind by the mangled pair. */
-export function stripStateFences(text: string): string {
-  let out = sansFences(text);
+export function stripStateFences(text: string): string {  let out = sansFences(text);
   const parts: string[] = [];
   let cursor = 0;
+  // A {"state_patch"…} inside a ```json fence is a quotation (recall W2) — must survive in
+  // the answer, not be scrubbed as leaked bookkeeping.
+  const jsonSpans = jsonFenceSpans(out);
   BARE_PATCH.lastIndex = 0;
   let m: RegExpExecArray | null;
   while ((m = BARE_PATCH.exec(out)) !== null) {
+    if (insideSpans(jsonSpans, m.index)) continue;
     const obj = balancedJsonObject(out, m.index);
     if (obj === undefined) continue;
     // Glom any immediately-preceding bare `state`/`.state` token (prose like "...report.state {").
@@ -156,6 +188,105 @@ export function stripStateFences(text: string): string {
   return parts.join("").replace(ORPHAN_FENCE, "").trim();
 }
 
+// ── RLM-paper finalize repair (App. A: programmatic template-mistake fixing) ────────────────
+
+/** A finalize the model wrote in the paper's surface syntax instead of flipping `answer`. */
+export type FinalTag =
+  | { readonly kind: "final"; readonly value: string }
+  | { readonly kind: "final_var"; readonly value: string };
+
+const FINAL_TAG_RE = /\b(FINAL_VAR|FINAL)\s*\(/g;
+const IDENTIFIER_RE = /^[A-Za-z_][A-Za-z0-9_]*$/;
+const FINAL_TAG_VALUE_MAX = 20_000;
+
+/** All match spans of a global regex — used to excise fenced regions before prose scans. */
+function collectFenceSpans(text: string, re: RegExp): readonly { readonly start: number; readonly end: number }[] {
+  const spans: { start: number; end: number }[] = [];
+  let m: RegExpExecArray | null;
+  re.lastIndex = 0;
+  while ((m = re.exec(text)) !== null) {
+    spans.push({ start: m.index, end: m.index + m[0].length });
+  }
+  return spans;
+}
+
+/** String-aware balanced-paren scan from `start` (an index of `(`). Returns the inner slice. */
+function balancedParens(text: string, start: number): string | undefined {
+  let depth = 0;
+  let inStr: string | undefined;
+  let esc = false;
+  for (let i = start; i < text.length; i++) {
+    const ch = text.charAt(i);
+    if (inStr !== undefined) {
+      if (esc) esc = false;
+      else if (ch === "\\") esc = true;
+      else if (ch === inStr) inStr = undefined;
+      continue;
+    }
+    if (ch === '"' || ch === "'") inStr = ch;
+    else if (ch === "(") depth += 1;
+    else if (ch === ")") {
+      depth -= 1;
+      if (depth === 0) return text.slice(start + 1, i);
+    }
+  }
+  return undefined;
+}
+
+/**
+ * Detect `FINAL(x)` / `FINAL_VAR(v)` in a turn's PROSE (the RLM paper's finalize surface).
+ * The paper found 16%/13% of small-model turns misuse these tags instead of setting
+ * `answer` — and that a programmatic fix measurably improved the post-trained RLM. Fenced
+ * code is excluded (those blocks already executed as code this turn); the LAST tag wins.
+ * Returns undefined when the value is empty or dump-sized (a plan/dump is not an answer).
+ */
+export function findFinalTag(text: string): FinalTag | undefined {
+  const spans: readonly { readonly start: number; readonly end: number }[] = [
+    ...collectFenceSpans(text, FENCE),
+    ...collectFenceSpans(text, FALLBACK_FENCE),
+  ];
+  let prose = text;
+  if (spans.length > 0) {
+    const ordered = [...spans].sort((a, b) => a.start - b.start);
+    const parts: string[] = [];
+    let cursor = 0;
+    for (const s of ordered) {
+      parts.push(text.slice(cursor, s.start));
+      cursor = s.end;
+    }
+    parts.push(text.slice(cursor));
+    prose = parts.join("");
+  }
+  let found: FinalTag | undefined;
+  FINAL_TAG_RE.lastIndex = 0;
+  let m: RegExpExecArray | null;
+  while ((m = FINAL_TAG_RE.exec(prose)) !== null) {
+    const open = m.index + m[0].length - 1;
+    const inner = balancedParens(prose, open);
+    if (inner === undefined) continue;
+    const value = inner.trim().replace(/^"([\s\S]*)"$|^'([\s\S]*)'$/, "$1$2");
+    if (value === "" || value.length > FINAL_TAG_VALUE_MAX) continue;
+    found =
+      m[1] === "FINAL_VAR" && IDENTIFIER_RE.test(value)
+        ? { kind: "final_var", value }
+        : { kind: "final", value: value };
+  }
+  return found;
+}
+
+/** Python snippet that completes a FINAL_VAR(v) repair inside the sandbox — never raises. */
+export function finalVarRepairCode(name: string): string {
+  const lit = JSON.stringify(name); // identifier chars only — JSON literal == Python literal
+  return [
+    `_final_name = ${lit}`,
+    `try:`,
+    `    answer["content"] = str(globals()[_final_name])`,
+    `except KeyError:`,
+    `    answer["content"] = "Error: FINAL_VAR variable " + _final_name + " is not defined in the REPL"`,
+    `answer["ready"] = True`,
+  ].join("\n");
+}
+
 /** Truncate REPL stdout for the model's context window (head + tail, with an elision note).
  *  `mark` lets callers specialize the wording (root elision cites the session log) while the
  *  head/tail math stays the one implementation. */

package/src/text/tokens.ts

@@ -7,6 +7,9 @@
  */
 
 const CHARS_PER_TOKEN = 4;
+// CJK ideographs/kana/hangul encode ~1.5 chars per token — the flat /4 heuristic under-counts
+// CJK-heavy content ~2.7x, delaying compaction and budget walls until near overflow.
+const CJK_CHARS_PER_TOKEN = 1.5;
 
 /** Rough token count for a character length (≈4 chars/token). Always ≥ 1 for non-empty text. */
 export function estimateTokens(charCount: number): number {
@@ -14,11 +17,43 @@ export function estimateTokens(charCount: number): number {
   return Math.ceil(charCount / CHARS_PER_TOKEN);
 }
 
-/** Rough token count for a list of role/content messages. */
+/** CJK codepoints (kana, ideographs, hangul, compatibility forms) in `text`. */
+function cjkChars(text: string): number {
+  let n = 0;
+  for (let i = 0; i < text.length; i++) {
+    const c = text.charCodeAt(i);
+    if (
+      (c >= 0x3040 && c <= 0x30ff) || // kana
+      (c >= 0x3400 && c <= 0x9fff) || // ideograph extensions + unified ideographs
+      (c >= 0xac00 && c <= 0xd7af) || // hangul syllables
+      (c >= 0xf900 && c <= 0xfaff) // compatibility ideographs
+    ) {
+      n += 1;
+    }
+  }
+  return n;
+}
+
+/** Blended token estimate for actual TEXT: ASCII at 4 chars/token, CJK at 1.5. Pure-ASCII
+ *  input is byte-identical to estimateTokens(length) — only non-English content moves. */
+export function estimateTextTokens(text: string): number {
+  const cjk = cjkChars(text);
+  if (cjk === 0) return estimateTokens(text.length);
+  const ascii = text.length - cjk;
+  return Math.ceil(ascii / CHARS_PER_TOKEN + cjk / CJK_CHARS_PER_TOKEN);
+}
+
+/** Rough token count for a list of role/content messages (script-aware — see estimateTextTokens). */
 export function estimateMessageTokens(messages: { content: string }[]): number {
-  let chars = 0;
-  for (const m of messages) chars += m.content.length + 8; // small per-message overhead
-  return estimateTokens(chars);
+  let ascii = 0;
+  let cjk = 0;
+  for (const m of messages) {
+    ascii += m.content.length + 8; // small per-message overhead
+    cjk += cjkChars(m.content);
+  }
+  if (cjk === 0) return estimateTokens(ascii);
+  const nonCjk = ascii - cjk;
+  return nonCjk <= 0 ? 0 : Math.ceil(nonCjk / CHARS_PER_TOKEN + cjk / CJK_CHARS_PER_TOKEN);
 }
 
 /**