bootstrap.py 128 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109211021112112211321142115211621172118211921202121212221232124212521262127212821292130213121322133213421352136213721382139214021412142214321442145214621472148214921502151215221532154215521562157215821592160216121622163216421652166216721682169217021712172217321742175217621772178217921802181218221832184218521862187218821892190219121922193219421952196219721982199220022012202220322042205220622072208220922102211221222132214221522162217221822192220222122222223222422252226222722282229223022312232223322342235223622372238223922402241224222432244224522462247224822492250225122522253225422552256225722582259226022612262226322642265226622672268226922702271227222732274227522762277227822792280228122822283228422852286228722882289229022912292229322942295229622972298229923002301230223032304230523062307230823092310231123122313231423152316231723182319232023212322232323242325232623272328232923302331233223332334233523362337233823392340234123422343234423452346234723482349235023512352235323542355235623572358235923602361236223632364236523662367236823692370237123722373237423752376237723782379238023812382238323842385238623872388238923902391239223932394239523962397239823992400240124022403240424052406240724082409241024112412241324142415241624172418241924202421242224232424242524262427242824292430243124322433243424352436243724382439244024412442244324442445244624472448244924502451245224532454245524562457245824592460246124622463246424652466
  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. # Module-level captures for the done-frame LAST-resort fallback. This bootstrap IS
  37. # ``__main__``, so ``import __main__; __main__.os = ...`` would rebind ``os.write``
  38. # at call time inside ``ProtocolChannel.write_encoded``. These module-level names
  39. # are the RAW primitive for the fallback; they are BOUND INTO ``_run`` LOCALS
  40. # before the program runs (see ``send_done``), which is what makes a one-line
  41. # rebind unable to change which write the fallback uses — the module global here
  42. # is itself reachable by ``__main__._os_write = boom``, so the immunity lives in
  43. # the ``_run`` frame-local binding, not in the module global.
  44. _os_write = os.write
  45. _memoryview = memoryview
  46. # A fixed, pre-encoded done frame for the fallback. It carries no live model
  47. # value, so it can always be written even when a transitive name (a ``_dump_*``
  48. # helper or ``os``) has been rebound and the normal encode/write threw. The
  49. # message is the fixed literal ``<unrenderable>`` — distinct from the failure
  50. # reporter's ``_UNRENDERABLE_DIAGNOSTIC`` text; the host renders the run as an
  51. # exception rather than a worker-exit, which is the honest verdict for a settled
  52. # run whose reporting was sabotaged. The bytes are JSON-valid and
  53. # newline-terminated.
  54. _FALLBACK_DONE_FRAME = b'{"type":"done","error":{"kind":"exception","message":"<unrenderable>"}}\n'
  55. # Code-unit ceiling on the exception class name interpolated into the LAST-resort
  56. # failure diagnostic. A metaclass `__name__` property can return any length, and
  57. # that construction runs outside the guard that would otherwise absorb a
  58. # MemoryError, so the name is sliced before it is copied. Generous enough that no
  59. # real class name is touched.
  60. _MAX_FALLBACK_NAME_CHARS = 200
  61. # Mirror of the host's output-budget/address-space gate (src/index.ts's
  62. # OUTPUT_BUDGET_WORST_CASE_ADDRESS_SPACE_MULTIPLE and INTERPRETER_BASELINE_BYTES),
  63. # re-applied against the EFFECTIVE RLIMIT_AS after inheritance clamping. An astral
  64. # character is one character but ~4 bytes of str storage and ~4 UTF-8 bytes, and
  65. # three such copies are live at the peak — the caller's write argument, the line
  66. # slice or joined pending handed to push, and the encode copy push takes — so a
  67. # budget's worst-case peak is twelve times its byte count; the interpreter's own
  68. # footprint is reserved on top. Kept in sync with the host constants by the shared
  69. # reasoning, not a wire field.
  70. _OUTPUT_BUDGET_WORST_CASE_MULTIPLE = 12
  71. _INTERPRETER_BASELINE_BYTES = 64 * 1024 * 1024
  72. # ---------------------------------------------------------------------------
  73. # Log buffer — Python-side ledger for captured text.
  74. # ---------------------------------------------------------------------------
  75. class LogBuffer:
  76. """Ordered text capture under one shared byte budget.
  77. Once the budget is exhausted the buffer emits exactly one in-band
  78. truncation marker via ``sink`` and silently drops everything after. The
  79. cap is a blast-radius bound; "how much was lost" intentionally stays
  80. unmeasured.
  81. """
  82. def __init__(self, max_bytes: int, sink) -> None:
  83. self._max_bytes = max_bytes
  84. # The ledger starts one byte below max_bytes: each entry is charged its
  85. # JSON-string cost plus one separator byte, and the serialized outer
  86. # logs array adds one more byte of envelope (two brackets and n-1 commas
  87. # over n entries' separators), so a result that exactly exhausts the
  88. # ledger serializes to exactly max_bytes; WITHOUT the reserved byte it
  89. # would serialize to max_bytes + 1. Reserving that byte keeps an
  90. # admitted result within the configured cap; the truncation-marker entry
  91. # is envelope, not payload, and rides uncharged (``_max_bytes`` stays the
  92. # configured value for the marker's message text).
  93. self._remaining = max_bytes - 1
  94. self._truncated = False
  95. # Re-entrant so a caller may hold it across a compound read-modify-write
  96. # (``_LogStream.write`` reads ``remaining`` several times and then calls
  97. # ``push`` while still holding it). One lock is shared by this buffer and
  98. # every stream that funnels into it: model code may start daemon threads
  99. # that keep calling ``print`` after the program body returns, and the
  100. # settlement ``flush_line`` on the main coroutine reads and mutates the
  101. # same ``_pending``/ledger state. Without a shared lock the flush could
  102. # interleave with a concurrent ``write`` — dropping or double-counting a
  103. # line, or costing the ``done`` frame on a mangled ledger. Fixing which
  104. # callable runs (binding ``out_stream.flush_line``) does not fix what it
  105. # reads.
  106. self._lock = threading.RLock()
  107. # ``sink(text, truncated=False)``. The marker is emitted with
  108. # ``truncated=True`` so the host can stop its own capture at the same
  109. # point rather than treating the marker as ordinary program output: the
  110. # two ledgers exhaust independently, and one entry larger than
  111. # ``max_bytes`` sends only the marker while the host budget is still
  112. # nearly empty.
  113. self._sink = sink
  114. @property
  115. def lock(self) -> "threading.RLock":
  116. """The shared re-entrant lock guarding this ledger and its streams' buffers."""
  117. return self._lock
  118. @property
  119. def remaining(self) -> int:
  120. """Serialized bytes still admissible; zero once truncated (streams use this to bound their own buffering, where a character count is a valid lower bound)."""
  121. return 0 if self._truncated else self._remaining
  122. def push(self, text: str) -> None:
  123. with self._lock:
  124. self._push_locked(text)
  125. def _push_locked(self, text: str) -> None:
  126. if self._truncated:
  127. return
  128. # Cheap lower bound FIRST: one char is at least one UTF-8 byte and the
  129. # JSON form adds two quotes plus the separator, so a single print() far
  130. # above the budget truncates without ever encoding it — the full encode
  131. # would allocate a second equally large string and could turn a
  132. # truncatable log into an RLIMIT_AS death.
  133. if len(text) + 3 > self._remaining:
  134. self._truncated = True
  135. self._sink(log_truncation_marker(self._max_bytes), truncated=True)
  136. return
  137. # A model print() can emit a lone surrogate; strict UTF-8 throws on it
  138. # here. Replace it rather than escaping it the way :func:`_dump_string`
  139. # preserves one inside a completion VALUE: log text is already a
  140. # truncatable, substituting channel (the byte cap replaces the tail with
  141. # a marker), and the ledger below charges the RAW UTF-8 bytes, which
  142. # would undercharge the six-byte escape by half. Bounded: the text
  143. # passed the length check, so this encodes at most ~4x remaining.
  144. try:
  145. raw = text.encode("utf-8")
  146. except UnicodeEncodeError:
  147. raw = text.encode("utf-8", errors="replace")
  148. text = raw.decode("utf-8")
  149. # Charge the SERIALIZED cost — the JSON string form's bytes plus one
  150. # separator byte — exactly as the host ledger does. Charging the raw
  151. # UTF-8 length instead undercharges control-heavy text, whose JSON
  152. # escaping expands it up to sixfold (a NUL costs one raw byte but six
  153. # as its ``\uXXXX`` escape): a NUL flood sized to fit ``maxLogBytes``
  154. # raw would serialize to roughly six times the shared cap, and the child
  155. # could then die on RLIMIT_AS (reported host-side as ``worker-exit``)
  156. # instead of emitting the truncation marker. The +1 also floors an empty
  157. # entry above zero, so a flood of blank ``print()`` lines exhausts the
  158. # budget instead of emitting unbounded zero-cost log frames.
  159. cost = _json_string_cost(raw) + 1
  160. if cost > self._remaining:
  161. self._truncated = True
  162. self._sink(log_truncation_marker(self._max_bytes), truncated=True)
  163. return
  164. self._remaining -= cost
  165. self._sink(text)
  166. class _LogStream(io.TextIOBase):
  167. """A newline-coalescing text stream backed by a :class:`LogBuffer`.
  168. Installed as ``sys.stdout`` / ``sys.stderr`` before executing the model
  169. program. ``print(...)`` calls ``write`` once per argument, separator, and
  170. newline, so a raw one-push-per-write stream would emit
  171. ``["a", " ", "b", "\\n"]`` for ``print("a", "b")`` — and Code Mode renders
  172. ``logs`` with ``join('\\n')``, turning that into spurious blank lines. This
  173. stream instead buffers writes and pushes one LogBuffer entry per completed
  174. LINE (the text up to each ``\\n``, newline stripped), so the rendered join
  175. reproduces ``a b``. Any unterminated tail is flushed by :meth:`flush_line`
  176. after the program settles.
  177. """
  178. def __init__(self, logs: LogBuffer) -> None:
  179. super().__init__()
  180. self._logs = logs
  181. # list-of-chunks, joined only at a newline or flush: repeated
  182. # ``print("x", end="")`` must not concatenate quadratically.
  183. self._pending: list[str] = []
  184. self._pending_blocks: list[str] = []
  185. self._pending_chars = 0
  186. # A newline-free drip must not accumulate one list slot per ``write``:
  187. # under a large ``maxLogBytes`` the list-of-fragments pointer array and
  188. # the per-fragment str objects cost host memory well before the byte
  189. # budget is reached, and a 25 M single-character drip would OOM on its
  190. # own accounting (plus the same-size list ``_push_bounded_prefix`` then
  191. # builds). Past this many fragments the chunks are SEALED into a
  192. # ``_pending_blocks`` entry (the character count is unchanged), bounding
  193. # the live fragment count exactly as the host-side ``captureStray`` does
  194. # with its ``MAX_PENDING_CHUNKS``. The seal is INCREMENTAL: only the
  195. # current ``_pending`` fragments (≤ cap) are joined into one block, not
  196. # the whole accumulated buffer, so a large drip stays O(B) rather than
  197. # re-copying the growing block O(B²/cap) times.
  198. self._PENDING_MAX_CHUNKS = 1024
  199. def _pending_parts(self) -> list[str]:
  200. """The sealed blocks followed by the current fragments, for joining."""
  201. return [*self._pending_blocks, *self._pending]
  202. def writable(self) -> bool: # noqa: D401 -- inherited contract
  203. return True
  204. def write(self, text: str) -> int: # noqa: D401 -- inherited contract
  205. # Serialize the whole read-modify-write against the settlement flush and
  206. # any other thread's write: model code may spawn daemon threads that keep
  207. # printing after the program body returns, and this method reads
  208. # ``remaining`` and mutates ``_pending``/the ledger across many steps. The
  209. # lock is the buffer's and is re-entrant, so the ``push`` calls below
  210. # (which re-acquire it) do not deadlock.
  211. with self._logs.lock:
  212. return self._write_locked(text)
  213. def _write_locked(self, text: str) -> int:
  214. # Drop an empty write instead of buffering it. An empty chunk adds no
  215. # character, so the budget check below can never fire on it:
  216. # ``while True: sys.stdout.write("")`` would append one list slot per
  217. # call with `_pending_chars` pinned at 0, growing unbounded long after
  218. # the log ledger was exhausted (about 3.7 M slots per CPU second here)
  219. # until RLIMIT_AS turned the allocation into a MemoryError — reported as
  220. # the program's own exception rather than the intended bounded-log
  221. # behavior. Returning here also keeps `flush_line` from pushing a
  222. # spurious empty log entry for a program whose only writes were empty.
  223. if not text:
  224. return 0
  225. if "\n" in text:
  226. # Scan `text` in place; the buffered chunks are joined ONLY into the
  227. # first line. Joining the pending chunks with the whole write first
  228. # made a second copy of that write, which an over-budget write
  229. # cannot afford (measured under a 400 MiB addressSpaceMb: one
  230. # buffered character followed by a 340 MiB write died on MemoryError
  231. # inside the join, reported as the program's own exception, and the
  232. # retained chunks made the settlement `flush_line` fail the same way
  233. # — costing the `done` frame and turning the run into a wall-clock
  234. # timeout instead of the promised truncation marker).
  235. length = len(text)
  236. pos = 0
  237. if self._pending or self._pending_blocks:
  238. newline = text.index("\n")
  239. if self._pending_chars + newline + 3 > self._logs.remaining:
  240. # The reconstructed first line cannot fit the ledger, so
  241. # LogBuffer would reject it whole: copy only the prefix that
  242. # fails its cheap bound and drop the chunks. The slice is
  243. # bounded HERE, not inside the helper: `text[:newline]` on a
  244. # 340 MiB newline-terminated write is the same full copy the
  245. # join was (measured: MemoryError inside `sys.stdout.write`
  246. # under a 400 MiB addressSpaceMb), and `remaining + 4`
  247. # characters are all the helper can use.
  248. self._push_bounded_prefix(text[: min(newline, self._logs.remaining + 4)])
  249. else:
  250. self._pending.append(text[:newline])
  251. line = "".join(self._pending_parts())
  252. self._pending = []
  253. self._pending_blocks = []
  254. self._pending_chars = 0
  255. self._logs.push(line)
  256. pos = newline + 1
  257. # Scan by offset and STOP once the ledger is exhausted: a single
  258. # write of many newlines (``print("\n" * 1000000)``) would otherwise
  259. # re-slice the tail once per line and keep pushing long after
  260. # LogBuffer truncated, burning the CPU budget on discarded lines.
  261. # `remaining` reads 0 the instant the buffer truncates, so the loop
  262. # exits immediately; the unscanned tail is simply dropped.
  263. while pos < length and self._logs.remaining > 0:
  264. newline = text.find("\n", pos)
  265. if newline < 0:
  266. break
  267. # Bound the SLICE the same way LogBuffer bounds the encode: a
  268. # first line far above the ledger would be copied whole before
  269. # push could reject it, and that copy is the allocation an
  270. # over-budget write cannot afford. Copy only a budget-sized
  271. # prefix, which push still rejects on its own cheap bound (the
  272. # prefix is longer than `remaining`), so the marker is emitted
  273. # and the oversized line is never materialized.
  274. if newline - pos + 3 > self._logs.remaining:
  275. self._logs.push(text[pos:pos + self._logs.remaining + 4])
  276. break
  277. self._logs.push(text[pos:newline])
  278. pos = newline + 1
  279. if pos < length:
  280. if self._logs.remaining > 0:
  281. # Buffer only a budget-sized PREFIX of the tail, not the whole
  282. # `text[pos:]`: an early newline followed by a huge unterminated
  283. # tail (`"\n" + "A" * 30 MiB`) would otherwise copy the entire
  284. # tail into `_pending` here — a second full copy of the model's
  285. # own string, the RLIMIT_AS death this path exists to avoid —
  286. # before the newline-free trigger below could bound it. Anything
  287. # past `remaining` characters cannot be admitted (the char count
  288. # is a lower bound on the serialized cost), so a
  289. # `remaining + 4`-character prefix is all that can ever survive;
  290. # the flush trigger below rejects it and emits the marker.
  291. tail = text[pos:pos + self._logs.remaining + 4]
  292. self._pending.append(tail)
  293. self._pending_chars = len(tail)
  294. else:
  295. # The ledger ran out with text still unscanned, so that text
  296. # IS being dropped and the run must say so. One push is
  297. # enough and is bounded: `remaining` is 0, so LogBuffer's
  298. # cheap length lower bound rejects immediately, emits the
  299. # marker, and never encodes the tail — and a push after the
  300. # marker is already out returns without emitting a second.
  301. # Reaching 0 EXACTLY (65 one-character lines against the
  302. # default 3-byte-per-entry serialized charge) leaves
  303. # `_truncated` unset, so without this the tail vanished with
  304. # no marker at all. Sliced to a budget-sized prefix, not the
  305. # whole tail: the tail can be hundreds of megabytes and the
  306. # copy would be the RLIMIT_AS death this bound exists to
  307. # avoid, while push only needs enough characters to fail its
  308. # own cheap length check.
  309. self._logs.push(text[pos:pos + self._logs.remaining + 4])
  310. else:
  311. self._pending.append(text)
  312. self._pending_chars += len(text)
  313. # Seal the fragment list past the chunk cap: a newline-free drip
  314. # appends one fragment per write, so a 25 M single-character flood
  315. # would accumulate that many list slots (and str objects) long before
  316. # the byte budget is met — the pointer array alone being ~25 M slots.
  317. # Past the cap the current fragments are joined into ONE block and
  318. # moved to `_pending_blocks` (character count unchanged), bounding the
  319. # live fragment count exactly as the host-side `captureStray` seal
  320. # does. The join is only the ≤cap current fragments, never the whole
  321. # accumulated buffer, so a large drip stays O(B) rather than
  322. # re-copying the growing block O(B²/cap) times; a newline never starts
  323. # a multi-byte sequence, so the fragments are un-sealable mid-line.
  324. if len(self._pending) >= self._PENDING_MAX_CHUNKS:
  325. self._pending_blocks.append("".join(self._pending))
  326. self._pending = []
  327. # A newline-free flood must hit the budget while running, not at
  328. # settlement: once the buffered tail alone can no longer fit the
  329. # ledger (chars lower-bound the serialized cost), push it through — LogBuffer
  330. # truncates, emits the marker once, and swallows everything after.
  331. if self._pending_chars > self._logs.remaining:
  332. self._push_bounded_prefix()
  333. return len(text)
  334. def _push_bounded_prefix(self, extra: str = "") -> None:
  335. # Reached only when the buffered characters already exceed what the
  336. # ledger admits, so LogBuffer is certain to reject on its cheap length
  337. # bound and emit the marker. Copy a budget-sized PREFIX rather than the
  338. # joined whole: ``sys.stdout.write("x")`` followed by one newline-free
  339. # 340 MiB write leaves two chunks whose join is a second copy of the
  340. # payload, and under a tight addressSpaceMb that join raises MemoryError
  341. # from inside `write` — surfacing as the program's own exception, or,
  342. # while the oversized chunks stayed retained, again from `flush_line`
  343. # after the program settled, which cost the `done` frame and turned the
  344. # run into a wall-clock timeout instead of the promised truncation
  345. # marker.
  346. #
  347. # The chunks are dropped BEFORE the push so neither this call nor the
  348. # settlement flush can repeat the allocation, and dropping the text is
  349. # exactly what the marker reports. `remaining + 4` is the shortest
  350. # prefix that still fails LogBuffer's ``len(text) + 3 > remaining``
  351. # check; the accumulation stops there, so the copy is bounded by the log
  352. # budget however large the pending chunks are.
  353. #
  354. # `self._pending` is iterated IN PLACE and `extra` handled after it:
  355. # `(*self._pending, extra)` would first copy every pending reference into
  356. # a same-size tuple, which for a single-character drip (millions of tiny
  357. # chunks) is a second pointer array as large as the list itself --
  358. # measured at +80 MiB of tuple on top of a 40 MiB list for 5.2M chunks,
  359. # the allocation this bounded prefix exists to avoid.
  360. limit = self._logs.remaining + 4
  361. parts: list[str] = []
  362. total = 0
  363. for chunk in self._pending_parts():
  364. parts.append(chunk[: limit - total])
  365. total += len(parts[-1])
  366. if total >= limit:
  367. break
  368. else:
  369. # Only reached when the pending chunks did not fill the prefix, so
  370. # `extra` is the one remaining source of text.
  371. parts.append(extra[: limit - total])
  372. self._pending = []
  373. self._pending_blocks = []
  374. self._pending_chars = 0
  375. self._logs.push("".join(parts))
  376. def flush(self) -> None: # noqa: D401 -- inherited contract
  377. # ``TextIOBase.flush`` is a no-op, so without this override an explicit
  378. # ``print(..., flush=True)`` or ``sys.stdout.flush()`` left the text in
  379. # `_pending` with nothing to drain it except `flush_line` after the
  380. # program settles. A run that then hangs or is killed never reaches that
  381. # call: ``print("before hang", end="", flush=True)`` followed by an
  382. # infinite loop returned `logs: []`, losing the one diagnostic the
  383. # program deliberately committed. Forwarding makes an explicit flush emit
  384. # the pending entry immediately, which is what the caller asked for; a
  385. # newline-terminated write already emitted on its own.
  386. self.flush_line()
  387. def flush_line(self) -> None:
  388. """Push any buffered text not terminated by a newline (also serves explicit flushes)."""
  389. # Same shared, re-entrant lock as ``write``: the settlement flush on the
  390. # main coroutine and a daemon thread's concurrent ``write`` both touch
  391. # ``_pending`` and the ledger, so this read-and-clear must be atomic
  392. # against them.
  393. with self._logs.lock:
  394. if self._pending or self._pending_blocks:
  395. # Join, drop the chunks, THEN push — the same join-clear-push order
  396. # as `_write_locked`'s newline branch. Pushing before the clear would keep the
  397. # pending chunks alive through `_push_locked`'s `text.encode`, so
  398. # the chunks, their join, and the encode copy would all be live at
  399. # once; dropping the chunks first leaves only the join and its
  400. # encode, matching that path's peak.
  401. line = "".join(self._pending_parts())
  402. self._pending = []
  403. self._pending_blocks = []
  404. self._pending_chars = 0
  405. self._logs.push(line)
  406. # ---------------------------------------------------------------------------
  407. # Fd-3 channel — line-framed JSON.
  408. # ---------------------------------------------------------------------------
  409. # Non-string scalars only. The string form is scanned by hand in
  410. # :func:`_decode_json_plain` because a ``(?:[^"\\]|\\.)*`` repetition makes
  411. # CPython's backtracking engine retain per-repetition state proportional to the
  412. # string's WIDTH: measured at ~146 MiB of engine state for a 1 MiB string and
  413. # ~558 MiB for 4 MiB, so a legitimate multi-megabyte binding reply raised
  414. # MemoryError out of ``_pump_replies``, leaving its future unsettled until the
  415. # wall clock reported a timeout.
  416. _SCALAR_RE = re.compile(
  417. r'-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?|true|false|null'
  418. )
  419. # A run of ordinary string body characters. The star applies to a CHARACTER
  420. # CLASS, which the engine matches in one linear pass with no backtracking state,
  421. # so the scanner's cost is the number of escapes, not the string's width.
  422. _STRING_CHUNK_RE = re.compile(r'[^"\\]*')
  423. def _decode_json_plain(
  424. text: str,
  425. # Def-time captures: the decoder runs AFTER the program (which is `__main__`)
  426. # may have rebound `__main__.json`, `__main__._SCALAR_RE`,
  427. # `__main__._STRING_CHUNK_RE`, or `__main__.len`; a call-time lookup would
  428. # let a one-line rebind kill the reply pump (a broken decode strands every
  429. # pending Future to the wall clock). Defaults are evaluated at def time.
  430. _json_loads: Any = json.loads,
  431. _scalar_re: Any = _SCALAR_RE,
  432. _string_chunk_re: Any = _STRING_CHUNK_RE,
  433. _len: Any = len,
  434. _isinstance: Any = isinstance,
  435. _str: Any = str,
  436. _list: Any = list,
  437. ) -> Any:
  438. """Parse one JSON document iteratively (no per-level recursion).
  439. ``json.loads`` recurses per nesting level and raises ``RecursionError``
  440. around ~10k levels, but a binding reply is depth-unbounded by the seam
  441. contract — the host's iterative encoder happily produces documents
  442. ``json.loads`` cannot read back. Scalars (numbers, strings with escapes)
  443. are delegated to ``json.loads`` one token at a time, so their grammar and
  444. semantics stay CPython's own; only the container structure is parsed here
  445. with an explicit stack. Raises ``ValueError`` on malformed input; frames
  446. come from the TRUSTED host, so strictness mirrors ``json.loads`` without
  447. extra hostile-input hardening.
  448. """
  449. length = _len(text)
  450. def skip_ws(i: int) -> int:
  451. while i < length and text[i] in " \t\n\r":
  452. i += 1
  453. return i
  454. def scan_string(i: int) -> int:
  455. # Walk chunk by chunk: each match consumes every character up to the next
  456. # quote or backslash, so an escape costs one extra step and a plain body
  457. # costs one pass. Returns the offset just past the closing quote.
  458. j = i + 1
  459. while True:
  460. j = _string_chunk_re.match(text, j).end()
  461. if j >= length:
  462. raise ValueError(f"unterminated string at offset {i}")
  463. char = text[j]
  464. if char == '"':
  465. return j + 1
  466. # text[j] is a backslash: skip it and the character it escapes. A
  467. # trailing backslash runs j past `length`, caught on the next pass.
  468. j += 2
  469. def scalar(i: int):
  470. if i < length and text[i] == '"':
  471. end = scan_string(i)
  472. return _json_loads(text[i:end]), end
  473. match = _scalar_re.match(text, i)
  474. if match is None:
  475. raise ValueError(f"invalid JSON at offset {i}")
  476. return _json_loads(match.group(0)), match.end()
  477. def string_key(i: int):
  478. key, end = scalar(i)
  479. if not _isinstance(key, _str):
  480. raise ValueError(f"object key must be a string at offset {i}")
  481. end = skip_ws(end)
  482. if end >= length or text[end] != ":":
  483. raise ValueError(f"expected ':' at offset {end}")
  484. return key, skip_ws(end + 1)
  485. # Frames: a list, or (dict, pending key). `value`/`have_value` carry each
  486. # completed value up to its parent frame.
  487. stack: list[Any] = []
  488. value: Any = None
  489. have_value = False
  490. i = skip_ws(0)
  491. while True:
  492. if not have_value:
  493. ch = text[i] if i < length else ""
  494. if ch == "[":
  495. i = skip_ws(i + 1)
  496. if i < length and text[i] == "]":
  497. i += 1
  498. value, have_value = [], True
  499. else:
  500. stack.append([])
  501. continue
  502. elif ch == "{":
  503. i = skip_ws(i + 1)
  504. if i < length and text[i] == "}":
  505. i += 1
  506. value, have_value = {}, True
  507. else:
  508. key, i = string_key(i)
  509. stack.append(({}, key))
  510. continue
  511. else:
  512. value, i = scalar(i)
  513. have_value = True
  514. if not stack:
  515. i = skip_ws(i)
  516. if i != length:
  517. raise ValueError(f"trailing data at offset {i}")
  518. return value
  519. top = stack[-1]
  520. i = skip_ws(i)
  521. ch = text[i] if i < length else ""
  522. if _isinstance(top, _list):
  523. top.append(value)
  524. if ch == ",":
  525. i = skip_ws(i + 1)
  526. have_value = False
  527. elif ch == "]":
  528. i += 1
  529. stack.pop()
  530. value = top
  531. else:
  532. raise ValueError(f"expected ',' or ']' at offset {i}")
  533. else:
  534. container, key = top
  535. container[key] = value
  536. if ch == ",":
  537. key, i = string_key(skip_ws(i + 1))
  538. stack[-1] = (container, key)
  539. have_value = False
  540. elif ch == "}":
  541. i += 1
  542. stack.pop()
  543. value = container
  544. else:
  545. raise ValueError(f"expected ',' or '}}' at offset {i}")
  546. class ProtocolChannel:
  547. """Blocking readers and synchronous writers over the fd-3 protocol pipe.
  548. Writes are unbuffered and go straight to the fd, so ``send_sync`` is safe
  549. from inside model code (which may run outside an asyncio task) and from
  550. background tasks alike. Concurrent writers are serialized by ``_write_lock``
  551. around a full-write loop (see ``send_sync``): ``os.write`` releases the GIL,
  552. a frame may exceed ``PIPE_BUF`` (logs up to ``maxLogBytes``, completions up
  553. to ``maxValueBytes``, uncapped ``call`` args), and one ``os.write`` may
  554. consume only part of a frame — so neither the GIL nor per-frame atomicity is
  555. relied on for framing.
  556. """
  557. def __init__(self, fd: int) -> None:
  558. self._fd = fd
  559. # Residual bytes read past a frame's newline, shared by the blocking and
  560. # async readers. Held here, not in the reading coroutine: the reply pump
  561. # is cancelled once `done` is posted, and read-ahead sitting in a local
  562. # would be lost with it. Both readers use `os.read(self._fd, ...)`
  563. # directly, so no buffered file object wraps the fd.
  564. self._pending = bytearray()
  565. # Serializes writers: os.write releases the GIL, and a frame larger
  566. # than PIPE_BUF is neither atomic nor guaranteed fully consumed by one
  567. # call — without the lock, model-created threads printing while a big
  568. # completion frame drains could interleave bytes mid-frame.
  569. self._write_lock = threading.Lock()
  570. def read_frame(
  571. self,
  572. # Def-time captures for the decode primitives (see _decode_json_plain):
  573. # this runs before the program, but the reply path does not — a rebind
  574. # of `__main__._decode_json_plain`/`__main__.os`/`__main__._READ_CHUNK_BYTES`
  575. # must not break the pump.
  576. _decode: Any = _decode_json_plain,
  577. _os_read: Any = os.read,
  578. _read_chunk: int = _READ_CHUNK_BYTES,
  579. _bytes: Any = bytes,
  580. _len: Any = len,
  581. ) -> dict[str, Any] | None:
  582. """Read one JSON-line frame (iteratively decoded). ``None`` on EOF.
  583. Blocking. Used for the two frames read BEFORE the model program starts
  584. (``boot`` and ``run``), where blocking is what the handshake wants. Reply
  585. frames arriving during the program go through :meth:`read_frame_async`,
  586. which must not occupy a thread.
  587. Reads in CHUNKS into the shared ``_pending`` buffer rather than through
  588. ``FileIO.readline()``: the fd is unbuffered (``buffering=0``), so
  589. ``readline`` issues one ``os.read(1)`` per byte, and a multi-megabyte
  590. ``run`` frame — RLIMIT_CPU already in force by then — would burn the
  591. budget in millions of syscalls before ``ast.parse`` even runs. The chunk
  592. reads and the same residual buffer the async path uses keep read-ahead
  593. past a newline for the next frame.
  594. """
  595. # Scan only the bytes not yet examined: `find` from a running offset so a
  596. # frame arriving in N chunks costs one linear pass total, not one rescan
  597. # of the whole buffer per chunk (which is quadratic in the frame size).
  598. scanned = 0
  599. while True:
  600. newline = self._pending.find(b"\n", scanned)
  601. if newline >= 0:
  602. line = _bytes(self._pending[:newline])
  603. del self._pending[: newline + 1]
  604. return _decode(line.decode("utf-8"))
  605. scanned = _len(self._pending)
  606. chunk = _os_read(self._fd, _read_chunk)
  607. if not chunk:
  608. # EOF before a newline: drop the partial line, as the host drops
  609. # a frame that never completed.
  610. return None
  611. self._pending.extend(chunk)
  612. async def read_frame_async(
  613. self,
  614. # Def-time captures, same rationale as read_frame.
  615. _decode: Any = _decode_json_plain,
  616. _os_read: Any = os.read,
  617. _read_chunk: int = _READ_CHUNK_BYTES,
  618. _bytes: Any = bytes,
  619. _get_event_loop: Any = asyncio.get_event_loop,
  620. _len: Any = len,
  621. ) -> dict[str, Any] | None:
  622. """Await one JSON-line frame without occupying a thread. ``None`` on EOF.
  623. ``loop.run_in_executor(None, read_frame)`` was the obvious spelling and
  624. the wrong one: the default executor spins up its first thread the moment
  625. the program awaits a binding, and on Linux/glibc that thread's 8 MiB
  626. stack plus a 64 MiB per-thread malloc arena reservation are charged to
  627. ``RLIMIT_AS`` — measured, the child's mappings went from 30.34 MiB to
  628. 102.39 MiB across one ``await tools.*``. Since the limit is already in
  629. force, that ~72 MiB comes straight out of the run's ``addressSpaceMb``:
  630. under a small limit the thread cannot start at all and a legitimate
  631. binding call hangs to ``maxWallMs``, and under a larger one an allocation
  632. that should have fit dies as ``MemoryError``. This is the same accounting
  633. the settlement-time CPU recheck was designed around, where a sampling
  634. thread cost the same 72 MiB.
  635. `loop.add_reader` watches the fd instead, so no thread exists.
  636. Bytes past a frame's newline belong to the next frame, so the residual
  637. lives on the CHANNEL rather than in this coroutine: the pump is cancelled
  638. once ``done`` is posted, and a local buffer would discard whatever it had
  639. read ahead.
  640. """
  641. loop = _get_event_loop()
  642. # Scan only the not-yet-examined bytes (running offset), so a frame
  643. # arriving across many reads costs one linear pass, not a quadratic
  644. # rescan of the whole buffer per read.
  645. scanned = 0
  646. while True:
  647. newline = self._pending.find(b"\n", scanned)
  648. if newline >= 0:
  649. line = _bytes(self._pending[:newline])
  650. del self._pending[: newline + 1]
  651. return _decode(line.decode("utf-8"))
  652. scanned = _len(self._pending)
  653. ready = loop.create_future()
  654. # `add_reader` only reports readability; the read itself happens here,
  655. # and `os.read` returns whatever is buffered without waiting for more.
  656. loop.add_reader(self._fd, lambda: ready.done() or ready.set_result(None))
  657. try:
  658. await ready
  659. finally:
  660. loop.remove_reader(self._fd)
  661. chunk = _os_read(self._fd, _read_chunk)
  662. if not chunk:
  663. # EOF. Any partial line is dropped, matching how the host drops a
  664. # frame that never completed.
  665. return None
  666. self._pending.extend(chunk)
  667. def send_sync(self, message: dict[str, Any]) -> None:
  668. """Post one frame synchronously.
  669. Encoded with the iterative :func:`_encode_json_plain` (not
  670. ``json.dumps``, whose per-level recursion would raise
  671. ``RecursionError`` on a deeply nested completion or call argument the
  672. depth-unbounded ``CodeJsonValue`` contract admits). NaN/Infinity still
  673. raise ``ValueError`` — they would serialize as non-standard tokens
  674. that Node's ``JSON.parse`` rejects, silently dropping the frame, and a
  675. call would then hang until the wall clock instead of failing fast.
  676. Callers turn the ``ValueError`` into their own contract error
  677. (dispatch raises the lossless-JSON message).
  678. """
  679. self.write_encoded(_encode_json_plain(message))
  680. def write_encoded(self, frame: str) -> None:
  681. """Write a frame that is ALREADY encoded to its JSON string form.
  682. Appends the frame's trailing newline and full-write-loops the bytes
  683. under the writer lock, identically to :meth:`send_sync`. The consumer
  684. supplies the encoded JSON (a ``"done"`` frame carrying a completion
  685. value that was serialized at its validation point — see
  686. :func:`_done_with_value`); the channel does not re-encode it, so the
  687. bytes written are exactly what was validated with no second traversal
  688. of a live object.
  689. """
  690. payload = (frame + "\n").encode("utf-8")
  691. # Full-write loop under the writer lock (same rationale as send_sync):
  692. # one os.write may consume only part of a frame beyond PIPE_BUF, and a
  693. # partial or interleaved frame is dropped host-side as malformed JSON.
  694. with self._write_lock:
  695. view = memoryview(payload)
  696. while view:
  697. view = view[os.write(self._fd, view):]
  698. # ---------------------------------------------------------------------------
  699. # Tools proxy — turns ``await tools.name(args)`` into a fd-3 call frame.
  700. # ---------------------------------------------------------------------------
  701. class _Namespace:
  702. """A proxy for one binding namespace: every declared name routes to the bridge.
  703. Names arrive from :class:`BootMessage.namespaces`. Both attribute access
  704. (``tools.name``) and subscript access (``tools["my-tool"]`` — the SDK's
  705. escape hatch for exotic or reserved names, which are legal function names
  706. on the wire) return a coroutine factory that posts a ``call`` frame and
  707. awaits the matching ``reply``. An undeclared name raises ``AttributeError``
  708. (attribute) or ``KeyError`` (subscript), matching the worker backend's
  709. own-property discipline.
  710. ``__getattribute__`` (not ``__getattr__``) intercepts attribute access so a
  711. declared name ALWAYS reaches the bridge — even one that collides with an
  712. inherited attribute like ``__class__``, which ordinary lookup would resolve
  713. on ``object`` before ``__getattr__`` ever ran. Internal state lives under
  714. name-mangled ``_Namespace__*`` attributes; a declared binding with such a
  715. name still wins (declared-names check runs first).
  716. """
  717. def __init__(self, global_name: str, names: list[str], dispatch) -> None:
  718. self.__global = global_name
  719. self.__names = set(names)
  720. self.__dispatch = dispatch
  721. def __call_for(self, name: str):
  722. dispatch = object.__getattribute__(self, "_Namespace__dispatch")
  723. global_name = object.__getattribute__(self, "_Namespace__global")
  724. async def call(args: Any) -> Any:
  725. return await dispatch(global_name, name, args)
  726. return call
  727. def __getattribute__(self, name: str):
  728. # Declared names route to the bridge unconditionally — before Python
  729. # can resolve an inherited attribute (``__class__``) or our own
  730. # internals. Everything else falls through to normal lookup so the
  731. # proxy machinery itself keeps working.
  732. names = object.__getattribute__(self, "_Namespace__names")
  733. if name in names:
  734. return object.__getattribute__(self, "_Namespace__call_for")(name)
  735. return object.__getattribute__(self, name)
  736. def __getattr__(self, name: str):
  737. # Reached only when normal lookup found nothing (declared names were
  738. # already intercepted above), so this is always an undeclared tool.
  739. raise AttributeError(
  740. f"tool {name!r} is not declared in namespace "
  741. f"{object.__getattribute__(self, '_Namespace__global')!r}"
  742. )
  743. def __getitem__(self, name: str):
  744. names = object.__getattribute__(self, "_Namespace__names")
  745. if name not in names:
  746. raise KeyError(
  747. f"tool {name!r} is not declared in namespace "
  748. f"{object.__getattribute__(self, '_Namespace__global')!r}"
  749. )
  750. return object.__getattribute__(self, "_Namespace__call_for")(name)
  751. class _BindingRejection(Exception):
  752. """Internal reply-pump rejection, converted by ``dispatch`` into the
  753. namespace's declared error class (or ``RuntimeError``) so the marker type
  754. itself never reaches model code."""
  755. def _make_error_class(
  756. name: str,
  757. member_name_property: str,
  758. # Def-time captures (see the __init__ body): the minted class runs AFTER
  759. # the program, so `__main__.Exception`/`__main__.setattr` rebinds must not
  760. # break the rejection constructor.
  761. _Exception_init: Any = Exception,
  762. _setattr: Any = setattr,
  763. ) -> type:
  764. """Mint one program-visible rejection class per the seam's
  765. ``CodeBindingErrorClass`` contract: instances carry the failed member name
  766. under ``member_name_property`` and render as their message."""
  767. def __init__(self, member_name: str, message: str) -> None: # noqa: N807
  768. # Exception.__init__ and setattr are captured as defaults at def time:
  769. # the class constructor runs model-visible code paths, and a rebind of
  770. # `__main__.Exception`/`__main__.setattr` must not break a rejection.
  771. _Exception_init.__init__(self, message)
  772. _setattr(self, member_name_property, member_name)
  773. return type(name, (Exception,), {"__init__": __init__})
  774. def _clamped(which: int, soft: int, hard: int) -> tuple[int, int]:
  775. """Bound a requested (soft, hard) rlimit pair by BOTH inherited limits.
  776. An unprivileged process may lower a hard limit but never raise it, so a
  777. harness already started under a tighter ceiling (``ulimit -v`` below
  778. ``addressSpaceBytes``, or a CPU cap below ``cpuSeconds`` + 1) would make
  779. ``setrlimit`` raise ``ValueError`` and fail every run — despite the
  780. inherited limit being STRONGER than the one requested. Clamping keeps the
  781. stricter of the two, which still satisfies the containment contract.
  782. Both inherited bounds matter, not just the hard one. A deployment that
  783. inherited a soft limit BELOW what is requested (e.g. inherited ``(100, 200)``,
  784. requested ``(150, 160)``) must keep the stricter soft — returning the
  785. requested ``150`` would RAISE the effective soft limit, loosening RLIMIT_AS
  786. memory or deferring the RLIMIT_CPU SIGXCPU, the opposite of "strictest of
  787. configured and inherited". So each side is clamped against its inherited
  788. counterpart. ``RLIM_INFINITY`` compares as -1, so an infinite inherited bound
  789. imposes no ceiling and the requested value stands.
  790. """
  791. inherited_soft, inherited_hard = resource.getrlimit(which)
  792. clamped_soft = soft if inherited_soft == resource.RLIM_INFINITY else min(soft, inherited_soft)
  793. clamped_hard = hard if inherited_hard == resource.RLIM_INFINITY else min(hard, inherited_hard)
  794. # setrlimit requires soft <= hard. Clamping the two sides independently can
  795. # invert them (a finite inherited soft below the clamped hard is fine, but a
  796. # requested hard below the inherited soft would leave soft > hard), so pin
  797. # soft under hard as the final step; the stricter hard ceiling wins.
  798. result_soft = min(clamped_soft, clamped_hard)
  799. result_hard = clamped_hard
  800. # For RLIMIT_CPU, a soft limit EQUAL to the hard limit leaves the kernel no
  801. # window to send SIGXCPU: it checks the hard limit in the same tick and
  802. # SIGKILLs directly (a `ulimit -t N` sets both, and a busy loop then dies by
  803. # SIGKILL, not SIGXCPU). The host classifies a CPU overrun ONLY on ``signal
  804. # === 'SIGXCPU'``, so a definite budget exhaustion would be misreported as a
  805. # `worker-exit`. Lowering the soft limit one unit below the hard (when the
  806. # hard is at least 2, so soft stays positive) keeps the stricter-of-the-two
  807. # containment semantics while giving SIGXCPU a window to fire — the CPU
  808. # overrun is then reported as a timeout, not a worker-exit. This is scoped to
  809. # RLIMIT_CPU: for RLIMIT_AS a one-byte soft differential would only misalign
  810. # the child's applied limit with the host-side budget gate, with no signal to
  811. # preserve. The ``hard >= 2`` guard leaves the ``hard == 1`` blind spot
  812. # (a 1-second dual limit cannot lower soft to 0) — a definite CPU overrun
  813. # there is still reported as `worker-exit`; see the settlement note.
  814. if (
  815. which == resource.RLIMIT_CPU
  816. and result_soft == result_hard
  817. and result_hard >= 2
  818. ):
  819. result_soft = result_hard - 1
  820. return (result_soft, result_hard)
  821. # ---------------------------------------------------------------------------
  822. # Main.
  823. # ---------------------------------------------------------------------------
  824. async def _run(channel: ProtocolChannel) -> None:
  825. # The exception class every `except` clause in this function catches is bound
  826. # into a LOCAL at the very top, before any model code runs. This bootstrap IS
  827. # `__main__`, so `__main__.BaseException = RuntimeError` would otherwise
  828. # rebind the module global the `except BaseException` clauses resolve at
  829. # runtime — and a program exception that no longer matches would escape `_run`
  830. # with no done frame, misreporting the run as a `worker-exit`. A frame local
  831. # is not reachable by `__main__._X = ...`, so the catch is immune.
  832. _BaseException = BaseException
  833. # `RuntimeError` and the `_BindingRejection` marker class are likewise bound
  834. # into locals: `dispatch`'s `call_failure` and its `except` clause resolve
  835. # them at call time, and the program (running as `__main__`) can rebind the
  836. # module globals — `__main__._BindingRejection = ValueError` would leak the
  837. # marker type into model code, violating the class's conversion contract.
  838. # The names differ from the module globals (`_RuntimeError_cls`) so the
  839. # assignment RHS resolves the module global, not an unbound local.
  840. _RuntimeError_cls = RuntimeError
  841. _BindingRejection_cls = _BindingRejection
  842. # `str` for dispatch's rejection conversion is likewise bound: a program
  843. # rebinding `__main__.str` would otherwise run a hostile callable when the
  844. # binding-rejection message is formatted. `isinstance` for `send_done`'s
  845. # frame-shape check is bound the same way.
  846. _str = str
  847. _isinstance = isinstance
  848. # dispatch's argument-validation, event-loop, and frame-send primitives are
  849. # bound here, before the program runs: a rebind of `__main__._lossless_json_violation`,
  850. # `__main__.asyncio`, or the channel's send method must not turn a legitimate
  851. # binding call into an exception or a wall-clock timeout.
  852. _lossless_json_violation_cls = _lossless_json_violation
  853. _get_event_loop_cls = asyncio.get_event_loop
  854. _send_sync_cls = channel.send_sync
  855. # The frame WRITE primitives are bound too: send_sync's body resolves
  856. # `_encode_json_plain` (module global) and `self.write_encoded` (class
  857. # attribute) at call time, so a program rebinding either before the first
  858. # binding call could turn a legitimate call into an exception. dispatch
  859. # and the log sink both write through these directly (the sink does NOT go
  860. # through send_sync, whose body resolves the same names at call time).
  861. _write_encoded_cls = channel.write_encoded
  862. _encode_plain_cls = _encode_json_plain
  863. # 1. Boot handshake.
  864. boot = channel.read_frame()
  865. if boot is None or boot.get("type") != "boot":
  866. raise RuntimeError("bootstrap: expected boot frame on fd 3")
  867. # A limit that cannot be applied must fail the run as a diagnosable done
  868. # frame, not a bare traceback + exit(1): running the program UNCAPPED would
  869. # silently void the containment contract, and the host can only relay what
  870. # rides the protocol.
  871. try:
  872. # SIGXCPU's default disposition (how the soft CPU limit stops the child)
  873. # dumps core, and the child inherits the host's RLIMIT_CORE — a CPU
  874. # timeout would otherwise write a large memory-bearing core file into
  875. # the workspace. Forbid core dumps first so the timeout path leaves none.
  876. resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
  877. # Soft limit at cpuSeconds fires SIGXCPU (its default disposition
  878. # terminates the child; the host classifies that close as a timeout).
  879. # Hard limit at +1s is a SIGKILL backstop for a program that traps
  880. # SIGXCPU and keeps burning CPU.
  881. cpu_soft, cpu_hard = _clamped(
  882. resource.RLIMIT_CPU, boot["cpuSeconds"], boot["cpuSeconds"] + 1
  883. )
  884. resource.setrlimit(resource.RLIMIT_CPU, (cpu_soft, cpu_hard))
  885. # Darwin maps the multi-GB dyld shared cache into every process at
  886. # exec, so any practical RLIMIT_AS cap sits below current usage and
  887. # the kernel rejects it — the child would die here on every run. Skip
  888. # the address-space cap there; RLIMIT_CPU and the host's wall-clock
  889. # ceiling still bound the run.
  890. if sys.platform != "darwin":
  891. addr_bytes = int(boot["addressSpaceBytes"])
  892. effective_as = _clamped(resource.RLIMIT_AS, addr_bytes, addr_bytes)
  893. resource.setrlimit(resource.RLIMIT_AS, effective_as)
  894. # The host rejected an output budget too large for the CONFIGURED
  895. # addressSpaceMb, but a launch environment can inherit a STRICTER
  896. # RLIMIT_AS (e.g. a `ulimit -v` wrapper below addressSpaceMb), which
  897. # `_clamped` correctly lowers the effective limit to — leaving the
  898. # budgets validated against a ceiling the child never gets. Re-check
  899. # both budgets against the EFFECTIVE soft limit here, mirroring the
  900. # host gate (each budget times the worst-case Unicode multiple must
  901. # fit the room left after the interpreter baseline), and fail loud at
  902. # boot rather than letting a near-budget output OOM mid-run. The
  903. # constants match src/index.ts's OUTPUT_BUDGET_WORST_CASE_ADDRESS_
  904. # SPACE_MULTIPLE and INTERPRETER_BASELINE_BYTES.
  905. # `effective_soft` is always finite, so the re-check is
  906. # unconditional: `_clamped` was asked for the finite `addr_bytes` on
  907. # both sides, and each of its branches returns either that value or a
  908. # `min` with an inherited bound -- RLIM_INFINITY is not reachable. A
  909. # guard here would have silently skipped the whole re-check on the
  910. # branch it claimed to protect.
  911. effective_soft = effective_as[0]
  912. budgetable = effective_soft - _INTERPRETER_BASELINE_BYTES
  913. for _budget_key in ("maxLogBytes", "maxValueBytes"):
  914. if int(boot[_budget_key]) * _OUTPUT_BUDGET_WORST_CASE_MULTIPLE >= budgetable:
  915. raise ValueError(
  916. "config.%s is too large for the inherited RLIMIT_AS of %d bytes "
  917. "(a near-budget output would breach it during encode); "
  918. "lower the budget or raise the inherited address-space limit"
  919. % (_budget_key, effective_soft)
  920. )
  921. except _BaseException as exc: # noqa: BLE001 -- report every failure to host; `_BaseException` is a pre-program local
  922. channel.send_sync(
  923. {
  924. "type": "done",
  925. "error": {
  926. "kind": "exception",
  927. # Exception-only rendering: format_exc() would embed the
  928. # absolute installed bootstrap.py path in model-visible
  929. # durable output, leaking host paths into transcripts.
  930. "message": "bootstrap: applying resource limits failed\n"
  931. + "".join(
  932. traceback.format_exception_only(type(exc), exc)
  933. ),
  934. },
  935. }
  936. )
  937. return
  938. logs = LogBuffer(
  939. int(boot["maxLogBytes"]),
  940. # The sink writes through the def-time bound encode+write primitives
  941. # (not _send_sync_cls, whose body still resolves _encode_json_plain and
  942. # self.write_encoded at call time) so a rebind cannot break a log frame.
  943. sink=lambda text, truncated=False: _write_encoded_cls(
  944. _encode_plain_cls(
  945. {"type": "log", "text": text, **({"truncated": True} if truncated else {})}
  946. )
  947. ),
  948. )
  949. # 2. Wire the tools proxies and the ack.
  950. #
  951. # Each entry records the reply Future AND the loop it was created on. Model
  952. # code may call a binding from a THREAD it started, spelled
  953. # ``asyncio.run(tools.x(...))`` or its own new loop in that thread, so a
  954. # Future here can belong to a loop other than the one ``_pump_replies`` runs
  955. # on. ``asyncio.Future`` is not thread-safe: completing it from another
  956. # thread does not wake its own loop, so the pump schedules the completion on
  957. # the owning loop via ``call_soon_threadsafe`` (see ``_pump_replies``) rather
  958. # than calling ``set_result`` directly.
  959. pending: dict[int, tuple[asyncio.AbstractEventLoop, asyncio.Future[Any]]] = {}
  960. next_id = 0
  961. # Serializes the id claim + write + counter advance in ``dispatch`` against
  962. # both other binding-calling threads and the pump's ``pop``. ``dispatch`` may
  963. # run concurrently on several loops/threads, and the host answers a ``call``
  964. # only when its id is the exact successor of the last one — so ids must reach
  965. # the wire in the order they are claimed. Holding this lock across the write
  966. # (not just the counter arithmetic) is what keeps two threads' frames from
  967. # interleaving on fd 3 out of id order, which the host would reject.
  968. pending_lock = threading.Lock()
  969. error_classes: dict[str, type] = {}
  970. async def dispatch(global_name: str, name: str, args: Any) -> Any:
  971. nonlocal next_id
  972. error_class = error_classes.get(global_name)
  973. def call_failure(message: str) -> BaseException:
  974. # The namespace's declared rejection contract (e.g. Code Mode's
  975. # ToolCallError with .toolName) when present; RuntimeError keeps
  976. # the pre-errorClass behavior for namespaces that declared none.
  977. if error_class is not None:
  978. return error_class(name, message)
  979. return _RuntimeError_cls(message)
  980. # Validate the argument shape before claiming an id, so a rejected call
  981. # leaves no gap in the sequence the host checks. json.dumps would coerce
  982. # a non-string dict key or non-finite float rather than raise (allow_nan
  983. # is off, but key coercion still slips through), silently corrupting what
  984. # the tool receives. Reject up front through the call's error contract.
  985. violation = _lossless_json_violation_cls(args)
  986. if violation is not None:
  987. raise call_failure(f"binding arguments must be lossless JSON ({violation})")
  988. # Ids are consecutive from 0 with NO gaps: the host answers a `call` only
  989. # when its id is the exact successor of the last one, which bounds the
  990. # state it retains to a single number. A frame that never reaches the
  991. # host must therefore not consume an id, so the counter advances only
  992. # once the write has succeeded.
  993. #
  994. # The whole claim-write-advance runs under ``pending_lock`` because a
  995. # binding may be called from more than one thread/loop at once (the model
  996. # can start a thread that runs ``asyncio.run(tools.x(...))``). Without
  997. # the lock two callers could claim the same id, or write their frames to
  998. # fd 3 in an order that does not match their ids — either of which the
  999. # host rejects as an out-of-sequence call. The Future's own loop is
  1000. # captured here so ``_pump_replies`` can complete it thread-safely.
  1001. loop = _get_event_loop_cls()
  1002. with pending_lock:
  1003. call_id = next_id
  1004. fut: asyncio.Future[Any] = loop.create_future()
  1005. pending[call_id] = (loop, fut)
  1006. try:
  1007. _write_encoded_cls(
  1008. _encode_plain_cls(
  1009. {
  1010. "type": "call",
  1011. "id": call_id,
  1012. "global": global_name,
  1013. "name": name,
  1014. "args": args,
  1015. }
  1016. )
  1017. )
  1018. except (TypeError, ValueError) as exc:
  1019. pending.pop(call_id, None)
  1020. raise call_failure(
  1021. f"binding arguments must be lossless JSON: {exc}"
  1022. ) from exc
  1023. next_id += 1
  1024. try:
  1025. return await fut
  1026. except _BindingRejection_cls as exc:
  1027. raise call_failure(_str(exc)) from None
  1028. namespaces: dict[str, Any] = {}
  1029. for entry in boot["namespaces"]:
  1030. namespaces[entry["global"]] = _Namespace(
  1031. entry["global"], entry["names"], dispatch
  1032. )
  1033. declared = entry.get("errorClass")
  1034. if declared:
  1035. error_class = _make_error_class(
  1036. declared["name"], declared["memberNameProperty"]
  1037. )
  1038. error_classes[entry["global"]] = error_class
  1039. # The class is program-visible under its own name so model code
  1040. # can `except ToolCallError as e:` and read the member property.
  1041. namespaces[declared["name"]] = error_class
  1042. channel.send_sync({"type": "boot-ack"})
  1043. # 3. Start a reply-pump task before the run message: replies can arrive
  1044. # interleaved with the run's own binding traffic.
  1045. # The pump's frame reader is bound here, before the program runs: the
  1046. # pump itself starts AFTER the program's top-level statements (no suspension
  1047. # point between create_task and `await __dsh_main__`), so a body-local
  1048. # `channel.read_frame_async` lookup would resolve a rebound class method.
  1049. pump_read = channel.read_frame_async
  1050. reply_task = asyncio.get_event_loop().create_task(
  1051. _pump_replies(channel, pending, pending_lock, pump_read)
  1052. )
  1053. # 4. Read the run message.
  1054. run = channel.read_frame()
  1055. if run is None or run.get("type") != "run":
  1056. reply_task.cancel()
  1057. raise RuntimeError("bootstrap: expected run frame on fd 3")
  1058. program: str = run["program"]
  1059. # 5. Install log capture — ``print``, tracebacks, and ordinary ``sys.stdout``
  1060. # writes funnel into the LogBuffer. The real fds stay open (host uses
  1061. # stderr for stray-byte accounting) but the Python-visible streams point
  1062. # at the buffer.
  1063. sys.stdout = _LogStream(logs) # type: ignore[assignment]
  1064. sys.stderr = _LogStream(logs) # type: ignore[assignment]
  1065. out_stream, err_stream = sys.stdout, sys.stderr
  1066. # The ORIGINAL std streams are bound here, before the program runs, so the
  1067. # settlement flush can drain bytes a program wrote through them without an
  1068. # explicit flush. The bootstrap only replaces `sys.stdout`/`sys.stderr` with
  1069. # the `_LogStream`; `sys.__stdout__`/`sys.__stderr__` (and C-ext stdio
  1070. # layered on the same fds) are untouched, and their block-buffered bytes are
  1071. # lost when the host SIGTERMs the child right after the done frame — the
  1072. # default SIGTERM disposition terminates without interpreter finalization.
  1073. # Binding the names here (before the program) makes them immune to a
  1074. # `sys.__stdout__ = boom` rebind in model code; `None` under `-S`-style
  1075. # redirects is guarded at flush time.
  1076. # Bind the FLUSH METHODS, not the stream objects: the settlement flush
  1077. # loop iterates callables, and a bare TextIOWrapper object is not callable —
  1078. # invoking it would raise TypeError and be swallowed by the loop's except,
  1079. # silently disabling the drain. A bound method captures its stream at
  1080. # binding time, so a later `sys.__stdout__ = boom` rebind cannot redirect
  1081. # it; `None` (stream absent) is guarded at flush time.
  1082. _stdout_orig = sys.__stdout__.flush if sys.__stdout__ is not None else None
  1083. _stderr_orig = sys.__stderr__.flush if sys.__stderr__ is not None else None
  1084. # 6. Compile the program as the body of an async function, matching the
  1085. # seam contract (`CodeRunRequest.program` is an async-function body: top-level
  1086. # `await` and `return` both work, and the returned value is the completion).
  1087. # AST-splicing the parsed body into an `async def` keeps every statement's
  1088. # original line number, so a traceback points at the model's own source.
  1089. ns: dict[str, Any] = {
  1090. "__name__": "__main__",
  1091. "__builtins__": __builtins__,
  1092. **namespaces,
  1093. }
  1094. # Read the enforcement callable and its budget into this frame's locals
  1095. # BEFORE the program runs: model code can rebind this module's globals
  1096. # (the bootstrap IS ``__main__``), and a frame local is not a module
  1097. # attribute, so a later ``__main__._DIE_IF_CPU_EXHAUSTED = ...`` cannot
  1098. # change which callable the post-check below invokes. This defeats the
  1099. # one-line rebind, not a determined `sys._getframe` walk; the unforgeable
  1100. # bounds are the RLIMIT_CPU hard limit and the host wall clock
  1101. # (see _make_cpu_enforcer).
  1102. die_if_cpu_exhausted = _DIE_IF_CPU_EXHAUSTED
  1103. # The settlement recheck compares against the EFFECTIVE soft CPU limit
  1104. # (`cpu_soft`, clamped to any stricter inherited limit above), NOT the
  1105. # configured `cpuSeconds`. When the deployment inherited a soft limit below
  1106. # the configured value, a program that traps SIGXCPU, burns past the
  1107. # inherited soft, and returns inside the soft-to-hard gap must be reported as
  1108. # a timeout — checking the configured value would falsely pass it and bypass
  1109. # the inherited limit.
  1110. cpu_seconds = cpu_soft
  1111. # Same capture, same reason, for the failure path and the send that follows
  1112. # it. The reporter was a module-global lookup inside the `except` block, so
  1113. # ``import __main__; __main__._SAFE_MODEL_TRACEBACK = ...`` put model code
  1114. # there with no guard around it; the flush and send were attribute lookups
  1115. # on the stream and channel CLASSES, which ``__main__._LogStream.flush_line
  1116. # = ...`` rebinds just as easily. All four run AFTER the handler, where a
  1117. # throw costs the `done` frame and the host reports a wall-clock timeout
  1118. # instead of the model's exception. Binding the callables now fixes what
  1119. # runs; what they in turn reach is closed over in _make_failure_reporter.
  1120. safe_model_traceback = _SAFE_MODEL_TRACEBACK
  1121. flush_out = out_stream.flush_line
  1122. flush_err = err_stream.flush_line
  1123. # `done` is either a pre-encoded frame STRING (a `_done_with_value` success:
  1124. # the completion value was serialized at its validation point, inside the try,
  1125. # so a later send never re-walks the live value a mutating daemon thread could
  1126. # have changed) or a dict ERROR frame (a rejection or the exception handler,
  1127. # which carry no live model value). `send_done` posts whichever form, going
  1128. # DIRECTLY through a bound `_encode_json_plain` and a bound `write_encoded` —
  1129. # never through `channel.send_sync`, whose body re-resolves `self.write_encoded`
  1130. # and `self`'s module-level `_encode_json_plain` at call time.
  1131. #
  1132. # The program runs as `__main__`, so `import __main__; __main__.ProtocolChannel
  1133. # .send_sync = boom` or `__main__._encode_json_plain = boom` would otherwise
  1134. # re-resolve the send/encode to a rebranded callable at call time and, when
  1135. # that replacement raises, skip the `done` frame and downgrade a settled
  1136. # verdict to a host-side worker-exit (the binding-all-names regression test
  1137. # pins this). Same reason `flush_out`/`flush_err`/`safe_model_traceback` are
  1138. # bound above.
  1139. encode_plain_bound = _encode_json_plain
  1140. write_encoded_bound = channel.write_encoded
  1141. # The completion-frame builder is bound into a LOCAL here, before the
  1142. # program runs: `done = _done_with_value(...)` below sits after the program
  1143. # (which is `__main__`) may have rebound `__main__._done_with_value`, so a
  1144. # module-global lookup at call time would let a one-line rebind rewrite a
  1145. # legitimate success into an `exception`. Binding it (with its own def-time
  1146. # default-captured `_check_done_value`/`_encode_json_plain`) makes the entry
  1147. # name immune.
  1148. done_with_value_bound = _done_with_value
  1149. # The fallback primitives are bound into LOCALS here, before the program
  1150. # runs, so `send_done`'s except arm does not read module globals at call
  1151. # time. This bootstrap is `__main__`, so `__main__._os_write = boom` (or
  1152. # `__main__._FALLBACK_DONE_FRAME`, `__main__._memoryview`) would otherwise
  1153. # rebind exactly the names the fallback reads, reopening the single-line-
  1154. # rebind hole the fallback exists to close.
  1155. _os_write_local = _os_write
  1156. _memoryview_local = _memoryview
  1157. _fallback_frame_local = _FALLBACK_DONE_FRAME
  1158. def send_done(payload: dict[str, Any] | str) -> None:
  1159. try:
  1160. if _isinstance(payload, _str):
  1161. write_encoded_bound(payload)
  1162. else:
  1163. write_encoded_bound(encode_plain_bound(payload))
  1164. except: # noqa: BLE001, E722 -- a rebind must not cost the done frame; bare except avoids naming BaseException
  1165. # `encode_plain_bound`/`write_encoded_bound` are bound callables, but
  1166. # their BODIES still resolve transitive module globals at call time —
  1167. # `_encode_json_plain` reaches `_dump_scalar`/`_dump_string`/`json.dumps`,
  1168. # `write_encoded` reaches `os.write` (via the `os` module). This
  1169. # bootstrap is `__main__`, so `__main__._dump_scalar = boom` (or
  1170. # `__main__.os = ...`) makes the error-frame encode/write throw AFTER
  1171. # the `except` block, which would drop the `done` frame and downgrade a
  1172. # settled `exception` verdict to a host-side `worker-exit`. Write a fixed
  1173. # literal done frame with the LOCALLY-BOUND `_os_write_local` and
  1174. # `_memoryview_local` (captured before the program runs, so a one-line
  1175. # rebind of the module global cannot change them) so the host still gets
  1176. # a verdict. The literal is JSON-valid and newline-terminated; the lock
  1177. # is the channel's, so the write is serialized against any concurrent
  1178. # writer.
  1179. with channel._write_lock:
  1180. view = _memoryview_local(_fallback_frame_local)
  1181. while view:
  1182. view = view[_os_write_local(channel._fd, view):]
  1183. max_value_bytes = int(boot["maxValueBytes"])
  1184. done: dict[str, Any] | str
  1185. try:
  1186. # filename="<model>" keeps the source label consistent with the later
  1187. # compile(wrapped, "<model>", ...) and the runtime traceback filtering
  1188. # (safe_model_traceback drops frames whose filename is not "<model>");
  1189. # the default "<unknown>" would leak a different label into model-visible
  1190. # syntax diagnostics.
  1191. module = ast.parse(program, filename="<model>")
  1192. wrapper = ast.AsyncFunctionDef(
  1193. name="__dsh_main__",
  1194. args=ast.arguments(
  1195. posonlyargs=[], args=[], vararg=None,
  1196. kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[],
  1197. ),
  1198. body=module.body or [ast.Pass()],
  1199. decorator_list=[],
  1200. returns=None,
  1201. )
  1202. # Anchor the synthetic wrapper on the first real statement (or line 1 for
  1203. # an empty program) so fix_missing_locations does not stamp it at 0.
  1204. anchor = module.body[0] if module.body else ast.parse("pass").body[0]
  1205. ast.copy_location(wrapper, anchor)
  1206. wrapped = ast.Module(body=[wrapper], type_ignores=[])
  1207. ast.fix_missing_locations(wrapped)
  1208. # `dont_inherit=True` stops this module's `from __future__ import
  1209. # annotations` (line 14) from leaking into the program's compile: PEP 563
  1210. # would otherwise stringify the program's type annotations, changing the
  1211. # semantics of a legal program that reads `f.__annotations__` at runtime.
  1212. code = compile(wrapped, "<model>", "exec", dont_inherit=True)
  1213. exec(code, ns) # noqa: S102 -- defines __dsh_main__; executing model code is the point
  1214. value = await ns["__dsh_main__"]()
  1215. die_if_cpu_exhausted(cpu_seconds)
  1216. # Flush the log buffers BEFORE metering and framing the completion value.
  1217. # `_done_with_value` materializes the value's escaped JSON form to meter
  1218. # it and then pre-encodes the admitted value into its frame (see its
  1219. # docstring for the TOCTOU rationale) — several copies of a near-budget
  1220. # value live at once (see OUTPUT_BUDGET_WORST_CASE_ADDRESS_SPACE_MULTIPLE).
  1221. # Any unflushed log pending would add its own bytes to that peak, so a
  1222. # `maxLogBytes` and a `maxValueBytes` each admitted alone by the load gate
  1223. # could together breach RLIMIT_AS. Flushing first frees the log pending so
  1224. # the value frame's peak stands alone against the address space.
  1225. flush_out()
  1226. flush_err()
  1227. done = done_with_value_bound(value, max_value_bytes)
  1228. except _BaseException as exc: # noqa: BLE001 -- report every failure to host; `_BaseException` is a pre-program local, not a rebindable module global
  1229. done = {
  1230. "type": "done",
  1231. "error": {
  1232. "kind": "exception",
  1233. # Cap the diagnostic BEFORE it crosses the wire: a program can
  1234. # raise with a gigabytes-long message, and formatting/sending
  1235. # it whole would allocate on both sides before the host's own
  1236. # cap runs. Byte-cap at maxValueBytes with the host's marker
  1237. # text so the truncated diagnostic reads identically wherever
  1238. # the cap was applied. The rendering is wrapped because the
  1239. # `done` send below sits outside this handler: a throw while
  1240. # formatting would skip it and strand the host on fd 3 until
  1241. # maxWallMs (see _make_failure_reporter).
  1242. "message": safe_model_traceback(exc, max_value_bytes),
  1243. },
  1244. }
  1245. # Flush any print output not terminated by a newline (a traceback always
  1246. # ends in one, but `print(x, end="")` or a bare write may not), so the
  1247. # final partial line is not silently dropped. The success path already
  1248. # flushed before framing the value; this is an idempotent no-op there and
  1249. # the flush the exception path needs.
  1250. #
  1251. # Guarded because `done` is ALREADY DECIDED here: on the exception path the
  1252. # handler above built it, and a flush that raises (its join/encode under
  1253. # memory pressure, after the program left a near-maxLogBytes pending and then
  1254. # allocated toward RLIMIT_AS) would skip `send_done` and downgrade a run the
  1255. # child already classified as `exception` into a host-side `worker-exit`.
  1256. # Losing the log tail is the lesser outcome, and the marker the ledger
  1257. # already pushed still reports the truncation. Same rule as
  1258. # `_make_failure_reporter`: a settled verdict must not be swallowed by the
  1259. # reporting that follows it.
  1260. for _flush in (flush_out, flush_err, _stdout_orig, _stderr_orig):
  1261. if _flush is None:
  1262. continue
  1263. try:
  1264. _flush()
  1265. except _BaseException: # noqa: BLE001 -- swallow ONLY the log tail; `done` must reach the host; `_BaseException` is a pre-program local
  1266. pass
  1267. reply_task.cancel()
  1268. send_done(done)
  1269. async def _pump_replies(
  1270. channel: ProtocolChannel,
  1271. pending: dict[int, tuple[asyncio.AbstractEventLoop, asyncio.Future[Any]]],
  1272. pending_lock: "threading.Lock",
  1273. # The frame reader is a bound method captured by _run BEFORE the program
  1274. # runs (see the create_task site), so a rebind of the class attribute cannot
  1275. # redirect it.
  1276. _read_frame: Any,
  1277. # Bound as DEFAULT ARGUMENTS so they are captured at def/import time, before
  1278. # ANY model code runs. This bootstrap IS `__main__`, so `__main__.RuntimeError
  1279. # = ...` (or `__main__._BindingRejection`, `__main__.str`, `__main__.bool`)
  1280. # as a program top-level statement would otherwise rebind the module globals
  1281. # these clauses resolve at runtime. A body-local `X = X` binding is too late:
  1282. # `_run` reaches `await __dsh_main__` (whose top-level statements run first)
  1283. # with no suspension point after `create_task`, so the model's rebind executes
  1284. # before the pump body. Defaults are evaluated in the enclosing scope at def
  1285. # time, truly before the program.
  1286. _RuntimeError: Any = RuntimeError,
  1287. _BindingRejection: Any = _BindingRejection,
  1288. _str: Any = str,
  1289. _bool: Any = bool,
  1290. ) -> None:
  1291. """Background task: read reply frames and settle pending futures.
  1292. Cancelled after ``done`` is posted. Unknown ids and post-settlement replies
  1293. are ignored (mirrors the worker backend's hostile-peer stance, though here
  1294. the host is the trusted side; the guards defend against races).
  1295. A pending Future may belong to a loop other than this pump's — the model can
  1296. call a binding from a thread running its own loop (``asyncio.run(tools.x())``).
  1297. ``asyncio.Future`` is not thread-safe, so the completion is scheduled on the
  1298. Future's OWN loop via ``call_soon_threadsafe`` rather than mutated here; a
  1299. direct ``set_result`` would never wake the waiting loop and the call would
  1300. hang to the wall clock. The ``pop`` shares ``pending_lock`` with ``dispatch``
  1301. so a reply cannot race the claim that registers its id.
  1302. """
  1303. def complete(fut: asyncio.Future[Any], ok: bool, value: Any, message: Any) -> None:
  1304. # Runs on the Future's own loop. `done()` re-checked here because
  1305. # cancellation or a duplicate reply may have settled it between the pop
  1306. # and this callback.
  1307. if fut.done():
  1308. return
  1309. if ok:
  1310. fut.set_result(value)
  1311. else:
  1312. fut.set_exception(_BindingRejection(_str(message)))
  1313. while True:
  1314. frame = await _read_frame()
  1315. if frame is None:
  1316. return
  1317. if frame.get("type") != "reply":
  1318. continue
  1319. with pending_lock:
  1320. entry = pending.pop(frame.get("id"), None)
  1321. if entry is None:
  1322. continue
  1323. loop, fut = entry
  1324. ok = _bool(frame.get("ok"))
  1325. value = frame.get("value")
  1326. message = frame.get("message")
  1327. try:
  1328. loop.call_soon_threadsafe(complete, fut, ok, value, message)
  1329. except _RuntimeError: # `_RuntimeError` is a pre-program local, not a rebindable module global
  1330. # The Future's loop has already closed — the thread that ran
  1331. # `asyncio.run(tools.x(...))` finished (its coroutine was cancelled
  1332. # or it exited) before this reply arrived, so nothing awaits the
  1333. # Future and the reply is moot. Drop it; scheduling onto a closed
  1334. # loop raises RuntimeError, and letting that escape would kill the
  1335. # pump and strand every later reply — the exact failure class this
  1336. # cross-loop delivery exists to prevent. An abandoned call's pending
  1337. # entry is not leaked: it is popped here when its reply arrives
  1338. # (dispatch's cancellation does not remove it), so stranded entries
  1339. # are bounded by the number of calls THIS run itself issued.
  1340. continue
  1341. def _encode_json_plain(value: Any) -> str:
  1342. """Encode JSON-plain data iteratively, byte-identical to compact ``json.dumps``.
  1343. ``json.dumps`` recurses one Python frame per nesting level and raises
  1344. ``RecursionError`` a few thousand levels deep, but the seam's
  1345. ``CodeJsonValue`` has no depth limit — a valid deeply nested completion or
  1346. call argument below the byte budget must cross intact (the host uses the
  1347. same iterative idiom in ``protocol.ts``). Accepts what the callers already
  1348. validated or constructed: ``None``/``bool``/``int``/finite ``float``/
  1349. ``str``, exact ``list``/``tuple``, and exact ``dict`` with ``str`` keys.
  1350. Scalar encoding delegates to ``json.dumps`` (string escaping, float repr)
  1351. so the bytes match; non-finite floats still raise ``ValueError`` exactly
  1352. like ``allow_nan=False``.
  1353. Containers are classified by EXACT type and traversed through the unbound
  1354. built-in methods rather than the instance's own: a ``dict``/``list``
  1355. subclass can override ``items``, ``keys``, ``__iter__``, ``__len__``, or
  1356. ``__getitem__``, and the validators only see the container it subclasses,
  1357. so an instance-method call here could emit different data than the walk
  1358. that metered and approved it. ``_check_done_value`` and
  1359. ``_lossless_json_violation`` reject subclasses outright, so this path only
  1360. ever sees exact containers; classifying on exact type keeps that agreement
  1361. checkable at one glance instead of resting on the caller.
  1362. """
  1363. # O(DEPTH) auxiliary space, not O(width). A container pushes ONE cursor frame
  1364. # that pulls its children one at a time and writes each into the shared buffer,
  1365. # rather than one stack entry (plus a separator marker) per child: a flat
  1366. # `[0] * 6_000_000` encodes to ~12 MB but per-element frames are ~400 MB — an
  1367. # RLIMIT_AS death on a value `_check_done_value` already admitted (it walks by
  1368. # depth as well). The output string is the only width-proportional
  1369. # allocation, and its size the caller metered within budget. `io.StringIO`
  1370. # accumulates without the intermediate `"".join(chunks)` second copy. A cursor
  1371. # frame is [kind, iterator, wrote_any]; a visit frame is (VISIT, value).
  1372. buffer = io.StringIO()
  1373. exhausted = object()
  1374. visit, list_cursor, dict_cursor = 0, 1, 2
  1375. stack: list[Any] = [(visit, value)]
  1376. while stack:
  1377. frame = stack.pop()
  1378. kind = frame[0]
  1379. if kind == list_cursor:
  1380. iterator, wrote_any = frame[1], frame[2]
  1381. child = next(iterator, exhausted)
  1382. if child is exhausted:
  1383. buffer.write("]")
  1384. continue
  1385. if wrote_any:
  1386. buffer.write(",")
  1387. else:
  1388. frame[2] = True
  1389. stack.append(frame)
  1390. stack.append((visit, child))
  1391. continue
  1392. if kind == dict_cursor:
  1393. iterator, wrote_any = frame[1], frame[2]
  1394. entry = next(iterator, exhausted)
  1395. if entry is exhausted:
  1396. buffer.write("}")
  1397. continue
  1398. key, item = entry
  1399. if wrote_any:
  1400. buffer.write(",")
  1401. else:
  1402. frame[2] = True
  1403. buffer.write(_dump_scalar(key))
  1404. buffer.write(":")
  1405. stack.append(frame)
  1406. stack.append((visit, item))
  1407. continue
  1408. current = frame[1]
  1409. current_type = type(current)
  1410. if current_type is list or current_type is tuple:
  1411. buffer.write("[")
  1412. stack.append([list_cursor, iter(current), False])
  1413. elif current_type is dict:
  1414. buffer.write("{")
  1415. stack.append([dict_cursor, iter(dict.items(current)), False])
  1416. else:
  1417. buffer.write(_dump_scalar(current))
  1418. return buffer.getvalue()
  1419. def _dump_scalar(value: Any) -> str:
  1420. """One scalar as compact JSON, byte-compatible with the host's encoder.
  1421. ``ensure_ascii=False`` keeps non-ASCII text as raw UTF-8 — the default
  1422. backslash-u escaping would make the child count ``"é"`` as 8 bytes where
  1423. the host meter (and the worker backend) count its UTF-8 JSON form as 4,
  1424. splitting the budget the two sides are supposed to share. Strings route
  1425. through :func:`_dump_string`, which restores the escaping for the one class
  1426. of character UTF-8 cannot hold. Floats route through :func:`_dump_float`
  1427. because CPython's ``repr`` and ECMAScript's Number-to-String disagree on
  1428. spelling.
  1429. Dispatch is on EXACT type, matching the validators: a ``float`` subclass
  1430. reaching :func:`_dump_float` would have its overridden ``__repr__`` read as
  1431. the number's digits, so ``F(2.5)`` whose ``__repr__`` returns ``"1.0"``
  1432. would serialize as ``1``. ``json.dumps`` then refuses any subclass by
  1433. ``TypeError`` instead of emitting a value nothing validated; the callers
  1434. reject subclasses first, so this is the encoder refusing to be the place a
  1435. validation gap turns into corrupted output.
  1436. """
  1437. if type(value) is float:
  1438. return _dump_float(value)
  1439. if type(value) is str:
  1440. return _dump_string(value)
  1441. if value is None or type(value) is bool or type(value) is int:
  1442. return json.dumps(value, ensure_ascii=False, allow_nan=False)
  1443. raise TypeError(f"unsupported type ({type(value).__name__})")
  1444. # A surrogate code unit, and an adjacent high-low pair. Python stores an astral
  1445. # character as ONE code point, so a surrogate reaching these patterns is either
  1446. # lone or half of a pair the program spelled out code unit by code unit.
  1447. _SURROGATE = re.compile("[\ud800-\udfff]")
  1448. _SURROGATE_PAIR = re.compile("[\ud800-\udbff][\udc00-\udfff]")
  1449. def _combine_surrogate_pair(match: re.Match[str]) -> str:
  1450. """Fold one spelled-out high-low pair into the astral code point it names."""
  1451. high, low = match.group(0)
  1452. return chr(0x10000 + ((ord(high) - 0xD800) << 10) + (ord(low) - 0xDC00))
  1453. def _dump_string(text: str) -> str:
  1454. """One string as compact JSON, byte-identical to the host's ``JSON.stringify``.
  1455. ``ensure_ascii=False`` cannot render a surrogate code unit: UTF-8 has no
  1456. encoding for one, so the frame write would raise and the run would strand
  1457. until the wall clock. JSON carries it as the ASCII escape ``\\ud800``, which
  1458. the host's ``JSON.parse`` reads back as the same UTF-16 code unit and its
  1459. ``JSON.stringify`` re-emits identically — so the shared seam
  1460. (``CodeJsonValue``, ``snapshotJsonValue``, the worker backend) keeps a
  1461. lone-surrogate string instead of failing the value. An adjacent high-low
  1462. pair is folded into its astral code point FIRST: the host holds strings as
  1463. UTF-16, where those two code units and the single character are the same
  1464. string, and the raw 4-byte form is what the host would emit — escaping the
  1465. halves separately would charge 12 bytes against a budget the host meters at
  1466. 4. Every remaining surrogate is lone and becomes six ASCII bytes, matching
  1467. the host exactly.
  1468. @param text: the string to encode.
  1469. @return: its compact JSON form, always UTF-8-encodable.
  1470. """
  1471. rendered = json.dumps(text, ensure_ascii=False)
  1472. if _SURROGATE.search(rendered) is None:
  1473. return rendered
  1474. return _SURROGATE.sub(
  1475. lambda match: "\\u%04x" % ord(match.group(0)),
  1476. _SURROGATE_PAIR.sub(_combine_surrogate_pair, rendered),
  1477. )
  1478. # How many bytes each byte that needs escaping adds beyond its raw self, as a
  1479. # ready-made (byte, surcharge) list so :func:`_json_string_cost` walks no
  1480. # branches per pass. ``"`` and ``\\`` take a one-character prefix; the five C0
  1481. # controls with a shorthand (``\\b\\f\\n\\r\\t``) likewise; every other C0
  1482. # control becomes a six-character ``\\uXXXX``.
  1483. _JSON_ESCAPE_SURCHARGES = [
  1484. (bytes((byte,)), 1 if byte in b'"\\\b\f\n\r\t' else 5)
  1485. for byte in [*range(0x20), ord('"'), ord("\\")]
  1486. ]
  1487. # Per-byte JSON-string serialized cost (the byte itself plus its escape
  1488. # surcharge), indexed by byte value. Lets :func:`_cap_message` accumulate the
  1489. # serialized cost of a growing prefix in one O(1) step per byte without building
  1490. # the escaped form. A non-ASCII byte stays raw (cost 1); a C0 control or ``"``/
  1491. # ``\\`` carries its surcharge from :data:`_JSON_ESCAPE_SURCHARGES`.
  1492. _JSON_BYTE_COST = [1] * 256
  1493. for _escaped_byte, _surcharge in _JSON_ESCAPE_SURCHARGES:
  1494. _JSON_BYTE_COST[_escaped_byte[0]] = 1 + _surcharge
  1495. def _json_str_cost(text: str) -> int:
  1496. """Byte length of ``text``'s JSON string form, WITHOUT building that form.
  1497. The str-side twin of :func:`_json_string_cost`, for the completion-value
  1498. meter. Measuring by materializing ``_dump_string(text).encode()`` allocates
  1499. the escaped copy plus its encode -- for a NUL-heavy string that is ~6x the
  1500. original each, so metering a value the budget would have REJECTED could
  1501. itself breach ``RLIMIT_AS`` and report ``exception`` where the contract
  1502. promises ``output-limit``.
  1503. The common case encodes once (~1x, well inside the load gate's envelope) and
  1504. counts escapes with the same C-level passes :func:`_json_string_cost` uses.
  1505. A string carrying surrogate code units has no UTF-8 form at all, so it takes
  1506. the exact path :func:`_dump_string` defines: fold each spelled-out high-low
  1507. pair into its astral character first (the host meters that as its raw 4-byte
  1508. form), then charge six ASCII bytes for every surviving lone surrogate and
  1509. count the rest from its encodable remainder.
  1510. @param text: the string to measure.
  1511. @return: the byte length of its JSON string form, quotes included.
  1512. """
  1513. try:
  1514. return _json_string_cost(text.encode("utf-8"))
  1515. except UnicodeEncodeError:
  1516. pass
  1517. folded = _SURROGATE_PAIR.sub(_combine_surrogate_pair, text)
  1518. lone = len(_SURROGATE.findall(folded))
  1519. # Six ASCII bytes per lone surrogate; the remainder is ordinary text whose
  1520. # own quotes are dropped here because the outer call adds them once.
  1521. without = _SURROGATE.sub("", folded)
  1522. return _json_string_cost(without.encode("utf-8")) + lone * 6
  1523. def _json_string_cost(raw: bytes) -> int:
  1524. """UTF-8 byte length of one string's JSON form, WITHOUT building that form.
  1525. Used by :class:`LogBuffer` to charge a log entry what it will actually cost
  1526. on the wire. Building ``json.dumps(text)`` to measure it would allocate a
  1527. second copy up to six times the original — the very allocation the ledger's
  1528. cheap pre-check exists to avoid, and enough to breach ``RLIMIT_AS`` on a
  1529. large control-heavy line. Counts exactly what :func:`_dump_scalar`'s
  1530. ``ensure_ascii=False`` output holds: the two quotes, each escaped byte's
  1531. surcharge from :data:`_JSON_ESCAPE_SURCHARGES`, and the raw bytes themselves
  1532. (non-ASCII stays raw, so its UTF-8 length already counts). Uses a fixed
  1533. number of C-level ``count`` passes — allocating nothing, unlike a
  1534. ``translate`` filter — because the caller admits up to ~4x the remaining
  1535. budget of bytes here and a per-byte Python loop over it would cost more than
  1536. the encode being avoided.
  1537. @param raw: the entry's UTF-8 bytes.
  1538. @return: the byte length of its JSON string form, quotes included.
  1539. """
  1540. extra = 0
  1541. for byte, surcharge in _JSON_ESCAPE_SURCHARGES:
  1542. extra += raw.count(byte) * surcharge
  1543. return len(raw) + 2 + extra
  1544. def _dump_float(value: float) -> str:
  1545. """One finite float in ECMAScript ``Number::toString`` spelling.
  1546. CPython's ``repr`` and the host's ``String(number)`` name the same double
  1547. differently: ``1.0`` is ``"1.0"`` here but ``"1"`` there, ``1e-07`` pads the
  1548. exponent the host writes as ``1e-7``, and ``1e+21``/``2**60`` differ again.
  1549. Since the child meters the completion value against ``maxValueBytes`` and
  1550. the host re-meters the frame it parses, any spelling difference splits the
  1551. shared budget: ``return 1.0`` under ``maxValueBytes: 1`` used to be reported
  1552. as ``output-limit`` by the child while the host would have counted the
  1553. one-byte ``1`` it actually receives. Both sides also emit these bytes (the
  1554. child through :func:`_encode_json_plain`, the host through
  1555. ``encodeJsonPlain``), so the fix has to be in the shared speller, not in the
  1556. meter.
  1557. Implements ECMA-262 ``Number::toString`` radix 10 directly: ``repr``
  1558. already yields the shortest round-tripping decimal digits, and ``Decimal``
  1559. splits them into the significand ``s`` (``digits``, ``k`` of them) and
  1560. decimal exponent ``n`` the spec's cases select on. The integral values above
  1561. the JS safe range take the host's BigInt branch, whose exact digits differ
  1562. from the shortest-round-trip form (``2**60`` prints ``...846976``, not
  1563. ``...847000``).
  1564. """
  1565. if value != value or value in (float("inf"), float("-inf")):
  1566. # json.dumps(allow_nan=False) raises the same way; the callers reject
  1567. # non-finite floats before metering, so this is unreachable defense.
  1568. raise ValueError("Out of range float values are not JSON compliant")
  1569. if value == 0.0:
  1570. # Covers -0.0 too; callers reject it as non-lossless before this point.
  1571. return "0"
  1572. if value < 0:
  1573. return "-" + _dump_float(-value)
  1574. if value.is_integer() and value > float(2**53 - 1):
  1575. # The host's BigInt branch: exact digits, not shortest-round-trip.
  1576. return str(int(value))
  1577. parts = Decimal(repr(value)).normalize().as_tuple()
  1578. digits = "".join(str(digit) for digit in parts.digits)
  1579. k = len(digits)
  1580. n = parts.exponent + k
  1581. if k <= n <= 21:
  1582. return digits + "0" * (n - k)
  1583. if 0 < n <= 21:
  1584. return digits[:n] + "." + digits[n:]
  1585. if -6 < n <= 0:
  1586. return "0." + "0" * -n + digits
  1587. exponent = ("+" if n - 1 >= 0 else "-") + str(abs(n - 1))
  1588. return (digits if k == 1 else digits[0] + "." + digits[1:]) + "e" + exponent
  1589. def _check_done_value(value: Any, max_bytes: int):
  1590. """Meter a completion value's JSON byte size AND validate its lossless-JSON
  1591. shape in one bounded post-order walk; return ``None`` when it passes.
  1592. Folds what was formerly a losslessness walk followed by a separate byte
  1593. meter into one pass. Running the losslessness walk first materialized one
  1594. traversal tuple per element before any size cap: ``return [0] * 2000000``
  1595. under ``maxValueBytes: 64`` allocated millions of frames (an RLIMIT_AS
  1596. death) before the meter could reject it. Folding the byte bound into the
  1597. walk rejects over-budget BEFORE enqueuing a container's children — every
  1598. element is at least one JSON byte — so the walk stays O(cap). Same
  1599. JS-double-exact integer boundary, cycle detection (a leave marker pops each
  1600. container off ``on_path``), and type rejections as
  1601. :func:`_lossless_json_violation`, and the same byte accounting as
  1602. :func:`_encode_json_plain`. :func:`_lossless_json_violation` stays for the
  1603. binding-argument path, which carries no size cap.
  1604. EVERY type here is matched EXACTLY, containers and scalars alike, so a
  1605. subclass is rejected as an unsupported type rather than admitted by
  1606. ``isinstance``. A subclass can override the operators and methods this walk
  1607. and the encoder call, and they need not agree: a populated ``dict``
  1608. subclass whose ``items()`` returns ``[]`` would meter as ``{}``; a ``float``
  1609. subclass overriding ``__repr__`` passes the non-finite and negative-zero
  1610. checks by its real value but serializes as whatever the override says, since
  1611. :func:`_dump_float` reads ``repr``; an ``int`` subclass overriding ``__gt__``
  1612. and ``__lt__`` slips past the JS-safe-range bound while ``json.dumps``
  1613. emits its true C-level digits, so ``2**53 + 1`` reaches the host as
  1614. ``...992``; a ``str`` subclass overriding ``__len__`` returns 0 from the
  1615. pre-encode lower bound and admits an arbitrarily large string. In each case
  1616. the value the host receives differs from the one this walk approved. The
  1617. worker backend rejects the equivalent shapes by prototype identity and
  1618. ``typeof`` (``hasPlainObjectPrototype`` in ``worker-json.ts``); a ``bool``
  1619. is checked before ``int`` because it is an ``int`` subclass that IS
  1620. lossless JSON.
  1621. Returns ``("invalid-output", message)`` for a non-lossless value,
  1622. ``("output-limit", message)`` once the size crosses ``max_bytes``, or
  1623. ``None`` when the value is lossless JSON within budget.
  1624. Metering and validation interleave in this single traversal: each member is
  1625. costed the moment it is visited, and it is rejected the moment it trips
  1626. either check. A value that holds BOTH an over-budget member and an
  1627. invalid-typed member therefore resolves to whichever tripped FIRST in
  1628. visit order — both are rejects, and neither kind claims priority over the
  1629. other, so that first-trip order is not part of the seam contract; the
  1630. host side independently re-measures the value it receives.
  1631. """
  1632. js_safe = 2**53 - 1
  1633. def invalid(reason: str):
  1634. return ("invalid-output", f"program completion must be lossless JSON ({reason})")
  1635. over_budget = ("output-limit", f"completion value exceeded {max_bytes} bytes")
  1636. total = 0
  1637. on_path: set[int] = set()
  1638. # The walk uses O(DEPTH) space, not O(width). A container pushes ONE cursor
  1639. # frame that pulls its children one at a time, rather than one traversal
  1640. # frame per child: a flat `[0] * 6_000_000` serializes to ~12 MB (well within
  1641. # a modest budget) but one tuple per element is ~380 MB — an RLIMIT_AS death
  1642. # on a value the byte meter would admit, the very inversion this meter exists
  1643. # to prevent. A cursor frame is (kind, container, iterator); a visit frame is
  1644. # (VISIT, value, None). The upfront structural bound still rejects a wide
  1645. # forgery before any iteration begins.
  1646. exhausted = object()
  1647. visit, list_cursor, dict_cursor = 0, 1, 2
  1648. stack: list[tuple[int, Any, Any]] = [(visit, value, None)]
  1649. while stack:
  1650. frame = stack.pop()
  1651. kind = frame[0]
  1652. if kind == list_cursor:
  1653. container, iterator = frame[1], frame[2]
  1654. child = next(iterator, exhausted)
  1655. if child is exhausted:
  1656. on_path.discard(id(container))
  1657. continue
  1658. # Resume this cursor after the child is fully walked; the child goes
  1659. # on top so it is visited next (order does not affect the byte total).
  1660. stack.append(frame)
  1661. stack.append((visit, child, None))
  1662. continue
  1663. if kind == dict_cursor:
  1664. container, iterator = frame[1], frame[2]
  1665. entry = next(iterator, exhausted)
  1666. if entry is exhausted:
  1667. on_path.discard(id(container))
  1668. continue
  1669. key, item = entry
  1670. # Only an EXACT str key survives: bool and int coerce or raise, and a
  1671. # str SUBCLASS can override the ``__len__`` the bound below reads while
  1672. # the encoder emits its real characters.
  1673. if type(key) is not str:
  1674. return invalid(f"non-string dict key ({type(key).__name__})")
  1675. # The same string lower bound, before escaping the key.
  1676. if total + len(key) + 3 > max_bytes:
  1677. return over_budget
  1678. # Same counting rule as the string branch: a control-heavy KEY
  1679. # expands just as far, and `_dump_scalar` on a str is `_dump_string`.
  1680. total += _json_str_cost(key) + 1
  1681. if total > max_bytes:
  1682. return over_budget
  1683. stack.append(frame)
  1684. stack.append((visit, item, None))
  1685. continue
  1686. current = frame[1]
  1687. if current is None or type(current) is bool:
  1688. total += len(_dump_scalar(current).encode("utf-8"))
  1689. elif type(current) is str:
  1690. # Lower-bound BEFORE materializing the escaped form: every character
  1691. # is at least one UTF-8 byte plus the two quotes, so a huge or
  1692. # control-heavy string (whose escaped copy expands severalfold) is
  1693. # rejected without allocating that copy.
  1694. if total + len(current) + 2 > max_bytes:
  1695. return over_budget
  1696. # A lone surrogate has no UTF-8 form but a lossless JSON one — the
  1697. # ASCII ``\uXXXX`` escape :func:`_dump_string` emits — so it is
  1698. # metered, not rejected, matching the shared seam. Metered by
  1699. # COUNTING, not by building the escaped form: that copy plus its
  1700. # encode is ~6x the original for a control-heavy string, so measuring
  1701. # a value the budget rejects could breach RLIMIT_AS and surface as
  1702. # `exception` instead of the promised `output-limit`.
  1703. total += _json_str_cost(current)
  1704. elif type(current) is int:
  1705. # The canonical boundary accepts every JS-double-exact value: an int
  1706. # outside +-2**53-1 is fine IFF the double round-trip is exact.
  1707. if current > js_safe or current < -js_safe:
  1708. try:
  1709. exact = int(float(current)) == current
  1710. except OverflowError:
  1711. exact = False
  1712. if not exact:
  1713. return invalid("integer not exactly representable as a JavaScript number")
  1714. total += len(_dump_scalar(current).encode("utf-8"))
  1715. elif type(current) is float:
  1716. if current != current or current in (float("inf"), float("-inf")):
  1717. return invalid("non-finite float")
  1718. # JSON turns -0.0 into a sign the host parses back to JS -0; the
  1719. # canonical boundary rejects it, so this side must too.
  1720. if current == 0.0 and math.copysign(1.0, current) < 0:
  1721. return invalid("negative zero")
  1722. total += len(_dump_scalar(current).encode("utf-8"))
  1723. elif type(current) is list:
  1724. if id(current) in on_path:
  1725. return invalid("circular reference")
  1726. count = len(current)
  1727. total += 2 + (count - 1 if count > 1 else 0)
  1728. # Reject over-budget BEFORE iterating: every element serializes to at
  1729. # least one byte, so a wide flat forgery fails here without pulling a
  1730. # single child.
  1731. if total + count > max_bytes:
  1732. return over_budget
  1733. on_path.add(id(current))
  1734. stack.append((list_cursor, current, iter(current)))
  1735. elif type(current) is dict:
  1736. if id(current) in on_path:
  1737. return invalid("circular reference")
  1738. # ``len`` without materializing ``current.items()``: that list
  1739. # allocates one tuple per member before the bound below could run,
  1740. # recreating the spike the bound exists to stop.
  1741. count = len(current)
  1742. total += 2 + (count - 1 if count > 1 else 0)
  1743. # Same pre-iterate bound: each entry contributes a quoted key
  1744. # (>= 2 bytes), a colon, and a >= 1-byte value. ``iter`` on the items
  1745. # view is O(1); the cursor meters each key as it is pulled.
  1746. if total + count * 4 > max_bytes:
  1747. return over_budget
  1748. on_path.add(id(current))
  1749. stack.append((dict_cursor, current, iter(current.items())))
  1750. else:
  1751. # tuple, set, or any other type: not round-trippable JSON.
  1752. return invalid(f"unsupported type ({type(current).__name__})")
  1753. if total > max_bytes:
  1754. return over_budget
  1755. return None
  1756. def _lossless_json_violation(value: Any) -> str | None:
  1757. """Return why ``value`` is not lossless JSON, or ``None`` when it is.
  1758. ``json.dumps`` succeeding is NOT proof of losslessness: it coerces a
  1759. non-string ``dict`` key to its string form (``{1: "a", "1": "b"}`` collapses
  1760. to one key, silently dropping data), emits non-standard ``NaN``/``Infinity``
  1761. tokens without ``allow_nan=False``, and accepts integers outside JavaScript's
  1762. safe range (``9007199254740993`` becomes ``...992`` once the host parses the
  1763. frame into a JS number). Validate the shape up front so a coercive or lossy
  1764. value fails as ``invalid-output`` instead of round-tripping to something the
  1765. program did not compute. Iterative so deep nesting cannot overflow the stack,
  1766. and it tracks the container ancestry on the current path so a cyclic value is
  1767. reported at once rather than spinning until the CPU budget. Only JSON-plain
  1768. types survive: ``None``/``bool``/JS-safe ``int``/finite ``float``/``str``,
  1769. exact ``list``, and exact ``dict`` with ``str`` keys. Every type matches
  1770. EXACTLY, containers and scalars alike, for the reason
  1771. :func:`_check_done_value` documents: a subclass can override the operators
  1772. and methods a traversal calls, so an ``isinstance`` admission here would
  1773. approve one shape and let the encoder emit another.
  1774. """
  1775. # The canonical boundary accepts every JS-double-exact value: an int
  1776. # outside +-2**53-1 is fine IFF the double round-trip is exact (2**53 or
  1777. # 2**60 survive; 2**53+1 rounds), matching the worker backend.
  1778. js_safe = 2**53 - 1
  1779. # Post-order walk with an explicit "leave" marker: a container's id is added
  1780. # to `on_path` when entered and removed when left, so a back-edge to an
  1781. # ancestor (a cycle) is detected without rejecting a legitimately shared
  1782. # acyclic subtree.
  1783. on_path: set[int] = set()
  1784. # O(DEPTH) auxiliary space, not O(width), for the reason
  1785. # :func:`_check_done_value` documents: this walk runs in ``dispatch`` on
  1786. # MODEL-CONSTRUCTED binding arguments, which no child-side byte budget
  1787. # bounds first (the frame ceiling is the host's, and it applies after this
  1788. # returns). Enqueueing one frame per member would let a legitimate
  1789. # ``[0] * 6_000_000`` argument -- ~17 MB of JSON -- allocate ~366 MB of
  1790. # traversal tuples and die as the program's own MemoryError instead of
  1791. # round-tripping. A container therefore pushes ONE cursor frame holding its
  1792. # iterator; children are pulled one at a time.
  1793. exhausted = object()
  1794. visit, container_cursor = 0, 1
  1795. # A visit frame is (visit, value); a cursor frame is (cursor, container, iterator).
  1796. stack: list[tuple[int, Any, Any]] = [(visit, value, None)]
  1797. while stack:
  1798. kind = stack[-1][0]
  1799. if kind == container_cursor:
  1800. _, container, iterator = stack[-1]
  1801. child = next(iterator, exhausted)
  1802. if child is exhausted:
  1803. # Leaving the container: it is no longer on the current path, so
  1804. # a legitimately shared acyclic subtree is not mistaken for a cycle.
  1805. on_path.discard(id(container))
  1806. stack.pop()
  1807. continue
  1808. if type(container) is dict:
  1809. # The dict cursor yields (key, value): check the key as it is
  1810. # pulled. Only an EXACT str key survives -- int, float, None, and
  1811. # tuple keys coerce or raise, and a str subclass can carry
  1812. # overrides the encoder does not honor.
  1813. key, child = child
  1814. if type(key) is not str:
  1815. return f"non-string dict key ({type(key).__name__})"
  1816. stack.append((visit, child, None))
  1817. continue
  1818. _, current, _unused = stack.pop()
  1819. if current is None or type(current) is bool:
  1820. continue
  1821. if type(current) is str:
  1822. # Every string is lossless JSON. A lone surrogate has no UTF-8 form,
  1823. # but JSON carries the code unit as its ASCII ``\uXXXX`` escape and
  1824. # :func:`_dump_string` emits exactly that, so the host receives the
  1825. # same code unit the program passed — the same acceptance
  1826. # ``CodeJsonValue``, ``snapshotJsonValue``, and the worker backend
  1827. # already give it.
  1828. continue
  1829. if type(current) is int:
  1830. if current > js_safe or current < -js_safe:
  1831. try:
  1832. exact = int(float(current)) == current
  1833. except OverflowError:
  1834. exact = False
  1835. if not exact:
  1836. return "integer not exactly representable as a JavaScript number"
  1837. continue
  1838. if type(current) is float:
  1839. if current != current or current in (float("inf"), float("-inf")):
  1840. return "non-finite float"
  1841. # JSON serialization turns -0.0 into 0 (or "-0.0" text that the
  1842. # host parses to JS -0), silently changing the sign bit either
  1843. # way; the repository's canonical lossless-JSON boundary and the
  1844. # worker backend both reject it, so this side must too.
  1845. if current == 0.0 and math.copysign(1.0, current) < 0:
  1846. return "negative zero"
  1847. continue
  1848. if type(current) is list or type(current) is dict:
  1849. if id(current) in on_path:
  1850. return "circular reference"
  1851. on_path.add(id(current))
  1852. if type(current) is dict:
  1853. # Keys are checked as the cursor pulls each entry, not in a
  1854. # separate pass: ``current.values()`` would need a second walk,
  1855. # and materializing ``items()`` up front allocates one tuple per
  1856. # member -- the very spike the cursor removes.
  1857. stack.append((container_cursor, current, iter(current.items())))
  1858. else:
  1859. stack.append((container_cursor, current, iter(current)))
  1860. continue
  1861. return f"unsupported type ({type(current).__name__})"
  1862. return None
  1863. def _make_cpu_enforcer() -> Any:
  1864. """Build the CPU post-check over closure-held primitives.
  1865. This bootstrap IS ``__main__``, so model code can reach every one of its
  1866. module globals: ``import __main__; __main__._X = ...`` rebinds the name the
  1867. enforcement would otherwise read at call time, which a plain module-level
  1868. function plus module-level captures made a one-line defeat. The primitives
  1869. therefore live in this factory's locals, which become closure cells of the
  1870. returned function, and :func:`_run` binds the returned function into a
  1871. local of its own frame BEFORE executing the program, so no assignment to
  1872. ``__main__`` changes which callable runs or what it calls. Capture happens
  1873. at import time, before model code runs, so the captured
  1874. ``resource.getrusage``/``signal.signal``/``os.kill`` are the real builtins.
  1875. This raises the cost of defeating the check; it does not make it
  1876. unreachable, and nothing in-process could. A cell is writable through
  1877. ``fn.__closure__[i].cell_contents``, and ``sys._getframe`` walks to
  1878. :func:`_run`'s frame and reads its locals, so a program determined to
  1879. tamper still can — consistent with this backend's documented posture, where
  1880. the in-process interpreter is containment rather than a security boundary
  1881. (§Trust posture in the Code Mode RFC). The bounds that model code cannot
  1882. forge are outside the interpreter: the RLIMIT_CPU HARD limit at
  1883. ``cpuSeconds + 1``, whose SIGKILL is undeliverable to a handler and
  1884. unraisable by a process that cannot raise its own hard limit, and the
  1885. host's wall-clock ceiling. This check exists to convert the two cases those
  1886. miss — a program that traps SIGXCPU and settles inside the soft-to-hard
  1887. gap, and a program that spends the budget in DESCENDANTS the kernel never
  1888. charged to this process — from a reported SUCCESS into the same `timeout`
  1889. an untrapped program gets.
  1890. @returns The one-argument enforcement callable, taking `cpuSeconds`.
  1891. """
  1892. getrusage = resource.getrusage
  1893. rusage_self = resource.RUSAGE_SELF
  1894. rusage_children = resource.RUSAGE_CHILDREN
  1895. set_signal = signal.signal
  1896. sig_dfl = signal.SIG_DFL
  1897. sigxcpu = signal.SIGXCPU
  1898. kill = os.kill
  1899. getpid = os.getpid
  1900. # SIGXCPU unmasking primitives for the re-raise below: a program can mask
  1901. # the signal and return past the soft limit, so the re-delivered signal
  1902. # must be unblocked first. Captured here (import time) so a rebind cannot
  1903. # defeat them. The ``getattr``/``None`` guard is defensive against a
  1904. # stripped CPython build (the host refuses win32 at construction, so every
  1905. # platform this backend actually starts on has ``pthread_sigmask``).
  1906. pthread_sigmask = getattr(signal, "pthread_sigmask", None)
  1907. sig_unblock = getattr(signal, "SIG_UNBLOCK", None)
  1908. def die_if_cpu_exhausted(cpu_seconds: int) -> None:
  1909. """Die by re-delivered SIGXCPU when the CPU budget is already spent.
  1910. Two cases reach here as a would-be SUCCESS. A model program can trap
  1911. SIGXCPU and return during the one-second soft-to-hard gap. And
  1912. ``RLIMIT_CPU`` is PER-PROCESS, inherited fresh by every child, so a
  1913. program calling ``subprocess`` or ``os.fork`` multiplies the run's CPU
  1914. budget by the number of descendants it starts: measured with
  1915. ``cpuSeconds: 1``, two sequential busy children burned 2.0
  1916. CPU-seconds and the parent, which had accrued almost no CPU of its own
  1917. while blocked in ``subprocess.wait``, still returned a completion.
  1918. The meter is therefore ``RUSAGE_SELF + RUSAGE_CHILDREN``, the kernel's
  1919. own aggregate, which accumulates the CPU of every REAPED descendant
  1920. (grandchildren included, verified).
  1921. ``getrusage`` is the kernel's own meter (unforgeable from model code),
  1922. and dying by SIGXCPU with the default disposition restored gives the
  1923. host the same kernel-authoritative close signal as the untrapped soft
  1924. limit — classified as `timeout`, after which the host's process-group
  1925. SIGTERM/SIGKILL teardown reaches any surviving descendants. Runs AFTER
  1926. the model program settled, so a program can re-trap SIGXCPU between
  1927. this SIG_DFL and the kill only by running more code, which it no longer
  1928. does. A program that tampers with this callable instead (see
  1929. :func:`_make_cpu_enforcer` on why in-process state cannot be hidden)
  1930. buys at most the remaining soft-to-hard gap: one more CPU second, after
  1931. which the hard limit's SIGKILL lands with no handler possible.
  1932. Checking at settle time rather than sampling mid-run is deliberate:
  1933. both mid-run designs perturb the run they measure. A sampling thread
  1934. cost 72 MiB of virtual address space in the child (8 MiB stack plus a
  1935. 64 MiB glibc per-thread malloc arena reservation; measured 30.23 MiB of
  1936. mappings without it against 102.37 MiB with it), and ``RLIMIT_AS``
  1937. counts reserved space, so it silently shrank every run's
  1938. `addressSpaceMb`. A ``SIGALRM`` interval timer costs no mappings but
  1939. makes the program's own syscalls return short under PEP 475 — measured
  1940. a 64 MiB ``os.write`` returning 65536 — which corrupts fd-3 framing.
  1941. The cost of checking only at settle time is that a descendant's CPU is
  1942. detected after it is spent, not while it runs; the host's wall-clock
  1943. ceiling bounds that interval, and a program that never reaps its child
  1944. is bounded by the wall clock alone, since ``RUSAGE_CHILDREN`` counts
  1945. only reaped descendants (verified: a still-running child contributes
  1946. 0.0).
  1947. @param cpu_seconds The `cpuSeconds` budget the soft RLIMIT_CPU used.
  1948. """
  1949. own = getrusage(rusage_self)
  1950. kids = getrusage(rusage_children)
  1951. spent = own.ru_utime + own.ru_stime + kids.ru_utime + kids.ru_stime
  1952. if spent >= cpu_seconds:
  1953. # A program can mask SIGXCPU (``pthread_sigmask(SIG_BLOCK, ...)``),
  1954. # burn past the soft limit, and return during the soft-to-hard gap;
  1955. # the re-delivered SIGXCPU below would then stay PENDING and the
  1956. # child would exit normally with a success result. Restore the
  1957. # default disposition BEFORE unblocking: a program that installed a
  1958. # custom handler AND masked the signal has that pending handler run
  1959. # the moment the signal is unblocked (CPython delivers it at the next
  1960. # eval-breaker checkpoint in model code), and it could re-mask or
  1961. # raise — so the disposition must already be SIG_DFL when the signal
  1962. # is released. With SIG_DFL restored first, the pending signal kills
  1963. # the process inside the kernel with no bytecode window; the ``kill``
  1964. # below is the fallback for the never-pending case. ``pthread_sigmask``
  1965. # is ``None``-guarded defensively (every platform this backend starts
  1966. # on has it; the host refuses win32 at construction).
  1967. set_signal(sigxcpu, sig_dfl)
  1968. if pthread_sigmask is not None:
  1969. pthread_sigmask(sig_unblock, (sigxcpu,))
  1970. kill(getpid(), sigxcpu)
  1971. return die_if_cpu_exhausted
  1972. _DIE_IF_CPU_EXHAUSTED = _make_cpu_enforcer()
  1973. _TRUNCATION_MARKER = "… [truncated]"
  1974. # The marker's own UTF-8 size, reserved out of the cap rather than added on top
  1975. # of it. Byte-identical to the host's TRUNCATION_MARKER_BYTES; the ellipsis is
  1976. # three bytes, so this is 15, not the string's 13 characters.
  1977. _TRUNCATION_MARKER_BYTES = len(_TRUNCATION_MARKER.encode("utf-8"))
  1978. def _cap_message(message: str, max_bytes: int) -> str:
  1979. """Cap a diagnostic by its SERIALIZED cost, appending the host's marker.
  1980. Metered by the JSON-string cost the ``done`` frame will actually carry, not
  1981. by raw UTF-8 length: the message crosses fd 3 inside a JSON frame where
  1982. control characters escape up to sixfold (a NUL is one raw byte but six as
  1983. ``\\u0000``), so a raw-length cap of ``maxValueBytes`` could serialize to
  1984. roughly six times that and breach the 64 MiB frame parse cap — the silent
  1985. ``worker-exit`` inversion the load-time cap check exists to prevent, and a
  1986. several-hundred-MiB escape allocation besides. The seam's load bound admits
  1987. ``maxValueBytes`` up to ``parse-cap - envelope`` on the premise that both the
  1988. completion value and the diagnostic are metered in serialized bytes, so this
  1989. honors that premise for the diagnostic.
  1990. Encoded with ``errors="replace"`` first: a model exception message can
  1991. contain an unpaired surrogate (``raise Exception("\\ud800")``), and a strict
  1992. encode would throw while BUILDING the failure frame — the run would then
  1993. strand until the wall clock instead of reporting the exception. The marker's
  1994. serialized cost comes OUT of ``max_bytes``, so the returned string's own
  1995. frame form honors the cap; the host meters the same field again on arrival.
  1996. A ``max_bytes`` below the marker's cost yields the marker alone.
  1997. This is the PRODUCING-side cap. The host's receive-side ``capMessage``
  1998. (``src/index.ts``) bills the same field by RAW bytes instead, because its
  1999. output goes into ``CodeRunResult.error.message`` and never re-crosses a
  2000. frame-bounded channel — see that function's JSDoc for the split.
  2001. """
  2002. raw = message.encode("utf-8", errors="replace")
  2003. if _json_string_cost(raw) <= max_bytes:
  2004. return raw.decode("utf-8")
  2005. # Truncating: the result is `prefix + marker`, whose serialized cost is
  2006. # `2 (quotes) + sum(prefix byte costs) + marker cost`. The marker is
  2007. # escape-free, so its cost is its UTF-8 length. Reserve that and the quotes,
  2008. # then take the longest raw prefix whose accumulated per-byte cost fits.
  2009. # `_JSON_BYTE_COST` is per-byte and additive, so the scan is exact and walks
  2010. # at most a budget's worth of bytes, allocating nothing (unlike building the
  2011. # escaped form). `max(0, ...)` handles a `max_bytes` below the marker's own
  2012. # cost, yielding the marker alone.
  2013. content_budget = max(0, max_bytes - 2 - _TRUNCATION_MARKER_BYTES)
  2014. cost = 0
  2015. end = 0
  2016. for end in range(len(raw)):
  2017. cost += _JSON_BYTE_COST[raw[end]]
  2018. if cost > content_budget:
  2019. break
  2020. else:
  2021. end = len(raw)
  2022. # Drop a trailing partial UTF-8 sequence the slice may have cut (continuation
  2023. # bytes are 0b10xxxxxx); `errors="ignore"` renders the clean prefix.
  2024. return raw[:end].decode("utf-8", errors="ignore") + _TRUNCATION_MARKER
  2025. # Fixed safety/liveness bound, not a tunable: a model can raise an exception
  2026. # with an arbitrarily deep __cause__/__context__ chain, and both the rendering
  2027. # walk and format() are linear in chain length. Capping how many links get
  2028. # RENDERED keeps traceback formatting from consuming the whole wall budget.
  2029. # 100 links is far beyond any legible human traceback.
  2030. _MAX_TRACEBACK_CHAIN = 100
  2031. # Diagnostic used when rendering the failure itself fails. Built from a fixed
  2032. # literal plus the exception CLASS name, never from the exception's own str.
  2033. _UNRENDERABLE_DIAGNOSTIC = "<diagnostic rendering failed>"
  2034. def _model_traceback(exc: BaseException, max_bytes: int) -> str:
  2035. """Format a model-program failure with only the MODEL's own frames.
  2036. Bootstrap frames carry host-absolute paths — meaningless to the model and
  2037. unstable across machines, so transcripts pinning them cannot replay. They
  2038. appear not only as a leading prefix (the bootstrap's ``exec``/``await``)
  2039. but also interleaved and trailing: an uncaught binding rejection re-raised
  2040. by ``dispatch`` puts bootstrap frames AFTER the model's, and chained
  2041. ``__cause__``/``__context__`` exceptions carry their own stacks. Filter
  2042. every non-``<model>`` frame across the whole chain rather than trimming a
  2043. prefix. A failure with no model frame anywhere (e.g. a SyntaxError raised
  2044. by ``compile``) keeps the standard exception-only rendering.
  2045. Rendering is bounded to ``_MAX_TRACEBACK_CHAIN`` links, cut on the
  2046. ``TracebackException`` COPY, and a marker line announces the truncation.
  2047. Nothing here touches the live exception: an exception class overriding
  2048. ``__setattr__`` would run MODEL code from inside the caller's failure
  2049. handler, and a throw there costs the ``done`` frame (see
  2050. ``_safe_model_traceback``). ``TracebackException`` instances hold no such
  2051. hooks, so clearing their links runs no model code. The walk is iterative,
  2052. so a deep chain cannot overflow the recursion limit.
  2053. ``from_exception`` still copies the WHOLE live chain, at a higher per-link
  2054. cost than building it took. That is bounded by the child's ``RLIMIT_AS``:
  2055. the model must materialize every link (exception object plus traceback)
  2056. before raising, so a chain long enough for the copy to matter is already
  2057. near the address-space cap, and a ``MemoryError`` in the copy lands in the
  2058. caller's fallback rather than stranding the run.
  2059. """
  2060. te = traceback.TracebackException.from_exception(exc)
  2061. # One iterative pass over the copy does both jobs: keep only <model> frames
  2062. # on every linked exception and group member, and cut the chain at the cap.
  2063. found = False
  2064. truncated = False
  2065. pending = [(te, 1)]
  2066. while pending:
  2067. entry, depth = pending.pop()
  2068. kept = [f for f in entry.stack if f.filename == "<model>"]
  2069. entry.stack = traceback.StackSummary.from_list(kept)
  2070. found = found or bool(kept)
  2071. # 3.11+ exception groups (a binding failure inside asyncio.TaskGroup)
  2072. # carry member stacks under `exceptions`, not the dunder links; a group
  2073. # member counts as a link so the cap bounds nesting through both edges.
  2074. members = getattr(entry, "exceptions", None) or ()
  2075. if depth >= _MAX_TRACEBACK_CHAIN:
  2076. if entry.__cause__ is not None or entry.__context__ is not None or members:
  2077. truncated = True
  2078. entry.__cause__ = None
  2079. entry.__context__ = None
  2080. if members:
  2081. entry.exceptions = None
  2082. continue
  2083. for linked in (entry.__cause__, entry.__context__):
  2084. if linked is not None:
  2085. pending.append((linked, depth + 1))
  2086. for member in members:
  2087. pending.append((member, depth + 1))
  2088. def emit():
  2089. if found:
  2090. yield from te.format()
  2091. else:
  2092. yield from traceback.format_exception_only(type(exc), exc)
  2093. if truncated:
  2094. yield f"[dsh-code-runtime-python] exception chain truncated at {_MAX_TRACEBACK_CHAIN} links\n"
  2095. return _join_bounded(emit(), max_bytes)
  2096. def _make_failure_reporter() -> Any:
  2097. """Build the failure-diagnostic renderer over closure-held primitives.
  2098. The returned callable renders a model failure diagnostic that cannot itself
  2099. raise. The caller sends the ``done`` frame AFTER its ``except BaseException``
  2100. block, so anything thrown while rendering the diagnostic skips the send
  2101. entirely: the host then blocks on fd 3 until ``maxWallMs`` and reports a
  2102. timeout instead of the exception that actually happened. Rendering runs
  2103. model code by design (``format()`` reaches ``__str__``, ``__repr__`` and
  2104. ``__notes__``) and allocates under ``RLIMIT_AS``, so it must be treated as
  2105. able to throw.
  2106. The fallback names the exception CLASS and a fixed literal — no ``str(exc)``
  2107. and no ``format_exception_only``, both of which reach the model's
  2108. ``__str__``. A ``__name__`` that is not exactly ``str`` (a metaclass
  2109. property can return anything, or raise) is discarded rather than
  2110. formatted, so no override runs on this path either.
  2111. The factory exists for the same reason :func:`_make_cpu_enforcer` does: this
  2112. bootstrap IS ``__main__``, so ``import __main__; __main__._X = ...`` rebinds
  2113. any module global a call-time lookup would read. On this path a rebind is
  2114. worst — the handler's own reporter, and everything the reporter reaches,
  2115. would run model code outside any guard, and a throw there costs the ``done``
  2116. frame. The traceback formatter, the byte cap and the fallback literal
  2117. therefore become closure cells captured at import time, before model code
  2118. runs, and :func:`_run` binds the returned callable into a local of its own
  2119. frame. A frame local is not a module attribute, so no assignment to
  2120. ``__main__`` changes which callable runs or what it calls. This defeats the
  2121. one-line rebind, not a determined ``sys._getframe`` walk; the unforgeable
  2122. bound is the host wall clock.
  2123. """
  2124. cap_message = _cap_message
  2125. model_traceback = _model_traceback
  2126. unrenderable = _UNRENDERABLE_DIAGNOSTIC
  2127. # The exception class the guards below catch is bound into a closure cell
  2128. # here, at import time, before model code runs. `safe_model_traceback` runs
  2129. # AFTER the program (which is `__main__`) may have rebound the module global
  2130. # `BaseException`, so `except BaseException` would resolve the rebound class
  2131. # and a render-time throw could escape — losing the done frame. A closure cell
  2132. # is not reachable by `__main__._X = ...`, so the catch is immune.
  2133. _BaseException = BaseException
  2134. def safe_model_traceback(exc: BaseException, max_bytes: int) -> str:
  2135. try:
  2136. return cap_message(model_traceback(exc, max_bytes), max_bytes)
  2137. except _BaseException: # noqa: BLE001 -- a throw here would cost the done frame
  2138. pass
  2139. try:
  2140. raw_name = type(exc).__name__
  2141. # Slice BEFORE interpolating. A metaclass `__name__` property can
  2142. # return an arbitrarily long string, and both the f-string and
  2143. # `cap_message`'s encode would copy it whole — under a tight
  2144. # RLIMIT_AS either allocation can raise MemoryError, and this is the
  2145. # LAST fallback, so a throw here costs the `done` frame outright and
  2146. # the run misreports as an exit or a timeout. The slice is a
  2147. # code-unit prefix, which bounds the bytes at 4x, and the following
  2148. # `cap_message` still applies the exact byte cap.
  2149. name = raw_name[:_MAX_FALLBACK_NAME_CHARS] if type(raw_name) is str else "<unknown>"
  2150. except _BaseException: # noqa: BLE001 -- a raising __name__ must not cost the done frame
  2151. name = "<unknown>"
  2152. # Wrapped for the same reason: `cap_message` encodes, and its allocation
  2153. # is the only step left that can still fail. The fixed literal needs no
  2154. # budget, so it can always be delivered.
  2155. try:
  2156. return cap_message(f"{name}: {unrenderable}", max_bytes)
  2157. except _BaseException: # noqa: BLE001 -- the done frame outranks the diagnostic's detail
  2158. return unrenderable
  2159. return safe_model_traceback
  2160. _SAFE_MODEL_TRACEBACK = _make_failure_reporter()
  2161. def _join_bounded(lines, max_bytes: int) -> str:
  2162. """Join formatter output, stopping once the budget is comfortably passed.
  2163. ``format()`` yields lines lazily; consuming it whole for an exception
  2164. carrying a huge message would materialize the full text only for
  2165. ``_cap_message`` to throw it away — enough over-shoot to exhaust
  2166. ``RLIMIT_AS``. Stop after the accumulated CHARACTER count passes the byte
  2167. budget (chars lower-bound UTF-8 bytes); the caller's ``_cap_message``
  2168. does the exact byte-level cut.
  2169. """
  2170. chunks: list[str] = []
  2171. total = 0
  2172. for line in lines:
  2173. # A single yielded line can itself dwarf the budget (the exception
  2174. # message rides in one line): keep only the prefix it can ever need.
  2175. if len(line) > max_bytes + 1:
  2176. line = line[: max_bytes + 1]
  2177. chunks.append(line)
  2178. total += len(line)
  2179. if total > max_bytes:
  2180. break
  2181. return "".join(chunks)
  2182. def _done_with_value(
  2183. value: Any,
  2184. max_value_bytes: int,
  2185. # Bound as DEFAULT ARGUMENTS so they are captured at import time, before
  2186. # model code runs: `_done_with_value` runs AFTER the program (which is
  2187. # `__main__`) may have rebound `__main__._check_done_value` or
  2188. # `__main__._encode_json_plain`, and a module-global lookup at call time
  2189. # would let a one-line rebind rewrite a legitimate success into an
  2190. # `exception`. Defaults are evaluated at def time, so they are the originals.
  2191. _check_done_value: Any = _check_done_value,
  2192. _encode_json_plain: Any = _encode_json_plain,
  2193. _cap_message: Any = _cap_message,
  2194. ) -> dict[str, Any] | str:
  2195. """Build the terminal done frame under the seam's lossless-JSON contract.
  2196. A completion value returned by the program (``None`` when it returns
  2197. nothing) that is not lossless JSON fails the run as ``invalid-output``; a
  2198. serialized value beyond ``max_value_bytes`` fails as ``output-limit``.
  2199. Substituting a ``repr`` or truncated string would be a silent lie about
  2200. what the program computed, so both paths refuse instead (mirroring the
  2201. worker backend's contract). ``None`` crosses as an exact JSON ``null``.
  2202. The SUCCESS path returns the whole ``"done"`` frame as an ALREADY-ENCODED
  2203. JSON string: the admitted value is serialized here, at its validation
  2204. point, rather than handed to ``send_sync`` to re-walk later. The program
  2205. can keep mutating the returned list/dict from a daemon thread or signal
  2206. handler after it returns, so a second traversal held at a later point
  2207. would be a TOCTOU — a mutation into a non-JSON type would let that later
  2208. encode throw outside the settlement handler and downgrade the run
  2209. host-side to ``worker-exit``. Serializing once, inside the try that wraps
  2210. this call, closes the window: if a concurrent mutation makes the encode
  2211. throw, the exception handler classifies it as ``exception``, and once the
  2212. string is produced the frame is sent verbatim with no further touching of
  2213. the live value. Returns a ``dict`` only for a rejection (an error frame
  2214. carries no live model value and is safe to send via ``send_sync``).
  2215. """
  2216. # One bounded walk folds the losslessness check and the byte meter (mirrors
  2217. # the host's checkDoneValue): the former split ran the full losslessness
  2218. # walk first, materializing one tuple per element for a wide completion
  2219. # before the size cap could reject it — an RLIMIT_AS death on a value the
  2220. # meter would have refused. The value's escaped JSON is then produced in the
  2221. # SAME call, so the admitted value is serialized exactly once (see above);
  2222. # its size the walk proved within budget. Iterative like the encoder, so a
  2223. # valid completion deeper than the recursion limit still checks.
  2224. rejection = _check_done_value(value, max_value_bytes)
  2225. if rejection is not None:
  2226. kind, message = rejection
  2227. # The rejection diagnostic is capped like an exception message: a
  2228. # reason embedding a hostile class name (a huge `type(value).__name__`)
  2229. # could otherwise make the done frame exceed the host's frame parse cap
  2230. # and be silently dropped — an invalid-output run misreported as a
  2231. # worker-exit.
  2232. return {"type": "done", "error": {"kind": kind, "message": _cap_message(message, max_value_bytes)}}
  2233. # Pre-encode the value at the validation point (not in `_run`'s later send,
  2234. # which is outside the try): see the TOCTOU note in the docstring. The value
  2235. # is JSON-plain by construction, so `_encode_json_plain` is the encoder.
  2236. return '{"type": "done", "value": ' + _encode_json_plain(value) + "}"
  2237. def main() -> None:
  2238. channel = ProtocolChannel(PROTOCOL_FD)
  2239. asyncio.run(_run(channel))
  2240. if __name__ == "__main__":
  2241. main()