Merge branch 'main' into codex/provider-prefixed-gemini-models

Author: kenrogers

src/google/adk/tools/mcp_tool/mcp_toolset.py

@@ -118,6 +118,7 @@ def __init__(
       use_mcp_resources: Optional[bool] = False,
       sampling_callback: Optional[SamplingFnT] = None,
       sampling_capabilities: Optional[SamplingCapability] = None,
+      credential_key: str | None = None,
   ):
     """Initializes the McpToolset.
 
@@ -157,6 +158,8 @@ def __init__(
       sampling_callback: Optional callback to handle sampling requests from the
         MCP server.
       sampling_capabilities: Optional capabilities for sampling.
+      credential_key: A user specified key used to load and save this credential
+        in a credential service. Used with auth_scheme.
     """
 
     # --- BEGIN BOUND TOKEN PATCH ---
@@ -197,6 +200,7 @@ def __init__(
         AuthConfig(
             auth_scheme=auth_scheme,
             raw_auth_credential=auth_credential,
+            credential_key=credential_key,
         )
         if auth_scheme
         else None
@@ -469,6 +473,7 @@ def from_config(
         tool_name_prefix=mcp_toolset_config.tool_name_prefix,
         auth_scheme=mcp_toolset_config.auth_scheme,
         auth_credential=mcp_toolset_config.auth_credential,
+        credential_key=mcp_toolset_config.credential_key,
         use_mcp_resources=mcp_toolset_config.use_mcp_resources,
     )
 
@@ -520,6 +525,8 @@ class McpToolsetConfig(BaseToolConfig):
 
   auth_credential: Optional[AuthCredential] = None
 
+  credential_key: str | None = None
+
   use_mcp_resources: bool = False
 
   @model_validator(mode="after")

tests/unittests/tools/mcp_tool/test_mcp_toolset.py

@@ -27,6 +27,7 @@
 from google.adk.auth.auth_credential import AuthCredentialTypes
 from google.adk.auth.auth_credential import HttpAuth
 from google.adk.auth.auth_credential import HttpCredentials
+from google.adk.auth.auth_credential import OAuth2Auth
 from google.adk.auth.auth_tool import AuthConfig
 from google.adk.tools.load_mcp_resource_tool import LoadMcpResourceTool
 from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager
@@ -35,6 +36,7 @@
 from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
 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_configs import ToolArgsConfig
 from mcp import StdioServerParameters
 from mcp.types import BlobResourceContents
 from mcp.types import ListResourcesResult
@@ -136,10 +138,8 @@ def test_init_with_tool_filter_list(self):
   def test_init_with_auth(self):
     """Test initialization with authentication."""
     # Create real auth scheme instances
-    from fastapi.openapi.models import OAuth2
 
     auth_scheme = OAuth2(flows={})
-    from google.adk.auth.auth_credential import OAuth2Auth
 
     auth_credential = AuthCredential(
         auth_type="oauth2",
@@ -155,6 +155,42 @@ def test_init_with_auth(self):
     assert toolset._auth_scheme == auth_scheme
     assert toolset._auth_credential == auth_credential
 
+  def test_init_with_auth_and_credential_key(self):
+    """Test initialization with authentication and a custom credential_key."""
+
+    auth_scheme = OAuth2(flows={})
+
+    auth_credential = AuthCredential(
+        auth_type="oauth2",
+        oauth2=OAuth2Auth(client_id="test_id", client_secret="test_secret"),
+    )
+
+    toolset = McpToolset(
+        connection_params=self.mock_stdio_params,
+        auth_scheme=auth_scheme,
+        auth_credential=auth_credential,
+        credential_key="my_custom_key",
+    )
+
+    assert toolset._auth_scheme == auth_scheme
+    assert toolset._auth_credential == auth_credential
+    assert toolset._auth_config.credential_key == "my_custom_key"
+
+  def test_from_config_with_credential_key(self):
+    """Test that from_config correctly parses credential_key."""
+
+    auth_scheme = OAuth2(flows={})
+
+    config = ToolArgsConfig(
+        stdio_server_params=self.mock_stdio_params,
+        auth_scheme=auth_scheme,
+        credential_key="my_custom_key",
+    )
+    toolset = McpToolset.from_config(config, "")
+
+    assert isinstance(toolset._auth_scheme, OAuth2)
+    assert toolset._auth_config.credential_key == "my_custom_key"
+
   def test_init_missing_connection_params(self):
     """Test initialization with missing connection params raises error."""
     with pytest.raises(ValueError, match="Missing connection params"):