fix: bound the token cache, and make peer message ids unique
Two bugs from review, the second in both agents because identity.py is a copy. The exchanged-token cache grew for the life of the process. It is keyed on (subject token, scope) and the subject token is minted fresh every turn, so every turn added keys nothing would ever read again — and expiry was checked on read but never collected. A slow leak that only shows in a pod up for days, which is the worst way to find one. It now drops expired entries first and evicts oldest beyond a bound. Expired first on purpose: an expired entry is dead weight wherever it sits in the ordering, so evicting by age alone would keep it while discarding a live one. The docstring was wrong before the code was: it claimed the cache lasted one turn. It never did. Corrected rather than quietly made true. message_id reused the context id, which is the conversation — and a conversation makes several calls to the same peer, relaying an approval being exactly that. A peer treating the second call as a replay of the first would surface as a stale answer rather than an error. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
@@ -19,6 +19,7 @@ from __future__ import annotations
|
||||
import base64
|
||||
import json
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass
|
||||
|
||||
import httpx2
|
||||
@@ -27,6 +28,11 @@ GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange"
|
||||
SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token"
|
||||
|
||||
|
||||
#: Enough for every scope of every conversation in flight, and small enough that
|
||||
#: the linear prune below costs nothing.
|
||||
MAX_CACHED_TOKENS = 500
|
||||
|
||||
|
||||
class TokenExchangeError(RuntimeError):
|
||||
"""Keycloak refused the exchange. Carries its response, which names the cause."""
|
||||
|
||||
@@ -46,9 +52,15 @@ class ExchangedToken:
|
||||
class TokenExchanger:
|
||||
"""Exchanges a user token for audience-scoped tokens.
|
||||
|
||||
Caches per (subject token, scope) for the lifetime of one turn. Deliberately
|
||||
not a long-lived cache keyed by user: a token that outlives the turn it was
|
||||
minted for is a token that can be replayed after the user's session ended.
|
||||
Caches per (subject token, scope), so the several exchanges one turn needs
|
||||
cost one round trip each rather than one per use. Deliberately not keyed by
|
||||
user: a cached token that outlives the turn it was minted for is a token
|
||||
that can be replayed after the user's session ended.
|
||||
|
||||
The cache is bounded and drops expired entries, which it has to be. The
|
||||
subject token changes every turn, so every turn adds keys that will never be
|
||||
read again — an unbounded dict here grows for the life of the process. An
|
||||
earlier docstring claimed the cache lasted one turn; the code never did.
|
||||
"""
|
||||
|
||||
def __init__(self, client: httpx2.AsyncClient, *, token_endpoint: str, client_id: str, client_secret: str) -> None:
|
||||
@@ -56,7 +68,7 @@ class TokenExchanger:
|
||||
self._endpoint = token_endpoint
|
||||
self._client_id = client_id
|
||||
self._client_secret = client_secret
|
||||
self._cache: dict[tuple[str, str], ExchangedToken] = {}
|
||||
self._cache: OrderedDict[tuple[str, str], ExchangedToken] = OrderedDict()
|
||||
|
||||
async def exchange(self, subject_token: str, audience_scope: str) -> str:
|
||||
key = (subject_token, audience_scope)
|
||||
@@ -87,8 +99,21 @@ class TokenExchanger:
|
||||
expires_at=time.monotonic() + float(payload.get("expires_in", 60)),
|
||||
)
|
||||
self._cache[key] = token
|
||||
self._prune()
|
||||
return token.access_token
|
||||
|
||||
def _prune(self) -> None:
|
||||
"""Drop what is dead, then what is oldest.
|
||||
|
||||
Expired first, because an expired entry is dead weight wherever it sits
|
||||
in the ordering — evicting by age alone would keep it while discarding a
|
||||
live one. The bound is the backstop for a burst that outruns expiry.
|
||||
"""
|
||||
for key in [k for k, v in self._cache.items() if v.expired]:
|
||||
del self._cache[key]
|
||||
while len(self._cache) > MAX_CACHED_TOKENS:
|
||||
self._cache.popitem(last=False)
|
||||
|
||||
|
||||
def caller_of(access_token: str) -> str:
|
||||
"""The username inside a token, for the prompt and nothing else.
|
||||
|
||||
Reference in New Issue
Block a user