[Bugfix] Align video frame consumption reporting with prompt sampling (#7329)

docs/serving/video_stream_api.md

@@ -50,6 +50,14 @@ WebSocket /v1/video/chat/stream
 | Server -> Client | `session.done` | none | Session closed. |
 | Server -> Client | `error` | `message` | Recoverable protocol or generation error. |
 
+### Frame Consumption Reporting
+
+When buffered frames include client-supplied `frame_id` values, the server emits `video.frames.consumed` after the first engine output. Its `frame_ids`, `frames`, and `latest_pts_ms` describe the images included in that query's prompt.
+
+Queries stride-sample the buffered frames, keeping the last frame, then exclude frames already known to have failed decoding. Excluded frames are not replaced with other buffered frames. Frames whose background decoding has not finished remain eligible through the image URL path.
+
+**Bugfix compatibility note:** `video.frames.consumed` now excludes known decode failures that were already excluded from the prompt. Older versions could report those frames and their timestamps as consumed. The event name and fields are unchanged; an empty selection reports empty lists and `latest_pts_ms: null`.
+
 ### `session.config` Fields
 
 | Field | Type | Default | Description |

tests/entrypoints/openai_api/test_serving_video_stream.py

@@ -375,6 +375,200 @@ async def test_video_frames_consumed_is_emitted_after_engine_uses_frame_prompt()
     )
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "num_frames,bad_index,expected_indices",
+    [
+        pytest.param(2, 4, [0], id="bad-last-no-refill"),
+        pytest.param(1, 4, [], id="all-selected-frames-bad"),
+        pytest.param(8, 4, [0, 1, 2, 3], id="below-sampling-limit"),
+        pytest.param(2, 1, [0, 4], id="bad-frame-not-selected"),
+        pytest.param(2, 0, [4], id="bad-first-frame"),
+    ],
+)
+async def test_video_frames_consumed_excludes_bad_frame_in_query_snapshot(
+    monkeypatch, num_frames, bad_index, expected_indices
+):
+    """A query queued before bad-frame cleanup must report only its prompt images."""
+    query_queued = asyncio.Event()
+    decode_failed = asyncio.Event()
+    release_decode = asyncio.Event()
+    pong_blocked = asyncio.Event()
+    release_pong = asyncio.Event()
+    turn_done = asyncio.Event()
+    original_to_thread = asyncio.to_thread
+    bad_image = b"not-a-jpeg"
+    frames = [_make_jpeg(r=index * 40) for index in range(5)]
+    frames[bad_index] = bad_image
+    engine_prompts = []
+
+    class ObservedQueue(asyncio.Queue):
+        def put_nowait(self, item):
+            super().put_nowait(item)
+            if isinstance(item, dict):
+                if item.get("type") == "video.query":
+                    query_queued.set()
+                elif item.get("type") == "_internal.frame_decode_failed":
+                    decode_failed.set()
+
+    async def gated_to_thread(function, *args, **kwargs):
+        if function is video_stream_base._decode_frame_bytes and args[0] == bad_image:
+            await release_decode.wait()
+        return await original_to_thread(function, *args, **kwargs)
+
+    class BackpressuredWebSocket(TimedWebSocket):
+        async def send_json(self, data):
+            await super().send_json(data)
+            if data["type"] == "pong":
+                pong_blocked.set()
+                await release_pong.wait()
+            elif data["type"] == "response.text.done":
+                turn_done.set()
+
+    class OneOutputEngine:
+        async def generate(self, **kwargs):
+            engine_prompts.append(kwargs["prompt"])
+            yield _text_result("visible")
+
+    class CapturingHandler(QwenOmniStreamingVideoHandler):
+        async def _preprocess_to_engine_prompt(self, request):
+            return {"messages": request.messages}
+
+    # Construct the client queue before observing the server's message queue.
+    ws = BackpressuredWebSocket()
+    monkeypatch.setattr(asyncio, "Queue", ObservedQueue)
+    monkeypatch.setattr(asyncio, "to_thread", gated_to_thread)
+    handler = CapturingHandler(chat_service=object(), engine_client=OneOutputEngine(), idle_timeout=5.0)
+    ws.put(
+        {
+            "type": "session.config",
+            "model": "test",
+            "modalities": ["text"],
+            "num_frames": num_frames,
+            "enable_frame_filter": False,
+        }
+    )
+    for index, frame in enumerate(frames):
+        ws.put({"type": "video.frame", "data": _b64(frame), "frame_id": f"frame-{index}", "pts_ms": index * 100})
+    ws.put({"type": "ping"})
+    ws.put({"type": "video.query", "text": "describe"})
+    task = asyncio.create_task(handler.handle_session(ws))
+    try:
+        # Hold the processor at a WebSocket send. Queue the query first, then
+        # finish the real failing decode so its cleanup lands behind the query.
+        await asyncio.wait_for(pong_blocked.wait(), timeout=5.0)
+        await asyncio.wait_for(query_queued.wait(), timeout=5.0)
+        release_decode.set()
+        await asyncio.wait_for(decode_failed.wait(), timeout=5.0)
+        release_pong.set()
+        await asyncio.wait_for(turn_done.wait(), timeout=5.0)
+        ws.put({"type": "video.done"})
+        await asyncio.wait_for(task, timeout=5.0)
+    finally:
+        release_decode.set()
+        release_pong.set()
+        if not task.done():
+            task.cancel()
+        await asyncio.gather(task, return_exceptions=True)
+
+    assert len(engine_prompts) == 1
+    images = [part for part in engine_prompts[0]["messages"][-1]["content"] if part["type"].startswith("image_")]
+    # Expected positions preserve sample-then-filter: do not refill a dropped
+    # position or shift the remaining metadata when a frame fails to decode.
+    assert len(images) == len(expected_indices)
+    for image, index in zip(images, expected_indices):
+        if image["type"] == "image_pil":
+            assert image["image_pil"].tobytes() == Image.open(io.BytesIO(frames[index])).convert("RGB").tobytes()
+        else:
+            assert image["image_url"]["url"] == f"data:image/jpeg;base64,{_b64(frames[index])}"
+    consumed_events = [message for message in ws.sent if message.get("type") == "video.frames.consumed"]
+    assert len(consumed_events) == 1
+    consumed = consumed_events[0]
+    expected_ids = [f"frame-{index}" for index in expected_indices]
+    assert consumed["frame_ids"] == expected_ids
+    assert [frame["frame_id"] for frame in consumed["frames"]] == expected_ids
+    assert [frame["pts_ms"] for frame in consumed["frames"]] == [index * 100 for index in expected_indices]
+    assert consumed["latest_pts_ms"] == (expected_indices[-1] * 100 if expected_indices else None)
+    assert {"type": "error", "message": "Frame decode failed"} in ws.sent
+    assert ws.sent[-1]["type"] == "session.done"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "num_frames,frame_ids,repeat_image,expected_indices",
+    [
+        pytest.param(2, ["f0", "f1", "f2", "f3", "f4"], False, [0, 4], id="stride-and-last"),
+        pytest.param(1, ["f0", "f1", "f2"], False, [2], id="last-only"),
+        pytest.param(8, ["f0", "f1", "f2"], False, [0, 1, 2], id="below-sampling-limit"),
+        pytest.param(2, [None, "f1", None, "f3", None], False, [0, 4], id="unidentified-selected-frames"),
+        pytest.param(2, [None, None, None], False, [0, 2], id="no-frame-ids"),
+        pytest.param(2, ["same", "same", "same"], False, [0, 2], id="repeated-frame-id"),
+        pytest.param(2, ["f0", "f1", "f2"], True, [0, 2], id="repeated-image"),
+    ],
+)
+async def test_video_frame_selection_preserves_uncached_frames(
+    monkeypatch, num_frames, frame_ids, repeat_image, expected_indices
+):
+    """Prewarm misses still use image URLs, with metadata paired by buffer position."""
+    engine_prompts = []
+    pending_decode = asyncio.Event()
+
+    async def blocked_to_thread(*args, **kwargs):
+        await pending_decode.wait()
+
+    monkeypatch.setattr(asyncio, "to_thread", blocked_to_thread)
+
+    class OneOutputEngine:
+        async def generate(self, **kwargs):
+            engine_prompts.append(kwargs["prompt"])
+            yield _text_result("visible")
+
+    class CapturingHandler(QwenOmniStreamingVideoHandler):
+        async def _preprocess_to_engine_prompt(self, request):
+            return {"messages": request.messages}
+
+    frames = [_b64(_make_jpeg(r=0 if repeat_image else index * 40)) for index in range(len(frame_ids))]
+    messages = [
+        {
+            "type": "session.config",
+            "model": "test",
+            "modalities": ["text"],
+            "num_frames": num_frames,
+            "enable_frame_filter": False,
+        },
+        *[
+            {"type": "video.frame", "data": frame, "frame_id": frame_ids[index], "pts_ms": index * 100}
+            for index, frame in enumerate(frames)
+        ],
+        {"type": "video.query", "text": "describe"},
+        {"type": "video.done"},
+    ]
+    ws = MockWebSocket([json.dumps(message) for message in messages])
+    handler = CapturingHandler(chat_service=object(), engine_client=OneOutputEngine())
+    await asyncio.wait_for(handler.handle_session(ws), timeout=5.0)
+
+    assert not [message for message in ws.sent if message["type"] == "error"]
+    assert len(engine_prompts) == 1
+    content = engine_prompts[0]["messages"][-1]["content"]
+    assert content == [
+        *[
+            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{frames[index]}"}}
+            for index in expected_indices
+        ],
+        {"type": "text", "text": "describe"},
+    ]
+    consumed_events = [message for message in ws.sent if message["type"] == "video.frames.consumed"]
+    if not any(frame_ids):
+        assert consumed_events == []
+    else:
+        assert len(consumed_events) == 1
+        consumed = consumed_events[0]
+        assert consumed["frame_ids"] == [frame_ids[index] for index in expected_indices if frame_ids[index] is not None]
+        assert [frame["frame_id"] for frame in consumed["frames"]] == [frame_ids[index] for index in expected_indices]
+        assert [frame["pts_ms"] for frame in consumed["frames"]] == [index * 100 for index in expected_indices]
+        assert consumed["latest_pts_ms"] == expected_indices[-1] * 100
+
+
 @pytest.mark.asyncio
 async def test_audio_in_video_sets_mm_processor_kwargs():
     captured_requests = []

