[Entrypoint][Refactor] Make field type hint more concrete (#3139)

Signed-off-by: wuhang <wuhang6@huawei.com>

tests/engine/test_async_omni_engine_input.py

@@ -4,7 +4,7 @@ from vllm.sampling_params import SamplingParams
 from vllm.v1.engine import EngineCoreRequest
 
 from vllm_omni.engine import OmniEngineCoreRequest
-from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -27,7 +27,7 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock
     engine = object.__new__(AsyncOmniEngine)
     params = SamplingParams(max_tokens=8)
     engine.default_sampling_params_list = [params]
-    engine.stage_metadata = [{"stage_type": "llm"}]
+    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
     engine.supported_tasks = ("speech",)
 
     input_processor = mocker.Mock()
@@ -53,7 +53,7 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock
         arrival_time=0.0,
     )
 
-    request = msg["prompt"]
+    request = msg.prompt
     assert isinstance(request, OmniEngineCoreRequest)
     assert request.external_req_id == "req-1"
     assert request.additional_information is not None
@@ -66,7 +66,7 @@ def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixtur
     engine = object.__new__(AsyncOmniEngine)
     params = SamplingParams(max_tokens=8)
     engine.default_sampling_params_list = [params]
-    engine.stage_metadata = [{"stage_type": "llm"}]
+    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
     engine.supported_tasks = ("generate",)
 
     input_processor = mocker.Mock()
@@ -85,6 +85,6 @@ def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixtur
         message_type="streaming_update",
     )
 
-    assert msg["type"] == "streaming_update"
+    assert msg.type == "streaming_update"
     input_processor.process_inputs.assert_called_once()
     assert input_processor.process_inputs.call_args.kwargs["resumable"] is True

tests/engine/test_async_omni_engine_outputs.py

@@ -10,6 +10,8 @@ import pytest
 from pytest_mock import MockerFixture
 
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.engine.messages import ErrorMessage, OutputMessage
+from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -30,14 +32,19 @@ def test_try_get_output_raises_after_orchestrator_dies(mocker: MockerFixture):
     mock_queue = mocker.MagicMock()
     # First call succeeds; second call finds the queue empty.
     mock_queue.sync_q.get.side_effect = [
-        {"type": "output", "request_id": "r1"},
+        OutputMessage(
+            request_id="r1",
+            stage_id=0,
+            engine_outputs=OmniRequestOutput(request_id="r1"),
+            finished=False,
+        ),
         queue.Empty,
     ]
 
     engine = _make_engine(mock_queue, mocker, thread_alive=True)
 
     # Collect the one buffered result.
-    assert engine.try_get_output()["request_id"] == "r1"
+    assert engine.try_get_output().request_id == "r1"
 
     # Orchestrator thread crashes between polls.
     engine.orchestrator_thread.is_alive.return_value = False
