feat: Add header_provider to OpenAPIToolset and RestApiTool

Fixes: #3782

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 843352147

Author: xuanyang15

src/google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py

@@ -18,6 +18,7 @@
 import logging
 import ssl
 from typing import Any
+from typing import Callable
 from typing import Dict
 from typing import Final
 from typing import List
@@ -71,6 +72,9 @@ def __init__(
       tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
       tool_name_prefix: Optional[str] = None,
       ssl_verify: Optional[Union[bool, str, ssl.SSLContext]] = None,
+      header_provider: Optional[
+          Callable[[ReadonlyContext], Dict[str, str]]
+      ] = None,
   ):
     """Initializes the OpenAPIToolset.
 
@@ -116,8 +120,14 @@ def __init__(
         - ssl.SSLContext: Custom SSL context for advanced configuration
         This is useful for enterprise environments where requests go through
         a TLS-intercepting proxy with a custom CA certificate.
+      header_provider: A callable that returns a dictionary of headers to be
+        included in API requests. The callable receives the ReadonlyContext as
+        an argument, allowing dynamic header generation based on the current
+        context. Useful for adding custom headers like correlation IDs,
+        authentication tokens, or other request metadata.
     """
     super().__init__(tool_filter=tool_filter, tool_name_prefix=tool_name_prefix)
+    self._header_provider = header_provider
     if not spec_dict:
       spec_dict = self._load_spec(spec_str, spec_str_type)
     self._ssl_verify = ssl_verify
@@ -189,7 +199,11 @@ def _parse(self, openapi_spec_dict: Dict[str, Any]) -> List[RestApiTool]:
 
     tools = []
     for o in operations:
-      tool = RestApiTool.from_parsed_operation(o, ssl_verify=self._ssl_verify)
+      tool = RestApiTool.from_parsed_operation(
+          o,
+          ssl_verify=self._ssl_verify,
+          header_provider=self._header_provider,
+      )
       logger.info("Parsed tool: %s", tool.name)
       tools.append(tool)
     return tools

src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py

@@ -16,6 +16,7 @@
 
 import ssl
 from typing import Any
+from typing import Callable
 from typing import Dict
 from typing import List
 from typing import Literal
@@ -29,6 +30,7 @@
 import requests
 from typing_extensions import override
 
+from ....agents.readonly_context import ReadonlyContext
 from ....auth.auth_credential import AuthCredential
 from ....auth.auth_schemes import AuthScheme
 from ..._gemini_schema_util import _to_gemini_schema
@@ -90,6 +92,9 @@ def __init__(
       auth_credential: Optional[Union[AuthCredential, str]] = None,
       should_parse_operation=True,
       ssl_verify: Optional[Union[bool, str, ssl.SSLContext]] = None,
+      header_provider: Optional[
+          Callable[[ReadonlyContext], Dict[str, str]]
+      ] = None,
   ):
     """Initializes the RestApiTool with the given parameters.
 
@@ -122,6 +127,11 @@ def __init__(
           - False: Disable SSL verification (insecure, not recommended)
           - str: Path to a CA bundle file or directory for custom CA
           - ssl.SSLContext: Custom SSL context for advanced configuration
+        header_provider: A callable that returns a dictionary of headers to be
+          included in API requests. The callable receives the ReadonlyContext as
+          an argument, allowing dynamic header generation based on the current
+          context. Useful for adding custom headers like correlation IDs,
+          authentication tokens, or other request metadata.
     """
     # Gemini restrict the length of function name to be less than 64 characters
     self.name = name[:60]
