fix: bound the async image download to the same 5s as sync

src/xai_sdk/aio/image.py

@@ -6,6 +6,7 @@
 from ..client import USER_AGENT
 from ..files import StorageOptions
 from ..image import (
+    IMAGE_DOWNLOAD_TIMEOUT_SECONDS,
     BaseClient,
     BaseImageResponse,
     ImageAspectRatio,
@@ -363,6 +364,7 @@ async def image(self) -> bytes:
                 "GET",
                 self.url,
                 headers={"User-Agent": USER_AGENT},
+                timeout=aiohttp.ClientTimeout(total=IMAGE_DOWNLOAD_TIMEOUT_SECONDS),
             ) as session:
                 session.raise_for_status()
                 return await session.read()

src/xai_sdk/image.py

@@ -11,6 +11,14 @@
 from .telemetry import should_disable_sensitive_attributes
 from .types import ImageAspectRatio, ImageFormat, ImageGenerationModel, ImageQuality, ImageResolution
 
+"""Seconds to wait when downloading an image returned as a URL.
+
+Shared by the sync and async `ImageResponse.image` accessors so the two
+transports cannot drift apart: `requests` has no default timeout and
+`aiohttp.request` defaults to 300s total, so each has to be told explicitly.
+"""
+IMAGE_DOWNLOAD_TIMEOUT_SECONDS = 5
+
 _IMAGE_ASPECT_RATIO_MAP: dict[ImageAspectRatio, image_pb2.ImageAspectRatio] = {
     "1:1": image_pb2.ImageAspectRatio.IMG_ASPECT_RATIO_1_1,
     "3:4": image_pb2.ImageAspectRatio.IMG_ASPECT_RATIO_3_4,

src/xai_sdk/sync/image.py

@@ -6,6 +6,7 @@
 from ..client import USER_AGENT
 from ..files import StorageOptions
 from ..image import (
+    IMAGE_DOWNLOAD_TIMEOUT_SECONDS,
     BaseClient,
     BaseImageResponse,
     ImageAspectRatio,
@@ -362,7 +363,7 @@ def image(self) -> bytes:
             response = requests.get(
                 self.url,
                 headers={"User-Agent": USER_AGENT},
-                timeout=5,  # 5 seconds
+                timeout=IMAGE_DOWNLOAD_TIMEOUT_SECONDS,
             )
             response.raise_for_status()
             return response.content

tests/aio/image_test.py

@@ -1,6 +1,7 @@
 import datetime
 from unittest import mock
 
+import aiohttp
 import pytest
 import pytest_asyncio
 from google.protobuf import timestamp_pb2
@@ -9,6 +10,7 @@
 from xai_sdk import AsyncClient
 from xai_sdk.cost import USD_PER_TICK
 from xai_sdk.image import (
+    IMAGE_DOWNLOAD_TIMEOUT_SECONDS,
     BaseImageResponse,
     ImageFormat,
     _make_generate_request,
@@ -48,6 +50,20 @@ async def test_url(client: AsyncClient, image_asset: bytes):
     assert image_asset == await response.image
 
 
+@pytest.mark.asyncio(loop_scope="session")
+async def test_url_download_uses_the_shared_timeout(client: AsyncClient, image_asset: bytes):
+    """The async download must use the same bound as the sync one. `aiohttp.request`
+    defaults to 300s total, so without an explicit timeout the two transports differ
+    by 60x for the same call."""
+    real_request = aiohttp.request
+    with mock.patch("xai_sdk.aio.image.aiohttp.request", wraps=real_request) as request:
+        response = await client.image.sample(prompt="foo", model="grok-2-image", image_format="url")
+        assert image_asset == await response.image
+
+    timeout = request.call_args.kwargs["timeout"]
+    assert timeout.total == IMAGE_DOWNLOAD_TIMEOUT_SECONDS
+
+
 @pytest.mark.asyncio(loop_scope="session")
 async def test_batch(client: AsyncClient, image_asset: bytes):
     responses = await client.image.sample_batch(prompt="foo", model="grok-2-image", n=2, image_format="base64")

tests/sync/image_test.py

@@ -2,12 +2,14 @@
 from unittest import mock
 
 import pytest
+import requests
 from google.protobuf import timestamp_pb2
 from opentelemetry.trace import SpanKind
 
 from xai_sdk import Client
 from xai_sdk.cost import USD_PER_TICK
 from xai_sdk.image import (
+    IMAGE_DOWNLOAD_TIMEOUT_SECONDS,
     BaseImageResponse,
     ImageFormat,
     _make_generate_request,
@@ -45,6 +47,16 @@ def test_url(client: Client, image_asset: bytes):
     assert image_asset == response.image
 
 
+def test_url_download_uses_the_shared_timeout(client: Client, image_asset: bytes):
+    """The URL download must be bounded. `requests` has no default timeout, so
+    omitting it would hang the caller indefinitely on a stalled connection."""
+    with mock.patch("xai_sdk.sync.image.requests.get", wraps=requests.get) as get:
+        response = client.image.sample(prompt="foo", model="grok-2-image", image_format="url")
+        assert image_asset == response.image
+
+    assert get.call_args.kwargs["timeout"] == IMAGE_DOWNLOAD_TIMEOUT_SECONDS
+
+
 def test_batch(client: Client, image_asset: bytes):
     responses = client.image.sample_batch(prompt="foo", model="grok-2-image", n=2, image_format="base64")