Blame for core/provider_availability.py

# -*- coding: utf-8 -*-
"""Provider-neutral availability taxonomy and benchmark-only health tracker.

A3-R5G scope is deliberately narrow: this module has no network calls, no
production route activation, no Worker/D1 access, and no product entitlement
logic.  Provider/model/endpoints and all retry/circuit thresholds remain
versioned policy data supplied by callers.
"""
from __future__ import annotations

from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from email.utils import parsedate_to_datetime
import re
import threading
import time
from typing import Any, Callable, Mapping


ERROR_KINDS = {
    "AUTHENTICATION",
    "PERMISSION",
    "BALANCE",
    "RATE_LIMIT",
    "OVERLOAD",
    "QUOTA_EXHAUSTED",
    "TIMEOUT",
    "CONNECTION",
    "SERVER_ERROR",
    "REQUEST_INVALID",
    "CANCELLED",
    "UNKNOWN",
}

AVAILABILITY_STATES = {
    "UNKNOWN",
    "HEALTHY",
    "DEGRADED",
    "RATE_LIMITED",
    "OVERLOADED",
    "UNREACHABLE",
    "AUTH_BLOCKED",
    "BALANCE_BLOCKED",
    "CIRCUIT_OPEN",
    "DISABLED",
}

CIRCUIT_STATES = {"CLOSED", "OPEN", "HALF_OPEN"}


@dataclass(frozen=True)
class HealthKey:
    provider_profile_id: str
    endpoint_profile_id: str
    runtime_egress_profile: str

    def as_string(self) -> str:
        return "::".join(
            (
                str(self.provider_profile_id or "").strip(),
                str(self.endpoint_profile_id or "").strip(),
                str(self.runtime_egress_profile or "").strip(),
            )
        )


@dataclass(frozen=True)
class ProviderErrorSignal:
    provider: str
    error_kind: str
    availability_state: str
    retryable: bool
    http_status: int | None = None
    provider_code: str = ""
    retry_after_seconds: float | None = None
    safe_reason_code: str = ""
    penalizes_health: bool = False

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True)
class HealthSnapshot:
    availability_state: str = "UNKNOWN"
    circuit_state: str = "CLOSED"
    eligible: bool = True
    recent_successes: int = 0
    recent_failures: int = 0
    success_rate: float = 0.0
    error_rate: float = 0.0
    latency_ms: float = 0.0
    blocked_until: float | None = None
    last_error_kind: str = ""
    last_provider_code: str = ""
    cause_state: str = "UNKNOWN"

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass
class _HealthRecord:
    availability_state: str = "UNKNOWN"
    circuit_state: str = "CLOSED"
    blocked_until: float | None = None
    last_error_kind: str = ""
    last_provider_code: str = ""
    consecutive_overloads: int = 0
    events: list[tuple[float, bool, bool, float]] = field(default_factory=list)


def _coerce_status(value: Any) -> int | None:
    try:
        number = int(value)
    except (TypeError, ValueError):
        return None
    return number if number > 0 else None


def _coerce_positive_float(value: Any) -> float | None:
    try:
        number = float(value)
    except (TypeError, ValueError):
        return None
    return number if number >= 0 else None


def _header_value(headers: Any, name: str) -> str:
    if not isinstance(headers, Mapping):
        return ""
    target = name.lower()
    for key, value in headers.items():
        if str(key).lower() == target:
            return str(value or "").strip()
    return ""


def _parse_retry_after(value: Any, *, now_epoch: float | None = None) -> float | None:
    text = str(value or "").strip()
    if not text:
        return None
    numeric = _coerce_positive_float(text)
    if numeric is not None:
        return numeric
    try:
        when = parsedate_to_datetime(text)
        if when.tzinfo is None:
            when = when.replace(tzinfo=timezone.utc)
        now = datetime.fromtimestamp(now_epoch if now_epoch is not None else time.time(), tz=timezone.utc)
        return max(0.0, (when.astimezone(timezone.utc) - now).total_seconds())
    except Exception:
        return None


