@hicaru/pi-rlm 0.3.1 → 0.3.2

package/src/context/refresh.ts

@@ -0,0 +1,141 @@
+/**
+ * Keep RLM `context` in sync with the disk after native edit/write.
+ *
+ * Seed packs file bodies once; without this, search/map_files/llm still see pre-edit text.
+ */
+
+import { readFile } from "node:fs/promises";
+import { isAbsolute, relative, resolve } from "node:path";
+import { estimateTokens } from "../text/tokens.ts";
+import type { ContextFile } from "./types.ts";
+
+/** Paths that look like tool file targets. */
+export function extractEditPaths(input: unknown): readonly string[] {
+  if (typeof input !== "object" || input === null) return Object.freeze([]);
+  const o = input as Record<string, unknown>;
+  const keys = ["path", "file_path", "filePath", "filename", "file"] as const;
+  const out: string[] = [];
+  for (const k of keys) {
+    const v = o[k];
+    if (typeof v === "string" && v.trim() !== "") out.push(v.trim());
+  }
+  // Some tools pass { path, oldText, newText } only — already covered.
+  return Object.freeze(out);
+}
+
+/**
+ * Normalize disk path to how cwd-seed entries usually appear (relative to cwd when under cwd).
+ */
+export function normalizeContextPath(filePath: string, cwd: string): string {
+  const abs = isAbsolute(filePath) ? resolve(filePath) : resolve(cwd, filePath);
+  const rel = relative(cwd, abs);
+  if (rel === "" || rel.startsWith("..") || isAbsolute(rel)) return abs;
+  return rel.split("\\").join("/");
+}
+
+function pathMatches(entryPath: string, target: string, cwd: string): boolean {
+  if (entryPath === target) return true;
+  const a = normalizeContextPath(entryPath, cwd);
+  const b = normalizeContextPath(target, cwd);
+  if (a === b) return true;
+  // suffix match for namespaced entries
+  return entryPath.endsWith("/" + target) || entryPath.endsWith(target);
+}
+
+/**
+ * Upsert one file into a context payload list. Returns a **new** array (identity change
+ * so BM25 stamp invalidates when the worker rebinds `context`).
+ */
+export function upsertContextFile(
+  payload: unknown,
+  filePath: string,
+  content: string,
+  cwd: string,
+): ContextFile[] {
+  const path = normalizeContextPath(filePath, cwd);
+  const tokens = estimateTokens(content.length);
+  const entry: ContextFile = Object.freeze({ path, content, tokens });
+
+  if (!Array.isArray(payload)) {
+    return [entry];
+  }
+
+  const next = new Array<ContextFile>(payload.length + 1);
+  let n = 0;
+  let replaced = false;
+  for (let i = 0; i < payload.length; i++) {
+    const item: unknown = payload[i];
+    if (
+      item !== null &&
+      typeof item === "object" &&
+      "path" in item &&
+      typeof (item as { path: unknown }).path === "string" &&
+      pathMatches((item as { path: string }).path, path, cwd)
+    ) {
+      next[n++] = entry;
+      replaced = true;
+    } else if (
+      item !== null &&
+      typeof item === "object" &&
+      "path" in item &&
+      "content" in item &&
+      typeof (item as { path: unknown }).path === "string" &&
+      typeof (item as { content: unknown }).content === "string"
+    ) {
+      const e = item as { path: string; content: string; tokens?: number };
+      next[n++] = Object.freeze({
+        path: e.path,
+        content: e.content,
+        tokens: typeof e.tokens === "number" ? e.tokens : estimateTokens(e.content.length),
+      });
+    }
+    // drop non-file entries silently (shouldn't appear in file bundles)
+  }
+  if (!replaced) next[n++] = entry;
+  next.length = n;
+  return next;
+}
+
+/** Read file from disk; return null if missing/unreadable. */
+export async function readDiskFile(filePath: string, cwd: string): Promise<string | null> {
+  const abs = isAbsolute(filePath) ? resolve(filePath) : resolve(cwd, filePath);
+  try {
+    return await readFile(abs, "utf8");
+  } catch {
+    return null;
+  }
+}
+
+/**
+ * Python snippet that rebinds `context` with updated path content and forces a new list id
+ * so BM25 rebuilds on next search.
+ */
+export function patchContextExecCode(filePath: string, content: string, cwd: string): string {
+  const path = normalizeContextPath(filePath, cwd);
+  // JSON for safe embedding in Python string literals
+  const pathLit = JSON.stringify(path);
+  const contentLit = JSON.stringify(content);
+  const tokens = estimateTokens(content.length);
+  return `
+_path = ${pathLit}
+_content = ${contentLit}
+_tokens = ${tokens}
+_old = context if isinstance(context, list) else []
+_next = []
+_found = False
+for _e in _old:
+    if isinstance(_e, dict) and str(_e.get("path", "")) in (_path, _path.replace("\\\\", "/")):
+        _next.append({"path": _path, "content": _content, "tokens": _tokens})
+        _found = True
+    elif isinstance(_e, dict) and (
+        str(_e.get("path", "")).endswith("/" + _path) or str(_e.get("path", "")).endswith(_path)
+    ):
+        _next.append({"path": str(_e.get("path")), "content": _content, "tokens": _tokens})
+        _found = True
+    else:
+        _next.append(_e)
+if not _found:
+    _next.append({"path": _path, "content": _content, "tokens": _tokens})
+context = _next
+`.trim();
+}

