fix(nous): pre-expiry adoption only ever swaps to a key for the SAME account

Independent review of the first fix found a credential-identity takeover:
_adopt_nous_key_before_expiry() called the singleton resolver
unconditionally, so an agent running on an explicitly supplied (or
pool-selected) account-A key near expiry was moved onto the logged-in
account-B key from auth.json before the A key had even failed, and the next
real SDK request went out as B. That silently changes who is billed.

The adoption now reads the `sub` claim of the key in hand and passes it as
require_account; _try_refresh_nous_client_credentials refuses any
replacement whose `sub` differs (logged at INFO, current key kept). A key
with no `sub` is not adopted proactively at all. The reactive 401 path is
unchanged, and same-account adoption (the keepalive's or a peer's fresh key)
still works.

Verified with the reviewer's own live probe (real AIAgent -> real
prepare_iteration -> real auth.json transaction -> real OpenAI SDK ->
loopback capture): explicit-account case account-A -> account-A,
adopted_fresh=False (was A -> B); same-account still adopts; 12 concurrent
near-expiry agents still 0 x 401.

Tests (2 new): a fresh key for a different account is never adopted; a key
without an account claim is left alone without touching the store.
This commit is contained in:
Teknium
2026-09-05 06:09:24 -07:00
parent 058ad0329e
commit ef897bfd7e
2 changed files with 45 additions and 8 deletions

View File

@@ -525,7 +525,7 @@ class ClientLifecycleMixin:
return False
return self._adopt_openai_credentials(api_key, base_url, reason=f"{self.provider}_credential_refresh")
def _try_refresh_nous_client_credentials(self, *, force: bool = True) -> bool:
def _try_refresh_nous_client_credentials(self, *, force: bool = True, require_account: str | None = None) -> bool:
# Portal serves anthropic/* on the native Messages route, so either client kind may hold the expiring JWT.
if self.provider != "nous" or self.api_mode not in ("chat_completions", "anthropic_messages"):
return False
@@ -545,6 +545,18 @@ class ClientLifecycleMixin:
return False
if str(api_key).strip() == str(self.api_key or "").strip():
return False # store holds the same key: nothing to adopt, no client rebuild
if require_account is not None:
try:
from hermes_cli.auth_constants import _decode_jwt_claims
new_account = _decode_jwt_claims(str(api_key)).get("sub")
except Exception:
new_account = None
if str(new_account or "") != require_account:
logger.info(
"Nous pre-expiry adoption skipped: the store's key belongs to a different account "
"than the one in hand; keeping the current credential."
)
return False
if self.api_mode == "anthropic_messages":
self.api_key, self.base_url = api_key.strip(), base_url.strip().rstrip("/")
self._anthropic_api_key, self._anthropic_base_url = self.api_key, self.base_url
@@ -567,17 +579,24 @@ class ClientLifecycleMixin:
every agent in a process learned about the hourly expiry from its own 401, all in the same
minute (620 in one 200-subagent run), and the pool benched the sole credential for all of
them. Returns True when a new key was adopted.
Identity guard: the replacement must belong to the SAME account (``sub`` claim) as the key
in hand. The store holds the logged-in singleton; an agent running on an explicitly supplied
or pool-selected key for a different account must never be silently moved onto it (that
changes who is billed). When either side lacks a ``sub`` nothing is adopted here; the
reactive 401 path is unchanged.
"""
if getattr(self, "provider", "") != "nous" or not getattr(self, "api_key", None):
return False
try:
from hermes_cli.auth_constants import _decode_jwt_claims
exp = _decode_jwt_claims(self.api_key).get("exp")
claims = _decode_jwt_claims(self.api_key)
except Exception:
return False
if not isinstance(exp, (int, float)) or exp - time.time() > self._NOUS_KEY_ADOPT_SKEW_S:
exp, account = claims.get("exp"), claims.get("sub")
if not account or not isinstance(exp, (int, float)) or exp - time.time() > self._NOUS_KEY_ADOPT_SKEW_S:
return False
return self._try_refresh_nous_client_credentials(force=False)
return self._try_refresh_nous_client_credentials(force=False, require_account=str(account))
def _resolve_env_credentials(self) -> Optional[tuple]:

View File

@@ -12,10 +12,10 @@ from unittest.mock import patch
from agent.client_lifecycle import ClientLifecycleMixin
def _jwt(exp: float) -> str:
def _jwt(exp: float, sub: str = "acct-A") -> str:
def b64(o):
return base64.urlsafe_b64encode(json.dumps(o).encode()).rstrip(b"=").decode()
return f"{b64({'alg': 'none'})}.{b64({'exp': exp})}.sig"
return f"{b64({'alg': 'none'})}.{b64({'exp': exp, 'sub': sub})}.sig"
class _Agent(ClientLifecycleMixin):
@@ -40,14 +40,32 @@ def test_key_inside_the_skew_adopts_the_stores_fresh_key_without_forcing_a_refre
agent = _Agent(_jwt(time.time() + 60))
calls = []
fresh = _jwt(time.time() + 3600) # same account A
def resolve(**kw):
calls.append(kw)
return {"api_key": "fresh-key", "base_url": agent.base_url}
return {"api_key": fresh, "base_url": agent.base_url}
with patch("hermes_cli.auth.resolve_nous_runtime_credentials", side_effect=resolve):
assert agent._adopt_nous_key_before_expiry() is True
assert calls[0]["force_refresh"] is False # the keepalive/peer refresh is adopted, never re-minted
assert agent.adopted == [("fresh-key", "nous_credential_refresh")]
assert agent.adopted == [(fresh, "nous_credential_refresh")]
def test_a_fresh_key_for_a_different_account_is_never_adopted():
"""Independent-review witness: an explicitly supplied account-A key near expiry was replaced by the
logged-in singleton's account-B key and the next real request went out as B. Identity is preserved."""
agent = _Agent(_jwt(time.time() + 60, sub="acct-A"))
other = _jwt(time.time() + 3600, sub="acct-B")
with patch("hermes_cli.auth.resolve_nous_runtime_credentials", return_value={"api_key": other, "base_url": agent.base_url}):
assert agent._adopt_nous_key_before_expiry() is False
assert agent.adopted == [] and agent.api_key != other
def test_a_key_without_an_account_claim_is_left_alone_proactively():
agent = _Agent(_jwt(time.time() + 60, sub=""))
with patch("hermes_cli.auth.resolve_nous_runtime_credentials", side_effect=AssertionError("must not hit the store")):
assert agent._adopt_nous_key_before_expiry() is False
def test_same_key_back_from_the_store_is_not_readopted():