"""OMP coding agent interface — an alternative to Pi for SSSF. Runs `omp -p --mode json` and tails its JSONL stdout line by line, forwarding each event to a callback WHILE the agent works. OMP's event stream is the same shape as Pi's (session / message_start / message_end / turn_end / tool_execution_start / tool_execution_end / agent_end), so the shared ToolCallTracker folds tool calls identically. The one real difference is session identity. Pi accepts `--session-id` and creates-or-continues that exact id. OMP mints its own session id and resumes by `--resume `. To keep SSSF's retry/correction loop (which re-enters the SAME context window) working, this adapter persists the real OMP session id in a `session_id.txt` next to the session dir: the first send creates a session and records its id; every later send with the same session_dir resumes it. """ from __future__ import annotations import json import os import subprocess from functools import lru_cache from pathlib import Path from typing import Callable, Optional from .agent_pi import (ToolCallTracker, _context_tokens, _label, _text_of) from .data_types import PiRequest, PiResult from .utils import now_iso, operator_env OMP_PATH = os.environ.get("OMP_PATH", "omp") RESULT_SNIPPET_CHARS = 20_000 # tool output rides along whole; clip only guards pathological cases ARG_VALUE_CHARS = 20_000 # args too — the UI scrolls, it must not be handed cut-off data LABEL_CHARS = 80 # "bash: " shown as the event name # The arg that identifies a call at a glance, in the order tools tend to use. PRIMARY_ARGS = ("command", "path", "file_path", "pattern", "query", "url") # The config's `tools` list is written in pi's vocabulary. OMP's tool set # overlaps but is not identical, so translate the names that differ and drop # the ones OMP has no equivalent for (the `writes` boundary, enforced in code, # is what actually keeps an agent read-only — the tool list is a hint). TOOL_TRANSLATION = { "find": "glob", # pi's find -> omp's glob "ls": None, # omp reads directories via `read` "subagent_create": "task", # omp's subagent tool "subagent_continue": "task", "subagent_list": "task", "subagent_remove": "task", } def _translate_tools(tools: list[str]) -> list[str]: """Map pi tool names to omp's vocabulary, dropping unmappable ones.""" out = [] for tool in tools: mapped = TOOL_TRANSLATION.get(tool, tool) if mapped and mapped not in out: out.append(mapped) return out @lru_cache(maxsize=1) def _omp_catalog() -> list[dict]: """Read omp's merged model catalog (`omp models --json`).""" try: result = subprocess.run( [OMP_PATH, "models", "--json"], capture_output=True, text=True, timeout=30, env=operator_env(), check=False, ) except (OSError, subprocess.TimeoutExpired): return [] if result.returncode != 0: return [] try: return json.loads(result.stdout).get("models", []) except json.JSONDecodeError: return [] def resolve_model(pattern: str) -> tuple[str, str]: """Resolve a model pattern to an explicit ``(provider, model_id)`` pair. OMP's catalog exposes a `selector` of the form ``provider/model-id``. A pattern with a slash must match a selector exactly; a bare pattern matches by model id (substring, then exact), failing on ambiguity. """ catalog = _omp_catalog() selectors = {m.get("selector"): (m.get("provider", ""), m.get("id", "")) for m in catalog if m.get("selector")} if "/" in pattern: if pattern in selectors: return selectors[pattern] raise ValueError(f"model pattern {pattern!r} not found in omp models — " "run `omp models` to see available selectors") matches = [(m.get("provider", ""), m.get("id", "")) for m in catalog if pattern == m.get("id") or pattern in m.get("id", "")] exact = [match for match in matches if match[1] == pattern] if len(exact) == 1: return exact[0] if len(matches) == 1: return matches[0] if not matches: raise ValueError(f"model pattern {pattern!r} not found in omp models — " "run `omp models` to see available models") raise ValueError(f"model pattern {pattern!r} is ambiguous: {matches}") def context_window(provider: str, model_id: str) -> int: """The model's context ceiling from omp's merged model catalog.""" for model in _omp_catalog(): if model.get("provider") == provider and model.get("id") == model_id: return int(model.get("contextWindow") or 0) return 0 def _session_id_file(session_dir: Path) -> Path: return session_dir / "session_id.txt" def run(request: PiRequest, on_event: Optional[Callable[[dict], None]] = None, on_spawn: Optional[Callable[[int], None]] = None, on_exit: Optional[Callable[[int], None]] = None) -> PiResult: """Run one non-interactive omp turn. `on_spawn(pid)` and `on_exit(pid)` bracket the child process so the caller can record it as killable — a hung coding agent is otherwise a pid you have to hunt for in `ps` while the run sits there. """ provider, model_id = resolve_model(request.model) selector = f"{provider}/{model_id}" session_dir = Path(request.session_dir) session_dir.mkdir(parents=True, exist_ok=True) sid_file = _session_id_file(session_dir) cmd = [ OMP_PATH, "-p", "--mode", "json", "--model", selector, "--thinking", request.thinking, "--system-prompt", request.system_prompt, "--session-dir", str(session_dir), "--auto-approve", # agents run tools autonomously; no approval prompts ] # Create-or-continue: the first send mints a session and records its id; # every later send with the same session_dir resumes that context window. if sid_file.exists(): cmd += ["--resume", sid_file.read_text().strip()] if request.tools: cmd += ["--tools", ",".join(_translate_tools(request.tools))] for extension in request.extensions: cmd += ["-e", extension] cmd.append(request.prompt) raw_path = Path(request.raw_output_path) raw_path.parent.mkdir(parents=True, exist_ok=True) result = PiResult(session_id="", context_window=context_window(provider, model_id)) # stdin is DEVNULL, deliberately. The prompt travels in argv, so the child # never needs stdin — but inheriting the parent's means omp sees a non-TTY # and can sit forever waiting for piped input that will never arrive or EOF. process = subprocess.Popen(cmd, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1, cwd=request.cwd, env=operator_env()) if on_spawn: on_spawn(process.pid) with raw_path.open("a") as raw: assert process.stdout is not None for line in process.stdout: raw.write(line) raw.flush() # events land on disk as they happen line = line.strip() if not line: continue try: event = json.loads(line) except json.JSONDecodeError: continue if event.get("type") == "session": sid = event.get("id") if sid: result.session_id = sid sid_file.write_text(sid) # persist for later resumes if event.get("type") == "message_end": message = event.get("message", {}) if message.get("role") == "assistant": text = _text_of(message) if text: result.text = text # last assistant message wins usage = message.get("usage", {}) or {} turn = _context_tokens(usage) result.tokens += turn result.usage.add_turn(usage, turn) # Occupancy is read off the last VALID assistant turn, the # way omp does it — an aborted or errored turn reports usage # you can't trust, so it must not overwrite a good reading. if turn and message.get("stopReason") not in ("aborted", "error"): result.context_tokens = turn result.cost += (usage.get("cost", {}) or {}).get("total", 0.0) or 0.0 if on_event: on_event(event) stderr = process.stderr.read() if process.stderr else "" result.returncode = process.wait() if on_exit: on_exit(process.pid) if result.returncode != 0 and not result.text: raise RuntimeError(f"omp exited {result.returncode}: {stderr.strip()[-800:]}") return result