tai/tai_pkg/sandbox.py

"""Container sandbox box and in-box command execution."""

import os
import shutil
import subprocess
import sys
import tempfile

from tai_pkg.consts import Consts
from tai_pkg.errors import BackendError
from tai_pkg.opencode import Opencode
from tai_pkg.render import Ansi
from tai_pkg.render import Render
from tai_pkg.text import Text

__all__ = ["Sandbox", "_BOX_PYTHON_CACHE", "_BOX_OPENCODE_CACHE"]


class Sandbox:
    """Sandbox."""

    @staticmethod
    def container_engine():
        if shutil.which("podman"):
            return "podman"
        if shutil.which("docker"):
            return "docker"
        return ""

    @staticmethod
    def box_state(engine):
        try:
            done = subprocess.run([engine, "inspect", "-f", "{{.State.Running}}", Consts.BOX_NAME], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=15)
        except (OSError, subprocess.SubprocessError):
            return "missing"
        if done.returncode != 0:
            return "missing"
        return "running" if done.stdout.strip() == "true" else "stopped"

    @staticmethod
    def ensure_box():
        engine = Sandbox.container_engine()
        if not engine:
            raise BackendError("no container engine found, install podman")
        state = Sandbox.box_state(engine)
        if state == "running":
            return engine
        if state == "stopped":
            done = subprocess.run([engine, "start", Consts.BOX_NAME], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=60)
            if done.returncode != 0:
                raise BackendError("cannot start box: " + (done.stderr or "").strip()[:200])
            return engine
        print("building sandbox image %s (one-time, takes minutes)..." % Consts.BOX_IMAGE)
        with tempfile.TemporaryDirectory() as build:
            for filename, content in (("Containerfile", Consts.BOX_CONTAINERFILE), ("stt.py", Consts.BOX_STT), ("tts.py", Consts.BOX_TTS)):
                with open(os.path.join(build, filename), "w", encoding="utf-8") as handle:
                    handle.write(content)
            done = subprocess.run([engine, "build", "-t", Consts.BOX_IMAGE, build], timeout=1800)
            if done.returncode != 0:
                raise BackendError("box image build failed")
        done = subprocess.run([engine, "run", "-d", "--name", Consts.BOX_NAME, "--restart", "unless-stopped", Consts.BOX_IMAGE], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=120)
        if done.returncode != 0:
            raise BackendError("cannot start box: " + (done.stderr or "").strip()[:200])
        return engine

    @staticmethod
    def box_exec(engine, argv, extra=(), input_data=None, timeout=180):
        return subprocess.run([engine, "exec", *extra, "-i", Consts.BOX_NAME, *argv], input=input_data, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=timeout)

    @staticmethod
    def box_python(engine):
        if engine not in _BOX_PYTHON_CACHE:
            try:
                probe = Sandbox.box_exec(engine, ["test", "-x", Consts.BOX_PYTHON], timeout=30)
                _BOX_PYTHON_CACHE[engine] = Consts.BOX_PYTHON if probe.returncode == 0 else "python3"
            except (OSError, subprocess.SubprocessError):
                _BOX_PYTHON_CACHE[engine] = "python3"
        return _BOX_PYTHON_CACHE[engine]

    @staticmethod
    def box_transcribe(engine, audio, timeout=300):
        done = Sandbox.box_exec(engine, [Sandbox.box_python(engine), "/box/stt.py"], input_data=audio, timeout=timeout)
        if done.returncode != 0:
            raise BackendError("box stt failed: " + done.stderr.decode("utf-8", "replace")[:300])
        return done.stdout.decode("utf-8", "replace").strip()

    @staticmethod
    def box_speak(engine, text, voice=""):
        argv = [Sandbox.box_python(engine), "/box/tts.py"] + ([voice] if voice else [])
        done = Sandbox.box_exec(engine, argv, input_data=text.encode("utf-8"), timeout=180)
        if done.returncode != 0 or not done.stdout:
            raise BackendError("box tts failed: " + done.stderr.decode("utf-8", "replace")[:300])
        return done.stdout

    @staticmethod
    def box_opencode_binary(engine):
        """True once the opencode CLI is confirmed present inside the sandbox
    box, installing Node.js + `npm install -g opencode-ai` into the box
    first if it isn't there yet. Only ever called as a last resort (see
    ChatClient.all_candidates()) and only against a box that is ALREADY
    running -- this never creates the box itself; that stays the explicit,
    approval-gated action it already is (ensure_box(), via --env sandbox or
    the install tool). Cached per engine for the process lifetime so normal
    turns that succeed through a real backend never pay this cost twice."""
        if engine in _BOX_OPENCODE_CACHE:
            return _BOX_OPENCODE_CACHE[engine]
        try:
            probe = Sandbox.box_exec(engine, ["which", "opencode"], timeout=10)
        except (OSError, subprocess.SubprocessError):
            _BOX_OPENCODE_CACHE[engine] = False
            return False
        if probe.returncode == 0:
            _BOX_OPENCODE_CACHE[engine] = True
            return True
        print(Render.paint(
            "opencode: no local CLI on this machine; installing it into the existing sandbox "
            "box (container '%s') instead, as a last-resort backend. Nothing outside that "
            "container is touched. Streaming the install below:" % Consts.BOX_NAME, Ansi.DIM
        ), file=sys.stderr)
        try:
            # Deliberately not box_exec(): this one streams straight to the
            # real stdout/stderr instead of capturing, so the person watching
            # sees the actual apt/curl/npm output live, not a silent pause.
            done = subprocess.run([engine, "exec", Consts.BOX_NAME, "sh", "-c", Consts.BOX_OPENCODE_INSTALL_SCRIPT], timeout=600)
        except (OSError, subprocess.SubprocessError) as exc:
            print(Render.paint("opencode box install failed: %s" % Text.short_error(exc), Ansi.DIM), file=sys.stderr)
            _BOX_OPENCODE_CACHE[engine] = False
            return False
        ok = done.returncode == 0
        print(Render.paint("opencode box install %s" % ("succeeded" if ok else "failed (exit %d)" % done.returncode), Ansi.DIM), file=sys.stderr)
        _BOX_OPENCODE_CACHE[engine] = ok
        return ok

    @staticmethod
    def box_opencode_run(engine, model, prompt, timeout):
        argv = ["opencode", "run", "--pure", "--agent", Consts.OPENCODE_AGENT, "--format", "json", "-m", "opencode/" + model]
        try:
            done = Sandbox.box_exec(engine, argv, input_data=prompt.encode("utf-8"), timeout=timeout)
        except subprocess.TimeoutExpired:
            raise BackendError("opencode (box) run timed out after %ds" % timeout)
        except (OSError, subprocess.SubprocessError) as exc:
            raise BackendError("opencode (box) run failed: %s" % Text.short_error(exc))
        return Opencode.opencode_parse_events(done.stdout.decode("utf-8", "replace"), done.returncode, done.stderr.decode("utf-8", "replace"))


_BOX_PYTHON_CACHE = {}


_BOX_OPENCODE_CACHE = {}