@hicaru/pi-rlm 0.3.1 → 0.3.3

package/src/bridge/handlers/task-registry.ts

@@ -0,0 +1,202 @@
+/**
+ * Single in-memory task registry for api_v5 async-by-default spawns.
+ *
+ * DRY: this is the ONLY place that assigns task_ids, parks await waiters, and
+ * resolves/rejects entries. Handlers and session UI both consume this type —
+ * never re-implement a Map of TaskEntry elsewhere.
+ */
+
+import type { AwaitResult, SpawnResult, TaskEntry } from "./types.ts";
+
+export const SPAWN_HINT =
+  "Call await_task(task_id=...) to get the result — this is NOT the answer.";
+
+export interface SpawnDeps {
+  nextId(): number;
+  register(kind: SpawnResult["kind"], n: number, taskId: string): SpawnResult;
+  resolve(taskId: string, result: string | readonly string[]): void;
+  reject(taskId: string, error: string): void;
+}
+
+export interface AwaitDeps {
+  get(taskId: string): TaskEntry | undefined;
+  wait(taskId: string, timeoutMs?: number): Promise<TaskEntry>;
+  unawaitedIds(): readonly string[];
+}
+
+export interface TaskRegistry {
+  readonly spawnDeps: SpawnDeps;
+  readonly awaitDeps: AwaitDeps;
+  entries(): ReadonlyMap<string, TaskEntry>;
+  toAwaitResult(taskId: string): AwaitResult;
+}
+
+interface Waiter {
+  resolve: (entry: TaskEntry) => void;
+  reject: (err: Error) => void;
+  timer?: ReturnType<typeof setTimeout>;
+}
+
+function notify(waiters: Map<string, Waiter>, taskId: string, entry: TaskEntry): void {
+  const w = waiters.get(taskId);
+  if (w === undefined) return;
+  if (w.timer !== undefined) clearTimeout(w.timer);
+  w.resolve(entry);
+  waiters.delete(taskId);
+}
+
+export function createTaskRegistry(): TaskRegistry {
+  const tasks = new Map<string, TaskEntry>();
+  const waiters = new Map<string, Waiter>();
+  let counter = 0;
+
+  const spawnDeps: SpawnDeps = {
+    nextId: () => {
+      counter += 1;
+      return counter;
+    },
+    register(kind, n, taskId) {
+      const entry: TaskEntry = {
+        taskId,
+        kind,
+        n,
+        status: "pending",
+        createdAt: Date.now(),
+      };
+      tasks.set(taskId, entry);
+      return Object.freeze({
+        ok: true,
+        task_id: taskId,
+        kind,
+        n,
+        status: "pending" as const,
+        hint: SPAWN_HINT,
+      });
+    },
+    resolve(taskId, result) {
+      const entry = tasks.get(taskId);
+      if (entry === undefined) return;
+      entry.status = "done";
+      if (typeof result === "string") {
+        entry.result = result;
+      } else {
+        entry.results = Object.freeze([...result]);
+      }
+      notify(waiters, taskId, entry);
+    },
+    reject(taskId, error) {
+      const entry = tasks.get(taskId);
+      if (entry === undefined) return;
+      entry.status = "error";
+      entry.error = error;
+      const w = waiters.get(taskId);
+      if (w !== undefined) {
+        if (w.timer !== undefined) clearTimeout(w.timer);
+        w.reject(new Error(error));
+        waiters.delete(taskId);
+      }
+    },
+  };
+
+  const awaitDeps: AwaitDeps = {
+    get: (taskId) => tasks.get(taskId),
+    wait(taskId, timeoutMs) {
+      return new Promise<TaskEntry>((resolve, reject) => {
+        const entry = tasks.get(taskId);
+        if (entry !== undefined && entry.status !== "pending") {
+          resolve(entry);
+          return;
+        }
+        const timer =
+          timeoutMs !== undefined
+            ? setTimeout(() => {
+                waiters.delete(taskId);
+                const e = tasks.get(taskId);
+                if (e !== undefined && e.status === "pending") {
+                  e.status = "timeout";
+                  e.error = `Timeout after ${timeoutMs}ms`;
+                }
+                reject(new Error(`Timeout waiting for task ${taskId}`));
+              }, timeoutMs)
+            : undefined;
+        waiters.set(taskId, { resolve, reject, timer });
+      });
+    },
+    unawaitedIds: () => {
+      const ids: string[] = [];
+      for (const [id, e] of tasks) {
+        if (e.status === "pending") ids.push(id);
+      }
+      return ids;
+    },
+  };
+
+  return {
+    spawnDeps,
+    awaitDeps,
+    entries: () => tasks,
+    toAwaitResult(taskId) {
+      const entry = tasks.get(taskId);
+      if (entry === undefined) {
+        return {
+          ok: false,
+          task_id: taskId,
+          kind: "unknown",
+          status: "error",
+          error: `Task ${taskId} not found`,
+        };
+      }
+      const status =
+        entry.status === "pending" ? "error" : entry.status;
+      return {
+        ok: entry.status === "done",
+        task_id: entry.taskId,
+        kind: entry.kind,
+        status,
+        result: entry.result,
+        results: entry.results,
+        error:
+          entry.error ??
+          (entry.status === "pending" ? "Task still pending" : undefined),
+      };
+    },
+  };
+}
+
+/**
+ * Register a spawn and start background work once.
+ * Returns SpawnResult immediately; never uses non-null assertions on task_id.
+ */
+export function spawnAndRun(
+  sd: SpawnDeps,
+  kind: SpawnResult["kind"],
+  n: number,
+  work: () => Promise<string | readonly string[]>,
+  trackDetached: (<T>(run: () => Promise<T>) => Promise<T>) | undefined,
+  detached: boolean,
+): SpawnResult {
+  const id = sd.nextId();
+  const taskId = `task_${id}`;
+  const spawned = sd.register(kind, n, taskId);
+
+  const run = async (): Promise<void> => {
+    try {
+      const result = await work();
+      sd.resolve(taskId, result);
+    } catch (err: unknown) {
+      const message = err instanceof Error ? err.message : String(err);
+      sd.reject(taskId, message.startsWith("Error:") ? message : `Error: ${message}`);
+    }
+  };
+
+  if (trackDetached !== undefined && detached) {
+    void trackDetached(run).catch((err: unknown) => {
+      const message = err instanceof Error ? err.message : String(err);
+      sd.reject(taskId, message.startsWith("Error:") ? message : `Error: ${message}`);
+    });
+  } else {
+    void run();
+  }
+
+  return spawned;
+}