vllm_omni/entrypoints/openai/serving_video_stream.py

@@ -30,7 +30,6 @@ from collections.abc import Mapping
 from typing import Any
 
 from vllm_omni.entrypoints.openai.video_stream_base import (
-    _BAD_FRAME,
     _DEFAULT_CONFIG_TIMEOUT,
     _DEFAULT_IDLE_TIMEOUT,
     PrewarmedFrame,
@@ -62,22 +61,19 @@ class QwenOmniStreamingVideoHandler(OmniStreamingVideoHandlerBase):
         message_history: list[dict[str, Any]],
         query_text: str,
         prewarmed_frames: Mapping[str, PrewarmedFrame],
+        *,
+        frame_indices: list[int] | None = None,
     ) -> tuple[list[dict[str, Any]], dict[str, Any]]:
-        n_buf = len(frame_buffer)
-        if n_buf <= config.num_frames:
-            frames = list(frame_buffer)
-        else:
-            stride = max(1, n_buf // config.num_frames)
-            idx = [i * stride for i in range(config.num_frames - 1)] + [n_buf - 1]
-            frames = [frame_buffer[i] for i in idx]
-
         prewarmed = prewarmed_frames or {}
+        if frame_indices is None:
+            frame_indices = self._sample_frame_indices(frame_buffer, config.num_frames, prewarmed)
         user_content: list[dict] = []
-        for frame_b64 in frames:
+        for index in frame_indices:
+            frame_b64 = frame_buffer[index]
             cached = prewarmed.get(frame_b64)
-            if cached is _BAD_FRAME:
-                continue
             if cached is not None:
+                # The shared selector excludes failed decodes before this point.
+                assert isinstance(cached, tuple)
                 pil, pil_uuid = cached
                 user_content.append(
                     {

vllm_omni/entrypoints/openai/video_stream_base.py

@@ -93,8 +93,10 @@ class VideoStreamPipelineHooks(Protocol):
         message_history: list[dict[str, Any]],
         query_text: str,
         prewarmed_frames: Mapping[str, PrewarmedFrame],
+        *,
+        frame_indices: list[int] | None = None,
     ) -> tuple[list[dict[str, Any]], dict[str, Any]]:
-        """Build OpenAI-style messages and the current user message."""
+        """Build messages using the supplied frame selection, or sample if omitted."""
         ...
 
     def on_turn_complete(
@@ -179,6 +181,8 @@ class OmniStreamingVideoHandler:
         message_history: list[dict[str, Any]],
         query_text: str,
         prewarmed_frames: Mapping[str, PrewarmedFrame],
+        *,
+        frame_indices: list[int] | None = None,
     ) -> tuple[list[dict[str, Any]], dict[str, Any]]:
         raise NotImplementedError
 
@@ -639,6 +643,9 @@ class OmniStreamingVideoHandler:
             ChatCompletionRequest,
         )
 
+        # Select once from this turn's snapshot so prompt images and the
+        # consumed event refer to exactly the same frames, including decode failures.
+        frame_indices = self._sample_frame_indices(frame_buffer, config.num_frames, prewarmed_frames)
         messages, user_message = self.build_engine_prompt(
             config,
             frame_buffer,
@@ -646,6 +653,7 @@ class OmniStreamingVideoHandler:
             message_history,
             query_text,
             prewarmed_frames,
+            frame_indices=frame_indices,
         )
 
         request_kwargs: dict[str, Any] = {
@@ -678,7 +686,7 @@ class OmniStreamingVideoHandler:
             await self._send_error(websocket, f"Preprocess failed: {e}")
             return
         decoded_ready_ts_ms = _time.monotonic() * 1000
-        selected_metadata = self._sample_frame_metadata(frame_metadata or [], config.num_frames)
+        selected_metadata = [frame_metadata[index] for index in frame_indices] if frame_metadata else []
         model_selected_ts_ms = _time.monotonic() * 1000
 
         await websocket.send_json({"type": "response.start"})
@@ -1071,15 +1079,19 @@ class OmniStreamingVideoHandler:
             pass
 
     @staticmethod
-    def _sample_frame_metadata(
-        frame_metadata: list[dict[str, Any]],
+    def _sample_frame_indices(
+        frame_buffer: list[str],
         num_frames: int,
-    ) -> list[dict[str, Any]]:
-        if len(frame_metadata) <= num_frames:
-            return list(frame_metadata)
-        stride = max(1, len(frame_metadata) // num_frames)
-        indices = [index * stride for index in range(num_frames - 1)] + [len(frame_metadata) - 1]
-        return [frame_metadata[index] for index in indices]
+        prewarmed_frames: Mapping[str, PrewarmedFrame],
+    ) -> list[int]:
+        """Stride-sample with the last frame, then drop known bad frames without refilling."""
+        n_buf = len(frame_buffer)
+        if n_buf <= num_frames:
+            indices = list(range(n_buf))
+        else:
+            stride = max(1, n_buf // num_frames)
+            indices = [index * stride for index in range(num_frames - 1)] + [n_buf - 1]
+        return [index for index in indices if prewarmed_frames.get(frame_buffer[index]) is not _BAD_FRAME]
 
     @staticmethod
     async def _send_frame_ack(