[Frontend] Add midway prompt update for streaming video generation (#4652)

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

tests/diffusion/test_diffusion_streaming_output.py

@@ -341,7 +341,9 @@ class TestPipelineStreamingOutputToEntrypoint:
                 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"
@@ -453,7 +455,6 @@ class TestPipelineStreamingOutputToEntrypoint:
         omni.final_output_task = None
         omni.event_resolver = AsyncEventResolver()
         omni._enable_ar_profiler = False
-        omni._is_sleeping = False
         omni.prom_metrics = MagicMock()
         omni.mod_metrics = MagicMock()
         omni.resolve_sampling_params_list = lambda params, allow_delta_coercion: params

tests/diffusion/test_prompt_update.py

@@ -0,0 +1,574 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for midway prompt update across runner, batch, pipeline, and entrypoint layers."""
+
+from __future__ import annotations
+
+import asyncio
+import queue
+import time
+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.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()
+    return pipeline
+
+
+def _make_diffusion_request_state(*, request_id: str = "req-1") -> StepRequestState:
+    state = StepRequestState(
+        request_id=request_id,
+        sampling=SimpleNamespace(num_outputs_per_prompt=1, max_sequence_length=226),  # pyright: ignore[reportArgumentType]
+        prompt="hello",
+    )
+    state.prompt_embeds = torch.zeros(1, 4, 2)
+    state.extra = {}
+    return state
+
+
+def _make_diffusion_model_runner(*, pipeline, streaming_output: bool = True) -> DiffusionModelRunner:
+    runner = object.__new__(DiffusionModelRunner)
+    runner.pipeline = pipeline
+    runner.state_cache = {}
+    runner.od_config = SimpleNamespace(  # pyright: ignore[reportAttributeAccessIssue]
+        model_class_name="HeliosPipeline",
+        streaming_output=streaming_output,
+    )
+    runner._supports_step_mode = lambda: True
+    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]
+        pending = state.extra["pending_prompt_update"]
+        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"}}}},
+        ],
+    )
+    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))
+
+    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.extra["prompt_update_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.extra["prompt_update_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_prepare_prompt_update_queues_pending_target(self, pipeline: HeliosPipeline) -> None:
+        """HeliosPipeline queues target embeds without mutating current prompt_embeds."""
+        state = _make_diffusion_request_state()
+
+        pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=2)
+
+        pending = state.extra["pending_prompt_update"]
+        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_prepare_prompt_update_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",
+        ):
+            pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=2)
+
+        assert "pending_prompt_update" not in state.extra
+
+    def test_helios_apply_prompt_update_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()
+        pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=2)
+
+        pipeline._apply_prompt_update_at_chunk_boundary(state)
+
+        assert state.extra.get("prompt_update_state") is not None
+        assert torch.equal(state.prompt_embeds, torch.ones(1, 4, 2))  # pyright: ignore[reportArgumentType]
+        assert state.extra["prompt_update_version"] == 1
+        assert state.extra["prompt_update_chunk_metadata"] == {
+            "started_event_ids": ["ui-update-1"],
+            "active_event_ids": ["ui-update-1"],
+            "completed_event_ids": [],
+        }
+
+    def test_helios_apply_prompt_update_advances_transition_over_chunks(self, pipeline: HeliosPipeline) -> None:
+        """At chunk boundary, interpolates embeds until the target prompt is reached."""
+        state = _make_diffusion_request_state()
+        pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=3)
+        pipeline._apply_prompt_update_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_prompt_update_at_chunk_boundary(state)
+
+        assert torch.allclose(state.prompt_embeds, torch.full((1, 4, 2), 2.0))  # pyright: ignore[reportArgumentType]
+        assert state.extra["prompt_update_version"] == 3
+
+        # Further chunk boundaries after completion must not bump the version.
+        pipeline._apply_prompt_update_at_chunk_boundary(state)
+        assert state.extra["prompt_update_version"] == 3
+
+
+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 prepare_prompt_update rejection is reported through the orchestrator."""
+        pipeline.prepare_prompt_update = MagicMock(  # pyright: ignore[reportAttributeAccessIssue, reportMethodAssignment]
+            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.prepare_prompt_update.assert_called_once_with(  # pyright: ignore[reportAttributeAccessIssue]
+            runner.state_cache["req-1"],
+            "new scene",
+            "ui-update-1",
+            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_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].extra["pending_prompt_update"]
+            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,
+                    batch_size=1,
+                )
+
+    @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 "pending_prompt_update" not in state.extra:
+                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

tests/entrypoints/openai_api/test_serving_video_output_stream.py

@@ -16,6 +16,7 @@ from fastapi import FastAPI, WebSocket
 from pytest_mock import MockerFixture
 from starlette.testclient import TestClient
 
+from tests.helpers.media import generate_synthetic_image
 from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -27,6 +28,12 @@ def _fake_video_frames(num_frames: int = 2) -> list[np.ndarray]:
     return [np.full((4, 4, 3), fill_value=i, dtype=np.uint8) for i in range(num_frames)]
 
 
+def _receive_video_chunk(ws) -> tuple[dict, bytes]:
+    metadata = ws.receive_json()
+    assert metadata["type"] == "video.chunk_metadata"
+    return metadata, ws.receive_bytes()
+
+
 def _build_test_app(
     *,
     mocker: MockerFixture,
@@ -37,6 +44,7 @@ def _build_test_app(
 ) -> tuple[FastAPI, OmniStreamingVideoOutputHandler, MagicMock]:
     engine_client = mocker.MagicMock()
     engine_client.abort = mocker.AsyncMock()
+    engine_client.submit_interaction_async = mocker.AsyncMock()
     engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()]
     engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
     engine_client.generate = mock_generate
@@ -126,8 +134,13 @@ class TestStreamingVideoOutputWebSocket:
                 assert "format" not in start["config"]
                 assert "request_id" in start
 
-                assert ws.receive_bytes() == b"mp4-chunk-0"
-                assert ws.receive_bytes() == b"mp4-chunk-1"
+                metadata, chunk = _receive_video_chunk(ws)
+                assert metadata["transport_chunk_index"] == 0
+                assert metadata["kind"] == "media"
+                assert chunk == b"mp4-chunk-0"
+                metadata, chunk = _receive_video_chunk(ws)
+                assert metadata["transport_chunk_index"] == 1
+                assert chunk == b"mp4-chunk-1"
 
                 done = ws.receive_json()
                 assert done["type"] == "session.done"
@@ -138,6 +151,48 @@ class TestStreamingVideoOutputWebSocket:
         engine_client.streaming_encoder_factory.assert_called_once()
         assert engine_client.streaming_encoder_factory.call_args.kwargs["output_format"] == "m4s"
 
+    def test_streaming_session_accepts_initial_image_reference(self, mocker: MockerFixture):
+        """session.start forwards image_reference as prompt multi_modal_data.image."""
+        captured_prompts: list[dict[str, object]] = []
+
+        async def mock_generate(*_args, **kwargs):
+            captured_prompts.append(kwargs["prompt"])
+            yield OmniRequestOutput.from_diffusion(
+                request_id="req-ws",
+                images=[_fake_video_frames(2)],
+                final_output_type="image",
+                finished=True,
+            )
+
+        app, _handler, engine_client = _build_test_app(
+            mocker=mocker,
+            streaming_chunks=[(b"mp4-chunk-0", True)],
+            mock_generate=mock_generate,
+        )
+
+        synth_image = generate_synthetic_image(4, 4)
+        image_url = f"data:image/jpeg;base64,{synth_image['base64']}"
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/realtime/video") as ws:
+                ws.send_json(
+                    {
+                        "type": "session.start",
+                        "prompt": "animate this image",
+                        "width": 4,
+                        "height": 4,
+                        "image_reference": {"image_url": image_url},
+                    }
+                )
+                assert ws.receive_json()["type"] == "video.start"
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
+                assert ws.receive_json()["type"] == "session.done"
+
+        prompt = captured_prompts[0]
+        image = prompt["multi_modal_data"]["image"]  # type: ignore[index]
+        assert image.size == (4, 4)
+        engine_client.abort.assert_not_called()
+
     def test_streaming_session_rejects_unknown_format(self, mocker: MockerFixture):
         async def mock_generate(*_args, **_kwargs):
             if False:
@@ -158,6 +213,33 @@ class TestStreamingVideoOutputWebSocket:
 
         engine_client.abort.assert_not_called()
 
+    def test_first_message_must_be_session_start(self, mocker: MockerFixture):
+        """Control messages sent before session.start are rejected at the handshake.
+
+        This test can practically ensure ``session.interaction`` before ``session.start`` is rejected,
+        but the implementation is more generic to also guard against other control messages.
+        """
+
+        async def mock_generate(*_args, **_kwargs):
+            if False:
+                yield None
+
+        app, _handler, engine_client = _build_test_app(
+            mocker=mocker,
+            streaming_chunks=[],
+            mock_generate=mock_generate,
+        )
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/realtime/video") as ws:
+                ws.send_json({"type": "session.interaction", "interaction": {"event": {"prompt": "too early"}}})
+                err = ws.receive_json()
+                assert err["type"] == "error"
+                assert err["message"] == "Expected session.start, got: session.interaction"
+
+        engine_client.abort.assert_not_called()
+        engine_client.submit_interaction_async.assert_not_called()
+
     def test_streaming_session_emits_final_encoder_delta_before_done(self, mocker: MockerFixture):
         """A non-empty encoder close delta is delivered as the final binary chunk."""
 
@@ -180,8 +262,13 @@ class TestStreamingVideoOutputWebSocket:
             with client.websocket_connect("/v1/realtime/video") as ws:
                 ws.send_json({"type": "session.start", "prompt": "A cat walking in the rain"})
                 assert ws.receive_json()["type"] == "video.start"
-                assert ws.receive_bytes() == b"mp4-chunk-0"
-                assert ws.receive_bytes() == b"mp4-trailer"
+                metadata, chunk = _receive_video_chunk(ws)
+                assert metadata["kind"] == "media"
+                assert chunk == b"mp4-chunk-0"
+                metadata, chunk = _receive_video_chunk(ws)
+                assert metadata["kind"] == "trailer"
+                assert metadata["generation_chunk_index"] is None
+                assert chunk == b"mp4-trailer"
 
                 done = ws.receive_json()
                 assert done["type"] == "session.done"
@@ -213,7 +300,7 @@ class TestStreamingVideoOutputWebSocket:
             with client.websocket_connect("/v1/realtime/video") as ws:
                 ws.send_json({"type": "session.start", "prompt": "hello"})
                 assert ws.receive_json()["type"] == "video.start"
-                assert ws.receive_bytes() == b"mp4-chunk-0"
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
                 err = ws.receive_json()
                 assert err["type"] == "error"
                 assert error_message in err["message"]
@@ -259,6 +346,7 @@ def _build_handler_for_async_tests(
 ) -> tuple[OmniStreamingVideoOutputHandler, MagicMock, _MockVideoOutputWebSocket]:
     engine_client = mocker.MagicMock()
     engine_client.abort = mocker.AsyncMock()
+    engine_client.submit_interaction_async = mocker.AsyncMock()
     engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()]
     engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
     engine_client.generate = mock_generate
