@hicaru/pi-rlm 0.3.1 → 0.3.2

package/src/sandbox/interrupts.ts

@@ -2,8 +2,15 @@
  * The interrupt surface: what the worker can ask the host for mid-exec, and how each request is
  * turned into a reply frame.
  *
- * Split from sandbox.ts, which owns the subprocess and the JSONL pump. Adding a sandbox function
- * touches this file and worker.py; the transport underneath does not change.
+ * Canonical api_v5 kinds only: llm_query | llm_batch | rlm_query | rlm_batch | await | finish | add_context.
+ *
+ * Wire reply shapes the Python worker reduces:
+ *   - single:  { response: string }  or { error }
+ *   - batch:   { responses: string[] } or { error }
+ *
+ * Host handlers may return either:
+ *   - plain string / string[] (tests, sync stubs)
+ *   - SpawnResult { task_id } (createSubcallHandlers) — this layer awaits to final content
  */
 
 import type { WorkerInterrupt } from "./protocol.ts";
@@ -12,47 +19,40 @@ import { errorMessage, formatError } from "../util/errors.ts";
 
 /** Result of a host-side pack requested by `add_context`. */
 export interface AddContextResult {
-  readonly payload: unknown;     // always ContextFile[] under ctx/<id>/ (or un-prefixed for cwd)
+  readonly payload: unknown;
   readonly files?: number;
   readonly chars: number;
   readonly sourceId: string;
   readonly pathPrefix: string;
-  /** Host already has this source — no pack, empty payload. */
   readonly alreadyLoaded?: boolean;
-  /** Document-type files in the payload (fresh + cache hits). */
   readonly documents?: number;
-  /** Documents freshly converted this call (cache hits excluded). */
   readonly converted?: number;
-  /** Paths skipped during packing (model-facing). */
   readonly skipped?: readonly { readonly path: string; readonly reason: string }[];
 }
 
-/**
- * Per-interrupt routing context for the sub-LLM handlers.
- *
- * Only the four sub-call kinds can be spawned, so only they carry it; add_context is
- * always synchronous within one exec.
- */
 export interface SubcallOpts {
-  /** Started via `spawn()` — route to session-scoped state, not the current invocation. */
   readonly detached: boolean;
-  /**
-   * `rlm_query(paths=[…])` — path prefixes narrowing the child's inherited context.
-   * Absent on every other path; `llm_query` never carries it.
-   */
+  /** Path prefixes for rlm_query / rlm_batch child context. */
   readonly paths?: readonly string[];
 }
 
-/** Handlers the bridge installs to service sub-LLM interrupts. Return the reply payload. */
+/** Handlers the bridge installs — canonical names only. */
 export interface SubLlmHandlers {
-  llmQuery(prompt: string, model: string | null, depth: number, opts: SubcallOpts): Promise<string>;
-  llmQueryBatched(prompts: readonly string[], model: string | null, depth: number, opts: SubcallOpts): Promise<string[]>;
-  rlmQuery(prompt: string, model: string | null, depth: number, opts: SubcallOpts): Promise<string>;
-  rlmQueryBatched(prompts: readonly string[], model: string | null, depth: number, opts: SubcallOpts): Promise<string[]>;
+  llmQuery(prompt: string, depth: number, opts: SubcallOpts): Promise<unknown>;
+  llmBatch(prompts: readonly string[], depth: number, opts: SubcallOpts): Promise<unknown>;
+  rlmQuery(task: string, depth: number, opts: SubcallOpts): Promise<unknown>;
+  rlmBatch(tasks: readonly string[], depth: number, opts: SubcallOpts): Promise<unknown>;
+  awaitTask(
+    taskId: string | undefined,
+    taskIds: readonly string[] | undefined,
+    timeoutS: number | undefined,
+    depth: number,
+    opts: SubcallOpts,
+  ): Promise<unknown>;
+  finishTask(summary: string, depth: number, opts: SubcallOpts): Promise<unknown>;
   addContext(source: string, depth: number): Promise<AddContextResult>;
 }
 
