chore: import upstream snapshot with attribution

tests/unittests/auth/__init__.py

@@ -0,0 +1,13 @@
+# 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.

tests/unittests/auth/credential_service/__init__.py

@@ -0,0 +1,13 @@
+# 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.

tests/unittests/auth/credential_service/test_in_memory_credential_service.py

@@ -0,0 +1,347 @@
+# 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.
+
+from unittest.mock import Mock
+
+from fastapi.openapi.models import OAuth2
+from fastapi.openapi.models import OAuthFlowAuthorizationCode
+from fastapi.openapi.models import OAuthFlows
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_credential import OAuth2Auth
+from google.adk.auth.auth_tool import AuthConfig
+from google.adk.auth.credential_service.in_memory_credential_service import InMemoryCredentialService
+import pytest
+
+
+class TestInMemoryCredentialService:
+  """Tests for the InMemoryCredentialService class."""
+
+  @pytest.fixture
+  def credential_service(self):
+    """Create an InMemoryCredentialService instance for testing."""
+    return InMemoryCredentialService()
+
+  @pytest.fixture
+  def oauth2_auth_scheme(self):
+    """Create an OAuth2 auth scheme for testing."""
+    flows = OAuthFlows(
+        authorizationCode=OAuthFlowAuthorizationCode(
+            authorizationUrl="https://example.com/oauth2/authorize",
+            tokenUrl="https://example.com/oauth2/token",
+            scopes={"read": "Read access", "write": "Write access"},
+        )
+    )
+    return OAuth2(flows=flows)
+
+  @pytest.fixture
+  def oauth2_credentials(self):
+    """Create OAuth2 credentials for testing."""
+    return AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="mock_client_id",
+            client_secret="mock_client_secret",
+            redirect_uri="https://example.com/callback",
+        ),
+    )
+
+  @pytest.fixture
+  def auth_config(self, oauth2_auth_scheme, oauth2_credentials):
+    """Create an AuthConfig for testing."""
+    exchanged_credential = oauth2_credentials.model_copy(deep=True)
+    return AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        raw_auth_credential=oauth2_credentials,
+        exchanged_auth_credential=exchanged_credential,
+    )
+
+  @pytest.fixture
+  def callback_context(self):
+    """Create a mock CallbackContext for testing."""
+    mock_context = Mock(spec=CallbackContext)
+    mock_invocation_context = Mock()
+    mock_invocation_context.app_name = "test_app"
+    mock_invocation_context.user_id = "test_user"
+    mock_context._invocation_context = mock_invocation_context
+    return mock_context
+
+  @pytest.fixture
+  def another_callback_context(self):
+    """Create another mock CallbackContext with different app/user for testing isolation."""
+    mock_context = Mock(spec=CallbackContext)
+    mock_invocation_context = Mock()
+    mock_invocation_context.app_name = "another_app"
+    mock_invocation_context.user_id = "another_user"
+    mock_context._invocation_context = mock_invocation_context
+    return mock_context
+
+  def test_init(self, credential_service):
+    """Test that the service initializes with an empty store."""
+    assert isinstance(credential_service._credentials, dict)
+    assert len(credential_service._credentials) == 0
+
+  @pytest.mark.asyncio
+  async def test_load_credential_not_found(
+      self, credential_service, auth_config, callback_context
+  ):
+    """Test loading a credential that doesn't exist returns None."""
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
+    assert result is None
+
+  @pytest.mark.asyncio
+  async def test_save_and_load_credential(
+      self, credential_service, auth_config, callback_context
+  ):
+    """Test saving and then loading a credential."""
+    # Save the credential
+    await credential_service.save_credential(auth_config, callback_context)
+
+    # Load the credential
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
+
+    # Verify the credential was saved and loaded correctly
+    assert result is not None
+    assert result == auth_config.exchanged_auth_credential
+    assert result.auth_type == AuthCredentialTypes.OAUTH2
+    assert result.oauth2.client_id == "mock_client_id"
+
+  @pytest.mark.asyncio
+  async def test_save_credential_updates_existing(
+      self,
+      credential_service,
+      auth_config,
+      callback_context,
+      oauth2_credentials,
+  ):
+    """Test that saving a credential updates an existing one."""
+    # Save initial credential
+    await credential_service.save_credential(auth_config, callback_context)
+
+    # Create a new credential and update the auth_config
+    new_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="updated_client_id",
+            client_secret="updated_client_secret",
+            redirect_uri="https://updated.com/callback",
+        ),
+    )
+    auth_config.exchanged_auth_credential = new_credential
+
+    # Save the updated credential
+    await credential_service.save_credential(auth_config, callback_context)
+
+    # Load and verify the credential was updated
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
+    assert result is not None
+    assert result.oauth2.client_id == "updated_client_id"
+    assert result.oauth2.client_secret == "updated_client_secret"
+
+  @pytest.mark.asyncio
+  async def test_credentials_isolated_by_context(
+      self,
+      credential_service,
+      auth_config,
+      callback_context,
+      another_callback_context,
+  ):
+    """Test that credentials are isolated between different app/user contexts."""
+    # Save credential in first context
+    await credential_service.save_credential(auth_config, callback_context)
+
+    # Try to load from another context
+    result = await credential_service.load_credential(
+        auth_config, another_callback_context
+    )
+    assert result is None
+
+    # Verify original context still has the credential
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
+    assert result is not None
+
+  @pytest.mark.asyncio
+  async def test_multiple_credentials_same_context(
+      self, credential_service, callback_context, oauth2_auth_scheme
+  ):
+    """Test storing multiple credentials in the same context with different keys."""
+    # Create two different auth configs with different credential keys
+    cred1 = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="client1",
+            client_secret="secret1",
+            redirect_uri="https://example1.com/callback",
+        ),
+    )
+
+    cred2 = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="client2",
+            client_secret="secret2",
+            redirect_uri="https://example2.com/callback",
+        ),
+    )
+
+    auth_config1 = AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        raw_auth_credential=cred1,
+        exchanged_auth_credential=cred1,
+        credential_key="key1",
+    )
+
+    auth_config2 = AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        raw_auth_credential=cred2,
+        exchanged_auth_credential=cred2,
+        credential_key="key2",
+    )
+
+    # Save both credentials
+    await credential_service.save_credential(auth_config1, callback_context)
+    await credential_service.save_credential(auth_config2, callback_context)
+
+    # Load and verify both credentials
+    result1 = await credential_service.load_credential(
+        auth_config1, callback_context
+    )
+    result2 = await credential_service.load_credential(
+        auth_config2, callback_context
+    )
+
+    assert result1 is not None
+    assert result2 is not None
+    assert result1.oauth2.client_id == "client1"
+    assert result2.oauth2.client_id == "client2"
+
+  def test_get_bucket_for_current_context_creates_nested_structure(
+      self, credential_service, callback_context
+  ):
+    """Test that _get_bucket_for_current_context creates the proper nested structure."""
+    storage = credential_service._get_bucket_for_current_context(
+        callback_context
+    )
+
+    # Verify the nested structure was created
+    assert "test_app" in credential_service._credentials
+    assert "test_user" in credential_service._credentials["test_app"]
+    assert isinstance(storage, dict)
+    assert storage is credential_service._credentials["test_app"]["test_user"]
+
+  def test_get_bucket_for_current_context_reuses_existing(
+      self, credential_service, callback_context
+  ):
+    """Test that _get_bucket_for_current_context reuses existing structure."""
+    # Create initial structure
+    storage1 = credential_service._get_bucket_for_current_context(
+        callback_context
+    )
+    storage1["test_key"] = "test_value"
+
+    # Get storage again
+    storage2 = credential_service._get_bucket_for_current_context(
+        callback_context
+    )
+
+    # Verify it's the same storage instance
+    assert storage1 is storage2
+    assert storage2["test_key"] == "test_value"
+
+  def test_get_storage_different_apps(
+      self, credential_service, callback_context, another_callback_context
+  ):
+    """Test that different apps get different storage instances."""
+    storage1 = credential_service._get_bucket_for_current_context(
+        callback_context
+    )
+    storage2 = credential_service._get_bucket_for_current_context(
+        another_callback_context
+    )
+
+    # Verify they are different storage instances
+    assert storage1 is not storage2
+
+    # Verify the structure
+    assert "test_app" in credential_service._credentials
+    assert "another_app" in credential_service._credentials
+    assert "test_user" in credential_service._credentials["test_app"]
+    assert "another_user" in credential_service._credentials["another_app"]
+
+  @pytest.mark.asyncio
+  async def test_same_user_different_apps(
+      self, credential_service, auth_config
+  ):
+    """Test that the same user in different apps get isolated storage."""
+    # Create two contexts with same user but different apps
+    context1 = Mock(spec=CallbackContext)
+    mock_invocation_context1 = Mock()
+    mock_invocation_context1.app_name = "app1"
+    mock_invocation_context1.user_id = "same_user"
+    context1._invocation_context = mock_invocation_context1
+
+    context2 = Mock(spec=CallbackContext)
+    mock_invocation_context2 = Mock()
+    mock_invocation_context2.app_name = "app2"
+    mock_invocation_context2.user_id = "same_user"
+    context2._invocation_context = mock_invocation_context2
+
+    # Save credential in first app
+    await credential_service.save_credential(auth_config, context1)
+
+    # Try to load from second app
+    result = await credential_service.load_credential(auth_config, context2)
+    assert result is None
+
+    # Verify first app still has the credential
+    result = await credential_service.load_credential(auth_config, context1)
+    assert result is not None
+
+  @pytest.mark.asyncio
+  async def test_same_app_different_users(
+      self, credential_service, auth_config
+  ):
+    """Test that different users in the same app get isolated storage."""
+    # Create two contexts with same app but different users
+    context1 = Mock(spec=CallbackContext)
+    mock_invocation_context1 = Mock()
+    mock_invocation_context1.app_name = "same_app"
+    mock_invocation_context1.user_id = "user1"
+    context1._invocation_context = mock_invocation_context1
+
+    context2 = Mock(spec=CallbackContext)
+    mock_invocation_context2 = Mock()
+    mock_invocation_context2.app_name = "same_app"
+    mock_invocation_context2.user_id = "user2"
+    context2._invocation_context = mock_invocation_context2
+
+    # Save credential for first user
+    await credential_service.save_credential(auth_config, context1)
+
+    # Try to load for second user
+    result = await credential_service.load_credential(auth_config, context2)
+    assert result is None
+
+    # Verify first user still has the credential
+    result = await credential_service.load_credential(auth_config, context1)
+    assert result is not None

tests/unittests/auth/credential_service/test_session_state_credential_service.py

@@ -0,0 +1,388 @@
+# 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.
+
+from unittest.mock import Mock
+
+from fastapi.openapi.models import OAuth2
+from fastapi.openapi.models import OAuthFlowAuthorizationCode
+from fastapi.openapi.models import OAuthFlows
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_credential import OAuth2Auth
+from google.adk.auth.auth_tool import AuthConfig
+from google.adk.auth.credential_service.session_state_credential_service import SessionStateCredentialService
+import pytest
+
+
+class TestSessionStateCredentialService:
+  """Tests for the SessionStateCredentialService class."""
+
+  @pytest.fixture
+  def credential_service(self):
+    """Create a SessionStateCredentialService instance for testing."""
+    return SessionStateCredentialService()
+
+  @pytest.fixture
+  def oauth2_auth_scheme(self):
+    """Create an OAuth2 auth scheme for testing."""
+    flows = OAuthFlows(
+        authorizationCode=OAuthFlowAuthorizationCode(
+            authorizationUrl="https://example.com/oauth2/authorize",
+            tokenUrl="https://example.com/oauth2/token",
+            scopes={"read": "Read access", "write": "Write access"},
+        )
+    )
+    return OAuth2(flows=flows)
+
+  @pytest.fixture
+  def oauth2_credentials(self):
+    """Create OAuth2 credentials for testing."""
+    return AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="mock_client_id",
+            client_secret="mock_client_secret",
+            redirect_uri="https://example.com/callback",
+        ),
+    )
+
+  @pytest.fixture
+  def auth_config(self, oauth2_auth_scheme, oauth2_credentials):
+    """Create an AuthConfig for testing."""
+    exchanged_credential = oauth2_credentials.model_copy(deep=True)
+    return AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        raw_auth_credential=oauth2_credentials,
+        exchanged_auth_credential=exchanged_credential,
+    )
+
+  @pytest.fixture
+  def callback_context(self):
+    """Create a mock CallbackContext for testing."""
+    mock_context = Mock(spec=CallbackContext)
+    # Create a state dictionary that behaves like session state
+    mock_context.state = {}
+    return mock_context
+
+  @pytest.fixture
+  def another_callback_context(self):
+    """Create another mock CallbackContext with different state for testing isolation."""
+    mock_context = Mock(spec=CallbackContext)
+    # Create a separate state dictionary to simulate different session
+    mock_context.state = {}
+    return mock_context
+
+  @pytest.mark.asyncio
+  async def test_load_credential_not_found(
+      self, credential_service, auth_config, callback_context
+  ):
+    """Test loading a credential that doesn't exist returns None."""
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
+    assert result is None
+
+  @pytest.mark.asyncio
+  async def test_save_and_load_credential(
+      self, credential_service, auth_config, callback_context
+  ):
+    """Test saving and then loading a credential."""
+    # Save the credential
+    await credential_service.save_credential(auth_config, callback_context)
+
+    # Load the credential
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
+
+    # Verify the credential was saved and loaded correctly
+    assert result is not None
+    assert result == auth_config.exchanged_auth_credential
+    assert result.auth_type == AuthCredentialTypes.OAUTH2
+    assert result.oauth2.client_id == "mock_client_id"
+
+  @pytest.mark.asyncio
+  async def test_save_credential_updates_existing(
+      self,
+      credential_service,
+      auth_config,
+      callback_context,
+      oauth2_credentials,
+  ):
+    """Test that saving a credential updates an existing one."""
+    # Save initial credential
+    await credential_service.save_credential(auth_config, callback_context)
+
+    # Create a new credential and update the auth_config
+    new_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="updated_client_id",
+            client_secret="updated_client_secret",
+            redirect_uri="https://updated.com/callback",
+        ),
+    )
+    auth_config.exchanged_auth_credential = new_credential
+
+    # Save the updated credential
+    await credential_service.save_credential(auth_config, callback_context)
+
+    # Load and verify the credential was updated
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
+    assert result is not None
+    assert result.oauth2.client_id == "updated_client_id"
+    assert result.oauth2.client_secret == "updated_client_secret"
+
+  @pytest.mark.asyncio
+  async def test_credentials_isolated_by_context(
+      self,
+      credential_service,
+      auth_config,
+      callback_context,
+      another_callback_context,
+  ):
+    """Test that credentials are isolated between different callback contexts."""
+    # Save credential in first context
+    await credential_service.save_credential(auth_config, callback_context)
+
+    # Try to load from another context (should not find it)
+    result = await credential_service.load_credential(
+        auth_config, another_callback_context
+    )
+    assert result is None
+
+    # Verify original context still has the credential
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
+    assert result is not None
+
+  @pytest.mark.asyncio
+  async def test_multiple_credentials_same_context(
+      self, credential_service, callback_context, oauth2_auth_scheme
+  ):
+    """Test storing multiple credentials in the same context with different keys."""
+    # Create two different auth configs with different credential keys
+    cred1 = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="client1",
+            client_secret="secret1",
+            redirect_uri="https://example1.com/callback",
+        ),
+    )
+
+    cred2 = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="client2",
+            client_secret="secret2",
+            redirect_uri="https://example2.com/callback",
+        ),
+    )
+
+    auth_config1 = AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        raw_auth_credential=cred1,
+        exchanged_auth_credential=cred1,
+        credential_key="key1",
+    )
+
+    auth_config2 = AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        raw_auth_credential=cred2,
+        exchanged_auth_credential=cred2,
+        credential_key="key2",
+    )
+
+    # Save both credentials
+    await credential_service.save_credential(auth_config1, callback_context)
+    await credential_service.save_credential(auth_config2, callback_context)
+
+    # Load and verify both credentials
+    result1 = await credential_service.load_credential(
+        auth_config1, callback_context
+    )
+    result2 = await credential_service.load_credential(
+        auth_config2, callback_context
+    )
+
+    assert result1 is not None
+    assert result2 is not None
+    assert result1.oauth2.client_id == "client1"
+    assert result2.oauth2.client_id == "client2"
+
+  @pytest.mark.asyncio
+  async def test_save_credential_with_none_exchanged_credential(
+      self, credential_service, auth_config, callback_context
+  ):
+    """Test that saving a credential with None exchanged_auth_credential stores None."""
+    # Set exchanged_auth_credential to None
+    auth_config.exchanged_auth_credential = None
+
+    # Save the credential
+    await credential_service.save_credential(auth_config, callback_context)
+
+    # Load and verify None was stored
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
+    assert result is None
+
+  @pytest.mark.asyncio
+  async def test_load_credential_with_empty_credential_key(
+      self, credential_service, auth_config, callback_context
+  ):
+    """Test that loading with an empty credential key returns None."""
+    # Set credential_key to empty string
+    auth_config.credential_key = ""
+
+    # Try to load credential
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
+    assert result is None
+
+  @pytest.mark.asyncio
+  async def test_state_persistence_across_operations(
+      self, credential_service, auth_config, callback_context
+  ):
+    """Test that state persists across multiple operations."""
+    # Save credential
+    await credential_service.save_credential(auth_config, callback_context)
+
+    # Verify state contains the credential
+    assert auth_config.credential_key in callback_context.state
+    assert (
+        callback_context.state[auth_config.credential_key]
+        == auth_config.exchanged_auth_credential
+    )
+
+    # Load credential
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
+    assert result is not None
+
+    # Verify state still contains the credential
+    assert auth_config.credential_key in callback_context.state
+    assert (
+        callback_context.state[auth_config.credential_key]
+        == auth_config.exchanged_auth_credential
+    )
+
+    # Update credential
+    new_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="updated_client_id",
+            client_secret="updated_client_secret",
+            redirect_uri="https://updated.com/callback",
+        ),
+    )
+    auth_config.exchanged_auth_credential = new_credential
+
+    # Save updated credential
+    await credential_service.save_credential(auth_config, callback_context)
+
+    # Verify state was updated
+    assert callback_context.state[auth_config.credential_key] == new_credential
+
+  @pytest.mark.asyncio
+  async def test_credential_key_uniqueness(
+      self, credential_service, oauth2_auth_scheme, callback_context
+  ):
+    """Test that different credential keys store different credentials."""
+    # Create credentials with different keys
+    cred1 = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="client1",
+            client_secret="secret1",
+            redirect_uri="https://example1.com/callback",
+        ),
+    )
+
+    cred2 = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="client2",
+            client_secret="secret2",
+            redirect_uri="https://example2.com/callback",
+        ),
+    )
+
+    auth_config1 = AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        raw_auth_credential=cred1,
+        exchanged_auth_credential=cred1,
+        credential_key="unique_key_1",
+    )
+
+    auth_config2 = AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        raw_auth_credential=cred2,
+        exchanged_auth_credential=cred2,
+        credential_key="unique_key_2",
+    )
+
+    # Save both credentials
+    await credential_service.save_credential(auth_config1, callback_context)
+    await credential_service.save_credential(auth_config2, callback_context)
+
+    # Verify both exist in state with different keys
+    assert "unique_key_1" in callback_context.state
+    assert "unique_key_2" in callback_context.state
+    assert (
+        callback_context.state["unique_key_1"]
+        != callback_context.state["unique_key_2"]
+    )
+
+    # Load and verify both credentials
+    result1 = await credential_service.load_credential(
+        auth_config1, callback_context
+    )
+    result2 = await credential_service.load_credential(
+        auth_config2, callback_context
+    )
+
+    assert result1 is not None
+    assert result2 is not None
+    assert result1.oauth2.client_id == "client1"
+    assert result2.oauth2.client_id == "client2"
+
+  @pytest.mark.asyncio
+  async def test_direct_state_access(
+      self, credential_service, auth_config, callback_context
+  ):
+    """Test that the service properly accesses callback context state."""
+    # Directly set a value in state
+    test_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="direct_client_id",
+            client_secret="direct_client_secret",
+            redirect_uri="https://direct.com/callback",
+        ),
+    )
+    callback_context.state[auth_config.credential_key] = test_credential
+
+    # Load using the service
+    result = await credential_service.load_credential(
+        auth_config, callback_context
+    )
+    assert result == test_credential

