Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 29 additions & 3 deletions google/genai/_api_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -865,6 +865,7 @@ def _use_google_auth_async(self) -> bool:
return bool(
has_aiohttp
and self.vertexai
and hasattr(mtls, 'should_use_client_cert')
and mtls.should_use_client_cert() # type: ignore[no-untyped-call]
and mtls.has_default_client_cert_source() # type: ignore[no-untyped-call]
and not self._http_options.httpx_async_client
Expand All @@ -877,11 +878,36 @@ async def _get_aiohttp_session(

if self._aiohttp_session is None and self._use_google_auth_async():
try:
from google.auth.aio.credentials import StaticCredentials
from google.auth.aio.credentials import Credentials as AsyncCredentials
from google.auth.aio.transport.sessions import AsyncAuthorizedSession

async_creds = StaticCredentials(token=self._access_token()) # type: ignore[no-untyped-call]
self._aiohttp_session = AsyncAuthorizedSession(async_creds) # type: ignore[no-untyped-call,assignment]
class _RefreshableAsyncCredentials(AsyncCredentials): # type: ignore[misc, valid-type]
"""Adapter to use the client's sync credentials in an AsyncAuthorizedSession."""

def __init__(self, client: 'BaseApiClient'):
super().__init__() # type: ignore[no-untyped-call]
self._client = client

async def before_request(
self, request: Any, method: str, url: str, headers: dict[str, str]
) -> None:
token = await self._client._async_access_token()
headers['Authorization'] = f'Bearer {token}'
if (
self._client._credentials
and self._client._credentials.quota_project_id
):
headers['x-goog-user-project'] = (
self._client._credentials.quota_project_id
)

@property
def valid(self) -> bool:
if not self._client._credentials:
return False
return not self._client._credentials.expired

self._aiohttp_session = AsyncAuthorizedSession(_RefreshableAsyncCredentials(self)) # type: ignore[no-untyped-call,assignment]
return self._aiohttp_session # type: ignore[return-value]
except ImportError:
pass
Expand Down
57 changes: 57 additions & 0 deletions google/genai/tests/client/test_client_initialization.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import logging
import os
import ssl
import sys
from unittest import mock

import certifi
Expand Down Expand Up @@ -1870,3 +1871,59 @@ async def test_get_aiohttp_session():
assert initial_session is not None
session = await client._api_client._get_aiohttp_session()
assert session is initial_session


@requires_aiohttp
@pytest.mark.asyncio
async def test_async_mtls_uses_refreshable_credentials(monkeypatch):
"""Tests that _RefreshableAsyncCredentials is used in async mTLS path."""
from google.genai import _api_client

# Ensure _use_google_auth_async returns True
monkeypatch.setattr(_api_client, "has_aiohttp", True)
monkeypatch.setattr(_api_client.mtls, "should_use_client_cert", lambda: True, raising=False)
monkeypatch.setattr(
_api_client.mtls, "has_default_client_cert_source", lambda: True
)

# Mock AsyncAuthorizedSession and google.auth.aio modules
mock_session = mock.MagicMock()
mock_auth_aio = mock.MagicMock()
monkeypatch.setitem(sys.modules, "google.auth.aio", mock_auth_aio)
monkeypatch.setitem(
sys.modules, "google.auth.aio.credentials", mock_auth_aio.credentials
)
monkeypatch.setitem(
sys.modules, "google.auth.aio.transport", mock_auth_aio.transport
)
monkeypatch.setitem(
sys.modules,
"google.auth.aio.transport.sessions",
mock_auth_aio.transport.sessions,
)
mock_auth_aio.transport.sessions.AsyncAuthorizedSession = mock_session
mock_auth_aio.credentials.Credentials = mock.MagicMock

# Mock credentials
mock_creds = mock.MagicMock()
mock_creds.expired = False
mock_creds.token = "initial_token"
monkeypatch.setattr(
google.auth, "default", lambda scopes=None: (mock_creds, "fake-project")
)

client = Client(vertexai=True, project="fake-project")
client._api_client._credentials = mock_creds

# Trigger session creation
await client._api_client._get_aiohttp_session()

# Verify AsyncAuthorizedSession was called with _RefreshableAsyncCredentials
assert mock_session.call_count == 1
passed_creds = mock_session.call_args[0][0]
assert type(passed_creds).__name__ == "_RefreshableAsyncCredentials"

# Verify valid property
assert passed_creds.valid == True
mock_creds.expired = True
assert passed_creds.valid == False
Loading