@@ -51,13 +58,18 @@ async def test_try_get_output_async_raises_after_orchestrator_dies(mocker: Mocke
     """Same scenario as above but for the async variant."""
     mock_queue = mocker.MagicMock()
     mock_queue.sync_q.get_nowait.side_effect = [
-        {"type": "output", "request_id": "r1"},
+        OutputMessage(
+            request_id="r1",
+            stage_id=0,
+            engine_outputs=OmniRequestOutput(request_id="r1"),
+            finished=False,
+        ),
         queue.Empty,
     ]
 
     engine = _make_engine(mock_queue, mocker, thread_alive=True)
 
-    assert (await engine.try_get_output_async())["request_id"] == "r1"
+    assert (await engine.try_get_output_async()).request_id == "r1"
 
     engine.orchestrator_thread.is_alive.return_value = False
 
@@ -71,7 +83,7 @@ def test_fatal_error_message_surfaces_through_try_get_output(mocker: MockerFixtu
     ``try_get_output`` must return this message so the caller
     (``OmniBase._handle_output_message``) can detect the fatal flag.
     """
-    fatal_msg = {"type": "error", "error": "Orchestrator thread crashed", "fatal": True}
+    fatal_msg = ErrorMessage(error="Orchestrator thread crashed", fatal=True)
 
     mock_queue = mocker.MagicMock()
     mock_queue.sync_q.get.return_value = fatal_msg
@@ -80,15 +92,15 @@ def test_fatal_error_message_surfaces_through_try_get_output(mocker: MockerFixtu
 
     msg = engine.try_get_output()
     assert msg is not None
-    assert msg["type"] == "error"
-    assert msg["fatal"] is True
-    assert "crashed" in msg["error"]
+    assert msg.type == "error"
+    assert msg.fatal is True
+    assert "crashed" in msg.error
 
 
 @pytest.mark.asyncio
 async def test_fatal_error_message_surfaces_through_try_get_output_async(mocker: MockerFixture):
     """Async variant of the fatal error message test."""
-    fatal_msg = {"type": "error", "error": "Orchestrator thread crashed", "fatal": True}
+    fatal_msg = ErrorMessage(error="Orchestrator thread crashed", fatal=True)
 
     mock_queue = mocker.MagicMock()
     mock_queue.sync_q.get_nowait.return_value = fatal_msg
@@ -97,5 +109,5 @@ async def test_fatal_error_message_surfaces_through_try_get_output_async(mocker:
 
     msg = await engine.try_get_output_async()
     assert msg is not None
-    assert msg["type"] == "error"
-    assert msg["fatal"] is True
+    assert msg.type == "error"
+    assert msg.fatal is True

tests/engine/test_async_omni_engine_stage_init.py

@@ -7,7 +7,7 @@ import types
 import pytest
 
 from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
-from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo
 from vllm_omni.engine.stage_init_utils import (
     LogicalStageInitPlan,
     ReplicaInitPlan,
@@ -332,8 +332,8 @@ def test_initialize_stages_exposes_logical_stage_views_and_builds_top_level_inpu
         stage1_client_r0.default_sampling_params,
     ]
     assert engine.stage_metadata == [
-        {"final_output": False, "final_output_type": None, "stage_type": "llm"},
-        {"final_output": True, "final_output_type": None, "stage_type": "llm"},
+        StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm"),
+        StageRuntimeInfo(final_output=True, final_output_type=None, stage_type="llm"),
     ]
 
 

tests/engine/test_orchestrator.py

@@ -17,6 +17,15 @@ from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.core_client import AsyncMPClient
 
+from vllm_omni.engine.messages import (
+    AbortRequestMessage,
+    AddCompanionRequestMessage,
+    CollectiveRPCRequestMessage,
+    CollectiveRPCResultMessage,
+    OutputMessage,
+    ShutdownRequestMessage,
+    StageSubmissionMessage,
+)
 from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
 from vllm_omni.engine.stage_pool import StagePool
@@ -44,14 +53,27 @@ class FakeStageClient:
         final_output: bool = False,
         final_output_type: str = "text",
         next_inputs: list[dict] | None = None,
+        engine_input_source: list[int] | None = None,
+        is_comprehension: bool = False,
+        model_stage: str | None = None,
+        kv_sender_info: dict[str, Any] | None = None,
     ) -> None:
+        self.stage_id = 0
+        self.replica_id = 0
         self.stage_type = stage_type
         self.final_output = final_output
         self.final_output_type = final_output_type
+        self.default_sampling_params = SamplingParams(max_tokens=1)
+        self.requires_multimodal_data = False
+        self.engine_input_source = list(engine_input_source or [0])
+        self.is_comprehension = is_comprehension
+        self.model_stage = model_stage
         self.next_inputs = list(next_inputs or [])
         self.custom_process_input_func = None
+        self._kv_sender_info = dict(kv_sender_info) if kv_sender_info is not None else None
         self.add_request_calls: list[tuple] = []
         self.abort_calls: list[list[str]] = []
+        self.collective_rpc_calls: list[tuple[str, float | None, tuple[Any, ...], dict[str, Any]]] = []
         self.shutdown_calls = 0
         self._engine_core_outputs = queue.Queue()
         self._diffusion_outputs = queue.Queue()
@@ -81,6 +103,30 @@ class FakeStageClient:
     async def abort_requests_async(self, request_ids: list[str]) -> None:
         self.abort_calls.append(list(request_ids))
 
+    async def collective_rpc_async(
+        self,
+        *,
+        method: str,
+        timeout: float | None = None,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> Any:
+        normalized_kwargs = dict(kwargs or {})
+        self.collective_rpc_calls.append((method, timeout, args, normalized_kwargs))
+        return {
+            "supported": False,
+            "todo": True,
+            "reason": f"{self.__class__.__name__}.collective_rpc_async is not implemented yet",
+        }
+
+    def get_kv_sender_info(self) -> dict[str, Any] | None:
+        if self._kv_sender_info is None:
+            return None
+        return dict(self._kv_sender_info)
+
+    def check_health(self) -> None:
+        return None
+
     def shutdown(self) -> None:
         self.shutdown_calls += 1
 
@@ -96,7 +142,6 @@ class FakeCollectiveRpcStageClient(FakeStageClient):
     def __init__(self, *args, rpc_result: Any = None, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self.rpc_result = rpc_result
-        self.collective_rpc_calls: list[tuple[str, float | None, tuple[Any, ...], dict[str, Any]]] = []
 
     async def collective_rpc_async(
         self,
@@ -318,7 +363,7 @@ def _build_harness(
 
 
 async def _shutdown_orchestrator(orchestrator_fixture: OrchestratorFixture) -> None:
-    orchestrator_fixture.request_sync_q.put_nowait({"type": "shutdown"})
+    orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
     await asyncio.to_thread(orchestrator_fixture.thread.join, 5)
     if orchestrator_fixture.thread.is_alive():
         raise AssertionError("Timed out waiting for orchestrator thread shutdown")
@@ -333,7 +378,7 @@ async def _wait_for(predicate, *, timeout: float = 2.0) -> None:
         await asyncio.sleep(0.01)
 
 
-async def _get_output_message(orchestrator_fixture: OrchestratorFixture, *, timeout: float = 2.0) -> dict:
+async def _get_output_message(orchestrator_fixture: OrchestratorFixture, *, timeout: float = 2.0) -> OutputMessage:
     deadline = time.monotonic() + timeout
     while True:
         if time.monotonic() >= deadline:
@@ -343,11 +388,15 @@ async def _get_output_message(orchestrator_fixture: OrchestratorFixture, *, time
         except queue.Empty:
             await asyncio.sleep(0.01)
             continue
-        if msg.get("type") == "output":
+        if isinstance(msg, OutputMessage):
             return msg
 
 
-async def _get_rpc_message(orchestrator_fixture: OrchestratorFixture, *, timeout: float = 2.0) -> dict:
+async def _get_rpc_message(
+    orchestrator_fixture: OrchestratorFixture,
+    *,
+    timeout: float = 2.0,
+) -> CollectiveRPCResultMessage:
     deadline = time.monotonic() + timeout
     rpc_sync_q = orchestrator_fixture.queues[2].sync_q
     while True:
@@ -369,24 +418,22 @@ async def _enqueue_add_request(
     final_stage_id: int,
 ) -> None:
     orchestrator_fixture.request_sync_q.put_nowait(
-        {
-            "type": "add_request",
-            "request_id": request_id,
-            "prompt": prompt,
-            "original_prompt": original_prompt,
-            "sampling_params_list": sampling_params_list,
-            "final_stage_id": final_stage_id,
-        }
+        StageSubmissionMessage(
+            type="add_request",
+            request_id=request_id,
+            prompt=prompt,
+            original_prompt=original_prompt,
+            output_prompt_text=None,
+            sampling_params_list=sampling_params_list,
+            final_stage_id=final_stage_id,
+            preprocess_ms=0.0,
+            enqueue_ts=time.perf_counter(),
+        )
     )
 
 
 async def _enqueue_abort_request(orchestrator_fixture: OrchestratorFixture, request_ids: list[str]) -> None:
-    orchestrator_fixture.request_sync_q.put_nowait(
-        {
-            "type": "abort",
-            "request_ids": request_ids,
-        }
-    )
+    orchestrator_fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=request_ids))
 
 
 def test_stage_engine_core_client_shutdown_cleans_children_if_base_shutdown_fails(monkeypatch):
@@ -457,7 +504,7 @@ def orchestrator_factory():
 
     for fixture in fixtures:
         if fixture.thread.is_alive():
-            fixture.request_sync_q.put_nowait({"type": "shutdown"})
+            fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
             fixture.thread.join(timeout=5)
         for q in fixture.queues:
             q.close()
@@ -505,10 +552,10 @@ async def test_run_two_stage_llm(orchestrator_factory) -> None:
 
         output_msg = await _get_output_message(orchestrator_fixture)
 
-        assert output_msg["request_id"] == "req-llm"
-        assert output_msg["stage_id"] == 1
-        assert output_msg["finished"] is True
-        assert output_msg["engine_outputs"].request_id == "req-llm"
+        assert output_msg.request_id == "req-llm"
+        assert output_msg.stage_id == 1
+        assert output_msg.finished is True
+        assert output_msg.engine_outputs.request_id == "req-llm"
         assert "req-llm" not in orchestrator_fixture.orchestrator.request_states
     finally:
         await _shutdown_orchestrator(orchestrator_fixture)
@@ -541,10 +588,10 @@ async def test_run_single_stage_diffusion(orchestrator_factory) -> None:
 
         output_msg = await _get_output_message(orchestrator_fixture)
 
-        assert output_msg["request_id"] == "req-diff"
-        assert output_msg["stage_id"] == 0
-        assert output_msg["finished"] is True
-        assert output_msg["engine_outputs"].request_id == "req-diff"
+        assert output_msg.request_id == "req-diff"
+        assert output_msg.stage_id == 0
+        assert output_msg.finished is True
+        assert output_msg.engine_outputs.request_id == "req-diff"
         assert "req-diff" not in orchestrator_fixture.orchestrator.request_states
     finally:
         await _shutdown_orchestrator(orchestrator_fixture)
@@ -589,10 +636,10 @@ async def test_run_llm_to_diffusion(orchestrator_factory) -> None:
 
         output_msg = await _get_output_message(orchestrator_fixture)
 
-        assert output_msg["request_id"] == "req-img"
-        assert output_msg["stage_id"] == 1
-        assert output_msg["finished"] is True
-        assert output_msg["engine_outputs"].request_id == "req-img"
+        assert output_msg.request_id == "req-img"
+        assert output_msg.stage_id == 1
+        assert output_msg.finished is True
+        assert output_msg.engine_outputs.request_id == "req-img"
         assert "req-img" not in orchestrator_fixture.orchestrator.request_states
     finally:
         await _shutdown_orchestrator(orchestrator_fixture)
@@ -633,9 +680,9 @@ async def test_run_async_chunk(orchestrator_factory) -> None:
 
         output_msg = await _get_output_message(orchestrator_fixture)
 
-        assert output_msg["request_id"] == "req-async"
-        assert output_msg["stage_id"] == 1
-        assert output_msg["finished"] is True
+        assert output_msg.request_id == "req-async"
+        assert output_msg.stage_id == 1
+        assert output_msg.finished is True
         assert "req-async" not in orchestrator_fixture.orchestrator.request_states
     finally:
         await _shutdown_orchestrator(orchestrator_fixture)
@@ -757,9 +804,9 @@ async def test_multi_replica_round_robin_distribution(orchestrator_factory) -> N
         stage1.push_engine_core_outputs(_engine_core_outputs("s1-raw", 2.0))
         output_msg = await _get_output_message(orchestrator_fixture)
 
-        assert output_msg["request_id"] == "req-0"
-        assert output_msg["stage_id"] == 1
-        assert output_msg["finished"] is True
+        assert output_msg.request_id == "req-0"
+        assert output_msg.stage_id == 1
+        assert output_msg.finished is True
         assert "req-0" not in orchestrator_fixture.orchestrator.request_states
     finally:
         await _shutdown_orchestrator(orchestrator_fixture)
@@ -927,12 +974,17 @@ async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> Non
     orchestrator.stage_pools = [pool]
 
     await orchestrator._handle_streaming_update(
-        {
-            "request_id": "req-stream",
-            "prompt": SimpleNamespace(request_id="req-stream", prompt_token_ids=[1], resumable=True),
-            "sampling_params_list": [_sampling_params()],
-            "output_prompt_text": "segment-2",
-        }
+        StageSubmissionMessage(
+            type="streaming_update",
+            request_id="req-stream",
+            prompt=SimpleNamespace(request_id="req-stream", prompt_token_ids=[1], resumable=True),
+            original_prompt={"prompt": "segment-2"},
+            output_prompt_text="segment-2",
+            sampling_params_list=[_sampling_params()],
+            final_stage_id=0,
+            preprocess_ms=0.0,
+            enqueue_ts=time.perf_counter(),
+        )
     )
 
     assert pool.calls == [("req-stream", "segment-2")]
@@ -1024,12 +1076,14 @@ async def test_collective_rpc_ignores_invalid_stage_ids(orchestrator_factory, ca
         target_logger.setLevel(logging.WARNING)
         try:
             orchestrator_fixture.request_sync_q.put_nowait(
-                {
-                    "type": "collective_rpc",
-                    "rpc_id": "rpc-1",
-                    "method": "list_loras",
-                    "stage_ids": [99, 1],
-                }
+                CollectiveRPCRequestMessage(
+                    rpc_id="rpc-1",
+                    method="list_loras",
+                    timeout=None,
+                    args=(),
+                    kwargs={},
+                    stage_ids=[99, 1],
+                )
             )
 
             msg = await _get_rpc_message(orchestrator_fixture)
@@ -1037,10 +1091,10 @@ async def test_collective_rpc_ignores_invalid_stage_ids(orchestrator_factory, ca
             target_logger.removeHandler(caplog.handler)
             target_logger.setLevel(prev_level)
 
-        assert msg["type"] == "collective_rpc_result"
-        assert msg["rpc_id"] == "rpc-1"
-        assert msg["stage_ids"] == [1]
-        assert msg["results"] == [{"stage": 1}]
+        assert msg.type == "collective_rpc_result"
+        assert msg.rpc_id == "rpc-1"
+        assert msg.stage_ids == [1]
+        assert msg.results == [{"stage": 1}]
         assert not stage0.collective_rpc_calls
         assert len(stage1.collective_rpc_calls) == 1
         assert "collective_rpc: ignoring invalid stage_id 99" in caplog.text
@@ -1084,15 +1138,14 @@ async def test_multi_replica_cfg_companion_inherits_parent_affinity(orchestrator
         await _wait_for(lambda: len(stage0_r1.add_request_calls) == 1)
 
         orchestrator_fixture.request_sync_q.put_nowait(
-            {
-                "type": "add_companion_request",
-                "companion_id": "parent-neg",
-                "parent_id": "parent",
-                "role": "negative",
-                "prompt": SimpleNamespace(request_id="parent-neg", prompt_token_ids=[9]),
-                "companion_prompt_text": {"prompt": "negative"},
-                "sampling_params_list": [_sampling_params()],
-            }
+            AddCompanionRequestMessage(
+                companion_id="parent-neg",
+                parent_id="parent",
+                role="negative",
+                prompt=SimpleNamespace(request_id="parent-neg", prompt_token_ids=[9]),
+                companion_prompt_text={"prompt": "negative"},
+                sampling_params_list=[_sampling_params()],
+            )
         )
         await _wait_for(lambda: len(stage0_r1.add_request_calls) == 2)
 

tests/engine/test_orchestrator_error_handling.py

@@ -17,6 +17,7 @@ from types import SimpleNamespace
 import pytest
 from vllm.v1.engine.exceptions import EngineDeadError
 
+from vllm_omni.engine.messages import EngineQueueMessage, ErrorMessage, ShutdownRequestMessage
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -37,7 +38,7 @@ def _sampling_params(max_tokens: int = 4):
     return SamplingParams(max_tokens=max_tokens)
 
 
-async def _get_any_output_message(fixture: OrchestratorFixture, *, timeout: float = 2.0) -> dict:
+async def _get_any_output_message(fixture: OrchestratorFixture, *, timeout: float = 2.0) -> EngineQueueMessage:
     """Like _get_output_message but returns any message type (including errors)."""
     deadline = time.monotonic() + timeout
     while True:
@@ -62,7 +63,7 @@ def orchestrator_factory():
 
     for fixture in fixtures:
         if fixture.thread.is_alive():
-            fixture.request_sync_q.put_nowait({"type": "shutdown"})
+            fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
             fixture.thread.join(timeout=5)
         for q in fixture.queues:
             q.close()
@@ -101,10 +102,11 @@ async def test_engine_dead_error_broadcasts_fatal_and_shuts_down(orchestrator_fa
         # Collect the fatal error message.
         msg = await _get_any_output_message(orchestrator_fixture)
 
-        assert msg["type"] == "error"
-        assert msg["fatal"] is True
-        assert msg["request_id"] == "req-dead"
-        assert "Stage-0 engine core is dead" in msg["error"]
+        assert isinstance(msg, ErrorMessage)
+        assert msg.type == "error"
+        assert msg.fatal is True
+        assert msg.request_id == "req-dead"
+        assert "Stage-0 engine core is dead" in msg.error
 
         # The orchestrator thread should exit after the fatal error.
         orchestrator_fixture.thread.join(timeout=5)
@@ -114,7 +116,7 @@ async def test_engine_dead_error_broadcasts_fatal_and_shuts_down(orchestrator_fa
         assert "req-dead" not in orchestrator_fixture.orchestrator.request_states
     finally:
         if orchestrator_fixture.thread.is_alive():
-            orchestrator_fixture.request_sync_q.put_nowait({"type": "shutdown"})
+            orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
             orchestrator_fixture.thread.join(timeout=5)
 
 
@@ -148,13 +150,14 @@ async def test_diffusion_error_output_routed_as_finished(orchestrator_factory) -
 
         msg = await _get_any_output_message(orchestrator_fixture)
 
-        assert msg["type"] == "error"
-        assert msg["request_id"] == "req-err"
-        assert msg["stage_id"] == 0
-        assert msg["error"] == "gpu fault"
+        assert isinstance(msg, ErrorMessage)
+        assert msg.type == "error"
+        assert msg.request_id == "req-err"
+        assert msg.stage_id == 0
+        assert msg.error == "gpu fault"
 
         # Request state should be cleaned up.
         await _wait_for(lambda: "req-err" not in orchestrator_fixture.orchestrator.request_states)
     finally:
-        orchestrator_fixture.request_sync_q.put_nowait({"type": "shutdown"})
+        orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
         orchestrator_fixture.thread.join(timeout=5)

tests/engine/test_orchestrator_kv_sender_info.py

@@ -5,6 +5,7 @@ import pytest
 from vllm import SamplingParams
 
 from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
+from vllm_omni.engine.messages import OutputMessage
 from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
 from vllm_omni.engine.stage_pool import StagePool
@@ -225,11 +226,12 @@ def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs():
     assert diffusion_stage.calls == []
     assert len(orchestrator.output_async_queue.items) == 1
     terminal_msg = orchestrator.output_async_queue.items[0]
-    assert terminal_msg["type"] == "output"
-    assert terminal_msg["request_id"] == "req-empty"
-    assert terminal_msg["stage_id"] == 1
-    assert terminal_msg["finished"] is True
-    assert "produced no valid inputs" in terminal_msg["engine_outputs"].error
+    assert isinstance(terminal_msg, OutputMessage)
+    assert terminal_msg.type == "output"
+    assert terminal_msg.request_id == "req-empty"
+    assert terminal_msg.stage_id == 1
+    assert terminal_msg.finished is True
+    assert "produced no valid inputs" in terminal_msg.engine_outputs.error
     assert "req-empty" not in orchestrator.request_states