tests/unittests/auth/exchanger/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+
+"""Tests for credential exchanger."""

tests/unittests/auth/exchanger/test_credential_exchanger_registry.py

@@ -0,0 +1,242 @@
+# 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.
+
+"""Unit tests for the CredentialExchangerRegistry."""
+
+from typing import Optional
+from unittest.mock import MagicMock
+
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_schemes import AuthScheme
+from google.adk.auth.exchanger.base_credential_exchanger import BaseCredentialExchanger
+from google.adk.auth.exchanger.credential_exchanger_registry import CredentialExchangerRegistry
+import pytest
+
+
+class MockCredentialExchanger(BaseCredentialExchanger):
+  """Mock credential exchanger for testing."""
+
+  def __init__(self, exchange_result: Optional[AuthCredential] = None):
+    self.exchange_result = exchange_result or AuthCredential(
+        auth_type=AuthCredentialTypes.HTTP
+    )
+
+  def exchange(
+      self,
+      auth_credential: AuthCredential,
+      auth_scheme: Optional[AuthScheme] = None,
+  ) -> AuthCredential:
+    """Mock exchange method."""
+    return self.exchange_result
+
+
+class TestCredentialExchangerRegistry:
+  """Test cases for CredentialExchangerRegistry."""
+
+  def test_initialization(self):
+    """Test that the registry initializes with an empty exchangers dictionary."""
+    registry = CredentialExchangerRegistry()
+
+    # Access the private attribute for testing
+    assert hasattr(registry, '_exchangers')
+    assert isinstance(registry._exchangers, dict)
+    assert len(registry._exchangers) == 0
+
+  def test_register_single_exchanger(self):
+    """Test registering a single exchanger."""
+    registry = CredentialExchangerRegistry()
+    mock_exchanger = MockCredentialExchanger()
+
+    registry.register(AuthCredentialTypes.API_KEY, mock_exchanger)
+
+    # Verify the exchanger was registered
+    retrieved_exchanger = registry.get_exchanger(AuthCredentialTypes.API_KEY)
+    assert retrieved_exchanger is mock_exchanger
+
+  def test_register_multiple_exchangers(self):
+    """Test registering multiple exchangers for different credential types."""
+    registry = CredentialExchangerRegistry()
+
+    api_key_exchanger = MockCredentialExchanger()
+    oauth2_exchanger = MockCredentialExchanger()
+    service_account_exchanger = MockCredentialExchanger()
+
+    registry.register(AuthCredentialTypes.API_KEY, api_key_exchanger)
+    registry.register(AuthCredentialTypes.OAUTH2, oauth2_exchanger)
+    registry.register(
+        AuthCredentialTypes.SERVICE_ACCOUNT, service_account_exchanger
+    )
+
+    # Verify all exchangers were registered correctly
+    assert (
+        registry.get_exchanger(AuthCredentialTypes.API_KEY) is api_key_exchanger
+    )
+    assert (
+        registry.get_exchanger(AuthCredentialTypes.OAUTH2) is oauth2_exchanger
+    )
+    assert (
+        registry.get_exchanger(AuthCredentialTypes.SERVICE_ACCOUNT)
+        is service_account_exchanger
+    )
+
+  def test_register_overwrites_existing_exchanger(self):
+    """Test that registering an exchanger for an existing type overwrites the previous one."""
+    registry = CredentialExchangerRegistry()
+
+    first_exchanger = MockCredentialExchanger()
+    second_exchanger = MockCredentialExchanger()
+
+    # Register first exchanger
+    registry.register(AuthCredentialTypes.API_KEY, first_exchanger)
+    assert (
+        registry.get_exchanger(AuthCredentialTypes.API_KEY) is first_exchanger
+    )
+
+    # Register second exchanger for the same type
+    registry.register(AuthCredentialTypes.API_KEY, second_exchanger)
+    assert (
+        registry.get_exchanger(AuthCredentialTypes.API_KEY) is second_exchanger
+    )
+    assert (
+        registry.get_exchanger(AuthCredentialTypes.API_KEY)
+        is not first_exchanger
+    )
+
+  def test_get_exchanger_returns_correct_instance(self):
+    """Test that get_exchanger returns the correct exchanger instance."""
+    registry = CredentialExchangerRegistry()
+    mock_exchanger = MockCredentialExchanger()
+
+    registry.register(AuthCredentialTypes.HTTP, mock_exchanger)
+
+    retrieved_exchanger = registry.get_exchanger(AuthCredentialTypes.HTTP)
+    assert retrieved_exchanger is mock_exchanger
+    assert isinstance(retrieved_exchanger, BaseCredentialExchanger)
+
+  def test_get_exchanger_nonexistent_type_returns_none(self):
+    """Test that get_exchanger returns None for nonexistent credential types."""
+    registry = CredentialExchangerRegistry()
+
+    # Try to get an exchanger that was never registered
+    result = registry.get_exchanger(AuthCredentialTypes.OAUTH2)
+    assert result is None
+
+  def test_get_exchanger_after_registration_and_removal(self):
+    """Test behavior when an exchanger is registered and then the registry is cleared indirectly."""
+    registry = CredentialExchangerRegistry()
+    mock_exchanger = MockCredentialExchanger()
+
+    # Register exchanger
+    registry.register(AuthCredentialTypes.API_KEY, mock_exchanger)
+    assert registry.get_exchanger(AuthCredentialTypes.API_KEY) is mock_exchanger
+
+    # Clear the internal dictionary (simulating some edge case)
+    registry._exchangers.clear()
+    assert registry.get_exchanger(AuthCredentialTypes.API_KEY) is None
+
+  def test_register_with_all_credential_types(self):
+    """Test registering exchangers for all available credential types."""
+    registry = CredentialExchangerRegistry()
+
+    exchangers = {}
+    credential_types = [
+        AuthCredentialTypes.API_KEY,
+        AuthCredentialTypes.HTTP,
+        AuthCredentialTypes.OAUTH2,
+        AuthCredentialTypes.OPEN_ID_CONNECT,
+        AuthCredentialTypes.SERVICE_ACCOUNT,
+    ]
+
+    # Register an exchanger for each credential type
+    for cred_type in credential_types:
+      exchanger = MockCredentialExchanger()
+      exchangers[cred_type] = exchanger
+      registry.register(cred_type, exchanger)
+
+    # Verify all exchangers can be retrieved
+    for cred_type in credential_types:
+      retrieved_exchanger = registry.get_exchanger(cred_type)
+      assert retrieved_exchanger is exchangers[cred_type]
+
+  def test_register_with_mock_exchanger_using_magicmock(self):
+    """Test registering with a MagicMock exchanger."""
+    registry = CredentialExchangerRegistry()
+    mock_exchanger = MagicMock(spec=BaseCredentialExchanger)
+
+    registry.register(AuthCredentialTypes.API_KEY, mock_exchanger)
+
+    retrieved_exchanger = registry.get_exchanger(AuthCredentialTypes.API_KEY)
+    assert retrieved_exchanger is mock_exchanger
+
+  def test_registry_isolation(self):
+    """Test that different registry instances are isolated from each other."""
+    registry1 = CredentialExchangerRegistry()
+    registry2 = CredentialExchangerRegistry()
+
+    exchanger1 = MockCredentialExchanger()
+    exchanger2 = MockCredentialExchanger()
+
+    # Register different exchangers in different registry instances
+    registry1.register(AuthCredentialTypes.API_KEY, exchanger1)
+    registry2.register(AuthCredentialTypes.API_KEY, exchanger2)
+
+    # Verify isolation
+    assert registry1.get_exchanger(AuthCredentialTypes.API_KEY) is exchanger1
+    assert registry2.get_exchanger(AuthCredentialTypes.API_KEY) is exchanger2
+    assert (
+        registry1.get_exchanger(AuthCredentialTypes.API_KEY) is not exchanger2
+    )
+    assert (
+        registry2.get_exchanger(AuthCredentialTypes.API_KEY) is not exchanger1
+    )
+
+  def test_exchanger_functionality_through_registry(self):
+    """Test that exchangers registered in the registry function correctly."""
+    registry = CredentialExchangerRegistry()
+
+    # Create a mock exchanger with specific return value
+    expected_result = AuthCredential(auth_type=AuthCredentialTypes.HTTP)
+    mock_exchanger = MockCredentialExchanger(exchange_result=expected_result)
+
+    registry.register(AuthCredentialTypes.API_KEY, mock_exchanger)
+
+    # Get the exchanger and test its functionality
+    retrieved_exchanger = registry.get_exchanger(AuthCredentialTypes.API_KEY)
+    input_credential = AuthCredential(auth_type=AuthCredentialTypes.API_KEY)
+
+    result = retrieved_exchanger.exchange(input_credential)
+    assert result is expected_result
+
+  def test_register_none_exchanger(self):
+    """Test that registering None as an exchanger works (edge case)."""
+    registry = CredentialExchangerRegistry()
+
+    # This should work but return None when retrieved
+    registry.register(AuthCredentialTypes.API_KEY, None)
+
+    result = registry.get_exchanger(AuthCredentialTypes.API_KEY)
+    assert result is None
+
+  def test_internal_dictionary_structure(self):
+    """Test the internal structure of the registry."""
+    registry = CredentialExchangerRegistry()
+    mock_exchanger = MockCredentialExchanger()
+
+    registry.register(AuthCredentialTypes.OAUTH2, mock_exchanger)
+
+    # Verify internal dictionary structure
+    assert AuthCredentialTypes.OAUTH2 in registry._exchangers
+    assert registry._exchangers[AuthCredentialTypes.OAUTH2] is mock_exchanger
+    assert len(registry._exchangers) == 1

tests/unittests/auth/exchanger/test_oauth2_credential_exchanger.py

