From 58940992ae203464dfbd5e410300cfcdabee0f45 Mon Sep 17 00:00:00 2001 From: Keith Date: Thu, 5 Nov 2026 10:14:53 -0800 Subject: [PATCH] [bugfix] Streaming: answer malformed client frames with invalid_message instead of dropping the session (#2073) --- docs/design/server_contracts/streaming.md | 5 ++ fastvideo/entrypoints/streaming/server.py | 36 ++++++++- .../entrypoints/streaming/test_server.py | 80 +++++++++++++++++++ 3 files changed, 119 insertions(+), 2 deletions(-) diff --git a/docs/design/server_contracts/streaming.md b/docs/design/server_contracts/streaming.md index 3184a8659..c6e866165 100644 --- a/docs/design/server_contracts/streaming.md +++ b/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 | diff --git a/fastvideo/entrypoints/streaming/server.py b/fastvideo/entrypoints/streaming/server.py index 707462b58..fcfb94bf7 100644 --- a/fastvideo/entrypoints/streaming/server.py +++ b/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) diff --git a/fastvideo/tests/entrypoints/streaming/test_server.py b/fastvideo/tests/entrypoints/streaming/test_server.py index 15c38b186..6e9280abb 100644 --- a/fastvideo/tests/entrypoints/streaming/test_server.py +++ b/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()