agent/nous_rate_guard.py
Rama main · réplica de lectura del repositorio original.
"""Cross-session rate limit guard for Nous Portal.
Writes rate limit state to a shared file so all sessions (CLI, gateway, cron,
auxiliary) can check whether Nous Portal is currently rate-limited before making
requests. Without it each 429 fans out into up to 9 calls per turn (3 SDK
retries x 3 Hermes retries), all counted against RPH.
"""
from __future__ import annotations
import contextlib
import json
import logging
import os
import time
from typing import Any, Mapping, Optional
from utils import atomic_write_text
from agent.retry_utils import parse_retry_after_seconds
from agent.rate_limit_tracker import (
_BUCKET_TAGS, _fmt_seconds, _safe_float, _safe_int, has_rate_limit_headers, lower_headers,
)
logger = logging.getLogger(__name__)
# Reset windows shorter than this are transient upstream jitter, not a quota
# exhaustion worth a cross-session breaker trip.
_MIN_RESET_FOR_BREAKER_SECONDS = 60.0
# The welcome tier's structured ``rate_limited`` refusal: a reset at or above this is an exhausted
# allowance (stop, tell the user when it refreshes and that signing in lifts it); below it the
# turn simply waits it out. The gateway's fairshare bucket names honest resets from a few seconds
# up to a minute, so the split sits where a quiet wait stops being quiet.
WELCOME_LONG_WAIT_SECONDS = 20.0
format_remaining = _fmt_seconds
def _state_path(*, anonymous: bool = False) -> str:
"""Path to the Nous rate limit state file."""
try:
from hermes_constants import get_hermes_home
base = get_hermes_home()
except ImportError:
base = os.path.join(os.path.expanduser("~"), ".hermes")
# Signing in must not inherit the anonymous allowance's cooldown (or clear it for
# another anonymous session). Keep the existing named-account file unchanged.
return os.path.join(base, "rate_limits", "nous-anonymous.json" if anonymous else "nous.json")
def _parse_reset_seconds(headers: Optional[Mapping[str, str]]) -> Optional[float]:
"""Best reset estimate (seconds from now) from hourly, per-minute, then retry-after headers."""
lowered = lower_headers(headers)
for key in ("x-ratelimit-reset-requests-1h", "x-ratelimit-reset-requests"):
val = _safe_float(lowered.get(key), 0.0)
if val > 0:
return val
retry_after = parse_retry_after_seconds(lowered.get("retry-after"))
return retry_after if retry_after else None
def record_nous_rate_limit(
*, headers: Optional[Mapping[str, str]] = None, error_context: Optional[dict[str, Any]] = None,
default_cooldown: float = 300.0,
anonymous: bool = False,
) -> None:
"""Record that Nous Portal is rate-limited in the shared state file.
Reset time comes from ``error_context["reset_at"]`` (the 429 body's own wait), then headers, then
``default_cooldown``. The body wins: a welcome-tier refusal says when to retry (410 s), while the
hourly bucket's reset header is that window's end (about an hour) even with requests left in it.
"""
now = time.time()
reset_at = None
if isinstance(error_context, dict):
ctx_reset = error_context.get("reset_at")
if isinstance(ctx_reset, (int, float)) and ctx_reset > now:
reset_at = float(ctx_reset)
if reset_at is None and (header_seconds := _parse_reset_seconds(headers)) is not None:
reset_at = now + header_seconds
if reset_at is None:
reset_at = now + default_cooldown
state = {"reset_at": reset_at, "recorded_at": now, "reset_seconds": reset_at - now}
try:
atomic_write_text(_state_path(anonymous=anonymous), json.dumps(state))
logger.info("Nous rate limit recorded: resets in %.0fs (at %.0f)", reset_at - now, reset_at)
except Exception as exc:
logger.debug("Failed to write Nous rate limit state: %s", exc)
def nous_rate_limit_remaining(*, anonymous: bool = False) -> Optional[float]:
"""Seconds remaining until reset, or None if not rate-limited (expired state is removed)."""
path = _state_path(anonymous=anonymous)
try:
with open(path, encoding="utf-8-sig") as f:
state = json.load(f)
remaining = state.get("reset_at", 0) - time.time()
if remaining > 0:
return remaining
with contextlib.suppress(OSError):
os.unlink(path)
return None
except (FileNotFoundError, json.JSONDecodeError, KeyError, TypeError):
return None
def clear_nous_rate_limit(*, anonymous: bool = False) -> None:
"""Clear the rate limit state (e.g., after a successful Nous request)."""
try:
os.unlink(_state_path(anonymous=anonymous))
except FileNotFoundError:
pass
except OSError as exc:
logger.debug("Failed to clear Nous rate limit state: %s", exc)
def _is_exhausted(remaining: Optional[int], reset: Optional[float]) -> bool:
"""remaining == 0 AND a reset window long enough to be a real quota exhaustion."""
return (
remaining is not None
and remaining <= 0
and reset is not None
and reset >= _MIN_RESET_FOR_BREAKER_SECONDS
)
def is_genuine_nous_rate_limit(
*, headers: Optional[Mapping[str, str]] = None, last_known_state: Optional[Any] = None,
) -> bool:
"""Decide whether a 429 from Nous Portal is a real account rate limit.
Nous multiplexes upstream providers behind one key, so a 429 may be one
upstream model out of capacity (clears in seconds) rather than our quota;
tripping the breaker on that would block every Nous model for minutes.
Only an exhausted bucket (remaining == 0 with reset >= 60s) in the 429's
own headers, or in the last-known-good ``RateLimitState``, is genuine.
"""
if _has_exhausted_bucket(_parse_buckets_from_headers(headers)):
return True
return last_known_state is not None and _has_exhausted_bucket_in_object(last_known_state)
def is_long_welcome_rate_limit(error_context: Any) -> bool:
"""True for a Nous welcome-tier ``rate_limited`` refusal whose reset is long enough to be an
exhausted allowance (``WELCOME_LONG_WAIT_SECONDS``), as parsed into ``error_context``
(``welcome_refusal`` from ``hermes_cli.anon_auth.parse_welcome_refusal``). Capacity refusals
(``at_capacity`` / ``admission_closed``) are never this: they are retried in place."""
if not isinstance(error_context, dict):
return False
refusal = error_context.get("welcome_refusal")
if not isinstance(refusal, dict) or refusal.get("reason") != "rate_limited":
return False
return _safe_float(refusal.get("retry_after"), 0.0) >= WELCOME_LONG_WAIT_SECONDS
def welcome_refusal_from_headers(
headers: Optional[Mapping[str, str]], *, body_wait: Optional[float] = None,
) -> dict[str, Any]:
"""A free-tier 429 that carried no structured ``reason``, read from its headers (else the wait
its body named) into the ``parse_welcome_refusal`` shape. Long wait or an empty bucket =
``rate_limited``; a short or unexplained one = ``at_capacity``. A missing wait stays 0 so the
breaker never trips on a guess."""
buckets = _parse_buckets_from_headers(headers).values()
# Only an empty bucket's reset is a wait; a healthy bucket's reset is just its window end.
exhausted_wait = max((reset for remaining, reset in buckets
if remaining is not None and remaining <= 0 and (reset or 0) > 0), default=None)
wait = (exhausted_wait or parse_retry_after_seconds(lower_headers(headers).get("retry-after"))
or body_wait or 0.0)
empty = any(remaining is not None and remaining <= 0 for remaining, _reset in buckets)
reason = "rate_limited" if wait >= WELCOME_LONG_WAIT_SECONDS or empty else "at_capacity"
return {"reason": reason, "retry_after": max(0, int(wait)), "alternates": [], "upgrade_url": ""}
def _parse_buckets_from_headers(
headers: Optional[Mapping[str, str]],
) -> dict[str, tuple[Optional[int], Optional[float]]]:
"""(remaining, reset_seconds) per bucket from x-ratelimit-* headers ({} if none)."""
lowered = lower_headers(headers)
if not has_rate_limit_headers(lowered):
return {}
result: dict[str, tuple[Optional[int], Optional[float]]] = {}
for _attr, tag in _BUCKET_TAGS:
remaining = _safe_int(lowered.get(f"x-ratelimit-remaining-{tag}"), None)
reset = _safe_float(lowered.get(f"x-ratelimit-reset-{tag}"), None)
if remaining is not None or reset is not None:
result[tag] = (remaining, reset)
return result
def _has_exhausted_bucket(buckets: Mapping[str, tuple[Optional[int], Optional[float]]]) -> bool:
return any(_is_exhausted(remaining, reset) for remaining, reset in buckets.values())
def _has_exhausted_bucket_in_object(state: Any) -> bool:
"""Check a RateLimitState-like object (duck-typed; missing attrs are skipped)."""
for attr, _tag in _BUCKET_TAGS:
bucket = getattr(state, attr, None)
if bucket is None or (getattr(bucket, "limit", 0) or 0) <= 0:
continue
remaining = getattr(bucket, "remaining", 0) or 0
reset = getattr(bucket, "remaining_seconds_now", None)
if reset is None:
reset = getattr(bucket, "reset_seconds", 0.0) or 0.0
if _is_exhausted(remaining, reset):
return True
return False