atlas-iac/services/hermes/scripts/claude_oauth_broker.py
2026-08-16 09:19:23 -03:00

660 lines
24 KiB
Python

#!/usr/bin/env python3
"""Expose the owner's native Claude Code subscription as an Anthropic lane."""
from __future__ import annotations
import hmac
import json
import os
import re
import shutil
import subprocess
import threading
import time
import uuid
from datetime import datetime, timezone
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any, Final
from routing_catalog import load_catalog, resolve_route
HOST: Final = os.environ.get("HERMES_CLAUDE_BROKER_HOST", "0.0.0.0")
PORT: Final = int(os.environ.get("HERMES_CLAUDE_BROKER_PORT", "9006"))
CLAUDE_BIN: Final = os.environ.get("HERMES_CLAUDE_BIN", "/opt/coordinator/claude")
MAX_BODY_BYTES: Final = int(
os.environ.get("HERMES_CLAUDE_BROKER_MAX_BODY", str(64 << 20))
)
TIMEOUT_SECONDS: Final = float(
os.environ.get("HERMES_CLAUDE_BROKER_READ_TIMEOUT", "1800")
)
MAX_CONCURRENCY: Final = int(os.environ.get("HERMES_CLAUDE_BROKER_CONCURRENCY", "4"))
HEALTH_PATH: Final = Path(
os.environ.get("HERMES_CLAUDE_HEALTH_PATH", "/opt/data/provider-health/claude.json")
)
ROUTED_MODEL_PREFIX: Final = "route/claude/"
EFFORTS: Final = {"low", "medium", "high", "xhigh"}
CAPACITY_PATTERN: Final = re.compile(
r"(?:rate.?limit|capacity|overload|usage.?limit|quota|credit|exhaust|429|529)",
re.I,
)
STRUCTURED_SCHEMA: Final = {
"type": "object",
"additionalProperties": False,
"required": ["type", "text", "tool_calls"],
"properties": {
"type": {"type": "string", "enum": ["final", "tool_calls"]},
"text": {"type": "string"},
"tool_calls": {
"type": "array",
"items": {
"type": "object",
"additionalProperties": False,
"required": ["name", "input"],
"properties": {
"name": {"type": "string"},
"input": {"type": "object"},
},
},
},
},
}
_slots = threading.BoundedSemaphore(MAX_CONCURRENCY)
_auth_probe_lock = threading.Lock()
_auth_probe_at = 0.0
_auth_probe_value: dict[str, Any] = {}
HEALTH_POLL_SECONDS: Final = max(
30, int(os.environ.get("HERMES_CLAUDE_HEALTH_POLL_SECONDS", "60"))
)
def _read_secret(file_env_name: str) -> str:
"""Read a secret only from its runtime-mounted file."""
path = os.environ.get(file_env_name, "").strip()
if not path:
return ""
try:
return Path(path).read_text(encoding="utf-8").strip()
except OSError:
return ""
def _relay_key() -> str:
"""Return the internal relay key shared with Switchyard."""
return _read_secret("HERMES_CLAUDE_BROKER_KEY_FILE")
def _authorized(authorization: str | None, api_key: str | None) -> bool:
"""Accept Switchyard's API-key or bearer-key header."""
expected = _relay_key()
if not expected:
return False
candidates = [api_key or ""]
if authorization and authorization.startswith("Bearer "):
candidates.append(authorization[7:].strip())
return any(
candidate and hmac.compare_digest(candidate, expected)
for candidate in candidates
)
def _route(model: str, payload: dict[str, Any]) -> tuple[str, str]:
"""Resolve one Switchyard model and effort into native Claude CLI values."""
effort = "medium"
if model.startswith(ROUTED_MODEL_PREFIX):
parts = model.split("/")
if parts[-1] in EFFORTS:
effort = parts[-1]
model = resolve_route(model)
output_config = payload.get("output_config")
if isinstance(output_config, dict) and output_config.get("effort") in EFFORTS:
effort = str(output_config["effort"])
if not model.startswith("claude-"):
raise ValueError("unsupported Claude model")
return model, effort
def _prompt(payload: dict[str, Any]) -> str:
"""Describe one Anthropic boundary without letting Claude run local tools."""
tools = payload.get("tools")
tools = tools if isinstance(tools, list) else []
contract = {
"system": payload.get("system") or "",
"messages": payload.get("messages") or [],
"tools": tools,
"tool_choice": payload.get("tool_choice") or {"type": "auto"},
}
return (
"You are serving one model boundary for Hermes. The JSON below is the "
"complete conversation and the only source of task context. Do not run "
"Claude Code tools or modify files yourself. Every tool listed in the "
"contract is an available external Hermes tool. If the requested work "
"needs one, return type=tool_calls with its exact name and a valid input "
"object; Hermes executes it after this boundary. Never simulate a tool "
"result or claim that a listed tool succeeded, failed, or is unavailable. "
"Do not return a plan or blocker when an available tool can advance the "
"request. Otherwise return type=final and place the complete user-facing "
"answer in text. Do not describe this envelope.\n\n"
+ json.dumps(contract, ensure_ascii=False, separators=(",", ":"))
)
def _tool_names(payload: dict[str, Any]) -> set[str]:
"""Return the external Hermes tool names advertised at this boundary."""
tools = payload.get("tools")
if not isinstance(tools, list):
return set()
return {
str(tool["name"])
for tool in tools
if isinstance(tool, dict) and isinstance(tool.get("name"), str)
}
def _latest_user_text(payload: dict[str, Any]) -> str:
"""Return the newest user text for false-positive-safe result validation."""
messages = payload.get("messages")
if not isinstance(messages, list):
return ""
for message in reversed(messages):
if not isinstance(message, dict) or message.get("role") != "user":
continue
content = message.get("content")
if isinstance(content, str):
return content
if isinstance(content, list):
return "\n".join(
str(block.get("text") or "")
for block in content
if isinstance(block, dict)
)
return ""
def _validate_structured(structured: dict[str, Any], payload: dict[str, Any]) -> None:
"""Reject fabricated tool outcomes so Switchyard can try another lane."""
response_type = structured.get("type")
tool_calls = structured.get("tool_calls")
if response_type not in {"final", "tool_calls"} or not isinstance(tool_calls, list):
raise RuntimeError("provider: Claude CLI returned an invalid result envelope")
available = _tool_names(payload)
for call in tool_calls:
if (
not isinstance(call, dict)
or call.get("name") not in available
or not isinstance(call.get("input"), dict)
):
raise RuntimeError("provider: Claude CLI returned an unavailable tool call")
if response_type == "tool_calls" and not tool_calls:
raise RuntimeError("provider: Claude CLI returned an empty tool request")
if response_type == "final" and tool_calls:
raise RuntimeError("provider: Claude CLI mixed a final answer with tool calls")
text = str(structured.get("text") or "")
latest_user = _latest_user_text(payload).lower()
lowered = text.lower()
unavailable_claims = (
"no such tool available",
"tools that worked earlier",
"not reachable at this boundary",
)
claims_tool_failure = any(claim in lowered for claim in unavailable_claims)
user_is_quoting_failure = any(claim in latest_user for claim in unavailable_claims)
names_tool = any(
re.search(rf"\b{re.escape(name.lower())}\b", lowered) for name in available
)
if (
response_type == "final"
and claims_tool_failure
and names_tool
and not user_is_quoting_failure
):
raise RuntimeError(
"provider: Claude CLI fabricated an external tool availability failure"
)
def _usage(event: dict[str, Any]) -> dict[str, int]:
"""Translate Claude Code's result accounting into Anthropic token fields."""
raw = event.get("usage")
raw = raw if isinstance(raw, dict) else {}
return {
"input_tokens": max(0, int(raw.get("input_tokens") or 0)),
"output_tokens": max(0, int(raw.get("output_tokens") or 0)),
"cache_creation_input_tokens": max(
0, int(raw.get("cache_creation_input_tokens") or 0)
),
"cache_read_input_tokens": max(0, int(raw.get("cache_read_input_tokens") or 0)),
}
def _atomic_health(value: dict[str, Any]) -> None:
"""Persist non-secret subscription/transport health for the owner dashboard."""
try:
HEALTH_PATH.parent.mkdir(parents=True, exist_ok=True)
temporary = HEALTH_PATH.with_name(f".{HEALTH_PATH.name}.{os.getpid()}.tmp")
temporary.write_text(
json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8"
)
os.replace(temporary, HEALTH_PATH)
except OSError:
pass
def _previous_health() -> dict[str, Any]:
"""Read the prior non-secret health snapshot when it is still valid JSON."""
try:
value = json.loads(HEALTH_PATH.read_text(encoding="utf-8"))
except (OSError, TypeError, ValueError, json.JSONDecodeError):
return {}
return value if isinstance(value, dict) else {}
def _claude_environment() -> dict[str, str]:
"""Return an environment that cannot silently select metered API billing."""
environment = os.environ.copy()
environment.pop("ANTHROPIC_API_KEY", None)
environment.pop("CLAUDE_API_KEY", None)
return environment
def _subscription_health(force: bool = False) -> dict[str, Any]:
"""Probe the native CLI login and cache the non-secret account result."""
global _auth_probe_at, _auth_probe_value
now = time.monotonic()
with _auth_probe_lock:
if not force and _auth_probe_value and now - _auth_probe_at < 60:
return dict(_auth_probe_value)
checked_at = datetime.now(timezone.utc).isoformat()
try:
completed = subprocess.run(
[CLAUDE_BIN, "auth", "status"],
text=True,
capture_output=True,
timeout=15,
env=_claude_environment(),
check=False,
)
raw = json.loads(completed.stdout) if completed.stdout.strip() else {}
except (OSError, subprocess.SubprocessError, ValueError, json.JSONDecodeError):
completed = None
raw = {}
authenticated = bool(
completed
and completed.returncode == 0
and isinstance(raw, dict)
and raw.get("loggedIn") is True
and raw.get("apiProvider") == "firstParty"
)
previous = _previous_health()
value = {
"transport": "claude-code-cli-subscription",
"state": "available" if authenticated else "unavailable",
"checked_at": checked_at,
"authenticated": authenticated,
"api_provider": raw.get("apiProvider") if isinstance(raw, dict) else None,
"auth_method": raw.get("authMethod") if isinstance(raw, dict) else None,
"subscription_type": raw.get("subscriptionType")
if isinstance(raw, dict)
else None,
}
# Readiness probes must not erase the most recently observed native
# usage window or successful route metadata. They only refresh auth.
for key in (
"rate_limit",
"last_success_at",
"last_error_at",
"latency_ms",
"model",
"effort",
):
if key in previous:
value[key] = previous[key]
_auth_probe_at = now
_auth_probe_value = value
_atomic_health(value)
return dict(value)
def _health_polling_loop() -> None:
"""Keep native auth health fresh even when no request reaches the broker."""
while True:
time.sleep(HEALTH_POLL_SECONDS)
_subscription_health(force=True)
def _invoke(payload: dict[str, Any]) -> tuple[dict[str, Any], dict[str, int], str]:
"""Run one native, first-party Claude Code subscription request."""
requested = payload.get("model")
if not isinstance(requested, str):
raise ValueError("Claude model is required")
model, effort = _route(requested, payload)
command = [
CLAUDE_BIN,
"-p",
"--output-format",
"stream-json",
"--verbose",
"--no-session-persistence",
"--tools",
"",
"--model",
model,
"--effort",
effort,
"--json-schema",
json.dumps(STRUCTURED_SCHEMA, separators=(",", ":")),
]
environment = _claude_environment()
started = time.monotonic()
with _slots:
completed = subprocess.run(
command,
input=_prompt(payload),
text=True,
capture_output=True,
timeout=TIMEOUT_SECONDS,
env=environment,
check=False,
)
result_event: dict[str, Any] = {}
rate_limit: dict[str, Any] = {}
for line in completed.stdout.splitlines():
try:
event = json.loads(line)
except (TypeError, ValueError, json.JSONDecodeError):
continue
if not isinstance(event, dict):
continue
if event.get("type") == "rate_limit_event":
raw = event.get("rate_limit_info")
if isinstance(raw, dict):
rate_limit = raw
if event.get("type") == "result":
result_event = event
error_text = "\n".join(
value for value in (completed.stderr.strip(), completed.stdout[-8000:]) if value
)
if completed.returncode or result_event.get("is_error"):
health = _subscription_health()
health.update(
{
"state": "capacity-limited"
if CAPACITY_PATTERN.search(error_text)
else "unavailable",
"last_error_at": datetime.now(timezone.utc).isoformat(),
"model": model,
"effort": effort,
"rate_limit": rate_limit,
}
)
_atomic_health(health)
kind = "capacity" if CAPACITY_PATTERN.search(error_text) else "provider"
raise RuntimeError(f"{kind}: {error_text[-1200:] or 'Claude CLI failed'}")
structured = result_event.get("structured_output")
if not isinstance(structured, dict):
raw_result = result_event.get("result")
try:
structured = json.loads(raw_result) if isinstance(raw_result, str) else None
except (TypeError, ValueError, json.JSONDecodeError):
structured = None
if not isinstance(structured, dict):
raise RuntimeError("provider: Claude CLI returned no structured result")
_validate_structured(structured, payload)
usage = _usage(result_event)
actual_model = str(
next(iter(result_event.get("modelUsage") or {}), model)
if isinstance(result_event.get("modelUsage"), dict)
else model
)
health = _subscription_health()
health.update(
{
"state": "available",
"last_success_at": datetime.now(timezone.utc).isoformat(),
"latency_ms": int((time.monotonic() - started) * 1000),
"model": actual_model,
"effort": effort,
"rate_limit": rate_limit,
}
)
_atomic_health(health)
return structured, usage, actual_model
def _message(
structured: dict[str, Any], usage: dict[str, int], model: str
) -> dict[str, Any]:
"""Build one Anthropic Messages response from the structured CLI result."""
content: list[dict[str, Any]] = []
text = structured.get("text")
if isinstance(text, str) and text:
content.append({"type": "text", "text": text})
tool_calls = structured.get("tool_calls")
if isinstance(tool_calls, list):
for raw in tool_calls:
if not isinstance(raw, dict) or not isinstance(raw.get("name"), str):
continue
tool_input = raw.get("input")
content.append(
{
"type": "tool_use",
"id": f"toolu_{uuid.uuid4().hex}",
"name": raw["name"],
"input": tool_input if isinstance(tool_input, dict) else {},
}
)
if not content:
content.append({"type": "text", "text": ""})
return {
"id": f"msg_{uuid.uuid4().hex}",
"type": "message",
"role": "assistant",
"model": model,
"content": content,
"stop_reason": "tool_use"
if any(item["type"] == "tool_use" for item in content)
else "end_turn",
"stop_sequence": None,
"usage": usage,
}
def _sse(message: dict[str, Any]) -> bytes:
"""Encode a completed message as a standards-compliant Anthropic SSE stream."""
events: list[tuple[str, dict[str, Any]]] = []
opening = dict(message)
opening["content"] = []
opening["stop_reason"] = None
opening["usage"] = {
"input_tokens": message["usage"]["input_tokens"],
"output_tokens": 0,
}
events.append(("message_start", {"type": "message_start", "message": opening}))
for index, block in enumerate(message["content"]):
if block["type"] == "text":
start = {"type": "text", "text": ""}
delta = {"type": "text_delta", "text": block["text"]}
else:
start = {key: block[key] for key in ("type", "id", "name")}
start["input"] = {}
delta = {
"type": "input_json_delta",
"partial_json": json.dumps(block["input"], separators=(",", ":")),
}
events.extend(
[
(
"content_block_start",
{
"type": "content_block_start",
"index": index,
"content_block": start,
},
),
(
"content_block_delta",
{
"type": "content_block_delta",
"index": index,
"delta": delta,
},
),
("content_block_stop", {"type": "content_block_stop", "index": index}),
]
)
events.append(
(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": message["stop_reason"], "stop_sequence": None},
"usage": {"output_tokens": message["usage"]["output_tokens"]},
},
)
)
events.append(("message_stop", {"type": "message_stop"}))
return "".join(
f"event: {name}\ndata: {json.dumps(value, separators=(',', ':'))}\n\n"
for name, value in events
).encode("utf-8")
class Handler(BaseHTTPRequestHandler):
"""Serve the subset of Anthropic Messages used by Switchyard."""
server_version = "HermesClaudeCodeBroker/2"
def log_message(self, format: str, *args: object) -> None:
return
def _json(self, status: int, value: dict[str, Any]) -> None:
body = json.dumps(value, separators=(",", ":")).encode("utf-8")
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.send_header("Cache-Control", "no-store")
self.end_headers()
try:
self.wfile.write(body)
except (BrokenPipeError, ConnectionResetError):
# Kubernetes and callers may close a health/request socket after
# their own deadline. The broker is still healthy and must not
# emit a noisy handler traceback for that normal disconnect.
return
def _error(self, status: int, error_type: str, message: str) -> None:
self._json(
status,
{"type": "error", "error": {"type": error_type, "message": message}},
)
def _check_auth(self) -> bool:
if _authorized(
self.headers.get("Authorization"), self.headers.get("x-api-key")
):
return True
self._error(401, "authentication_error", "unauthorized")
return False
def do_GET(self) -> None: # noqa: N802
if self.path == "/health":
executable = shutil.which(CLAUDE_BIN) or (
CLAUDE_BIN if Path(CLAUDE_BIN).is_file() else ""
)
health = _subscription_health()
status = 200 if executable and health.get("authenticated") else 503
self._json(
status,
{
"ok": status == 200,
"provider": "claude-code-subscription",
"subscription_type": health.get("subscription_type"),
"auth_method": health.get("auth_method"),
},
)
return
if self.path != "/v1/models":
self._error(404, "not_found", "not found")
return
if not self._check_auth():
return
catalog = load_catalog()
providers = catalog.get("providers", {})
claude = providers.get("claude", {}) if isinstance(providers, dict) else {}
raw_models = claude.get("models", []) if isinstance(claude, dict) else []
models = sorted(
model
for model in raw_models
if isinstance(model, str) and model.startswith("claude-")
)
self._json(200, {"data": [{"id": model, "type": "model"} for model in models]})
def do_POST(self) -> None: # noqa: N802
if self.path not in {"/v1/messages", "/v1/messages/count_tokens"}:
self._error(404, "not_found", "not found")
return
if not self._check_auth():
return
try:
length = int(self.headers.get("Content-Length", "0"))
except ValueError:
length = -1
if length < 0 or length > MAX_BODY_BYTES:
self._error(413, "request_too_large", "request too large")
return
try:
payload = json.loads(self.rfile.read(length))
except (TypeError, ValueError, json.JSONDecodeError):
self._error(400, "invalid_request_error", "JSON object required")
return
if not isinstance(payload, dict):
self._error(400, "invalid_request_error", "JSON object required")
return
if self.path.endswith("count_tokens"):
self._json(200, {"input_tokens": max(1, len(_prompt(payload)) // 4)})
return
try:
structured, usage, model = _invoke(payload)
message = _message(structured, usage, model)
except ValueError as exc:
self._error(400, "invalid_request_error", str(exc))
return
except subprocess.TimeoutExpired:
self._error(504, "timeout_error", "Claude Code request timed out")
return
except (OSError, RuntimeError) as exc:
detail = str(exc)
status = 429 if detail.startswith("capacity:") else 503
self._error(
status,
"rate_limit_error" if status == 429 else "api_error",
detail.partition(": ")[2] or detail,
)
return
if payload.get("stream"):
body = _sse(message)
self.send_response(200)
self.send_header("Content-Type", "text/event-stream")
self.send_header("Cache-Control", "no-store")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
try:
self.wfile.write(body)
except (BrokenPipeError, ConnectionResetError):
return
else:
self._json(200, message)
class Server(ThreadingHTTPServer):
"""Threaded broker with bounded provider-side concurrency."""
daemon_threads = True
if __name__ == "__main__":
_subscription_health(force=True)
threading.Thread(target=_health_polling_loop, daemon=True).start()
Server((HOST, PORT), Handler).serve_forever()