@@ -0,0 +1,507 @@
+# 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.
+
+import time
+from unittest.mock import Mock
+from unittest.mock import patch
+from urllib.parse import parse_qs
+
+from authlib.integrations.requests_client import OAuth2Session
+from authlib.oauth2.rfc6749 import OAuth2Token
+from fastapi.openapi.models import OAuth2
+from fastapi.openapi.models import OAuthFlowClientCredentials
+from fastapi.openapi.models import OAuthFlows
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_credential import OAuth2Auth
+from google.adk.auth.auth_schemes import OAuthGrantType
+from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
+from google.adk.auth.exchanger import oauth2_credential_exchanger
+from google.adk.auth.exchanger.base_credential_exchanger import CredentialExchangeError
+from google.adk.auth.exchanger.oauth2_credential_exchanger import OAuth2CredentialExchanger
+import pytest
+
+
+class _TokenBodyCapturingOAuth2Session(OAuth2Session):
+  """A mock OAuth2Session that captures the final token request body."""
+
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    self.request_body = ""
+
+  def _fetch_token(
+      self,
+      url=None,
+      body="",
+      auth=None,
+      method="POST",
+      headers=None,
+      **kwargs,
+  ):
+    if auth is not None:
+      _, _, body = auth.prepare(method, url, headers or {}, body)
+    self.request_body = body
+    return OAuth2Token({
+        "access_token": "new_access_token",
+        "refresh_token": "new_refresh_token",
+        "expires_at": int(time.time()) + 3600,
+        "expires_in": 3600,
+    })
+
+
+class TestOAuth2CredentialExchanger:
+  """Test suite for OAuth2CredentialExchanger."""
+
+  async def test_exchange_with_existing_token(self):
+    """Test exchange method when access token already exists."""
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            access_token="existing_token",
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+    exchange_result = await exchanger.exchange(credential, scheme)
+
+    # Should return the same credential since access token already exists
+    assert exchange_result.credential == credential
+    assert exchange_result.credential.oauth2.access_token == "existing_token"
+    assert not exchange_result.was_exchanged
+
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
+  async def test_exchange_success(self, mock_oauth2_session):
+    """Test successful token exchange."""
+    # Setup mock
+    mock_client = Mock()
+    mock_oauth2_session.return_value = mock_client
+    mock_tokens = OAuth2Token({
+        "access_token": "new_access_token",
+        "refresh_token": "new_refresh_token",
+        "expires_at": int(time.time()) + 3600,
+        "expires_in": 3600,
+    })
+    mock_client.fetch_token.return_value = mock_tokens
+
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            auth_response_uri="https://example.com/callback?code=auth_code",
+            auth_code="auth_code",
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+    exchange_result = await exchanger.exchange(credential, scheme)
+
+    # Verify token exchange was successful
+    assert exchange_result.credential.oauth2.access_token == "new_access_token"
+    assert (
+        exchange_result.credential.oauth2.refresh_token == "new_refresh_token"
+    )
+    assert exchange_result.was_exchanged
+    mock_client.fetch_token.assert_called_once()
+
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
+  async def test_exchange_success_pkce(self, mock_oauth2_session):
+    """Test successful token exchange with PKCE."""
+    # Setup mock
+    mock_client = Mock()
+    mock_oauth2_session.return_value = mock_client
+    mock_tokens = OAuth2Token({
+        "access_token": "new_access_token",
+        "refresh_token": "new_refresh_token",
+        "expires_at": int(time.time()) + 3600,
+        "expires_in": 3600,
+    })
+    mock_client.fetch_token.return_value = mock_tokens
+
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            auth_response_uri="https://example.com/callback?code=auth_code",
+            auth_code="auth_code",
+            code_verifier="mock_code_verifier",
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+    exchange_result = await exchanger.exchange(credential, scheme)
+
+    # Verify token exchange was successful
+    assert exchange_result.credential.oauth2.access_token == "new_access_token"
+    assert (
+        exchange_result.credential.oauth2.refresh_token == "new_refresh_token"
+    )
+    assert exchange_result.was_exchanged
+    mock_client.fetch_token.assert_called_once_with(
+        "https://example.com/token",
+        authorization_response="https://example.com/callback?code=auth_code",
+        code="auth_code",
+        grant_type=OAuthGrantType.AUTHORIZATION_CODE,
+        code_verifier="mock_code_verifier",
+    )
+
+  async def test_exchange_missing_auth_scheme(self):
+    """Test exchange with missing auth_scheme raises ValueError."""
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+    try:
+      await exchanger.exchange(credential, None)
+      assert False, "Should have raised ValueError"
+    except CredentialExchangeError as e:
+      assert "auth_scheme is required" in str(e)
+
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
+  async def test_exchange_no_session(self, mock_oauth2_session):
+    """Test exchange when OAuth2Session cannot be created."""
+    # Mock to return None for create_oauth2_session
+    mock_oauth2_session.return_value = None
+
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            # Missing client_secret to trigger session creation failure
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+    exchange_result = await exchanger.exchange(credential, scheme)
+
+    # Should return original credential when session creation fails
+    assert exchange_result.credential == credential
+    assert exchange_result.credential.oauth2.access_token is None
+    assert not exchange_result.was_exchanged
+
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
+  async def test_exchange_fetch_token_failure(self, mock_oauth2_session):
+    """Test exchange when fetch_token fails."""
+    # Setup mock to raise exception during fetch_token
+    mock_client = Mock()
+    mock_oauth2_session.return_value = mock_client
+    mock_client.fetch_token.side_effect = Exception("Token fetch failed")
+
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            auth_response_uri="https://example.com/callback?code=auth_code",
+            auth_code="auth_code",
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+    exchange_result = await exchanger.exchange(credential, scheme)
+
+    # Should return original credential when fetch_token fails
+    assert exchange_result.credential == credential
+    assert exchange_result.credential.oauth2.access_token is None
+    assert not exchange_result.was_exchanged
+    mock_client.fetch_token.assert_called_once()
+
+  async def test_exchange_authlib_not_available(self):
+    """Test exchange when authlib is not available."""
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            auth_response_uri="https://example.com/callback?code=auth_code",
+            auth_code="auth_code",
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+
+    # Mock AUTHLIB_AVAILABLE to False
+    with patch(
+        "google.adk.auth.exchanger.oauth2_credential_exchanger.AUTHLIB_AVAILABLE",
+        False,
+    ):
+      exchange_result = await exchanger.exchange(credential, scheme)
+
+    # Should return original credential when authlib is not available
+    assert exchange_result.credential == credential
+    assert exchange_result.credential.oauth2.access_token is None
+    assert not exchange_result.was_exchanged
+
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
+  async def test_exchange_client_credentials_success(self, mock_oauth2_session):
+    """Test successful client credentials exchange."""
+    # Setup mock
+    mock_client = Mock()
+    mock_oauth2_session.return_value = mock_client
+    mock_tokens = OAuth2Token({
+        "access_token": "client_access_token",
+        "expires_at": int(time.time()) + 3600,
+        "expires_in": 3600,
+    })
+    mock_client.fetch_token.return_value = mock_tokens
+
+    # Create OAuth2 scheme with client credentials flow
+    flows = OAuthFlows(
+        clientCredentials=OAuthFlowClientCredentials(
+            tokenUrl="https://example.com/token",
+            scopes={"read": "Read access", "write": "Write access"},
+        )
+    )
+    scheme = OAuth2(flows=flows)
+
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+    exchange_result = await exchanger.exchange(credential, scheme)
+
+    # Verify client credentials exchange was successful
+    assert (
+        exchange_result.credential.oauth2.access_token == "client_access_token"
+    )
+    assert exchange_result.was_exchanged
+    mock_client.fetch_token.assert_called_once_with(
+        "https://example.com/token",
+        grant_type="client_credentials",
+    )
+
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
+  async def test_exchange_client_credentials_failure(self, mock_oauth2_session):
+    """Test client credentials exchange failure."""
+    # Setup mock to raise exception during fetch_token
+    mock_client = Mock()
+    mock_oauth2_session.return_value = mock_client
+    mock_client.fetch_token.side_effect = Exception(
+        "Client credentials fetch failed"
+    )
+
+    # Create OAuth2 scheme with client credentials flow
+    flows = OAuthFlows(
+        clientCredentials=OAuthFlowClientCredentials(
+            tokenUrl="https://example.com/token", scopes={"read": "Read access"}
+        )
+    )
+    scheme = OAuth2(flows=flows)
+
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+    exchange_result = await exchanger.exchange(credential, scheme)
+
+    # Should return original credential when client credentials exchange fails
+    assert exchange_result.credential == credential
+    assert exchange_result.credential.oauth2.access_token is None
+    assert not exchange_result.was_exchanged
+    mock_client.fetch_token.assert_called_once()
+
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
+  async def test_exchange_normalize_uri(self, mock_oauth2_session):
+    """Test exchange method normalizes auth_response_uri."""
+    mock_client = Mock()
+    mock_oauth2_session.return_value = mock_client
+    mock_tokens = OAuth2Token({
+        "access_token": "new_access_token",
+        "refresh_token": "new_refresh_token",
+        "expires_at": int(time.time()) + 3600,
+        "expires_in": 3600,
+    })
+    mock_client.fetch_token.return_value = mock_tokens
+
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            auth_response_uri="https://example.com/callback?code=auth_code#",  # URI with trailing hash
+            auth_code="auth_code",
+        ),
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+    await exchanger.exchange(credential, scheme)
+
+    # Verify fetch_token was called with the normalized URI
+    mock_client.fetch_token.assert_called_once_with(
+        "https://example.com/token",
+        authorization_response="https://example.com/callback?code=auth_code",  # Normalized URI
+        code="auth_code",
+        grant_type=OAuthGrantType.AUTHORIZATION_CODE,
+    )
+
+  async def test_exchange_client_secret_post_has_single_client_id(self):
+    """Test exchange lets Authlib add client_id only once for body auth."""
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            token_endpoint_auth_method="client_secret_post",
+            auth_response_uri="https://example.com/callback?code=auth_code",
+            auth_code="auth_code",
+        ),
+    )
+
+    client = _TokenBodyCapturingOAuth2Session(
+        credential.oauth2.client_id,
+        credential.oauth2.client_secret,
+        token_endpoint_auth_method="client_secret_post",
+    )
+
+    with patch.object(
+        oauth2_credential_exchanger,
+        "create_oauth2_session",
+        autospec=True,
+        return_value=(client, "https://example.com/token"),
+    ):
+      exchanger = OAuth2CredentialExchanger()
+      exchange_result = await exchanger.exchange(credential, scheme)
+
+    request_params = parse_qs(client.request_body)
+
+    assert exchange_result.was_exchanged
+    assert request_params["grant_type"] == [OAuthGrantType.AUTHORIZATION_CODE]
+    assert request_params["code"] == ["auth_code"]
+    assert request_params["client_id"] == ["test_client_id"]
+    assert request_params["client_secret"] == ["test_client_secret"]
+
+  async def test_determine_grant_type_client_credentials(self):
+    """Test grant type determination for client credentials."""
+    flows = OAuthFlows(
+        clientCredentials=OAuthFlowClientCredentials(
+            tokenUrl="https://example.com/token", scopes={"read": "Read access"}
+        )
+    )
+    scheme = OAuth2(flows=flows)
+
+    exchanger = OAuth2CredentialExchanger()
+    grant_type = exchanger._determine_grant_type(scheme)
+
+    from google.adk.auth.auth_schemes import OAuthGrantType
+
+    assert grant_type == OAuthGrantType.CLIENT_CREDENTIALS
+
+  async def test_determine_grant_type_openid_connect(self):
+    """Test grant type determination for OpenID Connect (defaults to auth code)."""
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+
+    exchanger = OAuth2CredentialExchanger()
+    grant_type = exchanger._determine_grant_type(scheme)
+
+    from google.adk.auth.auth_schemes import OAuthGrantType
+
+    assert grant_type == OAuthGrantType.AUTHORIZATION_CODE

tests/unittests/auth/refresher/__init__.py

@@ -0,0 +1,13 @@
+# 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.

tests/unittests/auth/refresher/test_credential_refresher_registry.py

@@ -0,0 +1,174 @@
+# 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.
+
+"""Tests for CredentialRefresherRegistry."""
+
+from unittest.mock import Mock
+
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.refresher.base_credential_refresher import BaseCredentialRefresher
+from google.adk.auth.refresher.credential_refresher_registry import CredentialRefresherRegistry
+
+
+class TestCredentialRefresherRegistry:
+  """Tests for the CredentialRefresherRegistry class."""
+
+  def test_init(self):
+    """Test that registry initializes with empty refreshers dictionary."""
+    registry = CredentialRefresherRegistry()
+    assert registry._refreshers == {}
+
+  def test_register_refresher(self):
+    """Test registering a refresher instance for a credential type."""
+    registry = CredentialRefresherRegistry()
+    mock_refresher = Mock(spec=BaseCredentialRefresher)
+
+    registry.register(AuthCredentialTypes.OAUTH2, mock_refresher)
+
+    assert registry._refreshers[AuthCredentialTypes.OAUTH2] == mock_refresher
+
+  def test_register_multiple_refreshers(self):
+    """Test registering multiple refresher instances for different credential types."""
+    registry = CredentialRefresherRegistry()
+    mock_oauth2_refresher = Mock(spec=BaseCredentialRefresher)
+    mock_openid_refresher = Mock(spec=BaseCredentialRefresher)
+    mock_service_account_refresher = Mock(spec=BaseCredentialRefresher)
+
+    registry.register(AuthCredentialTypes.OAUTH2, mock_oauth2_refresher)
+    registry.register(
+        AuthCredentialTypes.OPEN_ID_CONNECT, mock_openid_refresher
+    )
+    registry.register(
+        AuthCredentialTypes.SERVICE_ACCOUNT, mock_service_account_refresher
+    )
+
+    assert (
+        registry._refreshers[AuthCredentialTypes.OAUTH2]
+        == mock_oauth2_refresher
+    )
+    assert (
+        registry._refreshers[AuthCredentialTypes.OPEN_ID_CONNECT]
+        == mock_openid_refresher
+    )
+    assert (
+        registry._refreshers[AuthCredentialTypes.SERVICE_ACCOUNT]
+        == mock_service_account_refresher
+    )
+
+  def test_register_overwrite_existing_refresher(self):
+    """Test that registering a refresher overwrites an existing one for the same credential type."""
+    registry = CredentialRefresherRegistry()
+    mock_refresher_1 = Mock(spec=BaseCredentialRefresher)
+    mock_refresher_2 = Mock(spec=BaseCredentialRefresher)
+
+    # Register first refresher
+    registry.register(AuthCredentialTypes.OAUTH2, mock_refresher_1)
+    assert registry._refreshers[AuthCredentialTypes.OAUTH2] == mock_refresher_1
+
+    # Register second refresher for same credential type
+    registry.register(AuthCredentialTypes.OAUTH2, mock_refresher_2)
+    assert registry._refreshers[AuthCredentialTypes.OAUTH2] == mock_refresher_2
+
+  def test_get_refresher_existing(self):
+    """Test getting a refresher instance for a registered credential type."""
+    registry = CredentialRefresherRegistry()
+    mock_refresher = Mock(spec=BaseCredentialRefresher)
+
+    registry.register(AuthCredentialTypes.OAUTH2, mock_refresher)
+    result = registry.get_refresher(AuthCredentialTypes.OAUTH2)
+
+    assert result == mock_refresher
+
+  def test_get_refresher_non_existing(self):
+    """Test getting a refresher instance for a non-registered credential type returns None."""
+    registry = CredentialRefresherRegistry()
+
+    result = registry.get_refresher(AuthCredentialTypes.OAUTH2)
+
+    assert result is None
+
+  def test_get_refresher_after_registration(self):
+    """Test getting refresher instances for multiple credential types."""
+    registry = CredentialRefresherRegistry()
+    mock_oauth2_refresher = Mock(spec=BaseCredentialRefresher)
+    mock_api_key_refresher = Mock(spec=BaseCredentialRefresher)
+
+    registry.register(AuthCredentialTypes.OAUTH2, mock_oauth2_refresher)
+    registry.register(AuthCredentialTypes.API_KEY, mock_api_key_refresher)
+
+    # Get registered refreshers
+    oauth2_result = registry.get_refresher(AuthCredentialTypes.OAUTH2)
+    api_key_result = registry.get_refresher(AuthCredentialTypes.API_KEY)
+
+    assert oauth2_result == mock_oauth2_refresher
+    assert api_key_result == mock_api_key_refresher
+
+    # Get non-registered refresher
+    http_result = registry.get_refresher(AuthCredentialTypes.HTTP)
+    assert http_result is None
+
+  def test_register_all_credential_types(self):
+    """Test registering refreshers for all available credential types."""
+    registry = CredentialRefresherRegistry()
+
+    refreshers = {}
+    for credential_type in AuthCredentialTypes:
+      mock_refresher = Mock(spec=BaseCredentialRefresher)
+      refreshers[credential_type] = mock_refresher
+      registry.register(credential_type, mock_refresher)
+
+    # Verify all refreshers are registered correctly
+    for credential_type in AuthCredentialTypes:
+      result = registry.get_refresher(credential_type)
+      assert result == refreshers[credential_type]
+
+  def test_empty_registry_get_refresher(self):
+    """Test getting refresher from empty registry returns None for any credential type."""
+    registry = CredentialRefresherRegistry()
+
+    for credential_type in AuthCredentialTypes:
+      result = registry.get_refresher(credential_type)
+      assert result is None
+
+  def test_registry_independence(self):
+    """Test that multiple registry instances are independent."""
+    registry1 = CredentialRefresherRegistry()
+    registry2 = CredentialRefresherRegistry()
+
+    mock_refresher1 = Mock(spec=BaseCredentialRefresher)
+    mock_refresher2 = Mock(spec=BaseCredentialRefresher)
+
+    registry1.register(AuthCredentialTypes.OAUTH2, mock_refresher1)
+    registry2.register(AuthCredentialTypes.OAUTH2, mock_refresher2)
+
+    # Verify registries are independent
+    assert (
+        registry1.get_refresher(AuthCredentialTypes.OAUTH2) == mock_refresher1
+    )
+    assert (
+        registry2.get_refresher(AuthCredentialTypes.OAUTH2) == mock_refresher2
+    )
+    assert registry1.get_refresher(
+        AuthCredentialTypes.OAUTH2
+    ) != registry2.get_refresher(AuthCredentialTypes.OAUTH2)
+
+  def test_register_with_none_refresher(self):
+    """Test registering None as a refresher instance."""
+    registry = CredentialRefresherRegistry()
+
+    # This should technically work as the registry accepts any value
+    registry.register(AuthCredentialTypes.OAUTH2, None)
+    result = registry.get_refresher(AuthCredentialTypes.OAUTH2)
+
+    assert result is None

tests/unittests/auth/refresher/test_oauth2_credential_refresher.py

@@ -0,0 +1,179 @@
+# 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.
+
+import time
+from unittest.mock import Mock
+from unittest.mock import patch
+
+from authlib.oauth2.rfc6749 import OAuth2Token
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_credential import OAuth2Auth
+from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
+from google.adk.auth.refresher.oauth2_credential_refresher import OAuth2CredentialRefresher
+import pytest
+
+
+class TestOAuth2CredentialRefresher:
+  """Test suite for OAuth2CredentialRefresher."""
+
+  @patch("google.adk.auth.refresher.oauth2_credential_refresher.OAuth2Token")
+  @pytest.mark.asyncio
+  async def test_needs_refresh_token_not_expired(self, mock_oauth2_token):
+    """Test needs_refresh when token is not expired."""
+    mock_token_instance = Mock()
+    mock_token_instance.is_expired.return_value = False
+    mock_oauth2_token.return_value = mock_token_instance
+
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            access_token="existing_token",
+            expires_at=int(time.time()) + 3600,
+        ),
+    )
+
+    refresher = OAuth2CredentialRefresher()
+    needs_refresh = await refresher.is_refresh_needed(credential, scheme)
+
+    assert not needs_refresh
+
+  @patch("google.adk.auth.refresher.oauth2_credential_refresher.OAuth2Token")
+  @pytest.mark.asyncio
+  async def test_needs_refresh_token_expired(self, mock_oauth2_token):
+    """Test needs_refresh when token is expired."""
+    mock_token_instance = Mock()
+    mock_token_instance.is_expired.return_value = True
+    mock_oauth2_token.return_value = mock_token_instance
+
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            access_token="existing_token",
+            expires_at=int(time.time()) - 3600,  # Expired
+        ),
+    )
+
+    refresher = OAuth2CredentialRefresher()
+    needs_refresh = await refresher.is_refresh_needed(credential, scheme)
+
+    assert needs_refresh
+
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Token")
+  @pytest.mark.asyncio
+  async def test_refresh_token_expired_success(
+      self, mock_oauth2_token, mock_oauth2_session
+  ):
+    """Test successful token refresh when token is expired."""
+    # Setup mock token
+    mock_token_instance = Mock()
+    mock_token_instance.is_expired.return_value = True
+    mock_oauth2_token.return_value = mock_token_instance
+
+    # Setup mock session
+    mock_client = Mock()
+    mock_oauth2_session.return_value = mock_client
+    mock_tokens = OAuth2Token({
+        "access_token": "refreshed_access_token",
+        "refresh_token": "refreshed_refresh_token",
+        "expires_at": int(time.time()) + 3600,
+        "expires_in": 3600,
+    })
+    mock_client.refresh_token.return_value = mock_tokens
+
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            access_token="old_token",
+            refresh_token="old_refresh_token",
+            expires_at=int(time.time()) - 3600,  # Expired
+        ),
+    )
+
+    refresher = OAuth2CredentialRefresher()
+    result = await refresher.refresh(credential, scheme)
+
+    # Verify token refresh was successful
+    assert result.oauth2.access_token == "refreshed_access_token"
+    assert result.oauth2.refresh_token == "refreshed_refresh_token"
+    mock_client.refresh_token.assert_called_once()
+
+  @pytest.mark.asyncio
+  async def test_refresh_no_oauth2_credential(self):
+    """Test refresh with no OAuth2 credential returns original."""
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        # No oauth2 field
+    )
+
+    refresher = OAuth2CredentialRefresher()
+    result = await refresher.refresh(credential, scheme)
+
+    assert result == credential
+
+  @pytest.mark.asyncio
+  async def test_needs_refresh_no_oauth2_credential(self):
+    """Test needs_refresh with no OAuth2 credential returns False."""
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.HTTP,
+        # No oauth2 field
+    )
+
+    refresher = OAuth2CredentialRefresher()
+    needs_refresh = await refresher.is_refresh_needed(credential, None)
+
+    assert not needs_refresh

tests/unittests/auth/test_auth_config.py

