api/opencode_runner.py

"""Run OpenCode headless in a throwaway workspace and collect the files it writes.

Each request gets its own temp directory::

    <tmp>/ws/        OpenCode's project directory; everything written here is the output
    <tmp>/state/     OpenCode's data/state (sessions), so --continue resumes *this* request only
    <tmp>/context.md the RAG context, attached to the message with -f

OpenCode is configured entirely through OPENCODE_CONFIG_CONTENT: one
OpenAI-compatible provider (LLM_BASE_URL / LLM_MODEL, shared platform key), file edits allowed
inside the workspace, shell/web access and edits outside the workspace denied.
"""

from __future__ import annotations

import json
import os
import shutil
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path

from .config import Settings

MAX_FILES = 60
MAX_FILE_BYTES = 200_000
IGNORED_DIRS = {".opencode", ".git", "__pycache__", "node_modules"}


class OpenCodeError(RuntimeError):
    pass


@dataclass
class OpenCodeRun:
    returncode: int
    events: list[dict]
    stderr_tail: str

    @property
    def error(self) -> str | None:
        for e in self.events:
            if e.get("type") == "error":
                err = e.get("error", {})
                return (err.get("data") or {}).get("message") or err.get("name") or "unknown OpenCode error"
        if self.returncode != 0:
            return f"opencode exited with {self.returncode}: {self.stderr_tail[-300:]}"
        return None

    def usage(self) -> dict:
        """Tokens and cost OpenCode reported for this run (its step_finish events), summed over its steps."""
        tin = tout = 0
        cost = None
        for e in self.events:
            part = e.get("part") or {}
            tokens = part.get("tokens") or e.get("tokens")
            if isinstance(tokens, dict):
                cache = tokens.get("cache") or {}
                tin += int(tokens.get("input") or 0) + int(cache.get("read") or 0)
                tout += int(tokens.get("output") or 0) + int(tokens.get("reasoning") or 0)
            c = part.get("cost", e.get("cost"))
            if isinstance(c, (int, float)):
                cost = (cost or 0.0) + float(c)
        out = {"input_tokens": tin, "output_tokens": tout}
        if cost is not None and (tin or tout):
            out["cost"] = cost
        return out


def opencode_config(settings: Settings) -> dict:
    providers = {
        "llm": {
            "npm": "@ai-sdk/openai-compatible",
            "name": "RosDiff LLM",
            "options": {"baseURL": settings.llm_base_url, "apiKey": "{env:LLM_API_KEY}"},
            # every model the router may pick, plus the configured default
            "models": {m: {"name": m} for m in _router_models(settings)},
        }
    }
    if settings.free_debug and settings.debug_model:  # the free debugging model, on OpenCode's own Zen provider
        zen, _, model = settings.debug_model.partition("/")
        options = {"apiKey": "{env:ZEN_API_KEY}"} if settings.zen_api_key else {}
        providers[zen] = {"options": options, "models": {model: {"name": model}}}
    return {
        "$schema": "https://opencode.ai/config.json",
        "model": f"llm/{settings.code_llm_model}",
        "provider": providers,
        "permission": {
            "edit": "allow",
            "bash": "deny",
            "webfetch": "deny",
            "websearch": "deny",
            "external_directory": "deny",
        },
        "autoupdate": False,
        "share": "disabled",
    }


def _router_models(settings: Settings) -> list[str]:
    from .routing import load_registry

    return list(dict.fromkeys([settings.code_llm_model, *(m.id for m in load_registry())]))


def find_opencode(settings: Settings) -> str | None:
    return shutil.which(settings.opencode_bin)


