refactor: Introduce HeaderProvider type and modularize header generation logic in MCP tools.

Author: timof1308

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

@@ -56,6 +56,7 @@
 from ..base_authenticated_tool import BaseAuthenticatedTool
 #  import
 from ..tool_context import ToolContext
+from .types import HeaderProvider
 
 logger = logging.getLogger("google_adk." + __name__)
 
@@ -78,9 +79,7 @@ def __init__(
       auth_scheme: Optional[AuthScheme] = None,
       auth_credential: Optional[AuthCredential] = None,
       require_confirmation: Union[bool, Callable[..., bool]] = False,
-      header_provider: Optional[
-          Callable[[ReadonlyContext], Dict[str, str]]
-      ] = None,
+      header_provider: Optional[HeaderProvider] = None,
   ):
     """Initializes an McpTool.
 

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

@@ -17,6 +17,7 @@
 import asyncio
 import logging
 import sys
+from typing import Any
 from typing import Callable
 from typing import Dict
 from typing import List
@@ -57,16 +58,30 @@
     raise e
 
 from .mcp_tool import MCPTool
+from .types import HeaderProvider
+
 
 logger = logging.getLogger("google_adk." + __name__)
 
 
+def _validate_header_value(state_key: str, value: Any) -> None:
+  """Validates that a state value is suitable for use in a header."""
+  if not isinstance(value, (str, int, float, bool)):
+    logger.warning(
+        'Value for state key "%s" is of type %s, which may not serialize'
+        ' correctly into a header. Consider pre-serializing complex values or'
+        ' using state_header_format.',
+        state_key,
+        type(value).__name__,
+    )
+
+
 def create_session_state_header_provider(
     state_key: str,
     header_name: str = "Authorization",
     header_format: str = "Bearer {value}",
     default_value: Optional[str] = None,
-) -> Callable[[ReadonlyContext], Dict[str, str]]:
+) -> HeaderProvider:
   """Creates a header provider that extracts values from session state.
 
   This utility function generates a header_provider callable that can be used
@@ -103,20 +118,34 @@ def provider(ctx: ReadonlyContext) -> Dict[str, str]:
     value = ctx.state.get(state_key, default_value)
     if value is None:
       return {}
-    if not isinstance(value, (str, int, float, bool)):
-      logger.warning(
-          'Value for state key "%s" is of type %s, which may not serialize'
-          ' correctly into a header. Consider pre-serializing complex values or'
-          ' using a different header_format.',
-          state_key,
-          type(value).__name__,
-      )
+    _validate_header_value(state_key, value)
     formatted_value = header_format.format(value=value)
     return {header_name: formatted_value}
 
   return provider
 
 
+def create_combined_header_provider(
+    providers: List[HeaderProvider],
+) -> HeaderProvider:
+  """Creates a header provider that combines multiple providers.
+
+  Args:
+      providers: A list of header providers to combine.
+
+  Returns:
+      A single header provider that merges the results of all input providers.
+  """
+
+  def combined_provider(ctx: ReadonlyContext) -> Dict[str, str]:
+    headers = {}
+    for provider in providers:
+      headers.update(provider(ctx))
+    return headers
+
+  return combined_provider
+
+
 class McpToolset(BaseToolset):
   """Connects to a MCP Server, and retrieves MCP Tools into ADK Tools.
 
@@ -162,9 +191,7 @@ def __init__(
       auth_scheme: Optional[AuthScheme] = None,
       auth_credential: Optional[AuthCredential] = None,
       require_confirmation: Union[bool, Callable[..., bool]] = False,
-      header_provider: Optional[
-          Callable[[ReadonlyContext], Dict[str, str]]
-      ] = None,
+      header_provider: Optional[HeaderProvider] = None,
   ):
     """Initializes the McpToolset.
 
@@ -298,30 +325,17 @@ def from_config(
       state_mapping = mcp_toolset_config.state_header_mapping
       state_format = mcp_toolset_config.state_header_format or {}
 
-      def config_based_header_provider(
-          ctx: ReadonlyContext,
-      ) -> Dict[str, str]:
-        headers = {}
-        for state_key, header_name in state_mapping.items():
-          value = ctx.state.get(state_key)
-          if value is not None:
-            if not isinstance(value, (str, int, float, bool)):
-              logger.warning(
-                  'Value for state key "%s" is of type %s, which may not'
-                  ' serialize correctly into a header. Consider pre-serializing'
-                  ' complex values or using a different header_format.',
-                  state_key,
-                  type(value).__name__,
-              )
-            # Apply formatting if specified for this header
-            if header_name in state_format:
-              formatted_value = state_format[header_name].format(value=value)
-            else:
-              formatted_value = str(value)
-            headers[header_name] = formatted_value
-        return headers
-
-      header_provider = config_based_header_provider
+      providers = [
+          create_session_state_header_provider(
+              state_key=state_key,
+              header_name=header_name,
+              header_format=state_format.get(header_name, "{value}"),
+              default_value=None,
+          )
+          for state_key, header_name in state_mapping.items()
+      ]
+
+      header_provider = create_combined_header_provider(providers)
 
     return cls(
         connection_params=connection_params,

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

@@ -0,0 +1,20 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Callable
+from typing import Dict
+
+from ...agents.readonly_context import ReadonlyContext
+
+HeaderProvider = Callable[[ReadonlyContext], Dict[str, str]]