agentsts: allow subject token lookup from session state (#1559)

Runtimes like Vertex do not pass along the HTTP headers to
ADK plugins, and instead we need to rely on passing custom
metadata such as application specific state via the session
state. This change allows retrieving the subject token using
a callback function that defaults to extracting it from the
Authorization header in the session state.

Signed-off-by: Shashank Ram <shashank.ram@solo.io>

Author: shashankram

python/packages/agentsts-adk/src/agentsts/adk/_base.py

@@ -3,7 +3,7 @@
 import inspect
 import logging
 import time
-from typing import Any, Awaitable, Callable, Dict, Optional, Union
+from typing import Awaitable, Callable, Dict, Optional, Union
 
 import jwt
 from google.adk.agents import BaseAgent, LlmAgent
@@ -33,8 +33,27 @@
 HEADERS_KEY = "headers"
 
 
+def _default_get_subject_token(state: dict) -> Optional[str]:
+    """Default subject token retrieval from Authorization header in session state."""
+    headers = state.get(HEADERS_KEY, None)
+    return _extract_jwt_from_headers(headers)
+
+
 class ADKSTSIntegration(STSIntegrationBase):
-    """Google ADK-specific STS integration."""
+    """Google ADK-specific STS integration.
+
+    By default, the subject token is read from the ``Authorization`` header
+    stored in the session state under the ``headers`` key.  To retrieve the
+    subject token from a custom source, pass a ``get_subject_token`` callback::
+
+        integration = ADKSTSIntegration(
+            well_known_uri="https://example.com/.well-known/sts",
+            get_subject_token=lambda state: state.get("my_custom_token_key"),
+        )
+
+    The callback receives ``session.state`` (a dict) and should return the
+    subject token string, or ``None`` if not available.
+    """
 
     def __init__(
         self,
@@ -44,7 +63,7 @@ def __init__(
         timeout: int = 5,
         verify_ssl: bool = True,
         use_issuer_host: bool = False,
-        additional_config: Optional[Dict[str, Any]] = None,
+        get_subject_token: Optional[Callable[[dict], Optional[str]]] = None,
     ):
         """Initialize the ADK STS integration.
 
@@ -55,7 +74,9 @@ def __init__(
             timeout: Request timeout in seconds
             verify_ssl: Whether to verify SSL certificates
             use_issuer_host: Replace the host:port in token_endpoint with the host:port from well_known_uri
-            additional_config: Additional configuration
+            get_subject_token: Optional callback that takes session.state (dict) and returns
+                the subject token string or None. If not set, defaults to extracting the
+                JWT from the Authorization header in session.state["headers"].
         """
         super().__init__(
             well_known_uri=well_known_uri,
@@ -64,7 +85,7 @@ def __init__(
             timeout=timeout,
             verify_ssl=verify_ssl,
             use_issuer_host=use_issuer_host,
-            additional_config=additional_config,
+            get_subject_token=get_subject_token or _default_get_subject_token,
         )
 
 
@@ -143,10 +164,14 @@ async def before_run_callback(
             return None
 
         # No valid cached token, need to get/exchange subject token
-        headers = invocation_context.session.state.get(HEADERS_KEY, None)
-        subject_token = _extract_jwt_from_headers(headers)
+        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("No subject token found in headers for token propagation")
+            logger.debug("subject token not found in session state for token propagation")
             return None
 
         if self.sts_integration:

python/packages/agentsts-adk/tests/test_adk_integration.py

@@ -18,12 +18,14 @@
 class TestADKTokenPropagationPlugin:
     """Unit tests for token propagation plugin covering: none, downstream, and STS exchange."""
 
-    def _make_invocation_context(self, session_id: str, headers: dict | None):
+    def _make_invocation_context(self, session_id: str, headers: dict | None, extra_state: dict | None = None):
         session = Mock()
         session.id = session_id
         session.state = {}
         if headers is not None:
             session.state[HEADERS_KEY] = headers
+        if extra_state is not None:
+            session.state.update(extra_state)
         invocation_context = Mock()
         invocation_context.session = session
         return invocation_context
@@ -48,9 +50,83 @@ async def test_before_run_callback_no_headers(self):
         with patch("agentsts.adk._base.logger") as mock_logger:
             result = await plugin.before_run_callback(invocation_context=ic)
             assert result is None
-            mock_logger.debug.assert_called_once_with("No subject token found in headers for token propagation")
+            mock_logger.debug.assert_called_once_with("subject token not found in session state for token propagation")
         assert plugin.token_cache == {}
 
+    @pytest.mark.asyncio
+    async def test_subject_token_from_callback(self):
+        """Case: get_subject_token callback set -> reads token from session state via callback."""
+        sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
+        sts.get_subject_token = lambda state: state.get("subject-token")
+        sts.fetch_actor_token = None
+        sts._actor_token = "actor-token"
+        sts.exchange_token = AsyncMock(return_value="exchanged-token")
+        plugin = ADKTokenPropagationPlugin(sts)
+        ic = self._make_invocation_context(
+            "sess-key-1",
+            headers=None,
+            extra_state={"subject-token": "subject-jwt-from-vertex"},
+        )
+        result = await plugin.before_run_callback(invocation_context=ic)
+        assert result is None
+        sts.exchange_token.assert_called_once_with(
+            subject_token="subject-jwt-from-vertex",
+            subject_token_type=TokenType.JWT,
+            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"
+
+    @pytest.mark.asyncio
+    async def test_subject_token_callback_returns_none(self):
+        """Case: get_subject_token callback returns None -> returns None."""
+        sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
+        sts.get_subject_token = lambda state: None
+        plugin = ADKTokenPropagationPlugin(sts)
+        ic = self._make_invocation_context("sess-key-2", headers=None)
+        with patch("agentsts.adk._base.logger") as mock_logger:
+            result = await plugin.before_run_callback(invocation_context=ic)
+            assert result is None
+            mock_logger.debug.assert_called_once_with("subject token not found in session state for token propagation")
+        assert plugin.token_cache == {}
+
+    @pytest.mark.asyncio
+    async def test_no_sts_no_headers_returns_none(self):
+        """Case: no STS integration, no headers -> default callback returns None, no token cached."""
+        plugin = ADKTokenPropagationPlugin(sts_integration=None)
+        ic = self._make_invocation_context(
+            "sess-key-3",
+            headers=None,
+        )
+        # Default callback looks for headers, finds none -> no token cached
+        result = await plugin.before_run_callback(invocation_context=ic)
+        assert result is None
+        assert plugin.token_cache == {}
+
+    @pytest.mark.asyncio
+    async def test_default_callback_extracts_from_headers(self):
+        """Case: no get_subject_token callback -> default extracts from headers."""
+        sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
+        sts.get_subject_token = None
+        sts.fetch_actor_token = None
+        sts._actor_token = "actor-token"
+        sts.exchange_token = AsyncMock(return_value="exchanged-via-headers")
+        plugin = ADKTokenPropagationPlugin(sts)
+        ic = self._make_invocation_context("sess-key-4", headers={"Authorization": "Bearer header-jwt"})
+        result = await plugin.before_run_callback(invocation_context=ic)
+        assert result is None
+        sts.exchange_token.assert_called_once_with(
+            subject_token="header-jwt",
+            subject_token_type=TokenType.JWT,
+            actor_token="actor-token",
+            actor_token_type=TokenType.JWT,
+        )
+        assert plugin.token_cache["sess-key-4"].token == "exchanged-via-headers"
+
     @pytest.mark.asyncio
     async def test_downstream_token_propagation_without_sts(self):
         """Case: headers present, no STS integration -> subject token cached and available via header_provider."""
@@ -83,6 +159,7 @@ async def test_downstream_token_propagation_without_sts(self):
     async def test_sts_token_exchange_success(self):
         """Case: STS integration exchanges token -> access token cached and returned by header provider."""
         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="access-token-XYZ")
@@ -115,6 +192,7 @@ async def test_sts_token_exchange_success(self):
     async def test_sts_token_exchange_failure(self):
         """Case: STS exchange raises -> no cache entry, graceful warning."""
         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=Exception("boom"))
@@ -161,6 +239,7 @@ async def test_dynamic_token_fetch_success_sync(self):
         """Case: sync fetch_actor_token is called successfully and token is exchanged."""
         fetch_token_mock = Mock(return_value="dynamic-actor-token")
         sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
         sts.fetch_actor_token = fetch_token_mock
         sts._actor_token = None
         sts.exchange_token = AsyncMock(return_value="access-token-dynamic")
@@ -200,6 +279,7 @@ async def async_fetch_token():
             return "dynamic-actor-token-async"
 
         sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
         sts.fetch_actor_token = async_fetch_token
         sts._actor_token = None
         sts.exchange_token = AsyncMock(return_value="access-token-dynamic-async")
@@ -233,6 +313,7 @@ async def test_dynamic_token_fetch_failure_sync(self):
         """Case: sync fetch_actor_token raises exception -> no token exchange, graceful handling."""
         fetch_token_mock = Mock(side_effect=Exception("Token fetch failed"))
         sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
         sts.fetch_actor_token = fetch_token_mock
         sts._actor_token = None
 
@@ -262,6 +343,7 @@ async def async_fetch_token_failing():
             raise Exception("Async token fetch failed")
 
         sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
         sts.fetch_actor_token = async_fetch_token_failing
         sts._actor_token = None
 
@@ -290,6 +372,7 @@ async def test_dynamic_token_preserved_when_not_expired(self):
         jwt_token = "header.payload.signature"  # Mock JWT
 
         sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
         sts.fetch_actor_token = Mock(return_value="dynamic-actor")
         sts.exchange_token = AsyncMock(return_value=jwt_token)
 
@@ -321,6 +404,7 @@ async def test_dynamic_token_removed_when_expired(self):
         jwt_token = "header.payload.signature"  # Mock JWT
 
         sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
         sts.fetch_actor_token = Mock(return_value="dynamic-actor")
         sts.exchange_token = AsyncMock(return_value=jwt_token)
 
@@ -344,6 +428,7 @@ async def test_dynamic_token_removed_when_expired(self):
     async def test_valid_token_preserved_in_cache(self):
         """Case: valid token (not expired) is preserved in after_run_callback."""
         sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
         sts.fetch_actor_token = None
         sts._actor_token = "static-actor"
         sts.exchange_token = AsyncMock(return_value="access-token-static")
@@ -378,6 +463,7 @@ def sync_fetch_token():
             return "dynamic-actor-token"
 
         sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
         sts.fetch_actor_token = sync_fetch_token
         sts._actor_token = None
         sts.exchange_token = AsyncMock(return_value="access-token")
@@ -417,6 +503,7 @@ def sync_fetch_token():
             return f"dynamic-actor-token-{fetch_count}"
 
         sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
         sts.fetch_actor_token = sync_fetch_token
         sts._actor_token = None
         sts.exchange_token = AsyncMock(return_value="access-token")
@@ -450,6 +537,7 @@ async def test_actor_token_cache_cleanup_on_expiry(self):
         past_expiry = int(time.time()) - 100
 
         sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
         sts.fetch_actor_token = Mock(return_value="dynamic-actor")
         sts._actor_token = None
         sts.exchange_token = AsyncMock(return_value="access-token")
@@ -482,6 +570,7 @@ async def test_actor_token_cache_preserved_when_not_expired(self):
         future_expiry = int(time.time()) + 3600
 
         sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
         sts.fetch_actor_token = Mock(return_value="dynamic-actor")
         sts._actor_token = None
         sts.exchange_token = AsyncMock(return_value="access-token")
@@ -514,6 +603,7 @@ def sync_fetch_token():
             return "actor-token-no-expiry"
 
         sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
         sts.fetch_actor_token = sync_fetch_token
         sts._actor_token = None
         sts.exchange_token = AsyncMock(return_value="access-token")
@@ -548,6 +638,7 @@ async def test_subject_token_cached_and_reused(self):
         future_expiry = int(time.time()) + 3600
 
         sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
         sts.fetch_actor_token = None
         sts._actor_token = "static-actor"
         sts.exchange_token = AsyncMock(return_value="exchanged-token")
@@ -582,6 +673,7 @@ async def test_subject_token_reexchanged_after_expiry(self):
         future_expiry = int(time.time()) + 3600
 
         sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
         sts.fetch_actor_token = None
         sts._actor_token = "static-actor"
         sts.exchange_token = AsyncMock(side_effect=["token-1", "token-2"])
@@ -613,6 +705,7 @@ async def test_subject_token_reexchanged_after_expiry(self):
     async def test_subject_token_cache_no_expiry(self):
         """Case: subject token without expiry is cached indefinitely and reused."""
         sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
         sts.fetch_actor_token = None
         sts._actor_token = "static-actor"
         sts.exchange_token = AsyncMock(return_value="exchanged-token-no-exp")
@@ -650,6 +743,7 @@ async def async_fetch_token():
             return "dynamic-actor-token-async"
 
         sts = Mock(spec=ADKSTSIntegration)
+        sts.get_subject_token = None
         sts.fetch_actor_token = async_fetch_token
         sts._actor_token = None
         sts.exchange_token = AsyncMock(return_value="access-token")

python/packages/agentsts-core/src/agentsts/core/_base.py

@@ -21,7 +21,7 @@ def __init__(
         timeout: int = 30,
         verify_ssl: bool = True,
         use_issuer_host: bool = False,
-        additional_config: Optional[Dict[str, Any]] = None,
+        get_subject_token: Optional[Callable[[dict], Optional[str]]] = None,
     ):
         """Initialize the STS integration.
 
@@ -32,13 +32,14 @@ def __init__(
             timeout: Request timeout in seconds
             verify_ssl: Whether to verify SSL certificates
             use_issuer_host: Replace the host:port in token_endpoint with the host:port from well_known_uri
-            additional_config: Additional configuration for the specific framework
+            get_subject_token: Optional callback that takes session state (dict) and returns
+                the subject token string or None
         """
         self.well_known_uri = well_known_uri
         self.timeout = timeout
         self.verify_ssl = verify_ssl
-        self.additional_config = additional_config or {}
         self.fetch_actor_token = fetch_actor_token
+        self.get_subject_token = get_subject_token
 
         # Initialize STS client
         config = STSConfig(