Merge pull request #2459 from alexxa/rspeed-3209-vertex-429-workaround

RSPEED-3209: Detect Vertex AI RESOURCE_EXHAUSTED wrapped as 500 by llama-stack

Author: tisnik

src/utils/agents/error_handler.py

@@ -22,6 +22,7 @@
 from utils.query import (
     handle_known_apistatus_errors,
     is_context_length_error,
+    is_resource_exhausted_error,
 )
 
 type AgentInferenceError = (
@@ -90,6 +91,13 @@ def map_pydantic_agent_run_error(  # pylint: disable=too-many-return-statements
             return PromptTooLongResponse(model=model_id)
         case ModelHTTPError(status_code=429):
             return QuotaExceededResponse.model(model_id)
+        case ModelHTTPError() as http_exc if is_resource_exhausted_error(str(http_exc)):
+            logger.warning(
+                "Detected RESOURCE_EXHAUSTED in ModelHTTPError with status %d; "
+                "treating as 429 (llama-stack wraps Vertex AI 429 as 500)",
+                http_exc.status_code,
+            )
+            return QuotaExceededResponse.model(model_id)
         case ModelHTTPError():
             return InternalServerErrorResponse.generic()
         case ModelAPIError() as api_exc:

src/utils/query.py

@@ -539,6 +539,23 @@ def normalize_vertex_ai_model_id(model_id: str) -> str:
     return model_id
 
 
+def is_resource_exhausted_error(error_message: str) -> bool:
+    """Detect Vertex AI RESOURCE_EXHAUSTED errors wrapped as 500 by llama-stack.
+
+    llama-stack's remote::vertexai provider translates Vertex AI's 429
+    RESOURCE_EXHAUSTED into a generic 500 InternalServerError, losing the
+    original status code.  The original gRPC status name is preserved in
+    the error message, so we match on that.
+
+    Args:
+        error_message: The error message to inspect.
+
+    Returns:
+        True if the message indicates a wrapped RESOURCE_EXHAUSTED error.
+    """
+    return "resource_exhausted" in error_message.lower()
+
+
 def handle_known_apistatus_errors(
     error: LLSApiStatusError | OpenAIAPIStatusError, model_id: str
 ) -> AbstractErrorResponse:
@@ -556,4 +573,11 @@ def handle_known_apistatus_errors(
         return PromptTooLongResponse(model=model_id)
     if error.status_code == 429:
         return QuotaExceededResponse.model(model_id)
+    if is_resource_exhausted_error(error_message):
+        logger.warning(
+            "Detected RESOURCE_EXHAUSTED in error message with status %d; "
+            "treating as 429 (llama-stack wraps Vertex AI 429 as 500)",
+            error.status_code,
+        )
+        return QuotaExceededResponse.model(model_id)
     return InternalServerErrorResponse.generic()

tests/unit/utils/agents/test_error_handler.py

@@ -0,0 +1,33 @@
+"""Tests for agent inference error mapping."""
+
+from pydantic_ai.exceptions import ModelHTTPError
+
+from models.api.responses.error import (
+    InternalServerErrorResponse,
+    QuotaExceededResponse,
+)
+from utils.agents.error_handler import map_pydantic_agent_run_error
+
+
+class TestMapPydanticAgentRunError:
+    """Tests for map_pydantic_agent_run_error with RESOURCE_EXHAUSTED workaround."""
+
+    def test_vertex_429_wrapped_as_500_model_http_error(self) -> None:
+        """Test that ModelHTTPError 500 with RESOURCE_EXHAUSTED is treated as 429."""
+        exc = ModelHTTPError(
+            status_code=500,
+            model_name="vertexai/gemini-2.5-flash",
+            body="RESOURCE_EXHAUSTED: Quota exceeded for model",
+        )
+        result = map_pydantic_agent_run_error(exc, "vertexai/gemini-2.5-flash")
+        assert isinstance(result, QuotaExceededResponse)
+
+    def test_generic_500_model_http_error(self) -> None:
+        """Test that a generic 500 without RESOURCE_EXHAUSTED stays as 500."""
+        exc = ModelHTTPError(
+            status_code=500,
+            model_name="vertexai/gemini-2.5-flash",
+            body="Internal server error",
+        )
+        result = map_pydantic_agent_run_error(exc, "vertexai/gemini-2.5-flash")
+        assert isinstance(result, InternalServerErrorResponse)

tests/unit/utils/test_query.py

@@ -358,6 +358,21 @@ def test_quota_exceeded(self) -> None:
         detail = result.model_dump()["detail"]
         assert "quota" in detail["response"].lower()
 
+    def test_vertex_429_wrapped_as_500(self) -> None:
+        """Test that Vertex AI RESOURCE_EXHAUSTED wrapped as 500 is treated as 429."""
+        error = type(
+            "APIStatusError",
+            (),
+            {
+                "status_code": 500,
+                "message": "RESOURCE_EXHAUSTED: Quota exceeded for model",
+            },
+        )()
+        result = handle_known_apistatus_errors(error, "model1")
+        assert isinstance(result, QuotaExceededResponse)
+        detail = result.model_dump()["detail"]
+        assert "quota" in detail["response"].lower()
+
     def test_generic_error(self) -> None:
         """Test handling generic error."""
         error = type(