[bugfix] Streaming: end sessions with the documented segment_cap, session_timeout and error frames

docs/design/server_contracts/streaming.md

@@ -92,7 +92,7 @@ closed.
 | binary frame | fMP4 fragment bytes | Subsequent media chunks; the protocol enforces that `media_init` precedes any binary frames. |
 | `media_segment_complete` | `segment_idx`, chunk count, byte count | Last media chunk for the segment. |
 | `ltx2_segment_complete` | `segment_idx`, segment summary | Segment fully shipped; ready for the next `segment_prompt_source`. |
-| `ltx2_stream_complete` | session summary | Session reached `generation_segment_cap` or client requested clean shutdown. |
+| `ltx2_stream_complete` | `reason` | Session reached `generation_segment_cap` (`segment_cap`, then close 1000), failed unexpectedly (`error`, then close 1011), or client requested clean shutdown (`stop_requested`). |
 | `session_timeout` | reason | Session hit `session_timeout_seconds`; immediately followed by close. |
 | `continuation_state_snapshot` | `kind`, `payload` | Reply to `snapshot_state`. The payload is the same shape produced by `LTX2ContinuationState.to_continuation_state(...)`. |
 | `error` | `code`, `message` | Any validation/runtime error. Non-fatal errors keep the connection open; fatal errors precede a `close`. |

fastvideo/entrypoints/streaming/server.py

@@ -38,6 +38,7 @@
     SeedPromptsUpdated,
     SegmentPromptSource,
     SessionInitV2,
+    SessionTimeout,
     SnapshotState,
     StepComplete,
     parse_client_message,
@@ -67,8 +68,12 @@
 logger = init_logger(__name__)
 
 # RFC 6455 WebSocket close codes used by the server.
+_WS_CLOSE_NORMAL = 1000
 _WS_CLOSE_UNSUPPORTED_DATA = 1003
+_WS_CLOSE_INTERNAL_ERROR = 1011
 _WS_CLOSE_TRY_AGAIN_LATER = 1013
