[Refactor][CI] Reorganize test tree, and move slow e2e to weekly (#4951)

Signed-off-by: wangyu <410167048@qq.com>

tests/e2e/features/comfyui/conftest.py

@@ -0,0 +1,89 @@
+"""
+Conftest for ComfyUI-vLLM-Omni tests.
+
+This module sets up the test environment by:
+1. Adding the ComfyUI plugin to Python path
+2. Mocking comfy_api.input module (AudioInput, VideoInput) since comfyui is not installed
+3. Mocking comfy_extras.nodes_audio module
+"""
+
+import os
+import sys
+from types import ModuleType, SimpleNamespace
+from typing import BinaryIO, TypedDict
+
+
+def pytest_configure(config):
+    """
+    Called after command line options have been parsed and before test collection.
+    This is the right place to set up sys.path and mock modules.
+    """
+    _setup_comfyui_test_environment()
+
+
+def _setup_comfyui_test_environment():
+    """Set up the test environment for ComfyUI plugin testing."""
+    # === Add ComfyUI plugin path to allow importing comfyui_vllm_omni ===
+    _COMFYUI_PLUGIN_PATH = os.path.abspath(
+        os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "apps", "ComfyUI-vLLM-Omni")
+    )
+    if not os.path.isdir(_COMFYUI_PLUGIN_PATH):
+        raise FileNotFoundError(
+            f"ComfyUI plugin not found at {_COMFYUI_PLUGIN_PATH}. "
+            "If it is moved elsewhere, please update the path in this conftest.py."
+        )
+    if _COMFYUI_PLUGIN_PATH not in sys.path:
+        sys.path.insert(0, _COMFYUI_PLUGIN_PATH)
+
+    # Import torch after changing import paths. (To be used later)
+    import torch
+
+    # === Mock ComfyUI internal modules (comfy_api & comfy_extras) and "import" them to sys.module ===
+    class AudioInput(TypedDict):
+        """Mock AudioInput TypedDict from comfy_api.input"""
+
+        waveform: torch.Tensor  # Shape: (B, C, T)
+        sample_rate: int
+
+    class VideoInput:
+        """Mock VideoInput class from comfy_api.input"""
+
+        def __init__(self, data: bytes = b"mock_video_data"):
+            self._data = data
+
+        def save_to(self, file: str | BinaryIO):
+            """Save video data to file or file-like object."""
+            if isinstance(file, str):
+                print("Called VideoInput.save_to with file path. Saving to a path is no-op in tests.")
+            else:
+                file.write(self._data)
+
+    mock_comfy_api = ModuleType("comfy_api")
+    mock_comfy_api_input = ModuleType("comfy_api.input")
+    mock_comfy_api_input.AudioInput = AudioInput
+    mock_comfy_api_input.VideoInput = VideoInput
+    mock_comfy_api.input = mock_comfy_api_input
+    mock_comfy_api_latest = ModuleType("comfy_api.latest")
+    mock_comfy_api_latest.Types = SimpleNamespace(VideoComponents=lambda **kwargs: kwargs)
+    mock_comfy_api_latest.InputImpl = SimpleNamespace(
+        VideoFromComponents=lambda _: VideoInput(b"mock_video_from_components")
+    )
+    mock_comfy_api.latest = mock_comfy_api_latest
+
+    def mock_load(_: str | BinaryIO):
+        """Mock nodes_audio.load that returns a waveform tensor (channels, samples) and sample rate."""
+        waveform = torch.zeros((1, 24000), dtype=torch.float32)
+        sample_rate = 24000
+        return waveform, sample_rate
+
+    mock_comfy_extras = ModuleType("comfy_extras")
+    mock_nodes_audio = ModuleType("comfy_extras.nodes_audio")
+    mock_nodes_audio.load = mock_load
+    mock_comfy_extras.nodes_audio = mock_nodes_audio
+
+    # Install mock modules BEFORE importing any comfyui_vllm_omni code
+    sys.modules["comfy_api"] = mock_comfy_api
+    sys.modules["comfy_api.input"] = mock_comfy_api_input
+    sys.modules["comfy_api.latest"] = mock_comfy_api_latest
+    sys.modules["comfy_extras"] = mock_comfy_extras
+    sys.modules["comfy_extras.nodes_audio"] = mock_nodes_audio

tests/e2e/features/comfyui/test_comfyui_integration.py

