From 0d9d57acd90f6b6418cf8ccc91c76991a84103e6 Mon Sep 17 00:00:00 2001
From: wuhang <wuhang6@huawei.com>
Date: Thu, 14 May 2026 11:25:55 +0800
Subject: [PATCH] [Entrypoint][Refactor] Make field type hint more concrete
 (#3139)

Signed-off-by: wuhang <wuhang6@huawei.com>
---
 tests/engine/test_async_omni_engine_input.py  |  10 +-
 .../engine/test_async_omni_engine_outputs.py  |  34 +-
 .../test_async_omni_engine_stage_init.py      |   6 +-
 tests/engine/test_orchestrator.py             | 179 +++++---
 .../test_orchestrator_error_handling.py       |  27 +-
 .../test_orchestrator_kv_sender_info.py       |  12 +-
 tests/entrypoints/test_omni_base_profiler.py  |  10 +-
 tests/entrypoints/test_omni_entrypoints.py    | 127 +++---
 .../inline_stage_diffusion_client.py          |   4 +-
 vllm_omni/diffusion/stage_diffusion_client.py |   4 +-
 vllm_omni/engine/async_omni_engine.py         | 197 +++++----
 vllm_omni/engine/messages.py                  |  90 ++++
 vllm_omni/engine/orchestrator.py              | 408 ++++++++++--------
 vllm_omni/engine/output_modality.py           |   3 +
 vllm_omni/engine/stage_client.py              | 111 +++++
 vllm_omni/engine/stage_engine_core_client.py  |   3 +-
 vllm_omni/engine/stage_pool.py                |  72 ++--
 vllm_omni/entrypoints/async_omni.py           |  68 +--
 vllm_omni/entrypoints/omni.py                 |   5 +-
 vllm_omni/entrypoints/omni_base.py            |  81 ++--
 20 files changed, 908 insertions(+), 543 deletions(-)
 create mode 100644 vllm_omni/engine/messages.py
 create mode 100644 vllm_omni/engine/stage_client.py

diff --git a/tests/engine/test_async_omni_engine_input.py b/tests/engine/test_async_omni_engine_input.py
index d86483f15..65b3a7ec8 100644
--- a/tests/engine/test_async_omni_engine_input.py
+++ b/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
diff --git a/tests/engine/test_async_omni_engine_outputs.py b/tests/engine/test_async_omni_engine_outputs.py
index 47b3d5e9f..ecf3ee9ee 100644
--- a/tests/engine/test_async_omni_engine_outputs.py
+++ b/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
diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py
index fcbf3788f..51ed6b46c 100644
--- a/tests/engine/test_async_omni_engine_stage_init.py
+++ b/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"),
     ]
 
 
diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py
index 93b2761ae..f0d45a1e6 100644
--- a/tests/engine/test_orchestrator.py
+++ b/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)
 
diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py
index 18099c016..c5cba1ea9 100644
--- a/tests/engine/test_orchestrator_error_handling.py
+++ b/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)
diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py
index 2bd9b0947..e3f58c129 100644
--- a/tests/engine/test_orchestrator_kv_sender_info.py
+++ b/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
 
 
diff --git a/tests/entrypoints/test_omni_base_profiler.py b/tests/entrypoints/test_omni_base_profiler.py
index ca10eed91..3f793e89e 100644
--- a/tests/entrypoints/test_omni_base_profiler.py
+++ b/tests/entrypoints/test_omni_base_profiler.py
@@ -1,5 +1,7 @@
 """Unit tests for OmniBase and AsyncOmni profiler methods."""
 
+from types import SimpleNamespace
+
 import pytest
 from pytest_mock import MockerFixture
 
@@ -16,10 +18,10 @@ class TestOmniBaseProfiler:
         engine.num_stages = 3
         engine.is_alive.return_value = True
         engine.default_sampling_params_list = [mocker.MagicMock() for _ in range(3)]
-        engine.get_stage_metadata.side_effect = lambda i: {
-            "final_output_type": "text" if i == 0 else "audio",
-            "final_output": True,
-        }
+        engine.get_stage_metadata.side_effect = lambda i: SimpleNamespace(
+            final_output_type="text" if i == 0 else "audio",
+            final_output=True,
+        )
         engine.collective_rpc.return_value = [None, None, None]
         return engine
 
diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py
index 251e54ab8..2360a68fc 100644
--- a/tests/entrypoints/test_omni_entrypoints.py
+++ b/tests/entrypoints/test_omni_entrypoints.py
@@ -13,44 +13,40 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 
+from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
+from vllm_omni.engine.messages import ErrorMessage, OutputMessage
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.entrypoints.omni_base import OmniEngineDeadError
+from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def _stage_meta(*, stage_type: str, final_output: bool, final_output_type: str | None) -> StageRuntimeInfo:
+    return StageRuntimeInfo(
+        stage_type=stage_type,
+        final_output=final_output,
+        final_output_type=final_output_type,
+    )
+
+
 THREE_STAGE_META = [
-    {"stage_type": "llm", "final_output": True, "final_output_type": "text"},
-    {"stage_type": "llm", "final_output": False, "final_output_type": None},
-    {"stage_type": "diffusion", "final_output": True, "final_output_type": "image"},
+    _stage_meta(stage_type="llm", final_output=True, final_output_type="text"),
+    _stage_meta(stage_type="llm", final_output=False, final_output_type=None),
+    _stage_meta(stage_type="diffusion", final_output=True, final_output_type="image"),
 ]
 
 DIFFUSION_ONLY_META = [
-    {"stage_type": "diffusion", "final_output": True, "final_output_type": "image"},
+    _stage_meta(stage_type="diffusion", final_output=True, final_output_type="image"),
 ]
 
 LLM_DIFFUSION_META = [
-    {"stage_type": "llm", "final_output": True, "final_output_type": "text"},
-    {"stage_type": "diffusion", "final_output": True, "final_output_type": "image"},
+    _stage_meta(stage_type="llm", final_output=True, final_output_type="text"),
+    _stage_meta(stage_type="diffusion", final_output=True, final_output_type="image"),
 ]
 
 
-class FakeEngineOutput:
-    def __init__(
-        self,
-        *,
-        payload: str,
-        finished: bool,
-        images: list[str] | None = None,
-        stage_durations: dict[str, float] | None = None,
-    ) -> None:
-        self.payload = payload
-        self.finished = finished
-        self.images = images or []
-        self.stage_durations = stage_durations or {}
-
-
 def make_output_msg(
     request_id: str,
     stage_id: int,
@@ -60,21 +56,25 @@ def make_output_msg(
     finished: bool | None = None,
     images: list[str] | None = None,
     metrics: Any = None,
-) -> dict[str, Any]:
+) -> OutputMessage:
     if finished is None:
         finished = output_finished
-    return {
-        "type": "output",
-        "request_id": request_id,
-        "stage_id": stage_id,
-        "engine_outputs": FakeEngineOutput(
-            payload=payload,
-            finished=output_finished,
-            images=images,
-        ),
-        "finished": finished,
-        "metrics": metrics,
-    }
+    final_output_type = "image" if images else "text"
+    engine_output = OmniRequestOutput(
+        request_id=request_id,
+        finished=output_finished,
+        final_output_type=final_output_type,
+        images=images or [],
+        stage_durations={},
+    )
+    engine_output.payload = payload
+    return OutputMessage(
+        request_id=request_id,
+        stage_id=stage_id,
+        engine_outputs=engine_output,
+        finished=finished,
+        metrics=metrics,
+    )
 
 
 class FakeAsyncOmniEngine:
@@ -82,7 +82,7 @@ class FakeAsyncOmniEngine:
         self,
         model: str = "dummy-model",
         *,
-        stage_metadata: list[dict[str, Any]] | None = None,
+        stage_metadata: list[StageRuntimeInfo] | None = None,
         default_sampling_params_list: list[Any] | None = None,
         on_add_request: Callable[[FakeAsyncOmniEngine, dict[str, Any]], None] | None = None,
         rpc_results: list[Any] | None = None,
@@ -102,7 +102,7 @@ class FakeAsyncOmniEngine:
         self.output_processors = [SimpleNamespace(tokenizer=None) for _ in range(self.num_stages)]
         self.input_processor = None
 
-        self.output_q: queue.Queue[dict[str, Any]] = queue.Queue()
+        self.output_q: queue.Queue[Any] = queue.Queue()
         self.submitted: list[dict[str, Any]] = []
         self.aborted: list[list[str]] = []
         self.rpc_results = rpc_results or []
@@ -133,16 +133,16 @@ class FakeAsyncOmniEngine:
     async def add_request_async(self, *args, **kwargs) -> None:
         self.add_request(*args, **kwargs)
 
-    def try_get_output(self, timeout: float = 0.001) -> dict[str, Any] | None:
+    def try_get_output(self, timeout: float = 0.001) -> Any | None:
         try:
             return self.output_q.get_nowait()
         except queue.Empty:
             return None
 
-    async def try_get_output_async(self) -> dict[str, Any] | None:
+    async def try_get_output_async(self) -> Any | None:
         return self.try_get_output()
 
-    def get_stage_metadata(self, stage_id: int) -> dict[str, Any]:
+    def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo:
         return self.stage_metadata[stage_id]
 
     def abort(self, request_ids: list[str]) -> None:
@@ -405,24 +405,22 @@ def _enqueue_async_llm_diffusion_outputs(engine: FakeAsyncOmniEngine, msg: dict[
 
 def _enqueue_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
     engine.output_q.put_nowait(
-        {
-            "type": "error",
-            "request_id": msg["request_id"],
-            "stage_id": 0,
-            "error": "engine boom",
-        }
+        ErrorMessage(
+            request_id=msg["request_id"],
+            stage_id=0,
+            error="engine boom",
+        )
     )
 
 
 def _enqueue_fatal_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
     engine.output_q.put_nowait(
-        {
-            "type": "error",
-            "fatal": True,
-            "request_id": msg["request_id"],
-            "stage_id": 2,
-            "error": "engine dead",
-        }
+        ErrorMessage(
+            fatal=True,
+            request_id=msg["request_id"],
+            stage_id=2,
+            error="engine dead",
+        )
     )
 
 
@@ -798,7 +796,7 @@ def test_omni_forces_final_only_on_llm_stages(monkeypatch: pytest.MonkeyPatch):
 
 def test_fatal_error_raises_engine_dead():
     base = _make_base()
-    msg = {"type": "error", "error": "orchestrator crashed", "fatal": True}
+    msg = ErrorMessage(error="orchestrator crashed", fatal=True)
 
     with pytest.raises(EngineDeadError, match="orchestrator crashed"):
         base._handle_output_message(msg)
@@ -806,7 +804,7 @@ def test_fatal_error_raises_engine_dead():
 
 def test_non_fatal_error_raises_runtime():
     base = _make_base()
-    msg = {"type": "error", "error": "something wrong"}
+    msg = ErrorMessage(error="something wrong")
 
     with pytest.raises(RuntimeError, match="something wrong"):
         base._handle_output_message(msg)
@@ -853,20 +851,15 @@ def _enqueue_stage_error(
     """Enqueue a stage error output, optionally killing the engine."""
     if kill_engine:
         engine._alive = False
+    engine_output = OmniRequestOutput.from_error(msg["request_id"], error_text)
+    engine_output.payload = ""
     engine.output_q.put_nowait(
-        {
-            "type": "output",
-            "request_id": msg["request_id"],
-            "stage_id": 0,
-            "engine_outputs": SimpleNamespace(
-                payload="",
-                finished=True,
-                images=[],
-                stage_durations={},
-                error=error_text,
-            ),
-            "finished": False,
-        }
+        OutputMessage(
+            request_id=msg["request_id"],
+            stage_id=0,
+            engine_outputs=engine_output,
+            finished=False,
+        )
     )
 
 
diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py
index 07256a264..62eff7ecc 100644
--- a/vllm_omni/diffusion/inline_stage_diffusion_client.py
+++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py
@@ -20,6 +20,7 @@ from vllm.v1.engine.exceptions import EngineDeadError
 from vllm_omni.diffusion.data import DiffusionRequestAbortedError
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.engine.stage_client import StageClientBase
 from vllm_omni.engine.stage_init_utils import StageMetadata
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -31,11 +32,12 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
-class InlineStageDiffusionClient:
+class InlineStageDiffusionClient(StageClientBase):
     """Runs DiffusionEngine in a thread executor inside the Orchestrator."""
 
     stage_type: str = "diffusion"
     replica_id: int = 0
+    is_comprehension: bool = False
 
     def __init__(
         self,
diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py
index e19db4536..11809f6ef 100644
--- a/vllm_omni/diffusion/stage_diffusion_client.py
+++ b/vllm_omni/diffusion/stage_diffusion_client.py
@@ -29,6 +29,7 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import (
     OmniMsgpackDecoder,
     OmniMsgpackEncoder,
 )
+from vllm_omni.engine.stage_client import StageClientBase
 from vllm_omni.engine.stage_init_utils import StageMetadata, terminate_alive_proc
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -57,7 +58,7 @@ def create_diffusion_client(
     )
 
 
-class StageDiffusionClient:
+class StageDiffusionClient(StageClientBase):
     """Communicates with StageDiffusionProc via ZMQ for use inside the Orchestrator.
 
     Exposes the same attributes and async methods the Orchestrator
@@ -68,6 +69,7 @@ class StageDiffusionClient:
 
     stage_type: str = "diffusion"
     replica_id: int = 0
+    is_comprehension: bool = False
 
     def __init__(
         self,
diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py
index ed4eed808..35b323487 100644
--- a/vllm_omni/engine/async_omni_engine.py
+++ b/vllm_omni/engine/async_omni_engine.py
@@ -20,8 +20,8 @@ import uuid
 import weakref
 from collections.abc import Mapping, Sequence
 from contextlib import ExitStack
-from dataclasses import asdict
-from typing import TYPE_CHECKING, Any
+from dataclasses import asdict, dataclass
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 import janus
 import torch
@@ -36,6 +36,7 @@ from vllm.v1.engine.input_processor import InputProcessor
 from vllm_omni.config.stage_config import strip_parent_engine_args
 from vllm_omni.diffusion.data import DiffusionParallelConfig, parse_attention_config
 from vllm_omni.diffusion.diffusion_engine import supports_audio_output
+from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
 from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
 from vllm_omni.diffusion.stage_diffusion_proc import (
     complete_diffusion_handshake,
@@ -45,11 +46,23 @@ from vllm_omni.distributed.omni_connectors.utils.initialization import (
     resolve_omni_kv_config_for_stage,
 )
 from vllm_omni.engine import OmniEngineCoreRequest
+from vllm_omni.engine.messages import (
+    AbortRequestMessage,
+    AddCompanionRequestMessage,
+    CollectiveRPCRequestMessage,
+    CollectiveRPCResultMessage,
+    EngineQueueMessage,
+    ErrorMessage,
+    ShutdownRequestMessage,
+    StageSubmissionMessage,
+)
 from vllm_omni.engine.orchestrator import Orchestrator
+from vllm_omni.engine.output_modality import FinalOutputModalityType
 from vllm_omni.engine.serialization import (
     deserialize_additional_information,
     serialize_additional_information,
 )
+from vllm_omni.engine.stage_client import StageClient
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase
 from vllm_omni.engine.stage_engine_core_proc import (
     complete_stage_handshake,
@@ -83,12 +96,13 @@ from vllm_omni.engine.stage_init_utils import (
     setup_stage_devices,
     terminate_alive_proc,
 )
-from vllm_omni.engine.stage_pool import StagePool
+from vllm_omni.engine.stage_pool import StagePool, StagePoolClient
 from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
 from vllm_omni.entrypoints.utils import (
     inject_omni_kv_config,
     load_and_resolve_stage_configs,
 )
+from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
 if TYPE_CHECKING:
@@ -99,6 +113,13 @@ logger = init_logger(__name__)
 _STARTUP_POLL_INTERVAL_S = 1.0
 
 
+@dataclass(frozen=True, slots=True)
+class StageRuntimeInfo:
+    final_output: bool
+    final_output_type: FinalOutputModalityType | None
+    stage_type: str
+
+
 # ============================================================================
 # Parent-EngineArgs field-routing contracts (consumed by
 # AsyncOmniEngine._strip_parent_engine_args when ``stage_configs_path`` is set).
@@ -190,14 +211,14 @@ def _apply_omni_final_stage_metadata(
 
 def _weak_shutdown_async_omni_engine(
     orchestrator_thread: threading.Thread | None,
-    request_queue: janus.Queue[dict[str, Any]] | None,
-    output_queue: janus.Queue[dict[str, Any]] | None,
-    rpc_output_queue: janus.Queue[dict[str, Any]] | None,
+    request_queue: janus.Queue[EngineQueueMessage] | None,
+    output_queue: janus.Queue[EngineQueueMessage] | None,
+    rpc_output_queue: janus.Queue[EngineQueueMessage] | None,
 ) -> None:
     """Best-effort orchestrator cleanup for GC finalization."""
     try:
         if request_queue is not None:
-            request_queue.sync_q.put_nowait({"type": "shutdown"})
+            request_queue.sync_q.put_nowait(ShutdownRequestMessage())
     except Exception:
         pass
 
@@ -295,14 +316,14 @@ class AsyncOmniEngine:
         stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None
         self.async_chunk = bool(getattr(stage0_args, "async_chunk", False))
         self.stage_pools: list[StagePool] = []
-        self.stage_clients: list[Any] = []  # logical-stage view for external readers
+        self.stage_clients: list[StageClient] = []  # logical-stage view for external readers
         self.input_processor: InputProcessor | None = None
         self.supported_tasks: tuple[str, ...] = ("generate",)
-        self.default_sampling_params_list: list[Any] = []
-        self.stage_metadata: list[dict[str, Any]] = []
-        self.request_queue: janus.Queue[dict[str, Any]] | None = None
-        self.output_queue: janus.Queue[dict[str, Any]] | None = None
-        self.rpc_output_queue: janus.Queue[dict[str, Any]] | None = None
+        self.default_sampling_params_list: list[OmniSamplingParams] = []
+        self.stage_metadata: list[StageRuntimeInfo] = []
+        self.request_queue: janus.Queue[EngineQueueMessage] | None = None
+        self.output_queue: janus.Queue[EngineQueueMessage] | None = None
+        self.rpc_output_queue: janus.Queue[EngineQueueMessage] | None = None
         self._shutdown_called = False
         self._weak_finalizer: weakref.finalize | None = None
         self._rpc_lock = threading.Lock()
@@ -379,15 +400,15 @@ class AsyncOmniEngine:
 
     @staticmethod
     def _collect_initialized_clients_for_cleanup(
-        stage_pools: Sequence[Any],
-        initialized_clients_by_stage: Mapping[int, Sequence[Any | None]],
-    ) -> list[Any]:
+        stage_pools: Sequence[StagePool],
+        initialized_clients_by_stage: Mapping[int, Sequence[StagePoolClient | None]],
+    ) -> list[StageClient]:
         """Collect initialized clients exactly once for failure cleanup."""
 
-        collected: list[Any] = []
+        collected: list[StageClient] = []
         seen: set[int] = set()
 
-        def _add_client(client: Any) -> None:
+        def _add_client(client: StageClient | None) -> None:
             if client is None:
                 return
             client_id = id(client)
@@ -407,7 +428,7 @@ class AsyncOmniEngine:
         return collected
 
     @staticmethod
-    def _shutdown_initialized_clients(clients: Sequence[Any]) -> None:
+    def _shutdown_initialized_clients(clients: Sequence[StageClient]) -> None:
         """Best-effort shutdown for attached clients after init failure."""
 
         for client in reversed(list(clients)):
@@ -585,7 +606,7 @@ class AsyncOmniEngine:
         plan: ReplicaInitPlan,
         stage_init_timeout: int,
         llm_stage_launch_lock: threading.Lock,
-    ) -> Any:
+    ) -> StageEngineCoreClientBase:
         """Initialize one LLM replica end-to-end."""
 
         proc = None
@@ -748,7 +769,7 @@ class AsyncOmniEngine:
         plan: ReplicaInitPlan,
         stage_init_timeout: int,
         stage_launch_lock: threading.Lock,
-    ) -> Any:
+    ) -> StageDiffusionClient | InlineStageDiffusionClient:
         """Initialize one diffusion replica end-to-end."""
 
         client = None
@@ -865,7 +886,7 @@ class AsyncOmniEngine:
         plan: ReplicaInitPlan,
         stage_init_timeout: int,
         stage_launch_lock: threading.Lock,
-    ) -> Any:
+    ) -> StagePoolClient:
         """Initialize one replica, regardless of backend type."""
 
         if plan.metadata.stage_type == "diffusion":
@@ -876,7 +897,7 @@ class AsyncOmniEngine:
         self,
         stage_plans: Sequence[LogicalStageInitPlan],
         stage_init_timeout: int,
-    ) -> dict[int, list[Any | None]]:
+    ) -> dict[int, list[StagePoolClient | None]]:
         """Initialize all stage replicas.
 
         Diffusion replicas are launched **inline on the orchestrator thread**
@@ -892,7 +913,7 @@ class AsyncOmniEngine:
         """
 
         stage_launch_lock = threading.Lock()
-        initialized_clients_by_stage: dict[int, list[Any | None]] = {
+        initialized_clients_by_stage: dict[int, list[StagePoolClient | None]] = {
             plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans
         }
         primary_exc: Exception | None = None
@@ -922,7 +943,7 @@ class AsyncOmniEngine:
 
         # --- 2) LLM replicas: parallel init via a scoped ThreadPoolExecutor. ---
         if primary_exc is None and llm_replicas:
-            future_to_replica: dict[concurrent.futures.Future[Any], tuple[int, int]] = {}
+            future_to_replica: dict[concurrent.futures.Future[StagePoolClient], tuple[int, int]] = {}
             with concurrent.futures.ThreadPoolExecutor(
                 max_workers=max(1, len(llm_replicas)),
                 thread_name_prefix="stage-init",
@@ -959,13 +980,13 @@ class AsyncOmniEngine:
     def _assemble_stage_pools(
         self,
         stage_plans: Sequence[LogicalStageInitPlan],
-        initialized_clients_by_stage: Mapping[int, Sequence[Any | None]],
+        initialized_clients_by_stage: Mapping[int, Sequence[StagePoolClient | None]],
     ) -> list[StagePool]:
         """Assemble logical stage pools and update top-level stage metadata."""
 
         stage_pools: list[StagePool] = []
-        default_sampling_params_list: list[Any] = []
-        stage_metadata_list: list[dict[str, Any]] = []
+        default_sampling_params_list: list[OmniSamplingParams] = []
+        stage_metadata_list: list[StageRuntimeInfo] = []
 
         for plan in stage_plans:
             replica_clients = initialized_clients_by_stage[plan.stage_idx]
@@ -973,7 +994,7 @@ class AsyncOmniEngine:
             if first_client is None:
                 raise RuntimeError(f"Stage {plan.stage_idx} initialization completed with a missing client")
 
-            clients = [client for client in replica_clients if client is not None]
+            clients: list[StagePoolClient] = [client for client in replica_clients if client is not None]
             stage_vllm_config = None
             output_processor = None
             if plan.replicas[0].metadata.stage_type != "diffusion":
@@ -991,11 +1012,11 @@ class AsyncOmniEngine:
             )
             default_sampling_params_list.append(first_client.default_sampling_params)
             stage_metadata_list.append(
-                {
-                    "final_output": first_client.final_output,
-                    "final_output_type": first_client.final_output_type,
-                    "stage_type": first_client.stage_type,
-                }
+                StageRuntimeInfo(
+                    final_output=first_client.final_output,
+                    final_output_type=first_client.final_output_type,
+                    stage_type=first_client.stage_type,
+                )
             )
 
         self.default_sampling_params_list = list(default_sampling_params_list)
@@ -1033,7 +1054,7 @@ class AsyncOmniEngine:
 
         stage_pools: list[StagePool] = []
         input_processor: InputProcessor | None = None
-        initialized_clients_by_stage: dict[int, list[Any | None]] = {
+        initialized_clients_by_stage: dict[int, list[StagePoolClient | None]] = {
             plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans
         }
 
@@ -1071,7 +1092,7 @@ class AsyncOmniEngine:
         self.prompt_expand_func = prompt_expand_func
 
         # Derive logical-stage views for external readers (entrypoints/async_omni.py).
-        self.stage_clients = [pool.stage_client for pool in self.stage_pools]
+        self.stage_clients = [cast(StageClient, pool.stage_client) for pool in self.stage_pools]
         self.stage_vllm_configs = [pool.stage_vllm_config for pool in self.stage_pools]
         self.output_processors = [pool.output_processor for pool in self.stage_pools]
 
@@ -1079,7 +1100,7 @@ class AsyncOmniEngine:
         supported_tasks: set[str] = set()
         if any(getattr(pool.stage_client, "is_comprehension", False) for pool in self.stage_pools):
             supported_tasks.add("generate")
-        if any(m.get("final_output_type") == "audio" for m in self.stage_metadata):
+        if any(m.final_output_type == "audio" for m in self.stage_metadata):
             supported_tasks.add("speech")
         self.supported_tasks = tuple(supported_tasks) if supported_tasks else ("generate",)
 
@@ -1127,7 +1148,7 @@ class AsyncOmniEngine:
             logger.exception("[AsyncOmniEngine] Orchestrator thread crashed")
             error_text = str(e) or "Orchestrator thread crashed"
             try:
-                error_msg = {"type": "error", "error": error_text, "fatal": True}
+                error_msg = ErrorMessage(error=error_text, fatal=True)
                 if self.output_queue is not None:
                     self.output_queue.sync_q.put_nowait(error_msg)
                 if self.rpc_output_queue is not None:
@@ -1194,11 +1215,13 @@ class AsyncOmniEngine:
         reasoning_ended: bool | None = None,
         *,
         resumable: bool = False,
-        message_type: str = "add_request",
-    ) -> dict[str, Any]:
+        message_type: Literal["add_request", "streaming_update"] = "add_request",
+    ) -> StageSubmissionMessage:
         """Build an add_request message after stage-0 preprocessing."""
-        effective_sampling_params_list = (
-            list(sampling_params_list) if sampling_params_list is not None else list(self.default_sampling_params_list)
+        effective_sampling_params_list: list[OmniSamplingParams] = (
+            list(cast(Sequence[OmniSamplingParams], sampling_params_list))
+            if sampling_params_list is not None
+            else list(self.default_sampling_params_list)
         )
         if not effective_sampling_params_list:
             raise ValueError(
@@ -1210,7 +1233,7 @@ class AsyncOmniEngine:
         # dict, e.g. for multi_modal_data).
         original_prompt = prompt
 
-        stage_type = self.stage_metadata[0].get("stage_type")
+        stage_type = self.stage_metadata[0].stage_type
         output_prompt_text: Any = None
         _preprocess_ms = 0.0
         if stage_type != "diffusion" and not isinstance(prompt, EngineCoreRequest):
@@ -1261,17 +1284,17 @@ class AsyncOmniEngine:
                 output_prompt_text = original_prompt.get("prompt")
             prompt = request
 
-        return {
-            "type": message_type,
-            "request_id": request_id,
-            "prompt": prompt,
-            "original_prompt": original_prompt,
-            "output_prompt_text": output_prompt_text,
-            "sampling_params_list": effective_sampling_params_list,
-            "final_stage_id": final_stage_id,
-            "preprocess_ms": _preprocess_ms,
-            "enqueue_ts": time.perf_counter(),
-        }
+        return StageSubmissionMessage(
+            type=message_type,
+            request_id=request_id,
+            prompt=prompt,
+            original_prompt=original_prompt,
+            output_prompt_text=output_prompt_text,
+            sampling_params_list=effective_sampling_params_list,
+            final_stage_id=final_stage_id,
+            preprocess_ms=_preprocess_ms,
+            enqueue_ts=time.perf_counter(),
+        )
 
     def _enqueue_cfg_companions(
         self,
@@ -1311,15 +1334,14 @@ class AsyncOmniEngine:
             # deferred to Orchestrator._handle_add_companion, which routes
             # admission through StagePool.submit_initial(..., affinity_request_id=...).
             self.request_queue.sync_q.put_nowait(
-                {
-                    "type": "add_companion_request",
-                    "companion_id": cid,
-                    "parent_id": parent_id,
-                    "role": ep.role,
-                    "prompt": request,
-                    "companion_prompt_text": companion_prompt,
-                    "sampling_params_list": companion_spl,
-                }
+                AddCompanionRequestMessage(
+                    companion_id=cid,
+                    parent_id=parent_id,
+                    role=ep.role,
+                    prompt=request,
+                    companion_prompt_text=companion_prompt,
+                    sampling_params_list=companion_spl,
+                )
             )
 
         logger.info(
@@ -1766,8 +1788,8 @@ class AsyncOmniEngine:
         # CFG companion expansion: create and enqueue companion requests
         # so the AR stage also generates their KV caches.
         if self.prompt_expand_func is not None and final_stage_id > 0:
-            original_prompt = msg.get("original_prompt", prompt)
-            effective_spl = msg.get("sampling_params_list", [])
+            original_prompt = msg.original_prompt
+            effective_spl = msg.sampling_params_list
             stage0_params = effective_spl[0] if effective_spl else None
             if stage0_params is not None:
                 self._enqueue_cfg_companions(request_id, original_prompt, stage0_params, effective_spl)
@@ -1854,7 +1876,7 @@ class AsyncOmniEngine:
             resumable=resumable,
         )
 
-    def try_get_output(self, timeout: float = 0.001) -> dict[str, Any] | None:
+    def try_get_output(self, timeout: float = 0.001) -> EngineQueueMessage | None:
         """Read one output message from the Orchestrator output queue."""
         if self.output_queue is None:
             return None
@@ -1865,7 +1887,7 @@ class AsyncOmniEngine:
                 raise RuntimeError("Orchestrator died unexpectedly. See logs above.")
             return None
 
-    async def try_get_output_async(self) -> dict[str, Any] | None:
+    async def try_get_output_async(self) -> EngineQueueMessage | None:
         """Async read from the Orchestrator output queue."""
         if self.output_queue is None:
             return None
@@ -1876,7 +1898,7 @@ class AsyncOmniEngine:
                 raise RuntimeError("Orchestrator died unexpectedly. See logs above.")
             return None
 
-    def get_stage_metadata(self, stage_id: int) -> dict[str, Any]:
+    def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo:
         """Get cached metadata for a stage."""
         return self.stage_metadata[stage_id]
 
@@ -1884,12 +1906,7 @@ class AsyncOmniEngine:
         """Send abort message to the Orchestrator."""
         if self.request_queue is None:
             raise RuntimeError("request_queue is not initialized")
-        self.request_queue.sync_q.put_nowait(
-            {
-                "type": "abort",
-                "request_ids": request_ids,
-            }
-        )
+        self.request_queue.sync_q.put_nowait(AbortRequestMessage(request_ids=request_ids))
 
     async def abort_async(self, request_ids: list[str]) -> None:
         """Async abort API."""
@@ -1914,14 +1931,14 @@ class AsyncOmniEngine:
             raise RuntimeError("rpc_output_queue is not initialized")
 
         rpc_id = uuid.uuid4().hex
-        msg = {
-            "type": "collective_rpc",
-            "rpc_id": rpc_id,
-            "method": method,
-            "args": tuple(args),
-            "kwargs": kwargs or {},
-            "stage_ids": stage_ids,
-        }
+        msg = CollectiveRPCRequestMessage(
+            rpc_id=rpc_id,
+            method=method,
+            timeout=timeout,
+            args=tuple(args),
+            kwargs=kwargs or {},
+            stage_ids=stage_ids,
+        )
 
         with self._rpc_lock:
             self.request_queue.sync_q.put_nowait(msg)
@@ -1934,25 +1951,25 @@ class AsyncOmniEngine:
                 except queue.Empty as exc:
                     raise TimeoutError(f"collective_rpc timed out after {timeout} seconds") from exc
 
-                if result_msg.get("type") == "error":
-                    raise RuntimeError(result_msg.get("error", "Orchestrator returned an error message"))
+                if isinstance(result_msg, ErrorMessage):
+                    raise RuntimeError(result_msg.error)
 
-                if result_msg.get("type") != "collective_rpc_result":
+                if not isinstance(result_msg, CollectiveRPCResultMessage):
                     logger.warning(
                         "[AsyncOmniEngine] Dropping unexpected rpc queue message type=%s",
-                        result_msg.get("type"),
+                        getattr(result_msg, "type", type(result_msg).__name__),
                     )
                     continue
 
-                if result_msg.get("rpc_id") != rpc_id:
+                if result_msg.rpc_id != rpc_id:
                     logger.warning(
                         "[AsyncOmniEngine] Dropping mismatched rpc result rpc_id=%s expected=%s",
-                        result_msg.get("rpc_id"),
+                        result_msg.rpc_id,
                         rpc_id,
                     )
                     continue
 
-                return list(result_msg.get("results", []))
+                return list(result_msg.results)
 
     async def collective_rpc_async(
         self,
@@ -1991,7 +2008,7 @@ class AsyncOmniEngine:
         logger.info("[AsyncOmniEngine] Shutting down Orchestrator")
         try:
             if self.request_queue is not None:
-                self.request_queue.sync_q.put_nowait({"type": "shutdown"})
+                self.request_queue.sync_q.put_nowait(ShutdownRequestMessage())
         except Exception:
             pass
         if self.is_alive():
diff --git a/vllm_omni/engine/messages.py b/vllm_omni/engine/messages.py
new file mode 100644
index 000000000..28a5c721c
--- /dev/null
+++ b/vllm_omni/engine/messages.py
@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+from typing import Literal
+
+import msgspec
+from vllm.inputs import PromptType
+from vllm.v1.engine import EngineCoreRequest
+
+from vllm_omni.inputs.data import OmniSamplingParams
+from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics
+from vllm_omni.outputs import OmniRequestOutput
+
+
+class EngineQueueMessage(msgspec.Struct, forbid_unknown_fields=True):
+    pass
+
+
+class StageSubmissionMessage(EngineQueueMessage, kw_only=True):
+    type: Literal["add_request", "streaming_update"]
+    request_id: str
+    prompt: EngineCoreRequest | PromptType
+    original_prompt: EngineCoreRequest | PromptType
+    output_prompt_text: object | None
+    sampling_params_list: list[OmniSamplingParams]
+    final_stage_id: int
+    preprocess_ms: float
+    enqueue_ts: float
+
+
+class AddCompanionRequestMessage(EngineQueueMessage, kw_only=True):
+    type: Literal["add_companion_request"] = "add_companion_request"
+    companion_id: str
+    parent_id: str
+    role: str
+    prompt: EngineCoreRequest
+    companion_prompt_text: object | None
+    sampling_params_list: list[OmniSamplingParams]
+
+
+class AbortRequestMessage(EngineQueueMessage, kw_only=True):
+    type: Literal["abort"] = "abort"
+    request_ids: list[str]
+
+
+class CollectiveRPCRequestMessage(EngineQueueMessage, kw_only=True):
+    type: Literal["collective_rpc"] = "collective_rpc"
+    rpc_id: str
+    method: str
+    timeout: float | None = None
+    args: tuple[object, ...]
+    kwargs: dict[str, object]
+    stage_ids: list[int] | None
+
+
+class ShutdownRequestMessage(EngineQueueMessage, kw_only=True):
+    type: Literal["shutdown"] = "shutdown"
+
+
+class ErrorMessage(EngineQueueMessage, kw_only=True):
+    type: Literal["error"] = "error"
+    error: str
+    fatal: bool = False
+    request_id: str | None = None
+    stage_id: int | None = None
+
+
+class OutputMessage(EngineQueueMessage, kw_only=True):
+    type: Literal["output"] = "output"
+    request_id: str
+    stage_id: int
+    engine_outputs: OmniRequestOutput
+    metrics: StageRequestMetrics | None = None
+    finished: bool
+    stage_submit_ts: float | None = None
+
+
+class StageMetricsMessage(EngineQueueMessage, kw_only=True):
+    type: Literal["stage_metrics"] = "stage_metrics"
+    request_id: str
+    stage_id: int
+    metrics: StageRequestMetrics
+    stage_submit_ts: float | None = None
+
+
+class CollectiveRPCResultMessage(EngineQueueMessage, kw_only=True):
+    type: Literal["collective_rpc_result"] = "collective_rpc_result"
+    rpc_id: str
+    method: str
+    stage_ids: list[int]
+    results: list[object]
diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py
index 2d2ac47cb..4bc032ae2 100644
--- a/vllm_omni/engine/orchestrator.py
+++ b/vllm_omni/engine/orchestrator.py
@@ -24,6 +24,18 @@ from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
+from vllm_omni.engine.messages import (
+    AbortRequestMessage,
+    AddCompanionRequestMessage,
+    CollectiveRPCRequestMessage,
+    CollectiveRPCResultMessage,
+    EngineQueueMessage,
+    ErrorMessage,
+    OutputMessage,
+    ShutdownRequestMessage,
+    StageMetricsMessage,
+    StageSubmissionMessage,
+)
 from vllm_omni.engine.serialization import serialize_additional_information
 from vllm_omni.engine.stage_pool import StagePool
 from vllm_omni.outputs import OmniRequestOutput
@@ -115,9 +127,9 @@ class Orchestrator:
 
     def __init__(
         self,
-        request_async_queue: janus.AsyncQueue[dict[str, Any]],
-        output_async_queue: janus.AsyncQueue[dict[str, Any]],
-        rpc_async_queue: janus.AsyncQueue[dict[str, Any]],
+        request_async_queue: janus.AsyncQueue[EngineQueueMessage],
+        output_async_queue: janus.AsyncQueue[EngineQueueMessage],
+        rpc_async_queue: janus.AsyncQueue[EngineQueueMessage],
         stage_pools: list[StagePool],
         *,
         async_chunk: bool = False,
@@ -196,36 +208,38 @@ class Orchestrator:
         """Read messages from the main thread via request_async_queue."""
         while True:
             msg = await self.request_async_queue.get()
-            msg_type = msg.get("type")
 
-            if msg_type == "add_request":
-                await self._handle_add_request(msg)
-            elif msg_type == "streaming_update":
-                await self._handle_streaming_update(msg)
-            elif msg_type == "add_companion_request":
+            if isinstance(msg, StageSubmissionMessage):
+                if msg.type == "add_request":
+                    await self._handle_add_request(msg)
+                elif msg.type == "streaming_update":
+                    await self._handle_streaming_update(msg)
+                else:
+                    logger.warning("[Orchestrator] Unknown stage submission type: %s", msg.type)
+            elif isinstance(msg, AddCompanionRequestMessage):
                 await self._handle_add_companion(msg)
-            elif msg_type == "abort":
+            elif isinstance(msg, AbortRequestMessage):
                 await self._handle_abort(msg)
-            elif msg_type == "collective_rpc":
+            elif isinstance(msg, CollectiveRPCRequestMessage):
                 await self._handle_collective_rpc(msg)
-            elif msg_type == "shutdown":
+            elif isinstance(msg, ShutdownRequestMessage):
                 logger.info("[Orchestrator] Received shutdown signal")
                 self._shutdown_event.set()
                 self._shutdown_stages()
                 break
             else:
-                logger.warning("[Orchestrator] Unknown message type: %s", msg_type)
+                logger.warning("[Orchestrator] Unknown message type: %s", type(msg).__name__)
 
-    async def _handle_add_request(self, msg: dict[str, Any]) -> None:
+    async def _handle_add_request(self, msg: StageSubmissionMessage) -> None:
         """Handle an add_request message from the main thread."""
         stage_id = 0
-        request_id = msg["request_id"]
-        prompt = msg["prompt"]
-        original_prompt = msg.get("original_prompt", prompt)
-        sampling_params_list = msg["sampling_params_list"]
+        request_id = msg.request_id
+        prompt = msg.prompt
+        original_prompt = msg.original_prompt
+        sampling_params_list = msg.sampling_params_list
         if not sampling_params_list:
             raise ValueError(f"Missing sampling params for stage 0. Got {len(sampling_params_list)} stage params.")
-        final_stage_id = msg["final_stage_id"]
+        final_stage_id = msg.final_stage_id
 
         logger.debug(
             "[Orchestrator] _handle_add_request: stage=%s req=%s "
@@ -249,27 +263,27 @@ class Orchestrator:
         self.request_states[request_id] = req_state
         req_state.streaming.enabled = bool(getattr(prompt, "resumable", False))
         req_state.stage_submit_ts[stage_id] = _time.time()
-        enqueue_ts = msg.get("enqueue_ts", 0.0)
+        enqueue_ts = msg.enqueue_ts
         if enqueue_ts > 0:
             req_state.pipeline_timings["queue_wait_ms"] = (_time.perf_counter() - enqueue_ts) * 1000.0
-        preprocess_ms = msg.get("preprocess_ms", 0.0)
+        preprocess_ms = msg.preprocess_ms
         if preprocess_ms > 0:
             req_state.pipeline_timings["preprocess_ms"] = preprocess_ms
         await self.stage_pools[stage_id].submit_initial(
             request_id,
             req_state,
             prompt,
-            prompt_text=msg.get("output_prompt_text"),
+            prompt_text=msg.output_prompt_text,
         )
 
         if self.async_chunk and stage_id == 0 and final_stage_id > 0:
             await self._prewarm_async_chunk_stages(request_id, prompt, req_state)
 
-    async def _handle_streaming_update(self, msg: dict[str, Any]) -> None:
+    async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None:
         """Handle a streaming_update message for an existing request."""
         stage_id = 0
-        request_id = msg["request_id"]
-        request = msg["prompt"]
+        request_id = msg.request_id
+        request = msg.prompt
 
         req_state = self.request_states.get(request_id)
         if req_state is None:
@@ -277,13 +291,22 @@ class Orchestrator:
                 "[Orchestrator] streaming_update for unknown req=%s, falling back to add_request",
                 request_id,
             )
-            fallback_msg = dict(msg)
-            fallback_msg["type"] = "add_request"
+            fallback_msg = StageSubmissionMessage(
+                type="add_request",
+                request_id=msg.request_id,
+                prompt=msg.prompt,
+                original_prompt=msg.original_prompt,
+                output_prompt_text=msg.output_prompt_text,
+                sampling_params_list=msg.sampling_params_list,
+                final_stage_id=msg.final_stage_id,
+                preprocess_ms=msg.preprocess_ms,
+                enqueue_ts=msg.enqueue_ts,
+            )
             await self._handle_add_request(fallback_msg)
             return
 
-        if "sampling_params_list" in msg and msg["sampling_params_list"]:
-            req_state.sampling_params_list = msg["sampling_params_list"]
+        if msg.sampling_params_list:
+            req_state.sampling_params_list = msg.sampling_params_list
 
         req_state.streaming.enabled = True
         req_state.stage_submit_ts[stage_id] = _time.time()
@@ -291,16 +314,16 @@ class Orchestrator:
             request_id,
             req_state,
             request,
-            prompt_text=msg.get("output_prompt_text"),
+            prompt_text=msg.output_prompt_text,
         )
 
-    async def _handle_add_companion(self, msg: dict[str, Any]) -> None:
+    async def _handle_add_companion(self, msg: AddCompanionRequestMessage) -> None:
         """Handle an add_companion_request message: submit companion to stage 0."""
-        companion_id = msg["companion_id"]
-        parent_id = msg["parent_id"]
-        role = msg["role"]
-        companion_prompt = msg["prompt"]
-        sampling_params_list = msg["sampling_params_list"]
+        companion_id = msg.companion_id
+        parent_id = msg.parent_id
+        role = msg.role
+        companion_prompt = msg.prompt
+        sampling_params_list = msg.sampling_params_list
 
         parent_state = self.request_states.get(parent_id)
         if parent_state is None:
@@ -326,7 +349,7 @@ class Orchestrator:
             companion_id,
             companion_state,
             companion_prompt,
-            prompt_text=msg.get("companion_prompt_text"),
+            prompt_text=msg.companion_prompt_text,
             affinity_request_id=parent_id,
         )
 
@@ -338,9 +361,9 @@ class Orchestrator:
             companion_replica_id,
         )
 
-    async def _handle_abort(self, msg: dict[str, Any]) -> None:
+    async def _handle_abort(self, msg: AbortRequestMessage) -> None:
         """Handle an abort message from the main thread."""
-        request_ids = msg["request_ids"]
+        request_ids = msg.request_ids
         await self._cleanup_request_ids(
             self._cfg_tracker.abort_parents(request_ids),
             abort=True,
@@ -359,14 +382,14 @@ class Orchestrator:
         for pool in self.stage_pools:
             pool.release_bindings(request_ids)
 
-    async def _handle_collective_rpc(self, msg: dict[str, Any]) -> None:
+    async def _handle_collective_rpc(self, msg: CollectiveRPCRequestMessage) -> None:
         """Handle a control-plane RPC request from the main thread."""
-        rpc_id = msg["rpc_id"]
-        method = msg["method"]
-        timeout = msg.get("timeout")
-        args = tuple(msg.get("args", ()))
-        kwargs = dict(msg.get("kwargs") or {})
-        requested_stage_ids = msg.get("stage_ids")
+        rpc_id = msg.rpc_id
+        method = msg.method
+        timeout = msg.timeout
+        args = msg.args
+        kwargs = dict(msg.kwargs)
+        requested_stage_ids = msg.stage_ids
 
         target_pools: list[StagePool] = []
         if requested_stage_ids is None:
@@ -393,13 +416,12 @@ class Orchestrator:
                 results.append(stage_result)
 
         await self.rpc_async_queue.put(
-            {
-                "type": "collective_rpc_result",
-                "rpc_id": rpc_id,
-                "method": method,
-                "stage_ids": stage_ids,
-                "results": results,
-            }
+            CollectiveRPCResultMessage(
+                rpc_id=rpc_id,
+                method=method,
+                stage_ids=stage_ids,
+                results=results,
+            )
         )
 
     # ---- Orchestration loop ----
@@ -469,13 +491,12 @@ class Orchestrator:
                             for req_id, req_state in list(self.request_states.items()):
                                 if stage_id in req_state.stage_submit_ts:
                                     await self.output_async_queue.put(
-                                        {
-                                            "type": "error",
-                                            "error": str(e),
-                                            "fatal": True,
-                                            "request_id": req_id,
-                                            "stage_id": stage_id,
-                                        }
+                                        ErrorMessage(
+                                            error=str(e),
+                                            fatal=True,
+                                            request_id=req_id,
+                                            stage_id=stage_id,
+                                        )
                                     )
                                     self.request_states.pop(req_id, None)
                             self._shutdown_event.set()
@@ -534,12 +555,11 @@ class Orchestrator:
         else:
             parent_id = output.request_id
         await self.output_async_queue.put(
-            {
-                "type": "error",
-                "request_id": parent_id,
-                "stage_id": stage_id,
-                "error": output.error,
-            }
+            ErrorMessage(
+                request_id=parent_id,
+                stage_id=stage_id,
+                error=output.error,
+            )
         )
         await self._cleanup_request_ids(
             [parent_id, *self._cfg_tracker.cleanup_parent(parent_id)],
@@ -596,25 +616,23 @@ class Orchestrator:
 
         if self.stage_pools[stage_id].final_output:
             await self.output_async_queue.put(
-                {
-                    "type": "output",
-                    "request_id": req_id,
-                    "stage_id": stage_id,
-                    "engine_outputs": output,
-                    "metrics": stage_metrics,
-                    "finished": finished and stage_id == req_state.final_stage_id,
-                    "stage_submit_ts": submit_ts,
-                }
+                OutputMessage(
+                    request_id=req_id,
+                    stage_id=stage_id,
+                    engine_outputs=output,
+                    metrics=stage_metrics,
+                    finished=finished and stage_id == req_state.final_stage_id,
+                    stage_submit_ts=submit_ts,
+                )
             )
         elif stage_metrics is not None:
             await self.output_async_queue.put(
-                {
-                    "type": "stage_metrics",
-                    "request_id": req_id,
-                    "stage_id": stage_id,
-                    "metrics": stage_metrics,
-                    "stage_submit_ts": submit_ts,
-                }
+                StageMetricsMessage(
+                    request_id=req_id,
+                    stage_id=stage_id,
+                    metrics=stage_metrics,
+                    stage_submit_ts=submit_ts,
+                )
             )
 
         if self._pd_pair is not None and finished and stage_id == self._pd_pair[0]:
@@ -775,79 +793,120 @@ class Orchestrator:
         """Forward output from the current logical stage to the next one."""
         next_logical = src_stage_id + 1
         next_pool = self.stage_pools[next_logical]
-        next_client = next_pool.stage_client
+
+        if next_pool.stage_type == "diffusion":
+            await self._forward_to_diffusion_stage(
+                req_id,
+                src_stage_id,
+                output,
+                req_state,
+            )
+        else:
+            await self._forward_to_llm_stage(
+                req_id,
+                src_stage_id,
+                output,
+                req_state,
+                is_streaming_session=is_streaming_session,
+                is_final_update=is_final_update,
+            )
+
+    async def _forward_to_diffusion_stage(
+        self,
+        req_id: str,
+        src_stage_id: int,
+        output: Any,
+        req_state: OrchestratorRequestState,
+    ) -> None:
+        """Forward output from an LLM stage to the next diffusion stage."""
+        next_logical = src_stage_id + 1
+        next_pool = self.stage_pools[next_logical]
+        params = req_state.sampling_params_list[next_logical]
+        source_outputs = [output]
+        already_submitted = self._next_stage_already_submitted(src_stage_id, req_state)
+
+        diffusion_client = next_pool.stage_client
+        if diffusion_client.custom_process_input_func is not None:
+            _t_ar2d = _time.perf_counter()
+            diffusion_prompt = diffusion_client.custom_process_input_func(
+                source_outputs,
+                req_state.prompt,
+                getattr(diffusion_client, "requires_multimodal_data", False),
+            )
+            _dt_ar2d = (_time.perf_counter() - _t_ar2d) * 1000
+            req_state.pipeline_timings["ar2diffusion_ms"] = _dt_ar2d
+            logger.info(
+                "[Orchestrator] ar2diffusion req=%s wall_time=%.3fms stage=%d->%d",
+                req_id,
+                _dt_ar2d,
+                src_stage_id,
+                next_logical,
+            )
+            if isinstance(diffusion_prompt, list):
+                if not diffusion_prompt:
+                    error_output = OmniRequestOutput.from_error(
+                        req_id,
+                        f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}",
+                    )
+                    logger.warning(
+                        "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; "
+                        "routing terminal error output",
+                        req_id,
+                        src_stage_id,
+                        next_logical,
+                    )
+                    await self.output_async_queue.put(
+                        OutputMessage(
+                            request_id=req_id,
+                            stage_id=next_logical,
+                            engine_outputs=error_output,
+                            metrics=None,
+                            finished=True,
+                        )
+                    )
+                    await self._cleanup_request_ids(
+                        [req_id, *self._cfg_tracker.cleanup_parent(req_id)],
+                    )
+                    return
+                if already_submitted and len(diffusion_prompt) == 1:
+                    diffusion_prompt = diffusion_prompt[0]
+        else:
+            diffusion_prompt = req_state.prompt
+
+        if already_submitted:
+            await next_pool.submit_update(req_id, req_state, diffusion_prompt)
+        else:
+            await next_pool.submit_initial(
+                req_id,
+                req_state,
+                diffusion_prompt,
+                submit_kwargs={
+                    "kv_sender_info": self._build_kv_sender_info(
+                        list(getattr(diffusion_client, "engine_input_source", None) or [src_stage_id]),
+                        request_id=req_id,
+                    )
+                },
+                params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params),
+            )
+        req_state.stage_submit_ts[next_logical] = _time.time()
+
+    async def _forward_to_llm_stage(
+        self,
+        req_id: str,
+        src_stage_id: int,
+        output: Any,
+        req_state: OrchestratorRequestState,
+        *,
+        is_streaming_session: bool = False,
+        is_final_update: bool = False,
+    ) -> None:
+        """Forward output from the current stage to the next LLM stage."""
+        next_logical = src_stage_id + 1
+        next_pool = self.stage_pools[next_logical]
         params = req_state.sampling_params_list[next_logical]
         source_outputs = [output]
         next_stage_resumable = is_streaming_session and not is_final_update
         already_submitted = self._next_stage_already_submitted(src_stage_id, req_state)
-        requires_multimodal_data = getattr(next_client, "requires_multimodal_data", False)
-
-        if next_pool.stage_type == "diffusion":
-            if next_client.custom_process_input_func is not None:
-                _t_ar2d = _time.perf_counter()
-                diffusion_prompt = next_client.custom_process_input_func(
-                    source_outputs,
-                    req_state.prompt,
-                    requires_multimodal_data,
-                )
-                _dt_ar2d = (_time.perf_counter() - _t_ar2d) * 1000
-                req_state.pipeline_timings["ar2diffusion_ms"] = _dt_ar2d
-                logger.info(
-                    "[Orchestrator] ar2diffusion req=%s wall_time=%.3fms stage=%d->%d",
-                    req_id,
-                    _dt_ar2d,
-                    src_stage_id,
-                    next_logical,
-                )
-                if isinstance(diffusion_prompt, list):
-                    if not diffusion_prompt:
-                        error_output = OmniRequestOutput.from_error(
-                            req_id,
-                            f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}",
-                        )
-                        logger.warning(
-                            "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; "
-                            "routing terminal error output",
-                            req_id,
-                            src_stage_id,
-                            next_logical,
-                        )
-                        await self.output_async_queue.put(
-                            {
-                                "type": "output",
-                                "request_id": req_id,
-                                "stage_id": next_logical,
-                                "engine_outputs": error_output,
-                                "metrics": None,
-                                "finished": True,
-                            }
-                        )
-                        await self._cleanup_request_ids(
-                            [req_id, *self._cfg_tracker.cleanup_parent(req_id)],
-                        )
-                        return
-                    if already_submitted and len(diffusion_prompt) == 1:
-                        diffusion_prompt = diffusion_prompt[0]
-            else:
-                diffusion_prompt = req_state.prompt
-
-            if already_submitted:
-                await next_pool.submit_update(req_id, req_state, diffusion_prompt)
-            else:
-                await next_pool.submit_initial(
-                    req_id,
-                    req_state,
-                    diffusion_prompt,
-                    submit_kwargs={
-                        "kv_sender_info": self._build_kv_sender_info(
-                            list(getattr(next_client, "engine_input_source", None) or [src_stage_id]),
-                            request_id=req_id,
-                        )
-                    },
-                    params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params),
-                )
-            req_state.stage_submit_ts[next_logical] = _time.time()
-            return
 
         # PD disaggregation: prefill → decode routing uses original prompt + KV transfer params
         if self._pd_pair is not None and (src_stage_id, next_logical) == self._pd_pair:
@@ -894,6 +953,7 @@ class Orchestrator:
                 {},
             )[req_id] = req_state.pd_prefill_multimodal_output
 
+        next_client = next_pool.llm_stage_client
         try:
             next_inputs = next_client.process_engine_inputs(
                 source_outputs,
@@ -912,7 +972,7 @@ class Orchestrator:
         for next_input in next_inputs:
             # Only AR thinker stages consume encoder mm_features; downstream
             # (talker/code2wav/…) must not see them (avoids encoder-cache misses).
-            model_stage = getattr(next_client, "model_stage", None)
+            model_stage = next_client.model_stage
             mm_features = req_state.mm_features if model_stage == "thinker" else None
             request = build_engine_core_request_from_tokens(
                 request_id=req_id,
@@ -962,7 +1022,7 @@ class Orchestrator:
                     req_state.prompt,
                     submit_kwargs={
                         "kv_sender_info": self._build_kv_sender_info(
-                            list(getattr(next_pool.stage_client, "engine_input_source", None) or [next_stage_id - 1]),
+                            list(next_pool.stage_client.engine_input_source or [next_stage_id - 1]),
                             request_id=request_id,
                         )
                     },
@@ -1009,14 +1069,14 @@ class Orchestrator:
                 continue
 
             sender_pool = self.stage_pools[sender_stage_id]
-            sender_stage = sender_pool.get_bound_client(request_id) if request_id is not None else None
-            if sender_stage is None:
-                sender_stage = sender_pool.stage_client
-            get_sender_info = getattr(sender_stage, "get_kv_sender_info", None)
-            if not callable(get_sender_info):
+            if sender_pool.stage_type == "diffusion":
                 continue
 
-            sender_info = get_sender_info()
+            sender_stage = sender_pool.get_bound_llm_client(request_id) if request_id is not None else None
+            if sender_stage is None:
+                sender_stage = sender_pool.llm_stage_client
+
+            sender_info = sender_stage.get_kv_sender_info()
             if not sender_info:
                 logger.warning(
                     "[Orchestrator] Stage-%s has no KV sender info available",
@@ -1050,16 +1110,15 @@ class Orchestrator:
                 msg = self.request_async_queue.get_nowait()
             except Exception:
                 break
-            if msg.get("type") == "add_request":
-                req_id = msg["request_id"]
+            if isinstance(msg, StageSubmissionMessage) and msg.type == "add_request":
+                req_id = msg.request_id
                 await self.output_async_queue.put(
-                    {
-                        "type": "error",
-                        "error": self._fatal_error,
-                        "fatal": True,
-                        "request_id": req_id,
-                        "stage_id": self._fatal_error_stage_id,
-                    }
+                    ErrorMessage(
+                        error=self._fatal_error,
+                        fatal=True,
+                        request_id=req_id,
+                        stage_id=self._fatal_error_stage_id,
+                    )
                 )
                 notified.add(req_id)
 
@@ -1069,13 +1128,12 @@ class Orchestrator:
         for req_id in list(self.request_states):
             if req_id not in notified:
                 await self.output_async_queue.put(
-                    {
-                        "type": "error",
-                        "error": self._fatal_error,
-                        "fatal": True,
-                        "request_id": req_id,
-                        "stage_id": self._fatal_error_stage_id,
-                    }
+                    ErrorMessage(
+                        error=self._fatal_error,
+                        fatal=True,
+                        request_id=req_id,
+                        stage_id=self._fatal_error_stage_id,
+                    )
                 )
             self.request_states.pop(req_id, None)
 
diff --git a/vllm_omni/engine/output_modality.py b/vllm_omni/engine/output_modality.py
index c2615ca8c..5887bb910 100644
--- a/vllm_omni/engine/output_modality.py
+++ b/vllm_omni/engine/output_modality.py
@@ -9,6 +9,9 @@ from __future__ import annotations
 
 import re
 from enum import Enum, Flag, StrEnum, auto
+from typing import Literal, TypeAlias
+
+FinalOutputModalityType: TypeAlias = Literal["text", "image", "audio", "video"]
 
 _MODALITY_ALIASES: dict[str, str] = {
     "speech": "audio",
diff --git a/vllm_omni/engine/stage_client.py b/vllm_omni/engine/stage_client.py
new file mode 100644
index 000000000..1f10a219c
--- /dev/null
+++ b/vllm_omni/engine/stage_client.py
@@ -0,0 +1,111 @@
+"""Shared stage-client typing for vLLM-Omni runtime surfaces."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any, Protocol
+
+if TYPE_CHECKING:
+    from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest
+
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniTokensPrompt
+    from vllm_omni.outputs import OmniRequestOutput
+
+from vllm_omni.engine.output_modality import FinalOutputModalityType
+from vllm_omni.inputs.data import OmniSamplingParams
+
+
+class StageClient(Protocol):
+    """Shared metadata and lifecycle surface for all stage clients.
+
+    This protocol intentionally covers the common attributes consumed by
+    orchestration and entrypoint code. Backend-specific request APIs stay
+    typed with the concrete client union.
+    """
+
+    stage_id: int
+    replica_id: int
+    stage_type: str
+    final_output: bool
+    final_output_type: FinalOutputModalityType | None
+    default_sampling_params: OmniSamplingParams
+    requires_multimodal_data: bool
+    custom_process_input_func: Any | None
+    engine_input_source: Sequence[int]
+    is_comprehension: bool
+
+    def check_health(self) -> None: ...
+
+    def shutdown(self) -> None: ...
+
+
+class StageClientBase:
+    """Runtime nominal base for stage clients.
+
+    Keeping the runtime base separate from the protocol avoids Protocol-related
+    MRO interference with concrete client super() calls.
+    """
+
+    pass
+
+
+class StagePoolClient(StageClient, Protocol):
+    """Common pool-facing client surface shared by every stage backend."""
+
+    async def abort_requests_async(self, request_ids: list[str]) -> None: ...
+
+    async def collective_rpc_async(
+        self,
+        method: str,
+        timeout: float | None = None,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> Any: ...
+
+
+class StagePoolLLMClient(StagePoolClient, Protocol):
+    """Pool-facing API for LLM-style stages."""
+
+    model_stage: str | None
+
+    async def add_request_async(self, request: EngineCoreRequest) -> None: ...
+
+    async def get_output_async(self) -> EngineCoreOutputs: ...
+
+    def set_engine_outputs(self, engine_outputs: EngineCoreOutput) -> None: ...
+
+    def process_engine_inputs(
+        self,
+        source_outputs: list[Any],
+        prompt: Any = None,
+        streaming_context: Any | None = None,
+    ) -> list[OmniTokensPrompt]: ...
+
+    def get_kv_sender_info(
+        self,
+        *,
+        base_port: int = ...,
+        kv_transfer_port_offset: int = ...,
+    ) -> dict[str, Any] | None: ...
+
+
+class StagePoolDiffusionClient(StagePoolClient, Protocol):
+    """Pool-facing API for diffusion stages."""
+
+    async def add_request_async(
+        self,
+        request_id: str,
+        prompt: OmniPromptType,
+        sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[int, dict[str, Any]] | None = None,
+    ) -> None: ...
+
+    async def add_batch_request_async(
+        self,
+        request_id: str,
+        prompts: list[OmniPromptType],
+        sampling_params: OmniDiffusionSamplingParams,
+        kv_sender_info: dict[int, dict[str, Any]] | None = None,
+    ) -> None: ...
+
+    def get_diffusion_output_nowait(self) -> OmniRequestOutput | None: ...
diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py
index 04a047fb3..f2d4779a9 100644
--- a/vllm_omni/engine/stage_engine_core_client.py
+++ b/vllm_omni/engine/stage_engine_core_client.py
@@ -23,6 +23,7 @@ from vllm.v1.engine.exceptions import EngineDeadError
 from vllm_omni.distributed.omni_connectors.utils.initialization import (
     KV_TRANSFER_PORT_OFFSET,
 )
+from vllm_omni.engine.stage_client import StageClientBase
 from vllm_omni.engine.stage_init_utils import StageMetadata
 
 if TYPE_CHECKING:
@@ -56,7 +57,7 @@ def _default_process_engine_inputs(
     ]
 
 
-class StageEngineCoreClientBase:
+class StageEngineCoreClientBase(StageClientBase):
     """Shared stage-aware behavior for async EngineCore clients.
 
     The concrete transport/load-balancing behavior is supplied by the
diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py
index 6f7454271..427368a02 100644
--- a/vllm_omni/engine/stage_pool.py
+++ b/vllm_omni/engine/stage_pool.py
@@ -5,11 +5,16 @@ from __future__ import annotations
 import asyncio
 import time as _time
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreOutputs
 
+from vllm_omni.engine.stage_client import (
+    StagePoolClient,
+    StagePoolDiffusionClient,
+    StagePoolLLMClient,
+)
 from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics
 from vllm_omni.metrics.stats import StageStats
 from vllm_omni.metrics.utils import count_tokens_from_outputs
@@ -35,20 +40,20 @@ class StagePool:
     def __init__(
         self,
         stage_id: int,
-        clients: Any | list[Any],
+        clients: StagePoolClient | list[StagePoolClient],
         *,
         output_processor: Any = None,
         stage_vllm_config: Any = None,
     ) -> None:
         if isinstance(clients, list):
-            normalized_clients = list(clients)
+            normalized_clients: list[StagePoolClient] = list(clients)
         else:
             normalized_clients = [clients]
 
         if not normalized_clients:
             raise ValueError(f"StagePool for stage {stage_id} has no replicas")
         self.stage_id = stage_id
-        self.clients: list[Any] = normalized_clients
+        self.clients: list[StagePoolClient] = normalized_clients
         self._output_processor = output_processor
         self._stage_vllm_config = stage_vllm_config
         self._next_replica_id = 0
@@ -63,16 +68,20 @@ class StagePool:
 
     @property
     def stage_type(self) -> str | None:
-        return getattr(self.stage_client, "stage_type", None)
+        return self.stage_client.stage_type
 
     @property
     def final_output(self) -> bool:
-        return bool(getattr(self.clients[0], "final_output", False))
+        return self.clients[0].final_output
 
     @property
-    def stage_client(self) -> Any:
+    def stage_client(self) -> StagePoolClient:
         return self.clients[0]
 
+    @property
+    def llm_stage_client(self) -> StagePoolLLMClient:
+        return cast(StagePoolLLMClient, self.stage_client)
+
     @property
     def stage_vllm_config(self) -> Any:
         return self._stage_vllm_config
@@ -87,13 +96,20 @@ class StagePool:
         """Return the currently bound replica id for *request_id* if present."""
         return self._request_bindings.get(request_id)
 
-    def get_bound_client(self, request_id: str) -> Any | None:
+    def get_bound_client(self, request_id: str) -> StagePoolClient | None:
         """Return the currently bound client for *request_id* if present."""
         replica_id = self.get_bound_replica_id(request_id)
         if replica_id is None:
             return None
         return self.clients[replica_id]
 
+    def get_bound_llm_client(self, request_id: str) -> StagePoolLLMClient | None:
+        """Return the currently bound LLM client for *request_id* if present."""
+        client = self.get_bound_client(request_id)
+        if client is None:
+            return None
+        return cast(StagePoolLLMClient, client)
+
     def release_binding(self, request_id: str) -> None:
         """Drop the route binding for *request_id* in this stage."""
         self._request_bindings.pop(request_id, None)
@@ -125,6 +141,12 @@ class StagePool:
         self._request_bindings[request_id] = chosen
         return chosen
 
+    def _llm_client(self, replica_id: int) -> StagePoolLLMClient:
+        return cast(StagePoolLLMClient, self.clients[replica_id])
+
+    def _diffusion_client(self, replica_id: int) -> StagePoolDiffusionClient:
+        return cast(StagePoolDiffusionClient, self.clients[replica_id])
+
     # ---- Metrics ----
 
     def build_stage_metrics(
@@ -188,7 +210,7 @@ class StagePool:
                 request_id,
                 affinity_request_id=affinity_request_id,
             )
-            client = self.clients[replica_id]
+            client = self._diffusion_client(replica_id)
             if isinstance(request, list):
                 await client.add_batch_request_async(request_id, request, params, **submit_kwargs)
             else:
@@ -212,7 +234,7 @@ class StagePool:
             raise
 
         try:
-            await self.clients[replica_id].add_request_async(request, **submit_kwargs)
+            await self._llm_client(replica_id).add_request_async(request, **submit_kwargs)
         except Exception:
             self.release_binding(request_id)
             rollback = getattr(self.output_processor, "remove_request", None)
@@ -244,7 +266,7 @@ class StagePool:
             replica_id = self.select_replica_id(request_id)
 
         if self.stage_type == "diffusion":
-            await self.clients[replica_id].add_request_async(request_id, request, params)
+            await self._diffusion_client(replica_id).add_request_async(request_id, request, params)
         else:
             # Refresh the shared output-processor state before yielding to the
             # stage client so streaming segments are merged against the latest
@@ -256,12 +278,12 @@ class StagePool:
                 request_index=0,
                 queue=None,
             )
-            await self.clients[replica_id].add_request_async(request)
+            await self._llm_client(replica_id).add_request_async(request)
         return replica_id
 
     # ---- Stage-local polling ----
 
-    async def _poll_stage_raw(self, client: Any) -> EngineCoreOutputs | None:
+    async def _poll_stage_raw(self, client: StagePoolLLMClient) -> EngineCoreOutputs | None:
         """Pull raw EngineCoreOutputs from a stage replica without processing."""
         outputs = await client.get_output_async()
         if not outputs.outputs:
@@ -274,7 +296,7 @@ class StagePool:
         raw_outputs: EngineCoreOutputs,
     ) -> list[Any]:
         """Run the shared LLM output processor on one raw poll result."""
-        client = self.clients[replica_id]
+        client = self._llm_client(replica_id)
         processor = self.output_processor
         processed = processor.process_outputs(
             raw_outputs.outputs,
@@ -297,7 +319,7 @@ class StagePool:
         timeout_s: float = 0.001,
     ) -> EngineCoreOutputs | None:
         """Poll raw EngineCore outputs from one LLM replica once."""
-        client = self.clients[replica_id]
+        client = self._llm_client(replica_id)
         try:
             return await asyncio.wait_for(
                 self._poll_stage_raw(client),
@@ -317,7 +339,7 @@ class StagePool:
 
     def poll_diffusion_output(self, replica_id: int) -> Any | None:
         """Drain one ready diffusion output from the given replica if present."""
-        return self.clients[replica_id].get_diffusion_output_nowait()
+        return self._diffusion_client(replica_id).get_diffusion_output_nowait()
 
     # ---- Stage-local control plane ----
 
@@ -360,18 +382,12 @@ class StagePool:
         kwargs = dict(kwargs or {})
         client = self.clients[replica_id]
         try:
-            if hasattr(client, "collective_rpc_async"):
-                return await client.collective_rpc_async(
-                    method=method,
-                    timeout=timeout,
-                    args=args,
-                    kwargs=kwargs,
-                )
-            return {
-                "supported": False,
-                "todo": True,
-                "reason": f"{client.__class__.__name__}.collective_rpc_async is not implemented yet",
-            }
+            return await client.collective_rpc_async(
+                method=method,
+                timeout=timeout,
+                args=args,
+                kwargs=kwargs,
+            )
         except Exception as exc:
             logger.exception(
                 "[StagePool] collective_rpc failed: stage=%s replica=%s method=%s",
diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py
index 056f56c00..c22e78096 100644
--- a/vllm_omni/entrypoints/async_omni.py
+++ b/vllm_omni/entrypoints/async_omni.py
@@ -26,6 +26,7 @@ from vllm.tasks import SupportedTask
 from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask
+from vllm_omni.engine.messages import ErrorMessage, OutputMessage
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni_base import (
     OmniBase,
@@ -416,7 +417,12 @@ class AsyncOmni(EngineClient, OmniBase):
             except (asyncio.CancelledError, GeneratorExit):
                 cancelled = True
             except Exception as error:
-                await req_state.queue.put({"request_id": request_id, "error": error})
+                await req_state.queue.put(
+                    ErrorMessage(
+                        request_id=request_id,
+                        error=str(error),
+                    )
+                )
             finally:
                 if not cancelled:
                     # Send empty final request to indicate that inputs have
@@ -501,23 +507,25 @@ class AsyncOmni(EngineClient, OmniBase):
         while True:
             result = await req_state.queue.get()
 
-            stage_id = result.get("stage_id", 0)
-
-            if result.get("type") == "error" and result.get("fatal"):
-                raise OmniEngineDeadError(
-                    result.get("error", ""),
-                    error_stage_id=result.get("stage_id"),
-                )
-
-            # Check for errors
-            if "error" in result:
+            if isinstance(result, ErrorMessage):
                 logger.error(
                     "[AsyncOmni] Orchestrator error for req=%s stage-%s: %s",
                     request_id,
-                    stage_id,
-                    result["error"],
+                    result.stage_id,
+                    result.error,
                 )
-                raise RuntimeError(result)
+                if result.fatal:
+                    raise OmniEngineDeadError(
+                        result.error,
+                        error_stage_id=result.stage_id,
+                    )
+                raise RuntimeError(result.error)
+
+            if not isinstance(result, OutputMessage):
+                logger.warning("[AsyncOmni] Dropping unexpected per-request message %r", result)
+                continue
+
+            stage_id = result.stage_id
 
             self._check_engine_output_error(result, request_id, stage_id)
 
@@ -541,7 +549,7 @@ class AsyncOmni(EngineClient, OmniBase):
                 yield output_to_yield
 
             # The Orchestrator sets "finished" when the final stage is done
-            if result.get("finished"):
+            if result.finished:
                 break
 
     # ==================== Output Handler ====================
@@ -596,29 +604,29 @@ class AsyncOmni(EngineClient, OmniBase):
             except OmniEngineDeadError as e:
                 logger.error("[AsyncOmni] Engine dead: %s", e)
                 for req_state in list(self.request_states.values()):
-                    error_msg = {
-                        "type": "error",
-                        "error": str(e),
-                        "fatal": True,
-                        "request_id": req_state.request_id,
-                    }
-                    if e.error_stage_id is not None:
-                        error_msg["stage_id"] = e.error_stage_id
+                    error_msg = ErrorMessage(
+                        error=str(e),
+                        fatal=True,
+                        request_id=req_state.request_id,
+                        stage_id=e.error_stage_id,
+                    )
                     await req_state.queue.put(error_msg)
             except EngineDeadError as e:
                 logger.error("[AsyncOmni] Engine dead: %s", e)
                 for req_state in list(self.request_states.values()):
-                    error_msg = {
-                        "type": "error",
-                        "error": str(e),
-                        "fatal": True,
-                        "request_id": req_state.request_id,
-                    }
+                    error_msg = ErrorMessage(
+                        error=str(e),
+                        fatal=True,
+                        request_id=req_state.request_id,
+                    )
                     await req_state.queue.put(error_msg)
             except Exception as e:
                 logger.exception("[AsyncOmni] final_output_loop failed.")
                 for req_state in list(self.request_states.values()):
-                    error_msg = {"request_id": req_state.request_id, "error": str(e)}
+                    error_msg = ErrorMessage(
+                        request_id=req_state.request_id,
+                        error=str(e),
+                    )
                     await req_state.queue.put(error_msg)
                 self.final_output_task = None
 
diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py
index 223c208af..7b1b15e48 100644
--- a/vllm_omni/entrypoints/omni.py
+++ b/vllm_omni/entrypoints/omni.py
@@ -10,6 +10,7 @@ from tqdm.auto import tqdm
 from vllm.logger import init_logger
 from vllm.sampling_params import RequestOutputKind
 
+from vllm_omni.engine.messages import OutputMessage
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni_base import OmniBase
 from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics
@@ -33,7 +34,7 @@ class Omni(OmniBase):
         for stage_id, params in enumerate(sampling_params_list):
             sp = copy.deepcopy(params)
             stage_meta = self.engine.get_stage_metadata(stage_id)
-            if stage_meta.get("stage_type") != "diffusion" and hasattr(sp, "output_kind"):
+            if stage_meta.stage_type != "diffusion" and hasattr(sp, "output_kind"):
                 sp.output_kind = RequestOutputKind.FINAL_ONLY
             effective_params.append(sp)
         return effective_params
@@ -181,7 +182,7 @@ class Omni(OmniBase):
                 if output_to_yield is not None:
                     yield output_to_yield
 
-                if msg.get("finished"):
+                if isinstance(msg, OutputMessage) and msg.finished:
                     active_reqs.discard(req_id)
                     if pbar is not None:
                         pbar.update(1)
diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py
index 4e84d6200..78f2b7985 100644
--- a/vllm_omni/entrypoints/omni_base.py
+++ b/vllm_omni/entrypoints/omni_base.py
@@ -4,7 +4,6 @@ import argparse
 import os
 import sys
 import time
-import types
 import warnings
 import weakref
 from collections.abc import Sequence
@@ -15,6 +14,12 @@ from vllm.logger import init_logger
 from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.engine.messages import (
+    EngineQueueMessage,
+    ErrorMessage,
+    OutputMessage,
+    StageMetricsMessage,
+)
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
 from vllm_omni.entrypoints.utils import coerce_param_message_types, get_final_stage_id_for_e2e
@@ -189,12 +194,10 @@ class OmniBase(PDDisaggregationMixin):
         self.default_sampling_params_list = self.engine.default_sampling_params_list
         if not self.output_modalities:
             self.output_modalities = [
-                self.engine.get_stage_metadata(i).get("final_output_type") for i in range(self.engine.num_stages)
+                self.engine.get_stage_metadata(i).final_output_type for i in range(self.engine.num_stages)
             ]
 
-        self._stage_meta_list = [
-            types.SimpleNamespace(**self.engine.get_stage_metadata(i)) for i in range(self.engine.num_stages)
-        ]
+        self._stage_meta_list = [self.engine.get_stage_metadata(i) for i in range(self.engine.num_stages)]
 
         logger.info(
             "[%s] Initialized with %s stages for model %s",
@@ -289,17 +292,15 @@ class OmniBase(PDDisaggregationMixin):
             self._stage_meta_list,
         )
 
-    def _process_stage_metrics_message(self, msg: dict[str, Any]) -> None:
-        req_id = msg.get("request_id")
+    def _process_stage_metrics_message(self, msg: StageMetricsMessage) -> None:
+        req_id = msg.request_id
         req_state = self.request_states.get(req_id)
         if req_state is None or req_state.metrics is None:
             return
-        _m = msg.get("metrics")
-        if _m is None:
-            return
-        stage_id = msg.get("stage_id", 0)
+        _m = msg.metrics
+        stage_id = msg.stage_id
         req_state.metrics.on_stage_metrics(stage_id, req_id, _m)
-        submit_ts = msg.get("stage_submit_ts")
+        submit_ts = msg.stage_submit_ts
         now = time.time()
         if req_state.metrics.stage_first_ts[stage_id] is None:
             req_state.metrics.stage_first_ts[stage_id] = submit_ts if submit_ts is not None else now
@@ -307,40 +308,30 @@ class OmniBase(PDDisaggregationMixin):
 
     def _handle_output_message(
         self,
-        msg: dict[str, Any] | None,
+        msg: EngineQueueMessage | None,
     ) -> OutputMessageHandleResult:
         """Handle one Orchestrator output-queue message."""
         if msg is None:
             return True, None, None, None
 
-        msg_type = msg.get("type")
-        if msg_type == "stage_metrics":
+        if isinstance(msg, StageMetricsMessage):
             self._process_stage_metrics_message(msg)
             return True, None, None, None
 
-        if msg_type == "error":
-            error_text = msg.get("error", "Orchestrator returned an error message")
-            stage_id = msg.get("stage_id")
-            if msg.get("fatal"):
+        if isinstance(msg, ErrorMessage):
+            if msg.fatal:
                 raise OmniEngineDeadError(
-                    error_text,
-                    error_stage_id=stage_id,
+                    msg.error,
+                    error_stage_id=msg.stage_id,
                 )
-            raise RuntimeError(error_text)
+            raise RuntimeError(msg.error)
 
-        if msg_type != "output":
-            logger.warning("[%s] got unexpected msg type: %s", self.__class__.__name__, msg_type)
+        if not isinstance(msg, OutputMessage):
+            logger.warning("[%s] got unexpected msg type: %s", self.__class__.__name__, msg.type)
             return True, None, None, None
 
-        req_id = msg.get("request_id")
-        if req_id is None:
-            logger.warning("[%s] got output message without request_id", self.__class__.__name__)
-            return True, None, None, None
-
-        stage_id = msg.get("stage_id")
-        if stage_id is None:
-            logger.warning("[%s] got output message without stage_id for req=%s", self.__class__.__name__, req_id)
-            return True, None, None, None
+        req_id = msg.request_id
+        stage_id = msg.stage_id
 
         req_state = self.request_states.get(req_id)
         if req_state is None:
@@ -357,7 +348,7 @@ class OmniBase(PDDisaggregationMixin):
 
     def _check_engine_output_error(
         self,
-        result: dict[str, Any],
+        result: OutputMessage,
         request_id: str,
         stage_id: int,
     ) -> None:
@@ -367,7 +358,7 @@ class OmniBase(PDDisaggregationMixin):
         engine is unrecoverable, otherwise raises :class:`EngineGenerateError`
         (recoverable, single-request failure).
         """
-        engine_outputs = result.get("engine_outputs")
+        engine_outputs = result.engine_outputs
         error_text = getattr(engine_outputs, "error", None)
         if error_text is None:
             return
@@ -388,17 +379,17 @@ class OmniBase(PDDisaggregationMixin):
 
     def _process_single_result(
         self,
-        result: dict[str, Any],
+        result: OutputMessage,
         stage_id: int,
         metrics: OrchestratorMetrics,
         req_start_ts: dict[str, float],
         wall_start_ts: float,
         final_stage_id_for_e2e: int,
     ) -> OmniRequestOutput | None:
-        req_id = result.get("request_id")
-        engine_outputs = result.get("engine_outputs")
-        stage_durations = getattr(result["engine_outputs"], "stage_durations", {})
-        peak_memory_mb = getattr(result["engine_outputs"], "peak_memory_mb", 0.0)
+        req_id = result.request_id
+        engine_outputs = result.engine_outputs
+        stage_durations = getattr(engine_outputs, "stage_durations", {})
+        peak_memory_mb = getattr(engine_outputs, "peak_memory_mb", 0.0)
 
         # Merge AR stage timing from OrchestratorAggregator.stage_events
         if self._enable_ar_profiler:
@@ -408,7 +399,7 @@ class OmniBase(PDDisaggregationMixin):
                     stage_durations[f"ar_stage_{evt.stage_id}"] = evt.stage_gen_time_ms / 1000.0
 
         # Merge pipeline timings from Orchestrator into stage_durations
-        _m = result.get("metrics")
+        _m = result.metrics
         if _m is not None and hasattr(_m, "pipeline_timings") and _m.pipeline_timings:
             for key, value in _m.pipeline_timings.items():
                 if key not in stage_durations:
@@ -425,18 +416,18 @@ class OmniBase(PDDisaggregationMixin):
 
         finished = engine_outputs.finished
 
-        submit_ts = result.get("stage_submit_ts")
+        submit_ts = result.stage_submit_ts
         now = time.time()
         if metrics.stage_first_ts[stage_id] is None:
             metrics.stage_first_ts[stage_id] = submit_ts if submit_ts is not None else now
         metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, now)
 
-        _m = result.get("metrics")
+        _m = result.metrics
         if finished and _m is not None:
             metrics.on_stage_metrics(stage_id, req_id, _m)
 
         stage_meta = self.engine.get_stage_metadata(stage_id)
-        if not stage_meta["final_output"]:
+        if not stage_meta.final_output:
             return None
 
         try:
@@ -450,7 +441,7 @@ class OmniBase(PDDisaggregationMixin):
         except Exception:
             logger.exception("[%s] Finalize request handling error", self.__class__.__name__)
 
-        output_type = getattr(engine_outputs, "final_output_type", stage_meta["final_output_type"])
+        output_type = getattr(engine_outputs, "final_output_type", stage_meta.final_output_type)
         images = getattr(engine_outputs, "images", []) if output_type == "image" else []
         return OmniRequestOutput(
             request_id=req_id or "",