diff --git a/k8s/deployment.yaml b/k8s/deployment.yaml index 0bbfc46..c2ba892 100644 --- a/k8s/deployment.yaml +++ b/k8s/deployment.yaml @@ -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 diff --git a/pyproject.toml b/pyproject.toml index e7529e3..879b890 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/src/catalogue_agent/identity.py b/src/catalogue_agent/identity.py index 98f1cc7..f8cc800 100644 --- a/src/catalogue_agent/identity.py +++ b/src/catalogue_agent/identity.py @@ -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. diff --git a/tests/test_token_cache.py b/tests/test_token_cache.py new file mode 100644 index 0000000..da79006 --- /dev/null +++ b/tests/test_token_cache.py @@ -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