@@ -0,0 +1,788 @@
+"""
+Integration tests for ComfyUI nodes that use the Omni API client, with a mocked AsyncOmni and a real API server running in a background process.
+These tests cover the integration between ComfyUI node and the API server, without actual model inference logic.
+It ensures that
+1. Changes made to the API (e.g., request and response formats) do not break the ComfyUI frontend that use it.
+2. The sampling parameters are correctly passed from the node to AsyncOmni through the API layer.
+"""
+
+from __future__ import annotations
+
+import multiprocessing
+import time
+import traceback
+from collections.abc import Iterable, Sequence
+from enum import Enum
+from types import SimpleNamespace
+from typing import Any, NamedTuple
+
+import pytest
+import requests
+import torch
+from comfy_api.input import AudioInput, VideoInput
+from comfyui_vllm_omni.nodes import (
+    VLLMOmniGenerateImage,
+    VLLMOmniGenerateVideo,
+    VLLMOmniTTS,
+    VLLMOmniUnderstanding,
+    VLLMOmniVoiceClone,
+)
+from comfyui_vllm_omni.utils.types import AutoregressionSamplingParams, DiffusionSamplingParams, WanModelSpecificParams
+from PIL import Image
+from pytest_mock import MockerFixture
+from vllm import SamplingParams
+from vllm.outputs import CompletionOutput, RequestOutput
+
+from vllm_omni.entrypoints.async_omni import AsyncOmni as RealAsyncOmni
+from vllm_omni.entrypoints.cli.serve import OmniServeCommand
+from vllm_omni.inputs.data import OmniSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class ServerCase(NamedTuple):
+    """Parametrizing the model to serve."""
+
+    served_model: str
+    stage_list: list
+    stage_configs: list[Any]
+    outputs: list[OmniRequestOutput]
+
+
+class SamplingCase(NamedTuple):
+    """Parametrizing the input sampling parameters."""
+
+    kind: SamplingKind
+    sampling_params: dict | list[dict] | None
+    lora: dict | None = None
+
+
+class SamplingKind(str, Enum):
+    IMAGE_NONE = "image_none"
+    IMAGE_DIFFUSION_SINGLE = "image_diffusion_single"
+    UNDERSTANDING_NONE = "understanding_none"
+    UNDERSTANDING_AR_LIST = "understanding_ar_list"
+    TTS_NONE = "tts_none"
+    TTS_DIFFUSION_SINGLE = "tts_diffusion_single"
+    VIDEO_NONE = "video_none"
+    VIDEO_DIFFUSION_SINGLE = "video_diffusion_single"
+
+
+# Pre-defined arguments to be used in function calls during the tests
+IMAGE_WIDTH = 64
+IMAGE_HEIGHT = 64
+VIDEO_WIDTH = 32
+VIDEO_HEIGHT = 32
+VIDEO_FPS = 8
+VIDEO_NUM_FRAMES = 5
+DIFFUSION_SINGLE_SAMPLING_PARAMS = DiffusionSamplingParams(
+    {
+        "n": 2,
+        "num_inference_steps": 30,
+        "guidance_scale": 6.0,
+        "true_cfg_scale": 1.5,
+    }
+)
+
+DIFFUSION_VIDEO_SINGLE_SAMPLING_PARAMS = DiffusionSamplingParams(
+    {
+        "num_inference_steps": 30,
+        "guidance_scale": 6.0,
+        "true_cfg_scale": 1.5,
+    }
+)
+
+
+AR_LIST_SAMPLING_PARAMS = [
+    AutoregressionSamplingParams(
+        {
+            "max_tokens": 64,
+            "temperature": 0.6,
+            "top_p": 0.9,
+            "repetition_penalty": 1.0,
+            "seed": 21,
+        }
+    ),
+    AutoregressionSamplingParams(
+        {
+            "max_tokens": 96,
+            "temperature": 0.75,
+            "top_p": 0.85,
+            "repetition_penalty": 1.05,
+            "seed": 22,
+        }
+    ),
+    AutoregressionSamplingParams(
+        {
+            "max_tokens": 128,
+            "temperature": 0.8,
+            "top_p": 0.8,
+            "repetition_penalty": 1.1,
+            "seed": 23,
+        }
+    ),
+]
+
+VIDEO_MODEL_PARAMS = WanModelSpecificParams({"guidance_scale_2": 5.0, "boundary_ratio": 0.98, "flow_shift": 12.0})
+
+LORA_PARAMS = {"local_path": "test_lora_path", "name": "test_name", "scale": 0.7, "int_id": 10}
+
+
+def _build_image_output(size: tuple[int, int] = (IMAGE_WIDTH, IMAGE_HEIGHT), color: str = "red") -> Image.Image:
+    return Image.new("RGB", size, color=color)
+
+
+def _build_text_output(text: str = "This is a test response.") -> OmniRequestOutput:
+    completion_output = CompletionOutput(
+        index=0,
+        text=text,
+        token_ids=[1, 2, 3],
+        cumulative_logprob=0.0,
+        logprobs=None,
+        finish_reason="stop",
+        stop_reason=None,
+    )
+    request_output = RequestOutput(
+        request_id="test_req_text",
+        prompt="test prompt",
+        prompt_token_ids=[1, 2, 3],
+        prompt_logprobs=None,
+        outputs=[completion_output],
+        finished=True,
+        metrics=None,
+        lora_request=None,
+    )
+    return OmniRequestOutput(
+        request_id="test_req_text",
+        finished=True,
+        final_output_type="text",
+        request_output=request_output,
+    )
+
+
+def _build_audio_chat_output(num_samples: int = 24000) -> OmniRequestOutput:
+    completion_output = CompletionOutput(
+        index=0,
+        text="",
+        token_ids=[],
+        cumulative_logprob=0.0,
+        logprobs=None,
+        finish_reason="stop",
+        stop_reason=None,
+    )
+    completion_output.multimodal_output = {"audio": [torch.zeros(1, num_samples)]}
+    request_output = RequestOutput(
+        request_id="test_req_audio_chat",
+        prompt="test prompt",
+        prompt_token_ids=[1, 2, 3],
+        prompt_logprobs=None,
+        outputs=[completion_output],
+        finished=True,
+        metrics=None,
+        lora_request=None,
+    )
+    return OmniRequestOutput(
+        request_id="test_req_audio_chat",
+        finished=True,
+        final_output_type="audio",
+        request_output=request_output,
+    )
+
+
+def _build_audio_speech_output(num_samples: int = 24000) -> OmniRequestOutput:
+    return OmniRequestOutput.from_diffusion(
+        request_id="test_req_audio_speech",
+        images=[],
+        multimodal_output={"audio": torch.zeros(num_samples), "sr": 24000},
+        final_output_type="audio",
+    )
+
+
+def _build_diffusion_image_output_for_images_endpoint() -> OmniRequestOutput:
+    return OmniRequestOutput.from_diffusion(
+        request_id="test_req_img_dalle",
+        images=[_build_image_output()],
+        final_output_type="image",
+    )
+
+
+def _build_diffusion_video_output() -> OmniRequestOutput:
+    # Small video: VIDEO_NUM_FRAMES frames of (VIDEO_HEIGHT x VIDEO_WIDTH) RGB, shape (F, H, W, C)
+    video_frames = torch.zeros((VIDEO_NUM_FRAMES, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.float32)
+    return OmniRequestOutput.from_diffusion(
+        request_id="test_req_video",
+        images=[video_frames],
+        final_output_type="video",
+    )
+
+
+def _build_diffusion_image_output_for_chat_endpoint() -> OmniRequestOutput:
+    request_output = SimpleNamespace(
+        images=[_build_image_output(color="blue")],
+        finished=True,
+    )
+    return OmniRequestOutput(
+        request_id="test_req_img_chat",
+        finished=True,
+        final_output_type="image",
+        request_output=request_output,
+    )
+
+
+def _assert_sampling_param_values(
+    received: OmniSamplingParams, expected: dict[str, Any], expected_lora: dict | None = None
+):
+    for key, expected_value in expected.items():
+        actual_value = getattr(received, key)
+        assert actual_value == expected_value, (
+            f"Expected sampling param '{key}'={expected_value}, got {actual_value}. The received sampling params: {received}"
+        )
+    if expected_lora:
+        assert received.lora_request.lora_name == expected_lora["name"], (
+            f"Expected lora name={(expected_lora['name'])}, got {received.lora_request.lora_name}. The received sampling params: {received}"
+        )
+        assert received.lora_request.lora_int_id == expected_lora["int_id"], (
+            f"Expected lora int_id={expected_lora['int_id']}, got {received.lora_request.lora_int_id}. The received sampling params: {received}"
+        )
+        assert received.lora_request.lora_path == expected_lora["local_path"], (
+            f"Expected lora path={expected_lora['local_path']}, got {received.lora_request.lora_path}. The received sampling params: {received}"
+        )
+        assert received.lora_scale == expected_lora["scale"], (
+            f"Expected lora scale={expected_lora['scale']}, got {received.lora_scale}. The received sampling params: {received}"
+        )
+
+
+def _assert_model_param_values(received: OmniSamplingParams, expected: dict):
+    for key, expected_value in expected.items():
+        try:
+            actual_value = getattr(received, key)
+            expected_param_name = key
+        except AttributeError:
+            actual_value = received.extra_args.get(key, None)
+            expected_param_name = f'extra_args["{key}"]'
+        assert actual_value == expected_value, (
+            f"Expected model param '{expected_param_name}'={expected_value}, got {actual_value}. The received sampling params: {received}"
+        )
+
+
+def _make_stage_config(
+    stage_type: str,
+    *,
+    is_comprehension: bool = False,
+    model_stage: str | None = None,
+):
+    engine_args = SimpleNamespace()
+    if model_stage is not None:
+        engine_args.model_stage = model_stage
+    return SimpleNamespace(
+        stage_type=stage_type,
+        is_comprehension=is_comprehension,
+        engine_args=engine_args,
+    )
+
+
+def _stage_type(stage: Any) -> str | None:
+    return getattr(stage, "stage_type", None)
+
+
+def _build_output_modalities(stage_configs: list[Any]) -> list[str]:
+    modalities: list[str] = []
+    for stage in stage_configs:
+        final_output = getattr(stage, "final_output", False)
+        final_output_type = getattr(stage, "final_output_type", None)
+        if final_output and isinstance(final_output_type, str):
+            modalities.append(final_output_type)
+    if modalities:
+        return modalities
+    if any(_stage_type(stage) == "diffusion" for stage in stage_configs):
+        return ["image"]
+    return ["text", "audio"]
+
+
+def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: SamplingCase, server_case: ServerCase):
+    async def _mock_generate(*args, **kwargs):
+        received_sampling_params_list: Sequence[OmniSamplingParams] | None = (
+            args[2] if len(args) > 2 else kwargs.get("sampling_params_list")
+        )
+
+        assert received_sampling_params_list is not None, (
+            "In the current codebase, the API layer always provides not-None sampling parameter list when calling `AsyncOmni.generate`"
+            "This test also uses this assumption for now."
+            "If this assertion fails, it means the API layer has changed and this test needs to be updated accordingly."
+            "It does not necessarily mean there is a bug, because `AsyncOmni.generate` does allow sampling_params_list to be None."
+        )
+        assert isinstance(received_sampling_params_list, Sequence), "sampling_params_list should be a Sequence"
+
+        if sampling_case.kind is SamplingKind.IMAGE_NONE:
+            assert len(received_sampling_params_list) == 1
+            _assert_sampling_param_values(
+                received_sampling_params_list[0],
+                {
+                    "width": IMAGE_WIDTH,
+                    "height": IMAGE_HEIGHT,
+                },
+            )
+        elif sampling_case.kind is SamplingKind.IMAGE_DIFFUSION_SINGLE:
+            assert len(received_sampling_params_list) == 1
+            expected = DIFFUSION_SINGLE_SAMPLING_PARAMS.copy()
+            expected["num_outputs_per_prompt"] = expected.pop("n")  # convert from n to num_outputs_per_prompt
+            _assert_sampling_param_values(
+                received_sampling_params_list[0],
+                {
+                    "width": IMAGE_WIDTH,
+                    "height": IMAGE_HEIGHT,
+                    **expected,
+                },
+                LORA_PARAMS,
+            )
+        elif sampling_case.kind is SamplingKind.UNDERSTANDING_NONE:
+            assert len(received_sampling_params_list) == 3
+        elif sampling_case.kind is SamplingKind.UNDERSTANDING_AR_LIST:
+            assert len(received_sampling_params_list) == 3
+            for i, expected in enumerate(AR_LIST_SAMPLING_PARAMS):
+                _assert_sampling_param_values(received_sampling_params_list[i], expected)
+        elif sampling_case.kind in {SamplingKind.TTS_NONE, SamplingKind.TTS_DIFFUSION_SINGLE}:
+            assert len(received_sampling_params_list) == 1
+        elif sampling_case.kind is SamplingKind.VIDEO_NONE:
+            assert len(received_sampling_params_list) == 1
+            _assert_sampling_param_values(
+                received_sampling_params_list[0],
+                {
+                    "width": VIDEO_WIDTH,
+                    "height": VIDEO_HEIGHT,
+                    "num_frames": VIDEO_NUM_FRAMES,
+                    "fps": VIDEO_FPS,
+                },
+            )
+        elif sampling_case.kind is SamplingKind.VIDEO_DIFFUSION_SINGLE:
+            assert len(received_sampling_params_list) == 1
+            expected = DIFFUSION_VIDEO_SINGLE_SAMPLING_PARAMS.copy()
+            # expected["num_outputs_per_prompt"] = expected.pop("n")  # convert from n to num_outputs_per_prompt
+            _assert_sampling_param_values(
+                received_sampling_params_list[0],
+                {
+                    "width": VIDEO_WIDTH,
+                    "height": VIDEO_HEIGHT,
+                    "num_frames": VIDEO_NUM_FRAMES,
+                    "fps": VIDEO_FPS,
+                    **expected,
+                },
+                LORA_PARAMS,
+            )
+            _assert_model_param_values(received_sampling_params_list[0], VIDEO_MODEL_PARAMS)
+        else:
+            raise AssertionError(f"Unknown sampling case: {sampling_case.kind}")
+
+        for output in outputs:
+            yield output
+
+    return _mock_generate
+
+
+@pytest.fixture
+def server_case(request) -> ServerCase:
+    return request.param
+
+
+@pytest.fixture
+def sampling_case(request) -> SamplingCase:
+    return request.param
+
+
+@pytest.fixture
+def mock_async_omni(
+    server_case: ServerCase,
+    sampling_case: SamplingCase,
+    monkeypatch: pytest.MonkeyPatch,
+    mocker: MockerFixture,
+):
+    async def _mock_preprocess_chat(self, *args, **kwargs):
+        return ([{"role": "user", "content": "test"}], [{"prompt": "test prompt"}])
+
+    # Need to mock AsyncOmni itself (not only its generate method) because
+    # 1. The API layer uses its stage_list and stage_configs attributes
+    # 2. Its __init__ method has slow side effects (model & config loading).
+    mock_async_omni_cls = mocker.patch("vllm_omni.entrypoints.openai.api_server.AsyncOmni")
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.openai.serving_chat.OmniOpenAIServingChat._preprocess_chat",
+        _mock_preprocess_chat,
+    )
+
+    mock_instance = mocker.AsyncMock(spec=RealAsyncOmni)
+    mock_instance.generate = _build_mock_outputs(server_case.outputs, sampling_case, server_case)
+
+    mock_instance.stage_list = server_case.stage_list
+    mock_instance.stage_configs = server_case.stage_configs
+    mock_instance.output_modalities = _build_output_modalities(server_case.stage_configs)
+    mock_instance.default_sampling_params_list = [
+        SamplingParams() if _stage_type(stage) != "diffusion" else mocker.MagicMock()
+        for stage in server_case.stage_configs
+    ]
+    mock_instance.errored = False
+    mock_instance.dead_error = RuntimeError("Mock engine error")
+    mock_instance.model_config = mocker.MagicMock(
+        max_model_len=4096,
+        io_processor_plugin=None,
+        allowed_local_media_path=None,
+        allowed_media_domains=None,
+    )
+    # Mimic Qwen3-TTS talker speaker config so CustomVoice validation passes.
+    mock_instance.model_config.hf_config = mocker.MagicMock()
+    mock_instance.model_config.hf_config.talker_config = mocker.MagicMock()
+    mock_instance.model_config.hf_config.talker_config.speaker_id = {"Vivian": 0}
+    mock_instance.io_processor = mocker.MagicMock()
+    mock_instance.input_processor = mocker.MagicMock()
+    mock_instance.shutdown = mocker.MagicMock()
+    mock_instance.get_vllm_config = mocker.AsyncMock(return_value=None)
+    mock_instance.get_supported_tasks = mocker.AsyncMock(return_value=["generate"])
+    mock_instance.get_tokenizer = mocker.AsyncMock(return_value=None)
+
+    mock_async_omni_cls.return_value = mock_instance
+    yield mock_async_omni_cls
+
+
+@pytest.fixture
+def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni):
+    """Set up a API server in background process from command line with parametrized model name and mocked AsyncOmni."""
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="command")
+    cmd = OmniServeCommand()
+    cmd.subparser_init(subparsers)
+
+    port = unused_tcp_port_factory()
+    args = parser.parse_args(["serve", server_case.served_model, "--omni", "--port", str(port)])
+
+    def run_server():
+        try:
+            cmd.cmd(args)
+        except Exception:
+            traceback.print_exc()
+
+    server_process = multiprocessing.Process(target=run_server)
+    server_process.start()
+
+    # Wait for the server to be ready by polling the health endpoint.
+    wait_time = 30
+    wait_poll_interval = 1
+    for _ in range(wait_time // wait_poll_interval):
+        try:
+            response = requests.get(f"http://127.0.0.1:{port}/health")
+            if response.status_code == 200:
+                break
+        except requests.ConnectionError:
+            time.sleep(wait_poll_interval)
+    else:
+        if server_process.is_alive():
+            server_process.terminate()
+            server_process.join(timeout=5)
+            if server_process.is_alive():
+                server_process.kill()
+                server_process.join(timeout=5)
+        pytest.fail(f"API server failed to start within {wait_time} seconds")
+
+    yield f"http://127.0.0.1:{port}/v1"
+
+    if server_process.is_alive():
+        server_process.terminate()
+    server_process.join(timeout=10)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "server_case,model,image_input",
+    [
+        pytest.param(
+            ServerCase(
+                served_model="Tongyi-MAI/Z-Image-Turbo",
+                stage_list=["diffusion"],
+                stage_configs=[_make_stage_config("diffusion")],
+                outputs=[_build_diffusion_image_output_for_images_endpoint()],
+            ),
+            "Tongyi-MAI/Z-Image-Turbo",
+            False,
+            id="text-to-image-dalle-endpoint",
+        ),
+        pytest.param(
+            ServerCase(
+                served_model="ByteDance-Seed/BAGEL-7B-MoT",
+                stage_list=["diffusion"],
+                stage_configs=[_make_stage_config("diffusion")],
+                outputs=[_build_diffusion_image_output_for_chat_endpoint()],
+            ),
+            "ByteDance-Seed/BAGEL-7B-MoT",
+            False,
+            id="text-to-image-bagel-chat-endpoint",
+        ),
+        pytest.param(
+            ServerCase(
+                served_model="Qwen/Qwen-Image-Edit",
+                stage_list=["diffusion"],
+                stage_configs=[_make_stage_config("diffusion")],
+                outputs=[_build_diffusion_image_output_for_images_endpoint()],
+            ),
+            "Qwen/Qwen-Image-Edit",
+            True,
+            id="image-to-image-dalle-endpoint",
+        ),
+        pytest.param(
+            ServerCase(
+                served_model="ByteDance-Seed/BAGEL-7B-MoT",
+                stage_list=["diffusion"],
+                stage_configs=[_make_stage_config("diffusion")],
+                outputs=[_build_diffusion_image_output_for_chat_endpoint()],
+            ),
+            "ByteDance-Seed/BAGEL-7B-MoT",
+            True,
+            id="image-to-image-bagel-chat-endpoint",
+        ),
+    ],
+    indirect=["server_case"],
+)
+@pytest.mark.parametrize(
+    "sampling_case",
+    [
+        pytest.param(
+            SamplingCase(kind=SamplingKind.IMAGE_NONE, sampling_params=None, lora=None), id="no-sampling-params"
+        ),
+        pytest.param(
+            SamplingCase(
+                kind=SamplingKind.IMAGE_DIFFUSION_SINGLE,
+                sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS,
+                lora=LORA_PARAMS,
+            ),
+            id="single-diffusion-sampling-params",
+        ),
+    ],
+    indirect=["sampling_case"],
+)
+async def test_image_generation_node(api_server: str, model: str, image_input: bool, sampling_case: SamplingCase):
+    node = VLLMOmniGenerateImage()
+
+    kwargs = {
+        "url": api_server,
+        "model": model,
+        "prompt": "A beautiful sunset",
+        "width": IMAGE_WIDTH,
+        "height": IMAGE_HEIGHT,
+    }
+    if image_input:
+        kwargs["image"] = torch.zeros((1, IMAGE_WIDTH, IMAGE_HEIGHT, 3), dtype=torch.float32)
+    if sampling_case.sampling_params is not None:
+        kwargs["sampling_params"] = sampling_case.sampling_params
+    if sampling_case.lora:
+        kwargs["lora"] = sampling_case.lora
+
+    result = await node.generate(**kwargs)
+
+    assert isinstance(result, tuple)
+    assert len(result) == 1
+    assert isinstance(result[0], torch.Tensor)
+    assert result[0].shape == (1, IMAGE_WIDTH, IMAGE_HEIGHT, 3)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "server_case",
+    [
+        pytest.param(
+            ServerCase(
+                served_model="Qwen/Qwen2.5-Omni-7B",
+                stage_list=[
+                    SimpleNamespace(is_comprehension=True, model_stage="llm"),
+                    SimpleNamespace(is_comprehension=False, model_stage="llm"),
+                    SimpleNamespace(is_comprehension=False, model_stage="llm"),
+                ],
+                stage_configs=[
+                    _make_stage_config("llm", is_comprehension=True, model_stage="thinker"),
+                    _make_stage_config("llm", is_comprehension=False, model_stage="talker"),
+                    _make_stage_config("llm", is_comprehension=False, model_stage="code2wav"),
+                ],
+                outputs=[_build_audio_chat_output(), _build_text_output("Understanding response")],
+            ),
+            id="multimodal-understanding",
+        )
+    ],
+    indirect=["server_case"],
+)
+@pytest.mark.parametrize(
+    "sampling_case",
+    [
+        pytest.param(SamplingCase(kind=SamplingKind.UNDERSTANDING_NONE, sampling_params=None), id="no-sampling-params"),
+        pytest.param(
+            SamplingCase(kind=SamplingKind.UNDERSTANDING_AR_LIST, sampling_params=AR_LIST_SAMPLING_PARAMS),
+            id="ar-sampling-params-list",
+        ),
+    ],
+    indirect=["sampling_case"],
+)
+async def test_understanding_node(api_server: str, sampling_case: SamplingCase):
+    node = VLLMOmniUnderstanding()
+
+    image = torch.zeros((1, IMAGE_WIDTH, IMAGE_HEIGHT, 3), dtype=torch.float32)
+    video = VideoInput(b"mock_video_for_test")  # type: ignore[reportAbstractUsage]
+    audio: AudioInput = {"waveform": torch.zeros((1, 1, 24000), dtype=torch.float32), "sample_rate": 24000}
+
+    text_response, audio_response = await node.generate(
+        url=api_server,
+        model="Qwen/Qwen2.5-Omni-7B",
+        prompt="Describe all modalities.",
+        image=image,
+        audio=audio,
+        video=video,
+        sampling_params=sampling_case.sampling_params,
+        output_text=True,
+        output_audio=True,
+        use_audio_in_video=True,
+    )
+
+    assert text_response == "Understanding response"
+    assert isinstance(audio_response, dict)
+    assert audio_response["sample_rate"] == 24000
+    assert isinstance(audio_response["waveform"], torch.Tensor)
+    assert audio_response["waveform"].shape == (1, 1, 24000)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "server_case,node_cls,call_kwargs",
+    [
+        pytest.param(
+            ServerCase(
+                served_model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+                stage_list=["llm"],
+                stage_configs=[_make_stage_config("llm", model_stage="qwen3_tts")],
+                outputs=[_build_audio_speech_output()],
+            ),
+            VLLMOmniTTS,
+            {
+                "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+                "input": "Hello from TTS test",
+                "voice": "Vivian",
+                "response_format": "wav",
+                "speed": 1.0,
+                "model_specific_params": None,
+            },
+            id="tts",
+        ),
+        pytest.param(
+            ServerCase(
+                served_model="Qwen/Qwen3-TTS-12Hz-1.7B-Base",
+                stage_list=["llm"],
+                stage_configs=[_make_stage_config("llm", model_stage="qwen3_tts")],
+                outputs=[_build_audio_speech_output()],
+            ),
+            VLLMOmniVoiceClone,
+            {
+                "model": "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
+                "input": "Hello from voice clone test",
+                "voice": "Vivian",
+                "response_format": "wav",
+                "speed": 1.0,
+                "ref_audio": {"waveform": torch.zeros((1, 1, 24000), dtype=torch.float32), "sample_rate": 24000},
+                "ref_text": "Reference transcript",
+                "x_vector_only_mode": False,
+                "model_specific_params": None,
+            },
+            id="tts-voice-clone",
+        ),
+    ],
+    indirect=["server_case"],
+)
+@pytest.mark.parametrize(
+    "sampling_case",
+    [
+        pytest.param(SamplingCase(kind=SamplingKind.TTS_NONE, sampling_params=None), id="no-sampling-params"),
+        pytest.param(
+            SamplingCase(kind=SamplingKind.TTS_DIFFUSION_SINGLE, sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS),
+            id="single-diffusion-sampling-params",
+        ),
+    ],
+    indirect=["sampling_case"],
+)
+async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_case: SamplingCase):
+    node = node_cls()
+    actual_kwargs = dict(call_kwargs)
+    if sampling_case.sampling_params is not None:
+        actual_kwargs["model_specific_params"] = sampling_case.sampling_params
+    result = await node.generate(url=api_server, **actual_kwargs)
+
+    assert isinstance(result, tuple)
+    assert len(result) == 1
+    assert isinstance(result[0], dict)
+    assert result[0]["sample_rate"] == 24000
+    assert isinstance(result[0]["waveform"], torch.Tensor)
+    assert result[0]["waveform"].shape == (1, 1, 24000)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "server_case,model,image_input",
+    [
+        pytest.param(
+            ServerCase(
+                served_model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+                stage_list=["diffusion"],
+                stage_configs=[{"stage_type": "diffusion"}],
+                outputs=[_build_diffusion_video_output()],
+            ),
+            "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+            False,
+            id="text-to-video",
+        ),
+        pytest.param(
+            ServerCase(
+                served_model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+                stage_list=["diffusion"],
+                stage_configs=[{"stage_type": "diffusion"}],
+                outputs=[_build_diffusion_video_output()],
+            ),
+            "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+            True,
+            id="image-to-video",
+        ),
+    ],
+    indirect=["server_case"],
+)
+@pytest.mark.parametrize(
+    "sampling_case",
+    [
+        pytest.param(SamplingCase(kind=SamplingKind.VIDEO_NONE, sampling_params=None), id="no-sampling-params"),
+        pytest.param(
+            SamplingCase(
+                kind=SamplingKind.VIDEO_DIFFUSION_SINGLE,
+                sampling_params=DIFFUSION_VIDEO_SINGLE_SAMPLING_PARAMS,
+                lora=LORA_PARAMS,
+            ),
+            id="single-diffusion-sampling-params",
+        ),
+    ],
+    indirect=["sampling_case"],
+)
+async def test_video_generation_node(api_server: str, model: str, image_input: bool, sampling_case: SamplingCase):
+    node = VLLMOmniGenerateVideo()
+
+    kwargs = {
+        "url": api_server,
+        "model": model,
+        "prompt": "A beautiful sunset timelapse",
+        "negative_prompt": "",
+        "width": VIDEO_WIDTH,
+        "height": VIDEO_HEIGHT,
+        "fps": VIDEO_FPS,
+        "num_frames": VIDEO_NUM_FRAMES,
+        "model_params": VIDEO_MODEL_PARAMS,
+    }
+    if image_input:
+        kwargs["image"] = torch.zeros((1, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.float32)
+    if sampling_case.sampling_params is not None:
+        kwargs["sampling_params"] = sampling_case.sampling_params
+    if sampling_case.lora:
+        kwargs["lora"] = sampling_case.lora
+
+    result = await node.generate(**kwargs)
+
+    assert isinstance(result, tuple)
+    assert len(result) == 1
+    assert isinstance(result[0], VideoInput)

tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py

@@ -0,0 +1,248 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Regression tests for AsyncOmni collective_rpc in inline diffusion mode.
+
+When AsyncOmni runs a single diffusion stage it activates "inline diffusion
+mode", which skips stage worker subprocess creation and therefore never
+attaches IPC queues (_in_q / _out_q) to the OmniStage.  Methods like
+list_loras(), add_lora(), sleep(), wake_up() all delegate to
+collective_rpc(), which must handle this mode correctly instead of
+trying to use the non-existent queues.
+
+This is the same code path that verl's vLLMOmniHttpServer.generate()
+exercises when it calls ``await self.engine.list_loras()`` before
+dispatching a generation request.
+
+Usage:
+    pytest tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py -v -s
+"""
+
+from __future__ import annotations
+
+import asyncio
+import uuid
+from contextlib import ExitStack
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+MODEL = "tiny-random/Qwen-Image"
+CUSTOM_PIPELINE_CLASS = "tests.e2e.features.helpers.custom_pipeline.QwenImagePipelineWithLogProbForTest"
+WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest"
+
+
+def _create_inline_engine() -> AsyncOmni:
+    """Create an AsyncOmni instance that uses inline diffusion mode.
+
+    A single diffusion stage triggers inline mode automatically.
+    """
+    engine = AsyncOmni(
+        model=MODEL,
+        custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+        worker_extension_cls=WORKER_EXTENSION_CLASS,
+        enforce_eager=True,
+    )
+
+    return engine
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_list_loras_inline_mode():
+    """list_loras() must not crash in inline diffusion mode.
+
+    This is the exact call that vLLMOmniHttpServer.generate() makes
+    before every generation request.
+    """
+    with ExitStack() as after:
+        engine = _create_inline_engine()
+        after.callback(engine.shutdown)
+
+        result = await engine.list_loras()
+        assert isinstance(result, list), f"Expected list, got {type(result)}"
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_collective_rpc_inline_mode():
+    """collective_rpc() must delegate to the inline engine, not stage queues."""
+    with ExitStack() as after:
+        engine = _create_inline_engine()
+        after.callback(engine.shutdown)
+
+        result = await engine.collective_rpc(method="list_loras")
+        assert isinstance(result, list), f"Expected list, got {type(result)}"
+        assert len(result) == 1, "Inline mode has exactly one stage"
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_sleep_wake_up_inline_mode():
+    """sleep() and wake_up() must work in inline diffusion mode."""
+    with ExitStack() as after:
+        engine = _create_inline_engine()
+        after.callback(engine.shutdown)
+
+        await engine.sleep(level=1)
+        assert await engine.is_sleeping()
+
+        await engine.wake_up()
+        assert not await engine.is_sleeping()
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_generate_after_list_loras_inline_mode():
+    """Full flow: list_loras() then generate(), matching vLLMOmniHttpServer.
+
+    This reproduces the exact sequence that caused the original crash:
+    1. list_loras() (was crashing with AssertionError on _out_q)
+    2. generate() (should succeed)
+    """
+    with ExitStack() as after:
+        engine = _create_inline_engine()
+        after.callback(engine.shutdown)
+
+        # Step 1: list_loras (the call that was crashing)
+        loras = await engine.list_loras()
+        assert isinstance(loras, list)
+
+        # Step 2: generate (should still work after list_loras)
+        sampling_params = OmniDiffusionSamplingParams(
+            num_inference_steps=2,
+            guidance_scale=0.0,
+            height=256,
+            width=256,
+            seed=42,
+        )
+
+        last_output = None
+        async for output in engine.generate(
+            prompt={"prompt_ids": list(range(50))},
+            request_id=f"test_after_lora_{uuid.uuid4().hex[:8]}",
+            sampling_params_list=[sampling_params],
+            output_modalities=["image"],
+        ):
+            last_output = output
+
+        assert last_output is not None
+        assert isinstance(last_output, OmniRequestOutput)
+        assert last_output.images, "Expected at least one generated image"
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_sleep_memory_reclaimed_custom_pipeline():
+    """sleep(level=1) must physically reclaim CuMemAllocator-tracked memory for
+    custom_pipeline.
+
+    Regression test for: custom pipelines constructed under ``with target_device:``
+    (CUDA default-device context) caused safetensors >=0.20.0 to use a
+    direct-to-GPU fast path (cudaMalloc via the driver API) that bypasses
+    CuMemAllocator, leaving weights invisible to sleep() and pinned in GPU
+    memory after the call.
+
+    The fix moves custom_pipeline init outside the CUDA context so all weights
+    go through the caching allocator and are therefore fully reclaimed by
+    sleep(level=1).  A non-zero ``CuMemAllocator.get_current_usage()`` after
+    sleep is the direct signal that the bypass is still occurring.
+    """
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+            enable_sleep_mode=True,
+        )
+        after.callback(engine.shutdown)
+
+        assert not await engine.is_sleeping(), "Engine should be awake after creation"
+
+        # Measure global VRAM before sleep (driver view; includes inline worker
+        # thread since inline mode runs in the same process).
+        torch.accelerator.synchronize()
+        free_before, total = torch.cuda.mem_get_info()
+        used_before_gib = (total - free_before) / 1024**3
+
+        # Measure CuMemAllocator-tracked usage before sleep.  In inline mode
+        # the worker runs in a thread pool inside this process, so the allocator
+        # singleton is shared and can be read directly.
+        allocator = None
+        tracked_before = 0
+        try:
+            from vllm.device_allocator.cumem import CuMemAllocator
+
+            allocator = CuMemAllocator.get_instance()
+            tracked_before = allocator.get_current_usage()
+        except Exception:
+            pass
+
+        # Put the engine to sleep; all weights should be offloaded via the pool.
+        acks = await engine.sleep(level=1)
+        await asyncio.sleep(0.5)  # allow the CUDA driver to settle
+        torch.accelerator.synchronize()
+
+        # Measure after sleep.
+        free_after, _ = torch.cuda.mem_get_info()
+        used_after_gib = (total - free_after) / 1024**3
+        drop_gib = used_before_gib - used_after_gib
+
+        # --- Primary assertion: allocator reports zero tracked memory. ---
+        # If this fails it means weights were allocated outside the CuMem pool
+        # (safetensors direct-to-GPU bypass) — the exact regression this test
+        # is designed to catch.
+        if allocator is not None:
+            tracked_after = allocator.get_current_usage()
+            assert tracked_after == 0, (
+                f"CuMemAllocator still tracks {tracked_after / 1024**3:.3f} GiB "
+                f"after sleep(level=1) on custom_pipeline path "
+                f"(was {tracked_before / 1024**3:.3f} GiB before sleep). "
+                "Weights were allocated outside the CuMem pool via the "
+                "safetensors direct-to-GPU fast path — loader-context fix "
+                "may not be applied."
+            )
+
+        # --- Secondary assertion: physical VRAM or ACK freed_bytes confirms
+        # reclamation at the driver level.
+        total_freed_bytes = sum(
+            (ack.freed_bytes if hasattr(ack, "freed_bytes") else ack.get("freed_bytes", 0))
+            for ack in acks
+            if ack is not None
+        )
+        freed_gib = total_freed_bytes / 1024**3
+        assert freed_gib > 0 or drop_gib > 0, (
+            f"Expected GPU memory to be reclaimed after sleep(level=1) on "
+            f"custom_pipeline + enable_sleep_mode=True. "
+            f"CuMemAllocator tracked before={tracked_before / 1024**3:.3f} GiB, "
+            f"ACK freed={freed_gib:.3f} GiB, global VRAM drop={drop_gib:.3f} GiB."
+        )
+
+        # Engine must report it is sleeping.
+        assert await engine.is_sleeping()
+
+        # Wake up and confirm the engine is functional again.
+        await engine.wake_up()
+        assert not await engine.is_sleeping()

tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py

@@ -0,0 +1,250 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""E2E tests for AsyncOmni Qwen-Image generation with trajectory_* fields.
+
+Validates that the custom Qwen-Image pipeline returns structured trajectory
+outputs (latents, timesteps, log_probs) via OmniRequestOutput's trajectory_*
+fields instead of the legacy custom_output dict.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import uuid
+from contextlib import ExitStack
+from functools import lru_cache
+
+import numpy as np
+import pytest
+from transformers import AutoTokenizer
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+CUSTOM_PIPELINE_CLASS = "tests.e2e.features.helpers.custom_pipeline.QwenImagePipelineWithLogProbForTest"
+WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest"
+
+# Use your specified HF repo name/path directly here
+MODEL = "tiny-random/Qwen-Image"
+
+TOKENIZER_MODEL = "Qwen/Qwen2-1.5B-Instruct"
+
+
+# ---------------------------------------------------------------------
+#                👇 Test Helper Functions & Fixtures 👇
+# ---------------------------------------------------------------------
+
+_MIN_PROMPT_TOKENS = 35
+
+
+def normalize_token_ids(tokenized_output) -> list[int]:
+    """Normalize tokenizer outputs into a flat ``list[int]``.
+
+    This handles Transformers 4/5 differences where ``apply_chat_template(tokenize=True)``
+    may return either ``list[int]`` or a ``BatchEncoding``/mapping with ``input_ids``.
+    """
+
+    token_ids = tokenized_output
+    if isinstance(tokenized_output, dict):
+        if "input_ids" in tokenized_output:
+            token_ids = tokenized_output["input_ids"]
+    elif hasattr(tokenized_output, "input_ids"):
+        token_ids = tokenized_output.input_ids
+
+    if hasattr(token_ids, "tolist"):
+        token_ids = token_ids.tolist()
+
+    if isinstance(token_ids, tuple):
+        token_ids = list(token_ids)
+
+    if isinstance(token_ids, list) and len(token_ids) == 1 and isinstance(token_ids[0], list | tuple):
+        token_ids = list(token_ids[0])
+
+    if not isinstance(token_ids, list):
+        raise TypeError(f"token_ids must be list-like token ids, got {type(token_ids).__name__}: {token_ids!r}")
+
+    normalized_ids = []
+    for idx, token_id in enumerate(token_ids):
+        if hasattr(token_id, "item"):
+            token_id = token_id.item()
+        try:
+            normalized_ids.append(int(token_id))
+        except (TypeError, ValueError) as e:
+            raise TypeError(f"token_id must be int-convertible, got {type(token_id).__name__}: {token_id!r}") from e
+    return normalized_ids
+
+
+@lru_cache(maxsize=1)
+def _tokenize_prompt(text: str) -> list[int]:
+    """Tokenize a text prompt into valid token IDs for the model."""
+    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_MODEL, trust_remote_code=True)
+    messages = [{"role": "user", "content": text}]
+    token_ids = normalize_token_ids(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False))
+    assert len(token_ids) > _MIN_PROMPT_TOKENS, (
+        f"Prompt too short ({len(token_ids)} tokens, need >{_MIN_PROMPT_TOKENS}). "
+        f"The pipeline drops the first 34 chat‑template prefix tokens; "
+        f"use a longer prompt so content tokens remain after the drop."
+    )
+    return token_ids
+
+
+def _sampling_params(*, logprobs: bool = False, seed: int = 42) -> OmniDiffusionSamplingParams:
+    return OmniDiffusionSamplingParams(
+        num_inference_steps=2,
+        guidance_scale=0.0,
+        height=256,
+        width=256,
+        seed=seed,
+        extra_args={"logprobs": logprobs},
+    )
+
+
+async def _generate_once(
+    engine: AsyncOmni,
+    prompt: str | list[int] | dict,
+    *,
+    request_id: str,
+    sampling_params: OmniDiffusionSamplingParams,
+) -> OmniRequestOutput:
+    # Convert text prompt to dict with tokenized prompt_ids
+    if isinstance(prompt, str):
+        prompt_ids = _tokenize_prompt(prompt)
+        prompt = {"prompt_ids": prompt_ids}
+    elif isinstance(prompt, list):
+        prompt = {"prompt_ids": prompt}
+    # else: assume it's already a dict with prompt_ids
+
+    last_output = None
+    async for output in engine.generate(
+        prompt=prompt,
+        request_id=request_id,
+        sampling_params_list=[sampling_params],
+        output_modalities=["image"],
+    ):
+        last_output = output
+
+    assert last_output is not None
+    assert isinstance(last_output, OmniRequestOutput)
+    return last_output
+
+
+def _assert_valid_image_output(output: OmniRequestOutput) -> None:
+    assert output.final_output_type == "image"
+    assert output.images, "Expected at least one generated image"
+
+    image = output.images[0]
+    arr = np.asarray(image, dtype=np.float32) / 255.0
+
+    assert arr.ndim == 3 and arr.shape[2] == 3, f"Expected HWC RGB image, got shape={arr.shape}"
+    assert arr.shape[0] > 0 and arr.shape[1] > 0
+    assert 0.0 <= float(arr[0, 0, 0]) <= 1.0
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_async_omni_generate():
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+            max_num_seqs=1,
+        )
+        after.callback(engine.shutdown)
+
+        output = await _generate_once(
+            engine,
+            "a beautiful sunset over the ocean with vibrant orange and purple clouds "
+            "reflecting on the calm water surface near a rocky coastline",
+            request_id=f"test_{uuid.uuid4().hex[:8]}",
+            sampling_params=_sampling_params(logprobs=False, seed=42),
+        )
+
+        _assert_valid_image_output(output)
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_async_omni_generate_with_logprobs():
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+            max_num_seqs=1,
+        )
+        after.callback(engine.shutdown)
+
+        output = await _generate_once(
+            engine,
+            "a futuristic city at night with neon lights glowing on tall glass "
+            "skyscrapers and flying vehicles soaring between the buildings",
+            request_id=f"test_lp_{uuid.uuid4().hex[:8]}",
+            sampling_params=_sampling_params(logprobs=True, seed=123),
+        )
+
+        _assert_valid_image_output(output)
+
+        assert output.trajectory_latents is not None, "trajectory_latents should be present"
+        assert hasattr(output.trajectory_latents, "shape")
+        assert output.trajectory_latents.numel() > 0
+
+        assert output.trajectory_timesteps is not None, "trajectory_timesteps should be present"
+        assert hasattr(output.trajectory_timesteps, "shape")
+        assert output.trajectory_timesteps.numel() > 0
+
+        assert output.trajectory_log_probs is not None, "trajectory_log_probs should be present when logprobs=True"
+        assert hasattr(output.trajectory_log_probs, "shape")
+        assert output.trajectory_log_probs.numel() > 0
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_async_omni_generate_concurrent():
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+            max_num_seqs=1,
+        )
+        after.callback(engine.shutdown)
+
+        prompts = [
+            "a beautiful sunset over the ocean with vibrant orange and purple clouds "
+            "reflecting on the calm water surface near a rocky coastline",
+            "a fluffy orange cat sitting on a wooden windowsill looking outside at "
+            "a garden full of colorful flowers on a bright sunny afternoon",
+            "a majestic mountain landscape covered with fresh white snow under a "
+            "clear blue sky with pine trees in the foreground and a frozen lake",
+            "a futuristic city at night with neon lights glowing on tall glass "
+            "skyscrapers and flying vehicles soaring between the buildings",
+        ]
+
+        tasks = [
+            _generate_once(
+                engine,
+                prompt,
+                request_id=f"concurrent_{i}_{uuid.uuid4().hex[:8]}",
+                sampling_params=_sampling_params(logprobs=False, seed=100 + i),
+            )
+            for i, prompt in enumerate(prompts)
+        ]
+
+        outputs = await asyncio.gather(*tasks)
+
+        assert len(outputs) == len(prompts)
+        for output in outputs:
+            _assert_valid_image_output(output)

tests/e2e/features/custom_pipeline/test_worker_extension.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from contextlib import ExitStack
+
+import pytest
+
+from tests.e2e.features.helpers.custom_pipeline import (
+    QwenImagePipelineWithLogProbForTest,
+)
+from tests.helpers.mark import hardware_test
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+CUSTOM_PIPELINE_CLASS = QwenImagePipelineWithLogProbForTest
+WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest"
+MODEL = "tiny-random/Qwen-Image"
+
+pytestmark = [pytest.mark.core_model]
+
+
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.asyncio
+async def test_collective_rpc_test_extension_name():
+    with ExitStack() as after:
+        engine = AsyncOmni(
+            model=MODEL,
+            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
+            worker_extension_cls=WORKER_EXTENSION_CLASS,
+            enforce_eager=True,
+        )
+        after.callback(engine.shutdown)
+
+        result = await engine.collective_rpc(method="test_extension_name")
+
+    assert isinstance(result, list), f"Expected list, got {type(result)}"
+    assert result == [["vllm-omni-colocate-worker-extension-for-test"]]

tests/e2e/features/fullduplex/engine/test_duplex_control_client.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.experimental.fullduplex.engine.duplex_control_client import (
+    DuplexControlClient,
+    DuplexControlRequestError,
+)
+from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseActivity
+from vllm_omni.experimental.fullduplex.engine.messages import (
+    DuplexControlError,
+    DuplexControlResultMessage,
+    DuplexFence,
+    ResumeDuplexSessionMessage,
+    TouchDuplexSessionMessage,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Transport:
+    def __init__(self) -> None:
+        self.calls = []
+
+    def execute(self, key, message, **kwargs):
+        self.calls.append((key, message, kwargs))
+        return DuplexControlResultMessage(
+            control_id=message.control_id,
+            fence=message.fence,
+            operation="touch" if isinstance(message, TouchDuplexSessionMessage) else "resume",
+            session_id=message.session_id,
+            ok=True,
+            stage_results=[],
+        )
+
+
+class _UnsupportedTransport:
+    def execute(self, key, message, **kwargs):
+        del key, kwargs
+        return DuplexControlResultMessage(
+            control_id=message.control_id,
+            fence=message.fence,
+            operation="touch",
+            session_id=message.session_id,
+            ok=False,
+            stage_results=[{"result": {"supported": False}}],
+            unsupported_count=1,
+            error=DuplexControlError(code="invalid_capability", message="unsupported", retryable=False),
+            accepted_fence=DuplexFence(message.session_id, epoch=2),
+            lease_generation=7,
+        )
+
+
+def test_control_client_routes_touch_and_resume_by_control_id() -> None:
+    transport = _Transport()
+    control_ids = iter(("touch-id", "resume-id"))
+    client = DuplexControlClient(transport, control_id_factory=lambda: next(control_ids))
+    fence = DuplexFence("sid-client")
+
+    assert (
+        client.touch(
+            fence.session_id,
+            fence=fence,
+            activity=DuplexLeaseActivity.PLAYBACK_ACK,
+            timeout=2.0,
+        )["ok"]
+        is True
+    )
+    assert (
+        client.resume(
+            fence.session_id,
+            fence=fence,
+            expected_lease_generation=7,
+            timeout=3.0,
+        )["ok"]
+        is True
+    )
+
+    touch_key, touch_message, touch_kwargs = transport.calls[0]
+    assert touch_key == ("duplex", "touch-id")
+    assert isinstance(touch_message, TouchDuplexSessionMessage)
+    assert touch_message.activity == DuplexLeaseActivity.PLAYBACK_ACK.value
+    assert touch_kwargs["timeout"] == 2.0
+
+    resume_key, resume_message, resume_kwargs = transport.calls[1]
+    assert resume_key == ("duplex", "resume-id")
+    assert isinstance(resume_message, ResumeDuplexSessionMessage)
+    assert resume_message.expected_lease_generation == 7
+    assert resume_kwargs["timeout"] == 3.0
+
+
+def test_control_client_treats_ok_false_as_authoritative_failure() -> None:
+    client = DuplexControlClient(_UnsupportedTransport(), control_id_factory=lambda: "unsupported")
+    fence = DuplexFence("sid-unsupported")
+
+    with pytest.raises(DuplexControlRequestError, match="duplex touch failed") as exc_info:
+        client.touch(
+            fence.session_id,
+            fence=fence,
+            activity=DuplexLeaseActivity.HEARTBEAT,
+            timeout=1.0,
+        )
+
+    assert exc_info.value.code == "invalid_capability"
+    assert exc_info.value.retryable is False
+    assert exc_info.value.accepted_fence == DuplexFence(fence.session_id, epoch=2)
+    assert exc_info.value.lease_generation == 7

tests/e2e/features/fullduplex/engine/test_duplex_deploy_config.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Focused tests for typed deployment-only runtime configuration."""
+
+from __future__ import annotations
+
+from dataclasses import FrozenInstanceError
+
+import pytest
+
+from vllm_omni.config.stage_config import (
+    DuplexSessionRuntimeConfig,
+    load_deploy_config,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_duplex_session_runtime_defaults_are_typed_and_immutable(tmp_path) -> None:
+    deploy_path = tmp_path / "duplex.yaml"
+    deploy_path.write_text("duplex_session: {}\nstages: []\n", encoding="utf-8")
+
+    deploy = load_deploy_config(deploy_path)
+
+    assert isinstance(deploy.duplex_session, DuplexSessionRuntimeConfig)
+    assert deploy.duplex_session.idle_ttl_s == 300.0
+    assert deploy.duplex_session.disconnect_grace_s == 30.0
+    assert deploy.duplex_session.reaper_interval_s == 5.0
+    assert deploy.duplex_session.resume_replay_ttl_s == 60.0
+    assert deploy.duplex_session.resume_replay_max_bytes_per_session == 8 * 1024 * 1024
+    assert deploy.duplex_session.max_pending_input_bytes_per_session == 16 * 1024 * 1024
+    assert deploy.duplex_session.max_pending_turns_per_session == 4
+    assert deploy.duplex_session.max_sessions == 1
+    assert deploy.duplex_session.completed_append_cache_size == 256
+    with pytest.raises(FrozenInstanceError):
+        deploy.duplex_session.idle_ttl_s = 1.0  # type: ignore[misc]
+
+
+def test_duplex_session_runtime_accepts_disabled_idle_expiry(tmp_path) -> None:
+    deploy_path = tmp_path / "duplex.yaml"
+    deploy_path.write_text("duplex_session:\n  idle_ttl_s: null\nstages: []\n", encoding="utf-8")
+
+    deploy = load_deploy_config(deploy_path)
+
+    assert deploy.duplex_session.idle_ttl_s is None
+
+
+@pytest.mark.parametrize(
+    ("name", "value"),
+    [
+        ("idle_ttl_s", 0),
+        ("idle_ttl_s", -1),
+        ("disconnect_grace_s", 0),
+        ("reaper_interval_s", -1),
+        ("resume_replay_ttl_s", 0),
+        ("resume_replay_max_bytes_per_session", -1),
+        ("max_pending_input_bytes_per_session", 0),
+        ("max_pending_turns_per_session", -1),
+        ("max_sessions", 0),
+        ("completed_append_cache_size", 0),
+    ],
+)
+def test_duplex_session_runtime_rejects_non_positive_values(tmp_path, name: str, value: int) -> None:
+    deploy_path = tmp_path / "duplex.yaml"
+    deploy_path.write_text(f"duplex_session:\n  {name}: {value}\nstages: []\n", encoding="utf-8")
+
+    with pytest.raises(ValueError, match=rf"duplex_session\.{name} must be positive"):
+        load_deploy_config(deploy_path)

tests/e2e/features/fullduplex/engine/test_duplex_intermediate.py

@@ -0,0 +1,14 @@
+import pytest
+
+from vllm_omni.experimental.fullduplex.engine.intermediate import get_stream_request_key
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_get_stream_request_key_requires_stable_identifier():
+    with pytest.raises(ValueError, match="stable request id"):
+        get_stream_request_key({"ids": {"tts": [1, 2, 3]}})
+
+
+def test_get_stream_request_key_accepts_global_request_id():
+    assert get_stream_request_key({"global_request_id": ["duplex-sid-stage0"]}) == "duplex-sid-stage0"

tests/e2e/features/fullduplex/engine/test_duplex_lease.py

@@ -0,0 +1,224 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from dataclasses import FrozenInstanceError
+
+import pytest
+
+from vllm_omni.experimental.fullduplex.engine.duplex_runtime import DuplexInputMode
+from vllm_omni.experimental.fullduplex.engine.duplex_session import DuplexSessionRuntimeManager
+from vllm_omni.experimental.fullduplex.engine.lease import (
+    DuplexLeaseActivity,
+    DuplexLeaseConfig,
+)
+from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeMonotonicClock:
+    def __init__(self, value: float) -> None:
+        self.value = value
+
+    def __call__(self) -> float:
+        return self.value
+
+    def advance(self, seconds: float) -> None:
+        self.value += seconds
+
+
+def _lease_config(*, idle_ttl_s: float | None = 300.0) -> DuplexLeaseConfig:
+    return DuplexLeaseConfig(
+        idle_ttl_s=idle_ttl_s,
+        disconnect_grace_s=30.0,
+    )
+
+
+def test_open_touch_detach_and_idle_expiry_use_monotonic_time() -> None:
+    clock = FakeMonotonicClock(100.0)
+    manager = DuplexSessionRuntimeManager(clock=clock)
+    fence = DuplexFence("sid-expiry")
+    session = manager.open_session(fence, lease_config=_lease_config())
+
+    assert session.lease.last_activity == 100.0
+    clock.advance(10.0)
+    session.touch(fence, DuplexLeaseActivity.HEARTBEAT)
+    session.detach(fence)
+    clock.advance(29.0)
+    assert session.lease.disconnect_grace_expired(clock()) is False
+    assert manager.collect_expired() == []
+
+    clock.advance(272.0)
+    expired = manager.collect_expired()
+
+    assert [item.session_id for item in expired] == ["sid-expiry"]
+    assert expired[0].reason == "disconnect_grace_expired"
+    assert manager.get("sid-expiry") is session
+    manager.finalize_close_session(session)
+    assert manager.get("sid-expiry") is None
+
+
+def test_detach_grace_and_resume_advance_token_independent_lease_generation() -> None:
+    clock = FakeMonotonicClock(10.0)
+    manager = DuplexSessionRuntimeManager(clock=clock)
+    fence = DuplexFence("sid-resume")
+    session = manager.open_session(fence, lease_config=_lease_config())
+
+    session.detach(fence)
+    clock.advance(29.0)
+    assert session.lease.disconnect_grace_expired(clock()) is False
+    clock.advance(2.0)
+    assert session.lease.disconnect_grace_expired(clock()) is True
+
+    generation = session.resume(fence, expected_lease_generation=0)
+
+    assert generation == 1
+    assert session.lease.detached_at is None
+    assert session.lease.disconnect_grace_expired(clock()) is False
+    with pytest.raises(ValueError, match="lease generation mismatch"):
+        session.resume(fence, expected_lease_generation=0)
+
+
+def test_active_operation_prevents_mid_transaction_expiry() -> None:
+    clock = FakeMonotonicClock(0.0)
+    manager = DuplexSessionRuntimeManager(clock=clock)
+    fence = DuplexFence("sid-operation")
+    session = manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=5.0))
+
+    session.begin_operation(fence, "append-1")
+    clock.advance(10.0)
+    assert manager.collect_expired() == []
+
+    session.end_operation(fence, "append-1")
+    assert manager.collect_expired() == []
+    clock.advance(6.0)
+    assert [item.session_id for item in manager.collect_expired()] == ["sid-operation"]
+
+
+def test_sessions_have_independent_activity_deadlines_and_resources() -> None:
+    clock = FakeMonotonicClock(0.0)
+    manager = DuplexSessionRuntimeManager(clock=clock)
+    fence_a = DuplexFence("sid-a")
+    fence_b = DuplexFence("sid-b")
+    session_a = manager.open_session(fence_a, lease_config=_lease_config(idle_ttl_s=10.0))
+    session_b = manager.open_session(fence_b, lease_config=_lease_config(idle_ttl_s=10.0))
+    session_a.reserve_stage_request(0, "req-a-reserved", fence=fence_a)
+    session_a.bind_stage_request(1, "req-a-submitted", fence=fence_a)
+    session_b.bind_stage_request(0, "req-b", fence=fence_b)
+
+    clock.advance(6.0)
+    session_b.touch(fence_b, DuplexLeaseActivity.MODEL_OUTPUT)
+    clock.advance(5.0)
+    expired = manager.collect_expired()
+
+    assert len(expired) == 1
+    assert expired[0].session_id == "sid-a"
+    assert expired[0].reserved_request_ids == ("req-a-reserved",)
+    assert expired[0].submitted_request_ids == ("req-a-submitted",)
+    assert manager.get("sid-b") is session_b
+    assert session_b.resource_request_ids() == ["req-b"]
+
+
+def test_close_and_reaper_select_exactly_one_terminal_transition() -> None:
+    clock = FakeMonotonicClock(0.0)
+    manager = DuplexSessionRuntimeManager(clock=clock)
+    fence = DuplexFence("sid-race")
+    manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=1.0))
+    clock.advance(2.0)
+
+    assert len(manager.collect_expired()) == 1
+    assert manager.close_session(fence, reason="explicit_close") is None
+    assert manager.collect_expired() == []
+
+    fence_2 = DuplexFence("sid-race-2")
+    manager.open_session(fence_2, lease_config=_lease_config(idle_ttl_s=1.0))
+    assert manager.close_session(fence_2, reason="explicit_close") is not None
+    clock.advance(2.0)
+    assert manager.collect_expired() == []
+
+
+def test_stale_fence_cannot_touch_detach_or_resume_lease() -> None:
+    clock = FakeMonotonicClock(0.0)
+    manager = DuplexSessionRuntimeManager(clock=clock)
+    current = DuplexFence("sid-fence", epoch=1)
+    stale = DuplexFence("sid-fence", epoch=0)
+    session = manager.open_session(current, lease_config=_lease_config())
+
+    with pytest.raises(RuntimeError, match="fence mismatch"):
+        session.touch(stale, DuplexLeaseActivity.APPEND)
+    with pytest.raises(RuntimeError, match="fence mismatch"):
+        session.detach(stale)
+    with pytest.raises(RuntimeError, match="fence mismatch"):
+        session.resume(stale, expected_lease_generation=0)
+
+
+def test_disabled_idle_expiry_never_collects_session() -> None:
+    clock = FakeMonotonicClock(0.0)
+    manager = DuplexSessionRuntimeManager(clock=clock)
+    manager.open_session(DuplexFence("sid-no-expiry"), lease_config=_lease_config(idle_ttl_s=None))
+
+    clock.advance(1_000_000.0)
+
+    assert manager.collect_expired() == []
+
+
+def test_detached_session_expires_at_disconnect_grace_when_idle_ttl_is_disabled() -> None:
+    clock = FakeMonotonicClock(0.0)
+    manager = DuplexSessionRuntimeManager(clock=clock)
+    fence = DuplexFence("sid-disconnect-grace")
+    session = manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=None))
+
+    session.detach(fence)
+    clock.advance(29.0)
+    assert manager.collect_expired() == []
+
+    clock.advance(1.0)
+    expired = manager.collect_expired()
+
+    assert [item.session_id for item in expired] == [fence.session_id]
+    assert expired[0].reason == "disconnect_grace_expired"
+
+
+def test_runtime_manager_enforces_server_owned_session_admission() -> None:
+    manager = DuplexSessionRuntimeManager(max_sessions=2)
+    manager.open_session(DuplexFence("sid-admission-a"))
+    manager.open_session(DuplexFence("sid-admission-b"))
+
+    with pytest.raises(RuntimeError, match="duplex_session_capacity_exhausted"):
+        manager.open_session(DuplexFence("sid-admission-c"))
+
+
+def test_completed_append_cache_is_bounded() -> None:
+    manager = DuplexSessionRuntimeManager(completed_append_limit=2)
+    fence = DuplexFence("sid-completed-cache")
+    session = manager.open_session(fence)
+
+    for index in range(3):
+        session.record_completed_append(
+            f"operation-{index}",
+            fence=fence,
+            mode=DuplexInputMode.TURN_COMMIT_ONLY,
+            final=True,
+            stage_results=[{"index": index}],
+        )
+
+    assert list(session.completed_appends) == ["operation-1", "operation-2"]
+
+    session.accept_fence(DuplexFence(fence.session_id, epoch=1))
+    assert session.completed_appends == {}
+
+
+def test_lease_config_and_expiry_record_are_immutable() -> None:
+    config = _lease_config()
+    with pytest.raises(FrozenInstanceError):
+        config.idle_ttl_s = 1.0  # type: ignore[misc]
+
+    clock = FakeMonotonicClock(0.0)
+    manager = DuplexSessionRuntimeManager(clock=clock)
+    manager.open_session(DuplexFence("sid-record"), lease_config=_lease_config(idle_ttl_s=1.0))
+    clock.advance(2.0)
+    record = manager.collect_expired()[0]
+    with pytest.raises(FrozenInstanceError):
+        record.reason = "changed"  # type: ignore[misc]

