app/agy_bridge.py

"""Bridge to the Antigravity CLI (`agy`) WITHOUT any direct API calls.

`agy -p` emits nothing to stdout; it persists the answer under
<home>/.gemini/antigravity-cli/brain/<conv-id>/.system_generated/logs/transcript.jsonl
and maps <working-dir> -> <conv-id> in <home>/.gemini/antigravity-cli/cache/last_conversations.json.

ISOLATION: every call takes a `home` directory. We point the agy subprocess's
HOME/USERPROFILE at it, so each API key's conversations, scratch, and config live
in their OWN folder. Combined with allowNonWorkspaceAccess=false + sandbox, the
agent cannot read or write outside that home (no cross-tenant access).
"""
from __future__ import annotations

import json
import os
import re
import shutil
import subprocess
import tempfile
import time
from pathlib import Path

AGY_BIN = os.environ.get(
    "AGY_BIN", str(Path.home() / "AppData" / "Local" / "agy" / "bin" / "agy.exe")
)
DEFAULT_TIMEOUT = float(os.environ.get("AGY_TIMEOUT", "180"))

# Fallback home when no per-key home is supplied (local dev / no auth).
_SERVER_HOME = os.environ.get("AGY_SERVER_HOME", "").strip()
DEFAULT_HOME = Path(_SERVER_HOME) if _SERVER_HOME else Path.home()

# Hardened, isolated config written into EVERY home: web search allowed; commands,
# files, image-gen, subagents denied; no access outside the workspace.
_HARDENED_SETTINGS = {
    "toolPermission": "strict",
    "allowNonWorkspaceAccess": False,
    "enableTerminalSandbox": True,
    "runningLightSpeed": "off",
    "verbosity": "low",
    "altScreenMode": "never",
    "showTips": False,
    "notifications": False,
    "showFeedbackSurvey": False,
    "model": os.environ.get("AGY_MODEL", "gemini-3.5-flash"),
    "permissions": {
        "allow": ["search_web(*)", "read_url_content(*)", "ask_question(*)", "finish(*)"],
        "deny": [
            "run_command(*)", "edit_file(*)", "create_file(*)", "view_file(*)",
            "find_file(*)", "list_directory(*)", "search_directory(*)",
            "start_subagent(*)", "generate_image(*)",
        ],
        "ask": [],
    },
}

_ready_homes: set[str] = set()

# On Linux/HF, agy reads creds from a file relative to HOME (no keyring). Since we
# redirect HOME per key, each key's home needs its own copy of the credential.
# No-op on Windows (AGY_CREDS_JSON unset -> uses the shared OS keyring).
_CREDS_JSON = os.environ.get("AGY_CREDS_JSON", "").strip()
# agy reads its OAuth token from this file (relative to HOME), confirmed on Linux.
_CREDS_FILENAME = os.environ.get("AGY_CREDS_FILENAME", "antigravity-oauth-token")

# Run agy with terminal restrictions so the agent can't run commands that read our
# home internals (OAuth token, trajectory DB, other sessions). Web search still works.
_SANDBOX = os.environ.get("AGY_SANDBOX", "true").strip().lower() == "true"

# Steps that are actually tool output / internal dumps, not a chat answer. We never
# return these as the assistant's reply.
_INTERNAL_STEP = re.compile(
    r"(tables in sqlite|trajectory_meta|executor_metadata|gen_metadata|parent_references|"
    r"trajectory_metadata_blob|battle_mode_infos|the command completed successfully|"
    r"antigravity-cli|\.system_generated|antigravity-oauth|last_conversations|"
    r"created at:\s*\d{4}-\d\d-\d\d|completed at:\s*\d{4}-\d\d-\d\d)",
    re.I | re.S,
)
_FALLBACK_ANSWER = "Sorry, I couldn't produce a proper response to that. Could you rephrase?"


def _pick_answer(steps: list[str]) -> str | None:
    """Return the most recent step that is a real chat reply (not a tool/internal dump)."""
    for s in reversed(steps):
        if s and s.strip() and not _INTERNAL_STEP.search(s):
            return s
    return None


class AgyError(RuntimeError):
    pass


def agy_dir(home: Path) -> Path:
    return home / ".gemini" / "antigravity-cli"


def ensure_home(home: Path) -> None:
    """Write the hardened config (+ creds on Linux/HF) into this home, once per home."""
    key = os.path.normcase(str(home))
    if key in _ready_homes:
        return
    cfg_dir = agy_dir(home)
    cfg_dir.mkdir(parents=True, exist_ok=True)
    (cfg_dir / "settings.json").write_text(json.dumps(_HARDENED_SETTINGS, indent=2), encoding="utf-8")
    if _CREDS_JSON:  # Linux/HF: each isolated home gets the credential file
        cf = cfg_dir / _CREDS_FILENAME
        cf.write_text(_CREDS_JSON, encoding="utf-8")
        try:
            os.chmod(cf, 0o600)
        except OSError:
            pass
    _ready_homes.add(key)


def _write_agents_md(workdir: Path, agents_md: str | None) -> None:
    """Drop (or clear) a per-user AGENTS.md in the workspace so agy loads it natively.
    Called per run so an admin's later edit takes effect on the next turn."""
    target = Path(workdir) / "AGENTS.md"
    try:
        if agents_md and agents_md.strip():
            target.write_text(agents_md.strip() + "\n", encoding="utf-8")
        elif target.exists():
            target.unlink()
    except OSError:
        pass


def _child_env(home: Path) -> dict:
    """Env for the agy subprocess; HOME/USERPROFILE -> this key's isolated home."""
    env = os.environ.copy()
    env["USERPROFILE"] = str(home)
    env["HOME"] = str(home)
    return env


