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:
+1
-1
@@ -30,7 +30,7 @@ spec:
|
||||
# repo is public, so the pull needs no credentials — which removes an
|
||||
# imagePullSecret that would otherwise have to be created by hand in
|
||||
# this namespace and kept in step with a GitHub token.
|
||||
image: gitea.sttlab.eu/stt/classic-models-agent-catalogue:v0.3.1
|
||||
image: gitea.sttlab.eu/stt/classic-models-agent-catalogue:v0.3.2
|
||||
ports:
|
||||
- name: http
|
||||
containerPort: 8080
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "classic-models-agent-catalogue"
|
||||
version = "0.3.1"
|
||||
version = "0.3.2"
|
||||
description = "Product management agent for the Classic Models catalog domain"
|
||||
requires-python = ">=3.12"
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
"""The token cache must not grow for the life of the process.
|
||||
|
||||
The subject token changes every turn, so every turn adds keys that will never be
|
||||
read again. An unbounded dict here is a slow leak that only shows up in a pod
|
||||
that has been up for days — which is the worst way to find it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
import httpx
|
||||
import httpx2
|
||||
import pytest
|
||||
|
||||
from catalogue_agent.identity import MAX_CACHED_TOKENS, ExchangedToken, TokenExchanger
|
||||
|
||||
|
||||
def exchanger(handler) -> TokenExchanger:
|
||||
transport = httpx2.MockTransport(handler)
|
||||
return TokenExchanger(
|
||||
httpx2.AsyncClient(transport=transport),
|
||||
token_endpoint="https://keycloak/token",
|
||||
client_id="agent",
|
||||
client_secret="secret",
|
||||
)
|
||||
|
||||
|
||||
def ok(_request: httpx2.Request) -> httpx2.Response:
|
||||
return httpx2.Response(200, json={"access_token": "a-token", "expires_in": 300})
|
||||
|
||||
|
||||
class TestExpiry:
|
||||
def test_an_expired_entry_is_reported_expired(self) -> None:
|
||||
assert ExchangedToken("t", time.monotonic() - 1).expired is True
|
||||
|
||||
def test_a_fresh_one_is_not(self) -> None:
|
||||
assert ExchangedToken("t", time.monotonic() + 300).expired is False
|
||||
|
||||
def test_the_margin_covers_a_turn_that_stalls(self) -> None:
|
||||
# A token valid for another 10s is treated as spent: a turn that stalls
|
||||
# on a slow model must not send one that dies mid-flight.
|
||||
assert ExchangedToken("t", time.monotonic() + 10).expired is True
|
||||
|
||||
|
||||
class TestPruning:
|
||||
async def test_a_hit_costs_no_round_trip(self) -> None:
|
||||
calls = 0
|
||||
|
||||
def counting(request: httpx2.Request) -> httpx2.Response:
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
return ok(request)
|
||||
|
||||
ex = exchanger(counting)
|
||||
await ex.exchange("subject", "scope")
|
||||
await ex.exchange("subject", "scope")
|
||||
assert calls == 1
|
||||
|
||||
async def test_expired_entries_are_dropped(self) -> None:
|
||||
ex = exchanger(ok)
|
||||
await ex.exchange("subject", "scope")
|
||||
# Age the only entry, then exchange something else: the prune that runs
|
||||
# on the way out must collect it.
|
||||
for key in list(ex._cache):
|
||||
ex._cache[key] = ExchangedToken("stale", time.monotonic() - 1)
|
||||
await ex.exchange("other", "scope")
|
||||
assert all(not v.expired for v in ex._cache.values())
|
||||
assert len(ex._cache) == 1
|
||||
|
||||
async def test_the_cache_is_bounded(self) -> None:
|
||||
# Every turn brings a new subject token, so without a bound this grows
|
||||
# until the pod does.
|
||||
ex = exchanger(ok)
|
||||
for i in range(MAX_CACHED_TOKENS + 50):
|
||||
await ex.exchange(f"subject-{i}", "scope")
|
||||
assert len(ex._cache) <= MAX_CACHED_TOKENS
|
||||
|
||||
async def test_eviction_drops_the_oldest_first(self) -> None:
|
||||
ex = exchanger(ok)
|
||||
for i in range(MAX_CACHED_TOKENS + 10):
|
||||
await ex.exchange(f"subject-{i}", "scope")
|
||||
assert ("subject-0", "scope") not in ex._cache
|
||||
assert (f"subject-{MAX_CACHED_TOKENS + 9}", "scope") in ex._cache
|
||||
Reference in New Issue
Block a user