#!/usr/bin/env python3 """Whalent-owned Jupyter kernel sidecar. Stdin/stdout are NDJSON. Stdout is protocol-only; diagnostics go to stderr. The process is launched from the selected Python environment so imports and the kernel executable are guaranteed to belong to that environment. """ from __future__ import annotations import base64 import json import os import queue import signal import sys import threading import time import traceback from typing import Any from jupyter_client import KernelManager from jupyter_client.kernelspec import KernelSpec WRITE_LOCK = threading.Lock() KERNEL_LOCK = threading.RLock() # Serializes every raw operation on the shell zmq socket: the execute worker # polls replies while the stdin thread may send comm messages, and zmq sockets # are not thread-safe. SHELL_LOCK = threading.Lock() # Guards EXECUTIONS, the parent-msg_id routing table the IOPub pump consults. EXEC_LOCK = threading.Lock() KM: KernelManager | None = None KC: Any = None GENERATION = 0 STATE = "dead" CURRENT_CELL: str | None = None SHUTTING_DOWN = False MAX_FRAME_BYTES = 8 * 1024 * 1024 # IOPub message types routed to the owning execution (same set the old # in-worker polling loop handled); everything else with an unknown parent is # dropped, exactly like the old parent-msg_id filter did. ROUTED_IOPUB_TYPES = {"stream", "display_data", "execute_result", "error", "clear_output", "status"} COMM_FRAME_TYPES = {"comm_open": "comm.open", "comm_msg": "comm.msg", "comm_close": "comm.close"} def emit(payload: dict[str, Any]) -> None: encoded = json.dumps(payload, ensure_ascii=False, separators=(",", ":")) if len(encoded.encode("utf-8")) > MAX_FRAME_BYTES: encoded = json.dumps( {"type": "protocol.error", "error": "output frame exceeds 8 MiB"}, ensure_ascii=False, separators=(",", ":"), ) with WRITE_LOCK: sys.stdout.write(encoded + "\n") sys.stdout.flush() def reply(request_id: str, ok: bool, result: Any = None, error: str = "") -> None: payload: dict[str, Any] = {"type": "rpc.reply", "request_id": request_id, "ok": ok} if ok: payload["result"] = result else: payload["error"] = {"message": error or "sidecar request failed"} emit(payload) def safe_status() -> dict[str, Any]: pid = None with KERNEL_LOCK: if KM is not None: provisioner = getattr(KM, "provisioner", None) pid = getattr(provisioner, "pid", None) return { "state": STATE, "generation": GENERATION, "current_cell": CURRENT_CELL, "pid": int(pid) if isinstance(pid, int) else None, } def start_kernel() -> dict[str, Any]: global KM, KC, GENERATION, STATE with KERNEL_LOCK: if KM is not None and KM.is_alive(): return safe_status() # A kernel that died on its own leaves a pump attached to dead # channels; retire it before wiring up the replacement. stop_iopub_pump() STATE = "starting" manager = KernelManager() # The sidecar itself runs under the selected interpreter. Pin ipykernel # to the same executable via an explicit KernelSpec: assigning # `manager.kernel_cmd` is a silent no-op on jupyter_client >= 7 (the # trait is gone), while `_kernel_spec` is the cache slot behind the # KernelManager.kernel_spec property and keeps its meaning on 5.x-8.x. manager._kernel_spec = KernelSpec( argv=[sys.executable, "-m", "ipykernel_launcher", "-f", "{connection_file}"], display_name="Whalent (pinned)", language="python", ) # Kernel stdout/stderr go to the sidecar's stderr so raw fd writes from # user code can never corrupt the NDJSON protocol on stdout. manager.start_kernel(stdout=sys.stderr, stderr=sys.stderr) client = manager.client() client.start_channels() client.wait_for_ready(timeout=30) KM = manager KC = client GENERATION += 1 # The pump outlives individual executes: it delivers kernel-initiated # comm messages even while the kernel is idle. Started after the # generation bump so its comm frames carry the new generation. start_iopub_pump(client) STATE = "idle" emit({"type": "kernel.status", **safe_status()}) return safe_status() def output_from_message(message: dict[str, Any]) -> dict[str, Any] | None: msg_type = str(message.get("header", {}).get("msg_type", "")) content = message.get("content", {}) if msg_type == "stream": return { "outputType": "stream", "name": str(content.get("name", "stdout")), "text": str(content.get("text", "")), } if msg_type in ("display_data", "execute_result"): data = content.get("data", {}) safe_data = { str(key): value if isinstance(value, str) else json.dumps(value, ensure_ascii=False) for key, value in data.items() if isinstance(key, str) } transient = content.get("transient", {}) return { "outputType": msg_type, "data": safe_data, "executionCount": content.get("execution_count"), "displayId": transient.get("display_id") if isinstance(transient, dict) else None, } if msg_type == "error": return { "outputType": "error", "name": str(content.get("ename", "Error")), "value": str(content.get("evalue", "")), "traceback": [str(line) for line in content.get("traceback", [])], } if msg_type == "clear_output": return {"outputType": "clear_output", "wait": bool(content.get("wait"))} return None class ExecutionContext: """Mailbox for one in-flight execute; the IOPub pump routes matching-parent messages here and a ``None`` sentinel means the kernel went away.""" __slots__ = ("cell_id", "generation", "events") def __init__(self, cell_id: str, generation: int) -> None: self.cell_id = cell_id self.generation = generation self.events: queue.Queue[dict[str, Any] | None] = queue.Queue() EXECUTIONS: dict[str, ExecutionContext] = {} PUMP_STOP: threading.Event | None = None PUMP_THREAD: threading.Thread | None = None def encode_buffers(buffers: Any) -> list[str]: encoded: list[str] = [] for item in buffers or []: try: encoded.append(base64.b64encode(bytes(item)).decode("ascii")) except Exception: continue return encoded def iopub_pump(client: Any, stop: threading.Event) -> None: """Resident IOPub reader for one kernel incarnation. Runs from right after start_kernel() until shutdown/restart stops it, so kernel-initiated comm traffic is delivered even while no cell executes. Execution outputs keep their old semantics: they are routed by parent msg_id to the owning execute worker, which emits them itself. """ while not stop.is_set(): try: message = client.get_iopub_msg(timeout=0.2) except queue.Empty: continue except Exception: if stop.is_set(): return time.sleep(0.05) continue msg_type = str(message.get("header", {}).get("msg_type", "")) if msg_type in COMM_FRAME_TYPES: content = message.get("content", {}) if isinstance(message.get("content"), dict) else {} frame: dict[str, Any] = { "type": COMM_FRAME_TYPES[msg_type], "comm_id": str(content.get("comm_id", "")), "data": content.get("data", {}), "buffers": encode_buffers(message.get("buffers")), "generation": GENERATION, } if msg_type == "comm_open": frame["target_name"] = str(content.get("target_name", "")) emit(frame) continue if msg_type not in ROUTED_IOPUB_TYPES: continue parent = str(message.get("parent_header", {}).get("msg_id", "")) with EXEC_LOCK: context = EXECUTIONS.get(parent) if context is not None: context.events.put(message) # No registered parent: drop, matching the old execute-window filter. def start_iopub_pump(client: Any) -> None: global PUMP_STOP, PUMP_THREAD stop = threading.Event() thread = threading.Thread(target=iopub_pump, args=(client, stop), daemon=True) PUMP_STOP = stop PUMP_THREAD = thread thread.start() def stop_iopub_pump() -> None: global PUMP_STOP, PUMP_THREAD stop, thread = PUMP_STOP, PUMP_THREAD PUMP_STOP = None PUMP_THREAD = None if stop is not None: stop.set() if thread is not None and thread is not threading.current_thread(): thread.join(timeout=2.0) # Wake blocked execute workers so they fail fast instead of sitting out # their full timeout (and later interrupting an innocent new kernel). with EXEC_LOCK: contexts = list(EXECUTIONS.values()) for context in contexts: context.events.put(None) def execute_worker(request_id: str, cell_id: str, code: str, timeout_seconds: float) -> None: global STATE, CURRENT_CELL generation = GENERATION msg_id = "" context: ExecutionContext | None = None try: start_kernel() with KERNEL_LOCK: # Capture the generation this execute belongs to (start_kernel may # have revived a dead kernel and bumped it); every frame below is # stamped with it so the daemon can fence stale generations. generation = GENERATION client = KC if client is None: raise RuntimeError("kernel client unavailable") # Register the routing entry atomically around the send: the pump # must never observe an IOPub reply whose parent is not yet routed. with EXEC_LOCK: msg_id = client.execute(code, stop_on_error=True, allow_stdin=False) context = ExecutionContext(cell_id, generation) EXECUTIONS[msg_id] = context STATE = "busy" CURRENT_CELL = cell_id emit({"type": "execution.started", "request_id": request_id, "cell_id": cell_id, **safe_status(), "generation": generation}) deadline = time.monotonic() + timeout_seconds execution_count = None shell_done = False idle = False while time.monotonic() < deadline and not (shell_done and idle): try: message = context.events.get(timeout=0.2) if message is None: raise RuntimeError("kernel was shut down or restarted during execution") msg_type = message.get("header", {}).get("msg_type") if msg_type == "status" and message.get("content", {}).get("execution_state") == "idle": idle = True output = output_from_message(message) if output: if output.get("outputType") == "clear_output": emit({"type": "output.clear", "cell_id": cell_id, "generation": generation, "wait": output.get("wait", False)}) else: emit({"type": "output", "cell_id": cell_id, "generation": generation, "output": output}) except queue.Empty: pass if not shell_done: try: with SHELL_LOCK: shell = client.get_shell_msg(timeout=0.01) if shell.get("parent_header", {}).get("msg_id") == msg_id: execution_count = shell.get("content", {}).get("execution_count") shell_done = True except queue.Empty: pass if not (shell_done and idle): try: with KERNEL_LOCK: # Interrupt only the kernel this execute ran on; after a # restart the same handle points at an innocent successor. if KM is not None and GENERATION == generation: KM.interrupt_kernel() except Exception: pass raise TimeoutError(f"cell execution exceeded {int(timeout_seconds)} seconds") with KERNEL_LOCK: if CURRENT_CELL == cell_id: STATE = "idle" CURRENT_CELL = None emit({ "type": "execution.reply", "request_id": request_id, "ok": True, "cell_id": cell_id, "execution_count": execution_count, **safe_status(), "generation": generation, }) except Exception as exc: # Only release the claim if this worker still owns it: a restart may # already have handed the kernel to a newer execute. with KERNEL_LOCK: if CURRENT_CELL == cell_id: STATE = "idle" if KM is not None and KM.is_alive() else "dead" CURRENT_CELL = None emit({ "type": "execution.reply", "request_id": request_id, "ok": False, "cell_id": cell_id, "error": {"message": str(exc)}, **safe_status(), "generation": generation, }) finally: if msg_id: with EXEC_LOCK: EXECUTIONS.pop(msg_id, None) def send_comm(params: dict[str, Any]) -> dict[str, Any]: """Forward one browser-originated comm_msg to the kernel's shell channel.""" comm_id = str(params.get("comm_id", "")) if not comm_id: raise ValueError("comm_id required") data = params.get("data") if data is None: data = {} if not isinstance(data, dict): raise ValueError("comm data must be an object") buffers = [base64.b64decode(str(item)) for item in params.get("buffers") or []] with KERNEL_LOCK: client = KC if client is None or KM is None or not KM.is_alive(): raise RuntimeError("kernel is not running") msg = client.session.msg("comm_msg", {"comm_id": comm_id, "data": data}) # jupyter_client 8.9.1: ZMQSocketChannel.send() takes no buffers # argument, but Session.send() does — target the shell socket directly. with SHELL_LOCK: client.session.send(client.shell_channel.socket, msg, buffers=buffers or None) return {"sent": True, "comm_id": comm_id, "buffers": len(buffers)} def interrupt_kernel() -> dict[str, Any]: global STATE with KERNEL_LOCK: if KM is None or not KM.is_alive(): raise RuntimeError("kernel is not running") KM.interrupt_kernel() STATE = "idle" emit({"type": "execution.interrupted", **safe_status()}) return safe_status() def shutdown_kernel() -> dict[str, Any]: global KM, KC, STATE, CURRENT_CELL with KERNEL_LOCK: client, manager = KC, KM KC = None KM = None # Stop the pump before closing the channels it reads from. stop_iopub_pump() if client is not None: try: client.stop_channels() except Exception: pass if manager is not None: try: manager.shutdown_kernel(now=True) except Exception: try: manager.cleanup_resources(restart=False) except Exception: pass STATE = "dead" CURRENT_CELL = None emit({"type": "kernel.status", **safe_status()}) return safe_status() def handle(message: dict[str, Any]) -> None: global STATE, CURRENT_CELL request_id = str(message.get("id", "")) method = str(message.get("method", "")) params = message.get("params") if isinstance(message.get("params"), dict) else {} try: if method == "start": reply(request_id, True, start_kernel()) elif method == "execute": cell_id = str(params.get("cell_id", "")) code = str(params.get("code", "")) if not cell_id: raise ValueError("cell_id required") with KERNEL_LOCK: if CURRENT_CELL is not None: raise ValueError(f"kernel busy: cell {CURRENT_CELL} still running") # Claim the kernel before the worker thread exists so two # racing execute requests can never drain IOPub concurrently. STATE = "busy" CURRENT_CELL = cell_id try: threading.Thread( target=execute_worker, args=(request_id, cell_id, code, max(1.0, min(float(params.get("timeout_seconds", 1800)), 3600.0))), daemon=True, ).start() except BaseException: with KERNEL_LOCK: STATE = "idle" if KM is not None and KM.is_alive() else "dead" CURRENT_CELL = None raise elif method == "comm": reply(request_id, True, send_comm(params)) elif method == "interrupt": reply(request_id, True, interrupt_kernel()) elif method == "restart": shutdown_kernel() reply(request_id, True, start_kernel()) elif method == "shutdown": reply(request_id, True, shutdown_kernel()) elif method == "status": reply(request_id, True, safe_status()) else: raise ValueError(f"unsupported method: {method}") except Exception as exc: reply(request_id, False, error=str(exc)) def cleanup(*_: Any) -> None: global SHUTTING_DOWN if SHUTTING_DOWN: return SHUTTING_DOWN = True try: shutdown_kernel() finally: raise SystemExit(0) def main() -> None: signal.signal(signal.SIGTERM, cleanup) if hasattr(signal, "SIGINT"): signal.signal(signal.SIGINT, cleanup) emit({"type": "sidecar.ready", "pid": os.getpid(), "protocol": 1}) for line in sys.stdin: try: message = json.loads(line) if not isinstance(message, dict): raise ValueError("request must be an object") handle(message) except Exception as exc: print(f"[notebook-sidecar] {exc}", file=sys.stderr, flush=True) cleanup() if __name__ == "__main__": main()