def _parse_duration_seconds(value: Any) -> float | None:
    text = str(value or "").strip().lower()
    match = re.fullmatch(r"([0-9]+(?:\.[0-9]+)?)s", text)
    if match:
        return float(match.group(1))
    return _coerce_positive_float(text)


def safe_response_metadata(response: Any, *, provider: str) -> dict[str, Any]:
    """Extract only non-sensitive operational metadata from a provider response."""
    status = _coerce_status(getattr(response, "status_code", None))
    retry_after = _parse_retry_after(
        _header_value(getattr(response, "headers", None), "Retry-After")
    )
    provider_code = ""
    provider_reason = ""
    try:
        payload = response.json()
    except Exception:
        payload = None

    if isinstance(payload, Mapping):
        error = payload.get("error")
        error_map = error if isinstance(error, Mapping) else payload
        raw_code = error_map.get("code") if isinstance(error_map, Mapping) else None
        if raw_code is None:
            raw_code = payload.get("code")
        if raw_code is not None:
            provider_code = str(raw_code).strip()[:80]
        if isinstance(error_map, Mapping):
            raw_reason = error_map.get("status") or error_map.get("type") or error_map.get("reason")
            if raw_reason:
                provider_reason = str(raw_reason).strip()[:120]
            if retry_after is None:
                details = error_map.get("details")
                if isinstance(details, list):
                    for detail in details:
                        if not isinstance(detail, Mapping):
                            continue
                        retry_after = _parse_duration_seconds(
                            detail.get("retryDelay") or detail.get("retry_delay")
                        )
                        if retry_after is not None:
                            break

    return {
        "provider": str(provider or "").strip().lower(),
        "http_status": status,
        "provider_code": provider_code,
        "provider_reason": provider_reason,
        "retry_after_seconds": retry_after,
    }


def attach_safe_error_metadata(error: BaseException, response: Any, *, provider: str) -> BaseException:
    metadata = safe_response_metadata(response, provider=provider)
    for key, value in metadata.items():
        try:
            setattr(error, key, value)
        except Exception:
            pass
    return error


def _signal(
    *,
    provider: str,
    error_kind: str,
    availability_state: str,
    retryable: bool,
    http_status: int | None,
    provider_code: str,
    retry_after_seconds: float | None,
    safe_reason_code: str,
    penalizes_health: bool,
) -> ProviderErrorSignal:
    kind = str(error_kind or "UNKNOWN").upper()
    state = str(availability_state or "UNKNOWN").upper()
    if kind not in ERROR_KINDS:
        kind = "UNKNOWN"
    if state not in AVAILABILITY_STATES:
        state = "UNKNOWN"
    return ProviderErrorSignal(
        provider=str(provider or "").strip().lower(),
        error_kind=kind,
        availability_state=state,
        retryable=bool(retryable),
        http_status=_coerce_status(http_status),
        provider_code=str(provider_code or "").strip()[:80],
        retry_after_seconds=_coerce_positive_float(retry_after_seconds),
        safe_reason_code=str(safe_reason_code or "").strip()[:120],
        penalizes_health=bool(penalizes_health),
    )


