can1357
9a2cc3bda0
- Added `cwd` and `env` optional parameters to kernel execution API for runtime working directory and environment variable control. - Implemented runtime environment setup in Python runner with `_apply_request_runtime()` to apply cwd and env from request before code execution. - Enhanced SIGINT handler management with `active_executions` counter and `_begin_exec_sigint()` / `_end_exec_sigint()` functions to prevent state mutation during concurrent execution. - Changed `SearchRenderArgs.paths` parameter type from `string[]` to `string | string[]` to accept single string paths. - Added comprehensive test coverage for kernel cwd updates, timeout interruption safety, and SystemExit handling in shared executor sessions.
1029 lines
32 KiB
Python
1029 lines
32 KiB
Python
"""OMP Python runner — subprocess wrapper used by the coding-agent host.
|
|
|
|
NDJSON protocol over stdin/stdout. Host writes one JSON object per line;
|
|
wrapper writes typed frames back.
|
|
|
|
Host -> wrapper:
|
|
{"id": str, "code": str, "silent": bool?, "storeHistory": bool?}
|
|
{"id": str, "code": str, "silent": bool?, "storeHistory": bool?, "cwd": str?, "env": dict?}
|
|
{"type": "exit"} # graceful shutdown
|
|
|
|
Wrapper -> host:
|
|
{"type": "started", "id": ...}
|
|
{"type": "stdout", "id": ..., "data": str}
|
|
{"type": "stderr", "id": ..., "data": str}
|
|
{"type": "display", "id": ..., "bundle": {<mime>: <value>}}
|
|
{"type": "result", "id": ..., "bundle": {<mime>: <value>}}
|
|
{"type": "error", "id": ..., "ename": str, "evalue": str, "traceback": [str]}
|
|
{"type": "done", "id": ..., "status": "ok"|"error",
|
|
"executionCount": int, "cancelled": bool}
|
|
|
|
The runner is intentionally self-contained: no third-party imports, no IPython.
|
|
Magics are translated by a small line-scanner before AST parsing; rich display
|
|
falls back through `_repr_*_` methods so pandas/PIL/plotly etc. still render
|
|
when installed.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import ast
|
|
import contextvars
|
|
import base64
|
|
import builtins
|
|
import inspect
|
|
import io
|
|
import json
|
|
import os
|
|
import re
|
|
import runpy
|
|
import shlex
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
import traceback
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Frame writer
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_RAW_STDOUT = sys.__stdout__
|
|
_RAW_STDERR = sys.__stderr__
|
|
_OUT_LOCK = threading.Lock()
|
|
|
|
|
|
def _json_default(o: Any) -> Any:
|
|
try:
|
|
return repr(o)
|
|
except Exception:
|
|
return f"<unrepr {type(o).__name__}>"
|
|
|
|
|
|
def _emit(frame: dict) -> None:
|
|
"""Serialize a frame and write it to the host as a single NDJSON line."""
|
|
line = json.dumps(frame, ensure_ascii=False, default=_json_default)
|
|
with _OUT_LOCK:
|
|
_RAW_STDOUT.write(line)
|
|
_RAW_STDOUT.write("\n")
|
|
_RAW_STDOUT.flush()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# User stdout/stderr proxies
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _StreamProxy(io.TextIOBase):
|
|
"""Emit each ``write()`` as a typed frame tied to the current request."""
|
|
|
|
def __init__(self, kind: str) -> None:
|
|
super().__init__()
|
|
self._kind = kind
|
|
|
|
def writable(self) -> bool: # noqa: D401 - protocol method
|
|
return True
|
|
|
|
def isatty(self) -> bool: # noqa: D401 - protocol method
|
|
return False
|
|
|
|
def write(self, data: Any) -> int: # type: ignore[override]
|
|
if not isinstance(data, str):
|
|
data = str(data)
|
|
if not data:
|
|
return 0
|
|
rid = _CURRENT_RID.get()
|
|
if rid is None:
|
|
_RAW_STDERR.write(data)
|
|
_RAW_STDERR.flush()
|
|
return len(data)
|
|
_emit({"type": self._kind, "id": rid, "data": data})
|
|
return len(data)
|
|
|
|
def flush(self) -> None: # noqa: D401 - protocol method
|
|
return None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Runner state
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _RunnerState:
|
|
def __init__(self) -> None:
|
|
self.execution_count: int = 0
|
|
self.cancel_requested: bool = False
|
|
# User globals — kept across requests when running in session mode.
|
|
self.user_ns: dict[str, Any] = {
|
|
"__name__": "__main__",
|
|
"__doc__": None,
|
|
"__builtins__": builtins,
|
|
}
|
|
self.last_install_marker: int = 0
|
|
self.loop: asyncio.AbstractEventLoop | None = None
|
|
self.active_executions: int = 0
|
|
|
|
|
|
_CURRENT_RID: contextvars.ContextVar[str | None] = contextvars.ContextVar("omp_current_rid", default=None)
|
|
|
|
_STATE = _RunnerState()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Magic source transformer
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_MAGIC_LINE_RE = re.compile(r"^(?P<indent>[ \t]*)(?P<name>[A-Za-z_][A-Za-z_0-9]*)(?:[ \t]+(?P<args>.*))?$")
|
|
_ASSIGN_LINE_RE = re.compile(
|
|
r"^(?P<indent>[ \t]*)(?P<lhs>[A-Za-z_][A-Za-z_0-9.\[\], ]*?)\s*=\s*(?P<rhs>.+)$"
|
|
)
|
|
|
|
|
|
def _fold_continuations(lines: list[str], start: int) -> tuple[str, int]:
|
|
"""Fold trailing backslash continuations starting at ``start``. Returns
|
|
``(folded_text, lines_consumed)``."""
|
|
parts: list[str] = []
|
|
i = start
|
|
while i < len(lines):
|
|
line = lines[i]
|
|
if line.endswith("\\"):
|
|
parts.append(line[:-1])
|
|
i += 1
|
|
continue
|
|
parts.append(line)
|
|
i += 1
|
|
break
|
|
return ("".join(parts), i - start)
|
|
|
|
|
|
def _quote_arg(text: str) -> str:
|
|
"""Return a Python string literal that round-trips ``text`` exactly."""
|
|
return json.dumps(text, ensure_ascii=False)
|
|
|
|
|
|
def transform_cell(source: str) -> str:
|
|
"""Translate IPython-style magics + shell escapes into plain Python.
|
|
|
|
Rules
|
|
-----
|
|
* ``%name args`` -> ``__omp_magic("name", "args")``
|
|
* ``var = %name args`` -> ``var = __omp_magic("name", "args")``
|
|
* ``!cmd`` -> ``__omp_shell("cmd")``
|
|
* ``var = !cmd`` -> ``var = __omp_shell("cmd")``
|
|
* ``%%name args\\n<body>`` -> ``__omp_magic_cell("name", "args", "<body>")``
|
|
(cell magic must be the first non-whitespace token of a top-level line and
|
|
consumes the remainder of the cell)
|
|
|
|
Lines inside strings or comments are left alone — we operate on the raw
|
|
text before parsing, but the scanner only fires on the first token of each
|
|
physical line and never touches the body of triple-quoted strings because
|
|
those bodies are never first tokens themselves.
|
|
"""
|
|
|
|
if "%" not in source and "!" not in source:
|
|
return source
|
|
|
|
lines = source.splitlines()
|
|
out: list[str] = []
|
|
i = 0
|
|
while i < len(lines):
|
|
line = lines[i]
|
|
stripped = line.lstrip()
|
|
indent = line[: len(line) - len(stripped)]
|
|
|
|
# Cell magic — consumes from here to EOF.
|
|
if stripped.startswith("%%"):
|
|
head, _ = _split_magic_head(stripped[2:])
|
|
name, args = head
|
|
body_lines = lines[i + 1 :]
|
|
body = "\n".join(body_lines)
|
|
out.append(
|
|
f"{indent}__omp_magic_cell({_quote_arg(name)}, {_quote_arg(args)}, {_quote_arg(body)})"
|
|
)
|
|
return "\n".join(out)
|
|
|
|
# Line magic / shell at start of line.
|
|
if stripped.startswith("%") and not stripped.startswith("%%"):
|
|
folded, consumed = _fold_continuations(lines, i)
|
|
stripped_folded = folded.lstrip()
|
|
indent = folded[: len(folded) - len(stripped_folded)]
|
|
head, _ = _split_magic_head(stripped_folded[1:])
|
|
name, args = head
|
|
out.append(f"{indent}__omp_magic({_quote_arg(name)}, {_quote_arg(args)})")
|
|
i += consumed
|
|
continue
|
|
|
|
if stripped.startswith("!"):
|
|
folded, consumed = _fold_continuations(lines, i)
|
|
stripped_folded = folded.lstrip()
|
|
indent = folded[: len(folded) - len(stripped_folded)]
|
|
cmd = stripped_folded[1:].strip()
|
|
out.append(f"{indent}__omp_shell({_quote_arg(cmd)})")
|
|
i += consumed
|
|
continue
|
|
|
|
# Assignment forms: var = %magic / var = !cmd
|
|
m = _ASSIGN_LINE_RE.match(line)
|
|
if m:
|
|
rhs = m.group("rhs").strip()
|
|
if rhs.startswith("!"):
|
|
cmd = rhs[1:].strip()
|
|
out.append(f"{m.group('indent')}{m.group('lhs').rstrip()} = __omp_shell({_quote_arg(cmd)})")
|
|
i += 1
|
|
continue
|
|
if rhs.startswith("%") and not rhs.startswith("%%"):
|
|
head, _ = _split_magic_head(rhs[1:])
|
|
name, args = head
|
|
out.append(
|
|
f"{m.group('indent')}{m.group('lhs').rstrip()} = __omp_magic({_quote_arg(name)}, {_quote_arg(args)})"
|
|
)
|
|
i += 1
|
|
continue
|
|
|
|
out.append(line)
|
|
i += 1
|
|
|
|
return "\n".join(out)
|
|
|
|
|
|
def _split_magic_head(text: str) -> tuple[tuple[str, str], str]:
|
|
"""Split ``"name rest"`` into ``("name", "rest")``."""
|
|
text = text.lstrip()
|
|
if not text:
|
|
return ("", ""), ""
|
|
m = re.match(r"([A-Za-z_][A-Za-z_0-9]*)(?:\s+(.*))?$", text)
|
|
if not m:
|
|
return ("", text), ""
|
|
return (m.group(1), (m.group(2) or "").rstrip()), ""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Magic registry
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_LINE_MAGICS: dict[str, Callable[[str], Any]] = {}
|
|
_CELL_MAGICS: dict[str, Callable[[str, str], Any]] = {}
|
|
|
|
|
|
def line_magic(name: str) -> Callable[[Callable[[str], Any]], Callable[[str], Any]]:
|
|
def decorator(fn: Callable[[str], Any]) -> Callable[[str], Any]:
|
|
_LINE_MAGICS[name] = fn
|
|
return fn
|
|
|
|
return decorator
|
|
|
|
|
|
def cell_magic(name: str) -> Callable[[Callable[[str, str], Any]], Callable[[str, str], Any]]:
|
|
def decorator(fn: Callable[[str, str], Any]) -> Callable[[str, str], Any]:
|
|
_CELL_MAGICS[name] = fn
|
|
return fn
|
|
|
|
return decorator
|
|
|
|
|
|
def _emit_status(op: str, **data: Any) -> None:
|
|
bundle = {"application/x-omp-status": {"op": op, **data}}
|
|
rid = _CURRENT_RID.get()
|
|
if rid is None:
|
|
return
|
|
_emit({"type": "display", "id": rid, "bundle": bundle})
|
|
|
|
|
|
@line_magic("pip")
|
|
def _magic_pip(args: str) -> None:
|
|
argv = shlex.split(args) if args else ["--help"]
|
|
cmd = [sys.executable, "-m", "pip", *argv]
|
|
proc = subprocess.Popen(
|
|
cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
bufsize=1,
|
|
)
|
|
installed_packages: list[str] = []
|
|
assert proc.stdout is not None
|
|
for raw_line in proc.stdout:
|
|
sys.stdout.write(raw_line)
|
|
m = re.search(r"Successfully installed\s+(.+)$", raw_line)
|
|
if m:
|
|
for token in m.group(1).split():
|
|
# Token is name-version; drop the version suffix.
|
|
pkg = token.rsplit("-", 1)[0]
|
|
installed_packages.append(pkg.replace("_", "-"))
|
|
proc.wait()
|
|
if installed_packages:
|
|
import importlib
|
|
|
|
importlib.invalidate_caches()
|
|
prefixes = {pkg.lower().replace("-", "_") for pkg in installed_packages}
|
|
for mod_name in list(sys.modules):
|
|
head = mod_name.split(".", 1)[0].lower()
|
|
if head in prefixes:
|
|
sys.modules.pop(mod_name, None)
|
|
_emit_status("pip", args=args, installed=installed_packages, exit_code=proc.returncode)
|
|
|
|
|
|
@line_magic("cd")
|
|
def _magic_cd(args: str) -> str:
|
|
path = os.path.expanduser(args.strip()) or os.path.expanduser("~")
|
|
os.chdir(path)
|
|
cwd = os.getcwd()
|
|
_emit_status("cd", path=cwd)
|
|
return cwd
|
|
|
|
|
|
@line_magic("pwd")
|
|
def _magic_pwd(_args: str) -> str:
|
|
cwd = os.getcwd()
|
|
_emit_status("pwd", path=cwd)
|
|
return cwd
|
|
|
|
|
|
@line_magic("ls")
|
|
def _magic_ls(args: str) -> list[str]:
|
|
target = os.path.expanduser(args.strip()) or "."
|
|
entries = sorted(os.listdir(target))
|
|
_emit_status("ls", path=os.path.abspath(target), count=len(entries))
|
|
return entries
|
|
|
|
|
|
@line_magic("env")
|
|
def _magic_env(args: str) -> Any:
|
|
args = args.strip()
|
|
if not args:
|
|
return dict(sorted(os.environ.items()))
|
|
if "=" in args:
|
|
key, value = args.split("=", 1)
|
|
os.environ[key.strip()] = value.strip()
|
|
return value.strip()
|
|
return os.environ.get(args)
|
|
|
|
|
|
@line_magic("set_env")
|
|
def _magic_set_env(args: str) -> str:
|
|
parts = args.split(None, 1)
|
|
if len(parts) != 2:
|
|
raise ValueError("Usage: %set_env KEY VALUE")
|
|
key, value = parts
|
|
os.environ[key] = value
|
|
return value
|
|
|
|
|
|
@line_magic("time")
|
|
def _magic_time(args: str) -> Any:
|
|
start = time.perf_counter()
|
|
result = eval(args, _STATE.user_ns)
|
|
elapsed = time.perf_counter() - start
|
|
sys.stdout.write(f"Wall time: {elapsed * 1000:.2f} ms\n")
|
|
_emit_status("time", elapsed_ms=round(elapsed * 1000, 3))
|
|
return result
|
|
|
|
|
|
@line_magic("timeit")
|
|
def _magic_timeit(args: str) -> None:
|
|
import timeit as _timeit
|
|
|
|
timer = _timeit.Timer(stmt=args, globals=_STATE.user_ns)
|
|
iters, total = timer.autorange()
|
|
per = total / iters
|
|
sys.stdout.write(f"{iters} loops, best of 1: {per * 1e6:.2f} us per loop\n")
|
|
_emit_status("timeit", loops=iters, total_ms=round(total * 1000, 3))
|
|
|
|
|
|
@line_magic("who")
|
|
def _magic_who(_args: str) -> list[str]:
|
|
names = sorted(
|
|
name
|
|
for name, value in _STATE.user_ns.items()
|
|
if not name.startswith("_") and not callable(value) or hasattr(value, "__class__")
|
|
)
|
|
return [n for n in names if not n.startswith("__")]
|
|
|
|
|
|
@line_magic("whos")
|
|
def _magic_whos(_args: str) -> list[tuple[str, str]]:
|
|
rows = []
|
|
for name in sorted(_STATE.user_ns):
|
|
if name.startswith("__"):
|
|
continue
|
|
value = _STATE.user_ns[name]
|
|
rows.append((name, type(value).__name__))
|
|
return rows
|
|
|
|
|
|
@line_magic("reset")
|
|
def _magic_reset(_args: str) -> None:
|
|
_STATE.user_ns.clear()
|
|
_STATE.user_ns.update({"__name__": "__main__", "__doc__": None, "__builtins__": builtins})
|
|
_install_builtins(_STATE.user_ns)
|
|
_emit_status("reset")
|
|
|
|
|
|
@line_magic("load")
|
|
def _magic_load(args: str) -> None:
|
|
path = Path(os.path.expanduser(args.strip()))
|
|
source = path.read_text(encoding="utf-8")
|
|
_emit({"type": "display", "id": _CURRENT_RID.get(), "bundle": {"text/plain": source}})
|
|
_exec_source(source, _STATE.user_ns)
|
|
|
|
|
|
@line_magic("run")
|
|
def _magic_run(args: str) -> None:
|
|
parts = shlex.split(args) if args else []
|
|
if not parts:
|
|
raise ValueError("Usage: %run <path>")
|
|
target = os.path.expanduser(parts[0])
|
|
saved_argv = sys.argv
|
|
try:
|
|
sys.argv = [target, *parts[1:]]
|
|
result_ns = runpy.run_path(target, run_name="__main__")
|
|
finally:
|
|
sys.argv = saved_argv
|
|
for name, value in result_ns.items():
|
|
if name.startswith("__"):
|
|
continue
|
|
_STATE.user_ns[name] = value
|
|
|
|
|
|
@cell_magic("bash")
|
|
def _magic_cell_bash(args: str, body: str) -> int:
|
|
return _run_shell_body(body, shell_arg="/bin/bash")
|
|
|
|
|
|
@cell_magic("sh")
|
|
def _magic_cell_sh(args: str, body: str) -> int:
|
|
return _run_shell_body(body, shell_arg="/bin/sh")
|
|
|
|
|
|
@cell_magic("capture")
|
|
def _magic_cell_capture(args: str, body: str) -> str:
|
|
"""Capture stdout/stderr of body; bind to ``args`` (a name) if provided."""
|
|
captured = io.StringIO()
|
|
saved_stdout, saved_stderr = sys.stdout, sys.stderr
|
|
sys.stdout = sys.stderr = captured
|
|
try:
|
|
_exec_source(body, _STATE.user_ns)
|
|
finally:
|
|
sys.stdout, sys.stderr = saved_stdout, saved_stderr
|
|
text = captured.getvalue()
|
|
name = args.strip()
|
|
if name:
|
|
_STATE.user_ns[name] = text
|
|
return text
|
|
|
|
|
|
@cell_magic("timeit")
|
|
def _magic_cell_timeit(args: str, body: str) -> None:
|
|
import timeit as _timeit
|
|
|
|
timer = _timeit.Timer(stmt=body, globals=_STATE.user_ns)
|
|
iters, total = timer.autorange()
|
|
per = total / iters
|
|
sys.stdout.write(f"{iters} loops, best of 1: {per * 1e6:.2f} us per loop\n")
|
|
_emit_status("timeit", loops=iters, total_ms=round(total * 1000, 3))
|
|
|
|
|
|
@cell_magic("writefile")
|
|
def _magic_cell_writefile(args: str, body: str) -> str:
|
|
path = Path(os.path.expanduser(args.strip()))
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(body, encoding="utf-8")
|
|
_emit_status("writefile", path=str(path), bytes=len(body))
|
|
return str(path)
|
|
|
|
|
|
def _run_shell_body(body: str, *, shell_arg: str) -> int:
|
|
proc = subprocess.Popen(
|
|
[shell_arg, "-c", body],
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
bufsize=1,
|
|
)
|
|
assert proc.stdout is not None
|
|
for raw_line in proc.stdout:
|
|
sys.stdout.write(raw_line)
|
|
proc.wait()
|
|
return proc.returncode
|
|
|
|
|
|
def __omp_magic(name: str, args: str) -> Any:
|
|
fn = _LINE_MAGICS.get(name)
|
|
if fn is None:
|
|
raise NameError(f"UsageError: Line magic function '%{name}' not found.")
|
|
return fn(args)
|
|
|
|
|
|
def __omp_magic_cell(name: str, args: str, body: str) -> Any:
|
|
fn = _CELL_MAGICS.get(name)
|
|
if fn is None:
|
|
raise NameError(f"UsageError: Cell magic function '%%{name}' not found.")
|
|
return fn(args, body)
|
|
|
|
|
|
class _ShellResult(list):
|
|
"""Result of ``!cmd`` — list of stripped output lines."""
|
|
|
|
def __init__(self, lines: list[str], returncode: int) -> None:
|
|
super().__init__(lines)
|
|
self.returncode = returncode
|
|
|
|
@property
|
|
def n(self) -> str: # IPython compat
|
|
return "\n".join(self)
|
|
|
|
@property
|
|
def s(self) -> str: # IPython compat
|
|
return " ".join(self)
|
|
|
|
|
|
def __omp_shell(cmd: str) -> _ShellResult:
|
|
proc = subprocess.run(
|
|
cmd,
|
|
shell=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
)
|
|
if proc.stdout:
|
|
sys.stdout.write(proc.stdout)
|
|
lines = [line for line in (proc.stdout or "").splitlines()]
|
|
return _ShellResult(lines, proc.returncode)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Display dispatch
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_REPR_MIMES = [
|
|
("_repr_html_", "text/html"),
|
|
("_repr_markdown_", "text/markdown"),
|
|
("_repr_svg_", "image/svg+xml"),
|
|
("_repr_png_", "image/png"),
|
|
("_repr_jpeg_", "image/jpeg"),
|
|
("_repr_json_", "application/json"),
|
|
("_repr_latex_", "text/latex"),
|
|
]
|
|
|
|
|
|
def _coerce_image_bytes(value: Any) -> str:
|
|
if isinstance(value, (bytes, bytearray)):
|
|
return base64.b64encode(bytes(value)).decode("ascii")
|
|
if isinstance(value, str):
|
|
return value
|
|
return base64.b64encode(repr(value).encode("utf-8")).decode("ascii")
|
|
|
|
|
|
def _mime_bundle(value: Any) -> dict:
|
|
"""Build a Jupyter-style MIME bundle for ``value``.
|
|
|
|
Honors ``_repr_mimebundle_`` first, falls back to individual ``_repr_*_``
|
|
accessors, and always provides ``text/plain``.
|
|
"""
|
|
bundle: dict[str, Any] = {}
|
|
|
|
mimebundle = getattr(value, "_repr_mimebundle_", None)
|
|
if callable(mimebundle):
|
|
try:
|
|
data = mimebundle()
|
|
except Exception:
|
|
data = None
|
|
if isinstance(data, tuple):
|
|
data = data[0]
|
|
if isinstance(data, dict):
|
|
bundle.update({str(k): v for k, v in data.items()})
|
|
|
|
for attr, mime in _REPR_MIMES:
|
|
if mime in bundle:
|
|
continue
|
|
repr_fn = getattr(value, attr, None)
|
|
if not callable(repr_fn):
|
|
continue
|
|
try:
|
|
data = repr_fn()
|
|
except Exception:
|
|
continue
|
|
if data is None:
|
|
continue
|
|
if mime in ("image/png", "image/jpeg"):
|
|
bundle[mime] = _coerce_image_bytes(data)
|
|
else:
|
|
bundle[mime] = data
|
|
|
|
if "text/plain" not in bundle:
|
|
try:
|
|
bundle["text/plain"] = repr(value)
|
|
except Exception:
|
|
bundle["text/plain"] = f"<unrepr {type(value).__name__}>"
|
|
|
|
return bundle
|
|
|
|
|
|
def _emit_display(bundle: dict, *, kind: str = "display") -> None:
|
|
rid = _CURRENT_RID.get()
|
|
if rid is None:
|
|
return
|
|
_emit({"type": kind, "id": rid, "bundle": bundle})
|
|
|
|
|
|
def __omp_display(value: Any, *, raw: bool = False, kind: str = "display") -> None:
|
|
if raw:
|
|
if not isinstance(value, dict):
|
|
raise TypeError("display(..., raw=True) requires a MIME bundle dict")
|
|
bundle = {str(k): v for k, v in value.items()}
|
|
if "text/plain" not in bundle:
|
|
bundle["text/plain"] = ""
|
|
_emit_display(bundle, kind=kind)
|
|
return
|
|
_emit_display(_mime_bundle(value), kind=kind)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Matplotlib post-cell flush
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _flush_matplotlib_figures() -> None:
|
|
plt = sys.modules.get("matplotlib.pyplot")
|
|
if plt is None:
|
|
return
|
|
try:
|
|
fignums = list(plt.get_fignums())
|
|
except Exception:
|
|
return
|
|
for num in fignums:
|
|
try:
|
|
fig = plt.figure(num)
|
|
buf = io.BytesIO()
|
|
fig.savefig(buf, format="png", bbox_inches="tight")
|
|
data = base64.b64encode(buf.getvalue()).decode("ascii")
|
|
_emit_display({"image/png": data, "text/plain": f"<Figure {num}>"})
|
|
plt.close(fig)
|
|
except Exception:
|
|
continue
|
|
|
|
|
|
# Force a non-interactive backend before user code imports matplotlib. Set as
|
|
# environ default so the user can still override it explicitly.
|
|
os.environ.setdefault("MPLBACKEND", "Agg")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Builtin injection
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _install_builtins(ns: dict) -> None:
|
|
ns["display"] = __omp_display
|
|
ns["__omp_display"] = __omp_display
|
|
ns["__omp_magic"] = __omp_magic
|
|
ns["__omp_magic_cell"] = __omp_magic_cell
|
|
ns["__omp_shell"] = __omp_shell
|
|
ns["__omp_current_run_id__"] = lambda: _CURRENT_RID.get()
|
|
|
|
|
|
_install_builtins(_STATE.user_ns)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Source execution (split last expression for rich display)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_TLA_FLAG = getattr(ast, "PyCF_ALLOW_TOP_LEVEL_AWAIT", 0x2000)
|
|
|
|
|
|
def _await_sync(coro) -> Any:
|
|
try:
|
|
running_loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
running_loop = None
|
|
if running_loop is not None and running_loop.is_running():
|
|
raise RuntimeError("top-level await is not supported from synchronous magic execution")
|
|
return asyncio.run(coro)
|
|
|
|
|
|
def _run_compiled_sync(code, ns: dict, *, want_value: bool) -> Any:
|
|
"""Synchronous execution path used by nested magic helpers."""
|
|
if code.co_flags & inspect.CO_COROUTINE:
|
|
result = _await_sync(eval(code, ns))
|
|
return result if want_value else None
|
|
if want_value:
|
|
return eval(code, ns)
|
|
exec(code, ns)
|
|
return None
|
|
|
|
|
|
|
|
async def _run_compiled_async(code, ns: dict, *, want_value: bool) -> Any:
|
|
"""Execute a code object in the persistent event loop.
|
|
|
|
Coroutine code is awaited in this task so top-level ``await`` interleaves
|
|
with sibling requests. Plain statement/expression code runs on the main
|
|
runner thread so SIGINT can interrupt it reliably.
|
|
"""
|
|
if code.co_flags & inspect.CO_COROUTINE:
|
|
result = await eval(code, ns)
|
|
return result if want_value else None
|
|
if want_value:
|
|
return eval(code, ns)
|
|
exec(code, ns)
|
|
return None
|
|
|
|
|
|
def _compile_source(source: str) -> tuple[Any, Any | None, bool]:
|
|
module = ast.parse(source, mode="exec")
|
|
if not module.body:
|
|
return None, None, False
|
|
|
|
last = module.body[-1]
|
|
if isinstance(last, ast.Expr):
|
|
body_module = ast.Module(body=module.body[:-1], type_ignores=[])
|
|
expr_module = ast.Expression(body=last.value)
|
|
ast.copy_location(expr_module, last)
|
|
body_code = compile(body_module, "<cell>", "exec", flags=_TLA_FLAG)
|
|
expr_code = compile(expr_module, "<cell>", "eval", flags=_TLA_FLAG)
|
|
return body_code, expr_code, True
|
|
|
|
return compile(module, "<cell>", "exec", flags=_TLA_FLAG), None, False
|
|
|
|
|
|
def _exec_source(source: str, ns: dict) -> None:
|
|
"""Synchronous source execution for legacy magic helpers."""
|
|
body_code, expr_code, has_expr = _compile_source(source)
|
|
if body_code is None:
|
|
return
|
|
_run_compiled_sync(body_code, ns, want_value=False)
|
|
if has_expr and expr_code is not None:
|
|
value = _run_compiled_sync(expr_code, ns, want_value=True)
|
|
if value is not None:
|
|
__omp_display(value, kind="result")
|
|
|
|
|
|
async def _exec_source_async(source: str, ns: dict) -> None:
|
|
"""Compile + execute ``source``; if the last node is an expression, route
|
|
its value through ``__omp_display`` so dataframes/figures render rich.
|
|
Top-level ``await`` / ``async for`` / ``async with`` is permitted; awaited
|
|
regions yield to other requests in the runner's persistent event loop."""
|
|
body_code, expr_code, has_expr = _compile_source(source)
|
|
if body_code is None:
|
|
return
|
|
await _run_compiled_async(body_code, ns, want_value=False)
|
|
if has_expr and expr_code is not None:
|
|
value = await _run_compiled_async(expr_code, ns, want_value=True)
|
|
if value is not None:
|
|
__omp_display(value, kind="result")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Signal handling
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _install_idle_sigint() -> None:
|
|
try:
|
|
signal.signal(signal.SIGINT, signal.SIG_IGN)
|
|
except (OSError, ValueError):
|
|
# Some platforms (Windows in non-console mode) reject this; fine.
|
|
pass
|
|
|
|
|
|
def _install_exec_sigint() -> None:
|
|
try:
|
|
signal.signal(signal.SIGINT, signal.default_int_handler)
|
|
except (OSError, ValueError):
|
|
pass
|
|
|
|
|
|
def _begin_exec_sigint() -> None:
|
|
_STATE.active_executions += 1
|
|
_install_exec_sigint()
|
|
|
|
|
|
def _end_exec_sigint() -> None:
|
|
if _STATE.active_executions > 0:
|
|
_STATE.active_executions -= 1
|
|
if _STATE.active_executions == 0:
|
|
_install_idle_sigint()
|
|
|
|
|
|
_MANAGED_ENV_KEYS = (
|
|
"PI_SESSION_FILE",
|
|
"PI_ARTIFACTS_DIR",
|
|
"PI_TOOL_BRIDGE_URL",
|
|
"PI_TOOL_BRIDGE_TOKEN",
|
|
"PI_TOOL_BRIDGE_SESSION",
|
|
)
|
|
|
|
|
|
def _apply_request_runtime(req: dict) -> None:
|
|
cwd = req.get("cwd")
|
|
if isinstance(cwd, str) and cwd:
|
|
os.chdir(cwd)
|
|
try:
|
|
sys.path.remove(cwd)
|
|
except ValueError:
|
|
pass
|
|
sys.path.insert(0, cwd)
|
|
|
|
env = req.get("env")
|
|
if isinstance(env, dict):
|
|
for key in _MANAGED_ENV_KEYS:
|
|
value = env.get(key)
|
|
if isinstance(value, str):
|
|
os.environ[key] = value
|
|
elif value is None:
|
|
os.environ.pop(key, None)
|
|
|
|
def _start_parent_watchdog() -> None:
|
|
"""Self-terminate when the host process dies.
|
|
|
|
The main loop only exits when stdin EOFs, which only happens once user
|
|
code finishes and the next ``readline`` call returns. If the host gets
|
|
SIGKILL mid-execution (or any way that skips graceful shutdown) the
|
|
runner would otherwise outlive its parent and keep holding kernel
|
|
state. Poll ``os.getppid()`` instead and ``os._exit`` the moment we get
|
|
reparented \u2014 covers POSIX hosts. Windows has no reliable ppid
|
|
equivalent; there we still bail out on the next stdin read.
|
|
"""
|
|
if os.name != "posix":
|
|
return
|
|
original_ppid = os.getppid()
|
|
if original_ppid <= 1:
|
|
return
|
|
|
|
def watch() -> None:
|
|
while True:
|
|
try:
|
|
if os.getppid() != original_ppid:
|
|
os._exit(0)
|
|
except Exception:
|
|
return
|
|
time.sleep(10)
|
|
|
|
thread = threading.Thread(target=watch, name="omp-parent-watchdog", daemon=True)
|
|
thread.start()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Request dispatch
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def _handle_request_async(req: dict) -> None:
|
|
rid = str(req.get("id"))
|
|
token = _CURRENT_RID.set(rid)
|
|
_STATE.user_ns["__omp_run_id__"] = rid
|
|
_STATE.cancel_requested = False
|
|
_STATE.execution_count += 1
|
|
execution_count = _STATE.execution_count
|
|
_emit({"type": "started", "id": rid})
|
|
|
|
status: str = "ok"
|
|
cancelled = False
|
|
|
|
try:
|
|
try:
|
|
_apply_request_runtime(req)
|
|
transformed = transform_cell(req.get("code", ""))
|
|
except SyntaxError as exc:
|
|
_emit_error(rid, exc)
|
|
_emit({
|
|
"type": "done",
|
|
"id": rid,
|
|
"status": "error",
|
|
"executionCount": execution_count,
|
|
"cancelled": False,
|
|
})
|
|
return
|
|
except BaseException as exc: # noqa: BLE001 - runtime setup errors must settle the request
|
|
_emit_error(rid, exc)
|
|
_emit({
|
|
"type": "done",
|
|
"id": rid,
|
|
"status": "error",
|
|
"executionCount": execution_count,
|
|
"cancelled": False,
|
|
})
|
|
return
|
|
|
|
_begin_exec_sigint()
|
|
try:
|
|
await _exec_source_async(transformed, _STATE.user_ns)
|
|
except KeyboardInterrupt:
|
|
cancelled = True
|
|
status = "error"
|
|
_emit_error(rid, KeyboardInterrupt("Execution interrupted"))
|
|
except SystemExit as exc:
|
|
status = "error"
|
|
_emit_error(rid, exc)
|
|
except BaseException as exc: # noqa: BLE001 - we want to surface every user error
|
|
status = "error"
|
|
_emit_error(rid, exc)
|
|
finally:
|
|
_end_exec_sigint()
|
|
try:
|
|
_flush_matplotlib_figures()
|
|
except Exception:
|
|
pass
|
|
|
|
_emit({
|
|
"type": "done",
|
|
"id": rid,
|
|
"status": status,
|
|
"executionCount": execution_count,
|
|
"cancelled": cancelled,
|
|
})
|
|
finally:
|
|
_CURRENT_RID.reset(token)
|
|
|
|
|
|
def _emit_error(rid: str, exc: BaseException) -> None:
|
|
tb_lines = traceback.format_exception(type(exc), exc, exc.__traceback__)
|
|
_emit({
|
|
"type": "error",
|
|
"id": rid,
|
|
"ename": type(exc).__name__,
|
|
"evalue": str(exc),
|
|
"traceback": [line.rstrip("\n") for line in tb_lines],
|
|
})
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main loop
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _read_stdin(loop: asyncio.AbstractEventLoop, queue: asyncio.Queue, stdin) -> None:
|
|
for raw_line in stdin:
|
|
line = raw_line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
req = json.loads(line)
|
|
except json.JSONDecodeError as exc:
|
|
_emit({
|
|
"type": "error",
|
|
"id": "",
|
|
"ename": "ProtocolError",
|
|
"evalue": f"Invalid JSON request: {exc}",
|
|
"traceback": [],
|
|
})
|
|
continue
|
|
loop.call_soon_threadsafe(queue.put_nowait, req)
|
|
loop.call_soon_threadsafe(queue.put_nowait, {"type": "exit"})
|
|
|
|
|
|
async def _main_async() -> None:
|
|
sys.stdout = _StreamProxy("stdout")
|
|
sys.stderr = _StreamProxy("stderr")
|
|
_install_idle_sigint()
|
|
_start_parent_watchdog()
|
|
|
|
stdin = sys.__stdin__
|
|
if stdin is None:
|
|
return
|
|
|
|
loop = asyncio.get_running_loop()
|
|
_STATE.loop = loop
|
|
queue: asyncio.Queue = asyncio.Queue()
|
|
reader = threading.Thread(target=_read_stdin, args=(loop, queue, stdin), name="omp-stdin-reader", daemon=True)
|
|
reader.start()
|
|
|
|
tasks: set[asyncio.Task] = set()
|
|
def _task_done(task: asyncio.Task) -> None:
|
|
tasks.discard(task)
|
|
try:
|
|
exc = task.exception()
|
|
except asyncio.CancelledError:
|
|
return
|
|
if exc is not None:
|
|
_emit_error("", exc)
|
|
try:
|
|
while True:
|
|
req = await queue.get()
|
|
if req.get("type") == "exit":
|
|
break
|
|
task = asyncio.create_task(_handle_request_async(req))
|
|
tasks.add(task)
|
|
task.add_done_callback(_task_done)
|
|
finally:
|
|
for task in tasks:
|
|
task.cancel()
|
|
if tasks:
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
|
|
def main() -> None:
|
|
asyncio.run(_main_async())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|