def _transcript_path(conv_id: str, home: Path) -> Path:
    return agy_dir(home) / "brain" / conv_id / ".system_generated" / "logs" / "transcript.jsonl"


def _conv_id_for_dir(workdir: Path, home: Path) -> str | None:
    lc = agy_dir(home) / "cache" / "last_conversations.json"
    try:
        mapping = json.loads(lc.read_text(encoding="utf-8"))
    except (FileNotFoundError, json.JSONDecodeError):
        return None
    target = os.path.normcase(os.path.abspath(str(workdir)))
    for k, conv_id in mapping.items():
        if os.path.normcase(os.path.abspath(k)) == target:
            return conv_id
    return None


def _iter_model(conv_id: str, home: Path):
    path = _transcript_path(conv_id, home)
    if not path.is_file():
        return
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
        except json.JSONDecodeError:
            continue
        yield obj


def _model_steps(conv_id: str, home: Path) -> list[str]:
    return [
        o["content"] for o in _iter_model(conv_id, home)
        if o.get("source") == "MODEL" and isinstance(o.get("content"), str) and o["content"].strip()
    ]


def _extract_answer(conv_id: str, home: Path) -> str:
    steps = _model_steps(conv_id, home)
    if not steps:
        raise AgyError(f"no MODEL text response in transcript for {conv_id}")
    return _pick_answer(steps) or _FALLBACK_ANSWER


def session_stats(conv_id: str, home: Path) -> dict:
    turns = chars = 0
    for o in _iter_model(conv_id, home):
        c = o.get("content")
        if isinstance(c, str):
            chars += len(c)
            if o.get("source") in ("USER_EXPLICIT", "MODEL"):
                turns += 1
    return {"turns": turns, "context_chars": chars}


def _cmd(base: list[str], model: str | None, image_dir: str | None) -> list[str]:
    cmd = list(base)
    if _SANDBOX:
        cmd += ["--sandbox"]  # terminal restrictions: no reading our home internals
    if model:
        cmd += ["--model", model]
    if image_dir:
        cmd += ["--add-dir", image_dir]  # grant read access to ONLY the attached image
    return cmd


def _run(cmd: list[str], cwd: str, home: Path, timeout: float) -> subprocess.CompletedProcess:
    if not Path(AGY_BIN).is_file():
        raise AgyError(f"agy binary not found at {AGY_BIN}")
    ensure_home(home)
    try:
        return subprocess.run(
            cmd, cwd=cwd, timeout=timeout, capture_output=True, text=True, env=_child_env(home)
        )
    except subprocess.TimeoutExpired as e:
        raise AgyError(f"agy timed out after {timeout}s") from e


def run_agy(prompt: str, model: str | None = None, image_dir: str | None = None,
            home: Path | None = None, timeout: float = DEFAULT_TIMEOUT,
            agents_md: str | None = None) -> dict:
    """Stateless one-shot in the key's isolated home."""
    home = home or DEFAULT_HOME
    workdir = Path(tempfile.mkdtemp(prefix="agy_bridge_"))
    started = time.time()
    try:
        _write_agents_md(workdir, agents_md)
        cmd = _cmd([AGY_BIN, "-p", prompt], model, image_dir)
        proc = _run(cmd, str(workdir), home, timeout)
        conv_id = _conv_id_for_dir(workdir, home)
        if not conv_id:
            raise AgyError(f"could not resolve conversation id (exit={proc.returncode}, stderr={proc.stderr[:200]!r})")
        return {"answer": _extract_answer(conv_id, home), "conversation_id": conv_id,
                "elapsed": round(time.time() - started, 2), "exit_code": proc.returncode}
    finally:
        shutil.rmtree(workdir, ignore_errors=True)


def start_session(prompt: str, session_dir: Path, model: str | None = None,
                  image_dir: str | None = None, home: Path | None = None,
                  timeout: float = DEFAULT_TIMEOUT, agents_md: str | None = None) -> dict:
    """First turn of a session in the key's isolated home."""
    home = home or DEFAULT_HOME
    session_dir.mkdir(parents=True, exist_ok=True)
    _write_agents_md(session_dir, agents_md)
    started = time.time()
    proc = _run(_cmd([AGY_BIN, "-p", prompt], model, image_dir), str(session_dir), home, timeout)
    conv_id = _conv_id_for_dir(session_dir, home)
    if not conv_id:
        raise AgyError(f"could not resolve conversation id (exit={proc.returncode}, stderr={proc.stderr[:200]!r})")
    steps = _model_steps(conv_id, home)
    if not steps:
        raise AgyError(f"no MODEL response for new session {conv_id}")
    answer = _pick_answer(steps) or _FALLBACK_ANSWER
    return {"answer": answer, "conversation_id": conv_id, "elapsed": round(time.time() - started, 2)}


def continue_session(prompt: str, conv_id: str, session_dir: Path, model: str | None = None,
                     image_dir: str | None = None, home: Path | None = None,
                     timeout: float = DEFAULT_TIMEOUT, agents_md: str | None = None) -> dict:
    """Resume a conversation in the key's isolated home; agy retains context."""
    home = home or DEFAULT_HOME
    _write_agents_md(session_dir, agents_md)
    before = len(_model_steps(conv_id, home))
    started = time.time()
    _run(_cmd([AGY_BIN, "--conversation", conv_id, "-p", prompt], model, image_dir),
         str(session_dir), home, timeout)
    steps = _model_steps(conv_id, home)
    # Only consider steps produced by THIS turn — never fall back to a prior turn's
    # step (that caused a stale tool-dump to repeat across unrelated questions).
    answer = _pick_answer(steps[before:]) or _FALLBACK_ANSWER
    return {"answer": answer, "conversation_id": conv_id, "elapsed": round(time.time() - started, 2)}