tai/tai_pkg/opencode.py

"""Opencode CLI backend, prompts, and response parsing."""

import json
import os
import random
import re
import shutil
import subprocess
import tempfile

from tai_pkg.consts import Consts
from tai_pkg.errors import BackendError
from tai_pkg.muse import Muse
from tai_pkg.text import Text

__all__ = ["Opencode", "_OPENCODE_ID_ALPHABET", "_OPENCODE_SESSION"]


class Opencode:
    """Opencode."""

    @staticmethod
    def opencode_binary():
        return os.environ.get("OPENCODE_BIN") or shutil.which("opencode")

    @staticmethod
    def opencode_prompt(messages, tools):
        schema = json.dumps(Muse.muse_output_schema())
        rules = (
            "You are the model behind another agent. The transcript above is the whole conversation; lines starting with TOOL <id> are results of tool calls that already ran. "
            "Do not use any of your own tools and do not repeat a tool call whose result is already shown. "
            "If the results answer the user, put the final answer in content with an empty tool_calls list; otherwise request the next tool calls from TOOLS. "
            "Reply with one JSON object only, no code fence, matching this schema: "
        )
        return Muse.muse_prompt(messages, tools) + "\n\n" + rules + schema

    @staticmethod
    def opencode_parse(text):
        text = (text or "").strip()
        fenced = re.match(r"^```(?:json)?\s*(.*?)\s*```$", text, re.S)
        if fenced:
            text = fenced.group(1)
        # Scan every "{" last-first and decode one value at each spot: free
        # models often echo the prompt schema or tool JSON before the real
        # answer, and a first-"{"-to-last-"}" slice spans both into garbage.
        # Last-first means the model's own trailing answer wins over echoes.
        decoder = json.JSONDecoder()
        starts = [pos for pos, char in enumerate(text) if char == "{"]
        for pos in reversed(starts[-Consts.OPENCODE_PARSE_MAX_SCANS:]):
            try:
                answer, _end = decoder.raw_decode(text, pos)
            except ValueError:
                continue
            if isinstance(answer, dict) and ("content" in answer or "tool_calls" in answer):
                return answer
        return {"content": text, "tool_calls": []}

    @staticmethod
    def opencode_parse_events(stdout_text, returncode, stderr_text):
        """Shared by opencode_run() (local CLI) and box_opencode_run() (same CLI,
    executed inside the sandbox box) -- the JSON-lines event stream is
    identical either way, only how the process gets launched differs."""
        texts = []
        for line in stdout_text.splitlines():
            try:
                event = json.loads(line)
            except ValueError:
                continue
            if not isinstance(event, dict):
                continue
            if event.get("type") == "error":
                data = (event.get("error") or {}).get("data") or {}
                raise BackendError("opencode: %s" % (data.get("message") or json.dumps(event.get("error"))[:300]), data.get("statusCode") or 0)
            part = event.get("part") or {}
            if event.get("type") == "text" and part.get("text"):
                texts.append(part["text"])
        if not texts:
            raise BackendError("opencode returned no text (exit %d): %s" % (returncode, (stderr_text or stdout_text)[-300:]))
        return "".join(texts)

    @staticmethod
    def opencode_run(model, prompt, timeout):
        binary = Opencode.opencode_binary()
        if not binary:
            raise BackendError("opencode CLI not found (OPENCODE_BIN or PATH)")
        argv = [binary, "run", "--pure", "--agent", Consts.OPENCODE_AGENT, "--format", "json", "-m", "opencode/" + model]
        with tempfile.TemporaryDirectory(prefix="tai-opencode-") as folder:
            try:
                done = subprocess.run(argv, input=prompt, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, cwd=folder, timeout=timeout)
            except subprocess.TimeoutExpired:
                raise BackendError("opencode run timed out after %ds" % timeout)
            except OSError as exc:
                raise BackendError("opencode run failed: %s" % Text.short_error(exc))
        return Opencode.opencode_parse_events(done.stdout, done.returncode, done.stderr)

    @staticmethod
    def opencode_free_models():
        binary = Opencode.opencode_binary()
        if not binary:
            return None
        try:
            done = subprocess.run([binary, "models", "opencode"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=Consts.MODEL_LIST_TIMEOUT * 4)
        except (OSError, subprocess.TimeoutExpired):
            return None
        models = [line.strip().split("/", 1)[1] for line in done.stdout.splitlines() if line.strip().startswith("opencode/")]
        return [item for item in models if Text.is_free_model(item)] or None

    @staticmethod
    def opencode_headers():
        return {
            "User-Agent": os.environ.get("OPENCODE_CLIENT_USER_AGENT") or Consts.OPENCODE_USER_AGENT,
            "x-opencode-client": os.environ.get("OPENCODE_CLIENT_NAME") or Consts.OPENCODE_CLIENT_NAME,
            "x-opencode-project": "global",
            "x-opencode-session": _OPENCODE_SESSION,
            "x-opencode-request": "msg_" + "".join(random.choice(_OPENCODE_ID_ALPHABET) for _ in range(26)),
        }


_OPENCODE_ID_ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"


_OPENCODE_SESSION = "ses_" + "".join(random.choice(_OPENCODE_ID_ALPHABET) for _ in range(26))