def classify_provider_error(
    error: BaseException | None = None,
    *,
    provider: str = "",
    http_status: int | None = None,
    provider_code: Any = None,
    provider_reason: Any = None,
    retry_after_seconds: float | None = None,
) -> ProviderErrorSignal:
    """Normalize provider/network failures without retaining response bodies."""
    provider_name = str(provider or getattr(error, "provider", "") or "").strip().lower()
    status = _coerce_status(http_status if http_status is not None else getattr(error, "http_status", None))
    code = str(
        provider_code if provider_code is not None else getattr(error, "provider_code", "") or ""
    ).strip()
    reason = str(
        provider_reason if provider_reason is not None else getattr(error, "provider_reason", "") or ""
    ).strip()
    retry_after = _coerce_positive_float(
        retry_after_seconds if retry_after_seconds is not None else getattr(error, "retry_after_seconds", None)
    )
    class_name = error.__class__.__name__.lower() if error is not None else ""
    safe_reason = (
        reason if provider_name == "gemini" and reason
        else code or reason or (f"http_{status}" if status else class_name or "unknown")
    )

    # Provider-specific business codes take precedence over outer HTTP status.
    if provider_name == "zhipu":
        if code in {"1000", "1001", "1003"}:
            return _signal(provider=provider_name, error_kind="AUTHENTICATION", availability_state="AUTH_BLOCKED", retryable=False, http_status=status, provider_code=code, retry_after_seconds=retry_after, safe_reason_code=safe_reason, penalizes_health=False)
        if code == "1113":
            return _signal(provider=provider_name, error_kind="BALANCE", availability_state="BALANCE_BLOCKED", retryable=False, http_status=status, provider_code=code, retry_after_seconds=retry_after, safe_reason_code=safe_reason, penalizes_health=False)
        if code == "1302":
            return _signal(provider=provider_name, error_kind="RATE_LIMIT", availability_state="RATE_LIMITED", retryable=True, http_status=status, provider_code=code, retry_after_seconds=retry_after, safe_reason_code=safe_reason, penalizes_health=True)
        if code == "1305":
            return _signal(provider=provider_name, error_kind="OVERLOAD", availability_state="OVERLOADED", retryable=True, http_status=status, provider_code=code, retry_after_seconds=retry_after, safe_reason_code=safe_reason, penalizes_health=True)
        if code == "1308":
            return _signal(provider=provider_name, error_kind="QUOTA_EXHAUSTED", availability_state="RATE_LIMITED", retryable=False, http_status=status, provider_code=code, retry_after_seconds=retry_after, safe_reason_code=safe_reason, penalizes_health=False)

    if "cancel" in class_name:
        return _signal(provider=provider_name, error_kind="CANCELLED", availability_state="UNKNOWN", retryable=False, http_status=status, provider_code=code, retry_after_seconds=retry_after, safe_reason_code=safe_reason, penalizes_health=False)
    if "timeout" in class_name or isinstance(error, TimeoutError):
        return _signal(provider=provider_name, error_kind="TIMEOUT", availability_state="DEGRADED", retryable=True, http_status=status, provider_code=code, retry_after_seconds=retry_after, safe_reason_code=safe_reason, penalizes_health=True)
    if "connection" in class_name or "chunked" in class_name:
        return _signal(provider=provider_name, error_kind="CONNECTION", availability_state="DEGRADED", retryable=True, http_status=status, provider_code=code, retry_after_seconds=retry_after, safe_reason_code=safe_reason, penalizes_health=True)

    if status == 401:
        kind, state, retryable, penalty = "AUTHENTICATION", "AUTH_BLOCKED", False, False
    elif status == 403:
        kind, state, retryable, penalty = "PERMISSION", "AUTH_BLOCKED", False, False
    elif status == 402:
        kind, state, retryable, penalty = "BALANCE", "BALANCE_BLOCKED", False, False
    elif status == 429:
        kind, state, retryable, penalty = "RATE_LIMIT", "RATE_LIMITED", True, True
    elif status == 503:
        kind, state, retryable, penalty = "OVERLOAD", "OVERLOADED", True, True
    elif status == 504:
        kind, state, retryable, penalty = "TIMEOUT", "DEGRADED", True, True
    elif status is not None and status >= 500:
        kind, state, retryable, penalty = "SERVER_ERROR", "DEGRADED", True, True
    elif status is not None and 400 <= status < 500:
        kind, state, retryable, penalty = "REQUEST_INVALID", "UNKNOWN", False, False
    elif "authentication" in class_name or "permission" in class_name:
        kind, state, retryable, penalty = "AUTHENTICATION", "AUTH_BLOCKED", False, False
    elif "ratelimit" in class_name or "rate_limit" in class_name:
        kind, state, retryable, penalty = "RATE_LIMIT", "RATE_LIMITED", True, True
    elif "request" in class_name:
        kind, state, retryable, penalty = "REQUEST_INVALID", "UNKNOWN", False, False
    else:
        kind, state, retryable, penalty = "UNKNOWN", "UNKNOWN", False, False

    return _signal(
        provider=provider_name,
        error_kind=kind,
        availability_state=state,
        retryable=retryable,
        http_status=status,
        provider_code=code,
        retry_after_seconds=retry_after,
        safe_reason_code=safe_reason,
        penalizes_health=penalty,
    )


