Handle ValueError safely on websockets to avoid ASGI crash

Author: Dieg0Code

src/api/error_handling.py

@@ -4,7 +4,7 @@
 from http import HTTPStatus
 from uuid import uuid4
 
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException, Request, WebSocket
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse
 from sqlalchemy.exc import OperationalError
@@ -14,7 +14,7 @@
 logger = logging.getLogger("api.errors")
 
 
-def _resolve_request_id(request: Request) -> str:
+def _resolve_request_id(request: Request | WebSocket) -> str:
     request_id = getattr(request.state, "request_id", None)
     if isinstance(request_id, str) and request_id:
         return request_id
@@ -128,17 +128,23 @@ async def validation_exception_handler(
         return JSONResponse(status_code=422, content=body)
 
     @app.exception_handler(ValueError)
-    async def value_error_handler(request: Request, exc: ValueError) -> JSONResponse:
+    async def value_error_handler(request: Request | WebSocket, exc: ValueError) -> JSONResponse | None:
         request_id = _resolve_request_id(request)
+        path = request.url.path if request.url is not None else "<unknown>"
+        method = request.method if isinstance(request, Request) else "WEBSOCKET"
         logger.warning(
             "value_error",
             extra={
                 "request_id": request_id,
-                "method": request.method,
-                "path": request.url.path,
+                "method": method,
+                "path": path,
                 "error_detail": str(exc),
             },
         )
+        if isinstance(request, WebSocket):
+            reason = str(exc) or "Invalid value"
+            await request.close(code=1008, reason=reason[:120])
+            return None
         body = _build_error_body(
             status_code=422,
             request_id=request_id,

tests/test_api_error_handling.py

@@ -4,10 +4,11 @@
 import unittest
 from pathlib import Path
 
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, WebSocket
 from fastapi.testclient import TestClient
 from pydantic import BaseModel
 from sqlalchemy.exc import DBAPIError, OperationalError
+from starlette.websockets import WebSocketDisconnect
 
 sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
 
@@ -46,6 +47,11 @@ def crash() -> None:
     def bad_value() -> None:
         raise ValueError("legacy game row has invalid enum")
 
+    @app.websocket("/ws-bad-value")
+    async def ws_bad_value(websocket: WebSocket) -> None:
+        await websocket.accept()
+        raise ValueError("invalid ws payload")
+
     return app
 
 
@@ -123,6 +129,12 @@ def test_value_error_maps_to_422_standard_shape(self) -> None:
         self.assertEqual(payload["detail"], "legacy game row has invalid enum")
         self.assertIn("request_id", payload)
 
+    def test_value_error_on_websocket_closes_with_policy_violation(self) -> None:
+        client = TestClient(build_test_app(), raise_server_exceptions=False)
+        with self.assertRaises(WebSocketDisconnect) as ctx, client.websocket_connect("/ws-bad-value") as websocket:
+            websocket.receive_json()
+        self.assertEqual(ctx.exception.code, 1008)
+
 
 if __name__ == "__main__":
     unittest.main()