feat: mock timeout

src/mlpa/core/classes.py

@@ -17,6 +17,7 @@ class ChatRequest(BaseModel):
     max_completion_tokens: Optional[int] = env.MAX_COMPLETION_TOKENS
     top_p: Optional[float] = env.TOP_P
     mock_response: Optional[str] = None
+    mock_timeout: Optional[bool] = None
     tools: Optional[list] = None
     tool_choice: Optional[str | dict] = None
     # Optional OpenAI params

src/mlpa/core/completions.py

@@ -10,6 +10,7 @@
 from mlpa.core.classes import AuthorizedChatRequest, AuthorizedSearchRequest
 from mlpa.core.config import (
     ERROR_CODE_MAX_USERS_REACHED,
+    ERROR_CODE_UPSTREAM_TIMEOUT,
     LITELLM_COMPLETIONS_URL,
     LITELLM_VIRTUAL_AUTH_HEADERS,
     env,
@@ -288,6 +289,8 @@ async def _read_next_chunk(
         result = PrometheusResult.ABORT
         log.info(_client_disconnected_msg)
         raise
+    except httpx.TimeoutException as e:
+        yield raise_and_log(e, True, 502, "Upstream request timed out", log=log)
     except httpx.ReadError as e:
         if disconnect_event.is_set() or await request.is_disconnected():
             disconnect_event.set()
@@ -401,6 +404,14 @@ async def _get_completion(authorized_chat_request: AuthorizedChatRequest):
         return data
     except HTTPException:
         raise
+    except httpx.TimeoutException as e:
+        try:
+            raise_and_log(e, False, 502, "Upstream request timed out")
+        except HTTPException as exc:
+            raise HTTPException(
+                status_code=exc.status_code,
+                detail={"error": ERROR_CODE_UPSTREAM_TIMEOUT},
+            ) from e
     except Exception as e:
         raise_and_log(e, False, 502, "Failed to proxy request")
     finally:

src/mlpa/core/config.py

@@ -395,6 +395,7 @@ def __init__(self):
 
 ERROR_CODE_INVALID_MODEL_NAME: int = 8
 ERROR_CODE_INVALID_REQUEST: int = 9
+ERROR_CODE_UPSTREAM_TIMEOUT: int = 10
 
 ERROR_RESPONSES: dict[int | str, dict[str, Any]] = {
     429: {
@@ -557,6 +558,46 @@ def __init__(self):
             }
         },
     },
+    502: {
+        "description": "Bad Gateway - upstream proxy request failed or timed out",
+        "content": {
+            "application/json": {
+                "schema": {
+                    "type": "object",
+                    "properties": {
+                        "detail": {
+                            "type": "object",
+                            "properties": {
+                                "error": {
+                                    "anyOf": [
+                                        {"type": "integer"},
+                                        {"type": "string"},
+                                    ],
+                                    "description": (
+                                        "Error code: 10 upstream timeout. "
+                                        "Unclassified proxy failures may return a string."
+                                    ),
+                                }
+                            },
+                            "required": ["error"],
+                        }
+                    },
+                    "required": ["detail"],
+                },
+                "examples": {
+                    "upstream_timeout": {
+                        "summary": "Upstream timeout",
+                        "value": {"detail": {"error": ERROR_CODE_UPSTREAM_TIMEOUT}},
+                        "description": "MLPA timed out waiting for LiteLLM or the upstream model provider.",
+                    },
+                    "proxy_failure": {
+                        "summary": "Unclassified proxy failure",
+                        "value": {"detail": {"error": "Failed to proxy request"}},
+                    },
+                },
+            }
+        },
+    },
 }
 
 # Reusable response docs for admin endpoints, which carry their own auth header

src/mlpa/core/search.py

@@ -5,6 +5,7 @@
 
 from mlpa.core.classes import AuthorizedSearchRequest
 from mlpa.core.config import (
+    ERROR_CODE_UPSTREAM_TIMEOUT,
     LITELLM_SEARCH_URL,
     LITELLM_VIRTUAL_AUTH_HEADERS,
 )
@@ -67,6 +68,14 @@ async def _get_search(authorized_search_request: AuthorizedSearchRequest):
         return data
     except HTTPException:
         raise
