[Bugfix][MiniCPM-o] Prevent native duplex Stage-1 prompt overflow

Signed-off-by: Sy03 <1370724210@qq.com>

tests/engine/test_orchestrator.py

@@ -30,6 +30,7 @@ from vllm_omni.engine.messages import (
 from vllm_omni.engine.orchestrator import (
     Orchestrator,
     OrchestratorRequestState,
+    StreamingSegmentState,
     _build_terminal_empty_output,
 )
 from vllm_omni.engine.stage_pool import StagePool
@@ -1951,7 +1952,7 @@ async def test_resumable_segment_boundary_builds_stage_metrics() -> None:
         final_stage_id=0,
     )
     req_state.streaming.enabled = True
-    req_state.streaming.segment_finished = True
+    req_state.streaming.segments[0] = StreamingSegmentState(finished=True)
     req_state.stage_submit_ts[0] = time.time()
     orchestrator.request_states = {"req-stream": req_state}
     orchestrator.stage_pools = [pool]

tests/engine/test_orchestrator_error_handling.py

@@ -880,6 +880,39 @@ async def test_diffusion_client_error_output_propagates_status_code(orchestrator
         orchestrator_fixture.thread.join(timeout=5)
 
 
+@pytest.mark.asyncio
+async def test_duplex_input_processor_failure_is_request_scoped(orchestrator_factory, monkeypatch) -> None:
+    class FailingStage(FakeStageClient):
+        def process_engine_inputs(self, source_outputs, prompt=None, streaming_context=None):
+            del source_outputs, prompt, streaming_context
+            raise ValueError("No latent or hidden_states found in thinker output")
+
+    fixture = orchestrator_factory(
+        [FakeStageClient(stage_type="llm"), FailingStage(stage_type="llm", final_output=True)]
+    )
+    state = OrchestratorRequestState(
+        request_id="bad",
+        prompt=SimpleNamespace(request_id="bad", prompt_token_ids=[1]),
+        sampling_params_list=[_sampling_params(), _sampling_params()],
+        final_stage_id=1,
+        duplex_identity=SimpleNamespace(),
+    )
+
+    async def no_cleanup(*_args: object, **_kwargs: object) -> None:
+        return None
+
+    monkeypatch.setattr(fixture.orchestrator, "_cleanup_request_ids", no_cleanup)
+    try:
+        await fixture.orchestrator._forward_to_next_stage_unguarded(
+            "bad", 0, _build_request_output("raw", finished=True), state
+        )
+        error = await _wait_for_error_message(fixture, request_id="bad")
+        assert error.fatal is False
+    finally:
+        fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
+        fixture.thread.join(timeout=5)
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     ("status_code", "error_type"),

tests/engine/test_orchestrator_segment_stage_keying.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Streaming segment-boundary state must be keyed by stage.
+
+``Orchestrator._orchestration_loop`` polls every stage into the same
+``OrchestratorRequestState``. While the segment-boundary fields lived in a single
+flat slot per request, whichever stage polled most recently overwrote the others,
+and the per-stage consumers then read back a different stage's boundary.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import pytest
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.engine.orchestrator import (
+    Orchestrator,
+    OrchestratorRequestState,
+    StreamingSegmentState,
+)
+from vllm_omni.experimental.fullduplex.engine.contracts import DuplexRequestIdentity
+from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+THINKER_STAGE = 0
+TALKER_STAGE = 1
+
+
+@dataclass
+class _StageMetrics:
+    pipeline_timings: dict[str, float] = field(default_factory=dict)
+
+
+@dataclass
+class _StageOutput:
+    """Stand-in for a processed stage output, with the fields the router reads."""
+
+    request_id: str
+    finished: bool = False
+    error: object | None = None
+
+
+class _RecordingPool:
+    """Stage pool that records the requests it was asked to build metrics for."""
+
+    def __init__(self) -> None:
+        self.built_for: list[str] = []
+
+    def build_stage_metrics(self, outputs: list[_StageOutput], **_kwargs: object) -> _StageMetrics:
+        self.built_for.extend(output.request_id for output in outputs)
+        return _StageMetrics()
+
+
+def _duplex_req_state(request_id: str = "req-duplex") -> OrchestratorRequestState:
+    req_state = OrchestratorRequestState(
+        request_id=request_id,
+        sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)],
+        final_stage_id=TALKER_STAGE,
+    )
+    req_state.streaming.enabled = True
+    req_state.duplex_identity = DuplexRequestIdentity(
+        session_id="session-1",
+        fence=DuplexFence(session_id="session-1"),
+    )
+    return req_state
+
+
+def test_duplex_output_context_reads_only_its_own_stage_segment() -> None:
+    req_state = _duplex_req_state()
+    req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState(
+        finished=True,
+        token_ids=[11, 22],
+        output_metadata={"stage": "thinker"},
+    )
+    # The talker polls after the thinker and is still mid-segment.
+    req_state.streaming.segments[TALKER_STAGE] = StreamingSegmentState(finished=False)
+
+    thinker = Orchestrator._duplex_output_context(req_state, stage_id=THINKER_STAGE)
+    talker = Orchestrator._duplex_output_context(req_state, stage_id=TALKER_STAGE)
+
+    assert thinker is not None
+    assert talker is not None
+
+    # These are the assertions that fail on the flat field: the talker is written
+    # second, so the thinker loses its boundary to whichever stage polled last.
+    assert thinker.segment_finished is True
+    assert thinker.segment_token_ids == (11, 22)
+    assert thinker.segment_output_metadata == {"stage": "thinker"}
+
+    assert talker.segment_finished is False
+    assert talker.segment_token_ids == ()
+    assert talker.segment_output_metadata == {}
+
+
+@pytest.mark.asyncio
+async def test_segment_boundary_on_one_stage_does_not_build_metrics_on_another() -> None:
+    req_state = _duplex_req_state()
+    req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState(finished=True)
+
+    talker_pool = _RecordingPool()
+    orchestrator = object.__new__(Orchestrator)
+    orchestrator.request_states = {req_state.request_id: req_state}
+    orchestrator.stage_pools = [_RecordingPool(), talker_pool]
+
+    routed: list[_StageMetrics | None] = []
+
+    async def record_route(
+        _stage: int,
+        _replica: int,
+        _output: _StageOutput,
+        _state: OrchestratorRequestState,
+        metrics: _StageMetrics | None,
+    ) -> None:
+        routed.append(metrics)
+
+    orchestrator._route_output = record_route
+
+    # The talker's own output is neither finished nor at a segment boundary, so
+    # it must not be treated as one just because the thinker reached one.
+    await orchestrator._handle_processed_outputs(
+        TALKER_STAGE,
+        0,
+        [_StageOutput(request_id=req_state.request_id)],
+    )
+
+    assert talker_pool.built_for == []
+    assert routed == [None]

tests/engine/test_orchestrator_stage_input_bridge.py

@@ -17,6 +17,7 @@ from vllm.sampling_params import SamplingParams
 from vllm_omni.engine.orchestrator import (
     Orchestrator,
     OrchestratorRequestState,
+    StreamingSegmentState,
     _OrchestratorDuplexStagePort,
 )
 from vllm_omni.engine.stage_pool import StagePool
@@ -358,7 +359,7 @@ async def test_streaming_segment_does_not_complete_final_output_stage() -> None:
         final_output_stage_ids={0},
     )
     req_state.streaming.enabled = True
-    req_state.streaming.segment_finished = True
+    req_state.streaming.segments[0] = StreamingSegmentState(finished=True)
     output = SimpleNamespace(
         request_id=req_state.request_id,
         finished=True,