@hicaru/pi-rlm 0.3.14 → 0.3.16

package/src/core/engine.ts

@@ -20,12 +20,14 @@ import { TaskLedger, contextSig, taskKey } from "./ledger.ts";
 import { type MemoryStore, rootContextPaths } from "./memory.ts";
 import { type ChatMsg, modelComplete } from "../bridge/model.ts";
 import { buildRlmSystemPrompt } from "../prompts/system.ts";
-import { buildTurnPrompt, FINALIZE_PROMPT } from "../prompts/user.ts";
+import { buildTurnPrompt, FINALIZE_PROMPT, RETRIEVAL_NUDGE, REASONING_BUDGET_HINT, VERIFICATION_NUDGE } from "../prompts/user.ts";
 import type { RlmEmitter } from "../tool/rlm-events.ts";
 import type { SubcallPhase } from "../tool/rlm-details.ts";
 import { PythonSandbox, SANDBOX_WATCHDOG_HEARTBEAT_MS } from "../sandbox/sandbox.ts";
+import type { ReplResult } from "../sandbox/protocol.ts";
 import { pinContext, type PinnedContext } from "../sandbox/context-file.ts";
 import { previewStdout, previewText } from "../text/preview.ts";
+import { findReplBlocks } from "../text/parsing.ts";
 import { contextLength, contextSizeStats, contextTypeLabel } from "../text/tokens.ts";
 import { finalAnswerOf, formatReplOutputs, latestAnswerContentOf, turnHadError } from "./answer.ts";
 import { compactHistory, elideOldToolPayloads, shouldCompact } from "./compaction.ts";
@@ -44,6 +46,9 @@ import { createSubcallGates, type SubcallGates } from "../util/concurrency.ts";
  * hold a finished run open on work whose result nobody can receive.
  */
 const DETACHED_SETTLE_MS = 5_000;
+/** Verification nudge (enableVerificationNudge): only an EARLY finalize is suspicious — from
+ *  turn 4 on, a bare answer is just... an answer. "Before iteration ~4", per the bench data. */
+const VERIFICATION_NUDGE_TURN_CAP = 4;
 /** H6 (audit): root episodes snapshot at most this many real files — replay invalidation for
  *  the disk-backed slice of the context without hashing an unbounded repository. */
 const ROOT_HASH_MAX = 64;