class OpenCodeWorkspace:
    """One request's workspace. Use as a context manager so it is always cleaned up.

    `config` replaces the code-generation config (the prompt enhancer uses a free Zen model with no tools)."""

    def __init__(self, settings: Settings, config: dict | None = None, timeout: float | None = None):
        self.settings = settings
        self.config = config
        self.timeout = timeout or settings.opencode_timeout
        self.root = Path(tempfile.mkdtemp(prefix="rosdiff_opencode_"))
        self.ws = self.root / "ws"
        self.state = self.root / "state"
        self.ws.mkdir()
        self.state.mkdir()
        self.runs = 0

    def __enter__(self) -> OpenCodeWorkspace:
        return self

    def __exit__(self, *exc) -> None:
        shutil.rmtree(self.root, ignore_errors=True)

    def _env(self) -> dict[str, str]:
        env = {k: v for k, v in os.environ.items() if not k.startswith("OPENCODE_")}
        env.update(
            LLM_API_KEY=self.settings.llm_api_key,
            ZEN_API_KEY=self.settings.zen_api_key,
            OPENCODE_CONFIG_CONTENT=json.dumps(self.config or opencode_config(self.settings)),
            OPENCODE_DISABLE_PROJECT_CONFIG="1",
            OPENCODE_DISABLE_AUTOUPDATE="1",
            OPENCODE_DISABLE_SHARE="1",
            OPENCODE_DISABLE_CLAUDE_CODE="1",
            OPENCODE_DISABLE_LSP_DOWNLOAD="1",
            OPENCODE_DISABLE_MODELS_FETCH="1",
            OPENCODE_DISABLE_DEFAULT_PLUGINS="1",
            OPENCODE_DISABLE_EXTERNAL_SKILLS="1",
            DO_NOT_TRACK="1",
            XDG_DATA_HOME=str(self.state),
            XDG_STATE_HOME=str(self.state),
        )
        return env

    def run(self, message: str, context: str, model: str | None = None, provider: str = "llm") -> OpenCodeRun:
        """First call starts a session; later calls continue it (used for repair rounds)."""
        binary = shutil.which(self.settings.opencode_bin) or self.settings.opencode_bin
        if not shutil.which(binary):
            raise OpenCodeError(f"OpenCode binary {self.settings.opencode_bin!r} not found (npm i -g opencode-ai)")
        cmd = [
            binary,
            "run",
            message,
            "--dir",
            str(self.ws),
            "--model",
            f"{provider}/{model or self.settings.code_llm_model}",
            "--format",
            "json",
        ]
        if self.runs:
            cmd.append("--continue")
        if context:
            context_file = self.root / f"context_{self.runs}.md"
            context_file.write_text(context)
            cmd += ["-f", str(context_file)]
        self.runs += 1
        try:
            proc = subprocess.run(
                cmd,
                cwd=self.ws,
                env=self._env(),
                capture_output=True,
                text=True,
                timeout=self.timeout,
            )
        except subprocess.TimeoutExpired as e:
            raise OpenCodeError(f"OpenCode did not finish within {self.timeout:.0f}s") from e
        events = []
        for line in proc.stdout.splitlines():
            line = line.strip()
            if line.startswith("{"):
                try:
                    events.append(json.loads(line))
                except json.JSONDecodeError:
                    pass
        return OpenCodeRun(proc.returncode, events, proc.stderr[-2000:])

    def collect(self) -> dict[str, str]:
        """All text files OpenCode wrote, keyed by path relative to the workspace."""
        files: dict[str, str] = {}
        for p in sorted(self.ws.rglob("*")):
            rel = p.relative_to(self.ws)
            if not p.is_file() or p.is_symlink() or IGNORED_DIRS.intersection(rel.parts):
                continue
            if len(files) >= MAX_FILES:
                break
            if p.stat().st_size > MAX_FILE_BYTES:
                continue
            try:
                files[rel.as_posix()] = p.read_text()
            except UnicodeDecodeError:
                continue
        return files

    def write(self, files: dict[str, str]) -> None:
        """Replace the workspace contents (unused by the default flow; handy for tests)."""
        for rel, content in files.items():
            dest = self.ws / rel
            dest.parent.mkdir(parents=True, exist_ok=True)
            dest.write_text(content)