diff --git a/google/genai/_gaos/google_genai.py b/google/genai/_gaos/google_genai.py index 95c9fdf7c..7dfa5c666 100644 --- a/google/genai/_gaos/google_genai.py +++ b/google/genai/_gaos/google_genai.py @@ -1021,7 +1021,8 @@ def upload( or 'v1alpha' ).lstrip('/') - handshake_res = self._internal.start_upload( + handshake_res = wrap_sdk_call( + self._internal.start_upload, environment=clean_env, path=clean_path, x_goog_upload_header_content_length=size_bytes, @@ -1179,7 +1180,8 @@ async def upload( or 'v1alpha' ).lstrip('/') - handshake_res = await self._internal.start_upload( + handshake_res = await async_wrap_sdk_call( + self._internal.start_upload, environment=clean_env, path=clean_path, x_goog_upload_header_content_length=size_bytes, diff --git a/google/genai/_gaos/lib/compat_errors.py b/google/genai/_gaos/lib/compat_errors.py index 7bd524927..c25a70db2 100644 --- a/google/genai/_gaos/lib/compat_errors.py +++ b/google/genai/_gaos/lib/compat_errors.py @@ -350,9 +350,9 @@ def after_error( response: Optional[httpx.Response], error: Optional[Exception], ) -> tuple[Optional[httpx.Response], Optional[Exception]]: - if error is None: - return response, error - return response, wrap_sdk_error(error) # type: ignore[return-value] + # Runs once per attempt inside the retry loop, which only retries raw + # httpx errors, so `wrap_sdk_call` translates after the last attempt. + return response, error def after_parse_error( self, diff --git a/google/genai/tests/interactions/test_httpx_compat.py b/google/genai/tests/interactions/test_httpx_compat.py index ea867f6a0..f6f6fdfb7 100644 --- a/google/genai/tests/interactions/test_httpx_compat.py +++ b/google/genai/tests/interactions/test_httpx_compat.py @@ -113,6 +113,67 @@ async def failing_agen(): assert "async stream network broken" in str(exc_info.value) +def test_interactions_retries_connection_errors(): + recorded_requests = [] + + def handler(request: httpx.Request) -> httpx.Response: + recorded_requests.append(request) + raise httpx.ConnectError("connection refused", request=request) + + client = client_lib.Client( + api_key="fake-key", + http_options={ + "client_args": {"transport": httpx.MockTransport(handler)}, + "retry_options": {"attempts": 2, "initial_delay": 0, "max_delay": 0}, + }, + ) + with pytest.raises(compat_errors.APIConnectionError) as exc_info: + client.interactions.create(model="gemini-1.5-flash", input="Hello") + assert len(recorded_requests) > 1 + assert isinstance(exc_info.value.__cause__, httpx.ConnectError) + + +@pytest.mark.asyncio +async def test_async_interactions_retries_connection_errors(): + recorded_requests = [] + + def handler(request: httpx.Request) -> httpx.Response: + recorded_requests.append(request) + raise httpx.ConnectError("connection refused", request=request) + + client = client_lib.Client( + api_key="fake-key", + http_options={ + "async_client_args": {"transport": httpx.MockTransport(handler)}, + "retry_options": {"attempts": 2, "initial_delay": 0, "max_delay": 0}, + }, + ) + with pytest.raises(compat_errors.APIConnectionError) as exc_info: + await client.aio.interactions.create( + model="gemini-1.5-flash", input="Hello" + ) + assert len(recorded_requests) > 1 + assert isinstance(exc_info.value.__cause__, httpx.ConnectError) + + +def test_environment_file_upload_wraps_connection_errors(): + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("connection refused", request=request) + + client = client_lib.Client( + api_key="fake-key", + http_options={ + "client_args": {"transport": httpx.MockTransport(handler)}, + "retry_options": {"attempts": 1, "initial_delay": 0, "max_delay": 0}, + }, + ) + with pytest.raises(compat_errors.APIConnectionError) as exc_info: + client.environments.files.upload( + environment="env_123", path="src/main.py", file=b"print('hello')" + ) + assert isinstance(exc_info.value.__cause__, httpx.ConnectError) + + # --- httpx2 tests (run in CI when httpx2 is installed) ---