bootstrap.py 82 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711
  1. """CPython bootstrap for dsh-code-runtime-python.
  2. Reads a :class:`BootMessage` on fd 3, applies resource limits and log capture,
  3. reads a :class:`RunMessage`, runs the model program as the body of an async
  4. function (top-level ``await`` and ``return`` both work; the returned value is
  5. the completion), and posts a terminal :class:`DoneMessage`. The program calls
  6. host functions through the ``tools`` (or other namespace) proxy, whose attribute
  7. and subscript access return awaitables that ride binding messages over fd 3.
  8. This module runs under ``python3 -I`` with an empty environment and
  9. ``sys.path`` containing only its own directory.
  10. """
  11. from __future__ import annotations
  12. import asyncio
  13. import ast
  14. import io
  15. import json
  16. import math
  17. import os
  18. import re
  19. import resource
  20. import signal
  21. import sys
  22. import threading
  23. import traceback
  24. from decimal import Decimal
  25. from pathlib import Path
  26. from typing import Any
  27. # ``python3 -I`` (isolated) drops the script directory from ``sys.path`` so
  28. # the sibling ``protocol.py`` is invisible by default. Restore it explicitly
  29. # before importing.
  30. sys.path.insert(0, str(Path(__file__).resolve().parent))
  31. from protocol import PROTOCOL_FD, log_truncation_marker # noqa: E402
  32. # Read size for the async fd-3 reader. One `os.read` returns whatever the pipe
  33. # holds, so this only bounds a single syscall's copy, not a frame: a larger frame
  34. # simply takes more reads. 64 KiB matches the usual pipe capacity.
  35. _READ_CHUNK_BYTES = 65536
  36. # Code-unit ceiling on the exception class name interpolated into the LAST-resort
  37. # failure diagnostic. A metaclass `__name__` property can return any length, and
  38. # that construction runs outside the guard that would otherwise absorb a
  39. # MemoryError, so the name is sliced before it is copied. Generous enough that no
  40. # real class name is touched.
  41. _MAX_FALLBACK_NAME_CHARS = 200
  42. # ---------------------------------------------------------------------------
  43. # Log buffer — Python-side ledger for captured text.
  44. # ---------------------------------------------------------------------------
  45. class LogBuffer:
  46. """Ordered text capture under one shared byte budget.
  47. Once the budget is exhausted the buffer emits exactly one in-band
  48. truncation marker via ``sink`` and silently drops everything after. The
  49. cap is a blast-radius bound; "how much was lost" intentionally stays
  50. unmeasured.
  51. """
  52. def __init__(self, max_bytes: int, sink) -> None:
  53. self._max_bytes = max_bytes
  54. self._remaining = max_bytes
  55. self._truncated = False
  56. # Re-entrant so a caller may hold it across a compound read-modify-write
  57. # (``_LogStream.write`` reads ``remaining`` several times and then calls
  58. # ``push`` while still holding it). One lock is shared by this buffer and
  59. # every stream that funnels into it: model code may start daemon threads
  60. # that keep calling ``print`` after the program body returns, and the
  61. # settlement ``flush_line`` on the main coroutine reads and mutates the
  62. # same ``_pending``/ledger state. Without a shared lock the flush could
  63. # interleave with a concurrent ``write`` — dropping or double-counting a
  64. # line, or costing the ``done`` frame on a mangled ledger. Fixing which
  65. # callable runs (binding ``out_stream.flush_line``) does not fix what it
  66. # reads.
  67. self._lock = threading.RLock()
  68. # ``sink(text, truncated=False)``. The marker is emitted with
  69. # ``truncated=True`` so the host can stop its own capture at the same
  70. # point rather than treating the marker as ordinary program output: the
  71. # two ledgers exhaust independently, and one entry larger than
  72. # ``max_bytes`` sends only the marker while the host budget is still
  73. # nearly empty.
  74. self._sink = sink
  75. @property
  76. def lock(self) -> "threading.RLock":
  77. """The shared re-entrant lock guarding this ledger and its streams' buffers."""
  78. return self._lock
  79. @property
  80. def remaining(self) -> int:
  81. """Serialized bytes still admissible; zero once truncated (streams use this to bound their own buffering, where a character count is a valid lower bound)."""
  82. return 0 if self._truncated else self._remaining
  83. def push(self, text: str) -> None:
  84. with self._lock:
  85. self._push_locked(text)
  86. def _push_locked(self, text: str) -> None:
  87. if self._truncated:
  88. return
  89. # Cheap lower bound FIRST: one char is at least one UTF-8 byte and the
  90. # JSON form adds two quotes plus the separator, so a single print() far
  91. # above the budget truncates without ever encoding it — the full encode
  92. # would allocate a second equally large string and could turn a
  93. # truncatable log into an RLIMIT_AS death.
  94. if len(text) + 3 > self._remaining:
  95. self._truncated = True
  96. self._sink(log_truncation_marker(self._max_bytes), truncated=True)
  97. return
  98. # A model print() can emit a lone surrogate; strict UTF-8 throws on it
  99. # here. Replace it rather than escaping it the way :func:`_dump_string`
  100. # preserves one inside a completion VALUE: log text is already a
  101. # truncatable, substituting channel (the byte cap replaces the tail with
  102. # a marker), and the ledger below charges the RAW UTF-8 bytes, which
  103. # would undercharge the six-byte escape by half. Bounded: the text
  104. # passed the length check, so this encodes at most ~4x remaining.
  105. try:
  106. raw = text.encode("utf-8")
  107. except UnicodeEncodeError:
  108. raw = text.encode("utf-8", errors="replace")
  109. text = raw.decode("utf-8")
  110. # Charge the SERIALIZED cost — the JSON string form's bytes plus one
  111. # separator byte — exactly as the host ledger does. Charging the raw
  112. # UTF-8 length instead undercharges control-heavy text, whose JSON
  113. # escaping expands it up to sixfold (a NUL costs one raw byte but six
  114. # as its ``\uXXXX`` escape): a NUL flood sized to fit ``maxLogBytes``
  115. # raw would serialize to roughly six times the shared cap, and the child
  116. # could then die on RLIMIT_AS (reported host-side as ``worker-exit``)
  117. # instead of emitting the truncation marker. The +1 also floors an empty
  118. # entry above zero, so a flood of blank ``print()`` lines exhausts the
  119. # budget instead of emitting unbounded zero-cost log frames.
  120. cost = _json_string_cost(raw) + 1
  121. if cost > self._remaining:
  122. self._truncated = True
  123. self._sink(log_truncation_marker(self._max_bytes), truncated=True)
  124. return
  125. self._remaining -= cost
  126. self._sink(text)
  127. class _LogStream(io.TextIOBase):
  128. """A newline-coalescing text stream backed by a :class:`LogBuffer`.
  129. Installed as ``sys.stdout`` / ``sys.stderr`` before executing the model
  130. program. ``print(...)`` calls ``write`` once per argument, separator, and
  131. newline, so a raw one-push-per-write stream would emit
  132. ``["a", " ", "b", "\\n"]`` for ``print("a", "b")`` — and Code Mode renders
  133. ``logs`` with ``join('\\n')``, turning that into spurious blank lines. This
  134. stream instead buffers writes and pushes one LogBuffer entry per completed
  135. LINE (the text up to each ``\\n``, newline stripped), so the rendered join
  136. reproduces ``a b``. Any unterminated tail is flushed by :meth:`flush_line`
  137. after the program settles.
  138. """
  139. def __init__(self, logs: LogBuffer) -> None:
  140. super().__init__()
  141. self._logs = logs
  142. # list-of-chunks, joined only at a newline or flush: repeated
  143. # ``print("x", end="")`` must not concatenate quadratically.
  144. self._pending: list[str] = []
  145. self._pending_chars = 0
  146. def writable(self) -> bool: # noqa: D401 -- inherited contract
  147. return True
  148. def write(self, text: str) -> int: # noqa: D401 -- inherited contract
  149. # Serialize the whole read-modify-write against the settlement flush and
  150. # any other thread's write: model code may spawn daemon threads that keep
  151. # printing after the program body returns, and this method reads
  152. # ``remaining`` and mutates ``_pending``/the ledger across many steps. The
  153. # lock is the buffer's and is re-entrant, so the ``push`` calls below
  154. # (which re-acquire it) do not deadlock.
  155. with self._logs.lock:
  156. return self._write_locked(text)
  157. def _write_locked(self, text: str) -> int:
  158. # Drop an empty write instead of buffering it. An empty chunk adds no
  159. # character, so the budget check below can never fire on it:
  160. # ``while True: sys.stdout.write("")`` would append one list slot per
  161. # call with `_pending_chars` pinned at 0, growing unbounded long after
  162. # the log ledger was exhausted (about 3.7 M slots per CPU second here)
  163. # until RLIMIT_AS turned the allocation into a MemoryError — reported as
  164. # the program's own exception rather than the intended bounded-log
  165. # behavior. Returning here also keeps `flush_line` from pushing a
  166. # spurious empty log entry for a program whose only writes were empty.
  167. if not text:
  168. return 0
  169. if "\n" in text:
  170. # Scan `text` in place; the buffered chunks are joined ONLY into the
  171. # first line. Joining the pending chunks with the whole write first
  172. # made a second copy of that write, which an over-budget write
  173. # cannot afford (measured under a 400 MiB addressSpaceMb: one
  174. # buffered character followed by a 340 MiB write died on MemoryError
  175. # inside the join, reported as the program's own exception, and the
  176. # retained chunks made the settlement `flush_line` fail the same way
  177. # — costing the `done` frame and turning the run into a wall-clock
  178. # timeout instead of the promised truncation marker).
  179. length = len(text)
  180. pos = 0
  181. if self._pending:
  182. newline = text.index("\n")
  183. if self._pending_chars + newline + 3 > self._logs.remaining:
  184. # The reconstructed first line cannot fit the ledger, so
  185. # LogBuffer would reject it whole: copy only the prefix that
  186. # fails its cheap bound and drop the chunks. The slice is
  187. # bounded HERE, not inside the helper: `text[:newline]` on a
  188. # 340 MiB newline-terminated write is the same full copy the
  189. # join was (measured: MemoryError inside `sys.stdout.write`
  190. # under a 400 MiB addressSpaceMb), and `remaining + 4`
  191. # characters are all the helper can use.
  192. self._push_bounded_prefix(text[: min(newline, self._logs.remaining + 4)])
  193. else:
  194. self._pending.append(text[:newline])
  195. line = "".join(self._pending)
  196. self._pending = []
  197. self._pending_chars = 0
  198. self._logs.push(line)
  199. pos = newline + 1
  200. # Scan by offset and STOP once the ledger is exhausted: a single
  201. # write of many newlines (``print("\n" * 1000000)``) would otherwise
  202. # re-slice the tail once per line and keep pushing long after
  203. # LogBuffer truncated, burning the CPU budget on discarded lines.
  204. # `remaining` reads 0 the instant the buffer truncates, so the loop
  205. # exits immediately; the unscanned tail is simply dropped.
  206. while pos < length and self._logs.remaining > 0:
  207. newline = text.find("\n", pos)
  208. if newline < 0:
  209. break
  210. # Bound the SLICE the same way LogBuffer bounds the encode: a
  211. # first line far above the ledger would be copied whole before
  212. # push could reject it, and that copy is the allocation an
  213. # over-budget write cannot afford. Copy only a budget-sized
  214. # prefix, which push still rejects on its own cheap bound (the
  215. # prefix is longer than `remaining`), so the marker is emitted
  216. # and the oversized line is never materialized.
  217. if newline - pos + 3 > self._logs.remaining:
  218. self._logs.push(text[pos:pos + self._logs.remaining + 4])
  219. break
  220. self._logs.push(text[pos:newline])
  221. pos = newline + 1
  222. if pos < length:
  223. if self._logs.remaining > 0:
  224. tail = text[pos:]
  225. self._pending.append(tail)
  226. self._pending_chars = len(tail)
  227. else:
  228. # The ledger ran out with text still unscanned, so that text
  229. # IS being dropped and the run must say so. One push is
  230. # enough and is bounded: `remaining` is 0, so LogBuffer's
  231. # cheap length lower bound rejects immediately, emits the
  232. # marker, and never encodes the tail — and a push after the
  233. # marker is already out returns without emitting a second.
  234. # Reaching 0 EXACTLY (65 one-character lines against the
  235. # default 3-byte-per-entry serialized charge) leaves
  236. # `_truncated` unset, so without this the tail vanished with
  237. # no marker at all. Sliced to a budget-sized prefix, not the
  238. # whole tail: the tail can be hundreds of megabytes and the
  239. # copy would be the RLIMIT_AS death this bound exists to
  240. # avoid, while push only needs enough characters to fail its
  241. # own cheap length check.
  242. self._logs.push(text[pos:pos + self._logs.remaining + 4])
  243. else:
  244. self._pending.append(text)
  245. self._pending_chars += len(text)
  246. # A newline-free flood must hit the budget while running, not at
  247. # settlement: once the buffered tail alone can no longer fit the
  248. # ledger (chars lower-bound the serialized cost), push it through — LogBuffer
  249. # truncates, emits the marker once, and swallows everything after.
  250. if self._pending_chars > self._logs.remaining:
  251. self._push_bounded_prefix()
  252. return len(text)
  253. def _push_bounded_prefix(self, extra: str = "") -> None:
  254. # Reached only when the buffered characters already exceed what the
  255. # ledger admits, so LogBuffer is certain to reject on its cheap length
  256. # bound and emit the marker. Copy a budget-sized PREFIX rather than the
  257. # joined whole: ``sys.stdout.write("x")`` followed by one newline-free
  258. # 340 MiB write leaves two chunks whose join is a second copy of the
  259. # payload, and under a tight addressSpaceMb that join raises MemoryError
  260. # from inside `write` — surfacing as the program's own exception, or,
  261. # while the oversized chunks stayed retained, again from `flush_line`
  262. # after the program settled, which cost the `done` frame and turned the
  263. # run into a wall-clock timeout instead of the promised truncation
  264. # marker.
  265. #
  266. # The chunks are dropped BEFORE the push so neither this call nor the
  267. # settlement flush can repeat the allocation, and dropping the text is
  268. # exactly what the marker reports. `remaining + 4` is the shortest
  269. # prefix that still fails LogBuffer's ``len(text) + 3 > remaining``
  270. # check; the accumulation stops there, so the copy is bounded by the log
  271. # budget however large the pending chunks are.
  272. limit = self._logs.remaining + 4
  273. parts: list[str] = []
  274. total = 0
  275. for chunk in (*self._pending, extra):
  276. parts.append(chunk[: limit - total])
  277. total += len(parts[-1])
  278. if total >= limit:
  279. break
  280. self._pending = []
  281. self._pending_chars = 0
  282. self._logs.push("".join(parts))
  283. def flush(self) -> None: # noqa: D401 -- inherited contract
  284. # ``TextIOBase.flush`` is a no-op, so without this override an explicit
  285. # ``print(..., flush=True)`` or ``sys.stdout.flush()`` left the text in
  286. # `_pending` with nothing to drain it except `flush_line` after the
  287. # program settles. A run that then hangs or is killed never reaches that
  288. # call: ``print("before hang", end="", flush=True)`` followed by an
  289. # infinite loop returned `logs: []`, losing the one diagnostic the
  290. # program deliberately committed. Forwarding makes an explicit flush emit
  291. # the pending entry immediately, which is what the caller asked for; a
  292. # newline-terminated write already emitted on its own.
  293. self.flush_line()
  294. def flush_line(self) -> None:
  295. """Push any buffered text not terminated by a newline (also serves explicit flushes)."""
  296. # Same shared, re-entrant lock as ``write``: the settlement flush on the
  297. # main coroutine and a daemon thread's concurrent ``write`` both touch
  298. # ``_pending`` and the ledger, so this read-and-clear must be atomic
  299. # against them.
  300. with self._logs.lock:
  301. if self._pending:
  302. self._logs.push("".join(self._pending))
  303. self._pending = []
  304. self._pending_chars = 0
  305. # ---------------------------------------------------------------------------
  306. # Fd-3 channel — line-framed JSON.
  307. # ---------------------------------------------------------------------------
  308. class ProtocolChannel:
  309. """Blocking readers and synchronous writers over the fd-3 protocol pipe.
  310. Writes are unbuffered and go straight to the fd, so ``send_sync`` is safe
  311. from inside model code (which may run outside an asyncio task) and from
  312. background tasks alike. The single writer is serialized by CPython's GIL
  313. plus one os.write per frame (POSIX guarantees atomicity for writes below
  314. ``PIPE_BUF``, and our frames are short JSON lines).
  315. """
  316. def __init__(self, fd: int) -> None:
  317. # Unbuffered binary I/O so we never lose frames to an idle flush.
  318. self._reader = os.fdopen(fd, "rb", buffering=0, closefd=False)
  319. self._fd = fd
  320. # Residual bytes read past a frame's newline. Held here, not in the
  321. # reading coroutine: the reply pump is cancelled once `done` is posted,
  322. # and read-ahead sitting in a local would be lost with it.
  323. self._pending = bytearray()
  324. # Serializes writers: os.write releases the GIL, and a frame larger
  325. # than PIPE_BUF is neither atomic nor guaranteed fully consumed by one
  326. # call — without the lock, model-created threads printing while a big
  327. # completion frame drains could interleave bytes mid-frame.
  328. self._write_lock = threading.Lock()
  329. def read_frame(self) -> dict[str, Any] | None:
  330. """Read one JSON-line frame (iteratively decoded). ``None`` on EOF.
  331. Blocking. Used for the two frames read BEFORE the model program starts
  332. (``boot`` and ``run``), where blocking is what the handshake wants. Reply
  333. frames arriving during the program go through :meth:`read_frame_async`,
  334. which must not occupy a thread.
  335. """
  336. line = self._reader.readline()
  337. if not line:
  338. return None
  339. return _decode_json_plain(line.decode("utf-8"))
  340. async def read_frame_async(self) -> dict[str, Any] | None:
  341. """Await one JSON-line frame without occupying a thread. ``None`` on EOF.
  342. ``loop.run_in_executor(None, read_frame)`` was the obvious spelling and
  343. the wrong one: the default executor spins up its first thread the moment
  344. the program awaits a binding, and on Linux/glibc that thread's 8 MiB
  345. stack plus a 64 MiB per-thread malloc arena reservation are charged to
  346. ``RLIMIT_AS`` — measured, the child's mappings went from 30.34 MiB to
  347. 102.39 MiB across one ``await tools.*``. Since the limit is already in
  348. force, that ~72 MiB comes straight out of the run's ``addressSpaceMb``:
  349. under a small limit the thread cannot start at all and a legitimate
  350. binding call hangs to ``maxWallMs``, and under a larger one an allocation
  351. that should have fit dies as ``MemoryError``. This is the same accounting
  352. the settlement-time CPU recheck was designed around, where a sampling
  353. thread cost the same 72 MiB.
  354. `loop.add_reader` watches the fd instead, so no thread exists.
  355. Bytes past a frame's newline belong to the next frame, so the residual
  356. lives on the CHANNEL rather than in this coroutine: the pump is cancelled
  357. once ``done`` is posted, and a local buffer would discard whatever it had
  358. read ahead.
  359. """
  360. loop = asyncio.get_event_loop()
  361. while True:
  362. newline = self._pending.find(b"\n")
  363. if newline >= 0:
  364. line = bytes(self._pending[:newline])
  365. del self._pending[: newline + 1]
  366. return _decode_json_plain(line.decode("utf-8"))
  367. ready = loop.create_future()
  368. # `add_reader` only reports readability; the read itself happens here,
  369. # and `os.read` returns whatever is buffered without waiting for more.
  370. loop.add_reader(self._fd, lambda: ready.done() or ready.set_result(None))
  371. try:
  372. await ready
  373. finally:
  374. loop.remove_reader(self._fd)
  375. chunk = os.read(self._fd, _READ_CHUNK_BYTES)
  376. if not chunk:
  377. # EOF. Any partial line is dropped, matching how the host drops a
  378. # frame that never completed.
  379. return None
  380. self._pending.extend(chunk)
  381. def send_sync(self, message: dict[str, Any]) -> None:
  382. """Post one frame synchronously.
  383. Encoded with the iterative :func:`_encode_json_plain` (not
  384. ``json.dumps``, whose per-level recursion would raise
  385. ``RecursionError`` on a deeply nested completion or call argument the
  386. depth-unbounded ``CodeJsonValue`` contract admits). NaN/Infinity still
  387. raise ``ValueError`` — they would serialize as non-standard tokens
  388. that Node's ``JSON.parse`` rejects, silently dropping the frame, and a
  389. call would then hang until the wall clock instead of failing fast.
  390. Callers turn the ``ValueError`` into their own contract error
  391. (dispatch raises the lossless-JSON message).
  392. """
  393. payload = (_encode_json_plain(message) + "\n").encode("utf-8")
  394. # Full-write loop under the writer lock: one os.write may consume only
  395. # part of a frame beyond PIPE_BUF (64 KiB logs / 32 KiB completions /
  396. # uncapped call args exceed it), and a partial or interleaved frame is
  397. # dropped host-side as malformed JSON — the run would then hang to the
  398. # wall clock.
  399. with self._write_lock:
  400. view = memoryview(payload)
  401. while view:
  402. view = view[os.write(self._fd, view):]
  403. # ---------------------------------------------------------------------------
  404. # Tools proxy — turns ``await tools.name(args)`` into a fd-3 call frame.
  405. # ---------------------------------------------------------------------------
  406. class _Namespace:
  407. """A proxy for one binding namespace: every declared name routes to the bridge.
  408. Names arrive from :class:`BootMessage.namespaces`. Both attribute access
  409. (``tools.name``) and subscript access (``tools["my-tool"]`` — the SDK's
  410. escape hatch for exotic or reserved names, which are legal function names
  411. on the wire) return a coroutine factory that posts a ``call`` frame and
  412. awaits the matching ``reply``. An undeclared name raises ``AttributeError``
  413. (attribute) or ``KeyError`` (subscript), matching the worker backend's
  414. own-property discipline.
  415. ``__getattribute__`` (not ``__getattr__``) intercepts attribute access so a
  416. declared name ALWAYS reaches the bridge — even one that collides with an
  417. inherited attribute like ``__class__``, which ordinary lookup would resolve
  418. on ``object`` before ``__getattr__`` ever ran. Internal state lives under
  419. name-mangled ``_Namespace__*`` attributes; a declared binding with such a
  420. name still wins (declared-names check runs first).
  421. """
  422. def __init__(self, global_name: str, names: list[str], dispatch) -> None:
  423. self.__global = global_name
  424. self.__names = set(names)
  425. self.__dispatch = dispatch
  426. def __call_for(self, name: str):
  427. dispatch = object.__getattribute__(self, "_Namespace__dispatch")
  428. global_name = object.__getattribute__(self, "_Namespace__global")
  429. async def call(args: Any) -> Any:
  430. return await dispatch(global_name, name, args)
  431. return call
  432. def __getattribute__(self, name: str):
  433. # Declared names route to the bridge unconditionally — before Python
  434. # can resolve an inherited attribute (``__class__``) or our own
  435. # internals. Everything else falls through to normal lookup so the
  436. # proxy machinery itself keeps working.
  437. names = object.__getattribute__(self, "_Namespace__names")
  438. if name in names:
  439. return object.__getattribute__(self, "_Namespace__call_for")(name)
  440. return object.__getattribute__(self, name)
  441. def __getattr__(self, name: str):
  442. # Reached only when normal lookup found nothing (declared names were
  443. # already intercepted above), so this is always an undeclared tool.
  444. raise AttributeError(
  445. f"tool {name!r} is not declared in namespace "
  446. f"{object.__getattribute__(self, '_Namespace__global')!r}"
  447. )
  448. def __getitem__(self, name: str):
  449. names = object.__getattribute__(self, "_Namespace__names")
  450. if name not in names:
  451. raise KeyError(
  452. f"tool {name!r} is not declared in namespace "
  453. f"{object.__getattribute__(self, '_Namespace__global')!r}"
  454. )
  455. return object.__getattribute__(self, "_Namespace__call_for")(name)
  456. class _BindingRejection(Exception):
  457. """Internal reply-pump rejection, converted by ``dispatch`` into the
  458. namespace's declared error class (or ``RuntimeError``) so the marker type
  459. itself never reaches model code."""
  460. def _make_error_class(name: str, member_name_property: str) -> type:
  461. """Mint one program-visible rejection class per the seam's
  462. ``CodeBindingErrorClass`` contract: instances carry the failed member name
  463. under ``member_name_property`` and render as their message."""
  464. def __init__(self, member_name: str, message: str) -> None: # noqa: N807
  465. Exception.__init__(self, message)
  466. setattr(self, member_name_property, member_name)
  467. return type(name, (Exception,), {"__init__": __init__})
  468. def _clamped(which: int, soft: int, hard: int) -> tuple[int, int]:
  469. """Bound a requested (soft, hard) rlimit pair by the inherited hard limit.
  470. An unprivileged process may lower a hard limit but never raise it, so a
  471. harness already started under a tighter ceiling (``ulimit -v`` below
  472. ``addressSpaceBytes``, or a CPU cap below ``cpuSeconds`` + 1) would make
  473. ``setrlimit`` raise ``ValueError`` and fail every run — despite the
  474. inherited limit being STRONGER than the one requested. Clamping keeps the
  475. stricter of the two, which still satisfies the containment contract.
  476. ``RLIM_INFINITY`` compares as -1, so it is special-cased rather than
  477. treated as the smallest bound.
  478. """
  479. inherited = resource.getrlimit(which)[1]
  480. if inherited == resource.RLIM_INFINITY:
  481. return (soft, hard)
  482. return (min(soft, inherited), min(hard, inherited))
  483. # ---------------------------------------------------------------------------
  484. # Main.
  485. # ---------------------------------------------------------------------------
  486. async def _run(channel: ProtocolChannel) -> None:
  487. # 1. Boot handshake.
  488. boot = channel.read_frame()
  489. if boot is None or boot.get("type") != "boot":
  490. raise RuntimeError("bootstrap: expected boot frame on fd 3")
  491. # A limit that cannot be applied must fail the run as a diagnosable done
  492. # frame, not a bare traceback + exit(1): running the program UNCAPPED would
  493. # silently void the containment contract, and the host can only relay what
  494. # rides the protocol.
  495. try:
  496. # SIGXCPU's default disposition (how the soft CPU limit stops the child)
  497. # dumps core, and the child inherits the host's RLIMIT_CORE — a CPU
  498. # timeout would otherwise write a large memory-bearing core file into
  499. # the workspace. Forbid core dumps first so the timeout path leaves none.
  500. resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
  501. # Soft limit at cpuSeconds fires SIGXCPU (its default disposition
  502. # terminates the child; the host classifies that close as a timeout).
  503. # Hard limit at +1s is a SIGKILL backstop for a program that traps
  504. # SIGXCPU and keeps burning CPU.
  505. cpu_soft, cpu_hard = _clamped(
  506. resource.RLIMIT_CPU, boot["cpuSeconds"], boot["cpuSeconds"] + 1
  507. )
  508. resource.setrlimit(resource.RLIMIT_CPU, (cpu_soft, cpu_hard))
  509. # Darwin maps the multi-GB dyld shared cache into every process at
  510. # exec, so any practical RLIMIT_AS cap sits below current usage and
  511. # the kernel rejects it — the child would die here on every run. Skip
  512. # the address-space cap there; RLIMIT_CPU and the host's wall-clock
  513. # ceiling still bound the run.
  514. if sys.platform != "darwin":
  515. addr_bytes = int(boot["addressSpaceBytes"])
  516. resource.setrlimit(
  517. resource.RLIMIT_AS, _clamped(resource.RLIMIT_AS, addr_bytes, addr_bytes)
  518. )
  519. except BaseException as exc: # noqa: BLE001 -- report every failure to host
  520. channel.send_sync(
  521. {
  522. "type": "done",
  523. "error": {
  524. "kind": "exception",
  525. # Exception-only rendering: format_exc() would embed the
  526. # absolute installed bootstrap.py path in model-visible
  527. # durable output, leaking host paths into transcripts.
  528. "message": "bootstrap: applying resource limits failed\n"
  529. + "".join(
  530. traceback.format_exception_only(type(exc), exc)
  531. ),
  532. },
  533. }
  534. )
  535. return
  536. logs = LogBuffer(
  537. int(boot["maxLogBytes"]),
  538. sink=lambda text, truncated=False: channel.send_sync(
  539. {"type": "log", "text": text, **({"truncated": True} if truncated else {})}
  540. ),
  541. )
  542. # 2. Wire the tools proxies and the ack.
  543. pending: dict[int, asyncio.Future[Any]] = {}
  544. next_id = 0
  545. error_classes: dict[str, type] = {}
  546. async def dispatch(global_name: str, name: str, args: Any) -> Any:
  547. nonlocal next_id
  548. error_class = error_classes.get(global_name)
  549. def call_failure(message: str) -> BaseException:
  550. # The namespace's declared rejection contract (e.g. Code Mode's
  551. # ToolCallError with .toolName) when present; RuntimeError keeps
  552. # the pre-errorClass behavior for namespaces that declared none.
  553. if error_class is not None:
  554. return error_class(name, message)
  555. return RuntimeError(message)
  556. # Validate the argument shape before claiming an id, so a rejected call
  557. # leaves no gap in the sequence the host checks. json.dumps would coerce
  558. # a non-string dict key or non-finite float rather than raise (allow_nan
  559. # is off, but key coercion still slips through), silently corrupting what
  560. # the tool receives. Reject up front through the call's error contract.
  561. violation = _lossless_json_violation(args)
  562. if violation is not None:
  563. raise call_failure(f"binding arguments must be lossless JSON ({violation})")
  564. # Ids are consecutive from 0 with NO gaps: the host answers a `call` only
  565. # when its id is the exact successor of the last one, which bounds the
  566. # state it retains to a single number. A frame that never reaches the
  567. # host must therefore not consume an id, so the counter advances only
  568. # once the write has succeeded.
  569. call_id = next_id
  570. fut: asyncio.Future[Any] = asyncio.get_event_loop().create_future()
  571. pending[call_id] = fut
  572. try:
  573. channel.send_sync(
  574. {
  575. "type": "call",
  576. "id": call_id,
  577. "global": global_name,
  578. "name": name,
  579. "args": args,
  580. }
  581. )
  582. except (TypeError, ValueError) as exc:
  583. pending.pop(call_id, None)
  584. raise call_failure(
  585. f"binding arguments must be lossless JSON: {exc}"
  586. ) from exc
  587. next_id += 1
  588. try:
  589. return await fut
  590. except _BindingRejection as exc:
  591. raise call_failure(str(exc)) from None
  592. namespaces: dict[str, Any] = {}
  593. for entry in boot["namespaces"]:
  594. namespaces[entry["global"]] = _Namespace(
  595. entry["global"], entry["names"], dispatch
  596. )
  597. declared = entry.get("errorClass")
  598. if declared:
  599. error_class = _make_error_class(
  600. declared["name"], declared["memberNameProperty"]
  601. )
  602. error_classes[entry["global"]] = error_class
  603. # The class is program-visible under its own name so model code
  604. # can `except ToolCallError as e:` and read the member property.
  605. namespaces[declared["name"]] = error_class
  606. channel.send_sync({"type": "boot-ack"})
  607. # 3. Start a reply-pump task before the run message: replies can arrive
  608. # interleaved with the run's own binding traffic.
  609. reply_task = asyncio.get_event_loop().create_task(_pump_replies(channel, pending))
  610. # 4. Read the run message.
  611. run = channel.read_frame()
  612. if run is None or run.get("type") != "run":
  613. reply_task.cancel()
  614. raise RuntimeError("bootstrap: expected run frame on fd 3")
  615. program: str = run["program"]
  616. # 5. Install log capture — ``print``, tracebacks, and ordinary ``sys.stdout``
  617. # writes funnel into the LogBuffer. The real fds stay open (host uses
  618. # stderr for stray-byte accounting) but the Python-visible streams point
  619. # at the buffer.
  620. sys.stdout = _LogStream(logs) # type: ignore[assignment]
  621. sys.stderr = _LogStream(logs) # type: ignore[assignment]
  622. out_stream, err_stream = sys.stdout, sys.stderr
  623. # 6. Compile the program as the body of an async function, matching the
  624. # seam contract (`CodeRunRequest.program` is an async-function body: top-level
  625. # `await` and `return` both work, and the returned value is the completion).
  626. # AST-splicing the parsed body into an `async def` keeps every statement's
  627. # original line number, so a traceback points at the model's own source.
  628. ns: dict[str, Any] = {
  629. "__name__": "__main__",
  630. "__builtins__": __builtins__,
  631. **namespaces,
  632. }
  633. # Read the enforcement callable and its budget into this frame's locals
  634. # BEFORE the program runs: model code can rebind this module's globals
  635. # (the bootstrap IS ``__main__``), and a frame local is not a module
  636. # attribute, so a later ``__main__._DIE_IF_CPU_EXHAUSTED = ...`` cannot
  637. # change which callable the post-check below invokes. This defeats the
  638. # one-line rebind, not a determined `sys._getframe` walk; the unforgeable
  639. # bounds are the RLIMIT_CPU hard limit and the host wall clock
  640. # (see _make_cpu_enforcer).
  641. die_if_cpu_exhausted = _DIE_IF_CPU_EXHAUSTED
  642. cpu_seconds = int(boot["cpuSeconds"])
  643. # Same capture, same reason, for the failure path and the send that follows
  644. # it. The reporter was a module-global lookup inside the `except` block, so
  645. # ``import __main__; __main__._SAFE_MODEL_TRACEBACK = ...`` put model code
  646. # there with no guard around it; the flush and send were attribute lookups
  647. # on the stream and channel CLASSES, which ``__main__._LogStream.flush_line
  648. # = ...`` rebinds just as easily. All four run AFTER the handler, where a
  649. # throw costs the `done` frame and the host reports a wall-clock timeout
  650. # instead of the model's exception. Binding the callables now fixes what
  651. # runs; what they in turn reach is closed over in _make_failure_reporter.
  652. safe_model_traceback = _SAFE_MODEL_TRACEBACK
  653. flush_out = out_stream.flush_line
  654. flush_err = err_stream.flush_line
  655. send_done = channel.send_sync
  656. max_value_bytes = int(boot["maxValueBytes"])
  657. done: dict[str, Any]
  658. try:
  659. module = ast.parse(program)
  660. wrapper = ast.AsyncFunctionDef(
  661. name="__dsh_main__",
  662. args=ast.arguments(
  663. posonlyargs=[], args=[], vararg=None,
  664. kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[],
  665. ),
  666. body=module.body or [ast.Pass()],
  667. decorator_list=[],
  668. returns=None,
  669. )
  670. # Anchor the synthetic wrapper on the first real statement (or line 1 for
  671. # an empty program) so fix_missing_locations does not stamp it at 0.
  672. anchor = module.body[0] if module.body else ast.parse("pass").body[0]
  673. ast.copy_location(wrapper, anchor)
  674. wrapped = ast.Module(body=[wrapper], type_ignores=[])
  675. ast.fix_missing_locations(wrapped)
  676. code = compile(wrapped, "<model>", "exec")
  677. exec(code, ns) # noqa: S102 -- defines __dsh_main__; executing model code is the point
  678. value = await ns["__dsh_main__"]()
  679. die_if_cpu_exhausted(cpu_seconds)
  680. done = _done_with_value(value, max_value_bytes)
  681. except BaseException as exc: # noqa: BLE001 -- report every failure to host
  682. done = {
  683. "type": "done",
  684. "error": {
  685. "kind": "exception",
  686. # Cap the diagnostic BEFORE it crosses the wire: a program can
  687. # raise with a gigabytes-long message, and formatting/sending
  688. # it whole would allocate on both sides before the host's own
  689. # cap runs. Byte-cap at maxValueBytes with the host's marker
  690. # text so the truncated diagnostic reads identically wherever
  691. # the cap was applied. The rendering is wrapped because the
  692. # `done` send below sits outside this handler: a throw while
  693. # formatting would skip it and strand the host on fd 3 until
  694. # maxWallMs (see _make_failure_reporter).
  695. "message": safe_model_traceback(exc, max_value_bytes),
  696. },
  697. }
  698. # Flush any print output not terminated by a newline (a traceback always
  699. # ends in one, but `print(x, end="")` or a bare write may not), so the
  700. # final partial line is not silently dropped.
  701. flush_out()
  702. flush_err()
  703. reply_task.cancel()
  704. send_done(done)
  705. async def _pump_replies(
  706. channel: ProtocolChannel, pending: dict[int, asyncio.Future[Any]]
  707. ) -> None:
  708. """Background task: read reply frames and settle pending futures.
  709. Cancelled after ``done`` is posted. Unknown ids and post-settlement replies
  710. are ignored (mirrors the worker backend's hostile-peer stance, though here
  711. the host is the trusted side; the guards defend against races).
  712. """
  713. while True:
  714. frame = await channel.read_frame_async()
  715. if frame is None:
  716. return
  717. if frame.get("type") != "reply":
  718. continue
  719. fut = pending.pop(frame.get("id"), None)
  720. if fut is None or fut.done():
  721. continue
  722. if frame.get("ok"):
  723. fut.set_result(frame.get("value"))
  724. else:
  725. fut.set_exception(_BindingRejection(str(frame.get("message"))))
  726. _SCALAR_RE = re.compile(
  727. r'"(?:[^"\\]|\\.)*"|-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?|true|false|null'
  728. )
  729. def _decode_json_plain(text: str) -> Any:
  730. """Parse one JSON document iteratively (no per-level recursion).
  731. ``json.loads`` recurses per nesting level and raises ``RecursionError``
  732. around ~10k levels, but a binding reply is depth-unbounded by the seam
  733. contract — the host's iterative encoder happily produces documents
  734. ``json.loads`` cannot read back. Scalars (numbers, strings with escapes)
  735. are delegated to ``json.loads`` one token at a time, so their grammar and
  736. semantics stay CPython's own; only the container structure is parsed here
  737. with an explicit stack. Raises ``ValueError`` on malformed input; frames
  738. come from the TRUSTED host, so strictness mirrors ``json.loads`` without
  739. extra hostile-input hardening.
  740. """
  741. length = len(text)
  742. def skip_ws(i: int) -> int:
  743. while i < length and text[i] in " \t\n\r":
  744. i += 1
  745. return i
  746. def scalar(i: int):
  747. match = _SCALAR_RE.match(text, i)
  748. if match is None:
  749. raise ValueError(f"invalid JSON at offset {i}")
  750. return json.loads(match.group(0)), match.end()
  751. def string_key(i: int):
  752. key, end = scalar(i)
  753. if not isinstance(key, str):
  754. raise ValueError(f"object key must be a string at offset {i}")
  755. end = skip_ws(end)
  756. if end >= length or text[end] != ":":
  757. raise ValueError(f"expected ':' at offset {end}")
  758. return key, skip_ws(end + 1)
  759. # Frames: a list, or (dict, pending key). `value`/`have_value` carry each
  760. # completed value up to its parent frame.
  761. stack: list[Any] = []
  762. value: Any = None
  763. have_value = False
  764. i = skip_ws(0)
  765. while True:
  766. if not have_value:
  767. ch = text[i] if i < length else ""
  768. if ch == "[":
  769. i = skip_ws(i + 1)
  770. if i < length and text[i] == "]":
  771. i += 1
  772. value, have_value = [], True
  773. else:
  774. stack.append([])
  775. continue
  776. elif ch == "{":
  777. i = skip_ws(i + 1)
  778. if i < length and text[i] == "}":
  779. i += 1
  780. value, have_value = {}, True
  781. else:
  782. key, i = string_key(i)
  783. stack.append(({}, key))
  784. continue
  785. else:
  786. value, i = scalar(i)
  787. have_value = True
  788. if not stack:
  789. i = skip_ws(i)
  790. if i != length:
  791. raise ValueError(f"trailing data at offset {i}")
  792. return value
  793. top = stack[-1]
  794. i = skip_ws(i)
  795. ch = text[i] if i < length else ""
  796. if isinstance(top, list):
  797. top.append(value)
  798. if ch == ",":
  799. i = skip_ws(i + 1)
  800. have_value = False
  801. elif ch == "]":
  802. i += 1
  803. stack.pop()
  804. value = top
  805. else:
  806. raise ValueError(f"expected ',' or ']' at offset {i}")
  807. else:
  808. container, key = top
  809. container[key] = value
  810. if ch == ",":
  811. key, i = string_key(skip_ws(i + 1))
  812. stack[-1] = (container, key)
  813. have_value = False
  814. elif ch == "}":
  815. i += 1
  816. stack.pop()
  817. value = container
  818. else:
  819. raise ValueError(f"expected ',' or '}}' at offset {i}")
  820. def _encode_json_plain(value: Any) -> str:
  821. """Encode JSON-plain data iteratively, byte-identical to compact ``json.dumps``.
  822. ``json.dumps`` recurses one Python frame per nesting level and raises
  823. ``RecursionError`` a few thousand levels deep, but the seam's
  824. ``CodeJsonValue`` has no depth limit — a valid deeply nested completion or
  825. call argument below the byte budget must cross intact (the host uses the
  826. same iterative idiom in ``protocol.ts``). Accepts what the callers already
  827. validated or constructed: ``None``/``bool``/``int``/finite ``float``/
  828. ``str``, exact ``list``/``tuple``, and exact ``dict`` with ``str`` keys.
  829. Scalar encoding delegates to ``json.dumps`` (string escaping, float repr)
  830. so the bytes match; non-finite floats still raise ``ValueError`` exactly
  831. like ``allow_nan=False``.
  832. Containers are classified by EXACT type and traversed through the unbound
  833. built-in methods rather than the instance's own: a ``dict``/``list``
  834. subclass can override ``items``, ``keys``, ``__iter__``, ``__len__``, or
  835. ``__getitem__``, and the validators only see the container it subclasses,
  836. so an instance-method call here could emit different data than the walk
  837. that metered and approved it. ``_check_done_value`` and
  838. ``_lossless_json_violation`` reject subclasses outright, so this path only
  839. ever sees exact containers; classifying on exact type keeps that agreement
  840. checkable at one glance instead of resting on the caller.
  841. """
  842. chunks: list[str] = []
  843. # Each frame is either a literal string to emit or a value to expand.
  844. stack: list[Any] = [value]
  845. while stack:
  846. current = stack.pop()
  847. current_type = type(current)
  848. if current_type is _Emit:
  849. chunks.append(current.text)
  850. elif current_type is list or current_type is tuple:
  851. count = len(current)
  852. chunks.append("[")
  853. stack.append(_Emit("]"))
  854. for index in range(count - 1, -1, -1):
  855. if index < count - 1:
  856. stack.append(_Emit(","))
  857. stack.append(current[index])
  858. elif current_type is dict:
  859. chunks.append("{")
  860. stack.append(_Emit("}"))
  861. items = list(dict.items(current))
  862. for index in range(len(items) - 1, -1, -1):
  863. key, item = items[index]
  864. if index < len(items) - 1:
  865. stack.append(_Emit(","))
  866. stack.append(item)
  867. stack.append(_Emit(_dump_scalar(key) + ":"))
  868. else:
  869. chunks.append(_dump_scalar(current))
  870. return "".join(chunks)
  871. def _dump_scalar(value: Any) -> str:
  872. """One scalar as compact JSON, byte-compatible with the host's encoder.
  873. ``ensure_ascii=False`` keeps non-ASCII text as raw UTF-8 — the default
  874. backslash-u escaping would make the child count ``"é"`` as 8 bytes where
  875. the host meter (and the worker backend) count its UTF-8 JSON form as 4,
  876. splitting the budget the two sides are supposed to share. Strings route
  877. through :func:`_dump_string`, which restores the escaping for the one class
  878. of character UTF-8 cannot hold. Floats route through :func:`_dump_float`
  879. because CPython's ``repr`` and ECMAScript's Number-to-String disagree on
  880. spelling.
  881. Dispatch is on EXACT type, matching the validators: a ``float`` subclass
  882. reaching :func:`_dump_float` would have its overridden ``__repr__`` read as
  883. the number's digits, so ``F(2.5)`` whose ``__repr__`` returns ``"1.0"``
  884. would serialize as ``1``. ``json.dumps`` then refuses any subclass by
  885. ``TypeError`` instead of emitting a value nothing validated; the callers
  886. reject subclasses first, so this is the encoder refusing to be the place a
  887. validation gap turns into corrupted output.
  888. """
  889. if type(value) is float:
  890. return _dump_float(value)
  891. if type(value) is str:
  892. return _dump_string(value)
  893. if value is None or type(value) is bool or type(value) is int:
  894. return json.dumps(value, ensure_ascii=False, allow_nan=False)
  895. raise TypeError(f"unsupported type ({type(value).__name__})")
  896. # A surrogate code unit, and an adjacent high-low pair. Python stores an astral
  897. # character as ONE code point, so a surrogate reaching these patterns is either
  898. # lone or half of a pair the program spelled out code unit by code unit.
  899. _SURROGATE = re.compile("[\ud800-\udfff]")
  900. _SURROGATE_PAIR = re.compile("[\ud800-\udbff][\udc00-\udfff]")
  901. def _combine_surrogate_pair(match: re.Match[str]) -> str:
  902. """Fold one spelled-out high-low pair into the astral code point it names."""
  903. high, low = match.group(0)
  904. return chr(0x10000 + ((ord(high) - 0xD800) << 10) + (ord(low) - 0xDC00))
  905. def _dump_string(text: str) -> str:
  906. """One string as compact JSON, byte-identical to the host's ``JSON.stringify``.
  907. ``ensure_ascii=False`` cannot render a surrogate code unit: UTF-8 has no
  908. encoding for one, so the frame write would raise and the run would strand
  909. until the wall clock. JSON carries it as the ASCII escape ``\\ud800``, which
  910. the host's ``JSON.parse`` reads back as the same UTF-16 code unit and its
  911. ``JSON.stringify`` re-emits identically — so the shared seam
  912. (``CodeJsonValue``, ``snapshotJsonValue``, the worker backend) keeps a
  913. lone-surrogate string instead of failing the value. An adjacent high-low
  914. pair is folded into its astral code point FIRST: the host holds strings as
  915. UTF-16, where those two code units and the single character are the same
  916. string, and the raw 4-byte form is what the host would emit — escaping the
  917. halves separately would charge 12 bytes against a budget the host meters at
  918. 4. Every remaining surrogate is lone and becomes six ASCII bytes, matching
  919. the host exactly.
  920. @param text: the string to encode.
  921. @return: its compact JSON form, always UTF-8-encodable.
  922. """
  923. rendered = json.dumps(text, ensure_ascii=False)
  924. if _SURROGATE.search(rendered) is None:
  925. return rendered
  926. return _SURROGATE.sub(
  927. lambda match: "\\u%04x" % ord(match.group(0)),
  928. _SURROGATE_PAIR.sub(_combine_surrogate_pair, rendered),
  929. )
  930. # How many bytes each byte that needs escaping adds beyond its raw self, as a
  931. # ready-made (byte, surcharge) list so :func:`_json_string_cost` walks no
  932. # branches per pass. ``"`` and ``\\`` take a one-character prefix; the five C0
  933. # controls with a shorthand (``\\b\\f\\n\\r\\t``) likewise; every other C0
  934. # control becomes a six-character ``\\uXXXX``.
  935. _JSON_ESCAPE_SURCHARGES = [
  936. (bytes((byte,)), 1 if byte in b'"\\\b\f\n\r\t' else 5)
  937. for byte in [*range(0x20), ord('"'), ord("\\")]
  938. ]
  939. def _json_string_cost(raw: bytes) -> int:
  940. """UTF-8 byte length of one string's JSON form, WITHOUT building that form.
  941. Used by :class:`LogBuffer` to charge a log entry what it will actually cost
  942. on the wire. Building ``json.dumps(text)`` to measure it would allocate a
  943. second copy up to six times the original — the very allocation the ledger's
  944. cheap pre-check exists to avoid, and enough to breach ``RLIMIT_AS`` on a
  945. large control-heavy line. Counts exactly what :func:`_dump_scalar`'s
  946. ``ensure_ascii=False`` output holds: the two quotes, each escaped byte's
  947. surcharge from :data:`_JSON_ESCAPE_SURCHARGES`, and the raw bytes themselves
  948. (non-ASCII stays raw, so its UTF-8 length already counts). Uses a fixed
  949. number of C-level ``count`` passes — allocating nothing, unlike a
  950. ``translate`` filter — because the caller admits up to ~4x the remaining
  951. budget of bytes here and a per-byte Python loop over it would cost more than
  952. the encode being avoided.
  953. @param raw: the entry's UTF-8 bytes.
  954. @return: the byte length of its JSON string form, quotes included.
  955. """
  956. extra = 0
  957. for byte, surcharge in _JSON_ESCAPE_SURCHARGES:
  958. extra += raw.count(byte) * surcharge
  959. return len(raw) + 2 + extra
  960. def _dump_float(value: float) -> str:
  961. """One finite float in ECMAScript ``Number::toString`` spelling.
  962. CPython's ``repr`` and the host's ``String(number)`` name the same double
  963. differently: ``1.0`` is ``"1.0"`` here but ``"1"`` there, ``1e-07`` pads the
  964. exponent the host writes as ``1e-7``, and ``1e+21``/``2**60`` differ again.
  965. Since the child meters the completion value against ``maxValueBytes`` and
  966. the host re-meters the frame it parses, any spelling difference splits the
  967. shared budget: ``return 1.0`` under ``maxValueBytes: 1`` used to be reported
  968. as ``output-limit`` by the child while the host would have counted the
  969. one-byte ``1`` it actually receives. Both sides also emit these bytes (the
  970. child through :func:`_encode_json_plain`, the host through
  971. ``encodeJsonPlain``), so the fix has to be in the shared speller, not in the
  972. meter.
  973. Implements ECMA-262 ``Number::toString`` radix 10 directly: ``repr``
  974. already yields the shortest round-tripping decimal digits, and ``Decimal``
  975. splits them into the significand ``s`` (``digits``, ``k`` of them) and
  976. decimal exponent ``n`` the spec's cases select on. The integral values above
  977. the JS safe range take the host's BigInt branch, whose exact digits differ
  978. from the shortest-round-trip form (``2**60`` prints ``...846976``, not
  979. ``...847000``).
  980. """
  981. if value != value or value in (float("inf"), float("-inf")):
  982. # json.dumps(allow_nan=False) raises the same way; the callers reject
  983. # non-finite floats before metering, so this is unreachable defense.
  984. raise ValueError("Out of range float values are not JSON compliant")
  985. if value == 0.0:
  986. # Covers -0.0 too; callers reject it as non-lossless before this point.
  987. return "0"
  988. if value < 0:
  989. return "-" + _dump_float(-value)
  990. if value.is_integer() and value > float(2**53 - 1):
  991. # The host's BigInt branch: exact digits, not shortest-round-trip.
  992. return str(int(value))
  993. parts = Decimal(repr(value)).normalize().as_tuple()
  994. digits = "".join(str(digit) for digit in parts.digits)
  995. k = len(digits)
  996. n = parts.exponent + k
  997. if k <= n <= 21:
  998. return digits + "0" * (n - k)
  999. if 0 < n <= 21:
  1000. return digits[:n] + "." + digits[n:]
  1001. if -6 < n <= 0:
  1002. return "0." + "0" * -n + digits
  1003. exponent = ("+" if n - 1 >= 0 else "-") + str(abs(n - 1))
  1004. return (digits if k == 1 else digits[0] + "." + digits[1:]) + "e" + exponent
  1005. def _check_done_value(value: Any, max_bytes: int):
  1006. """Meter a completion value's JSON byte size AND validate its lossless-JSON
  1007. shape in one bounded post-order walk; return ``None`` when it passes.
  1008. Folds what was formerly a losslessness walk followed by a separate byte
  1009. meter into one pass. Running the losslessness walk first materialized one
  1010. traversal tuple per element before any size cap: ``return [0] * 2000000``
  1011. under ``maxValueBytes: 64`` allocated millions of frames (an RLIMIT_AS
  1012. death) before the meter could reject it. Folding the byte bound into the
  1013. walk rejects over-budget BEFORE enqueuing a container's children — every
  1014. element is at least one JSON byte — so the walk stays O(cap). Same
  1015. JS-double-exact integer boundary, cycle detection (a leave marker pops each
  1016. container off ``on_path``), and type rejections as
  1017. :func:`_lossless_json_violation`, and the same byte accounting as
  1018. :func:`_encode_json_plain`. :func:`_lossless_json_violation` stays for the
  1019. binding-argument path, which carries no size cap.
  1020. EVERY type here is matched EXACTLY, containers and scalars alike, so a
  1021. subclass is rejected as an unsupported type rather than admitted by
  1022. ``isinstance``. A subclass can override the operators and methods this walk
  1023. and the encoder call, and they need not agree: a populated ``dict``
  1024. subclass whose ``items()`` returns ``[]`` would meter as ``{}``; a ``float``
  1025. subclass overriding ``__repr__`` passes the non-finite and negative-zero
  1026. checks by its real value but serializes as whatever the override says, since
  1027. :func:`_dump_float` reads ``repr``; an ``int`` subclass overriding ``__gt__``
  1028. and ``__lt__`` slips past the JS-safe-range bound while ``json.dumps``
  1029. emits its true C-level digits, so ``2**53 + 1`` reaches the host as
  1030. ``...992``; a ``str`` subclass overriding ``__len__`` returns 0 from the
  1031. pre-encode lower bound and admits an arbitrarily large string. In each case
  1032. the value the host receives differs from the one this walk approved. The
  1033. worker backend rejects the equivalent shapes by prototype identity and
  1034. ``typeof`` (``hasPlainObjectPrototype`` in ``worker-json.ts``); a ``bool``
  1035. is checked before ``int`` because it is an ``int`` subclass that IS
  1036. lossless JSON.
  1037. Returns ``("invalid-output", message)`` for a non-lossless value,
  1038. ``("output-limit", message)`` once the size crosses ``max_bytes``, or
  1039. ``None`` when the value is lossless JSON within budget.
  1040. """
  1041. js_safe = 2**53 - 1
  1042. def invalid(reason: str):
  1043. return ("invalid-output", f"program completion must be lossless JSON ({reason})")
  1044. over_budget = ("output-limit", f"completion value exceeded {max_bytes} bytes")
  1045. total = 0
  1046. on_path: set[int] = set()
  1047. # Each frame is (value, is_leave): a leave frame pops its container off the path.
  1048. stack: list[tuple[Any, bool]] = [(value, False)]
  1049. while stack:
  1050. current, is_leave = stack.pop()
  1051. if is_leave:
  1052. on_path.discard(id(current))
  1053. continue
  1054. if current is None or type(current) is bool:
  1055. total += len(_dump_scalar(current).encode("utf-8"))
  1056. elif type(current) is str:
  1057. # Lower-bound BEFORE materializing the escaped form: every character
  1058. # is at least one UTF-8 byte plus the two quotes, so a huge or
  1059. # control-heavy string (whose escaped copy expands severalfold) is
  1060. # rejected without allocating that copy.
  1061. if total + len(current) + 2 > max_bytes:
  1062. return over_budget
  1063. # A lone surrogate has no UTF-8 form but a lossless JSON one — the
  1064. # ASCII ``\uXXXX`` escape :func:`_dump_string` emits — so it is
  1065. # metered, not rejected, matching the shared seam.
  1066. total += len(_dump_string(current).encode("utf-8"))
  1067. elif type(current) is int:
  1068. # The canonical boundary accepts every JS-double-exact value: an int
  1069. # outside +-2**53-1 is fine IFF the double round-trip is exact.
  1070. if current > js_safe or current < -js_safe:
  1071. try:
  1072. exact = int(float(current)) == current
  1073. except OverflowError:
  1074. exact = False
  1075. if not exact:
  1076. return invalid("integer not exactly representable as a JavaScript number")
  1077. total += len(_dump_scalar(current).encode("utf-8"))
  1078. elif type(current) is float:
  1079. if current != current or current in (float("inf"), float("-inf")):
  1080. return invalid("non-finite float")
  1081. # JSON turns -0.0 into a sign the host parses back to JS -0; the
  1082. # canonical boundary rejects it, so this side must too.
  1083. if current == 0.0 and math.copysign(1.0, current) < 0:
  1084. return invalid("negative zero")
  1085. total += len(_dump_scalar(current).encode("utf-8"))
  1086. elif type(current) is list:
  1087. if id(current) in on_path:
  1088. return invalid("circular reference")
  1089. count = len(current)
  1090. total += 2 + (count - 1 if count > 1 else 0)
  1091. # Reject over-budget BEFORE enqueuing children: every element
  1092. # serializes to at least one byte, so a wide flat forgery fails here
  1093. # without materializing millions of leave frames first.
  1094. if total + count > max_bytes:
  1095. return over_budget
  1096. on_path.add(id(current))
  1097. stack.append((current, True))
  1098. stack.extend((child, False) for child in current)
  1099. elif type(current) is dict:
  1100. if id(current) in on_path:
  1101. return invalid("circular reference")
  1102. # ``len`` without materializing ``current.items()``: that list
  1103. # allocates one tuple per member before the bound below could run,
  1104. # recreating the spike the bound exists to stop.
  1105. count = len(current)
  1106. total += 2 + (count - 1 if count > 1 else 0)
  1107. # Same pre-enqueue bound: each entry contributes a quoted key
  1108. # (>= 2 bytes), a colon, and a >= 1-byte value.
  1109. if total + count * 4 > max_bytes:
  1110. return over_budget
  1111. on_path.add(id(current))
  1112. stack.append((current, True))
  1113. for key, item in current.items():
  1114. # Only an EXACT str key survives: bool and int coerce or raise,
  1115. # and a str SUBCLASS can override the ``__len__`` the bound
  1116. # below reads while the encoder emits its real characters.
  1117. if type(key) is not str:
  1118. return invalid(f"non-string dict key ({type(key).__name__})")
  1119. # The same string lower bound, before escaping the key.
  1120. if total + len(key) + 3 > max_bytes:
  1121. return over_budget
  1122. total += len(_dump_scalar(key).encode("utf-8")) + 1
  1123. stack.append((item, False))
  1124. else:
  1125. # tuple, set, or any other type: not round-trippable JSON.
  1126. return invalid(f"unsupported type ({type(current).__name__})")
  1127. if total > max_bytes:
  1128. return over_budget
  1129. return None
  1130. class _Emit:
  1131. """A pre-rendered fragment on :func:`_encode_json_plain`'s explicit stack."""
  1132. __slots__ = ("text",)
  1133. def __init__(self, text: str) -> None:
  1134. self.text = text
  1135. def _lossless_json_violation(value: Any) -> str | None:
  1136. """Return why ``value`` is not lossless JSON, or ``None`` when it is.
  1137. ``json.dumps`` succeeding is NOT proof of losslessness: it coerces a
  1138. non-string ``dict`` key to its string form (``{1: "a", "1": "b"}`` collapses
  1139. to one key, silently dropping data), emits non-standard ``NaN``/``Infinity``
  1140. tokens without ``allow_nan=False``, and accepts integers outside JavaScript's
  1141. safe range (``9007199254740993`` becomes ``...992`` once the host parses the
  1142. frame into a JS number). Validate the shape up front so a coercive or lossy
  1143. value fails as ``invalid-output`` instead of round-tripping to something the
  1144. program did not compute. Iterative so deep nesting cannot overflow the stack,
  1145. and it tracks the container ancestry on the current path so a cyclic value is
  1146. reported at once rather than spinning until the CPU budget. Only JSON-plain
  1147. types survive: ``None``/``bool``/JS-safe ``int``/finite ``float``/``str``,
  1148. exact ``list``, and exact ``dict`` with ``str`` keys. Every type matches
  1149. EXACTLY, containers and scalars alike, for the reason
  1150. :func:`_check_done_value` documents: a subclass can override the operators
  1151. and methods a traversal calls, so an ``isinstance`` admission here would
  1152. approve one shape and let the encoder emit another.
  1153. """
  1154. # The canonical boundary accepts every JS-double-exact value: an int
  1155. # outside +-2**53-1 is fine IFF the double round-trip is exact (2**53 or
  1156. # 2**60 survive; 2**53+1 rounds), matching the worker backend.
  1157. js_safe = 2**53 - 1
  1158. # Post-order walk with an explicit "leave" marker: a container's id is added
  1159. # to `on_path` when entered and removed when left, so a back-edge to an
  1160. # ancestor (a cycle) is detected without rejecting a legitimately shared
  1161. # acyclic subtree.
  1162. on_path: set[int] = set()
  1163. # Each frame is (value, is_leave): a leave frame pops its container off the path.
  1164. stack: list[tuple[Any, bool]] = [(value, False)]
  1165. while stack:
  1166. current, is_leave = stack.pop()
  1167. if is_leave:
  1168. on_path.discard(id(current))
  1169. continue
  1170. if current is None or type(current) is bool:
  1171. continue
  1172. if type(current) is str:
  1173. # Every string is lossless JSON. A lone surrogate has no UTF-8 form,
  1174. # but JSON carries the code unit as its ASCII ``\uXXXX`` escape and
  1175. # :func:`_dump_string` emits exactly that, so the host receives the
  1176. # same code unit the program passed — the same acceptance
  1177. # ``CodeJsonValue``, ``snapshotJsonValue``, and the worker backend
  1178. # already give it.
  1179. continue
  1180. if type(current) is int:
  1181. if current > js_safe or current < -js_safe:
  1182. try:
  1183. exact = int(float(current)) == current
  1184. except OverflowError:
  1185. exact = False
  1186. if not exact:
  1187. return "integer not exactly representable as a JavaScript number"
  1188. continue
  1189. if type(current) is float:
  1190. if current != current or current in (float("inf"), float("-inf")):
  1191. return "non-finite float"
  1192. # JSON serialization turns -0.0 into 0 (or "-0.0" text that the
  1193. # host parses to JS -0), silently changing the sign bit either
  1194. # way; the repository's canonical lossless-JSON boundary and the
  1195. # worker backend both reject it, so this side must too.
  1196. if current == 0.0 and math.copysign(1.0, current) < 0:
  1197. return "negative zero"
  1198. continue
  1199. if type(current) is list or type(current) is dict:
  1200. if id(current) in on_path:
  1201. return "circular reference"
  1202. on_path.add(id(current))
  1203. stack.append((current, True))
  1204. if type(current) is dict:
  1205. for key in current:
  1206. # Only an EXACT str key survives: int, float, None, and
  1207. # tuple keys coerce or raise, and a str subclass can carry
  1208. # overrides the encoder does not honor.
  1209. if type(key) is not str:
  1210. return f"non-string dict key ({type(key).__name__})"
  1211. stack.extend((child, False) for child in current.values())
  1212. else:
  1213. stack.extend((child, False) for child in current)
  1214. continue
  1215. return f"unsupported type ({type(current).__name__})"
  1216. return None
  1217. def _make_cpu_enforcer() -> Any:
  1218. """Build the CPU post-check over closure-held primitives.
  1219. This bootstrap IS ``__main__``, so model code can reach every one of its
  1220. module globals: ``import __main__; __main__._X = ...`` rebinds the name the
  1221. enforcement would otherwise read at call time, which a plain module-level
  1222. function plus module-level captures made a one-line defeat. The primitives
  1223. therefore live in this factory's locals, which become closure cells of the
  1224. returned function, and :func:`_run` binds the returned function into a
  1225. local of its own frame BEFORE executing the program, so no assignment to
  1226. ``__main__`` changes which callable runs or what it calls. Capture happens
  1227. at import time, before model code runs, so the captured
  1228. ``resource.getrusage``/``signal.signal``/``os.kill`` are the real builtins.
  1229. This raises the cost of defeating the check; it does not make it
  1230. unreachable, and nothing in-process could. A cell is writable through
  1231. ``fn.__closure__[i].cell_contents``, and ``sys._getframe`` walks to
  1232. :func:`_run`'s frame and reads its locals, so a program determined to
  1233. tamper still can — consistent with this backend's documented posture, where
  1234. the in-process interpreter is containment rather than a security boundary
  1235. (§Trust posture in the Code Mode RFC). The bounds that model code cannot
  1236. forge are outside the interpreter: the RLIMIT_CPU HARD limit at
  1237. ``cpuSeconds + 1``, whose SIGKILL is undeliverable to a handler and
  1238. unraisable by a process that cannot raise its own hard limit, and the
  1239. host's wall-clock ceiling. This check exists to convert the two cases those
  1240. miss — a program that traps SIGXCPU and settles inside the soft-to-hard
  1241. gap, and a program that spends the budget in DESCENDANTS the kernel never
  1242. charged to this process — from a reported SUCCESS into the same `timeout`
  1243. an untrapped program gets.
  1244. @returns The one-argument enforcement callable, taking `cpuSeconds`.
  1245. """
  1246. getrusage = resource.getrusage
  1247. rusage_self = resource.RUSAGE_SELF
  1248. rusage_children = resource.RUSAGE_CHILDREN
  1249. set_signal = signal.signal
  1250. sig_dfl = signal.SIG_DFL
  1251. sigxcpu = signal.SIGXCPU
  1252. kill = os.kill
  1253. getpid = os.getpid
  1254. def die_if_cpu_exhausted(cpu_seconds: int) -> None:
  1255. """Die by re-delivered SIGXCPU when the CPU budget is already spent.
  1256. Two cases reach here as a would-be SUCCESS. A model program can trap
  1257. SIGXCPU and return during the one-second soft-to-hard gap. And
  1258. ``RLIMIT_CPU`` is PER-PROCESS, inherited fresh by every child, so a
  1259. program calling ``subprocess`` or ``os.fork`` multiplies the run's CPU
  1260. budget by the number of descendants it starts: measured with
  1261. ``cpuSeconds: 1``, two sequential busy children burned 2.0
  1262. CPU-seconds and the parent, which had accrued almost no CPU of its own
  1263. while blocked in ``subprocess.wait``, still returned a completion.
  1264. The meter is therefore ``RUSAGE_SELF + RUSAGE_CHILDREN``, the kernel's
  1265. own aggregate, which accumulates the CPU of every REAPED descendant
  1266. (grandchildren included, verified).
  1267. ``getrusage`` is the kernel's own meter (unforgeable from model code),
  1268. and dying by SIGXCPU with the default disposition restored gives the
  1269. host the same kernel-authoritative close signal as the untrapped soft
  1270. limit — classified as `timeout`, after which the host's process-group
  1271. SIGTERM/SIGKILL teardown reaches any surviving descendants. Runs AFTER
  1272. the model program settled, so a program can re-trap SIGXCPU between
  1273. this SIG_DFL and the kill only by running more code, which it no longer
  1274. does. A program that tampers with this callable instead (see
  1275. :func:`_make_cpu_enforcer` on why in-process state cannot be hidden)
  1276. buys at most the remaining soft-to-hard gap: one more CPU second, after
  1277. which the hard limit's SIGKILL lands with no handler possible.
  1278. Checking at settle time rather than sampling mid-run is deliberate:
  1279. both mid-run designs perturb the run they measure. A sampling thread
  1280. cost 72 MiB of virtual address space in the child (8 MiB stack plus a
  1281. 64 MiB glibc per-thread malloc arena reservation; measured 30.23 MiB of
  1282. mappings without it against 102.37 MiB with it), and ``RLIMIT_AS``
  1283. counts reserved space, so it silently shrank every run's
  1284. `addressSpaceMb`. A ``SIGALRM`` interval timer costs no mappings but
  1285. makes the program's own syscalls return short under PEP 475 — measured
  1286. a 64 MiB ``os.write`` returning 65536 — which corrupts fd-3 framing.
  1287. The cost of checking only at settle time is that a descendant's CPU is
  1288. detected after it is spent, not while it runs; the host's wall-clock
  1289. ceiling bounds that interval, and a program that never reaps its child
  1290. is bounded by the wall clock alone, since ``RUSAGE_CHILDREN`` counts
  1291. only reaped descendants (verified: a still-running child contributes
  1292. 0.0).
  1293. @param cpu_seconds The `cpuSeconds` budget the soft RLIMIT_CPU used.
  1294. """
  1295. own = getrusage(rusage_self)
  1296. kids = getrusage(rusage_children)
  1297. spent = own.ru_utime + own.ru_stime + kids.ru_utime + kids.ru_stime
  1298. if spent >= cpu_seconds:
  1299. set_signal(sigxcpu, sig_dfl)
  1300. kill(getpid(), sigxcpu)
  1301. return die_if_cpu_exhausted
  1302. _DIE_IF_CPU_EXHAUSTED = _make_cpu_enforcer()
  1303. _TRUNCATION_MARKER = "… [truncated]"
  1304. # The marker's own UTF-8 size, reserved out of the cap rather than added on top
  1305. # of it. Byte-identical to the host's TRUNCATION_MARKER_BYTES; the ellipsis is
  1306. # three bytes, so this is 15, not the string's 13 characters.
  1307. _TRUNCATION_MARKER_BYTES = len(_TRUNCATION_MARKER.encode("utf-8"))
  1308. def _cap_message(message: str, max_bytes: int) -> str:
  1309. """Byte-cap a diagnostic, appending the same marker the host uses.
  1310. Encoded with ``errors="replace"`` first: a model exception message can
  1311. contain an unpaired surrogate (``raise Exception("\\ud800")``), and a
  1312. strict encode would throw while BUILDING the failure frame — the run
  1313. would then strand until the wall clock instead of reporting the
  1314. exception. Then a UTF-8 slice with a trailing partial sequence dropped
  1315. by ``errors="ignore"``; the marker text matches the host-side
  1316. ``capMessage`` so a truncated diagnostic reads identically wherever the
  1317. cap was applied.
  1318. The marker's bytes come OUT of ``max_bytes``, so the returned string as a
  1319. whole honors the cap; retaining a full cap of text and then appending the
  1320. marker would exceed the bound this function enforces, and the host meters
  1321. the same field again on arrival. A ``max_bytes`` below the marker's own
  1322. size leaves no room for message text and yields the marker alone, so the
  1323. true bound is ``max(max_bytes, 15)`` — reporting that truncation happened
  1324. is worth those 15 bytes.
  1325. """
  1326. raw = message.encode("utf-8", errors="replace")
  1327. if len(raw) <= max_bytes:
  1328. return raw.decode("utf-8")
  1329. budget = max(0, max_bytes - _TRUNCATION_MARKER_BYTES)
  1330. return raw[:budget].decode("utf-8", errors="ignore") + _TRUNCATION_MARKER
  1331. # Fixed safety/liveness bound, not a tunable: a model can raise an exception
  1332. # with an arbitrarily deep __cause__/__context__ chain, and both the rendering
  1333. # walk and format() are linear in chain length. Capping how many links get
  1334. # RENDERED keeps traceback formatting from consuming the whole wall budget.
  1335. # 100 links is far beyond any legible human traceback.
  1336. _MAX_TRACEBACK_CHAIN = 100
  1337. # Diagnostic used when rendering the failure itself fails. Built from a fixed
  1338. # literal plus the exception CLASS name, never from the exception's own str.
  1339. _UNRENDERABLE_DIAGNOSTIC = "<diagnostic rendering failed>"
  1340. def _model_traceback(exc: BaseException, max_bytes: int) -> str:
  1341. """Format a model-program failure with only the MODEL's own frames.
  1342. Bootstrap frames carry host-absolute paths — meaningless to the model and
  1343. unstable across machines, so transcripts pinning them cannot replay. They
  1344. appear not only as a leading prefix (the bootstrap's ``exec``/``await``)
  1345. but also interleaved and trailing: an uncaught binding rejection re-raised
  1346. by ``dispatch`` puts bootstrap frames AFTER the model's, and chained
  1347. ``__cause__``/``__context__`` exceptions carry their own stacks. Filter
  1348. every non-``<model>`` frame across the whole chain rather than trimming a
  1349. prefix. A failure with no model frame anywhere (e.g. a SyntaxError raised
  1350. by ``compile``) keeps the standard exception-only rendering.
  1351. Rendering is bounded to ``_MAX_TRACEBACK_CHAIN`` links, cut on the
  1352. ``TracebackException`` COPY, and a marker line announces the truncation.
  1353. Nothing here touches the live exception: an exception class overriding
  1354. ``__setattr__`` would run MODEL code from inside the caller's failure
  1355. handler, and a throw there costs the ``done`` frame (see
  1356. ``_safe_model_traceback``). ``TracebackException`` instances hold no such
  1357. hooks, so clearing their links runs no model code. The walk is iterative,
  1358. so a deep chain cannot overflow the recursion limit.
  1359. ``from_exception`` still copies the WHOLE live chain, at a higher per-link
  1360. cost than building it took. That is bounded by the child's ``RLIMIT_AS``:
  1361. the model must materialize every link (exception object plus traceback)
  1362. before raising, so a chain long enough for the copy to matter is already
  1363. near the address-space cap, and a ``MemoryError`` in the copy lands in the
  1364. caller's fallback rather than stranding the run.
  1365. """
  1366. te = traceback.TracebackException.from_exception(exc)
  1367. # One iterative pass over the copy does both jobs: keep only <model> frames
  1368. # on every linked exception and group member, and cut the chain at the cap.
  1369. found = False
  1370. truncated = False
  1371. pending = [(te, 1)]
  1372. while pending:
  1373. entry, depth = pending.pop()
  1374. kept = [f for f in entry.stack if f.filename == "<model>"]
  1375. entry.stack = traceback.StackSummary.from_list(kept)
  1376. found = found or bool(kept)
  1377. # 3.11+ exception groups (a binding failure inside asyncio.TaskGroup)
  1378. # carry member stacks under `exceptions`, not the dunder links; a group
  1379. # member counts as a link so the cap bounds nesting through both edges.
  1380. members = getattr(entry, "exceptions", None) or ()
  1381. if depth >= _MAX_TRACEBACK_CHAIN:
  1382. if entry.__cause__ is not None or entry.__context__ is not None or members:
  1383. truncated = True
  1384. entry.__cause__ = None
  1385. entry.__context__ = None
  1386. if members:
  1387. entry.exceptions = None
  1388. continue
  1389. for linked in (entry.__cause__, entry.__context__):
  1390. if linked is not None:
  1391. pending.append((linked, depth + 1))
  1392. for member in members:
  1393. pending.append((member, depth + 1))
  1394. def emit():
  1395. if found:
  1396. yield from te.format()
  1397. else:
  1398. yield from traceback.format_exception_only(type(exc), exc)
  1399. if truncated:
  1400. yield f"[dsh-code-runtime-python] exception chain truncated at {_MAX_TRACEBACK_CHAIN} links\n"
  1401. return _join_bounded(emit(), max_bytes)
  1402. def _make_failure_reporter() -> Any:
  1403. """Build the failure-diagnostic renderer over closure-held primitives.
  1404. The returned callable renders a model failure diagnostic that cannot itself
  1405. raise. The caller sends the ``done`` frame AFTER its ``except BaseException``
  1406. block, so anything thrown while rendering the diagnostic skips the send
  1407. entirely: the host then blocks on fd 3 until ``maxWallMs`` and reports a
  1408. timeout instead of the exception that actually happened. Rendering runs
  1409. model code by design (``format()`` reaches ``__str__``, ``__repr__`` and
  1410. ``__notes__``) and allocates under ``RLIMIT_AS``, so it must be treated as
  1411. able to throw.
  1412. The fallback names the exception CLASS and a fixed literal — no ``str(exc)``
  1413. and no ``format_exception_only``, both of which reach the model's
  1414. ``__str__``. A ``__name__`` that is not exactly ``str`` (a metaclass
  1415. property can return anything, or raise) is discarded rather than
  1416. formatted, so no override runs on this path either.
  1417. The factory exists for the same reason :func:`_make_cpu_enforcer` does: this
  1418. bootstrap IS ``__main__``, so ``import __main__; __main__._X = ...`` rebinds
  1419. any module global a call-time lookup would read. On this path a rebind is
  1420. worst — the handler's own reporter, and everything the reporter reaches,
  1421. would run model code outside any guard, and a throw there costs the ``done``
  1422. frame. The traceback formatter, the byte cap and the fallback literal
  1423. therefore become closure cells captured at import time, before model code
  1424. runs, and :func:`_run` binds the returned callable into a local of its own
  1425. frame. A frame local is not a module attribute, so no assignment to
  1426. ``__main__`` changes which callable runs or what it calls. This defeats the
  1427. one-line rebind, not a determined ``sys._getframe`` walk; the unforgeable
  1428. bound is the host wall clock.
  1429. """
  1430. cap_message = _cap_message
  1431. model_traceback = _model_traceback
  1432. unrenderable = _UNRENDERABLE_DIAGNOSTIC
  1433. def safe_model_traceback(exc: BaseException, max_bytes: int) -> str:
  1434. try:
  1435. return cap_message(model_traceback(exc, max_bytes), max_bytes)
  1436. except BaseException: # noqa: BLE001 -- a throw here would cost the done frame
  1437. pass
  1438. try:
  1439. raw_name = type(exc).__name__
  1440. # Slice BEFORE interpolating. A metaclass `__name__` property can
  1441. # return an arbitrarily long string, and both the f-string and
  1442. # `cap_message`'s encode would copy it whole — under a tight
  1443. # RLIMIT_AS either allocation can raise MemoryError, and this is the
  1444. # LAST fallback, so a throw here costs the `done` frame outright and
  1445. # the run misreports as an exit or a timeout. The slice is a
  1446. # code-unit prefix, which bounds the bytes at 4x, and the following
  1447. # `cap_message` still applies the exact byte cap.
  1448. name = raw_name[:_MAX_FALLBACK_NAME_CHARS] if type(raw_name) is str else "<unknown>"
  1449. except BaseException: # noqa: BLE001 -- a raising __name__ must not cost the done frame
  1450. name = "<unknown>"
  1451. # Wrapped for the same reason: `cap_message` encodes, and its allocation
  1452. # is the only step left that can still fail. The fixed literal needs no
  1453. # budget, so it can always be delivered.
  1454. try:
  1455. return cap_message(f"{name}: {unrenderable}", max_bytes)
  1456. except BaseException: # noqa: BLE001 -- the done frame outranks the diagnostic's detail
  1457. return unrenderable
  1458. return safe_model_traceback
  1459. _SAFE_MODEL_TRACEBACK = _make_failure_reporter()
  1460. def _join_bounded(lines, max_bytes: int) -> str:
  1461. """Join formatter output, stopping once the budget is comfortably passed.
  1462. ``format()`` yields lines lazily; consuming it whole for an exception
  1463. carrying a huge message would materialize the full text only for
  1464. ``_cap_message`` to throw it away — enough over-shoot to exhaust
  1465. ``RLIMIT_AS``. Stop after the accumulated CHARACTER count passes the byte
  1466. budget (chars lower-bound UTF-8 bytes); the caller's ``_cap_message``
  1467. does the exact byte-level cut.
  1468. """
  1469. chunks: list[str] = []
  1470. total = 0
  1471. for line in lines:
  1472. # A single yielded line can itself dwarf the budget (the exception
  1473. # message rides in one line): keep only the prefix it can ever need.
  1474. if len(line) > max_bytes + 1:
  1475. line = line[: max_bytes + 1]
  1476. chunks.append(line)
  1477. total += len(line)
  1478. if total > max_bytes:
  1479. break
  1480. return "".join(chunks)
  1481. def _done_with_value(value: Any, max_value_bytes: int) -> dict[str, Any]:
  1482. """Build the terminal done frame under the seam's lossless-JSON contract.
  1483. A completion value returned by the program (``None`` when it returns
  1484. nothing) that is not lossless JSON fails the run as ``invalid-output``; a
  1485. serialized value beyond ``max_value_bytes`` fails as ``output-limit``.
  1486. Substituting a ``repr`` or truncated string would be a silent lie about
  1487. what the program computed, so both paths refuse instead (mirroring the
  1488. worker backend's contract). ``None`` crosses as an exact JSON ``null``.
  1489. """
  1490. # One bounded walk folds the losslessness check and the byte meter (mirrors
  1491. # the host's checkDoneValue): the former split ran the full losslessness
  1492. # walk first, materializing one tuple per element for a wide completion
  1493. # before the size cap could reject it — an RLIMIT_AS death on a value the
  1494. # meter would have refused. send_sync later encodes the admitted value,
  1495. # whose size the walk proved within budget. Iterative like the encoder, so a
  1496. # valid completion deeper than the recursion limit still checks.
  1497. rejection = _check_done_value(value, max_value_bytes)
  1498. if rejection is not None:
  1499. kind, message = rejection
  1500. return {"type": "done", "error": {"kind": kind, "message": message}}
  1501. return {"type": "done", "value": value}
  1502. def main() -> None:
  1503. channel = ProtocolChannel(PROTOCOL_FD)
  1504. asyncio.run(_run(channel))
  1505. if __name__ == "__main__":
  1506. main()