[Bugfix] Preserve prompt token usage details (#5181)

Signed-off-by: Ian Eaves <ian.k.eaves@gmail.com>

.buildkite/cuda/test-merge.yml

@@ -119,6 +119,8 @@ steps:
         source_file_dependencies:
           - tests/e2e/offline_inference/test_qwen3_omni.py
           - tests/e2e/online_serving/test_qwen3_omni.py
+          - tests/helpers/runtime.py
+          - vllm_omni/entrypoints/openai/serving_chat.py
           - vllm_omni/model_executor/models/qwen3_omni/
           - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
           - vllm_omni/model_executor/models/common/snake_activation.py

.buildkite/cuda/test-ready.yml

@@ -168,6 +168,8 @@ steps:
       - label: "Omni · Qwen3-Omni Test"
         source_file_dependencies:
           - tests/e2e/online_serving/test_qwen3_omni.py
+          - tests/helpers/runtime.py
+          - vllm_omni/entrypoints/openai/serving_chat.py
           - vllm_omni/model_executor/models/qwen3_omni/
           - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
           - vllm_omni/model_executor/models/common/snake_activation.py

tests/e2e/online_serving/test_qwen3_omni.py

@@ -149,10 +149,19 @@ def _run_prefix_cache_check(openai_client, request_config: dict):
     2. The number of cached tokens is divisible by the block size used in
     test_params, because currently upstream vLLM does not cache partial
     blocks.
+    3. The first request also reports prompt token details; details must be
+    present whenever --enable-prompt-tokens-details is on, even when zero
+    tokens are cached. (An exact cached_tokens == 0 check would be
+    order-dependent here: the module-scoped server may have already cached
+    the shared system-prompt blocks from earlier tests.)
+
+    Returns the (first, second) responses for additional per-test assertions.
     """
-    openai_client.send_omni_request(request_config, request_num=1)[0]
+    first_response = openai_client.send_omni_request(request_config, request_num=1)[0]
     cached_response = openai_client.send_omni_request(request_config, request_num=1)[0]
 
+    assert first_response.cached_tokens is not None
+
     # Ensure that we have a prefix cache hit on the second request and that only the last
     # partial block is uncached (since currently we don't cache partial blocks).
     num_cached_tokens = cached_response.cached_tokens
@@ -163,6 +172,8 @@ def _run_prefix_cache_check(openai_client, request_config: dict):
     assert num_cached_tokens % BLOCK_SIZE == 0
     assert (num_cached_tokens + num_uncached_tokens) == num_prompt_tokens
 
+    return first_response, cached_response
+
 
 @pytest.mark.advanced_model
 @pytest.mark.core_model
@@ -195,7 +206,10 @@ def test_thinker_prefix_caching_text_output(omni_server, openai_client) -> None:
         "stream": False,
         "modalities": ["text"],
     }
-    _run_prefix_cache_check(openai_client, request_config)
+    first_response, _ = _run_prefix_cache_check(openai_client, request_config)
+
+    assert first_response.multimodal_tokens is not None
+    assert first_response.multimodal_tokens.get("image", 0) > 0
 
 
 @pytest.mark.advanced_model

tests/entrypoints/openai_api/test_serving_chat_metrics.py

@@ -2,9 +2,18 @@
 """Unit tests for OmniChatCompletionResponse/StreamResponse metrics field."""
 
 from types import SimpleNamespace
+from unittest.mock import MagicMock
 
 import pytest
 
+from tests.helpers.serving_chat import (
+    build_serving_chat,
+    collect_stream,
+    make_request,
+    make_text_omni_output,
+    parse_sse_chunks,
+)
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
@@ -54,6 +63,65 @@ def test_omni_chat_completion_stream_response_metrics():
     assert response.metrics == {"stage_latency": 0.5}
 
 
+@pytest.mark.asyncio
+async def test_stream_prompt_token_details_include_zero_cached_and_multimodal_tokens():
+    serving_chat = build_serving_chat()
+    serving_chat.enable_prompt_tokens_details = True
+    request = make_request(modalities=["text"], include_usage=True)
+
+    async def result_generator():
+        yield make_text_omni_output(
+            token_ids=[10],
+            finish_reason="stop",
+            num_cached_tokens=0,
+        )
+
+    raw_lines = await collect_stream(
+        serving_chat.chat_completion_stream_generator(
+            request=request,
+            result_generator=result_generator(),
+            request_id="test-req",
+            model_name="test-model",
+            conversation=[],
+            tokenizer=MagicMock(),
+            request_metadata=MagicMock(),
+            mm_token_counts={"image": 7},
+        )
+    )
+
+    chunks = parse_sse_chunks(raw_lines)
+    usage = next(chunk["usage"] for chunk in chunks if chunk.get("usage"))
+
+    assert usage["prompt_tokens_details"] == {
+        "cached_tokens": 0,
+        "multimodal_tokens": {"image": 7},
+    }
+
+
+def test_non_stream_prompt_token_details_include_zero_cached_and_multimodal_tokens():
+    serving_chat = build_serving_chat()
+    serving_chat.enable_prompt_tokens_details = True
+    request = make_request(modalities=["text"], stream=False)
+    omni_output = make_text_omni_output(
+        token_ids=[10],
+        finish_reason="stop",
+        num_cached_tokens=0,
+    )
+
+    _, usage, _, _, _ = serving_chat._create_text_choice(
+        request=request,
+        omni_outputs=omni_output,
+        tokenizer=MagicMock(),
+        conversation=[],
+        role="assistant",
+        mm_token_counts={"audio": 5},
+    )
+
+    assert usage.prompt_tokens_details is not None
+    assert usage.prompt_tokens_details.cached_tokens == 0
+    assert usage.prompt_tokens_details.multimodal_tokens == {"audio": 5}
+
+
 def test_create_image_choice_exposes_diffusion_metrics():
     """Ensure image chat content exposes profiler metrics for clients."""
     from PIL import Image

tests/entrypoints/openai_api/test_serving_chat_sampling_params.py

@@ -14,6 +14,8 @@ from pytest_mock import MockerFixture
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
 from vllm.sampling_params import SamplingParams
 
+from tests.helpers.serving_chat import build_serving_chat
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
@@ -77,18 +79,8 @@ def mock_engine_client(
     return engine_client
 
 
-@pytest.fixture
-def serving_chat(mock_engine_client):
-    """Create OmniOpenAIServingChat instance with mocked dependencies."""
-    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
-
-    # Create instance without calling __init__
-    instance = object.__new__(OmniOpenAIServingChat)
-    instance.engine_client = mock_engine_client
-    return instance
-
-
-def test_serving_boundary_normalizes_declared_root_and_nested_extras(serving_chat):
+def test_serving_boundary_normalizes_declared_root_and_nested_extras(mock_engine_client):
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     serving_chat._diffusion_extra_body_params = frozenset({"cfg_text_scale", "negative_prompt"})
     request = ChatCompletionRequest(
         model="test",
@@ -118,7 +110,8 @@ def test_serving_boundary_normalizes_declared_root_and_nested_extras(serving_cha
     assert diffusion_request_args["modalities"] == ["image"]
 
 
-def test_unknown_root_extra_does_not_claim_canonical_extra(serving_chat):
+def test_unknown_root_extra_does_not_claim_canonical_extra(mock_engine_client):
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     serving_chat._diffusion_extra_body_params = frozenset()
     request = ChatCompletionRequest(
         model="test",
@@ -135,7 +128,8 @@ def test_unknown_root_extra_does_not_claim_canonical_extra(serving_chat):
     assert "pipeline_option" not in diffusion_request_args
 
 
-def test_serving_boundary_rejects_invalid_quality(serving_chat):
+def test_serving_boundary_rejects_invalid_quality(mock_engine_client):
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     serving_chat._diffusion_extra_body_params = frozenset()
     request = ChatCompletionRequest(
         model="test",
@@ -148,7 +142,8 @@ def test_serving_boundary_rejects_invalid_quality(serving_chat):
         serving_chat._normalize_diffusion_request_args(request)
 
 
-def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(serving_chat):
+def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(mock_engine_client):
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     serving_chat._diffusion_extra_body_params = frozenset()
     request = ChatCompletionRequest(
         model="test",
@@ -170,7 +165,8 @@ def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(serving_chat):
     ],
     ids=["registered-model-extra", "unregistered-common-alias"],
 )
-def test_cfg_scale_owner_conflicts_are_rejected(serving_chat, registered, request_kwargs):
+def test_cfg_scale_owner_conflicts_are_rejected(mock_engine_client, registered, request_kwargs):
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     serving_chat._diffusion_extra_body_params = frozenset(registered)
     request = ChatCompletionRequest(model="test", messages=[], **request_kwargs)
 
@@ -180,10 +176,11 @@ def test_cfg_scale_owner_conflicts_are_rejected(serving_chat, registered, reques
 
 @pytest.mark.parametrize("diffusion_mode", [True, False], ids=["pure", "mixed"])
 def test_duplicate_extras_return_the_same_bad_request_before_dispatch(
-    serving_chat,
+    mock_engine_client,
     mocker: MockerFixture,
     diffusion_mode: bool,
 ):
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     serving_chat._diffusion_mode = diffusion_mode
     serving_chat._diffusion_extra_body_params = frozenset({"sample_solver"})
     serving_chat.engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
@@ -211,9 +208,10 @@ def test_duplicate_extras_return_the_same_bad_request_before_dispatch(
     pure_dispatch.assert_not_awaited()
 
 
-def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(serving_chat, mocker: MockerFixture):
+def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(mock_engine_client, mocker: MockerFixture):
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     captured: dict[str, object] = {}
 
     async def generate(**kwargs):
@@ -261,19 +259,20 @@ def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(serving_chat,
     assert response.error.message == "No output generated from AsyncOmni"
 
 
-def test_mixed_consumer_keeps_root_common_args_with_nested_extras(serving_chat, mocker: MockerFixture):
+def test_mixed_consumer_keeps_root_common_args_with_nested_extras(mock_engine_client, mocker: MockerFixture):
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
-    serving_chat.engine_client.stage_configs = [
+    mock_engine_client.stage_configs = [
         SimpleNamespace(stage_type="llm", is_comprehension=True),
         SimpleNamespace(stage_type="diffusion", is_comprehension=False),
     ]
-    serving_chat.engine_client.default_sampling_params_list = [
+    mock_engine_client.default_sampling_params_list = [
         SamplingParams(),
         OmniDiffusionSamplingParams(),
     ]
-    serving_chat.engine_client.output_modalities = ["image"]
-    serving_chat.engine_client.errored = False
+    mock_engine_client.output_modalities = ["image"]
+    mock_engine_client.errored = False
+    mock_engine_client.renderer = SimpleNamespace(get_tokenizer=lambda: object())
 
     captured: dict[str, object] = {}
 
@@ -285,21 +284,15 @@ def test_mixed_consumer_keeps_root_common_args_with_nested_extras(serving_chat,
         captured.update(kwargs)
         return results()
 
-    serving_chat.engine_client.generate = generate
-    serving_chat.__dict__.update(
-        _diffusion_mode=False,
-        _diffusion_extra_body_params=frozenset(),
+    mock_engine_client.generate = generate
+    serving_chat = build_serving_chat(
+        engine_client=mock_engine_client,
         models=SimpleNamespace(model_name=lambda _: "test"),
-        renderer=SimpleNamespace(get_tokenizer=lambda: object()),
         online_renderer=SimpleNamespace(validate_chat_template=lambda **_: None),
-        parser_cls=None,
-        use_harmony=False,
-        enable_auto_tools=False,
-        exclude_tools_when_tool_choice_none=False,
         trust_request_chat_template=True,
-        chat_template=None,
-        chat_template_content_format="auto",
     )
+    serving_chat._diffusion_mode = False
+    serving_chat._diffusion_extra_body_params = frozenset()
     mocker.patch.multiple(
         serving_chat,
         _check_model=mocker.AsyncMock(return_value=None),
@@ -386,8 +379,9 @@ def test_openai_sampling_fields_contains_expected_fields():
 # =============================================================================
 
 
-def test_preserves_yaml_defaults_when_no_request_params(serving_chat, mock_request):
+def test_preserves_yaml_defaults_when_no_request_params(mock_engine_client, mock_request):
     """Test that YAML defaults are preserved when request has no params."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     result = serving_chat._build_sampling_params_list_from_request(mock_request)
 
     assert len(result) == 2
@@ -400,8 +394,9 @@ def test_preserves_yaml_defaults_when_no_request_params(serving_chat, mock_reque
     assert comprehension_params.repetition_penalty == 1.05  # YAML custom param preserved
 
 
-def test_request_temperature_overrides_yaml_default(serving_chat, mock_request):
+def test_request_temperature_overrides_yaml_default(mock_engine_client, mock_request):
     """Test that request temperature overrides YAML default."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     mock_request.temperature = 0.8
     mock_request.model_fields_set = {"temperature"}
 
@@ -413,8 +408,9 @@ def test_request_temperature_overrides_yaml_default(serving_chat, mock_request):
     assert comprehension_params.top_k == 1  # YAML custom param preserved
 
 
-def test_request_top_p_overrides_yaml_default(serving_chat, mock_request):
+def test_request_top_p_overrides_yaml_default(mock_engine_client, mock_request):
     """Test that request top_p overrides YAML default."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     mock_request.top_p = 0.95
     mock_request.model_fields_set = {"top_p"}
 
@@ -425,8 +421,9 @@ def test_request_top_p_overrides_yaml_default(serving_chat, mock_request):
     assert comprehension_params.temperature == 0.4  # Preserved from YAML
 
 
-def test_request_max_tokens_overrides_yaml_default(serving_chat, mock_request):
+def test_request_max_tokens_overrides_yaml_default(mock_engine_client, mock_request):
     """Test that request max_tokens overrides YAML default."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     mock_request.max_tokens = 100
     mock_request.model_fields_set = {"max_tokens"}
 
@@ -435,15 +432,17 @@ def test_request_max_tokens_overrides_yaml_default(serving_chat, mock_request):
     assert result[0].max_tokens == 100
 
 
-def test_max_tokens_uses_yaml_default_when_not_specified(serving_chat, mock_request):
+def test_max_tokens_uses_yaml_default_when_not_specified(mock_engine_client, mock_request):
     """Test that max_tokens falls back to YAML default when not in request."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     result = serving_chat._build_sampling_params_list_from_request(mock_request)
 
     assert result[0].max_tokens == 4353
 
 
-def test_request_seed_overrides_yaml_default(serving_chat, mock_request):
+def test_request_seed_overrides_yaml_default(mock_engine_client, mock_request):
     """Test that request seed overrides YAML default."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     mock_request.seed = 123
     mock_request.model_fields_set = {"seed"}
 
@@ -454,8 +453,9 @@ def test_request_seed_overrides_yaml_default(serving_chat, mock_request):
     assert comprehension_params.temperature == 0.4  # Preserved from YAML
 
 
-def test_request_frequency_penalty_overrides(serving_chat, mock_request):
+def test_request_frequency_penalty_overrides(mock_engine_client, mock_request):
     """Test that request frequency_penalty is applied."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     mock_request.frequency_penalty = 0.5
     mock_request.model_fields_set = {"frequency_penalty"}
 
@@ -464,8 +464,9 @@ def test_request_frequency_penalty_overrides(serving_chat, mock_request):
     assert result[0].frequency_penalty == 0.5
 
 
-def test_request_presence_penalty_overrides(serving_chat, mock_request):
+def test_request_presence_penalty_overrides(mock_engine_client, mock_request):
     """Test that request presence_penalty is applied."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     mock_request.presence_penalty = 0.3
     mock_request.model_fields_set = {"presence_penalty"}
 
@@ -474,8 +475,9 @@ def test_request_presence_penalty_overrides(serving_chat, mock_request):
     assert result[0].presence_penalty == 0.3
 
 
-def test_non_comprehension_stages_use_cloned_defaults(serving_chat, mock_request):
+def test_non_comprehension_stages_use_cloned_defaults(mock_engine_client, mock_request):
     """Test that non-comprehension stages always use cloned YAML defaults."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     mock_request.max_tokens = 50
     mock_request.temperature = 0.1
 
@@ -488,8 +490,9 @@ def test_non_comprehension_stages_use_cloned_defaults(serving_chat, mock_request
     assert other_params.seed == 42  # YAML default
 
 
-def test_multiple_params_override_together(serving_chat, mock_request):
+def test_multiple_params_override_together(mock_engine_client, mock_request):
     """Test that multiple request params can override together."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     mock_request.max_tokens = 200
     mock_request.temperature = 0.7
     mock_request.top_p = 0.85
@@ -514,15 +517,17 @@ def test_multiple_params_override_together(serving_chat, mock_request):
 # =============================================================================
 
 
-def test_apply_request_overrides_clones_params(serving_chat, mock_request, default_comprehension_params):
+def test_apply_request_overrides_clones_params(mock_engine_client, mock_request, default_comprehension_params):
     """Test that _apply_request_overrides returns a cloned object."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     result = serving_chat._apply_request_overrides(default_comprehension_params, mock_request)
 
     assert result is not default_comprehension_params  # Different object
 
 
-def test_apply_request_overrides_preserves_defaults(serving_chat, mock_request, default_comprehension_params):
+def test_apply_request_overrides_preserves_defaults(mock_engine_client, mock_request, default_comprehension_params):
     """Test that _apply_request_overrides preserves defaults when request has None."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     result = serving_chat._apply_request_overrides(default_comprehension_params, mock_request)
 
     assert result.temperature == 0.4
@@ -531,8 +536,9 @@ def test_apply_request_overrides_preserves_defaults(serving_chat, mock_request,
     assert result.top_k == 1  # YAML custom param
 
 
-def test_apply_request_overrides_applies_values(serving_chat, mock_request, default_comprehension_params):
+def test_apply_request_overrides_applies_values(mock_engine_client, mock_request, default_comprehension_params):
     """Test that _apply_request_overrides applies non-None request values."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     mock_request.temperature = 0.8
     mock_request.seed = 123
     mock_request.model_fields_set = {"temperature", "seed"}
@@ -550,8 +556,9 @@ def test_apply_request_overrides_applies_values(serving_chat, mock_request, defa
 # =============================================================================
 
 
-def test_apply_overrides_empty_stop_list_preserves_default(serving_chat, mocker):
+def test_apply_overrides_empty_stop_list_preserves_default(mock_engine_client, mocker):
     """Test that request.stop=[] does NOT override YAML default stop words."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     default_params = SamplingParams(temperature=0.5, stop=["<|im_end|>"])
     request = mocker.MagicMock()
     request.temperature = None
@@ -573,8 +580,9 @@ def test_apply_overrides_empty_stop_list_preserves_default(serving_chat, mocker)
     assert result.stop == ["<|im_end|>"]  # YAML default preserved
 
 
-def test_apply_overrides_nonempty_stop_list_overrides_default(serving_chat, mocker):
+def test_apply_overrides_nonempty_stop_list_overrides_default(mock_engine_client, mocker):
     """Test that request.stop=["\\n"] overrides YAML default stop words."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     default_params = SamplingParams(temperature=0.5, stop=["<|im_end|>"])
     request = mocker.MagicMock()
     request.temperature = None
@@ -596,8 +604,9 @@ def test_apply_overrides_nonempty_stop_list_overrides_default(serving_chat, mock
     assert result.stop == ["\n"]  # Overridden by request
 
 
-def test_apply_overrides_empty_stop_token_ids_preserves_default(serving_chat, mocker):
+def test_apply_overrides_empty_stop_token_ids_preserves_default(mock_engine_client, mocker):
     """Test that request.stop_token_ids=[] does NOT override YAML default."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     default_params = SamplingParams(temperature=0.5, stop_token_ids=[2, 3])
     request = mocker.MagicMock()
     request.temperature = None
@@ -617,8 +626,9 @@ def test_apply_overrides_empty_stop_token_ids_preserves_default(serving_chat, mo
     assert result.stop_token_ids == [2, 3]  # YAML default preserved
 
 
-def test_apply_overrides_nonempty_stop_token_ids_overrides_default(serving_chat, mocker):
+def test_apply_overrides_nonempty_stop_token_ids_overrides_default(mock_engine_client, mocker):
     """Test that request.stop_token_ids=[100] overrides YAML default."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     default_params = SamplingParams(temperature=0.5, stop_token_ids=[2, 3])
     request = mocker.MagicMock()
     request.temperature = None
@@ -640,8 +650,9 @@ def test_apply_overrides_nonempty_stop_token_ids_overrides_default(serving_chat,
     assert result.stop_token_ids == [100]  # Overridden by request
 
 
-def test_apply_overrides_mixed_empty_and_nonempty_lists(serving_chat, mocker):
+def test_apply_overrides_mixed_empty_and_nonempty_lists(mock_engine_client, mocker):
     """Test mixing empty and non-empty list fields with scalar fields."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     default_params = SamplingParams(
         temperature=0.4,
         stop=["<|end|>"],
@@ -669,8 +680,9 @@ def test_apply_overrides_mixed_empty_and_nonempty_lists(serving_chat, mocker):
     assert result.stop_token_ids == [100, 200]  # Non-empty list DID override
 
 
-def test_apply_overrides_none_scalar_still_preserves_default(serving_chat, mocker):
+def test_apply_overrides_none_scalar_still_preserves_default(mock_engine_client, mocker):
     """Regression: ensure None scalar values still don't override defaults."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     default_params = SamplingParams(temperature=0.5, max_tokens=100, seed=42)
     request = mocker.MagicMock()
     request.temperature = None
@@ -694,8 +706,9 @@ def test_apply_overrides_none_scalar_still_preserves_default(serving_chat, mocke
     assert result.seed == 42
 
 
-def test_apply_overrides_both_lists_empty_preserves_defaults(serving_chat, mocker):
+def test_apply_overrides_both_lists_empty_preserves_defaults(mock_engine_client, mocker):
     """Test that both stop=[] and stop_token_ids=[] preserve YAML defaults."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     default_params = SamplingParams(
         temperature=0.5,
         stop=["<|end|>", "\\n"],
@@ -722,9 +735,10 @@ def test_apply_overrides_both_lists_empty_preserves_defaults(serving_chat, mocke
     assert result.stop_token_ids == [2, 32000]
 
 
-def test_build_sampling_params_list_empty_stop_preserves_yaml(serving_chat, mock_request):
+def test_build_sampling_params_list_empty_stop_preserves_yaml(mock_engine_client, mock_request):
     """Test that empty stop list in request preserves YAML defaults via
     _build_sampling_params_list_from_request."""
+    serving_chat = build_serving_chat(engine_client=mock_engine_client)
     mock_request.stop = []
     mock_request.stop_token_ids = []
 
@@ -736,21 +750,18 @@ def test_build_sampling_params_list_empty_stop_preserves_yaml(serving_chat, mock
     assert comprehension_params.stop_token_ids == []
 
 
-def test_to_sampling_params_list_pads_missing_tail_stage_with_defaults():
+def test_to_sampling_params_list_pads_missing_tail_stage_with_defaults(mocker: MockerFixture):
     """AURA callers may pass 3 semantic model params for a 4-stage engine pipeline."""
-    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
-
-    instance = object.__new__(OmniOpenAIServingChat)
     default_params = [
         SamplingParams(max_tokens=10),
         SamplingParams(max_tokens=20),
         SamplingParams(max_tokens=30),
         SamplingParams(max_tokens=40),
     ]
-    instance.engine_client = SimpleNamespace(
-        stage_configs=[SimpleNamespace(stage_type="llm") for _ in range(4)],
-        default_sampling_params_list=default_params,
-    )
+    engine_client = mocker.MagicMock()
+    engine_client.stage_configs = [SimpleNamespace(stage_type="llm") for _ in range(4)]
+    engine_client.default_sampling_params_list = default_params
+    instance = build_serving_chat(engine_client=engine_client)
 
     result = instance._to_sampling_params_list(
         [
@@ -772,44 +783,35 @@ def test_to_sampling_params_list_pads_missing_tail_stage_with_defaults():
 
 def test_get_comprehension_stage_index_finds_first_stage(mock_engine_client):
     """Test finding comprehension stage when it's at index 0."""
-    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
-
-    instance = object.__new__(OmniOpenAIServingChat)
-    instance.engine_client = mock_engine_client
+    instance = build_serving_chat(engine_client=mock_engine_client)
 
     assert instance._get_comprehension_stage_index() == 0
 
 
 def test_get_comprehension_stage_index_finds_second_stage(mocker: MockerFixture):
     """Test finding comprehension stage when it's at index 1."""
-    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
-
-    instance = object.__new__(OmniOpenAIServingChat)
-
     other = mocker.MagicMock()
     other.is_comprehension = False
     comprehension = mocker.MagicMock()
     comprehension.is_comprehension = True
 
-    instance.engine_client = mocker.MagicMock()
-    instance.engine_client.stage_configs = [other, comprehension]
+    engine_client = mocker.MagicMock()
+    engine_client.stage_configs = [other, comprehension]
+    instance = build_serving_chat(engine_client=engine_client)
 
     assert instance._get_comprehension_stage_index() == 1
 
 
 def test_get_comprehension_stage_index_raises_when_not_found(mocker: MockerFixture):
     """Test that ValueError is raised when no comprehension stage exists."""
-    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
-
-    instance = object.__new__(OmniOpenAIServingChat)
-
     stage1 = mocker.MagicMock()
     stage1.is_comprehension = False
     stage2 = mocker.MagicMock()
     stage2.is_comprehension = False
 
-    instance.engine_client = mocker.MagicMock()
-    instance.engine_client.stage_configs = [stage1, stage2]
+    engine_client = mocker.MagicMock()
+    engine_client.stage_configs = [stage1, stage2]
+    instance = build_serving_chat(engine_client=engine_client)
 
     with pytest.raises(ValueError, match="No comprehension stage"):
         instance._get_comprehension_stage_index()

tests/entrypoints/test_stream_finish_reason.py

@@ -16,19 +16,20 @@ Key invariants tested:
   - voice/speaker parameter compatibility in chat completions
 """
 
-import json
 import time
 from unittest.mock import MagicMock
 
 import pytest
-from vllm.entrypoints.openai.chat_completion.protocol import (
-    ChatCompletionRequest,
-    ChatCompletionResponseStreamChoice,
-)
-from vllm.entrypoints.openai.engine.protocol import DeltaMessage
-from vllm.entrypoints.openai.models.serving import OpenAIServingModels
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
 from vllm.outputs import CompletionOutput, RequestOutput
 
+from tests.helpers.serving_chat import (
+    build_serving_chat,
+    collect_stream,
+    make_request,
+    make_text_omni_output,
+    parse_sse_chunks,
+)
 from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -40,43 +41,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 # ---------------------------------------------------------------------------
 
 
-def _make_text_omni_output(
-    request_id: str = "test-req",
-    text: str = "hello",
-    token_ids: list[int] | None = None,
-    finish_reason: str | None = None,
-    index: int = 0,
-    num_prompt_tokens: int = 3,
-) -> OmniRequestOutput:
-    """Build an OmniRequestOutput wrapping a text RequestOutput."""
-    if token_ids is None:
-        token_ids = [10, 11, 12]
-    res = RequestOutput(
-        request_id=request_id,
-        prompt="test",
-        prompt_token_ids=list(range(num_prompt_tokens)),
-        prompt_logprobs=None,
-        outputs=[
-            CompletionOutput(
-                index=index,
-                text=text,
-                token_ids=token_ids,
-                cumulative_logprob=0.0,
-                logprobs=None,
-                finish_reason=finish_reason,
-                stop_reason=None,
-            )
-        ],
-        finished=finish_reason is not None,
-    )
-    return OmniRequestOutput.from_stage_output(
-        res,
-        request_id=request_id,
-        final_output_type="text",
-        finished=finish_reason is not None,
-    )
-
-
 def _make_audio_omni_output(
     request_id: str = "test-req",
     index: int = 0,
@@ -118,82 +82,6 @@ def _make_audio_omni_output(
     )
 
 
-def _mock_audio_choices(index: int = 0, role: str = "assistant"):
-    return [
-        ChatCompletionResponseStreamChoice(
-            index=index,
-            delta=DeltaMessage(role=role, content="dGVzdA=="),
-            logprobs=None,
-            finish_reason="stop",
-        )
-    ]
-
-
-def _build_serving_chat():
-    """Create a minimal OmniOpenAIServingChat for testing."""
-    mock_engine = MagicMock()
-    mock_engine.errored = False
-
-    models = OpenAIServingModels(
-        engine_client=mock_engine,
-        base_model_paths=[],
-    )
-    mock_render = MagicMock()
-
-    instance = OmniOpenAIServingChat(
-        engine_client=mock_engine,
-        models=models,
-        response_role="assistant",
-        online_renderer=mock_render,
-        request_logger=None,
-        chat_template=None,
-        chat_template_content_format="auto",
-    )
-    instance._create_audio_choice = MagicMock(
-        side_effect=lambda omni_res, role, request, stream=False: _mock_audio_choices(
-            index=omni_res.outputs[0].index,
-            role=role,
-        )
-    )
-    return instance
-
-
-def _make_request(modalities: list[str], n: int = 1) -> ChatCompletionRequest:
-    req = ChatCompletionRequest(
-        model="test-model",
-        messages=[{"role": "user", "content": "hello"}],
-        n=n,
-        stream=True,
-    )
-    req.modalities = modalities  # type: ignore[attr-defined]
-    return req
-
-
-def _parse_sse_chunks(lines: list[str]) -> list[dict]:
-    """Parse SSE lines into JSON dicts."""
-    prefix = "data: "
-    chunks = []
-    for line in lines:
-        line = line.strip()
-        if not line.startswith(prefix):
-            continue
-        payload = line[len(prefix) :].strip()
-        if payload == "[DONE]":
-            continue
-        try:
-            chunks.append(json.loads(payload))
-        except json.JSONDecodeError:
-            pass
-    return chunks
-
-
-async def _collect_stream(gen):
-    result = []
-    async for item in gen:
-        result.append(item)
-    return result
-
-
 # ---------------------------------------------------------------------------
 # Tests: finish_reason correctness
 # ---------------------------------------------------------------------------
@@ -202,14 +90,14 @@ async def _collect_stream(gen):
 @pytest.mark.asyncio
 async def test_single_modality_text_only_one_stop():
     """Text-only streaming: exactly one chunk has finish_reason='stop'."""
-    serving_chat = _build_serving_chat()
-    request = _make_request(modalities=["text"])
+    serving_chat = build_serving_chat()
+    request = make_request(modalities=["text"])
 
     async def result_generator():
-        yield _make_text_omni_output(text="he", token_ids=[10, 11], finish_reason=None)
-        yield _make_text_omni_output(text="llo", token_ids=[12], finish_reason="stop")
+        yield make_text_omni_output(text="he", token_ids=[10, 11], finish_reason=None)
+        yield make_text_omni_output(text="llo", token_ids=[12], finish_reason="stop")
 
-    raw_lines = await _collect_stream(
+    raw_lines = await collect_stream(
         serving_chat.chat_completion_stream_generator(
             request=request,
             result_generator=result_generator(),
@@ -221,7 +109,7 @@ async def test_single_modality_text_only_one_stop():
         )
     )
 
-    chunks = _parse_sse_chunks(raw_lines)
+    chunks = parse_sse_chunks(raw_lines)
     finish_reasons = [c["choices"][0]["finish_reason"] for c in chunks if c.get("choices")]
 
     assert finish_reasons[-1] == "stop"
@@ -233,15 +121,15 @@ async def test_single_modality_text_only_one_stop():
 @pytest.mark.asyncio
 async def test_multi_modal_text_audio_only_last_stop():
     """text+audio: text finish sends finish_reason=null, audio sends stop."""
-    serving_chat = _build_serving_chat()
-    request = _make_request(modalities=["text", "audio"])
+    serving_chat = build_serving_chat()
+    request = make_request(modalities=["text", "audio"])
 
     async def result_generator():
-        yield _make_text_omni_output(text="he", token_ids=[10, 11], finish_reason=None)
-        yield _make_text_omni_output(text="llo", token_ids=[12], finish_reason="stop")
+        yield make_text_omni_output(text="he", token_ids=[10, 11], finish_reason=None)
+        yield make_text_omni_output(text="llo", token_ids=[12], finish_reason="stop")
         yield _make_audio_omni_output()
 
-    raw_lines = await _collect_stream(
+    raw_lines = await collect_stream(
         serving_chat.chat_completion_stream_generator(
             request=request,
             result_generator=result_generator(),
@@ -253,7 +141,7 @@ async def test_multi_modal_text_audio_only_last_stop():
         )
     )
 
-    chunks = _parse_sse_chunks(raw_lines)
+    chunks = parse_sse_chunks(raw_lines)
     finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])]
 
     assert finish_reasons.count("stop") == 1
@@ -269,18 +157,18 @@ async def test_multi_modal_text_audio_only_last_stop():
 @pytest.mark.asyncio
 async def test_multi_modal_n2_independent_per_choice():
     """n=2 with text+audio: each choice gets exactly one stop, at the end."""
-    serving_chat = _build_serving_chat()
-    request = _make_request(modalities=["text", "audio"], n=2)
+    serving_chat = build_serving_chat()
+    request = make_request(modalities=["text", "audio"], n=2)
 
     async def result_generator():
-        yield _make_text_omni_output(text="A", token_ids=[10], finish_reason=None, index=0)
-        yield _make_text_omni_output(text="B", token_ids=[20], finish_reason=None, index=1)
-        yield _make_text_omni_output(text="", token_ids=[11], finish_reason="stop", index=0)
-        yield _make_text_omni_output(text="", token_ids=[21], finish_reason="stop", index=1)
+        yield make_text_omni_output(text="A", token_ids=[10], finish_reason=None, index=0)
+        yield make_text_omni_output(text="B", token_ids=[20], finish_reason=None, index=1)
+        yield make_text_omni_output(text="", token_ids=[11], finish_reason="stop", index=0)
+        yield make_text_omni_output(text="", token_ids=[21], finish_reason="stop", index=1)
         yield _make_audio_omni_output(index=0)
         yield _make_audio_omni_output(index=1)
 
-    raw_lines = await _collect_stream(
+    raw_lines = await collect_stream(
         serving_chat.chat_completion_stream_generator(
             request=request,
             result_generator=result_generator(),
@@ -292,7 +180,7 @@ async def test_multi_modal_n2_independent_per_choice():
         )
     )
 
-    chunks = _parse_sse_chunks(raw_lines)
+    chunks = parse_sse_chunks(raw_lines)
     per_choice: dict[int, list] = {}
     for c in chunks:
         for ch in c.get("choices", []):
@@ -306,13 +194,13 @@ async def test_multi_modal_n2_independent_per_choice():
 @pytest.mark.asyncio
 async def test_single_modality_audio_only_one_stop():
     """Audio-only streaming: the audio chunk carries finish_reason='stop'."""
-    serving_chat = _build_serving_chat()
-    request = _make_request(modalities=["audio"])
+    serving_chat = build_serving_chat()
+    request = make_request(modalities=["audio"])
 
     async def result_generator():
         yield _make_audio_omni_output()
 
-    raw_lines = await _collect_stream(
+    raw_lines = await collect_stream(
         serving_chat.chat_completion_stream_generator(
             request=request,
             result_generator=result_generator(),
@@ -324,7 +212,7 @@ async def test_single_modality_audio_only_one_stop():
         )
     )
 
-    chunks = _parse_sse_chunks(raw_lines)
+    chunks = parse_sse_chunks(raw_lines)
     finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])]
 
     assert finish_reasons.count("stop") == 1
@@ -349,8 +237,8 @@ async def test_streaming_audio_metrics_resolve_replica_id(
     import vllm_omni.entrypoints.openai.serving_chat as serving_chat_mod
     from vllm_omni.entrypoints.client_request_state import ClientRequestState
 
-    serving_chat = _build_serving_chat()
-    request = _make_request(modalities=["audio"])
+    serving_chat = build_serving_chat()
+    request = make_request(modalities=["audio"])
     req_state = ClientRequestState(
         request_id="internal-req",
         external_request_id="test-req",
@@ -377,7 +265,7 @@ async def test_streaming_audio_metrics_resolve_replica_id(
     async def result_generator():
         yield _make_audio_omni_output(stage_id=2, replica_id=output_replica_id, audio_samples=2400)
 
-    await _collect_stream(
+    await collect_stream(
         serving_chat.chat_completion_stream_generator(
             request=request,
             result_generator=result_generator(),
@@ -408,15 +296,15 @@ async def test_streaming_audio_metrics_resolve_replica_id(
 async def test_declared_modality_not_produced_emits_fallback_stop():
     """If request.modalities declares ["text","audio"] but engine only produces
     text, a fallback stop chunk is emitted at stream end."""
-    serving_chat = _build_serving_chat()
-    request = _make_request(modalities=["text", "audio"])
+    serving_chat = build_serving_chat()
+    request = make_request(modalities=["text", "audio"])
 
     async def result_generator():
         # Engine only produces text, no audio output at all
-        yield _make_text_omni_output(text="hi", token_ids=[10], finish_reason=None)
-        yield _make_text_omni_output(text="!", token_ids=[11], finish_reason="stop")
+        yield make_text_omni_output(text="hi", token_ids=[10], finish_reason=None)
+        yield make_text_omni_output(text="!", token_ids=[11], finish_reason="stop")
 
-    raw_lines = await _collect_stream(
+    raw_lines = await collect_stream(
         serving_chat.chat_completion_stream_generator(
             request=request,
             result_generator=result_generator(),
@@ -428,7 +316,7 @@ async def test_declared_modality_not_produced_emits_fallback_stop():
         )
     )
 
-    chunks = _parse_sse_chunks(raw_lines)
+    chunks = parse_sse_chunks(raw_lines)
     finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])]
 
     # Text finish is suppressed (audio not seen yet), but fallback stop
@@ -441,15 +329,15 @@ async def test_declared_modality_not_produced_emits_fallback_stop():
 async def test_declared_modality_not_produced_text_finish_suppressed():
     """When text finishes but audio (declared in modalities) never appears,
     the text finish chunk has finish_reason=null (suppressed)."""
-    serving_chat = _build_serving_chat()
-    request = _make_request(modalities=["text", "audio"])
+    serving_chat = build_serving_chat()
+    request = make_request(modalities=["text", "audio"])
 
     async def result_generator():
-        yield _make_text_omni_output(text="hi", token_ids=[10], finish_reason=None)
-        yield _make_text_omni_output(text="!", token_ids=[11], finish_reason="stop")
+        yield make_text_omni_output(text="hi", token_ids=[10], finish_reason=None)
+        yield make_text_omni_output(text="!", token_ids=[11], finish_reason="stop")
         # No audio output — stream ends
 
-    raw_lines = await _collect_stream(
+    raw_lines = await collect_stream(
         serving_chat.chat_completion_stream_generator(
             request=request,
             result_generator=result_generator(),
@@ -461,7 +349,7 @@ async def test_declared_modality_not_produced_text_finish_suppressed():
         )
     )
 
-    chunks = _parse_sse_chunks(raw_lines)
+    chunks = parse_sse_chunks(raw_lines)
 
     # Find the text finish chunk (content "!")
     for c in chunks:
@@ -485,8 +373,8 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive():
     yields (field, value) tuples, so the loop raised
     ``AttributeError: 'tuple' object has no attribute 'finish_reason'``.
     """
-    serving_chat = _build_serving_chat()
-    request = _make_request(modalities=["text", "audio"])
+    serving_chat = build_serving_chat()
+    request = make_request(modalities=["text", "audio"])
 
     empty_audio = _make_audio_omni_output()
     empty_audio.outputs[0].multimodal_output = {"audio": []}
@@ -497,11 +385,11 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive():
     serving_chat._create_audio_choice = create_audio_choice
 
     async def result_generator():
-        yield _make_text_omni_output(text="hi", token_ids=[10], finish_reason=None)
-        yield _make_text_omni_output(text="!", token_ids=[11], finish_reason="stop")
+        yield make_text_omni_output(text="hi", token_ids=[10], finish_reason=None)
+        yield make_text_omni_output(text="!", token_ids=[11], finish_reason="stop")
         yield empty_audio
 
-    raw_lines = await _collect_stream(
+    raw_lines = await collect_stream(
         serving_chat.chat_completion_stream_generator(
             request=request,
             result_generator=result_generator(),
@@ -514,7 +402,7 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive():
     )
 
     assert not any("Error in chat completion stream generator" in line for line in raw_lines)
-    chunks = _parse_sse_chunks(raw_lines)
+    chunks = parse_sse_chunks(raw_lines)
     finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])]
     assert finish_reasons.count("stop") == 1, f"Expected 1 stop, got {finish_reasons}"
     assert finish_reasons[-1] == "stop"
@@ -525,18 +413,18 @@ async def test_audio_choice_error_response_is_not_iterated_as_choices():
     """Any ErrorResponse from the audio path is skipped, not iterated."""
     from vllm.entrypoints.openai.engine.protocol import ErrorResponse
 
-    serving_chat = _build_serving_chat()
-    request = _make_request(modalities=["text", "audio"])
+    serving_chat = build_serving_chat()
+    request = make_request(modalities=["text", "audio"])
 
     serving_chat._create_audio_choice = MagicMock(
         side_effect=lambda omni_res, role, request, stream=False: serving_chat._create_error_response("boom")
     )
 
     async def result_generator():
-        yield _make_text_omni_output(text="hi", token_ids=[10], finish_reason="stop")
+        yield make_text_omni_output(text="hi", token_ids=[10], finish_reason="stop")
         yield _make_audio_omni_output()
 
-    raw_lines = await _collect_stream(
+    raw_lines = await collect_stream(
         serving_chat.chat_completion_stream_generator(
             request=request,
             result_generator=result_generator(),
@@ -550,7 +438,7 @@ async def test_audio_choice_error_response_is_not_iterated_as_choices():
 
     assert isinstance(serving_chat._create_error_response("boom"), ErrorResponse)
     assert not any("AttributeError" in line for line in raw_lines)
-    chunks = _parse_sse_chunks(raw_lines)
+    chunks = parse_sse_chunks(raw_lines)
     finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])]
     assert finish_reasons.count("stop") == 1, f"Expected 1 stop, got {finish_reasons}"
 

tests/helpers/runtime.py

@@ -770,6 +770,7 @@ class OmniResponse:
     success: bool = False
     prompt_tokens: int | None = None
     cached_tokens: int | None = None
+    multimodal_tokens: dict[str, int] | None = None
     logprobs: list | None = None
     #: HTTP status + error text for the error-handling path (e.g. validator
     #: rejections); populated when the OpenAI client raises an APIError.
@@ -1089,6 +1090,7 @@ class OpenAIClientHandler:
                     result.prompt_tokens = chunk.usage.prompt_tokens
                     if details := getattr(chunk.usage, "prompt_tokens_details", None):
                         result.cached_tokens = details.cached_tokens
+                        result.multimodal_tokens = getattr(details, "multimodal_tokens", None)
 
             if audio_data:
                 merged_seg = _merge_base64_audio_to_segment(audio_data)
@@ -1121,6 +1123,7 @@ class OpenAIClientHandler:
                 result.prompt_tokens = usage.prompt_tokens
                 if details := getattr(usage, "prompt_tokens_details", None):
                     result.cached_tokens = details.cached_tokens
+                    result.multimodal_tokens = getattr(details, "multimodal_tokens", None)
             if audio_data:
                 result.audio_bytes = base64.b64decode(audio_data)
             result.text_content = text_content

tests/helpers/serving_chat.py

@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Shared scaffolding for OmniOpenAIServingChat unit tests.
+
+Provides factories for a minimal serving-chat instance, chat requests, and
+OmniRequestOutput objects, plus SSE stream collection/parsing utilities.
+"""
+
+import json
+from unittest.mock import MagicMock
+
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    ChatCompletionRequest,
+    ChatCompletionResponseStreamChoice,
+)
+from vllm.entrypoints.openai.engine.protocol import DeltaMessage
+from vllm.entrypoints.openai.models.serving import OpenAIServingModels
+from vllm.outputs import CompletionOutput, RequestOutput
+
+from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+from vllm_omni.outputs import OmniRequestOutput
+
+
+def make_text_omni_output(
+    request_id: str = "test-req",
+    text: str = "hello",
+    token_ids: list[int] | None = None,
+    finish_reason: str | None = None,
+    index: int = 0,
+    num_prompt_tokens: int = 3,
+    num_cached_tokens: int | None = None,
+) -> OmniRequestOutput:
+    """Build an OmniRequestOutput wrapping a text RequestOutput."""
+    if token_ids is None:
+        token_ids = [10, 11, 12]
+    res = RequestOutput(
+        request_id=request_id,
+        prompt="test",
+        prompt_token_ids=list(range(num_prompt_tokens)),
+        prompt_logprobs=None,
+        outputs=[
+            CompletionOutput(
+                index=index,
+                text=text,
+                token_ids=token_ids,
+                cumulative_logprob=0.0,
+                logprobs=None,
+                finish_reason=finish_reason,
+                stop_reason=None,
+            )
+        ],
+        finished=finish_reason is not None,
+        num_cached_tokens=num_cached_tokens,
+    )
+    return OmniRequestOutput.from_stage_output(
+        res,
+        request_id=request_id,
+        final_output_type="text",
+        finished=finish_reason is not None,
+    )
+
+
+def mock_audio_choices(index: int = 0, role: str = "assistant"):
+    return [
+        ChatCompletionResponseStreamChoice(
+            index=index,
+            delta=DeltaMessage(role=role, content="dGVzdA=="),
+            logprobs=None,
+            finish_reason="stop",
+        )
+    ]
+
+
+def build_serving_chat(
+    *,
+    engine_client=None,
+    models=None,
+    online_renderer=None,
+    response_role: str = "assistant",
+    request_logger=None,
+    chat_template=None,
+    chat_template_content_format="auto",
+    **init_overrides,
+) -> OmniOpenAIServingChat:
+    """Create a real OmniOpenAIServingChat via its actual __init__.
+
+    engine_client/models/online_renderer default to minimal mocks; pass your
+    own to wire in test-specific behavior (e.g. engine_client.stage_configs).
+    Any other OmniOpenAIServingChat.__init__ keyword — enable_auto_tools,
+    enable_prompt_tokens_details, trust_request_chat_template, etc. — can be
+    passed through **init_overrides instead of set on the instance after
+    construction.
+    """
+    if engine_client is None:
+        engine_client = MagicMock()
+        engine_client.errored = False
+    if models is None:
+        models = OpenAIServingModels(
+            engine_client=engine_client,
+            base_model_paths=[],
+        )
+    if online_renderer is None:
+        online_renderer = MagicMock()
+
+    instance = OmniOpenAIServingChat(
+        engine_client=engine_client,
+        models=models,
+        response_role=response_role,
+        online_renderer=online_renderer,
+        request_logger=request_logger,
+        chat_template=chat_template,
+        chat_template_content_format=chat_template_content_format,
+        **init_overrides,
+    )
+    instance._create_audio_choice = MagicMock(
+        side_effect=lambda omni_res, role, request, stream=False: mock_audio_choices(
+            index=omni_res.outputs[0].index,
+            role=role,
+        )
+    )
+    return instance
+
+
+def make_request(
+    modalities: list[str],
+    n: int = 1,
+    *,
+    stream: bool = True,
+    include_usage: bool = False,
+) -> ChatCompletionRequest:
+    req = ChatCompletionRequest(
+        model="test-model",
+        messages=[{"role": "user", "content": "hello"}],
+        n=n,
+        stream=stream,
+        stream_options={"include_usage": True} if include_usage else None,
+    )
+    req.modalities = modalities  # type: ignore[attr-defined]
+    return req
+
+
+def parse_sse_chunks(lines: list[str]) -> list[dict]:
+    """Parse SSE lines into JSON dicts."""
+    prefix = "data: "
+    chunks = []
+    for line in lines:
+        line = line.strip()
+        if not line.startswith(prefix):
+            continue
+        payload = line[len(prefix) :].strip()
+        if payload == "[DONE]":
+            continue
+        chunks.append(json.loads(payload))
+    return chunks
+
+
+async def collect_stream(gen):
+    result = []
+    async for item in gen:
+        result.append(item)
+    return result

vllm_omni/entrypoints/openai/serving_chat.py

@@ -58,7 +58,11 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionResponseStreamChoice,
     ChatMessage,
 )
-from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
+from vllm.entrypoints.openai.chat_completion.serving import (
+    OpenAIServingChat,
+    _get_mm_token_counts,
+    _make_prompt_tokens_details,
+)
 from vllm.entrypoints.openai.engine.protocol import (
     DeltaFunctionCall,
     DeltaMessage,
@@ -67,7 +71,6 @@ from vllm.entrypoints.openai.engine.protocol import (
     ErrorResponse,
     FunctionCall,
     FunctionDefinition,
-    PromptTokenUsageInfo,
     RequestResponseMetadata,
     ToolCall,
     UsageInfo,
@@ -730,8 +733,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
         # Schedule the request and get the result generator.
         generators: list[AsyncGenerator[RequestOutput, None]] = []
+        mm_token_counts: dict[str, int] | None = None
         try:
             for i, engine_prompt in enumerate(engine_prompts):
+                if self.enable_prompt_tokens_details:
+                    mm_token_counts = _get_mm_token_counts(engine_prompt)
                 if hasattr(request, "sampling_params_list"):
                     sampling_params_list = self._to_sampling_params_list(request.sampling_params_list)
                 else:
@@ -814,6 +820,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 tokenizer,
                 request_metadata,
                 reasoning_parser,
+                mm_token_counts=mm_token_counts,
                 raw_request=raw_request,
             )
 
@@ -827,6 +834,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 tokenizer,
                 request_metadata,
                 reasoning_parser,
+                mm_token_counts=mm_token_counts,
             )
         except ValueError as e:
             return self.create_error_response(e)
@@ -1387,6 +1395,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         tokenizer: AnyTokenizer,
         request_metadata: RequestResponseMetadata,
         reasoning_parser: ReasoningParser | None = None,
+        mm_token_counts: dict[str, int] | None = None,
         raw_request: Request | None = None,
     ):
         created_time = int(time.time())
@@ -1407,6 +1416,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         stop_reason_emitted: list[bool] = [False] * num_choices
         num_prompt_tokens = 0
         num_cached_tokens = None
+        num_cache_creation_tokens = None
         if self.use_harmony:
             harmony_parsers = [get_streamable_parser_for_assistant() for _ in range(num_choices)]
             harmony_tools_streamed = [False] * num_choices
@@ -1503,6 +1513,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                             num_prompt_tokens += len(res.encoder_prompt_token_ids)
 
                     num_cached_tokens = res.num_cached_tokens
+                    num_cache_creation_tokens = res.num_cache_creation_tokens
                     # Send first response for each choice with role
                     # NOTE: num_choices defaults to 1 so this usually executes once per request
                     for i in range(num_choices):
@@ -2234,8 +2245,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     completion_tokens=completion_tokens,
                     total_tokens=num_prompt_tokens + completion_tokens,
                 )
-                if self.enable_prompt_tokens_details and num_cached_tokens:
-                    final_usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=num_cached_tokens)
+                final_usage.prompt_tokens_details = _make_prompt_tokens_details(
+                    self.enable_prompt_tokens_details,
+                    num_cached_tokens,
+                    num_cache_creation_tokens,
+                    mm_token_counts,
+                )
 
                 final_usage_chunk = OmniChatCompletionStreamResponse(
                     id=request_id,
@@ -2307,6 +2322,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         tokenizer: TokenizerLike,
         request_metadata: RequestResponseMetadata,
         reasoning_parser: ReasoningParser | None = None,
+        mm_token_counts: dict[str, int] | None = None,
     ) -> ErrorResponse | OmniChatCompletionResponse:
         created_time = int(time.time())
         final_res: RequestOutput | None = None
@@ -2361,6 +2377,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         conversation,
                         role,
                         reasoning_parser,
+                        mm_token_counts=mm_token_counts,
                     )
                     final_res = omni_outputs
                 else:
@@ -2462,6 +2479,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         conversation: list[ConversationMessage],
         role: str,
         reasoning_parser: ReasoningParser | None = None,
+        mm_token_counts: dict[str, int] | None = None,
     ):
         final_res = omni_outputs
         if self.tool_call_id_type == "kimi_k2":
@@ -2698,8 +2716,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             completion_tokens=num_generated_tokens,
             total_tokens=num_prompt_tokens + num_generated_tokens,
         )
-        if self.enable_prompt_tokens_details and final_res.num_cached_tokens:
-            usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=final_res.num_cached_tokens)
+        usage.prompt_tokens_details = _make_prompt_tokens_details(
+            self.enable_prompt_tokens_details,
+            final_res.num_cached_tokens,
+            final_res.num_cache_creation_tokens,
+            mm_token_counts,
+        )
 
         prompt_logprobs = clamp_prompt_logprobs(final_res.prompt_logprobs)
         prompt_token_ids = final_res.prompt_token_ids if request.return_token_ids else None