tests/e2e/features/fullduplex/engine/test_duplex_runtime.py

@@ -0,0 +1,480 @@
+from dataclasses import FrozenInstanceError
+from types import SimpleNamespace
+
+import pytest
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.experimental.fullduplex.engine import duplex_runtime
+from vllm_omni.experimental.fullduplex.engine.contracts import (
+    duplex_data_plane_request_info,
+    duplex_resource_request_belongs_to_session,
+    duplex_resource_request_id,
+)
+from vllm_omni.experimental.fullduplex.engine.duplex_runtime import (
+    DuplexInputMode,
+    DuplexOutputAction,
+    DuplexOutputDecision,
+    DuplexRuntimeCapabilities,
+    DuplexSessionRuntimeManager,
+)
+from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+from vllm_omni.experimental.fullduplex.minicpmo45.runtime import (
+    MiniCPMO45DuplexRuntimeExtension,
+    build_duplex_data_plane_prompt,
+    duplex_scheduler_token_budget,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_duplex_fence_is_immutable():
+    fence = DuplexFence("session")
+
+    with pytest.raises(FrozenInstanceError):
+        fence.epoch = 1  # type: ignore[misc]
+
+    assert not hasattr(fence, "__dict__")
+
+
+def test_duplex_runtime_tracks_stage_bindings_and_barge_in_epoch():
+    manager = DuplexSessionRuntimeManager()
+    session = manager.open_session(
+        DuplexFence("sid-1"),
+        capabilities=DuplexRuntimeCapabilities(
+            input_modes={DuplexInputMode.APPEND_TOKENS},
+        ),
+    )
+    session.bind_stage_request(stage_id=0, request_id="req-stage0", fence=session.fence)
+    session.bind_stage_request(stage_id=1, request_id="req-stage1", fence=session.fence)
+
+    update = session.append_input(mode=DuplexInputMode.APPEND_TOKENS, fence=session.fence)
+    next_fence = DuplexFence("sid-1", epoch=1)
+    stale_request_ids = session.release_fence(session.fence)
+    session.accept_fence(next_fence)
+
+    assert update.seq == 1
+    assert session.fence == next_fence
+    assert stale_request_ids == ["req-stage0", "req-stage1"]
+    assert session.stage_bindings == {}
+
+
+def test_duplex_runtime_tracks_same_request_id_for_each_pipeline_stage():
+    manager = DuplexSessionRuntimeManager()
+    fence = DuplexFence("sid-shared-pipeline-request")
+    session = manager.open_session(fence)
+
+    session.reserve_stage_request(0, "req-shared", fence=fence)
+    session.bind_stage_request(0, "req-shared", fence=fence)
+    session.bind_stage_request(1, "req-shared", fence=fence)
+
+    assert session.stage_bindings[0].request_id == "req-shared"
+    assert session.stage_bindings[1].request_id == "req-shared"
+    assert session.resource_request_ids() == ["req-shared"]
+    assert session.input_seq == 0
+
+
+def test_duplex_runtime_extension_validation_rejects_missing_methods():
+    class IncompleteExtension:
+        def configure_sampling_params(self, *, runtime_config, defaults):
+            del runtime_config
+            return defaults
+
+    with pytest.raises(TypeError, match="plan_append"):
+        duplex_runtime.validate_duplex_runtime_extension(IncompleteExtension())
+
+
+def test_duplex_runtime_extension_validation_rejects_stage_count_mismatch():
+    class WrongStageCountExtension(MiniCPMO45DuplexRuntimeExtension):
+        def configure_sampling_params(self, *, runtime_config, defaults):
+            del runtime_config
+            return defaults[:1]
+
+    with pytest.raises(ValueError, match="one sampling parameter per stage"):
+        duplex_runtime.validate_duplex_runtime_extension(
+            WrongStageCountExtension(),
+            sampling_defaults=(SamplingParams(), SamplingParams()),
+        )
+
+
+def test_duplex_runtime_extension_validation_rejects_sampling_type_mismatch():
+    class WrongSamplingTypeExtension(MiniCPMO45DuplexRuntimeExtension):
+        def configure_sampling_params(self, *, runtime_config, defaults):
+            del runtime_config
+            return tuple(object() for _ in defaults)
+
+    with pytest.raises(TypeError, match="sampling parameter type"):
+        duplex_runtime.validate_duplex_runtime_extension(
+            WrongSamplingTypeExtension(),
+            sampling_defaults=(SamplingParams(), SamplingParams()),
+        )
+
+
+def test_duplex_output_decision_metadata_is_immutable():
+    decision = DuplexOutputDecision(
+        action=DuplexOutputAction.DIRECT_RESPONSE,
+        metadata={"model_listen": True},
+    )
+
+    with pytest.raises(TypeError):
+        decision.metadata["model_listen"] = False
+
+
+def _decide_minicpmo_output(
+    output: object,
+    *,
+    segment_token_ids: tuple[int, ...] = (),
+    segment_output_metadata: dict | None = None,
+):
+    return MiniCPMO45DuplexRuntimeExtension().decide_output(
+        stage_id=0,
+        final_stage_id=1,
+        segment_finished=True,
+        segment_token_ids=segment_token_ids,
+        segment_output_metadata=segment_output_metadata or {},
+        output=output,
+    )
+
+
+def test_minicpmo_extension_owns_stage_sampling_overrides():
+    defaults = (
+        SamplingParams(max_tokens=4),
+        SamplingParams(max_tokens=8),
+    )
+
+    configured = MiniCPMO45DuplexRuntimeExtension().configure_sampling_params(
+        runtime_config={
+            "duplex_stage_max_tokens": {"0": 20},
+            "duplex_stage_sampling_params": {"1": {"stop_token_ids": [151645]}},
+        },
+        defaults=defaults,
+    )
+
+    assert configured[0].max_tokens == 20
+    assert configured[1].stop_token_ids == [151645]
+    assert defaults[0].max_tokens == 4
+    assert 151645 not in (defaults[1].stop_token_ids or [])
+
+
+def test_minicpmo_output_decision_uses_raw_streaming_token_snapshot():
+    decision = _decide_minicpmo_output(
+        SimpleNamespace(outputs=[SimpleNamespace()]),
+        segment_token_ids=(151705,),
+        segment_output_metadata={"special_token_ids": {"listen_token_id": 151705}},
+    )
+
+    assert decision is not None
+    assert decision.action is DuplexOutputAction.DIRECT_RESPONSE
+    assert decision.metadata["duplex_native_decision"] == "listen"
+    assert decision.metadata["model_listen"] is True
+
+
+@pytest.mark.parametrize("attr", ["token_ids", "cumulative_token_ids"])
+def test_minicpmo_output_decision_ignores_output_level_token_history(attr):
+    output = SimpleNamespace(
+        multimodal_output={"special_token_ids": {"listen_token_id": 151705}},
+        outputs=[SimpleNamespace()],
+        **{attr: [42, 151705]},
+    )
+
+    assert _decide_minicpmo_output(output) is None
+
+
+@pytest.mark.parametrize("attr", ["token_ids", "cumulative_token_ids"])
+def test_minicpmo_output_decision_uses_completion_token_ids(attr):
+    output = SimpleNamespace(
+        multimodal_output={"special_token_ids": {"listen_token_id": 151705}},
+        outputs=[SimpleNamespace(**{attr: [42, 151705]})],
+    )
+
+    assert _decide_minicpmo_output(output) is not None
+
+
+def test_minicpmo_output_decision_uses_completion_stop_reason():
+    output = SimpleNamespace(
+        multimodal_output={"special_token_ids": {"listen_token_id": 151705}},
+        outputs=[SimpleNamespace(stop_reason=151705)],
+    )
+
+    assert _decide_minicpmo_output(output) is not None
+
+
+def test_duplex_runtime_cancel_fence_rejects_late_append_and_accepts_next_epoch():
+    manager = DuplexSessionRuntimeManager()
+    cancelled_fence = DuplexFence("sid-cancel-race")
+    next_fence = DuplexFence("sid-cancel-race", epoch=1)
+    session = manager.open_session(
+        cancelled_fence,
+        capabilities=DuplexRuntimeCapabilities(
+            input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
+        ),
+    )
+    session.bind_stage_request(
+        stage_id=0,
+        request_id="req-cancelled",
+        fence=cancelled_fence,
+    )
+
+    stale_request_ids = session.cancel_fence(cancelled_fence, next_fence)
+
+    assert stale_request_ids == ["req-cancelled"]
+    assert session.fence == next_fence
+    assert session.stage_bindings == {}
+    with pytest.raises(RuntimeError, match="fence mismatch"):
+        session.append_input(
+            mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
+            fence=cancelled_fence,
+        )
+    update = session.append_input(
+        mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
+        fence=next_fence,
+    )
+    assert update.seq == 1
+
+
+def test_duplex_runtime_stale_close_preserves_live_session_and_bindings():
+    manager = DuplexSessionRuntimeManager()
+    current_fence = DuplexFence("sid-stale-close", epoch=1)
+    session = manager.open_session(
+        current_fence,
+        capabilities=DuplexRuntimeCapabilities(
+            input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
+        ),
+    )
+    session.bind_stage_request(0, "req-live", fence=current_fence)
+
+    with pytest.raises(RuntimeError, match="fence mismatch"):
+        manager.close_session(DuplexFence("sid-stale-close"))
+
+    assert manager.get("sid-stale-close") is session
+    assert session.stage_request_ids() == ["req-live"]
+
+
+def test_duplex_runtime_reopen_rejects_late_append_from_old_incarnation():
+    manager = DuplexSessionRuntimeManager()
+    old_fence = DuplexFence("sid-reopen", incarnation=0)
+    old_session = manager.open_session(
+        old_fence,
+        capabilities=DuplexRuntimeCapabilities(
+            input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
+        ),
+    )
+    manager.close_session(old_fence)
+
+    new_fence = DuplexFence("sid-reopen", incarnation=1)
+    new_session = manager.open_session(
+        new_fence,
+        capabilities=old_session.capabilities,
+    )
+
+    with pytest.raises(RuntimeError, match="fence mismatch"):
+        new_session.append_input(
+            mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
+            fence=old_fence,
+        )
+    assert (
+        new_session.append_input(
+            mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
+            fence=new_fence,
+        ).seq
+        == 1
+    )
+
+
+def test_duplex_prompt_expands_incarnation_metadata():
+    fence = DuplexFence("sid-incarnation", incarnation=3)
+
+    prompt = build_duplex_data_plane_prompt(
+        request_id=duplex_resource_request_id(fence, "stage0"),
+        fence=fence,
+        session_config={},
+        runtime_config={},
+        seq=1,
+        turn_seq=1,
+        mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
+        payload={"is_speech": True},
+        final=False,
+    )
+
+    assert prompt["model_intermediate_buffer"]["duplex"]["incarnation"] == 3
+
+
+def test_duplex_runtime_tracks_turn_local_append_sequence():
+    manager = DuplexSessionRuntimeManager()
+    session = manager.open_session(
+        DuplexFence("sid-turn-seq"),
+        capabilities=DuplexRuntimeCapabilities(
+            input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK},
+        ),
+    )
+
+    first = session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK, fence=session.fence)
+    second = session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK, fence=session.fence)
+    next_turn = DuplexFence("sid-turn-seq", turn_id=1, response_seq=1)
+    third = session.append_input(
+        mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
+        fence=next_turn,
+    )
+    fourth = session.append_input(
+        mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
+        fence=next_turn,
+    )
+
+    assert [first.seq, second.seq, third.seq, fourth.seq] == [1, 2, 3, 4]
+    assert [first.turn_seq, second.turn_seq, third.turn_seq, fourth.turn_seq] == [1, 2, 1, 2]
+    assert [first.turn_id, second.turn_id, third.turn_id, fourth.turn_id] == [0, 0, 1, 1]
+
+
+def test_duplex_runtime_rejects_unsupported_append_mode():
+    manager = DuplexSessionRuntimeManager()
+    session = manager.open_session(
+        DuplexFence("sid-2"),
+        capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.TURN_COMMIT_ONLY}),
+    )
+
+    with pytest.raises(ValueError, match="not supported"):
+        session.append_input(mode=DuplexInputMode.APPEND_TOKENS, fence=session.fence)
+
+
+def test_duplex_data_plane_request_info_extracts_structured_stage_result():
+    request_id, response_stage_id = duplex_data_plane_request_info(
+        {
+            "stage_results": [
+                {"result": {"supported": True}},
+                {
+                    "result": {
+                        "data_plane_append": True,
+                        "request_id": "duplex-sid-e0-stage0-s1",
+                        "response_stage_id": 1,
+                    }
+                },
+            ]
+        }
+    )
+
+    assert request_id == "duplex-sid-e0-stage0-s1"
+    assert response_stage_id == 1
+
+
+def test_duplex_data_plane_request_info_rejects_missing_request_id():
+    assert duplex_data_plane_request_info(
+        {
+            "stage_results": [
+                {
+                    "result": {
+                        "data_plane_append": True,
+                        "request_id": "",
+                        "response_stage_id": 1,
+                    }
+                }
+            ]
+        }
+    ) == (None, None)
+
+
+def test_duplex_scheduler_token_budget_estimates_pcm_slots():
+    assert (
+        duplex_scheduler_token_budget(
+            {
+                "audio": "AAAAAA==",
+                "format": "pcm_f32le",
+                "sample_rate_hz": 16000,
+            }
+        )
+        == 16
+    )
+
+
+def test_duplex_scheduler_token_budget_ignores_client_budget_fields():
+    assert (
+        duplex_scheduler_token_budget(
+            {
+                "audio": "AAAAAA==",
+                "format": "pcm_f32le",
+                "duplex_num_input_tokens": 999,
+                "num_input_tokens": 999,
+            }
+        )
+        == 16
+    )
+
+
+def test_resource_state_rejects_fence_regression_and_requires_explicit_fence():
+    current = DuplexFence("sid", epoch=2, turn_id=3, response_seq=4)
+    manager = DuplexSessionRuntimeManager()
+    session = manager.open_session(
+        current,
+        capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}),
+    )
+
+    for stale in (
+        DuplexFence("sid", epoch=1, turn_id=99, response_seq=99),
+        DuplexFence("sid", epoch=2, turn_id=2, response_seq=4),
+        DuplexFence("sid", epoch=2, turn_id=3, response_seq=3),
+    ):
+        with pytest.raises(RuntimeError, match="fence mismatch"):
+            session.accept_fence(stale)
+        assert session.fence == current
+
+    with pytest.raises(TypeError, match="fence"):
+        session.bind_stage_request(0, "request")
+    with pytest.raises(TypeError, match="fence"):
+        session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK)
+    with pytest.raises(TypeError, match="DuplexFence"):
+        manager.open_session("legacy-session")
+    with pytest.raises(TypeError, match="DuplexFence"):
+        manager.close_session("legacy-session")
+
+
+def test_resource_request_id_is_derived_from_fence_and_role():
+    fence = DuplexFence("sid-with-dashes", epoch=7, turn_id=11, response_seq=13)
+
+    assert duplex_resource_request_id(fence, "stage0") == "duplex-s.c2lkLXdpdGgtZGFzaGVz.i.0.e.7.r.stage0"
+    assert duplex_resource_request_id(fence, "stage1") == "duplex-s.c2lkLXdpdGgtZGFzaGVz.i.0.e.7.r.stage1"
+
+
+def test_resource_request_id_codec_separates_session_id_from_incarnation():
+    embedded_incarnation = duplex_resource_request_id(
+        DuplexFence("foo-i1", incarnation=0),
+        "stage0",
+    )
+    actual_incarnation = duplex_resource_request_id(
+        DuplexFence("foo", incarnation=1),
+        "stage0",
+    )
+
+    assert embedded_incarnation != actual_incarnation
+
+
+def test_resource_request_id_session_membership_uses_encoded_identity():
+    request_id = duplex_resource_request_id(
+        DuplexFence("sid-with-dashes", incarnation=2, epoch=7),
+        "stage0",
+    )
+
+    assert duplex_resource_request_belongs_to_session(request_id, "sid-with-dashes") is True
+    assert duplex_resource_request_belongs_to_session(request_id, "sid") is False
+    assert duplex_resource_request_belongs_to_session("duplex-s.invalid.i.x.e.7.r.stage0", "sid") is False
+
+
+def test_placeholder_budget_is_planned_inside_omni_engine_boundary():
+    fence = DuplexFence("sid", turn_id=1, response_seq=1)
+    prompt = build_duplex_data_plane_prompt(
+        request_id=duplex_resource_request_id(fence, "stage0"),
+        fence=fence,
+        session_config={},
+        runtime_config={},
+        seq=2,
+        turn_seq=1,
+        mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
+        payload={
+            "audio": "AAAAAA==",
+            "format": "pcm_f32le",
+            "duplex_num_input_tokens": 999,
+            "num_input_tokens": 999,
+        },
+        final=False,
+    )
+
+    assert len(prompt["prompt_token_ids"]) == 16
+    assert prompt["model_intermediate_buffer"]["duplex"]["fence"] == fence
+    assert prompt["model_intermediate_buffer"]["duplex"]["scheduler_token_budget"] == 16

