feat: Add header_provider to OpenAPIToolset and RestApiTool

Fixes: https://github.com/google/adk-python/issues/3782

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

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,8 +150,7 @@ 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]):
@@ -160,15 +158,13 @@ def test_openapi_toolset_configure_verify_all(openapi_spec: Dict[str, Any]):
   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 @@ class TestRestApiTool:
       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"