Retry bounded remote transport failures including response read errors

great_ai/remote/call_remote_great_ai_async.py

@@ -1,3 +1,4 @@
+import asyncio
 from typing import Any, Mapping, Optional, Type, TypeVar
 
 import httpx
@@ -19,7 +20,8 @@ async def call_remote_great_ai_async(
     """Communicate with a GreatAI object through an HTTP request.
 
     Send a POST request using [httpx](https://www.python-httpx.org/) to implement a
-    remote call. Error-handling and retries are provided by httpx.
+    remote call. Transport failures are retried up to `retry_count` times. A retry
+    may repeat a prediction when its previous response was lost.
 
     The return value is inflated into a Trace. If `model_class` is specified, the
     original output is deserialised.
@@ -40,13 +42,22 @@ async def call_remote_great_ai_async(
     if not base_uri.endswith("/predict"):
         base_uri = f"{base_uri}/predict"
 
-    transport = httpx.AsyncHTTPTransport(retries=retry_count)
+    if retry_count < 0:
+        raise ValueError("retry_count must be non-negative")
+    transport = httpx.AsyncHTTPTransport(retries=0)
 
     try:
         async with httpx.AsyncClient(
             transport=transport, timeout=timeout_in_seconds
         ) as client:
-            response = await client.post(base_uri, json=data)
+            for attempt in range(retry_count + 1):
+                try:
+                    response = await client.post(base_uri, json=data)
+                    break
+                except httpx.TransportError:
+                    if attempt == retry_count:
+                        raise
+                    await asyncio.sleep(min(0.1 * 2**attempt, 1))
             try:
                 response.raise_for_status()
             except Exception:
@@ -54,7 +65,7 @@ async def call_remote_great_ai_async(
                     f"Unexpected status code, reason: {response.text}"
                 )
     except Exception as e:
-        raise RemoteCallError from e
+        raise RemoteCallError(str(e)) from e
 
     try:
         trace = response.json()

tests/test_remote_regressions.py

@@ -0,0 +1,58 @@
+from typing import Any
+from unittest.mock import patch
+
+import httpx
+import pytest
+
+from great_ai import RemoteCallError, call_remote_great_ai_async
+
+
+def response_trace() -> Any:
+    return dict(
+        trace_id="remote",
+        created="2026-01-01",
+        original_execution_time_ms=0,
+        models=[],
+        tags=[],
+        logged_values={},
+        output=7,
+    )
+
+
+@pytest.mark.asyncio
+async def test_remote_call_retries_response_read_failures() -> None:
+    calls = 0
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        if calls == 1:
+            raise httpx.ReadError("response disconnected")
+        return httpx.Response(200, json=response_trace())
+
+    with patch(
+        "great_ai.remote.call_remote_great_ai_async.httpx.AsyncHTTPTransport",
+        return_value=httpx.MockTransport(handler),
+    ):
+        assert (
+            await call_remote_great_ai_async("http://test", {}, retry_count=2)
+        ).output == 7
+    assert calls == 2
+
+
+@pytest.mark.asyncio
+async def test_remote_retry_limit_is_respected() -> None:
+    calls = 0
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        raise httpx.ReadError("response disconnected")
+
+    with patch(
+        "great_ai.remote.call_remote_great_ai_async.httpx.AsyncHTTPTransport",
+        return_value=httpx.MockTransport(handler),
+    ):
+        with pytest.raises(RemoteCallError, match="response disconnected"):
+            await call_remote_great_ai_async("http://test", {}, retry_count=1)
+    assert calls == 2