@hicaru/pi-rlm 0.1.5 → 0.1.6

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@hicaru/pi-rlm",
-  "version": "0.1.5",
+  "version": "0.1.6",
   "type": "module",
   "description": "Save 99% tokens, Recursive Language Model (RLM) for the Pi",
   "license": "MIT",

package/src/index.ts

@@ -7,6 +7,8 @@ import { registerRlmCommand } from "./commands/rlm.ts";
 import { registerRlmConfigCommand } from "./commands/rlm-config.ts";
 import { createRlmTool } from "./tool/rlm-tool.ts";
 import { createReplTool } from "./tool/repl-tool.ts";
+import { createApplyEditsTool } from "./tool/apply-edits-tool.ts";
+import { EditRegistry } from "./registry/edit-registry.ts";
 import { loadSettings, mergeConfig, resolveModelId } from "./config/settings.ts";
 import { RlmController, cheapestModel } from "./mode/rlm-mode.ts";
 import { postRlmGuide } from "./ui/intro.ts";
@@ -24,12 +26,14 @@ export default function rlmExtension(pi: ExtensionAPI): void {
   // Init synchronously with defaults — ensures commands/tools/handlers register before session_start
   const config = mergeConfig({});
   const controller = new RlmController(config);
+  const editRegistry = new EditRegistry();
   const sandboxManager = new SandboxManager({
     execTimeoutS: config.execTimeoutS,
     requestTimeoutMs: config.requestTimeoutMs,
     python: config.python,
     sandboxInitTimeoutMs: config.sandboxInitTimeoutMs,
     maxPromptChars: config.maxPromptChars,
+    onSandboxDiscarded: () => { editRegistry.clear(); },
   });
   let packedContextText: string | undefined;
   let contextPackPromise: Promise<string | undefined> | undefined;
@@ -77,6 +81,7 @@ export default function rlmExtension(pi: ExtensionAPI): void {
 
   // ── Tool registration ──
   pi.registerTool(createRlmTool(controller));
+  pi.registerTool(createApplyEditsTool(editRegistry));
   let guidePosted = false;
 
   pi.on("session_start", async (_event, ctx) => {
@@ -100,6 +105,7 @@ export default function rlmExtension(pi: ExtensionAPI): void {
           getModel: () => controller.resolveModels(ctx)?.model,
           getWorkerModel: () => controller.resolveModels(ctx)?.worker,
           registry: ctx.modelRegistry,
+          editRegistry,
           config: controller.config,
           ensureContext: async () => {
             const contextText = await ensureRepositoryContext(ctx.cwd ?? process.cwd());

package/src/prompts/system.ts

@@ -233,8 +233,8 @@ function nativeReplGlossary(): string {
     "",
     "- `todo(action, **kwargs) -> str` — manage a task list. Actions: create, update, list, get, delete, clear. Statuses: pending → in_progress → completed.",
     "- `SHOW_VARS() -> str` — list all variables currently in the REPL.",
-    "- `stage_edit(path, old_text, new_text)`: stage exact edits from `context`; apply returned STAGED_EDITS verbatim with edit().",
-    "- `answer`: dict `{\"content\": \"\", \"ready\": False}`. To submit: `answer[\"content\"] = \"...\"; answer[\"ready\"] = True`.",
+    "- `stage_edit(path, old_text, new_text) -> str`: stages an edit and returns an edit ID; apply IDs with `apply_edits`.",
+    "- `answer`: dict `{\"content\": \"\", \"ready\": False}`. Setting `answer[\"ready\"] = True` delivers it to the user; do not restate it.",
     "",
     "### Orchestrator Pattern",
     "You are an **orchestrator, not a solver**. After probing `context`, decompose the task into sub-LLM / REPL steps,",
@@ -267,13 +267,13 @@ function nativeReplGlossary(): string {
     "| `llm_query` (inside repl) | Extract, summarize, or classify a chunk of text |",
     "| `rlm_query` (inside repl) | Complex sub-task needing iterative reasoning with its own REPL |",
     "| `todo` (inside repl) | Track multi-step progress visibly to the user |",
-    "| `stage_edit(path, old, new)` (inside repl) | Sub-agent stages exact edit params; relay STAGED_EDITS to `edit` |",
+    "| `stage_edit(path, old, new)` (inside repl) | Stage exact edit params; apply returned IDs with `apply_edits` |",
     "",
     "### Workflow",
     "1. **Plan**: Create todos for the multi-step analysis. Probe `context` — print length, inspect a few entries.",
     "2. **Chunk & Delegate**: Slice `context` into batches, delegate each batch to sub-LLMs via `llm_query_batched`.",
     "3. **Aggregate**: Collect results in Python, pass aggregated results to a final `llm_query` or produce the answer directly.",
-    "4. **Finalize**: For file changes, stage them inside repl() via stage_edit(path, old, new), then relay the STAGED_EDITS from the result to `edit`. For analysis tasks, write a normal message.",
+    "4. **Finalize**: For file changes, stage them inside repl() via `stage_edit(path, old, new)`, then apply the returned IDs with `apply_edits({ ids })`. For analysis tasks, write a normal message.",
     "",
     "### Task-Specific Patterns",
     LARGE_FILE_RULE_NATIVE,
@@ -305,12 +305,9 @@ export function buildNativeSystemPrompt(): string {
     "All file content is pre-loaded in the REPL `context` variable. Use ONLY `repl({code})`.",
     "If sub-LLM credits are exhausted → report the error to the user and stop.",
     "",
-    "For file changes, use `edit` (modify existing) or `write` (create new) — these route through",
-    "Pi's native tool flow, visible to all plugins with a `+/-` diff preview.",
-    "",
-    "When repl() returns a STAGED_EDITS block, apply each entry by calling `edit` verbatim:",
-    "  edit({ path: entry.path, edits: [{ oldText: entry.oldText, newText: entry.newText }] })",
-    "Do not analyze or modify the parameters — relay them exactly as provided by the sub-agent.",
+    "For file changes, prefer `stage_edit()` inside repl(); it returns edit IDs and keeps edit bodies out of your output.",
+    "When repl() returns STAGED_EDITS, apply them with `apply_edits({ ids: [\"e1\", ...] })`.",
+    "Never re-type file paths, oldText, newText, file bodies, or `answer[\"content\"]` in your own output.",
     "",
     nativeReplGlossary(),
   ].join("\n");

package/src/registry/edit-registry.ts

@@ -0,0 +1,22 @@
+import type { ProposedEdit } from "../sandbox/protocol.ts";
+
+export class EditRegistry {
+  private readonly edits = new Map<string, ProposedEdit>();
+
+  registerAll(edits: readonly ProposedEdit[] | undefined): void {
+    if (edits === undefined) return;
+    for (const edit of edits) this.edits.set(edit.id, edit);
+  }
+
+  get(id: string): ProposedEdit | undefined {
+    return this.edits.get(id);
+  }
+
+  delete(id: string): boolean {
+    return this.edits.delete(id);
+  }
+
+  clear(): void {
+    this.edits.clear();
+  }
+}

package/src/sandbox/protocol.ts

@@ -27,6 +27,7 @@ export interface LlmReply {
 export type ParentMessage = WorkerRequest | LlmReply;
 
 export interface ProposedEdit {
+  readonly id: string;
   readonly path: string;
   readonly oldText: string;
   readonly newText: string;

package/src/sandbox/sandbox-manager.ts

@@ -15,6 +15,7 @@ export interface SandboxManagerConfig {
   readonly sandboxInitTimeoutMs: number;
   readonly maxPromptChars: number;
   readonly signal?: AbortSignal;
+  readonly onSandboxDiscarded?: () => void;
 }
 
 export class SandboxManager {
@@ -118,6 +119,7 @@ export class SandboxManager {
         try { await this.sandbox.dispose(); } catch { /* already dead */ }
         this.sandbox = null;
         this.contextLoaded = false;
+        this.config.onSandboxDiscarded?.();
       }
       throw err;
     } finally {
@@ -141,7 +143,10 @@ export class SandboxManager {
     if (this.disposed) return;
     this.disposed = true;
     await this.sandbox?.dispose();
-    this.sandbox = null;
-    this.contextLoaded = false;
+    if (this.sandbox !== null) {
+      this.sandbox = null;
+      this.contextLoaded = false;
+      this.config.onSandboxDiscarded?.();
+    }
   }
 }

package/src/sandbox/worker.py

@@ -132,6 +132,7 @@ class Worker:
         self.ns = {"__builtins__": _SAFE_BUILTINS.copy(), "__name__": "__main__"}
         self._ctx_payloads: dict[int, Any] = {}
         self._staged_edits: list[dict[str, str]] = []
+        self._edit_counter = 0
         self._nudged: set[str] = set()
         self._restore_scaffold()
 
@@ -325,8 +326,10 @@ class Worker:
     def _stage_edit(self, path: str, old_text: str, new_text: str) -> str:
         if not isinstance(path, str) or not isinstance(old_text, str) or not isinstance(new_text, str):
             return "Error: path, old_text, new_text must be strings"
-        self._staged_edits.append({"path": path, "oldText": old_text, "newText": new_text})
-        return f"Staged edit for {path} ({len(old_text)} → {len(new_text)} chars)"
+        self._edit_counter += 1
+        edit_id = f"e{self._edit_counter}"
+        self._staged_edits.append({"id": edit_id, "path": path, "oldText": old_text, "newText": new_text})
+        return edit_id
 
     def _advance_phase(self, phase: str, summary: str | None = None) -> str:
         """Transition the root RLM pipeline to a new phase.

package/src/tool/apply-edits-tool.ts

@@ -0,0 +1,126 @@
+import { readFile } from "node:fs/promises";
+import { resolve } from "node:path";
+import { createEditToolDefinition, type AgentToolResult, type ToolDefinition } from "@earendil-works/pi-coding-agent";
+import { Text } from "@earendil-works/pi-tui";
+import { Type } from "typebox";
+import type { EditToolDetails } from "@earendil-works/pi-coding-agent";
+import type { EditRegistry } from "../registry/edit-registry.ts";
+import { countOccurrences } from "../text/edits.ts";
+import { errorMessage, formatError } from "../util/errors.ts";
+
+export const ApplyEditsToolParams = Object.freeze(Type.Object({
+  ids: Type.Array(Type.String({ description: "A staged edit ID returned by stage_edit()." }), {
+    description: "Staged edit IDs to apply.",
+  }),
+}));
+
+export interface ApplyEditsFailure {
+  readonly id: string;
+  readonly error: string;
+}
+
+export interface ApplyEditsDetails {
+  readonly status: "done" | "partial" | "error";
+  readonly appliedIds: readonly string[];
+  readonly errors: readonly ApplyEditsFailure[];
+  readonly editDetails: readonly EditToolDetails[];
+}
+
+function statusFor(appliedCount: number, errorCount: number): ApplyEditsDetails["status"] {
+  if (errorCount === 0) return "done";
+  return appliedCount > 0 ? "partial" : "error";
+}
+
+function summarize(details: ApplyEditsDetails): string {
+  const head = details.errors.length > 0
+    ? `apply_edits: ${details.appliedIds.length} applied, ${details.errors.length} failed`
+    : `apply_edits: ${details.appliedIds.length} applied`;
+  if (details.errors.length === 0) return `${head}.`;
+  const rows = new Array<string>(details.errors.length);
+  for (let i = 0; i < details.errors.length; i++) {
+    const error = details.errors[i];
+    rows[i] = `${error.id}: ${error.error}`;
+  }
+  return `${head}.\n${rows.join("\n")}`;
+}
+
+export function createApplyEditsTool(editRegistry: EditRegistry): ToolDefinition<typeof ApplyEditsToolParams, ApplyEditsDetails> {
+  return {
+    name: "apply_edits",
+    label: "Apply Edits",
+    description: "Apply staged REPL edits by ID without re-typing file paths or edit bodies.",
+    parameters: ApplyEditsToolParams,
+
+    async execute(toolCallId, params, signal, _onUpdate, ctx): Promise<AgentToolResult<ApplyEditsDetails>> {
+      const appliedIds = new Array<string>(params.ids.length);
+      const errors = new Array<ApplyEditsFailure>(params.ids.length);
+      const editDetails = new Array<EditToolDetails>(params.ids.length);
+      let appliedCount = 0;
+      let errorCount = 0;
+      let detailCount = 0;
+
+      const editTool = createEditToolDefinition(ctx.cwd);
+      for (let i = 0; i < params.ids.length; i++) {
+        const id = params.ids[i];
+        const edit = editRegistry.get(id);
+        if (edit === undefined) {
+          errors[errorCount] = { id, error: formatError("unknown edit id") };
+          errorCount++;
+          continue;
+        }
+
+        try {
+          const fullPath = resolve(ctx.cwd, edit.path);
+          const content = await readFile(fullPath, "utf8");
+          const occurrences = countOccurrences(content, edit.oldText);
+          if (occurrences !== 1) {
+            errors[errorCount] = { id, error: formatError(`anchor occurs ${occurrences} times in ${edit.path}`) };
+            errorCount++;
+            continue;
+          }
+
+          const result = await editTool.execute(
+            toolCallId,
+            { path: edit.path, edits: [{ oldText: edit.oldText, newText: edit.newText }] },
+            signal,
+            undefined,
+            ctx,
+          );
+          if (result.details !== undefined) {
+            editDetails[detailCount] = result.details;
+            detailCount++;
+          }
+          editRegistry.delete(id);
+          appliedIds[appliedCount] = id;
+          appliedCount++;
+        } catch (error) {
+          errors[errorCount] = { id, error: formatError(errorMessage(error)) };
+          errorCount++;
+        }
+      }
+
+      const details: ApplyEditsDetails = {
+        status: statusFor(appliedCount, errorCount),
+        appliedIds: appliedIds.slice(0, appliedCount),
+        errors: errors.slice(0, errorCount),
+        editDetails: editDetails.slice(0, detailCount),
+      };
+      return { content: [{ type: "text", text: summarize(details) }], details };
+    },
+
+    renderCall(args, theme) {
+      return new Text(
+        theme.fg("toolTitle", theme.bold("apply_edits ")) + theme.fg("dim", args.ids.join(", ")),
+        0,
+        0,
+      );
+    },
+
+    renderResult(result, _options, theme) {
+      const details = result.details;
+      if (details === undefined) return new Text("(no apply_edits details)", 0, 0);
+      const summary = summarize(details);
+      return new Text(theme.fg(details.status === "error" ? "error" : "success", summary), 0, 0);
+    },
+  };
+}

package/src/tool/repl-details.ts

@@ -21,6 +21,8 @@ export interface ReplDetails {
   readonly subcalls: readonly RlmSubcall[];
   /** Running totals for this repl() call (cost + tokens from sub-LLM calls). */
   readonly totals: { readonly costUsd: number; readonly tokens: number };
-  /** File edits staged inside the REPL for native relay through edit(). */
+  /** Final answer submitted through answer["ready"] without echoing it to the model. */
+  readonly finalAnswer?: string;
+  /** File edits staged inside the REPL for native relay through apply_edits(). */
   readonly edits?: readonly ProposedEdit[];
 }

package/src/tool/repl-tool.ts

@@ -34,6 +34,7 @@ import type { ReplDetails } from "./repl-details.ts";
 import type { RlmSubcall } from "./rlm-details.ts";
 import { createEngine } from "../core/engine.ts";
 import { formatCost, formatTokens, spinnerFrame } from "../ui/theme.ts";
+import type { EditRegistry } from "../registry/edit-registry.ts";
 import { errorMessage, formatError, isErrorText } from "../util/errors.ts";
 import {
   headlineStatusGlyph,
@@ -59,29 +60,49 @@ export interface ReplResultText {
   readonly surfacedEdits: readonly ProposedEdit[] | undefined;
 }
 
+function countLines(text: string): number {
+  if (text.length === 0) return 0;
+  let count = 1;
+  for (const ch of text) if (ch === "\n") count++;
+  return count;
+}
+
+function stagedEditSummary(edits: readonly ProposedEdit[]): string {
+  const rows = new Array<string>(edits.length);
+  for (let i = 0; i < edits.length; i++) {
+    const edit = edits[i];
+    rows[i] = `  ${edit.id}  ${edit.path}  (-${countLines(edit.oldText)}/+${countLines(edit.newText)} lines)`;
+  }
+  return [
+    "STAGED_EDITS (apply by id with apply_edits; do NOT re-type content):",
+    ...rows,
+  ].join("\n");
+}
+
 /**
  * Assemble the model-visible text for a repl() result: cap stdout, append a zero-subcall
- * delegation nudge (suppressed when edits were staged), and append the STAGED_EDITS block
- * AFTER capping so edit JSON is never truncated. Extracted as a pure function so the
- * capping/ordering invariants are testable independently of the sandbox.
+ * delegation nudge (suppressed when edits were staged), and summarize staged edits by ID
+ * without exposing oldText/newText bodies to the root model.
  */
 export function buildReplResultText(
   stdout: string,
-  answerContent: string | undefined,
+  finalAnswer: string | undefined,
   edits: readonly ProposedEdit[],
   raised: boolean,
   subcalls: readonly RlmSubcall[],
 ): ReplResultText {
-  const rawText = stdout || answerContent || "(no output)";
+  const answerSubmitted = finalAnswer !== undefined;
+  const rawText = answerSubmitted
+    ? `ANSWER_SUBMITTED (${finalAnswer.length} chars) — delivered to user. Do not restate it.`
+    : stdout || "(no output)";
   const surfacedEdits = surfaceReplEdits(edits, raised);
-  const editsBlock = surfacedEdits
-    ? `\n\nSTAGED_EDITS:\n${JSON.stringify(surfacedEdits)}`
-    : "";
-  // Model-visible text is capped; the caller keeps full stdout in `details.output` for the TUI.
-  const cappedText = capReplResultText(rawText) ?? rawText;
+  const editsBlock = surfacedEdits ? `\n\n${stagedEditSummary(surfacedEdits)}` : "";
+  const modelText = rawText + editsBlock;
+  // Model-visible text is capped; the caller keeps full stdout/final answer in `details` for the TUI.
+  const cappedText = capReplResultText(modelText) ?? modelText;
   const delegated = subcalls.some((s) => s.kind === "llm" || s.kind === "batch" || s.kind === "rlm");
-  const nudge = surfacedEdits ? undefined : replDelegationNudge(rawText.length, delegated);
-  return { text: cappedText + (nudge ?? "") + editsBlock, surfacedEdits };
+  const nudge = surfacedEdits || answerSubmitted ? undefined : replDelegationNudge(rawText.length, delegated);
+  return { text: cappedText + (nudge ?? ""), surfacedEdits };
 }
 
 // ── Mutable bridge state (handler indirection) ──
@@ -315,6 +336,7 @@ export interface ReplToolDeps {
   readonly getModel?: () => Model<Api> | undefined;
   readonly getWorkerModel?: () => Model<Api> | undefined;
   readonly registry: ModelRegistry;
+  readonly editRegistry?: EditRegistry;
   readonly config: RlmConfig;
   readonly signal?: AbortSignal;
   readonly onUsage?: (usage: Usage, role: "sub") => void;
@@ -322,7 +344,7 @@ export interface ReplToolDeps {
 }
 
 export function createReplTool(deps: ReplToolDeps): ToolDefinition<typeof ReplToolParams, ReplDetails> {
-  const { sandboxManager, workerModel, registry, config, signal, onUsage } = deps;
+  const { sandboxManager, workerModel, registry, editRegistry, config, signal, onUsage } = deps;
   const bridgeState = new NativeBridgeState();
 
   // Build handlers once — llm/rlm use mutable refs, interactive is session-stable
@@ -464,13 +486,15 @@ export function createReplTool(deps: ReplToolDeps): ToolDefinition<typeof ReplTo
 
         if (queuedId) emitter.emitSubcallUpdated({ id: queuedId, status: "done" });
 
+        const finalAnswer = result.finalAnswer ?? undefined;
         const { text: resultText, surfacedEdits } = buildReplResultText(
           result.stdout,
-          result.answerContent,
+          finalAnswer,
           result.edits,
           result.raised,
           store.getSubcalls(),
         );
+        editRegistry?.registerAll(surfacedEdits);
 
         const details: ReplDetails = {
           status: "done",
@@ -479,10 +503,14 @@ export function createReplTool(deps: ReplToolDeps): ToolDefinition<typeof ReplTo
           executionTimeMs: elapsed,
           subcalls: store.getSubcalls(),
           totals: store.getTotals(),
+          finalAnswer,
           edits: surfacedEdits,
         };
+        const progressText = finalAnswer !== undefined
+          ? `ANSWER_SUBMITTED (${finalAnswer.length} chars)`
+          : result.stdout.slice(0, 500) || "(no output)";
         // Final progressive update
-        onUpdate?.({ content: [{ type: "text", text: result.stdout.slice(0, 500) || "(no output)" }], details });
+        onUpdate?.({ content: [{ type: "text", text: progressText }], details });
         return { content: [{ type: "text", text: resultText }], details };
       } catch (e) {
         progressStatus = "error";