[bugfix] Streaming: answer malformed client frames with invalid_message instead of dropping the session (#2073)

docs/design/server_contracts/streaming.md

@@ -79,6 +79,11 @@ The opening frame must be `session_init_v2`. Any other first frame is
 rejected with an `error` (code `invalid_message`) and the WebSocket is
 closed.
 
+Client frames are JSON text. A binary frame, or text that is not JSON, is
+answered with an `error` (code `invalid_message`) as well: as the opening
+frame it closes the WebSocket with 1003 `invalid_init`; later in the session
+the error is `retryable` and the session continues.
+
 ### Server → client
 
 | `type` | Carries | When emitted |

fastvideo/entrypoints/streaming/server.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import asyncio
 import contextlib
+import json
 import os
 import time
 from dataclasses import dataclass
@@ -249,7 +250,7 @@ async def _read_init_message(
 ) -> SessionInitV2 | None:
     try:
         raw = await asyncio.wait_for(
-            websocket.receive_json(),
+            _receive_client_json(websocket),
             timeout=state.sessions.session_timeout_seconds,
         )
     except asyncio.TimeoutError:
@@ -259,6 +260,9 @@ async def _read_init_message(
         return None
     except WebSocketDisconnect:
         return None
+    except _MalformedClientFrame as exc:
+        await _reject_init(websocket, session, f"opening frame failed validation: {exc}", "invalid_init")
+        return None
     try:
         parsed = parse_client_message(raw)
     except Exception as exc:
@@ -322,7 +326,7 @@ async def _run_segment_loop(
 
         try:
             raw = await asyncio.wait_for(
-                websocket.receive_json(),
+                _receive_client_json(websocket),
                 timeout=state.sessions.session_timeout_seconds,
             )
         except asyncio.TimeoutError:
@@ -332,6 +336,10 @@ async def _run_segment_loop(
             return
         except WebSocketDisconnect:
             return
+        except _MalformedClientFrame as exc:
+            session.touch()
+            await _send_error(websocket, "invalid_message", str(exc), retryable=True)
+            continue
         session.touch()
 
         try:
@@ -545,6 +553,30 @@ def _extract_state(result: Any) -> ContinuationState | None:
     return None
 
 
+class _MalformedClientFrame(ValueError):
+    """A client frame that is not a JSON text frame."""
+
+
+async def _receive_client_json(websocket: WebSocket) -> Any:
+    """Receive one client frame and decode it as JSON.
+
+    Client frames are JSON text. ``WebSocket.receive_json`` raises ``KeyError``
+    for a binary frame and ``JSONDecodeError`` for text that is not JSON, and
+    neither reached the ``invalid_message`` handling, so one bad frame closed
+    the session with no error. Raise ``_MalformedClientFrame`` for both instead.
+    """
+    message = await websocket.receive()
+    if message["type"] == "websocket.disconnect":
+        raise WebSocketDisconnect(message.get("code", 1000), message.get("reason"))
+    text = message.get("text")
+    if text is None:
+        raise _MalformedClientFrame("client frames must be JSON text, got a binary frame")
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError as exc:
+        raise _MalformedClientFrame(f"client frame is not valid JSON: {exc}") from exc
+
+
 async def _send_json(websocket: WebSocket, message: Any) -> None:
     payload = (message.model_dump(mode="json", exclude_none=True) if hasattr(message, "model_dump") else message)
     await websocket.send_json(payload)

fastvideo/tests/entrypoints/streaming/test_server.py

@@ -6,7 +6,9 @@ weights). Skips the fMP4 assertions when ``ffmpeg`` is missing.
 """
 from __future__ import annotations
 
+import contextlib
 import shutil
+import signal
 from dataclasses import dataclass
 from typing import Any
 
@@ -101,6 +103,25 @@ class TestSessionHandshake:
             err = ws.receive_json()
             assert err["type"] == "error"
 
+    def test_rejects_opening_frame_that_is_not_json(self):
+        client, _ = _build_client()
+        with _deadline(), client.websocket_connect("/v1/stream") as ws:
+            ws.send_text("{not json")
+            err = ws.receive_json()
+            assert err["type"] == "error"
+            assert err["code"] == "invalid_message"
+            assert err["retryable"] is False
+            assert ws.receive() == {"type": "websocket.close", "code": 1003, "reason": "invalid_init"}
+
+    def test_rejects_binary_opening_frame(self):
+        client, _ = _build_client()
+        with _deadline(), client.websocket_connect("/v1/stream") as ws:
+            ws.send_bytes(b"\x00\x01")
+            err = ws.receive_json()
+            assert err["type"] == "error"
+            assert err["code"] == "invalid_message"
+            assert ws.receive() == {"type": "websocket.close", "code": 1003, "reason": "invalid_init"}
+
     def test_emits_queue_and_gpu_assigned_on_valid_init(self):
         client, _ = _build_client()
         with client.websocket_connect("/v1/stream") as ws:
@@ -139,6 +160,42 @@ class TestSessionHandshake:
             assert snap["state"]["payload"]["segment_index"] == 3
 
 
+class TestMalformedFramesMidSession:
+
+    @pytest.mark.parametrize("send", [
+        lambda ws: ws.send_text("{not json"),
+        lambda ws: ws.send_bytes(b"\x00\x01"),
+    ],
+                             ids=["text_not_json", "binary"])
+    def test_malformed_frame_is_reported_and_session_stays_open(self, send):
+        client, _ = _build_client()
+        with _deadline(), client.websocket_connect("/v1/stream") as ws:
+            ws.send_json({
+                "type": "session_init_v2",
+                "preset": "ltx2_two_stage",
+                "continuation_state": {
+                    "kind": "ltx2.v1",
+                    "payload": {
+                        "schema_version": 1,
+                        "segment_index": 3
+                    },
+                },
+            })
+            for _ in range(3):
+                ws.receive_json()  # queue_status + gpu_assigned + stream_start
+
+            send(ws)
+            err = ws.receive_json()
+            assert err["type"] == "error"
+            assert err["code"] == "invalid_message"
+            assert err["retryable"] is True
+
+            ws.send_json({"type": "snapshot_state"})
+            snap = ws.receive_json()
+            assert snap["type"] == "continuation_state_snapshot"
+            assert snap["state"]["payload"]["segment_index"] == 3
+
+
 @pytest.mark.skipif(not _FFMPEG_AVAILABLE, reason="ffmpeg not installed")
 class TestSegmentFlow:
 
@@ -207,6 +264,29 @@ class TestContinuationStatePersistence:
 # ----------------------------------------------------------------------
 
 
+@contextlib.contextmanager
+def _deadline(seconds: float = 10.0):
+    """Fail instead of hanging when the server never sends the frame a test waits for.
+
+    ``TestClient`` websocket receives block with no timeout, and the server
+    sending nothing is exactly the regression these tests guard against.
+    """
+    if not hasattr(signal, "SIGALRM"):
+        yield
+        return
+
+    def _expired(signum, frame):
+        raise TimeoutError(f"no server frame within {seconds:g} s")
+
+    previous = signal.signal(signal.SIGALRM, _expired)
+    signal.setitimer(signal.ITIMER_REAL, seconds)
+    try:
+        yield
+    finally:
+        signal.setitimer(signal.ITIMER_REAL, 0)
+        signal.signal(signal.SIGALRM, previous)
+
+
 def _drain_until(ws, target_type: str) -> dict[str, Any]:
     while True:
         msg = ws.receive()