package/src/core/engine.ts

@@ -13,10 +13,10 @@ import { buildAddContextHandler } from "../bridge/add-context.ts";
 import { mergeIntoContext } from "../context/merge.ts";
 import {
   createSubcallHandlers,
+  createTaskRegistry,
   type Invocation,
-} from "../bridge/subcall-handlers.ts";
+} from "../bridge/handlers/index.ts";
 import { type ChatMsg, modelComplete } from "../bridge/model.ts";
-import { resolveModelId } from "../config/settings.ts";
 import { buildRlmSystemPrompt } from "../prompts/system.ts";
 import { buildTurnPrompt, FINALIZE_PROMPT } from "../prompts/user.ts";
 import type { RlmEmitter } from "../tool/rlm-events.ts";
@@ -30,7 +30,6 @@ import { appendUserMessage } from "./history.ts";
 import { runTurn } from "./iteration.ts";
 import { type Limits, LimitError, LimitGuard } from "./limits.ts";
 import type { RlmConfig, RlmInput, RlmResult, RunRlm, Sampling } from "./types.ts";
-import { formatError } from "../util/errors.ts";
 import { createSubcallGates, type SubcallGates } from "../util/concurrency.ts";
 
 /**
@@ -71,20 +70,7 @@ export function createEngine(deps: EngineDeps): RunRlm {
       emitter.emitTurn(0, deps.config.maxIterations);
     }
 
-    const overrideModel = input.modelOverride ? resolveModelId(deps.registry, input.modelOverride) : undefined;
-    if (input.modelOverride && !overrideModel) {
-      if (selfReportId) emitter.emitSubcallUpdated({ id: selfReportId, status: "error", detail: "unknown model override" });
-      else emitter.emitStatus("error");
-      return {
-        answer: formatError(`unknown model override '${input.modelOverride}'`),
-        iterations: 0,
-        costUsd: 0,
-        inputTokens: 0,
-        outputTokens: 0,
-        durationMs: 0,
-      };
-    }
-    const model = overrideModel ?? deps.model;
+    const model = deps.model;
 
     // Create LimitGuard BEFORE the bridge so sub-LLM usage feeds into it.
     // Children inherit the parent's remaining timeout (propagated as remaining amount, not
@@ -117,6 +103,8 @@ export function createEngine(deps: EngineDeps): RunRlm {
     // run can settle or abort it first (a child engine left running would keep spending).
     let detachedInFlight = 0;
     let detachedIdle: (() => void) | undefined;
+    // One registry per run — unawaited task reminders share the same map as await handlers.
+    const taskRegistry = createTaskRegistry();
     const subcalls = createSubcallHandlers({
       resolve: () => invocation,
       gates: deps.gates
@@ -140,7 +128,7 @@ export function createEngine(deps: EngineDeps): RunRlm {
           if (detachedInFlight === 0) detachedIdle?.();
         }
       },
-    });
+    }, taskRegistry);
     /** Wait (bounded) for detached work before the sandbox goes away. */
     const settleDetached = async (): Promise<void> => {
       if (detachedInFlight === 0) return;
@@ -189,6 +177,7 @@ export function createEngine(deps: EngineDeps): RunRlm {
         maxPromptChars: deps.config.maxPromptChars,
         contextLoader: deps.config.contextLoader,
         child: input.depth > 0,
+        depth: input.depth,
       });
 
       const contextHandlers = deps.config.contextLoader
@@ -251,6 +240,15 @@ export function createEngine(deps: EngineDeps): RunRlm {
           pendingReplOutputs = undefined;
         }
 
+        // Soft runtime nudge (rlm_test parity): remind the model to await pending host tasks.
+        const pendingIds = taskRegistry.awaitDeps.unawaitedIds();
+        if (pendingIds.length > 0) {
+          appendUserMessage(
+            history,
+            `[runtime] Unawaited task_ids: ${pendingIds.join(", ")} — call await before finish.`,
+          );
+        }
+
         appendUserMessage(history, buildTurnPrompt(i, deps.config.maxIterations));
 
         // rootSampling fields win; smartReasoning is the default reasoning when not overridden.

package/src/core/types.ts

@@ -55,7 +55,7 @@ export interface RlmConfig {
    *  Keeps each turn short so the next turn's input stays manageable.
    *  `reasoning` is read from `smartReasoning` if omitted here. */
   readonly rootSampling?: Readonly<Sampling>;
-  /** System prompt injected into every llm_query / llm_query_batched sub-call.
+  /** System prompt injected into every llm_query / llm_batch sub-call.
    *  Instructs the worker model to respond concisely.
    *  undefined = no system prompt (raw completion). */
   readonly subSystemPrompt?: string;
@@ -73,8 +73,6 @@ export interface RlmInput {
   readonly depth: number;
   /** AgentTree node to attach this run's node under (set when recursing). */
   readonly parentNodeId?: string;
-  /** "provider/id" — overrides the root model for this run (set by recursive rlm_query). */
-  readonly modelOverride?: string;
   /** Remaining timeout for this subtree (set by parent from its LimitGuard). */
   readonly remainingTimeoutMs?: number;
 }

package/src/index.ts

@@ -18,9 +18,10 @@ import { BackgroundTasks } from "./tool/background-tasks.ts";
 import { resolve } from "node:path";
 import { resolveSource } from "./context/resolve.ts";
 import { formatContextListing } from "./context/listing.ts";
+import { extractEditPaths, readDiskFile } from "./context/refresh.ts";
 import type { AddContextHandlerBundle } from "./bridge/add-context.ts";
-import { buildNativeSystemPrompt, NATIVE_TURN_REMINDER } from "./prompts/native.ts";
-import { bashCommandFromInput, isFileReadingCommand, capToolResultText, BASH_BLOCK_REASON } from "./mode/native-guards.ts";
+import { buildNativeSystemPrompt } from "./prompts/native.ts";
+import { capToolResultText } from "./mode/native-guards.ts";
 import {
   isSubagentChildBypass,
   commitSubagentForceActivation,
@@ -37,14 +38,13 @@ export {
   processRlmDepth,
 } from "./mode/subagent.ts";
 
-const BLOCKED_NATIVE_TOOLS = Object.freeze(new Set(["read", "grep"]));
 /** How often to keep the parent sandbox's request watchdog alive during detached work. */
 const WATCHDOG_HEARTBEAT_MS = 30_000;
-const CAPPED_RESULT_TOOLS = Object.freeze(new Set(["bash", "find", "ls"]));
+/** Soft token guard — cap bulk tool stdout; do NOT hard-block read/grep/bash readers. */
+const CAPPED_RESULT_TOOLS = Object.freeze(new Set(["bash", "find", "ls", "read", "grep"]));
 
 export default function rlmExtension(pi: ExtensionAPI): void {
-  // Subagent children run a native tool flow; RLM's contract is the opposite
-  // (block read/grep, route through repl). Env fast path: full bypass when
+  // Subagent children may bypass full RLM registration. Env fast path when
   // PI_SUBAGENT_CHILD=1 (unless force-in under the depth cap). See mode/subagent.ts.
   if (isSubagentChildBypass()) {
     if (traceEnabled) {
@@ -182,7 +182,23 @@ export default function rlmExtension(pi: ExtensionAPI): void {
 
     if (controller.savedLlmRef) {
       const resolved = resolveModelId(ctx.modelRegistry, controller.savedLlmRef);
-      if (resolved) controller.llmModel = resolved;
+      if (resolved) {
+        controller.llmModel = resolved;
+      } else {
+        // Keep the pin on disk/controller — do not fall back permanently. Runtime uses
+        // cheapest until the catalog has the model again; surface that once per session.
+        console.warn(
+          `[rlm] pinned sub-LLM ${controller.savedLlmRef} not in registry; using cheapest until it reappears`,
+        );
+        try {
+          ctx.ui.notify(
+            `RLM: pinned llm=${controller.savedLlmRef} unavailable — using cheapest until it is`,
+            "warning",
+          );
+        } catch {
+          // Some hosts have no UI at session_start.
+        }
+      }
     }
 
     // Re-register repl tool each session to pick up model provider changes
@@ -254,7 +270,7 @@ export default function rlmExtension(pi: ExtensionAPI): void {
     const filtered = event.messages.filter(
       (message) =>
         !(message.role === "custom" && message.customType === "rlm-intro")
-        && !(message.role === "user" && typeof message.content === "string" && message.content === NATIVE_TURN_REMINDER),
+
     );
     if (!nativeTradeHolds()) return { messages: filtered };
 
@@ -266,11 +282,11 @@ export default function rlmExtension(pi: ExtensionAPI): void {
       listingPayloadRef = payload;
       const listing = formatContextListing(payload);
       const instruction = [
-        "ANALYZE with repl({code}) — read/grep are DISABLED.",
-        "Files you have loaded live in the Python REPL `context` variable (starts empty; cwd seeds on first repl()).",
-        "Locate with search()/grep_context()/outline() (free), then delegate bulk reading to",
-        "map_files()/llm_query_batched(). Use add_context(path) for external dirs/files/docs/git URLs.",
-        "If credits exhausted → report and stop.",
+        "Prefer repl({code}) for bulk analysis: free search/grep/outline, then fan-out Tasks.",
+        "Multi-module work → rlm_batch (or rlm_query); one-shot extracts → map_files/llm_batch.",
+        "Always-spawn returns Task (↯bg); only await_task has content — fire-all then await.",
+        "Large tool/repl outputs are capped. Files live in REPL `context` (cwd seeds first repl()).",
+        "add_context(path) for external dirs/files/docs/git. Credits exhausted → report and stop.",
         "",
       ].join("\n");
       filtered.unshift({
@@ -280,42 +296,38 @@ export default function rlmExtension(pi: ExtensionAPI): void {
       } as PiMessage);
     }
 
-    // Per-turn last-position reminder (not persisted — context hook rebuilds every request)
-    filtered.push({
-      role: "user" as const,
-      content: NATIVE_TURN_REMINDER,
-      timestamp: 0,
-    } as PiMessage);
 
     return { messages: filtered };
   });
 
-  // ── Native mode restrictions: keep bulk file content out of root-model context ──
-  // `edit`/`write` stay unblocked so the agent modifies files through Pi's native
-  // tool flow (visible to all plugins, +/- diff preview). File reading/searching
-  // belongs in the REPL, and bash output is capped as a backstop.
-  // Fail-open when repl is not active (e.g. --tools allowlist without repl): never
-  // confiscate readers without a working substitute (RLM paper §2 trade).
-  pi.on("tool_call", async (event) => {
-    if (!nativeTradeHolds()) {
-      if (traceEnabled && controller.enabled) {
-        trace("native.block_skip", { toolName: event.toolName, reason: "no_active_repl" });
+  // Soft token guard only — never hard-block read/grep/bash. Large tool results are capped.
+  // After edit/write, re-read disk into RLM context so search/llm see fresh content.
+  const MUTATING_FILE_TOOLS = Object.freeze(new Set(["edit", "write"]));
+
+  pi.on("tool_result", async (event, ctx) => {
+    // ── Keep RLM context fresh after native file mutations ──
+    if (
+      nativeTradeHolds()
+      && MUTATING_FILE_TOOLS.has(event.toolName)
+      && event.isError !== true
+    ) {
+      const cwd = resolve(ctx?.cwd ?? process.cwd());
+      const paths = extractEditPaths(event.input);
+      for (const p of paths) {
+        const body = await readDiskFile(p, cwd);
+        if (body === null) continue;
+        try {
+          await sandboxManager.refreshFileFromDisk(p, body, cwd);
+          // Listing must re-inject if we rewrote payload identity
+          listingPayloadRef = undefined;
+        } catch (err) {
+          if (traceEnabled) {
+            trace("context.refresh_fail", { path: p, error: errorMessage(err) });
+          }
+        }
       }
-      return;
-    }
-    if (BLOCKED_NATIVE_TOOLS.has(event.toolName)) {
-      return {
-        block: true,
-        reason: "RLM mode active. Use repl({code}) to read files and search the repository — loaded files live in the REPL `context` variable (cwd seeds on first call). Use `edit`/`write` for file changes. If sub-LLM credits are exhausted, report to the user.",
-      };
     }
-    const bashCommand = event.toolName === "bash" ? bashCommandFromInput(event.input) : undefined;
-    if (bashCommand !== undefined && isFileReadingCommand(bashCommand)) {
-      return { block: true, reason: BASH_BLOCK_REASON };
-    }
-  });
 
-  pi.on("tool_result", async (event) => {
     if (!nativeTradeHolds() || !CAPPED_RESULT_TOOLS.has(event.toolName)) return;
     let changed = false;
     const content = event.content.map((c) => {

package/src/mode/native-guards.ts

@@ -43,7 +43,7 @@ export function isFileReadingCommand(command: string): boolean {
 export const BASH_BLOCK_REASON =
   "RLM mode: reading files via bash is blocked — that dumps file content into the root model's " +
   "context. All files are pre-loaded in the REPL `context` variable: use repl({code}) with Python " +
-  "string/regex search, and delegate bulk analysis to llm_query / llm_query_batched / " +
+  "string/regex search, and delegate bulk analysis to llm_query / llm_batch / " +
   "llm_query_chunked. bash is for RUNNING things (tests, builds, git).";
 
 /** Max chars of tool output forwarded to the root model (≈1K tokens). */
@@ -51,12 +51,12 @@ export const TOOL_RESULT_CAP = 4_000;
 
 const CAP_NOTE =
   `\n[RLM: tool output capped at ${TOOL_RESULT_CAP.toLocaleString()} chars to protect the root ` +
-  "model's context — route bulk text through repl() + llm_query_chunked / llm_query_batched.]";
+  "model's context — route bulk text through repl() + llm_query_chunked / llm_batch.]";
 
 const REPL_CAP_NOTE =
   `\n[RLM: repl() stdout capped at ${TOOL_RESULT_CAP.toLocaleString()} chars — printing bulk text ` +
   "is useless. Keep results in REPL variables and delegate semantic reading to llm_query / " +
-  "llm_query_batched / llm_query_chunked.]";
+  "llm_batch / llm_query_chunked.]";
 
 /** Shared truncation core. Returns undefined when under the cap (leave the text untouched). */
 function capText(text: string, note: string): string | undefined {
@@ -86,7 +86,7 @@ export function replDelegationNudge(stdoutChars: number, delegated: boolean): st
   if (delegated || stdoutChars <= NUDGE_STDOUT_CHARS) return undefined;
   return (
     `\n[RLM: this repl() printed ${stdoutChars.toLocaleString()} chars with 0 sub-LLM calls — ` +
-    "if you were READING, delegate via llm_query / llm_query_batched / llm_query_chunked. " +
+    "if you were READING, delegate via llm_query / llm_batch / llm_query_chunked. " +
     "Authoring an edit body yourself is correct and needs no delegation.]"
   );
 }

package/src/mode/subagent.ts

@@ -6,7 +6,7 @@
  *   2. Capability gate — never confiscate native readers unless `repl` is in the
  *      active tool set (paper trade: scaffold only if the REPL substitute exists).
  *
- * In-process rlm_query depth is handled by subcall-handlers.childRun; this module
+ * In-process rlm_query depth is handled by bridge/handlers childRun; this module
  * only covers OS-process children (pi subagents), which restart at depth 0.
  */
 import { DEFAULT_CONFIG } from "../config/defaults.ts";