tests/e2e/features/fullduplex/minicpmo45/test_commit_policy.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.experimental.fullduplex.openai.commit_policy import (
+    CommitAction,
+    CommitSnapshot,
+    decide_commit_action,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    ("case", "snapshot", "expected"),
+    [
+        (
+            "first idle commit",
+            CommitSnapshot(
+                auto_responds=True,
+                speech_since_commit=True,
+                active_response_id=None,
+                overlap_speech_ms=0,
+                native_response_in_progress=False,
+                playback_active=False,
+            ),
+            CommitAction.START_AUTO_RESPONSE,
+        ),
+        (
+            "native auto-response remains model-owned during playback",
+            CommitSnapshot(
+                auto_responds=True,
+                speech_since_commit=True,
+                active_response_id=None,
+                overlap_speech_ms=800,
+                native_response_in_progress=True,
+                playback_active=True,
+            ),
+            CommitAction.START_AUTO_RESPONSE,
+        ),
+        (
+            "native auto-response commit without speech",
+            CommitSnapshot(
+                auto_responds=True,
+                speech_since_commit=False,
+                active_response_id=None,
+                overlap_speech_ms=0,
+                native_response_in_progress=True,
+                playback_active=False,
+            ),
+            CommitAction.COMMIT_ONLY,
+        ),
+        (
+            "non-auto response overlap",
+            CommitSnapshot(
+                auto_responds=False,
+                speech_since_commit=True,
+                active_response_id=None,
+                overlap_speech_ms=800,
+                native_response_in_progress=True,
+                playback_active=False,
+            ),
+            CommitAction.DEFER_ACTIVE_RESPONSE,
+        ),
+        (
+            "non-auto playback overlap",
+            CommitSnapshot(
+                auto_responds=False,
+                speech_since_commit=True,
+                active_response_id=None,
+                overlap_speech_ms=800,
+                native_response_in_progress=False,
+                playback_active=True,
+            ),
+            CommitAction.DEFER_ACTIVE_RESPONSE,
+        ),
+        (
+            "non-auto commit without overlap",
+            CommitSnapshot(
+                auto_responds=False,
+                speech_since_commit=True,
+                active_response_id=None,
+                overlap_speech_ms=0,
+                native_response_in_progress=False,
+                playback_active=False,
+            ),
+            CommitAction.COMMIT_ONLY,
+        ),
+    ],
+    ids=lambda value: value if isinstance(value, str) else None,
+)
+def test_decide_commit_action_covers_realtime_lifecycle_sequences(
+    case: str,
+    snapshot: CommitSnapshot,
+    expected: CommitAction,
+) -> None:
+    del case
+
+    assert decide_commit_action(snapshot) is expected

tests/e2e/features/fullduplex/minicpmo45/test_input.py

@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import base64
+
+import numpy as np
+import pytest
+
+from vllm_omni.experimental.fullduplex.minicpmo45.input import MiniCPMO45PcmAppendBuffer
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def pcm_payload(samples: int, *, speech: bool = True) -> dict[str, object]:
+    audio = np.ones(samples, dtype=np.float32).tobytes()
+    return {
+        "type": "audio",
+        "audio": base64.b64encode(audio).decode("ascii"),
+        "format": "pcm_f32le",
+        "sample_rate_hz": 16_000,
+        "is_speech": speech,
+    }
+
+
+def test_commit_does_not_add_silence_after_incremental_audio_was_drained():
+    buffer = MiniCPMO45PcmAppendBuffer()
+
+    emitted = buffer.append(pcm_payload(16_000), chunk_period_ms=1_000)
+    committed = buffer.commit(chunk_period_ms=1_000)
+
+    assert emitted is not None
+    assert not buffer.has_pending()
+    assert committed is None
+
+
+def test_append_emits_one_model_unit_when_multiple_units_are_buffered():
+    buffer = MiniCPMO45PcmAppendBuffer()
+
+    emitted = buffer.append(pcm_payload(32_000), chunk_period_ms=1_000)
+
+    assert emitted is not None
+    assert len(base64.b64decode(emitted["audio"])) == 16_000 * 4
+    assert buffer.pending_byte_count == 16_000 * 4
+
+
+def test_speech_marker_does_not_leak_across_irregular_chunk_boundaries():
+    buffer = MiniCPMO45PcmAppendBuffer()
+
+    assert buffer.append(pcm_payload(15_000), chunk_period_ms=1_000) is None
+    speech_unit = buffer.append(pcm_payload(2_000, speech=False), chunk_period_ms=1_000)
+    silence_unit = buffer.append(pcm_payload(15_000, speech=False), chunk_period_ms=1_000)
+
+    assert speech_unit is not None
+    assert speech_unit["is_speech"] is True
+    assert silence_unit is not None
+    assert silence_unit["is_speech"] is False
+
+
+def test_force_listen_marker_does_not_leak_across_irregular_chunk_boundaries():
+    buffer = MiniCPMO45PcmAppendBuffer()
+    forced = pcm_payload(15_000, speech=False)
+    forced["force_listen"] = True
+
+    assert buffer.append(forced, chunk_period_ms=1_000) is None
+    forced_unit = buffer.append(pcm_payload(2_000, speech=False), chunk_period_ms=1_000)
+    unforced_unit = buffer.append(pcm_payload(15_000, speech=False), chunk_period_ms=1_000)
+
+    assert forced_unit is not None
+    assert forced_unit["force_listen"] is True
+    assert unforced_unit is not None
+    assert unforced_unit["force_listen"] is False
+
+
+def test_pcm_append_rollback_restores_per_span_speech_markers():
+    buffer = MiniCPMO45PcmAppendBuffer()
+    assert buffer.append(pcm_payload(15_000), chunk_period_ms=1_000) is None
+    reservation = buffer.prepare_append(
+        pcm_payload(2_000, speech=False),
+        operation_id="mixed-unit",
+        chunk_period_ms=1_000,
+    )
+
+    assert reservation is not None
+    reservation.rollback()
+    first = buffer.append(pcm_payload(15_000, speech=False), chunk_period_ms=1_000)
+    second = buffer.flush(chunk_period_ms=1_000)
+
+    assert first is not None
+    assert first["is_speech"] is True
+    assert second is not None
+    assert second["is_speech"] is False
+
+
+def test_commit_without_speech_does_not_synthesize_terminal_audio():
+    buffer = MiniCPMO45PcmAppendBuffer()
+    buffer.append(pcm_payload(8_000, speech=False), chunk_period_ms=1_000)
+
+    committed = buffer.commit(chunk_period_ms=1_000)
+
+    assert committed is None
+
+
+def test_commit_resets_cumulative_turn_accounting():
+    buffer = MiniCPMO45PcmAppendBuffer()
+    buffer.append(pcm_payload(16_000), chunk_period_ms=1_000)
+    buffer.commit(chunk_period_ms=1_000)
+
+    empty = buffer.commit(chunk_period_ms=1_000)
+
+    assert empty is None
+
+
+def test_pcm_append_reservation_rollback_restores_emitted_audio():
+    buffer = MiniCPMO45PcmAppendBuffer()
+    original = pcm_payload(16_000)
+
+    reservation = buffer.prepare_append(
+        original,
+        operation_id="append-1",
+        chunk_period_ms=1_000,
+    )
+
+    assert reservation is not None
+    assert not buffer.has_pending()
+    reservation.rollback()
+    assert buffer.has_pending()
+
+    retried = buffer.flush(chunk_period_ms=1_000)
+    assert retried is not None
+    assert base64.b64decode(retried["audio"]) == base64.b64decode(original["audio"])
+
+
+def test_pcm_commit_keeps_prior_append_reservation_active():
+    buffer = MiniCPMO45PcmAppendBuffer()
+    append_reservation = buffer.prepare_append(
+        pcm_payload(16_000),
+        operation_id="append-before-commit",
+        chunk_period_ms=1_000,
+    )
+
+    commit_reservation = buffer.prepare_commit(
+        operation_id="commit-after-append",
+        chunk_period_ms=1_000,
+    )
+
+    assert append_reservation is not None
+    assert append_reservation.active
+    assert commit_reservation.payload is None
+    append_reservation.commit()
+    commit_reservation.commit()
+
+
+def test_pcm_commit_reservation_rollback_restores_residual_audio():
+    buffer = MiniCPMO45PcmAppendBuffer()
+    original = pcm_payload(8_000)
+    assert (
+        buffer.prepare_append(
+            original,
+            operation_id="buffer-half-chunk",
+            chunk_period_ms=1_000,
+        )
+        is None
+    )
+
+    reservation = buffer.prepare_commit(
+        operation_id="final-half-chunk",
+        chunk_period_ms=1_000,
+    )
+
+    assert reservation.payload is not None
+    assert reservation.payload["final"] is True
+    reservation.rollback()
+
+    retried = buffer.prepare_commit(
+        operation_id="retry-final-half-chunk",
+        chunk_period_ms=1_000,
+    )
+    assert retried.payload is not None
+    assert base64.b64decode(retried.payload["audio"]) == (base64.b64decode(original["audio"]) + b"\x00" * (8_000 * 4))

tests/e2e/features/fullduplex/openai/test_duplex_audio.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import base64
+
+import numpy as np
+import pytest
+
+from vllm_omni.experimental.fullduplex.openai.audio import convert_input_audio_with_rate
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_pcm16_input_conversion_is_pure_and_reports_target_rate():
+    pcm16 = np.array([-32768, 0, 32767], dtype="<i2").tobytes()
+    encoded = base64.b64encode(pcm16).decode("ascii")
+
+    converted, fmt, sample_rate = convert_input_audio_with_rate(
+        encoded,
+        "pcm16",
+        sample_rate_hz=16_000,
+    )
+
+    assert encoded == base64.b64encode(pcm16).decode("ascii")
+    assert fmt == "pcm_f32le"
+    assert sample_rate == 16_000
+    samples = np.frombuffer(base64.b64decode(converted), dtype="<f4")
+    assert np.allclose(samples, [-1.0, 0.0, 32767 / 32768])
+
+
+@pytest.mark.parametrize("sample_rate_hz", [8_000, 192_000])
+def test_pcm16_input_conversion_accepts_supported_sample_rate_boundaries(sample_rate_hz: int):
+    encoded = base64.b64encode(np.zeros(8, dtype="<i2").tobytes()).decode("ascii")
+
+    _, fmt, converted_rate = convert_input_audio_with_rate(
+        encoded,
+        "pcm16",
+        sample_rate_hz=sample_rate_hz,
+    )
+
+    assert fmt == "pcm_f32le"
+    assert converted_rate == 16_000
+
+
+@pytest.mark.parametrize("sample_rate_hz", [7_999, 192_001, 10**1000])
+def test_pcm16_input_conversion_rejects_sample_rates_outside_supported_range(sample_rate_hz: int):
+    encoded = base64.b64encode(np.zeros(8, dtype="<i2").tobytes()).decode("ascii")
+
+    with pytest.raises(ValueError, match="sample_rate_hz"):
+        convert_input_audio_with_rate(
+            encoded,
+            "pcm16",
+            sample_rate_hz=sample_rate_hz,
+        )

tests/e2e/features/fullduplex/openai/test_runtime_adapter_boundary.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import subprocess
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.experimental.fullduplex.openai.runtime_adapter import (
+    validate_serving_runtime_adapter,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    "module_name",
+    [
+        "vllm_omni.experimental.fullduplex.openai.session_runner",
+        "vllm_omni.experimental.fullduplex.openai.runtime_bridge",
+        "vllm_omni.experimental.fullduplex.openai.serving",
+    ],
+)
+def test_generic_openai_runtime_import_does_not_load_minicpmo45(module_name: str) -> None:
+    repo_root = Path(__file__).resolve().parents[3]
+    script = f"""
+import sys
+
+import {module_name}
+
+loaded = sorted(
+    name
+    for name in sys.modules
+    if name == "vllm_omni.experimental.fullduplex.minicpmo45"
+    or name.startswith("vllm_omni.experimental.fullduplex.minicpmo45.")
+)
+if loaded:
+    raise SystemExit("generic OpenAI runtime loaded MiniCPM-o: " + ", ".join(loaded))
+"""
+
+    result = subprocess.run(
+        [sys.executable, "-c", script],
+        cwd=repo_root,
+        capture_output=True,
+        text=True,
+        timeout=60,
+        check=False,
+    )
+
+    assert result.returncode == 0, result.stderr
+
+
+def test_generic_handler_requires_explicit_serving_runtime_adapter() -> None:
+    repo_root = Path(__file__).resolve().parents[3]
+    script = """
+from types import SimpleNamespace
+
+from vllm_omni.experimental.fullduplex.openai.serving import OmniDuplexSessionHandler
+
+chat_service = SimpleNamespace(engine_client=SimpleNamespace())
+try:
+    OmniDuplexSessionHandler(chat_service=chat_service)
+except ValueError as exc:
+    if "serving runtime adapter" not in str(exc).lower():
+        raise
+else:
+    raise SystemExit("generic handler silently selected a model serving adapter")
+"""
+
+    result = subprocess.run(
+        [sys.executable, "-c", script],
+        cwd=repo_root,
+        capture_output=True,
+        text=True,
+        timeout=60,
+        check=False,
+    )
+
+    assert result.returncode == 0, result.stderr
+
+
+def _valid_runtime_adapter() -> SimpleNamespace:
+    data_plane = SimpleNamespace(
+        begin_request=lambda request_id: None,
+        is_terminal=lambda request_id: False,
+        mark_terminal=lambda request_id: None,
+        close_stream=lambda request_id: None,
+        close_session=lambda session_id, **kwargs: None,
+        project=lambda result, **kwargs: (),
+    )
+    return SimpleNamespace(
+        adapter_id="test-adapter",
+        session_states={},
+        data_plane=data_plane,
+        clean_response_done_prefix="clean",
+        interrupted_tts_prefix="interrupted",
+        private_runtime_config_keys=frozenset({"private_key"}),
+        create_session_state=lambda: object(),
+        session_state=lambda session_id: object(),
+        remove_session_state=lambda session_id: None,
+        is_enabled=lambda config: True,
+        capabilities=lambda **kwargs: object(),
+        validate_client_extra_body=lambda extra_body: None,
+        prepare_runtime_config=lambda config, **kwargs: {},
+        runtime_config_for_update=lambda config, current: {},
+        data_plane_context=lambda **kwargs: object(),
+    )
+
+
+@pytest.mark.parametrize(
+    ("attribute", "invalid_value"),
+    [
+        ("session_states", None),
+        ("clean_response_done_prefix", None),
+        ("interrupted_tts_prefix", None),
+        ("private_runtime_config_keys", {"private_key"}),
+    ],
+)
+def test_runtime_adapter_validator_rejects_invalid_protocol_attributes(
+    attribute: str,
+    invalid_value: object,
+) -> None:
+    adapter = _valid_runtime_adapter()
+    setattr(adapter, attribute, invalid_value)
+
+    with pytest.raises(TypeError, match=attribute):
+        validate_serving_runtime_adapter(adapter)
+
+
+def test_runtime_adapter_validator_rejects_incomplete_data_plane() -> None:
+    adapter = _valid_runtime_adapter()
+    adapter.data_plane.close_session = None
+
+    with pytest.raises(TypeError, match="data_plane.*close_session"):
+        validate_serving_runtime_adapter(adapter)

tests/e2e/features/fullduplex/openai/test_websocket_actor.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+
+import pytest
+
+from vllm_omni.experimental.fullduplex.openai.websocket import (
+    DuplexWebSocketActor,
+    normalize_duplex_input_event,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeWebSocket:
+    def __init__(self) -> None:
+        self.sent: list[dict[str, object]] = []
+
+    async def send_json(self, payload: dict[str, object]) -> None:
+        self.sent.append(dict(payload))
+
+
+@pytest.mark.parametrize(
+    ("event", "expected"),
+    [
+        ({"type": "signal_turn", "event": "barge_in"}, {"type": "turn.signal", "event": "barge_in"}),
+        ({"type": "close_session"}, {"type": "session.close"}),
+        ({"type": "audio.playback_ack", "played_ms": 1}, {"type": "playback.ack", "played_ms": 1}),
+        ({"type": "input_text.append", "text": "a"}, {"type": "input.text.append", "text": "a"}),
+        ({"type": "push_text", "text": "b"}, {"type": "input.text.append", "text": "b"}),
+        (
+            {"type": "input.audio.append", "audio": "wav"},
+            {"type": "input_audio_buffer.append", "audio": "wav", "format": "wav"},
+        ),
+        (
+            {"type": "push_chunk", "audio": "wav", "format": "pcm_f32le"},
+            {"type": "input_audio_buffer.append", "audio": "wav", "format": "pcm_f32le"},
+        ),
+    ],
+)
+def test_input_aliases_normalize_once_at_mailbox_boundary(event, expected):
+    assert normalize_duplex_input_event(event) == expected
+
+
+def test_actor_uses_one_fifo_mailbox_for_inbound_events():
+    actor = DuplexWebSocketActor(FakeWebSocket())
+
+    assert isinstance(actor.mailbox, asyncio.Queue)
+    assert not isinstance(actor.mailbox, asyncio.PriorityQueue)
+    assert not hasattr(actor, "control_queue")
+    assert not hasattr(actor, "input_queue")
+    assert not hasattr(actor, "event_queue")
+    assert not hasattr(actor, "lifecycle_state")
+    assert not hasattr(actor, "last_response_id")
+    assert not hasattr(actor, "overlap_speech_ms")
+    assert not hasattr(actor, "transition")
+    assert not hasattr(actor, "output_generation_in_flight")
+    assert not hasattr(actor, "session")
+    assert not hasattr(actor, "runtime_opened")
+    assert not hasattr(actor, "runtime_closed")
+    assert not hasattr(actor, "drain_input_queue")
+    assert not hasattr(actor, "control_events_seen")
+    assert not hasattr(actor, "input_events_seen")
+    assert not hasattr(actor, "cancel_count")
+
+
+@pytest.mark.asyncio
+async def test_control_event_does_not_overtake_earlier_audio_input():
+    actor = DuplexWebSocketActor(FakeWebSocket())
+    await actor.enqueue_event({"type": "input_audio_buffer.append", "audio": "pcm"})
+    await actor.enqueue_event({"type": "response.cancel"})
+
+    assert await actor.next_event() == {"type": "input_audio_buffer.append", "audio": "pcm"}
+    assert await actor.next_event() == {"type": "response.cancel"}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("event_type", ["input.cancel", "session.close", "close_session"])
+async def test_terminal_control_preserves_wire_order(event_type: str):
+    actor = DuplexWebSocketActor(FakeWebSocket())
+    await actor.enqueue_event({"type": "input_audio_buffer.append", "audio": "pcm"})
+    await actor.enqueue_event({"type": event_type})
+
+    assert await actor.next_event() == {"type": "input_audio_buffer.append", "audio": "pcm"}
+    assert await actor.next_event() == {"type": event_type}
+
+
+@pytest.mark.asyncio
+async def test_mailbox_delivers_every_enqueued_event_exactly_once():
+    actor = DuplexWebSocketActor(FakeWebSocket())
+    events = [
+        {"type": "input_audio_buffer.append", "audio": str(index)}
+        if index % 2
+        else {"type": "response.created", "response_id": f"resp-{index}"}
+        for index in range(20)
+    ]
+
+    await asyncio.gather(*(actor.enqueue_event(dict(event)) for event in events))
+    received = [await actor.next_event() for _ in events]
+
+    assert sorted(received, key=repr) == sorted(events, key=repr)
+
+
+@pytest.mark.asyncio
+async def test_writer_is_single_owner_of_websocket_send():
+    websocket = FakeWebSocket()
+    actor = DuplexWebSocketActor(websocket)
+    writer = asyncio.create_task(actor.writer_loop())
+
+    await actor.send_json({"type": "one"})
+    await actor.send_json({"type": "two"})
+    await actor.close_writer()
+    await writer
+
+    assert websocket.sent == [{"type": "one"}, {"type": "two"}]
+
+
+@pytest.mark.asyncio
+async def test_stale_fence_payload_is_dropped_before_websocket_send():
+    websocket = FakeWebSocket()
+    actor = DuplexWebSocketActor(websocket, current_epoch=lambda: 2)
+    writer = asyncio.create_task(actor.writer_loop())
+
+    await actor.send_json({"type": "response.audio.delta", "epoch": 1})
+    await actor.send_json({"type": "response.audio.delta", "epoch": 2})
+    await actor.close_writer()
+    await writer
+
+    assert websocket.sent == [{"type": "response.audio.delta", "epoch": 2}]
+    assert actor.stale_output_dropped == 1
+
+
+@pytest.mark.asyncio
+async def test_writer_does_not_revoke_accepted_terminal_after_close_starts():
+    websocket = FakeWebSocket()
+    actor = DuplexWebSocketActor(websocket, current_epoch=lambda: 1)
+    await actor.send_json({"type": "response.done", "epoch": 1, "response_id": "resp-1"})
+    actor.closing = True
+    await actor.close_writer()
+
+    await actor.writer_loop()
+
+    assert websocket.sent == [{"type": "response.done", "epoch": 1, "response_id": "resp-1"}]
+    assert actor.stale_output_dropped == 0

tests/e2e/features/fullduplex/test_client.py

@@ -0,0 +1,240 @@
+import base64
+import wave
+from urllib.parse import parse_qs, urlsplit
+
+import pytest
+
+from vllm_omni.experimental.fullduplex.client import (
+    RealtimeDuplexClient,
+    RealtimeEventCollector,
+    build_realtime_url,
+    read_pcm16_wav,
+    write_pcm16_wav,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_realtime_client_builds_explicit_native_duplex_url():
+    url = build_realtime_url(
+        "ws://localhost:8099/v1/realtime?custom=1",
+        "openbmb/MiniCPM-o-4_5",
+        session_id="session-a",
+    )
+
+    query = parse_qs(urlsplit(url).query)
+    assert query == {
+        "custom": ["1"],
+        "duplex": ["1"],
+        "model": ["openbmb/MiniCPM-o-4_5"],
+        "minicpmo45_native_duplex": ["1"],
+        "session_id": ["session-a"],
+    }
+
+
+def test_realtime_client_builds_resume_only_url_when_autostart_disabled():
+    url = build_realtime_url(
+        "ws://localhost:8099/v1/realtime?duplex=1",
+        "openbmb/MiniCPM-o-4_5",
+        autostart=False,
+    )
+
+    query = parse_qs(urlsplit(url).query)
+    assert query["autostart"] == ["0"]
+    assert query["minicpmo45_native_duplex"] == ["1"]
+
+
+@pytest.mark.asyncio
+async def test_realtime_client_configure_omits_ref_audio_by_default():
+    class Client(RealtimeDuplexClient):
+        def __init__(self):
+            super().__init__("ws://unused")
+            self.sent = []
+
+        async def send(self, event):
+            self.sent.append(event)
+            self.events.add({"type": "session.created"})
+
+    client = Client()
+
+    await client.configure("openbmb/MiniCPM-o-4_5", timeout_s=1)
+
+    session = client.sent[0]["session"]
+    assert "ref_audio" not in session
+
+
+@pytest.mark.asyncio
+async def test_realtime_client_configure_sends_explicit_ref_audio():
+    class Client(RealtimeDuplexClient):
+        def __init__(self):
+            super().__init__("ws://unused")
+            self.sent = []
+
+        async def send(self, event):
+            self.sent.append(event)
+            self.events.add({"type": "session.created"})
+
+    client = Client()
+
+    await client.configure(
+        "openbmb/MiniCPM-o-4_5",
+        ref_audio="data:audio/wav;base64,AAAA",
+        timeout_s=1,
+    )
+
+    session = client.sent[0]["session"]
+    assert session["ref_audio"] == "data:audio/wav;base64,AAAA"
+
+
+def test_realtime_event_collector_partitions_audio_by_response():
+    collector = RealtimeEventCollector()
+    collector.add({"type": "response.created", "response": {"id": "resp-a"}})
+    collector.add(
+        {
+            "type": "response.audio.delta",
+            "response_id": "resp-a",
+            "delta": base64.b64encode(b"audio-a").decode("ascii"),
+            "sample_rate_hz": 16_000,
+        }
+    )
+
+    assert collector.response_ids == ["resp-a"]
+    assert collector.audio_bytes("resp-a") == b"audio-a"
+    assert collector.output_sample_rate_hz == 16_000
+    assert collector.first_received_at("response.created") is not None
+    assert collector.last_received_at("response.audio.delta") is not None
+
+
+def test_realtime_event_collector_reports_engine_token_and_audio_intervals():
+    collector = RealtimeEventCollector()
+    collector.add(
+        {"type": "response.created", "response": {"id": "resp-a"}},
+        received_at_s=10.0,
+    )
+    stage_metrics = {
+        "0": {
+            "num_tokens_out": 4,
+            "vllm_ttft_ms": 120.0,
+            "vllm_tpot_ms": 15.0,
+            "vllm_itl_ms": 14.0,
+            "vllm_itls_ms": [10.0, 14.0, 18.0],
+        }
+    }
+    for received_at_s, cumulative_audio_ms in ((10.2, 80), (10.25, 160), (10.36, 240)):
+        collector.add(
+            {
+                "type": "response.audio.delta",
+                "response_id": "resp-a",
+                "delta": base64.b64encode(b"audio").decode("ascii"),
+                "sample_rate_hz": 16_000,
+                "metadata": {
+                    "audio_duration_ms": cumulative_audio_ms,
+                    "vllm_omni": {"stage_metrics": stage_metrics},
+                },
+            },
+            received_at_s=received_at_s,
+        )
+
+    timing = collector.timing_summary(
+        after_s=10.0,
+        input_committed_at_s=9.9,
+        response_id="resp-a",
+    )
+
+    assert timing["stage0_tokens"] == {
+        "source": "engine_stage_metrics",
+        "output_token_count": 4,
+        "ttft_ms": 120.0,
+        "tpot_ms": 15.0,
+        "inter_token_interval_ms": {
+            "count": 3,
+            "mean": 14.0,
+            "p50": 14.0,
+            "p95": 18.0,
+            "max": 18.0,
+        },
+    }
+    assert timing["audio_output"] == {
+        "source": "client_monotonic_receive",
+        "chunk_count": 3,
+        "response_created_to_first_audio_ms": 200.0,
+        "commit_to_first_audio_ms": 300.0,
+        "inter_chunk_interval_ms": {
+            "count": 2,
+            "mean": 80.0,
+            "p50": 50.0,
+            "p95": 110.0,
+            "max": 110.0,
+        },
+        "chunk_duration_ms": {
+            "count": 3,
+            "mean": 80.0,
+            "p50": 80.0,
+            "p95": 80.0,
+            "max": 80.0,
+        },
+        "max_chunk_gap_ms": 110.0,
+    }
+
+
+def test_response_timing_ignores_unowned_session_level_metrics():
+    collector = RealtimeEventCollector()
+    collector.add(
+        {"type": "response.created", "response": {"id": "resp-a"}},
+        received_at_s=10.0,
+    )
+    collector.add(
+        {
+            "type": "response.audio.delta",
+            "response_id": "resp-a",
+            "delta": base64.b64encode(b"audio").decode("ascii"),
+            "metadata": {
+                "vllm_omni": {
+                    "stage_metrics": {
+                        "0": {
+                            "num_tokens_out": 20,
+                            "vllm_ttft_ms": 157.0,
+                            "vllm_tpot_ms": 16.0,
+                            "vllm_itls_ms": [15.0, 17.0],
+                        }
+                    }
+                }
+            },
+        },
+        received_at_s=10.2,
+    )
+    collector.add(
+        {
+            "type": "response.listen",
+            "metadata": {
+                "vllm_omni": {
+                    "stage_metrics": {
+                        "0": {
+                            "num_tokens_out": 2,
+                            "vllm_ttft_ms": 106.0,
+                            "vllm_tpot_ms": 0.0,
+                            "vllm_itls_ms": [],
+                        }
+                    }
+                }
+            },
+        },
+        received_at_s=10.3,
+    )
+
+    timing = collector.timing_summary(after_s=10.0, response_id="resp-a")
+
+    assert timing["stage0_tokens"]["output_token_count"] == 20
+    assert timing["stage0_tokens"]["ttft_ms"] == 157.0
+
+
+def test_realtime_client_pcm16_wav_round_trip(tmp_path):
+    path = tmp_path / "audio.wav"
+    pcm16 = b"\x01\x00\x02\x00"
+
+    write_pcm16_wav(path, pcm16, sample_rate_hz=16_000)
+
+    with wave.open(str(path), "rb") as wav_file:
+        assert wav_file.getnchannels() == 1
+        assert wav_file.getframerate() == 16_000
+    assert read_pcm16_wav(path) == pcm16

tests/e2e/features/fullduplex/test_joyvl_brain.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm_omni.experimental.fullduplex.joyvl.decision.output_parser import parse_action
+from vllm_omni.experimental.fullduplex.joyvl.memory.brain import InteractionBrain
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeSummarizer:
+    def __init__(self):
+        self.chunks = 0
+
+    async def summarize_chunk(self, chunk_index, frame_range, frames):
+        self.chunks += 1
+        return f"summary-of-chunk-{chunk_index}"
+
+    async def compress_to_long_term(self, mids):
+        return f"compressed({len(mids)})"
+
+
+class _FakeDelegation:
+    def __init__(self):
+        self._ready = {}
+
+    async def submit(self, question, note, frames):
+        self._ready["t1"] = question
+        return "t1"
+
+    async def poll(self, task_id):
+        from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import DelegationResult
+
+        return DelegationResult(task_id, "ready", digest=f"answer to {self._ready[task_id]}")
+
+
+def test_query_freshness_and_archive_on_change():
+    b = InteractionBrain()
+    assert b.update_query("count bottles") is True
+    assert b.update_query("count bottles") is False
+    b.record_response("1 bottle")
+    assert b.update_query("now describe") is True
+    assert b.memory.qa_history[-1].query == "count bottles"
+    assert b.memory.qa_history[-1].responses == [("0.0 seconds", "1 bottle")]
+
+
+def test_query_in_frame_on_arrival_then_prefix_on_later_chunks():
+    b = InteractionBrain(frame_seconds=1.0)
+    b.tick(3)
+    b.update_query("alert me if a fire breaks out")
+    assert b.query_in_current_chunk is True
+    assert "alert me if a fire breaks out" not in b.build_prefix()
+    b.record_response("a fire is breaking out")
+    b.close_chunk()
+    assert b.query_in_current_chunk is False
+    prefix = b.build_prefix()
+    assert "alert me if a fire breaks out" in prefix
+    assert "a fire is breaking out" in prefix
+
+
+@pytest.mark.asyncio
+async def test_flush_archives_qa_and_summarizes():
+    summ = _FakeSummarizer()
+    b = InteractionBrain(summarizer=summ, chunk_frames=3, long_term_every_n_chunks=2, frame_seconds=1.0)
+    b.update_query("watch the room")
+    b.tick(3)
+    assert b.should_flush() is True
+    b.record_response("a person enters")
+    await b.flush([("0.0s", "u0"), ("2.0s", "u2")])
+    assert summ.chunks == 1
+    assert b.chunk_index == 2
+    assert b.should_flush() is False
+    assert b.memory.qa_history[-1].responses == [("2.0 seconds", "a person enters")]
+    assert len(b.memory.mid_term_summaries) == 1
+
+
+@pytest.mark.asyncio
+async def test_long_term_compression_after_n_chunks():
+    summ = _FakeSummarizer()
+    b = InteractionBrain(summarizer=summ, chunk_frames=1, long_term_every_n_chunks=2, frame_seconds=1.0)
+    for _ in range(2):
+        await b.flush([("0.0s", "u")])
+    assert b.memory.long_term_memory.startswith("compressed(")
+    assert b.memory.mid_term_summaries == []
+
+
+@pytest.mark.asyncio
+async def test_long_term_window_caps_blocks():
+    # Long-term memory must not grow unbounded over hours: oldest blocks are dropped
+    # once the sliding window is full (reference window = 15; here 2 for the test).
+    summ = _FakeSummarizer()
+    b = InteractionBrain(
+        summarizer=summ, chunk_frames=1, long_term_every_n_chunks=1, long_term_window=2, frame_seconds=1.0
+    )
+    for _ in range(4):
+        await b.flush([("0.0s", "u")])
+    assert len(b.memory.long_term_blocks) == 2  # capped, not 4
+    assert b.memory.long_term_memory.count("compressed(") == 2
+
+
+@pytest.mark.asyncio
+async def test_delegation_submit_and_fold():
+    b = InteractionBrain(delegation=_FakeDelegation())
+    action = parse_action("</response> hold on <delegation> solve this")
+    info = await b.submit_delegation(action, [])
+    assert info["status"] == "submitted"
+    folded = await b.fold_delegations()
+    assert folded["status"] == "ready"
+    assert b.memory.qa_history[-1].responses[0][1] == "answer to solve this"

tests/e2e/features/fullduplex/test_joyvl_policy.py

@@ -0,0 +1,282 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+
+import pytest
+
+from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import (
+    DelegationResult,
+    ImageEditDelegationBridge,
+    ImageGenDelegationBridge,
+    OpenAIDelegationBridge,
+)
+from vllm_omni.experimental.fullduplex.joyvl.decision.policy import JoyVLPolicy, sample_frames
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_sample_frames_keeps_recent():
+    assert sample_frames(["a", "b"], 4) == ["a", "b"]
+    out = sample_frames([str(i) for i in range(10)], 4)
+    assert out[-1] == "9" and len(out) == 4
+
+
+def test_build_messages_stable_head_and_query():
+    p = JoyVLPolicy(num_frames=4)
+    p.tick(2)
+    p.set_query("count the bottles")
+    msgs, user = p.build_messages([{"type": "image_url", "image_url": {"url": "x"}}])
+    assert msgs[0]["role"] == "system" and "</silence>" in msgs[0]["content"]
+
+    assert any("count the bottles" in part.get("text", "") for part in user["content"])
+
+
+def test_commit_records_only_spoken():
+    p = JoyVLPolicy()
+    p.set_query("q")
+    assert p.commit("</silence>").action.value == "silence"
+    assert p.brain.response_records == []
+    assert p.commit("</response> hello").action.value == "response"
+    assert p.brain.response_records == [("0.0 seconds", "hello")]
+
+
+def test_commit_dedup_default_is_exact_match_only():
+    p = JoyVLPolicy()  # threshold 1.0 == reference behavior
+    assert p.commit("</response> a red ball appears").action.value == "response"
+    assert p.commit("</response> a red ball appears").action.value == "silence"  # exact repeat dropped
+    assert p.commit("</response> a red ball is appearing").action.value == "response"  # near-dup kept
+
+
+def test_commit_dedup_threshold_drops_near_duplicates():
+    p = JoyVLPolicy(response_dedup_threshold=0.8)
+    assert p.commit("</response> a red ball appears").action.value == "response"
+    assert p.commit("</response> a red ball appears!").action.value == "silence"  # near-dup dropped
+    assert p.commit("</response> the sky is now blue").action.value == "response"  # distinct kept
+
+
+def test_commit_dedup_rearms_after_silence():
+    p = JoyVLPolicy(response_dedup_threshold=0.85)
+    assert p.commit("</response> someone raises a hand").action.value == "response"
+    # no gap: an immediate repeat is still suppressed
+    assert p.commit("</response> someone raises a hand").action.value == "silence"
+    # a genuine silent tick (hand lowered) re-arms; the recurrence speaks again
+    assert p.commit("</silence>").action.value == "silence"
+    assert p.commit("</response> someone raises a hand").action.value == "response"
+
+
+class _FakeBrain(OpenAIDelegationBridge):
+    """Override only the network call so submit/poll plumbing is exercised for real."""
+
+    def __init__(self, answer, gate=None):
+        super().__init__("http://unused/v1", "background-model")
+        self._answer = answer
+        self._gate = gate
+
+    async def _complete(self, messages):
+        if self._gate is not None:
+            await self._gate.wait()
+        return self._answer
+
+
+@pytest.mark.asyncio
+async def test_openai_delegation_bridge_is_nonblocking_and_resolves():
+    gate = asyncio.Event()
+    bridge = _FakeBrain("the bottle is on the left shelf", gate=gate)
+    frames = [("0.0s", "data:image/jpeg;base64,AAA")]
+    task_id = await bridge.submit("where is the bottle?", "let me check", frames)
+
+    # poll must not block on the slow background brain
+    assert (await bridge.poll(task_id)).status == "pending"
+
+    gate.set()
+    await asyncio.sleep(0)
+    for _ in range(10):
+        result = await bridge.poll(task_id)
+        if result.status != "pending":
+            break
+        await asyncio.sleep(0)
+    assert result.is_ready
+    assert "bottle is on the left shelf" in result.digest
+    assert (await bridge.poll("missing")).status == "error"
+
+
+class _FakeImageBrain(ImageGenDelegationBridge):
+    async def _call_images(self, payload):
+        # echo the prompt back so the test can assert it was passed through
+        return {"data": [{"b64_json": "QUJD"}]}  # "ABC"
+
+
+@pytest.mark.asyncio
+async def test_image_delegation_bridge_returns_media_not_text():
+    bridge = _FakeImageBrain("http://unused/v1")
+    task_id = await bridge.submit("draw a red dragon", "", [])
+    for _ in range(10):
+        result = await bridge.poll(task_id)
+        if result.status != "pending":
+            break
+        await asyncio.sleep(0)
+    assert result.is_ready
+    # the image rides in media as a data URL; digest is only a short text placeholder
+    assert result.media == "data:image/png;base64,QUJD"
+    assert "QUJD" not in result.digest and "red dragon" in result.digest
+
+
+class _FakeEditBrain(ImageEditDelegationBridge):
+    def __init__(self, base_url, model):
+        super().__init__(base_url, model)
+        self.seen_frame = None
+
+    async def _call_chat(self, frame, instruction):
+        self.seen_frame = frame
+        return {"choices": [{"message": {"content": [{"image_url": {"url": "data:image/png;base64,Q0FSVA=="}}]}}]}
+
+
+@pytest.mark.asyncio
+async def test_image_edit_delegation_conditions_on_frame_and_returns_media():
+    bridge = _FakeEditBrain("http://unused/v1", "Qwen-Image-Edit")
+    frames = [("0.0s", "data:image/jpeg;base64,OLD"), ("1.0s", "data:image/jpeg;base64,NOW")]
+    task_id = await bridge.submit("把画面变成卡通风格", "", frames)
+    for _ in range(10):
+        result = await bridge.poll(task_id)
+        if result.status != "pending":
+            break
+        await asyncio.sleep(0)
+    assert result.is_ready
+    assert bridge.seen_frame == "data:image/jpeg;base64,NOW"  # latest frame, not the old one
+    assert result.media == "data:image/png;base64,Q0FSVA=="
+    assert result.media not in result.digest  # image rides in media, not context
+
+
+@pytest.mark.asyncio
+async def test_image_edit_delegation_errors_without_frame():
+    bridge = _FakeEditBrain("http://unused/v1", "Qwen-Image-Edit")
+    task_id = await bridge.submit("cartoonify", "", [])
+    for _ in range(10):
+        result = await bridge.poll(task_id)
+        if result.status != "pending":
+            break
+        await asyncio.sleep(0)
+    assert result.status == "error"
+
+
+class _NamedBridge:
+    def __init__(self, name):
+        self.name = name
+
+    async def submit(self, question, note, frames):
+        return f"{self.name}-1"
+
+    async def poll(self, task_id):
+        return DelegationResult(task_id, "ready", digest=self.name)
+
+
+@pytest.mark.asyncio
+async def test_routing_delegation_dispatches_by_request():
+    from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import RoutingDelegationBridge
+
+    chat, image, edit = _NamedBridge("chat"), _NamedBridge("image"), _NamedBridge("edit")
+    r = RoutingDelegationBridge(chat=chat, image=image, edit=edit)
+
+    async def route(q):
+        tid = await r.submit(q, "", [])
+        return (await r.poll(tid)).digest
+
+    assert await route("把你看到的画面变成卡通风格") == "edit"  # restyle current view
+    assert await route("画一只红色的龙") == "image"  # generate a new image
+    assert await route("这道数学题的答案是什么") == "chat"  # hard question -> chat brain
+
+
+class _Delegation:
+    async def submit(self, question, note, frames):
+        return "t1"
+
+    async def poll(self, task_id):
+        return DelegationResult(task_id, "ready", digest="background answer")
+
+
+def test_build_delegation_off_unless_backend_configured():
+    from vllm_omni.experimental.fullduplex.joyvl.serving.config import InteractionConfig
+    from vllm_omni.experimental.fullduplex.joyvl.serving.server import SessionManager
+
+    # No backend URL -> delegation disabled (None), NOT a silent stub that folds fake answers.
+    assert SessionManager._build_delegation(InteractionConfig()) is None
+    assert SessionManager._build_delegation(InteractionConfig(delegation_kind="router")) is None
+    # Stub only on explicit opt-in.
+    assert SessionManager._build_delegation(InteractionConfig(delegation_kind="stub")) is not None
+    # A real chat backend -> a live bridge.
+    cfg = InteractionConfig(delegation_backend_url="http://x/v1", delegation_model="m")
+    assert SessionManager._build_delegation(cfg) is not None
+
+
+@pytest.mark.asyncio
+async def test_delegation_cancel_drops_task():
+    from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import StubDelegationBridge
+
+    b = StubDelegationBridge(ready_after_ticks=1)
+    task_id = await b.submit("q", "", [])
+    b.cancel(task_id)
+    # after cancel the task is gone -> poll reports it as unknown rather than leaking
+    assert (await b.poll(task_id)).status == "error"
+
+
+@pytest.mark.asyncio
+async def test_delegation_aclose_clears_pending():
+    from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import StubDelegationBridge
+
+    b = StubDelegationBridge(ready_after_ticks=5)
+    task_id = await b.submit("q", "", [])
+    await b.aclose()
+    assert (await b.poll(task_id)).status == "error"
+
+
+@pytest.mark.asyncio
+async def test_session_reset_cancels_and_awaits_consolidation():
+    from vllm_omni.experimental.fullduplex.joyvl.serving.config import InteractionConfig
+    from vllm_omni.experimental.fullduplex.joyvl.serving.session import InteractionSession
+
+    class _Backend:
+        async def generate(self, *a, **k):
+            return "</silence>", None
+
+        async def aclose(self):
+            pass
+
+    class _Summarizer:
+        async def summarize_chunk(self, *a, **k):
+            await asyncio.sleep(100)  # hang so the consolidation task is in-flight
+            return "x"
+
+        async def compress_to_long_term(self, *a, **k):
+            return ""
+
+        async def aclose(self):
+            pass
+
+    sess = InteractionSession("s", InteractionConfig(), _Backend(), summarizer=_Summarizer())
+    sess._spawn_consolidation(1, [("0.0s", "u")])
+    assert len(sess._consolidating) == 1
+    await sess.reset()  # must cancel AND await the in-flight task, not just clear the set
+    assert len(sess._consolidating) == 0
+
+
+@pytest.mark.asyncio
+async def test_manager_aclose_closes_without_error():
+    from vllm_omni.experimental.fullduplex.joyvl.serving.config import InteractionConfig
+    from vllm_omni.experimental.fullduplex.joyvl.serving.server import SessionManager
+
+    mgr = SessionManager(InteractionConfig(enable_memory=False, enable_delegation=False))
+    await mgr.aclose()
+    await mgr.aclose()  # idempotent
+
+
+@pytest.mark.asyncio
+async def test_delegation_submit_and_fold():
+    p = JoyVLPolicy(delegation=_Delegation())
+    action = p.commit("</response> hold on <delegation> hard question")
+    assert action.action.value == "delegate"
+    info = await p.submit_if_delegate(action)
+    assert info["status"] == "submitted"
+    folded = await p.fold_delegations()
+    assert folded["status"] == "ready"
+    assert p.brain.memory.qa_history[-1].responses[0][1] == "background answer"

tests/e2e/features/fullduplex/test_runtime.py

@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+from collections.abc import AsyncIterator
+
+import pytest
+
+from vllm_omni.experimental.fullduplex.core import protocol as ev
+from vllm_omni.experimental.fullduplex.core.adapter import DuplexCapability, OutputChunk
+from vllm_omni.experimental.fullduplex.core.runtime import DuplexRuntime
+from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexSessionConfig
+from vllm_omni.experimental.fullduplex.joyvl.adapter import JoyVLDuplexAdapter
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+async def _feed(events):
+    for e in events:
+        yield e
+
+
+def _collector():
+    out: list[dict] = []
+
+    async def emit(event: dict) -> None:
+        out.append(event)
+
+    return out, emit
+
+
+class _FakeAdapter:
+    def __init__(self, chunks, barge_after=None):
+        self._chunks = chunks
+        self._barge_after = barge_after
+
+    def capabilities(self):
+        return DuplexCapability(frozenset({"text"}), frozenset({"text"}), proactive=False)
+
+    async def on_input(self, session, modality, data):
+        pass
+
+    def should_respond(self, session):
+        return True
+
+    async def respond(self, session) -> AsyncIterator[OutputChunk]:
+        for i, c in enumerate(self._chunks):
+            if self._barge_after is not None and i == self._barge_after:
+                session.barge_in()
+            yield OutputChunk("text", c)
+
+    async def on_barge_in(self, session):
+        pass
+
+    async def on_playback_ack(self, session, cursor):
+        pass
+
+
+@pytest.mark.asyncio
+async def test_runtime_basic_response():
+    session = DuplexSession("s", DuplexSessionConfig(output_modalities=("text",)))
+    rt = DuplexRuntime(session, _FakeAdapter(["a", "b"]))
+    out, emit = _collector()
+    await rt.run(_feed([{"type": ev.INPUT_COMMIT}, {"type": ev.CLOSE}]), emit)
+    types = [e["type"] for e in out]
+    assert types == [ev.RESPONSE_CREATED, ev.RESPONSE_DELTA, ev.RESPONSE_DELTA, ev.RESPONSE_DONE]
+    assert [e["data"] for e in out if e["type"] == ev.RESPONSE_DELTA] == ["a", "b"]
+
+
+@pytest.mark.asyncio
+async def test_runtime_barge_in_drops_stale_output():
+    session = DuplexSession("s")
+    rt = DuplexRuntime(session, _FakeAdapter(["a", "b", "c"], barge_after=1))
+    out, emit = _collector()
+    await rt.run(_feed([{"type": ev.INPUT_COMMIT}, {"type": ev.CLOSE}]), emit)
+    data = [e["data"] for e in out if e["type"] == ev.RESPONSE_DELTA]
+    assert data == ["a"]
+    assert ev.RESPONSE_DONE not in [e["type"] for e in out]
+
+
+class _SlowAdapter:
+    def capabilities(self):
+        return DuplexCapability(frozenset({"text"}), frozenset({"text"}), proactive=False)
+
+    async def on_input(self, session, modality, data):
+        pass
+
+    def should_respond(self, session):
+        return True
+
+    async def respond(self, session) -> AsyncIterator[OutputChunk]:
+        for c in ["a", "b", "c", "d", "e"]:
+            await asyncio.sleep(0.02)
+            yield OutputChunk("text", c)
+
+    async def on_barge_in(self, session):
+        pass
+
+    async def on_playback_ack(self, session, cursor):
+        pass
+
+
+@pytest.mark.asyncio
+async def test_runtime_cancel_event_interrupts_active_response():
+    session = DuplexSession("s")
+    rt = DuplexRuntime(session, _SlowAdapter())
+    out, emit = _collector()
+
+    async def feed():
+        yield {"type": ev.INPUT_COMMIT}
+        await asyncio.sleep(0.03)
+        yield {"type": ev.RESPONSE_CANCEL}
+        yield {"type": ev.CLOSE}
+
+    await rt.run(feed(), emit)
+    deltas = [e["data"] for e in out if e["type"] == ev.RESPONSE_DELTA]
+    types = [e["type"] for e in out]
+    assert len(deltas) < 5
+    assert ev.RESPONSE_CANCELLED in types
+    assert ev.RESPONSE_DONE not in types
+
+
+@pytest.mark.asyncio
+async def test_runtime_new_response_supersedes_inflight_without_blocking():
+    # A second response trigger arriving while the first is still streaming must
+    # supersede it (barge-in) rather than block the input loop until it finishes.
+    session = DuplexSession("s")
+    rt = DuplexRuntime(session, _SlowAdapter())
+    out, emit = _collector()
+
+    async def feed():
+        yield {"type": ev.INPUT_COMMIT}
+        await asyncio.sleep(0.03)  # let the first response emit ~one chunk
+        yield {"type": ev.INPUT_COMMIT}  # supersedes the in-flight first response
+        await asyncio.sleep(0.2)  # let the second response finish
+        yield {"type": ev.CLOSE}
+
+    await rt.run(feed(), emit)
+    created = [e for e in out if e["type"] == ev.RESPONSE_CREATED]
+    done = [e for e in out if e["type"] == ev.RESPONSE_DONE]
+    # two responses created; only the latest one completes (the first was cancelled)
+    assert len(created) == 2
+    assert len(done) == 1
+    assert done[0]["response_index"] == created[-1]["response_index"]
+
+
+@pytest.mark.asyncio
+async def test_joyvl_adapter_speaks_then_silences():
+    replies = iter(["</response> a fire is breaking out", "</silence>"])
+
+    async def fake_generate(messages):
+        return next(replies)
+
+    adapter = JoyVLDuplexAdapter(fake_generate, num_frames=4)
+    cfg = DuplexSessionConfig(input_modalities=("video", "text"), output_modalities=("text",), proactive=True)
+    rt = DuplexRuntime(DuplexSession("s", cfg), adapter)
+    out, emit = _collector()
+    await rt.run(
+        _feed(
+            [
+                {"type": ev.INPUT_APPEND, "modality": "text", "data": "alert me if a fire breaks out"},
+                {"type": ev.INPUT_APPEND, "modality": "video", "data": "data:image/jpeg;base64,AAA"},
+                {"type": ev.INPUT_APPEND, "modality": "video", "data": "data:image/jpeg;base64,BBB"},
+                {"type": ev.CLOSE},
+            ]
+        ),
+        emit,
+    )
+    deltas = [e["data"] for e in out if e["type"] == ev.RESPONSE_DELTA]
+    assert "a fire is breaking out" in deltas
+
+    assert deltas.count("a fire is breaking out") == 1

tests/e2e/features/helpers/custom_pipeline.py

@@ -0,0 +1,688 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Test doubles for custom-pipeline E2E and RLHF integration.
+
+Three layers:
+
+* ``FlowMatchSDEDiscreteSchedulerForTest`` — SDE scheduler with per-step logprobs.
+* ``vLLMOmniColocateWorkerExtensionForTest`` — verl-omni colocate worker extension mirror.
+* ``QwenImagePipelineWithLogProbForTest`` — Qwen-Image pipeline returning trajectory_* fields.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import logging
+import math
+import os
+import platform
+import signal
+from dataclasses import dataclass
+from typing import Any, Literal
+
+import torch
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import BaseOutput
+from diffusers.utils.torch_utils import randn_tensor
+from msgspec import field
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.models.qwen_image import QwenImagePipeline
+from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.lora.request import LoRARequest as OmniLoRARequest
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Scheduler layer
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class FlowMatchSDEDiscreteSchedulerOutputForTest(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+        log_prob (`torch.FloatTensor` of shape `(batch_size,)`, *optional*):
+            The log probability of the previous sample.
+        prev_sample_mean (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images):
+            The mean of the computed sample of previous timestep.
+        std_dev_t (`torch.FloatTensor` of shape `(batch_size, 1, 1)`):
+            The standard deviation used to compute `prev_sample`.
+    """
+
+    prev_sample: torch.FloatTensor
+    log_prob: torch.FloatTensor | None
+    prev_sample_mean: torch.FloatTensor
+    std_dev_t: torch.FloatTensor
+
+
+class FlowMatchSDEDiscreteSchedulerForTest(FlowMatchEulerDiscreteScheduler):
+    """SDE version of the FlowMatchEulerDiscreteScheduler.
+    The implementation is based on FlowGRPO paper (https://arxiv.org/abs/2505.05470)
+    and diffusers v0.37 branch.
+    """
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: float | torch.FloatTensor,
+        sample: torch.FloatTensor,
+        s_churn: float = 0.0,
+        s_tmin: float = 0.0,
+        s_tmax: float = float("inf"),
+        s_noise: float = 1.0,
+        generator: torch.Generator | None = None,
+        per_token_timesteps: torch.Tensor | None = None,
+        return_dict: bool = True,
+        noise_level: float = 0.7,
+        prev_sample: torch.FloatTensor | None = None,
+        sde_type: Literal["sde", "cps"] = "sde",
+        return_logprobs: bool = True,
+    ) -> FlowMatchSDEDiscreteSchedulerOutputForTest | tuple:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Modified from https://github.com/yifan123/flow_grpo/blob/main/flow_grpo/diffusers_patch/sd3_sde_with_logprob.py
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            s_churn (`float`):
+            s_tmin  (`float`):
+            s_tmax  (`float`):
+            s_noise (`float`, defaults to 1.0):
+                Scaling factor for noise added to the sample.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            per_token_timesteps (`torch.Tensor`, *optional*):
+                The timesteps for each token in the sample.
+            return_dict (`bool`):
+                Whether or not to return a
+                [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchSDEDiscreteSchedulerOutput`] or tuple.
+            noise_level (`float`, *optional*, defaults to 0.7):
+                The noise level used in the SDE.
+            prev_sample (`torch.FloatTensor`, *optional*):
+                The sample from the previous timestep. If not provided, it will be sampled inside the function.
+            sde_type (`str`, *optional*, defaults to "sde"):
+                The type of SDE to use. Choose between "sde" and "cps".
+            return_logprobs (`bool`, *optional*, defaults to True):
+                Whether to return log probabilities of the previous sample.
+        """
+
+        if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
+            raise ValueError(
+                (
+                    "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+                    " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                    " one of the `scheduler.timesteps` as a timestep."
+                ),
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Upcast to avoid precision issues when computing prev_sample
+        sample = sample.to(torch.float32)
+        if prev_sample is not None:
+            prev_sample = prev_sample.to(torch.float32)
+
+        prev_sample, log_prob, prev_sample_mean, std_dev_t = self.sample_previous_step(
+            sample=sample,
+            model_output=model_output,
+            generator=generator,
+            per_token_timesteps=per_token_timesteps,
+            noise_level=noise_level,
+            prev_sample=prev_sample,
+            sde_type=sde_type,
+            return_logprobs=return_logprobs,
+        )
+
+        # upon completion increase step index by one
+        self._step_index += 1
+        if per_token_timesteps is None:
+            # Cast sample back to model compatible dtype
+            prev_sample = prev_sample.to(model_output.dtype)
+
+        if not return_dict:
+            return (prev_sample, log_prob, prev_sample_mean, std_dev_t)
+
+        return FlowMatchSDEDiscreteSchedulerOutputForTest(
+            prev_sample=prev_sample, log_prob=log_prob, prev_sample_mean=prev_sample_mean, std_dev_t=std_dev_t
+        )
+
+    def sample_previous_step(
+        self,
+        sample: torch.Tensor,
+        model_output: torch.Tensor,
+        timestep: torch.FloatTensor | None = None,
+        generator: torch.Generator | None = None,
+        per_token_timesteps: torch.Tensor | None = None,
+        noise_level: float = 0.7,
+        prev_sample: torch.Tensor | None = None,
+        sde_type: Literal["cps", "sde"] = "sde",
+        return_logprobs: bool = True,
+    ):
+        assert sde_type in ["sde", "cps"]
+        assert sample.dtype == torch.float32
+        if prev_sample is not None:
+            assert prev_sample.dtype == torch.float32
+
+        if per_token_timesteps is not None:
+            raise NotImplementedError("per_token_timesteps is not supported yet for FlowMatchSDEDiscreteScheduler.")
+        else:
+            if timestep is None:
+                sigma_idx = self.step_index
+                sigma = self.sigmas[sigma_idx]
+                sigma_prev = self.sigmas[sigma_idx + 1]
+            else:
+                sigma_idx = torch.tensor([self.index_for_timestep(t) for t in timestep])
+                sigma = self.sigmas[sigma_idx].view(-1, *([1] * (len(sample.shape) - 1)))
+                sigma_prev = self.sigmas[sigma_idx + 1].view(-1, *([1] * (len(sample.shape) - 1)))
+
+            sigma_max = self.sigmas[1]
+            dt = sigma_prev - sigma
+
+        if sde_type == "sde":
+            std_dev_t = torch.sqrt(sigma / (1 - torch.where(sigma == 1, sigma_max, sigma))) * noise_level
+
+            prev_sample_mean = (
+                sample * (1 + std_dev_t**2 / (2 * sigma) * dt)
+                + model_output * (1 + std_dev_t**2 * (1 - sigma) / (2 * sigma)) * dt
+            )
+
+            if prev_sample is None:
+                variance_noise = randn_tensor(
+                    model_output.shape,
+                    generator=generator,
+                    device=model_output.device,
+                    dtype=model_output.dtype,
+                )
+                prev_sample = prev_sample_mean + std_dev_t * torch.sqrt(-1 * dt) * variance_noise
+
+            if return_logprobs:
+                log_prob = (
+                    -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * ((std_dev_t * torch.sqrt(-1 * dt)) ** 2))
+                    - torch.log(std_dev_t * torch.sqrt(-1 * dt))
+                    - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi)))
+                )
+            else:
+                log_prob = None
+
+        elif sde_type == "cps":
+            std_dev_t = sigma_prev * math.sin(noise_level * math.pi / 2)
+            pred_original_sample = sample - sigma * model_output
+            noise_estimate = sample + model_output * (1 - sigma)
+            prev_sample_mean = pred_original_sample * (1 - sigma_prev) + noise_estimate * torch.sqrt(
+                sigma_prev**2 - std_dev_t**2
+            )
+
+            if prev_sample is None:
+                variance_noise = randn_tensor(
+                    model_output.shape,
+                    generator=generator,
+                    device=model_output.device,
+                    dtype=model_output.dtype,
+                )
+                prev_sample = prev_sample_mean + std_dev_t * variance_noise
+
+            if return_logprobs:
+                log_prob = -((prev_sample.detach() - prev_sample_mean) ** 2)
+            else:
+                log_prob = None
+
+        # mean along all but batch dimension
+        log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) if log_prob is not None else None
+        return prev_sample, log_prob, prev_sample_mean, std_dev_t
+
+
+# ---------------------------------------------------------------------------
+# Worker extension layer
+# ---------------------------------------------------------------------------
+
+
+class OmniTensorLoRARequestForTest(OmniLoRARequest):
+    peft_config: dict = field(default=None)
+    lora_tensors: dict = field(default=None)
+
+
+def set_death_signal() -> None:
+    """Verbatim port of ``verl.workers.rollout.vllm_rollout.utils.set_death_signal``."""
+    if platform.system() != "Linux":
+        return
+    try:
+        libc = ctypes.CDLL("libc.so.6")
+        libc.prctl(1, signal.SIGKILL)
+        if os.getppid() == 1:
+            os.kill(os.getpid(), signal.SIGKILL)
+    except Exception:  # noqa: BLE001
+        # libc.so.6 may not be present (musl, alpine) — best-effort only.
+        pass
+
+
+class vLLMOmniColocateWorkerExtensionForTest(CustomPipelineWorkerExtension):
+    """Mirror of ``vLLMOmniColocateWorkerExtension`` (verl-omni).
+
+    The production ``__new__`` runs ``set_death_signal`` + ``VLLMOmniHijack.hijack``
+    on every vLLM worker process. Replicating both ensures this test
+    reproduces the same monkey-patched environment.
+    """
+
+    def __new__(cls, **kwargs):
+        set_death_signal()
+        return super().__new__(cls)
+
+    @staticmethod
+    def test_extension_name() -> str:
+        """Return a stable identifier for assertions in unit tests."""
+        return "vllm-omni-colocate-worker-extension-for-test"
+
+
+# ---------------------------------------------------------------------------
+# Pipeline layer
+# ---------------------------------------------------------------------------
+
+
+def _maybe_to_cpu(v):
+    if isinstance(v, torch.Tensor):
+        return v.detach().cpu()
+    return v
+
+
+class QwenImagePipelineWithLogProbForTest(QwenImagePipeline):
+    supports_request_batch = False
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__(od_config=od_config, prefix=prefix)
+        self.device = get_local_device()
+        model = od_config.model
+        # Check if model is a local path
+        local_files_only = os.path.exists(model)
+
+        self.scheduler = FlowMatchSDEDiscreteSchedulerForTest.from_pretrained(
+            model, subfolder="scheduler", local_files_only=local_files_only
+        )
+
+    def _get_qwen_prompt_embeds(
+        self,
+        prompt_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        dtype: torch.dtype | None = None,
+    ):
+        dtype = dtype or self.text_encoder.dtype
+
+        if attention_mask is None:
+            attention_mask = torch.ones_like(prompt_ids, dtype=torch.long)
+
+        prompt_ids = prompt_ids.unsqueeze(0) if prompt_ids.ndim == 1 else prompt_ids
+        attention_mask = attention_mask.unsqueeze(0) if attention_mask.ndim == 1 else attention_mask
+        drop_idx = self.prompt_template_encode_start_idx
+        encoder_hidden_states = self.text_encoder(
+            input_ids=prompt_ids.to(self.device),
+            attention_mask=attention_mask.to(self.device),
+            output_hidden_states=True,
+        )
+        hidden_states = encoder_hidden_states.hidden_states[-1]
+        split_hidden_states = self._extract_masked_hidden(hidden_states, attention_mask)
+        split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
+        attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states]
+        max_seq_len = max([e.size(0) for e in split_hidden_states])
+        prompt_embeds = torch.stack(
+            [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]
+        )
+        encoder_attention_mask = torch.stack(
+            [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]
+        )
+
+        prompt_embeds = prompt_embeds.to(dtype=dtype)
+
+        return prompt_embeds, encoder_attention_mask
+
+    def encode_prompt(
+        self,
+        prompt_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        num_images_per_prompt: int = 1,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_embeds_mask: torch.Tensor | None = None,
+        max_sequence_length: int = 1024,
+    ):
+        prompt_ids = prompt_ids.unsqueeze(0) if prompt_ids.ndim == 1 else prompt_ids
+        attention_mask = (
+            attention_mask.unsqueeze(0) if attention_mask is not None and attention_mask.ndim == 1 else attention_mask
+        )
+        batch_size = prompt_ids.shape[0] if prompt_embeds is None else prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt_ids, attention_mask=attention_mask)
+
+        prompt_embeds = prompt_embeds[:, :max_sequence_length]
+        prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
+
+        _, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+        prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+
+        return prompt_embeds, prompt_embeds_mask
+
+    def diffuse(
+        self,
+        prompt_embeds,
+        prompt_embeds_mask,
+        negative_prompt_embeds,
+        negative_prompt_embeds_mask,
+        latents,
+        img_shapes,
+        txt_seq_lens,
+        negative_txt_seq_lens,
+        timesteps,
+        do_true_cfg,
+        guidance,
+        true_cfg_scale,
+        noise_level,
+        sde_window,
+        sde_type,
+        generator,
+        logprobs,
+    ):
+        all_latents = []
+        all_log_probs = []
+        all_timesteps = []
+        self.scheduler.set_begin_index(0)
+        for i, t in enumerate(timesteps):
+            if self.interrupt:
+                continue
+
+            if i < sde_window[0]:
+                cur_noise_level = 0.0
+            elif i == sde_window[0]:
+                cur_noise_level = noise_level
+                all_latents.append(latents)
+            elif i > sde_window[0] and i < sde_window[1]:
+                cur_noise_level = noise_level
+            else:
+                cur_noise_level = 0.0
+
+            self._current_timestep = t
+
+            # Broadcast timestep to match batch size
+            timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
+
+            # Forward pass for positive prompt (or unconditional if no CFG)
+            self.transformer.do_true_cfg = do_true_cfg
+            noise_pred = self.transformer(
+                hidden_states=latents,
+                timestep=timestep / 1000,
+                guidance=guidance,
+                encoder_hidden_states_mask=prompt_embeds_mask,
+                encoder_hidden_states=prompt_embeds,
+                img_shapes=img_shapes,
+                txt_seq_lens=txt_seq_lens,
+                attention_kwargs=self.attention_kwargs,
+                return_dict=False,
+            )[0]
+            # Forward pass for negative prompt (CFG)
+            if do_true_cfg:
+                neg_noise_pred = self.transformer(
+                    hidden_states=latents,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    encoder_hidden_states_mask=negative_prompt_embeds_mask,
+                    encoder_hidden_states=negative_prompt_embeds,
+                    img_shapes=img_shapes,
+                    txt_seq_lens=negative_txt_seq_lens,
+                    attention_kwargs=self.attention_kwargs,
+                    return_dict=False,
+                )[0]
+                comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+                cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
+                noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
+                noise_pred = comb_pred * (cond_norm / noise_norm)
+            # compute the previous noisy sample x_t -> x_t-1
+            latents, log_prob, _, _ = self.scheduler.step(
+                noise_pred,
+                t,
+                latents,
+                generator=generator,
+                noise_level=cur_noise_level,
+                sde_type=sde_type,
+                return_logprobs=logprobs,
+                return_dict=False,
+            )
+
+            if i >= sde_window[0] and i < sde_window[1]:
+                all_latents.append(latents)
+                all_log_probs.append(log_prob)
+                all_timesteps.append(t)
+
+        all_latents = torch.stack(all_latents, dim=1)
+
+        if all_log_probs[0] is not None:
+            all_log_probs = torch.stack(all_log_probs, dim=1)
+        else:
+            all_log_probs = None
+
+        all_timesteps = torch.stack(all_timesteps).unsqueeze(0).expand(latents.shape[0], -1)
+
+        return latents, all_latents, all_log_probs, all_timesteps
+
+    def forward(
+        self,
+        req: DiffusionRequestBatch,
+        prompt_ids: torch.Tensor | list[int] | None = None,
+        prompt_mask: torch.Tensor | None = None,
+        negative_prompt_ids: torch.Tensor | list[int] | None = None,
+        negative_prompt_mask: torch.Tensor | None = None,
+        true_cfg_scale: float = 4.0,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int = 50,
+        sigmas: list[float] | None = None,
+        guidance_scale: float = 1.0,
+        num_images_per_prompt: int = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_embeds_mask: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds_mask: torch.Tensor | None = None,
+        output_type: str | None = "pil",
+        attention_kwargs: dict[str, Any] | None = None,
+        callback_on_step_end_tensor_inputs: tuple[str, ...] = ("latents",),
+        max_sequence_length: int = 512,
+        noise_level: float = 0.7,
+        sde_window_size: int | None = None,
+        sde_window_range: tuple[int, int] = (0, 5),
+        sde_type: Literal["sde", "cps"] = "sde",
+        logprobs: bool = True,
+    ) -> DiffusionOutput:
+        if req.num_reqs != 1:
+            raise ValueError("QwenImagePipelineWithLogProbForTest supports only single-request forward.")
+        request = req.requests[0]
+
+        # Extract prompt data from OmniCustomPrompt in req.prompt
+        custom_prompt = request.prompt if request.prompt is not None else {}
+        if isinstance(custom_prompt, dict):
+            prompt_ids = custom_prompt.get("prompt_ids", prompt_ids)
+            prompt_mask = custom_prompt.get("prompt_mask", prompt_mask)
+            negative_prompt_ids = custom_prompt.get("negative_prompt_ids", negative_prompt_ids)
+            negative_prompt_mask = custom_prompt.get("negative_prompt_mask", negative_prompt_mask)
+
+        # Read sampling params from req.sampling_params
+        sp = request.sampling_params
+        height = sp.height or self.default_sample_size * self.vae_scale_factor
+        width = sp.width or self.default_sample_size * self.vae_scale_factor
+        num_inference_steps = sp.num_inference_steps or num_inference_steps
+        max_sequence_length = sp.max_sequence_length or max_sequence_length
+
+        noise_level = sp.extra_args.get("noise_level", None) or noise_level
+        sde_window_size = sp.extra_args.get("sde_window_size", None) or sde_window_size
+        sde_window_range = sp.extra_args.get("sde_window_range", None) or sde_window_range
+        sde_type = sp.extra_args.get("sde_type", None) or sde_type
+        logprobs = sp.extra_args.get("logprobs", None) or logprobs
+
+        generator = sp.generator or generator
+        if generator is None and sp.seed is not None:
+            generator = torch.Generator(device=self.device).manual_seed(sp.seed)
+        true_cfg_scale = sp.true_cfg_scale or true_cfg_scale
+        req_num_outputs = getattr(sp, "num_outputs_per_prompt", None)
+        if req_num_outputs and req_num_outputs > 0:
+            num_images_per_prompt = req_num_outputs
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        if prompt_ids is not None:
+            if isinstance(prompt_ids, list):
+                prompt_ids = torch.tensor(prompt_ids, device=self.device)
+            batch_size = prompt_ids.shape[0] if prompt_ids.ndim == 2 else 1
+        elif prompt_embeds is not None:
+            batch_size = prompt_embeds.shape[0]
+        else:
+            # Both prompt_ids and prompt_embeds are None (e.g. during warmup/dummy run).
+            # Return a minimal dummy output to avoid crashing.
+            return DiffusionOutput(output=None)
+
+        if isinstance(negative_prompt_ids, list):
+            negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device)
+
+        has_neg_prompt = negative_prompt_ids is not None or (
+            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
+        )
+
+        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+            prompt_ids=prompt_ids,
+            attention_mask=prompt_mask,
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+        if do_true_cfg:
+            negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
+                prompt_ids=negative_prompt_ids,
+                attention_mask=negative_prompt_mask,
+                prompt_embeds=negative_prompt_embeds,
+                prompt_embeds_mask=negative_prompt_embeds_mask,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+            )
+
+        num_channels_latents = self.transformer.in_channels // 4
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            self.device,
+            generator,
+            latents,
+        )
+        img_shapes = [[(1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2)]] * batch_size
+
+        timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1])
+        # num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # handle guidance
+        if self.transformer.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0])
+        else:
+            guidance = None
+
+        if self.attention_kwargs is None:
+            self._attention_kwargs = {}
+
+        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
+        negative_txt_seq_lens = (
+            negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None
+        )
+
+        if sde_window_size is not None:
+            start = torch.randint(
+                sde_window_range[0],
+                sde_window_range[1] - sde_window_size + 1,
+                (1,),
+                generator=generator,
+                device=self.device,
+            ).item()
+            end = start + sde_window_size
+            sde_window = (start, end)
+        else:
+            sde_window = (0, len(timesteps) - 1)
+
+        latents, all_latents, all_log_probs, all_timesteps = self.diffuse(
+            prompt_embeds,
+            prompt_embeds_mask,
+            negative_prompt_embeds,
+            negative_prompt_embeds_mask,
+            latents,
+            img_shapes,
+            txt_seq_lens,
+            negative_txt_seq_lens,
+            timesteps,
+            do_true_cfg,
+            guidance,
+            true_cfg_scale,
+            noise_level,
+            sde_window,
+            sde_type,
+            generator,
+            logprobs,
+        )
+
+        self._current_timestep = None
+        if output_type == "latent":
+            image = latents
+        else:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = latents.to(self.vae.dtype)
+            latents_mean = (
+                torch.tensor(self.vae.config.latents_mean)
+                .view(1, self.vae.config.z_dim, 1, 1, 1)
+                .to(latents.device, latents.dtype)
+            )
+            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
+                latents.device, latents.dtype
+            )
+            latents = latents / latents_std + latents_mean
+            image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
+
+        return DiffusionOutput(
+            output={
+                "payload": {
+                    "image": _maybe_to_cpu(image),
+                },
+                "metadata": {
+                    "prompt_embeddings": {
+                        "prompt_embeds": _maybe_to_cpu(prompt_embeds),
+                        "prompt_embeds_mask": _maybe_to_cpu(prompt_embeds_mask),
+                        "negative_prompt_embeds": _maybe_to_cpu(negative_prompt_embeds),
+                        "negative_prompt_embeds_mask": _maybe_to_cpu(negative_prompt_embeds_mask),
+                    },
+                },
+            },
+            trajectory_latents=_maybe_to_cpu(all_latents),
+            trajectory_log_probs=_maybe_to_cpu(all_log_probs),
+            trajectory_timesteps=_maybe_to_cpu(all_timesteps),
+        )

tests/e2e/features/helpers/verl_omni_server.py

@@ -0,0 +1,320 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Ray-serializable verl-omni HTTP server test double for RLHF E2E tests.
+
+The actor class must live in an importable module (not a pytest test file) so
+Ray workers can unpickle ``vLLMOmniHttpServerLocal`` via
+``tests.e2e.features.helpers.verl_omni_server``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
+from dataclasses import asdict
+from enum import Enum
+from typing import Any
+
+import torchvision.transforms as T
+from omegaconf import DictConfig, OmegaConf
+from pydantic import BaseModel, ConfigDict
+from vllm.entrypoints.openai.api_server import build_app
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+import vllm_omni.entrypoints.cli.serve
+from vllm_omni.engine.arg_utils import OmniEngineArgs
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.openai.api_server import omni_init_app_state
+from vllm_omni.inputs.data import OmniCustomPrompt, OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+logger = logging.getLogger(__name__)
+
+CUSTOM_PIPELINE_CLASS = "tests.e2e.features.helpers.custom_pipeline.QwenImagePipelineWithLogProbForTest"
+WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest"
+
+
+class RolloutMode(Enum):
+    HYBRID = "hybrid"
+    COLOCATED = "colocated"
+    STANDALONE = "standalone"
+
+
+class DiffusionOutput(BaseModel):
+    """Pydantic mirror of ``verl_omni.workers.rollout.replica.DiffusionOutput``."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    diffusion_output: Any
+    log_probs: Any | None = None
+    stop_reason: str | None = None
+    num_preempted: int | None = None
+    extra_fields: dict[str, Any] = {}
+
+
+def normalize_token_ids(tokenized_output) -> list[int]:
+    token_ids = tokenized_output
+    if isinstance(tokenized_output, dict):
+        if "input_ids" in tokenized_output:
+            token_ids = tokenized_output["input_ids"]
+    elif hasattr(tokenized_output, "input_ids"):
+        token_ids = tokenized_output.input_ids
+
+    if hasattr(token_ids, "tolist"):
+        token_ids = token_ids.tolist()
+    if isinstance(token_ids, tuple):
+        token_ids = list(token_ids)
+    if isinstance(token_ids, list) and len(token_ids) == 1 and isinstance(token_ids[0], (list, tuple)):
+        token_ids = list(token_ids[0])
+    if not isinstance(token_ids, list):
+        raise TypeError(f"token_ids must be list-like, got {type(token_ids).__name__}")
+
+    out: list[int] = []
+    for tid in token_ids:
+        if hasattr(tid, "item"):
+            tid = tid.item()
+        out.append(int(tid))
+    return out
+
+
+def build_cli_args_from_config(config: dict[str, Any]) -> list[str]:
+    cli_args: list[str] = []
+    for k, v in config.items():
+        if v is None:
+            continue
+        if isinstance(v, bool):
+            if v:
+                cli_args.append(f"--{k}")
+        elif isinstance(v, list):
+            if not v:
+                continue
+            cli_args.append(f"--{k}")
+            cli_args.extend([str(item) for item in v])
+        else:
+            cli_args.append(f"--{k}")
+            cli_args.append(json.dumps(v) if isinstance(v, dict) else str(v))
+    return cli_args
+
+
+def _cfg_get(cfg: Any, key: str, default: Any = None) -> Any:
+    """Uniform getter that works for both DictConfig and dataclass-like objects."""
+    if isinstance(cfg, DictConfig):
+        return cfg.get(key, default)
+    return getattr(cfg, key, default)
+
+
+class vLLMOmniHttpServerLocal:
+    """Self-contained simulation of ``vLLMOmniHttpServer`` for diffusion mode."""
+
+    def __init__(
+        self,
+        config,
+        model_config,
+        rollout_mode: RolloutMode,
+        workers: list,
+        replica_rank: int,
+        node_rank: int,
+        gpus_per_node: int,
+        nnodes: int,
+        cuda_visible_devices: str,
+    ):
+        os.environ.setdefault("CUDA_VISIBLE_DEVICES", cuda_visible_devices)
+        os.environ["VERL_REPLICA_RANK"] = str(replica_rank)
+
+        self.config = OmegaConf.create(config) if not isinstance(config, DictConfig) else config
+        self.model_config = OmegaConf.create(model_config) if not isinstance(model_config, DictConfig) else model_config
+
+        self.rollout_mode = rollout_mode
+        self.workers = workers
+        self.replica_rank = replica_rank
+        self.node_rank = node_rank
+        self.gpus_per_node = gpus_per_node
+        self.nnodes = nnodes
+        self.global_steps = 0
+        self.engine: AsyncOmni | None = None
+        self._to_tensor = T.PILToTensor()
+
+        logger.info(
+            "%s, replica_rank: %d, node_rank: %d, CUDA_VISIBLE_DEVICES: %s",
+            self.__class__.__name__,
+            replica_rank,
+            node_rank,
+            cuda_visible_devices,
+        )
+
+    async def launch_server(self) -> None:
+        engine_kwargs_root = _cfg_get(self.config, "engine_kwargs", {}) or {}
+        if isinstance(engine_kwargs_root, (dict, DictConfig)):
+            engine_kwargs = engine_kwargs_root.get("vllm_omni", {}) or {}
+        else:
+            engine_kwargs = {}
+        engine_kwargs = {k: v for k, v in dict(engine_kwargs).items() if v is not None}
+
+        engine_kwargs.pop("output_mode", None)
+
+        compilation_config = engine_kwargs.pop("compilation_config", None) or {}
+        if isinstance(compilation_config, str):
+            compilation_config = json.loads(compilation_config)
+        if isinstance(compilation_config, DictConfig):
+            compilation_config = OmegaConf.to_container(compilation_config, resolve=True)
+        compilation_config.setdefault("cudagraph_mode", "FULL_AND_PIECEWISE")
+        compilation_config = json.dumps(compilation_config)
+
+        args = {
+            "dtype": _cfg_get(self.config, "dtype", "bfloat16"),
+            "load_format": _cfg_get(self.config, "load_format", "auto"),
+            "skip_tokenizer_init": False,
+            "distributed_executor_backend": "mp",
+            "worker_extension_cls": WORKER_EXTENSION_CLASS,
+            "trust_remote_code": _cfg_get(self.model_config, "trust_remote_code", True),
+            "max_model_len": _cfg_get(self.config, "max_model_len", 1058),
+            "max_num_seqs": _cfg_get(self.config, "max_num_seqs", 1),
+            "enable_chunked_prefill": _cfg_get(self.config, "enable_chunked_prefill", False),
+            "max_num_batched_tokens": _cfg_get(self.config, "max_num_batched_tokens", 8192),
+            "enable_prefix_caching": _cfg_get(self.config, "enable_prefix_caching", False),
+            "enable_sleep_mode": _cfg_get(self.config, "enable_sleep_mode", False),
+            "logprobs_mode": _cfg_get(self.config, "logprobs_mode", "processed_logprobs"),
+            "enforce_eager": _cfg_get(self.config, "enforce_eager", True),
+            "gpu_memory_utilization": _cfg_get(self.config, "gpu_memory_utilization", 0.8),
+            "disable_log_stats": _cfg_get(self.config, "disable_log_stats", True),
+            "tensor_parallel_size": _cfg_get(self.config, "tensor_model_parallel_size", 1),
+            "seed": self.replica_rank + int(_cfg_get(self.config, "seed", 0) or 0),
+            "override_generation_config": json.dumps({}),
+            "scheduling_policy": _cfg_get(self.config, "scheduling_policy", "fcfs"),
+            "compilation_config": compilation_config,
+            **engine_kwargs,
+        }
+
+        model_path = str(_cfg_get(self.model_config, "local_path", None) or _cfg_get(self.model_config, "path"))
+        server_args_list = ["serve", model_path] + build_cli_args_from_config(args)
+
+        parser = FlexibleArgumentParser(description="vLLM-Omni CLI")
+        subparsers = parser.add_subparsers(required=False, dest="subparser")
+        cmds: dict[str, Any] = {}
+        for cmd in vllm_omni.entrypoints.cli.serve.cmd_init():
+            cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
+            cmds[cmd.name] = cmd
+        parsed = parser.parse_args(args=server_args_list)
+        parsed.model = parsed.model_tag
+        if parsed.subparser in cmds:
+            cmds[parsed.subparser].validate(parsed)
+
+        await self.run_server(parsed)
+
+    async def run_server(self, args: argparse.Namespace) -> None:
+        engine_args = OmniEngineArgs.from_cli_args(args)
+        engine_args = asdict(engine_args)
+
+        engine_args["enable_dummy_pipeline"] = True
+        engine_args["custom_pipeline_args"] = {"pipeline_class": CUSTOM_PIPELINE_CLASS}
+
+        engine_client = AsyncOmni(**engine_args)
+        app = build_app(args)
+        await omni_init_app_state(engine_client, app.state, args)
+
+        self.engine = engine_client
+
+    async def generate(
+        self,
+        prompt_ids: list[int],
+        sampling_params: dict[str, Any],
+        request_id: str,
+        image_data: list[Any] | None = None,
+        video_data: list[Any] | None = None,
+        negative_prompt_ids: list[int] | None = None,
+        priority: int = 0,  # noqa: ARG002 (signature parity)
+    ) -> DiffusionOutput:
+        prompt_ids = normalize_token_ids(prompt_ids)
+
+        multi_modal_data: dict[str, Any] = {}
+        if image_data is not None:
+            multi_modal_data["image"] = image_data
+        if video_data is not None:
+            multi_modal_data["video"] = video_data
+
+        custom_prompt: OmniCustomPrompt = {"prompt_ids": prompt_ids}
+        if negative_prompt_ids is not None:
+            custom_prompt["negative_prompt_ids"] = negative_prompt_ids
+        if multi_modal_data:
+            custom_prompt["extra_args"] = {"multi_modal_data": multi_modal_data}
+
+        sampling_kwargs: dict[str, Any] = {}
+        extra_args: dict[str, Any] = {}
+        for k, v in sampling_params.items():
+            if hasattr(OmniDiffusionSamplingParams, k):
+                sampling_kwargs[k] = v
+            else:
+                extra_args[k] = v
+        sampling_kwargs["extra_args"] = extra_args
+        diffusion_sampling_params = OmniDiffusionSamplingParams(**sampling_kwargs)
+
+        generator = self.engine.generate(
+            prompt=custom_prompt,
+            request_id=request_id,
+            sampling_params_list=[diffusion_sampling_params],
+        )
+
+        final_res: OmniRequestOutput | None = None
+        async for output in generator:
+            final_res = output
+        assert final_res is not None
+
+        diffusion_output = self._to_tensor(final_res.images[0]).float() / 255.0
+
+        metadata = (
+            final_res.multimodal_output.get("metadata", {}) if isinstance(final_res.multimodal_output, dict) else {}
+        )
+        prompt_embeddings = metadata.get("prompt_embeddings", {}) if isinstance(metadata, dict) else {}
+
+        if sampling_params.get("logprobs", False):
+            all_log_probs = final_res.trajectory_log_probs
+            log_probs = all_log_probs[0] if all_log_probs is not None else None
+        else:
+            log_probs = None
+
+        all_latents = final_res.trajectory_latents
+        all_timesteps = final_res.trajectory_timesteps
+        prompt_embeds = prompt_embeddings.get("prompt_embeds")
+        prompt_embeds_mask = prompt_embeddings.get("prompt_embeds_mask")
+        negative_prompt_embeds = prompt_embeddings.get("negative_prompt_embeds")
+        negative_prompt_embeds_mask = prompt_embeddings.get("negative_prompt_embeds_mask")
+
+        extra_fields = {
+            "all_latents": all_latents[0] if all_latents is not None else None,
+            "all_timesteps": all_timesteps[0] if all_timesteps is not None else None,
+            "prompt_embeds": prompt_embeds[0] if prompt_embeds is not None else None,
+            "prompt_embeds_mask": prompt_embeds_mask[0] if prompt_embeds_mask is not None else None,
+            "negative_prompt_embeds": negative_prompt_embeds[0] if negative_prompt_embeds is not None else None,
+            "negative_prompt_embeds_mask": (
+                negative_prompt_embeds_mask[0] if negative_prompt_embeds_mask is not None else None
+            ),
+            "global_steps": self.global_steps,
+        }
+
+        req_output = final_res.request_output
+        if req_output is not None and hasattr(req_output, "finish_reason"):
+            finish_reason = req_output.finish_reason or "stop"
+        else:
+            finish_reason = "stop"
+
+        if finish_reason == "abort":
+            stop_reason = "aborted"
+        elif finish_reason in ("stop", "length"):
+            stop_reason = "completed"
+        else:
+            stop_reason = finish_reason
+
+        num_preempted = None
+        if req_output is not None and hasattr(req_output, "num_preempted"):
+            num_preempted = req_output.num_preempted
+
+        return DiffusionOutput(
+            diffusion_output=diffusion_output,
+            log_probs=log_probs,
+            stop_reason=stop_reason,
+            num_preempted=num_preempted,
+            extra_fields=extra_fields,
+        )

tests/e2e/features/rlhf_test/test_verl_omni_e2e.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""E2E test that follows the EXACT same flow as
+``verl-omni/tests/workers/rollout/rollout_vllm/test_vllm_omni_generate.py``
+but inlines all the ``verl`` / ``verl_omni`` symbols the flow depends on
+(stripped down to only what the diffusion-mode test actually exercises).
+
+Flow (identical to the verl-omni test):
+
+1. Build two ``OmegaConf`` configs (``DiffusionRolloutConfig`` /
+   ``DiffusionModelConfig`` shape).
+2. ``ServerCls = ray.remote(vLLMOmniHttpServerLocal)`` then
+   ``ServerCls.options(...).remote(config=..., model_config=...,
+   rollout_mode=RolloutMode.STANDALONE, workers=[], replica_rank=0,
+   node_rank=0, gpus_per_node=1, nnodes=1, cuda_visible_devices="0")``.
+3. ``ray.get(server.launch_server.remote())``.
+4. ``server.generate.remote(prompt_ids=..., sampling_params=..., request_id=...)``
+   and assert the returned ``DiffusionOutput``.
+
+The Ray actor lives in ``tests.e2e.features.helpers.verl_omni_server`` so
+workers can import it after directory migration (pytest test modules are not
+importable as ``test_verl_omni_e2e`` in Ray child processes).
+
+Usage:
+    pytest tests/e2e/features/rlhf_test/test_verl_omni_e2e.py -v -s
+"""
+
+from __future__ import annotations
+
+import os
+from functools import lru_cache
+from uuid import uuid4
+
+import pytest
+import ray
+import torch
+from huggingface_hub import snapshot_download
+from omegaconf import OmegaConf
+from transformers import AutoTokenizer
+
+from tests.e2e.features.helpers.verl_omni_server import (
+    DiffusionOutput,
+    RolloutMode,
+    normalize_token_ids,
+    vLLMOmniHttpServerLocal,
+)
+from tests.helpers.mark import hardware_test
+
+MODEL = "tiny-random/Qwen-Image"
+TOKENIZER_MODEL = "Qwen/Qwen2-1.5B-Instruct"
+_MIN_PROMPT_TOKENS = 35
+
+
+def _resolve_model_path(repo_id: str) -> str:
+    if os.path.isdir(repo_id):
+        return repo_id
+    return snapshot_download(repo_id=repo_id)
+
+
+@lru_cache(maxsize=1)
+def _get_tokenizer():
+    return AutoTokenizer.from_pretrained(TOKENIZER_MODEL, trust_remote_code=True)
+
+
+def _tokenize_prompt(text: str) -> list[int]:
+    tokenizer = _get_tokenizer()
+    messages = [{"role": "user", "content": text}]
+    token_ids = normalize_token_ids(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False))
+    assert len(token_ids) > _MIN_PROMPT_TOKENS, (
+        f"Prompt too short ({len(token_ids)} tokens, need >{_MIN_PROMPT_TOKENS}). "
+        "The pipeline drops the first 34 chat-template prefix tokens; "
+        "use a longer prompt so content tokens remain after the drop."
+    )
+    return token_ids
+
+
+@pytest.fixture
+def init_server():
+    """Create and launch a vLLMOmniHttpServerLocal Ray actor with Qwen-Image."""
+    model_path = _resolve_model_path(MODEL)
+
+    ray.init(
+        runtime_env={
+            "env_vars": {
+                "TOKENIZERS_PARALLELISM": "true",
+                "NCCL_DEBUG": "WARN",
+                "VLLM_LOGGING_LEVEL": "INFO",
+            }
+        },
+        ignore_reinit_error=True,
+    )
+
+    rollout_cfg = OmegaConf.create(
+        {
+            "name": "vllm_omni",
+            "mode": "async",
+            "tensor_model_parallel_size": 1,
+            "data_parallel_size": 1,
+            "pipeline_model_parallel_size": 1,
+            "gpu_memory_utilization": 0.8,
+            "max_num_batched_tokens": 8192,
+            # QwenImagePipeline (and the logprob custom subclass) only supports
+            # serial request execution; max_num_seqs>1 enables request-level batching.
+            "max_num_seqs": 1,
+            "max_model_len": 1058,
+            "dtype": "bfloat16",
+            "load_format": "auto",
+            "enforce_eager": True,
+            "enable_chunked_prefill": False,
+            "enable_prefix_caching": False,
+            "enable_sleep_mode": False,
+            "free_cache_engine": True,
+            "disable_log_stats": True,
+            "logprobs_mode": "processed_logprobs",
+            "scheduling_policy": "fcfs",
+            "seed": 0,
+            "n": 4,
+            "pipeline": {
+                "height": 512,
+                "width": 512,
+                "num_inference_steps": 10,
+            },
+        }
+    )
+
+    model_cfg = OmegaConf.create(
+        {
+            "path": str(model_path),
+            "local_path": str(model_path),
+            "tokenizer_path": TOKENIZER_MODEL,
+            "trust_remote_code": True,
+            "load_tokenizer": True,
+        }
+    )
+
+    ServerCls = ray.remote(vLLMOmniHttpServerLocal)
+    server = ServerCls.options(
+        runtime_env={
+            "env_vars": {
+                "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1",
+                "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1",
+                "NCCL_CUMEM_ENABLE": "0",
+            }
+        },
+        max_concurrency=10,
+        num_gpus=1,
+    ).remote(
+        config=rollout_cfg,
+        model_config=model_cfg,
+        rollout_mode=RolloutMode.STANDALONE,
+        workers=[],
+        replica_rank=0,
+        node_rank=0,
+        gpus_per_node=1,
+        nnodes=1,
+        cuda_visible_devices="0",
+    )
+
+    ray.get(server.launch_server.remote())
+
+    yield server
+
+    ray.shutdown()
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_generate(init_server):
+    """Concurrent generate() calls covering basic output, logprobs, and multi-request correctness."""
+    server = init_server
+
+    prompts = [
+        "a beautiful sunset over the ocean with vibrant orange and purple clouds "
+        "reflecting on the calm water surface near a rocky coastline",
+        "a fluffy orange cat sitting on a wooden windowsill looking outside at "
+        "a garden full of colorful flowers on a bright sunny afternoon",
+        "a majestic mountain landscape covered with fresh white snow under a "
+        "clear blue sky with pine trees in the foreground and a frozen lake",
+        "a futuristic city at night with neon lights glowing on tall glass "
+        "skyscrapers and flying vehicles soaring between the buildings",
+    ]
+
+    refs = []
+    for i, prompt in enumerate(prompts):
+        rid = f"test_{i}_{uuid4().hex[:8]}"
+        ref = server.generate.remote(
+            prompt_ids=_tokenize_prompt(prompt),
+            sampling_params={
+                "num_inference_steps": 10,
+                "true_cfg_scale": 4.0,
+                "height": 512,
+                "width": 512,
+                "logprobs": i == 0,
+            },
+            request_id=rid,
+        )
+        refs.append(ref)
+
+    results = ray.get(refs, timeout=600)
+
+    for i, output in enumerate(results):
+        assert isinstance(output, DiffusionOutput), f"Request {i}: expected DiffusionOutput"
+        img = output.diffusion_output
+        assert isinstance(img, torch.Tensor) and img.ndim == 3, (
+            f"Request {i}: expected CHW torch.Tensor, got {type(img).__name__} ndim={getattr(img, 'ndim', None)}"
+        )
+        c, h, w = img.shape
+        assert c == 3, f"Request {i}: expected 3 channels (CHW), got {c}"
+        assert h > 0 and w > 0, f"Request {i}: image dimensions must be positive"
+        assert output.stop_reason in ("completed", "aborted", None), (
+            f"Request {i}: unexpected stop_reason {output.stop_reason!r}"
+        )
+        assert 0.0 <= float(img[0, 0, 0]) <= 1.0, f"Request {i}: pixel values must be in [0, 1]"
+
+    print(f"All {len(prompts)} concurrent requests returned valid DiffusionOutput")