diff --git a/k8s/deployment.yaml b/k8s/deployment.yaml index 3712502..2d1a72d 100644 --- a/k8s/deployment.yaml +++ b/k8s/deployment.yaml @@ -28,7 +28,7 @@ spec: # gitea's own registry: the image sits beside the source it was built # from and inherits the repo's public visibility, so the pull needs no # credentials at all. - image: gitea.sttlab.eu/stt/classic-models-agent-assistant:v0.2.0 + image: gitea.sttlab.eu/stt/classic-models-agent-assistant:v0.2.1 ports: - name: http containerPort: 8080 diff --git a/pyproject.toml b/pyproject.toml index 4ae0577..d38f996 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "assistant-agent" -version = "0.2.0" +version = "0.2.1" description = "Assistant — the orchestrating agent for the classic-models estate" readme = "README.md" requires-python = ">=3.12" diff --git a/src/assistant_agent/__pycache__/delegation.cpython-312.pyc b/src/assistant_agent/__pycache__/delegation.cpython-312.pyc index 44e80db..69d2b9e 100644 Binary files a/src/assistant_agent/__pycache__/delegation.cpython-312.pyc and b/src/assistant_agent/__pycache__/delegation.cpython-312.pyc differ diff --git a/src/assistant_agent/__pycache__/identity.cpython-312.pyc b/src/assistant_agent/__pycache__/identity.cpython-312.pyc index 692e52d..85f6494 100644 Binary files a/src/assistant_agent/__pycache__/identity.cpython-312.pyc and b/src/assistant_agent/__pycache__/identity.cpython-312.pyc differ diff --git a/src/assistant_agent/delegation.py b/src/assistant_agent/delegation.py index ef491fc..491adfb 100644 --- a/src/assistant_agent/delegation.py +++ b/src/assistant_agent/delegation.py @@ -13,6 +13,7 @@ client is a long-lived token. from __future__ import annotations import logging +import uuid from dataclasses import dataclass import httpx @@ -72,7 +73,13 @@ class Delegator: try: request = a2a_types.SendMessageRequest( message=a2a_types.Message( - message_id=f"{context_id}-{specialist.name}", + # Unique per message. The context id is the + # conversation, and a conversation makes several calls to + # the same specialist — relaying an approval is exactly + # that. Reusing an id invites a peer to treat the second + # call as a replay of the first, which surfaces as a + # stale answer rather than as an error. + message_id=f"{context_id}-{specialist.name}-{uuid.uuid4().hex[:8]}", # Carried across so the specialist keeps one conversation # per user conversation rather than one per question — # which is what lets its own approval flow resume. diff --git a/src/assistant_agent/identity.py b/src/assistant_agent/identity.py index 7893a2c..53a7c58 100644 --- a/src/assistant_agent/identity.py +++ b/src/assistant_agent/identity.py @@ -24,6 +24,7 @@ from __future__ import annotations import base64 import json import time +from collections import OrderedDict from dataclasses import dataclass import httpx @@ -32,6 +33,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.""" @@ -51,9 +57,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: httpx.AsyncClient, *, token_endpoint: str, client_id: str, client_secret: str) -> None: @@ -61,7 +73,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) @@ -92,8 +104,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/__pycache__/test_token_cache.cpython-312-pytest-9.1.1.pyc b/tests/__pycache__/test_token_cache.cpython-312-pytest-9.1.1.pyc new file mode 100644 index 0000000..43f0b30 Binary files /dev/null and b/tests/__pycache__/test_token_cache.cpython-312-pytest-9.1.1.pyc differ diff --git a/tests/test_token_cache.py b/tests/test_token_cache.py new file mode 100644 index 0000000..6d4ac9e --- /dev/null +++ b/tests/test_token_cache.py @@ -0,0 +1,83 @@ +"""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 pytest + +from assistant_agent.identity import MAX_CACHED_TOKENS, ExchangedToken, TokenExchanger + + +def exchanger(handler) -> TokenExchanger: + transport = httpx.MockTransport(handler) + return TokenExchanger( + httpx.AsyncClient(transport=transport), + token_endpoint="https://keycloak/token", + client_id="agent", + client_secret="secret", + ) + + +def ok(_request: httpx.Request) -> httpx.Response: + return httpx.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: httpx.Request) -> httpx.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