Merge remote-tracking branch 'upstream/main' into topic/bump-litellm-cap

Author: cwest

pyproject.toml

@@ -102,6 +102,7 @@ optional-dependencies.docs = [
   "myst-parser",
   "sphinx<9",
   "sphinx-autodoc-typehints",
+  "sphinx-click",
   "sphinx-rtd-theme",
 ]
 optional-dependencies.eval = [

src/google/adk/integrations/agent_identity/gcp_auth_provider.py

@@ -82,7 +82,7 @@ def _construct_auth_credential(
     return AuthCredential(
         auth_type=AuthCredentialTypes.HTTP,
         http=HttpAuth(
-            scheme="bearer",
+            scheme="Bearer",
             credentials=HttpCredentials(token=response.token),
         ),
     )

src/google/adk/plugins/reflect_retry_tool_plugin.py

@@ -242,8 +242,8 @@ async def _handle_tool_error(
     """
     if self.max_retries == 0:
       if self.throw_exception_if_retry_exceeded:
-        raise error
-      return self._get_tool_retry_exceed_msg(tool, error, tool_args)
+        raise self._ensure_exception(error)
+      return self._get_tool_retry_exceed_msg(tool, tool_args, error)
 
     scope_key = self._get_scope_key(tool_context)
     async with self._lock:
@@ -260,7 +260,7 @@ async def _handle_tool_error(
 
       # Max Retry exceeded
       if self.throw_exception_if_retry_exceeded:
-        raise error
+        raise self._ensure_exception(error)
       else:
         return self._get_tool_retry_exceed_msg(tool, tool_args, error)
 

tests/integration/integrations/agent_identity/test_2lo_flow.py

@@ -18,13 +18,6 @@
 from typing import Any
 from unittest import mock
 
-import pytest
-
-pytest.importorskip(
-    "google.cloud.iamconnectorcredentials_v1alpha",
-    reason="Requires google-cloud-iamconnectorcredentials",
-)
-
 from google.adk import Agent
 from google.adk import Runner
 from google.adk.auth.auth_tool import AuthConfig
@@ -34,9 +27,12 @@
 from google.adk.integrations.agent_identity import GcpAuthProviderScheme
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
 from google.adk.tools.base_authenticated_tool import BaseAuthenticatedTool
+from google.adk.tools.mcp_tool.mcp_tool import McpTool
 from google.cloud.iamconnectorcredentials_v1alpha import RetrieveCredentialsRequest
 from google.cloud.iamconnectorcredentials_v1alpha import RetrieveCredentialsResponse
 from google.genai import types
+from mcp.types import Tool as McpBaseTool
+import pytest
 
 from tests.unittests import testing_utils
 
@@ -200,3 +196,73 @@ async def test_gcp_agent_identity_2lo_gets_token() -> None:
 
     assert function_response.name == "dummy_tool"
     assert DUMMY_TOKEN in str(function_response.response)
+
+
+@pytest.mark.parametrize("llm_backend", ["GOOGLE_AI"], indirect=True)
+@pytest.mark.asyncio
+async def test_gcp_agent_identity_2lo_sends_authorization_header_to_mcp_session(
+    llm_backend: Any,
+) -> None:
+  """Ensures a 2LO token from GCP is correctly passed into the outbound MCP session headers."""
+  CredentialManager._auth_provider_registry._providers.clear()
+  CredentialManager.register_auth_provider(GcpAuthProvider())
+
+  mock_operation = _DummyOperation()
+  with mock.patch.object(
+      gcp_auth_provider, "Client", autospec=True
+  ) as mock_gcp:
+    mock_gcp.return_value.retrieve_credentials.return_value = mock_operation
+
+    mock_session_mgr = mock.AsyncMock()
+    mock_session_mgr.create_session.return_value.call_tool.return_value = (
+        mock.MagicMock()
+    )
+
+    mcp_tool = McpTool(
+        mcp_tool=McpBaseTool(
+            name="dummy_mcp",
+            description="Dummy MCP tool for testing.",
+            inputSchema={"type": "object", "properties": {}},
+        ),
+        mcp_session_manager=mock_session_mgr,
+        auth_scheme=GcpAuthProviderScheme(
+            name=TEST_CONNECTOR_2LO, scopes=["test-scope"]
+        ),
+    )
+
+    agent = Agent(
+        name="test_agent",
+        model=testing_utils.MockModel.create(
+            responses=[
+                types.Part.from_function_call(name="dummy_mcp", args={}),
+                "Tool executed successfully.",
+            ]
+        ),
+        instruction="Use dummy_mcp tool.",
+        tools=[mcp_tool],
+    )
+
+    async for _ in Runner(
+        app_name="test_mcp_header_app",
+        agent=agent,
+        session_service=InMemorySessionService(),
+        auto_create_session=True,
+    ).run_async(
+        user_id="test_user",
+        session_id="session-id-2",
+        new_message=types.UserContent(parts=[types.Part(text="Run tool.")]),
+    ):
+      pass
+
+    mock_gcp.return_value.retrieve_credentials.assert_called_once_with(
+        RetrieveCredentialsRequest(
+            connector=TEST_CONNECTOR_2LO,
+            user_id="test_user",
+            scopes=["test-scope"],
+            force_refresh=False,
+        )
+    )
+
+    assert mock_session_mgr.create_session.call_args.kwargs.get("headers") == {
+        "Authorization": f"Bearer {DUMMY_TOKEN}"
+    }

tests/integration/integrations/agent_identity/test_3lo_flow.py

@@ -18,13 +18,6 @@
 from typing import Any
 from unittest import mock
 
-import pytest
-
-pytest.importorskip(
-    "google.cloud.iamconnectorcredentials_v1alpha",
-    reason="Requires google-cloud-iamconnectorcredentials",
-)
-
 from google.adk import Agent
 from google.adk import Runner
 from google.adk.auth.auth_tool import AuthConfig
@@ -38,6 +31,7 @@
 from google.cloud.iamconnectorcredentials_v1alpha import RetrieveCredentialsRequest
 from google.cloud.iamconnectorcredentials_v1alpha import RetrieveCredentialsResponse
 from google.genai import types
+import pytest
 
 from tests.unittests import testing_utils
 

tests/unittests/integrations/agent_identity/test_gcp_auth_provider.py

@@ -169,7 +169,7 @@ async def test_get_auth_credential_returns_credential_if_available_immediately(
   auth_credential = await provider.get_auth_credential(auth_config, context)
 
   assert auth_credential.auth_type == AuthCredentialTypes.HTTP
-  assert auth_credential.http.scheme == "bearer"
+  assert auth_credential.http.scheme == "Bearer"
   assert auth_credential.http.credentials.token == "test-token"
   mock_client.retrieve_credentials.assert_called_once()
 
@@ -433,7 +433,7 @@ async def test_get_auth_credential_returns_token_if_consent_was_completed(
   # Verify
   assert auth_credential is not None
   assert auth_credential.auth_type == AuthCredentialTypes.HTTP
-  assert auth_credential.http.scheme == "bearer"
+  assert auth_credential.http.scheme == "Bearer"
   assert auth_credential.http.credentials.token == "test-token"
 
 

tests/unittests/plugins/test_reflect_retry_tool_plugin.py

@@ -168,6 +168,57 @@ async def test_on_tool_error_callback_max_retries_zero(self):
     # Should re-raise the original exception when max_retries is 0
     self.assertIs(cm.exception, error)
 
+  async def test_on_tool_error_callback_max_retries_zero_without_exception(
+      self,
+  ):
+    """Test error callback when max_retries is 0 and exception is disabled."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = ReflectAndRetryToolPlugin(
+        max_retries=0, throw_exception_if_retry_exceeded=False
+    )
+    error = ValueError("Test error")
+
+    result = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+
+    # Should return a retry exceeded message instead of raising
+    self.assertIsNotNone(result)
+    self.assertEqual(result["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
+    self.assertEqual(result["error_type"], "ValueError")
+    self.assertEqual(result["retry_count"], 0)
+    self.assertIn(
+        "the retry limit has been exceeded", result["reflection_guidance"]
+    )
+
+  async def test_on_tool_error_callback_max_retries_zero_with_dict_error(self):
+    """Test error callback when max_retries is 0 and error is a dict."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = CustomErrorExtractionPlugin(
+        max_retries=0, throw_exception_if_retry_exceeded=True
+    )
+    dict_error = {"status": "error", "message": "Custom dict error"}
+    plugin.set_error_condition(lambda result: dict_error)
+
+    with self.assertRaises(Exception) as cm:
+      await plugin.after_tool_callback(
+          tool=mock_tool,
+          tool_args=sample_tool_args,
+          tool_context=mock_tool_context,
+          result={"some": "result"},
+      )
+
+    # Should raise an Exception wrapping the dict
+    self.assertNotIsInstance(cm.exception, TypeError)
+    self.assertIn("Custom dict error", str(cm.exception))
+
   async def test_on_tool_error_callback_first_failure(self):
     """Test first tool failure creates reflection response."""
     plugin = self.get_plugin()
@@ -280,6 +331,40 @@ async def test_max_retries_exceeded_with_exception(self):
     # Verify exception properties
     self.assertIs(cm.exception, error)
 
+  async def test_max_retries_exceeded_with_dict_error(self):
+    """Test that Exception is raised when max retries exceeded with dict error."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = CustomErrorExtractionPlugin(
+        max_retries=1, throw_exception_if_retry_exceeded=True
+    )
+    dict_error = {"status": "error", "message": "Custom dict error"}
+    plugin.set_error_condition(lambda result: dict_error)
+
+    # First call should fail and return a retry response
+    result1 = await plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result={"some": "result"},
+    )
+    self.assertIsNotNone(result1)
+    self.assertEqual(result1["retry_count"], 1)
+
+    # Second call should exceed max_retries and raise
+    with self.assertRaises(Exception) as cm:
+      await plugin.after_tool_callback(
+          tool=mock_tool,
+          tool_args=sample_tool_args,
+          tool_context=mock_tool_context,
+          result={"some": "result"},
+      )
+
+    # Verify exception properties
+    self.assertNotIsInstance(cm.exception, TypeError)
+    self.assertIn("Custom dict error", str(cm.exception))
+
   async def test_max_retries_exceeded_without_exception(self):
     """Test max retries exceeded returns failure message when exception is disabled."""
     mock_tool = self.get_mock_tool()