fix: remove dead code and make mock STS server audience-aware

- Remove unused imports from _base.py (BaseAgent, LlmAgent, AuthCredential, etc.)
- Remove add_to_agent method that would overwrite audience-aware closures
  created by create_header_provider in types.py
- Remove LlmAgent import and add_to_agent test code from test_adk_integration.py
- Make mock STS server include 'aud' claim in generated tokens when audience
  is provided, improving E2E test coverage for per-audience token exchange

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Simon Zhu <simon.zhu@mongodb.com>

Author: syn-zhu

go/core/test/e2e/mocks/mock_sts_server.go

@@ -178,7 +178,7 @@ func (m *MockSTSServer) handleTokenExchange(w http.ResponseWriter, r *http.Reque
 		return
 	}
 
-	accessToken, err := m.generateMockAccessToken(req.SubjectToken)
+	accessToken, err := m.generateMockAccessToken(req.SubjectToken, req.Audience)
 	if err != nil {
 		http.Error(w, fmt.Sprintf("Error generating mock access token: %v", err), http.StatusBadRequest)
 		return
@@ -200,7 +200,7 @@ func (m *MockSTSServer) handleTokenExchange(w http.ResponseWriter, r *http.Reque
 	m.requests = append(m.requests, req)
 }
 
-func (m *MockSTSServer) generateMockAccessToken(subjectToken string) (string, error) {
+func (m *MockSTSServer) generateMockAccessToken(subjectToken string, audience string) (string, error) {
 	// Try to parse JWT token to extract subject claim
 	subject, err := extractSubjectFromJWT(subjectToken)
 	if err != nil {
@@ -219,6 +219,10 @@ func (m *MockSTSServer) generateMockAccessToken(subjectToken string) (string, er
 		"iss":   "mock-sts-server",
 	}
 
+	if audience != "" {
+		tokenData["aud"] = audience
+	}
+
 	// For testing purposes, we'll return a simple JSON string
 	// In a real implementation, this would be a signed JWT
 	tokenBytes, err := json.Marshal(tokenData)

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

@@ -3,15 +3,9 @@
 import logging
 from typing import Any, Dict, Optional
 
-from google.adk.agents import BaseAgent, LlmAgent
 from google.adk.agents.invocation_context import InvocationContext
 from google.adk.agents.readonly_context import ReadonlyContext
-from google.adk.auth.auth_credential import AuthCredential, AuthCredentialTypes, HttpAuth, HttpCredentials
-from google.adk.events.event import Event
 from google.adk.plugins.base_plugin import BasePlugin
-from google.adk.runners import Runner
-from google.adk.sessions import BaseSessionService
-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_toolset import MCPToolset
@@ -67,23 +61,6 @@ def register_toolset(self, toolset: MCPToolset, audience: Optional[str]) -> None
         if audience and hasattr(toolset, "_mcp_session_manager"):
             self._audience_map[id(toolset._mcp_session_manager)] = audience
 
-    def add_to_agent(self, agent: BaseAgent):
-        """
-        Add the plugin to an ADK LLM agent by updating its MCP toolset
-        Call this once when setting up the agent; do not call it at runtime.
-        """
-        if not isinstance(agent, LlmAgent):
-            return
-
-        if not agent.tools:
-            return
-
-        for tool in agent.tools:
-            if isinstance(tool, MCPToolset):
-                mcp_toolset = tool
-                mcp_toolset._header_provider = self.header_provider
-                logger.debug("Updated tool connection params to include access token from STS server")
-
     def header_provider(
         self, readonly_context: Optional[ReadonlyContext], audience: Optional[str] = None
     ) -> Dict[str, str]:

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

@@ -3,7 +3,6 @@
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
-from google.adk.agents import LlmAgent
 from google.adk.tools.mcp_tool import MCPTool
 from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
 
@@ -98,14 +97,6 @@ async def test_downstream_token_propagation_without_sts(self):
         # Without STS, subject token is directly cached under session key
         assert plugin.token_cache["sess-2"] == "subj-token-123"
 
-        # propagate toolset
-        mcp_toolset = Mock(spec=MCPToolset)
-        agent = Mock(spec=LlmAgent)
-        agent.tools = [mcp_toolset]
-        plugin.add_to_agent(agent)
-        # The toolset._header_provider should be callable
-        assert callable(mcp_toolset._header_provider)
-
         # header provider should return subject token
         ro_ctx = self._make_readonly_context(ic)
         headers = plugin.header_provider(ro_ctx)