fix: Pass in auth headers with header provider instead of connection params

This allows headers to refresh on each request. Also only add auth headers if no auth_scheme or auth_credential is specified

Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 900900323

Author: wukath

src/google/adk/integrations/agent_registry/agent_registry.py

@@ -329,7 +329,6 @@ def get_mcp_toolset(
           f"MCP Server endpoint URI not found for: {mcp_server_name}"
       )
 
-    headers = self._get_auth_headers() if _is_google_api(endpoint_uri) else None
     if mcp_server_id and not auth_scheme:
       try:
         bindings_data = self._make_request("bindings")
@@ -349,13 +348,25 @@ def get_mcp_toolset(
 
     connection_params = StreamableHTTPConnectionParams(
         url=endpoint_uri,
-        headers=headers,
     )
+
+    def combined_header_provider(context: ReadonlyContext) -> Dict[str, str]:
+      headers = {}
+      if (
+          not auth_scheme
+          and not auth_credential
+          and _is_google_api(endpoint_uri)
+      ):
+        headers.update(self._get_auth_headers())
+      if self._header_provider:
+        headers.update(self._header_provider(context))
+      return headers
+
     return AgentRegistrySingleMcpToolset(
         destination_resource_id=mcp_server_id,
         connection_params=connection_params,
         tool_name_prefix=name,
-        header_provider=self._header_provider,
+        header_provider=combined_header_provider,
         auth_scheme=auth_scheme,
         auth_credential=auth_credential,
     )

tests/unittests/integrations/agent_registry/test_agent_registry.py

@@ -321,16 +321,17 @@ def test_get_endpoint(self, mock_httpx, registry):
     assert server == {"name": "test-endpoint"}
 
   @pytest.mark.parametrize(
-      "url, expected_auth",
+      "url, expected_auth, use_custom_provider",
       [
-          ("https://mcp.com", False),
-          ("https://mcp.googleapis.com/v1", True),
-          ("https://example.com/googleapis/v1", False),
+          ("https://mcp.com", False, False),
+          ("https://mcp.googleapis.com/v1", True, False),
+          ("https://example.com/googleapis/v1", False, False),
+          ("https://mcp.googleapis.com/v1", True, True),
       ],
   )
   @patch("httpx.Client")
   def test_get_mcp_toolset_auth_headers(
-      self, mock_httpx, registry, url, expected_auth
+      self, mock_httpx, registry, url, expected_auth, use_custom_provider
   ):
     mock_response = MagicMock()
     mock_response.json.return_value = {
@@ -345,19 +346,34 @@ def test_get_mcp_toolset_auth_headers(
         mock_response
     )
 
+    if use_custom_provider:
+      custom_header_provider = lambda context: {
+          "Authorization": "Bearer custom_token"
+      }
+      with patch(
+          "google.auth.default", return_value=(MagicMock(), "project-id")
+      ):
+        registry = AgentRegistry(
+            project_id="test-project",
+            location="global",
+            header_provider=custom_header_provider,
+        )
+
     registry._credentials.token = "token"
     registry._credentials.refresh = MagicMock()
 
     toolset = registry.get_mcp_toolset("test-mcp")
     assert isinstance(toolset, McpToolset)
     assert toolset.tool_name_prefix == "TestPrefix"
-    if expected_auth:
-      assert toolset._connection_params.headers is not None
-      assert (
-          toolset._connection_params.headers["Authorization"] == "Bearer token"
-      )
+    assert toolset._connection_params.headers is None
+    headers = toolset._header_provider(MagicMock())
+
+    if use_custom_provider:
+      assert headers.get("Authorization") == "Bearer custom_token"
+    elif expected_auth:
+      assert headers.get("Authorization") == "Bearer token"
     else:
-      assert toolset._connection_params.headers is None
+      assert "Authorization" not in headers
 
   @patch("httpx.Client")
   def test_get_mcp_toolset_with_auth(self, mock_httpx, registry):
@@ -392,6 +408,40 @@ def test_get_mcp_toolset_with_auth(self, mock_httpx, registry):
     assert auth_config.auth_scheme == auth_scheme
     assert auth_config.raw_auth_credential == auth_credential
 
+  @patch("httpx.Client")
+  def test_get_mcp_toolset_with_auth_blocks_gcp_headers(
+      self, mock_httpx, registry
+  ):
+    mock_response = MagicMock()
+    mock_response.json.return_value = {
+        "displayName": "TestPrefix",
+        "interfaces": [{
+            "url": "https://mcp.googleapis.com/v1",
+            "protocolBinding": "JSONRPC",
+        }],
+    }
+    mock_response.raise_for_status = MagicMock()
+    mock_httpx.return_value.__enter__.return_value.get.return_value = (
+        mock_response
+    )
+
+    registry._credentials.token = "token"
+    registry._credentials.refresh = MagicMock()
+
+    auth_scheme = OAuth2(flows={})
+    auth_credential = AuthCredential(
+        auth_type="oauth2",
+        oauth2=OAuth2Auth(client_id="test_id", client_secret="test_secret"),
+    )
+
+    toolset = registry.get_mcp_toolset(
+        "test-mcp", auth_scheme=auth_scheme, auth_credential=auth_credential
+    )
+    assert isinstance(toolset, McpToolset)
+
+    headers = toolset._header_provider(MagicMock())
+    assert "Authorization" not in headers
+
   @patch("httpx.Client")
   def test_get_remote_a2a_agent(self, mock_httpx, registry):
     mock_response = MagicMock()