308 lines
12 KiB
Python
308 lines
12 KiB
Python
#!/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
|
|
from routing_catalog import catalog_contains
|
|
|
|
_FALLBACK_RATIONALE = re.compile(
|
|
r"(?P<source>worker/(?:codex|claude)/auto(?:-(?:economy|balanced|advanced|frontier))?/(?:low|medium|high|xhigh)) "
|
|
r"(?:was unavailable|exceeded its context window); fell back to "
|
|
r"(?P<replacement>worker/(?:codex|claude)/auto(?:-(?:economy|balanced|advanced|frontier))?/(?:low|medium|high|xhigh))"
|
|
)
|
|
|
|
|
|
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)-(?:(?:economy|balanced|advanced|frontier)-)?(low|medium|high|xhigh)",
|
|
value,
|
|
)
|
|
if not match:
|
|
raise ValueError(f"unsupported external lane assignee: {assignee}")
|
|
return match.group(1), match.group(2)
|
|
|
|
|
|
def parse_assignee_capability(assignee: str) -> tuple[str | None, str | None, str | None]:
|
|
"""Return optional provider, effort, and capability lane constraints."""
|
|
value = str(assignee or "").strip().lower()
|
|
if value == "cli-auto":
|
|
return None, None, None
|
|
match = re.fullmatch(
|
|
r"cli-(codex|claude)-(?:(economy|balanced|advanced|frontier)-)?(low|medium|high|xhigh)",
|
|
value,
|
|
)
|
|
if not match:
|
|
raise ValueError(f"unsupported external lane assignee: {assignee}")
|
|
return match.group(1), match.group(3), 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 _target_capability(model: str, effort: str) -> str:
|
|
"""Read the stable capability selector from a worker target header."""
|
|
if model == "auto":
|
|
# Legacy AUTO encoded only the effort tier. Preserve its established
|
|
# mapping instead of treating a low/medium task as Sol-class work.
|
|
return {"low": "economy", "medium": "balanced"}.get(effort, "advanced")
|
|
match = re.fullmatch(r"auto-(economy|balanced|advanced|frontier)", model)
|
|
if match:
|
|
return match.group(1)
|
|
return "advanced"
|
|
|
|
|
|
def _fallback_floor(rationale: str) -> tuple[str, str, str] | None:
|
|
"""Return the originally selected worker floor from a Switchyard fallback."""
|
|
match = _FALLBACK_RATIONALE.search(rationale)
|
|
if not match:
|
|
return None
|
|
provider, model, effort = _decode_worker_target(match.group("source"))
|
|
return provider, _target_capability(model, effort), 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,
|
|
catalog_model_allowed: Callable[[str, str], bool] = catalog_contains,
|
|
) -> Route:
|
|
"""Ask Switchyard to select one native CLI worker at this boundary."""
|
|
started = time.monotonic()
|
|
manual_provider, manual_effort, manual_capability = parse_assignee_capability(assignee)
|
|
if manual_provider and manual_effort:
|
|
route_id = f"atlas/worker/manual/{manual_provider}/"
|
|
route_id += f"{manual_capability}/" if manual_capability else ""
|
|
route_id += 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)
|
|
capability = _target_capability(model, effort)
|
|
fallback_floor = _fallback_floor(rationale)
|
|
# 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.
|
|
resolved_valid = False
|
|
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-"
|
|
resolved_valid = (
|
|
resolved_provider == provider
|
|
and resolved_effort == effort
|
|
and (
|
|
resolved_model.startswith(required_prefix)
|
|
or (provider == "claude" and catalog_model_allowed(provider, resolved_model))
|
|
)
|
|
)
|
|
if resolved_valid:
|
|
model = resolved_model
|
|
except (AttributeError, IndexError, KeyError, RuntimeError, TypeError, ValueError, json.JSONDecodeError):
|
|
resolved_valid = False
|
|
if model.startswith("auto") and not resolved_valid:
|
|
raise RuntimeError("Switchyard did not return a current provider model")
|
|
if (
|
|
assignee == "cli-auto"
|
|
and fallback_floor
|
|
and (capability, effort) != fallback_floor[1:]
|
|
):
|
|
# FallThrough retries the route's entire target set after an upstream
|
|
# error. Its static list cannot retain the classifier's dynamic floor,
|
|
# so perform one exact alternate-provider selection before a lower role
|
|
# can become a CLI receipt.
|
|
failed_provider, required_capability, required_effort = fallback_floor
|
|
alternate = "claude" if failed_provider == "codex" else "codex"
|
|
guarded = select_route(
|
|
prompt,
|
|
f"cli-{alternate}-{required_capability}-{required_effort}",
|
|
switchyard_url=switchyard_url,
|
|
open_request=open_request,
|
|
catalog_model_allowed=catalog_model_allowed,
|
|
)
|
|
if (
|
|
guarded.provider == failed_provider
|
|
or guarded.capability != required_capability
|
|
or guarded.effort != required_effort
|
|
):
|
|
raise RuntimeError(
|
|
"Switchyard request fallback did not preserve the requested "
|
|
f"{required_capability}/{required_effort} route floor"
|
|
)
|
|
return Route(
|
|
provider=guarded.provider,
|
|
model=guarded.model,
|
|
effort=guarded.effort,
|
|
profile=guarded.profile,
|
|
classifier=f"{source}-fallback-floor-guard",
|
|
reason=(
|
|
f"Switchyard request fallback changed {required_capability}/"
|
|
f"{required_effort}; retried the alternate provider at the "
|
|
f"original floor. {guarded.reason}"
|
|
),
|
|
latency_ms=int((time.monotonic() - started) * 1000),
|
|
fallback_chain=(),
|
|
capability=guarded.capability,
|
|
)
|
|
if exclude_provider and provider == exclude_provider:
|
|
alternate = "claude" if provider == "codex" else "codex"
|
|
guarded = select_route(
|
|
prompt,
|
|
f"cli-{alternate}-{capability}-{effort}",
|
|
switchyard_url=switchyard_url,
|
|
open_request=open_request,
|
|
catalog_model_allowed=catalog_model_allowed,
|
|
)
|
|
if (
|
|
guarded.provider == exclude_provider
|
|
or guarded.capability != capability
|
|
or guarded.effort != effort
|
|
):
|
|
raise RuntimeError(
|
|
"Switchyard health guard did not preserve the requested "
|
|
f"{capability}/{effort} route floor"
|
|
)
|
|
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 {capability}/{effort} floor on a healthy-provider route. "
|
|
f"{guarded.reason}"
|
|
),
|
|
latency_ms=int((time.monotonic() - started) * 1000),
|
|
fallback_chain=(),
|
|
capability=guarded.capability,
|
|
)
|
|
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=(),
|
|
capability=capability,
|
|
)
|