@@ -340,8 +428,8 @@ class TestStreamingVideoOutputStallTimeout:
             with client.websocket_connect("/v1/realtime/video") as ws:
                 ws.send_json({"type": "session.start", "prompt": "slow generation"})
                 assert ws.receive_json()["type"] == "video.start"
-                assert ws.receive_bytes() == b"mp4-chunk-0"
-                assert ws.receive_bytes() == b"mp4-chunk-1"
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-1"
                 done = ws.receive_json()
                 assert done["type"] == "session.done"
                 assert done["stopped"] is False
@@ -444,3 +532,374 @@ class TestStreamingVideoOutputStallTimeout:
         assert any(m.get("type") == "session.pong" for m in ws.sent_json)
         assert any(m.get("type") == "session.done" for m in ws.sent_json)
         engine_client.abort.assert_not_called()
+
+
+class TestStreamingVideoOutputPromptUpdate:
+    """session.interaction control messages during streaming."""
+
+    def test_prompt_update_accepted_during_streaming(self, mocker: MockerFixture):
+        """A valid interaction during generation is forwarded to the engine and acknowledged."""
+
+        async def mock_generate(*_args, **_kwargs):
+            yield OmniRequestOutput.from_diffusion(
+                request_id="req-ws",
+                images=[_fake_video_frames(2)],
+                final_output_type="image",
+                finished=False,
+            )
+            await asyncio.sleep(0.3)
+            yield OmniRequestOutput.from_diffusion(
+                request_id="req-ws",
+                images=[_fake_video_frames(2)],
+                final_output_type="image",
+                finished=True,
+            )
+
+        app, _handler, engine_client = _build_test_app(
+            mocker=mocker,
+            streaming_chunks=[(b"mp4-chunk-0", False), (b"mp4-chunk-1", True)],
+            mock_generate=mock_generate,
+        )
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/realtime/video") as ws:
+                ws.send_json({"type": "session.start", "prompt": "initial scene"})
+                start = ws.receive_json()
+                assert start["type"] == "video.start"
+                request_id = start["request_id"]
+
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
+
+                ws.send_json(
+                    {
+                        "type": "session.interaction",
+                        "interaction": {
+                            "event_id": "ui-update-42",
+                            "event": {"prompt": "new scene"},
+                        },
+                    }
+                )
+                queued = ws.receive_json()
+                assert queued["type"] == "session.interaction.queued"
+                assert queued["request_id"] == request_id
+                assert queued["event_id"] == "ui-update-42"
+
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-1"
+                done = ws.receive_json()
+                assert done["type"] == "session.done"
+                assert done["stopped"] is False
+
+        engine_client.submit_interaction_async.assert_awaited_once_with(
+            request_id,
+            interaction={"event_id": "ui-update-42", "event": {"prompt": "new scene"}},
+        )
+
+    def test_prompt_update_rejects_duplicate_event_id(self, mocker: MockerFixture):
+        """Any repeated event_id is rejected at the session API boundary."""
+
+        async def mock_generate(*_args, **_kwargs):
+            yield OmniRequestOutput.from_diffusion(
+                request_id="req-ws",
+                images=[_fake_video_frames(2)],
+                final_output_type="image",
+                finished=False,
+            )
+            await asyncio.sleep(0.3)
+            yield OmniRequestOutput.from_diffusion(
+                request_id="req-ws",
+                images=[_fake_video_frames(2)],
+                final_output_type="image",
+                finished=True,
+            )
+
+        app, _handler, engine_client = _build_test_app(
+            mocker=mocker,
+            streaming_chunks=[(b"mp4-chunk-0", False), (b"mp4-chunk-1", True)],
+            mock_generate=mock_generate,
+        )
+
+        interaction = {
+            "type": "session.interaction",
+            "interaction": {
+                "event_id": "ui-update-42",
+                "event": {"prompt": "new scene"},
+            },
+        }
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/realtime/video") as ws:
+                ws.send_json({"type": "session.start", "prompt": "initial scene"})
+                request_id = ws.receive_json()["request_id"]
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
+
+                ws.send_json(interaction)
+                queued = ws.receive_json()
+                assert queued["type"] == "session.interaction.queued"
+
+                ws.send_json(interaction)
+                err = ws.receive_json()
+                assert err["type"] == "error"
+                assert err["code"] == "event_id_conflict"
+                assert err["request_id"] == request_id
+                assert err["event_id"] == "ui-update-42"
+
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-1"
+                assert ws.receive_json()["type"] == "session.done"
+
+        engine_client.submit_interaction_async.assert_awaited_once_with(
+            request_id,
+            interaction={"event_id": "ui-update-42", "event": {"prompt": "new scene"}},
+        )
+
+    def test_prompt_update_forwards_transition_chunks(self, mocker: MockerFixture):
+        """transition_chunks is passed through to the engine client."""
+
+        async def mock_generate(*_args, **_kwargs):
+            yield OmniRequestOutput.from_diffusion(
+                request_id="req-ws",
+                images=[_fake_video_frames(2)],
+                final_output_type="image",
+                finished=False,
+            )
+            await asyncio.sleep(0.3)
+            yield OmniRequestOutput.from_diffusion(
+                request_id="req-ws",
+                images=[_fake_video_frames(2)],
+                final_output_type="image",
+                finished=True,
+            )
+
+        app, _handler, engine_client = _build_test_app(
+            mocker=mocker,
+            streaming_chunks=[(b"mp4-chunk-0", False), (b"mp4-chunk-1", True)],
+            mock_generate=mock_generate,
+        )
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/realtime/video") as ws:
+                ws.send_json({"type": "session.start", "prompt": "initial"})
+                request_id = ws.receive_json()["request_id"]
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
+
+                ws.send_json(
+                    {
+                        "type": "session.interaction",
+                        "interaction": {
+                            "event_id": "ui-update-42",
+                            "event": {"prompt": "fade to sunset"},
+                            "transition_chunks": 5,
+                        },
+                    }
+                )
+                assert ws.receive_json()["type"] == "session.interaction.queued"
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-1"
+                assert ws.receive_json()["type"] == "session.done"
+
+        engine_client.submit_interaction_async.assert_awaited_once_with(
+            request_id,
+            interaction={
+                "event_id": "ui-update-42",
+                "event": {"prompt": "fade to sunset"},
+                "transition_chunks": 5,
+            },
+        )
+
+    def test_prompt_update_rejects_empty_prompt(self, mocker: MockerFixture):
+        """An empty prompt is rejected before calling the engine."""
+
+        async def mock_generate(*_args, **_kwargs):
+            yield OmniRequestOutput.from_diffusion(
+                request_id="req-ws",
+                images=[_fake_video_frames(2)],
+                final_output_type="image",
+                finished=False,
+            )
+            await asyncio.sleep(0.3)
+            yield OmniRequestOutput.from_diffusion(
+                request_id="req-ws",
+                images=[_fake_video_frames(2)],
+                final_output_type="image",
+                finished=True,
+            )
+
+        app, _handler, engine_client = _build_test_app(
+            mocker=mocker,
+            streaming_chunks=[(b"mp4-chunk-0", False), (b"mp4-chunk-1", True)],
+            mock_generate=mock_generate,
+        )
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/realtime/video") as ws:
+                ws.send_json({"type": "session.start", "prompt": "initial"})
+                assert ws.receive_json()["type"] == "video.start"
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
+
+                ws.send_json({"type": "session.interaction", "interaction": {"event": {"prompt": "   "}}})
+                err = ws.receive_json()
+                assert err["type"] == "error"
+                assert "non-empty event.prompt" in err["message"]
+
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-1"
+                assert ws.receive_json()["type"] == "session.done"
+
+        engine_client.submit_interaction_async.assert_not_called()
+
+    @pytest.mark.parametrize(
+        ("transition_chunks"),
+        ["not-an-int", 2.9, True, -1],
+    )
+    def test_prompt_update_rejects_invalid_transition_chunks(
+        self,
+        mocker: MockerFixture,
+        transition_chunks: object,
+    ):
+        """Non-integral or negative transition_chunks returns an error."""
+
+        async def mock_generate(*_args, **_kwargs):
+            yield OmniRequestOutput.from_diffusion(
+                request_id="req-ws",
+                images=[_fake_video_frames(2)],
+                final_output_type="image",
+                finished=False,
+            )
+            await asyncio.sleep(0.3)
+            yield OmniRequestOutput.from_diffusion(
+                request_id="req-ws",
+                images=[_fake_video_frames(2)],
+                final_output_type="image",
+                finished=True,
+            )
+
+        app, _handler, engine_client = _build_test_app(
+            mocker=mocker,
+            streaming_chunks=[(b"mp4-chunk-0", False), (b"mp4-chunk-1", True)],
+            mock_generate=mock_generate,
+        )
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/realtime/video") as ws:
+                ws.send_json({"type": "session.start", "prompt": "initial"})
+                assert ws.receive_json()["type"] == "video.start"
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
+
+                ws.send_json(
+                    {
+                        "type": "session.interaction",
+                        "interaction": {
+                            "event": {"prompt": "new scene"},
+                            "transition_chunks": transition_chunks,
+                        },
+                    }
+                )
+                err = ws.receive_json()
+                assert err["type"] == "error"
+
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-1"
+                assert ws.receive_json()["type"] == "session.done"
+
+        engine_client.submit_interaction_async.assert_not_called()
+
+    def test_prompt_update_engine_failure_returns_error(self, mocker: MockerFixture):
+        """Engine failures during interaction are surfaced to the client."""
+
+        async def mock_generate(*_args, **_kwargs):
+            yield OmniRequestOutput.from_diffusion(
+                request_id="req-ws",
+                images=[_fake_video_frames(2)],
+                final_output_type="image",
+                finished=False,
+            )
+            await asyncio.sleep(0.3)
+            yield OmniRequestOutput.from_diffusion(
+                request_id="req-ws",
+                images=[_fake_video_frames(2)],
+                final_output_type="image",
+                finished=True,
+            )
+
+        app, _handler, engine_client = _build_test_app(
+            mocker=mocker,
+            streaming_chunks=[(b"mp4-chunk-0", False), (b"mp4-chunk-1", True)],
+            mock_generate=mock_generate,
+        )
+        engine_client.submit_interaction_async.side_effect = RuntimeError("engine unavailable")
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/realtime/video") as ws:
+                ws.send_json({"type": "session.start", "prompt": "initial"})
+                assert ws.receive_json()["type"] == "video.start"
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
+
+                ws.send_json(
+                    {
+                        "type": "session.interaction",
+                        "interaction": {
+                            "event_id": "ui-update-42",
+                            "event": {"prompt": "new scene"},
+                        },
+                    }
+                )
+                err = ws.receive_json()
+                assert err["type"] == "error"
+                assert "Failed to apply interaction" in err["message"]
+                assert err["code"] == "interaction_failed"
+                assert err["event_id"] == "ui-update-42"
+
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-1"
+                assert ws.receive_json()["type"] == "session.done"
+
+        engine_client.submit_interaction_async.assert_awaited_once()
+
+    @pytest.mark.asyncio
+    async def test_prompt_update_accepted_via_async_handler(self, mocker: MockerFixture):
+        """interaction acknowledgement is sent while generation is in progress."""
+
+        release = asyncio.Event()
+
+        async def mock_generate(*_args, **_kwargs):
+            await release.wait()
+            yield OmniRequestOutput.from_diffusion(
+                request_id="req-ws",
+                images=[_fake_video_frames(2)],
+                final_output_type="image",
+                finished=True,
+            )
+
+        handler, engine_client, ws = _build_handler_for_async_tests(
+            mocker,
+            mock_generate=mock_generate,
+        )
+
+        async def _send_interaction_after_delay() -> None:
+            await asyncio.sleep(0.05)
+            ws.enqueue(
+                json.dumps(
+                    {
+                        "type": "session.interaction",
+                        "interaction": {
+                            "event_id": "ui-update-42",
+                            "event": {"prompt": "mid-stream update"},
+                            "transition_chunks": 2,
+                        },
+                    }
+                )
+            )
+
+        prompt_task = asyncio.create_task(_send_interaction_after_delay())
+        session_task = asyncio.create_task(handler.handle_session(ws))  # type: ignore[arg-type]
+        await asyncio.sleep(0.15)
+        release.set()
+        await session_task
+        await prompt_task
+
+        start = next(m for m in ws.sent_json if m.get("type") == "video.start")
+        request_id = start["request_id"]
+        assert any(m.get("type") == "session.interaction.queued" for m in ws.sent_json)
+        engine_client.submit_interaction_async.assert_awaited_once_with(
+            request_id,
+            interaction={
+                "event_id": "ui-update-42",
+                "event": {"prompt": "mid-stream update"},
+                "transition_chunks": 2,
+            },
+        )