@@ -0,0 +1,216 @@
+# 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.
+
+import os
+from pathlib import Path
+import subprocess
+import sys
+
+from fastapi.openapi.models import OAuth2
+from fastapi.openapi.models import OAuthFlowAuthorizationCode
+from fastapi.openapi.models import OAuthFlows
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_credential import OAuth2Auth
+from google.adk.auth.auth_schemes import CustomAuthScheme
+from google.adk.auth.auth_tool import AuthConfig
+import pytest
+
+
+class TestAuthConfig:
+  """Tests for the AuthConfig method."""
+
+
+@pytest.fixture
+def oauth2_auth_scheme():
+  """Create an OAuth2 auth scheme for testing."""
+  # Create the OAuthFlows object first
+  flows = OAuthFlows(
+      authorizationCode=OAuthFlowAuthorizationCode(
+          authorizationUrl="https://example.com/oauth2/authorize",
+          tokenUrl="https://example.com/oauth2/token",
+          scopes={"read": "Read access", "write": "Write access"},
+      )
+  )
+
+  # Then create the OAuth2 object with the flows
+  return OAuth2(flows=flows)
+
+
+@pytest.fixture
+def oauth2_credentials():
+  """Create OAuth2 credentials for testing."""
+  return AuthCredential(
+      auth_type=AuthCredentialTypes.OAUTH2,
+      oauth2=OAuth2Auth(
+          client_id="mock_client_id",
+          client_secret="mock_client_secret",
+          redirect_uri="https://example.com/callback",
+      ),
+  )
+
+
+@pytest.fixture
+def auth_config(oauth2_auth_scheme, oauth2_credentials):
+  """Create an AuthConfig for testing."""
+  # Create a copy of the credentials for the exchanged_auth_credential
+  exchanged_credential = oauth2_credentials.model_copy(deep=True)
+
+  return AuthConfig(
+      auth_scheme=oauth2_auth_scheme,
+      raw_auth_credential=oauth2_credentials,
+      exchanged_auth_credential=exchanged_credential,
+  )
+
+
+@pytest.fixture
+def auth_config_with_key(oauth2_auth_scheme, oauth2_credentials):
+  """Create an AuthConfig for testing."""
+
+  return AuthConfig(
+      auth_scheme=oauth2_auth_scheme,
+      raw_auth_credential=oauth2_credentials,
+      credential_key="test_key",
+  )
+
+
+def test_custom_credential_key(auth_config_with_key):
+  """Test using custom credential key."""
+
+  key = auth_config_with_key.credential_key
+  assert key == "test_key"
+
+
+def test_credential_key(auth_config):
+  """Test generating a unique credential key."""
+
+  key = auth_config.credential_key
+  assert key.startswith("adk_oauth2_")
+  assert "_oauth2_" in key
+
+
+def test_get_credential_key_with_extras(auth_config):
+  """Test generating a key when model_extra exists."""
+  # Add model_extra to test cleanup
+
+  original_key = auth_config.credential_key
+  key = auth_config.credential_key
+
+  auth_config.auth_scheme.model_extra["extra_field"] = "value"
+  auth_config.raw_auth_credential.model_extra["extra_field"] = "value"
+
+  assert original_key == key
+  assert "extra_field" in auth_config.auth_scheme.model_extra
+  assert "extra_field" in auth_config.raw_auth_credential.model_extra
+
+
+def test_credential_key_is_stable_across_python_hash_seed():
+  """Test AuthConfig key generation does not depend on PYTHONHASHSEED."""
+  repo_root = Path(__file__).resolve().parents[3]
+  pythonpath = str(repo_root / "src")
+  code = "\n".join([
+      "from fastapi.openapi.models import OAuth2",
+      "from fastapi.openapi.models import OAuthFlowAuthorizationCode",
+      "from fastapi.openapi.models import OAuthFlows",
+      "from google.adk.auth.auth_credential import AuthCredential",
+      "from google.adk.auth.auth_credential import AuthCredentialTypes",
+      "from google.adk.auth.auth_credential import OAuth2Auth",
+      "from google.adk.auth.auth_tool import AuthConfig",
+      "",
+      "auth_scheme = OAuth2(",
+      "  flows=OAuthFlows(",
+      "    authorizationCode=OAuthFlowAuthorizationCode(",
+      "      authorizationUrl='https://example.com/oauth2/authorize',",
+      "      tokenUrl='https://example.com/oauth2/token',",
+      "      scopes={'read': 'Read access'},",
+      "    )",
+      "  )",
+      ")",
+      "auth_cred = AuthCredential(",
+      "  auth_type=AuthCredentialTypes.OAUTH2,",
+      "  oauth2=OAuth2Auth(",
+      "    client_id='mock_client_id',",
+      "    client_secret='mock_client_secret',",
+      "  ),",
+      ")",
+      "print(AuthConfig(",
+      "  auth_scheme=auth_scheme,",
+      "  raw_auth_credential=auth_cred,",
+      ").credential_key)",
+  ])
+
+  def _run_with_seed(seed: str) -> str:
+    env = os.environ.copy()
+    env["PYTHONHASHSEED"] = seed
+    env["PYTHONPATH"] = os.pathsep.join(
+        [pythonpath, env.get("PYTHONPATH", "")]
+    ).strip(os.pathsep)
+    return subprocess.check_output(
+        [sys.executable, "-c", code],
+        env=env,
+        text=True,
+    ).strip()
+
+  assert _run_with_seed("0") == _run_with_seed("1")
+
+
+def test_credential_key_with_custom_auth_scheme():
+  """Test generating a credential key when the auth scheme is a CustomAuthScheme (type_ is a string)."""
+  custom_scheme = CustomAuthScheme.model_validate({"type": "mock_custom_type"})
+
+  custom_config = AuthConfig(
+      auth_scheme=custom_scheme,
+  )
+
+  key = custom_config.credential_key
+  assert key.startswith("adk_mock_custom_type_")
+  assert len(key) > len("adk_mock_custom_type_")
+
+
+def test_credential_key_is_stable_across_redirect_uri(oauth2_auth_scheme):
+  """AuthConfig.credential_key should be invariant under redirect_uri changes.
+
+  redirect_uri is deployment configuration (which callback URL the auth
+  server should redirect to), not part of the credential identity. Two
+  AuthConfig instances built from credentials that share the same client_id,
+  client_secret, and scopes but differ only in redirect_uri should produce
+  the same credential_key.
+  """
+  credential_local = AuthCredential(
+      auth_type=AuthCredentialTypes.OAUTH2,
+      oauth2=OAuth2Auth(
+          client_id="client",
+          client_secret="secret",
+          redirect_uri="http://localhost:8001/oauth2callback",
+      ),
+  )
+  credential_deployed = AuthCredential(
+      auth_type=AuthCredentialTypes.OAUTH2,
+      oauth2=OAuth2Auth(
+          client_id="client",
+          client_secret="secret",
+          redirect_uri="https://deployed.example.com/oauth2callback",
+      ),
+  )
+
+  config_local = AuthConfig(
+      auth_scheme=oauth2_auth_scheme,
+      raw_auth_credential=credential_local,
+  )
+  config_deployed = AuthConfig(
+      auth_scheme=oauth2_auth_scheme,
+      raw_auth_credential=credential_deployed,
+  )
+
+  assert config_local.credential_key == config_deployed.credential_key

tests/unittests/auth/test_auth_handler.py