+# Why the segment loop stopped; decides the session's last frame.
+_SessionEnd = Literal["segment_cap", "idle_timeout", "disconnected"]
 _ErrorCode = Literal[
     "session_rejected",
     "invalid_message",
@@ -236,10 +241,18 @@ async def _handle_session(
     await _send_json(websocket, _build_stream_start(session, state))
 
     try:
-        await _run_segment_loop(websocket, session, state)
-    finally:
-        with contextlib.suppress(RuntimeError):
-            await _send_json(websocket, Ltx2StreamComplete(reason="stop_requested"))
+        end = await _run_segment_loop(websocket, session, state)
+    except WebSocketDisconnect:
+        raise
+    except Exception:
+        await _send_final_frame(websocket, Ltx2StreamComplete(reason="error"), _WS_CLOSE_INTERNAL_ERROR, "error")
+        raise
+    if end == "segment_cap":
+        with contextlib.suppress(InvalidSessionTransition):
+            session.transition(SessionState.COMPLETE)
+        await _send_final_frame(websocket, Ltx2StreamComplete(reason="segment_cap"), _WS_CLOSE_NORMAL, "segment_cap")
+    elif end == "idle_timeout":
+        await _send_session_timeout(websocket, state)
 
 
 async def _read_init_message(
@@ -256,6 +269,7 @@ async def _read_init_message(
         logger.info("session %s: init timeout", session.id[:8])
         with contextlib.suppress(InvalidSessionTransition):
             session.transition(SessionState.TIMEOUT)
+        await _send_session_timeout(websocket, state)
         return None
     except WebSocketDisconnect:
         return None
@@ -313,12 +327,12 @@ async def _run_segment_loop(
     websocket: WebSocket,
     session: Session,
     state: ServerState,
-) -> None:
+) -> _SessionEnd:
     cap = state.sessions.segment_cap
     while True:
         if session.segment_cap_reached(cap):
             logger.info("session %s: segment cap (%d) reached", session.id[:8], cap)
-            return
+            return "segment_cap"
 
         try:
             raw = await asyncio.wait_for(
@@ -329,9 +343,9 @@ async def _run_segment_loop(
             logger.info("session %s: idle timeout", session.id[:8])
             with contextlib.suppress(InvalidSessionTransition):
                 session.transition(SessionState.TIMEOUT)
-            return
+            return "idle_timeout"
         except WebSocketDisconnect:
-            return
+            return "disconnected"
         session.touch()
 
         try:
@@ -563,6 +577,22 @@ async def _send_error(
     )
 
 
+async def _send_final_frame(websocket: WebSocket, message: Any, close_code: int, close_reason: str) -> None:
+    """Send the session's last frame and close; the client may already be gone."""
+    with contextlib.suppress(RuntimeError, WebSocketDisconnect):
+        await _send_json(websocket, message)
+        await websocket.close(code=close_code, reason=close_reason)
+
+
+async def _send_session_timeout(websocket: WebSocket, state: ServerState) -> None:
+    await _send_final_frame(
+        websocket,
+        SessionTimeout(timeout_seconds=state.sessions.session_timeout_seconds),
+        _WS_CLOSE_NORMAL,
+        "session_timeout",
+    )
+
+
 def _cleanup_session(session: Session, state: ServerState) -> None:
     state.sessions.close(session.id)
     state.session_store.drop(session.id)

fastvideo/tests/entrypoints/streaming/test_server.py

@@ -49,7 +49,7 @@ def generate(self, request: GenerationRequest) -> dict[str, Any]:
         }
 
 
-def _build_serve_config() -> ServeConfig:
+def _build_serve_config(session_timeout_seconds: int = 60) -> ServeConfig:
     return ServeConfig(
         generator=GeneratorConfig(model_path="/models/fake"),
         default_request=GenerationRequest(sampling=SamplingConfig(
@@ -60,15 +60,15 @@ def _build_serve_config() -> ServeConfig:
             num_inference_steps=1,
         ), ),
         streaming=StreamingConfig(
-            session_timeout_seconds=60,
+            session_timeout_seconds=session_timeout_seconds,
             generation_segment_cap=2,
         ),
     )
 
 
-def _build_client() -> tuple[TestClient, _MockGenerator]:
+def _build_client(session_timeout_seconds: int = 60) -> tuple[TestClient, _MockGenerator]:
     generator = _MockGenerator()
-    app = build_app(_build_serve_config(), generator)
+    app = build_app(_build_serve_config(session_timeout_seconds), generator)
     return TestClient(app), generator
 
 
@@ -179,6 +179,56 @@ def test_segment_generates_media_init_plus_complete(self):
             assert final["segment_idx"] == 0
 
 
+class TestSessionEnd:
+    """The session's last frame says why it ended (docs/design/server_contracts/streaming.md)."""
+
+    def test_segment_cap_ends_with_segment_cap_reason(self):
+        if not _FFMPEG_AVAILABLE:
+            pytest.skip("ffmpeg not installed")
+        client, _ = _build_client()
+        with client.websocket_connect("/v1/stream") as ws:
+            ws.send_json({"type": "session_init_v2", "preset": "ltx2_two_stage"})
+            for _ in range(3):
+                ws.receive_json()
+            for i in range(2):  # generation_segment_cap=2
+                ws.send_json({"type": "segment_prompt_source", "prompt": f"segment {i}", "num_inference_steps": 1})
+                _drain_until(ws, "ltx2_segment_complete")
+            final = ws.receive_json()
+            assert final == {"type": "ltx2_stream_complete", "reason": "segment_cap"}
+            assert ws.receive() == {"type": "websocket.close", "code": 1000, "reason": "segment_cap"}
+
+    def test_idle_session_is_told_it_timed_out(self):
+        client, _ = _build_client(session_timeout_seconds=1)
+        with client.websocket_connect("/v1/stream") as ws:
+            ws.send_json({"type": "session_init_v2", "preset": "ltx2_two_stage"})
+            for _ in range(3):
+                ws.receive_json()
+            assert ws.receive_json() == {"type": "session_timeout", "timeout_seconds": 1}
+            assert ws.receive() == {"type": "websocket.close", "code": 1000, "reason": "session_timeout"}
+
+    def test_missing_opening_frame_is_told_it_timed_out(self):
+        client, _ = _build_client(session_timeout_seconds=1)
+        with client.websocket_connect("/v1/stream") as ws:
+            assert ws.receive_json() == {"type": "session_timeout", "timeout_seconds": 1}
+            assert ws.receive() == {"type": "websocket.close", "code": 1000, "reason": "session_timeout"}
+
+    def test_unexpected_failure_ends_with_error_reason(self, monkeypatch):
+        import fastvideo.entrypoints.streaming.server as server
+
+        async def _fail(*args, **kwargs):
+            raise RuntimeError("boom")
+
+        monkeypatch.setattr(server, "_run_segment", _fail)
+        client, _ = _build_client()
+        with client.websocket_connect("/v1/stream") as ws:
+            ws.send_json({"type": "session_init_v2", "preset": "ltx2_two_stage"})
+            for _ in range(3):
+                ws.receive_json()
+            ws.send_json({"type": "segment_prompt_source", "prompt": "x", "num_inference_steps": 1})
+            assert ws.receive_json() == {"type": "ltx2_stream_complete", "reason": "error"}
+            assert ws.receive() == {"type": "websocket.close", "code": 1011, "reason": "error"}
+
+
 class TestContinuationStatePersistence:
 
     def test_snapshot_after_segment_carries_generator_state(self):