chore: reformat

This commit is contained in:
can1357
2026-07-15 00:08:50 +02:00
parent 7d02778c60
commit 55f5ebec49
97 changed files with 12237 additions and 3179 deletions
+89 -41
View File
@@ -1,10 +1,12 @@
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, math, re
from urllib.parse import unquote
INTENT_FIELD = "i"
# __omp_display is injected by runner.py before the prelude executes; it
@@ -40,7 +42,6 @@ if "__omp_prelude_loaded__" not in globals():
"""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:
@@ -150,18 +151,18 @@ if "__omp_prelude_loaded__" not in globals():
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
@@ -180,33 +181,35 @@ if "__omp_prelude_loaded__" not in globals():
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)
_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:
@@ -214,39 +217,60 @@ if "__omp_prelude_loaded__" not in globals():
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"
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
_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}
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),
"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:
@@ -256,12 +280,10 @@ if "__omp_prelude_loaded__" not in globals():
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")]
)
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])}"
@@ -269,7 +291,7 @@ if "__omp_prelude_loaded__" not in globals():
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":
@@ -277,13 +299,13 @@ if "__omp_prelude_loaded__" not in globals():
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"]})
@@ -295,13 +317,13 @@ if "__omp_prelude_loaded__" not in globals():
"""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 = ""
@@ -320,7 +342,7 @@ if "__omp_prelude_loaded__" not in globals():
j = i + 1
while j < len(query) and query[j] != "]":
j += 1
bracket_content = query[i+1:j]
bracket_content = query[i + 1 : j]
if bracket_content.startswith('"') and bracket_content.endswith('"'):
tokens.append(("key", bracket_content[1:-1]))
else:
@@ -331,7 +353,7 @@ if "__omp_prelude_loaded__" not in globals():
i += 1
if current_token:
tokens.append(("key", current_token))
# Apply tokens
current = data
for token_type, value in tokens:
@@ -343,9 +365,8 @@ if "__omp_prelude_loaded__" not in globals():
if not isinstance(current, dict) or value not in current:
return None
current = current[value]
return current
return current
def _tool_proxy_from_env() -> tuple[str, str, str]:
base = os.environ.get("PI_TOOL_BRIDGE_URL")
@@ -358,9 +379,14 @@ if "__omp_prelude_loaded__" not in globals():
def _bridge_call(name: str, args: dict):
"""POST one request to the host tool bridge and return its `value`."""
import urllib.request, urllib.error
base, token, session = _tool_proxy_from_env()
_run_id_getter = globals().get("__omp_current_run_id__")
_run_id = _run_id_getter() if callable(_run_id_getter) else globals().get("__omp_run_id__")
_run_id = (
_run_id_getter()
if callable(_run_id_getter)
else globals().get("__omp_run_id__")
)
payload = json.dumps(
{"session": session, "run": _run_id, "name": name, "args": args}
).encode("utf-8")
@@ -429,7 +455,11 @@ if "__omp_prelude_loaded__" not in globals():
def __repr__(self) -> str:
session = os.environ.get("PI_TOOL_BRIDGE_SESSION")
return f"<tool proxy session={session}>" if session else "<tool proxy unavailable>"
return (
f"<tool proxy session={session}>"
if session
else "<tool proxy unavailable>"
)
tool = _ToolProxy()
@@ -450,7 +480,18 @@ if "__omp_prelude_loaded__" not in globals():
text = res.get("text") if isinstance(res, dict) else res
return json.loads(text) if schema is not None else text
def agent(prompt, *, agent="task", model=None, label=None, schema=None, isolated=None, apply=None, merge=None, handle=False):
def agent(
prompt,
*,
agent="task",
model=None,
label=None,
schema=None,
isolated=None,
apply=None,
merge=None,
handle=False,
):
"""Run a subagent and return its final output.
`agent` selects the subagent definition (default "task"). Pass
@@ -513,7 +554,13 @@ if "__omp_prelude_loaded__" not in globals():
return parsed
details = res.get("details") if isinstance(res, dict) else None
if not isinstance(details, dict) or details.get("id") is None:
return {"text": text, "output": text, "handle": None, "id": None, "agent": None}
return {
"text": text,
"output": text,
"handle": None,
"id": None,
"agent": None,
}
node = {
"text": text,
"output": text,
@@ -559,6 +606,7 @@ if "__omp_prelude_loaded__" not in globals():
pool width tracks ``task.maxConcurrency`` (0 = run every item at once).
"""
import concurrent.futures, contextvars
items = list(items)
if not items:
return []
+109 -57
View File
@@ -191,10 +191,14 @@ class _RunnerState:
self.capture_rid: str | None = None
_CURRENT_RID: contextvars.ContextVar[str | None] = contextvars.ContextVar("omp_current_rid", default=None)
_CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS: contextvars.ContextVar[set[int] | None] = contextvars.ContextVar(
"omp_displayed_matplotlib_figure_ids",
default=None,
_CURRENT_RID: contextvars.ContextVar[str | None] = contextvars.ContextVar(
"omp_current_rid", default=None
)
_CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS: contextvars.ContextVar[set[int] | None] = (
contextvars.ContextVar(
"omp_displayed_matplotlib_figure_ids",
default=None,
)
)
@@ -233,7 +237,9 @@ def _drain_captured_stdout() -> None:
def _start_capture_drain() -> None:
if _CAPTURE_READ_FD is None:
return
thread = threading.Thread(target=_drain_captured_stdout, name="omp-fd1-capture", daemon=True)
thread = threading.Thread(
target=_drain_captured_stdout, name="omp-fd1-capture", daemon=True
)
thread.start()
@@ -242,7 +248,9 @@ def _start_capture_drain() -> None:
# ---------------------------------------------------------------------------
_MAGIC_LINE_RE = re.compile(r"^(?P<indent>[ \t]*)(?P<name>[A-Za-z_][A-Za-z_0-9]*)(?:[ \t]+(?P<args>.*))?$")
_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>.+)$"
)
@@ -337,7 +345,9 @@ def transform_cell(source: str) -> str:
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)})")
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("%%"):
@@ -383,7 +393,9 @@ def line_magic(name: str) -> Callable[[Callable[[str], Any]], Callable[[str], An
return decorator
def cell_magic(name: str) -> Callable[[Callable[[str, str], Any]], Callable[[str, str], Any]]:
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
@@ -398,6 +410,7 @@ def _emit_status(op: str, **data: Any) -> None:
return
_emit({"type": "display", "id": rid, "bundle": bundle})
_SHELL_READ_CHUNK_BYTES = 8192
_SHELL_OUTPUT_MAX_BYTES = 1024 * 1024
_SHELL_OUTPUT_MAX_LINES = 3000
@@ -458,12 +471,16 @@ class _ShellOutputLimiter:
return
limited = _take_prefix_by_lines(text, self._remaining_lines)
truncated = limited != text
byte_limited = _take_prefix_by_encoded_bytes(limited, self._remaining_bytes, self._encoding)
byte_limited = _take_prefix_by_encoded_bytes(
limited, self._remaining_bytes, self._encoding
)
truncated = truncated or byte_limited != limited
if byte_limited:
sys.stdout.write(byte_limited)
sys.stdout.flush()
self._remaining_bytes -= len(byte_limited.encode(self._encoding, errors="strict"))
self._remaining_bytes -= len(
byte_limited.encode(self._encoding, errors="strict")
)
self._remaining_lines -= byte_limited.count("\n")
self._at_line_start = byte_limited.endswith("\n")
if truncated:
@@ -478,7 +495,9 @@ class _ShellOutputLimiter:
self._truncated = True
def _stream_process_output(proc: subprocess.Popen, on_text: Callable[[str], None] | None = None) -> None:
def _stream_process_output(
proc: subprocess.Popen, on_text: Callable[[str], None] | None = None
) -> None:
assert proc.stdout is not None
encoding = _process_output_encoding()
decoder = _process_output_decoder(encoding)
@@ -514,7 +533,9 @@ class _BoundedTextCapture:
if self._remaining_bytes <= 0 or self._remaining_lines <= 0:
return
line_limited = _take_prefix_by_lines(text, self._remaining_lines)
part = _take_prefix_by_encoded_bytes(line_limited, self._remaining_bytes, self._encoding)
part = _take_prefix_by_encoded_bytes(
line_limited, self._remaining_bytes, self._encoding
)
if not part:
return
self._parts.append(part)
@@ -583,7 +604,9 @@ def _magic_pip(args: str) -> None:
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)
_emit_status(
"pip", args=args, installed=installed_packages, exit_code=proc.returncode
)
@line_magic("cd")
@@ -658,7 +681,9 @@ 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__")
if not name.startswith("_")
and not callable(value)
or hasattr(value, "__class__")
)
return [n for n in names if not n.startswith("__")]
@@ -677,7 +702,9 @@ def _magic_whos(_args: str) -> list[tuple[str, str]]:
@line_magic("reset")
def _magic_reset(_args: str) -> None:
_STATE.user_ns.clear()
_STATE.user_ns.update({"__name__": "__main__", "__doc__": None, "__builtins__": builtins})
_STATE.user_ns.update(
{"__name__": "__main__", "__doc__": None, "__builtins__": builtins}
)
_install_builtins(_STATE.user_ns)
_emit_status("reset")
@@ -686,7 +713,9 @@ def _magic_reset(_args: str) -> None:
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}})
_emit(
{"type": "display", "id": _CURRENT_RID.get(), "bundle": {"text/plain": source}}
)
_exec_source(source, _STATE.user_ns)
@@ -712,6 +741,7 @@ def _magic_run(args: str) -> None:
def _magic_cell_bash(args: str, body: str) -> int:
return _run_shell_body(body, shell_arg="/bin/bash")
@cell_magic("capture")
def _magic_cell_capture(args: str, body: str) -> str:
"""Capture stdout/stderr of body; bind to ``args`` (a name) if provided."""
@@ -797,7 +827,9 @@ def __omp_shell(cmd: str) -> _ShellResult:
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
)
capture = _BoundedTextCapture(_SHELL_RESULT_CAPTURE_BYTES, _SHELL_OUTPUT_MAX_LINES, _process_output_encoding())
capture = _BoundedTextCapture(
_SHELL_RESULT_CAPTURE_BYTES, _SHELL_OUTPUT_MAX_LINES, _process_output_encoding()
)
_stream_process_output(proc, capture.add)
proc.wait()
lines = [line for line in capture.text().splitlines()]
@@ -827,7 +859,9 @@ def _is_matplotlib_figure(value: Any) -> bool:
return True
value_type = type(value)
return value_type.__module__ == "matplotlib.figure" and value_type.__name__ == "Figure"
return (
value_type.__module__ == "matplotlib.figure" and value_type.__name__ == "Figure"
)
def _matplotlib_figure_png(value: Any) -> str | None:
@@ -869,7 +903,6 @@ def _mime_bundle(value: Any) -> dict:
if matplotlib_png is not None:
bundle["image/png"] = matplotlib_png
mimebundle = getattr(value, "_repr_mimebundle_", None)
if callable(mimebundle):
try:
@@ -990,7 +1023,9 @@ def _await_sync(coro) -> Any:
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")
raise RuntimeError(
"top-level await is not supported from synchronous magic execution"
)
return asyncio.run(coro)
@@ -1005,7 +1040,6 @@ def _run_compiled_sync(code, ns: dict, *, want_value: bool) -> Any:
return None
async def _run_compiled_async(code, ns: dict, *, want_value: bool) -> Any:
"""Execute a code object in the persistent event loop.
@@ -1127,6 +1161,7 @@ def _apply_request_runtime(req: dict) -> None:
elif value is None:
os.environ.pop(key, None)
def _start_parent_watchdog() -> None:
"""Self-terminate when the host process dies.
@@ -1182,23 +1217,27 @@ async def _handle_request_async(req: dict) -> None:
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,
})
_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,
})
_emit(
{
"type": "done",
"id": rid,
"status": "error",
"executionCount": execution_count,
"cancelled": False,
}
)
return
_begin_exec_sigint()
@@ -1222,13 +1261,15 @@ async def _handle_request_async(req: dict) -> None:
pass
_flush_stream_proxies(rid)
_emit({
"type": "done",
"id": rid,
"status": status,
"executionCount": execution_count,
"cancelled": cancelled,
})
_emit(
{
"type": "done",
"id": rid,
"status": status,
"executionCount": execution_count,
"cancelled": cancelled,
}
)
finally:
if _STATE.capture_rid == rid:
_STATE.capture_rid = None
@@ -1239,13 +1280,15 @@ async def _handle_request_async(req: dict) -> None:
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],
})
_emit(
{
"type": "error",
"id": rid,
"ename": type(exc).__name__,
"evalue": str(exc),
"traceback": [line.rstrip("\n") for line in tb_lines],
}
)
# ---------------------------------------------------------------------------
@@ -1261,13 +1304,15 @@ def _read_stdin(loop: asyncio.AbstractEventLoop, queue: asyncio.Queue, stdin) ->
try:
req = json.loads(line)
except json.JSONDecodeError as exc:
_emit({
"type": "error",
"id": "",
"ename": "ProtocolError",
"evalue": f"Invalid JSON request: {exc}",
"traceback": [],
})
_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"})
@@ -1287,10 +1332,16 @@ async def _main_async() -> None:
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 = 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:
@@ -1299,6 +1350,7 @@ async def _main_async() -> None:
return
if exc is not None:
_emit_error("", exc)
try:
while True:
req = await queue.get()