atlas-iac/services/hermes/scripts/suite_backends.py

"""Fixed-destination transports and an isolated, tool-free Claude invocation."""
from __future__ import annotations

import json
import os
from pathlib import Path
import signal
import subprocess
import tempfile
import time
from urllib.error import HTTPError, URLError
from urllib.request import HTTPRedirectHandler, ProxyHandler, Request, build_opener

from suite_contract import CLAUDE_MAX_TURNS, CLAUDE_MODELS, MODELS, SCHEMA, SYSTEM, Problem, encoded, prompt
from suite_cli_diagnostics import snapshot

SWITCHYARD = "http://hermes-switchyard.hermes.svc.cluster.local:9005/v1/chat/completions"
LOCAL = "http://hermes-model-gate-lan-api.hermes.svc.cluster.local:8082"
CLAUDE_BIN = "/opt/cli/claude"
OUTPUT_BYTES = 4 << 20


class NoRedirect(HTTPRedirectHandler):
    """Never forward credentials or source material to a redirect destination."""

    def redirect_request(self, req, fp, code, msg, headers, newurl):
        raise Problem("upstream_redirect", 502)


def post(url, value, timeout, headers=None):
    """One fixed HTTP request, without proxies, redirects, retries, or body logs."""
    request = Request(url, data=encoded(value), method="POST",
                      headers={"Content-Type": "application/json", **(headers or {})})
    try:
        with build_opener(ProxyHandler({}), NoRedirect()).open(request, timeout=timeout) as response:
            raw = response.read(OUTPUT_BYTES + 1)
    except HTTPError as exc:
        status = exc.code
        exc.close()
        raise Problem("rate_limit" if status == 429 else "backend_unavailable", 503,
                      upstream_status=status) from None
    except (URLError, OSError, TimeoutError):
        raise Problem("backend_timeout_or_unavailable", 504) from None
    if len(raw) > OUTPUT_BYTES:
        raise Problem("response_too_large", 502)
    try:
        return json.loads(raw)
    except ValueError:
        raise Problem("invalid_json_result", 502) from None


def switchyard_decision(provider):
    """Switchyard sees only a fixed route label, never the suite or its profiles."""
    result = post(SWITCHYARD, {"model": "atlas/planning/" + provider,
                             "messages": [{"role": "user", "content": "select"}],
                             "stream": False}, 20)
    try:
        decision = json.loads(result["choices"][0]["message"]["content"])
    except (KeyError, IndexError, TypeError, ValueError):
        raise Problem("routing_decision_invalid", 502) from None
    expected = {"provider": provider, "model": MODELS[provider]["model"]}
    if decision != expected:
        raise Problem("routing_policy_violation", 502)
    return expected


def local_generate(request, cancel, client_ip, *, invocation=None):
    """Reuse the unchanged RTX API and its pinned-model, local-only safeguards."""
    headers = {"Authorization": "Bearer " + Path("/vault/secrets/local-token").read_text().strip(),
               "X-Forwarded-For": client_ip}
    schema = json.loads(encoded(invocation["schema"] if invocation else SCHEMA))
    # Constrain the local decoder to aliases, excluding copied case descriptions.
    schema["properties"]["groups"]["items"]["properties"]["members"]["items"]["enum"] = [
        case["alias"] for case in request["cases"]]
    response = post(LOCAL + "/api/generate", {
        "model": MODELS["local"]["model"],
        "prompt": invocation["system"] + "\n" + invocation["input"] if invocation else SYSTEM + "\n" + prompt(request),
        "stream": False, "format": schema,
        "options": {"num_predict": 2048, "temperature": 0, "seed": 0}},
        request["execution"]["max_seconds"], headers)
    if cancel.is_set():
        raise Problem("cancelled", 409)
    if not response.get("done") or response.get("done_reason") == "length":
        raise Problem("incomplete_generation", 502)
    if response.get("model") != MODELS["local"]["model"]:
        raise Problem("model_changed", 502)
    try:
        result = json.loads(response["response"])
    except (KeyError, TypeError, ValueError):
        raise Problem("invalid_json_result", 502) from None
    return result, {"model": response["model"], "compaction": False,
                    "truncation": False, "usage": {k: response.get(k) for k in (
                        "prompt_eval_count", "eval_count", "total_duration",
                        "load_duration", "prompt_eval_duration", "eval_duration")},
                    "provenance": response.get("inference_provenance")}


