feat(agentsts): move STS token exchange to before_tool_callback

Move STS token exchange from before_run_callback (eager, once per
invocation) to before_tool_callback (lazy, per MCP tool call). This
avoids the sync/async problem in header_provider and only performs the
exchange when an MCP tool is actually invoked.

- before_run_callback now only extracts and stores the subject token
- before_tool_callback exchanges on first McpTool call per session
- Non-MCP tools (memory, AgentTool, etc.) are skipped
- Cached tokens are reused for subsequent MCP calls in same session
- Sets the stage for per-audience token exchange

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: syn-zhu

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

@@ -14,6 +14,7 @@
 from google.adk.sessions.session import Session
 from google.adk.tools.base_tool import BaseTool
 from google.adk.tools.mcp_tool import MCPTool
+from google.adk.tools.mcp_tool.mcp_tool import McpTool
 from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
 from google.adk.tools.tool_context import ToolContext
 from typing_extensions import override
@@ -40,7 +41,18 @@ def __init__(
 
 
 class ADKTokenPropagationPlugin(BasePlugin):
-    """Plugin for propagating STS tokens to ADK tools."""
+    """Plugin for propagating STS tokens to ADK tools.
+
+    Token exchange lifecycle:
+    1. before_run_callback: extracts the subject token from request headers
+       and stores it for the duration of the invocation.
+    2. before_tool_callback: when an MCP tool is about to be called and STS
+       is configured, exchanges the subject token (async) and caches the
+       access token so header_provider can read it.
+    3. header_provider (sync): returns cached access token as Authorization
+       header -- called by McpToolset/McpTool during MCP session setup.
+    4. after_run_callback: cleans up all cached state.
+    """
 
     def __init__(self, sts_integration: Optional[STSIntegrationBase] = None):
         """Initialize the token propagation plugin.
@@ -51,6 +63,7 @@ def __init__(self, sts_integration: Optional[STSIntegrationBase] = None):
         super().__init__("ADKTokenPropagationPlugin")
         self.sts_integration = sts_integration
         self.token_cache: Dict[str, str] = {}
+        self._subject_tokens: Dict[str, str] = {}
 
     def add_to_agent(self, agent: BaseAgent):
         """
@@ -70,7 +83,6 @@ def add_to_agent(self, agent: BaseAgent):
                 logger.debug("Updated tool connection params to include access token from STS server")
 
     def header_provider(self, readonly_context: Optional[ReadonlyContext]) -> Dict[str, str]:
-        # access save token
         access_token = self.token_cache.get(self.cache_key(readonly_context._invocation_context), "")
         if not access_token:
             return {}
@@ -85,25 +97,58 @@ async def before_run_callback(
         *,
         invocation_context: InvocationContext,
     ) -> Optional[dict]:
-        """Propagate token to model before execution."""
+        """Extract and store the subject token for later exchange."""
         headers = invocation_context.session.state.get(HEADERS_KEY, None)
         subject_token = _extract_jwt_from_headers(headers)
         if not subject_token:
             logger.debug("No subject token found in headers for token propagation")
             return None
-        if self.sts_integration:
-            try:
-                subject_token = await self.sts_integration.exchange_token(
-                    subject_token=subject_token,
-                    subject_token_type=TokenType.JWT,
-                    actor_token=self.sts_integration._actor_token,
-                    actor_token_type=TokenType.JWT if self.sts_integration._actor_token else None,
-                )
-            except Exception as e:
-                logger.warning(f"STS token exchange failed: {e}")
-                return None
-        # no sts, just propagate the subject token upstream
-        self.token_cache[self.cache_key(invocation_context)] = subject_token
+        key = self.cache_key(invocation_context)
+        self._subject_tokens[key] = subject_token
+        if not self.sts_integration:
+            # No STS -- propagate the subject token directly so
+            # header_provider can return it on the first tool call.
+            self.token_cache[key] = subject_token
+        return None
+
+    @override
+    async def before_tool_callback(
+        self,
+        *,
+        tool: BaseTool,
+        tool_args: dict[str, Any],
+        tool_context: ToolContext,
+    ) -> Optional[dict]:
+        """Exchange the subject token via STS before each MCP tool call."""
+        if not self.sts_integration:
+            return None
+        # Only exchange tokens for MCP tool calls. Other tool types
+        # (memory tools, AgentTool, etc.) don't use header_provider and
+        # have their own auth mechanisms, so exchanging here would be a
+        # wasted HTTP round-trip to the STS.
+        if not isinstance(tool, McpTool):
+            return None
+
+        key = self.cache_key(tool_context._invocation_context)
+        # Already exchanged for this session
+        if key in self.token_cache:
+            return None
+
+        subject_token = self._subject_tokens.get(key)
+        if not subject_token:
+            return None
+
+        try:
+            access_token = await self.sts_integration.exchange_token(
+                subject_token=subject_token,
+                subject_token_type=TokenType.JWT,
+                actor_token=self.sts_integration._actor_token,
+                actor_token_type=TokenType.JWT if self.sts_integration._actor_token else None,
+            )
+            self.token_cache[key] = access_token
+        except Exception as e:
+            logger.warning(f"STS token exchange failed: {e}")
+
         return None
 
     def cache_key(self, invocation_context: InvocationContext) -> str:
@@ -116,8 +161,9 @@ async def after_run_callback(
         *,
         invocation_context: InvocationContext,
     ) -> Optional[dict]:
-        # delete token after run
-        self.token_cache.pop(self.cache_key(invocation_context), None)
+        key = self.cache_key(invocation_context)
+        self.token_cache.pop(key, None)
+        self._subject_tokens.pop(key, None)
         return None
 
 

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

@@ -4,6 +4,7 @@
 
 import pytest
 from google.adk.agents import LlmAgent
+from google.adk.tools.mcp_tool.mcp_tool import McpTool
 from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
 
 from agentsts.adk import ADKSTSIntegration, ADKTokenPropagationPlugin
@@ -31,6 +32,16 @@ def _make_readonly_context(self, invocation_context):
         readonly_context._invocation_context = invocation_context
         return readonly_context
 
+    def _make_tool_context(self, invocation_context):
+        tool_context = Mock()
+        tool_context._invocation_context = invocation_context
+        return tool_context
+
+    def _make_mcp_tool(self):
+        tool = Mock(spec=McpTool)
+        tool.name = "test_tool"
+        return tool
+
     def test_init(self):
         mock_sts_integration = Mock()
         plugin = ADKTokenPropagationPlugin(mock_sts_integration)
@@ -77,49 +88,102 @@ async def test_downstream_token_propagation_without_sts(self):
 
     @pytest.mark.asyncio
     async def test_sts_token_exchange_success(self):
-        """Case: STS integration exchanges token -> access token cached and returned by header provider."""
+        """Case: STS integration -- before_run stores subject token, before_tool_callback exchanges it."""
         sts = Mock(spec=ADKSTSIntegration)
         sts._actor_token = "actor-token"
         sts.exchange_token = AsyncMock(return_value="access-token-XYZ")
         plugin = ADKTokenPropagationPlugin(sts)
         ic = self._make_invocation_context("sess-3", headers={"Authorization": "Bearer original-subject"})
-        with patch("agentsts.adk._base.logger") as mock_logger:
-            result = await plugin.before_run_callback(invocation_context=ic)
-            assert result is None
-            sts.exchange_token.assert_called_once_with(
-                subject_token="original-subject",
-                subject_token_type=TokenType.JWT,
-                actor_token="actor-token",
-                actor_token_type=TokenType.JWT,
-            )
-            # optional debug log length check
-            mock_logger.debug.assert_called()  # at least one debug log
+
+        # before_run_callback should store the subject token but NOT exchange
+        result = await plugin.before_run_callback(invocation_context=ic)
+        assert result is None
+        sts.exchange_token.assert_not_called()
+        assert "sess-3" not in plugin.token_cache
+        assert plugin._subject_tokens["sess-3"] == "original-subject"
+
+        # before_tool_callback should exchange on first MCP tool call
+        tool = self._make_mcp_tool()
+        tc = self._make_tool_context(ic)
+        result = await plugin.before_tool_callback(tool=tool, tool_args={}, tool_context=tc)
+        assert result is None
+        sts.exchange_token.assert_called_once_with(
+            subject_token="original-subject",
+            subject_token_type=TokenType.JWT,
+            actor_token="actor-token",
+            actor_token_type=TokenType.JWT,
+        )
         assert plugin.token_cache["sess-3"] == "access-token-XYZ"
 
+        # header_provider should return the exchanged token
         ro_ctx = self._make_readonly_context(ic)
         headers = plugin.header_provider(ro_ctx)
         assert headers == {"Authorization": "Bearer access-token-XYZ"}
 
+        # second tool call should not exchange again (cached)
+        sts.exchange_token.reset_mock()
+        result = await plugin.before_tool_callback(tool=tool, tool_args={}, tool_context=tc)
+        assert result is None
+        sts.exchange_token.assert_not_called()
+
         await plugin.after_run_callback(invocation_context=ic)
         assert "sess-3" not in plugin.token_cache
+        assert "sess-3" not in plugin._subject_tokens
 
     @pytest.mark.asyncio
     async def test_sts_token_exchange_failure(self):
-        """Case: STS exchange raises -> no cache entry, graceful warning."""
+        """Case: STS exchange raises in before_tool_callback -> no cache entry, graceful warning."""
         sts = Mock(spec=ADKSTSIntegration)
         sts._actor_token = "actor-token"
         sts.exchange_token = AsyncMock(side_effect=Exception("boom"))
         plugin = ADKTokenPropagationPlugin(sts)
         ic = self._make_invocation_context("sess-4", headers={"Authorization": "Bearer original-subject"})
+
+        await plugin.before_run_callback(invocation_context=ic)
+        assert plugin._subject_tokens["sess-4"] == "original-subject"
+
+        tool = self._make_mcp_tool()
+        tc = self._make_tool_context(ic)
         with patch("agentsts.adk._base.logger") as mock_logger:
-            result = await plugin.before_run_callback(invocation_context=ic)
+            result = await plugin.before_tool_callback(tool=tool, tool_args={}, tool_context=tc)
             assert result is None
             mock_logger.warning.assert_called_once()
         assert "sess-4" not in plugin.token_cache
+
         # header provider should yield empty dict
         ro_ctx = self._make_readonly_context(ic)
         assert plugin.header_provider(ro_ctx) == {}
 
+    @pytest.mark.asyncio
+    async def test_before_tool_callback_skips_non_mcp_tools(self):
+        """Case: before_tool_callback ignores non-MCP tools."""
+        sts = Mock(spec=ADKSTSIntegration)
+        sts._actor_token = "actor-token"
+        sts.exchange_token = AsyncMock(return_value="access-token")
+        plugin = ADKTokenPropagationPlugin(sts)
+        ic = self._make_invocation_context("sess-7", headers={"Authorization": "Bearer subj"})
+        await plugin.before_run_callback(invocation_context=ic)
+
+        non_mcp_tool = Mock()  # not a McpTool
+        tc = self._make_tool_context(ic)
+        result = await plugin.before_tool_callback(tool=non_mcp_tool, tool_args={}, tool_context=tc)
+        assert result is None
+        sts.exchange_token.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_before_tool_callback_no_sts(self):
+        """Case: before_tool_callback is a no-op without STS integration."""
+        plugin = ADKTokenPropagationPlugin(sts_integration=None)
+        ic = self._make_invocation_context("sess-8", headers={"Authorization": "Bearer subj"})
+        await plugin.before_run_callback(invocation_context=ic)
+
+        tool = self._make_mcp_tool()
+        tc = self._make_tool_context(ic)
+        result = await plugin.before_tool_callback(tool=tool, tool_args={}, tool_context=tc)
+        assert result is None
+        # token_cache should still have the subject token from before_run
+        assert plugin.token_cache["sess-8"] == "subj"
+
     def test_header_provider_no_entry(self):
         """Case: header_provider called with no cached token -> returns empty dict."""
         plugin = ADKTokenPropagationPlugin()
@@ -131,13 +195,15 @@ def test_header_provider_no_entry(self):
 
     @pytest.mark.asyncio
     async def test_after_run_callback_removes_token(self):
-        """Case: after_run_callback removes cached token."""
+        """Case: after_run_callback removes cached token and subject token."""
         plugin = ADKTokenPropagationPlugin()
         ic = self._make_invocation_context("sess-6", headers={"Authorization": "Bearer AAA"})
         await plugin.before_run_callback(invocation_context=ic)
         assert "sess-6" in plugin.token_cache
+        assert "sess-6" in plugin._subject_tokens
         await plugin.after_run_callback(invocation_context=ic)
         assert "sess-6" not in plugin.token_cache
+        assert "sess-6" not in plugin._subject_tokens
 
     def test_extract_jwt_from_headers_success(self):
         """Test successful JWT extraction from headers."""