def snapshot_from_legacy_health(
    value: Any,
    *,
    latency_degraded_threshold_ms: float = 0.0,
) -> HealthSnapshot:
    """Compatibility bridge for the pre-R5G benchmark health shape."""
    if isinstance(value, Mapping):
        available = bool(value.get("available", False))
        rate_limited = bool(value.get("rate_limited", False))
        latency_ms = float(value.get("latency_ms") or 0.0)
        error_rate = max(0.0, min(float(value.get("error_rate") or 0.0), 1.0))
    else:
        available = bool(value)
        rate_limited = False
        latency_ms = 0.0
        error_rate = 0.0

    if not available:
        state, eligible = "UNREACHABLE", False
    elif rate_limited:
        state, eligible = "RATE_LIMITED", False
    elif latency_degraded_threshold_ms > 0 and latency_ms > latency_degraded_threshold_ms:
        state, eligible = "DEGRADED", True
    else:
        state, eligible = "HEALTHY", True
    return HealthSnapshot(
        availability_state=state,
        circuit_state="CLOSED",
        eligible=eligible,
        recent_successes=1 if available else 0,
        recent_failures=0 if available else 1,
        success_rate=1.0 if available else 0.0,
        error_rate=error_rate,
        latency_ms=latency_ms,
        cause_state=state,
    )