package/src/bridge/handlers/types.ts

@@ -0,0 +1,136 @@
+/**
+ * Shared types for the async-by-default subcall handler set.
+ *
+ * Every subcall (llm_query, llm_batch, rlm_query, rlm_batch) returns a SpawnResult
+ * immediately with a task_id. The model must call await(task_id) to collect the
+ * real answer. This contract is proven in rlm_test (api_v5 + batch, scores 0.89–1.0).
+ */
+
+import type { Api, Model, Usage } from "@earendil-works/pi-ai";
+import type { ModelRegistry } from "@earendil-works/pi-coding-agent";
+import type { RlmInput, RlmResult, Sampling } from "../../core/types.ts";
+import type { SubcallGates } from "../../util/concurrency.ts";
+import type { SubcallOpts } from "../../sandbox/interrupts.ts";
+import type { RlmEmitter } from "../../tool/rlm-events.ts";
+
+// ---------------------------------------------------------------------------
+// Spawn / Await / Finish — the three shapes the model sees
+// ---------------------------------------------------------------------------
+
+/** Returned by every spawn: llm_query, llm_batch, rlm_query, rlm_batch. */
+export interface SpawnResult {
+  readonly ok: boolean;
+  readonly task_id: string | null;
+  readonly kind: "llm" | "rlm" | "llm_batch" | "rlm_batch";
+  /** Number of sub-tasks in this batch (1 for singles). */
+  readonly n: number;
+  readonly status: "pending";
+  /** Hint text reminding the model to await. */
+  readonly hint: string;
+  readonly error?: string;
+}
+
+/** Returned by await(task_id) or await(task_ids=[...]). */
+export interface AwaitResult {
+  readonly ok: boolean;
+  readonly task_id: string;
+  readonly kind: string;
+  readonly status: "done" | "error" | "timeout";
+  /** Single result (llm_query, rlm_query). */
+  readonly result?: string;
+  /** Ordered batch results (llm_batch, rlm_batch). */
+  readonly results?: readonly string[];
+  readonly error?: string;
+}
+
+/** Returned by finish(summary) — the contract boundary. */
+export interface FinishResult {
+  readonly ok: true;
+  readonly finished: true;
+}
+
+// ---------------------------------------------------------------------------
+// Invocation — per-call context (emitter, limits, depth)
+// ---------------------------------------------------------------------------
+
+export interface InvocationLimits {
+  remainingTimeoutMs(): number | undefined;
+  addUsage(usage: Usage): void;
+  addRaw(costUsd: number, inputTokens: number, outputTokens: number): void;
+}
+
+export function limitsFromRemaining(
+  remaining?: () => { readonly timeoutMs?: number },
+): InvocationLimits {
+  return Object.freeze({
+    remainingTimeoutMs: () => remaining?.().timeoutMs,
+    addUsage: (_usage: Usage) => {},
+    addRaw: (_costUsd: number, _inputTokens: number, _outputTokens: number) => {},
+  });
+}
+
+export interface Invocation {
+  readonly emitter: RlmEmitter;
+  readonly parentId: string | undefined;
+  readonly depth: number;
+  readonly limits: InvocationLimits;
+}
+
+// ---------------------------------------------------------------------------
+// Handler dependencies — supplied by the engine or repl() tool
+// ---------------------------------------------------------------------------
+
+export interface SubcallConfig {
+  readonly maxPromptChars: number;
+  readonly maxDepth: number;
+  readonly subSampling?: Sampling;
+  readonly subSystemPrompt?: string;
+}
+
+export interface SubcallHandlerDeps {
+  /** Resolve the invocation for a given interrupt (emitter + limits). */
+  readonly resolve: (opts: SubcallOpts, depth: number) => Invocation | null;
+  /** Session-wide concurrency gates. */
+  readonly gates: SubcallGates;
+  readonly registry: ModelRegistry;
+  readonly getLlmModel: () => Model<Api>;
+  readonly getConfig: () => SubcallConfig;
+  readonly signal?: AbortSignal;
+  readonly onUsage?: (usage: Usage, role: "sub") => void;
+  // Recursion
+  readonly runChild?: (input: RlmInput, inv: Invocation) => Promise<RlmResult>;
+  readonly getChildContext?: () => unknown;
+  readonly getModel?: () => Model<Api>;
+  readonly degrade?: (prompt: string, depth: number) => Promise<string>;
+  readonly onChildUsage?: (costUsd: number, inputTokens: number, outputTokens: number) => void;
+  readonly trackDetached?: <T>(run: () => Promise<T>) => Promise<T>;
+}
+
+// ---------------------------------------------------------------------------
+// Handler set returned to the sandbox (canonical names only)
+// ---------------------------------------------------------------------------
+
+/** Canonical api_v5 handler set — no legacy aliases. */
+export interface SubcallHandlers {
+  readonly llmQuery: (prompt: string, depth: number, opts: SubcallOpts) => Promise<SpawnResult>;
+  readonly llmBatch: (prompts: readonly string[], depth: number, opts: SubcallOpts) => Promise<SpawnResult>;
+  readonly rlmQuery: (task: string, depth: number, opts: SubcallOpts) => Promise<SpawnResult>;
+  readonly rlmBatch: (tasks: readonly string[], depth: number, opts: SubcallOpts) => Promise<SpawnResult>;
+  readonly awaitTask: (taskId: string | undefined, taskIds: readonly string[] | undefined, timeoutS: number | undefined, depth: number, opts: SubcallOpts) => Promise<AwaitResult>;
+  readonly finishTask: (summary: string, depth: number, opts: SubcallOpts) => Promise<FinishResult>;
+}
+
+// ---------------------------------------------------------------------------
+// Task registry entry — tracks in-flight spawns
+// ---------------------------------------------------------------------------
+
+export interface TaskEntry {
+  readonly taskId: string;
+  readonly kind: "llm" | "rlm" | "llm_batch" | "rlm_batch";
+  readonly n: number;
+  status: "pending" | "done" | "error" | "timeout";
+  result?: string;
+  results?: readonly string[];
+  error?: string;
+  readonly createdAt: number;
+}