def claude_command(model, max_cost, *, reasoning=None):
    """Use the native pinned binary, never the privileged Hermes shell wrapper."""
    if model not in CLAUDE_MODELS:
        raise Problem("unsupported_backend", 422)
    reasoning = MODELS["claude"]["reasoning"] if reasoning is None else reasoning
    if reasoning not in {"high", "xhigh"}:
        raise Problem("unsupported_reasoning", 422)
    settings = {"enabledPlugins": {"agents-md@builtin": False,
                                    "cc-plugin-agents-md@builtin": False},
                "disableAllHooks": True, "disableBundledSkills": True,
                "disableClaudeAiConnectors": True, "syncClaudeAiPlugins": False,
                "availableModels": [model], "switchModelsOnFlag": False}
    return [CLAUDE_BIN, "-p", "--output-format", "stream-json", "--verbose",
            "--no-session-persistence", "--safe-mode", "--tools", "",
            "--strict-mcp-config", "--mcp-config", '{"mcpServers":{}}',
            "--setting-sources", "", "--settings", encoded(settings).decode(), "--disable-slash-commands",
            "--permission-mode", "dontAsk", "--no-chrome",
            "--model", model + "[1m]", "--effort", reasoning,
            "--max-budget-usd", str(max_cost),
            "--max-turns", str(CLAUDE_MAX_TURNS), "--system-prompt", SYSTEM,
            "--json-schema", encoded(SCHEMA).decode()]


def claude_environment(directory, token):
    """Construct an allowlisted environment with no inherited tools or API keys."""
    env = {"PATH": "/usr/bin:/bin", "HOME": directory,
           "CLAUDE_CONFIG_DIR": directory + "/config", "TMPDIR": directory,
           "CLAUDE_CODE_OAUTH_TOKEN": token,
           "CLAUDE_CODE_MAX_OUTPUT_TOKENS": "64000", "CLAUDE_CODE_MAX_RETRIES": "0"}
    for key in ("DISABLE_COMPACT", "DISABLE_AUTO_COMPACT", "DISABLE_TELEMETRY",
                "DISABLE_ERROR_REPORTING", "DISABLE_AUTOUPDATER", "DISABLE_UPDATES",
                "DISABLE_PROMPT_CACHING", "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC",
                "CLAUDE_CODE_DISABLE_BACKGROUND_TASKS", "CLAUDE_CODE_DISABLE_TERMINAL_TITLE",
                "CLAUDE_CODE_DISABLE_AUTO_MEMORY", "CLAUDE_CODE_SKIP_PROMPT_HISTORY",
                "CLAUDE_AGENT_SDK_DISABLE_BUILTIN_AGENTS", "CLAUDE_CODE_DISABLE_GIT_INSTRUCTIONS",
                "CLAUDE_CODE_DISABLE_BUNDLED_SKILLS"):
        env[key] = "1"
    return env


