agent/credential_auth_recovery.py
Rama main · réplica de lectura del repositorio original.
"""Credential-pool recovery for an auth failure (401/403) on the dispatched credential.
Nous answers bursts of ``401 {"error": "authentication required"}`` for agent keys that are
still unexpired and valid a few seconds later. Re-minting cannot fix a server-side blip, and
each forced refresh redeems the single-use refresh token, so a Nous 401 episode gets exactly
ONE forced refresh. Any later 401 in the same turn (or a refresh that handed back the bearer
that just failed) waits out ``TRANSIENT_AUTH_BACKOFF_S`` and retries with the store's current
key, adopting a peer's rotation but never POSTing again. A dead grant (refresh rejected)
still fails fast through rotation, and so does an entitlement 403.
"""
from __future__ import annotations
import time
from typing import Any, Optional
# Waits before each retry of an unexpired Nous key that keeps answering 401 (~50 s in total).
TRANSIENT_AUTH_BACKOFF_S = (5.0, 15.0, 30.0)
# Cap same-entry OAuth refreshes on a persistent auth failure, else a single-entry pool re-mints forever.
_MAX_AUTH_REFRESH_ATTEMPTS = 2
_ACTIVITY_TOUCH_S = 30.0
# A 401 this long after the episode's last one starts a new episode (one more refresh allowed):
# hour-long turns must still recover from a key revoked later in the same turn.
_EPISODE_IDLE_S = 120.0
def _ra():
import run_agent
return run_agent
def _episode(agent: Any, provider: str, entry_id: Optional[str]) -> dict[str, Any]:
"""Per-turn auth-failure state for one pool entry (the turn loop resets the map)."""
if getattr(agent, "_auth_pool_refresh_counts", None) is None:
agent._auth_pool_refresh_counts = {}
now = time.monotonic()
episode = agent._auth_pool_refresh_counts.get((provider, entry_id))
if episode is None or now - episode["last"] > _EPISODE_IDLE_S:
episode = agent._auth_pool_refresh_counts[(provider, entry_id)] = {"refreshes": 0, "waits": 0}
episode["last"] = now
return episode
def nous_pool_refreshed_this_turn(agent: Any) -> bool:
"""True once the pool spent this turn's forced Nous refresh (the singleton path must not re-POST)."""
counts = getattr(agent, "_auth_pool_refresh_counts", None) or {}
return any(key[0] == "nous" and state.get("refreshes") for key, state in counts.items())
def _interruptible_wait(agent: Any, seconds: float) -> bool:
"""Sleep in 200 ms slices; False when the user interrupted."""
deadline = time.monotonic() + seconds
next_touch = time.monotonic() + _ACTIVITY_TOUCH_S
while time.monotonic() < deadline:
if getattr(agent, "_interrupt_requested", False):
return False
time.sleep(min(0.2, max(0.0, deadline - time.monotonic())))
if time.monotonic() >= next_touch:
agent._touch_activity("waiting out a Nous 401 on a valid key")
next_touch = time.monotonic() + _ACTIVITY_TOUCH_S
return True
def _entry_for(pool: Any, credential_id: Optional[str], api_key_hint: Optional[str]) -> Any:
entries = pool.entries()
if credential_id:
found = next((e for e in entries if e.id == credential_id), None)
if found is not None:
return found
return next((e for e in entries if api_key_hint and e.runtime_api_key == api_key_hint), None)
def _retry_after_backoff(agent: Any, pool: Any, episode: dict[str, Any], *, failed_key: Optional[str],
entry_id: Optional[str], credential_id: Optional[str], has_retried_429: bool):
"""Wait out a server-side 401 on an unexpired key, then retry with the store's current key."""
log = _ra().logger
step = episode["waits"]
if step >= len(TRANSIENT_AUTH_BACKOFF_S):
log.warning(
"Credential auth failure persists for pool entry %s after one refresh and %s "
"backoff retries — treating as unrecoverable and allowing fallback to activate.",
entry_id, step)
return False, has_retried_429
wait = TRANSIENT_AUTH_BACKOFF_S[step]
episode["waits"] = step + 1
log.info(
"Credential auth failure on an unexpired key for pool entry %s — transient 401, "
"retrying in %.0fs (backoff %s/%s, no further refresh)",
entry_id, wait, step + 1, len(TRANSIENT_AUTH_BACKOFF_S))
agent._vprint(
f"{agent.log_prefix}🔐 Nous rejected a still-valid key (transient); retrying in {wait:.0f}s...",
force=True, diagnostic=True)
if not _interruptible_wait(agent, wait):
# Not recovered: the caller's interrupt check then aborts (or redirects) the turn. A
# True here re-sends the request, and the retry loop never re-checks the interrupt.
return False, has_retried_429
# A peer process may have rotated meanwhile: adopt what the store holds now (no POST).
entry = _entry_for(pool, credential_id or entry_id, failed_key)
if entry is not None:
if agent._swap_credential(pool._sync_nous_entry_from_auth_store(entry)) is False:
return False, has_retried_429
return True, has_retried_429
def recover_auth_failure(agent, pool, *, status_code, has_retried_429, error_context, api_key_hint,
credential_id, rotate_and_swap):
from agent.agent_runtime_helpers import _is_entitlement_403
log = _ra().logger
if _is_entitlement_403(agent, status_code, error_context):
log.info(
"Credential %s — entitlement-shaped 403 from %s; "
"skipping pool refresh (account lacks subscription, "
"not a transient auth failure).", status_code if status_code is not None else "auth",
agent.provider or "provider",
)
return False, has_retried_429
provider = agent.provider
# Only a 401 is the "valid key refused" shape worth backing off on; a non-entitlement 403
# keeps the capped-refresh path and fails fast.
nous = provider == "nous" and status_code == 401
entry_id = credential_id
if nous:
found = _entry_for(pool, credential_id, api_key_hint)
entry_id = getattr(found, "id", None) or credential_id
episode = _episode(agent, provider, entry_id)
if episode["refreshes"] >= 1:
# The one forced refresh of this episode already happened: never re-redeem the grant.
return _retry_after_backoff(
agent, pool, episode, failed_key=api_key_hint, entry_id=entry_id,
credential_id=credential_id, has_retried_429=has_retried_429)
# Refresh the entry that supplied the failing key, not current(): refreshing a healthy entry
# burns its single-use refresh token for a failure it never had.
refresh_kwargs = {"api_key_hint": api_key_hint}
if credential_id:
refresh_kwargs["credential_id"] = credential_id
refreshed = pool.try_refresh_matching(**refresh_kwargs)
if refreshed is None:
# Refresh failed; rotate (the failed entry is already marked exhausted).
return (True, False) if rotate_and_swap(401, "auth refresh failed") else (False, has_retried_429)
refreshed_id = getattr(refreshed, "id", None)
if nous:
episode = _episode(agent, provider, refreshed_id or entry_id)
episode["refreshes"] += 1
if api_key_hint and getattr(refreshed, "runtime_api_key", None) == api_key_hint:
# The "refresh" handed back the bearer that just failed: replaying it at once is
# the 1.5 s give-up; it is a valid key the server is refusing, so back off instead.
return _retry_after_backoff(
agent, pool, episode, failed_key=api_key_hint, entry_id=refreshed_id,
credential_id=credential_id, has_retried_429=has_retried_429)
elif refreshed_id is not None:
# try_refresh_matching() reports success even when upstream keeps rejecting; cap same-entry
# refreshes so a single-entry pool falls through to fallback.
state = _episode(agent, provider, refreshed_id)
state["refreshes"] += 1
if state["refreshes"] > _MAX_AUTH_REFRESH_ATTEMPTS:
log.warning(
"Credential auth failure persists after %s refreshes for "
"pool entry %s — treating as unrecoverable and allowing "
"fallback to activate.", state["refreshes"] - 1, refreshed_id,
)
return False, has_retried_429
log.info("Credential auth failure — refreshed pool entry %s", refreshed_id or "?")
if agent._swap_credential(refreshed) is False:
return False, has_retried_429
return True, has_retried_429