diff --git a/python/packages/agentsts-adk/src/agentsts/adk/_base.py b/python/packages/agentsts-adk/src/agentsts/adk/_base.py index 9c8497549..1eb3ee918 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,31 @@ 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. 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 "" + 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 +182,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 +202,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 +231,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 +302,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..669c26c79 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,108 @@ 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}, "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): + """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 + + @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."""