def stop(process):
    """Terminate the complete CLI process group and reap it on every exit path."""
    if process.poll() is None:
        try:
            os.killpg(process.pid, signal.SIGTERM)
        except ProcessLookupError:
            pass  # The process can exit between poll and kill; still reap it.
        try:
            process.wait(timeout=2)
        except subprocess.TimeoutExpired:
            try:
                os.killpg(process.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
            process.wait(timeout=5)


def parse_claude(raw, expected_model, **process_info):
    """Normalize only a completed, uncompacted CLI result with a pinned model."""
    diagnostics = snapshot(raw, **process_info)

    def fail(code, stage):
        raise Problem(code, 502, **diagnostics, failure_stage=stage)

    final, initialized = None, False
    aliases = {expected_model, expected_model + "[1m]"}
    for line in raw.splitlines():
        try:
            event = json.loads(line)
        except ValueError:
            fail("invalid_json_result", "cli_event_json")
        if not isinstance(event, dict):
            fail("invalid_json_result", "cli_event_json")
        if "compact" in str(event.get("subtype", "")):
            fail("compaction_detected", "cli_events")
        if event.get("type") == "system" and event.get("subtype") == "init":
            initialized = True
            if (event.get("model") not in aliases or event.get("mcp_servers") or
                    event.get("plugins") or set(event.get("tools", [])) - {"StructuredOutput"}):
                fail("worker_isolation_failed", "cli_initialization")
        if event.get("type") == "result":
            final = event
    if not initialized or not final:
        fail("incomplete_generation", "missing_init_or_final_event")
    if final.get("is_error") or final.get("subtype") != "success":
        if final.get("subtype") == "error_max_budget_usd":
            fail("job_cost_budget_exhausted", "cli_final_result")
        status = final.get("api_error_status")
        assistant_error = diagnostics["assistant_error_code"]
        transport_error = diagnostics["cli_transport_error"]
        code = "incomplete_generation"
        if status == 429 or assistant_error == "rate_limit":
            code = "rate_limit"
        elif status in (401, 403) or assistant_error in {"authentication_failed", "oauth_org_not_allowed"}:
            code = "provider_authentication"
        elif transport_error == "request_timeout":
            code = "backend_timeout_or_unavailable"
        elif (transport_error in {"connection_error", "connection_refused"}
              or assistant_error in {"overloaded", "server_error"}
              or (type(status) is int and 500 <= status <= 599)):
            code = "backend_unavailable"
        fail(code, "cli_final_result")
    if final.get("stop_reason") in {"max_tokens", "model_context_window_exceeded"}:
        fail("incomplete_generation", "provider_stop")
    models = final.get("modelUsage", {})
    if not isinstance(models, dict) or not models or set(models) - aliases:
        fail("model_changed", "runtime_model")
    for limits in models.values():
        if (not isinstance(limits, dict) or limits.get("contextWindow") != 1000000
                or limits.get("maxOutputTokens") != CLAUDE_MODELS.get(expected_model)
                or limits.get("canonicalModel") != expected_model
                or limits.get("provider") != "firstParty"):
            fail("backend_capabilities_changed", "runtime_model_limits")
    usage = final.get("usage", {})
    if not isinstance(usage, dict) or not isinstance(usage.get("server_tool_use", {}), dict):
        fail("invalid_json_result", "cli_usage")
    if any(usage.get("server_tool_use", {}).values()):
        fail("worker_isolation_failed", "provider_tools")
    result = final.get("structured_output")
    if not isinstance(result, dict):
        fail("invalid_json_result", "structured_output_extraction")
    if process_info.get("exit_code"):
        fail("incomplete_generation", "process_exit")
    # Persist measured metadata only, including on later schema/coverage failure.
    safe_models = {name: {"canonicalModel": expected_model, "provider": "firstParty",
                         "contextWindow": 1000000, "maxOutputTokens": CLAUDE_MODELS[expected_model]}
                   for name in models}
    return result, {"model": expected_model, "compaction": False,
                    "truncation": False, "compaction_signal": "CLI events and disabled compaction",
                    "usage": diagnostics["usage"], "model_usage": safe_models,
                    "duration_api_ms": diagnostics["duration_api_ms"],
                    "cost_usd_estimate": diagnostics["cost_usd_estimate"],
                    "turns": diagnostics["turns"], "cli_diagnostics": diagnostics}


def claude_generate(request, cancel, *, invocation=None, progress=None):
    """Run one fresh job in tmpfs; input, output, configuration and caches expire together."""
    token = Path("/vault/secrets/claude-token").read_text().strip()
    if not token:
        raise Problem("provider_authentication", 503)
    model = MODELS["claude"]["model"]
    effort = invocation.get("reasoning", MODELS["claude"]["reasoning"]) if invocation else MODELS["claude"]["reasoning"]
    with tempfile.TemporaryDirectory(prefix="suite-", dir="/jobs") as directory:
        root = Path(directory)
        (root / "input").write_text(invocation["input"] if invocation else prompt(request))
        command = claude_command(model, request["execution"]["max_cost_usd"],
                                 reasoning=effort)
        if invocation:
            command[command.index("--system-prompt") + 1] = invocation["system"]
            command[command.index("--json-schema") + 1] = encoded(invocation["schema"]).decode()
        env = claude_environment(directory, token)
        seconds = request["execution"]["max_seconds"]
        failure = None
        with (root / "input").open("rb") as source, (root / "output").open("wb") as output:
            try:
                process = subprocess.Popen(command,
                                           stdin=source, stdout=output, stderr=subprocess.DEVNULL,
                                           env=env, cwd=directory, start_new_session=True)
            except OSError:
                raise Problem("backend_unavailable", 503, failure_stage="process_start",
                              **snapshot("", subprocess_timeout_seconds=seconds,
                                         reasoning_effort=effort)) from None
            deadline = time.monotonic() + seconds
            next_progress = 0.0
            try:
                while process.poll() is None:
                    if cancel.wait(0.1):
                        raise Problem("cancelled", 409)
                    if time.monotonic() > deadline:
                        raise Problem("timeout", 504)
                    activity = (root / "output").stat()
                    if progress and time.monotonic() >= next_progress:
                        # File activity proves CLI activity, not semantic completion.
                        # Do not read event bodies or emit model reasoning as progress.
                        progress({"cli_running": True, "cli_output_bytes": activity.st_size,
                                  "last_cli_activity_seconds_ago": round(max(0, time.time() - activity.st_mtime), 1)})
                        next_progress = time.monotonic() + 5
                    if activity.st_size > OUTPUT_BYTES:
                        raise Problem("response_too_large", 502)
            except Problem as exc:
                failure = exc
            finally:
                stop(process)
        if (root / "output").stat().st_size > OUTPUT_BYTES:
            failure = failure or Problem("response_too_large", 502)
        # Read only a bounded prefix even if the CLI exits between size checks.
        with (root / "output").open("rb") as output:
            raw = output.read(OUTPUT_BYTES).decode("utf-8", errors="replace")
        info = {"exit_code": process.returncode, "subprocess_timeout_seconds": seconds,
                "termination_reason": failure.code if failure else None, "reasoning_effort": effort}
        if failure:
            raise Problem(failure.code, failure.status, failure_stage="subprocess",
                          **snapshot(raw, **info))
        if process.returncode and not raw.strip():
            raise Problem("backend_unavailable", 503, failure_stage="process_exit",
                          **snapshot(raw, **info))
        result, metadata = parse_claude(raw, model, **info)
        metadata["temporary_files_deleted"] = True
        return result, metadata