can1357
8d144e17ec
- Replaced Python execution with a local `python -u runner.py` subprocess and NDJSON stdin/stdout framing. - Removed shared-gateway architecture, including coordinator lifecycle APIs, `useSharedGateway` wiring, and `jupyter` CLI/actions. - Simplified setup checks to a plain Python 3 availability probe and removed automatic dependency-install fallbacks. - Updated kernel cancellation and display processing to use status frames, SIGINT/SIGTERM escalation, and normalized output coercion. - Added `python-runner` integration and display tests while deleting legacy websocket and kernel lifecycle test suites.
463 lines
19 KiB
Python
463 lines
19 KiB
Python
from __future__ import annotations
|
|
# OMP prelude helpers (loaded once into the runner namespace)
|
|
if "__omp_prelude_loaded__" not in globals():
|
|
__omp_prelude_loaded__ = True
|
|
from pathlib import Path
|
|
import os, json
|
|
|
|
# __omp_display is injected by runner.py before the prelude executes; it
|
|
# mirrors IPython's display() semantics with the same MIME bundle output.
|
|
_omp_display = __omp_display # type: ignore[name-defined]
|
|
|
|
_PRESENTABLE_REPRS = (
|
|
"_repr_mimebundle_",
|
|
"_repr_html_",
|
|
"_repr_json_",
|
|
"_repr_markdown_",
|
|
"_repr_png_",
|
|
"_repr_jpeg_",
|
|
"_repr_svg_",
|
|
"_repr_latex_",
|
|
)
|
|
|
|
def display(value):
|
|
"""Render a value. Falls back to a JSON+text/plain bundle for plain dict/list/tuple."""
|
|
if any(hasattr(value, attr) for attr in _PRESENTABLE_REPRS):
|
|
_omp_display(value)
|
|
return
|
|
if isinstance(value, (dict, list, tuple)):
|
|
try:
|
|
bundle = {"application/json": value, "text/plain": repr(value)}
|
|
_omp_display(bundle, raw=True)
|
|
return
|
|
except Exception:
|
|
pass
|
|
_omp_display(value)
|
|
|
|
def _emit_status(op: str, **data):
|
|
"""Emit structured status event for TUI rendering."""
|
|
_omp_display({"application/x-omp-status": {"op": op, **data}}, raw=True)
|
|
|
|
|
|
def env(key: str | None = None, value: str | None = None):
|
|
"""Get/set environment variables."""
|
|
if key is None:
|
|
items = dict(sorted(os.environ.items()))
|
|
_emit_status("env", count=len(items), keys=list(items.keys())[:20])
|
|
return items
|
|
if value is not None:
|
|
os.environ[key] = value
|
|
_emit_status("env", key=key, value=value, action="set")
|
|
return value
|
|
val = os.environ.get(key)
|
|
_emit_status("env", key=key, value=val, action="get")
|
|
return val
|
|
|
|
def read(path: str | Path, *, offset: int = 1, limit: int | None = None) -> str:
|
|
"""Read file contents. offset/limit are 1-indexed line numbers."""
|
|
p = Path(path)
|
|
data = p.read_text(encoding="utf-8")
|
|
lines = data.splitlines(keepends=True)
|
|
if offset > 1 or limit is not None:
|
|
start = max(0, offset - 1)
|
|
end = start + limit if limit else len(lines)
|
|
lines = lines[start:end]
|
|
data = "".join(lines)
|
|
preview = data[:500]
|
|
_emit_status("read", path=str(p), chars=len(data), preview=preview)
|
|
return data
|
|
|
|
def write(path: str | Path, content: str) -> Path:
|
|
"""Write file contents (create parents)."""
|
|
p = Path(path)
|
|
p.parent.mkdir(parents=True, exist_ok=True)
|
|
p.write_text(content, encoding="utf-8")
|
|
_emit_status("write", path=str(p), chars=len(content))
|
|
return p
|
|
|
|
def append(path: str | Path, content: str) -> Path:
|
|
"""Append to file."""
|
|
p = Path(path)
|
|
p.parent.mkdir(parents=True, exist_ok=True)
|
|
with p.open("a", encoding="utf-8") as f:
|
|
f.write(content)
|
|
_emit_status("append", path=str(p), chars=len(content))
|
|
return p
|
|
|
|
def sort(text: str, *, reverse: bool = False, unique: bool = False) -> str:
|
|
"""Sort lines of text."""
|
|
lines = text.splitlines()
|
|
if unique:
|
|
lines = list(dict.fromkeys(lines))
|
|
lines = sorted(lines, reverse=reverse)
|
|
out = "\n".join(lines)
|
|
_emit_status("sort", lines=len(lines), unique=unique, reverse=reverse)
|
|
return out
|
|
|
|
def uniq(text: str, *, count: bool = False) -> str | list[tuple[int, str]]:
|
|
"""Remove duplicate adjacent lines (like uniq)."""
|
|
lines = text.splitlines()
|
|
if not lines:
|
|
_emit_status("uniq", groups=0)
|
|
return [] if count else ""
|
|
groups: list[tuple[int, str]] = []
|
|
current = lines[0]
|
|
current_count = 1
|
|
for line in lines[1:]:
|
|
if line == current:
|
|
current_count += 1
|
|
continue
|
|
groups.append((current_count, current))
|
|
current = line
|
|
current_count = 1
|
|
groups.append((current_count, current))
|
|
_emit_status("uniq", groups=len(groups), count_mode=count)
|
|
if count:
|
|
return groups
|
|
return "\n".join(line for _, line in groups)
|
|
|
|
def counter(
|
|
items: str | list,
|
|
*,
|
|
limit: int | None = None,
|
|
reverse: bool = True,
|
|
) -> list[tuple[int, str]]:
|
|
"""Count occurrences and sort by frequency. Like sort | uniq -c | sort -rn.
|
|
|
|
items: text (splits into lines) or list of strings
|
|
reverse: True for descending (most common first), False for ascending
|
|
Returns: [(count, item), ...] sorted by count
|
|
"""
|
|
from collections import Counter
|
|
if isinstance(items, str):
|
|
items = items.splitlines()
|
|
counts = Counter(items)
|
|
sorted_items = sorted(counts.items(), key=lambda x: (x[1], x[0]), reverse=reverse)
|
|
if limit is not None:
|
|
sorted_items = sorted_items[:limit]
|
|
result = [(count, item) for item, count in sorted_items]
|
|
_emit_status("counter", unique=len(counts), total=sum(counts.values()), top=result[:10])
|
|
return result
|
|
def tree(path: str | Path = ".", *, max_depth: int = 3, show_hidden: bool = False) -> str:
|
|
"""Return directory tree."""
|
|
base = Path(path)
|
|
lines = []
|
|
def walk(p: Path, prefix: str, depth: int):
|
|
if depth > max_depth:
|
|
return
|
|
items = sorted(p.iterdir(), key=lambda x: (not x.is_dir(), x.name.lower()))
|
|
items = [i for i in items if show_hidden or not i.name.startswith(".")]
|
|
for i, item in enumerate(items):
|
|
is_last = i == len(items) - 1
|
|
connector = "└── " if is_last else "├── "
|
|
suffix = "/" if item.is_dir() else ""
|
|
lines.append(f"{prefix}{connector}{item.name}{suffix}")
|
|
if item.is_dir():
|
|
ext = " " if is_last else "│ "
|
|
walk(item, prefix + ext, depth + 1)
|
|
lines.append(str(base) + "/")
|
|
walk(base, "", 1)
|
|
out = "\n".join(lines)
|
|
_emit_status("tree", path=str(base), entries=len(lines) - 1, preview=out[:1000])
|
|
return out
|
|
|
|
def diff(a: str | Path, b: str | Path) -> str:
|
|
"""Compare two files, return unified diff."""
|
|
import difflib
|
|
path_a, path_b = Path(a), Path(b)
|
|
lines_a = path_a.read_text(encoding="utf-8").splitlines(keepends=True)
|
|
lines_b = path_b.read_text(encoding="utf-8").splitlines(keepends=True)
|
|
result = difflib.unified_diff(lines_a, lines_b, fromfile=str(path_a), tofile=str(path_b))
|
|
out = "".join(result)
|
|
_emit_status("diff", file_a=str(path_a), file_b=str(path_b), identical=not out, preview=out[:500])
|
|
return out
|
|
|
|
def output(
|
|
*ids: str,
|
|
format: str = "raw",
|
|
query: str | None = None,
|
|
offset: int | None = None,
|
|
limit: int | None = None,
|
|
) -> str | dict | list[dict]:
|
|
"""Read task/agent output by ID. Returns text or JSON depending on format.
|
|
|
|
Args:
|
|
*ids: Output IDs to read (e.g., 'explore_0', 'reviewer_1')
|
|
format: 'raw' (default), 'json' (dict with metadata), 'stripped' (no ANSI)
|
|
query: jq-like query for JSON outputs (e.g., '.endpoints[0].file')
|
|
offset: Line number to start reading from (1-indexed)
|
|
limit: Maximum number of lines to read
|
|
|
|
Returns:
|
|
Single ID: str (format='raw'/'stripped') or dict (format='json')
|
|
Multiple IDs: list of dict with 'id' and 'content'/'data' keys
|
|
|
|
Examples:
|
|
output('explore_0') # Read as raw text
|
|
output('reviewer_0', format='json') # Read with metadata
|
|
output('explore_0', query='.files[0]') # Extract JSON field
|
|
output('explore_0', offset=10, limit=20) # Lines 10-29
|
|
output('explore_0', 'reviewer_1') # Read multiple outputs
|
|
"""
|
|
# Prefer PI_ARTIFACTS_DIR so subagents resolve through the parent's
|
|
# shared artifacts dir; fall back to deriving from PI_SESSION_FILE
|
|
# for legacy callers / top-level sessions where the two coincide.
|
|
artifacts_dir = os.environ.get("PI_ARTIFACTS_DIR")
|
|
if not artifacts_dir:
|
|
session_file = os.environ.get("PI_SESSION_FILE")
|
|
if not session_file:
|
|
_emit_status("output", error="No session file available")
|
|
raise RuntimeError("No session - output artifacts unavailable")
|
|
artifacts_dir = session_file.rsplit(".", 1)[0] # Strip .jsonl extension
|
|
if not Path(artifacts_dir).exists():
|
|
_emit_status("output", error="Artifacts directory not found", path=artifacts_dir)
|
|
raise RuntimeError(f"No artifacts directory found: {artifacts_dir}")
|
|
|
|
if not ids:
|
|
_emit_status("output", error="No IDs provided")
|
|
raise ValueError("At least one output ID is required")
|
|
|
|
if query and (offset is not None or limit is not None):
|
|
_emit_status("output", error="query cannot be combined with offset/limit")
|
|
raise ValueError("query cannot be combined with offset/limit")
|
|
|
|
results: list[dict] = []
|
|
not_found: list[str] = []
|
|
|
|
for output_id in ids:
|
|
output_path = Path(artifacts_dir) / f"{output_id}.md"
|
|
if not output_path.exists():
|
|
not_found.append(output_id)
|
|
continue
|
|
|
|
raw_content = output_path.read_text(encoding="utf-8")
|
|
raw_lines = raw_content.splitlines()
|
|
total_lines = len(raw_lines)
|
|
|
|
selected_content = raw_content
|
|
range_info: dict | None = None
|
|
|
|
# Handle query
|
|
if query:
|
|
try:
|
|
json_value = json.loads(raw_content)
|
|
except json.JSONDecodeError as e:
|
|
_emit_status("output", id=output_id, error=f"Not valid JSON: {e}")
|
|
raise ValueError(f"Output {output_id} is not valid JSON: {e}")
|
|
|
|
# Apply jq-like query
|
|
result_value = _apply_query(json_value, query)
|
|
try:
|
|
selected_content = json.dumps(result_value, indent=2) if result_value is not None else "null"
|
|
except (TypeError, ValueError):
|
|
selected_content = str(result_value)
|
|
|
|
# Handle offset/limit
|
|
elif offset is not None or limit is not None:
|
|
start_line = max(1, offset or 1)
|
|
if start_line > total_lines:
|
|
_emit_status("output", id=output_id, error=f"Offset {start_line} beyond end ({total_lines} lines)")
|
|
raise ValueError(f"Offset {start_line} is beyond end of output ({total_lines} lines) for {output_id}")
|
|
|
|
effective_limit = limit if limit is not None else total_lines - start_line + 1
|
|
end_line = min(total_lines, start_line + effective_limit - 1)
|
|
selected_lines = raw_lines[start_line - 1 : end_line]
|
|
selected_content = "\n".join(selected_lines)
|
|
range_info = {"start_line": start_line, "end_line": end_line, "total_lines": total_lines}
|
|
|
|
# Strip ANSI codes if requested
|
|
if format == "stripped":
|
|
import re
|
|
selected_content = re.sub(r"\x1b\[[0-9;]*m", "", selected_content)
|
|
|
|
# Build result
|
|
if format == "json":
|
|
result_data = {
|
|
"id": output_id,
|
|
"path": str(output_path),
|
|
"line_count": total_lines if not query else len(selected_content.splitlines()),
|
|
"char_count": len(raw_content) if not query else len(selected_content),
|
|
"content": selected_content,
|
|
}
|
|
if range_info:
|
|
result_data["range"] = range_info
|
|
if query:
|
|
result_data["query"] = query
|
|
results.append(result_data)
|
|
else:
|
|
results.append({"id": output_id, "content": selected_content})
|
|
|
|
# Handle not found
|
|
if not_found:
|
|
available = sorted(
|
|
[f.stem for f in Path(artifacts_dir).glob("*.md")]
|
|
)
|
|
error_msg = f"Output not found: {', '.join(not_found)}"
|
|
if available:
|
|
error_msg += f"\n\nAvailable outputs: {', '.join(available[:20])}"
|
|
if len(available) > 20:
|
|
error_msg += f" (and {len(available) - 20} more)"
|
|
_emit_status("output", not_found=not_found, available_count=len(available))
|
|
raise FileNotFoundError(error_msg)
|
|
|
|
# Return format
|
|
if len(ids) == 1:
|
|
if format == "json":
|
|
_emit_status("output", id=ids[0], chars=results[0]["char_count"])
|
|
return results[0]
|
|
_emit_status("output", id=ids[0], chars=len(results[0]["content"]))
|
|
return results[0]["content"]
|
|
|
|
# Multiple IDs
|
|
if format == "json":
|
|
total_chars = sum(r["char_count"] for r in results)
|
|
_emit_status("output", count=len(results), total_chars=total_chars)
|
|
return results
|
|
|
|
combined_output: list[dict] = []
|
|
for r in results:
|
|
combined_output.append({"id": r["id"], "content": r["content"]})
|
|
total_chars = sum(len(r["content"]) for r in combined_output)
|
|
_emit_status("output", count=len(combined_output), total_chars=total_chars)
|
|
return combined_output
|
|
|
|
def _apply_query(data: any, query: str) -> any:
|
|
"""Apply jq-like query to data. Supports .key, [index], and chaining."""
|
|
if not query:
|
|
return data
|
|
|
|
query = query.strip()
|
|
if query.startswith("."):
|
|
query = query[1:]
|
|
if not query:
|
|
return data
|
|
|
|
# Parse query into tokens
|
|
tokens = []
|
|
current_token = ""
|
|
i = 0
|
|
while i < len(query):
|
|
ch = query[i]
|
|
if ch == ".":
|
|
if current_token:
|
|
tokens.append(("key", current_token))
|
|
current_token = ""
|
|
elif ch == "[":
|
|
if current_token:
|
|
tokens.append(("key", current_token))
|
|
current_token = ""
|
|
# Find matching ]
|
|
j = i + 1
|
|
while j < len(query) and query[j] != "]":
|
|
j += 1
|
|
bracket_content = query[i+1:j]
|
|
if bracket_content.startswith('"') and bracket_content.endswith('"'):
|
|
tokens.append(("key", bracket_content[1:-1]))
|
|
else:
|
|
tokens.append(("index", int(bracket_content)))
|
|
i = j
|
|
else:
|
|
current_token += ch
|
|
i += 1
|
|
if current_token:
|
|
tokens.append(("key", current_token))
|
|
|
|
# Apply tokens
|
|
current = data
|
|
for token_type, value in tokens:
|
|
if token_type == "index":
|
|
if not isinstance(current, list) or value >= len(current):
|
|
return None
|
|
current = current[value]
|
|
elif token_type == "key":
|
|
if not isinstance(current, dict) or value not in current:
|
|
return None
|
|
current = current[value]
|
|
|
|
return current
|
|
|
|
|
|
class _ToolCallable:
|
|
"""Invokes one host-side tool via the loopback HTTP bridge."""
|
|
|
|
__slots__ = ("_proxy", "_name")
|
|
|
|
def __init__(self, proxy: "_ToolProxy", name: str):
|
|
self._proxy = proxy
|
|
self._name = name
|
|
|
|
def __repr__(self) -> str:
|
|
return f"<tool.{self._name}>"
|
|
|
|
def __call__(self, args=None, /, **kwargs):
|
|
import urllib.request, urllib.error
|
|
if args is None:
|
|
merged: dict = {}
|
|
elif isinstance(args, dict):
|
|
merged = dict(args)
|
|
else:
|
|
raise TypeError(
|
|
f"tool.{self._name}(...) expects a dict of arguments (got {type(args).__name__})"
|
|
)
|
|
merged.update(kwargs)
|
|
if "_i" not in merged:
|
|
merged["_i"] = "py prelude"
|
|
payload = json.dumps(
|
|
{"session": self._proxy._session, "name": self._name, "args": merged}
|
|
).encode("utf-8")
|
|
req = urllib.request.Request(
|
|
f"{self._proxy._base}/v1/tool",
|
|
data=payload,
|
|
method="POST",
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {self._proxy._token}",
|
|
},
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(req) as resp:
|
|
body = resp.read()
|
|
except urllib.error.HTTPError as exc:
|
|
body = exc.read()
|
|
try:
|
|
data = json.loads(body)
|
|
except json.JSONDecodeError:
|
|
raise RuntimeError(
|
|
f"tool.{self._name}: bridge returned non-JSON response: {body[:200]!r}"
|
|
) from None
|
|
if not isinstance(data, dict) or not data.get("ok"):
|
|
msg = (data or {}).get("error") if isinstance(data, dict) else None
|
|
raise RuntimeError(msg or f"tool.{self._name} failed")
|
|
return data.get("value")
|
|
|
|
class _ToolProxy:
|
|
"""`tool.<name>(args)` proxy mirroring the JS runtime bridge."""
|
|
|
|
__slots__ = ("_base", "_token", "_session")
|
|
|
|
def __init__(self, base: str, token: str, session: str):
|
|
self._base = base.rstrip("/")
|
|
self._token = token
|
|
self._session = session
|
|
|
|
def __getattr__(self, name: str) -> _ToolCallable:
|
|
if name.startswith("_"):
|
|
raise AttributeError(name)
|
|
return _ToolCallable(self, name)
|
|
|
|
def __getitem__(self, name: str) -> _ToolCallable:
|
|
return _ToolCallable(self, name)
|
|
|
|
def __repr__(self) -> str:
|
|
return f"<tool proxy session={self._session}>"
|
|
|
|
if all(
|
|
_k in os.environ
|
|
for _k in ("PI_TOOL_BRIDGE_URL", "PI_TOOL_BRIDGE_TOKEN", "PI_TOOL_BRIDGE_SESSION")
|
|
):
|
|
tool = _ToolProxy(
|
|
os.environ["PI_TOOL_BRIDGE_URL"],
|
|
os.environ["PI_TOOL_BRIDGE_TOKEN"],
|
|
os.environ["PI_TOOL_BRIDGE_SESSION"],
|
|
)
|