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

255 lines
12 KiB
Python

"""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, 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):
"""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(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": 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):
"""Use the native pinned binary, never the privileged Hermes shell wrapper."""
return [CLAUDE_BIN, "-p", "--output-format", "stream-json", "--verbose",
"--no-session-persistence", "--safe-mode", "--tools", "",
"--strict-mcp-config", "--mcp-config", '{"mcpServers":{}}',
"--setting-sources", "", "--disable-slash-commands",
"--permission-mode", "dontAsk", "--no-chrome",
"--model", MODELS["claude"]["cli_model"], "--effort", "medium", "--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"):
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":
status = final.get("api_error_status")
code = "rate_limit" if status == 429 else "incomplete_generation"
if status in (401, 403):
code = "provider_authentication"
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") != 64000
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": 64000}
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):
"""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"]
with tempfile.TemporaryDirectory(prefix="suite-", dir="/jobs") as directory:
root = Path(directory)
(root / "input").write_text(prompt(request))
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(claude_command(model, request["execution"]["max_cost_usd"]),
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)) from None
deadline = time.monotonic() + seconds
try:
while process.poll() is None:
if cancel.wait(0.1):
raise Problem("cancelled", 409)
if time.monotonic() > deadline:
raise Problem("timeout", 504)
if (root / "output").stat().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}
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