@hicaru/pi-rlm 0.1.0 → 0.1.2

package/src/core/engine.ts

@@ -21,7 +21,7 @@ import { PythonSandbox } from "../sandbox/sandbox.ts";
 import { advancePhase as validatePhaseTransition, phaseGatePrompt, type PhaseState } from "./pipeline.ts";
 import { previewStdout, previewText } from "../text/preview.ts";
 import { contextLength, contextSizeStats, contextTypeLabel } from "../text/tokens.ts";
-import { collectDiffs, collectEdits, finalAnswerOf, formatReplOutputs, latestAnswerContentOf, turnHadError } from "./answer.ts";
+import { collectEdits, finalAnswerOf, formatReplOutputs, latestAnswerContentOf, turnHadError } from "./answer.ts";
 import { compactHistory, shouldCompact } from "./compaction.ts";
 import { appendUserMessage } from "./history.ts";
 import { runTurn } from "./iteration.ts";
@@ -32,7 +32,7 @@ import { appendRow, appendTodoRow, generateRunId, pruneRuns, snapshotPath, write
 import { STATE_SCHEMA_VERSION } from "../state/rows.ts";
 import type { PhaseRow, RunHeader } from "../state/rows.ts";
 import { serializeForSandbox, type ContextBundle } from "../context/repomix-context.ts";
-import type { ProposedDiffEdit, ProposedEdit } from "../sandbox/protocol.ts";
+import type { ProposedEdit } from "../sandbox/protocol.ts";
 import { formatError } from "../util/errors.ts";
 
 
@@ -57,7 +57,7 @@ export function createEngine(deps: EngineDeps): RunRlm {
   const run: RunRlm = async (input: RlmInput): Promise<RlmResult> => {
     const nowIso = (): string => new Date().toISOString(); // local helper — 4 call sites below
     const persist = input.depth === 0 && deps.runState !== undefined;
-    // Compute runId early so it can tag the MLflow root span (Ops: trace correlation on resume).
+    // Compute runId early for run-state correlation on resume.
     const runId = persist
       ? (input.resume ? input.resume.header.runId : generateRunId())
       : undefined;
@@ -136,7 +136,6 @@ export function createEngine(deps: EngineDeps): RunRlm {
     let compactions = 0;
     let completedTurns = 0;
     let editsAcc: ProposedEdit[] = [];
-    let diffsAcc: ProposedDiffEdit[] = [];
     let phaseState: PhaseState | undefined;
     let nodeStatus: "done" | "error" = "done";
     let persistOn = persist;
@@ -228,7 +227,6 @@ export function createEngine(deps: EngineDeps): RunRlm {
         limits.addRaw(input.resume.usageSeed.costUsd, input.resume.usageSeed.inputTokens, input.resume.usageSeed.outputTokens);
         best = input.resume.best;
         editsAcc = [];
-        diffsAcc = [];
         compactions = input.resume.compactions;
         completedTurns = input.resume.completedTurns;
         if (input.resume.phase) {
@@ -312,12 +310,9 @@ export function createEngine(deps: EngineDeps): RunRlm {
         completedTurns = i + 1;
         const proposedEdits = collectEdits(turn.results);
         if (proposedEdits.length > 0) editsAcc = proposedEdits;
-        const proposedDiffs = collectDiffs(turn.results);
-        if (proposedDiffs.length > 0) diffsAcc = proposedDiffs;
-
         const final = finalAnswerOf(turn.results);
         if (final != null) {
-          const done = result(final, i + 1, limits, editsAcc, diffsAcc);
+          const done = result(final, i + 1, limits, editsAcc);
           await recordTerminal("completed", done);
           lastAnswer = done.answer;
           return done;
@@ -348,21 +343,21 @@ export function createEngine(deps: EngineDeps): RunRlm {
         }
       }
       if (pendingReplOutputs) appendUserMessage(history, pendingReplOutputs);
-      const finalized = result(await finalize(history, deps, limits), deps.config.maxIterations, limits, editsAcc, diffsAcc);
+      const finalized = result(await finalize(history, deps, limits), deps.config.maxIterations, limits, editsAcc);
       await recordTerminal("finalized", finalized);
       lastAnswer = finalized.answer;
       return finalized;
     } catch (err) {
       // Abort is a user action — resolve with the best partial, not an error.
       if (deps.signal?.aborted) {
-        const aborted = result(best.trim() || "(aborted)", completedTurns, limits, editsAcc, diffsAcc);
+        const aborted = result(best.trim() || "(aborted)", completedTurns, limits, editsAcc);
         await recordTerminal("aborted", aborted);
         lastAnswer = aborted.answer;
         return aborted;
       }
       if (err instanceof LimitError) {
         nodeStatus = "error";
-        const stopped = result(best.trim() || `(stopped: ${err.message})`, completedTurns, limits, editsAcc, diffsAcc);
+        const stopped = result(best.trim() || `(stopped: ${err.message})`, completedTurns, limits, editsAcc);
         await recordTerminal("stopped", stopped);
         lastAnswer = stopped.answer;
         return stopped;
@@ -388,9 +383,9 @@ export function createEngine(deps: EngineDeps): RunRlm {
   return run;
 }
 
-function result(answer: string, iterations: number, limits: LimitGuard, edits: ProposedEdit[] = [], diffs: ProposedDiffEdit[] = []): RlmResult {
+function result(answer: string, iterations: number, limits: LimitGuard, edits: ProposedEdit[] = []): RlmResult {
   const u = limits.usage();
-  return { answer, edits, diffs, iterations, costUsd: u.costUsd, inputTokens: u.inputTokens, outputTokens: u.outputTokens, durationMs: u.durationMs };
+  return { answer, edits, iterations, costUsd: u.costUsd, inputTokens: u.inputTokens, outputTokens: u.outputTokens, durationMs: u.durationMs };
 }
 
 /** Out of turns: ask the model for its best final answer (plain text). */

package/src/core/types.ts

@@ -1,7 +1,7 @@
 /** Shared configuration + runtime types for the RLM engine. */
 
 import type { ThinkingLevel } from "@earendil-works/pi-ai";
-import type { AskAnswer, AskQuestion, ProposedDiffEdit, ProposedEdit } from "../sandbox/protocol.ts";
+import type { AskAnswer, AskQuestion, ProposedEdit } from "../sandbox/protocol.ts";
 import type { ReconstructResult } from "../state/resume.ts";
 
 export interface Sampling {
@@ -12,15 +12,6 @@ export interface Sampling {
 
 type MutableSampling = { -readonly [Key in keyof Sampling]?: Sampling[Key] };
 
-export interface TelemetryConfig {
-  /** Default: enabled iff a tracking URI resolves from config or MLFLOW_TRACKING_URI. */
-  readonly enabled?: boolean;
-  readonly trackingUri?: string;
-  readonly experimentId?: string;
-  /** Bearer token is env-only via MLFLOW_TRACKING_TOKEN; never persisted in rlm.json. */
-  readonly maxQueueSize?: number;
-}
-
 export interface RunLogConfig {
   /** Default: true — always-on, opt-out. */
   readonly enabled?: boolean;
@@ -81,8 +72,6 @@ export interface RlmConfig {
   subSystemPrompt?: string;
   /** Sampling for sub-LLM (worker) calls. */
   subSampling: MutableSampling;
-  /** Optional MLflow telemetry export configuration. Omitted by default. */
-  readonly telemetry?: TelemetryConfig;
   /** Optional run-state persistence configuration. Enabled by default. */
   readonly runLog?: RunLogConfig;
 }
@@ -112,7 +101,6 @@ export interface RlmResult {
   readonly answer: string;
   /** Legacy anchor edits retained for compatibility while older run-state rows exist. */
   readonly edits?: readonly ProposedEdit[];
-  readonly diffs?: readonly ProposedDiffEdit[];
   readonly iterations: number;
   readonly costUsd: number;
   readonly inputTokens: number;

package/src/index.ts

@@ -6,7 +6,6 @@ import { Markdown } from "@earendil-works/pi-tui";
 import { registerRlmCommand } from "./commands/rlm.ts";
 import { registerRlmConfigCommand } from "./commands/rlm-config.ts";
 import { createRlmTool } from "./tool/rlm-tool.ts";
-import { createApplyDiffTool } from "./tool/apply-diff-tool.ts";
 import { createReplTool } from "./tool/repl-tool.ts";
 import { loadSettings, mergeConfig, resolveModelId } from "./config/settings.ts";
 import { RlmController, cheapestModel } from "./mode/rlm-mode.ts";
@@ -17,27 +16,45 @@ import { packRepository, formatForLLM, serializeForSandbox } from "./context/rep
 import { buildNativeSystemPrompt } from "./prompts/system.ts";
 import { errorMessage } from "./util/errors.ts";
 
-const BLOCKED_NATIVE_TOOLS = Object.freeze(new Set(["read", "grep", "bash", "write", "edit"]));
+const BLOCKED_NATIVE_TOOLS = Object.freeze(new Set(["read", "grep"]));
 
 export default function rlmExtension(pi: ExtensionAPI): void {
-  void setupRlmExtension(pi).catch((error) => {
-    console.warn(`[rlm] extension setup failed: ${errorMessage(error)}`);
-  });
-}
-
-async function setupRlmExtension(pi: ExtensionAPI): Promise<void> {
-  const persisted = await loadSettings();
-  const config = mergeConfig(persisted.config);
+  // Init synchronously with defaults — ensures commands/tools/handlers register before session_start
+  const config = mergeConfig({});
   const controller = new RlmController(config);
-  controller.savedWorkerRef = persisted.worker;
-
-  // ── SandboxManager — persistent singleton for native-mode repl() ──
   const sandboxManager = new SandboxManager({
     execTimeoutS: config.execTimeoutS,
     requestTimeoutMs: config.requestTimeoutMs,
     python: config.python,
     sandboxInitTimeoutMs: config.sandboxInitTimeoutMs,
   });
+  let packedContextText: string | undefined;
+  let contextPackPromise: Promise<string | undefined> | undefined;
+  const ensureRepositoryContext = async (cwd: string): Promise<string | undefined> => {
+    if (packedContextText !== undefined && sandboxManager.contextPayload !== null) return packedContextText;
+    contextPackPromise ??= packRepository(cwd)
+      .then((result) => {
+        if (!result.ok) {
+          console.warn(`[rlm] repository context pack failed: ${result.error}`);
+          return undefined;
+        }
+        sandboxManager.contextPayload = serializeForSandbox(result.value);
+        packedContextText = formatForLLM(result.value);
+        return packedContextText;
+      })
+      .finally(() => { contextPackPromise = undefined; });
+    return contextPackPromise;
+  };
+
+  // Load persisted settings async — applied before session_start handler reads controller state
+  const settingsReady = loadSettings()
+    .then((persisted) => {
+      controller.config = mergeConfig(persisted.config);
+      controller.savedWorkerRef = persisted.worker;
+    })
+    .catch((err) => {
+      console.warn(`[rlm] settings load failed: ${errorMessage(err)}`);
+    });
 
   // ── Message renderers ──
   pi.registerMessageRenderer(
@@ -56,23 +73,19 @@ async function setupRlmExtension(pi: ExtensionAPI): Promise<void> {
   registerRlmConfigCommand(pi, controller);
 
   // ── Tool registration ──
-  // Existing rlm tool (stays for backward compat with /rlm mode)
   pi.registerTool(createRlmTool(controller));
-  pi.registerTool(createApplyDiffTool());
-
-  // Native repl tool — re-registered each session to pick up model provider changes
   let guidePosted = false;
 
   pi.on("session_start", async (_event, ctx) => {
-    // Restore saved worker ref — retry lazily if registry/auth is not ready yet.
-    // session_start can fire before the provider/auth registry is fully loaded, so a miss
-    // here does NOT mean the ref is stale: RlmController.resolveModels() retries later.
+    // Wait for persisted settings before reading controller state
+    await settingsReady;
+
     if (controller.savedWorkerRef) {
       const resolved = resolveModelId(ctx.modelRegistry, controller.savedWorkerRef);
       if (resolved) controller.workerModel = resolved;
     }
 
-    // Register repl tool with current models (re-registers each session for provider changes)
+    // Re-register repl tool each session to pick up model provider changes
     const workerModel = controller.workerModel ?? cheapestModel(ctx.modelRegistry) ?? ctx.model;
     const model = ctx.model;
     if (workerModel && model) {
@@ -84,7 +97,11 @@ async function setupRlmExtension(pi: ExtensionAPI): Promise<void> {
           getModel: () => controller.resolveModels(ctx)?.model,
           getWorkerModel: () => controller.resolveModels(ctx)?.worker,
           registry: ctx.modelRegistry,
-          config,
+          config: controller.config,
+          ensureContext: async () => {
+            const contextText = await ensureRepositoryContext(ctx.cwd ?? process.cwd());
+            if (contextText === undefined) throw new Error("repository context could not be loaded into RLM sandbox");
+          },
         }));
       } catch { /* re-registration on provider change — ignore if already registered */ }
     }
@@ -111,14 +128,13 @@ async function setupRlmExtension(pi: ExtensionAPI): Promise<void> {
 
     // Inject repository context as a compact listing (once per session, only when RLM is enabled)
     if (controller.enabled && !contextInjected) {
-      contextInjected = true;
       const cwd = ctx.cwd ?? process.cwd();
-      const result = await packRepository(cwd);
-      if (result.ok) {
-        const contextText = formatForLLM(result.value);
+      const contextText = await ensureRepositoryContext(cwd);
+      if (contextText !== undefined) {
+        contextInjected = true;
         const instruction = [
-          "ANALYZE THIS REPOSITORY using repl({code}) — read/grep/bash are DISABLED.",
-          `Total: ${result.value.totalFiles} files, ${result.value.totalChars.toLocaleString()} chars — must use repl().`,
+          "ANALYZE THIS REPOSITORY using repl({code}) — read/grep are DISABLED.",
+          "Repository contents are pre-loaded in the Python REPL `context` variable.",
           "Chunk context via Python, delegate to llm_query. If credits exhausted → report and stop.",
           "",
         ].join("\n");
@@ -128,9 +144,6 @@ async function setupRlmExtension(pi: ExtensionAPI): Promise<void> {
           timestamp: 0,
         } as (typeof filtered)[number];
 
-        // Store context for sandbox loading on first repl() call
-        sandboxManager.contextPayload = serializeForSandbox(result.value);
-
         return { messages: [contextMsg, ...filtered] };
       }
     }
@@ -139,19 +152,20 @@ async function setupRlmExtension(pi: ExtensionAPI): Promise<void> {
   });
 
   // ── Input routing: native mode — agent decides whether to use repl() or other tools ──
-  // The old black-box rlm() routing is removed; the main agent receives messages normally
-  // and chooses natively when to call repl(), read, grep, zebra-mcp, etc.
   pi.on("input", async (_event, _ctx) => {
     return { action: "continue" };
   });
 
-  // ── Tool restriction: block read/grep/bash when RLM is ON ──
+  // ── Tool restriction: block analysis tools when RLM is ON ──
+  // `edit`/`write` stay unblocked so the agent modifies files through Pi's native
+  // tool flow (visible to all plugins, +/- diff preview). Only read/grep are
+  // blocked — the repository is pre-loaded in the REPL `context` variable.
   pi.on("tool_call", async (event) => {
     if (!controller.enabled) return;
     if (BLOCKED_NATIVE_TOOLS.has(event.toolName)) {
       return {
         block: true,
-        reason: "RLM mode active. Use repl({code}) to read files and apply_diff({diff}) to modify them. All files are pre-loaded in the REPL. If sub-LLM credits are exhausted, report to the user.",
+        reason: "RLM mode active. Use repl({code}) to read files and search the repository — all files are pre-loaded in the REPL `context` variable. Use `edit`/`write` for file changes. If sub-LLM credits are exhausted, report to the user.",
       };
     }
   });
@@ -161,5 +175,8 @@ async function setupRlmExtension(pi: ExtensionAPI): Promise<void> {
     controller.abort();
     await sandboxManager.dispose();
     contextInjected = false;
+    packedContextText = undefined;
+    contextPackPromise = undefined;
+    sandboxManager.contextPayload = null;
   });
 }

package/src/prompts/system.ts

@@ -162,11 +162,13 @@ function nativeReplGlossary(): string {
     "- `rlm_query(prompt, model=None) -> str` — recursive RLM with its own REPL for complex sub-tasks needing iterative reasoning.",
     "- `rlm_query_batched(prompts, model=None) -> list[str]` — concurrent recursive RLM calls.",
     "",
-    "**Choosing between `llm_query` and `rlm_query`:** default to `llm_query` (fast/cheap) for one-shot tasks",
-    "and fan out with `llm_query_batched`; reach for `rlm_query` only when a sub-task needs its own iterative",
-    "reasoning. Avoid excessive recursive sub-calls when a batched one-shot suffices.",
+    "**Choosing between `llm_query` and `rlm_query`:** default to `llm_query`/batched; use `rlm_query` only for iterative sub-tasks.",
     "- `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) -> str`: stage a file edit computed inside the REPL.",
+    "  Read the file from `context`, compute the exact change in Python, then call",
+    "  stage_edit once per file. The repl() result will include a STAGED_EDITS JSON block.",
+    "  The main agent must then call `edit` for each entry verbatim — zero analysis needed.",
     "- `answer`: dict `{\"content\": \"\", \"ready\": False}`. To submit: `answer[\"content\"] = \"...\"; answer[\"ready\"] = True`.",
     "",
     "### Orchestrator Pattern",
@@ -198,45 +200,23 @@ function nativeReplGlossary(): string {
     "| `repl({code})` | Need to chunk/delegate `context` to sub-LLMs; need Python scripting; need REPL state across calls |",
     "| `read` / `grep` | Inspect a few specific files directly; small codebase |",
     "| `zebra-mcp` | Semantic search over the codebase |",
-    "| `apply_diff({diff})` | Apply a unified diff to any file — shows patch preview before writing |",
+    "| `edit` | Modify an existing file with exact text replacement (native Pi flow, visible to all plugins) |",
+    "| `write` | Create a new file (native Pi flow, visible to all plugins) |",
     "| `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` |",
     "",
     "### 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**: Set `answer[\"content\"]` and `answer[\"ready\"] = True`, or just write your final answer as a normal message.",
+    "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.",
     "",
     "### Task-Specific Patterns",
-    "",
-    "**Architecture analysis / \"learn this project\" / diagram**:",
-    "1. Probe → `repl()` to list all files grouped by directory",
-    "2. Chunk → split files into module batches (~10-15 files each)",
-    "3. **DELEGATE ALL** → `llm_query_batched` on EVERY module: \"Summarize each file's role, what it exports, and how it connects\". Send ALL batches.",
-    "4. Aggregate → collect all sub-LLM summaries, synthesize diagram from them.",
-    "5. If sub-LLM credits exhausted → report to user: \"Credits exhausted after N batches. Results so far: ...\"",
-    "",
-    "**Bug investigation / \"find the issue\"**:",
-    "1. `repl()` → grep context for keywords (use Python re/in operators)",
-    "2. `llm_query` on matching files: \"Is there a bug here? What could cause X?\"",
-    "",
-    "**Full code review / audit**:",
-    "1. `repl()` → chunk all files, delegate ALL to `llm_query_batched` with review criteria",
-    "2. Aggregate findings, report to user",
-    "",
-    "CRITICAL: Never read files directly. If sub-LLMs fail → report, don't fall back to read.",
-    "",
-    "### Handling Sub-LLM Failures",
-    "Sub-LLM calls can fail (credit limits, rate limits, timeouts). Handle gracefully:",
-    "",
-    "| Failure | Action |",
-    "|---------|--------|",
-    "| `llm_query_batched` all fail | Reduce batch size (try 3-5 instead of 10+). If still failing, use individual `llm_query` calls. |",
-    "| Individual `llm_query` fails | Check error message. If credit/rate-limit, wait and retry once. If still failing, read files directly with `read`/`grep`. |",
-    "| `rlm_query` fails | Fall back to `llm_query` — it's a one-shot call that uses fewer resources. |",
-    "| All sub-LLMs exhausted | Read key files directly. For small repos (<20 files), direct reading is fine. For large repos, prioritize the most important files. |",
+    "- Architecture/code review: chunk relevant files and delegate summaries or review to `llm_query_batched`.",
+    "- Bug investigation: use Python string/regex search over `context`; delegate matching files for analysis.",
+    "- If sub-LLM credits are exhausted, report partial results and stop — do not bypass REPL restrictions.",
     "",
     "Reserve your own tokens for high-level decisions: what to ask next, how to combine sub-LLM outputs, when to finalize.",
     "Delegate everything else. Do not submit a final answer before inspecting `context`.",
@@ -250,14 +230,16 @@ export function buildNativeSystemPrompt(): string {
     "║  NATIVE RLM MODE — YOU ARE AN ORCHESTRATOR, NOT A READER      ║",
     "╚══════════════════════════════════════════════════════════════════╝",
     "",
-    "ABSOLUTE RESTRICTION: Do NOT use `read`, `grep`, or `bash` to access files.",
+    "ABSOLUTE RESTRICTION: Do NOT use `read` or `grep` to access files.",
     "All file content is pre-loaded in the REPL `context` variable. Use ONLY `repl({code})`.",
-    "You may read at most 2 hub files directly (README.md, package.json) for quick orientation.",
     "If sub-LLM credits are exhausted → report the error to the user and stop.",
     "",
-    "ABSOLUTE RESTRICTION: Do NOT use `write` or `edit` to modify files directly.",
-    "All file modifications MUST go through apply_diff({diff}).",
-    "The diff MUST be a complete unified diff with --- a/<path> / +++ b/<path> header and @@ hunk markers.",
+    "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.",
     "",
     nativeReplGlossary(),
   ].join("\n");

package/src/sandbox/protocol.ts

@@ -32,10 +32,6 @@ export interface ProposedEdit {
   readonly newText: string;
 }
 
-export interface ProposedDiffEdit {
-  readonly diff: string;
-}
-
 /** A normal response to a request (keyed by the request `id`). */
 export interface WorkerResponse {
   readonly id: string;
@@ -47,7 +43,6 @@ export interface WorkerResponse {
   readonly final_answer?: string | null;
   readonly answer_content?: string;
   readonly edits?: readonly ProposedEdit[];
-  readonly diffs?: readonly ProposedDiffEdit[];
   readonly raised?: boolean;
   readonly execution_time?: number;
   // user-created variable names after this exec (filters builtins/context) — Metadata(stdout) for history orientation
@@ -183,7 +178,6 @@ export interface ReplResult {
   readonly finalAnswer: string | null;
   readonly answerContent: string;
   readonly edits: readonly ProposedEdit[];
-  readonly diffs: readonly ProposedDiffEdit[];
   readonly raised: boolean;
   readonly executionTimeMs: number;
   /** User-created variable names after this exec (builtins/context filtered out). */

package/src/sandbox/sandbox-manager.ts

@@ -35,8 +35,7 @@ export class SandboxManager {
    * given handlers. Subsequent calls return the existing sandbox immediately.
    * Deduplicates concurrent calls via initPromise.
    *
-   * The caller is responsible for calling loadContext() on the returned sandbox
-   * before first use.
+   * If contextPayload is set, it is loaded before the sandbox is returned.
    */
   async getOrCreate(handlers: Partial<SubLlmHandlers>): Promise<PythonSandbox> {
     if (this.disposed) throw new Error("SandboxManager disposed");
@@ -45,8 +44,8 @@ export class SandboxManager {
       // "context" event's async packRepository resolves after the first repl() call).
       // Load it into the live sandbox now if still pending.
       if (this.contextPayload !== null && !this.contextLoaded) {
+        await this.sandbox.loadContext(this.contextPayload);
         this.contextLoaded = true;
-        try { await this.sandbox.loadContext(this.contextPayload); } catch { /* best-effort */ }
       }
       return this.sandbox;
     }
@@ -61,15 +60,16 @@ export class SandboxManager {
       initTimeoutMs: this.config.sandboxInitTimeoutMs,
       handlers,
     }).then(async (s) => {
-      // Load context on first creation if available
+      // Load context on first creation if available.
       if (this.contextPayload !== null) {
+        await s.loadContext(this.contextPayload);
         this.contextLoaded = true;
-        try { await s.loadContext(this.contextPayload); } catch { /* best-effort */ }
       }
       this.sandbox = s;
       this.initPromise = null;
       return s;
     }).catch((err) => {
+      this.contextLoaded = false;
       this.initPromise = null;
       throw err;
     });
@@ -115,6 +115,7 @@ export class SandboxManager {
         // Best-effort dispose of the dead sandbox
         try { await this.sandbox.dispose(); } catch { /* already dead */ }
         this.sandbox = null;
+        this.contextLoaded = false;
       }
       throw err;
     } finally {
@@ -139,5 +140,6 @@ export class SandboxManager {
     this.disposed = true;
     await this.sandbox?.dispose();
     this.sandbox = null;
+    this.contextLoaded = false;
   }
 }

package/src/sandbox/sandbox.ts

@@ -161,6 +161,7 @@ export class PythonSandbox {
     try {
       path = await this.writeContextFile(payload, isJson);
       const res = await this.request({ type: "load_context", path, index, json: isJson });
+      if (!res.ok) throw new Error(res.error ?? "load_context failed");
       return res.index ?? 0;
     } finally {
       if (path) await unlink(path).catch(() => {});
@@ -183,13 +184,13 @@ export class PythonSandbox {
 
   async exec(code: string): Promise<ReplResult> {
     const res = await this.request({ type: "exec", code });
+    if (!res.ok) throw new Error(res.error ?? "exec failed");
     return {
       stdout: res.stdout ?? "",
       stderr: res.stderr ?? "",
       finalAnswer: res.final_answer ?? null,
       answerContent: res.answer_content ?? "",
       edits: res.edits ?? [],
-      diffs: res.diffs ?? [],
       raised: res.raised ?? false,
       executionTimeMs: Math.round((res.execution_time ?? 0) * 1000),
       varNames: res.var_names ?? [],
@@ -213,6 +214,7 @@ export class PythonSandbox {
   async snapshot(path: string, nonce: string): Promise<boolean> {
     try {
       const res = await this.request({ type: "snapshot", path, nonce });
+      if (!res.ok) return false;
       return res.ok;
     } catch {
       return false;
@@ -223,6 +225,7 @@ export class PythonSandbox {
   async restore(path: string, nonce: string): Promise<boolean> {
     try {
       const res = await this.request({ type: "restore", path, nonce });
+      if (!res.ok) return false;
       return res.ok;
     } catch {
       return false;

package/src/sandbox/worker.py

@@ -65,7 +65,8 @@ RESERVED = frozenset(
         "llm_query", "llm_query_batched", "rlm_query", "rlm_query_batched",
         "advance_phase",
         "ask_user_question", "todo",
-        "SHOW_EDITS", "SHOW_DIFFS", "SHOW_VARS", "answer", "context",
+        "stage_edit",
+        "SHOW_VARS", "answer", "context",
     }
 )
 
@@ -103,6 +104,7 @@ class Worker:
     def _setup(self) -> None:
         self.ns = {"__builtins__": _SAFE_BUILTINS.copy(), "__name__": "__main__"}
         self._ctx_payloads: dict[int, Any] = {}
+        self._staged_edits: list[dict[str, str]] = []
         self._restore_scaffold()
 
     def _capture_answer(self, content: Any) -> None:
@@ -118,8 +120,7 @@ class Worker:
         ns["advance_phase"] = self._advance_phase
         ns["ask_user_question"] = self._ask_user_question
         ns["todo"] = self._todo
-        ns["SHOW_EDITS"] = self._show_edits
-        ns["SHOW_DIFFS"] = self._show_diffs
+        ns["stage_edit"] = self._stage_edit
         ns["SHOW_VARS"] = self._show_vars
         if not isinstance(ns.get("answer"), _AnswerDict):
             cur = ns.get("answer")
@@ -153,12 +154,6 @@ class Worker:
         avail = {k: type(self.ns[k]).__name__ for k in self._user_var_names()}
         return f"Available variables: {avail}" if avail else "No variables created yet."
 
-    def _show_edits(self) -> str:
-        return "No edits — edit tools are not available in this run."
-
-    def _show_diffs(self) -> str:
-        return "No diffs — edit tools are not available in this run."
-
     # ---- sub-LLM bridge over stdio --------------------------------------------------------
 
     def _rpc(self, kind: str, payload: dict[str, Any]) -> dict[str, Any]:
@@ -263,6 +258,12 @@ class Worker:
             return f"Error: {r['error']}"
         return str(r.get("response", "ok"))
 
+    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)"
+
     def _advance_phase(self, phase: str, summary: str | None = None) -> str:
         """Transition the root RLM pipeline to a new phase.
 
@@ -348,6 +349,7 @@ class Worker:
                 stdout = out.getvalue()
                 stderr = err.getvalue() + f"\n{type(e).__name__}: {e}\n" + traceback.format_exc()
         final, self._final_answer = self._final_answer, None
+        edits, self._staged_edits = self._staged_edits, []
         answer = self.ns.get("answer")
         answer_content = answer.get("content", "") if isinstance(answer, dict) else ""
         return {
@@ -355,8 +357,7 @@ class Worker:
             "stderr": stderr,
             "final_answer": final,
             "answer_content": str(answer_content),
-            "edits": [],
-            "diffs": [],
+            "edits": edits,
             "raised": raised,
             "execution_time": time.perf_counter() - start,
             "var_names": self._user_var_names(),

package/src/tool/repl-details.ts

@@ -6,6 +6,7 @@
  * accumulated into the subcalls array for tree rendering.
  */
 
+import type { ProposedEdit } from "../sandbox/protocol.ts";
 import type { RlmSubcall } from "./rlm-details.ts";
 
 export interface ReplDetails {
@@ -20,4 +21,6 @@ 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(). */
+  readonly edits?: readonly ProposedEdit[];
 }