-/** Narrow an unknown JSON value to a frozen string array. Non-strings and blanks are dropped. */
 function toStringArray(value: unknown): readonly string[] | undefined {
   if (!Array.isArray(value)) return undefined;
   const out = new Array<string>(value.length);
@@ -65,16 +65,23 @@ function toStringArray(value: unknown): readonly string[] | undefined {
   return n > 0 ? Object.freeze(out) : undefined;
 }
 
-/** Default handlers — every sandbox function refuses until a bridge installs a real one. */
-export const REJECT: SubLlmHandlers = {
-  llmQuery: async () => formatError("sub-LLM bridge not configured"),
-  llmQueryBatched: async (p) => p.map(() => formatError("sub-LLM bridge not configured")),
-  rlmQuery: async () => formatError("sub-LLM bridge not configured"),
-  rlmQueryBatched: async (p) => p.map(() => formatError("sub-LLM bridge not configured")),
-  addContext: async () => { throw new Error("add_context not configured"); },
-};
+const UNCONFIGURED = formatError("sub-LLM bridge not configured");
+
+const rejectBatch = async (items: readonly string[]): Promise<readonly string[]> =>
+  Object.freeze(items.map(() => UNCONFIGURED));
+
+export const REJECT: SubLlmHandlers = Object.freeze({
+  llmQuery: async () => UNCONFIGURED,
+  llmBatch: rejectBatch,
+  rlmQuery: async () => UNCONFIGURED,
+  rlmBatch: rejectBatch,
+  awaitTask: async () => UNCONFIGURED,
+  finishTask: async () => UNCONFIGURED,
+  addContext: async () => {
+    throw new Error("add_context not configured");
+  },
+});
 
-/** Body of a reply frame — the union of every handler's payload shape. */
 export interface ReplyBody {
   response?: string;
   responses?: string[];
@@ -91,11 +98,138 @@ export interface ReplyBody {
   error?: string;
 }
 
+const RLM_PATH_TYPES = new Set(["rlm_query", "rlm_batch"]);
+
+/** Narrow unknown to SpawnResult-shaped object from createSubcallHandlers. */
+function isSpawnResult(
+  value: unknown,
+): value is {
+  readonly ok: boolean;
+  readonly task_id: string | null;
+  readonly kind: string;
+  readonly error?: string;
+} {
+  if (typeof value !== "object" || value === null) return false;
+  const o = value as Record<string, unknown>;
+  return (
+    typeof o.ok === "boolean" &&
+    (typeof o.task_id === "string" || o.task_id === null) &&
+    typeof o.kind === "string" &&
+    typeof o.status === "string"
+  );
+}
+
+/** Narrow unknown to AwaitResult-shaped object. */
+function isAwaitResult(
+  value: unknown,
+): value is {
+  readonly ok: boolean;
+  readonly result?: string;
+  readonly results?: readonly string[];
+  readonly error?: string;
+} {
+  if (typeof value !== "object" || value === null) return false;
+  const o = value as Record<string, unknown>;
+  return typeof o.ok === "boolean" && typeof o.task_id === "string" && typeof o.status === "string";
+}
+
+/**
+ * Resolve a handler return value to a single `response` string for the worker.
+ * Accepts plain string stubs OR SpawnResult (awaits to completion).
+ */
+async function resolveSingle(
+  h: SubLlmHandlers,
+  raw: unknown,
+  depth: number,
+  opts: SubcallOpts,
+): Promise<ReplyBody> {
+  if (typeof raw === "string") {
+    return { response: raw };
+  }
+  if (isSpawnResult(raw)) {
+    if (!raw.ok || raw.task_id === null) {
+      const err = raw.error ?? "spawn failed";
+      return { error: err, response: formatError(err) };
+    }
+    const collected = await h.awaitTask(raw.task_id, undefined, undefined, depth, opts);
+    if (typeof collected === "string") {
+      return { response: collected };
+    }
+    if (isAwaitResult(collected)) {
+      if (!collected.ok && collected.error !== undefined) {
+        return { error: collected.error, response: formatError(collected.error) };
+      }
+      return { response: collected.result ?? "" };
+    }
+    return { response: String(collected ?? "") };
+  }
+  // Unexpected shape — surface as text rather than crash the worker.
+  return { response: String(raw ?? "") };
+}
+
+/**
+ * Resolve a handler return value to `responses: string[]` for the worker batch reducer.
+ */
+async function resolveBatch(
+  h: SubLlmHandlers,
+  raw: unknown,
+  expectedN: number,
+  depth: number,
+  opts: SubcallOpts,
+): Promise<ReplyBody> {
+  if (Array.isArray(raw)) {
+    const responses = raw.map((x) => (typeof x === "string" ? x : String(x)));
+    return { responses };
+  }
+  if (isSpawnResult(raw)) {
+    if (!raw.ok || raw.task_id === null) {
+      const err = raw.error ?? "spawn failed";
+      const msg = formatError(err);
+      return {
+        error: err,
+        responses: Array.from({ length: Math.max(1, expectedN) }, () => msg),
+      };
+    }
+    const collected = await h.awaitTask(raw.task_id, undefined, undefined, depth, opts);
+    if (Array.isArray(collected)) {
+      return { responses: collected.map(String) };
+    }
+    if (isAwaitResult(collected)) {
+      if (collected.results !== undefined) {
+        return { responses: [...collected.results] };
+      }
+      if (!collected.ok && collected.error !== undefined) {
+        const msg = formatError(collected.error);
+        return {
+          error: collected.error,
+          responses: Array.from({ length: Math.max(1, expectedN) }, () => msg),
+        };
+      }
+      if (collected.result !== undefined) {
+        return { responses: [collected.result] };
+      }
+    }
+    return {
+      error: "malformed batch await result",
+      responses: Array.from({ length: Math.max(1, expectedN) }, () =>
+        formatError("malformed batch await result"),
+      ),
+    };
+  }
+  if (typeof raw === "string") {
+    return { responses: [raw] };
+  }
+  return {
+    error: "malformed batch handler result",
+    responses: Array.from({ length: Math.max(1, expectedN) }, () =>
+      formatError("malformed batch handler result"),
+    ),
+  };
+}
+
 /**
  * Service one interrupt and hand the reply body to `reply`.
- *
- * Errors are replied, never thrown: the caller invokes this from the stdio pump, where a
- * rejection would surface as an unhandled promise and leave the worker parked forever.
+ * Errors are replied, never thrown.
  */
 export async function serviceInterrupt(
   msg: WorkerInterrupt,
@@ -103,44 +237,93 @@ export async function serviceInterrupt(
   reply: (rid: string, body: ReplyBody) => void,
 ): Promise<void> {
   const d = msg.depth;
+  const paths =
+    "paths" in msg && RLM_PATH_TYPES.has(msg.type)
+      ? toStringArray(msg.paths)
+      : undefined;
   const opts: SubcallOpts = Object.freeze({
     detached: msg.detached === true,
-    // Only the recursive kinds carry a context slice; the value crossed JSON, so guard it.
-    paths: msg.type === "rlm_query" || msg.type === "rlm_query_batched"
-      ? toStringArray(msg.paths)
-      : undefined,
+    paths,
   });
+
   try {
-    if (msg.type === "llm_query") {
-      const response = await h.llmQuery(msg.prompt ?? "", msg.model ?? null, d, opts);
-      reply(msg.rid, { response });
-    } else if (msg.type === "rlm_query") {
-      const response = await h.rlmQuery(msg.prompt ?? "", msg.model ?? null, d, opts);
-      reply(msg.rid, { response });
-    } else if (msg.type === "llm_query_batched") {
-      const responses = await h.llmQueryBatched(msg.prompts ?? [], msg.model ?? null, d, opts);
-      reply(msg.rid, { responses });
-    } else if (msg.type === "rlm_query_batched") {
-      const responses = await h.rlmQueryBatched(msg.prompts ?? [], msg.model ?? null, d, opts);
-      reply(msg.rid, { responses });
-    } else if (msg.type === "add_context") {
-      const lib = await h.addContext(msg.source ?? "", d);
-      if (lib.alreadyLoaded) {
-        // No temp file — worker short-circuits on already_loaded.
+    switch (msg.type) {
+      case "llm_query": {
+        const raw = await h.llmQuery(msg.prompt ?? "", d, opts);
+        reply(msg.rid, await resolveSingle(h, raw, d, opts));
+        return;
+      }
+      case "rlm_query": {
+        const raw = await h.rlmQuery(msg.prompt ?? "", d, opts);
+        reply(msg.rid, await resolveSingle(h, raw, d, opts));
+        return;
+      }
+      case "llm_batch": {
+        const prompts = msg.prompts ?? [];
+        const raw = await h.llmBatch(prompts, d, opts);
+        reply(msg.rid, await resolveBatch(h, raw, prompts.length, d, opts));
+        return;
+      }
+      case "rlm_batch": {
+        const tasks = msg.tasks ?? msg.prompts ?? [];
+        const raw = await h.rlmBatch(tasks, d, opts);
+        reply(msg.rid, await resolveBatch(h, raw, tasks.length, d, opts));
+        return;
+      }
+      case "await": {
+        // Host-level await (orchestrator tools). Worker uses Task + await_task in-process.
+        const result = await h.awaitTask(
+          msg.task_id,
+          msg.task_ids,
+          msg.timeout_s,
+          d,
+          opts,
+        );
+        if (typeof result === "string") {
+          reply(msg.rid, { response: result });
+          return;
+        }
+        if (isAwaitResult(result)) {
+          if (result.results !== undefined) {
+            reply(msg.rid, { responses: [...result.results] });
+            return;
+          }
+          reply(msg.rid, {
+            response: result.result ?? "",
+            error: result.ok ? undefined : result.error,
+          });
+          return;
+        }
+        reply(msg.rid, { response: String(result ?? "") });
+        return;
+      }
+      case "finish": {
+        const result = await h.finishTask(msg.summary ?? "", d, opts);
+        // finish is not reduced by the worker as content — stringify is fine
         reply(msg.rid, {
-          already_loaded: true,
-          files: 0,
-          chars: lib.chars,
-          source_id: lib.sourceId,
-          path_prefix: lib.pathPrefix,
-          documents: lib.documents ?? 0,
-          converted: lib.converted ?? 0,
-          skipped: lib.skipped,
+          response:
+            typeof result === "string"
+              ? result
+              : JSON.stringify(result ?? { ok: true, finished: true }),
         });
-      } else {
+        return;
+      }
+      case "add_context": {
+        const lib = await h.addContext(msg.source ?? "", d);
+        if (lib.alreadyLoaded) {
+          reply(msg.rid, {
+            already_loaded: true,
+            files: 0,
+            chars: lib.chars,
+            source_id: lib.sourceId,
+            path_prefix: lib.pathPrefix,
+            documents: lib.documents ?? 0,
+            converted: lib.converted ?? 0,
+            skipped: lib.skipped,
+          });
+          return;
+        }
         const { path, json: isJson } = await writeContextTempFile(lib.payload);
-        // Worker reads then unlinks (worker._add_context). Host must not unlink here —
-        // if the worker is SIGKILLed before os.remove, the temp file leaks in tmpdir (acceptable).
         reply(msg.rid, {
           path,
           json: isJson,
@@ -152,9 +335,16 @@ export async function serviceInterrupt(
           converted: lib.converted ?? 0,
           skipped: lib.skipped,
         });
+        return;
+      }
+      default: {
+        const _exhaustive: never = msg;
+        reply((_exhaustive as WorkerInterrupt).rid, {
+          error: "unknown interrupt type",
+        });
       }
     }
-  } catch (err) {
+  } catch (err: unknown) {
     reply(msg.rid, { error: errorMessage(err) });
   }
 }

package/src/sandbox/protocol.ts

@@ -4,12 +4,20 @@
  * Newline-delimited JSON over the worker's stdin/stdout — no sockets, no HTTP.
  * Parent -> worker: requests (exec/load_context/shutdown) and llm replies.
  * Worker -> parent: request responses and mid-exec sub-LLM interrupts.
+ *
+ * Canonical api_v5 kinds only — no legacy `*_query_batched` wire names.
  */
 
 /** Requests the parent sends to the worker. */
 export type WorkerRequest =
   | { readonly id: string; readonly type: "exec"; readonly code: string }
-  | { readonly id: string; readonly type: "load_context"; readonly path: string; readonly index?: number; readonly json: boolean }
+  | {
+      readonly id: string;
+      readonly type: "load_context";
+      readonly path: string;
+      readonly index?: number;
+      readonly json: boolean;
+    }
   | { readonly id: string; readonly type: "shutdown" };
 
 /** Reply the parent sends to satisfy a sub-LLM interrupt. */
@@ -50,27 +58,27 @@ export interface WorkerResponse {
   readonly answer_content?: string;
   readonly raised?: boolean;
   readonly execution_time?: number;
-  // user-created variable names after this exec (filters builtins/context) — Metadata(stdout) for history orientation
+  // user-created variable names after this exec
   readonly var_names?: readonly string[];
   // load_context:
   readonly index?: number;
 }
 
-/** Kinds of sub-LLM interrupt the worker can raise mid-exec. */
+/** Canonical interrupt kinds (api_v5). */
 export type InterruptKind =
   | "llm_query"
-  | "llm_query_batched"
   | "rlm_query"
-  | "rlm_query_batched"
+  | "llm_batch"
+  | "rlm_batch"
+  | "await"
+  | "finish"
   | "add_context";
 
 interface InterruptBase {
   readonly rid: string;
   readonly depth: number;
   /**
-   * Started via `spawn()`: the request may outlive the `exec` that issued it, so the host
-   * must not attach it to that invocation's emitter or LimitGuard. Absent on the
-   * synchronous path.
+   * Detached work may outlive the exec that issued it.
    */
   readonly detached?: boolean;
 }
@@ -78,22 +86,29 @@ interface InterruptBase {
 interface PromptInterrupt extends InterruptBase {
   readonly type: "llm_query" | "rlm_query";
   readonly prompt?: string;
-  readonly model?: string | null;
-  /**
-   * `rlm_query` only — path prefixes narrowing the child's inherited context. Never sent for
-   * `llm_query`, whose frame stays byte-identical to before.
-   */
+  /** `rlm_query` only — path prefixes narrowing the child's inherited context. */
   readonly paths?: readonly string[];
 }
 
-interface BatchedPromptInterrupt extends InterruptBase {
-  readonly type: "llm_query_batched" | "rlm_query_batched";
+interface BatchInterrupt extends InterruptBase {
+  readonly type: "llm_batch" | "rlm_batch";
   readonly prompts?: readonly string[];
-  readonly model?: string | null;
-  /** `rlm_query_batched` only — one prefix set shared by every prompt in the batch. */
+  readonly tasks?: readonly string[];
   readonly paths?: readonly string[];
 }
 
+interface AwaitInterrupt extends InterruptBase {
+  readonly type: "await";
+  readonly task_id?: string;
+  readonly task_ids?: readonly string[];
+  readonly timeout_s?: number;
+}
+
+interface FinishInterrupt extends InterruptBase {
+  readonly type: "finish";
+  readonly summary?: string;
+}
+
 export interface AddContextInterrupt extends InterruptBase {
   readonly type: "add_context";
   readonly source?: string;
@@ -102,18 +117,24 @@ export interface AddContextInterrupt extends InterruptBase {
 /** A mid-exec sub-LLM/tool request from the worker. */
 export type WorkerInterrupt =
   | PromptInterrupt
-  | BatchedPromptInterrupt
+  | BatchInterrupt
+  | AwaitInterrupt
+  | FinishInterrupt
   | AddContextInterrupt;
 
 export type WorkerMessage = WorkerResponse | WorkerInterrupt;
 
-export const INTERRUPT_KINDS = Object.freeze(new Set<InterruptKind>([
-  "llm_query",
-  "llm_query_batched",
-  "rlm_query",
-  "rlm_query_batched",
-  "add_context",
-]));
+export const INTERRUPT_KINDS = Object.freeze(
+  new Set<InterruptKind>([
+    "llm_query",
+    "rlm_query",
+    "llm_batch",
+    "rlm_batch",
+    "await",
+    "finish",
+    "add_context",
+  ]),
+);
 
 function isRecord(value: unknown): value is Record<string, unknown> {
   return typeof value === "object" && value !== null;
@@ -124,11 +145,13 @@ function isWorkerResponse(value: unknown): value is WorkerResponse {
 }
 
 export function isInterrupt(msg: unknown): msg is WorkerInterrupt {
-  return isRecord(msg)
-    && typeof msg.type === "string"
-    && INTERRUPT_KINDS.has(msg.type as InterruptKind)
-    && typeof msg.rid === "string"
-    && typeof msg.depth === "number";
+  return (
+    isRecord(msg) &&
+    typeof msg.type === "string" &&
+    INTERRUPT_KINDS.has(msg.type as InterruptKind) &&
+    typeof msg.rid === "string" &&
+    typeof msg.depth === "number"
+  );
 }
 
 export function isWorkerMessage(msg: unknown): msg is WorkerMessage {

package/src/sandbox/py/guards.py

@@ -88,10 +88,13 @@ for _blocked in ("eval", "exec", "compile", "input", "globals", "locals"):
 
 RESERVED = frozenset(
     {
-        "llm_query", "llm_query_batched", "llm_query_chunked",
-        "rlm_query", "rlm_query_batched",
-        "spawn", "rlm_await", "rlm_await_all",
-        "map_files", "llm_map_reduce",
+        # Canonical api_v5
+        "llm_query", "llm_batch",
+        "rlm_query", "rlm_batch",
+        "await_task", "finish",
+        "spawn",
+        # Helpers (not the old *_query_batched API)
+        "llm_query_chunked", "map_files", "llm_map_reduce",
         "search", "grep_context", "outline",
         "add_context",
         "SHOW_VARS", "answer", "context",
@@ -126,7 +129,7 @@ def _stall_alarm(exec_timeout_s: float, stall_timeout_s: float):
     def _fire(signum, frame):  # noqa: ARG001
         raise _StallTimeout(
             f"sub-call stalled — no reply from the host for {stall_timeout_s:g}s "
-            "(the task may still be running; rlm_await it again in a later block)"
+            "(the task may still be running; await_task it again in a later block)"
         )
 
     old = signal.signal(signal.SIGALRM, _fire) if use else None

package/src/sandbox/py/retrieval.py

@@ -15,7 +15,7 @@ from typing import Any
 
 
 _CHUNK_HEADER_OVERHEAD = 64
-_MAX_CHUNK_BATCH = 20          # fan-out per llm_query_batched call (matches prompt guidance)
+_MAX_CHUNK_BATCH = 20          # fan-out per llm_batch call (matches prompt guidance)
 _MAX_CHUNKS = 500              # ceiling: above this, force pre-filtering in Python
 _NUDGE_CHARS = 500_000         # str/bytes vars above this trigger a one-time stdout hint
 
@@ -167,11 +167,14 @@ class _Bm25Index:
         out: list[dict[str, Any]] = [None] * len(top)  # type: ignore[list-item]
         for i, (idx, score) in enumerate(top):
             text = self.texts[idx]
+            snip = text[:_SNIPPET_CHARS]
+            # Both `snippet` and `text` so agents never KeyError mixing search vs grep shapes.
             out[i] = {
                 "path": self.paths[idx],
                 "line": self.starts[idx],
                 "score": round(score, 3),
-                "snippet": text[:_SNIPPET_CHARS],
+                "snippet": snip,
+                "text": snip,
             }
         return out
 
@@ -179,8 +182,8 @@ class _Bm25Index:
 def search(entries: list[tuple[str, str]], index: _Bm25Index, query: str, k: int = 10, path_glob: str | None = None) -> list[dict[str, Any]]:
     """Rank `context` windows against a natural-language query (BM25).
 
-    Returns [{path, line, score, snippet}] — pointers, not bodies. Follow up by slicing the
-    named files out of `context` and delegating them to llm_query / map_files.
+    Returns [{path, line, score, snippet, text}] — pointers, not bodies.
+    `text` is an alias of `snippet` (same as grep_context hits).
     """
     terms = _tokenize(str(query))
     if not terms:
@@ -201,9 +204,9 @@ def grep_context(
 ) -> dict[str, Any]:
     """Regex over `context`, capped and shaped.
 
-    Returns {"hits": [{path, line, text}], "counts": {path: n}, "total": n, "truncated": bool}.
-    `counts` is complete even when `hits` is capped, so a wide pattern reports its shape
-    instead of flooding stdout.
+    Returns {"hits": [{path, line, text, snippet}], "counts": {path: n}, "total": n, "truncated": bool}.
+    `snippet` is an alias of `text` (same as search hits) to avoid KeyError footguns.
+    `counts` is complete even when `hits` is capped.
     """
     try:
         rx = re.compile(pattern)
@@ -234,7 +237,13 @@ def grep_context(
                 continue
             lo = max(0, i - pad_before)
             hi = min(len(lines), i + pad_after + 1)
-            hits.append({"path": path, "line": i + 1, "text": "\n".join(lines[lo:hi])[:_SNIPPET_CHARS]})
+            body = "\n".join(lines[lo:hi])[:_SNIPPET_CHARS]
+            hits.append({
+                "path": path,
+                "line": i + 1,
+                "text": body,
+                "snippet": body,
+            })
     return {"hits": hits, "counts": counts, "total": total, "truncated": total > len(hits)}
 
 def outline(entries: list[tuple[str, str]], path: str) -> str:

package/src/sandbox/py/tasks.py

@@ -42,7 +42,7 @@ def _reduce_batch(n: int):
 
 
 def _reduce_chunked(sizes: list[int], drop_empty: bool = True):
-    """Concatenate several llm_query_batched replies back into one flat chunk list.
+    """Concatenate several llm_batch replies back into one flat chunk list.
 
     drop_empty=True (llm_query_chunked): filter "" replies so results never degrade
     to blank entries; the flattened list may then be shorter than the chunk count,