@whalent/agent 0.3.265 → 0.3.267

package/dist/jupyter_sidecar.py

@@ -8,6 +8,7 @@ kernel executable are guaranteed to belong to that environment.
 
 from __future__ import annotations
 
+import base64
 import json
 import os
 import queue
@@ -24,6 +25,12 @@ 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
@@ -32,6 +39,12 @@ 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=(",", ":"))
@@ -74,6 +87,9 @@ def start_kernel() -> dict[str, Any]:
     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
@@ -95,6 +111,10 @@ def start_kernel() -> dict[str, Any]:
         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()
@@ -135,9 +155,106 @@ def output_from_message(message: dict[str, Any]) -> dict[str, Any] | None:
     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:
@@ -148,7 +265,12 @@ def execute_worker(request_id: str, cell_id: str, code: str, timeout_seconds: fl
             client = KC
             if client is None:
                 raise RuntimeError("kernel client unavailable")
-            msg_id = client.execute(code, stop_on_error=True, allow_stdin=False)
+            # 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})
@@ -158,22 +280,24 @@ def execute_worker(request_id: str, cell_id: str, code: str, timeout_seconds: fl
         idle = False
         while time.monotonic() < deadline and not (shell_done and idle):
             try:
-                message = client.get_iopub_msg(timeout=0.2)
-                if message.get("parent_header", {}).get("msg_id") == msg_id:
-                    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})
+                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:
-                    shell = client.get_shell_msg(timeout=0.01)
+                    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
@@ -182,7 +306,9 @@ def execute_worker(request_id: str, cell_id: str, code: str, timeout_seconds: fl
         if not (shell_done and idle):
             try:
                 with KERNEL_LOCK:
-                    if KM is not None:
+                    # 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
@@ -216,6 +342,33 @@ def execute_worker(request_id: str, cell_id: str, code: str, timeout_seconds: fl
             **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]:
@@ -235,6 +388,8 @@ def shutdown_kernel() -> dict[str, Any]:
         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()
@@ -285,6 +440,8 @@ def handle(message: dict[str, Any]) -> None:
                     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":

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@whalent/agent",
-  "version": "0.3.265",
+  "version": "0.3.267",
   "description": "Stable wrapper for Whalent Agent core runtime",
   "type": "module",
   "main": "dist/index.cjs",
@@ -54,7 +54,7 @@
     "yaml": "^2.4.0"
   },
   "optionalDependencies": {
-    "@whalent/agent-core": "0.3.265",
+    "@whalent/agent-core": "0.3.267",
     "node-pty": "^1.1.0"
   },
   "devDependencies": {