package/src/commands/rlm-config.ts

@@ -8,7 +8,7 @@ import type { RlmController } from "../mode/rlm-mode.ts";
 import { cheapestModel } from "../mode/llm-model.ts";
 import { setRlmModeStatus } from "../ui/status.ts";
 import { showConfigPanel } from "../ui/config-panel.ts";
-import { pickableModels, selectModel } from "../ui/model-picker.ts";
+import { pickableModels, selectModel, type ModelSelection } from "../ui/model-picker.ts";
 
 /** Newer Pi hosts expose session-scoped models; 0.79 peers do not — duck-type safely. */
 function sessionScopedModels(
@@ -18,6 +18,36 @@ function sessionScopedModels(
   return Array.isArray(scoped) ? scoped as readonly { readonly model: Model<Api> }[] : undefined;
 }
 
+/**
+ * Apply a model-picker result to controller pin state.
+ *
+ * - `null` → explicit "cheapest (auto)" (clear pin; leave reasoning alone)
+ * - `ModelSelection` → pin that model (and its thinking level, which may be undefined)
+ * - `undefined` → ESC / no change
+ */
+export function applyLlmSelection(
+  controller: RlmController,
+  llm: ModelSelection | null | undefined,
+): void {
+  if (llm === undefined) return;
+  if (llm === null) {
+    controller.llmModel = undefined;
+    controller.savedLlmRef = undefined;
+    controller.explicitClearPin = true;
+    return;
+  }
+  controller.llmModel = llm.model;
+  controller.savedLlmRef = modelRef(llm.model);
+  controller.explicitClearPin = false;
+  controller.setConfig(Object.freeze({
+    ...controller.config,
+    subSampling: Object.freeze({
+      ...controller.config.subSampling,
+      reasoning: llm.thinkingLevel,
+    }),
+  }));
+}
+
 export async function runRlmConfig(controller: RlmController, ctx: ExtensionContext): Promise<boolean> {
   // Match Pi's native list: refresh so a just-added key appears, then use scoped models when
   // the session narrowed them, else every available (auth-configured) model. Never getAll().
@@ -34,23 +64,21 @@ export async function runRlmConfig(controller: RlmController, ctx: ExtensionCont
     models,
     controller.llmModel,
     controller.config.subSampling.reasoning,
+    controller.savedLlmRef,
   );
+  // Only an explicit choice touches the pin. ESC leaves model + reasoning alone.
+  // Choosing cheapest must NOT wipe subSampling.reasoning (null !== undefined used to).
+  applyLlmSelection(controller, llm);
+
+  // Persist model choice immediately — if showConfigPanel throws or process exits before it
+  // returns, the pin survives (Root Cause #2, v0.3.2).
   if (llm !== undefined) {
-    controller.llmModel = llm?.model;
-    controller.setConfig(Object.freeze({
-      ...controller.config,
-      subSampling: Object.freeze({ ...controller.config.subSampling, reasoning: llm?.thinkingLevel }),
-    }));
+    const saved = await controller.persist();
+    if (!saved) ctx.ui.notify("RLM: failed to save llm setting", "error");
   }
 
   controller.setConfig(await showConfigPanel(ctx, controller.config));
 
-  // Only an explicit choice touches the persisted pin. ESC (`undefined`) used to fall through
-  // here and freeze whatever cheapest resolved to at that moment, which silently ended
-  // "cheapest (auto)" for every later session — including once a cheaper model appeared.
-  if (llm === null) controller.savedLlmRef = undefined;          // "⟳ cheapest (auto)"
-  else if (llm !== undefined) controller.savedLlmRef = modelRef(llm.model);
-
   const persisted = await controller.persist();
   if (!persisted) ctx.ui.notify("RLM: failed to save settings to ~/.pi/agent/rlm.json", "error");
   setRlmModeStatus(ctx.ui, controller, ctx.getContextUsage());

package/src/config/settings.ts

@@ -9,8 +9,9 @@ import { DEFAULT_CONFIG } from "./defaults.ts";
 
 export interface PersistedSettings {
   readonly config: Partial<RlmConfig>;
-  /** "provider/id" of the pinned sub-LLM, or undefined for "cheapest (auto)". */
-  readonly llm?: string;
+  /** "provider/id" of the pinned sub-LLM, or undefined for "cheapest (auto)".
+   *  `null` = explicit "cheapest" clear (omit key on disk). */
+  readonly llm?: string | null;
 }
 
 type MutablePartialRlmConfig = { -readonly [K in keyof RlmConfig]?: RlmConfig[K] };
@@ -133,7 +134,16 @@ export async function saveSettings(s: PersistedSettings): Promise<boolean> {
   try {
     const p = settingsPath();
     await mkdir(dirname(p), { recursive: true });
-    await writeFile(p, `${JSON.stringify(s, null, 2)}\n`);
+    const body: Record<string, unknown> = { config: s.config };
+    if (s.llm !== undefined) {
+      // Explicit: string → write pin, null → omit key (cheapest).
+      if (s.llm !== null) body.llm = s.llm;
+    } else {
+      // Merge: preserve existing disk pin so config-only saves never strip it.
+      const existing = await loadSettings();
+      if (existing.llm) body.llm = existing.llm;
+    }
+    await writeFile(p, `${JSON.stringify(body, null, 2)}\n`);
     return true;
   } catch {
     return false;

package/src/context/listing.ts

@@ -23,7 +23,7 @@ export function formatContextListing(context: unknown): string {
       'Use `add_context("/path/to/dir")` / `add_context("docs.pdf")` / `add_context("https://…")` for external sources.',
       "Documents (PDF, DOCX, XLSX, PPTX, CSV, …) are converted to Markdown on the way in.",
       "",
-      "Use repl({code}) and delegate semantic reading to llm_query / llm_query_batched / llm_query_chunked.",
+      "Use repl({code}) and delegate semantic reading to llm_query / llm_batch / llm_query_chunked.",
     ].join("\n");
   }
 
@@ -53,7 +53,7 @@ export function formatContextListing(context: unknown): string {
     truncated,
     "",
     "File contents are loaded in the REPL `context` variable — file-reading tools are disabled.",
-    "Use repl({code}) and delegate semantic reading to llm_query / llm_query_batched / llm_query_chunked.",
+    "Use repl({code}) and delegate semantic reading to llm_query / llm_batch / llm_query_chunked.",
   ].join("\n");
 }
 

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.