fix: sibling Nous 401 recovery adopts a peer's refresh instead of rotating again
The stampede fix added a `stale_access_token` hint to resolve_nous_runtime_credentials() so a process whose bearer just 401'd adopts a token a sibling already rotated instead of re-POSTing the shared grant — but only the credential-pool caller passed it. The main agent's 401 path (run_agent._try_refresh_nous_client_credentials), the auxiliary client rebuild, and the proxy adapter all called force_refresh=True with no hint, so `_already_rotated_by_peer` could never fire: N subagents hitting hourly expiry still issued N serialized refreshes, each one invalidating the token a sibling had just adopted. Live 12-process A/B against a fake Portal: 12 refresh POSTs / 9 distinct final tokens before, 1 POST / 1 token after.
This commit is contained in:
@@ -3031,13 +3031,19 @@ def _resolve_nous_pool_runtime_api(*, force_refresh: bool = False) -> Optional[t
|
||||
return api_key, base_url
|
||||
|
||||
|
||||
def _resolve_nous_runtime_api(*, force_refresh: bool = False) -> Optional[tuple[str, str]]:
|
||||
def _resolve_nous_runtime_api(
|
||||
*, force_refresh: bool = False, stale_access_token: Optional[str] = None
|
||||
) -> Optional[tuple[str, str]]:
|
||||
"""Return fresh Nous runtime credentials when available.
|
||||
|
||||
This mirrors the main agent's 401 recovery path and keeps auxiliary
|
||||
clients aligned with the singleton auth store + JWT refresh flow instead of
|
||||
relying only on whatever raw tokens happen to be sitting in auth.json
|
||||
or the credential pool.
|
||||
|
||||
``stale_access_token`` is the bearer that just 401'd; with ``force_refresh``
|
||||
it lets the auth store adopt a sibling process's rotation instead of
|
||||
re-POSTing the shared grant.
|
||||
"""
|
||||
pooled = _resolve_nous_pool_runtime_api(force_refresh=force_refresh)
|
||||
if pooled is not None:
|
||||
@@ -3049,6 +3055,7 @@ def _resolve_nous_runtime_api(*, force_refresh: bool = False) -> Optional[tuple[
|
||||
creds = resolve_nous_runtime_credentials(
|
||||
timeout_seconds=env_float("HERMES_NOUS_TIMEOUT_SECONDS", 15),
|
||||
force_refresh=force_refresh,
|
||||
stale_access_token=stale_access_token or None,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.debug("Auxiliary Nous runtime credential resolution failed: %s", exc)
|
||||
@@ -8298,7 +8305,7 @@ def _refresh_nous_auxiliary_client(
|
||||
401 lands under the ``task=""`` key while the stale entry survives under the
|
||||
task-scoped key (#58894).
|
||||
"""
|
||||
runtime = _resolve_nous_runtime_api(force_refresh=True)
|
||||
runtime = _resolve_nous_runtime_api(force_refresh=True, stale_access_token=api_key)
|
||||
if runtime is None:
|
||||
return None, model
|
||||
|
||||
|
||||
@@ -82,18 +82,19 @@ class NousPortalAdapter(UpstreamAdapter):
|
||||
failed_credential: UpstreamCredential,
|
||||
status_code: int,
|
||||
) -> Optional[UpstreamCredential]:
|
||||
_ = failed_credential
|
||||
if status_code != 401:
|
||||
return None
|
||||
logger.info("proxy: Nous upstream rejected bearer; force-refreshing invoke JWT")
|
||||
return self._get_credential(
|
||||
force_refresh=True,
|
||||
stale_access_token=failed_credential.bearer,
|
||||
)
|
||||
|
||||
def _get_credential(
|
||||
self,
|
||||
*,
|
||||
force_refresh: bool = False,
|
||||
stale_access_token: Optional[str] = None,
|
||||
) -> UpstreamCredential:
|
||||
with self._lock:
|
||||
state = self._read_state()
|
||||
@@ -105,6 +106,7 @@ class NousPortalAdapter(UpstreamAdapter):
|
||||
try:
|
||||
refreshed = resolve_nous_runtime_credentials(
|
||||
force_refresh=force_refresh,
|
||||
stale_access_token=stale_access_token or None,
|
||||
)
|
||||
except AuthError as exc:
|
||||
if _is_terminal_nous_refresh_error(exc):
|
||||
|
||||
@@ -6423,9 +6423,12 @@ class AIAgent:
|
||||
try:
|
||||
from hermes_cli.auth import resolve_nous_runtime_credentials
|
||||
|
||||
# Pass the bearer that just 401'd so a refresh already done by a
|
||||
# sibling process is adopted instead of rotating the grant again.
|
||||
creds = resolve_nous_runtime_credentials(
|
||||
timeout_seconds=env_float("HERMES_NOUS_TIMEOUT_SECONDS", 15),
|
||||
force_refresh=force,
|
||||
stale_access_token=self.api_key or None,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.debug("Nous credential refresh failed: %s", exc)
|
||||
|
||||
@@ -74,7 +74,7 @@ def test_call_llm_auto_provider_evicts_stale_client_end_to_end(monkeypatch):
|
||||
# The 401 refresh rebuilds a fresh client from refreshed runtime creds.
|
||||
monkeypatch.setattr(
|
||||
ac, "_resolve_nous_runtime_api",
|
||||
lambda *, force_refresh=False: ("fresh-key", NOUS_BASE_URL),
|
||||
lambda *, force_refresh=False, stale_access_token=None: ("fresh-key", NOUS_BASE_URL),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ac, "_create_openai_client",
|
||||
@@ -123,7 +123,7 @@ async def test_async_call_llm_auto_provider_evicts_stale_client_end_to_end(monke
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ac, "_resolve_nous_runtime_api",
|
||||
lambda *, force_refresh=False: ("fresh-key", NOUS_BASE_URL),
|
||||
lambda *, force_refresh=False, stale_access_token=None: ("fresh-key", NOUS_BASE_URL),
|
||||
)
|
||||
# Async refresh builds a sync client then wraps it; patch the wrap to `fresh`.
|
||||
monkeypatch.setattr(
|
||||
|
||||
@@ -99,3 +99,32 @@ def test_lock_timeout_during_nous_refresh_does_not_bench_entry(monkeypatch, capl
|
||||
|
||||
assert result is entry
|
||||
assert pool._entries[0].last_status is None, "lock contention is not a credential failure"
|
||||
|
||||
|
||||
def test_agent_401_refresh_passes_failed_bearer_as_stale_hint(monkeypatch):
|
||||
"""Every 401-recovery caller must hand the auth store the bearer that
|
||||
failed — without it ``_already_rotated_by_peer`` can never fire and each
|
||||
subagent rotates the shared grant again (the "no crash, N refreshes"
|
||||
variant of the Sep 2 stampede).
|
||||
"""
|
||||
from run_agent import AIAgent
|
||||
|
||||
agent = AIAgent.__new__(AIAgent)
|
||||
agent.provider = "nous"
|
||||
agent.api_mode = "chat_completions"
|
||||
agent.api_key = "jwt-that-just-401d"
|
||||
agent.base_url = "https://inference-api.nousresearch.com/v1"
|
||||
agent._client_kwargs = {}
|
||||
monkeypatch.setattr(agent, "_replace_primary_openai_client", lambda **k: True)
|
||||
|
||||
seen = {}
|
||||
|
||||
def _fake_resolve(**kwargs):
|
||||
seen.update(kwargs)
|
||||
return {"api_key": "fresh", "base_url": agent.base_url}
|
||||
|
||||
monkeypatch.setattr(auth_mod, "resolve_nous_runtime_credentials", _fake_resolve)
|
||||
|
||||
assert agent._try_refresh_nous_client_credentials(force=True) is True
|
||||
assert seen["force_refresh"] is True
|
||||
assert seen["stale_access_token"] == "jwt-that-just-401d"
|
||||
|
||||
Reference in New Issue
Block a user