@@ -0,0 +1,726 @@
+# 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.
+
+import copy
+import time
+from unittest.mock import Mock
+from unittest.mock import patch
+
+from authlib.oauth2.rfc6749 import OAuth2Token
+from fastapi.openapi.models import APIKey
+from fastapi.openapi.models import APIKeyIn
+from fastapi.openapi.models import OAuth2
+from fastapi.openapi.models import OAuthFlowAuthorizationCode
+from fastapi.openapi.models import OAuthFlowClientCredentials
+from fastapi.openapi.models import OAuthFlows
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_credential import OAuth2Auth
+from google.adk.auth.auth_handler import AuthHandler
+from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
+from google.adk.auth.auth_tool import AuthConfig
+import pytest
+
+
+# Mock classes for testing
+class MockState(dict):
+  """Mock State class for testing."""
+
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+
+  def get(self, key, default=None):
+    return super().get(key, default)
+
+
+class MockOAuth2Session:
+  """Mock OAuth2Session for testing."""
+
+  def __init__(
+      self,
+      client_id=None,
+      client_secret=None,
+      scope=None,
+      redirect_uri=None,
+      state=None,
+      **kwargs,
+  ):
+    self.client_id = client_id
+    self.client_secret = client_secret
+    self.scope = scope
+    self.redirect_uri = redirect_uri
+    self.state = state
+    self.extra_kwargs = kwargs
+
+  def create_authorization_url(self, url, **kwargs):
+    params = f"client_id={self.client_id}&scope={self.scope}"
+    if kwargs.get("audience"):
+      params += f"&audience={kwargs.get('audience')}"
+    if kwargs.get("prompt"):
+      params += f"&prompt={kwargs.get('prompt')}"
+    return f"{url}?{params}", "mock_state"
+
+  def fetch_token(
+      self,
+      token_endpoint,
+      authorization_response=None,
+      code=None,
+      grant_type=None,
+  ):
+    return {
+        "access_token": "mock_access_token",
+        "token_type": "bearer",
+        "expires_in": 3600,
+        "refresh_token": "mock_refresh_token",
+    }
+
+
+# Fixtures for common test objects
+@pytest.fixture
+def oauth2_auth_scheme():
+  """Create an OAuth2 auth scheme for testing."""
+  # Create the OAuthFlows object first
+  flows = OAuthFlows(
+      authorizationCode=OAuthFlowAuthorizationCode(
+          authorizationUrl="https://example.com/oauth2/authorize",
+          tokenUrl="https://example.com/oauth2/token",
+          scopes={"read": "Read access", "write": "Write access"},
+      )
+  )
+
+  # Then create the OAuth2 object with the flows
+  return OAuth2(flows=flows)
+
+
+@pytest.fixture
+def openid_auth_scheme():
+  """Create an OpenID Connect auth scheme for testing."""
+  return OpenIdConnectWithConfig(
+      openIdConnectUrl="https://example.com/.well-known/openid-configuration",
+      authorization_endpoint="https://example.com/oauth2/authorize",
+      token_endpoint="https://example.com/oauth2/token",
+      scopes=["openid", "profile", "email"],
+  )
+
+
+@pytest.fixture
+def oauth2_credentials():
+  """Create OAuth2 credentials for testing."""
+  return AuthCredential(
+      auth_type=AuthCredentialTypes.OAUTH2,
+      oauth2=OAuth2Auth(
+          client_id="mock_client_id",
+          client_secret="mock_client_secret",
+          redirect_uri="https://example.com/callback",
+      ),
+  )
+
+
+@pytest.fixture
+def oauth2_credentials_with_token():
+  """Create OAuth2 credentials with a token for testing."""
+  return AuthCredential(
+      auth_type=AuthCredentialTypes.OAUTH2,
+      oauth2=OAuth2Auth(
+          client_id="mock_client_id",
+          client_secret="mock_client_secret",
+          redirect_uri="https://example.com/callback",
+          access_token="mock_access_token",
+          refresh_token="mock_refresh_token",
+      ),
+  )
+
+
+@pytest.fixture
+def oauth2_credentials_with_auth_uri():
+  """Create OAuth2 credentials with an auth URI for testing."""
+  return AuthCredential(
+      auth_type=AuthCredentialTypes.OAUTH2,
+      oauth2=OAuth2Auth(
+          client_id="mock_client_id",
+          client_secret="mock_client_secret",
+          redirect_uri="https://example.com/callback",
+          auth_uri="https://example.com/oauth2/authorize?client_id=mock_client_id&scope=read,write",
+          state="mock_state",
+      ),
+  )
+
+
+@pytest.fixture
+def oauth2_credentials_with_auth_code():
+  """Create OAuth2 credentials with an auth code for testing."""
+  return AuthCredential(
+      auth_type=AuthCredentialTypes.OAUTH2,
+      oauth2=OAuth2Auth(
+          client_id="mock_client_id",
+          client_secret="mock_client_secret",
+          redirect_uri="https://example.com/callback",
+          auth_uri="https://example.com/oauth2/authorize?client_id=mock_client_id&scope=read,write",
+          state="mock_state",
+          auth_code="mock_auth_code",
+          auth_response_uri="https://example.com/callback?code=mock_auth_code&state=mock_state",
+      ),
+  )
+
+
+@pytest.fixture
+def auth_config(oauth2_auth_scheme, oauth2_credentials):
+  """Create an AuthConfig for testing."""
+  # Create a copy of the credentials for the exchanged_auth_credential
+  exchanged_credential = oauth2_credentials.model_copy(deep=True)
+
+  return AuthConfig(
+      auth_scheme=oauth2_auth_scheme,
+      raw_auth_credential=oauth2_credentials,
+      exchanged_auth_credential=exchanged_credential,
+  )
+
+
+@pytest.fixture
+def auth_config_with_exchanged(
+    oauth2_auth_scheme, oauth2_credentials, oauth2_credentials_with_auth_uri
+):
+  """Create an AuthConfig with exchanged credentials for testing."""
+  return AuthConfig(
+      auth_scheme=oauth2_auth_scheme,
+      raw_auth_credential=oauth2_credentials,
+      exchanged_auth_credential=oauth2_credentials_with_auth_uri,
+  )
+
+
+@pytest.fixture
+def auth_config_with_auth_code(
+    oauth2_auth_scheme, oauth2_credentials, oauth2_credentials_with_auth_code
+):
+  """Create an AuthConfig with auth code for testing."""
+  return AuthConfig(
+      auth_scheme=oauth2_auth_scheme,
+      raw_auth_credential=oauth2_credentials,
+      exchanged_auth_credential=oauth2_credentials_with_auth_code,
+  )
+
+
+class TestAuthHandlerInit:
+  """Tests for the AuthHandler initialization."""
+
+  def test_init(self, auth_config):
+    """Test the initialization of AuthHandler."""
+    handler = AuthHandler(auth_config)
+    assert handler.auth_config == auth_config
+
+
+class TestGenerateAuthUri:
+  """Tests for the generate_auth_uri method."""
+
+  @patch("google.adk.auth.auth_handler.OAuth2Session", MockOAuth2Session)
+  def test_generate_auth_uri_oauth2(self, auth_config):
+    """Test generating an auth URI for OAuth2."""
+    handler = AuthHandler(auth_config)
+    result = handler.generate_auth_uri()
+
+    assert result.oauth2.auth_uri.startswith(
+        "https://example.com/oauth2/authorize"
+    )
+    assert "client_id=mock_client_id" in result.oauth2.auth_uri
+    assert "audience" not in result.oauth2.auth_uri
+    assert result.oauth2.state == "mock_state"
+
+  @patch("google.adk.auth.auth_handler.OAuth2Session", MockOAuth2Session)
+  def test_generate_auth_uri_with_audience_and_prompt(
+      self, openid_auth_scheme, oauth2_credentials
+  ):
+    """Test generating an auth URI with audience and prompt."""
+    oauth2_credentials.oauth2.audience = "test_audience"
+    exchanged = oauth2_credentials.model_copy(deep=True)
+
+    config = AuthConfig(
+        auth_scheme=openid_auth_scheme,
+        raw_auth_credential=oauth2_credentials,
+        exchanged_auth_credential=exchanged,
+    )
+    handler = AuthHandler(config)
+    result = handler.generate_auth_uri()
+
+    assert "audience=test_audience" in result.oauth2.auth_uri
+    assert "prompt=consent" in result.oauth2.auth_uri
+
+  @patch("google.adk.auth.auth_handler.OAuth2Session", MockOAuth2Session)
+  def test_generate_auth_uri_with_custom_prompt(
+      self, openid_auth_scheme, oauth2_credentials
+  ):
+    """Test generating an auth URI with a custom prompt override."""
+    oauth2_credentials.oauth2.prompt = "none"
+    exchanged = oauth2_credentials.model_copy(deep=True)
+
+    config = AuthConfig(
+        auth_scheme=openid_auth_scheme,
+        raw_auth_credential=oauth2_credentials,
+        exchanged_auth_credential=exchanged,
+    )
+    handler = AuthHandler(config)
+    result = handler.generate_auth_uri()
+
+    assert "prompt=none" in result.oauth2.auth_uri
+
+  @patch("google.adk.auth.auth_handler.OAuth2Session", MockOAuth2Session)
+  def test_generate_auth_uri_openid(
+      self, openid_auth_scheme, oauth2_credentials
+  ):
+    """Test generating an auth URI for OpenID Connect."""
+    # Create a copy for the exchanged credential
+    exchanged = oauth2_credentials.model_copy(deep=True)
+
+    config = AuthConfig(
+        auth_scheme=openid_auth_scheme,
+        raw_auth_credential=oauth2_credentials,
+        exchanged_auth_credential=exchanged,
+    )
+    handler = AuthHandler(config)
+    result = handler.generate_auth_uri()
+
+    assert result.oauth2.auth_uri.startswith(
+        "https://example.com/oauth2/authorize"
+    )
+    assert "client_id=mock_client_id" in result.oauth2.auth_uri
+    assert result.oauth2.state == "mock_state"
+
+  @patch("google.adk.auth.auth_handler.OAuth2Session", MockOAuth2Session)
+  def test_generate_auth_uri_client_credentials_with_missing_scopes(
+      self, oauth2_credentials
+  ):
+    """Test client credentials flow tolerates missing scopes."""
+    auth_scheme = OAuth2(
+        flows=OAuthFlows(
+            clientCredentials=OAuthFlowClientCredentials(
+                tokenUrl="https://example.com/oauth2/token"
+            )
+        )
+    )
+    auth_scheme.flows.clientCredentials.scopes = None
+
+    config = AuthConfig(
+        auth_scheme=auth_scheme,
+        raw_auth_credential=oauth2_credentials,
+        exchanged_auth_credential=oauth2_credentials.model_copy(deep=True),
+    )
+
+    handler = AuthHandler(config)
+    result = handler.generate_auth_uri()
+
+    assert (
+        result.oauth2.auth_uri
+        == "https://example.com/oauth2/token?client_id=mock_client_id&scope=&prompt=consent"
+    )
+    assert result.oauth2.state == "mock_state"
+
+  @patch("google.adk.auth.auth_handler.OAuth2Session")
+  def test_generate_auth_uri_pkce(
+      self, mock_oauth2_session, oauth2_auth_scheme, oauth2_credentials
+  ):
+    """Test generating an auth URI with PKCE."""
+    oauth2_credentials.oauth2.code_challenge_method = "S256"
+    exchanged = oauth2_credentials.model_copy(deep=True)
+
+    config = AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        raw_auth_credential=oauth2_credentials,
+        exchanged_auth_credential=exchanged,
+    )
+
+    mock_client = Mock()
+    mock_oauth2_session.return_value = mock_client
+    mock_client.create_authorization_url.return_value = (
+        "https://example.com/oauth2/authorize?code_challenge=...&code_challenge_method=S256",
+        "mock_state",
+    )
+
+    handler = AuthHandler(config)
+    result = handler.generate_auth_uri()
+
+    assert result.oauth2.code_verifier is not None
+    assert len(result.oauth2.code_verifier) == 48
+    mock_client.create_authorization_url.assert_called_once()
+    _, kwargs = mock_client.create_authorization_url.call_args
+    assert "code_verifier" in kwargs
+    assert kwargs["code_verifier"] == result.oauth2.code_verifier
+
+  def test_generate_auth_uri_unsupported_pkce_method(
+      self, oauth2_auth_scheme, oauth2_credentials
+  ):
+    """Test generating an auth URI with unsupported PKCE method."""
+    oauth2_credentials.oauth2.code_challenge_method = "plain"
+    exchanged = oauth2_credentials.model_copy(deep=True)
+
+    config = AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        raw_auth_credential=oauth2_credentials,
+        exchanged_auth_credential=exchanged,
+    )
+
+    handler = AuthHandler(config)
+    with pytest.raises(ValueError, match="Unsupported code_challenge_method"):
+      handler.generate_auth_uri()
+
+
+class TestGenerateAuthRequest:
+  """Tests for the generate_auth_request method."""
+
+  def test_non_oauth_scheme(self):
+    """Test with a non-OAuth auth scheme."""
+    # Use a SecurityBase instance without using APIKey which has validation issues
+    api_key_scheme = APIKey(**{"name": "test_api_key", "in": APIKeyIn.header})
+
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.API_KEY, api_key="test_api_key"
+    )
+
+    # Create a copy for the exchanged credential
+    exchanged = credential.model_copy(deep=True)
+
+    config = AuthConfig(
+        auth_scheme=api_key_scheme,
+        raw_auth_credential=credential,
+        exchanged_auth_credential=exchanged,
+    )
+
+    handler = AuthHandler(config)
+    result = handler.generate_auth_request()
+
+    assert result == config
+
+  def test_with_existing_auth_uri(self, auth_config_with_exchanged):
+    """Test when auth_uri already exists in exchanged credential."""
+    handler = AuthHandler(auth_config_with_exchanged)
+    result = handler.generate_auth_request()
+
+    assert (
+        result.exchanged_auth_credential.oauth2.auth_uri
+        == auth_config_with_exchanged.exchanged_auth_credential.oauth2.auth_uri
+    )
+
+  def test_missing_raw_credential(self, oauth2_auth_scheme):
+    """Test when raw_auth_credential is missing."""
+
+    config = AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+    )
+    handler = AuthHandler(config)
+
+    with pytest.raises(ValueError, match="requires auth_credential"):
+      handler.generate_auth_request()
+
+  def test_missing_oauth2_in_raw_credential(self, oauth2_auth_scheme):
+    """Test when oauth2 is missing in raw_auth_credential."""
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.API_KEY, api_key="test_api_key"
+    )
+
+    # Create a copy for the exchanged credential
+    exchanged = credential.model_copy(deep=True)
+
+    config = AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        raw_auth_credential=credential,
+        exchanged_auth_credential=exchanged,
+    )
+    handler = AuthHandler(config)
+
+    with pytest.raises(ValueError, match="requires oauth2 in auth_credential"):
+      handler.generate_auth_request()
+
+  def test_auth_uri_in_raw_credential(
+      self, oauth2_auth_scheme, oauth2_credentials_with_auth_uri
+  ):
+    """Test when auth_uri exists in raw_credential."""
+    config = AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        raw_auth_credential=oauth2_credentials_with_auth_uri,
+        exchanged_auth_credential=oauth2_credentials_with_auth_uri.model_copy(
+            deep=True
+        ),
+        credential_key="my_tool_tokens",
+    )
+    handler = AuthHandler(config)
+    result = handler.generate_auth_request()
+
+    assert result.credential_key == "my_tool_tokens"
+    assert (
+        result.exchanged_auth_credential.oauth2.auth_uri
+        == oauth2_credentials_with_auth_uri.oauth2.auth_uri
+    )
+
+  def test_missing_client_credentials(self, oauth2_auth_scheme):
+    """Test when client_id or client_secret is missing."""
+    bad_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(redirect_uri="https://example.com/callback"),
+    )
+
+    # Create a copy for the exchanged credential
+    exchanged = bad_credential.model_copy(deep=True)
+
+    config = AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        raw_auth_credential=bad_credential,
+        exchanged_auth_credential=exchanged,
+    )
+    handler = AuthHandler(config)
+
+    with pytest.raises(
+        ValueError, match="requires both client_id and client_secret"
+    ):
+      handler.generate_auth_request()
+
+  @patch("google.adk.auth.auth_handler.AuthHandler.generate_auth_uri")
+  def test_generate_new_auth_uri(self, mock_generate_auth_uri, auth_config):
+    """Test generating a new auth URI."""
+    mock_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="mock_client_id",
+            client_secret="mock_client_secret",
+            redirect_uri="https://example.com/callback",
+            auth_uri="https://example.com/generated",
+            state="generated_state",
+        ),
+    )
+    mock_generate_auth_uri.return_value = mock_credential
+
+    handler = AuthHandler(auth_config)
+    result = handler.generate_auth_request()
+
+    assert mock_generate_auth_uri.called
+    assert result.exchanged_auth_credential == mock_credential
+
+  @patch("google.adk.auth.auth_handler.AuthHandler.generate_auth_uri")
+  def test_preserves_credential_key_on_generated_request(
+      self, mock_generate_auth_uri, oauth2_auth_scheme, oauth2_credentials
+  ):
+    """Test that AuthHandler preserves an explicit credential_key."""
+    mock_generate_auth_uri.return_value = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="mock_client_id",
+            client_secret="mock_client_secret",
+            auth_uri="https://example.com/generated",
+            state="generated_state",
+        ),
+    )
+
+    config = AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        raw_auth_credential=oauth2_credentials,
+        credential_key="my_tool_tokens",
+    )
+    handler = AuthHandler(config)
+    result = handler.generate_auth_request()
+
+    assert result.credential_key == "my_tool_tokens"
+
+
+class TestGetAuthResponse:
+  """Tests for the get_auth_response method."""
+
+  def test_get_auth_response_exists(
+      self, auth_config, oauth2_credentials_with_auth_uri
+  ):
+    """Test retrieving an existing auth response from state."""
+    handler = AuthHandler(auth_config)
+    state = MockState()
+
+    # Store a credential in the state
+    credential_key = auth_config.credential_key
+    state["temp:" + credential_key] = oauth2_credentials_with_auth_uri
+
+    result = handler.get_auth_response(state)
+    assert result == oauth2_credentials_with_auth_uri
+
+  def test_get_auth_response_not_exists(self, auth_config):
+    """Test retrieving a nonexistent auth response from state."""
+    handler = AuthHandler(auth_config)
+    state = MockState()
+
+    result = handler.get_auth_response(state)
+    assert result is None
+
+
+class TestParseAndStoreAuthResponse:
+  """Tests for the parse_and_store_auth_response method."""
+
+  @pytest.mark.asyncio
+  async def test_non_oauth_scheme(self, auth_config_with_exchanged):
+    """Test with a non-OAuth auth scheme."""
+    # Modify the auth scheme type to be non-OAuth
+    auth_config = copy.deepcopy(auth_config_with_exchanged)
+    auth_config.auth_scheme = APIKey(
+        **{"name": "test_api_key", "in": APIKeyIn.header}
+    )
+
+    handler = AuthHandler(auth_config)
+    state = MockState()
+
+    await handler.parse_and_store_auth_response(state)
+
+    credential_key = auth_config.credential_key
+    assert (
+        state["temp:" + credential_key] == auth_config.exchanged_auth_credential
+    )
+
+  @patch("google.adk.auth.auth_handler.AuthHandler.exchange_auth_token")
+  @pytest.mark.asyncio
+  async def test_oauth_scheme(
+      self, mock_exchange_token, auth_config_with_exchanged
+  ):
+    """Test with an OAuth auth scheme."""
+    mock_exchange_token.return_value = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(access_token="exchanged_token"),
+    )
+
+    handler = AuthHandler(auth_config_with_exchanged)
+    state = MockState()
+
+    await handler.parse_and_store_auth_response(state)
+
+    credential_key = auth_config_with_exchanged.credential_key
+    assert state["temp:" + credential_key] == mock_exchange_token.return_value
+    assert mock_exchange_token.called
+
+
+class TestExchangeAuthToken:
+  """Tests for the exchange_auth_token method."""
+
+  @pytest.mark.asyncio
+  async def test_token_exchange_not_supported(
+      self, auth_config_with_auth_code, monkeypatch
+  ):
+    """Test when token exchange is not supported."""
+    monkeypatch.setattr(
+        "google.adk.auth.exchanger.oauth2_credential_exchanger.AUTHLIB_AVAILABLE",
+        False,
+    )
+
+    handler = AuthHandler(auth_config_with_auth_code)
+    result = await handler.exchange_auth_token()
+
+    assert result == auth_config_with_auth_code.exchanged_auth_credential
+
+  @pytest.mark.asyncio
+  async def test_openid_missing_token_endpoint(
+      self, openid_auth_scheme, oauth2_credentials_with_auth_code
+  ):
+    """Test OpenID Connect without a token endpoint."""
+    # Create a scheme without token_endpoint
+    scheme_without_token = copy.deepcopy(openid_auth_scheme)
+    delattr(scheme_without_token, "token_endpoint")
+
+    config = AuthConfig(
+        auth_scheme=scheme_without_token,
+        raw_auth_credential=oauth2_credentials_with_auth_code,
+        exchanged_auth_credential=oauth2_credentials_with_auth_code,
+    )
+
+    handler = AuthHandler(config)
+    result = await handler.exchange_auth_token()
+
+    assert result == oauth2_credentials_with_auth_code
+
+  @pytest.mark.asyncio
+  async def test_oauth2_missing_token_url(
+      self, oauth2_auth_scheme, oauth2_credentials_with_auth_code
+  ):
+    """Test OAuth2 without a token URL."""
+    # Create a scheme without tokenUrl
+    scheme_without_token = copy.deepcopy(oauth2_auth_scheme)
+    scheme_without_token.flows.authorizationCode.tokenUrl = None
+
+    config = AuthConfig(
+        auth_scheme=scheme_without_token,
+        raw_auth_credential=oauth2_credentials_with_auth_code,
+        exchanged_auth_credential=oauth2_credentials_with_auth_code,
+    )
+
+    handler = AuthHandler(config)
+    result = await handler.exchange_auth_token()
+
+    assert result == oauth2_credentials_with_auth_code
+
+  @pytest.mark.asyncio
+  async def test_non_oauth_scheme(self, auth_config_with_auth_code):
+    """Test with a non-OAuth auth scheme."""
+    # Modify the auth scheme type to be non-OAuth
+    auth_config = copy.deepcopy(auth_config_with_auth_code)
+    auth_config.auth_scheme = APIKey(
+        **{"name": "test_api_key", "in": APIKeyIn.header}
+    )
+
+    handler = AuthHandler(auth_config)
+    result = await handler.exchange_auth_token()
+
+    assert result == auth_config.exchanged_auth_credential
+
+  @pytest.mark.asyncio
+  async def test_missing_credentials(self, oauth2_auth_scheme):
+    """Test with missing credentials."""
+    empty_credential = AuthCredential(auth_type=AuthCredentialTypes.OAUTH2)
+
+    config = AuthConfig(
+        auth_scheme=oauth2_auth_scheme,
+        exchanged_auth_credential=empty_credential,
+    )
+
+    handler = AuthHandler(config)
+    result = await handler.exchange_auth_token()
+
+    assert result == empty_credential
+
+  @pytest.mark.asyncio
+  async def test_credentials_with_token(
+      self, auth_config, oauth2_credentials_with_token
+  ):
+    """Test when credentials already have a token."""
+    config = AuthConfig(
+        auth_scheme=auth_config.auth_scheme,
+        raw_auth_credential=auth_config.raw_auth_credential,
+        exchanged_auth_credential=oauth2_credentials_with_token,
+    )
+
+    handler = AuthHandler(config)
+    result = await handler.exchange_auth_token()
+
+    assert result == oauth2_credentials_with_token
+
+  @patch("google.adk.auth.oauth2_credential_util.OAuth2Session")
+  @pytest.mark.asyncio
+  async def test_successful_token_exchange(
+      self, mock_oauth2_session, auth_config_with_auth_code
+  ):
+    """Test a successful token exchange."""
+    # Setup mock OAuth2Session
+    mock_client = Mock()
+    mock_oauth2_session.return_value = mock_client
+    mock_tokens = OAuth2Token({
+        "access_token": "mock_access_token",
+        "refresh_token": "mock_refresh_token",
+        "expires_at": int(time.time()) + 3600,
+        "expires_in": 3600,
+    })
+    mock_client.fetch_token.return_value = mock_tokens
+
+    handler = AuthHandler(auth_config_with_auth_code)
+    result = await handler.exchange_auth_token()
+
+    assert result.oauth2.access_token == "mock_access_token"
+    assert result.oauth2.refresh_token == "mock_refresh_token"
+    assert result.auth_type == AuthCredentialTypes.OAUTH2

tests/unittests/auth/test_auth_preprocessor.py

