[Core][Diffusion] Refactor prompt updates into a modality-ready interaction framework (#6294)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>

tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py

@@ -0,0 +1,555 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Integration tests for diffusion pipeline streaming output mode.
+
+Covers:
+- Mock pipeline -> StageDiffusionClient (ZMQ subprocess path)
+- Mock pipeline -> InlineStageDiffusionClient -> Orchestrator -> AsyncOmni
+- Mock pipeline -> InlineStageDiffusionClient -> Orchestrator -> AsyncOmni -> `/v1/realtime/video`
+"""
+
+import asyncio
+import contextlib
+import queue
+import time
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+import torch
+import zmq
+from fastapi import FastAPI, WebSocket
+from pytest_mock import MockerFixture
+from starlette.testclient import TestClient
+from vllm.utils.network_utils import get_open_zmq_ipc_path
+
+import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
+from tests.engine.test_orchestrator import OrchestratorFixture, _build_harness, _wait_for
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+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 StageDiffusionProc
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.distributed.omni_connectors.utils.serialization import (
+    OmniMsgpackDecoder,
+    OmniMsgpackEncoder,
+)
+from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
+from vllm_omni.engine.messages import ShutdownRequestMessage, StageSubmissionMessage
+from vllm_omni.engine.stage_init_utils import StageMetadata
+from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni
+from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _StepStreamingPipeline:
+    supports_step_execution = True
+
+    def __init__(self, outputs: list[DiffusionOutput]) -> None:
+        self.outputs = outputs
+        self.requests: list[Any] = []
+
+    def step_outputs(self, request):
+        self.requests.append(request)
+        return list(self.outputs)
+
+
+class _FailingStreamingPipeline:
+    supports_step_execution = True
+
+    def __init__(self) -> None:
+        self.requests: list[Any] = []
+
+    def step_outputs(self, request):
+        self.requests.append(request)
+        raise RuntimeError("stream failed midway")
+
+
+class _PipelineBackedEngine:
+    """DiffusionEngine stand-in that exposes step-streaming chunks from a mock pipeline."""
+
+    def __init__(self, pipeline: _StepStreamingPipeline | _FailingStreamingPipeline) -> None:
+        self.pipeline = pipeline
+        self.executor = SimpleNamespace(
+            register_failure_callback=MagicMock(),
+            check_health=MagicMock(),
+        )
+
+    async def step_streaming(self, request):
+        try:
+            outputs = self.pipeline.step_outputs(request)
+        except Exception as exc:
+            outputs = [
+                _streaming_diffusion_output(chunk=0, finished=False),
+                DiffusionOutput(error=str(exc), finished=True),
+            ]
+        for output in outputs:
+            if output.error is not None:
+                yield [OmniRequestOutput.from_error(request_id=request.request_id, error_message=output.error)]
+                continue
+            payload, metadata = _streaming_payload_and_metadata(output)
+            stream_metadata = metadata.get("stream", {})
+            yield [
+                OmniRequestOutput.from_diffusion(
+                    request_id=request.request_id,
+                    images=list(payload.get("image") or []),
+                    multimodal_output={"metadata": {"stream": stream_metadata}},
+                    finished=output.finished,
+                )
+            ]
+
+    def abort(self, request_id: str) -> None:
+        del request_id
+
+
+def _streaming_diffusion_output(
+    *,
+    chunk: int,
+    finished: bool,
+    images: list | None = None,
+) -> DiffusionOutput:
+    payload = {"image": images or []}
+    metadata = {"stream": {"chunk": chunk}}
+    return DiffusionOutput(
+        output={
+            "payload": payload,
+            "metadata": metadata,
+        },
+        finished=finished,
+    )
+
+
+def _streaming_payload_and_metadata(output: DiffusionOutput) -> tuple[dict, dict]:
+    envelope = output.output if isinstance(output.output, dict) else {}
+    payload = envelope.get("payload") if isinstance(envelope.get("payload"), dict) else {}
+    metadata = envelope.get("metadata") if isinstance(envelope.get("metadata"), dict) else {}
+    return payload, metadata
+
+
+class TestPipelineStreamingOutputToStageDiffusionClient:
+    """Streaming pipeline output over ZMQ into ``StageDiffusionClient``."""
+
+    @pytest.mark.asyncio
+    async def test_streaming_output_reaches_stage_client_from_mock_pipeline(self) -> None:
+        """Mock pipeline chunks over ZMQ reach StageDiffusionClient with correct finished flags."""
+        pipeline = _StepStreamingPipeline(
+            [
+                _streaming_diffusion_output(chunk=0, finished=False),
+                _streaming_diffusion_output(chunk=1, finished=True),
+            ]
+        )
+
+        outputs = await self._run_streaming_client_proc(pipeline, request_id="req-stream")
+
+        assert [output.request_id for output in outputs] == ["req-stream", "req-stream"]
+        assert [output.finished for output in outputs] == [False, True]
+        assert [output.multimodal_output["metadata"]["stream"]["chunk"] for output in outputs] == [0, 1]
+        assert pipeline.requests[0].request_id == "req-stream"
+
+    @pytest.mark.asyncio
+    async def test_streaming_midway_error_reaches_stage_client_after_prior_chunk(self) -> None:
+        """A pipeline error after the first chunk still delivers that chunk then a terminal error output."""
+        pipeline = _FailingStreamingPipeline()
+
+        outputs = await self._run_streaming_client_proc(pipeline, request_id="req-error")
+
+        assert len(outputs) == 2
+        assert outputs[0].request_id == "req-error"
+        assert outputs[0].finished is False
+        assert outputs[0].multimodal_output["metadata"]["stream"]["chunk"] == 0
+        assert outputs[1].request_id == "req-error"
+        assert outputs[1].finished is True
+        assert outputs[1].error is not None
+        assert "stream failed midway" in outputs[1].error
+
+    @classmethod
+    async def _run_streaming_client_proc(
+        cls,
+        pipeline: _StepStreamingPipeline | _FailingStreamingPipeline,
+        *,
+        request_id: str,
+    ) -> list[OmniRequestOutput]:
+        request_address = get_open_zmq_ipc_path()
+        response_address = get_open_zmq_ipc_path()
+        proc = object.__new__(StageDiffusionProc)
+        proc._od_config = SimpleNamespace(streaming_output=True)
+        proc._engine = _PipelineBackedEngine(pipeline)
+        proc._closed = False
+        client = cls._make_client(request_address, response_address)
+        proc_task = asyncio.create_task(proc.run_loop(request_address, response_address))
+        try:
+            await asyncio.sleep(0.05)
+            await client.add_request_async(
+                request_id,
+                "prompt",
+                OmniDiffusionSamplingParams(),
+            )
+            deadline = asyncio.get_running_loop().time() + 2.0
+            outputs: list[OmniRequestOutput] = []
+            while (not outputs or not outputs[-1].finished) and asyncio.get_running_loop().time() < deadline:
+                output = client.get_diffusion_output_nowait()
+                if output is not None:
+                    outputs.append(output)
+                else:
+                    await asyncio.sleep(0.01)
+            return outputs
+        finally:
+            try:
+                client._request_socket.send(client._encoder.encode({"type": "shutdown"}))
+            except Exception:
+                pass
+            try:
+                await asyncio.wait_for(proc_task, timeout=2)
+            finally:
+                cls._close_client(client)
+
+    @staticmethod
+    def _make_client(request_address: str, response_address: str) -> StageDiffusionClient:
+        client = object.__new__(StageDiffusionClient)
+        client.stage_id = 0
+        client.replica_id = 0
+        client.final_output = True
+        client.final_output_type = "image"
+        client.default_sampling_params = {}
+        client.requires_multimodal_data = False
+        client.custom_process_input_func = None
+        client.engine_input_source = []
+        client._proc = None
+        client._proc_manager = None
+        client._owns_process = False
+        client._zmq_ctx = zmq.Context()
+        client._request_socket = client._zmq_ctx.socket(zmq.PUSH)
+        client._request_socket.bind(request_address)
+        client._response_socket = client._zmq_ctx.socket(zmq.PULL)
+        client._response_socket.bind(response_address)
+        client._encoder = OmniMsgpackEncoder()
+        client._decoder = OmniMsgpackDecoder()
+        client._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue()
+        client._rpc_results = {}
+        client._pending_rpcs = set()
+        client._tasks = {}
+        client._shutting_down = False
+        client._engine_dead = False
+        return client
+
+    @staticmethod
+    def _close_client(client: StageDiffusionClient) -> None:
+        client._request_socket.close(linger=0)
+        client._response_socket.close(linger=0)
+        client._zmq_ctx.term()
+
+
+class TestPipelineStreamingOutputToEntrypoint:
+    """Streaming pipeline output through inline stage, orchestrator, and entrypoints."""
+
+    @pytest.mark.asyncio
+    async def test_streaming_output_reaches_async_omni_from_mock_pipeline(self) -> None:
+        """Mock pipeline streaming chunks reach AsyncOmni.generate() via inline stage and orchestrator."""
+        pipeline = _StepStreamingPipeline(
+            [
+                _streaming_diffusion_output(chunk=0, finished=False),
+                _streaming_diffusion_output(chunk=1, finished=True),
+            ]
+        )
+        inline_client = self._make_inline_pipeline_client(pipeline)
+        fixture = _build_harness([inline_client])
+        omni = self._make_async_omni(self._OrchestratorBridgeEngine(fixture))
+
+        try:
+            outputs: list[OmniRequestOutput] = []
+            async for output in omni.generate(
+                prompt={"prompt": "a cat"},
+                request_id="req-omni",
+                sampling_params_list=[OmniDiffusionSamplingParams()],
+                output_modalities=["image"],
+            ):
+                outputs.append(output)
+
+            await _wait_for(lambda: len(pipeline.requests) == 1)
+            assert pipeline.requests[0].request_id.startswith("req-omni-")
+            assert [output.multimodal_output["metadata"]["stream"]["chunk"] for output in outputs] == [0, 1]
+            assert [output.finished for output in outputs] == [False, True]
+        finally:
+            await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client)
+
+    @pytest.mark.asyncio
+    async def test_streaming_output_reaches_api_websocket_from_mock_pipeline(self, mocker: MockerFixture) -> None:
+        """Mock pipeline streaming chunks reach the video WebSocket handler through AsyncOmni.generate()."""
+        frames = [np.full((4, 4, 3), fill_value=i, dtype=np.uint8) for i in range(2)]
+        pipeline = _StepStreamingPipeline(
+            [
+                _streaming_diffusion_output(chunk=0, images=[frames], finished=False),
+                _streaming_diffusion_output(chunk=1, images=[frames], finished=True),
+            ]
+        )
+        inline_client = self._make_inline_pipeline_client(pipeline)
+        fixture = _build_harness([inline_client])
+        bridge_engine = self._OrchestratorBridgeEngine(fixture)
+        omni = self._make_async_omni(bridge_engine)
+
+        encoded = [b"pipeline-fmp4-0", b"pipeline-fmp4-1"]
+
+        class FakeStreamingVideoEncoder:
+            encode_calls: list[int] = []
+
+            def encode(self, video):
+                del video
+                idx = len(self.encode_calls)
+                self.encode_calls.append(idx)
+                return encoded[idx]
+
+            def close(self):
+                return b""
+
+        mocker.patch(
+            "vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder",
+            lambda *, output_format, fps, video_codec_options=None: FakeStreamingVideoEncoder(),
+        )
+        mocker.patch(
+            "vllm_omni.entrypoints.openai.serving_video_output_stream.is_video_generation_pipeline",
+            return_value=True,
+        )
+        mocker.patch(
+            "vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
+            return_value=[OmniDiffusionSamplingParams()],
+        )
+        mocker.patch(
+            "vllm_omni.entrypoints.openai.serving_video_output_stream.get_default_sampling_params_list",
+            return_value=[OmniDiffusionSamplingParams()],
+        )
+
+        handler = OmniStreamingVideoOutputHandler(
+            engine_client=omni,
+            model_name="test-model",
+            stage_configs=bridge_engine.stage_configs,
+            stall_timeout=5.0,
+            start_timeout=5.0,
+        )
+        app = FastAPI()
+
+        @app.websocket("/v1/realtime/video")
+        async def ws_endpoint(websocket: WebSocket):
+            await handler.handle_session(websocket)
+
+        try:
+            with TestClient(app) as client:
+                with client.websocket_connect("/v1/realtime/video") as ws:
+                    ws.send_json({"type": "session.start", "prompt": "integration test"})
+                    assert ws.receive_json()["type"] == "video.start"
+                    assert ws.receive_json()["type"] == "video.chunk_metadata"
+                    assert ws.receive_bytes() == b"pipeline-fmp4-0"
+                    assert ws.receive_json()["type"] == "video.chunk_metadata"
+                    assert ws.receive_bytes() == b"pipeline-fmp4-1"
+                    done = ws.receive_json()
+                    assert done["type"] == "session.done"
+                    assert done["chunks"] == 2
+
+            await _wait_for(lambda: len(pipeline.requests) == 1)
+            assert pipeline.requests[0].request_id
+        finally:
+            await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client)
+
+    class _OrchestratorBridgeEngine:
+        """Minimal AsyncOmni engine facade backed by a live Orchestrator harness."""
+
+        def __init__(self, fixture: OrchestratorFixture) -> None:
+            self._fixture = fixture
+            self.stage_metadata = [
+                StageRuntimeInfo(
+                    stage_type="diffusion",
+                    final_output=True,
+                    final_output_type="image",
+                )
+            ]
+            self.stage_configs: list[Any] = [SimpleNamespace(stage_type="diffusion")]
+            self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
+            self.num_stages = 1
+            self.supported_tasks = ("generate",)
+            self._alive = True
+
+        async def add_request_async(
+            self,
+            *,
+            request_id: str,
+            prompt: Any,
+            sampling_params_list: list[Any],
+            final_stage_id: int,
+            **kwargs: Any,
+        ) -> None:
+            self._fixture.request_sync_q.put_nowait(
+                StageSubmissionMessage(
+                    type="add_request",
+                    request_id=request_id,
+                    prompt=prompt,
+                    original_prompt=prompt,
+                    output_prompt_text=None,
+                    sampling_params_list=sampling_params_list,
+                    final_stage_id=final_stage_id,
+                    final_output_stage_ids=kwargs.get("final_output_stage_ids"),
+                    preprocess_ms=0.0,
+                    request_timestamp=kwargs.get("arrival_time", time.time()),
+                    enqueue_ts=time.perf_counter(),
+                )
+            )
+
+        async def try_get_output_async(self) -> Any | None:
+            try:
+                return self._fixture.output_sync_q.get_nowait()
+            except queue.Empty:
+                return None
+
+        def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo:
+            return self.stage_metadata[stage_id]
+
+        def is_alive(self) -> bool:
+            return self._fixture.thread.is_alive()
+
+        async def abort_async(self, request_ids: list[str]) -> None:
+            del request_ids
+
+    @classmethod
+    def _make_inline_pipeline_client(cls, pipeline: _StepStreamingPipeline) -> InlineStageDiffusionClient:
+        """Build an inline diffusion stage client whose engine runs the given mock pipeline."""
+        metadata = StageMetadata(
+            stage_id=0,
+            stage_type="diffusion",
+            engine_output_type="image",
+            is_comprehension=False,
+            requires_multimodal_data=False,
+            engine_input_source=[],
+            final_output=True,
+            final_output_type="image",
+            default_sampling_params=OmniDiffusionSamplingParams(),
+            custom_process_input_func=None,
+            model_stage=None,
+            runtime_cfg=None,
+        )
+        pipeline_engine = _PipelineBackedEngine(pipeline)
+        with patch.object(InlineStageDiffusionClient, "_enrich_config"):
+            with patch(
+                "vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine.make_engine",
+                return_value=pipeline_engine,
+            ):
+                od_config = MagicMock(spec=OmniDiffusionConfig)
+                od_config.streaming_output = True
+                return InlineStageDiffusionClient(
+                    model="test_model",
+                    od_config=od_config,
+                    metadata=metadata,
+                )
+
+    @staticmethod
+    def _make_async_omni(engine: _OrchestratorBridgeEngine) -> AsyncOmni:
+        omni = object.__new__(AsyncOmni)
+        omni.engine = engine
+        omni.log_stats = False
+        omni._pause_cond = asyncio.Condition()
+        omni._paused = False
+        omni.request_states = {}
+        omni.final_output_task = None
+        omni.event_resolver = AsyncEventResolver()
+        omni._enable_ar_profiler = False
+        omni.prom_metrics = MagicMock()
+        omni.mod_metrics = MagicMock()
+        omni.resolve_sampling_params_list = lambda params, allow_delta_coercion: params
+        omni._compute_final_stage_id = lambda output_modalities: 0
+        omni._compute_final_output_stage_ids = lambda output_modalities: [0]
+        omni.default_sampling_params_list = engine.default_sampling_params_list
+        omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None)
+        return omni
+
+    @staticmethod
+    async def _shutdown_pipeline_omni_harness(
+        omni: AsyncOmni,
+        fixture: OrchestratorFixture,
+        inline_client: InlineStageDiffusionClient,
+    ) -> None:
+        if omni.final_output_task is not None:
+            omni.final_output_task.cancel()
+            await asyncio.gather(omni.final_output_task, return_exceptions=True)
+        inline_client.shutdown()
+        fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
+        await asyncio.to_thread(fixture.thread.join, 5)
+
+
+def _make_vllm_config():
+    @contextlib.contextmanager
+    def set_priority(*args, **kwargs):
+        del args, kwargs
+        yield
+
+    return SimpleNamespace(
+        kernel_config=SimpleNamespace(ir_op_priority=SimpleNamespace(set_priority=set_priority)),
+        compilation_config=SimpleNamespace(ir_enable_torch_wrap=True),
+    )
+
+
+class TestSupportedPipelines:
+    """Streaming-output protocol checks for supported pipelines."""
+
+    def test_helios_supports_step_execution_for_streaming_output(self) -> None:
+        from vllm_omni.diffusion.models.helios.pipeline_helios import HeliosPipeline
+        from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution
+
+        # Avoid loading model weights; protocol membership depends on the class contract.
+        pipeline = object.__new__(HeliosPipeline)
+
+        assert pipeline.supports_step_execution is True
+        assert supports_step_execution(pipeline) is True
+        assert isinstance(pipeline, SupportsStepExecution) is True
+
+    def test_load_model_rejects_streaming_output_without_step_execution(self, monkeypatch) -> None:
+        class _NoStepPipeline:
+            def forward(self): ...
+
+        class _FakeDiffusersPipelineLoader:
+            def __init__(self, *args, **kwargs): ...
+
+            def load_model(self, **kwargs):
+                return _NoStepPipeline()
+
+        class _FakeDeviceMemoryProfiler:
+            consumed_memory = 0
+
+            def __enter__(self):
+                return self
+
+            def __exit__(self, exc_type, exc, tb):
+                return False
+
+        runner = object.__new__(DiffusionModelRunner)
+        runner.vllm_config = _make_vllm_config()
+        runner.od_config = SimpleNamespace(
+            max_num_seqs=1,
+            enable_cpu_offload=False,
+            enable_layerwise_offload=False,
+            enforce_eager=True,
+            cache_backend=None,
+            cache_config=None,
+            streaming_output=True,
+            step_execution=False,
+            model_class_name="NoStepPipeline",
+            parallel_config=SimpleNamespace(use_hsdp=False),
+        )
+        runner.device = torch.device("cpu")
+        runner.pipeline = None
+        runner.cache_backend = None
+        runner.offload_backend = None
+        runner.kv_transfer_manager = SimpleNamespace()
+
+        monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeDiffusersPipelineLoader)
+        monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeDeviceMemoryProfiler)
+        monkeypatch.setattr(
+            model_runner_module,
+            "enable_offload_backend",
+            lambda _config, pipeline, **_kwargs: (pipeline, None),
+        )
+        monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None)
+
+        with pytest.raises(ValueError, match="NoStepPipeline"):
+            DiffusionModelRunner.load_model(runner)

tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py

@@ -0,0 +1,812 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Tests for midway prompt update across runner, batch, pipeline, and entrypoint layers."""
+
+from __future__ import annotations
+
+import asyncio
+import queue
+import time
+from contextlib import contextmanager
+from types import SimpleNamespace
+from typing import Any, cast
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from tests.engine.test_orchestrator import (
+    OrchestratorFixture,
+    _build_harness,
+    _wait_for,
+)
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
+from vllm_omni.diffusion.interaction.coordinator import InteractionCoordinator
+from vllm_omni.diffusion.interaction.modality_handlers.prompt import (
+    PromptInteractionHandler,
+    PromptSession,
+)
+from vllm_omni.diffusion.models.helios.pipeline_helios import HeliosPipeline
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.diffusion.worker.input_batch import InputBatch
+from vllm_omni.diffusion.worker.utils import StepRequestState
+from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
+from vllm_omni.engine.messages import (
+    ErrorMessage,
+    InteractionMessage,
+    ShutdownRequestMessage,
+    StageSubmissionMessage,
+)
+from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
+from vllm_omni.engine.stage_client import StagePoolClient
+from vllm_omni.engine.stage_init_utils import StageMetadata
+from vllm_omni.engine.stage_pool import StagePool
+from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniInteractionPrompt
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.fixture
+def pipeline() -> HeliosPipeline:
+    pipeline = object.__new__(HeliosPipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.transformer = SimpleNamespace(dtype=torch.float32)  # pyright: ignore[reportAttributeAccessIssue]
+    pipeline.encode_prompt = MagicMock(
+        return_value=(
+            torch.full((1, 4, 2), 2.0),
+            None,
+        )
+    )
+    pipeline.prepare_next_chunk = MagicMock()
+    od_config = SimpleNamespace(model_class_name="HeliosPipeline")
+    pipeline._interaction_coordinator = InteractionCoordinator.build(pipeline, od_config)
+    return pipeline
+
+
+def _make_diffusion_request_state(*, request_id: str = "req-1", fps: int | None = None) -> StepRequestState:
+    state = StepRequestState(
+        request_id=request_id,
+        sampling=SimpleNamespace(  # pyright: ignore[reportArgumentType]
+            num_outputs_per_prompt=1,
+            max_sequence_length=226,
+            fps=fps,
+        ),
+        prompt="hello",
+    )
+    state.prompt_embeds = torch.zeros(1, 4, 2)
+    state.extra = {"window_num_frames": 8}
+    return state
+
+
+def _make_diffusion_model_runner(
+    *,
+    pipeline,
+    streaming_output: bool = True,
+    model_class_name: str = "HeliosPipeline",
+) -> DiffusionModelRunner:
+    runner = object.__new__(DiffusionModelRunner)
+    runner.pipeline = pipeline
+    runner.state_cache = {}
+    runner.od_config = SimpleNamespace(  # pyright: ignore[reportAttributeAccessIssue]
+        model_class_name=model_class_name,
+        streaming_output=streaming_output,
+        cache_backend=None,
+        parallel_config=SimpleNamespace(use_hsdp=False),
+    )
+    runner.vllm_config = object()
+    runner.device = torch.device("cpu")
+    runner.cache_backend = None
+    runner.offload_backend = None
+    runner.prompt_embed_cache = None
+    runner.kv_transfer_manager = SimpleNamespace(
+        receive_multi_kv_cache_distributed=lambda *a, **k: None,
+    )
+    # HeliosPipeline skeletons from object.__new__ are not full SupportsStepExecution
+    # instances; unit tests of submit_interaction stub this. Stepwise path tests
+    # rebind to the real method.
+    runner._supports_step_mode = lambda: True
+    runner._interaction_coordinator = None
+    return runner
+
+
+def _prompt_interaction(
+    prompt: str = "updated",
+    *,
+    event_id: str = "ui-update-1",
+    transition_chunks: int | None = None,
+) -> OmniInteractionPrompt:
+    interaction: dict[str, Any] = {
+        "event_id": event_id,
+        "event": {"prompt": prompt},
+    }
+    if transition_chunks is not None:
+        interaction["transition_chunks"] = transition_chunks
+    return cast(OmniInteractionPrompt, interaction)
+
+
+class TestPromptUpdateExecution:
+    """Runner, InputBatch cache, and pipeline chunk-boundary behavior."""
+
+    def test_runner_interaction_delegates_to_helios_pipeline(self, pipeline: HeliosPipeline) -> None:
+        """Runner encodes the new prompt and queues pending embeds on the request state."""
+        runner = _make_diffusion_model_runner(pipeline=pipeline)
+        state = _make_diffusion_request_state()
+        runner.state_cache["req-1"] = state
+
+        runner.submit_interaction("req-1", _prompt_interaction(transition_chunks=2))
+
+        pipeline.encode_prompt.assert_called_once()  # pyright: ignore[reportAttributeAccessIssue]
+        session = state.interaction_sessions["prompt"]
+        assert isinstance(session, PromptSession)
+        pending = session.pending_event
+        assert pending is not None
+        assert pending.event_id == "ui-update-1"
+        assert pending.transition_chunks == 2
+        assert torch.equal(pending.target_prompt_embeds, torch.full((1, 4, 2), 2.0))
+
+    def test_runner_prompt_update_rejects_unsupported_pipeline(self) -> None:
+        """Runner rejects prompt updates when the pipeline lacks prompt-update support."""
+
+        class _UnsupportedPipeline:
+            supports_step_execution = True
+
+        runner = _make_diffusion_model_runner(pipeline=_UnsupportedPipeline())
+        with pytest.raises(ValueError, match="not supported"):
+            runner.submit_interaction("req-1", _prompt_interaction())
+
+    def test_runner_prompt_update_rejects_missing_request(self, pipeline) -> None:
+        """Runner rejects prompt updates when no active request state exists."""
+        runner = _make_diffusion_model_runner(pipeline=pipeline)
+        with pytest.raises(ValueError, match="No active request state"):
+            runner.submit_interaction("missing", _prompt_interaction())
+
+    @pytest.mark.parametrize(
+        "interaction",
+        [
+            {},
+            {"multi_modal_data": {"camera": {"type": "pose"}}},
+            {"event": {"prompt": "updated", "multi_modal_data": {"camera": {"type": "pose"}}}},
+            {"event": {"prompt": "updated", "multi_modal_data": {}}},
+            {"event": {"prompt": "updated", "multi_modal_data": None}},
+            {"event": {"prompt": "updated", "multi_modal_data": "bad"}},
+        ],
+    )
+    def test_runner_interaction_rejects_structural_payloads_until_implemented(
+        self,
+        pipeline: HeliosPipeline,
+        interaction: dict[str, Any],
+    ) -> None:
+        """Unsupported interaction dict shapes are preserved to and rejected by the runner."""
+        runner = _make_diffusion_model_runner(pipeline=pipeline)
+        runner.state_cache["req-1"] = _make_diffusion_request_state()
+
+        with pytest.raises(NotImplementedError, match="Only text-only prompt update interactions"):
+            runner.submit_interaction("req-1", cast(Any, interaction))
+
+    @pytest.mark.parametrize("model_class_name", ["HeliosPipeline", "HeliosPyramidPipeline"])
+    def test_coordinator_registers_prompt_for_helios_aliases(
+        self,
+        pipeline: HeliosPipeline,
+        model_class_name: str,
+    ) -> None:
+        """Both Helios architecture names resolve the prompt handler registry."""
+        od_config = SimpleNamespace(model_class_name=model_class_name)
+        coordinator = InteractionCoordinator.build(pipeline, od_config)
+        assert coordinator.has_modality("prompt")
+
+        runner = _make_diffusion_model_runner(pipeline=pipeline, model_class_name=model_class_name)
+        runner.state_cache["req-1"] = _make_diffusion_request_state()
+        runner.submit_interaction("req-1", _prompt_interaction())
+        session = runner.state_cache["req-1"].interaction_sessions["prompt"]
+        assert isinstance(session, PromptSession)
+        assert session.pending_event is not None
+
+    def test_helios_peek_chunk_media_requires_positive_fps(self, pipeline: HeliosPipeline) -> None:
+        """Media-timeline peek must raise on missing/invalid fps instead of coercing to 0.0."""
+        state = _make_diffusion_request_state(fps=None)
+        with pytest.raises(ValueError, match="sampling.fps is required and must be > 0"):
+            pipeline.peek_chunk_media(state)
+
+        state.sampling.fps = 0  # pyright: ignore[reportAttributeAccessIssue]
+        with pytest.raises(ValueError, match="sampling.fps is required and must be > 0"):
+            pipeline.peek_chunk_media(state)
+
+    def test_input_batch_refreshes_prompt_embeds_on_version_change(self) -> None:
+        """InputBatch rebuilds prompt_embeds when prompt_update_version changes."""
+        state = StepRequestState(
+            request_id="req-1",
+            sampling=SimpleNamespace(),  # pyright: ignore[reportArgumentType]
+            prompt="hello",
+        )
+        state.prompt_embeds = torch.zeros(1, 2, 3)
+        state.latents = torch.zeros(1, 2)
+        state.timesteps = torch.tensor([1.0])
+        state.interaction_sessions["prompt"] = PromptSession(version=0)
+
+        batch = InputBatch.make_batch([state])
+        assert torch.equal(batch.prompt_embeds, torch.zeros(1, 2, 3))  # pyright: ignore[reportArgumentType]
+
+        state.prompt_embeds = torch.ones(1, 2, 3)
+        state.interaction_sessions["prompt"] = PromptSession(version=1)
+        refreshed = InputBatch.make_batch([state], cached_batch=batch)
+        assert torch.equal(refreshed.prompt_embeds, torch.ones(1, 2, 3))  # pyright: ignore[reportArgumentType]
+
+    def test_helios_prompt_handler_enqueue_queues_pending_target(self, pipeline: HeliosPipeline) -> None:
+        """Prompt handler queues target embeds without mutating current prompt_embeds."""
+        state = _make_diffusion_request_state()
+
+        PromptInteractionHandler.from_pipeline(pipeline).enqueue(
+            state,
+            event_id="ui-update-1",
+            received_at=0.0,
+            payload={"prompt": "new scene"},
+            transition_chunks=2,
+        )
+
+        session = state.interaction_sessions["prompt"]
+        assert isinstance(session, PromptSession)
+        pending = session.pending_event
+        assert pending is not None
+        assert torch.equal(pending.target_prompt_embeds, torch.full((1, 4, 2), 2.0))
+        assert pending.transition_chunks == 2
+        assert torch.equal(state.prompt_embeds, torch.zeros(1, 4, 2))  # pyright: ignore[reportArgumentType]
+
+    def test_helios_prompt_handler_enqueue_rejects_before_initial_generation(self, pipeline: HeliosPipeline) -> None:
+        """Reject prompt updates submitted before initial prompt embeds exist."""
+        state = _make_diffusion_request_state()
+        state.prompt_embeds = None
+
+        with pytest.raises(
+            ValueError,
+            match="prompt_update is not allowed before initial generation has started",
+        ):
+            PromptInteractionHandler.from_pipeline(pipeline).enqueue(
+                state,
+                event_id="ui-update-1",
+                received_at=0.0,
+                payload={"prompt": "new scene"},
+                transition_chunks=2,
+            )
+
+        session = state.interaction_sessions.get("prompt")
+        assert session is None or (isinstance(session, PromptSession) and session.pending_event is None)
+
+    def test_helios_apply_interaction_at_chunk_boundary_starts_transition(self, pipeline: HeliosPipeline) -> None:
+        """At chunk boundary, starts transition state and bumps prompt_update_version."""
+        state = _make_diffusion_request_state(fps=None)
+        PromptInteractionHandler.from_pipeline(pipeline).enqueue(
+            state,
+            event_id="ui-update-1",
+            received_at=0.0,
+            payload={"prompt": "new scene"},
+            transition_chunks=2,
+        )
+
+        pipeline.apply_interaction_at_chunk_boundary(state)
+
+        session = state.interaction_sessions["prompt"]
+        assert isinstance(session, PromptSession)
+        assert session.active_event is not None
+        assert torch.equal(state.prompt_embeds, torch.ones(1, 4, 2))  # pyright: ignore[reportArgumentType]
+        assert session.version == 1
+        assert state.interaction_chunk_metadata is not None
+        assert state.interaction_chunk_metadata.as_dict() == {
+            "started_event_ids": ["ui-update-1"],
+            "active_event_ids": ["ui-update-1"],
+            "completed_event_ids": [],
+        }
+
+    def test_helios_apply_interaction_advances_transition_over_chunks(self, pipeline: HeliosPipeline) -> None:
+        """At chunk boundary, interpolates embeds until the target prompt is reached."""
+        state = _make_diffusion_request_state(fps=None)
+        PromptInteractionHandler.from_pipeline(pipeline).enqueue(
+            state,
+            event_id="ui-update-1",
+            received_at=0.0,
+            payload={"prompt": "new scene"},
+            transition_chunks=3,
+        )
+        pipeline.apply_interaction_at_chunk_boundary(state)
+
+        assert torch.allclose(state.prompt_embeds, torch.full((1, 4, 2), 2.0 / 3.0))  # pyright: ignore[reportArgumentType]
+
+        for _ in range(2):
+            pipeline.apply_interaction_at_chunk_boundary(state)
+
+        assert torch.allclose(state.prompt_embeds, torch.full((1, 4, 2), 2.0))  # pyright: ignore[reportArgumentType]
+        session = state.interaction_sessions["prompt"]
+        assert isinstance(session, PromptSession)
+        assert session.version == 3
+
+        # Further chunk boundaries after completion must not bump the version.
+        pipeline.apply_interaction_at_chunk_boundary(state)
+        assert session.version == 3
+
+    def test_stepwise_runner_applies_prompt_update_and_acks(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """Drive the production runner loop: submit → apply(no media args) → ACK on next chunk.
+
+        Also regresses ``sampling.fps is None`` on the prompt-only path (no peek).
+        ACK ids land on the *next* chunk's ``DiffusionOutput`` via
+        ``_attach_stepwise_metadata`` consuming the prior boundary metadata.
+        """
+        import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
+        from vllm_omni.diffusion.interaction.mixin import InteractionMixin
+        from vllm_omni.diffusion.interaction.types import ChunkMediaSpec
+
+        chunks = [
+            DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2),  # pyright: ignore[reportArgumentType]
+            DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2),  # pyright: ignore[reportArgumentType]
+        ]
+
+        class _InteractiveChunkPipeline(InteractionMixin):
+            device = torch.device("cpu")
+            supports_step_execution = True
+
+            def __init__(self) -> None:
+                self._outputs = chunks
+                self.decode_calls = 0
+                self.transformer = SimpleNamespace(dtype=torch.float32)
+                self.encode_prompt = MagicMock(return_value=(torch.full((1, 4, 2), 2.0), None))
+                self.prepare_next_chunk = MagicMock()
+                od_config = SimpleNamespace(model_class_name="HeliosPipeline")
+                self._interaction_coordinator = InteractionCoordinator.build(self, od_config)  # pyright: ignore[reportArgumentType]
+
+            def peek_chunk_media(self, state: StepRequestState) -> ChunkMediaSpec:
+                # Same contract as Helios: never coerce missing fps to 0.0.
+                fps = state.sampling.fps
+                if fps is None or float(fps) <= 0:
+                    raise ValueError(
+                        "sampling.fps is required and must be > 0 for interaction modalities "
+                        f"that use the chunk media timeline, got {fps!r} "
+                        f"(request_id={state.request_id!r})"
+                    )
+                return ChunkMediaSpec(num_frames=int(state.extra["window_num_frames"]), fps=float(fps))
+
+            def prepare_encode(self, state: StepRequestState) -> StepRequestState:
+                state.prompt_embeds = torch.zeros(1, 4, 2)
+                state.latents = torch.zeros(1, 1)
+                state.timesteps = torch.tensor([1.0, 0.0, 1.0, 0.0])
+                state.step_index = 0
+                state.step_in_chunk = 0
+                state.chunk_num_steps = 2
+                state.total_chunks = 2
+                state.extra = {"window_num_frames": 8}
+                return state
+
+            def denoise_step(self, input_batch: Any, states: Any) -> torch.Tensor:
+                del states
+                return torch.ones_like(input_batch.latents)
+
+            def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor | None) -> None:
+                state.latents = noise_pred
+                state.step_index += 1
+                state.step_in_chunk += 1
+
+            def post_decode(self, state: StepRequestState) -> DiffusionOutput:
+                output = self._outputs[self.decode_calls]
+                self.decode_calls += 1
+                state.chunk_index += 1
+                state.step_index = 0
+                state.step_in_chunk = 0
+                if not state.request_denoise_completed:
+                    state.latents = torch.zeros(1, 1)
+                return output
+
+        pipeline = _InteractiveChunkPipeline()
+        runner = _make_diffusion_model_runner(pipeline=pipeline)
+        runner._supports_step_mode = lambda: DiffusionModelRunner._supports_step_mode(runner)
+        runner.od_config.streaming_output = True
+        runner.od_config.step_execution = True
+        runner.diffusion_kv_backend = SimpleNamespace(remove_diffusion_kv_requests=lambda *_: 0)  # pyright: ignore[reportAttributeAccessIssue]
+
+        sampling = SimpleNamespace(
+            generator=None,
+            seed=None,
+            generator_device=None,
+            num_inference_steps=4,
+            num_outputs_per_prompt=1,
+            max_sequence_length=226,
+            fps=None,
+        )
+        req = SimpleNamespace(
+            request_id="req",
+            prompt="a prompt",
+            sampling_params=sampling,
+            kv_sender_info=None,
+        )
+
+        @contextmanager
+        def _noop_forward_context(*args: Any, **kwargs: Any):
+            del args, kwargs
+            yield
+
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+        monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
+        monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+        monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+        monkeypatch.setattr(model_runner_module.current_omni_platform, "is_available", lambda: False)
+
+        scheduler_output = SimpleNamespace(
+            finished_req_ids=set(),
+            scheduled_new_reqs=[SimpleNamespace(request_id="req", req=req, diffusion_kv_metadata=None)],
+            scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
+        )
+        # Admit request and run first denoise step (no chunk boundary yet).
+        DiffusionModelRunner.execute_stepwise(runner, scheduler_output)  # pyright: ignore[reportArgumentType]
+        assert "req" in runner.state_cache
+        assert runner.state_cache["req"].sampling.fps is None
+
+        runner.submit_interaction("req", _prompt_interaction("new scene", transition_chunks=2))
+
+        cached = SimpleNamespace(
+            finished_req_ids=set(),
+            scheduled_new_reqs=[],
+            scheduled_cached_reqs=SimpleNamespace(request_ids=["req"]),
+        )
+        # Second step completes chunk 0: apply runs (no peek), metadata staged.
+        chunk0 = DiffusionModelRunner.execute_stepwise(runner, cached).get_request_output("req")  # pyright: ignore[reportArgumentType]
+        assert chunk0 is not None and chunk0.result is not None
+        assert chunk0.result.started_event_ids == []
+        state = runner.state_cache["req"]
+        assert torch.equal(state.prompt_embeds, torch.ones(1, 4, 2))  # pyright: ignore[reportArgumentType]
+        assert state.interaction_chunk_metadata is not None
+        assert state.interaction_chunk_metadata.started_event_ids == ["ui-update-1"]
+        pipeline.prepare_next_chunk.assert_called()  # pyright: ignore[reportAttributeAccessIssue]
+
+        # Next chunk decode attaches prior-boundary ACK ids onto DiffusionOutput.
+        DiffusionModelRunner.execute_stepwise(runner, cached)  # pyright: ignore[reportArgumentType]
+        chunk1 = DiffusionModelRunner.execute_stepwise(runner, cached).get_request_output("req")  # pyright: ignore[reportArgumentType]
+        assert chunk1 is not None and chunk1.result is not None
+        assert chunk1.result.started_event_ids == ["ui-update-1"]
+        assert chunk1.result.active_event_ids == ["ui-update-1"]
+        assert chunk1.finished is True
+
+
+class TestPromptUpdateIntegration:
+    """AsyncOmni prompt update through orchestrator, inline client, and runner."""
+
+    @pytest.mark.asyncio
+    async def test_orchestrator_prompt_update_failure_surfaces_non_fatal_error(
+        self,
+    ) -> None:
+        """Worker-side prompt_update rejection is reported without crashing the orchestrator."""
+
+        class _RejectingStagePool:
+            async def submit_interaction(self, *args: Any, **kwargs: Any) -> None:
+                del args, kwargs
+                raise ValueError("prompt embeds are not ready")
+
+        output_queue: asyncio.Queue[ErrorMessage] = asyncio.Queue()
+        orchestrator = object.__new__(Orchestrator)
+        orchestrator.stage_pools = [_RejectingStagePool()]  # pyright: ignore[reportAttributeAccessIssue]
+        orchestrator.output_async_queue = output_queue  # pyright: ignore[reportAttributeAccessIssue]
+        orchestrator.request_states = {"req-1": OrchestratorRequestState(request_id="req-1")}
+
+        await orchestrator._handle_interaction(
+            InteractionMessage(
+                request_id="req-1",
+                interaction=_prompt_interaction("new scene", event_id="ui-update-1", transition_chunks=2),
+            )
+        )
+
+        msg = output_queue.get_nowait()
+        assert msg == ErrorMessage(
+            error="Failed interaction for request req-1: prompt embeds are not ready",
+            fatal=False,
+            request_id="req-1",
+            event_id="ui-update-1",
+            stage_id=0,
+        )
+
+    @pytest.mark.asyncio
+    async def test_runner_prompt_update_failure_surfaces_non_fatal_error(self, pipeline: HeliosPipeline) -> None:
+        """Runner-side prompt encode rejection is reported through the orchestrator."""
+        pipeline.encode_prompt = MagicMock(  # pyright: ignore[reportAttributeAccessIssue]
+            side_effect=ValueError("prompt embeds are not ready")
+        )
+        runner = _make_diffusion_model_runner(pipeline=pipeline)
+        runner.state_cache["req-1"] = _make_diffusion_request_state(request_id="req-1")
+        prompt_update_engine = self._PromptUpdateEngine(self._IncrementalStreamingPipeline(), runner)
+        inline_client = self._make_inline_pipeline_client(prompt_update_engine)
+
+        output_queue: asyncio.Queue[ErrorMessage] = asyncio.Queue()
+        orchestrator = object.__new__(Orchestrator)
+        orchestrator.stage_pools = [  # pyright: ignore[reportAttributeAccessIssue]
+            StagePool(0, cast(StagePoolClient, inline_client))
+        ]
+        orchestrator.output_async_queue = output_queue  # pyright: ignore[reportAttributeAccessIssue]
+        orchestrator.request_states = {"req-1": OrchestratorRequestState(request_id="req-1")}
+
+        try:
+            await orchestrator._handle_interaction(
+                InteractionMessage(
+                    request_id="req-1",
+                    interaction=_prompt_interaction("new scene", event_id="ui-update-1", transition_chunks=2),
+                )
+            )
+        finally:
+            inline_client.shutdown()
+
+        pipeline.encode_prompt.assert_called_once()  # pyright: ignore[reportAttributeAccessIssue]
+        msg = output_queue.get_nowait()
+        assert msg == ErrorMessage(
+            error="Failed interaction for request req-1: prompt embeds are not ready",
+            fatal=False,
+            request_id="req-1",
+            event_id="ui-update-1",
+            stage_id=0,
+        )
+
+    @pytest.mark.asyncio
+    async def test_prompt_update_reaches_runner_from_async_omni(self, pipeline: HeliosPipeline) -> None:
+        """Midway prompt update submitted via AsyncOmni reaches the diffusion runner."""
+        streaming_pipeline = self._IncrementalStreamingPipeline()
+        runner = _make_diffusion_model_runner(pipeline=pipeline)
+        prompt_update_engine = self._PromptUpdateEngine(streaming_pipeline, runner)
+        inline_client = self._make_inline_pipeline_client(prompt_update_engine)
+        fixture = _build_harness([inline_client])
+        omni = self._make_async_omni(self._OrchestratorBridgeEngine(fixture))
+
+        generate_task: asyncio.Task[list[OmniRequestOutput]] | None = None
+        try:
+            generate_task = asyncio.create_task(self._collect_generate_outputs(omni))
+
+            await _wait_for(
+                lambda: (
+                    len(runner.state_cache) > 0
+                    and any(state.external_request_id == "req-omni" for state in omni.request_states.values())
+                )
+            )
+            internal_request_id = next(iter(runner.state_cache))
+
+            await omni.submit_interaction_async(
+                "req-omni",
+                interaction=_prompt_interaction("new scene", event_id="ui-update-1", transition_chunks=2),
+            )
+
+            outputs = await generate_task
+
+            pending = runner.state_cache[internal_request_id].interaction_sessions["prompt"].pending_event
+            assert pending is not None
+            assert pending.transition_chunks == 2
+            assert torch.equal(pending.target_prompt_embeds, torch.full((1, 4, 2), 2.0))
+            pipeline.encode_prompt.assert_called_once()  # pyright: ignore[reportAttributeAccessIssue]
+            assert [output.custom_output["chunk"] for output in outputs] == [0, 1]
+        finally:
+            if generate_task is not None and not generate_task.done():
+                generate_task.cancel()
+                await asyncio.gather(generate_task, return_exceptions=True)
+            await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client)
+
+    @staticmethod
+    async def _collect_generate_outputs(omni: AsyncOmni) -> list[OmniRequestOutput]:
+        outputs: list[OmniRequestOutput] = []
+        async for output in omni.generate(
+            prompt={"prompt": "a cat"},
+            request_id="req-omni",
+            sampling_params_list=[OmniDiffusionSamplingParams()],
+            output_modalities=["image"],
+        ):
+            outputs.append(output)
+        return outputs
+
+    @classmethod
+    def _make_inline_pipeline_client(cls, engine: _PromptUpdateEngine) -> InlineStageDiffusionClient:
+        metadata = StageMetadata(
+            stage_id=0,
+            stage_type="diffusion",
+            engine_output_type="image",
+            is_comprehension=False,
+            requires_multimodal_data=False,
+            engine_input_source=[],
+            final_output=True,
+            final_output_type="image",
+            default_sampling_params=OmniDiffusionSamplingParams(),
+            custom_process_input_func=None,
+            model_stage=None,
+            runtime_cfg=None,
+        )
+        with patch.object(InlineStageDiffusionClient, "_enrich_config"):
+            with patch(
+                "vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine.make_engine",
+                return_value=engine,
+            ):
+                od_config = MagicMock(spec=OmniDiffusionConfig)
+                od_config.streaming_output = True
+                return InlineStageDiffusionClient(
+                    model="test_model",
+                    od_config=od_config,
+                    metadata=metadata,
+                )
+
+    @staticmethod
+    def _make_async_omni(engine: _OrchestratorBridgeEngine) -> AsyncOmni:
+        omni = object.__new__(AsyncOmni)
+        omni.engine = engine  # pyright: ignore[reportAttributeAccessIssue]
+        omni.log_stats = False
+        omni._pause_cond = asyncio.Condition()
+        omni._paused = False
+        omni.request_states = {}
+        omni.final_output_task = None
+        omni.event_resolver = AsyncEventResolver()
+        omni._enable_ar_profiler = False
+        omni.prom_metrics = MagicMock()
+        omni.mod_metrics = MagicMock()
+        omni.resolve_sampling_params_list = lambda params, allow_delta_coercion: params  # pyright: ignore[reportAttributeAccessIssue]
+        omni._compute_final_stage_id = lambda output_modalities: 0
+        omni._compute_final_output_stage_ids = lambda output_modalities: [0]
+        omni.default_sampling_params_list = engine.default_sampling_params_list  # pyright: ignore[reportAttributeAccessIssue]
+        omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None)  # pyright: ignore[reportAttributeAccessIssue]
+        return omni
+
+    @staticmethod
+    async def _shutdown_pipeline_omni_harness(
+        omni: AsyncOmni,
+        fixture: OrchestratorFixture,
+        inline_client: InlineStageDiffusionClient,
+    ) -> None:
+        if omni.final_output_task is not None:
+            omni.final_output_task.cancel()
+            await asyncio.gather(omni.final_output_task, return_exceptions=True)
+        inline_client.shutdown()
+        fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
+        await asyncio.to_thread(fixture.thread.join, 5)
+
+    class _IncrementalStreamingPipeline:
+        """Returns one streaming chunk per ``step_outputs`` call."""
+
+        supports_step_execution = True
+
+        def __init__(self) -> None:
+            self.requests: list[Any] = []
+            self._step_index = 0
+
+        def step_outputs(self, request):
+            self.requests.append(request)
+            if self._step_index == 0:
+                self._step_index += 1
+                return [DiffusionOutput(output={"chunk": 0}, finished=False)]
+            return [DiffusionOutput(output={"chunk": 1}, finished=True)]
+
+    class _OrchestratorBridgeEngine:
+        """Minimal AsyncOmni engine facade backed by a live Orchestrator harness."""
+
+        def __init__(self, fixture: OrchestratorFixture) -> None:
+            self._fixture = fixture
+            self.stage_metadata = [
+                StageRuntimeInfo(
+                    stage_type="diffusion",
+                    final_output=True,
+                    final_output_type="image",
+                )
+            ]
+            self.stage_configs: list[Any] = [SimpleNamespace(stage_type="diffusion")]
+            self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
+            self.num_stages = 1
+            self.supported_tasks = ("generate",)
+            self._alive = True
+
+        async def add_request_async(
+            self,
+            *,
+            request_id: str,
+            prompt: Any,
+            sampling_params_list: list[Any],
+            final_stage_id: int,
+            **kwargs: Any,
+        ) -> None:
+            self._fixture.request_sync_q.put_nowait(
+                StageSubmissionMessage(
+                    type="add_request",
+                    request_id=request_id,
+                    prompt=prompt,
+                    original_prompt=prompt,
+                    output_prompt_text=None,
+                    sampling_params_list=sampling_params_list,
+                    final_stage_id=final_stage_id,
+                    final_output_stage_ids=kwargs.get("final_output_stage_ids"),
+                    preprocess_ms=0.0,
+                    request_timestamp=kwargs.get("arrival_time", time.time()),
+                    enqueue_ts=time.perf_counter(),
+                )
+            )
+
+        async def submit_interaction_async(
+            self,
+            request_id: str,
+            *,
+            interaction: dict[str, Any],
+        ) -> None:
+            self._fixture.request_sync_q.put_nowait(
+                InteractionMessage(
+                    request_id=request_id,
+                    interaction=cast(Any, interaction),
+                )
+            )
+
+        async def try_get_output_async(self) -> Any | None:
+            try:
+                return self._fixture.output_sync_q.get_nowait()
+            except queue.Empty:
+                return None
+
+        def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo:
+            return self.stage_metadata[stage_id]
+
+        def is_alive(self) -> bool:
+            return self._fixture.thread.is_alive()
+
+        async def abort_async(self, request_ids: list[str]) -> None:
+            del request_ids
+
+    class _PromptUpdateEngine:
+        """DiffusionEngine stand-in that streams chunks and routes prompt_update to a runner."""
+
+        def __init__(
+            self,
+            streaming_pipeline: TestPromptUpdateIntegration._IncrementalStreamingPipeline,
+            runner: DiffusionModelRunner,
+        ) -> None:
+            self.streaming_pipeline = streaming_pipeline
+            self._runner = runner
+            self.executor = SimpleNamespace(
+                register_failure_callback=MagicMock(),
+                check_health=MagicMock(),
+            )
+
+        async def step_streaming(self, request):
+            if request.request_id not in self._runner.state_cache:
+                state = _make_diffusion_request_state(request_id=request.request_id)
+                self._runner.state_cache[request.request_id] = state
+
+            state = self._runner.state_cache[request.request_id]
+
+            outputs = self.streaming_pipeline.step_outputs(request)
+            output = outputs[0]
+            custom_output = cast(dict[str, Any], output.output) or {}
+            yield [
+                OmniRequestOutput.from_diffusion(
+                    request_id=request.request_id,
+                    images=list(custom_output.get("images") or []),
+                    custom_output=custom_output,
+                    finished=output.finished,
+                )
+            ]
+
+            deadline = time.monotonic() + 5.0
+            while (
+                not isinstance(state.interaction_sessions.get("prompt"), PromptSession)
+                or state.interaction_sessions["prompt"].pending_event is None
+            ):
+                if time.monotonic() >= deadline:
+                    raise TimeoutError("timed out waiting for prompt update during streaming")
+                await asyncio.sleep(0.01)
+
+            outputs = self.streaming_pipeline.step_outputs(request)
+            output = outputs[0]
+            custom_output = cast(dict[str, Any], output.output) or {}
+            yield [
+                OmniRequestOutput.from_diffusion(
+                    request_id=request.request_id,
+                    images=list(custom_output.get("images") or []),
+                    custom_output=custom_output,
+                    finished=output.finished,
+                )
+            ]
+
+        def collective_rpc(self, method, timeout, args, kwargs, unique_reply_rank):
+            del timeout, kwargs, unique_reply_rank
+            if method == "submit_interaction":
+                request_id, interaction = args
+                self._runner.submit_interaction(request_id, interaction)
+                return None
+            raise NotImplementedError(f"collective_rpc not mocked for {method!r}")
+
+        def abort(self, request_id: str) -> None:
+            del request_id