[Refactor][CI] Reorganize test tree, and move slow e2e to weekly (#4951)

Signed-off-by: wangyu <410167048@qq.com>

tests/engine/test_arg_utils_shared_fields.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for vllm_omni.engine.arg_utils — invariants that must
+hold for the orchestrator/engine/server CLI flag partition."""
+
+from __future__ import annotations
+
+from dataclasses import fields
+
+import pytest
+
+from vllm_omni.engine.arg_utils import (
+    SHARED_FIELDS,
+    OmniEngineArgs,
+    internal_blacklist_keys,
+    orchestrator_field_names,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_no_ambiguous_overlap_with_real_engine():
+    """OrchestratorArgs ∩ OmniEngineArgs must be ⊆ SHARED_FIELDS ∪ orchestrator-captured.
+
+    Fields that the orchestrator intentionally captures to prevent uniform
+    propagation (async_chunk, tokenizer) are excluded — they exist in both
+    classes by design but flow through DeployConfig, not kwargs forwarding.
+    """
+    # Fields on both classes by design: orchestrator captures them to prevent
+    # uniform per-stage propagation; redistributed via DeployConfig.
+    orchestrator_captured = {"async_chunk", "tokenizer"}
+
+    orch = orchestrator_field_names()
+    engine = {f.name for f in fields(OmniEngineArgs)}
+    overlap = orch & engine
+    unexpected = overlap - SHARED_FIELDS - orchestrator_captured
+    assert not unexpected, (
+        f"OmniEngineArgs has ambiguous overlap with OrchestratorArgs: "
+        f"{sorted(unexpected)}. Update SHARED_FIELDS or remove duplication."
+    )
+
+
+def test_internal_blacklist_keys_derived_from_orchestrator():
+    """Blacklist is exactly OrchestratorArgs fields minus SHARED_FIELDS.
+
+    This function replaces the old hardcoded INTERNAL_STAGE_OVERRIDE_KEYS
+    frozenset. Asserts the contract so future changes to OrchestratorArgs
+    automatically propagate to the blacklist.
+    """
+    blacklist = internal_blacklist_keys()
+    assert blacklist == orchestrator_field_names() - SHARED_FIELDS

tests/engine/test_data_entry_keys.py

@@ -0,0 +1,396 @@
+"""Tests for data_entry_keys."""
+
+import msgspec
+import pytest
+import torch
+
+from vllm_omni.data_entry_keys import (
+    CodesStruct,
+    EmbeddingsStruct,
+    HiddenStatesStruct,
+    IdsStruct,
+    MetaStruct,
+    OmniPayload,
+    OmniPayloadStruct,
+    deserialize_payload,
+    flatten_payload,
+    serialize_payload,
+    to_dict,
+    to_struct,
+    unflatten_payload,
+)
+from vllm_omni.engine import AdditionalInformationPayload
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestOmniPayloadStruct:
+    """Runtime-validated mirror of OmniPayload (msgspec.Struct)."""
+
+    def test_to_struct_validates_dict(self):
+        d = {"meta": {"left_context_size": 25, "finished": torch.tensor(False)}}
+        s = to_struct(d)
+        assert s.meta.left_context_size == 25
+
+    def test_to_struct_rejects_legacy_flat_top_level(self):
+        with pytest.raises(msgspec.ValidationError, match="unknown field"):
+            to_struct({"code_predictor_codes": torch.zeros(3, 8)})
+
+    def test_to_struct_rejects_legacy_flat_meta_field(self):
+        # `left_context_size` at top level (legacy) instead of under `meta`
+        with pytest.raises(msgspec.ValidationError, match="unknown field"):
+            to_struct({"left_context_size": 25})
+
+    def test_to_struct_rejects_typo_in_subkey(self):
+        with pytest.raises(msgspec.ValidationError, match="unknown field"):
+            to_struct({"meta": {"finisheed": True}})
+
+    def test_to_struct_rejects_minicpmo_duplex_runtime_handoff_fields(self):
+        with pytest.raises(msgspec.ValidationError, match="unknown field"):
+            to_struct({"hidden_states": {"tts": torch.zeros(1)}})
+        with pytest.raises(msgspec.ValidationError, match="unknown field"):
+            to_struct({"ids": {"tts": [1]}})
+        with pytest.raises(msgspec.ValidationError, match="unknown field"):
+            to_struct({"meta": {"ref_audio_sr": 16000}})
+        with pytest.raises(msgspec.ValidationError, match="unknown field"):
+            to_struct({"meta": {"listen_token_id": 151685}})
+
+    def test_to_struct_rejects_wrong_type(self):
+        with pytest.raises(msgspec.ValidationError, match="Expected"):
+            to_struct({"meta": {"left_context_size": "not_an_int"}})
+
+    def test_round_trip_dict_struct_dict(self):
+        original = {
+            "meta": {"left_context_size": 7, "finished": torch.tensor(True)},
+            "codes": {"audio": torch.zeros(2, 8)},
+            "hidden_states": {"output": torch.zeros(4, 16)},
+        }
+        s = to_struct(original)
+        d = to_dict(s)
+        assert sorted(d.keys()) == sorted(original.keys())
+        for top, sub in original.items():
+            assert sorted(d[top].keys()) == sorted(sub.keys())
+
+    def test_to_dict_drops_unset_fields(self):
+        s = OmniPayloadStruct(meta=MetaStruct(left_context_size=10))
+        d = to_dict(s)
+        assert d == {"meta": {"left_context_size": 10}}
+
+    def test_struct_with_all_categories(self):
+        d = {
+            "hidden_states": {"output": torch.zeros(1), "last": torch.ones(2, 4)},
+            "embed": {"prefill": torch.zeros(1), "tts_bos": torch.zeros(1)},
+            "ids": {"all": [1, 2], "prompt": [1], "speech_token": [3, 4]},
+            "codes": {"audio": torch.zeros(1), "ref": torch.ones(4)},
+            "meta": {"left_context_size": 3, "num_processed_tokens": 7},
+        }
+        s = to_struct(d)
+        assert isinstance(s.hidden_states, HiddenStatesStruct)
+        assert isinstance(s.embed, EmbeddingsStruct)
+        assert isinstance(s.ids, IdsStruct)
+        assert isinstance(s.codes, CodesStruct)
+        assert isinstance(s.meta, MetaStruct)
+        assert s.ids.all == [1, 2]
+        assert s.ids.speech_token == [3, 4]
+        assert torch.equal(s.hidden_states.last, torch.ones(2, 4))
+        assert s.meta.num_processed_tokens == 7
+        assert torch.equal(s.codes.ref, torch.ones(4))
+
+
+class TestValidatePayload:
+    def test_raises_on_unknown_top_level(self):
+        from vllm_omni.data_entry_keys import validate_payload
+
+        with pytest.raises(msgspec.ValidationError, match="unknown field"):
+            validate_payload({"code_predictor_codes": torch.zeros(3, 8)}, context="test_boundary")
+
+    def test_raises_on_unknown_sub_key(self):
+        from vllm_omni.data_entry_keys import validate_payload
+
+        with pytest.raises(msgspec.ValidationError, match="unknown field"):
+            validate_payload({"meta": {"finisheed": True}})
+
+    def test_none_is_ok(self):
+        from vllm_omni.data_entry_keys import validate_payload
+
+        validate_payload(None)  # should not raise
+
+    def test_valid_payload_passes(self):
+        from vllm_omni.data_entry_keys import validate_payload
+
+        validate_payload({"meta": {"left_context_size": 5}})
+
+    def test_context_in_error_message(self):
+        from vllm_omni.data_entry_keys import validate_payload
+
+        with pytest.raises(msgspec.ValidationError, match="my_call_site"):
+            validate_payload({"bad": 1}, context="my_call_site")
+
+
+class TestWireEquivalenceStructVsDict:
+    """Producer return-side migration invariant: encoding an ``OmniPayloadStruct``
+    via the connector serializer must decode to the same payload as encoding
+    the equivalent ``to_dict(struct)``.
+
+    Guards against regressions where the wire format diverges between the two
+    paths (e.g. msgspec adds a struct tag, or ``to_dict`` drops a non-default
+    sub-field that ``omit_defaults`` retains).
+    """
+
+    @staticmethod
+    def _round_trip(obj):
+        from vllm_omni.distributed.omni_connectors.utils.serialization import (
+            OmniMsgpackDecoder,
+            OmniMsgpackEncoder,
+        )
+
+        return OmniMsgpackDecoder().decode(OmniMsgpackEncoder().encode(obj))
+
+    @staticmethod
+    def _assert_decoded_equal(a, b):
+        if isinstance(a, dict):
+            assert isinstance(b, dict)
+            assert sorted(a.keys()) == sorted(b.keys())
+            for k in a:
+                TestWireEquivalenceStructVsDict._assert_decoded_equal(a[k], b[k])
+        elif isinstance(a, torch.Tensor):
+            assert isinstance(b, torch.Tensor)
+            assert a.dtype == b.dtype
+            assert a.shape == b.shape
+            assert torch.equal(a, b)
+        elif isinstance(a, list):
+            assert isinstance(b, list) and len(a) == len(b)
+            for x, y in zip(a, b, strict=True):
+                TestWireEquivalenceStructVsDict._assert_decoded_equal(x, y)
+        else:
+            assert a == b
+
+    def test_basic_payload(self):
+        struct = OmniPayloadStruct(
+            codes=CodesStruct(audio=torch.tensor([1, 2, 3], dtype=torch.long)),
+            meta=MetaStruct(left_context_size=10, finished=torch.tensor(True, dtype=torch.bool)),
+        )
+        self._assert_decoded_equal(self._round_trip(struct), self._round_trip(to_dict(struct)))
+
+    def test_nested_sub_structs(self):
+        # Exercises depth-2 sub-struct encoding (embed.*) which was the case that
+        # exposed schema drift in the #1829 migration.
+        struct = OmniPayloadStruct(
+            codes=CodesStruct(audio=torch.tensor([5, 6], dtype=torch.long)),
+            meta=MetaStruct(finished=torch.tensor(False, dtype=torch.bool)),
+            embed=EmbeddingsStruct(
+                speech_token=torch.tensor([[1, 2, 3]]),
+                speech_feat=torch.tensor([[[0.1, 0.2], [0.3, 0.4]]]),
+                embedding=torch.tensor([[0.5, 0.6]]),
+            ),
+        )
+        self._assert_decoded_equal(self._round_trip(struct), self._round_trip(to_dict(struct)))
+
+
+class TestFlattenPayload:
+    def test_basic_nested_to_dotted(self):
+        nested = {
+            "codes": {"audio": torch.tensor([1.0])},
+            "meta": {"finished": torch.tensor(True, dtype=torch.bool), "left_context_size": 5},
+        }
+        flat = flatten_payload(nested)
+        assert torch.equal(flat["codes.audio"], torch.tensor([1.0]))
+        assert flat["meta.finished"].item() is True
+        assert flat["meta.left_context_size"] == 5
+        assert "codes" not in flat
+        assert "meta" not in flat
+
+    def test_top_level_keys_preserved(self):
+        nested = {
+            "latent": torch.tensor([9.0]),
+            "generated_len": 42,
+        }
+        flat = flatten_payload(nested)
+        assert torch.equal(flat["latent"], torch.tensor([9.0]))
+        assert flat["generated_len"] == 42
+
+    def test_hidden_states_layers_expanded(self):
+        nested = {
+            "hidden_states": {
+                "output": torch.tensor([1.0]),
+                "layers": {
+                    0: torch.tensor([2.0]),
+                    24: torch.tensor([3.0]),
+                },
+            },
+        }
+        flat = flatten_payload(nested)
+        assert torch.equal(flat["hidden_states.output"], torch.tensor([1.0]))
+        assert torch.equal(flat["hidden_states.layer_0"], torch.tensor([2.0]))
+        assert torch.equal(flat["hidden_states.layer_24"], torch.tensor([3.0]))
+        assert "hidden_states.layers" not in flat
+
+    def test_mixed_nested_and_top_level(self):
+        nested: OmniPayload = {
+            "codes": {"audio": torch.tensor([1.0])},
+            "latent": torch.tensor([2.0]),
+            "meta": {"finished": torch.tensor(False, dtype=torch.bool)},
+        }
+        flat = flatten_payload(nested)
+        assert set(flat.keys()) == {"codes.audio", "latent", "meta.finished"}
+
+
+class TestUnflattenPayload:
+    def test_basic_dotted_to_nested(self):
+        flat = {
+            "codes.audio": torch.tensor([1.0]),
+            "meta.finished": torch.tensor(True, dtype=torch.bool),
+            "meta.left_context_size": 5,
+        }
+        nested = unflatten_payload(flat)
+        assert torch.equal(nested["codes"]["audio"], torch.tensor([1.0]))
+        assert nested["meta"]["finished"].item() is True
+        assert nested["meta"]["left_context_size"] == 5
+
+    def test_top_level_keys_preserved(self):
+        flat = {"latent": torch.tensor([9.0]), "generated_len": 42}
+        nested = unflatten_payload(flat)
+        assert torch.equal(nested["latent"], torch.tensor([9.0]))
+        assert nested["generated_len"] == 42
+
+    def test_hidden_states_layers_collected(self):
+        flat = {
+            "hidden_states.output": torch.tensor([1.0]),
+            "hidden_states.layer_0": torch.tensor([2.0]),
+            "hidden_states.layer_24": torch.tensor([3.0]),
+        }
+        nested = unflatten_payload(flat)
+        assert torch.equal(nested["hidden_states"]["output"], torch.tensor([1.0]))
+        assert torch.equal(nested["hidden_states"]["layers"][0], torch.tensor([2.0]))
+        assert torch.equal(nested["hidden_states"]["layers"][24], torch.tensor([3.0]))
+
+
+class TestFlattenUnflattenRoundTrip:
+    def test_round_trip_simple(self):
+        original: OmniPayload = {
+            "codes": {"audio": torch.tensor([1.0, 2.0])},
+            "meta": {"finished": torch.tensor(True, dtype=torch.bool), "left_context_size": 10},
+            "ids": {"prompt": [1, 2, 3]},
+            "latent": torch.tensor([5.0]),
+        }
+        restored = unflatten_payload(flatten_payload(original))
+        assert torch.equal(restored["codes"]["audio"], original["codes"]["audio"])
+        assert restored["meta"]["finished"].item() is True
+        assert restored["meta"]["left_context_size"] == 10
+        assert restored["ids"]["prompt"] == [1, 2, 3]
+        assert torch.equal(restored["latent"], original["latent"])
+
+    def test_round_trip_with_layers(self):
+        original = {
+            "hidden_states": {
+                "output": torch.tensor([1.0]),
+                "layers": {0: torch.tensor([2.0]), 24: torch.tensor([3.0])},
+            },
+        }
+        restored = unflatten_payload(flatten_payload(original))
+        assert torch.equal(restored["hidden_states"]["output"], torch.tensor([1.0]))
+        assert torch.equal(restored["hidden_states"]["layers"][0], torch.tensor([2.0]))
+        assert torch.equal(restored["hidden_states"]["layers"][24], torch.tensor([3.0]))
+
+    def test_round_trip_all_categories(self):
+        original: OmniPayload = {
+            "hidden_states": {"output": torch.tensor([1.0]), "last": torch.tensor([2.0])},
+            "embed": {"prefill": torch.tensor([3.0]), "tts_bos": torch.tensor([4.0])},
+            "codes": {"audio": torch.tensor([5.0]), "ref": torch.tensor([6.0])},
+            "ids": {"all": [1, 2], "prompt": [3, 4]},
+            "meta": {"finished": torch.tensor(False, dtype=torch.bool), "ar_width": 8},
+        }
+        restored = unflatten_payload(flatten_payload(original))
+        assert torch.equal(restored["hidden_states"]["output"], torch.tensor([1.0]))
+        assert torch.equal(restored["hidden_states"]["last"], torch.tensor([2.0]))
+        assert torch.equal(restored["embed"]["prefill"], torch.tensor([3.0]))
+        assert torch.equal(restored["embed"]["tts_bos"], torch.tensor([4.0]))
+        assert torch.equal(restored["codes"]["audio"], torch.tensor([5.0]))
+        assert torch.equal(restored["codes"]["ref"], torch.tensor([6.0]))
+        assert restored["ids"]["all"] == [1, 2]
+        assert restored["ids"]["prompt"] == [3, 4]
+        assert restored["meta"]["finished"].item() is False
+        assert restored["meta"]["ar_width"] == 8
+
+
+class TestSerializeDeserializePayload:
+    def test_tensor_round_trip(self):
+        original: OmniPayload = {
+            "hidden_states": {"output": torch.tensor([[1.0, 2.0], [3.0, 4.0]])},
+        }
+        wire = serialize_payload(original)
+        assert isinstance(wire, AdditionalInformationPayload)
+        restored = deserialize_payload(wire)
+        assert torch.equal(restored["hidden_states"]["output"], original["hidden_states"]["output"])
+
+    def test_list_round_trip(self):
+        original: OmniPayload = {
+            "ids": {"prompt": [10, 20, 30]},
+        }
+        wire = serialize_payload(original)
+        restored = deserialize_payload(wire)
+        assert restored["ids"]["prompt"] == [10, 20, 30]
+
+    def test_finished_tensor_round_trip(self):
+        original: OmniPayload = {
+            "meta": {"finished": torch.tensor(True, dtype=torch.bool), "left_context_size": 5},
+        }
+        wire = serialize_payload(original)
+        restored = deserialize_payload(wire)
+        assert isinstance(restored["meta"]["finished"], torch.Tensor)
+        assert restored["meta"]["finished"].dtype == torch.bool
+        assert restored["meta"]["finished"].item() is True
+        assert restored["meta"]["left_context_size"] == 5
+
+    def test_mixed_types_round_trip(self):
+        original: OmniPayload = {
+            "hidden_states": {"output": torch.tensor([1.0, 2.0])},
+            "ids": {"all": [1, 2, 3]},
+            "meta": {"finished": torch.tensor(False, dtype=torch.bool), "ar_width": 4},
+            "codes": {"audio": torch.tensor([3.0]), "ref": torch.tensor([0.1, 0.2])},
+        }
+        wire = serialize_payload(original)
+        restored = deserialize_payload(wire)
+        assert torch.equal(restored["hidden_states"]["output"], original["hidden_states"]["output"])
+        assert restored["ids"]["all"] == [1, 2, 3]
+        assert restored["meta"]["finished"].item() is False
+        assert restored["meta"]["ar_width"] == 4
+        assert torch.equal(restored["codes"]["audio"], original["codes"]["audio"])
+        assert torch.equal(restored["codes"]["ref"], original["codes"]["ref"])
+
+    def test_hidden_states_layers_round_trip(self):
+        original = {
+            "hidden_states": {
+                "output": torch.tensor([1.0]),
+                "layers": {0: torch.tensor([2.0]), 24: torch.tensor([3.0])},
+            },
+        }
+        wire = serialize_payload(original)
+        restored = deserialize_payload(wire)
+        assert torch.equal(restored["hidden_states"]["output"], torch.tensor([1.0]))
+        assert torch.equal(restored["hidden_states"]["layers"][0], torch.tensor([2.0]))
+        assert torch.equal(restored["hidden_states"]["layers"][24], torch.tensor([3.0]))
+
+    def test_tensor_dtype_preserved(self):
+        # bfloat16 excluded: numpy() doesn't support it; callers must cast before serializing.
+        for dtype in [torch.float16, torch.float32, torch.int64, torch.int32, torch.bool]:
+            original: OmniPayload = {"codes": {"audio": torch.tensor([1], dtype=dtype)}}
+            wire = serialize_payload(original)
+            restored = deserialize_payload(wire)
+            assert restored["codes"]["audio"].dtype == dtype, f"dtype mismatch for {dtype}"
+
+    def test_tensor_shape_preserved(self):
+        t = torch.randn(3, 4, 5)
+        original: OmniPayload = {"hidden_states": {"output": t}}
+        wire = serialize_payload(original)
+        restored = deserialize_payload(wire)
+        assert restored["hidden_states"]["output"].shape == (3, 4, 5)
+        assert torch.allclose(restored["hidden_states"]["output"], t)
+
+    def test_empty_payload_returns_none(self):
+        assert serialize_payload({}) is None
+
+    def test_none_values_skipped(self):
+        original: OmniPayload = {"meta": {"finished": None}}
+        wire = serialize_payload(original)
+        assert wire is None

tests/engine/test_omni_request.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for OmniRequest substitutability with the base Request.
+
+vllm-omni rebinds ``vllm.v1.request.Request`` to ``OmniRequest`` at import
+time (see ``vllm_omni/patch.py``). vLLM core constructs ``Request`` positionally
+in some paths (notably ``vllm/v1/worker/gpu/warmup.py::warmup_kernels`` for V2
+model-runner architectures like Qwen3ForCausalLM). The omni-specific params
+must therefore be keyword-only so positional construction stays Liskov-
+substitutable with the base class — including base-style calls that pass
+``prompt_embeds`` (itself a positional-capable base param) positionally.
+"""
+
+import inspect
+
+import numpy as np
+import pytest
+import torch
+from vllm.sampling_params import SamplingParams
+from vllm.v1.request import Request
+
+from vllm_omni.engine import PromptEmbedsPayload
+from vllm_omni.request import OmniRequest, OmniStreamingUpdate
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_omni_params_are_keyword_only():
+    """The three omni params must be keyword-only after ``*args``.
+
+    Guards against re-introducing the import-time rebind bug where omni params
+    came first positionally and broke positional ``Request(...)`` construction.
+    """
+    params = inspect.signature(OmniRequest.__init__).parameters
+    for name in (
+        "prompt_embeds",
+        "external_req_id",
+        "additional_information",
+        "model_intermediate_buffer",
+    ):
+        assert params[name].kind is inspect.Parameter.KEYWORD_ONLY, name
+
+
+def test_positional_construction_matches_base_request():
+    """Reproduces the warmup_kernels call: Request(id, tokens, sp, pp)."""
+    req = OmniRequest("req-0", [1, 2, 3], SamplingParams(), None)
+
+    assert isinstance(req, Request)
+    assert req.request_id == "req-0"
+    assert req.prompt_token_ids == [1, 2, 3]
+    # Omni params default cleanly when constructed positionally.
+    assert req.external_req_id is None
+    assert req.additional_information is None
+    assert req.prompt_embeds_payload is None
+
+
+def test_positional_prompt_embeds_does_not_collide():
+    """Base-style positional call that includes ``prompt_embeds``.
+
+    ``prompt_embeds`` is a positional-capable base param (after
+    ``arrival_time``). A base-style call must not trip
+    ``got multiple values for argument 'prompt_embeds'`` — the subclass must
+    not also inject its keyword override when the value arrived positionally.
+    Positional order: request_id, prompt_token_ids, sampling_params,
+    pooling_params, client_index, arrival_time, prompt_embeds.
+    """
+    embeds = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+    req = OmniRequest("req-pe", [1, 2], SamplingParams(), None, 0, None, embeds)
+
+    assert isinstance(req, Request)
+    assert torch.equal(req.prompt_embeds, embeds)
+    # Positional tensor is not a serialized payload.
+    assert req.prompt_embeds_payload is None
+    assert req.external_req_id is None
+
+
+def test_keyword_omni_params_round_trip():
+    """Keyword omni params are preserved; serialized embeds are decoded."""
+    arr = np.arange(6, dtype=np.float32).reshape(2, 3)
+    payload = PromptEmbedsPayload(data=arr.tobytes(), shape=[2, 3], dtype="float32")
+
+    req = OmniRequest(
+        request_id="req-1",
+        prompt_token_ids=[7, 8],
+        sampling_params=SamplingParams(),
+        pooling_params=None,
+        prompt_embeds=payload,
+        external_req_id="ext-1",
+    )
+
+    assert req.external_req_id == "ext-1"
+    # The serialized payload is retained and decoded into a tensor on the base.
+    assert req.prompt_embeds_payload is payload
+    assert torch.equal(req.prompt_embeds, torch.from_numpy(arr))
+
+
+def test_model_intermediate_buffer_round_trips_to_streaming_update():
+    buffer = {"duplex": {"turn_id": 3}}
+    request = OmniRequest(
+        request_id="req-model-buffer",
+        prompt_token_ids=[1, 2],
+        sampling_params=SamplingParams(max_tokens=16),
+        pooling_params=None,
+        resumable=True,
+        model_intermediate_buffer=buffer,
+    )
+
+    update = OmniStreamingUpdate.from_request(request)
+
+    assert update is not None
+    assert request.model_intermediate_buffer is buffer
+    assert update.model_intermediate_buffer is buffer

tests/engine/test_omni_request_output.py

@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for OmniRequestOutput class."""
+
+import pytest
+from PIL import Image
+from pytest_mock import MockerFixture
+
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestOmniRequestOutput:
+    """Tests for OmniRequestOutput class."""
+
+    def test_from_diffusion(self):
+        """Test creating output from diffusion model."""
+        images = [Image.new("RGB", (64, 64), color="red")]
+        output = OmniRequestOutput.from_diffusion(
+            request_id="test-123",
+            images=images,
+            prompt="a cat",
+            metrics={"steps": 50},
+        )
+        assert output.request_id == "test-123"
+        assert output.images == images
+        assert output.prompt == "a cat"
+        assert output.metrics == {"steps": 50}
+        assert output.is_diffusion_output
+        assert output.num_images == 1
+
+    def test_from_pipeline(self, mocker: MockerFixture):
+        """Test creating output from pipeline stage."""
+        mock_request_output = mocker.Mock()
+        mock_request_output.request_id = "pipeline-123"
+        mock_request_output.prompt_token_ids = [1, 2, 3]
+        mock_request_output.outputs = [mocker.Mock()]
+        mock_request_output.encoder_prompt_token_ids = None
+        mock_request_output.prompt_logprobs = None
+        mock_request_output.num_cached_tokens = 10
+        mock_request_output.kv_transfer_params = None
+        mock_request_output.multimodal_output = {"image": mocker.Mock()}
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+
+        assert output.request_id == "pipeline-123"
+        assert output.stage_id == 0
+        assert output.final_output_type == "text"
+        assert output.is_pipeline_output
+
+    def test_prompt_token_ids_property(self, mocker: MockerFixture):
+        """Test prompt_token_ids property for streaming compatibility."""
+        mock_request_output = mocker.Mock()
+        mock_request_output.prompt_token_ids = [1, 2, 3, 4, 5]
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+
+        assert output.prompt_token_ids == [1, 2, 3, 4, 5]
+
+    def test_prompt_token_ids_none_when_no_request_output(self):
+        """Test prompt_token_ids returns None when no request_output."""
+        output = OmniRequestOutput.from_diffusion(
+            request_id="test-123",
+            images=[],
+            prompt="a cat",
+        )
+        assert output.prompt_token_ids is None
+
+    def test_outputs_property(self, mocker: MockerFixture):
+        """Test outputs property for chat completion compatibility."""
+        mock_output = mocker.Mock()
+        mock_request_output = mocker.Mock()
+        mock_request_output.outputs = [mock_output]
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+
+        assert output.outputs == [mock_output]
+
+    def test_outputs_empty_when_no_request_output(self):
+        """Test outputs returns empty list when no request_output."""
+        output = OmniRequestOutput.from_diffusion(
+            request_id="test-123",
+            images=[],
+            prompt="a cat",
+        )
+        assert output.outputs == []
+
+    def test_encoder_prompt_token_ids_property(self, mocker: MockerFixture):
+        """Test encoder_prompt_token_ids property."""
+        mock_request_output = mocker.Mock()
+        mock_request_output.encoder_prompt_token_ids = [10, 20, 30]
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+
+        assert output.encoder_prompt_token_ids == [10, 20, 30]
+
+    def test_num_cached_tokens_property(self, mocker: MockerFixture):
+        """Test num_cached_tokens property."""
+        mock_request_output = mocker.Mock()
+        mock_request_output.num_cached_tokens = 42
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+
+        assert output.num_cached_tokens == 42
+
+    def test_multimodal_output_property(self, mocker: MockerFixture):
+        """Test multimodal_output property."""
+        mock_request_output = mocker.Mock()
+        mock_audio = mocker.Mock()
+        expected_output = {"audio": mock_audio}
+        mock_request_output.outputs = []
+        mock_request_output.multimodal_output = expected_output
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="audio",
+            request_output=mock_request_output,
+        )
+
+        assert output.multimodal_output is expected_output
+
+    def test_to_dict_diffusion(self):
+        """Test to_dict for diffusion output."""
+        output = OmniRequestOutput.from_diffusion(
+            request_id="test-123",
+            images=[Image.new("RGB", (64, 64), color="red")],
+            prompt="a cat",
+            metrics={"steps": 50},
+        )
+        result = output.to_dict()
+
+        assert result["request_id"] == "test-123"
+        assert result["finished"] is True
+        assert result["final_output_type"] == "image"
+        assert result["num_images"] == 1
+        assert result["prompt"] == "a cat"
+
+    def test_to_dict_pipeline(self, mocker: MockerFixture):
+        """Test to_dict for pipeline output."""
+        mock_request_output = mocker.Mock()
+        mock_request_output.request_id = "pipeline-123"
+
+        output = OmniRequestOutput.from_pipeline(
+            stage_id=0,
+            final_output_type="text",
+            request_output=mock_request_output,
+        )
+        result = output.to_dict()
+
+        assert result["request_id"] == "pipeline-123"
+        assert result["finished"] is True
+        assert result["final_output_type"] == "text"
+        assert result["stage_id"] == 0

tests/engine/test_orchestrator_error_handling.py

@@ -73,9 +73,17 @@ def orchestrator_factory():
 
 
 class FakeDeadLLMStageClient(FakeStageClient):
-    """LLM stage client that raises EngineDeadError on get_output_async."""
+    """LLM stage client that raises EngineDeadError once a request is in-flight.
+
+    Raising on the first poll races with ``add_request``: under suite load the
+    orchestrator can die and drain before the test enqueues anything, so no
+    fatal ErrorMessage is ever published. Keep returning empty outputs until
+    ``add_request_async`` has been called, then fail on the next poll.
+    """
 
     async def get_output_async(self):
+        if not self.add_request_calls:
+            return SimpleNamespace(outputs=[])
         raise EngineDeadError("Stage-0 engine core is dead")