Something went wrong. Try again.
Personal fork of prime agent. Mirror of: https://mnemosyne.sigilzero.dev/sigilzero/prime-agent
Something went wrong. Try again.
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771"""Minimal CPython REPL runtime speaking newline-delimited JSON over stdio.
Entry point: ``python -m rlm.repl``. The protocol is documented in repl.mdnext to this file. Cells execute with top-level await in one persistent``__main__`` namespace on a single asyncio event loop."""
from __future__ import annotations
import astimport asyncioimport codecsimport contextvarsimport ctypesimport functoolsimport inspectimport ioimport jsonimport linecacheimport osimport pickleimport platformimport selectimport signalimport sysimport tempfileimport threadingimport timeimport tracebackimport typesimport uuidfrom collections.abc import Awaitable, Callablefrom typing import Any
from .bash import _kill_live_handles
PROTOCOL_VERSION = 3
DEFAULT_SNAPSHOT_MAX_BYTES = 256 * 1024 * 1024DEFAULT_SNAPSHOT_MAX_VARIABLE_BYTES = 16 * 1024 * 1024
# Plain ASCII, never a pickle start: _restore_state sniffs it to tell v2 framed# payloads from legacy (single dill-pickled dict) ones._SNAPSHOT_MAGIC = b"PRIME-AGENT-KERNEL-SNAPSHOT-V2\n"
# Stream writes must fit one protocol frame: the host buffers whole lines# before its per-execution truncation, and raw fd writes already arrive as# 64 KiB pump chunks._STREAM_FRAME_TEXT_CAP = 64 * 1024# The host truncates results at a smaller per-execution maxChars, so this only# bounds a pathological repr or exception text in transit._RESULT_TEXT_CAP = 1_048_576_RESULT_TRUNCATION_MARKER = f"\n[... result truncated at {_RESULT_TEXT_CAP} characters ...]"# Oversized display and host_request payloads fail the cell instead of wedging host memory._PAYLOAD_CAP = 16 * 1024 * 1024
# Names the session bootstrap re-creates on every start; never snapshotted._ALWAYS_SKIP = {"rlm", "mcp", "bash", "plugin", "asyncio", "In", "Out", "get_ipython", "exit", "quit", "open"}# Legacy IPython names and live plugin bindings are never restored from snapshot data._RESTORE_SKIP = {"In", "Out", "get_ipython", "plugin", "_prime_agent_plugin_namespace"}# The target of the last successful snapshot, remembered so an EOF shutdown# (the host process died without a graceful dispose) can flush the final# namespace before exit. `None` until this process has committed a snapshot:# an EOF before that must not overwrite the on-disk payload with a namespace# the host never considered durable._last_snapshot_target: dict[str, Any] | None = None
_protocol_fd: int = -1# The host's stderr before _setup_fds captures fd 2: the only channel that can# report a dropped protocol frame without feeding it back into the protocol._host_stderr_fd: int = -1_write_lock = threading.Lock()_loop: asyncio.AbstractEventLoop | None = None_serve_task: asyncio.Task[Any] | None = None
class _CellExecution: def __init__(self) -> None: self.finished = asyncio.Event() self.owner: asyncio.Task[Any] | None = None
# Asyncio tasks copy cell context at creation, so detached tasks retain their# output attribution and completion barrier. Threads start with a fresh context._current_cell: contextvars.ContextVar[str | None] = contextvars.ContextVar("_current_cell", default=None)_current_cell_execution: contextvars.ContextVar[_CellExecution | None] = contextvars.ContextVar( "_current_cell_execution", default=None)_active: dict[str, Any] = {"task": None, "rid": None, "interrupted": False}_cell_counter = 0_pending_host: dict[str, "asyncio.Future[dict[str, Any]]"] = {}# Set on the loop thread once stdin hits EOF or a shutdown request arrives; no# host reply can arrive after that, so waiting (and future) host_request calls fail._host_closed = False
# In-flight dedicated-lane MCP status requests (the eager settle and the# connections view ride this lane): the serve loop hands them off instead of# awaiting them inline, so a hanging server bounds its own request while the# execution queue stays free for the user's cells. Shutdown cancels the set# before closing the MCP registry._mcp_status_tasks: "set[asyncio.Task[None]]" = set()
# Interrupt bookkeeping shared between the reader thread and the loop thread._interrupt_lock = threading.Lock()_inflight: set[str] = set()_pending_interrupts: dict[str, Any] = {"ids": set(), "any": False}_sigint_target: str | None = None_finishing_rid: str | None = None_handoff_interrupted = False
def _send(event: dict[str, Any]) -> None: """Write one protocol frame; the locked single write keeps frames atomic.""" data = (json.dumps(event, separators=(",", ":")) + "\n").encode() with _write_lock: view = memoryview(data) try: while view: view = view[os.write(_protocol_fd, view) :] except OSError as err: try: os.write(_host_stderr_fd, f"rlm.repl: dropped protocol frame: {err}\n".encode()) except OSError: pass
def _check_payload(event: str, data: dict[str, Any]) -> None: """Fail the calling cell when a `data` payload would not fit one protocol frame.
Strict-dumps validation: default allow_nan=True would let NaN/Infinity serialize as non-JSON text and tear the host's protocol framing (a non-serializable value raises TypeError here before any bytes are written, so NaN is the only corruption vector). The encoded length enforces the frame cap; _send re-serializes. """ if len(json.dumps(data, allow_nan=False)) > _PAYLOAD_CAP: raise ValueError(f"{event} payload exceeds the {_PAYLOAD_CAP}-character frame cap")
def emit(data: dict[str, Any]) -> None: """Ship one display event carrying a dict of MIME type -> JSON payload.
Thread-safe; the event is tagged with the cell running at call time. """ if not isinstance(data, dict) or not data or not all(isinstance(k, str) for k in data): raise TypeError("emit() requires a non-empty dict keyed by MIME type strings") _check_payload("display", data) _send({"event": "display", "id": _current_cell.get(), "data": data})
def is_active() -> bool: """True when this process serves the repl protocol (not merely imported).""" return _protocol_fd >= 0
def current_cell_completion_context() -> tuple[asyncio.Event, asyncio.Task[Any] | None] | None: """Return the calling cell's completion barrier and owning execution task.""" execution = _current_cell_execution.get() if execution is None: return None return execution.finished, execution.owner
def active_cell_task() -> asyncio.Task[Any] | None: """The cell body task executing right now, or None between cells (global state, not the cell contextvar — detached tasks keep stale context copies).""" with _interrupt_lock: task = _active["task"] return task if isinstance(task, asyncio.Task) and not task.done() else None
async def host_request(data: dict[str, Any]) -> dict[str, Any]: """Send one typed request to the host and await its raw reply dict.""" if _loop is None: raise RuntimeError("repl runtime is not serving") if _host_closed: raise RuntimeError("host connection closed; host_request cannot be answered") _check_payload("host_request", data) rid = uuid.uuid4().hex future: asyncio.Future[dict[str, Any]] = _loop.create_future() _pending_host[rid] = future try: _send({"event": "host_request", "id": rid, "data": data}) return await future except BaseException: # cancellation/KeyboardInterrupt must retire native work too _send({"event": "host_cancel", "id": rid}) raise finally: _pending_host.pop(rid, None)
def _fail_pending_host_requests() -> None: """Loop-thread half of teardown: no host reply can arrive anymore, so every awaiting cell must unblock or the queued shutdown would never be served.""" global _host_closed _host_closed = True for rid, future in _pending_host.items(): if not future.done(): _send({"event": "host_cancel", "id": rid}) future.set_exception(RuntimeError("host connection closed; host_request cannot be answered"))
def _resolve_host_reply(rid: str, data: dict[str, Any]) -> None: """Reader-thread half of the host bridge; late/unknown replies are dropped.""" assert _loop is not None
def deliver() -> None: future = _pending_host.get(rid) if future is not None and not future.done(): future.set_result(data)
_loop.call_soon_threadsafe(deliver)
class _Pump: """Reads one captured-output pipe and ships its bytes as stream events."""
def __init__(self, read_fd: int, write_fd: int, stream: str) -> None: self._read_fd = read_fd # Private write end: a cell closing/reclaiming fd 1/2 cannot hijack drain tokens. self._token_fd = os.dup(write_fd) self._stream = stream self._decoder = codecs.getincrementaldecoder("utf-8")("replace") self._lock = threading.Lock() self._watch: tuple[bytes, threading.Event] | None = None self._buf = b"" self._thread = threading.Thread(target=self._run, daemon=True) self._thread.start()
def drain(self) -> None: """Block until every byte written to the fd so far has been shipped.""" token = b"\xff<drain:" + uuid.uuid4().hex.encode() + b">\xff" seen = threading.Event() with self._lock: self._watch = (token, seen) try: os.write(self._token_fd, token) # Backstop only: a dead pump (read end closed under it) can never set seen. while not seen.wait(0.1): if not self._thread.is_alive(): return except OSError: return finally: with self._lock: self._watch = None
def _run(self) -> None: while True: try: chunk = os.read(self._read_fd, 65536) except OSError: break if not chunk: break self._feed(chunk)
def _feed(self, chunk: bytes) -> None: data = self._buf + chunk self._buf = b"" with self._lock: watch = self._watch if watch is None: self._emit(data) return token, seen = watch while True: i = data.find(token) if i == -1: break self._emit(data[:i]) self._finish_decode() seen.set() data = data[i + len(token) :] # Hold back a tail that could be the start of a token split across reads. hold = 0 for k in range(min(len(data), len(token) - 1), 0, -1): if data.endswith(token[:k]): hold = k break if hold: self._buf = data[len(data) - hold :] data = data[: len(data) - hold] self._emit(data)
def _emit(self, data: bytes) -> None: if not data: return text = self._decoder.decode(data) if text: # Raw fd bytes have no provable owner (os.write, C extensions, # subprocesses, threads from earlier cells): never credit a cell. _send({"event": self._stream, "id": None, "text": text})
def _finish_decode(self) -> None: text = self._decoder.decode(b"", final=True) self._decoder = codecs.getincrementaldecoder("utf-8")("replace") if text: _send({"event": self._stream, "id": None, "text": text})
class _TaggedBuffer(io.RawIOBase): """Binary proxy for _TaggedWriter.buffer: bytes go to the raw fd channel.
Byte ownership cannot be proven at this layer, so buffer writes ride the captured pipe and surface as id:null stream events (drained before done like any other raw fd write). """
def __init__(self, fallback_fd: int) -> None: self._fallback_fd = fallback_fd
def write(self, data: Any) -> int: # memoryview (unlike bytes()) rejects int, matching a real buffer's TypeError. view = memoryview(data).cast("B") total = len(view) # Pipe writes can be short for payloads above the pipe capacity. while view: view = view[os.write(self._fallback_fd, view) :] return total
def flush(self) -> None: pass
def fileno(self) -> int: return self._fallback_fd
def writable(self) -> bool: return True
class _TaggedWriter(io.TextIOBase): """sys.stdout/sys.stderr replacement tagging writes with the writer's cell id.
Python-level writes carry write-time provenance from the _current_cell contextvar (asyncio tasks inherit the spawning cell's id; user threads see None) and ship straight to the protocol, bypassing the fd pipe. fileno() and .buffer expose the captured pipe so subprocesses, C-level writers, and sys.stdout.buffer.write() keep working through the raw channel (null-attributed). """
def __init__(self, stream: str, fallback_fd: int) -> None: self._stream = stream self._fallback_fd = fallback_fd # Keeps one write()'s frames contiguous under concurrent writers. self._frame_lock = threading.Lock() self._buffer = _TaggedBuffer(fallback_fd)
def write(self, text: str) -> int: if not isinstance(text, str): raise TypeError(f"write() argument must be str, not {type(text).__name__}") if text: cell_id = _current_cell.get() with self._frame_lock: for start in range(0, len(text), _STREAM_FRAME_TEXT_CAP): _send( { "event": self._stream, "id": cell_id, "text": text[start : start + _STREAM_FRAME_TEXT_CAP], } ) return len(text)
def flush(self) -> None: pass
def fileno(self) -> int: return self._fallback_fd
def writable(self) -> bool: return True
@property def buffer(self) -> _TaggedBuffer: return self._buffer
@property def encoding(self) -> str: return "utf-8"
@property def errors(self) -> str: return "replace"
def _consume_task_exception(task: asyncio.Task[Any]) -> None: """Retrieve a killed task's exception so no never-retrieved noise is logged.""" if not task.cancelled(): task.exception()
def _sigint_handler(signum: int, frame: types.FrameType | None) -> None: global _handoff_interrupted task = _active["task"] # No lock (the main thread may hold it): the rid equality revalidates the # target so a SIGINT delayed past its request's finish cannot hit a later cell. if task is None or task.done() or _active["rid"] != _sigint_target: if _sigint_target is not None and _sigint_target == _active["rid"]: # Handoff: the task is done but _run_guarded's finally has not run # yet, so the main thread may be inside loop internals where raising # would kill the serve loop. Record it; the finishing phase consumes it. _handoff_interrupted = True return # Post-run repr/drain is synchronous main-thread work: raise into it. # The equality revalidation drops a SIGINT delayed past the done send. if _sigint_target is not None and _sigint_target == _finishing_rid: raise KeyboardInterrupt return _active["interrupted"] = True # Handler runs in the main (loop) thread: current_task is whose step the signal interrupted. running = asyncio.current_task(_loop) if _loop is not None else None if running is task: # The active request's own step (sync bytecode or an EINTR-woken syscall): raise into it. raise KeyboardInterrupt # Loop idle in select() or another task mid-step: cancel the active task (same thread, safe). task.cancel() if running is not None and running is not _serve_task: # A background task blocked in sync code occupies the only thread and would keep the # cancel from ever running: raise into it to unwind its step; it dies with the KI. running.add_done_callback(_consume_task_exception) raise KeyboardInterrupt
def _request_interrupt(target: str | None) -> None: """Deliver an interrupt now, or park it for the request it targets.
Runs on the reader thread. Without a target id the interrupt applies to the running request, else to the next queued one; with a target id it applies to that request only. A request finishing its post-run repr/drain is still interrupted (never parked: parking would hit the NEXT request). Interrupts for finished or unknown requests are dropped. """ global _sigint_target with _interrupt_lock: task = _active["task"] rid = _active["rid"] if rid is not None and (target is None or target == rid): # Active, or in the done-task handoff before _run_guarded's finally: # either way the rid still owns the interrupt (parking here would # leak it onto the next request); the handler decides delivery. _sigint_target = rid elif _finishing_rid is not None and (target is None or target == _finishing_rid): _sigint_target = _finishing_rid elif target is not None: if target in _inflight: _pending_interrupts["ids"].add(target) return elif _inflight: _pending_interrupts["any"] = True return else: return # SIGINT must land on the main thread, where cells execute. Windows has no # signal.pthread_kill: fall back to cancelling the active task on the loop # (sync-blocked cells and the finishing repr/drain cannot be broken there; # best-effort parity). if hasattr(signal, "pthread_kill"): signal.pthread_kill(threading.main_thread().ident, signal.SIGINT) if _loop is not None: # Wake the selector so a cancel scheduled by the handler runs promptly. _loop.call_soon_threadsafe(lambda: None) return if _loop is not None:
def cancel_active() -> None: current = _active["task"] if current is task and current is not None and not current.done(): _active["interrupted"] = True current.cancel()
_loop.call_soon_threadsafe(cancel_active)
def _consume_pending_interrupt(rid: str) -> bool: """Check-and-clear any interrupt parked for this request.""" pending = _pending_interrupts["any"] or rid in _pending_interrupts["ids"] _pending_interrupts["any"] = False _pending_interrupts["ids"].discard(rid) return pending
def _consume_handoff_interrupt() -> bool: """Check-and-clear an interrupt that landed in the done-task handoff.""" global _handoff_interrupted with _interrupt_lock: pending = _handoff_interrupted _handoff_interrupted = False return pending
def _finish_locked(rid: str) -> None: """Drop a finished request; a parked untargeted interrupt survives while others are inflight.""" global _finishing_rid, _handoff_interrupted, _sigint_target if _finishing_rid == rid: # An unconsumed handoff interrupt dies with its request (state requests # have no cancellable post-run work); it must never hit the next request. _finishing_rid = None _handoff_interrupted = False if _sigint_target == rid: # The target dies with its request: a later request reusing this id must # not match a stale target when a delayed/external SIGINT arrives. _sigint_target = None _inflight.discard(rid) _pending_interrupts["ids"].discard(rid) if not _inflight: _pending_interrupts["any"] = False
def _finish_request(rid: str) -> None: with _interrupt_lock: _finish_locked(rid)
_RUNTIME_FILE = __file__
def _cell_stack(stack: traceback.StackSummary) -> traceback.StackSummary | None: """Frames from the first cell frame on, minus runtime-internal frames.
Returns None when no cell frame exists (e.g. a compile-time SyntaxError). """ start = next((i for i, f in enumerate(stack) if f.filename.startswith("<cell-")), None) if start is None: return None return traceback.StackSummary.from_list([f for f in stack[start:] if f.filename != _RUNTIME_FILE])
def _safe_str(exc: BaseException) -> str: try: return str(exc) except BaseException: # noqa: BLE001 - a broken __str__ must not kill the runtime return "<exception str() failed>"
def _cap_text(text: str) -> str: if len(text) > _RESULT_TEXT_CAP: return text[:_RESULT_TEXT_CAP] + _RESULT_TRUNCATION_MARKER return text
def _cap_traceback_lines(lines: list[str]) -> list[str]: """Bound the aggregate, not just each entry: an exception chain can carry thousands of entries, and per-entry caps alone would still let one error event exceed the host's protocol line limit. Keep the newest entries — the outermost exception carries the actionable failure — and lead with a truncation marker.""" total = sum(len(line) for line in lines) if total <= _RESULT_TEXT_CAP: return lines kept: list[str] = [] remaining = _RESULT_TEXT_CAP for line in reversed(lines): if len(line) > remaining: if not kept: # Never drop the newest entry: it names the raised exception. kept.append(line) break kept.append(line) remaining -= len(line) kept.reverse() kept.insert( 0, f"[... traceback truncated: kept the newest {len(kept)} of {len(lines)} entries " f"to fit {_RESULT_TEXT_CAP} characters ...]\n", ) return kept
def _error_event(cell_id: str, exc: BaseException) -> dict[str, Any]: # No cell frame (e.g. SyntaxError): exception-only keeps filename, source, and caret. te = traceback.TracebackException.from_exception(exc) stack = _cell_stack(te.stack) if stack is None: lines = traceback.format_exception_only(type(exc), exc) else: te.stack = stack lines = list(te.format()) return { "event": "error", "id": cell_id, "ename": type(exc).__name__, "evalue": _cap_text(_safe_str(exc)), "traceback": _cap_traceback_lines([_cap_text(line) for line in lines]), }
def _interrupt_event(cell_id: str, exc: BaseException) -> dict[str, Any]: """Report a cancelled await-suspended cell as a KeyboardInterrupt.""" stack = _cell_stack(traceback.extract_tb(exc.__traceback__)) lines = [] if stack: lines = ["Traceback (most recent call last):\n"] lines.extend(stack.format()) lines.append("KeyboardInterrupt\n") return {"event": "error", "id": cell_id, "ename": "KeyboardInterrupt", "evalue": "", "traceback": lines}
def _compile_cell(code: str, filename: str) -> tuple[list[types.CodeType], bool]: """Compile a cell; a trailing expression compiles separately in eval mode.""" linecache.cache[filename] = (len(code), None, code.splitlines(keepends=True), filename) tree = ast.parse(code, filename) trailing: ast.Expression | None = None if tree.body and isinstance(tree.body[-1], ast.Expr): trailing = ast.Expression(tree.body.pop().value) flags = ast.PyCF_ALLOW_TOP_LEVEL_AWAIT codes: list[types.CodeType] = [] if tree.body: codes.append(compile(tree, filename, "exec", flags=flags, dont_inherit=True)) if trailing is not None: codes.append(compile(trailing, filename, "eval", flags=flags, dont_inherit=True)) return codes, trailing is not None
async def _run_codes( codes: list[types.CodeType], ns: dict[str, Any], *, refresh_plugins: bool = False) -> Any: if refresh_plugins: # Keep the bootstrap-owned namespace, not a user-rebound alias. Refresh is # part of the guarded execution so host errors/interrupts prevent the cell. namespace = ns.get("_prime_agent_plugin_namespace") if namespace is None: raise RuntimeError("plugin namespace was not initialized by the session bootstrap") await namespace.refresh() ns["plugin"] = namespace value: Any = None for code_obj in codes: value = eval(code_obj, ns) # noqa: S307 - executing the model's cell is the runtime's job if code_obj.co_flags & inspect.CO_COROUTINE: value = await value return value
async def _run_guarded(task: asyncio.Task[Any], rid: str) -> tuple[str, Any, dict[str, Any] | None]: """Await a request task; returns (status, value, error event or None).""" with _interrupt_lock: _active["interrupted"] = False _active["rid"] = rid _active["task"] = task if _consume_pending_interrupt(rid): # Interrupt parked before activation: cancel before the first step. _active["interrupted"] = True task.cancel() try: value = await task return "ok", value, None except asyncio.CancelledError as exc: if _active["interrupted"]: return "error", None, _interrupt_event(rid, exc) return "error", None, _error_event(rid, exc) except BaseException as exc: # noqa: BLE001 - every cell failure becomes an error event return "error", None, _error_event(rid, exc) finally: with _interrupt_lock: global _finishing_rid # The rid stays inflight and interrupt-targetable through the # post-run repr/drain; the handler closes the window via _finish_request. # Set before clearing _active: the lock-free handler must always see # the rid in one of the two slots, never a torn in-between state. _finishing_rid = rid _active["task"] = None _active["rid"] = None
async def _handle_execute(req: dict[str, Any], ns: dict[str, Any]) -> None: global _cell_counter cell_id = req["id"] _cell_counter += 1 filename = f"<cell-{_cell_counter}>" execution = _CellExecution() cell_token = _current_cell.set(cell_id) execution_token = _current_cell_execution.set(execution) try: codes, has_trailing = _compile_cell(req["code"], filename) assert _loop is not None task = _loop.create_task(_run_codes(codes, ns, refresh_plugins=req.get("refresh_plugins") is True)) execution.owner = task status, value, error = await _run_guarded(task, cell_id) result_text: str | None = None try: if _consume_handoff_interrupt() and status == "ok": # SIGINT landed between the task's completion and the finishing # phase: it targeted this request, so cancel its remaining work. status, error = "error", _error_event(cell_id, KeyboardInterrupt()) if status == "ok" and has_trailing and value is not None: try: ns["_"] = value result_text = repr(value) except BaseException as exc: # noqa: BLE001 - a broken __repr__ is a cell error status, error = "error", _error_event(cell_id, exc) if result_text is not None: result_text = _cap_text(result_text) _drain_output() finally: # Close the interrupt window before the protocol sends so a # handler-raised KeyboardInterrupt can never tear a frame mid-_send. _finish_request(cell_id) if result_text is not None: _send({"event": "result", "id": cell_id, "text": result_text}) if error is not None: _send(error) _send({"event": "done", "id": cell_id, "status": status}) finally: execution.owner = None execution.finished.set() _current_cell_execution.reset(execution_token) _current_cell.reset(cell_token)
def _drain_output() -> None: # Per-stream, and ValueError too: a cell may close sys.stdout/sys.stderr, and # flushing a closed file raises ValueError, or rebind them to a flush-less # object (AttributeError); neither may kill the serve loop nor skip flushing # the other stream. for stream in (sys.stdout, sys.stderr): try: stream.flush() except (OSError, ValueError, AttributeError): pass _pump_out.drain() _pump_err.drain()
class _SnapshotSizeLimitExceeded(Exception): pass
class _CappedWriter: def __init__(self, sink: Any, limit: int) -> None: self._sink = sink self._limit = limit self.written = 0
def write(self, chunk: Any) -> int: size = len(chunk) if self.written + size > self._limit: raise _SnapshotSizeLimitExceeded() self._sink.write(chunk) self.written += size return size
def _read_snapshot_records(fh: Any, max_bytes: int, max_variable_bytes: int) -> dict[str, bytes]: """Framing damage is a corrupt snapshot: a restore error, never a partial namespace. Length fields are bounds-checked before their reads, so a corrupt header cannot force a huge allocation; the writer's own per-record and aggregate caps bound every blob read, so a sparse multi-gigabyte file cannot OOM the process either.""" fh.seek(0, os.SEEK_END) size = fh.tell() fh.seek(len(_SNAPSHOT_MAGIC)) records: dict[str, bytes] = {} total = 0 while fh.tell() < size: header = fh.read(4) if len(header) < 4: raise ValueError("truncated snapshot record") name_len = int.from_bytes(header, "little") if fh.tell() + name_len + 8 > size: raise ValueError("truncated snapshot record") # The name is bounded by the same aggregate cap as the blobs: a # corrupt or sparse snapshot declaring a multi-gigabyte name must # fail the cap BEFORE the read allocates it (the same OOM class # the blob caps close). if name_len > max_bytes: raise ValueError("snapshot record name exceeds the aggregate byte cap") name = fh.read(name_len) raw_len = fh.read(8) blob_len = int.from_bytes(raw_len, "little") if len(raw_len) < 8 or fh.tell() + blob_len > size: raise ValueError("truncated snapshot record") if blob_len > max_variable_bytes: raise ValueError("snapshot record exceeds the per-variable byte cap") total += blob_len if total > max_bytes: raise ValueError("snapshot payload exceeds the aggregate byte cap") blob = fh.read(blob_len) if len(blob) < blob_len: raise ValueError("truncated snapshot record") records[name.decode("utf-8")] = blob return records
def _reject_plugin_snapshot_reference(value: Any) -> None: """Plugin authority is live host state, never portable Python namespace data.
Pickler persistent_id visits nested containers, instance attributes, closures, defaults and dill's captured globals too; checking only top-level names would let a saved alias smuggle a retired generation into the next kernel. """ cls = value if isinstance(value, type) else type(value) if any(base.__module__ == "rlm.plugins" for base in cls.__mro__): raise pickle.PicklingError("live plugin references are not restored from snapshots")
def _snapshot_state( ns: dict[str, Any], path: str, manifest_path: str, max_bytes: int, max_variable_bytes: int, prune_oversized: bool, committed: list[dict[str, Any]] | None = None,) -> dict[str, Any]: import datetime
try: import dill except Exception as err: # noqa: BLE001 - dill is provisioned by the host, not a hard dep return {"error": f"dill unavailable: {err}"} dill.settings["recurse"] = True
class SnapshotPickler(dill.Pickler): def persistent_id(self, value: Any) -> None: _reject_plugin_snapshot_reference(value) return None
saved: list[str] = [] skipped: list[dict[str, str]] = [] oversized: list[str] = [] missing = object()
os.makedirs(os.path.dirname(path) or ".", exist_ok=True) temps: list[str] = []
def stage_temp(target: str, mode: str): # Unique same-directory temps: a fixed '.tmp' name could alias the other # final path (clobbering it) or collide with a concurrent snapshot. fd, name = tempfile.mkstemp( dir=os.path.dirname(target) or ".", prefix=os.path.basename(target) + ".", suffix=".tmp" ) temps.append(name) try: return os.fdopen(fd, mode), name except BaseException: os.close(fd) # fdopen never took ownership: the raw fd would leak raise
def discard_temps() -> None: for stale in temps: try: os.remove(stale) except OSError: pass
# Stage both temps before replacing anything: any failure up to the first # replace leaves the previous payload+manifest pair fully intact. stage = "write" parked: list[int] = [] handler_installed = False previous = None try: try: if max_bytes < len(_SNAPSHOT_MAGIC): # Even the header alone busts the cap: keep the committed-payload <= cap invariant. return {"error": "write failed: snapshot exceeds aggregate snapshot size cap"} fh, tmp = stage_temp(path, "wb") with fh: # Single pass: each variable is dill-serialized exactly once, streamed # into the staged temp. The record header is charged against the aggregate # cap up front, so a completed record can never overflow it (no prefix re-dump). total = fh.write(_SNAPSHOT_MAGIC) for name in list(ns.keys()): if name.startswith("_") or name in _ALWAYS_SKIP: continue value = ns.get(name, missing) if value is missing: # A background thread deleted the name after the key listing. skipped.append({"name": name, "reason": "deleted during snapshot"}) continue try: encoded = name.encode("utf-8") except UnicodeEncodeError as err: # A lone-surrogate name (e.g. "\ud800") cannot ride the # v2 record header; skip it like any other unserializable # name instead of failing the whole snapshot. skipped.append({"name": name, "reason": f"{type(err).__name__}: {_safe_str(err)[:200]}"}) continue # Record header: 4-byte name length + 8-byte blob length, plus the name itself. budget = max_bytes - total - 12 - len(encoded) # Prune mode measures at the full per-variable cap: only that cap decides # pruned-ness, and the write always re-measures — in-place mutation # defeats any name-based size tracking from an earlier dump. limit = max_variable_bytes if prune_oversized else min(max_variable_bytes, budget) buffer = io.BytesIO() try: SnapshotPickler(_CappedWriter(buffer, limit)).dump(value) blob = buffer.getvalue() except _SnapshotSizeLimitExceeded: if not prune_oversized and budget < max_variable_bytes: skipped.append({"name": name, "reason": "exceeds aggregate snapshot size cap"}) else: skipped.append({"name": name, "reason": "exceeds per-variable snapshot size cap"}) oversized.append(name) continue except Exception as err: # noqa: BLE001 - one unpicklable name must not abort the snapshot skipped.append({"name": name, "reason": f"{type(err).__name__}: {_safe_str(err)[:200]}"}) continue if total + 12 + len(encoded) + len(blob) > max_bytes: # Only reachable in prune mode, where the measurement cap ignores the budget. skipped.append({"name": name, "reason": "exceeds aggregate snapshot size cap"}) continue fh.write(len(encoded).to_bytes(4, "little")) fh.write(encoded) fh.write(len(blob).to_bytes(8, "little")) fh.write(blob) total += 12 + len(encoded) + len(blob) saved.append(name) saved.sort() pruned = sorted(name for name in oversized if name in ns) if prune_oversized else [] manifest = { "version": 1, "savedNames": saved, "skipped": skipped, "pruned": pruned, "bytes": total, "pythonVersion": sys.version.split()[0], "timestamp": datetime.datetime.now(datetime.timezone.utc).isoformat(), } stage = "manifest write" fh, manifest_tmp = stage_temp(manifest_path, "w") with fh: json.dump(manifest, fh) except BaseException as err: # noqa: BLE001 - Exception -> error dict, rest propagates if not isinstance(err, Exception): raise # e.g. KeyboardInterrupt: clean up (outer finally), then propagate return {"error": f"{stage} failed: {err}"}
# A SIGINT-raised KeyboardInterrupt anywhere from the first commit through the # last cleanup removal would desync payload/manifest/namespace or misreport a # committed snapshot: park SIGINT until the end; it is consumed, see below. previous = signal.signal(signal.SIGINT, lambda signum, frame: parked.append(signum)) handler_installed = True try: os.replace(tmp, path) except OSError as err: return {"error": f"write failed: {err}"} try: os.replace(manifest_tmp, manifest_path) except OSError as err: # Fail before the prune deletions so a bad manifest path never destroys state. return {"error": f"manifest write failed: {err}"} for name in pruned: ns.pop(name, None) result = {"saved": saved, "skipped": skipped, "pruned": pruned, "bytes": total} # Publish while still parked: a later KeyboardInterrupt into this task finds the committed result (see _handle_state). if committed is not None: committed.append(result) finally: # The one guaranteed cleanup point (unique owned names: after a successful # commit the renamed temps no longer exist, so this is a no-op). It runs with # SIGINT still parked; the nested finally makes the restore the guaranteed # last action even when cleanup itself fails. try: discard_temps() finally: if handler_installed: signal.signal(signal.SIGINT, previous) # The parked SIGINT is consumed, not re-raised: with the manifest committed and # the namespace pruned, the destructive snapshot has succeeded, and re-raising # would misreport it as failed and risk the host discarding the only copy of # the pruned variables. The interrupt targeted this now-complete request. return result
def _revive_with_live_globals( value: Any, ns: dict[str, Any], backfill: list[tuple[str, Any]] | None = None, memo: dict[int, Any] | None = None,) -> Any: """Rebind restored __main__ callables onto the live namespace, collecting names their saved globals carry but ns lacks as backfill for the caller to apply at commit (live ns values always win).""" import functools
_reject_plugin_snapshot_reference(value) if memo is None: memo = {} if id(value) in memo: return memo[id(value)]
def revive(dep: Any) -> Any: return _revive_with_live_globals(dep, ns, backfill, memo)
if isinstance(value, functools.partial): # No placeholder memo entry: a partial is immutable, so it could never be patched; # every cycle passes through a function, which is memoized before recursing. rebuilt = revive(value.func) changed = rebuilt is not value.func args = [] keywords = {} for arg in value.args: revived = revive(arg) changed = changed or revived is not arg args.append(revived) for key, arg in value.keywords.items(): revived = revive(arg) changed = changed or revived is not arg keywords[key] = revived if not changed: # An unchanged partial still carries its original attributes: # __main__ callables there would keep frozen snapshot globals, so # revive them in place. Memoize first — an attribute can cycle # back to this partial. memo[id(value)] = value value.__dict__.update({key: revive(attr) for key, attr in value.__dict__.items()}) return value rebuilt_partial = functools.partial(rebuilt, *args, **keywords) # Memoize before the attribute walk: attributes can hold the partial # itself (a self-cycle or mutual partials), and unlike the function # branch below there is no outer memo entry yet. Once set, the # not-changed fast path also returns the rebuilt one. memo[id(value)] = rebuilt_partial rebuilt_partial.__dict__.update({key: revive(attr) for key, attr in value.__dict__.items()}) return rebuilt_partial atoms = (int, float, str, bytes, bool, type(None)) # dill loads __main__.__dict__ by reference, so a saved globals() IS the live ns: never walk it. if value is ns: return value if isinstance(value, (list, dict, set)): # Memoized before recursing and revived in place: cycles and identity come for free. # Skipping atoms keeps the walk over million-element containers near dill.loads cost. memo[id(value)] = value if isinstance(value, set): # Iterate a snapshot: discard+add during iteration would skip members. for item in list(value): revived = item if type(item) in atoms else revive(item) if revived is not item: value.discard(item) value.add(revived) elif isinstance(value, list): for key, item in enumerate(value): revived = item if type(item) in atoms else revive(item) if revived is not item: value[key] = revived else: # Keys can be __main__ callables too: revive them, or lookups # through the dict keep observing frozen globals. The snapshot # tolerates the delete+reinsert a rebuilt key needs. for key, item in list(value.items()): revived = item if type(item) in atoms else revive(item) revived_key = key if type(key) in atoms else revive(key) if revived_key is not key: del value[key] value[revived_key] = revived elif revived is not item: value[key] = revived return value if type(value) is tuple: items = tuple(item if type(item) in atoms else revive(item) for item in value) if all(new is old for new, old in zip(items, value)): items = value return memo.setdefault(id(value), items) if type(value) is frozenset: # Immutable: rebuild when any member revived (identity equality makes # the comparison exact — a rebuilt function never equals the original). items = frozenset(item if type(item) in atoms else revive(item) for item in value) if items == value: items = value return memo.setdefault(id(value), items) if not isinstance(value, types.FunctionType) or value.__module__ != "__main__": return value # Defaults and cell contents are revived only after the rebound function is memoized, so a # function reachable from its own defaults or closure resolves to it. Cells are revived in # place: holders this walk never sees (attribute-held siblings) must keep sharing them. rebound = types.FunctionType(value.__code__, ns, value.__name__, None, value.__closure__) memo[id(value)] = rebound if backfill is not None: for name, dep in value.__globals__.items(): # Snapshots never save _-prefixed or skip-listed names; backfill must not smuggle them past that policy. if name in ns or name.startswith("_") or name in _ALWAYS_SKIP or name in _RESTORE_SKIP: continue backfill.append((name, revive(dep))) if value.__defaults__: rebound.__defaults__ = tuple(revive(dep) for dep in value.__defaults__) if value.__kwdefaults__: rebound.__kwdefaults__ = {key: revive(dep) for key, dep in value.__kwdefaults__.items()} for cell in value.__closure__ or (): if id(cell) in memo: continue memo[id(cell)] = cell try: contents = cell.cell_contents except ValueError: continue cell.cell_contents = revive(contents) rebound.__doc__ = value.__doc__ rebound.__dict__.update({key: revive(attr) for key, attr in value.__dict__.items()}) rebound.__annotations__ = value.__annotations__ rebound.__qualname__ = value.__qualname__ rebound.__module__ = value.__module__ # PEP 695 generics carry their type params here on 3.12+; plain 3.11 # functions lack the attribute entirely, hence the getattr guard. params = getattr(value, "__type_params__", None) if params is not None: rebound.__type_params__ = params return rebound
def _snapshot_unpickler(dill: Any) -> type: """Unpickler that refuses to reopen a pickled raw fd number in this kernel.
dill serializes a pipe/socket-backed file as its fd NUMBER; restoring a saved closed one reopens that number here and closes it again, killing whatever owns the fd now (e.g. the event loop's self-pipe). The refusal fails just that record; every other name still restores. """ create_filehandle = dill._dill._create_filehandle
def refuse_raw_fd(name: Any, *args: Any) -> Any: if isinstance(name, int): raise pickle.UnpicklingError( f"refusing to reopen raw file descriptor {name} from a snapshot" ) return create_filehandle(name, *args)
class GuardedUnpickler(dill.Unpickler): def find_class(self, module: str, name: str) -> Any: target = super().find_class(module, name) _reject_plugin_snapshot_reference(target) return refuse_raw_fd if target is create_filehandle else target
return GuardedUnpickler
def _restore_state( ns: dict[str, Any], path: str, committed: list[dict[str, Any]] | None = None, max_bytes: int | None = None, max_variable_bytes: int | None = None,) -> dict[str, Any]: if not os.path.exists(path): return {"restored": [], "failed": [], "reason": "snapshot not found"} try: import dill except Exception as err: # noqa: BLE001 return {"error": f"dill unavailable: {err}"} unpickler = _snapshot_unpickler(dill) try: with open(path, "rb") as fh: if fh.read(len(_SNAPSHOT_MAGIC)) == _SNAPSHOT_MAGIC: payload = _read_snapshot_records( fh, max_bytes if max_bytes is not None else DEFAULT_SNAPSHOT_MAX_BYTES, ( max_variable_bytes if max_variable_bytes is not None else DEFAULT_SNAPSHOT_MAX_VARIABLE_BYTES ), ) else: # Legacy: one dill-pickled dict; old snapshot files must keep restoring. fh.seek(0) payload = unpickler(fh).load() except Exception as err: # noqa: BLE001 - a corrupt snapshot yields an empty restore return {"error": f"load failed: {_safe_str(err)}"} if not isinstance(payload, dict): return {"error": "corrupt snapshot: not a dict"}
staged: dict[str, Any] = {} failed: list[dict[str, str]] = [] for name, blob in payload.items(): if name in _RESTORE_SKIP or name.startswith("_prime_agent_plugin_"): continue try: staged[name] = unpickler(io.BytesIO(blob)).load() except Exception as err: # noqa: BLE001 - revive every other name regardless failed.append({"name": name, "reason": f"{type(err).__name__}: {_safe_str(err)[:200]}"}) # Revive every staged name before parking: a failure must never abort the # apply halfway and leave the namespace half old, half new. prepared: dict[str, Any] = {} backfill: list[tuple[str, Any]] = [] revive_failed: list[dict[str, str]] = [] for name, value in staged.items(): # The backfill entries a name's revival produced merge only if THAT # name revives: a failed revival adds nothing (its saved globals # would partially restore state the failure report says failed). name_backfill: list[tuple[str, Any]] = [] try: prepared[name] = _revive_with_live_globals(value, ns, name_backfill) except Exception as err: # noqa: BLE001 - one broken revival must not abort the restore revive_failed.append({"name": name, "reason": f"{type(err).__name__}: {_safe_str(err)[:200]}"}) continue backfill.extend(name_backfill) result = {"restored": sorted(prepared), "failed": failed + revive_failed} # Park SIGINT across the whole apply so it is all-or-nothing; the parked interrupt is consumed by the commit (as in snapshot). previous = signal.signal(signal.SIGINT, lambda signum, frame: None) try: for name, value in prepared.items(): ns[name] = value for name, value in backfill: # prepared names already sit in ns here: a restored value always beats backfill. if name not in ns: ns[name] = value # Publish while still parked: a later KeyboardInterrupt into this task finds the committed result (see _handle_state). if committed is not None: committed.append(result) finally: signal.signal(signal.SIGINT, previous) return result
async def _handle_state(req: dict[str, Any], ns: dict[str, Any]) -> None: """Run snapshot/restore as an interruptible task and reply in the done event.""" rid = req["id"] committed: list[dict[str, Any]] = []
async def run() -> dict[str, Any]: if req["type"] == "snapshot": prune = req.get("prune_oversized", False) if not isinstance(prune, bool): return {"error": "prune_oversized must be a boolean"} for field in ("max_bytes", "max_variable_bytes"): # Any present value must be a non-negative int; a JSON null is not a valid way to ask # for the default, and a negative cap would prune every user variable from ns. if field in req and ( isinstance(req[field], bool) or not isinstance(req[field], int) or req[field] < 0 ): return {"error": f"{field} must be a non-negative integer"} # realpath resolves symlinks, so aliased paths cannot silently clobber the payload. if os.path.realpath(req["path"]) == os.path.realpath(req["manifest_path"]): return {"error": "path and manifest_path must differ"} return _snapshot_state( ns, req["path"], req["manifest_path"], req.get("max_bytes", DEFAULT_SNAPSHOT_MAX_BYTES), req.get("max_variable_bytes", DEFAULT_SNAPSHOT_MAX_VARIABLE_BYTES), prune, committed, ) return _restore_state( ns, req["path"], committed, req.get("max_bytes", DEFAULT_SNAPSHOT_MAX_BYTES), req.get("max_variable_bytes", DEFAULT_SNAPSHOT_MAX_VARIABLE_BYTES), )
assert _loop is not None task = _loop.create_task(run()) outcome: tuple[str, Any, dict[str, Any] | None] | None = None try: outcome = await _run_guarded(task, rid) _finish_request(rid) # no post-run repr/drain: close the interrupt window now except KeyboardInterrupt: # A finishing-targeted SIGINT can raise anywhere between _run_guarded's # finally publishing _finishing_rid and _finish_request clearing it; the # handler only raises once _finishing_rid is set, so the task is already # complete (destructively so for a pruning snapshot). Consume the # interrupt and report the task's real outcome; escaping to the backstop # would misreport a committed snapshot as failed. _finish_request(rid) if outcome is None: # The KeyboardInterrupt pre-empted _run_guarded's return: recover # the completed task's outcome with _run_guarded's failure mapping. try: outcome = ("ok", task.result(), None) except asyncio.CancelledError as exc: event = _interrupt_event(rid, exc) if _active["interrupted"] else _error_event(rid, exc) outcome = ("error", None, event) except BaseException as exc: # noqa: BLE001 - every request failure becomes an error event outcome = ("error", None, _error_event(rid, exc)) status, result, error = outcome if ( committed and _active["interrupted"] and error is not None and error.get("ename") == "KeyboardInterrupt" ): # Recover only a protocol interrupt that landed after the commit; a user KeyboardInterrupt keeps interrupted reporting. status, result, error = "ok", committed[0], None if status != "ok": reason = "interrupted" if error and error.get("ename") == "KeyboardInterrupt" else ( f"{error.get('ename')}: {error.get('evalue')}" if error else "failed" ) _send({"event": "done", "id": rid, "status": "error", "reason": reason}) return if "error" in result: _send({"event": "done", "id": rid, "status": "error", "reason": result["error"]}) return if req["type"] == "snapshot": global _last_snapshot_target _last_snapshot_target = { "path": req["path"], "manifest_path": req["manifest_path"], "max_bytes": req.get("max_bytes", DEFAULT_SNAPSHOT_MAX_BYTES), "max_variable_bytes": req.get("max_variable_bytes", DEFAULT_SNAPSHOT_MAX_VARIABLE_BYTES), } _send({"event": "done", "id": rid, "status": "ok", **result})
def _flush_final_snapshot(ns: dict[str, Any]) -> None: """EOF-only best-effort final snapshot before exit.
The host died without a graceful dispose (crash, SIGKILL, worker exit path that skipped `shutdown`), so its debounced snapshots stop at the last one. Flush the current namespace to the last-known target so a resume revives state up to the EOF instead of up to the debounce. Atomic staging (temp + rename) means an interrupted flush leaves the previous payload intact, never a torn one. """ target = _last_snapshot_target if target is None: return try: _snapshot_state( ns, target["path"], target["manifest_path"], target["max_bytes"], target["max_variable_bytes"], False, None, ) except BaseException: # noqa: BLE001 - never block or crash the shutdown path pass
def _list_names(ns: dict[str, Any]) -> list[str]: """User-defined top-level names, filtered like the snapshot.""" # Non-string keys (globals()[1] = 1) are not user-listable names. return sorted( name for name in ns if isinstance(name, str) and not name.startswith("_") and name not in _ALWAYS_SKIP )
async def _handle_list_names(req: dict[str, Any], ns: dict[str, Any]) -> None: _send({"event": "done", "id": req["id"], "status": "ok", "names": _list_names(ns)})
async def _handle_mcp_status(req: dict[str, Any], ns: dict[str, Any]) -> None: from . import mcp as mcp_mod
servers = req.get("servers") if not isinstance(servers, list) or not all(isinstance(name, str) for name in servers): raise ValueError("mcp_status requires a list of server names") timeout_ms = req.get("timeout_ms", 10_000) if isinstance(timeout_ms, bool) or not isinstance(timeout_ms, (int, float)) or timeout_ms <= 0: raise ValueError("mcp_status timeout_ms must be a positive number") connections = await mcp_mod.status(servers, float(timeout_ms)) _send({"event": "done", "id": req["id"], "status": "ok", "connections": connections})
async def _handle_request( handler: Callable[[dict[str, Any], dict[str, Any]], Awaitable[None]], req: dict[str, Any], ns: dict[str, Any],) -> None: # Backstop: one broken request (e.g. RecursionError in compile) fails alone, never the serve loop. try: await handler(req, ns) except BaseException as exc: # noqa: BLE001 - any per-request failure becomes error+done rid = req["id"] with _interrupt_lock: # Only a still-inflight request (never reached _run_guarded, e.g. compile # failure) owns a parked interrupt; after _run_guarded finished it, a parked # "any" belongs to the next request and must survive. if rid in _inflight: _consume_pending_interrupt(rid) _finish_locked(rid) _send(_error_event(rid, exc)) _send({"event": "done", "id": rid, "status": "error"})
async def _mcp_status_lane(req: dict[str, Any], ns: dict[str, Any]) -> None: """One dedicated-lane `mcp_status` request.
The failure shape mirrors `_handle_request`'s backstop (one broken request fails alone, never the lane, never the serve loop); a shutdown cancellation is silent by design — the host is already tearing the kernel down and owns the deadline. """ try: await _handle_mcp_status(req, ns) except asyncio.CancelledError: raise except BaseException as exc: # noqa: BLE001 - one broken lane request fails alone rid = req["id"] try: _send(_error_event(rid, exc)) _send({"event": "done", "id": rid, "status": "error"}) except BaseException: pass
def _spawn_mcp_status_lane(req: dict[str, Any], ns: dict[str, Any]) -> None: """Hand one `mcp_status` request to the dedicated lane.
`mcp_status` opens and lists the configured generic servers (the eager background settle and the connections view both ride it), and a slow or unreachable server can hold it for the whole per-server timeout. Awaiting it inline on the execution queue would park the user's first python cell behind a server handshake that no cell asked for (finding: the eager settle must never contend with the user's first-turn path; the join point stays at first use through the registry's per-server locks). The lane task runs on the same loop (the registry dispatch joins concurrent listings and an early user call onto the same in-flight open), so ordering between cells is untouched — only status work stops occupying the cell queue. """ task = asyncio.get_running_loop().create_task(_mcp_status_lane(req, ns)) _mcp_status_tasks.add(task) task.add_done_callback(_mcp_status_tasks.discard)
async def _serve(queue: asyncio.Queue[dict[str, Any]], ns: dict[str, Any]) -> None: while True: req = await queue.get() # A cell (or a snapshot-restored prior handler) may have rebound SIGINT; the # protocol handler must own it before each request. Mid-cell rebinds remain # that cell's own problem for that cell only. signal.signal(signal.SIGINT, _sigint_handler) rtype = req.get("type") if rtype == "shutdown": rid = req.get("id") if req.get("eof"): # Host stdin closed without a shutdown request: the host # process is gone, so this is the last chance to persist. _flush_final_snapshot(ns) # The dedicated MCP status lane goes first: a still-listing # server is cancelled (bounded, silent) so the registry close # below cannot race an in-flight open under the host's 5s # shutdown deadline. if _mcp_status_tasks: for task in list(_mcp_status_tasks): task.cancel() await asyncio.wait(set(_mcp_status_tasks), timeout=1.0) _mcp_status_tasks.clear() # MCP children must close before the loop dies; close() is internally bounded under the host's 5s deadline. mcp_mod = sys.modules.get("rlm.mcp") if mcp_mod is not None: try: await mcp_mod.close() except BaseException as exc: print(f"MCP shutdown failed: {type(exc).__name__}: {exc}", file=sys.stderr) # Kill live bash children now; atexit would wait on parked executor threads. _kill_live_handles() if isinstance(rid, str): _send({"event": "done", "id": rid, "status": "ok"}) return if rtype == "execute": await _handle_request(_handle_execute, req, ns) elif rtype in ("snapshot", "restore"): await _handle_request(_handle_state, req, ns) elif rtype == "list_names": await _handle_request(_handle_list_names, req, ns) elif rtype == "mcp_status": # The dedicated lane (see `_spawn_mcp_status_lane`): status/open # work never occupies the execution queue. _spawn_mcp_status_lane(req, ns)
def _handle_bash_activity(req: dict[str, Any]) -> None: """Out-of-band: a running cell must not block inspection or cancellation.""" from .bash import activity_request
rid = req["id"] try: response = activity_request(req["action"], req.get("activityId"), req.get("lines", 50)) frame = {"event": "done", "id": rid, "status": "ok", **response} _cap_bash_activity_frame(frame) _send(frame) except (KeyError, ValueError) as exc: _send({"event": "done", "id": rid, "status": "error", "reason": str(exc)})
def _cap_bash_activity_frame(frame: dict[str, Any]) -> None: """Keep the serialized response under the 16 KiB wire cap.
json escaping can expand one character to six bytes (uXXXX-style), so the byte slices in `activity_request` cannot bound the frame alone. Trim from the oldest end: a tail keeps its newest lines, a list keeps its newest rows. """ tail = frame.get("tail") if isinstance(tail, str): while len(json.dumps(frame)) > 16_384: excess = len(json.dumps(frame)) - 16_384 keep = max(0, len(tail) - excess // 6 - 1) if keep >= len(tail): # The frame cannot fit no matter how the payload shrinks # (oversized request metadata): emit the smallest frame # instead of looping forever on the reader thread. frame["tail"] = "" break tail = tail[-keep:] if keep else "" frame["tail"] = tail return rows = frame.get("activities") while len(json.dumps(frame)) > 16_384 and isinstance(rows, list) and len(rows) > 1: victim = next( (index for index, row in enumerate(rows) if row.get("status") != "running"), 0, ) rows.pop(victim)
_REQUIRED_FIELDS = { "execute": ("id", "code"), "snapshot": ("id", "path", "manifest_path"), "restore": ("id", "path"), "list_names": ("id",), # mcp_status's server list is a JSON array, so only its id is a # string-required field; the handler validates the list itself. "mcp_status": ("id",), "bash_activity": ("id", "action"), "shutdown": (),}
def _protocol_error(message: str) -> None: _send({"event": "error", "id": None, "ename": "ProtocolError", "evalue": message, "traceback": []})
def _handle_request_line(raw: bytes, queue: asyncio.Queue[dict[str, Any]]) -> None: assert _loop is not None req = json.loads(raw) if not isinstance(req, dict): raise ValueError("request is not a JSON object") rtype = req.get("type") if rtype == "interrupt": if "id" in req and not isinstance(req["id"], str): _protocol_error("interrupt request id must be a string") return _request_interrupt(req.get("id")) return if rtype == "host_reply": # Bypass the FIFO queue: the awaiting cell IS the in-flight # execute, so a queued reply would deadlock behind it. rid = req.get("id") data = req.get("data") if isinstance(rid, str) and isinstance(data, dict): _resolve_host_reply(rid, data) else: _protocol_error("host_reply request needs string id and dict data") return if not isinstance(rtype, str) or rtype not in _REQUIRED_FIELDS: _protocol_error(f"unknown request type: {rtype!r}") return missing = [f for f in _REQUIRED_FIELDS[rtype] if not isinstance(req.get(f), str)] if missing: _protocol_error(f"{rtype} request needs string fields: {', '.join(missing)}") return if rtype == "bash_activity": if req["action"] not in ("list", "tail", "kill"): _protocol_error("unknown bash activity action") return if req["action"] != "list" and not isinstance(req.get("activityId"), str): _protocol_error("bash activity tail/kill requires string activityId") return if len(req["id"]) > 256 or len(req.get("activityId") or "") > 256: # Frame metadata rides every response: an unbounded id would # leave no room for the capped payload. _protocol_error("bash activity ids must stay under 256 characters") return # Like host_reply, this bypasses the cell FIFO. Handles remain owned # by the runtime, not by an arbitrary PID supplied by the client. _handle_bash_activity(req) return if rtype in ("execute", "snapshot", "restore"): with _interrupt_lock: # A reused in-flight id would corrupt interrupt/finish bookkeeping. duplicate = req["id"] in _inflight if not duplicate: _inflight.add(req["id"]) # The protocol write can block on backpressure: never send under the lock. if duplicate: _protocol_error(f"duplicate in-flight request id: {req['id']!r}") return if rtype == "shutdown": # No host reply follows a shutdown; a cell awaiting host_request # must fail now or it would block _serve from ever consuming this. _loop.call_soon_threadsafe(_fail_pending_host_requests) _loop.call_soon_threadsafe(queue.put_nowait, req)
def _read_requests(stdin_fd: int, queue: asyncio.Queue[dict[str, Any]]) -> None: assert _loop is not None with os.fdopen(stdin_fd, "rb") as stream: for raw in stream: raw = raw.strip() if not raw: continue try: # The whole per-line handling sits inside the backstop: hostile # input (RecursionError from pathological nesting, unhashable # field types, ...) must never kill the reader thread. _handle_request_line(raw, queue) except BaseException as err: # noqa: BLE001 _protocol_error(f"{type(err).__name__}: {_safe_str(err)}") # Host closed stdin: shut the runtime down. The marker distinguishes # this from the host's explicit shutdown request (which runs after the # host flushed its own final snapshot, so no runtime-side flush runs). _loop.call_soon_threadsafe(_fail_pending_host_requests) _loop.call_soon_threadsafe(queue.put_nowait, {"type": "shutdown", "eof": True})
def _resolve_owner_pid() -> int: raw = os.environ.get("PRIME_AGENT_KERNEL_OWNER_PID", "") try: owner = int(raw) except ValueError: owner = 0 return owner if owner > 0 else os.getppid()
def _owner_alive_posix(owner: int, initial_ppid: int) -> bool: # Reparenting is the race-free parent-death signal when the owner is the # parent; the kill-0 probe covers an env-designated non-parent owner. if initial_ppid == owner and os.getppid() != initial_ppid: return False try: os.kill(owner, 0) except ProcessLookupError: return False except OSError: pass # EPERM etc.: alive but unprobeable return True
def _wait_owner_posix(owner: int, initial_ppid: int) -> None: # Blocks until the owner exits: the OS exit notification (kqueue NOTE_EXIT # on macOS/BSD, a pidfd on Linux) wakes this thread once, with no polling. # Registration binds whichever process holds the pid right now, so the # liveness check after it catches a parent owner that died (and whose pid # may be reused) before the watch existed. try: if hasattr(select, "kqueue"): kq = select.kqueue() # max_events=0 makes a registration error raise (ESRCH for a gone # owner) instead of arriving as an EV_ERROR event. kq.control( [select.kevent(owner, select.KQ_FILTER_PROC, select.KQ_EV_ADD, select.KQ_NOTE_EXIT)], 0, 0, ) wait_for_exit = functools.partial(kq.control, None, 1) else: poller = select.poll() poller.register(os.pidfd_open(owner), select.POLLIN) wait_for_exit = poller.poll except ProcessLookupError: return # already gone except (AttributeError, OSError): # Slow fallback, only where exit notification is unavailable: Linux # before 5.3, a seccomp filter denying pidfd_open, a CPython built # without os.pidfd_open. while _owner_alive_posix(owner, initial_ppid): time.sleep(30.0) return if _owner_alive_posix(owner, initial_ppid): wait_for_exit()
def _wait_owner_windows(owner: int) -> None: # Blocks until the owner exits. os.kill(pid, 0) on Windows TERMINATES the # target, so a SYNCHRONIZE handle wait is the only sound probe. from ctypes import wintypes
SYNCHRONIZE = 0x00100000 INFINITE = 0xFFFFFFFF k32 = ctypes.WinDLL("kernel32", use_last_error=True) k32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] k32.OpenProcess.restype = wintypes.HANDLE k32.WaitForSingleObject.argtypes = [wintypes.HANDLE, wintypes.DWORD] k32.WaitForSingleObject.restype = wintypes.DWORD k32.CloseHandle.argtypes = [wintypes.HANDLE] k32.CloseHandle.restype = wintypes.BOOL handle = k32.OpenProcess(SYNCHRONIZE, False, owner) if not handle: return # already gone (or unprobeable): exit rather than run ownerless try: k32.WaitForSingleObject(handle, INFINITE) finally: k32.CloseHandle(handle)
def _owner_watchdog(owner: int, initial_ppid: int) -> None: if os.name == "nt": _wait_owner_windows(owner) else: _wait_owner_posix(owner, initial_ppid) # Event-loop-independent by design: a synchronous cell monopolizes the # loop, so the queued EOF shutdown can never run; hard-exit from here. try: _kill_live_handles() except BaseException: # noqa: BLE001 pass os._exit(1)
def _start_owner_watchdog() -> None: threading.Thread( target=_owner_watchdog, args=(_resolve_owner_pid(), os.getppid()), daemon=True ).start()
_pump_out: _Pump_pump_err: _Pump
def _setup_fds() -> int: """Reserve stdout for the protocol; route fds 1/2 through captured pipes.""" global _protocol_fd, _pump_out, _pump_err, _host_stderr_fd _protocol_fd = os.dup(1) os.set_inheritable(_protocol_fd, False) _host_stderr_fd = os.dup(2) os.set_inheritable(_host_stderr_fd, False) out_r, out_w = os.pipe() err_r, err_w = os.pipe() os.dup2(out_w, 1) os.dup2(err_w, 2) os.close(out_w) os.close(err_w) sys.stdout = _TaggedWriter("stdout", fallback_fd=os.dup(1)) sys.stderr = _TaggedWriter("stderr", fallback_fd=os.dup(2)) stdin_fd = os.dup(0) devnull = os.open(os.devnull, os.O_RDONLY) os.dup2(devnull, 0) os.close(devnull) sys.stdin = open(os.devnull, "r") # user input() sees EOF, never protocol frames _pump_out = _Pump(out_r, 1, "stdout") _pump_err = _Pump(err_r, 2, "stderr") return stdin_fd
def main() -> None: global _loop, _serve_task stdin_fd = _setup_fds() _start_owner_watchdog()
# Alias the executing module so an in-cell `from rlm.repl import emit` # binds the live module, not a second copy. sys.modules.setdefault("rlm.repl", sys.modules[__name__]) # A real __main__ module makes dill pickle user functions/classes by value. user_module = types.ModuleType("__main__") user_module.__dict__["__builtins__"] = __builtins__ sys.modules["__main__"] = user_module
_loop = asyncio.new_event_loop() asyncio.set_event_loop(_loop) queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue() signal.signal(signal.SIGINT, _sigint_handler) threading.Thread(target=_read_requests, args=(stdin_fd, queue), daemon=True).start()
_send({"event": "ready", "protocol": PROTOCOL_VERSION, "python": platform.python_version()})
_serve_task = _loop.create_task(_serve(queue, user_module.__dict__)) # A KeyboardInterrupt escaping a cell or background task stops # run_until_complete; the interrupt is already recorded, so resume serving. # SystemExit is the same kind of escape, not a shutdown: asyncio re-raises it # from the cell task into the loop, and letting it out of main() ends the # interpreter before the awaiting _handle_execute can report the cell. (A # non-daemon thread left by an earlier cell then blocks interpreter exit # forever, leaving a live kernel that answers nothing.) Resuming the loop # delivers the stored SystemExit to `await task`, which reports it as that # cell's error and keeps the kernel serving. while not _serve_task.done(): try: _loop.run_until_complete(_serve_task) except (KeyboardInterrupt, SystemExit): continue _loop.close()
if __name__ == "__main__": main()