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

docs/design/server_contracts/streaming.md

@@ -41,7 +41,7 @@ through the states in `SessionState` (`fastvideo/entrypoints/streaming/session.p
                 ┌───────────┼───────────┐
                 ▼           ▼           ▼
             COMPLETE      ERROR      TIMEOUT
-        (clean leave)  (any failure)  (idle / segment_cap reached)
+        (segment_cap)  (any failure)     (idle)
 ```
 
 Terminal states (`COMPLETE`, `ERROR`, `TIMEOUT`, `REJECTED`) are sinks —
@@ -92,10 +92,10 @@ 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. |
-| `session_timeout` | reason | Session hit `session_timeout_seconds`; immediately followed by close. |
+| `ltx2_stream_complete` | `reason` | Session reached `generation_segment_cap` (`reason: "segment_cap"`); followed by close 1000 `segment_cap`. `stop_requested` and `error` are reserved: no client message requests a stop yet, and a server failure is reported with an `error` frame instead. |
+| `session_timeout` | `timeout_seconds` | Session hit `session_timeout_seconds`, idle or before its opening frame; followed by close 1000 `session_timeout`. |
 | `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`. |
+| `error` | `code`, `message`, `retryable` | Any validation/runtime error. Non-fatal errors keep the connection open; fatal errors precede a `close` (an unexpected server failure sends `code: "internal_error"`, then close 1011). |
 
 ## Continuation state
 
@@ -161,7 +161,7 @@ WS /v1/stream  ─────── connect ───────────
                        "payload": {"schema_version": 1, ...}}
 
 (close) ──────────────────────────────────────────────────►
-                                                          (session → COMPLETE)
+                                                          (connection closed)
 ```
 
 ## Backward / forward compatibility

fastvideo/entrypoints/streaming/server.py

@@ -38,6 +38,7 @@
     SeedPromptsUpdated,
     SegmentPromptSource,
     SessionInitV2,
+    SessionTimeout,
     SnapshotState,
     StepComplete,
     parse_client_message,
@@ -67,8 +68,11 @@
 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
+_SessionEnd = Literal["segment_cap", "idle_timeout", "disconnected"]
 _ErrorCode = Literal[
     "session_rejected",
     "invalid_message",
@@ -236,10 +240,25 @@ 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 as exc:
+        await _send_final_frame(
+            websocket,
+            ErrorMessage(code="internal_error",
+                         message=f"unexpected server error ({type(exc).__name__})",
+                         retryable=False),
+            _WS_CLOSE_INTERNAL_ERROR,
+            "internal_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 +275,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 +333,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 +349,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 +583,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

@@ -6,7 +6,9 @@
 """
 from __future__ import annotations
 
+import contextlib
 import shutil
+import signal
 from dataclasses import dataclass
 from typing import Any
 
@@ -49,7 +51,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 +62,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 +181,61 @@ 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 _deadline(), 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):
+                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 _deadline(), 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 _deadline(), 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_sends_a_fatal_error_frame(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 _deadline(), 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": "error",
+                "code": "internal_error",
+                "message": "unexpected server error (RuntimeError)",
+                "retryable": False,
+            }
+            assert ws.receive() == {"type": "websocket.close", "code": 1011, "reason": "internal_error"}
+
+
 class TestContinuationStatePersistence:
 
     def test_snapshot_after_segment_carries_generator_state(self):
@@ -207,6 +264,29 @@ def test_snapshot_after_segment_carries_generator_state(self):
 # ----------------------------------------------------------------------
 
 
+@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()