tests/entrypoints/test_async_omni_prompt_update.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Entrypoint contract tests for ``AsyncOmni.submit_interaction_async``."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.client_request_state import ClientRequestState
+from vllm_omni.inputs.data import OmniInteractionPrompt
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_async_omni(*, num_stages: int = 1, stage_type: str = "diffusion") -> AsyncOmni:
+    omni = object.__new__(AsyncOmni)
+    omni.log_stats = False
+    omni.request_states = {
+        "external-abc-uuid-1": ClientRequestState(
+            request_id="external-abc-uuid-1",
+            external_request_id="external-abc",
+            queue=AsyncMock(),
+        ),
+        "external-abc-uuid-2": ClientRequestState(
+            request_id="external-abc-uuid-2",
+            external_request_id="external-abc",
+            queue=AsyncMock(),
+        ),
+        "other-req-uuid": ClientRequestState(
+            request_id="other-req-uuid",
+            external_request_id="other-req",
+            queue=AsyncMock(),
+        ),
+    }
+    omni.engine = SimpleNamespace(  # pyright: ignore[reportAttributeAccessIssue]
+        num_stages=num_stages,
+        get_stage_metadata=lambda stage_id: SimpleNamespace(stage_type=stage_type),
+        submit_interaction_async=AsyncMock(),
+    )
+    return omni
+
+
+def _prompt_interaction(prompt: str = "new prompt", transition_chunks: int | None = None) -> OmniInteractionPrompt:
+    interaction: OmniInteractionPrompt = {"event_id": "ui-update-1", "event": {"prompt": prompt}}
+    if transition_chunks is not None:
+        interaction["transition_chunks"] = transition_chunks
+    return interaction
+
+
+@pytest.mark.asyncio
+async def test_submit_interaction_async_maps_external_to_internal_id() -> None:
+    omni = _make_async_omni()
+    with pytest.raises(ValueError, match="exactly one active request"):
+        await omni.submit_interaction_async("external-abc", interaction=_prompt_interaction(transition_chunks=2))
+
+    omni.request_states.pop("external-abc-uuid-2")
+    await omni.submit_interaction_async(
+        "external-abc",
+        interaction=_prompt_interaction(transition_chunks=2),
+    )
+    omni.engine.submit_interaction_async.assert_awaited_once_with(  # pyright: ignore[reportAttributeAccessIssue]
+        "external-abc-uuid-1",
+        interaction=_prompt_interaction(transition_chunks=2),
+    )
+
+
+@pytest.mark.asyncio
+async def test_submit_interaction_async_passes_missing_transition_chunks() -> None:
+    omni = _make_async_omni()
+    omni.request_states.pop("external-abc-uuid-2")
+    await omni.submit_interaction_async("external-abc", interaction=_prompt_interaction())
+    omni.engine.submit_interaction_async.assert_awaited_once_with(  # pyright: ignore[reportAttributeAccessIssue]
+        "external-abc-uuid-1",
+        interaction=_prompt_interaction(),
+    )
+
+
+@pytest.mark.asyncio
+async def test_submit_interaction_async_rejects_empty_prompt() -> None:
+    omni = _make_async_omni()
+    omni.request_states.pop("external-abc-uuid-2")
+    with pytest.raises(ValueError, match="prompt must be non-empty"):
+        await omni.submit_interaction_async("external-abc", interaction=_prompt_interaction(""))
+
+
+@pytest.mark.asyncio
+async def test_submit_interaction_async_rejects_negative_transition_chunks() -> None:
+    omni = _make_async_omni()
+    omni.request_states.pop("external-abc-uuid-2")
+    with pytest.raises(ValueError, match="transition_chunks must be >= 0"):
+        await omni.submit_interaction_async(
+            "external-abc",
+            interaction=_prompt_interaction(transition_chunks=-1),
+        )
+
+
+@pytest.mark.asyncio
+async def test_submit_interaction_async_rejects_inactive_request() -> None:
+    omni = _make_async_omni()
+    omni.request_states.clear()
+    with pytest.raises(ValueError, match="No active request"):
+        await omni.submit_interaction_async("missing", interaction=_prompt_interaction(transition_chunks=2))
+
+
+@pytest.mark.asyncio
+async def test_submit_interaction_async_rejects_non_diffusion() -> None:
+    omni = _make_async_omni(stage_type="llm")
+    omni.request_states = {
+        "req-uuid": ClientRequestState(
+            request_id="req-uuid",
+            external_request_id="req",
+            queue=AsyncMock(),
+        ),
+    }
+    with pytest.raises(ValueError, match="requires a diffusion stage"):
+        await omni.submit_interaction_async("req", interaction=_prompt_interaction(transition_chunks=2))
+
+
+@pytest.mark.asyncio
+async def test_submit_interaction_async_forwards_unsupported_dict_shape_to_engine() -> None:
+    omni = _make_async_omni()
+    omni.request_states.pop("external-abc-uuid-2")
+    payload: OmniInteractionPrompt = {  # pyright: ignore[reportAssignmentType]  # deliberately wrong
+        "event": {"multi_modal_data": {"unsupported_modality": {"type": "pose"}}},
+        "transition_chunks": 2,
+    }
+
+    await omni.submit_interaction_async(
+        "external-abc",
+        interaction=payload,
+    )
+
+    omni.engine.submit_interaction_async.assert_awaited_once_with(  # pyright: ignore[reportAttributeAccessIssue]
+        "external-abc-uuid-1",
+        interaction=payload,
+    )