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

#!/usr/bin/env python3
"""Switchyard route selection and provider-health exclusion for CLI workers."""

from __future__ import annotations

import json
import re
import time
import urllib.error
import urllib.request
from typing import Any, Callable

from cli_lane_config import (
    EFFORTS,
    PROVIDER_AUTH_FAILURE_MAX_AGE_SECONDS,
    PROVIDER_HEALTH_MAX_AGE_SECONDS,
    PROVIDER_HEALTH_PATHS,
    SWITCHYARD_URL,
    Route,
)
from cli_lane_files import load_json


def parse_assignee(assignee: str) -> tuple[str | None, str | None]:
    """Parse external lane overrides while leaving cli-auto fully automatic."""
    value = str(assignee or "").strip().lower()
    if value == "cli-auto":
        return None, None
    match = re.fullmatch(r"cli-(codex|claude)-(low|medium|high|xhigh)", value)
    if not match:
        raise ValueError(f"unsupported external lane assignee: {assignee}")
    return match.group(1), match.group(2)

def _health_number(value: Any) -> float | None:
    """Return a plain numeric health field, rejecting booleans and strings."""
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        return None
    return float(value)

def _quota_reset_readmits(
    health: dict[str, Any], current: float, reset_timestamp: float | None
) -> bool:
    """A passed provider quota reset ends a capacity cooldown early."""
    if health.get("authenticated") is False:
        # An authentication failure is not repaired by a quota window reset.
        return False
    failed_at = _health_number(health.get("failed_at"))
    return (
        reset_timestamp is not None
        and failed_at is not None
        and failed_at < reset_timestamp <= current
    )

def fresh_unavailable_provider(
    now: float | None = None,
    quota_resets: dict[str, float] | None = None,
) -> str | None:
    """Return one recently proven-down provider for automatic route exclusion."""
    current = time.time() if now is None else now
    resets = quota_resets or {}
    unavailable: list[str] = []
    for provider, path in PROVIDER_HEALTH_PATHS.items():
        health = load_json(path)
        if health.get("state") not in ("unavailable", "capacity-limited"):
            continue
        cooldown_until = _health_number(health.get("cooldown_until"))
        if cooldown_until is not None:
            # Lane-recorded failures re-enter only after the full cooldown has
            # elapsed (hysteresis), a quota reset passed, or fresh success
            # rewrote the state above.
            if current >= cooldown_until or _quota_reset_readmits(
                health, current, resets.get(provider)
            ):
                continue
            unavailable.append(provider)
            continue
        try:
            age = current - path.stat().st_mtime
        except OSError:
            continue
        max_age = (
            PROVIDER_AUTH_FAILURE_MAX_AGE_SECONDS
            if health.get("authenticated") is False
            else PROVIDER_HEALTH_MAX_AGE_SECONDS
        )
        if 0 <= age <= max_age:
            unavailable.append(provider)
    # If both providers are down, keep the normal attempt/fallback path so the
    # card records authoritative current errors instead of trusting snapshots.
    return unavailable[0] if len(unavailable) == 1 else None

def _decode_worker_target(value: str) -> tuple[str, str, str]:
    """Decode the selected model header emitted by a worker decision route."""
    parts = value.split("/", 3)
    if len(parts) != 4 or parts[0] != "worker":
        raise RuntimeError(f"invalid Switchyard worker target: {value}")
    provider, model, effort = parts[1:]
    if provider not in {"codex", "claude"} or effort not in EFFORTS:
        raise RuntimeError(f"unsupported Switchyard worker target: {value}")
    return provider, model, effort

def select_route(
    prompt: str,
    assignee: str,
    *,
    exclude_provider: str | None = None,
    exclude_reason: str | None = None,
    switchyard_url: str = SWITCHYARD_URL,
    open_request: Callable[..., Any] = urllib.request.urlopen,
) -> Route:
    """Ask Switchyard to select one native CLI worker at this boundary."""
    started = time.monotonic()
    manual_provider, manual_effort = parse_assignee(assignee)
    if manual_provider and manual_effort:
        route_id = f"atlas/worker/manual/{manual_provider}/{manual_effort}"
        source = "switchyard-manual"
    else:
        route_id = "atlas/worker/auto/maximum"
        source = "switchyard-classifier"
    context = prompt
    if exclude_provider:
        reason = exclude_reason or "failed or exhausted capacity at this boundary"
        context += (
            f"\n\nRouting constraint: the {exclude_provider} provider {reason}. "
            "Do not select it."
        )
    payload = json.dumps(
        {
            "model": route_id,
            "messages": [{"role": "user", "content": context}],
            "stream": False,
            "max_tokens": 1,
        }
    ).encode("utf-8")
    request = urllib.request.Request(
        switchyard_url,
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with open_request(request, timeout=60) as response:
            selected = str(response.headers.get("x-model-router-selected-model") or "")
            rationale = str(response.headers.get("x-model-router-rationale") or "")
            response_body = response.read()
    except (OSError, urllib.error.URLError) as exc:
        raise RuntimeError(f"Switchyard worker routing failed: {exc}") from exc
    provider, model, effort = _decode_worker_target(selected)
    # Switchyard preserves the stable tier target in the selection header and
    # top-level response model. The worker broker's assistant content contains
    # the steward-resolved provider model required by the native CLI.
    try:
        response_document = json.loads(response_body)
        resolved_target = str(response_document["choices"][0]["message"]["content"] or "")
        resolved_provider, resolved_model, resolved_effort = _decode_worker_target(resolved_target)
        required_prefix = "gpt-" if provider == "codex" else "claude-"
        if (
            resolved_provider == provider
            and resolved_effort == effort
            and resolved_model.startswith(required_prefix)
        ):
            model = resolved_model
    except (AttributeError, IndexError, KeyError, RuntimeError, TypeError, ValueError, json.JSONDecodeError):
        pass
    if exclude_provider and provider == exclude_provider:
        alternate = "claude" if provider == "codex" else "codex"
        guarded = select_route(
            prompt,
            f"cli-{alternate}-{effort}",
            switchyard_url=switchyard_url,
            open_request=open_request,
        )
        return Route(
            provider=guarded.provider,
            model=guarded.model,
            effort=guarded.effort,
            profile=guarded.profile,
            classifier=f"{source}-health-guard",
            reason=(
                f"Switchyard selected excluded {exclude_provider}; preserved "
                f"its {effort} effort on healthy-provider route. {guarded.reason}"
            ),
            latency_ms=int((time.monotonic() - started) * 1000),
            fallback_chain=(),
        )
    return Route(
        provider=provider,
        model=model,
        effort=effort,
        profile=f"{provider}-{effort}",
        classifier=source,
        reason=rationale or f"Switchyard selected {selected}",
        latency_ms=int((time.monotonic() - started) * 1000),
        fallback_chain=(),
    )