@@ -0,0 +1,581 @@
+# 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.
+
+"""Unit tests for auth_preprocessor module."""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock
+from unittest.mock import Mock
+from unittest.mock import patch
+
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.auth.auth_handler import AuthHandler
+from google.adk.auth.auth_preprocessor import _AuthLlmRequestProcessor
+from google.adk.auth.auth_tool import AuthConfig
+from google.adk.auth.auth_tool import AuthToolArguments
+from google.adk.events.event import Event
+from google.adk.flows.llm_flows.functions import REQUEST_EUC_FUNCTION_CALL_NAME
+from google.adk.models.llm_request import LlmRequest
+import pytest
+
+
+class TestAuthLlmRequestProcessor:
+  """Tests for _AuthLlmRequestProcessor class."""
+
+  @pytest.fixture
+  def processor(self):
+    """Create an _AuthLlmRequestProcessor instance."""
+    return _AuthLlmRequestProcessor()
+
+  @pytest.fixture
+  def mock_llm_agent(self):
+    """Create a mock LlmAgent."""
+    from google.adk.agents.llm_agent import LlmAgent
+
+    agent = Mock(spec=LlmAgent)
+    agent.canonical_tools = AsyncMock(return_value=[])
+    return agent
+
+  @pytest.fixture
+  def mock_non_llm_agent(self):
+    """Create a mock non-LLM agent."""
+    agent = Mock()
+    agent.__class__.__name__ = 'BaseAgent'
+    return agent
+
+  @pytest.fixture
+  def mock_session(self):
+    """Create a mock session."""
+    session = Mock()
+    session.state = {}
+    session.events = []
+    return session
+
+  @pytest.fixture
+  def mock_invocation_context(self, mock_llm_agent, mock_session):
+    """Create a mock invocation context."""
+    context = Mock(spec=InvocationContext)
+    context.agent = mock_llm_agent
+    context.session = mock_session
+    context._get_events.side_effect = lambda **_: context.session.events
+    return context
+
+  @pytest.fixture
+  def mock_llm_request(self):
+    """Create a mock LlmRequest."""
+    return Mock(spec=LlmRequest)
+
+  @pytest.fixture
+  def mock_auth_config(self):
+    """Create a mock AuthConfig."""
+    config = Mock(spec=AuthConfig)
+    config.credential_key = None
+    return config
+
+  @pytest.fixture
+  def mock_function_response_with_auth(self, mock_auth_config):
+    """Create a mock function response with auth data."""
+    function_response = Mock()
+    function_response.name = REQUEST_EUC_FUNCTION_CALL_NAME
+    function_response.id = 'auth_response_id'
+    function_response.response = mock_auth_config
+    return function_response
+
+  @pytest.fixture
+  def mock_function_response_without_auth(self):
+    """Create a mock function response without auth data."""
+    function_response = Mock()
+    function_response.name = 'some_other_function'
+    function_response.id = 'other_response_id'
+    return function_response
+
+  @pytest.fixture
+  def mock_user_event_with_auth_response(
+      self, mock_function_response_with_auth
+  ):
+    """Create a mock user event with auth response."""
+    event = Mock(spec=Event)
+    event.author = 'user'
+    event.content = Mock()  # Non-None content
+    event.get_function_responses.return_value = [
+        mock_function_response_with_auth
+    ]
+    return event
+
+  @pytest.fixture
+  def mock_user_event_without_auth_response(
+      self, mock_function_response_without_auth
+  ):
+    """Create a mock user event without auth response."""
+    event = Mock(spec=Event)
+    event.author = 'user'
+    event.content = Mock()  # Non-None content
+    event.get_function_responses.return_value = [
+        mock_function_response_without_auth
+    ]
+    return event
+
+  @pytest.fixture
+  def mock_user_event_no_responses(self):
+    """Create a mock user event with no responses."""
+    event = Mock(spec=Event)
+    event.author = 'user'
+    event.content = Mock()  # Non-None content
+    event.get_function_responses.return_value = []
+    return event
+
+  @pytest.fixture
+  def mock_agent_event(self):
+    """Create a mock agent-authored event."""
+    event = Mock(spec=Event)
+    event.author = 'test_agent'
+    event.content = Mock()  # Non-None content
+    return event
+
+  @pytest.fixture
+  def mock_event_no_content(self):
+    """Create a mock event with no content."""
+    event = Mock(spec=Event)
+    event.author = 'user'
+    event.content = None
+    return event
+
+  @pytest.fixture
+  def mock_agent_event_with_content(self):
+    """Create a mock agent event with content."""
+    event = Mock(spec=Event)
+    event.author = 'test_agent'
+    event.content = Mock()  # Non-None content
+    return event
+
+  @pytest.mark.asyncio
+  async def test_non_llm_agent_returns_early(
+      self, processor, mock_llm_request, mock_session
+  ):
+    """Test that non-LLM agents return early."""
+    mock_context = Mock(spec=InvocationContext)
+    # Using spec=[] ensures hasattr(agent, 'canonical_tools') returns False.
+    mock_context.agent = Mock(spec=[])
+    mock_context.agent.__class__.__name__ = 'BaseAgent'
+    mock_context.session = mock_session
+
+    result = []
+    async for event in processor.run_async(mock_context, mock_llm_request):
+      result.append(event)
+
+    assert result == []
+
+  @pytest.mark.asyncio
+  async def test_empty_events_returns_early(
+      self, processor, mock_invocation_context, mock_llm_request
+  ):
+    """Test that empty events list returns early."""
+    mock_invocation_context.session.events = []
+
+    result = []
+    async for event in processor.run_async(
+        mock_invocation_context, mock_llm_request
+    ):
+      result.append(event)
+
+    assert result == []
+
+  @pytest.mark.asyncio
+  async def test_no_events_with_content_returns_early(
+      self,
+      processor,
+      mock_invocation_context,
+      mock_llm_request,
+      mock_event_no_content,
+  ):
+    """Test that no events with content returns early."""
+    mock_invocation_context.session.events = [mock_event_no_content]
+
+    result = []
+    async for event in processor.run_async(
+        mock_invocation_context, mock_llm_request
+    ):
+      result.append(event)
+
+    assert result == []
+
+  @pytest.mark.asyncio
+  async def test_last_event_with_content_not_user_authored_returns_early(
+      self,
+      processor,
+      mock_invocation_context,
+      mock_llm_request,
+      mock_event_no_content,
+      mock_agent_event_with_content,
+  ):
+    """Test that last event with content not user-authored returns early."""
+    # Mix of events: user event with no content, then agent event with content
+    mock_invocation_context.session.events = [
+        mock_event_no_content,
+        mock_agent_event_with_content,
+    ]
+
+    result = []
+    async for event in processor.run_async(
+        mock_invocation_context, mock_llm_request
+    ):
+      result.append(event)
+
+    assert result == []
+
+  @pytest.mark.asyncio
+  async def test_last_event_no_responses_returns_early(
+      self,
+      processor,
+      mock_invocation_context,
+      mock_llm_request,
+      mock_user_event_no_responses,
+  ):
+    """Test that user event with no responses returns early."""
+    mock_invocation_context.session.events = [mock_user_event_no_responses]
+
+    result = []
+    async for event in processor.run_async(
+        mock_invocation_context, mock_llm_request
+    ):
+      result.append(event)
+
+    assert result == []
+
+  @pytest.mark.asyncio
+  async def test_last_event_no_auth_responses_returns_early(
+      self,
+      processor,
+      mock_invocation_context,
+      mock_llm_request,
+      mock_user_event_without_auth_response,
+  ):
+    """Test that user event with non-auth responses returns early."""
+    mock_invocation_context.session.events = [
+        mock_user_event_without_auth_response
+    ]
+
+    result = []
+    async for event in processor.run_async(
+        mock_invocation_context, mock_llm_request
+    ):
+      result.append(event)
+
+    assert result == []
+
+  @pytest.mark.asyncio
+  @patch('google.adk.auth.auth_preprocessor.AuthHandler')
+  @patch('google.adk.auth.auth_tool.AuthConfig.model_validate')
+  async def test_ignores_auth_responses_outside_current_branch(
+      self,
+      mock_auth_config_validate,
+      mock_auth_handler_class,
+      processor,
+      mock_invocation_context,
+      mock_llm_request,
+      mock_user_event_with_auth_response,
+  ):
+    """Test auth responses hidden by branch filtering are ignored."""
+    mock_invocation_context.session.events = [
+        mock_user_event_with_auth_response
+    ]
+    mock_invocation_context._get_events.side_effect = None
+    mock_invocation_context._get_events.return_value = []
+
+    result = []
+    async for event in processor.run_async(
+        mock_invocation_context, mock_llm_request
+    ):
+      result.append(event)
+
+    mock_invocation_context._get_events.assert_called_once_with(
+        current_branch=True
+    )
+    mock_auth_config_validate.assert_not_called()
+    mock_auth_handler_class.assert_not_called()
+    assert result == []
+
+  @pytest.mark.asyncio
+  @patch('google.adk.auth.auth_preprocessor.AuthHandler')
+  @patch('google.adk.auth.auth_tool.AuthConfig.model_validate')
+  async def test_processes_auth_response_successfully(
+      self,
+      mock_auth_config_validate,
+      mock_auth_handler_class,
+      processor,
+      mock_invocation_context,
+      mock_llm_request,
+      mock_user_event_with_auth_response,
+      mock_auth_config,
+  ):
+    """Test successful processing of auth response in last event."""
+    # Setup mocks
+    mock_auth_config_validate.return_value = mock_auth_config
+    mock_auth_handler = Mock(spec=AuthHandler)
+    mock_auth_handler.parse_and_store_auth_response = AsyncMock()
+    mock_auth_handler_class.return_value = mock_auth_handler
+
+    mock_invocation_context.session.events = [
+        mock_user_event_with_auth_response
+    ]
+
+    result = []
+    async for event in processor.run_async(
+        mock_invocation_context, mock_llm_request
+    ):
+      result.append(event)
+
+    # Verify auth config validation was called
+    mock_auth_config_validate.assert_called_once()
+
+    # Verify auth handler was created with the config
+    mock_auth_handler_class.assert_called_once_with(
+        auth_config=mock_auth_config
+    )
+
+    # Verify parse_and_store_auth_response was called
+    mock_auth_handler.parse_and_store_auth_response.assert_called_once_with(
+        state=mock_invocation_context.session.state
+    )
+
+  @pytest.mark.asyncio
+  @patch('google.adk.auth.auth_preprocessor.AuthHandler')
+  @patch('google.adk.auth.auth_tool.AuthConfig.model_validate')
+  @patch('google.adk.auth.auth_preprocessor.handle_function_calls_async')
+  async def test_processes_multiple_auth_responses_and_resumes_tools(
+      self,
+      mock_handle_function_calls,
+      mock_auth_config_validate,
+      mock_auth_handler_class,
+      processor,
+      mock_invocation_context,
+      mock_llm_request,
+      mock_auth_config,
+  ):
+    """Test processing multiple auth responses and resuming tools."""
+    # Create multiple auth responses
+    auth_response_1 = Mock()
+    auth_response_1.name = REQUEST_EUC_FUNCTION_CALL_NAME
+    auth_response_1.id = 'auth_id_1'
+    auth_response_1.response = mock_auth_config
+
+    auth_response_2 = Mock()
+    auth_response_2.name = REQUEST_EUC_FUNCTION_CALL_NAME
+    auth_response_2.id = 'auth_id_2'
+    auth_response_2.response = mock_auth_config
+
+    user_event_with_multiple_responses = Mock(spec=Event)
+    user_event_with_multiple_responses.author = 'user'
+    user_event_with_multiple_responses.content = Mock()  # Non-None content
+    user_event_with_multiple_responses.get_function_responses.return_value = [
+        auth_response_1,
+        auth_response_2,
+    ]
+    user_event_with_multiple_responses.get_function_calls.return_value = []
+
+    # Create system function call events
+    system_function_call_1 = Mock()
+    system_function_call_1.id = 'auth_id_1'
+    system_function_call_1.name = REQUEST_EUC_FUNCTION_CALL_NAME
+    system_function_call_1.args = {
+        'function_call_id': 'tool_id_1',
+        'auth_config': mock_auth_config,
+    }
+
+    system_function_call_2 = Mock()
+    system_function_call_2.id = 'auth_id_2'
+    system_function_call_2.name = REQUEST_EUC_FUNCTION_CALL_NAME
+    system_function_call_2.args = {
+        'function_call_id': 'tool_id_2',
+        'auth_config': mock_auth_config,
+    }
+
+    system_event = Mock(spec=Event)
+    system_event.content = Mock()  # Non-None content
+    system_event.get_function_calls.return_value = [
+        system_function_call_1,
+        system_function_call_2,
+    ]
+
+    # Create original function call event
+    original_function_call_1 = Mock()
+    original_function_call_1.id = 'tool_id_1'
+
+    original_function_call_2 = Mock()
+    original_function_call_2.id = 'tool_id_2'
+
+    original_event = Mock(spec=Event)
+    original_event.content = Mock()  # Non-None content
+    original_event.get_function_calls.return_value = [
+        original_function_call_1,
+        original_function_call_2,
+    ]
+
+    # Setup events in order: original -> system -> user_with_responses
+    mock_invocation_context.session.events = [
+        original_event,
+        system_event,
+        user_event_with_multiple_responses,
+    ]
+
+    # Setup mocks
+    mock_auth_config_validate.return_value = mock_auth_config
+    mock_auth_handler = Mock(spec=AuthHandler)
+    mock_auth_handler.parse_and_store_auth_response = AsyncMock()
+    mock_auth_handler_class.return_value = mock_auth_handler
+
+    mock_function_response_event = Mock(spec=Event)
+    mock_handle_function_calls.return_value = mock_function_response_event
+
+    result = []
+    async for event in processor.run_async(
+        mock_invocation_context, mock_llm_request
+    ):
+      result.append(event)
+
+    # Verify auth responses were processed
+    assert mock_auth_handler.parse_and_store_auth_response.call_count == 2
+
+    # Verify function calls were resumed
+    mock_handle_function_calls.assert_called_once()
+    call_args = mock_handle_function_calls.call_args
+    assert call_args[0][1] == original_event  # The original event
+    assert call_args[0][3] == {'tool_id_1', 'tool_id_2'}  # Tools to resume
+
+    # Verify the function response event was yielded
+    assert result == [mock_function_response_event]
+
+  @pytest.mark.asyncio
+  @patch('google.adk.auth.auth_preprocessor.AuthHandler')
+  @patch('google.adk.auth.auth_tool.AuthConfig.model_validate')
+  async def test_no_matching_system_function_calls_returns_early(
+      self,
+      mock_auth_config_validate,
+      mock_auth_handler_class,
+      processor,
+      mock_invocation_context,
+      mock_llm_request,
+      mock_user_event_with_auth_response,
+      mock_auth_config,
+  ):
+    """Test that missing matching system function calls returns early."""
+    # Setup mocks
+    mock_auth_config_validate.return_value = mock_auth_config
+    mock_auth_handler = Mock(spec=AuthHandler)
+    mock_auth_handler.parse_and_store_auth_response = AsyncMock()
+    mock_auth_handler_class.return_value = mock_auth_handler
+
+    # Create a non-matching system event
+    non_matching_function_call = Mock()
+    non_matching_function_call.id = (  # Different from 'auth_response_id'
+        'different_id'
+    )
+
+    system_event = Mock(spec=Event)
+    system_event.content = Mock()  # Non-None content
+    system_event.get_function_calls.return_value = [non_matching_function_call]
+
+    mock_invocation_context.session.events = [
+        system_event,
+        mock_user_event_with_auth_response,
+    ]
+
+    result = []
+    async for event in processor.run_async(
+        mock_invocation_context, mock_llm_request
+    ):
+      result.append(event)
+
+    # Should process auth response but not resume any tools
+    mock_auth_handler.parse_and_store_auth_response.assert_called_once()
+    assert result == []
+
+  @pytest.mark.asyncio
+  @patch('google.adk.auth.auth_preprocessor.AuthHandler')
+  @patch('google.adk.auth.auth_tool.AuthConfig.model_validate')
+  @patch('google.adk.auth.auth_tool.AuthToolArguments.model_validate')
+  async def test_handles_missing_original_function_calls(
+      self,
+      mock_auth_tool_args_validate,
+      mock_auth_config_validate,
+      mock_auth_handler_class,
+      processor,
+      mock_invocation_context,
+      mock_llm_request,
+      mock_user_event_with_auth_response,
+      mock_auth_config,
+  ):
+    """Test handling when original function calls are not found."""
+    # Setup mocks
+    mock_auth_config_validate.return_value = mock_auth_config
+    mock_auth_handler = Mock(spec=AuthHandler)
+    mock_auth_handler.parse_and_store_auth_response = AsyncMock()
+    mock_auth_handler_class.return_value = mock_auth_handler
+
+    # Create matching system function call
+    auth_tool_args = Mock(spec=AuthToolArguments)
+    auth_tool_args.function_call_id = 'tool_id_1'
+    mock_auth_tool_args_validate.return_value = auth_tool_args
+
+    system_function_call = Mock()
+    system_function_call.id = 'auth_response_id'  # Matches the response ID
+    system_function_call.args = {
+        'function_call_id': 'tool_id_1',
+        'auth_config': mock_auth_config,
+    }
+
+    system_event = Mock(spec=Event)
+    system_event.content = Mock()  # Non-None content
+    system_event.get_function_calls.return_value = [system_function_call]
+
+    # Create event with no function calls (original function calls missing)
+    empty_event = Mock(spec=Event)
+    empty_event.content = Mock()  # Non-None content
+    empty_event.get_function_calls.return_value = []
+
+    mock_invocation_context.session.events = [
+        empty_event,
+        system_event,
+        mock_user_event_with_auth_response,
+    ]
+
+    result = []
+    async for event in processor.run_async(
+        mock_invocation_context, mock_llm_request
+    ):
+      result.append(event)
+
+    # Should process auth response but not find original function calls
+    mock_auth_handler.parse_and_store_auth_response.assert_called_once()
+    assert result == []
+
+  @pytest.mark.asyncio
+  async def test_isinstance_check_for_llm_agent(
+      self, processor, mock_llm_request, mock_session
+  ):
+    """Test that isinstance check works correctly for LlmAgent."""
+    # This test ensures the isinstance check work as expected
+
+    # Create a mock that fails isinstance check
+    mock_context = Mock(spec=InvocationContext)
+    # This will fail isinstance(agent, LlmAgent)
+    mock_context.agent = Mock(spec=[])
+    mock_context.session = mock_session
+
+    result = []
+    async for event in processor.run_async(mock_context, mock_llm_request):
+      result.append(event)
+
+    assert result == []

tests/unittests/auth/test_auth_provider_registry.py

@@ -0,0 +1,66 @@
+# 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.
+
+"""Unit tests for the AuthProviderRegistry."""
+
+from google.adk.auth.auth_provider_registry import AuthProviderRegistry
+from google.adk.auth.auth_schemes import CustomAuthScheme
+from google.adk.auth.base_auth_provider import BaseAuthProvider
+from pydantic import Field
+
+
+class SchemeA(CustomAuthScheme):
+  type_: str = Field(default="scheme_a")
+
+
+class SchemeB(CustomAuthScheme):
+  type_: str = Field(default="scheme_b")
+
+
+class TestAuthProviderRegistry:
+  """Test cases for AuthProviderRegistry."""
+
+  def test_register_and_get_provider(self, mocker):
+    """Test registering and retrieving providers for different auth scheme types."""
+    registry = AuthProviderRegistry()
+    provider_a = mocker.create_autospec(BaseAuthProvider, instance=True)
+    provider_b = mocker.create_autospec(BaseAuthProvider, instance=True)
+
+    registry.register(SchemeA, provider_a)
+    registry.register(SchemeB, provider_b)
+
+    assert registry.get_provider(SchemeA()) is provider_a
+    assert registry.get_provider(SchemeB()) is provider_b
+
+    # Test getting by scheme type
+    assert registry.get_provider(SchemeA) is provider_a
+    assert registry.get_provider(SchemeB) is provider_b
+
+  def test_get_unregistered_provider_returns_none(self):
+    """Test that get_provider returns None for unregistered scheme types."""
+    registry = AuthProviderRegistry()
+    assert registry.get_provider(SchemeA()) is None
+    assert registry.get_provider(SchemeA) is None
+
+  def test_register_duplicate_type_overwrites_existing(self, mocker):
+    """Test that registering a provider for an existing type overwrites the previous one."""
+    registry = AuthProviderRegistry()
+    provider_1 = mocker.create_autospec(BaseAuthProvider, instance=True)
+    provider_2 = mocker.create_autospec(BaseAuthProvider, instance=True)
+
+    registry.register(SchemeA, provider_1)
+    registry.register(SchemeA, provider_2)
+
+    assert registry.get_provider(SchemeA()) is provider_2
+    assert registry.get_provider(SchemeA) is provider_2

tests/unittests/auth/test_oauth2_credential_util.py