@@ -145,6 +155,7 @@ def __init__(
     self.credential_exchanger = AutoAuthCredentialExchanger()
     self._default_headers: Dict[str, str] = {}
     self._ssl_verify = ssl_verify
+    self._header_provider = header_provider
     if should_parse_operation:
       self._operation_parser = OperationParser(self.operation)
 
@@ -153,12 +164,20 @@ def from_parsed_operation(
       cls,
       parsed: ParsedOperation,
       ssl_verify: Optional[Union[bool, str, ssl.SSLContext]] = None,
+      header_provider: Optional[
+          Callable[[ReadonlyContext], Dict[str, str]]
+      ] = None,
   ) -> "RestApiTool":
     """Initializes the RestApiTool from a ParsedOperation object.
 
     Args:
         parsed: A ParsedOperation object.
         ssl_verify: SSL certificate verification option.
+        header_provider: A callable that returns a dictionary of headers to be
+          included in API requests. The callable receives the ReadonlyContext as
+          an argument, allowing dynamic header generation based on the current
+          context. Useful for adding custom headers like correlation IDs,
+          authentication tokens, or other request metadata.
 
     Returns:
         A RestApiTool object.
@@ -178,6 +197,7 @@ def from_parsed_operation(
         auth_scheme=parsed.auth_scheme,
         auth_credential=parsed.auth_credential,
         ssl_verify=ssl_verify,
+        header_provider=header_provider,
     )
     generated._operation_parser = operation_parser
     return generated
@@ -450,6 +470,13 @@ async def call(
     request_params = self._prepare_request_params(api_params, api_args)
     if self._ssl_verify is not None:
       request_params["verify"] = self._ssl_verify
+
+    # Add headers from header_provider if configured
+    if self._header_provider is not None and tool_context is not None:
+      provider_headers = self._header_provider(tool_context)
+      if provider_headers:
+        request_params.setdefault("headers", {}).update(provider_headers)
+
     response = requests.request(**request_params)
 
     # Parse API response

tests/unittests/tools/openapi_tool/openapi_spec_parser/test_openapi_toolset.py

@@ -135,9 +135,8 @@ def test_openapi_toolset_configure_auth_on_init(openapi_spec: Dict):
       auth_scheme=auth_scheme,
       auth_credential=auth_credential,
   )
-  for tool in toolset._tools:
-    assert tool.auth_scheme == auth_scheme
-    assert tool.auth_credential == auth_credential
+  assert all(tool.auth_scheme == auth_scheme for tool in toolset._tools)
+  assert all(tool.auth_credential == auth_credential for tool in toolset._tools)
 
 
 @pytest.mark.parametrize(
@@ -151,24 +150,21 @@ def test_openapi_toolset_verify_on_init(
       spec_dict=openapi_spec,
       ssl_verify=verify_value,
   )
-  for tool in toolset._tools:
-    assert tool._ssl_verify == verify_value
+  assert all(tool._ssl_verify == verify_value for tool in toolset._tools)
 
 
 def test_openapi_toolset_configure_verify_all(openapi_spec: Dict[str, Any]):
   """Test configure_verify_all method."""
   toolset = OpenAPIToolset(spec_dict=openapi_spec)
 
   # Initially verify should be None
-  for tool in toolset._tools:
-    assert tool._ssl_verify is None
+  assert all(tool._ssl_verify is None for tool in toolset._tools)
 
   # Configure verify for all tools
   ca_bundle_path = "/path/to/custom-ca.crt"
   toolset.configure_ssl_verify_all(ca_bundle_path)
 
-  for tool in toolset._tools:
-    assert tool._ssl_verify == ca_bundle_path
+  assert all(tool._ssl_verify == ca_bundle_path for tool in toolset._tools)
 
 
 async def test_openapi_toolset_tool_name_prefix(openapi_spec: Dict[str, Any]):
@@ -183,10 +179,42 @@ async def test_openapi_toolset_tool_name_prefix(openapi_spec: Dict[str, Any]):
   assert len(prefixed_tools) == 5
 
   # Verify all tool names are prefixed
-  for tool in prefixed_tools:
-    assert tool.name.startswith(f"{prefix}_")
+  assert all(tool.name.startswith(f"{prefix}_") for tool in prefixed_tools)
 
   # Verify specific tool name is prefixed
   expected_prefixed_name = "my_api_calendar_calendars_insert"
   prefixed_tool_names = [t.name for t in prefixed_tools]
   assert expected_prefixed_name in prefixed_tool_names
+
+
+def test_openapi_toolset_header_provider(openapi_spec: Dict[str, Any]):
+  """Test header_provider parameter is passed to tools."""
+
+  def my_header_provider(context):
+    return {"X-Custom-Header": "custom-value", "X-Request-ID": "12345"}
+
+  toolset = OpenAPIToolset(
+      spec_dict=openapi_spec,
+      header_provider=my_header_provider,
+  )
+
+  # Verify the toolset has the header_provider set
+  assert toolset._header_provider is my_header_provider
+
+  # Verify all tools have the header_provider
+  assert all(
+      tool._header_provider is my_header_provider for tool in toolset._tools
+  )
+
+
+def test_openapi_toolset_header_provider_none_by_default(
+    openapi_spec: Dict[str, Any],
+):
+  """Test that header_provider is None by default."""
+  toolset = OpenAPIToolset(spec_dict=openapi_spec)
+
+  # Verify the toolset has no header_provider by default
+  assert toolset._header_provider is None
+
+  # Verify all tools have no header_provider
+  assert all(tool._header_provider is None for tool in toolset._tools)

tests/unittests/tools/openapi_tool/openapi_spec_parser/test_rest_api_tool.py

@@ -1036,6 +1036,149 @@ async def test_call_with_configure_verify(
       call_kwargs = mock_request.call_args[1]
       assert call_kwargs["verify"] == ca_bundle_path
 
+  def test_init_with_header_provider(
+      self,
+      sample_endpoint,
+      sample_operation,
+  ):
+    """Test that header_provider is stored correctly."""
+
+    def my_header_provider(context):
+      return {"X-Custom": "value"}
+
+    tool = RestApiTool(
+        name="test_tool",
+        description="Test Tool",
+        endpoint=sample_endpoint,
+        operation=sample_operation,
+        header_provider=my_header_provider,
+    )
+    assert tool._header_provider is my_header_provider
+
+  def test_init_header_provider_none_by_default(
+      self,
+      sample_endpoint,
+      sample_operation,
+  ):
+    """Test that header_provider is None by default."""
+    tool = RestApiTool(
+        name="test_tool",
+        description="Test Tool",
+        endpoint=sample_endpoint,
+        operation=sample_operation,
+    )
+    assert tool._header_provider is None
+
+  @pytest.mark.asyncio
+  async def test_call_with_header_provider(
+      self,
+      mock_tool_context,
+      sample_endpoint,
+      sample_operation,
+      sample_auth_scheme,
+      sample_auth_credential,
+  ):
+    """Test that header_provider adds headers to the request."""
+    mock_response = mock.create_autospec(
+        requests.Response, instance=True, spec_set=True
+    )
+    mock_response.json.return_value = {"result": "success"}
+
+    def my_header_provider(context):
+      return {"X-Custom-Header": "custom-value", "X-Request-ID": "12345"}
+
+    tool = RestApiTool(
+        name="test_tool",
+        description="Test Tool",
+        endpoint=sample_endpoint,
+        operation=sample_operation,
+        auth_scheme=sample_auth_scheme,
+        auth_credential=sample_auth_credential,
+        header_provider=my_header_provider,
+    )
+
+    with patch.object(
+        requests, "request", return_value=mock_response, autospec=True
+    ) as mock_request:
+      await tool.call(args={}, tool_context=mock_tool_context)
+
+      # Verify the headers were added to the request
+      assert mock_request.called
+      _, call_kwargs = mock_request.call_args
+      assert call_kwargs["headers"]["X-Custom-Header"] == "custom-value"
+      assert call_kwargs["headers"]["X-Request-ID"] == "12345"
+
+  @pytest.mark.asyncio
+  async def test_call_header_provider_receives_tool_context(
+      self,
+      mock_tool_context,
+      sample_endpoint,
+      sample_operation,
+      sample_auth_scheme,
+      sample_auth_credential,
+  ):
+    """Test that header_provider receives the tool_context."""
+    mock_response = mock.create_autospec(
+        requests.Response, instance=True, spec_set=True
+    )
+    mock_response.json.return_value = {"result": "success"}
+
+    received_context = []
+
+    def my_header_provider(context):
+      received_context.append(context)
+      return {"X-Test": "test"}
+
+    tool = RestApiTool(
+        name="test_tool",
+        description="Test Tool",
+        endpoint=sample_endpoint,
+        operation=sample_operation,
+        auth_scheme=sample_auth_scheme,
+        auth_credential=sample_auth_credential,
+        header_provider=my_header_provider,
+    )
+
+    with patch.object(
+        requests, "request", return_value=mock_response, autospec=True
+    ):
+      await tool.call(args={}, tool_context=mock_tool_context)
+
+      # Verify header_provider was called with the tool_context
+      assert len(received_context) == 1
+      assert received_context[0] is mock_tool_context
+
+  @pytest.mark.asyncio
+  async def test_call_without_header_provider(
+      self,
+      mock_tool_context,
+      sample_endpoint,
+      sample_operation,
+      sample_auth_scheme,
+      sample_auth_credential,
+  ):
+    """Test that call works without header_provider."""
+    mock_response = mock.create_autospec(
+        requests.Response, instance=True, spec_set=True
+    )
+    mock_response.json.return_value = {"result": "success"}
+
+    tool = RestApiTool(
+        name="test_tool",
+        description="Test Tool",
+        endpoint=sample_endpoint,
+        operation=sample_operation,
+        auth_scheme=sample_auth_scheme,
+        auth_credential=sample_auth_credential,
+    )
+
+    with patch.object(
+        requests, "request", return_value=mock_response, autospec=True
+    ):
+      result = await tool.call(args={}, tool_context=mock_tool_context)
+
+      assert result == {"result": "success"}
+
 
 def test_snake_to_lower_camel():
   assert snake_to_lower_camel("single") == "single"