@hicaru/pi-rlm 0.2.1 → 0.3.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (79) hide show
  1. package/README.md +28 -47
  2. package/README.ru.md +18 -23
  3. package/README.zh-CN.md +17 -28
  4. package/package.json +22 -19
  5. package/src/bridge/add-context.ts +322 -0
  6. package/src/bridge/subcall-handlers.ts +63 -17
  7. package/src/commands/rlm-config.ts +47 -18
  8. package/src/commands/rlm.ts +3 -152
  9. package/src/config/defaults.ts +8 -18
  10. package/src/config/settings.ts +13 -34
  11. package/src/context/anydoc.ts +67 -0
  12. package/src/context/listing.ts +70 -0
  13. package/src/context/md-cache.ts +112 -0
  14. package/src/context/merge.ts +97 -0
  15. package/src/context/namespace.ts +180 -0
  16. package/src/context/resolve.ts +122 -0
  17. package/src/context/source-dir.ts +166 -0
  18. package/src/context/source-doc.ts +71 -0
  19. package/src/context/source-git.ts +51 -0
  20. package/src/context/source-text.ts +45 -0
  21. package/src/context/types.ts +88 -0
  22. package/src/context/walk.ts +250 -0
  23. package/src/core/engine.ts +61 -345
  24. package/src/core/history.ts +1 -1
  25. package/src/core/limits.ts +5 -12
  26. package/src/core/resource-limits.ts +0 -2
  27. package/src/core/types.ts +10 -38
  28. package/src/index.ts +92 -54
  29. package/src/mode/llm-model.ts +54 -0
  30. package/src/mode/rlm-mode.ts +28 -58
  31. package/src/prompts/glossary.ts +290 -0
  32. package/src/prompts/native.ts +127 -0
  33. package/src/prompts/system.ts +15 -408
  34. package/src/sandbox/context-file.ts +154 -0
  35. package/src/sandbox/interrupts.ts +160 -0
  36. package/src/sandbox/protocol.ts +20 -75
  37. package/src/sandbox/py/__pycache__/guards.cpython-314.pyc +0 -0
  38. package/src/sandbox/py/__pycache__/retrieval.cpython-314.pyc +0 -0
  39. package/src/sandbox/py/__pycache__/tasks.cpython-314.pyc +0 -0
  40. package/src/sandbox/py/guards.py +150 -0
  41. package/src/sandbox/py/retrieval.py +265 -0
  42. package/src/sandbox/py/tasks.py +129 -0
  43. package/src/sandbox/py/worker.py +856 -0
  44. package/src/sandbox/sandbox-manager.ts +24 -9
  45. package/src/sandbox/sandbox.ts +99 -193
  46. package/src/text/tokens.ts +31 -5
  47. package/src/tool/repl-details.ts +2 -2
  48. package/src/tool/repl-render.ts +58 -0
  49. package/src/tool/repl-result.ts +70 -0
  50. package/src/tool/repl-tool.ts +60 -170
  51. package/src/tool/rlm-aggregator.ts +2 -10
  52. package/src/tool/rlm-details.ts +0 -2
  53. package/src/tool/rlm-events.ts +0 -14
  54. package/src/tool/rlm-tool.ts +2 -13
  55. package/src/ui/config-panel.ts +12 -20
  56. package/src/ui/intro.ts +1 -2
  57. package/src/ui/model-picker.ts +34 -10
  58. package/src/ui/status.ts +3 -7
  59. package/src/util/concurrency.ts +9 -5
  60. package/src/bridge/fallback-todo.ts +0 -148
  61. package/src/bridge/interactive.ts +0 -65
  62. package/src/bridge/library.ts +0 -155
  63. package/src/bridge/pi-interactive.ts +0 -41
  64. package/src/context/library-context.ts +0 -266
  65. package/src/context/repomix-context.ts +0 -204
  66. package/src/core/artifacts.ts +0 -89
  67. package/src/core/critique.ts +0 -92
  68. package/src/core/gates.ts +0 -301
  69. package/src/core/pipeline-handlers.ts +0 -319
  70. package/src/core/pipeline.ts +0 -268
  71. package/src/prompts/phases.ts +0 -104
  72. package/src/sandbox/worker.py +0 -1456
  73. package/src/state/index.ts +0 -24
  74. package/src/state/internal.ts +0 -46
  75. package/src/state/paths.ts +0 -44
  76. package/src/state/reads.ts +0 -133
  77. package/src/state/resume.ts +0 -173
  78. package/src/state/rows.ts +0 -123
  79. package/src/state/writes.ts +0 -58
