fix: retry transient model failures

agent_platform/runtime/provider.py

@@ -1,11 +1,16 @@
 from __future__ import annotations
 
+import asyncio
 from typing import Any
 
 import httpx
 
 from agent_platform.config import Settings
 
+TRANSIENT_COMPLETE_STATUS_CODES = frozenset({408, 425, 429, 500, 502, 503, 504, 524})
+COMPLETE_MAX_ATTEMPTS = 2
+MAX_RETRY_DELAY_SECONDS = 5.0
+
 
 def completions_url(base_url: str) -> str:
     base = base_url.rstrip("/")
@@ -14,6 +19,17 @@ def completions_url(base_url: str) -> str:
     return f"{base}/v1/chat/completions"
 
 
+def retry_delay_seconds(response: httpx.Response | None, attempt: int) -> float:
+    if response is not None:
+        retry_after = response.headers.get("Retry-After")
+        if retry_after:
+            try:
+                return min(MAX_RETRY_DELAY_SECONDS, max(0.0, float(retry_after)))
+            except ValueError:
+                pass
+    return min(MAX_RETRY_DELAY_SECONDS, float(2 ** (attempt - 1)))
+
+
 class ModelProvider:
     def __init__(self, settings: Settings, client: httpx.AsyncClient | None = None) -> None:
         self.settings = settings
@@ -52,16 +68,31 @@ class ModelProvider:
             payload["tool_choice"] = "auto"
         if temperature is not None:
             payload["temperature"] = temperature
-        response = await self.client.post(
+        for attempt in range(1, COMPLETE_MAX_ATTEMPTS + 1):
-            completions_url(self.settings.model_api_base_url),
+            try:
-            headers=self.headers,
+                response = await self.client.post(
-            json=payload,
+                    completions_url(self.settings.model_api_base_url),
-        )
+                    headers=self.headers,
-        response.raise_for_status()
+                    json=payload,
-        data = response.json()
+                )
-        if not data.get("choices"):
+            except httpx.TransportError:
-            raise RuntimeError("Model provider returned no choices")
+                if attempt >= COMPLETE_MAX_ATTEMPTS:
-        return data
+                    raise
+                await asyncio.sleep(retry_delay_seconds(None, attempt))
+                continue
+
+            if response.status_code in TRANSIENT_COMPLETE_STATUS_CODES and attempt < COMPLETE_MAX_ATTEMPTS:
+                await response.aread()
+                await asyncio.sleep(retry_delay_seconds(response, attempt))
+                continue
+
+            response.raise_for_status()
+            data = response.json()
+            if not data.get("choices"):
+                raise RuntimeError("Model provider returned no choices")
+            return data
+
+        raise RuntimeError("Model provider retry loop exited unexpectedly")
 
     async def forward(self, payload: dict[str, Any]) -> httpx.Response:
         request = self.client.build_request(

tests/test_provider.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+
+import httpx
+import pytest
+
+from agent_platform.config import Settings
+from agent_platform.runtime.provider import ModelProvider
+
+
+@pytest.fixture(autouse=True)
+def disable_retry_sleep(monkeypatch: pytest.MonkeyPatch) -> None:
+    async def no_sleep(_: float) -> None:
+        return None
+
+    monkeypatch.setattr("agent_platform.runtime.provider.asyncio.sleep", no_sleep)
+
+
+async def _complete_with_handler(
+    settings: Settings,
+    handler: Callable[[httpx.Request], httpx.Response | Awaitable[httpx.Response]],
+) -> dict:
+    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
+        provider = ModelProvider(settings, client)
+        return await provider.complete(model="test-model", messages=[{"role": "user", "content": "hello"}])
+
+
+async def test_complete_retries_transient_gateway_response(settings: Settings) -> None:
+    calls = 0
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        if calls == 1:
+            return httpx.Response(524, request=request, headers={"Retry-After": "0"})
+        return httpx.Response(
+            200,
+            request=request,
+            json={"choices": [{"message": {"role": "assistant", "content": "ok"}}]},
+        )
+
+    result = await _complete_with_handler(settings, handler)
+
+    assert calls == 2
+    assert result["choices"][0]["message"]["content"] == "ok"
+
+
+async def test_complete_retries_transport_error(settings: Settings) -> None:
+    calls = 0
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        if calls == 1:
+            raise httpx.ConnectError("temporary disconnect", request=request)
+        return httpx.Response(
+            200,
+            request=request,
+            json={"choices": [{"message": {"role": "assistant", "content": "recovered"}}]},
+        )
+
+    result = await _complete_with_handler(settings, handler)
+
+    assert calls == 2
+    assert result["choices"][0]["message"]["content"] == "recovered"
+
+
+async def test_complete_does_not_retry_client_error(settings: Settings) -> None:
+    calls = 0
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        return httpx.Response(400, request=request, json={"error": "bad request"})
+
+    with pytest.raises(httpx.HTTPStatusError):
+        await _complete_with_handler(settings, handler)
+
+    assert calls == 1