+    except httpx.TimeoutException as e:
+        try:
+            raise_and_log(e, False, 502, "Upstream request timed out")
+        except HTTPException as exc:
+            raise HTTPException(
+                status_code=exc.status_code,
+                detail={"error": ERROR_CODE_UPSTREAM_TIMEOUT},
+            ) from e
     except Exception as e:
         raise_and_log(e, False, 502, "Failed to proxy request")
     finally:

src/tests/unit/test_completions.py

@@ -21,6 +21,7 @@
     ERROR_CODE_RATE_LIMIT_EXCEEDED,
     ERROR_CODE_REQUEST_TOO_LARGE,
     ERROR_CODE_UPSTREAM_RATE_LIMIT_EXCEEDED,
+    ERROR_CODE_UPSTREAM_TIMEOUT,
     LITELLM_COMPLETIONS_URL,
     LITELLM_HEADER_ATTEMPTED_FALLBACKS,
     LITELLM_HEADER_ATTEMPTED_RETRIES,
@@ -346,20 +347,40 @@ async def test_get_completion_http_error(mocker, metrics_spy):
 
 async def test_get_completion_network_error(mocker, metrics_spy):
     mock_client = AsyncMock()
-    mock_client.post.side_effect = httpx.TimeoutException("Connection timed out")
+    mock_client.post.side_effect = httpx.TimeoutException("")
     mocker.patch("mlpa.core.completions.get_http_client", return_value=mock_client)
     mocker.patch.object(env, "MLPA_DEBUG", True)
 
     with pytest.raises(HTTPException) as exc_info:
         await get_completion(SAMPLE_REQUEST)
 
     assert exc_info.value.status_code == 502
-    assert exc_info.value.detail["error"] == "Connection timed out"
+    assert exc_info.value.detail == {"error": ERROR_CODE_UPSTREAM_TIMEOUT}
 
     metrics_spy.assert_only({"chat_completion_latency", "chat_availability"})
     assert _latency_count(metrics_spy, PrometheusResult.ERROR) == 1
 
 
+async def test_stream_completion_timeout_uses_standard_proxy_error(
+    mocker, mock_request, metrics_spy
+):
+    class TimeoutClient:
+        def stream(self, *args, **kwargs):
+            raise httpx.ReadTimeout("")
+
+    mocker.patch("mlpa.core.completions.get_http_client", return_value=TimeoutClient())
+
+    received_chunks = [
+        chunk async for chunk in stream_completion(SAMPLE_REQUEST, mock_request)
+    ]
+
+    assert received_chunks == [
+        b'data: {"code": 502, "error": "Upstream service returned an error"}\n\n'
+    ]
+    metrics_spy.assert_only({"chat_completion_latency", "chat_availability"})
+    assert _latency_count(metrics_spy, PrometheusResult.ERROR) == 1
+
+
 async def test_stream_completion_success(
     httpx_mock: HTTPXMock, mock_request, metrics_spy
 ):

src/tests/unit/test_search.py

@@ -9,6 +9,7 @@
 from mlpa.core.config import (
     ERROR_CODE_BUDGET_LIMIT_EXCEEDED,
     ERROR_CODE_REQUEST_TOO_LARGE,
+    ERROR_CODE_UPSTREAM_TIMEOUT,
 )
 from mlpa.core.metrics import SEARCH_MODEL
 from mlpa.core.prometheus_metrics import PrometheusRejectionReason, PrometheusResult
@@ -79,6 +80,27 @@ async def test_get_search_sanitizes_response_surrogates(mocker):
     _httpx_encode_json(data)  # must not raise
 
 
+async def test_get_search_timeout_returns_custom_error_code(mocker, metrics_spy):
+    mock_client = AsyncMock()
+    mock_client.post.side_effect = httpx.ReadTimeout("")
+    mocker.patch("mlpa.core.search.get_http_client", return_value=mock_client)
+
+    req = AuthorizedSearchRequest(
+        user="test-user:search",
+        service_type="search",
+        query="weather in tokyo",
+        max_results=5,
+    )
+
+    with pytest.raises(HTTPException) as exc_info:
+        await get_search(req)
+
+    assert exc_info.value.status_code == 502
+    assert exc_info.value.detail == {"error": ERROR_CODE_UPSTREAM_TIMEOUT}
+    metrics_spy.assert_only({"search_latency"})
+    assert _search_latency_count(metrics_spy, PrometheusResult.ERROR) == 1
+
+
 def _search_rejection_count(
     spy,
     reason: PrometheusRejectionReason,