mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-20 15:33:54 +00:00
fix(cron/chronos): cache PyJWKClient across fires to stop JWKS fetch storm (#64641)
The inbound cron-fire verifier constructed a fresh PyJWKClient on every fire, discarding the client's key cache and forcing a synchronous JWKS HTTP GET to the portal on each fire. Under a burst of concurrent fires (a hosted instance with several cron jobs firing in the same window) this fanned out into N simultaneous JWKS fetches that the portal rate-limited (HTTP 403 -> verification fails -> agent 401), or that blocked the event loop long enough that the fire webhook could not return its 202 before the relay's 30s timeout (observed in prod as relay 504s concentrated on high-job-count instances). Cache one PyJWKClient per JWKS URL at module scope (double-checked lock) so the signing keys are reused across fires; NAS keys rotate rarely, so the steady state is zero JWKS fetches per fire. Regression test proves 5 fires -> 1 client construction (was 5).
This commit is contained in:
parent
df5700ebe3
commit
9884b4faad
2 changed files with 97 additions and 2 deletions
|
|
@ -19,6 +19,7 @@ hand-roll JWT verification.
|
|||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from typing import Any, Callable, Dict, Optional
|
||||
|
||||
logger = logging.getLogger("cron.chronos.verify")
|
||||
|
|
@ -27,6 +28,44 @@ logger = logging.getLogger("cron.chronos.verify")
|
|||
# JWT (without this claim) must NOT be replayable against /api/cron/fire.
|
||||
_FIRE_PURPOSE = "cron_fire"
|
||||
|
||||
# Process-wide cache of PyJWKClient instances, keyed by JWKS URL.
|
||||
#
|
||||
# WHY THIS EXISTS: a PyJWKClient caches the fetched JWKS (signing keys) on the
|
||||
# INSTANCE. Constructing a fresh client per fire therefore threw that cache
|
||||
# away and forced a synchronous JWKS HTTP GET to the portal on EVERY fire. Under
|
||||
# a burst of concurrent fires (an instance with several cron jobs firing in the
|
||||
# same window) that fanned out into N simultaneous JWKS fetches, which the
|
||||
# portal rate-limited (HTTP 403) — verification then failed and the agent
|
||||
# answered 401. When a fetch was merely slow rather than rate-limited, it blocked
|
||||
# the event loop long enough that the fire webhook could not return its 202
|
||||
# before the relay's 30s timeout (observed in prod as relay 504s concentrated on
|
||||
# high-job-count instances). Reusing one client per URL keeps the signing keys
|
||||
# cached (NAS keys rotate rarely), so the steady state is zero JWKS fetches per
|
||||
# fire. See docs/chronos-managed-cron-contract.md and the betterstack triage.
|
||||
_JWK_CLIENTS: Dict[str, Any] = {}
|
||||
_JWK_CLIENTS_LOCK = threading.Lock()
|
||||
|
||||
|
||||
def _get_jwk_client(jwks_url: str) -> Any:
|
||||
"""Return a process-cached PyJWKClient for ``jwks_url`` (one per URL).
|
||||
|
||||
PyJWKClient does its own key caching internally (``cache_keys``/``lifespan``);
|
||||
the whole point here is to reuse the SAME instance across fires so that cache
|
||||
is actually hit instead of discarded. Double-checked-locked so concurrent
|
||||
fires resolve to a single shared client without racing.
|
||||
"""
|
||||
client = _JWK_CLIENTS.get(jwks_url)
|
||||
if client is not None:
|
||||
return client
|
||||
with _JWK_CLIENTS_LOCK:
|
||||
client = _JWK_CLIENTS.get(jwks_url)
|
||||
if client is None:
|
||||
from jwt import PyJWKClient
|
||||
|
||||
client = PyJWKClient(jwks_url)
|
||||
_JWK_CLIENTS[jwks_url] = client
|
||||
return client
|
||||
|
||||
|
||||
def verify_nas_fire_token(
|
||||
*,
|
||||
|
|
@ -60,12 +99,15 @@ def verify_nas_fire_token(
|
|||
|
||||
try:
|
||||
import jwt
|
||||
from jwt import PyJWKClient
|
||||
|
||||
# Resolve the signing key from the JWKS endpoint by the token's kid.
|
||||
signing_key = None
|
||||
if jwks_or_key.startswith("http://") or jwks_or_key.startswith("https://"):
|
||||
jwk_client = PyJWKClient(jwks_or_key)
|
||||
# Reuse a process-cached client so the JWKS fetch is amortised across
|
||||
# fires (a fresh client per fire re-fetched the JWKS every time and,
|
||||
# under concurrent fires, tripped the portal's rate limit → 403 →
|
||||
# 401, or blocked the event loop past the relay's 30s timeout → 504).
|
||||
jwk_client = _get_jwk_client(jwks_or_key)
|
||||
signing_key = jwk_client.get_signing_key_from_jwt(token).key
|
||||
else:
|
||||
# A PEM public key passed inline (test / pinned-key deployments).
|
||||
|
|
|
|||
|
|
@ -152,6 +152,7 @@ def test_empty_token_refused(rsa_keys):
|
|||
|
||||
def test_jwks_url_path_resolves_key(rsa_keys, monkeypatch):
|
||||
"""The JWKS-URL branch resolves the signing key via PyJWKClient."""
|
||||
from plugins.cron_providers.chronos import verify as verify_mod
|
||||
from plugins.cron_providers.chronos.verify import verify_nas_fire_token
|
||||
|
||||
priv, pub = rsa_keys
|
||||
|
|
@ -168,6 +169,8 @@ def test_jwks_url_path_resolves_key(rsa_keys, monkeypatch):
|
|||
return FakeKey()
|
||||
|
||||
monkeypatch.setattr("jwt.PyJWKClient", FakeJWKClient)
|
||||
# Isolate from the process-wide client cache (other tests may have populated it).
|
||||
monkeypatch.setattr(verify_mod, "_JWK_CLIENTS", {})
|
||||
claims = verify_nas_fire_token(
|
||||
token=token, expected_audience=AUD,
|
||||
jwks_or_key="https://portal.nousresearch.com/.well-known/jwks.json",
|
||||
|
|
@ -176,6 +179,56 @@ def test_jwks_url_path_resolves_key(rsa_keys, monkeypatch):
|
|||
assert claims is not None and claims["purpose"] == "cron_fire"
|
||||
|
||||
|
||||
def test_jwks_client_is_cached_across_calls(rsa_keys, monkeypatch):
|
||||
"""Regression (Chronos relay 403/401 + 504 storm): the JWKS client must be
|
||||
constructed ONCE per URL and reused across fires, NOT rebuilt per call.
|
||||
|
||||
A fresh PyJWKClient per fire discards its key cache and forces a synchronous
|
||||
JWKS HTTP GET on every fire; a burst of concurrent fires then fans out into N
|
||||
simultaneous fetches that the portal rate-limits (403 → agent 401) or that
|
||||
block the event loop past the relay's 30s timeout (504). Reusing one cached
|
||||
client keeps the steady state at zero fetches per fire. This test fails
|
||||
against the pre-fix code (construct_count == N) and passes with the cache
|
||||
(construct_count == 1).
|
||||
"""
|
||||
from plugins.cron_providers.chronos import verify as verify_mod
|
||||
from plugins.cron_providers.chronos.verify import verify_nas_fire_token
|
||||
|
||||
priv, pub = rsa_keys
|
||||
url = "https://portal.nousresearch.com/.well-known/jwks.json"
|
||||
|
||||
counters = {"construct": 0, "fetch": 0}
|
||||
|
||||
class FakeKey:
|
||||
key = pub
|
||||
|
||||
class CountingJWKClient:
|
||||
def __init__(self, u):
|
||||
counters["construct"] += 1
|
||||
|
||||
def get_signing_key_from_jwt(self, tok):
|
||||
counters["fetch"] += 1
|
||||
return FakeKey()
|
||||
|
||||
monkeypatch.setattr("jwt.PyJWKClient", CountingJWKClient)
|
||||
# Start from an empty cache so this test's count is deterministic.
|
||||
monkeypatch.setattr(verify_mod, "_JWK_CLIENTS", {})
|
||||
|
||||
for _ in range(5):
|
||||
token = _mint(priv, _base_claims())
|
||||
claims = verify_nas_fire_token(
|
||||
token=token, expected_audience=AUD, jwks_or_key=url, issuer=ISS,
|
||||
)
|
||||
assert claims is not None
|
||||
|
||||
# The client is built once and reused; only the fetch (served from the
|
||||
# client's own cache in production) is per-call.
|
||||
assert counters["construct"] == 1, (
|
||||
f"expected 1 PyJWKClient construction, got {counters['construct']} "
|
||||
"(a fresh client per fire is the bug this guards against)"
|
||||
)
|
||||
|
||||
|
||||
def test_get_fire_verifier_returns_nas_verifier():
|
||||
from plugins.cron_providers.chronos.verify import get_fire_verifier, verify_nas_fire_token
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue