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

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

docs/contributing/model/adding_diffusion_model.md

@@ -779,6 +779,31 @@ turning them on. For Qwen-Image-style serving examples, document
 `--step-execution` as the feature gate and `--max-num-seqs N` as the
 companion batching knob.
 
+### Streaming Output and Mid-Stream Interactions
+
+To add streaming generation with optional mid-stream prompt updates to a new diffusion pipeline (often a video generation pipeline):
+
+0. **Model should internally support** chunked video emission and bounded VAE decode of a chunk.
+1. **Implement stepwise execution** so interaction RPCs can be consumed at chunk
+   boundaries (see [Step Execution](#step-execution) above).
+2. **Mix in** ``InteractionMixin`` on the pipeline class.
+3. **Implement** ``peek_chunk_media(StepRequestState) -> ChunkMediaSpec`` to offer information
+   about the next chunk to generate.
+   This is called whenever one supported interaction modality needs frame-level calculation,
+   e.g., per-frame camera trajectory. If only chunk-level calculation is involved, peeking is skipped.
+4. **Optionally implement** ``prepare_next_chunk(StepRequestState)`` if the
+   pipeline must rebuild something (e.g., per-chunk state) after an interaction is applied.
+   This function is called at chunk boundary.
+5. **Register** supported modalities and handlers in ``vllm_omni/diffusion/interaction/registry.py``
+   (keyed by ``od_config.model_class_name`` → modality key as in ``req["multi_modal_data"]``).
+
+Serve with ``--diffusion-streaming-output`` (auto-enables step execution). See
+[Diffusion Execution Modes](../../user_guide/diffusion/execution_modes.md#streaming-output)
+and the example script at `examples/online_serving/streaming_video_generation/`.
+
+**Interaction payload contract:** Although ``OmniInteractionPrompt`` type annocation requires ``event_id``,
+the WebSocket API payload may omit ``event_id`` and let the API layer assigns one.
+
 ### Cache Acceleration
 
 #### TeaCache

docs/user_guide/diffusion/execution_modes.md

@@ -8,7 +8,7 @@ are exposed.
 ## Choose a Mode
 
 | Goal | CLI configuration |
-|---|---|
+| --- | --- |
 | Serial request execution | `--max-num-seqs 1` |
 | Fused request-level batching | `--max-num-seqs N` |
 | Single-request step execution | `--step-execution --max-num-seqs 1` |
@@ -177,8 +177,10 @@ wait
 The scheduler may batch these requests when their sampling parameters are
 compatible. See the
 [Image Generation API](../../serving/image_generation_api.md) for response
-formats and additional client examples. Streaming models can use
-model-specific streaming endpoints documented by their serving guide.
+formats and additional client examples.
+
+Streaming models can use model-specific streaming endpoints documented at
+[Streaming Video Output API](../../serving/streaming_video_output_api.md).
 
 ## Python API
 
@@ -241,7 +243,7 @@ For step execution, set `step_execution: true` and remove
 ## CLI Reference
 
 | Flag | Default | Effect |
-|---|---:|---|
+| --- | ---: | --- |
 | `--step-execution` | disabled | Select step-wise scheduling |
 | `--max-num-seqs` | `1` for diffusion stages | Set request- or step-scheduler capacity |
 | `--request-batch-max-wait-ms` | `0` | Wait for burst coalescing in request mode |

tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# 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
@@ -7,6 +7,7 @@ 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
@@ -21,6 +22,11 @@ from tests.engine.test_orchestrator import (
 )
 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
@@ -54,30 +60,55 @@ def pipeline() -> HeliosPipeline:
             None,
         )
     )
-    pipeline._prepare_next_chunk = MagicMock()
+    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") -> StepRequestState:
+def _make_diffusion_request_state(*, request_id: str = "req-1", fps: int | None = None) -> StepRequestState:
     state = StepRequestState(
         request_id=request_id,
-        sampling=SimpleNamespace(num_outputs_per_prompt=1, max_sequence_length=226),  # pyright: ignore[reportArgumentType]
+        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 = {}
+    state.extra = {"window_num_frames": 8}
     return state
 
 
-def _make_diffusion_model_runner(*, pipeline, streaming_output: bool = True) -> DiffusionModelRunner:
+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="HeliosPipeline",
+        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
 
 
@@ -108,10 +139,13 @@ class TestPromptUpdateExecution:
         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))
+        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."""
@@ -135,6 +169,9 @@ class TestPromptUpdateExecution:
             {},
             {"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(
@@ -149,6 +186,34 @@ class TestPromptUpdateExecution:
         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(
@@ -159,28 +224,37 @@ class TestPromptUpdateExecution:
         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
+        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.extra["prompt_update_version"] = 1
+        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_prepare_prompt_update_queues_pending_target(self, pipeline: HeliosPipeline) -> None:
-        """HeliosPipeline queues target embeds without mutating current prompt_embeds."""
+    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()
 
-        pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=2)
+        PromptInteractionHandler.from_pipeline(pipeline).enqueue(
+            state,
+            event_id="ui-update-1",
+            received_at=0.0,
+            payload={"prompt": "new scene"},
+            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
+        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_prepare_prompt_update_rejects_before_initial_generation(self, pipeline: HeliosPipeline) -> None:
+    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
@@ -189,43 +263,209 @@ class TestPromptUpdateExecution:
             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)
+            PromptInteractionHandler.from_pipeline(pipeline).enqueue(
+                state,
+                event_id="ui-update-1",
+                received_at=0.0,
+                payload={"prompt": "new scene"},
+                transition_chunks=2,
+            )
 
-        assert "pending_prompt_update" not in state.extra
+        session = state.interaction_sessions.get("prompt")
+        assert session is None or (isinstance(session, PromptSession) and session.pending_event is None)
 
-    def test_helios_apply_prompt_update_at_chunk_boundary_starts_transition(self, pipeline: HeliosPipeline) -> 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()
-        pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=2)
+        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_prompt_update_at_chunk_boundary(state)
+        pipeline.apply_interaction_at_chunk_boundary(state)
 
-        assert state.extra.get("prompt_update_state") is not None
+        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 state.extra["prompt_update_version"] == 1
-        assert state.extra["prompt_update_chunk_metadata"] == {
+        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_prompt_update_advances_transition_over_chunks(self, pipeline: HeliosPipeline) -> None:
+    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()
-        pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=3)
-        pipeline._apply_prompt_update_at_chunk_boundary(state)
+        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_prompt_update_at_chunk_boundary(state)
+            pipeline.apply_interaction_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
+        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_prompt_update_at_chunk_boundary(state)
-        assert state.extra["prompt_update_version"] == 3
+        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:
@@ -266,8 +506,8 @@ class TestPromptUpdateIntegration:
 
     @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]
+        """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)
@@ -293,12 +533,7 @@ class TestPromptUpdateIntegration:
         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,
-        )
+        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",
@@ -337,9 +572,10 @@ class TestPromptUpdateIntegration:
 
             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))
+            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:
@@ -544,7 +780,10 @@ class TestPromptUpdateIntegration:
             ]
 
             deadline = time.monotonic() + 5.0
-            while "pending_prompt_update" not in state.extra:
+            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)

vllm_omni/diffusion/interaction/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/diffusion/interaction/coordinator.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Interaction coordinator to be co-owned by a pipeline and a runner."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from vllm_omni.diffusion.interaction.registry import STRUCTURED_HANDLER_REGISTRY
+from vllm_omni.diffusion.interaction.types import (
+    InteractionChunkMetadata,
+    InteractionPayload,
+    merge_interaction_metadata,
+)
+from vllm_omni.diffusion.models.interface import supports_interaction_apply
+from vllm_omni.diffusion.worker.utils import StepRequestState
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.interaction.modality_handlers.base import InteractionHandler
+
+
+class InteractionCoordinator:
+    """Handles mid-way interaction inputs for one loaded pipeline.
+
+    Request-local session state lives on ``StepRequestState.interaction_sessions``;
+    this object is request-agnostic.
+    """
+
+    def __init__(
+        self,
+        handlers: dict[str, InteractionHandler],
+        *,
+        model_class_name: str | None = None,
+    ) -> None:
+        self._handlers = handlers
+        self.model_class_name = model_class_name
+
+    @classmethod
+    def build(cls, pipeline: object, od_config: OmniDiffusionConfig) -> InteractionCoordinator:
+        model_class_name = od_config.model_class_name
+        handlers: dict[str, InteractionHandler] = {}
+
+        if supports_interaction_apply(pipeline):
+            for modality, handler_cls in STRUCTURED_HANDLER_REGISTRY.get(model_class_name or "", {}).items():
+                handlers[modality] = handler_cls.from_pipeline(pipeline)
+
+        return cls(handlers, model_class_name=model_class_name)
+
+    def has_modality(self, modality: str) -> bool:
+        return modality in self._handlers
+
+    @property
+    def needs_chunk_media(self) -> bool:
+        """Whether any registered handler needs chunk ``num_frames``/``fps``."""
+        return any(handler.needs_chunk_media for handler in self._handlers.values())
+
+    def get_handler(self, modality: str) -> InteractionHandler:
+        handler = self._handlers.get(modality)
+        if handler is None:
+            raise ValueError(
+                f"interaction modality {modality!r} is not supported by pipeline {self.model_class_name!r}"
+            )
+        return handler
+
+    def enqueue(
+        self,
+        state: StepRequestState,
+        *,
+        modality: str,
+        event_id: str,
+        received_at: float,
+        payload: InteractionPayload,
+        transition_chunks: int | None,
+    ) -> None:
+        """Fan out enqueue to handlers in stable order (prompt first)."""
+        handler = self.get_handler(modality)
+        handler.enqueue(
+            state,
+            event_id=event_id,
+            received_at=received_at,
+            payload=payload,
+            transition_chunks=transition_chunks,
+        )
+
+    def apply_at_chunk_boundary(
+        self,
+        state: StepRequestState,
+        *,
+        boundary_at: float,
+        chunk_index: int | None = None,
+        num_frames: int | None = None,
+        fps: float | None = None,
+    ) -> InteractionChunkMetadata:
+        """Fan out chunk-boundary apply to handlers in stable order (prompt first)."""
+        if chunk_index is None:
+            chunk_index = state.chunk_index
+        metas: list[InteractionChunkMetadata] = []
+        for handler in self._handlers_in_apply_order():
+            if handler.modality not in state.interaction_sessions:
+                # This function is always called at chunk boundary without knowing if an interaction is really enqueued.
+                # `state.interaction_sessions[this modality]` is lazy-populated after the first interaction is enqueued.
+                # So, check here and skip a modality handler if this modality's session is not yet created.
+                continue
+            meta = handler.apply_at_chunk_boundary(
+                state,
+                boundary_at=boundary_at,
+                chunk_index=chunk_index,
+                num_frames=num_frames,
+                fps=fps,
+            )
+            if meta is not None:
+                metas.append(meta)
+        return merge_interaction_metadata(metas)
+
+    def _handlers_in_apply_order(self) -> list[InteractionHandler]:
+        ordered: list[InteractionHandler] = []
+        if "prompt" in self._handlers:
+            ordered.append(self._handlers["prompt"])
+        for modality, handler in self._handlers.items():
+            if modality == "prompt":
+                continue
+            ordered.append(handler)
+        return ordered

vllm_omni/diffusion/interaction/mixin.py

@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Interaction-handling mixin for diffusion pipelines. Process interactions at chunk boundaries."""
+
+from __future__ import annotations
+
+import time
+from typing import TYPE_CHECKING, cast
+
+from vllm_omni.diffusion.models.interface import SupportsInteractionApply
+from vllm_omni.diffusion.worker.utils import StepRequestState
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.interaction.coordinator import InteractionCoordinator
+
+
+class InteractionMixin:
+    """Unified chunk-boundary interaction hook.
+
+    Pipelines inherit this mixin and implement ``peek_chunk_media`` to offer playback information about the next chunk.
+    Diffusion runner calls ``apply_interaction_at_chunk_boundary`` without knowing which modality is being applied.
+
+    Per-modality session state lives on ``StepRequestState.interaction_sessions``;
+    handler strategy objects live on the pipeline/runner ``InteractionCoordinator``.
+    """
+
+    _interaction_coordinator: InteractionCoordinator | None = None
+
+    def apply_interaction_at_chunk_boundary(self, state: StepRequestState) -> None:
+        """Advance all active interaction tracks before the next chunk."""
+        if self._interaction_coordinator is None:
+            raise RuntimeError(
+                "interaction coordinator is not initialized; "
+                "DiffusionModelRunner.load_model must wire InteractionCoordinator "
+                "onto the pipeline before chunked generation"
+            )
+
+        num_frames: int | None = None
+        fps: float | None = None
+        if self._interaction_coordinator.needs_chunk_media:
+            media = cast(SupportsInteractionApply, self).peek_chunk_media(state)
+            num_frames = media.num_frames
+            fps = media.fps
+
+        merged = self._interaction_coordinator.apply_at_chunk_boundary(
+            state,
+            boundary_at=time.monotonic(),
+            num_frames=num_frames,
+            fps=fps,
+        )
+        state.interaction_chunk_metadata = merged
+
+    def prepare_next_chunk(self, state: StepRequestState) -> None:
+        """Set up pipeline state for the next chunk after interaction apply.
+
+        Default no-op. Pipelines may override this if needed.
+        """
+        pass

vllm_omni/diffusion/interaction/modality_handlers/base.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Base interaction handler interface."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import ClassVar, Self
+
+from vllm_omni.diffusion.interaction.types import (
+    InteractionChunkMetadata,
+    InteractionPayload,
+)
+from vllm_omni.diffusion.worker.utils import StepRequestState
+
+
+class InteractionHandler(ABC):
+    """Strategy object for one interaction modality.
+
+    Handler instances are pipeline/runner-owned and request-agnostic. Per-request
+    session state lives on ``StepRequestState.interaction_sessions``.
+    Each concrete modality handler decides how to handle timing from ``received_at``.
+    """
+
+    modality: ClassVar[str]
+    # When True, chunk-boundary apply needs ``ChunkMediaSpec`` (num_frames/fps)
+    # Those information are useful when interaction handler needs interpolation/integration on a frame-by-frame basis
+    needs_chunk_media: ClassVar[bool] = False
+
+    @classmethod
+    def from_pipeline(cls, pipeline: object) -> Self:
+        """Bind a request-agnostic handler to ``pipeline`` when construction needs it.
+
+        Default construction ignores the pipeline. Prompt handlers override this to
+        capture ``encode_prompt`` / device / dtype.
+        """
+        del pipeline
+        return cls()
+
+    @abstractmethod
+    def enqueue(
+        self,
+        state: StepRequestState,
+        *,
+        event_id: str,
+        received_at: float,
+        payload: InteractionPayload,
+        transition_chunks: int | None,
+    ) -> None:
+        """Validate and queue this track on request-local state."""
+
+    @abstractmethod
+    def apply_at_chunk_boundary(
+        self,
+        state: StepRequestState,
+        *,
+        boundary_at: float,
+        chunk_index: int | None = None,  # defaults to state.chunk_index when omitted
+        num_frames: int | None = None,  # only set when at least one interaction handler needs it
+        fps: float | None = None,  # only set when at least one interaction handler needs it
+    ) -> InteractionChunkMetadata | None:
+        """Advance request-local state and materialize this chunk's effects."""

vllm_omni/diffusion/interaction/modality_handlers/prompt.py

@@ -0,0 +1,266 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Prompt-track interaction handler (midway prompt updates)."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+from typing import Any, ClassVar
+
+import torch
+from typing_extensions import override
+
+from vllm_omni.diffusion.interaction.modality_handlers.base import InteractionHandler
+from vllm_omni.diffusion.interaction.types import (
+    InteractionChunkMetadata,
+    InteractionEvent,
+    InteractionMode,
+    InteractionPayload,
+    InteractionSession,
+)
+from vllm_omni.diffusion.worker.utils import StepRequestState
+
+DEFAULT_TRANSITION_CHUNKS = 3
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(kw_only=True)
+class QueuedPromptEvent(InteractionEvent):
+    """Last-write-wins prompt update waiting for / active at a chunk boundary.
+
+    ``source_prompt_embeds`` is the lerp start endpoint, staged when this event
+    is activated (not at enqueue). Live blended embeds live on
+    ``StepRequestState.prompt_embeds``.
+    """
+
+    mode: InteractionMode = "target"
+    prompt: str
+    target_prompt_embeds: torch.Tensor
+    source_prompt_embeds: torch.Tensor | None = None
+
+    def blended_prompt_embeds(self) -> torch.Tensor:
+        assert self.source_prompt_embeds is not None
+        alpha = self._current_alpha()
+        if alpha >= 1.0:
+            return self.target_prompt_embeds
+        return (1.0 - alpha) * self.source_prompt_embeds + alpha * self.target_prompt_embeds
+
+    def advance_transition(self) -> None:
+        assert self.source_prompt_embeds is not None
+        if self.transition_chunks <= 0:
+            self.source_prompt_embeds = self.target_prompt_embeds
+            return
+        self.elapsed_transition_chunks += 1
+        if self.elapsed_transition_chunks >= self.transition_chunks:
+            self.source_prompt_embeds = self.target_prompt_embeds
+            self.elapsed_transition_chunks = float(self.transition_chunks)
+
+    def _current_alpha(self) -> float:
+        if self.transition_chunks <= 0:
+            return 1.0
+        return min(1.0, self.elapsed_transition_chunks / self.transition_chunks)
+
+
+@dataclass
+class PromptSession(InteractionSession):
+    """Request-local prompt-update session under ``state.interaction_sessions['prompt']``."""
+
+    pending_event: QueuedPromptEvent | None = None
+    active_event: QueuedPromptEvent | None = None
+    version: int = 0
+
+
+def prompt_update_versions(states: Sequence[StepRequestState]) -> tuple[int, ...]:
+    """Return per-request prompt-update versions for batch cache comparison."""
+    versions: list[int] = []
+    for state in states:
+        session = state.interaction_sessions.get("prompt")
+        versions.append(session.version if isinstance(session, PromptSession) else 0)
+    return tuple(versions)
+
+
+class PromptInteractionHandler(InteractionHandler):
+    """Encode-and-queue prompt updates; lerp embeds at chunk boundaries.
+
+    Mirrors the former ``PromptUpdateMixin`` behavior. Encoding happens at
+    enqueue time via an injected ``encode_prompt`` callable.
+    """
+
+    modality: ClassVar[str] = "prompt"
+
+    def __init__(
+        self,
+        *,
+        encode_prompt: Callable[..., tuple[torch.Tensor, ...]],
+        device: torch.device | str,
+        dtype: torch.dtype,
+    ) -> None:
+        self._encode_prompt = encode_prompt
+        self._device = device
+        self._dtype = dtype
+
+    @classmethod
+    @override
+    def from_pipeline(cls, pipeline: Any) -> PromptInteractionHandler:
+        """Build a handler from a diffusion pipeline that exposes ``encode_prompt``."""
+        return cls(
+            encode_prompt=pipeline.encode_prompt,
+            device=pipeline.device,
+            dtype=pipeline.transformer.dtype,
+        )
+
+    @override
+    def enqueue(
+        self,
+        state: StepRequestState,
+        *,
+        event_id: str,
+        received_at: float,
+        payload: InteractionPayload,
+        transition_chunks: int | None,
+    ) -> None:
+        """Prompt updates are last-write-win and unbuffered at chunk boundary."""
+        prompt = payload.get("prompt")
+        if not isinstance(prompt, str) or not prompt:
+            raise ValueError("prompt must be non-empty")
+        if not event_id:
+            raise ValueError("event_id must be non-empty")
+        if state.prompt_embeds is None:
+            raise ValueError(
+                f"prompt_update is not allowed before initial generation has started (request_id={state.request_id!r})"
+            )
+        duration = DEFAULT_TRANSITION_CHUNKS if transition_chunks is None else transition_chunks
+        if duration < 0:
+            raise ValueError("transition_chunks must be >= 0")
+
+        target_prompt_embeds, _ = self._encode_prompt(
+            prompt=prompt,
+            negative_prompt=None,
+            do_classifier_free_guidance=False,
+            num_videos_per_prompt=state.sampling.num_outputs_per_prompt,
+            max_sequence_length=state.sampling.max_sequence_length,
+            device=self._device,
+            dtype=self._dtype,
+        )
+        session = state.interaction_sessions.setdefault("prompt", PromptSession())
+        assert isinstance(session, PromptSession)
+        with session.lock:
+            # Chunk-level LWW: replace any prior pending event.
+            session.pending_event = QueuedPromptEvent(
+                event_id=event_id,
+                received_at=received_at,
+                transition_chunks=duration,
+                prompt=prompt,
+                target_prompt_embeds=target_prompt_embeds,
+            )
+
+    @override
+    def apply_at_chunk_boundary(
+        self,
+        state: StepRequestState,
+        *,
+        boundary_at: float,
+        chunk_index: int | None = None,
+        num_frames: int | None = None,
+        fps: float | None = None,
+    ) -> InteractionChunkMetadata | None:
+        """Advance or start prompt interpolation before the next chunk."""
+        # Prompt lerp is chunk-LWW; media timeline and caller chunk_index unused.
+        del chunk_index, num_frames, fps
+        session = state.interaction_sessions.get("prompt")
+        assert isinstance(session, PromptSession)
+
+        embeds_changed = False
+        next_chunk_index = state.chunk_index
+        started_event_ids: list[str] = []
+        active_event_ids: list[str] = []
+        completed_event_ids: list[str] = []
+
+        with session.lock:
+            # Prompt ignores frame scheduling but still records the boundary clock.
+            session.last_boundary_at = boundary_at
+            pending_event = session.pending_event
+            session.pending_event = None
+            active_event = session.active_event
+
+            # If current transition is not complete, advance it.
+            # After completion, leave active_event in place (so a later pending
+            # update can abort/overwrite), but do not keep bumping the version.
+            if active_event is not None:
+                in_transition = (
+                    active_event.transition_chunks > 0
+                    and active_event.elapsed_transition_chunks < active_event.transition_chunks
+                )
+                if in_transition:
+                    active_event.advance_transition()
+                    state.prompt_embeds = active_event.blended_prompt_embeds()
+                    embeds_changed = True
+                    active_event_ids.append(active_event.event_id)
+                    if active_event.elapsed_transition_chunks >= active_event.transition_chunks:
+                        state.prompt_embeds = active_event.target_prompt_embeds
+                        active_event.source_prompt_embeds = active_event.target_prompt_embeds
+                        completed_event_ids.append(active_event.event_id)
+                        logger.debug(
+                            "prompt_update transition complete request_id=%s next_chunk_index=%d prompt=%.20s...",
+                            state.request_id,
+                            next_chunk_index,
+                            active_event.prompt,
+                        )
+
+            # If a new prompt update is pending, start a new transition.
+            if pending_event is not None:
+                if state.prompt_embeds is None:
+                    raise RuntimeError(
+                        "internal error: trying to apply a pending prompt update but "
+                        f"current prompt_embeds is None (request_id={state.request_id!r})"
+                    )
+                # Stage lerp-start embeds at activation (after any same-boundary advance).
+                pending_event.source_prompt_embeds = state.prompt_embeds.detach().clone()
+                pending_event.elapsed_transition_chunks = 0.0
+                session.active_event = pending_event
+                active_event = pending_event
+                event_id = pending_event.event_id
+                duration = pending_event.transition_chunks
+                prompt = pending_event.prompt
+                target = pending_event.target_prompt_embeds
+                started_event_ids.append(event_id)
+                active_event_ids.append(event_id)
+                if duration <= 0:
+                    # A hard/immediate transition
+                    state.prompt_embeds = target
+                    pending_event.source_prompt_embeds = target
+                    completed_event_ids.append(event_id)
+                    logger.debug(
+                        "prompt_update sharp transition request_id=%s next_chunk_index=%d prompt=%.20s...",
+                        state.request_id,
+                        next_chunk_index,
+                        prompt,
+                    )
+                else:
+                    # A transition that really takes some time
+                    active_event.advance_transition()
+                    state.prompt_embeds = active_event.blended_prompt_embeds()
+                    if active_event.elapsed_transition_chunks >= active_event.transition_chunks:
+                        state.prompt_embeds = active_event.target_prompt_embeds
+                        active_event.source_prompt_embeds = active_event.target_prompt_embeds
+                        completed_event_ids.append(event_id)
+                    logger.debug(
+                        "prompt_update transition start request_id=%s next_chunk_index=%d prompt=%.20s...",
+                        state.request_id,
+                        next_chunk_index,
+                        prompt,
+                    )
+                embeds_changed = True
+
+            # Indicate that prompt embeddings have changed---clear input batch cache.
+            if embeds_changed:
+                session.version += 1
+
+        return InteractionChunkMetadata(
+            started_event_ids=started_event_ids,
+            active_event_ids=active_event_ids,
+            completed_event_ids=completed_event_ids,
+        )

vllm_omni/diffusion/interaction/registry.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Registry of mid-way interaction handlers during a diffusion generation.
+A pipeline only declares the modalities that it supports.
+And a pipeline can choose which handler for each modality (there can be more than one handler for a modality).
+
+Outer key: pipeline architecture name (``od_config.model_class_name``).
+Inner key: modality name (e.g. ``\"prompt\"``).
+"""
+
+from __future__ import annotations
+
+from vllm_omni.diffusion.interaction.modality_handlers.base import InteractionHandler
+from vllm_omni.diffusion.interaction.modality_handlers.prompt import PromptInteractionHandler
+
+STRUCTURED_HANDLER_REGISTRY: dict[str, dict[str, type[InteractionHandler]]] = {
+    # Pipeline class name -> modality -> handler class.
+    "HeliosPipeline": {
+        "prompt": PromptInteractionHandler,
+    },
+    "HeliosPyramidPipeline": {
+        "prompt": PromptInteractionHandler,
+    },
+}

vllm_omni/diffusion/interaction/types.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared types for midway interaction handlers."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from threading import Lock
+from typing import Literal
+
+# Wire payload for one modality track. Handlers validate keys they need.
+InteractionPayload = Mapping[str, object]
+
+# Shared command semantics across modalities that support them.
+# Prompt specialization uses ``"target"`` only.
+InteractionMode = Literal["target"]
+
+
+@dataclass(frozen=True)
+class ChunkMediaSpec:
+    """Media extent of one generation chunk, for interaction timeline mapping."""
+
+    num_frames: int
+    fps: float
+
+    def __post_init__(self) -> None:
+        if self.num_frames <= 0:
+            raise ValueError(f"ChunkMediaSpec.num_frames must be > 0, got {self.num_frames}")
+        if self.fps <= 0:
+            raise ValueError(f"ChunkMediaSpec.fps must be > 0, got {self.fps}")
+
+    @property
+    def duration_s(self) -> float:
+        return float(self.num_frames) / float(self.fps)
+
+
+@dataclass(kw_only=True)
+class InteractionEvent:
+    """Shared envelope for one queued or in-flight interaction command. **Intended to be subclassed**.
+
+    Example:
+    - ``QueuedPromptEvent`` in ``vllm_omni.diffusion.interaction.modality_handlers.prompt``
+    """
+
+    event_id: str
+    received_at: float
+    mode: InteractionMode
+    transition_chunks: int
+    elapsed_transition_chunks: float = 0.0
+
+
+@dataclass
+class InteractionSession:
+    """Request-local session base for one interaction modality. **Intended to be subclassed**.
+
+    Concrete modalities subclass this and store their own pending/active state.
+    ``last_boundary_at`` is available for frame-scheduled modalities; chunk-LWW
+    modalities may ignore it.
+
+    Example:
+    - ``PromptSession`` in ``vllm_omni.diffusion.interaction.modality_handlers.prompt``
+    """
+
+    lock: Lock = field(default_factory=Lock, repr=False)
+    last_boundary_at: float | None = None
+
+
+@dataclass
+class InteractionChunkMetadata:
+    """Per-modality (or merged) event-id bookkeeping for one chunk boundary."""
+
+    started_event_ids: list[str] = field(default_factory=list)
+    active_event_ids: list[str] = field(default_factory=list)
+    completed_event_ids: list[str] = field(default_factory=list)
+
+    def as_dict(self) -> dict[str, list[str]]:
+        return {
+            "started_event_ids": list(self.started_event_ids),
+            "active_event_ids": list(self.active_event_ids),
+            "completed_event_ids": list(self.completed_event_ids),
+        }
+
+
+def merge_interaction_metadata(
+    metas: list[InteractionChunkMetadata],
+) -> InteractionChunkMetadata:
+    """Merge per-handler metadata while preserving order and de-duplicating ids."""
+    started: list[str] = []
+    active: list[str] = []
+    completed: list[str] = []
+    seen_started: set[str] = set()
+    seen_active: set[str] = set()
+    seen_completed: set[str] = set()
+    for meta in metas:
+        for event_id in meta.started_event_ids:
+            if event_id not in seen_started:
+                seen_started.add(event_id)
+                started.append(event_id)
+        for event_id in meta.active_event_ids:
+            if event_id not in seen_active:
+                seen_active.add(event_id)
+                active.append(event_id)
+        for event_id in meta.completed_event_ids:
+            if event_id not in seen_completed:
+                seen_completed.add(event_id)
+                completed.append(event_id)
+    return InteractionChunkMetadata(
+        started_event_ids=started,
+        active_event_ids=active,
+        completed_event_ids=completed,
+    )

vllm_omni/diffusion/models/helios/pipeline_helios.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Adapted from Helios (https://github.com/BestWishYsh/Helios)
 
 from __future__ import annotations
@@ -18,18 +18,20 @@ from diffusers import AutoencoderKLWan
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel
+from typing_extensions import override
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.interaction.mixin import InteractionMixin
+from vllm_omni.diffusion.interaction.types import ChunkMediaSpec
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel
 from vllm_omni.diffusion.models.helios.scheduling_helios import HeliosScheduler
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
-from vllm_omni.diffusion.prompt_update import PromptUpdateMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.platforms import current_omni_platform
@@ -160,7 +162,7 @@ class HeliosPipeline(
     CFGParallelMixin,
     ProgressBarMixin,
     DiffusionPipelineProfilerMixin,
-    PromptUpdateMixin,
+    InteractionMixin,
     SupportsComponentDiscovery,
 ):
     """Helios text-to-video / image-to-video / video-to-video pipeline for vllm-omni.
@@ -533,10 +535,22 @@ class HeliosPipeline(
                 "zero_steps": int(extra.get("zero_steps", 1)),
             }
         )
-        self._prepare_next_chunk(state)
+        self.prepare_next_chunk(state)
         return state
 
-    def _prepare_next_chunk(self, state: StepRequestState) -> None:
+    def peek_chunk_media(self, state: StepRequestState) -> ChunkMediaSpec:
+        """Expose this chunk's decoded media extent for interaction timelines."""
+        num_frames = int(state.extra["window_num_frames"])
+        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 that use the "
+                f"chunk media timeline, got {fps!r} (request_id={state.request_id!r})"
+            )
+        return ChunkMediaSpec(num_frames=num_frames, fps=float(fps))
+
+    @override
+    def prepare_next_chunk(self, state: StepRequestState) -> None:
         extra = state.extra
         k = state.chunk_index
         is_first_chunk = k == 0
@@ -900,14 +914,9 @@ class HeliosPipeline(
 
         output = current_latents if extra["output_type"] == "latent" else current_video
         completed_chunk_index = state.chunk_index
-        prompt_update_metadata = state.extra.pop("prompt_update_chunk_metadata", {})
         state.chunk_index += 1
         finished = state.request_denoise_completed
-        if not finished:
-            # Apply queued/advancing prompt updates before preparing the next chunk.
-            self._apply_prompt_update_at_chunk_boundary(state)
-            self._prepare_next_chunk(state)
-        else:
+        if finished:
             self._current_timestep = None
             if current_omni_platform.is_available():
                 current_omni_platform.empty_cache()
@@ -917,9 +926,6 @@ class HeliosPipeline(
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else {},
             chunk_index=completed_chunk_index,
             total_chunks=state.total_chunks,
-            started_event_ids=prompt_update_metadata.get("started_event_ids", []),
-            active_event_ids=prompt_update_metadata.get("active_event_ids", []),
-            completed_event_ids=prompt_update_metadata.get("completed_event_ids", []),
             finished=finished,
         )
 

vllm_omni/diffusion/models/interface.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 from __future__ import annotations
 
 from dataclasses import dataclass
@@ -19,6 +19,7 @@ if TYPE_CHECKING:
 
     from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend
     from vllm_omni.diffusion.data import DiffusionOutput
+    from vllm_omni.diffusion.interaction.types import ChunkMediaSpec
     from vllm_omni.diffusion.worker.input_batch import InputBatch
     from vllm_omni.diffusion.worker.utils import StepRequestState
 
@@ -108,30 +109,34 @@ def supports_step_execution(pipeline: object) -> bool:
 
 
 @runtime_checkable
-class SupportsPromptUpdate(Protocol):
-    """Optional protocol for pipelines that support midway prompt updates.
+class SupportsInteractionApply(Protocol):
+    """Optional protocol for pipelines with unified mid-generation, chunk-boundary hooks."""
 
-    Pipelines typically implement this via
-    :class:`~vllm_omni.diffusion.prompt_update.PromptUpdateMixin`.
-    """
+    def peek_chunk_media(self, state: StepRequestState) -> ChunkMediaSpec:
+        """Return the media timeline represented by the upcoming/current chunk.
 
-    supports_prompt_update: ClassVar[bool] = True
+        Useful when interaction handler needs interpolation/integration on a frame-by-frame basis,
+        or for backpressure/pacing.
+        """
+        ...
 
-    def prepare_prompt_update(
-        self,
-        state: StepRequestState,
-        prompt: str,
-        event_id: str,
-        transition_chunks: int | None = None,
-    ) -> None:
-        """Encode and queue a prompt update on request-local state."""
+    def apply_interaction_at_chunk_boundary(self, state: StepRequestState) -> None:
+        """Advance queued interactions before the next generation chunk."""
+        ...
+
+    def prepare_next_chunk(self, state: StepRequestState) -> None:
+        """Set up pipeline state for the next chunk after interaction apply.
+
+        Default implementations is a no-op; model-specific pipelines override
+        when chunk transitions require latent/history bookkeeping.
+        """
         ...
 
 
-def supports_prompt_update(pipeline: object) -> bool:
-    """Return whether ``pipeline`` implements :class:`SupportsPromptUpdate`."""
+def supports_interaction_apply(pipeline: object) -> bool:
+    """Return whether ``pipeline`` implements :class:`SupportsInteractionApply`."""
 
-    return isinstance(pipeline, SupportsPromptUpdate)
+    return isinstance(pipeline, SupportsInteractionApply)
 
 
 @runtime_checkable

vllm_omni/diffusion/prompt_update.py

@@ -1,245 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Midway prompt updates for chunked diffusion step execution.
-
-``PromptUpdateMixin`` exposes
-- a pipeline interface for the external to queue a prompt update,
-- a pipeline internal method to apply prompt transitions at chunk boundaries.
-
-``prompt_update_versions``: a helper for ``InputBatch`` cache invalidation.
-"""
-
-from __future__ import annotations
-
-import logging
-from collections.abc import Sequence
-from dataclasses import dataclass
-from threading import Lock
-from typing import ClassVar, TypedDict, cast
-
-import torch
-
-from vllm_omni.diffusion.worker.utils import StepRequestState
-
-DEFAULT_TRANSITION_CHUNKS = 3
-
-logger = logging.getLogger(__name__)
-
-
-class PromptUpdateExtra(TypedDict, total=False):
-    """A "protocol" for ``StepRequestState.extra`` keys used by the prompt-update feature."""
-
-    pending_prompt_update: _PendingPromptUpdate
-    prompt_update_state: _PromptUpdateState
-    prompt_update_chunk_metadata: _PromptUpdateChunkMetadata
-    prompt_update_version: int  # Bumped when prompt_embeds change; invalidates InputBatch cache.
-
-    # The lock for storing/updating pending prompt updates.
-    # Saved in state.extra to be naturally cleaned up with the request state.
-    # As PromptUpdateMixin (the pipeline) currently lacks an obvious request-finish hook
-    prompt_update_lock: Lock
-
-
-def prompt_update_versions(states: Sequence[StepRequestState]) -> tuple[int, ...]:
-    """Return per-request prompt-update versions for batch cache comparison.
-
-    To be used by ``InputBatch`` or other external to determine if the cached batch is still valid.
-    """
-    return tuple(int(cast(PromptUpdateExtra, state.extra).get("prompt_update_version", 0)) for state in states)
-
-
-class PromptUpdateMixin:
-    """Mixin for chunked streaming pipelines that support midway prompt updates.
-
-    Implements the behavior expected by :class:`~vllm_omni.diffusion.models.interface.SupportsPromptUpdate`.
-    All per-request transition state lives on ``StepRequestState.extra``; this mixin itself is stateless.
-    """
-
-    supports_prompt_update: ClassVar[bool] = True
-
-    @staticmethod
-    def _prompt_update_lock(extra: PromptUpdateExtra) -> Lock:
-        return extra.setdefault("prompt_update_lock", Lock())
-
-    def prepare_prompt_update(
-        self,
-        state: StepRequestState,
-        prompt: str,
-        event_id: str,
-        transition_chunks: int | None = None,
-    ) -> None:
-        """Encode and queue a prompt update to apply at the next chunk boundary.
-        It does not actually apply the prompt transition.
-
-        This is an extended interface for the pipeline, to be called by the runner.
-        """
-        if not prompt:
-            raise ValueError("prompt must be non-empty")
-        if not event_id:
-            raise ValueError("event_id must be non-empty")
-        if state.prompt_embeds is None:
-            raise ValueError(
-                f"prompt_update is not allowed before initial generation has started (request_id={state.request_id!r})"
-            )
-        duration = DEFAULT_TRANSITION_CHUNKS if transition_chunks is None else transition_chunks
-        if duration < 0:
-            raise ValueError("transition_chunks must be >= 0")
-
-        target_prompt_embeds, _ = self.encode_prompt(
-            prompt=prompt,
-            negative_prompt=None,
-            do_classifier_free_guidance=False,
-            num_videos_per_prompt=state.sampling.num_outputs_per_prompt,
-            max_sequence_length=state.sampling.max_sequence_length,
-            device=self.device,
-            dtype=self.transformer.dtype,
-        )
-        extra = cast(PromptUpdateExtra, state.extra)
-        with self._prompt_update_lock(extra):
-            extra["pending_prompt_update"] = {
-                "event_id": event_id,
-                "prompt": prompt,
-                "target_prompt_embeds": target_prompt_embeds,
-                "transition_chunks": duration,
-            }
-
-    def _apply_prompt_update_at_chunk_boundary(self, state: StepRequestState) -> None:
-        """Advance or start prompt interpolation before the next chunk."""
-        extra = cast(PromptUpdateExtra, state.extra)
-        update_state: _PromptUpdateState | None = extra.get("prompt_update_state")
-        with self._prompt_update_lock(extra):
-            pending = extra.pop("pending_prompt_update", None)
-        embeds_changed = False
-        next_chunk_index = state.chunk_index
-        started_event_ids: list[str] = []
-        active_event_ids: list[str] = []
-        completed_event_ids: list[str] = []
-
-        # If current transition is not complete, advance it.
-        # After completion, leave prompt_update_state in place (so a later pending
-        # update can abort/overwrite), but do not keep bumping the version.
-        if update_state is not None:
-            in_transition = (
-                update_state.transition_chunks > 0
-                and update_state.elapsed_transition_chunks < update_state.transition_chunks
-            )
-            if in_transition:
-                update_state.advance_transition()
-                state.prompt_embeds = update_state.blended_prompt_embeds()
-                embeds_changed = True
-                active_event_ids.append(update_state.event_id)
-                if update_state.elapsed_transition_chunks >= update_state.transition_chunks:
-                    state.prompt_embeds = update_state.target_prompt_embeds
-                    update_state.source_prompt_embeds = update_state.target_prompt_embeds
-                    completed_event_ids.append(update_state.event_id)
-                    if update_state.target_prompt is not None:
-                        logger.debug(
-                            "prompt_update transition complete request_id=%s next_chunk_index=%d prompt=%.20s...",
-                            state.request_id,
-                            next_chunk_index,
-                            update_state.target_prompt,
-                        )
-
-        # If a new prompt update is pending, start a new transition.
-        if pending is not None:
-            if state.prompt_embeds is None:
-                raise RuntimeError(
-                    "internal error: trying to apply a pending prompt update but "
-                    f"current prompt_embeds is None (request_id={state.request_id!r})"
-                )
-            source = state.prompt_embeds.detach().clone()
-            target = pending["target_prompt_embeds"]
-            duration = int(pending["transition_chunks"])
-            prompt = str(pending.get("prompt", ""))
-            event_id = pending.get("event_id")
-            update_state = _PromptUpdateState(
-                source_prompt_embeds=source,
-                target_prompt_embeds=target,
-                transition_chunks=duration,
-                event_id=event_id,
-                target_prompt=prompt,
-            )
-            extra["prompt_update_state"] = update_state
-            started_event_ids.append(event_id)
-            active_event_ids.append(event_id)
-            if duration <= 0:
-                # A hard/immediate transition
-                state.prompt_embeds = target
-                completed_event_ids.append(event_id)
-                logger.debug(
-                    "prompt_update sharp transition request_id=%s next_chunk_index=%d prompt=%.20s...",
-                    state.request_id,
-                    next_chunk_index,
-                    prompt,
-                )
-            else:
-                # A transition that really takes some time
-                update_state.advance_transition()
-                state.prompt_embeds = update_state.blended_prompt_embeds()
-                if update_state.elapsed_transition_chunks >= update_state.transition_chunks:
-                    state.prompt_embeds = update_state.target_prompt_embeds
-                    update_state.source_prompt_embeds = update_state.target_prompt_embeds
-                    completed_event_ids.append(event_id)
-                logger.debug(
-                    "prompt_update transition start request_id=%s next_chunk_index=%d prompt=%.20s...",
-                    state.request_id,
-                    next_chunk_index,
-                    prompt,
-                )
-            embeds_changed = True
-
-        # Indicate that prompt embeddings have changed---clear input batch cache.
-        if embeds_changed:
-            new_version = int(extra.get("prompt_update_version", 0)) + 1
-            extra["prompt_update_version"] = new_version
-
-        extra["prompt_update_chunk_metadata"] = {
-            "started_event_ids": started_event_ids,
-            "active_event_ids": active_event_ids,
-            "completed_event_ids": completed_event_ids,
-        }
-
-
-@dataclass
-class _PromptUpdateState:
-    """Per-request prompt interpolation stored in ``state.extra["prompt_update_state"]``."""
-
-    source_prompt_embeds: torch.Tensor
-    target_prompt_embeds: torch.Tensor
-    transition_chunks: int
-    event_id: str
-    elapsed_transition_chunks: int = 0
-    target_prompt: str | None = None
-
-    def blended_prompt_embeds(self) -> torch.Tensor:
-        alpha = self._current_alpha()
-        if alpha >= 1.0:
-            return self.target_prompt_embeds
-        return (1.0 - alpha) * self.source_prompt_embeds + alpha * self.target_prompt_embeds
-
-    def advance_transition(self) -> None:
-        if self.transition_chunks <= 0:
-            self.source_prompt_embeds = self.target_prompt_embeds
-            return
-        self.elapsed_transition_chunks += 1
-        if self.elapsed_transition_chunks >= self.transition_chunks:
-            self.source_prompt_embeds = self.target_prompt_embeds
-            self.elapsed_transition_chunks = self.transition_chunks
-
-    def _current_alpha(self) -> float:
-        if self.transition_chunks <= 0:
-            return 1.0
-        return min(1.0, self.elapsed_transition_chunks / self.transition_chunks)
-
-
-class _PendingPromptUpdate(TypedDict):
-    event_id: str
-    prompt: str
-    target_prompt_embeds: torch.Tensor
-    transition_chunks: int
-
-
-class _PromptUpdateChunkMetadata(TypedDict):
-    started_event_ids: list[str]
-    active_event_ids: list[str]
-    completed_event_ids: list[str]

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -42,12 +42,13 @@ from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
 )
 from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_rank
 from vllm_omni.diffusion.forward_context import set_forward_context
+from vllm_omni.diffusion.interaction.coordinator import InteractionCoordinator
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import (
-    SupportsPromptUpdate,
+    SupportsInteractionApply,
     adopt_request_scoped_cache_dit,
     is_request_scoped_cache_dit_enabled,
-    supports_prompt_update,
+    supports_interaction_apply,
     supports_step_execution,
 )
 from vllm_omni.diffusion.offloader import enable_offload_backend
@@ -176,6 +177,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         self.offload_backend: Any | None = None
         self.prompt_embed_cache: Any | None = None
         self.input_batch: InputBatch | None = None
+        self._interaction_coordinator: InteractionCoordinator | None = None
         self.model_memory_usage = 0
         self.diffusion_kv_backend = DiffusionKVModelRunnerBackend(
             vllm_config=vllm_config,
@@ -427,6 +429,10 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 model_tag=self.od_config.model_class_name,
             )
 
+        self._interaction_coordinator = InteractionCoordinator.build(self.pipeline, self.od_config)
+        if hasattr(self.pipeline, "_interaction_coordinator"):
+            self.pipeline._interaction_coordinator = self._interaction_coordinator
+
         logger.info("Model runner: Initialization complete.")
 
     def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
@@ -789,7 +795,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
         return self._runner_output_from_outputs(reqs, outputs)
 
-    def _attach_stepwise_metrics(
+    def _attach_stepwise_metadata(
         self,
         state: StepRequestState,
         output: DiffusionOutput,
@@ -800,6 +806,14 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         )
         attach_stage_durations(state, output)
 
+        # In streaming output mode with interaction, acknowledge which events are handled in this chunk.
+        meta = state.interaction_chunk_metadata
+        state.interaction_chunk_metadata = None
+        if meta is not None:
+            output.started_event_ids = list(meta.started_event_ids)
+            output.active_event_ids = list(meta.active_event_ids)
+            output.completed_event_ids = list(meta.completed_event_ids)
+
     def execute_model(
         self,
         req: OmniDiffusionRequest,
@@ -1275,10 +1289,16 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                                 clear_pipeline_stage_durations(self.pipeline)
                                 result = self.pipeline.post_decode(req)
                                 if result is not None:
-                                    self._attach_stepwise_metrics(
+                                    self._attach_stepwise_metadata(
                                         req,
                                         result,
                                     )
+                                    # After consuming this chunk's interaction metadata, apply pending interactions and
+                                    # prepare the next chunk (prepare_next_chunk may be a no-op---depending on pipeline)
+                                    if supports_interaction_apply(self.pipeline) and not req.request_denoise_completed:
+                                        pipe = cast(SupportsInteractionApply, self.pipeline)
+                                        pipe.apply_interaction_at_chunk_boundary(req)
+                                        pipe.prepare_next_chunk(req)
                             else:
                                 result = None
                             # finished should be computed after post_decode() advanced chunk_index
@@ -1354,42 +1374,48 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         request_id: str,
         interaction: OmniInteractionPrompt,
     ) -> None:
-        """Route a midway interaction to the matching active stepwise request feature."""
+        """Route a midway interaction through the pipeline interaction coordinator."""
         assert self.pipeline is not None, "Model not loaded. Call load_model() first."
         if not self.od_config.streaming_output:
             raise ValueError("submit_interaction requires streaming_output=True")
         if not self._supports_step_mode():
             raise ValueError("submit_interaction requires step execution support")
 
+        coordinator = self._interaction_coordinator
+        if coordinator is None:
+            coordinator = InteractionCoordinator.build(self.pipeline, self.od_config)
+            self._interaction_coordinator = coordinator
+            if hasattr(self.pipeline, "_interaction_coordinator"):
+                self.pipeline._interaction_coordinator = coordinator
+
         event = interaction.get("event")
-        if isinstance(event, dict) and "prompt" in event and "multi_modal_data" not in event:
-            # Is a prompt update interaction.
-            self._submit_prompt_update_interaction(request_id, interaction)
-            return
+        has_mm = isinstance(event, dict) and "multi_modal_data" in event
+        has_prompt = isinstance(event, dict) and "prompt" in event and event.get("prompt") is not None
 
-        raise NotImplementedError(
-            "Only text-only prompt update interactions with 'event.prompt' and optional "
-            "'transition_chunks' are supported in this release"
-        )
-
-    def _submit_prompt_update_interaction(
-        self,
-        request_id: str,
-        interaction: OmniInteractionPrompt,
-    ) -> None:
-        """Queue a prompt-update interaction for an active stepwise request."""
-        if not supports_prompt_update(self.pipeline):
+        # Prompt-only interactions in this release; multi_modal_data lands with camera support.
+        if not isinstance(event, dict) or has_mm or not has_prompt:
+            raise NotImplementedError(
+                "Only text-only prompt update interactions with 'event.prompt' and optional "
+                "'transition_chunks' are supported in this release"
+            )
+        if not coordinator.has_modality("prompt"):
             raise ValueError(f"prompt_update is not supported by pipeline {self.od_config.model_class_name!r}")
 
         state = self.state_cache.get(request_id)
         if state is None:
-            raise ValueError(f"No active request state for prompt_update: {request_id!r}")
+            raise ValueError(f"No active request state for interaction: {request_id!r}")
 
-        event = cast(dict[str, Any], interaction.get("event"))
-        prompt = event["prompt"]
-        transition_chunks = interaction.get("transition_chunks")
         event_id = interaction.get("event_id")
         if not isinstance(event_id, str) or not event_id:
             raise ValueError("event_id must be non-empty")
-        pipeline = cast(SupportsPromptUpdate, self.pipeline)
-        pipeline.prepare_prompt_update(state, prompt, event_id, transition_chunks)
+        prompt = event["prompt"]
+        if not isinstance(prompt, str) or not prompt:
+            raise ValueError("prompt must be non-empty")
+        coordinator.enqueue(
+            state,
+            modality="prompt",
+            event_id=event_id,
+            received_at=time.monotonic(),
+            payload={"prompt": prompt},
+            transition_chunks=interaction.get("transition_chunks"),
+        )

vllm_omni/diffusion/worker/input_batch.py

@@ -17,7 +17,7 @@ from dataclasses import dataclass, field
 import numpy as np
 import torch
 
-from vllm_omni.diffusion.prompt_update import prompt_update_versions
+from vllm_omni.diffusion.interaction.modality_handlers.prompt import prompt_update_versions
 from vllm_omni.diffusion.worker.utils import StepRequestState
 
 

vllm_omni/diffusion/worker/utils.py

@@ -12,6 +12,10 @@ import torch
 
 if TYPE_CHECKING:
     from vllm_omni.diffusion.data import DiffusionOutput
+    from vllm_omni.diffusion.interaction.types import (
+        InteractionChunkMetadata,
+        InteractionSession,
+    )
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
 
 
@@ -100,6 +104,10 @@ class StepRequestState:
     total_chunks: int = 1
     chunk_num_steps: int | None = None
 
+    # ── Optional interaction information in streaming output mode ──
+    interaction_sessions: dict[str, InteractionSession] = field(default_factory=dict)  # Modality -> transition progress
+    interaction_chunk_metadata: InteractionChunkMetadata | None = None  # Acknowledging completion of each interaction
+
     # ── Per-request scheduler instance (set once by prepare_encode) ──
     scheduler: Any | None = None
 

vllm_omni/inputs/data.py

@@ -46,11 +46,14 @@ class OmniInteractionEvent(OmniTextPrompt, total=False):
     pass
 
 
-class OmniInteractionPrompt(TypedDict, total=False):
-    """Mid-stream interaction payload."""
+class OmniInteractionPrompt(TypedDict):
+    """Normalized mid-stream interaction payload passed to ``submit_interaction``.
 
-    event_id: NotRequired[str]
-    event: NotRequired[OmniInteractionEvent]
+    WebSocket clients payload may omit ``event_id``; the API layer assigns a fallback ID in this case.
+    """
+
+    event_id: str
+    event: OmniInteractionEvent
     transition_chunks: NotRequired[int]