@@ -0,0 +1,423 @@
+# 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.
+
+import time
+from typing import Optional
+from unittest.mock import Mock
+from unittest.mock import patch
+
+from authlib.oauth2.rfc6749 import OAuth2Token
+from fastapi.openapi.models import OAuth2
+from fastapi.openapi.models import OAuthFlowAuthorizationCode
+from fastapi.openapi.models import OAuthFlows
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_credential import OAuth2Auth
+from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
+from google.adk.auth.oauth2_credential_util import create_oauth2_session
+from google.adk.auth.oauth2_credential_util import update_credential_with_tokens
+import pytest
+
+
+@pytest.fixture
+def openid_connect_scheme() -> OpenIdConnectWithConfig:
+  """Fixture providing a standard OpenIdConnectWithConfig scheme."""
+  return OpenIdConnectWithConfig(
+      type_="openIdConnect",
+      openId_connect_url="https://example.com/.well-known/openid_configuration",
+      authorization_endpoint="https://example.com/auth",
+      token_endpoint="https://example.com/token",
+      scopes=["openid", "profile"],
+  )
+
+
+def create_oauth2_auth_credential(
+    auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+    token_endpoint_auth_method: Optional[str] = None,
+):
+  """Helper function to create OAuth2Auth credential with optional token_endpoint_auth_method."""
+  oauth2_auth = OAuth2Auth(
+      client_id="test_client_id",
+      client_secret="test_client_secret",
+      redirect_uri="https://example.com/callback",
+      state="test_state",
+  )
+  if token_endpoint_auth_method is not None:
+    oauth2_auth.token_endpoint_auth_method = token_endpoint_auth_method
+
+  return AuthCredential(
+      auth_type=auth_type,
+      oauth2=oauth2_auth,
+  )
+
+
+class TestOAuth2CredentialUtil:
+  """Test suite for OAuth2 credential utility functions."""
+
+  def test_create_oauth2_session_openid_connect(self):
+    """Test create_oauth2_session with OpenID Connect scheme."""
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid", "profile"],
+    )
+    credential = create_oauth2_auth_credential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        token_endpoint_auth_method="client_secret_jwt",
+    )
+
+    client, token_endpoint = create_oauth2_session(scheme, credential)
+
+    assert client is not None
+    assert token_endpoint == "https://example.com/token"
+    assert client.client_id == "test_client_id"
+    assert client.client_secret == "test_client_secret"
+
+  def test_create_oauth2_session_oauth2_scheme(self):
+    """Test create_oauth2_session with OAuth2 scheme."""
+    flows = OAuthFlows(
+        authorizationCode=OAuthFlowAuthorizationCode(
+            authorizationUrl="https://example.com/auth",
+            tokenUrl="https://example.com/token",
+            scopes={"read": "Read access", "write": "Write access"},
+        )
+    )
+    scheme = OAuth2(type_="oauth2", flows=flows)
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            redirect_uri="https://example.com/callback",
+        ),
+    )
+
+    client, token_endpoint = create_oauth2_session(scheme, credential)
+
+    assert client is not None
+    assert token_endpoint == "https://example.com/token"
+
+  def test_create_oauth2_session_invalid_scheme(self):
+    """Test create_oauth2_session with invalid scheme."""
+    scheme = Mock()  # Invalid scheme type
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+        ),
+    )
+
+    client, token_endpoint = create_oauth2_session(scheme, credential)
+
+    assert client is None
+    assert token_endpoint is None
+
+  def test_create_oauth2_session_missing_credentials(self):
+    """Test create_oauth2_session with missing credentials."""
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            # Missing client_secret
+        ),
+    )
+
+    client, token_endpoint = create_oauth2_session(scheme, credential)
+
+    assert client is None
+    assert token_endpoint is None
+
+  def _google_openid_scheme(self) -> OpenIdConnectWithConfig:
+    """OpenID Connect scheme that uses Google's OAuth2 token endpoint."""
+    return OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://accounts.google.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://accounts.google.com/o/oauth2/v2/auth",
+        token_endpoint="https://oauth2.googleapis.com/token",
+        scopes=["openid"],
+    )
+
+  @patch.dict("os.environ", {}, clear=True)
+  @patch("google.adk.utils._mtls_utils.configure_session_for_mtls")
+  @patch("google.adk.utils._mtls_utils.use_client_cert_effective")
+  def test_create_oauth2_session_google_endpoint_uses_mtls(
+      self, mock_use_cert, mock_configure
+  ):
+    """Google token endpoint is switched to mTLS when a cert is mounted."""
+    mock_use_cert.return_value = True
+    mock_configure.return_value = True
+    credential = create_oauth2_auth_credential(
+        auth_type=AuthCredentialTypes.OAUTH2
+    )
+
+    client, token_endpoint = create_oauth2_session(
+        self._google_openid_scheme(), credential
+    )
+
+    assert client is not None
+    assert token_endpoint == "https://oauth2.mtls.googleapis.com/token"
+    mock_configure.assert_called_once_with(client)
+
+  @patch.dict("os.environ", {}, clear=True)
+  @patch("google.adk.utils._mtls_utils.configure_session_for_mtls")
+  @patch("google.adk.utils._mtls_utils.use_client_cert_effective")
+  def test_create_oauth2_session_google_endpoint_no_cert_keeps_plain(
+      self, mock_use_cert, mock_configure
+  ):
+    """Without a client cert the plain Google endpoint is kept."""
+    mock_use_cert.return_value = False
+    credential = create_oauth2_auth_credential(
+        auth_type=AuthCredentialTypes.OAUTH2
+    )
+
+    _, token_endpoint = create_oauth2_session(
+        self._google_openid_scheme(), credential
+    )
+
+    assert token_endpoint == "https://oauth2.googleapis.com/token"
+    mock_configure.assert_not_called()
+
+  @patch.dict("os.environ", {}, clear=True)
+  @patch("google.adk.utils._mtls_utils.configure_session_for_mtls")
+  @patch("google.adk.utils._mtls_utils.use_client_cert_effective")
+  def test_create_oauth2_session_cert_unavailable_keeps_plain(
+      self, mock_use_cert, mock_configure
+  ):
+    """If the adapter cannot be mounted, the endpoint is not switched."""
+    mock_use_cert.return_value = True
+    mock_configure.return_value = False
+    credential = create_oauth2_auth_credential(
+        auth_type=AuthCredentialTypes.OAUTH2
+    )
+
+    client, token_endpoint = create_oauth2_session(
+        self._google_openid_scheme(), credential
+    )
+
+    assert token_endpoint == "https://oauth2.googleapis.com/token"
+    mock_configure.assert_called_once_with(client)
+
+  @patch.dict("os.environ", {}, clear=True)
+  @patch("google.adk.utils._mtls_utils.configure_session_for_mtls")
+  @patch("google.adk.utils._mtls_utils.use_client_cert_effective")
+  def test_create_oauth2_session_non_google_endpoint_skips_mtls(
+      self, mock_use_cert, mock_configure
+  ):
+    """Non-Google providers are never switched to an mTLS endpoint."""
+    mock_use_cert.return_value = True
+    credential = create_oauth2_auth_credential(
+        auth_type=AuthCredentialTypes.OAUTH2
+    )
+    scheme = OpenIdConnectWithConfig(
+        type_="openIdConnect",
+        openId_connect_url=(
+            "https://example.com/.well-known/openid_configuration"
+        ),
+        authorization_endpoint="https://example.com/auth",
+        token_endpoint="https://example.com/token",
+        scopes=["openid"],
+    )
+
+    _, token_endpoint = create_oauth2_session(scheme, credential)
+
+    assert token_endpoint == "https://example.com/token"
+    mock_configure.assert_not_called()
+
+  @pytest.mark.parametrize(
+      "token_endpoint_auth_method, expected_auth_method",
+      [
+          ("client_secret_post", "client_secret_post"),
+          (None, "client_secret_basic"),
+      ],
+  )
+  def test_create_oauth2_session_with_token_endpoint_auth_method(
+      self,
+      openid_connect_scheme,
+      token_endpoint_auth_method,
+      expected_auth_method,
+  ):
+    """Test create_oauth2_session with various token_endpoint_auth_method settings."""
+    credential = create_oauth2_auth_credential(
+        token_endpoint_auth_method=token_endpoint_auth_method
+    )
+
+    client, token_endpoint = create_oauth2_session(
+        openid_connect_scheme, credential
+    )
+
+    assert client is not None
+    assert token_endpoint == "https://example.com/token"
+    assert client.client_id == "test_client_id"
+    assert client.client_secret == "test_client_secret"
+    assert client.token_endpoint_auth_method == expected_auth_method
+
+  def test_create_oauth2_session_oauth2_scheme_with_token_endpoint_auth_method(
+      self,
+  ):
+    """Test create_oauth2_session with OAuth2 scheme and token_endpoint_auth_method."""
+    flows = OAuthFlows(
+        authorizationCode=OAuthFlowAuthorizationCode(
+            authorizationUrl="https://example.com/auth",
+            tokenUrl="https://example.com/token",
+            scopes={"read": "Read access", "write": "Write access"},
+        )
+    )
+    scheme = OAuth2(type_="oauth2", flows=flows)
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            redirect_uri="https://example.com/callback",
+            token_endpoint_auth_method="client_secret_jwt",
+        ),
+    )
+
+    client, token_endpoint = create_oauth2_session(scheme, credential)
+
+    assert client is not None
+    assert token_endpoint == "https://example.com/token"
+    assert client.token_endpoint_auth_method == "client_secret_jwt"
+
+  def _oauth2_scheme_with_scopes(self):
+    """Build an OAuth2 scheme that declares scopes."""
+    return OAuth2(
+        type_="oauth2",
+        flows=OAuthFlows(
+            authorizationCode=OAuthFlowAuthorizationCode(
+                authorizationUrl="https://example.com/auth",
+                tokenUrl="https://example.com/token",
+                scopes={"read": "Read access", "write": "Write access"},
+            )
+        ),
+    )
+
+  def _capturing_post(self, captured):
+    """Stub for OAuth2Session.post that records the token-request body."""
+
+    def _post(*args, **kwargs):
+      captured["data"] = kwargs.get("data")
+      response = Mock()
+      response.status_code = 200
+      response.json.return_value = {
+          "access_token": "new_access_token",
+          "token_type": "Bearer",
+          "expires_in": 3600,
+          "refresh_token": "new_refresh_token",
+      }
+      return response
+
+    return _post
+
+  def test_refresh_request_omits_scope(self):
+    """Refresh requests must not carry scope (some providers reject it)."""
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            redirect_uri="https://example.com/callback",
+        ),
+    )
+
+    client, token_endpoint = create_oauth2_session(
+        self._oauth2_scheme_with_scopes(), credential
+    )
+    assert client is not None
+
+    captured = {}
+    client.post = self._capturing_post(captured)
+    client.refresh_token(token_endpoint, refresh_token="old_refresh_token")
+
+    assert "scope" not in captured["data"]
+
+  def test_token_exchange_omits_scope(self):
+    """Authorization-code exchange must not carry scope (it is redundant)."""
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+            redirect_uri="https://example.com/callback",
+        ),
+    )
+
+    client, token_endpoint = create_oauth2_session(
+        self._oauth2_scheme_with_scopes(), credential
+    )
+    assert client is not None
+
+    captured = {}
+    client.post = self._capturing_post(captured)
+    client.fetch_token(
+        token_endpoint, grant_type="authorization_code", code="test_code"
+    )
+
+    assert "scope" not in captured["data"]
+
+  def test_update_credential_with_tokens(self):
+    """Test update_credential_with_tokens function."""
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
+        oauth2=OAuth2Auth(
+            client_id="test_client_id",
+            client_secret="test_client_secret",
+        ),
+    )
+
+    # Store the expected expiry time to avoid timing issues
+    expected_expires_at = int(time.time()) + 3600
+    tokens = OAuth2Token({
+        "access_token": "new_access_token",
+        "refresh_token": "new_refresh_token",
+        "id_token": "new_id_token",
+        "expires_at": expected_expires_at,
+        "expires_in": 3600,
+    })
+
+    assert credential.oauth2 is not None
+
+    update_credential_with_tokens(credential, tokens)
+
+    assert credential.oauth2.access_token == "new_access_token"
+    assert credential.oauth2.refresh_token == "new_refresh_token"
+    assert credential.oauth2.id_token == "new_id_token"
+    assert credential.oauth2.expires_at == expected_expires_at
+    assert credential.oauth2.expires_in == 3600
+
+  def test_update_credential_with_tokens_none(self) -> None:
+    credential = AuthCredential(
+        auth_type=AuthCredentialTypes.API_KEY,
+    )
+    tokens = OAuth2Token({"access_token": "new_access_token"})
+
+    # Should not raise any exceptions when oauth2 is None
+    update_credential_with_tokens(credential, tokens)
+    assert credential.oauth2 is None

tests/unittests/auth/test_oauth2_discovery.py

@@ -0,0 +1,285 @@
+# 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.
+
+import json
+from unittest.mock import call
+from unittest.mock import Mock
+from unittest.mock import patch
+
+from google.adk.auth.oauth2_discovery import AuthorizationServerMetadata
+from google.adk.auth.oauth2_discovery import OAuth2DiscoveryManager
+from google.adk.auth.oauth2_discovery import ProtectedResourceMetadata
+import httpx
+import pytest
+
+
+class TestOAuth2Discovery:
+  """Tests for the OAuth2DiscoveryManager class."""
+
+  @pytest.fixture
+  def auth_server_metadata(self):
+    """Create AuthorizationServerMetadata object."""
+    return AuthorizationServerMetadata(
+        issuer="https://auth.example.com",
+        authorization_endpoint="https://auth.example.com/authorize",
+        token_endpoint="https://auth.example.com/token",
+        scopes_supported=["read", "write"],
+    )
+
+  @pytest.fixture
+  def resource_metadata(self):
+    """Create ProtectedResourceMetadata object."""
+    return ProtectedResourceMetadata(
+        resource="https://resource.example.com",
+        authorization_servers=["https://auth.example.com"],
+    )
+
+  @pytest.fixture
+  def mock_failed_response(self):
+    """Create a mock HTTP response with a failure status."""
+    response = Mock()
+    response.raise_for_status.side_effect = httpx.HTTPError("Failed")
+    return response
+
+  @pytest.fixture
+  def mock_empty_response(self):
+    """Create a mock HTTP response with an empty JSON body."""
+    response = Mock()
+    response.json = lambda: {}
+    return response
+
+  @pytest.fixture
+  def mock_invalid_json_response(self):
+    """Create a mock HTTP response with an invalid JSON body."""
+    response = Mock()
+    response.json.side_effect = json.decoder.JSONDecodeError(
+        "Invalid JSON", "invalid_json", 0
+    )
+    return response
+
+  def mock_success_response(self, json_data):
+    """Create a mock HTTP successful response with auth server metadata."""
+    response = Mock()
+    response.json = json_data.model_dump
+    return response
+
+  @patch("httpx.AsyncClient.get")
+  @pytest.mark.asyncio
+  async def test_discover_auth_server_metadata_failed(
+      self,
+      mock_get,
+      mock_failed_response,
+  ):
+    """Test discovering auth server metadata with failed response."""
+
+    mock_get.side_effect = mock_failed_response
+    discovery_manager = OAuth2DiscoveryManager()
+    result = await discovery_manager.discover_auth_server_metadata(
+        "https://auth.example.com"
+    )
+    assert not result
+    mock_get.assert_has_calls([
+        call(
+            "https://auth.example.com/.well-known/oauth-authorization-server",
+            timeout=5,
+        ),
+        call(
+            "https://auth.example.com/.well-known/openid-configuration",
+            timeout=5,
+        ),
+    ])
+
+  @pytest.mark.asyncio
+  async def test_discover_metadata_invalid_url(self):
+    """Test discovering resource/auth metadata with an invalid URL."""
+    discovery_manager = OAuth2DiscoveryManager()
+    result = await discovery_manager.discover_auth_server_metadata("bad_url")
+    assert not result
+    result = await discovery_manager.discover_resource_metadata("bad_url")
+    assert not result
+
+  @patch("httpx.AsyncClient.get")
+  @pytest.mark.asyncio
+  async def test_discover_auth_server_metadata_without_path(
+      self,
+      mock_get,
+      auth_server_metadata,
+      mock_empty_response,
+  ):
+    """Test discovering auth server metadata with an issuer URL without a path."""
+
+    mock_get.side_effect = [
+        mock_empty_response,
+        self.mock_success_response(auth_server_metadata),
+    ]
+    discovery_manager = OAuth2DiscoveryManager()
+    result = await discovery_manager.discover_auth_server_metadata(
+        "https://auth.example.com/"
+    )
+    assert result == auth_server_metadata
+    mock_get.assert_has_calls([
+        call(
+            "https://auth.example.com/.well-known/oauth-authorization-server",
+            timeout=5,
+        ),
+        call(
+            "https://auth.example.com/.well-known/openid-configuration",
+            timeout=5,
+        ),
+    ])
+
+  @patch("httpx.AsyncClient.get")
+  @pytest.mark.asyncio
+  async def test_discover_auth_server_metadata_with_path(
+      self,
+      mock_get,
+      auth_server_metadata,
+      mock_failed_response,
+      mock_invalid_json_response,
+  ):
+    """Test discovering auth server metadata with an issuer URL with a path."""
+
+    auth_server_metadata.issuer = "https://auth.example.com/oauth"
+    mock_get.side_effect = [
+        mock_failed_response,
+        mock_invalid_json_response,
+        self.mock_success_response(auth_server_metadata),
+    ]
+    discovery_manager = OAuth2DiscoveryManager()
+    result = await discovery_manager.discover_auth_server_metadata(
+        "https://auth.example.com/oauth"
+    )
+    assert result == auth_server_metadata
+    mock_get.assert_has_calls([
+        call(
+            "https://auth.example.com/.well-known/oauth-authorization-server/oauth",
+            timeout=5,
+        ),
+        call(
+            "https://auth.example.com/.well-known/openid-configuration/oauth",
+            timeout=5,
+        ),
+        call(
+            "https://auth.example.com/oauth/.well-known/openid-configuration",
+            timeout=5,
+        ),
+    ])
+
+  @patch("httpx.AsyncClient.get")
+  @pytest.mark.asyncio
+  async def test_discover_auth_server_metadata_discard_mismatched_issuer(
+      self,
+      mock_get,
+      auth_server_metadata,
+  ):
+    """Test discover_auth_server_metadata() discards response with mismatched issuer."""
+
+    bad_auth_server_metadata = auth_server_metadata.model_copy(
+        update={"issuer": "https://bad.example.com"}
+    )
+    mock_get.side_effect = [
+        self.mock_success_response(bad_auth_server_metadata),
+        self.mock_success_response(auth_server_metadata),
+    ]
+    discovery_manager = OAuth2DiscoveryManager()
+    result = await discovery_manager.discover_auth_server_metadata(
+        "https://auth.example.com"
+    )
+    assert result == auth_server_metadata
+    mock_get.assert_has_calls([
+        call(
+            "https://auth.example.com/.well-known/oauth-authorization-server",
+            timeout=5,
+        ),
+        call(
+            "https://auth.example.com/.well-known/openid-configuration",
+            timeout=5,
+        ),
+    ])
+
+  @patch("httpx.AsyncClient.get")
+  @pytest.mark.asyncio
+  async def test_discover_resource_metadata_failed(
+      self,
+      mock_get,
+      mock_failed_response,
+  ):
+    """Test discovering resource metadata fails."""
+
+    mock_get.return_value = mock_failed_response
+    discovery_manager = OAuth2DiscoveryManager()
+    result = await discovery_manager.discover_resource_metadata(
+        "https://resource.example.com"
+    )
+    assert not result
+    mock_get.assert_called_once_with(
+        "https://resource.example.com/.well-known/oauth-protected-resource",
+        timeout=5,
+    )
+
+  @patch("httpx.AsyncClient.get")
+  @pytest.mark.asyncio
+  async def test_discover_resource_metadata_without_path(
+      self, mock_get, resource_metadata
+  ):
+    """Test discovering resource metadata with a resource URL without a path."""
+    mock_get.return_value = self.mock_success_response(resource_metadata)
+    discovery_manager = OAuth2DiscoveryManager()
+    result = await discovery_manager.discover_resource_metadata(
+        "https://resource.example.com/"
+    )
+    assert result == resource_metadata
+    mock_get.assert_called_once_with(
+        "https://resource.example.com/.well-known/oauth-protected-resource",
+        timeout=5,
+    )
+
+  @patch("httpx.AsyncClient.get")
+  @pytest.mark.asyncio
+  async def test_discover_resource_metadata_with_path(
+      self, mock_get, resource_metadata
+  ):
+    """Test discovering resource metadata with a resource URL with a path."""
+    resource_metadata.resource = "https://resource.example.com/tenant1"
+    mock_get.return_value = self.mock_success_response(resource_metadata)
+    discovery_manager = OAuth2DiscoveryManager()
+    result = await discovery_manager.discover_resource_metadata(
+        "https://resource.example.com/tenant1"
+    )
+    assert result == resource_metadata
+    mock_get.assert_called_once_with(
+        "https://resource.example.com/.well-known/oauth-protected-resource/tenant1",
+        timeout=5,
+    )
+
+  @patch("httpx.AsyncClient.get")
+  @pytest.mark.asyncio
+  async def test_discover_resource_metadata_discard_mismatched_resource(
+      self,
+      mock_get,
+      resource_metadata,
+  ):
+    """Test discover_resource_metadata() discards response with mismatched resource."""
+
+    resource_metadata.resource = "https://bad.example.com"
+    mock_get.return_value = self.mock_success_response(resource_metadata)
+    discovery_manager = OAuth2DiscoveryManager()
+    result = await discovery_manager.discover_resource_metadata(
+        "https://resource.example.com"
+    )
+    assert not result
+    mock_get.assert_called_once_with(
+        "https://resource.example.com/.well-known/oauth-protected-resource",
+        timeout=5,
+    )