@@ -1,1456 +0,0 @@
1
- """RLM sandbox worker: a persistent Python REPL driven over a JSONL stdio protocol.
2
-
3
- Executes model-authored Python with secrets stripped from the environment.
4
- This is NOT a security sandbox: __import__ and open are available, so code can import networking modules
5
- (socket, urllib, subprocess) and read/write local files. Trust the root model's code.
6
-
7
- Protocol (parent -> worker): {"id","type":"exec"|"load_context"|"shutdown", ...}
8
- Protocol (worker -> parent): {"id","ok",...result} # response to a request
9
- {"type":"llm_query"|"llm_query_batched"|"rlm_query"|...
10
- "advance_phase"|"save_artifact"|"ask_user_question"|"todo","rid",...}
11
- # mid-exec helper request
12
- When sandbox code calls llm_query/rlm_query/advance_phase/save_artifact/ask_user_question/todo, the worker writes a request line
13
- and BLOCKS reading stdin until the matching {"type":"llm_reply","rid",...} arrives. The parent
14
- services the request in-process (it holds API keys).
15
-
16
- Requests and replies are decoupled: `_post` writes a request and returns its rid without
17
- waiting, and replies are parked in `_inbox` keyed by rid until something asks for them. That
18
- is what makes `spawn()` / `rlm_await()` / `rlm_await_all()` possible — many requests can be
19
- in flight at once (the parent already services interrupts concurrently), and a task may be
20
- awaited in a LATER exec than the one that started it.
21
- """
22
-
23
- from __future__ import annotations
24
-
25
- import argparse
26
- import fnmatch
27
- import heapq
28
- import io
29
- import json
30
- import math
31
- import os
32
- import pickle
33
- import re
34
- import signal
35
- import sys
36
- import time
37
- import traceback
38
- from contextlib import contextmanager
39
- from typing import Any
40
-
41
- # Capture the REAL stdio before exec() redirects sys.stdout/sys.stderr into buffers.
42
- # All protocol writes must go to the real stdout even while user code's prints are captured.
43
- _REAL_STDOUT = sys.stdout
44
- _REAL_STDIN = sys.stdin
45
- _REAL_STDERR = sys.stderr
46
-
47
-
48
- def _builtin(name: str):
49
- return __builtins__[name] if isinstance(__builtins__, dict) else getattr(__builtins__, name, None)
50
-
51
-
52
- # Restricted builtins: enough for real data work, minus the dangerous reflection escapes.
53
- _SAFE_BUILTINS = {
54
- name: _builtin(name)
55
- for name in (
56
- "abs", "all", "any", "ascii", "bin", "bool", "bytearray", "bytes", "callable",
57
- "chr", "classmethod", "complex", "dict", "dir", "divmod", "enumerate", "filter",
58
- "float", "format", "frozenset", "getattr", "hasattr", "hash", "hex", "id", "int",
59
- "isinstance", "issubclass", "iter", "len", "list", "map", "max", "min", "next",
60
- "object", "oct", "ord", "pow", "print", "property", "range", "repr", "reversed",
61
- "round", "set", "setattr", "slice", "sorted", "staticmethod", "str", "sum", "super",
62
- "tuple", "type", "vars", "zip", "delattr", "memoryview", "__import__", "__build_class__",
63
- "Exception", "BaseException", "ValueError", "TypeError", "KeyError", "IndexError",
64
- "AttributeError", "FileNotFoundError", "OSError", "IOError", "RuntimeError",
65
- "NameError", "ImportError", "StopIteration", "AssertionError", "NotImplementedError",
66
- "ArithmeticError", "ZeroDivisionError", "LookupError", "Warning", "True", "False", "None",
67
- )
68
- }
69
- # `open` is allowed for data work; eval/exec/compile/input/globals/locals are not.
70
- # When read_only=True (pipeline runs), write modes raise PermissionError via
71
- # builtins.open, io.open (pathlib), and os.open. Steering, not a security sandbox.
72
- _WRITE_MODE_CHARS = frozenset("wax+")
73
- _OS_WRITE_FLAGS = os.O_WRONLY | os.O_RDWR | os.O_CREAT | os.O_APPEND | os.O_TRUNC
74
-
75
- _REAL_IO_OPEN = io.open
76
- _REAL_OS_OPEN = os.open
77
-
78
-
79
- def _install_read_only_guards():
80
- """Route every common file-open path through the read-only check.
81
-
82
- Steering, not a sandbox: closes builtins.open, io.open (hence pathlib), and
83
- os.open. A determined model can still reach the filesystem via ctypes or a
84
- subprocess — the goal is that ACCIDENTAL writes cannot pass silently.
85
- Worker-internal I/O keeps using _REAL_IO_OPEN / _REAL_OS_OPEN.
86
- """
87
- def guarded_io_open(file, mode="r", *args, **kwargs):
88
- if _WRITE_MODE_CHARS & set(str(mode)):
89
- raise PermissionError(
90
- f"read-only RLM run: refusing to open {file!r} with mode {mode!r}. "
91
- "This pipeline produces a plan; file changes go through the host edit tool."
92
- )
93
- return _REAL_IO_OPEN(file, mode, *args, **kwargs)
94
-
95
- def guarded_os_open(path, flags, *args, **kwargs):
96
- if flags & _OS_WRITE_FLAGS:
97
- raise PermissionError(
98
- f"read-only RLM run: refusing os.open({path!r}) with write flags."
99
- )
100
- return _REAL_OS_OPEN(path, flags, *args, **kwargs)
101
-
102
- io.open = guarded_io_open
103
- os.open = guarded_os_open
104
- return guarded_io_open
105
-
106
-
107
- # _builtin()'s getattr(..., None) fallback would silently inject None for a name this
108
- # interpreter lacks, surfacing much later as "'NoneType' object is not callable" inside model
109
- # code. Fail at startup instead. Note "None" is legitimately None, and the block-list below is
110
- # deliberate — which is why this check runs BEFORE it.
111
- _MISSING = sorted(name for name, value in _SAFE_BUILTINS.items() if value is None and name != "None")
112
- if _MISSING:
113
- raise RuntimeError(f"unsupported Python interpreter: missing builtins {_MISSING}")
114
-
115
- def _blocked_builtin(name: str):
116
- """Bind a disabled builtin to a callable that explains itself.
117
-
118
- Binding these to None made `eval(...)` fail with a bare "'NoneType' object is not callable",
119
- which reads as a broken sandbox rather than a deliberate block: an audit session spent six
120
- execs on it and filed a phantom "namespace corruption" bug. Saying so at the point of failure
121
- fixes it for every model without spending native-prompt budget on a rule most runs never hit.
122
- """
123
- def blocked(*_args, **_kwargs):
124
- raise PermissionError(
125
- f"{name}() is disabled in the RLM sandbox by design — it is not missing and the "
126
- "namespace is not corrupt. Names are already bound, so reference them directly; "
127
- "inspect `context` with search() / grep_context() / outline()."
128
- )
129
- blocked.__name__ = name
130
- return blocked
131
-
132
-
133
- # Blocked on purpose (NOT missing) — see _blocked_builtin. The _MISSING check above runs first,
134
- # so a genuinely absent builtin is still a startup failure rather than a silent None.
135
- for _blocked in ("eval", "exec", "compile", "input", "globals", "locals"):
136
- _SAFE_BUILTINS[_blocked] = _blocked_builtin(_blocked)
137
-
138
- RESERVED = frozenset(
139
- {
140
- "llm_query", "llm_query_batched", "llm_query_chunked",
141
- "rlm_query", "rlm_query_batched",
142
- "spawn", "rlm_await", "rlm_await_all",
143
- "map_files", "llm_map_reduce",
144
- "search", "grep_context", "outline",
145
- "advance_phase", "save_artifact",
146
- "ask_user_question", "todo",
147
- "load_library",
148
- "SHOW_VARS", "answer", "context",
149
- }
150
- )
151
- # NOTE: `answers` and `plan` are deliberately NOT reserved. They are seeded by the scaffold but
152
- # owned by the model, so they must appear in SHOW_VARS and be captured by snapshots — losing a
153
- # memoized answer across a resume is exactly the failure the memo exists to prevent.
154
- # Only the single name `context` is the packed world. Legacy context_N names are filtered out.
155
- _CONTEXT_NAME = re.compile(r"context(_\d+)?\Z")
156
-
157
- # Sizing for llm_query_chunked: leave room for the instruction and the chunk header.
158
- _CHUNK_HEADER_OVERHEAD = 64
159
- _MAX_CHUNK_BATCH = 20 # fan-out per llm_query_batched call (matches prompt guidance)
160
- _MAX_CHUNKS = 500 # ceiling: above this, force pre-filtering in Python
161
- _NUDGE_CHARS = 500_000 # str/bytes vars above this trigger a one-time stdout hint
162
-
163
-
164
- def _chunk_text(text: str, chunk_chars: int) -> list[str]:
165
- """Split text into <=chunk_chars pieces, preferring newline boundaries."""
166
- chunks: list[str] = []
167
- n = len(text)
168
- start = 0
169
- while start < n:
170
- end = min(start + chunk_chars, n)
171
- if end < n:
172
- nl = text.rfind("\n", start, end)
173
- if nl > start:
174
- end = nl + 1
175
- chunks.append(text[start:end])
176
- start = end
177
- return chunks
178
-
179
-
180
- # ---- deterministic retrieval over `context` -----------------------------------------------
181
- #
182
- # The RLM paper's trajectories retrieve by having the root model hand-write regex over the
183
- # context (App. E.1). Frontier models do that well; small/fast models guess keywords badly and
184
- # the first decomposition attempt disproportionately decides the outcome (paper §5, Fig. 4a).
185
- # These primitives make retrieval deterministic and token-free: no sub-LLM call, no root tokens
186
- # spent on printed file bodies — the model gets ranked pointers and decides what to delegate.
187
-
188
- _INDEX_WINDOW_LINES = 40 # a window is the retrieval unit: big enough to carry meaning
189
- _INDEX_MAX_WINDOWS = 20_000 # ceiling so a huge load_library() cannot exhaust worker memory
190
- _SNIPPET_CHARS = 400
191
- _GREP_HARD_CAP = 200 # absolute ceiling on returned grep hits, whatever k asks for
192
- _BM25_K1 = 1.2
193
- _BM25_B = 0.75
194
-
195
- _TOKEN_SPLIT = re.compile(r"[^0-9A-Za-z]+") # also splits snake_case and paths
196
- _CAMEL_SPLIT = re.compile(r"(?<=[a-z0-9])(?=[A-Z])")
197
-
198
- # Definition-ish lines across the languages this plugin is likely to meet. Deliberately
199
- # lexical: an outline is an orientation aid, not a parse tree.
200
- _OUTLINE_LINE = re.compile(
201
- r"^\s*(?:"
202
- r"(?:export\s+)?(?:default\s+)?(?:async\s+)?(?:function|class|interface|type|enum|struct|impl|trait|namespace)\s+\w+"
203
- r"|(?:export\s+)?(?:const|let|var)\s+\w+\s*[:=]\s*(?:async\s*)?(?:function|\(|<)"
204
- r"|(?:pub\s+)?(?:async\s+)?fn\s+\w+"
205
- r"|def\s+\w+|class\s+\w+"
206
- r"|func\s+\w+"
207
- r"|#{1,4}\s+\S"
208
- r")"
209
- )
210
-
211
-
212
- def _tokenize(text: str) -> list[str]:
213
- """Lowercased alphanumeric runs, plus camelCase parts so `resolveModelId` matches `model id`."""
214
- out: list[str] = []
215
- for raw in _TOKEN_SPLIT.split(text):
216
- if not raw:
217
- continue
218
- lowered = raw.lower()
219
- out.append(lowered)
220
- if len(raw) > 3:
221
- parts = _CAMEL_SPLIT.split(raw)
222
- if len(parts) > 1:
223
- for part in parts:
224
- piece = part.lower()
225
- if piece and piece != lowered:
226
- out.append(piece)
227
- return out
228
-
229
-
230
- def _context_entries(context: Any) -> list[tuple[str, str]]:
231
- """(path, content) pairs for either context shape: list[dict] bundles or a raw string."""
232
- if isinstance(context, str):
233
- return [("<context>", context)]
234
- if not isinstance(context, list):
235
- return []
236
- out: list[tuple[str, str]] = []
237
- for i, item in enumerate(context):
238
- if isinstance(item, dict):
239
- content = item.get("content", "")
240
- out.append((
241
- str(item.get("path", f"<context[{i}]>")),
242
- content if isinstance(content, str) else str(content),
243
- ))
244
- elif isinstance(item, str):
245
- out.append((f"<context[{i}]>", item))
246
- return out
247
-
248
-
249
- class _Bm25Index:
250
- """Okapi BM25 over fixed-line windows of `context`. Built lazily, discarded on change."""
251
-
252
- __slots__ = ("paths", "starts", "texts", "postings", "doc_len", "avg_len", "truncated")
253
-
254
- def __init__(self, entries: list[tuple[str, str]]) -> None:
255
- self.paths: list[str] = []
256
- self.starts: list[int] = []
257
- self.texts: list[str] = []
258
- self.postings: dict[str, list[tuple[int, int]]] = {}
259
- self.doc_len: list[int] = []
260
- self.truncated = False
261
-
262
- for path, content in entries:
263
- if not content:
264
- continue
265
- lines = content.split("\n")
266
- for start in range(0, len(lines), _INDEX_WINDOW_LINES):
267
- if len(self.texts) >= _INDEX_MAX_WINDOWS:
268
- self.truncated = True
269
- break
270
- window = "\n".join(lines[start:start + _INDEX_WINDOW_LINES])
271
- idx = len(self.texts)
272
- self.paths.append(path)
273
- self.starts.append(start + 1)
274
- self.texts.append(window)
275
- terms = _tokenize(window)
276
- self.doc_len.append(len(terms))
277
- freq: dict[str, int] = {}
278
- for term in terms:
279
- freq[term] = freq.get(term, 0) + 1
280
- for term, tf in freq.items():
281
- self.postings.setdefault(term, []).append((idx, tf))
282
- if self.truncated:
283
- break
284
-
285
- total = len(self.doc_len)
286
- self.avg_len = (sum(self.doc_len) / total) if total else 1.0
287
-
288
- def query(self, terms: list[str], k: int, path_glob: str | None) -> list[dict[str, Any]]:
289
- total = len(self.texts)
290
- if total == 0:
291
- return []
292
- scores: dict[int, float] = {}
293
- for term in set(terms):
294
- posting = self.postings.get(term)
295
- if not posting:
296
- continue
297
- df = len(posting)
298
- idf = math.log(1.0 + (total - df + 0.5) / (df + 0.5))
299
- for idx, tf in posting:
300
- norm = _BM25_K1 * (1.0 - _BM25_B + _BM25_B * self.doc_len[idx] / self.avg_len)
301
- scores[idx] = scores.get(idx, 0.0) + idf * (tf * (_BM25_K1 + 1.0)) / (tf + norm)
302
- if path_glob:
303
- scores = {i: s for i, s in scores.items() if fnmatch.fnmatch(self.paths[i], path_glob)}
304
- if not scores:
305
- return []
306
- top = heapq.nlargest(k, scores.items(), key=lambda kv: kv[1])
307
- out: list[dict[str, Any]] = [None] * len(top) # type: ignore[list-item]
308
- for i, (idx, score) in enumerate(top):
309
- text = self.texts[idx]
310
- out[i] = {
311
- "path": self.paths[idx],
312
- "line": self.starts[idx],
313
- "score": round(score, 3),
314
- "snippet": text[:_SNIPPET_CHARS],
315
- }
316
- return out
317
-
318
-
319
- class _AnswerDict(dict):
320
- """`answer` dict; flipping `ready` True captures the final answer for the parent."""
321
-
322
- def __init__(self, on_ready):
323
- super().__init__()
324
- super().__setitem__("content", "")
325
- super().__setitem__("ready", False)
326
- self._on_ready = on_ready
327
-
328
- def __setitem__(self, key, value):
329
- super().__setitem__(key, value)
330
- if key == "ready" and value:
331
- self._on_ready(self.get("content", ""))
332
-
333
-
334
- def _send(obj: dict[str, Any]) -> None:
335
- _REAL_STDOUT.write(json.dumps(obj, ensure_ascii=False) + "\n")
336
- _REAL_STDOUT.flush()
337
-
338
-
339
- class _StallTimeout(Exception):
340
- """No frame from the host while a sub-call was pending."""
341
-
342
-
343
- @contextmanager
344
- def _stall_alarm(exec_timeout_s: float, stall_timeout_s: float):
345
- """Swap the per-cell alarm for a stall alarm while blocked on the parent.
346
-
347
- Sub-LLM latency is network time, not cell compute time, so it must not count against the
348
- ```repl``` block timeout — but an unbounded wait is exactly how a lost reply turns into a
349
- dead session. The yielded `rearm()` restarts the stall clock on every frame, so a healthy
350
- long-running child never trips it.
351
- """
352
- use = hasattr(signal, "SIGALRM")
353
- remaining = signal.getitimer(signal.ITIMER_REAL)[0] if (use and exec_timeout_s > 0) else 0.0
354
-
355
- def _fire(signum, frame): # noqa: ARG001
356
- raise _StallTimeout(
357
- f"sub-call stalled — no reply from the host for {stall_timeout_s:g}s "
358
- "(the task may still be running; rlm_await it again in a later block)"
359
- )
360
-
361
- old = signal.signal(signal.SIGALRM, _fire) if use else None
362
-
363
- def rearm() -> None:
364
- if use and stall_timeout_s > 0:
365
- signal.setitimer(signal.ITIMER_REAL, stall_timeout_s)
366
-
367
- rearm()
368
- try:
369
- yield rearm
370
- finally:
371
- if use:
372
- signal.setitimer(signal.ITIMER_REAL, 0)
373
- if old is not None:
374
- signal.signal(signal.SIGALRM, old)
375
- if remaining > 0: # restore the cell's remaining budget
376
- signal.setitimer(signal.ITIMER_REAL, remaining)
377
-
378
-
379
- def _surfaced_error(message: str) -> str:
380
- """The "Error: …" contract value, ALSO written to the cell's stderr.
381
-
382
- A spawn/await misuse whose only trace is the returned value reads to the model as a random
383
- string much later — which is exactly how `tasks.items()` blew up on a str.
384
- """
385
- print(f"[rlm] {message}", file=sys.stderr)
386
- return f"Error: {message}"
387
-
388
-
389
- # ---- reply reducers: raw parent replies -> the value the scaffold fn returns ----------------
390
- # One reducer per result shape, shared by the sync helpers and their spawned equivalents.
391
-
392
-
393
- def _reduce_one(replies: list[dict[str, Any]]) -> str:
394
- r = replies[0]
395
- return f"Error: {r['error']}" if r.get("error") else r.get("response", "")
396
-
397
-
398
- def _reduce_batch(n: int):
399
- """Reducer for a single *_query_batched reply of n prompts."""
400
- def reduce(replies: list[dict[str, Any]]) -> list[str]:
401
- r = replies[0]
402
- if r.get("error"):
403
- return [f"Error: {r['error']}"] * n
404
- out = r.get("responses")
405
- if not isinstance(out, list) or len(out) != n:
406
- return ["Error: malformed batched response"] * n
407
- return [s if isinstance(s, str) else f"Error: {s}" for s in out]
408
- return reduce
409
-
410
-
411
- def _reduce_chunked(sizes: list[int]):
412
- """Concatenate several llm_query_batched replies back into one flat chunk list."""
413
- per = [_reduce_batch(n) for n in sizes]
414
-
415
- def reduce(replies: list[dict[str, Any]]) -> list[str]:
416
- out: list[str] = []
417
- for red, rep in zip(per, replies):
418
- out.extend(red([rep]))
419
- return out
420
- return reduce
421
-
422
-
423
- def _reduce_map_files(sizes: list[int], spans: list[tuple[str, int]]):
424
- """Flatten the batch replies (same as chunked), then regroup them per path.
425
-
426
- A file larger than the per-prompt budget contributed several requests; its answers rejoin
427
- in order, which is what makes map_files a {path: answer} dict rather than a flat list.
428
- """
429
- flatten = _reduce_chunked(sizes)
430
-
431
- def reduce(replies: list[dict[str, Any]]) -> dict[str, str]:
432
- responses = flatten(replies)
433
- out: dict[str, str] = {}
434
- cursor = 0
435
- for path, count in spans:
436
- part = responses[cursor:cursor + count]
437
- cursor += count
438
- out[path] = part[0] if count == 1 and part else "\n\n".join(part)
439
- return out
440
- return reduce
441
-
442
-
443
- def _spawnable(name: str):
444
- """Tag a sync scaffold fn with the request kind spawn() should route it to."""
445
- def mark(fn):
446
- fn._rlm_name = name
447
- return fn
448
- return mark
449
-
450
-
451
- class Task:
452
- """Handle for parent-side work in flight, returned by spawn().
453
-
454
- Opaque to model code apart from `done` and repr. A Task may be awaited in a later
455
- ```repl``` block than the one that created it.
456
- """
457
-
458
- __slots__ = ("kind", "label", "_worker", "_rids", "_reduce", "_value", "_settled")
459
-
460
- def __init__(self, worker: "Worker", kind: str, rids, reduce, label: str = ""):
461
- self.kind = kind
462
- self.label = label
463
- self._worker = worker
464
- self._rids = tuple(rids)
465
- self._reduce = reduce
466
- self._value: Any = None
467
- self._settled = False
468
-
469
- @staticmethod
470
- def resolved(worker: "Worker", kind: str, value: Any, label: str = "") -> "Task":
471
- """A Task that never hit the wire — validation errors and empty inputs."""
472
- task = Task(worker, kind, (), lambda _replies: value, label)
473
- task._value = value
474
- task._settled = True
475
- return task
476
-
477
- @property
478
- def done(self) -> bool:
479
- """True once every reply has landed — awaiting will not block."""
480
- return self._settled or all(r in self._worker.inbox for r in self._rids)
481
-
482
- def __repr__(self) -> str:
483
- return f"<Task {self.kind} {'done' if self.done else 'running'} {self.label}>"
484
-
485
-
486
- class Worker:
487
- def __init__(
488
- self,
489
- depth: int,
490
- exec_timeout_s: float,
491
- max_prompt_chars: int,
492
- read_only: bool = False,
493
- await_timeout_s: float = 600.0,
494
- ):
495
- self.depth = depth
496
- self.exec_timeout_s = exec_timeout_s
497
- self.max_prompt_chars = max_prompt_chars
498
- self.read_only = read_only
499
- self.await_timeout_s = await_timeout_s
500
- self._rid = 0
501
- self._final_answer: str | None = None
502
- # Replies parked by rid until something awaits them. Unbounded by design: a task
503
- # the model spawns and never awaits keeps its entry for the life of the process.
504
- # Bounded in practice by session length; evicting would silently hang a later
505
- # rlm_await, which is strictly worse than the memory.
506
- self.inbox: dict[str, dict[str, Any]] = {}
507
- self._inflight: set[str] = set()
508
- # Requests (exec/snapshot/shutdown) that arrived mid-exec; main() replays them.
509
- self._deferred: list[Any] = []
510
- # True only while spawn() runs a builder — marks requests that may outlive this exec.
511
- self._detached = False
512
- self.ns: dict[str, Any] = {}
513
- self._setup()
514
-
515
- def _setup(self) -> None:
516
- builtins = _SAFE_BUILTINS.copy()
517
- if self.read_only:
518
- builtins["open"] = _install_read_only_guards()
519
- else:
520
- builtins["open"] = open
521
- self.ns = {"__builtins__": builtins, "__name__": "__main__"}
522
- self._context_payload: Any | None = None # pristine restore for the single `context` var
523
- self._nudged: set[str] = set()
524
- self._index: _Bm25Index | None = None
525
- self._index_stamp: tuple[int, int] | None = None # (id(context), len(context))
526
- self._restore_scaffold()
527
-
528
- def _capture_answer(self, content: Any) -> None:
529
- self._final_answer = str(content)
530
-
531
- def _restore_scaffold(self) -> None:
532
- # Re-inject any scaffolding the user code clobbered.
533
- ns = self.ns
534
- ns["llm_query"] = self._llm_query
535
- ns["llm_query_batched"] = self._llm_query_batched
536
- ns["llm_query_chunked"] = self._llm_query_chunked
537
- ns["rlm_query"] = self._rlm_query
538
- ns["rlm_query_batched"] = self._rlm_query_batched
539
- ns["spawn"] = self._spawn
540
- ns["rlm_await"] = self._await_task
541
- ns["rlm_await_all"] = self._await_all
542
- ns["map_files"] = self._map_files
543
- ns["llm_map_reduce"] = self._llm_map_reduce
544
- ns["search"] = self._search
545
- ns["grep_context"] = self._grep_context
546
- ns["outline"] = self._outline
547
- # env_tips memo (paper App. C.3): "If a value isn't in `answers`, it doesn't exist."
548
- # Re-created only when deleted — contents must survive every turn.
549
- if not isinstance(ns.get("answers"), dict):
550
- ns["answers"] = {}
551
- if not isinstance(ns.get("plan"), dict):
552
- ns["plan"] = {}
553
- ns["advance_phase"] = self._advance_phase
554
- ns["save_artifact"] = self._save_artifact
555
- ns["ask_user_question"] = self._ask_user_question
556
- ns["todo"] = self._todo
557
- ns["load_library"] = self._load_library
558
- ns["SHOW_VARS"] = self._show_vars
559
- if not isinstance(ns.get("answer"), _AnswerDict):
560
- cur = ns.get("answer")
561
- ans = _AnswerDict(self._capture_answer)
562
- if isinstance(cur, dict):
563
- for k, v in cur.items():
564
- dict.__setitem__(ans, k, v)
565
- if cur.get("ready") and self._final_answer is None:
566
- self._final_answer = str(cur.get("content", ""))
567
- ns["answer"] = ans
568
- # Single context variable (RLM paper: the context lives in the environment and
569
- # the model may transform it in place). Re-inject only if the model deleted the
570
- # name entirely; mutations and re-binds persist within the run.
571
- if self._context_payload is not None:
572
- ns.setdefault("context", self._context_payload)
573
- # Scrub any legacy context_N names so the model never sees multi-slot APIs.
574
- for k in list(ns.keys()):
575
- if k != "context" and _CONTEXT_NAME.match(k):
576
- del ns[k]
577
-
578
- def _user_var_names(self) -> list[str]:
579
- """User-created variable names — filters builtins, scaffold, and `context`.
580
-
581
- Shared by SHOW_VARS() and the exec result so both expose the same namespace view.
582
- This is the cheap orientation hint that goes into history instead of full stdout.
583
- """
584
- return [
585
- k for k in self.ns
586
- if not k.startswith("_")
587
- and not _CONTEXT_NAME.match(k)
588
- and k not in RESERVED
589
- ]
590
-
591
- def _show_vars(self) -> str:
592
- avail = {k: type(self.ns[k]).__name__ for k in self._user_var_names()}
593
- return f"Available variables: {avail}" if avail else "No variables created yet."
594
-
595
- # ---- sub-LLM bridge over stdio --------------------------------------------------------
596
-
597
- def _post(self, kind: str, payload: dict[str, Any]) -> str:
598
- """Write one parent request and return its rid WITHOUT waiting for the reply."""
599
- self._rid += 1
600
- rid = f"q{self._rid}"
601
- # Register only after the write succeeds — a broken pipe must not leave an
602
- # _inflight entry that nothing will ever settle.
603
- _send({"type": kind, "rid": rid, "depth": self.depth,
604
- "detached": self._detached, **payload})
605
- self._inflight.add(rid)
606
- return rid
607
-
608
- def park_reply(self, msg: Any) -> bool:
609
- """File an llm_reply against its rid. True when the frame was a reply.
610
-
611
- Public because main() needs it too: a spawned task can settle while the worker
612
- sits idle between execs, and that reply must not fall through to "unknown type".
613
- """
614
- if not isinstance(msg, dict) or msg.get("type") != "llm_reply":
615
- return False
616
- rid = msg.get("rid")
617
- if isinstance(rid, str) and rid in self._inflight:
618
- self._inflight.discard(rid)
619
- self.inbox[rid] = msg
620
- else:
621
- # Late reply to an abandoned rid (e.g. a request from a discarded sandbox).
622
- print(f"[rlm-sandbox] dropping reply for unknown rid: {rid!r}", file=_REAL_STDERR)
623
- return True
624
-
625
- def take_deferred(self) -> Any | None:
626
- """Pop a request that arrived mid-exec, for main() to replay. None when empty."""
627
- return self._deferred.pop(0) if self._deferred else None
628
-
629
- def _pump(self) -> bool:
630
- """Read one frame from the parent into the inbox. False when the pipe closed."""
631
- line = _REAL_STDIN.readline()
632
- if not line:
633
- return False
634
- try:
635
- msg = json.loads(line)
636
- except ValueError:
637
- print(f"[rlm-sandbox] skipping non-JSON parent frame: {line[:200]}", file=_REAL_STDERR)
638
- return True
639
- if self.park_reply(msg):
640
- return True
641
- # A request (exec/snapshot/shutdown) arriving mid-exec: main() replays it.
642
- self._deferred.append(msg)
643
- return True
644
-
645
- def _drain_until(self, rids) -> None:
646
- """Block until every rid in `rids` has its reply parked in the inbox.
647
-
648
- Bounded: a host that goes silent raises inside the ```repl``` block instead of hanging
649
- the session forever.
650
- """
651
- if all(r in self.inbox for r in rids):
652
- return
653
- with _stall_alarm(self.exec_timeout_s, self.await_timeout_s) as rearm:
654
- while not all(r in self.inbox for r in rids):
655
- if not self._pump():
656
- raise RuntimeError("parent closed the pipe during a sub-LLM request")
657
- rearm()
658
-
659
- def _take(self, rids) -> list[dict[str, Any]]:
660
- return [self.inbox.pop(r) for r in rids]
661
-
662
- def _rpc(self, kind: str, payload: dict[str, Any]) -> dict[str, Any]:
663
- """Post one request and block for its reply — the synchronous single-shot path."""
664
- rid = self._post(kind, payload)
665
- self._drain_until((rid,))
666
- return self._take((rid,))[0]
667
-
668
- # ---- spawn / await ---------------------------------------------------------------------
669
-
670
- def _start_prompt(self, kind: str, prompt, model) -> Task:
671
- text = str(prompt)
672
- # A sub-LLM asked nothing answers something: the confabulation then sits in `answers`
673
- # looking exactly like data. Refuse instead of spending a call on it.
674
- if not text.strip():
675
- return Task.resolved(self, kind, _surfaced_error(
676
- f"{kind}() got an empty prompt — a sub-LLM would confabulate an answer to nothing"))
677
- rid = self._post(kind, {"prompt": text, "model": model})
678
- return Task(self, kind, (rid,), _reduce_one, text[:40])
679
-
680
- def _start_prompts(self, kind: str, prompts, model) -> Task:
681
- prompts = [str(p) for p in prompts]
682
- if not prompts:
683
- return Task.resolved(self, kind, [])
684
- # Only the all-blank case: one blank prompt among twenty is the caller's business.
685
- if not any(p.strip() for p in prompts):
686
- return Task.resolved(self, kind, [
687
- _surfaced_error(f"{kind}() got only empty prompts")
688
- ] * len(prompts))
689
- rid = self._post(kind, {"prompts": prompts, "model": model})
690
- return Task(self, kind, (rid,), _reduce_batch(len(prompts)), f"×{len(prompts)}")
691
-
692
- def _start_llm_query(self, prompt, model: str | None = None) -> Task:
693
- return self._start_prompt("llm_query", prompt, model)
694
-
695
- def _start_rlm_query(self, prompt, model: str | None = None) -> Task:
696
- return self._start_prompt("rlm_query", prompt, model)
697
-
698
- def _start_llm_query_batched(self, prompts, model: str | None = None) -> Task:
699
- return self._start_prompts("llm_query_batched", prompts, model)
700
-
701
- def _start_rlm_query_batched(self, prompts, model: str | None = None) -> Task:
702
- return self._start_prompts("rlm_query_batched", prompts, model)
703
-
704
- def _start_llm_query_chunked(self, text, prompt: str, model: str | None = None) -> Task:
705
- """Split oversized text into cap-sized chunks and post EVERY batch at once.
706
-
707
- One answer per chunk, order preserved. No exceptions escape: errors come back as
708
- "Error: ..." strings per chunk (same contract as llm_query_batched). Because all
709
- batches go on the wire together, a large input costs one round-trip of latency
710
- rather than one per 20 chunks.
711
-
712
- NOTE: budget uses Python code-point length (len) while the parent-side cap check counts
713
- UTF-16 units (JS string.length); astral/emoji-heavy text may be marginally larger on the
714
- parent and get per-chunk rejected. Acceptable trade-off for typical code/log/profile text.
715
- """
716
- text, prompt = str(text), str(prompt)
717
- if not text:
718
- return Task.resolved(self, "llm_query_chunked", [])
719
- budget = self.max_prompt_chars - len(prompt) - _CHUNK_HEADER_OVERHEAD
720
- if budget < 1_000:
721
- return Task.resolved(self, "llm_query_chunked", [
722
- f"Error: prompt leaves under 1,000 chars per chunk (cap {self.max_prompt_chars:,}) — shorten the instruction"
723
- ])
724
- chunks = _chunk_text(text, budget)
725
- total = len(chunks)
726
- if total > _MAX_CHUNKS:
727
- return Task.resolved(self, "llm_query_chunked", [
728
- f"Error: {total} chunks would be needed — filter/slice the text in Python first"
729
- ])
730
- rids: list[str] = []
731
- sizes: list[int] = []
732
- for i in range(0, total, _MAX_CHUNK_BATCH):
733
- batch = [
734
- f"{prompt}\n\n[chunk {i + j + 1}/{total} of the input]\n{c}"
735
- for j, c in enumerate(chunks[i:i + _MAX_CHUNK_BATCH])
736
- ]
737
- rids.append(self._post("llm_query_batched", {"prompts": batch, "model": model}))
738
- sizes.append(len(batch))
739
- return Task(self, "llm_query_chunked", tuple(rids), _reduce_chunked(sizes), f"{total} chunks")
740
-
741
- def _builder_for(self, name: str):
742
- # llm_map_reduce is deliberately absent: its reduce step is a SECOND sub-LLM call that
743
- # depends on its own map results, so it cannot be one (rids, pure reduce) Task.
744
- return {
745
- "llm_query": self._start_llm_query,
746
- "llm_query_batched": self._start_llm_query_batched,
747
- "llm_query_chunked": self._start_llm_query_chunked,
748
- "map_files": self._start_map_files,
749
- "rlm_query": self._start_rlm_query,
750
- "rlm_query_batched": self._start_rlm_query_batched,
751
- }.get(name)
752
-
753
- def _spawn(self, fn, *args, **kwargs) -> Task:
754
- """Start a sub-call without waiting for it. `fn` is the scaffold function itself.
755
-
756
- Returns a Task for rlm_await / rlm_await_all, possibly in a later ```repl``` block.
757
- Misuse returns an already-resolved error Task rather than raising, matching the
758
- "Error: ..." contract of the synchronous helpers.
759
- """
760
- name = getattr(fn, "_rlm_name", None)
761
- builder = self._builder_for(name) if isinstance(name, str) else None
762
- if builder is None:
763
- return Task.resolved(self, "spawn", _surfaced_error(
764
- "spawn() takes llm_query, llm_query_batched, llm_query_chunked, map_files, "
765
- "rlm_query or rlm_query_batched — not llm_map_reduce, whose reduce step depends "
766
- "on its own map results and so cannot be a single Task"))
767
- # Mark every request this builder posts as detached: the parent routes them to its
768
- # session-scoped registry, since they may outlive the exec that started them.
769
- self._detached = True
770
- try:
771
- return builder(*args, **kwargs)
772
- except TypeError as e:
773
- return Task.resolved(self, "spawn", _surfaced_error(f"bad spawn arguments — {e}"))
774
- finally:
775
- self._detached = False
776
-
777
- def _await_task(self, task) -> Any:
778
- """Block until `task` has its result. Idempotent — the value is memoized."""
779
- if not isinstance(task, Task):
780
- return _surfaced_error(
781
- f"rlm_await expects a Task from spawn(), got {type(task).__name__}"
782
- )
783
- if not task._settled:
784
- self._drain_until(task._rids)
785
- task._value = task._reduce(self._take(task._rids))
786
- task._settled = True
787
- return task._value
788
-
789
- def _await_all(self, tasks) -> list:
790
- """Block until every task has its result. Order matches the input."""
791
- tasks = list(tasks)
792
- # One union drain so the tasks overlap instead of settling one after another.
793
- union: list[str] = []
794
- seen: set[str] = set()
795
- for t in tasks:
796
- if not isinstance(t, Task) or t._settled:
797
- continue
798
- for rid in t._rids:
799
- if rid not in seen:
800
- seen.add(rid)
801
- union.append(rid)
802
- if union:
803
- self._drain_until(union)
804
- return [self._await_task(t) for t in tasks]
805
-
806
- # ---- deterministic retrieval (no sub-LLM calls, no root tokens) -----------------------
807
-
808
- def _entries(self) -> list[tuple[str, str]]:
809
- return _context_entries(self.ns.get("context"))
810
-
811
- def _get_index(self) -> _Bm25Index:
812
- """Build the BM25 index on first use; rebuild when `context` was replaced or resized.
813
-
814
- Identity+length is a cheap stamp that catches the two ways context actually changes:
815
- load_library() extending the list, and the model re-binding the name. In-place edits
816
- that preserve length are not detected — documented, and rare in practice.
817
- """
818
- ctx = self.ns.get("context")
819
- stamp = (id(ctx), len(ctx) if isinstance(ctx, (list, str)) else 0)
820
- if self._index is None or self._index_stamp != stamp:
821
- self._index = _Bm25Index(self._entries())
822
- self._index_stamp = stamp
823
- return self._index
824
-
825
- def _search(self, query: str, k: int = 10, path_glob: str | None = None) -> list[dict[str, Any]]:
826
- """Rank `context` windows against a natural-language query (BM25).
827
-
828
- Returns [{path, line, score, snippet}] — pointers, not bodies. Follow up by slicing the
829
- named files out of `context` and delegating them to llm_query / map_files.
830
- """
831
- terms = _tokenize(str(query))
832
- if not terms:
833
- return []
834
- try:
835
- limit = max(1, min(int(k), 100))
836
- except (TypeError, ValueError):
837
- limit = 10
838
- return self._get_index().query(terms, limit, path_glob)
839
-
840
- def _grep_context(
841
- self,
842
- pattern: str,
843
- k: int = 50,
844
- path_glob: str | None = None,
845
- before: int = 0,
846
- after: int = 0,
847
- ) -> dict[str, Any]:
848
- """Regex over `context`, capped and shaped.
849
-
850
- Returns {"hits": [{path, line, text}], "counts": {path: n}, "total": n, "truncated": bool}.
851
- `counts` is complete even when `hits` is capped, so a wide pattern reports its shape
852
- instead of flooding stdout.
853
- """
854
- try:
855
- rx = re.compile(pattern)
856
- except re.error as e:
857
- return {"hits": [], "counts": {}, "total": 0, "truncated": False, "error": f"bad regex: {e}"}
858
- try:
859
- limit = max(1, min(int(k), _GREP_HARD_CAP))
860
- except (TypeError, ValueError):
861
- limit = 50
862
- pad_before = max(0, min(int(before or 0), 10))
863
- pad_after = max(0, min(int(after or 0), 10))
864
-
865
- hits: list[dict[str, Any]] = []
866
- counts: dict[str, int] = {}
867
- total = 0
868
- for path, content in self._entries():
869
- if path_glob and not fnmatch.fnmatch(path, path_glob):
870
- continue
871
- if not rx.search(content):
872
- continue
873
- lines = content.split("\n")
874
- for i, line in enumerate(lines):
875
- if not rx.search(line):
876
- continue
877
- total += 1
878
- counts[path] = counts.get(path, 0) + 1
879
- if len(hits) >= limit:
880
- continue
881
- lo = max(0, i - pad_before)
882
- hi = min(len(lines), i + pad_after + 1)
883
- hits.append({"path": path, "line": i + 1, "text": "\n".join(lines[lo:hi])[:_SNIPPET_CHARS]})
884
- return {"hits": hits, "counts": counts, "total": total, "truncated": total > len(hits)}
885
-
886
- def _outline(self, path: str) -> str:
887
- """Definition/heading skeleton of one context file — orient in ~200 chars, not 20K.
888
-
889
- `path` matches exactly, then by suffix, then as a glob.
890
- """
891
- target = str(path)
892
- entries = self._entries()
893
- content: str | None = None
894
- for p, c in entries:
895
- if p == target:
896
- content = c
897
- break
898
- if content is None:
899
- for p, c in entries:
900
- if p.endswith(target) or fnmatch.fnmatch(p, target):
901
- content = c
902
- target = p
903
- break
904
- if content is None:
905
- return f"Error: no context file matching {path!r} — use search() or list paths from context"
906
- out: list[str] = [f"# {target}"]
907
- for i, line in enumerate(content.split("\n")):
908
- if _OUTLINE_LINE.match(line):
909
- out.append(f"{i + 1}: {line.strip()[:160]}")
910
- if len(out) == 1:
911
- return f"# {target}\n(no definition-like lines found)"
912
- return "\n".join(out)
913
-
914
- # ---- one-line delegation (structural: orchestrating must be easier than solving) -------
915
-
916
- def _start_map_files(self, files: Any, prompt: str, model: str | None = None) -> Task:
917
- """Post every batch map_files needs, WITHOUT waiting. Contract: see _map_files.
918
-
919
- All batches go on the wire together, so a 100-file map costs one round-trip of latency
920
- rather than one per 20 files.
921
- """
922
- prompt = str(prompt)
923
- by_path: list[tuple[str, str]] = []
924
- lookup: dict[str, str] | None = None
925
- for item in files if isinstance(files, (list, tuple)) else [files]:
926
- if isinstance(item, dict):
927
- content = item.get("content", "")
928
- by_path.append((str(item.get("path", "?")), content if isinstance(content, str) else str(content)))
929
- elif isinstance(item, str):
930
- if lookup is None:
931
- lookup = {p: c for p, c in self._entries()}
932
- if item in lookup:
933
- by_path.append((item, lookup[item]))
934
- else:
935
- by_path.append((item, ""))
936
- if not by_path:
937
- return Task.resolved(self, "map_files", {})
938
-
939
- # Per-file prompt budget; anything larger is chunked and its answers concatenated.
940
- budget = self.max_prompt_chars - len(prompt) - _CHUNK_HEADER_OVERHEAD - 256
941
- if budget < 1_000:
942
- return Task.resolved(self, "map_files", {
943
- p: "Error: prompt too long to leave room for file content" for p, _ in by_path
944
- })
945
-
946
- requests: list[str] = []
947
- spans: list[tuple[str, int]] = [] # (path, number of chunks contributed)
948
- for path, content in by_path:
949
- chunks = _chunk_text(content, budget) if len(content) > budget else [content]
950
- spans.append((path, len(chunks)))
951
- for j, chunk in enumerate(chunks):
952
- header = f"[file {path}" + (f", part {j + 1}/{len(chunks)}]" if len(chunks) > 1 else "]")
953
- requests.append(f"{prompt}\n\n{header}\n{chunk}")
954
-
955
- rids: list[str] = []
956
- sizes: list[int] = []
957
- for i in range(0, len(requests), _MAX_CHUNK_BATCH):
958
- batch = requests[i:i + _MAX_CHUNK_BATCH]
959
- rids.append(self._post("llm_query_batched", {"prompts": batch, "model": model}))
960
- sizes.append(len(batch))
961
- return Task(self, "map_files", tuple(rids),
962
- _reduce_map_files(sizes, spans), f"{len(by_path)} files")
963
-
964
- @_spawnable("map_files")
965
- def _map_files(self, files: Any, prompt: str, model: str | None = None) -> dict[str, str]:
966
- """Ask `prompt` of every given file, batched, and return {path: answer}.
967
-
968
- `files` accepts context entries (dicts), paths (strings), or a mix — the whole
969
- chunk/batch/collect loop the system prompt used to spell out, as one call.
970
- Oversized files are split and their per-chunk answers joined.
971
- """
972
- return self._await_task(self._start_map_files(files, prompt, model))
973
-
974
- def _llm_map_reduce(
975
- self,
976
- items: Any,
977
- map_prompt: str,
978
- reduce_prompt: str,
979
- model: str | None = None,
980
- ) -> str:
981
- """Map `map_prompt` over `items` in one batch, then reduce the answers with one call.
982
-
983
- The paper's canonical strategy ("query an LLM per chunk ... then query an LLM with all
984
- the buffers") as a single call, so the root never hand-rolls the loop.
985
- """
986
- map_prompt, reduce_prompt = str(map_prompt), str(reduce_prompt)
987
- seq = list(items) if isinstance(items, (list, tuple)) else [items]
988
- if not seq:
989
- return "Error: llm_map_reduce got no items"
990
- texts = [
991
- (str(it.get("content", "")) if isinstance(it, dict) else str(it))
992
- for it in seq
993
- ]
994
- labels = [
995
- (str(it.get("path", f"item {i + 1}")) if isinstance(it, dict) else f"item {i + 1}")
996
- for i, it in enumerate(seq)
997
- ]
998
- mapped: list[str] = []
999
- for i in range(0, len(texts), _MAX_CHUNK_BATCH):
1000
- batch = [
1001
- f"{map_prompt}\n\n[{labels[i + j]}]\n{t}"
1002
- for j, t in enumerate(texts[i:i + _MAX_CHUNK_BATCH])
1003
- ]
1004
- mapped.extend(self._llm_query_batched(batch, model))
1005
- joined = "\n\n".join(f"[{labels[i]}]\n{a}" for i, a in enumerate(mapped))
1006
- return self._llm_query(f"{reduce_prompt}\n\nPartial answers:\n{joined}", model)
1007
-
1008
- # ---- sync helpers: await(start(...)), so there is exactly one code path -----------------
1009
-
1010
- @_spawnable("llm_query")
1011
- def _llm_query(self, prompt: str, model: str | None = None) -> str:
1012
- return self._await_task(self._start_llm_query(prompt, model))
1013
-
1014
- @_spawnable("llm_query_batched")
1015
- def _llm_query_batched(self, prompts, model: str | None = None) -> list[str]:
1016
- return self._await_task(self._start_llm_query_batched(prompts, model))
1017
-
1018
- @_spawnable("llm_query_chunked")
1019
- def _llm_query_chunked(self, text, prompt: str, model: str | None = None) -> list[str]:
1020
- return self._await_task(self._start_llm_query_chunked(text, prompt, model))
1021
-
1022
- @_spawnable("rlm_query")
1023
- def _rlm_query(self, prompt: str, model: str | None = None) -> str:
1024
- return self._await_task(self._start_rlm_query(prompt, model))
1025
-
1026
- def _ask_user_question(self, questions: list[dict]) -> list[dict]:
1027
- """Present structured questions to the user; blocks until answered.
1028
-
1029
- Returns a list of {question, selected, custom?} dicts.
1030
- Each dict has: question (str), selected (list[str]), custom (str|None).
1031
- Only valid at root depth; sub-RLM calls return an error answer.
1032
- """
1033
- if self.depth > 0:
1034
- qlist = questions if isinstance(questions, list) else []
1035
- return [
1036
- {"question": str(q.get("question", "")) if isinstance(q, dict) else "",
1037
- "selected": [],
1038
- "custom": "Error: ask_user_question not available inside rlm_query sub-calls"}
1039
- for q in qlist
1040
- ] or [{"question": "", "selected": [],
1041
- "custom": "Error: ask_user_question not available inside rlm_query sub-calls"}]
1042
- if not isinstance(questions, list) or not questions:
1043
- return [{"question": "", "selected": [], "custom": "Error: questions must be a non-empty list"}]
1044
- cleaned = []
1045
- for q in questions:
1046
- if not isinstance(q, dict) or "question" not in q or "options" not in q:
1047
- return [{"question": "", "selected": [], "custom": "Error: each question needs 'question', 'header', 'options'"}]
1048
- opts = q.get("options")
1049
- if not isinstance(opts, list):
1050
- return [{"question": str(q.get("question", "")), "selected": [], "custom": "Error: options must be a list"}]
1051
- cleaned_opts = []
1052
- for o in opts:
1053
- if not isinstance(o, dict) or "label" not in o:
1054
- return [{"question": str(q.get("question", "")), "selected": [], "custom": "Error: each option needs 'label'"}]
1055
- item = {"label": str(o["label"]), "description": str(o.get("description", ""))}
1056
- if "preview" in o:
1057
- item["preview"] = str(o["preview"])
1058
- cleaned_opts.append(item)
1059
- cleaned.append({
1060
- "question": str(q["question"]),
1061
- "header": str(q.get("header", "Q")),
1062
- "multiSelect": bool(q.get("multiSelect", False)),
1063
- "options": cleaned_opts,
1064
- })
1065
- r = self._rpc("ask_user_question", {"questions": cleaned})
1066
- if r.get("error"):
1067
- return [{"question": q["question"], "selected": [], "custom": f"Error: {r['error']}"} for q in cleaned]
1068
- answers = r.get("answers", [])
1069
- return answers if isinstance(answers, list) else []
1070
-
1071
- def _todo(self, action: str, **kwargs) -> str:
1072
- """Manage the run's task list.
1073
-
1074
- action: "create" | "update" | "list" | "get" | "delete" | "clear"
1075
- kwargs: id, subject, description, status, activeForm, blockedBy, owner, filterStatus
1076
- Returns a human-readable string result.
1077
- """
1078
- params = {k: v for k, v in kwargs.items() if v is not None}
1079
- r = self._rpc("todo", {"action": str(action), **params})
1080
- if r.get("error"):
1081
- return f"Error: {r['error']}"
1082
- return str(r.get("response", "ok"))
1083
-
1084
- def _load_library(self, source: str) -> dict[str, Any] | str:
1085
- """Pack an external dir/file/git-URL on the host and append it into `context`.
1086
-
1087
- Paths are namespaced under lib/<source_id>/ (host). Content is always in the
1088
- single `context` list — never a new context_N variable.
1089
- Host-side idempotency may return already_loaded without a payload path.
1090
- """
1091
- r = self._rpc("load_library", {"source": str(source)})
1092
- if r.get("error"):
1093
- return f"Error: {r['error']}"
1094
- if r.get("already_loaded"):
1095
- source_id = r.get("source_id") if isinstance(r.get("source_id"), str) else "lib"
1096
- path_prefix = r.get("path_prefix") if isinstance(r.get("path_prefix"), str) else f"lib/{source_id}/"
1097
- ctx = self.ns.get("context")
1098
- ctx_len = len(ctx) if isinstance(ctx, list) else 0
1099
- print(
1100
- f"[rlm] load_library: already loaded {source_id} "
1101
- f"(paths under {path_prefix}, context len={ctx_len})"
1102
- )
1103
- return {
1104
- "source": str(source),
1105
- "source_id": source_id,
1106
- "path_prefix": path_prefix,
1107
- "files": 0,
1108
- "chars": r.get("chars"),
1109
- "context_len": ctx_len,
1110
- "already_loaded": True,
1111
- }
1112
- path = r.get("path")
1113
- if not isinstance(path, str):
1114
- return "Error: malformed load_library reply (no path)"
1115
- try:
1116
- # Worker-internal read — use real io.open so read-only guards never block us.
1117
- with _REAL_IO_OPEN(path, "r") as f:
1118
- payload = json.load(f) if r.get("json") else f.read()
1119
- finally:
1120
- try:
1121
- os.remove(path) # worker owns temp-file cleanup (host does NOT unlink)
1122
- except OSError:
1123
- pass
1124
- return self._append_library(str(source), payload, r)
1125
-
1126
- def _append_library(self, source: str, payload: Any, meta: dict[str, Any]) -> dict[str, Any] | str:
1127
- """Append host-packed library files into `context` (idempotent by path prefix)."""
1128
- ctx = self.ns.get("context")
1129
- if not isinstance(ctx, list):
1130
- kind = type(ctx).__name__ if ctx is not None else "None"
1131
- return f"Error: load_library requires list context (file bundle); got {kind}"
1132
-
1133
- source_id = meta.get("source_id")
1134
- if not isinstance(source_id, str) or not source_id:
1135
- source_id = "lib"
1136
- path_prefix = meta.get("path_prefix")
1137
- if not isinstance(path_prefix, str) or not path_prefix:
1138
- path_prefix = f"lib/{source_id}/"
1139
-
1140
- # Idempotent: already present if any path uses this library prefix.
1141
- for item in ctx:
1142
- if isinstance(item, dict) and str(item.get("path", "")).startswith(path_prefix):
1143
- print(
1144
- f"[rlm] load_library: already loaded {source_id} "
1145
- f"(paths under {path_prefix}, context len={len(ctx)})"
1146
- )
1147
- return {
1148
- "source": source,
1149
- "source_id": source_id,
1150
- "path_prefix": path_prefix,
1151
- "files": 0,
1152
- "chars": meta.get("chars"),
1153
- "context_len": len(ctx),
1154
- "already_loaded": True,
1155
- }
1156
-
1157
- files = self._library_file_entries(payload, path_prefix)
1158
- if not files:
1159
- return "Error: load_library produced no files"
1160
-
1161
- ctx.extend(files)
1162
- # Keep restore payload in sync with the live list.
1163
- self._context_payload = ctx
1164
- self.ns["context"] = ctx
1165
-
1166
- print(
1167
- f"[rlm] load_library: +{len(files)} files into context "
1168
- f"(len={len(ctx)}); paths under {path_prefix}"
1169
- )
1170
- return {
1171
- "source": source,
1172
- "source_id": source_id,
1173
- "path_prefix": path_prefix,
1174
- "files": len(files),
1175
- "chars": meta.get("chars"),
1176
- "context_len": len(ctx),
1177
- "already_loaded": False,
1178
- }
1179
-
1180
- @staticmethod
1181
- def _library_file_entries(payload: Any, path_prefix: str) -> list[dict[str, Any]]:
1182
- """Normalize host payload to list[dict]. Host already namespaces; string is fallback."""
1183
- if isinstance(payload, str):
1184
- return [{
1185
- "path": f"{path_prefix}content",
1186
- "content": payload,
1187
- "tokens": max(1, (len(payload) + 3) // 4),
1188
- }]
1189
- if not isinstance(payload, list):
1190
- return []
1191
- out: list[dict[str, Any]] = []
1192
- for item in payload:
1193
- if isinstance(item, dict) and "path" in item and "content" in item:
1194
- out.append(item)
1195
- return out
1196
-
1197
- def _advance_phase(self, phase: str, summary: str | None = None) -> str:
1198
- """Transition the root RLM pipeline to a new phase.
1199
-
1200
- Only callable at depth 0. The parent handler validates the transition
1201
- against the phase state machine (research → blueprint → validate)
1202
- and runs deterministic artifact gates before accepting the transition.
1203
- Returns a short confirmation, or an `Error: …` string the model can act on.
1204
- """
1205
- if self.depth > 0:
1206
- return "Error: advance_phase is only available at the root RLM depth"
1207
- r = self._rpc("advance_phase", {"phase": str(phase), "summary": summary})
1208
- if r.get("error"):
1209
- return f"Error: {r['error']}"
1210
- response = r.get("response", "ok")
1211
- if isinstance(response, str) and response.startswith("Error:"):
1212
- return response
1213
- return response if isinstance(response, str) else "ok"
1214
-
1215
- def _save_artifact(self, kind: str, content: str) -> str:
1216
- """Persist a stage artifact (research/plan/validation) under .rlm/artifacts/.
1217
-
1218
- Only callable at depth 0. The engine gates advance_phase against the latest
1219
- saved artifact for the current stage.
1220
- """
1221
- if self.depth > 0:
1222
- return "Error: save_artifact is only available at the root RLM depth"
1223
- r = self._rpc("save_artifact", {"artifactKind": str(kind), "content": str(content)})
1224
- if r.get("error"):
1225
- return f"Error: {r['error']}"
1226
- response = r.get("response", "ok")
1227
- if isinstance(response, str) and response.startswith("Error:"):
1228
- return response
1229
- return response if isinstance(response, str) else "ok"
1230
-
1231
- @_spawnable("rlm_query_batched")
1232
- def _rlm_query_batched(self, prompts, model: str | None = None) -> list[str]:
1233
- return self._await_task(self._start_rlm_query_batched(prompts, model))
1234
-
1235
- # ---- context + execution --------------------------------------------------------------
1236
-
1237
- def load_context(self, path: str, index: int | None = None, is_json: bool = False) -> int:
1238
- """Load the packed world into the single REPL variable `context`.
1239
-
1240
- `index` is accepted for protocol compatibility but ignored — there is only
1241
- one context slot. Libraries are merged on the host (or via load_library).
1242
- """
1243
- with open(path, "r") as f:
1244
- payload = json.load(f) if is_json else f.read()
1245
- self._context_payload = payload
1246
- self.ns["context"] = payload
1247
- # Drop legacy multi-slot names if present.
1248
- for k in list(self.ns.keys()):
1249
- if k != "context" and _CONTEXT_NAME.match(k):
1250
- del self.ns[k]
1251
- return 0
1252
-
1253
- @contextmanager
1254
- def _capture(self):
1255
- out, err = io.StringIO(), io.StringIO()
1256
- old_out, old_err = sys.stdout, sys.stderr
1257
- sys.stdout, sys.stderr = out, err
1258
- try:
1259
- yield out, err
1260
- finally:
1261
- sys.stdout, sys.stderr = old_out, old_err
1262
-
1263
- def _exec(self, code: str, ns: dict[str, Any]) -> None:
1264
- t = self.exec_timeout_s
1265
- if t <= 0 or not hasattr(signal, "SIGALRM"):
1266
- exec(compile(code, "<repl>", "exec"), ns, ns) # noqa: S102
1267
- return
1268
-
1269
- def _alarm(signum, frame): # noqa: ARG001
1270
- raise TimeoutError(f"```repl``` block exceeded {t:g}s timeout")
1271
-
1272
- old = signal.signal(signal.SIGALRM, _alarm)
1273
- signal.setitimer(signal.ITIMER_REAL, t)
1274
- try:
1275
- exec(compile(code, "<repl>", "exec"), ns, ns) # noqa: S102
1276
- finally:
1277
- signal.setitimer(signal.ITIMER_REAL, 0)
1278
- signal.signal(signal.SIGALRM, old)
1279
-
1280
- def _nudge_lines(self) -> list[str]:
1281
- """One-time hint for newly created huge raw-text variables (single line).
1282
-
1283
- Collapses to one line so it survives headless stdout elision (head 200 + tail 200).
1284
- """
1285
- names: list[str] = []
1286
- for k in self._user_var_names():
1287
- v = self.ns.get(k)
1288
- if isinstance(v, (str, bytes)) and len(v) > _NUDGE_CHARS and k not in self._nudged:
1289
- self._nudged.add(k)
1290
- names.append(f"{k} ({len(v):,} chars)")
1291
- if not names:
1292
- return []
1293
- return [
1294
- f"[rlm] huge raw-text variable(s): {', '.join(names)} — do NOT analyze them yourself; "
1295
- 'delegate with llm_query_chunked(name, "your question") or slice + llm_query_batched.'
1296
- ]
1297
-
1298
- def execute(self, code: str) -> dict[str, Any]:
1299
- start = time.perf_counter()
1300
- raised = False
1301
- with self._capture() as (out, err):
1302
- try:
1303
- self._restore_scaffold()
1304
- self._exec(code, self.ns)
1305
- self._restore_scaffold()
1306
- stdout, stderr = out.getvalue(), err.getvalue()
1307
- except BaseException as e: # noqa: BLE001
1308
- raised = True
1309
- self._restore_scaffold()
1310
- stdout = out.getvalue()
1311
- stderr = err.getvalue() + f"\n{type(e).__name__}: {e}\n" + traceback.format_exc()
1312
- final, self._final_answer = self._final_answer, None
1313
- answer = self.ns.get("answer")
1314
- answer_content = answer.get("content", "") if isinstance(answer, dict) else ""
1315
- # ready may have been flipped with empty content before content was assigned later
1316
- # in the same block; the dict's current content is the real submission.
1317
- if final is not None and not final.strip() and str(answer_content).strip():
1318
- final = str(answer_content)
1319
- nudges = self._nudge_lines()
1320
- if nudges:
1321
- parts = [stdout] if stdout else []
1322
- parts.extend(nudges)
1323
- stdout = "\n".join(parts) + "\n"
1324
- return {
1325
- "stdout": stdout,
1326
- "stderr": stderr,
1327
- "final_answer": final,
1328
- "answer_content": str(answer_content),
1329
- "raised": raised,
1330
- "execution_time": time.perf_counter() - start,
1331
- "var_names": self._user_var_names(),
1332
- }
1333
-
1334
- def _serializer(self):
1335
- try:
1336
- import dill as s
1337
- return s
1338
- except ImportError:
1339
- return pickle
1340
-
1341
- def snapshot(self, path: str, nonce: str) -> dict:
1342
- """Pickle user variables atomically to path. Stores session nonce for restore verification.
1343
-
1344
- Writes to path.tmp then os.rename — atomic on POSIX, so no .tmp leak and no
1345
- TypeScript-side finalize step needed. On resume (fresh session = different nonce),
1346
- restore fails — caller falls back to history-only replay.
1347
- """
1348
- s = self._serializer()
1349
- out, skipped = {}, []
1350
- MAX_VAR_BYTES = 50 * 1024 * 1024
1351
- for k, v in self.ns.items():
1352
- if k.startswith("_") or _CONTEXT_NAME.match(k) or k in RESERVED or k == "__builtins__":
1353
- continue
1354
- # A Task holds a back-reference to this Worker, so dill would happily pickle the
1355
- # whole process. Top-level guard only: a Task nested inside a list/dict still
1356
- # falls to the generic `except` below and skips the variable — which is why this
1357
- # guard is explicit rather than left to that fallback.
1358
- if isinstance(v, Task):
1359
- skipped.append(k)
1360
- continue
1361
- try:
1362
- blob = s.dumps(v)
1363
- if len(blob) > MAX_VAR_BYTES:
1364
- skipped.append(k)
1365
- continue
1366
- out[k] = v
1367
- except Exception:
1368
- skipped.append(k)
1369
- if skipped:
1370
- print(f"[rlm-sandbox] snapshot skipped {len(skipped)} unpicklable/oversized vars: {skipped}", file=_REAL_STDERR)
1371
- tmp = path + ".tmp"
1372
- with _REAL_IO_OPEN(tmp, "wb") as f:
1373
- s.dump({"nonce": nonce, "vars": out}, f)
1374
- os.rename(tmp, path) # atomic rename
1375
- return {"skipped": skipped}
1376
-
1377
- def restore(self, path: str, nonce: str) -> dict:
1378
- """Restore user variables from a pickle file. Verifies session nonce before deserializing.
1379
-
1380
- SECURITY: pickle.load executes arbitrary code. The session nonce check ensures the
1381
- .pkl was written by THIS engine session. Cross-session resume falls back to
1382
- history-only replay (caller skips restore when sessionNonce is undefined).
1383
- """
1384
- s = self._serializer()
1385
- with _REAL_IO_OPEN(path, "rb") as f:
1386
- data = s.load(f)
1387
- if not isinstance(data, dict) or data.get("nonce") != nonce:
1388
- raise ValueError("snapshot nonce mismatch — not from this session")
1389
- self.ns.update(data.get("vars", {}))
1390
- self._restore_scaffold()
1391
- return {"restored": list(data.get("vars", {}).keys())}
1392
-
1393
-
1394
- def main() -> None:
1395
- ap = argparse.ArgumentParser()
1396
- ap.add_argument("--depth", type=int, default=int(os.environ.get("RLM_DEPTH", "1")))
1397
- ap.add_argument("--timeout", type=float, default=float(os.environ.get("RLM_EXEC_TIMEOUT_S", "600")))
1398
- ap.add_argument("--await-timeout", type=float,
1399
- default=float(os.environ.get("RLM_AWAIT_TIMEOUT_S", "600")))
1400
- ap.add_argument("--max-prompt-chars", type=int,
1401
- default=int(os.environ.get("RLM_MAX_PROMPT_CHARS", "400000")))
1402
- ap.add_argument("--read-only", action="store_true",
1403
- default=os.environ.get("RLM_READ_ONLY", "").lower() in ("1", "true", "yes"),
1404
- help="Reject open() write modes (pipeline runs)")
1405
- args = ap.parse_args()
1406
-
1407
- worker = Worker(depth=args.depth, exec_timeout_s=args.timeout,
1408
- max_prompt_chars=args.max_prompt_chars, read_only=args.read_only,
1409
- await_timeout_s=args.await_timeout)
1410
- _send({"id": "_init", "ok": True})
1411
-
1412
- while True:
1413
- # Requests that arrived mid-exec were parked by _pump; replay them before reading.
1414
- req = worker.take_deferred()
1415
- if req is None:
1416
- line = _REAL_STDIN.readline()
1417
- if not line:
1418
- return
1419
- raw = line.strip()
1420
- if not raw:
1421
- continue
1422
- try:
1423
- req = json.loads(raw)
1424
- except json.JSONDecodeError as e:
1425
- _send({"id": "?", "ok": False, "error": f"bad json: {e}"})
1426
- continue
1427
- # A task spawned in an earlier exec settling while the worker is idle. Park it for
1428
- # a later rlm_await; without this it would fall through to "unknown type" and the
1429
- # result would be lost.
1430
- if worker.park_reply(req):
1431
- continue
1432
- if not isinstance(req, dict):
1433
- _send({"id": "?", "ok": False, "error": f"expected an object, got {type(req).__name__}"})
1434
- continue
1435
- rid, kind = req.get("id", "?"), req.get("type")
1436
- try:
1437
- if kind == "exec":
1438
- _send({"id": rid, "ok": True, **worker.execute(req.get("code", ""))})
1439
- elif kind == "load_context":
1440
- idx = worker.load_context(req.get("path"), req.get("index"), req.get("json"))
1441
- _send({"id": rid, "ok": True, "index": idx})
1442
- elif kind == "shutdown":
1443
- _send({"id": rid, "ok": True})
1444
- return
1445
- elif kind == "snapshot":
1446
- _send({"id": rid, "ok": True, **worker.snapshot(req.get("path", ""), req.get("nonce", ""))})
1447
- elif kind == "restore":
1448
- _send({"id": rid, "ok": True, **worker.restore(req.get("path", ""), req.get("nonce", ""))})
1449
- else:
1450
- _send({"id": rid, "ok": False, "error": f"unknown type: {kind!r}"})
1451
- except BaseException as e: # noqa: BLE001
1452
- _send({"id": rid, "ok": False, "error": f"{type(e).__name__}: {e}\n{traceback.format_exc()}"})
1453
-
1454
-
1455
- if __name__ == "__main__":
1456
- main()