chore: import upstream snapshot with attribution

tests/integration/integrations/agent_identity/test_2lo_flow.py

@@ -0,0 +1,268 @@
+# Copyright 2026 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.
+
+"""E2E Integration Test for GCP Agent Identity Auth Provider two-legged OAuth Flow."""
+
+import dataclasses
+from typing import Any
+from unittest import mock
+
+from google.adk import Agent
+from google.adk import Runner
+from google.adk.auth.auth_tool import AuthConfig
+from google.adk.auth.credential_manager import CredentialManager
+from google.adk.integrations.agent_identity import _iam_connector_credentials_provider
+from google.adk.integrations.agent_identity import GcpAuthProvider
+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
+
+DUMMY_TOKEN = "fake-gcp-2lo-token-123"
+TEST_CONNECTOR_2LO = (
+    "projects/test-project/locations/global/connectors/test-connector"
+)
+
+
+class DummyTool(BaseAuthenticatedTool):
+
+  def __init__(self, auth_config: AuthConfig) -> None:
+    super().__init__(
+        name="dummy_tool",
+        description="Dummy tool for testing 2LO.",
+        auth_config=auth_config,
+    )
+
+  def _get_declaration(self) -> types.FunctionDeclaration:
+    return types.FunctionDeclaration(
+        name=self.name,
+        description=self.description,
+        parameters=types.Schema(
+            type="OBJECT",
+            properties={},
+        ),
+    )
+
+  async def _run_async_impl(
+      self, *, args: dict[str, Any] | None, tool_context: Any, credential: Any
+  ) -> Any:
+    # Return the token to prove the provider gave the expected credential
+    if credential.http and credential.http.credentials:
+      return credential.http.credentials.token
+    if credential.oauth2 and credential.oauth2.access_token:
+      return credential.oauth2.access_token
+    return None
+
+
+@dataclasses.dataclass
+class _DummyOperation:
+  done: bool = True
+  error: Any = None
+  metadata: Any = None
+  response: Any = dataclasses.field(init=False)
+  operation: Any = dataclasses.field(init=False)
+
+  def __post_init__(self) -> None:
+    self.response = mock.Mock()
+    mock_credential = RetrieveCredentialsResponse(
+        header="Authorization: Bearer", token=DUMMY_TOKEN
+    )
+    self.response.value = RetrieveCredentialsResponse.serialize(mock_credential)
+    self.operation = self
+
+  def HasField(self, field_name: str) -> bool:
+    return getattr(self, field_name, None) is not None
+
+
+# Mocked execution; pin to a single LLM backend to avoid duplicate runs.
+@pytest.mark.parametrize("llm_backend", ["GOOGLE_AI"], indirect=True)
+@pytest.mark.asyncio
+async def test_gcp_agent_identity_2lo_gets_token() -> None:
+  """Test the end-to-end flow fetching 2LO OAuth token from GCP Agent Identity credentials service."""
+
+  # Clear registry to isolate tests
+  CredentialManager._auth_provider_registry._providers.clear()
+
+  # 1. Setup mocked GCP Client to return the fake Bearer token
+  with mock.patch.object(
+      _iam_connector_credentials_provider,
+      "Client",
+      autospec=True,
+  ) as mock_client_cls:
+
+    mock_operation = _DummyOperation()
+
+    mock_client_cls.return_value.retrieve_credentials.return_value = (
+        mock_operation
+    )
+
+    # 2. Configure Auth and DummyTool
+    auth_scheme = GcpAuthProviderScheme(
+        name=TEST_CONNECTOR_2LO,
+        scopes=["test-scope"],
+    )
+    auth_config = AuthConfig(auth_scheme=auth_scheme)
+    dummy_tool = DummyTool(auth_config=auth_config)
+
+    # 3. Setup LLM, Agent, and Runner
+    # We mock the LLM to just issue the tool call to 'dummy_tool'
+    mock_model = testing_utils.MockModel.create(
+        responses=[
+            types.Part.from_function_call(name="dummy_tool", args={}),
+            "Tool executed successfully.",
+        ]
+    )
+
+    agent = Agent(
+        name="test_agent",
+        model=mock_model,
+        instruction="You are an agent. Use the dummy_tool when needed.",
+        tools=[dummy_tool],
+    )
+
+    runner = Runner(
+        app_name="test_mcp_2lo_app",
+        agent=agent,
+        session_service=InMemorySessionService(),
+        auto_create_session=True,
+    )
+
+    # 4. Register Auth Provider
+    CredentialManager.register_auth_provider(GcpAuthProvider())
+
+    # 5. Execute Flow
+    event_list = []
+    async for event in runner.run_async(
+        user_id="test_user",
+        session_id="test_session1",
+        new_message=types.UserContent(
+            parts=[types.Part(text="Get me the token.")]
+        ),
+    ):
+      event_list.append(event)
+
+    # 6. Assertions
+
+    # Assert GCP Agent Identity client was invoked for credentials
+    expected_request = RetrieveCredentialsRequest(
+        connector=TEST_CONNECTOR_2LO,
+        user_id="test_user",
+        scopes=["test-scope"],
+        continue_uri="",
+        force_refresh=False,
+    )
+    mock_client_cls.return_value.retrieve_credentials.assert_called_once_with(
+        expected_request
+    )
+
+    # 3 Events: Model FunctionCall -> Tool FunctionResponse -> Final LLM Text
+    assert len(event_list) == 3
+    last_event = event_list[-1]
+    assert last_event.content.parts[0].text == "Tool executed successfully."
+
+    # Validate that the mock model received the query and the tool callback
+    requests = mock_model.requests
+    # 2 Events: User Input -> Tool FunctionResponse
+    assert len(requests) == 2
+
+    # Extract the function response from the prompt payload sent to the LLM
+    last_request = requests[-1]
+    function_response = next(
+        (
+            p.function_response
+            for p in last_request.contents[-1].parts
+            if p.function_response
+        ),
+        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(
+      _iam_connector_credentials_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}"
+    }