From f972971d65770666e854904609bbb8a6c76d0979 Mon Sep 17 00:00:00 2001 From: QuentinBisson Date: Sat, 25 Jul 2026 15:34:01 +0200 Subject: [PATCH 1/2] fix(sts): key the ADK token cache by acting subject, not session alone The Python token-propagation plugin cached the exchanged token under the session id, so a session carrying messages from more than one subject reused whichever subject's token seeded the cache first. Every later caller's MCP tool calls then ran under that first subject's delegated identity. The acting subject is now resolved from session state before the cache lookup and folded into the key as the issuer-scoped sub claim, falling back to a hash of the raw token for opaque or sub-less tokens. after_run_callback sweeps every expired entry, since a session now holds one per subject and only the acting subject's key is derivable there. Signed-off-by: QuentinBisson --- .../agentsts-adk/src/agentsts/adk/_base.py | 75 ++++++--- .../tests/test_adk_integration.py | 158 ++++++++++++++---- 2 files changed, 178 insertions(+), 55 deletions(-) diff --git a/python/packages/agentsts-adk/src/agentsts/adk/_base.py b/python/packages/agentsts-adk/src/agentsts/adk/_base.py index 9c8497549..af0a5f8b3 100644 --- a/python/packages/agentsts-adk/src/agentsts/adk/_base.py +++ b/python/packages/agentsts-adk/src/agentsts/adk/_base.py @@ -1,5 +1,6 @@ """Google ADK-specific STS integration.""" +import hashlib import inspect import logging import time @@ -38,6 +39,29 @@ def _default_get_subject_token(state: dict) -> Optional[str]: return _extract_jwt_from_headers(headers) +def _subject_key(token: Optional[str]) -> str: + """Derive a stable per-principal cache discriminator from a bearer token. + + Prefers the issuer-scoped ``sub`` claim. ``sub`` is unique only within an + issuer, so it is paired with ``iss``. Opaque or sub-less tokens fall back to + a hash of the raw token so they still partition per principal. + + NOTE: this parses the token without verification and uses it only to + partition the cache, never to gate a security decision. Tokens are validated + server-side during the STS exchange. + """ + if not token: + return "" + try: + claims = jwt.decode(token, options={"verify_signature": False}) + except Exception: + claims = {} + sub = claims.get("sub") + if sub: + return f"{claims.get('iss', '')}\0{sub}" + return "h:" + hashlib.sha256(token.encode()).hexdigest() + + class ADKSTSIntegration(STSIntegrationBase): """Google ADK-specific STS integration. @@ -156,7 +180,15 @@ async def before_run_callback( invocation_context: InvocationContext, ) -> Optional[dict]: """Propagate token to model before execution.""" - cache_key = self.cache_key(invocation_context) + # Resolve the acting subject before the cache lookup: the cache is keyed + # by subject, and a session carrying messages from several subjects would + # otherwise reuse whichever caller arrived first. + subject_token = self._read_subject_token(invocation_context.session.state) + if not subject_token: + logger.debug("subject token not found in session state for token propagation") + return None + + cache_key = self._cache_key_for(invocation_context.session.id, subject_token) # Check if we have a valid cached subject token cached_entry = self.token_cache.get(cache_key) @@ -168,17 +200,6 @@ async def before_run_callback( logger.debug("Using cached subject token (no expiry)") return None - # No valid cached token, need to get/exchange subject token - get_subject_token = ( - self.sts_integration.get_subject_token - if self.sts_integration and self.sts_integration.get_subject_token - else _default_get_subject_token - ) - subject_token = get_subject_token(invocation_context.session.state) - if not subject_token: - logger.debug("subject token not found in session state for token propagation") - return None - if self.sts_integration: # Get actor token (from cache or fetch dynamically) actor_token = await self._get_actor_token() @@ -208,9 +229,24 @@ async def before_run_callback( logger.debug("Cached new subject token") return None + def _read_subject_token(self, state: dict) -> Optional[str]: + """Resolve the acting caller's subject token from session state.""" + get_subject_token = ( + self.sts_integration.get_subject_token + if self.sts_integration and self.sts_integration.get_subject_token + else _default_get_subject_token + ) + return get_subject_token(state) + + def _cache_key_for(self, session_id: str, subject_token: Optional[str]) -> str: + return f"{session_id}\0{_subject_key(subject_token)}" + def cache_key(self, invocation_context: InvocationContext) -> str: - """Generate a cache key based on the session ID.""" - return invocation_context.session.id + """Key the cache on the session and the acting subject, so a session + carrying messages from several subjects keeps one token per subject + instead of collapsing onto whichever arrived first.""" + session = invocation_context.session + return self._cache_key_for(session.id, self._read_subject_token(session.state)) async def _get_actor_token(self) -> Optional[str]: """Get actor token from cache or fetch dynamically. @@ -264,13 +300,12 @@ async def after_run_callback( invocation_context: InvocationContext, ) -> Optional[dict]: """Clean up expired tokens after run, preserving valid tokens.""" - cache_key = self.cache_key(invocation_context) - cache_entry = self.token_cache.get(cache_key) - - # Clean up subject token cache - only remove if expired - if cache_entry and _has_token_expired(cache_entry.expiry): + # A session now holds one entry per subject, and only the acting + # subject's key is derivable here, so sweep every expired entry rather + # than leaving the other subjects' behind. + for key in [key for key, entry in self.token_cache.items() if _has_token_expired(entry.expiry)]: logger.debug("Removing expired subject token from cache") - self.token_cache.pop(cache_key, None) + self.token_cache.pop(key, None) # Clean up expired actor token cache if self.actor_token_cache and _has_token_expired(self.actor_token_cache.expiry): diff --git a/python/packages/agentsts-adk/tests/test_adk_integration.py b/python/packages/agentsts-adk/tests/test_adk_integration.py index 6915041c3..f226ac377 100644 --- a/python/packages/agentsts-adk/tests/test_adk_integration.py +++ b/python/packages/agentsts-adk/tests/test_adk_integration.py @@ -1,5 +1,6 @@ """Tests for ADK integration classes (STS + token propagation).""" +import time from unittest.mock import AsyncMock, Mock, patch import pytest @@ -13,6 +14,7 @@ from agentsts.adk._base import _extract_jwt_expiry as extract_jwt_expiry from agentsts.adk._base import _extract_jwt_from_headers as extract_jwt_from_headers from agentsts.adk._base import _has_token_expired as has_token_expired +from agentsts.adk._base import _subject_key as subject_key class TestADKTokenPropagationPlugin: @@ -76,8 +78,8 @@ async def test_subject_token_from_callback(self): actor_token="actor-token", actor_token_type=TokenType.JWT, ) - assert "sess-key-1" in plugin.token_cache - assert plugin.token_cache["sess-key-1"].token == "exchanged-token" + assert plugin.cache_key(ic) in plugin.token_cache + assert plugin.token_cache[plugin.cache_key(ic)].token == "exchanged-token" @pytest.mark.asyncio async def test_subject_token_callback_returns_none(self): @@ -125,7 +127,7 @@ async def test_default_callback_extracts_from_headers(self): actor_token="actor-token", actor_token_type=TokenType.JWT, ) - assert plugin.token_cache["sess-key-4"].token == "exchanged-via-headers" + assert plugin.token_cache[plugin.cache_key(ic)].token == "exchanged-via-headers" @pytest.mark.asyncio async def test_downstream_token_propagation_without_sts(self): @@ -134,8 +136,8 @@ async def test_downstream_token_propagation_without_sts(self): ic = self._make_invocation_context("sess-2", headers={"Authorization": "Bearer subj-token-123"}) result = await plugin.before_run_callback(invocation_context=ic) assert result is None - assert "sess-2" in plugin.token_cache - assert plugin.token_cache["sess-2"].token == "subj-token-123" + assert plugin.cache_key(ic) in plugin.token_cache + assert plugin.token_cache[plugin.cache_key(ic)].token == "subj-token-123" # propagate toolset mcp_toolset = Mock(spec=MCPToolset) @@ -153,7 +155,7 @@ async def test_downstream_token_propagation_without_sts(self): # cleanup - token should still be cached if not expired await plugin.after_run_callback(invocation_context=ic) # Token has no expiry, so it's preserved - assert "sess-2" in plugin.token_cache + assert plugin.cache_key(ic) in plugin.token_cache @pytest.mark.asyncio async def test_sts_token_exchange_success(self): @@ -176,8 +178,8 @@ async def test_sts_token_exchange_success(self): ) # optional debug log length check mock_logger.debug.assert_called() # at least one debug log - assert "sess-3" in plugin.token_cache - assert plugin.token_cache["sess-3"].token == "access-token-XYZ" + assert plugin.cache_key(ic) in plugin.token_cache + assert plugin.token_cache[plugin.cache_key(ic)].token == "access-token-XYZ" ro_ctx = self._make_readonly_context(ic) headers = plugin.header_provider(ro_ctx) @@ -186,7 +188,7 @@ async def test_sts_token_exchange_success(self): # cleanup - token should still be cached if not expired await plugin.after_run_callback(invocation_context=ic) # Token has no expiry, so it's preserved - assert "sess-3" in plugin.token_cache + assert plugin.cache_key(ic) in plugin.token_cache @pytest.mark.asyncio async def test_sts_token_exchange_failure(self): @@ -202,7 +204,7 @@ async def test_sts_token_exchange_failure(self): result = await plugin.before_run_callback(invocation_context=ic) assert result is None mock_logger.warning.assert_called_once() - assert "sess-4" not in plugin.token_cache + assert plugin.cache_key(ic) not in plugin.token_cache # header provider should yield empty dict ro_ctx = self._make_readonly_context(ic) assert plugin.header_provider(ro_ctx) == {} @@ -228,11 +230,11 @@ async def test_after_run_callback_removes_expired_token(self): # Mock expiry to return expired timestamp with patch("agentsts.adk._base._extract_jwt_expiry", return_value=past_expiry): await plugin.before_run_callback(invocation_context=ic) - assert "sess-6" in plugin.token_cache + assert plugin.cache_key(ic) in plugin.token_cache # Token is expired, should be removed await plugin.after_run_callback(invocation_context=ic) - assert "sess-6" not in plugin.token_cache + assert plugin.cache_key(ic) not in plugin.token_cache @pytest.mark.asyncio async def test_dynamic_token_fetch_success_sync(self): @@ -267,8 +269,8 @@ async def test_dynamic_token_fetch_success_sync(self): assert any("Fetched and cached new actor token" in call for call in debug_calls) # Verify token is cached - assert "sess-7" in plugin.token_cache - cache_entry = plugin.token_cache["sess-7"] + assert plugin.cache_key(ic) in plugin.token_cache + cache_entry = plugin.token_cache[plugin.cache_key(ic)] assert cache_entry.token == "access-token-dynamic" @pytest.mark.asyncio @@ -304,8 +306,8 @@ async def async_fetch_token(): assert any("Fetched and cached new actor token" in call for call in debug_calls) # Verify token is cached - assert "sess-7a" in plugin.token_cache - cache_entry = plugin.token_cache["sess-7a"] + assert plugin.cache_key(ic) in plugin.token_cache + cache_entry = plugin.token_cache[plugin.cache_key(ic)] assert cache_entry.token == "access-token-dynamic-async" @pytest.mark.asyncio @@ -333,7 +335,7 @@ async def test_dynamic_token_fetch_failure_sync(self): assert "Failed to fetch actor token dynamically" in warning_msg # No token should be cached - assert "sess-8" not in plugin.token_cache + assert plugin.cache_key(ic) not in plugin.token_cache @pytest.mark.asyncio async def test_dynamic_token_fetch_failure_async(self): @@ -360,7 +362,7 @@ async def async_fetch_token_failing(): assert "Failed to fetch actor token dynamically" in warning_msg # No token should be cached - assert "sess-8a" not in plugin.token_cache + assert plugin.cache_key(ic) not in plugin.token_cache @pytest.mark.asyncio async def test_dynamic_token_preserved_when_not_expired(self): @@ -384,15 +386,15 @@ async def test_dynamic_token_preserved_when_not_expired(self): await plugin.before_run_callback(invocation_context=ic) # Verify token is cached with expiry - assert "sess-9" in plugin.token_cache - cache_entry = plugin.token_cache["sess-9"] + assert plugin.cache_key(ic) in plugin.token_cache + cache_entry = plugin.token_cache[plugin.cache_key(ic)] assert cache_entry.expiry == future_expiry # Call after_run_callback - token should be preserved await plugin.after_run_callback(invocation_context=ic) # Verify token is still cached (not expired) - assert "sess-9" in plugin.token_cache + assert plugin.cache_key(ic) in plugin.token_cache @pytest.mark.asyncio async def test_dynamic_token_removed_when_expired(self): @@ -416,13 +418,13 @@ async def test_dynamic_token_removed_when_expired(self): await plugin.before_run_callback(invocation_context=ic) # Verify token is cached - assert "sess-10" in plugin.token_cache + assert plugin.cache_key(ic) in plugin.token_cache # Call after_run_callback - token should be removed (expired) await plugin.after_run_callback(invocation_context=ic) # Verify token is removed - assert "sess-10" not in plugin.token_cache + assert plugin.cache_key(ic) not in plugin.token_cache @pytest.mark.asyncio async def test_valid_token_preserved_in_cache(self): @@ -439,15 +441,15 @@ async def test_valid_token_preserved_in_cache(self): await plugin.before_run_callback(invocation_context=ic) # Verify token is cached with expected value - assert "sess-11" in plugin.token_cache - cache_entry = plugin.token_cache["sess-11"] + assert plugin.cache_key(ic) in plugin.token_cache + cache_entry = plugin.token_cache[plugin.cache_key(ic)] assert cache_entry.token == "access-token-static" # Call after_run_callback - token should still be in cache if not expired await plugin.after_run_callback(invocation_context=ic) # Verify the same cache entry is still present - assert plugin.token_cache.get("sess-11") is cache_entry + assert plugin.token_cache.get(plugin.cache_key(ic)) is cache_entry @pytest.mark.asyncio async def test_actor_token_cached_and_reused(self): @@ -655,8 +657,8 @@ async def test_subject_token_cached_and_reused(self): assert sts.exchange_token.call_count == 1 # Verify token is cached - assert "sess-18" in plugin.token_cache - assert plugin.token_cache["sess-18"].token == "exchanged-token" + assert plugin.cache_key(ic) in plugin.token_cache + assert plugin.token_cache[plugin.cache_key(ic)].token == "exchanged-token" # Second call with same session - should use cached token await plugin.before_run_callback(invocation_context=ic) @@ -685,12 +687,12 @@ async def test_subject_token_reexchanged_after_expiry(self): with patch("agentsts.adk._base._extract_jwt_expiry", return_value=past_expiry): await plugin.before_run_callback(invocation_context=ic) assert sts.exchange_token.call_count == 1 - assert plugin.token_cache["sess-19"].token == "token-1" - assert plugin.token_cache["sess-19"].expiry == past_expiry + assert plugin.token_cache[plugin.cache_key(ic)].token == "token-1" + assert plugin.token_cache[plugin.cache_key(ic)].expiry == past_expiry # Cleanup expired token await plugin.after_run_callback(invocation_context=ic) - assert "sess-19" not in plugin.token_cache + assert plugin.cache_key(ic) not in plugin.token_cache # Second call - should detect missing cache and re-exchange with patch("agentsts.adk._base._extract_jwt_expiry", return_value=future_expiry): @@ -698,8 +700,8 @@ async def test_subject_token_reexchanged_after_expiry(self): # Verify exchange was called again assert sts.exchange_token.call_count == 2 - assert plugin.token_cache["sess-19"].token == "token-2" - assert plugin.token_cache["sess-19"].expiry == future_expiry + assert plugin.token_cache[plugin.cache_key(ic)].token == "token-2" + assert plugin.token_cache[plugin.cache_key(ic)].expiry == future_expiry @pytest.mark.asyncio async def test_subject_token_cache_no_expiry(self): @@ -718,11 +720,11 @@ async def test_subject_token_cache_no_expiry(self): # First call await plugin.before_run_callback(invocation_context=ic) assert sts.exchange_token.call_count == 1 - assert plugin.token_cache["sess-20"].expiry is None + assert plugin.token_cache[plugin.cache_key(ic)].expiry is None # after_run_callback should preserve it (no expiry) await plugin.after_run_callback(invocation_context=ic) - assert "sess-20" in plugin.token_cache + assert plugin.cache_key(ic) in plugin.token_cache # Second call - should reuse cached token await plugin.before_run_callback(invocation_context=ic) @@ -978,6 +980,92 @@ def test_expiry_end_to_end_with_real_jwt(self): assert extract_jwt_expiry(no_exp_token) is None assert has_token_expired(extract_jwt_expiry(no_exp_token)) is False + @staticmethod + def _jwt(iss: str, sub: str) -> str: + import jwt as pyjwt + + return pyjwt.encode({"iss": iss, "sub": sub}, "secret", algorithm="HS256") + + @pytest.mark.asyncio + async def test_two_subjects_in_one_session_keep_separate_tokens(self): + """Case: one session, two callers -> each keeps and reuses its own exchanged token.""" + alice = self._jwt("https://dex.example", "alice") + bob = self._jwt("https://dex.example", "bob") + + sts = Mock(spec=ADKSTSIntegration) + sts.get_subject_token = None + sts.fetch_actor_token = None + sts._actor_token = "actor-token" + sts.exchange_token = AsyncMock(side_effect=lambda subject_token, **_: f"exchanged-for-{subject_token[-5:]}") + plugin = ADKTokenPropagationPlugin(sts) + + ic_alice = self._make_invocation_context("shared-sess", headers={"Authorization": f"Bearer {alice}"}) + ic_bob = self._make_invocation_context("shared-sess", headers={"Authorization": f"Bearer {bob}"}) + + await plugin.before_run_callback(invocation_context=ic_alice) + await plugin.before_run_callback(invocation_context=ic_bob) + + # Both callers share a session id, so a session-only key would have + # collapsed them onto whichever exchanged first. + assert plugin.cache_key(ic_alice) != plugin.cache_key(ic_bob) + assert len(plugin.token_cache) == 2 + assert sts.exchange_token.await_count == 2 + + # Each caller's tool invocation gets its own token back. + alice_headers = plugin.header_provider(self._make_readonly_context(ic_alice)) + bob_headers = plugin.header_provider(self._make_readonly_context(ic_bob)) + assert alice_headers["Authorization"] == f"Bearer {plugin.token_cache[plugin.cache_key(ic_alice)].token}" + assert alice_headers != bob_headers + + @pytest.mark.asyncio + async def test_same_subject_reuses_cached_token(self): + """Case: same caller twice in one session -> exchanged once, second call is a cache hit.""" + alice = self._jwt("https://dex.example", "alice") + + sts = Mock(spec=ADKSTSIntegration) + sts.get_subject_token = None + sts.fetch_actor_token = None + sts._actor_token = "actor-token" + sts.exchange_token = AsyncMock(return_value="exchanged-once") + plugin = ADKTokenPropagationPlugin(sts) + + ic = self._make_invocation_context("shared-sess", headers={"Authorization": f"Bearer {alice}"}) + await plugin.before_run_callback(invocation_context=ic) + await plugin.before_run_callback(invocation_context=ic) + + assert sts.exchange_token.await_count == 1 + assert len(plugin.token_cache) == 1 + + def test_subject_key_does_not_collide_across_issuers(self): + """Case: same sub at two issuers -> distinct keys, since sub is issuer-scoped.""" + assert subject_key(self._jwt("https://iss-a", "same")) != subject_key(self._jwt("https://iss-b", "same")) + + def test_subject_key_falls_back_to_hash_for_opaque_tokens(self): + """Case: non-JWT (or sub-less) tokens still partition per principal.""" + assert subject_key("opaque-a") != subject_key("opaque-b") + assert subject_key("opaque-a") == subject_key("opaque-a") + assert subject_key("opaque-a").startswith("h:") + assert subject_key(None) == "" + + @pytest.mark.asyncio + async def test_after_run_sweeps_other_subjects_expired_entries(self): + """Case: expired entry for a subject other than the acting one is still evicted.""" + alice = self._jwt("https://dex.example", "alice") + bob = self._jwt("https://dex.example", "bob") + + plugin = ADKTokenPropagationPlugin(sts_integration=None) + ic_alice = self._make_invocation_context("shared-sess", headers={"Authorization": f"Bearer {alice}"}) + ic_bob = self._make_invocation_context("shared-sess", headers={"Authorization": f"Bearer {bob}"}) + + past_expiry = int(time.time()) - 100 + with patch("agentsts.adk._base._extract_jwt_expiry", return_value=past_expiry): + await plugin.before_run_callback(invocation_context=ic_alice) + assert plugin.cache_key(ic_alice) in plugin.token_cache + + # bob runs; alice's entry has expired and must not survive the sweep. + await plugin.after_run_callback(invocation_context=ic_bob) + assert plugin.cache_key(ic_alice) not in plugin.token_cache + class TestADKSTSIntegration: """Test cases for ADKSTSIntegration.""" From be7a504b0fb0b886b1b880a01d14b7ea0d60277f Mon Sep 17 00:00:00 2001 From: QuentinBisson Date: Sat, 25 Jul 2026 15:42:43 +0200 Subject: [PATCH 2/2] test(sts): cover the fail-closed read path and tighten the subject-key note Signed-off-by: QuentinBisson --- .../agentsts-adk/src/agentsts/adk/_base.py | 6 ++++-- .../agentsts-adk/tests/test_adk_integration.py | 18 +++++++++++++++++- 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/python/packages/agentsts-adk/src/agentsts/adk/_base.py b/python/packages/agentsts-adk/src/agentsts/adk/_base.py index af0a5f8b3..1eb3ee918 100644 --- a/python/packages/agentsts-adk/src/agentsts/adk/_base.py +++ b/python/packages/agentsts-adk/src/agentsts/adk/_base.py @@ -47,8 +47,10 @@ def _subject_key(token: Optional[str]) -> str: a hash of the raw token so they still partition per principal. NOTE: this parses the token without verification and uses it only to - partition the cache, never to gate a security decision. Tokens are validated - server-side during the STS exchange. + partition the cache. On a cache hit the key selects which cached delegated + token the caller receives, so authenticating the inbound bearer remains the + caller's (upstream) responsibility; tokens are validated server-side during + the STS exchange on a miss. """ if not token: return "" diff --git a/python/packages/agentsts-adk/tests/test_adk_integration.py b/python/packages/agentsts-adk/tests/test_adk_integration.py index f226ac377..669c26c79 100644 --- a/python/packages/agentsts-adk/tests/test_adk_integration.py +++ b/python/packages/agentsts-adk/tests/test_adk_integration.py @@ -984,7 +984,7 @@ def test_expiry_end_to_end_with_real_jwt(self): def _jwt(iss: str, sub: str) -> str: import jwt as pyjwt - return pyjwt.encode({"iss": iss, "sub": sub}, "secret", algorithm="HS256") + return pyjwt.encode({"iss": iss, "sub": sub}, "test-signing-key-of-at-least-32-bytes!", algorithm="HS256") @pytest.mark.asyncio async def test_two_subjects_in_one_session_keep_separate_tokens(self): @@ -1066,6 +1066,22 @@ async def test_after_run_sweeps_other_subjects_expired_entries(self): await plugin.after_run_callback(invocation_context=ic_bob) assert plugin.cache_key(ic_alice) not in plugin.token_cache + @pytest.mark.asyncio + async def test_tokenless_caller_does_not_reuse_cached_session_token(self): + """Case: session holds a cached token, but a caller with no subject token gets nothing back.""" + alice = self._jwt("https://dex.example", "alice") + + plugin = ADKTokenPropagationPlugin(sts_integration=None) + ic_alice = self._make_invocation_context("shared-sess", headers={"Authorization": f"Bearer {alice}"}) + await plugin.before_run_callback(invocation_context=ic_alice) + assert len(plugin.token_cache) == 1 + + ic_anon = self._make_invocation_context("shared-sess", headers=None) + await plugin.before_run_callback(invocation_context=ic_anon) + assert len(plugin.token_cache) == 1 + + assert plugin.header_provider(self._make_readonly_context(ic_anon)) == {} + class TestADKSTSIntegration: """Test cases for ADKSTSIntegration."""