219 lines
10 KiB
Python
219 lines
10 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 re
|
|
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 MODELS, SCHEMA, SYSTEM, Problem, encoded, prompt
|
|
|
|
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}
|
|
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", model, "--effort", "medium", "--max-budget-usd", str(max_cost),
|
|
"--max-turns", "3", "--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_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:
|
|
os.killpg(process.pid, signal.SIGTERM)
|
|
try:
|
|
process.wait(timeout=2)
|
|
except subprocess.TimeoutExpired:
|
|
os.killpg(process.pid, signal.SIGKILL)
|
|
process.wait(timeout=5)
|
|
|
|
|
|
def parse_claude(raw, expected_model):
|
|
"""Normalize only a completed, uncompacted CLI result with a pinned model."""
|
|
final, initialized = None, False
|
|
aliases = {expected_model, expected_model + "[1m]"}
|
|
for line in raw.splitlines():
|
|
try:
|
|
event = json.loads(line)
|
|
except ValueError:
|
|
raise Problem("invalid_json_result", 502) from None
|
|
if not isinstance(event, dict):
|
|
raise Problem("invalid_json_result", 502)
|
|
if "compact" in str(event.get("subtype", "")):
|
|
raise Problem("compaction_detected", 502)
|
|
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"}):
|
|
raise Problem("worker_isolation_failed", 502)
|
|
if event.get("type") == "result":
|
|
final = event
|
|
if not initialized or not final:
|
|
raise Problem("incomplete_generation", 502)
|
|
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"
|
|
raise Problem(code, 502)
|
|
if final.get("stop_reason") in {"max_tokens", "model_context_window_exceeded"}:
|
|
raise Problem("incomplete_generation", 502)
|
|
models = final.get("modelUsage", {})
|
|
if not models or set(models) - aliases:
|
|
observed = [name if re.fullmatch(r"claude-[a-z0-9.-]+(?:\[1m\])?", name)
|
|
else "unrecognized" for name in models]
|
|
raise Problem("model_changed", 502, observed_models=observed)
|
|
for limits in models.values():
|
|
if (limits.get("contextWindow") != 1000000 or limits.get("maxOutputTokens") != 64000
|
|
or limits.get("canonicalModel") != expected_model
|
|
or limits.get("provider") != "firstParty"):
|
|
raise Problem("backend_capabilities_changed", 502)
|
|
usage = final.get("usage", {})
|
|
if any(usage.get("server_tool_use", {}).values()):
|
|
raise Problem("worker_isolation_failed", 502)
|
|
result = final.get("structured_output")
|
|
if not isinstance(result, dict):
|
|
raise Problem("invalid_json_result", 502)
|
|
return result, {"model": expected_model, "compaction": False,
|
|
"truncation": False, "compaction_signal": "CLI events and disabled compaction",
|
|
"usage": usage, "model_usage": models,
|
|
"duration_api_ms": final.get("duration_api_ms"),
|
|
"cost_usd_estimate": final.get("total_cost_usd"),
|
|
"turns": final.get("num_turns")}
|
|
|
|
|
|
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)
|
|
with (root / "input").open("rb") as source, (root / "output").open("wb") as output:
|
|
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)
|
|
deadline = time.monotonic() + request["execution"]["max_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)
|
|
finally:
|
|
stop(process)
|
|
if (root / "output").stat().st_size > OUTPUT_BYTES:
|
|
raise Problem("response_too_large", 502)
|
|
raw = (root / "output").read_text()
|
|
if process.returncode and not raw.strip():
|
|
raise Problem("backend_unavailable", 503)
|
|
result, metadata = parse_claude(raw, model)
|
|
if process.returncode:
|
|
raise Problem("incomplete_generation", 502)
|
|
metadata["temporary_files_deleted"] = True
|
|
return result, metadata
|