tests/unittests/auth/test_toolset_auth.py

@@ -0,0 +1,449 @@
+# 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.
+
+"""Tests for toolset authentication functionality."""
+
+from typing import Optional
+from unittest.mock import AsyncMock
+from unittest.mock import MagicMock
+from unittest.mock import Mock
+from unittest.mock import patch
+
+from fastapi.openapi.models import OAuth2
+from fastapi.openapi.models import OAuthFlowAuthorizationCode
+from fastapi.openapi.models import OAuthFlows
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.auth.auth_credential import OAuth2Auth
+from google.adk.auth.auth_preprocessor import TOOLSET_AUTH_CREDENTIAL_ID_PREFIX
+from google.adk.auth.auth_tool import AuthConfig
+from google.adk.auth.auth_tool import AuthToolArguments
+from google.adk.flows.llm_flows.base_llm_flow import _resolve_toolset_auth
+from google.adk.flows.llm_flows.base_llm_flow import TOOLSET_AUTH_CREDENTIAL_ID_PREFIX as FLOW_PREFIX
+from google.adk.flows.llm_flows.functions import build_auth_request_event
+from google.adk.flows.llm_flows.functions import REQUEST_EUC_FUNCTION_CALL_NAME
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.base_toolset import BaseToolset
+import pytest
+
+
+class MockToolset(BaseToolset):
+  """A mock toolset for testing."""
+
+  def __init__(
+      self,
+      auth_config: Optional[AuthConfig] = None,
+      tools: Optional[list[BaseTool]] = None,
+  ):
+    super().__init__()
+    self._auth_config = auth_config
+    self._tools = tools or []
+
+  def get_auth_config(self) -> Optional[AuthConfig]:
+    return self._auth_config
+
+  async def get_tools(self, readonly_context=None) -> list[BaseTool]:
+    return self._tools
+
+  async def close(self):
+    pass
+
+
+def create_oauth2_auth_config() -> AuthConfig:
+  """Create a sample OAuth2 auth config for testing."""
+  return AuthConfig(
+      auth_scheme=OAuth2(
+          flows=OAuthFlows(
+              authorizationCode=OAuthFlowAuthorizationCode(
+                  authorizationUrl="https://example.com/auth",
+                  tokenUrl="https://example.com/token",
+                  scopes={"read": "Read access"},
+              )
+          )
+      ),
+      raw_auth_credential=AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          oauth2=OAuth2Auth(
+              client_id="test_client_id",
+              client_secret="test_client_secret",
+          ),
+      ),
+  )
+
+
+class TestToolsetAuthPrefixConstant:
+  """Test that prefix constants are consistent."""
+
+  def test_prefix_constants_match(self):
+    """Ensure auth_preprocessor and _reasoning use the same prefix."""
+    assert TOOLSET_AUTH_CREDENTIAL_ID_PREFIX == FLOW_PREFIX
+    assert TOOLSET_AUTH_CREDENTIAL_ID_PREFIX == "_adk_toolset_auth_"
+
+
+class TestResolveToolsetAuth:
+  """Tests for _resolve_toolset_auth."""
+
+  @pytest.fixture
+  def mock_invocation_context(self):
+    """Create a mock invocation context."""
+    ctx = Mock(spec=InvocationContext)
+    ctx._state_schema = None
+    ctx.invocation_id = "test-invocation-id"
+    ctx.end_invocation = False
+    ctx.branch = None
+    ctx.session = Mock()
+    ctx.session.state = {}
+    ctx.session.id = "test-session-id"
+    ctx.credential_service = None
+    ctx.app_name = "test-app"
+    ctx.user_id = "test-user"
+    ctx.credential_by_key = {}
+    return ctx
+
+  @pytest.fixture
+  def mock_agent(self):
+    """Create a mock LLM agent."""
+    agent = Mock()
+    agent.name = "test-agent"
+    agent.tools = []
+    return agent
+
+  @pytest.mark.asyncio
+  async def test_no_tools_returns_no_events(
+      self, mock_invocation_context, mock_agent
+  ):
+    """Test that no events are yielded when agent has no tools."""
+    mock_agent.tools = []
+
+    events = []
+    async for event in _resolve_toolset_auth(
+        mock_invocation_context, mock_agent
+    ):
+      events.append(event)
+
+    assert len(events) == 0
+    assert mock_invocation_context.end_invocation is False
+
+  @pytest.mark.asyncio
+  async def test_toolset_without_auth_config_skipped(
+      self, mock_invocation_context, mock_agent
+  ):
+    """Test that toolsets without auth config are skipped."""
+    toolset = MockToolset(auth_config=None)
+    mock_agent.tools = [toolset]
+
+    events = []
+    async for event in _resolve_toolset_auth(
+        mock_invocation_context, mock_agent
+    ):
+      events.append(event)
+
+    assert len(events) == 0
+    assert mock_invocation_context.end_invocation is False
+
+  @pytest.mark.asyncio
+  async def test_toolset_with_credential_available_populates_context(
+      self, mock_invocation_context, mock_agent
+  ):
+    """Test that credential is stored in invocation context when available."""
+    auth_config = create_oauth2_auth_config()
+    toolset = MockToolset(auth_config=auth_config)
+    mock_agent.tools = [toolset]
+
+    # Mock CredentialManager to return a credential
+    mock_credential = AuthCredential(
+        auth_type=AuthCredentialTypes.OAUTH2,
+        oauth2=OAuth2Auth(access_token="test-token"),
+    )
+
+    with patch(
+        "google.adk.auth.credential_manager.CredentialManager"
+    ) as MockCredentialManager:
+      mock_manager = AsyncMock()
+      mock_manager.get_auth_credential = AsyncMock(return_value=mock_credential)
+      MockCredentialManager.return_value = mock_manager
+
+      events = []
+      async for event in _resolve_toolset_auth(
+          mock_invocation_context, mock_agent
+      ):
+        events.append(event)
+
+    # No auth request events - credential was available
+    assert len(events) == 0
+    assert mock_invocation_context.end_invocation is False
+    # Credential should be stored in invocation context, not auth_config
+    assert (
+        mock_invocation_context.credential_by_key[auth_config.credential_key]
+        == mock_credential
+    )
+    assert auth_config.exchanged_auth_credential is None
+
+  @pytest.mark.asyncio
+  async def test_toolset_auth_uses_copy_and_does_not_mutate_shared_config(
+      self, mock_invocation_context, mock_agent
+  ):
+    """Test that _resolve_toolset_auth uses a copy and does not mutate shared config."""
+    auth_config = create_oauth2_auth_config()
+    toolset = MockToolset(auth_config=auth_config)
+    mock_agent.tools = [toolset]
+
+    def create_mock_cm(cfg):
+      m = AsyncMock()
+      m._auth_config = cfg
+
+      async def get_cred(ctx):
+        cfg.exchanged_auth_credential = AuthCredential(
+            auth_type=AuthCredentialTypes.OAUTH2,
+            oauth2=OAuth2Auth(auth_uri="https://example.com/consent"),
+        )
+        return None
+
+      m.get_auth_credential = AsyncMock(side_effect=get_cred)
+      return m
+
+    with patch(
+        "google.adk.auth.credential_manager.CredentialManager",
+        side_effect=create_mock_cm,
+    ):
+      events = []
+      async for event in _resolve_toolset_auth(
+          mock_invocation_context, mock_agent
+      ):
+        events.append(event)
+
+    # Should yield one auth request event
+    assert len(events) == 1
+    assert mock_invocation_context.end_invocation is True
+
+    # The shared auth_config should NOT be mutated
+    assert auth_config.exchanged_auth_credential is None
+
+  @pytest.mark.asyncio
+  async def test_toolset_without_credential_yields_auth_event(
+      self, mock_invocation_context, mock_agent
+  ):
+    """Test that auth request event is yielded when credential not available."""
+    auth_config = create_oauth2_auth_config()
+    toolset = MockToolset(auth_config=auth_config)
+    mock_agent.tools = [toolset]
+
+    with patch(
+        "google.adk.auth.credential_manager.CredentialManager"
+    ) as MockCredentialManager:
+      mock_manager = AsyncMock()
+      mock_manager.get_auth_credential = AsyncMock(return_value=None)
+      MockCredentialManager.return_value = mock_manager
+
+      events = []
+      async for event in _resolve_toolset_auth(
+          mock_invocation_context, mock_agent
+      ):
+        events.append(event)
+
+    # Should yield one auth request event
+    assert len(events) == 1
+    assert mock_invocation_context.end_invocation is True
+
+    # Check event structure
+    event = events[0]
+    assert event.invocation_id == "test-invocation-id"
+    assert event.author == "test-agent"
+    assert event.content is not None
+    assert len(event.content.parts) == 1
+
+    # Check function call
+    fc = event.content.parts[0].function_call
+    assert fc.name == REQUEST_EUC_FUNCTION_CALL_NAME
+    # The args use camelCase aliases from the pydantic model
+    assert fc.args["functionCallId"].startswith(
+        TOOLSET_AUTH_CREDENTIAL_ID_PREFIX
+    )
+    assert "MockToolset" in fc.args["functionCallId"]
+
+  @pytest.mark.asyncio
+  async def test_multiple_toolsets_needing_auth(
+      self, mock_invocation_context, mock_agent
+  ):
+    """Test that multiple toolsets needing auth yield multiple function calls."""
+    auth_config1 = create_oauth2_auth_config()
+    auth_config2 = create_oauth2_auth_config()
+    toolset1 = MockToolset(auth_config=auth_config1)
+    toolset2 = MockToolset(auth_config=auth_config2)
+    mock_agent.tools = [toolset1, toolset2]
+
+    with patch(
+        "google.adk.auth.credential_manager.CredentialManager"
+    ) as MockCredentialManager:
+      mock_manager = AsyncMock()
+      mock_manager.get_auth_credential = AsyncMock(return_value=None)
+      MockCredentialManager.return_value = mock_manager
+
+      events = []
+      async for event in _resolve_toolset_auth(
+          mock_invocation_context, mock_agent
+      ):
+        events.append(event)
+
+    # Should yield one event with multiple function calls
+    # But since both toolsets have same class name, they'll have same ID
+    # and only one will be in pending_auth_requests (dict overwrites)
+    assert len(events) == 1
+    assert mock_invocation_context.end_invocation is True
+
+
+class TestAuthPreprocessorToolsetAuthSkip:
+  """Tests for auth preprocessor skipping toolset auth."""
+
+  def test_toolset_auth_prefix_skipped(self):
+    """Test that function calls with toolset auth prefix are skipped."""
+    from google.adk.auth.auth_preprocessor import TOOLSET_AUTH_CREDENTIAL_ID_PREFIX
+
+    # Verify the prefix is correct
+    assert TOOLSET_AUTH_CREDENTIAL_ID_PREFIX == "_adk_toolset_auth_"
+
+    # Test that a function_call_id starting with this prefix would be skipped
+    toolset_function_call_id = f"{TOOLSET_AUTH_CREDENTIAL_ID_PREFIX}McpToolset"
+    assert toolset_function_call_id.startswith(
+        TOOLSET_AUTH_CREDENTIAL_ID_PREFIX
+    )
+
+    # Regular tool auth function_call_id should NOT start with prefix
+    regular_function_call_id = "call_123"
+    assert not regular_function_call_id.startswith(
+        TOOLSET_AUTH_CREDENTIAL_ID_PREFIX
+    )
+
+
+class TestCallbackContextGetAuthResponse:
+  """Tests for CallbackContext.get_auth_response method."""
+
+  @pytest.fixture
+  def mock_invocation_context(self):
+    """Create a mock invocation context."""
+    ctx = Mock(spec=InvocationContext)
+    ctx._state_schema = None
+    ctx.session = Mock()
+    ctx.session.state = {}
+    return ctx
+
+  def test_get_auth_response_returns_none_when_no_response(
+      self, mock_invocation_context
+  ):
+    """Test that get_auth_response returns None when no auth response in state."""
+    callback_context = CallbackContext(mock_invocation_context)
+    auth_config = create_oauth2_auth_config()
+
+    result = callback_context.get_auth_response(auth_config)
+
+    # Should return None when no auth response is stored
+    assert result is None
+
+  def test_get_auth_response_delegates_to_auth_handler(
+      self, mock_invocation_context
+  ):
+    """Test that get_auth_response delegates to AuthHandler."""
+    callback_context = CallbackContext(mock_invocation_context)
+    auth_config = create_oauth2_auth_config()
+
+    # AuthHandler is imported inside the method, so we patch the module
+    with patch("google.adk.auth.auth_handler.AuthHandler") as MockAuthHandler:
+      mock_handler = Mock()
+      mock_handler.get_auth_response = Mock(return_value=None)
+      MockAuthHandler.return_value = mock_handler
+
+      callback_context.get_auth_response(auth_config)
+
+      MockAuthHandler.assert_called_once_with(auth_config)
+      mock_handler.get_auth_response.assert_called_once()
+
+
+class TestBuildAuthRequestEvent:
+  """Tests for build_auth_request_event helper function."""
+
+  @pytest.fixture
+  def mock_invocation_context(self):
+    """Create a mock invocation context."""
+    ctx = Mock(spec=InvocationContext)
+    ctx._state_schema = None
+    ctx.invocation_id = "test-invocation-id"
+    ctx.branch = None
+    ctx.agent = Mock()
+    ctx.agent.name = "test-agent"
+    return ctx
+
+  def test_builds_event_with_auth_requests(self, mock_invocation_context):
+    """Test that build_auth_request_event creates correct event."""
+    auth_requests = {
+        "call_123": create_oauth2_auth_config(),
+    }
+
+    event = build_auth_request_event(mock_invocation_context, auth_requests)
+
+    assert event.invocation_id == "test-invocation-id"
+    assert event.author == "test-agent"
+    assert event.content is not None
+    assert len(event.content.parts) == 1
+
+    fc = event.content.parts[0].function_call
+    assert fc.name == REQUEST_EUC_FUNCTION_CALL_NAME
+    assert fc.args["functionCallId"] == "call_123"
+
+  def test_multiple_auth_requests_create_multiple_parts(
+      self, mock_invocation_context
+  ):
+    """Test that multiple auth requests create multiple function call parts."""
+    auth_requests = {
+        "call_1": create_oauth2_auth_config(),
+        "call_2": create_oauth2_auth_config(),
+    }
+
+    event = build_auth_request_event(mock_invocation_context, auth_requests)
+
+    assert len(event.content.parts) == 2
+    function_call_ids = {
+        p.function_call.args["functionCallId"] for p in event.content.parts
+    }
+    assert function_call_ids == {"call_1", "call_2"}
+
+  def test_always_adds_long_running_tool_ids(self, mock_invocation_context):
+    """Test that long_running_tool_ids is always set."""
+    auth_requests = {"call_123": create_oauth2_auth_config()}
+
+    event = build_auth_request_event(mock_invocation_context, auth_requests)
+
+    assert event.long_running_tool_ids is not None
+    assert len(event.long_running_tool_ids) == 1
+
+  def test_custom_author_overrides_default(self, mock_invocation_context):
+    """Test that custom author overrides default agent name."""
+    auth_requests = {"call_123": create_oauth2_auth_config()}
+
+    event = build_auth_request_event(
+        mock_invocation_context, auth_requests, author="custom-author"
+    )
+
+    assert event.author == "custom-author"
+
+  def test_role_is_set_in_content(self, mock_invocation_context):
+    """Test that role is set in content."""
+    auth_requests = {"call_123": create_oauth2_auth_config()}
+
+    event = build_auth_request_event(
+        mock_invocation_context, auth_requests, role="model"
+    )
+
+    assert event.content.role == "model"