#!/usr/bin/env python3 """pi-repl bridge: owns ipykernel and the Jupyter protocol; talks to the TS host over one stdio pipe. Protocol (JSON lines; host -> bridge on stdin, bridge -> host on stdout): in: {"op":"boot","helpers":[{"name":str,"source":str}]} first message, before anything {"op":"exec","id":str,"code":str} run a cell (streams + one result) {"op":"snapshot","id":str,"path":str,"max_bytes":int} snapshot the namespace to path {"op":"restore","id":str,"path":str} revive path into the namespace {"op":"listNames","id":str} current public names {"op":"interrupt"} KeyboardInterrupt in the kernel {"op":"shutdown","id":str} graceful teardown, then exit 0 out: {"type":"ready","helpers":[{"name":str,"ok":bool,"error":str?}]} kernel up, helpers loaded {"type":"stream","id":str,"name":"stdout"|"stderr","text":str} cell output, streamed {"type":"result","id":str,"status":"ok"|"error"|"aborted", cell settled "result":str?,"error":{"name":str,"message":str,"stack":[str]}?} {"type":"reply","id":str,"..."} snapshot/restore/names/shutdown {"type":"error","id":str?,"message":str} protocol-level failure {"type":"dying","message":str?} best-effort before exit(1) The bridge is single-threaded over the kernel (one op at a time); the host serializes ops anyway. A stdin reader thread only relays lines; all kernel I/O stays on the main thread. When the kernel dies (os._exit, SIGKILL, crash) the bridge exits: the host treats bridge exit as kernel death and rebuilds from the last snapshot. stdin EOF (host gone) shuts the kernel down cleanly. """ import json import os import queue import sys import threading import time from jupyter_client import KernelManager from jupyter_client.kernelspec import NoSuchKernel SNAPSHOT_MIME = "application/vnd.pi-repl.snapshot+json" RESTORE_MIME = "application/vnd.pi-repl.restore+json" NAMES_MIME = "application/vnd.pi-repl.names+json" READY_TIMEOUT_MS = 30_000 # --- kernel-side programs. These run inside ipykernel, so they stay strings (now plain Python). --- def skip_set_literal(helper_names): skip = list(helper_names) + ["helper_description", "In", "Out", "get_ipython", "exit", "quit", "open"] return json.dumps(skip) def snapshot_code(helper_names, max_bytes): # v4: cloudpickle serializes functions and classes by value — no source-capture machinery. # Entries are zlib-compressed cloudpickle streams; one unpicklable binding costs itself only. return ("import cloudpickle as _cp, base64 as _b64, json as _js, zlib as _zl\n" + "__repl_skip = set(" + skip_set_literal(helper_names) + ")\n" + "__repl_max = " + str(max_bytes) + "\n" + """__repl_e = [] __repl_f = [] __repl_total = 0 for _k, _v in list(globals().items()): if _k.startswith('_') or _k in __repl_skip: continue try: __repl_p = _b64.b64encode(_zl.compress(_cp.dumps(_v), 1)).decode() __repl_b = len(__repl_p) if __repl_b > __repl_max: __repl_f.append({'name': _k, 'reason': 'exceeds per-entry snapshot cap'}) elif __repl_total + __repl_b > __repl_max: __repl_f.append({'name': _k, 'reason': 'exceeds total snapshot cap'}) else: __repl_e.append({'name': _k, 'kind': 'value', 'payload': __repl_p}) __repl_total += __repl_b except Exception as _e: __repl_f.append({'name': _k, 'reason': str(_e)}) """ + "get_ipython().display_pub.publish({" + json.dumps(SNAPSHOT_MIME) + ": _js.dumps({'version': 4, 'entries': __repl_e, 'failed': __repl_f})})\n") def _restore_body(name, kind, payload, version): n = json.dumps(name) pl = json.dumps(payload) if version >= 4: body = "globals()[" + n + "] = _cp.loads(_zl.decompress(_b64.b64decode(" + pl + ")))" elif version == 3 and kind == "def": # re-execute captured source and register it in linecache under the code # so a later snapshot can capture it again body = ( "__repl_src = _b64.b64decode(" + pl + ").decode()\n" + " exec(__repl_src, globals())\n" + " __repl_obj = globals().get(" + n + ")\n" + " if __repl_obj is not None:\n" + " __repl_fname = getattr(getattr(__repl_obj, '__code__', None), 'co_filename', None)\n" + " if __repl_fname is None:\n" + " __repl_init = getattr(__repl_obj, '__init__', None)\n" + " __repl_fname = getattr(getattr(__repl_init, '__code__', None), 'co_filename', None)\n" + " if __repl_fname:\n" + " _lc.cache[__repl_fname] = (len(__repl_src.splitlines()), None, __repl_src.splitlines(True), __repl_fname)" ) else: inner = "_b64.b64decode(" + pl + ")" if version == 3: inner = "_zl.decompress(" + inner + ")" body = "globals()[" + n + "] = _pk.loads(" + inner + ")" return ("try:\n " + body + "\n __repl_r['restored'].append(" + n + ")\n" + "except Exception as _e:\n __repl_r['failed'].append({'name': " + n + ", 'reason': str(_e)})\n") def restore_code(entries, version): if version >= 4: head = "import cloudpickle as _cp, base64 as _b64, json as _js, zlib as _zl\n" elif version == 3: head = "import pickle as _pk, base64 as _b64, json as _js, zlib as _zl, linecache as _lc\n" else: head = "import pickle as _pk, base64 as _b64, json as _js\n" per = "".join(_restore_body(e["name"], e.get("kind", "value"), e["payload"], version) for e in entries) return (head + "__repl_r = {'restored': [], 'failed': []}\n" + per + "get_ipython().display_pub.publish({" + json.dumps(RESTORE_MIME) + ": _js.dumps(__repl_r)})\n") def names_code(helper_names): return ("import json as _js\n" + "__repl_skip = set(" + skip_set_literal(helper_names) + ")\n" + "__repl_n = sorted(n for n in globals() if not n.startswith('_') and n not in __repl_skip)\n" + "get_ipython().display_pub.publish({" + json.dumps(NAMES_MIME) + ": _js.dumps(__repl_n)})\n") class Bridge: def __init__(self): self.km = None self.kc = None self.helper_names = [] self.requests = queue.Queue() self.op_id = None # id of the exec currently streaming, for stream routing self.op_payload_mime = None # private MIME awaited for the current internal cell # snapshot policy + name-diff gate: the bridge schedules snapshots itself in quiet gaps, # so a pickle can never land in front of a user cell (the bridge IS the queue) self.policy = None # {"path": str, "max_bytes": int, "period_ms": int} | None self.last_names = None # names at the last persisted write; None = never written self.last_persisted_at = 0.0 self.last_bytes = 0 self.pending_check = False # a cell just completed; the gate owes a name diff self.deferred_restore = None # background revive: runs at the first quiet gap # ----- IO ----- def emit(self, msg): sys.stdout.write(json.dumps(msg) + "\n") sys.stdout.flush() def _stdin_reader(self): # BufferedReader.read(n) blocks until n bytes arrive; line iteration returns per line for line in sys.stdin: line = line.strip() if not line: continue try: self.requests.put(json.loads(line)) except Exception as e: self.emit({"type": "error", "message": "bad request line: %s" % e}) self.requests.put({"op": "shutdown", "id": "bad-line"}) self.requests.put({"op": "__eof__"}) # ----- kernel lifecycle ----- def start_kernel(self): km = KernelManager(kernel_name="python3") try: km.start_kernel() except NoSuchKernel: km = KernelManager(kernel_name="python3", kernel_cmd=[ sys.executable, "-m", "ipykernel_launcher", "-f", "{connection_file}"]) km.start_kernel() kc = km.client() kc.start_channels() try: kc.wait_for_ready(timeout=READY_TIMEOUT_MS / 1000) except Exception: km.shutdown_kernel(now=True) raise self.km, self.kc = km, kc def stop_kernel(self): try: self.kc.stop_channels() except Exception: pass try: self.km.shutdown_kernel(now=True) except Exception: pass # ----- cells ----- def run_cell(self, code, emit_stream, payload_mime=None, emit_id=None): """Run one cell; settle on shell reply AND iopub idle (a reply alone can beat big output). Returns {status, result?, error?, payload?}. Raises when the kernel died.""" kc = self.kc msg_id = kc.execute(code, silent=False, store_history=False, user_expressions={}, allow_stdin=False, stop_on_error=True) self.op_id = msg_id self.op_payload_mime = payload_mime result = None error = None payload = None quiet_for = 0.0 try: while True: # pending interrupts/shutdowns act within one poll cadence self._drain_control_queue() try: msg = kc.get_iopub_msg(timeout=0.25) except Exception: # timeout: poll the heartbeat instead of timing out the whole bridge quiet_for += 0.25 if quiet_for >= 2.0: if not self.kernel_alive(): self.emit({"type": "dying", "message": "kernel heartbeat lost"}) sys.exit(1) quiet_for = 0.0 continue if msg.get("parent_header", {}).get("msg_id") != msg_id: continue quiet_for = 0.0 ctype = msg["msg_type"] content = msg["content"] if ctype == "stream": if emit_stream: self.emit({"type": "stream", "id": emit_id or msg_id, "name": "stderr" if content.get("name") == "stderr" else "stdout", "text": content.get("text") or ""}) elif ctype == "execute_result": data = content.get("data") or {} if "text/plain" in data: result = data["text/plain"] payload = self._take_payload(content) elif ctype == "display_data": payload = self._take_payload(content) elif ctype == "error": error = { "name": content.get("ename") or "Error", "message": content.get("evalue") or "", "stack": content.get("traceback") or [], } elif ctype == "status" and content.get("execution_state") == "idle": break except Exception as e: if not self.kernel_alive(): self.emit({"type": "dying", "message": "kernel died mid-cell"}) sys.exit(1) raise status = "ok" reply = None for _ in range(20): try: candidate = kc.get_shell_msg(timeout=10) except Exception: break if candidate.get("parent_header", {}).get("msg_id") == msg_id: reply = candidate break # stale reply from an earlier op (a late kernel_info, or an "incomplete # input" error flushed by this request): drain it, keep looking if reply: status = reply.get("content", {}).get("status", "ok") if status != "aborted" and error is not None: # iopub error wins over an "ok" reply (ipykernel: compile-phase failures reply ok) status = "error" if status == "error" and error is None: c = reply.get("content", {}) error = {"name": c.get("ename") or "Error", "message": c.get("evalue") or "cell failed", "stack": []} if status == "aborted" and error is None: error = {"name": "KeyboardInterrupt", "message": "cell interrupted", "stack": []} self.op_id = None self.op_payload_mime = None return {"status": status, "result": result, "error": error, "payload": payload} def _take_payload(self, content): data = content.get("data") or {} mime = self.op_payload_mime if mime is not None and mime in data: return data[mime] return None def _drain_control_queue(self): # control-plane ops may jump the queue mid-cell; anything else is left in order while True: try: op = self.requests.get_nowait() except queue.Empty: break if op["op"] == "interrupt": self.interrupt() elif op["op"] in ("shutdown", "__eof__"): self.stop_kernel() sys.exit(0) else: self.requests.put(op) break def interrupt(self): # BlockingKernelClient has no interrupt() in jupyter_client 8.x; the old host sent # interrupt_request on the control channel directly — do the same msg = self.kc.session.msg("interrupt_request", content={}) self.kc.control_channel.send(msg) def kernel_alive(self): try: return self.km.is_alive() except Exception: return False # ----- ops ----- FORCED_SNAPSHOT_MAX_BYTES = 8 * 1024 * 1024 QUIET_WINDOW_S = 0.15 def _names(self): res = self.run_cell(names_code(self.helper_names), emit_stream=False, payload_mime=NAMES_MIME) return json.loads(res["payload"]) if res["payload"] else [] def _snapshot_due(self, names): changed = self.last_names is None or names != self.last_names period = (self.policy or {}).get("period_ms") or 0 stale = (period > 0 and self.last_persisted_at > 0 and (time.time() - self.last_persisted_at) * 1000 >= period and self.last_bytes <= self.FORCED_SNAPSHOT_MAX_BYTES) return changed or stale def quiet_seconds(self): """How long the serve loop may block waiting for ops before the gap counts as quiet. None = nothing is due; block indefinitely.""" if self.deferred_restore is not None: return self.QUIET_WINDOW_S if self.pending_check and self.policy: return self.QUIET_WINDOW_S return None def quiet_tick(self): """The queue stayed empty for the quiet window: run the deferred revive, then the owed snapshot. The gate advances only on a persisted write, so a failed write retries at the next gap.""" if self.deferred_restore is not None: op, self.deferred_restore = self.deferred_restore, None self.op_restore(op) return if not (self.pending_check and self.policy): return self.pending_check = False names = self._names() if not self._snapshot_due(names): return self._write_snapshot(self.policy["path"], self.policy["max_bytes"], names) def op_boot(self, boot): self.helper_names = [h["name"] for h in boot.get("helpers") or []] self.policy = boot.get("snapshot") self.start_kernel() report = [] for h in boot.get("helpers") or []: res = self.run_cell(h["source"], emit_stream=False) if res["status"] == "ok": report.append({"name": h["name"], "ok": True}) else: err = res["error"] or {} report.append({"name": h["name"], "ok": False, "error": (err.get("name") + ": " + err.get("message", "")).strip() or "failed to load"}) self.emit({"type": "ready", "helpers": report}) def op_exec(self, op): res = self.run_cell(op["code"], emit_stream=True, emit_id=op["id"]) out = {"type": "result", "id": op["id"], "status": res["status"]} if res["result"] is not None: out["result"] = res["result"] if res["error"] is not None: out["error"] = res["error"] self.emit(out) def _write_snapshot(self, path, max_bytes, names=None): import cloudpickle # noqa: F401 — fail loudly here, not at bridge boot res = self.run_cell(snapshot_code(self.helper_names, max_bytes), emit_stream=False, payload_mime=SNAPSHOT_MIME) payload = res["payload"] if payload is None: return {"saved": [], "failed": [], "complete": False, "bytes": 0} body = json.loads(payload) entries, failed = body.get("entries", []), body.get("failed", []) # atomic write: temp + rename, so a crash can't corrupt the last good snapshot os.makedirs(os.path.dirname(path) or ".", exist_ok=True) tmp = path + ".tmp" with open(tmp, "w", encoding="utf-8") as fh: json.dump({"version": 4, "entries": entries, "failed": failed}, fh) os.replace(tmp, path) # only a persisted write advances the gate self.last_names = names if names is not None else [e["name"] for e in entries] self.last_persisted_at = time.time() self.last_bytes = sum(len(e["payload"]) for e in entries) return {"saved": [e["name"] for e in entries], "failed": failed, "complete": True, "bytes": self.last_bytes} def op_snapshot(self, op): try: import cloudpickle # noqa: F401 — fail loudly here, not at bridge boot except ImportError as e: self.emit({"type": "reply", "id": op["id"], "error": "cloudpickle is missing from the evaluator venv: %s" % e}) return out = self._write_snapshot(op["path"], op["max_bytes"]) self.emit({"type": "reply", "id": op["id"], **out}) def op_restore(self, op): try: with open(op["path"], encoding="utf-8") as fh: body = json.load(fh) except Exception as e: self.emit({"type": "reply", "id": op["id"], "error": "snapshot unreadable: %s" % e}) return version = body.get("version") if version is not None and version >= 2: entries = body.get("entries") or [] failed_at_save = body.get("failed") or [] else: # v1 files (pre-source-capture) restore via plain pickles entries = [{"name": n, "kind": "value", "payload": p} for n, p in (body.get("vars") or {}).items()] failed_at_save = [] res = self.run_cell(restore_code(entries, version or 4), emit_stream=False, payload_mime=RESTORE_MIME) restored, failed = [], [] if res["payload"] is not None: body2 = json.loads(res["payload"]) restored, failed = body2.get("restored", []), body2.get("failed", []) # merge save-time skips (oversized bindings) so the resume notice names every loss seen = {f["name"] for f in failed} for f in failed_at_save: if f["name"] not in seen: failed.append(f) seen.add(f["name"]) self.emit({"type": "reply", "id": op["id"], "restored": restored, "failed": failed}) def op_names(self, op): res = self.run_cell(names_code(self.helper_names), emit_stream=False, payload_mime=NAMES_MIME) names = json.loads(res["payload"]) if res["payload"] is not None else [] self.emit({"type": "reply", "id": op["id"], "names": names}) # ----- main loop ----- def serve(self): threading.Thread(target=self._stdin_reader, daemon=True).start() boot = self.requests.get() if boot.get("op") == "__eof__": return try: self.op_boot(boot) except Exception as e: self.emit({"type": "error", "message": "boot failed: %s" % e}) sys.exit(1) while True: try: op = self.requests.get(timeout=self.quiet_seconds()) except queue.Empty: self.quiet_tick() continue o = op.get("op") if o == "__eof__": self.stop_kernel() sys.exit(0) try: if o == "exec": self.op_exec(op) self.pending_check = True elif o == "snapshot": self.op_snapshot(op) elif o == "restore": if op.get("defer"): # background revive: run at the first quiet gap, never ahead of a cell; # the reply (same id) settles the host's restoreResult promise self.deferred_restore = op else: self.op_restore(op) elif o == "listNames": self.op_names(op) elif o == "interrupt": self.interrupt() elif o == "shutdown": self.stop_kernel() sys.exit(0) else: self.emit({"type": "error", "id": op.get("id"), "message": "unknown op: %s" % o}) except Exception as e: self.emit({"type": "error", "id": op.get("id"), "message": "%s" % e}) def main(): try: Bridge().serve() except KeyboardInterrupt: pass if __name__ == "__main__": main()