@@ -248,6 +253,13 @@ export function createEngine(deps: EngineDeps): RunRlm {
     // v5 budget cascade state: the wrap-up note fires for exactly ONE turn after crossing soft.
     let softFired = false;
     let softNoteTurn = -1;
+    // H3: retrieval-discipline coach — one-shot per run; children inherit it via the same loop.
+    let sawRetrieval = false;
+    let retrievalNudged = false;
+    // Verification-discipline coach (enableVerificationNudge, default OFF): one coached redo
+    // when an early finalize looks like the confident-wrong bench shape.
+    let verificationNudged = false;
+    let verificationNudgePending = false;
 
     try {
       const meta = {
@@ -315,6 +327,13 @@ export function createEngine(deps: EngineDeps): RunRlm {
       liveContext = input.context ?? [];
       contextPin = await pinContext(liveContext);
       await sandbox.loadContextPinned(contextPin);
+
+      // rootSampling fields win; smartReasoning is the default reasoning when not overridden.
+      // Loop-invariant — built once here; finalize() applies the same merge to its own turn.
+      const rootSampling: Sampling = {
+        reasoning: deps.config.smartReasoning,
+        ...deps.config.rootSampling,
+      };
       for (let i = 0; i < deps.config.maxIterations; i++) {
         limits.checkTimeout();
         if (selfReportId) emitter.emitSubcallUpdated({ id: selfReportId, detail: `turn ${i + 1}/${deps.config.maxIterations}` });
@@ -355,21 +374,27 @@ export function createEngine(deps: EngineDeps): RunRlm {
         // v5 [ledger] blackboard + [memory] notes — each silent ("") when it has nothing to say.
         const ledgerBlock = deps.config.enableLedger ? runLedger.injectBlock() : "";
         const memoryBlock = rootMemory !== undefined ? rootMemory.injectBlock(input.rootPrompt) : "";
+        // H3: after two retrieval-free turns, inject the coach nudge exactly once, for one turn.
+        const nudgeNow = i >= 2 && !sawRetrieval && !retrievalNudged;
+        if (nudgeNow) retrievalNudged = true;
         const notes =
           [
             i === softNoteTurn ? WRAP_UP_BUDGET : undefined,
             ledgerBlock === "" ? undefined : ledgerBlock,
             memoryBlock === "" ? undefined : memoryBlock,
+            nudgeNow ? RETRIEVAL_NUDGE : undefined,
+            verificationNudgePending ? VERIFICATION_NUDGE : undefined,
+            // One-shot (turn 0 only): thinking tokens share the completion budget — mirror of
+            // the bench's doubling rule. Advisory; never fatal, never repeated.
+            i === 0 && rootSampling.reasoning !== undefined && (rootSampling.maxTokens ?? 16_384) < 8_192
+              ? REASONING_BUDGET_HINT
+              : undefined,
           ]
             .filter((s): s is string => s !== undefined)
             .join("\n\n") || undefined;
+        verificationNudgePending = false;
         appendUserMessage(history, buildTurnPrompt(i, deps.config.maxIterations, notes));
 
-        // rootSampling fields win; smartReasoning is the default reasoning when not overridden.
-        const rootSampling: Sampling = {
-          reasoning: deps.config.smartReasoning,
-          ...deps.config.rootSampling,
-        };
         const turn = await runTurn(history, sandbox, {
           model: model,
           registry: deps.registry,
@@ -379,6 +404,7 @@ export function createEngine(deps: EngineDeps): RunRlm {
           complete: deps.complete,
           onPhase: reportPhase,
         });
+        if (turn.blocks.some((b) => /\b(?:search|grep_context)\s*\(/.test(b))) sawRetrieval = true;
         const allBlocks = turn.blocks.length > 0
           ? turn.blocks.map((b) => previewText(b, 400)).join("\n")
           : previewText(turn.response, 400);
@@ -386,8 +412,17 @@ export function createEngine(deps: EngineDeps): RunRlm {
           emitter.emitSubcallUpdated({ id: selfReportId, args: `▶ ${allBlocks}`, resultPreview: previewStdout(turn.results) });
         }
         limits.addUsage(turn.usage);
-        if (selfReportId) emitter.emitSubcallUpdated({ id: selfReportId, costUsd: turn.usage.cost.total, tokens: turn.usage.totalTokens });
-        else emitter.emitRootUsage(turn.usage.cost.total, turn.usage.totalTokens);
+        if (selfReportId) {
+          emitter.emitSubcallUpdated({
+            id: selfReportId,
+            costUsd: turn.usage.cost.total,
+            tokens: turn.usage.totalTokens,
+            tokensIn: turn.usage.input,
+            tokensOut: turn.usage.output,
+          });
+        } else {
+          emitter.emitRootUsage(turn.usage.cost.total, turn.usage.totalTokens, turn.usage.input, turn.usage.output);
+        }
         deps.onUsage?.(turn.usage, "root");
         const answerContent = latestAnswerContentOf(turn.results);
         if (answerContent) best = answerContent;
@@ -395,10 +430,19 @@ export function createEngine(deps: EngineDeps): RunRlm {
         completedTurns = i + 1;
         const final = finalAnswerOf(turn.results);
         if (final != null) {
-          const done = result(final, i + 1, limits);
-          persistRoot(done.answer);
-          lastAnswer = done.answer;
-          return done;
+          // Verification-discipline nudge (enableVerificationNudge, default OFF): an early
+          // finalize whose answer is a bare number / short label is the confident-wrong shape
+          // that dominated bench failures. ONE coached redo, then the answer is accepted.
+          if (deps.config.enableVerificationNudge === true && !verificationNudged
+            && completedTurns < VERIFICATION_NUDGE_TURN_CAP && isBareAnswer(final)) {
+            verificationNudged = true;
+            verificationNudgePending = true;
+          } else {
+            const done = result(final, i + 1, limits);
+            persistRoot(done.answer);
+            lastAnswer = done.answer;
+            return done;
+          }
         }
 
         limits.observe(turnHadError(turn.results));
@@ -458,7 +502,7 @@ export function createEngine(deps: EngineDeps): RunRlm {
         }
       }
       if (pendingReplOutputs) appendUserMessage(history, pendingReplOutputs);
-      const finalized = result(await finalize(history, model, deps, limits), deps.config.maxIterations, limits);
+      const finalized = result(await finalize(history, model, deps, limits, sandbox), deps.config.maxIterations, limits);
       persistRoot(finalized.answer);
       lastAnswer = finalized.answer;
       return finalized;
@@ -516,17 +560,48 @@ function contextWindowOrFallback(model: Model<Api>, registry: ModelContextRegist
   return registry.limitFor(`${model.provider}/${model.id}`);
 }
 
-/** Out of turns: ask the model for its best final answer (plain text). */
-async function finalize(history: ChatMsg[], model: Model<Api>, deps: EngineDeps, limits: LimitGuard): Promise<string> {
+/** Bare number / short label — the early-confident answer shape the verification nudge
+ *  targets (28/33 bench failures were early confident wrong answers). */
+function isBareAnswer(answer: string): boolean {
+  const t = answer.trim();
+  return t.length <= 12 || /^[-+$(€£¥]?\d+(?:[.,]\d+)*\s*%?$/.test(t);
+}
+
+/** Out of turns: ask the model for its best final answer. FINALIZE_PROMPT asks for a fenced
+ *  ```repl``` block, so execute it like any turn and prefer the captured answer (H2) — a raw
+ *  fence echoed verbatim must never become the run answer. Plain text stays the fallback. */
+async function finalize(
+  history: ChatMsg[],
+  model: Model<Api>,
+  deps: EngineDeps,
+  limits: LimitGuard,
+  sandbox: PythonSandbox,
+): Promise<string> {
   const finalHistory = [...history];
   appendUserMessage(finalHistory, FINALIZE_PROMPT);
   const complete = deps.complete ?? modelComplete;
+  // Same merge rule as the main loop (the `rootSampling` construction in run()): rootSampling
+  // wins, smartReasoning is the reasoning default. Finalize is a root turn — it must obey the
+  // user's sampling too, or the last turn of every run silently reverts to provider defaults.
+  const rootSampling: Sampling = {
+    reasoning: deps.config.smartReasoning,
+    ...deps.config.rootSampling,
+  };
   const { text, usage } = await complete(finalHistory, {
     model,
     registry: deps.registry,
-    reasoning: deps.config.smartReasoning,
+    maxTokens: rootSampling.maxTokens,
+    temperature: rootSampling.temperature,
+    reasoning: rootSampling.reasoning,
     signal: deps.signal,
   });
   limits.addUsage(usage);
+  const blocks = findReplBlocks(text);
+  const results = new Array<ReplResult>(blocks.length);
+  for (let i = 0; i < blocks.length; i++) {
+    results[i] = await sandbox.exec(blocks[i]);
+  }
+  const final = finalAnswerOf(results) ?? latestAnswerContentOf(results);
+  if (final !== null && final.trim() !== "") return final.trim();
   return text.trim();
 }

package/src/core/types.ts

@@ -110,6 +110,11 @@ export interface RlmConfig {
   /** v5 doctrine: "delegation" = child engines get llm/memory/ledger only (no repo retrieval);
    *  "legacy" keeps today's full child surface as a one-flip rollback. */
   readonly childSurface: "delegation" | "legacy";
+  /** Verification-discipline nudge (default OFF — it changes interactive behavior): when the
+   *  root finalizes before turn 4 with a bare number / short label, it gets ONE coached redo
+   *  ("recompute and sanity-check in Python") instead of accepting the answer. Opt-in via
+   *  rlm.json; evidence: 28/33 bench failures were early confident wrong answers. */
+  readonly enableVerificationNudge?: boolean;
 }
 
 /** Input to a (headless) RLM run. */

package/src/prompts/user.ts

@@ -26,3 +26,28 @@ export const FINALIZE_PROMPT =
   "You are out of turns. Finalize NOW: set `answer[\"content\"]` and `answer[\"ready\"] = True` " +
     "(fenced ```repl```) with your best final answer from everything you have gathered. " +
     "Only if the REPL is unavailable, answer as plain text.";
+
+/** One-shot retrieval-discipline nudge (the engine owns the when — see core/engine.ts). The
+ *  context is external by design, so a model that never calls search()/grep_context() is
+ *  guessing from padding vocabulary; after two retrieval-free turns it gets this once. */
+export const RETRIEVAL_NUDGE =
+  "[coach] You have not inspected the external context yet — it is NOT included in this " +
+  "chat, and guessing is useless: the text is padding. On THIS turn, call search(\"...\") " +
+  "or grep_context(\"...\") inside a ```repl block before answering.";
+
+/** One-shot budget hint (engine-owned, turn 0 only): reasoning tokens share the completion
+ *  budget with the answer, mirroring the bench's doubling rule — a reasoning root with a
+ *  small output cap risks truncated thought. Advisory only; never fatal. */
+export const REASONING_BUDGET_HINT =
+  "[budget] Reasoning is on while the root output cap is below 8192 tokens: thinking shares " +
+  "the completion budget with the answer, so long thought may be cut off mid-reasoning. " +
+  "Keep thought concise, or raise rootSampling.maxTokens.";
+
+/** One-shot verification-discipline nudge (default OFF — enableVerificationNudge): fired when
+ *  the root finalizes suspiciously early with a bare number / short label. The model gets ONE
+ *  coached redo instead of having the answer accepted. */
+export const VERIFICATION_NUDGE =
+  "[coach] That answer was submitted suspiciously early and looks under-verified. Before " +
+  "finalizing: recompute the key quantity inside a ```repl block (show the actual computation, " +
+  "not a restatement), sanity-check it against the source material, and only then set " +
+  "answer[\"content\"] again.";

package/src/sandbox/py/retrieval.py

@@ -36,6 +36,30 @@ def _chunk_text(text: str, chunk_chars: int) -> list[str]:
     return chunks
 
 
+def _snippet_window(text: str, terms: set[str]) -> str:
+    """Slice `text` around the earliest occurrence of any query term, capped at _SNIPPET_CHARS.
+
+    BM25 finds the right window chunk; the snippet must show the match, not the chunk head.
+    Clipped edges get "..." markers, which count toward the cap (the body is trimmed to fit).
+    Falls back to the chunk head when no term occurs (tokenize/camelCase mismatches).
+    """
+    lowered = text.lower()
+    hits = [p for p in (lowered.find(t) for t in terms) if p >= 0]
+    if not hits:
+        return text[:_SNIPPET_CHARS]
+    start = max(0, min(hits) - _SNIPPET_LEAD)
+    end = min(len(text), start + _SNIPPET_CHARS)
+    lead = "..." if start > 0 else ""
+    trail = "..." if end < len(text) else ""
+    body = text[start:end]
+    if len(body) > _SNIPPET_CHARS - len(lead) - len(trail):
+        body = body[:_SNIPPET_CHARS - len(lead) - len(trail)]
+    if start + len(body) < len(text):
+        trail = "..."  # trimming pulled the window edge back inside the chunk
+        body = body[:_SNIPPET_CHARS - len(lead) - len(trail)]
+    return lead + body + trail
+
+
 # ---- deterministic retrieval over `context` -----------------------------------------------
 #
 # The RLM paper's trajectories retrieve by having the root model hand-write regex over the
@@ -48,6 +72,7 @@ def _chunk_text(text: str, chunk_chars: int) -> list[str]:
 _INDEX_WINDOW_LINES = 40       # a window is the retrieval unit: big enough to carry meaning
 _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
 _GREP_HARD_CAP = 200           # absolute ceiling on returned grep hits, whatever k asks for
 _BM25_K1 = 1.2
 _BM25_B = 0.75
@@ -165,9 +190,10 @@ class _Bm25Index:
             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 = text[:_SNIPPET_CHARS]
+            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],
@@ -209,7 +235,10 @@ def grep_context(
     `counts` is complete even when `hits` is capped.
     """
     try:
-        rx = re.compile(pattern)
+        # MULTILINE: the doc-level gate below must not veto line-anchored patterns (^/$) —
+        # without it, `^foo` on a multi-line doc never matches outside position 0 and every
+        # line hit is silently filtered out (grep is line-oriented; match that).
+        rx = re.compile(pattern, re.MULTILINE)
     except re.error as e:
         return {"hits": [], "counts": {}, "total": 0, "truncated": False, "error": f"bad regex: {e}"}
     try:

package/src/sandbox/py/worker.py

@@ -22,6 +22,7 @@ awaited in a LATER exec than the one that started it.
 from __future__ import annotations
 
 import argparse
+import ast
 import io
 import json
 import os
@@ -61,7 +62,15 @@ class _AnswerDict(dict):
     def __setitem__(self, key, value):
         super().__setitem__(key, value)
         if key == "ready" and value:
-            self._on_ready(self.get("content", ""))
+            content = self.get("content", "")
+            # An empty ready-flip must not capture "": defer it — a later non-blank content
+            # assignment while ready stays True (branch below) is the real submission.
+            if str(content).strip():
+                self._on_ready(content)
+        elif key == "content" and self.get("ready"):
+            content = str(value)
+            if content.strip():
+                self._on_ready(content)
 
 
 
@@ -151,7 +160,7 @@ class Worker(WorkerScaffold):
             if isinstance(cur, dict):
                 for k, v in cur.items():
                     dict.__setitem__(ans, k, v)
-                if cur.get("ready") and self._final_answer is None:
+                if cur.get("ready") and self._final_answer is None and str(cur.get("content", "")).strip():
                     self._final_answer = str(cur.get("content", ""))
             ns["answer"] = ans
         # Single context variable (RLM paper: the context lives in the environment and
@@ -295,8 +304,25 @@ class Worker(WorkerScaffold):
 
     def _exec(self, code: str, ns: dict[str, Any]) -> None:
         t = self.exec_timeout_s
+        # Jupyter-style auto-echo: plain exec() discards a trailing bare expression's value
+        # (models write `search("...")` and never see the hits), so split the block: exec the
+        # head, eval the tail, and repr() the value when it is not None (None stays silent).
+        tree = ast.parse(code)
+        tail = tree.body[-1] if tree.body else None
+        tail_is_expr = isinstance(tail, ast.Expr)
+
+        def _run() -> None:
+            if tail_is_expr and tail is not None:
+                head = ast.Module(body=tree.body[:-1], type_ignores=[])
+                exec(compile(head, "<repl>", "exec"), ns, ns)  # noqa: S102
+                value = eval(compile(ast.Expression(tail.value), "<repl>", "eval"), ns, ns)  # noqa: S102
+                if value is not None:
+                    print(repr(value))
+            else:
+                exec(compile(tree, "<repl>", "exec"), ns, ns)  # noqa: S102
+
         if t <= 0 or not hasattr(signal, "SIGALRM"):
-            exec(compile(code, "<repl>", "exec"), ns, ns)  # noqa: S102
+            _run()
             return
 
         def _alarm(signum, frame):  # noqa: ARG001
@@ -305,7 +331,7 @@ class Worker(WorkerScaffold):
         old = signal.signal(signal.SIGALRM, _alarm)
         signal.setitimer(signal.ITIMER_REAL, t)
         try:
-            exec(compile(code, "<repl>", "exec"), ns, ns)  # noqa: S102
+            _run()
         finally:
             signal.setitimer(signal.ITIMER_REAL, 0)
             signal.signal(signal.SIGALRM, old)

package/src/sandbox/sandbox.ts

@@ -344,6 +344,8 @@ export class PythonSandbox {
         frame: msg.type,
         id: "id" in msg ? msg.id : undefined,
         rid: "rid" in msg ? msg.rid : undefined,
+        // Post-mortem needs the payload, not just the shape (mirrors repl-tool's 400-char cap).
+        ...(msg.type === "exec" ? { chars: msg.code.length, code: msg.code.slice(0, 400) } : {}),
       });
     }
     // Never write to a corpse. The write would fail asynchronously and, historically, take the
@@ -423,7 +425,18 @@ export class PythonSandbox {
           prompts: "prompts" in msg ? msg.prompts?.length : 1,
         });
       } else {
-        trace("frame.in", { frame: "response", id: msg.id, ok: msg.ok });
+        trace("frame.in", {
+          frame: "response",
+          id: msg.id,
+          ok: msg.ok,
+          ...(msg.ok
+            ? {
+                stdout: msg.stdout?.slice(0, 300),
+                finalAnswer: msg.final_answer ?? undefined,
+                vars: msg.var_names?.length,
+              }
+            : { error: msg.error?.slice(0, 200) }),
+        });
       }
     }
     if (isInterrupt(msg)) {

package/src/text/parsing.ts

@@ -2,23 +2,44 @@
  * Parsing helpers: extract ```repl``` code blocks from a model response.
  *
  * The RLM root model emits Python wrapped in fenced blocks tagged `repl`. We extract those
- * blocks in order; everything else is prose the model uses to think out loud.
+ * blocks in order; everything else is prose the model uses to think out loud. Small instruct
+ * models often finalize inside ```python / ```py fences (or bare ones) instead; when a response
+ * contains no `repl` block at all, those code-ish fences are used as a fallback so the engine
+ * still executes their code. Other language tags (```text, ```json, ...) are never executed —
+ * the sandbox only runs Python.
  */
 
 const FENCE = /(`{3,})[ \t]*repl[ \t]*\r?\n([\s\S]*?)\1/g;
+// The info string is captured so a rejected tag (e.g. ```text) still consumes its whole fence —
+// otherwise that fence's own closing ``` could later match as a bare opener and swallow code.
+const FALLBACK_FENCE = /(`{3,})[ \t]*([^`\r\n]*)[ \t]*\r?\n([\s\S]*?)\1/g;
+const PYTHON_TAG = /^py(thon)?$/i;
 
-/** Return every ```repl``` block body, in document order. */
-export function findReplBlocks(text: string): string[] {
+/** Shared fence scan: run `re` over `text`, keep bodies the selector accepts (same trimming). */
+function collectFences(text: string, re: RegExp, select: (m: RegExpExecArray) => string | null): string[] {
   const blocks: string[] = [];
   let m: RegExpExecArray | null;
-  FENCE.lastIndex = 0;
-  while ((m = FENCE.exec(text)) !== null) {
-    const code = m[2] ?? "";
-    if (code.trim()) blocks.push(code.replace(/\s+$/, ""));
+  re.lastIndex = 0;
+  while ((m = re.exec(text)) !== null) {
+    const code = select(m);
+    if (code !== null && code.trim()) blocks.push(code.replace(/\s+$/, ""));
   }
   return blocks;
 }
 
+/**
+ * Return every ```repl``` block body, in document order. If the response has none, fall back to
+ * ```python / ```py / untagged fences — never other language tags, and never a mix of both kinds.
+ */
+export function findReplBlocks(text: string): string[] {
+  const repl = collectFences(text, FENCE, (m) => m[2] ?? "");
+  if (repl.length > 0) return repl;
+  return collectFences(text, FALLBACK_FENCE, (m) => {
+    const tag = m[2] ?? "";
+    return tag === "" || PYTHON_TAG.test(tag) ? (m[3] ?? "") : null;
+  });
+}
+
 /** Truncate REPL stdout for the model's context window (head + tail, with an elision note). */
 export function truncateOutput(text: string, limit = 20_000): string {
   if (text.length <= limit) return text;

package/src/tool/repl-tool.ts

@@ -267,6 +267,8 @@ export function createReplTool(deps: ReplToolDeps): ToolDefinition<typeof ReplTo
           return modelRef(m) ?? m.id;
         },
         rootTokens: () => store.getRootUsage().tokens,
+        rootTokensIn: () => store.getRootUsage().tokensIn,
+        rootTokensOut: () => store.getRootUsage().tokensOut,
       });
       let capturedStdout = "";
       let capturedStderr = "";

package/src/tool/rlm-aggregator.ts

@@ -55,7 +55,7 @@ export class RlmEventAggregator extends EmitterListener {
   }
 
   private handleRootUsage(event: RootUsageEvent): void {
-    this.store.addRootUsage(event.costUsd, event.tokens);
+    this.store.addRootUsage(event.costUsd, event.tokens, event.tokensIn, event.tokensOut);
     this.notify();
   }
 

package/src/tool/rlm-details.ts

@@ -33,6 +33,9 @@ export interface RlmSubcall {
   readonly endedAt?: number;
   readonly costUsd: number;
   readonly tokens: number;
+  /** In/out split (input / output) — mirrors tokens. */
+  readonly tokensIn: number;
+  readonly tokensOut: number;
   /** For batch subcalls: failed prompt count (partial failure). */
   readonly failedCount?: number;
   /** For batch subcalls: total prompt count. */
@@ -46,7 +49,7 @@ export interface RlmDetails {
   readonly rootPrompt: string;
   readonly turns: { readonly current: number; readonly max: number };
   readonly subcalls: readonly RlmSubcall[];
-  readonly totals: { readonly costUsd: number; readonly tokens: number };
+  readonly totals: { readonly costUsd: number; readonly tokens: number; readonly tokensIn: number; readonly tokensOut: number };
   readonly answer?: string;
 }
 

package/src/tool/rlm-events.ts

@@ -41,6 +41,9 @@ export interface SubcallUpdatedEvent {
   readonly costUsd?: number;
   /** Delta — additive on both the subcall and running totals. */
   readonly tokens?: number;
+  /** Deltas for the in/out split shown in the tree (input / output). Additive like tokens. */
+  readonly tokensIn?: number;
+  readonly tokensOut?: number;
   /** For batch subcalls: failed prompt count. */
   readonly failedCount?: number;
   /** For batch subcalls: total prompt count. */
@@ -55,6 +58,8 @@ export interface TurnEvent {
 export interface RootUsageEvent {
   readonly costUsd: number;
   readonly tokens: number;
+  readonly tokensIn?: number;
+  readonly tokensOut?: number;
 }
 
 export interface AnswerEvent {
@@ -114,8 +119,8 @@ export class RlmEmitter {
   }
 
   /** Accumulate usage directly to root-level totals. */
-  emitRootUsage(costUsd: number, tokens: number): void {
-    this.ee.emit("root-usage", { costUsd, tokens } satisfies RootUsageEvent);
+  emitRootUsage(costUsd: number, tokens: number, tokensIn?: number, tokensOut?: number): void {
+    this.ee.emit("root-usage", { costUsd, tokens, tokensIn, tokensOut } satisfies RootUsageEvent);
   }
 
   /** Set the final answer text (root-only). */

package/src/tool/rlm-tool.ts

@@ -53,7 +53,7 @@ export function createRlmTool(controller: RlmController, runRegistry?: RunRegist
         rootPrompt: "",
         turns: { current: 0, max: 0 },
         subcalls: [],
-        totals: { costUsd: 0, tokens: 0 },
+        totals: { costUsd: 0, tokens: 0, tokensIn: 0, tokensOut: 0 },
       }));
       if (!validation.ok) return validation.error;
       const params = validation.value;
@@ -79,6 +79,8 @@ export function createRlmTool(controller: RlmController, runRegistry?: RunRegist
           return m === undefined ? undefined : modelRef(m) ?? m.id;
         },
         rootTokens: () => aggregator.getRootUsage().tokens,
+        rootTokensIn: () => aggregator.getRootUsage().tokensIn,
+        rootTokensOut: () => aggregator.getRootUsage().tokensOut,
       });
 
       // Wire abort signal to controller

package/src/tool/subcall-render.ts

@@ -9,7 +9,7 @@
 import { Text } from "@earendil-works/pi-tui";
 import { keyText } from "@earendil-works/pi-coding-agent";
 import type { SubcallStatus } from "./rlm-details.ts";
-import { formatTokens, spinnerFrame } from "../ui/theme.ts";
+import { formatTokens, formatTokensSplit, spinnerFrame } from "../ui/theme.ts";
 import type { Theme } from "@earendil-works/pi-coding-agent";
 
 // ── Glyphs ──
@@ -25,15 +25,21 @@ export function headlineStatusGlyph(status: SubcallStatus | "aborted" | "done",
 
 // ── Stats formatting ──
 
-/** The `4.2k tok · 812ms` run of a card header. Omits any zero component. */
+/** The `4.2k tok · 812ms` run of a card header. In/out split when known; omits zero parts. */
 export function cardStatsLine(
-  totals: { readonly tokens: number },
+  totals: { readonly tokens: number; readonly tokensIn?: number; readonly tokensOut?: number },
   theme: Theme,
   extra?: string,
   backgroundPending?: number,
 ): string {
   const parts: string[] = [];
-  if (totals.tokens > 0) parts.push(`${formatTokens(totals.tokens)} tok`);
+  if (totals.tokens > 0) {
+    parts.push(
+      totals.tokensOut !== undefined && totals.tokensOut > 0
+        ? formatTokensSplit(totals.tokensIn ?? totals.tokens, totals.tokensOut)
+        : `${formatTokens(totals.tokens)} tok`,
+    );
+  }
   if (extra) parts.push(extra);
   const line = theme.fg("dim", parts.join(" · "));
   // The one thing no single line can show: spawned work that may outlive this block.