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:
Stéphane Tailland
2026-08-30 08:03:33 +02:00
parent 60d25edb52
commit afe8f1103a
4 changed files with 115 additions and 6 deletions
+29 -4
View File
@@ -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.