diff --git a/google/genai/_api_client.py b/google/genai/_api_client.py index 4d90ba86e..c8e821ea2 100644 --- a/google/genai/_api_client.py +++ b/google/genai/_api_client.py @@ -806,7 +806,8 @@ def __init__( append_library_version_headers(self._http_options.headers) client_args, async_client_args = self._ensure_httpx_ssl_ctx( - self._http_options + self._http_options, + vertexai=bool(self.vertexai), ) self._async_httpx_client_args = async_client_args self._authorized_session: Optional[AuthorizedSession] = None @@ -837,7 +838,10 @@ def __init__( import aiohttp # pylint: disable=g-import-not-at-top # Do it once at the genai.Client level. Share among all requests. self._async_client_session_request_args = ( - self._ensure_aiohttp_ssl_ctx(self._http_options) + self._ensure_aiohttp_ssl_ctx( + self._http_options, + vertexai=bool(self.vertexai), + ) ) if self._use_google_auth_async(): self._async_client_session_request_args['ssl'] = True # type: ignore[no-untyped-call] @@ -851,7 +855,10 @@ def __init__( pass retry_kwargs = retry_args(self._http_options.retry_options) - self._websocket_ssl_ctx = self._ensure_websocket_ssl_ctx(self._http_options) + self._websocket_ssl_ctx = self._ensure_websocket_ssl_ctx( + self._http_options, + vertexai=bool(self.vertexai), + ) self._retry = tenacity.Retrying(**retry_kwargs) self._async_retry = tenacity.AsyncRetrying(**retry_kwargs) @@ -1014,6 +1021,7 @@ def __del__(self, _warnings: Any = warnings) -> None: @staticmethod def _ensure_httpx_ssl_ctx( options: HttpOptions, + vertexai: bool = False, ) -> Tuple[_common.StringDict, _common.StringDict]: """Ensures the SSL context is present in the HTTPX client args. @@ -1021,6 +1029,7 @@ def _ensure_httpx_ssl_ctx( Args: options: The http options to check for SSL context. + vertexai: Whether Vertex AI is enabled. Returns: A tuple of sync/async httpx client args. @@ -1037,15 +1046,24 @@ def _ensure_httpx_ssl_ctx( else None ) - if not ctx: + if ctx is None: # Initialize the SSL context for the httpx client. # Unlike requests, the httpx package does not automatically pull in the # environment variables SSL_CERT_FILE or SSL_CERT_DIR. They need to be # enabled explicitly. - ctx = ssl.create_default_context( - cafile=os.environ.get('SSL_CERT_FILE', certifi.where()), - capath=os.environ.get('SSL_CERT_DIR'), - ) + if vertexai: + get_ctx_fn = getattr(mtls, 'get_default_ssl_context', None) + if get_ctx_fn is not None: + try: + ctx = get_ctx_fn() + except Exception as e: # pylint: disable=broad-except + logger.warning('Failed to get default SSL context from google-auth: %s', e) + + if ctx is None: + ctx = ssl.create_default_context( + cafile=os.environ.get('SSL_CERT_FILE', certifi.where()), + capath=os.environ.get('SSL_CERT_DIR'), + ) def _maybe_set( args: Optional[_common.StringDict], @@ -1080,13 +1098,17 @@ def _maybe_set( ) @staticmethod - def _ensure_aiohttp_ssl_ctx(options: HttpOptions) -> _common.StringDict: + def _ensure_aiohttp_ssl_ctx( + options: HttpOptions, + vertexai: bool = False, + ) -> _common.StringDict: """Ensures the SSL context is present in the async client args. Creates a default SSL context if one is not provided. Args: options: The http options to check for SSL context. + vertexai: Whether Vertex AI is enabled. Returns: An async aiohttp ClientSession._request args. @@ -1095,11 +1117,20 @@ def _ensure_aiohttp_ssl_ctx(options: HttpOptions) -> _common.StringDict: async_args = options.async_client_args ctx = async_args.get(verify) if async_args else None - if not ctx: - ctx = ssl.create_default_context( - cafile=os.environ.get('SSL_CERT_FILE', certifi.where()), - capath=os.environ.get('SSL_CERT_DIR'), - ) + if ctx is None: + if vertexai: + get_ctx_fn = getattr(mtls, 'get_default_ssl_context', None) + if get_ctx_fn is not None: + try: + ctx = get_ctx_fn() + except Exception as e: # pylint: disable=broad-except + logger.warning('Failed to get default SSL context from google-auth: %s', e) + + if ctx is None: + ctx = ssl.create_default_context( + cafile=os.environ.get('SSL_CERT_FILE', certifi.where()), + capath=os.environ.get('SSL_CERT_DIR'), + ) def _maybe_set( args: Optional[_common.StringDict], @@ -1132,13 +1163,17 @@ def _maybe_set( return _maybe_set(async_args, ctx) @staticmethod - def _ensure_websocket_ssl_ctx(options: HttpOptions) -> _common.StringDict: + def _ensure_websocket_ssl_ctx( + options: HttpOptions, + vertexai: bool = False, + ) -> _common.StringDict: """Ensures the SSL context is present in the async client args. Creates a default SSL context if one is not provided. Args: options: The http options to check for SSL context. + vertexai: Whether Vertex AI is enabled. Returns: An async aiohttp ClientSession._request args. @@ -1148,16 +1183,25 @@ def _ensure_websocket_ssl_ctx(options: HttpOptions) -> _common.StringDict: async_args = options.async_client_args ctx = async_args.get(verify) if async_args else None - if not ctx: + if ctx is None: # Initialize the SSL context for the httpx client. # Unlike requests, the aiohttp package does not automatically pull in the # environment variables SSL_CERT_FILE or SSL_CERT_DIR. They need to be # enabled explicitly. Instead of 'verify' at client level in httpx, # aiohttp uses 'ssl' at request level. - ctx = ssl.create_default_context( - cafile=os.environ.get('SSL_CERT_FILE', certifi.where()), - capath=os.environ.get('SSL_CERT_DIR'), - ) + if vertexai: + get_ctx_fn = getattr(mtls, 'get_default_ssl_context', None) + if get_ctx_fn is not None: + try: + ctx = get_ctx_fn() + except Exception as e: # pylint: disable=broad-except + logger.warning('Failed to get default SSL context from google-auth: %s', e) + + if ctx is None: + ctx = ssl.create_default_context( + cafile=os.environ.get('SSL_CERT_FILE', certifi.where()), + capath=os.environ.get('SSL_CERT_DIR'), + ) def _maybe_set( args: Optional[_common.StringDict], @@ -1498,7 +1542,10 @@ async def _async_request_once( logger.info('Retrying due to aiohttp error: %s' % e) # Retrieve the SSL context from the session. self._async_client_session_request_args = ( - self._ensure_aiohttp_ssl_ctx(self._http_options) + self._ensure_aiohttp_ssl_ctx( + self._http_options, + vertexai=bool(self.vertexai), + ) ) # Instantiate a new session with the updated SSL context. session = await self._get_aiohttp_session() # type: ignore[assignment] @@ -1576,7 +1623,10 @@ async def _async_request_once( logger.info('Retrying due to aiohttp error: %s' % e) # Retrieve the SSL context from the session. self._async_client_session_request_args = ( - self._ensure_aiohttp_ssl_ctx(self._http_options) + self._ensure_aiohttp_ssl_ctx( + self._http_options, + vertexai=bool(self.vertexai), + ) ) # Instantiate a new session with the updated SSL context. session = await self._get_aiohttp_session() # type: ignore[assignment] diff --git a/google/genai/tests/client/test_client_initialization.py b/google/genai/tests/client/test_client_initialization.py index 9a99ddc37..d571c41ef 100644 --- a/google/genai/tests/client/test_client_initialization.py +++ b/google/genai/tests/client/test_client_initialization.py @@ -1929,3 +1929,54 @@ async def test_async_mtls_uses_refreshable_credentials(monkeypatch): assert passed_creds.valid == True mock_creds.expired = True assert passed_creds.valid == False + + +def test_ensure_httpx_ssl_ctx_mtls(): + with mock.patch.object(api_client, "mtls") as mock_mtls: + fake_ctx = mock.MagicMock(spec=ssl.SSLContext) + mock_mtls.get_default_ssl_context.return_value = fake_ctx + mock_mtls.should_use_client_cert.return_value = True + mock_mtls.has_default_client_cert_source.return_value = True + + + options = types.HttpOptions() + client_args, async_client_args = ( + api_client.BaseApiClient._ensure_httpx_ssl_ctx(options, vertexai=True) + ) + + assert client_args["verify"] == fake_ctx + assert async_client_args["verify"] == fake_ctx + mock_mtls.get_default_ssl_context.assert_called_once() + + +@requires_aiohttp +def test_ensure_aiohttp_ssl_ctx_mtls(): + with mock.patch.object(api_client, "mtls") as mock_mtls: + fake_ctx = mock.MagicMock(spec=ssl.SSLContext) + mock_mtls.get_default_ssl_context.return_value = fake_ctx + + options = types.HttpOptions() + async_client_args = ( + api_client.BaseApiClient._ensure_aiohttp_ssl_ctx(options, vertexai=True) + ) + + assert async_client_args["ssl"] == fake_ctx + mock_mtls.get_default_ssl_context.assert_called_once() + + +def test_ensure_websocket_ssl_ctx_mtls(): + with mock.patch.object(api_client, "mtls") as mock_mtls: + fake_ctx = mock.MagicMock(spec=ssl.SSLContext) + mock_mtls.get_default_ssl_context.return_value = fake_ctx + mock_mtls.should_use_client_cert.return_value = True + mock_mtls.has_default_client_cert_source.return_value = True + + options = types.HttpOptions() + async_client_args = ( + api_client.BaseApiClient._ensure_websocket_ssl_ctx( + options, vertexai=True + ) + ) + + assert async_client_args["ssl"] == fake_ctx + mock_mtls.get_default_ssl_context.assert_called_once() diff --git a/pyproject.toml b/pyproject.toml index aee0cc07d..5530ce1cc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ classifiers = [ ] dependencies = [ "anyio>=4.8.0, <5.0.0", - "google-auth[requests]>=2.48.1, <2.56.0", + "google-auth[requests]>=2.56.0, <3.0.0", "httpx>=0.28.1, <1.0.0", "pydantic>=2.12.5, <3.0.0", "requests>=2.28.1, <3.0.0", diff --git a/requirements.txt b/requirements.txt index 322d9fe27..b6c18ad3f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,7 +7,7 @@ charset-normalizer==3.4.0 coverage==7.6.9 distro==1.9.0 httpx==0.28.1 -google-auth==2.47.0 +google-auth==2.56.0 idna==3.10 iniconfig==2.0.0 packaging==24.2