class ProviderAvailabilityTracker:
    """Thread-safe, in-memory availability tracker for benchmark/testing only."""

    def __init__(self, policy: Mapping[str, Any], *, clock: Callable[[], float] | None = None):
        self.policy = dict(policy or {})
        self.clock = clock or time.monotonic
        self._lock = threading.RLock()
        self._records: dict[HealthKey, _HealthRecord] = {}

    def _now(self) -> float:
        return float(self.clock())

    def _record(self, key: HealthKey) -> _HealthRecord:
        return self._records.setdefault(key, _HealthRecord())

    def _window_seconds(self) -> float:
        return max(1.0, float(self.policy.get("failure_window_seconds") or 1.0))

    def _failure_threshold(self) -> int:
        return max(1, int(self.policy.get("failure_threshold") or 1))

    def _latency_threshold(self) -> float:
        return max(0.0, float(self.policy.get("latency_degraded_threshold_ms") or 0.0))

    def _prune(self, record: _HealthRecord, now: float) -> None:
        cutoff = now - self._window_seconds()
        record.events[:] = [row for row in record.events if row[0] >= cutoff]

    def _penalizing_failures(self, record: _HealthRecord) -> int:
        return sum(1 for _, success, penalty, _ in record.events if not success and penalty)

    def _cooldown(self, record: _HealthRecord, signal: ProviderErrorSignal) -> float | None:
        if signal.retry_after_seconds is not None:
            return max(0.0, float(signal.retry_after_seconds))
        if signal.error_kind == "OVERLOAD":
            schedule = [max(0.0, float(v)) for v in self.policy.get("overload_cooldown_schedule") or []]
            if schedule:
                index = min(max(record.consecutive_overloads - 1, 0), len(schedule) - 1)
                return schedule[index]
        if signal.error_kind in {"RATE_LIMIT", "QUOTA_EXHAUSTED"}:
            backoff = self.policy.get("rate_limit_backoff_policy")
            if isinstance(backoff, Mapping):
                base = max(0.0, float(backoff.get("base_seconds") or 0.0))
                cap = max(base, float(backoff.get("max_seconds") or base))
                attempts = max(1, self._penalizing_failures(record))
                return min(cap, base * (2 ** max(0, attempts - 1)))
        return None

    def _refresh(self, record: _HealthRecord, now: float) -> None:
        self._prune(record, now)
        if record.circuit_state == "OPEN" and record.blocked_until is not None and now >= record.blocked_until:
            record.circuit_state = "HALF_OPEN"
            record.availability_state = "DEGRADED"
            record.blocked_until = None
        elif (
            record.circuit_state == "CLOSED"
            and record.availability_state in {"RATE_LIMITED", "OVERLOADED"}
            and record.blocked_until is not None
            and now >= record.blocked_until
        ):
            record.availability_state = "DEGRADED"
            record.blocked_until = None

    def record_success(self, key: HealthKey, *, latency_ms: float = 0.0) -> HealthSnapshot:
        with self._lock:
            now = self._now()
            record = self._record(key)
            self._refresh(record, now)
            latency = max(0.0, float(latency_ms or 0.0))
            record.events.append((now, True, False, latency))
            record.circuit_state = "CLOSED"
            record.blocked_until = None
            record.last_error_kind = ""
            record.last_provider_code = ""
            record.consecutive_overloads = 0
            threshold = self._latency_threshold()
            record.availability_state = "DEGRADED" if threshold > 0 and latency > threshold else "HEALTHY"
            return self._snapshot_locked(record, now)

    def record_failure(
        self,
        key: HealthKey,
        signal: ProviderErrorSignal,
        *,
        latency_ms: float = 0.0,
    ) -> HealthSnapshot:
        with self._lock:
            now = self._now()
            record = self._record(key)
            self._refresh(record, now)
            latency = max(0.0, float(latency_ms or 0.0))
            record.events.append((now, False, bool(signal.penalizes_health), latency))
            record.last_error_kind = signal.error_kind
            record.last_provider_code = signal.provider_code
            record.availability_state = signal.availability_state
            if signal.error_kind == "OVERLOAD":
                record.consecutive_overloads += 1
            else:
                record.consecutive_overloads = 0

            cooldown = self._cooldown(record, signal)
            if signal.availability_state in {"AUTH_BLOCKED", "BALANCE_BLOCKED"}:
                record.circuit_state = "CLOSED"
                record.blocked_until = None
            elif signal.error_kind in {"RATE_LIMIT", "QUOTA_EXHAUSTED"}:
                record.circuit_state = "CLOSED"
                record.blocked_until = now + cooldown if cooldown is not None else None
            else:
                immediate = {
                    str(v or "").strip().upper()
                    for v in self.policy.get("immediate_open_error_kinds") or []
                }
                should_open = signal.error_kind in immediate
                if signal.penalizes_health and self._penalizing_failures(record) >= self._failure_threshold():
                    should_open = True
                if should_open:
                    record.circuit_state = "OPEN"
                    if cooldown is None:
                        cooldown = max(0.0, float(self.policy.get("half_open_probe_interval") or 0.0))
                    record.blocked_until = now + cooldown if cooldown is not None else None
            return self._snapshot_locked(record, now)

    def reset(self, key: HealthKey) -> None:
        with self._lock:
            self._records.pop(key, None)

    def snapshot(self, key: HealthKey) -> HealthSnapshot:
        with self._lock:
            now = self._now()
            record = self._record(key)
            self._refresh(record, now)
            return self._snapshot_locked(record, now)

    def _snapshot_locked(self, record: _HealthRecord, now: float) -> HealthSnapshot:
        self._prune(record, now)
        successes = sum(1 for _, success, _, _ in record.events if success)
        failures = sum(1 for _, success, penalty, _ in record.events if not success and penalty)
        denominator = successes + failures
        success_rate = successes / denominator if denominator else 0.0
        error_rate = failures / denominator if denominator else 0.0
        latencies = [latency for _, _, _, latency in record.events if latency > 0]
        latency_ms = sum(latencies) / len(latencies) if latencies else 0.0
        cause_state = record.availability_state
        state = "CIRCUIT_OPEN" if record.circuit_state == "OPEN" else cause_state
        eligible = (
            record.circuit_state != "OPEN"
            and cause_state not in {
                "AUTH_BLOCKED",
                "BALANCE_BLOCKED",
                "RATE_LIMITED",
                "OVERLOADED",
                "UNREACHABLE",
                "DISABLED",
            }
        )
        return HealthSnapshot(
            availability_state=state,
            circuit_state=record.circuit_state,
            eligible=eligible,
            recent_successes=successes,
            recent_failures=failures,
            success_rate=success_rate,
            error_rate=error_rate,
            latency_ms=latency_ms,
            blocked_until=record.blocked_until,
            last_error_kind=record.last_error_kind,
            last_provider_code=record.last_provider_code,
            cause_state=cause_state,
        )