skills/sssf/templates/adws/adw_modules/agent_omp.py

"""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 <id-prefix>`. 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: <command>" 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