[Core][Frontend] Bound duplex output delivery and reject cancelled audio (#7644)

Signed-off-by: Nolen Liang <nliang@nvidia.com>
Signed-off-by: bcsdhjew <nliang@nvidia.com>

tests/entrypoints/duplex/test_duplex_serving.py

@@ -18,6 +18,7 @@ from fastapi import WebSocketDisconnect
 from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
 from vllm_omni.engine.duplex import commands
 from vllm_omni.engine.duplex.config import DuplexCapabilities
+from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
 from vllm_omni.engine.duplex.events import AudioDelta, DuplexEvent, SessionClosed, SessionCreated
 from vllm_omni.engine.duplex.messages import DuplexSessionError
 from vllm_omni.entrypoints.duplex.realtime_input import RealtimeEnvelope, parse_resume_request
@@ -94,10 +95,13 @@ class FakeHandle:
         self.closed = False
         self.close_reasons: list[str] = []
         self.commands: list[commands.DuplexCommand] = []
-        self._outbox: asyncio.Queue[DuplexEvent] = asyncio.Queue()
+        self._outbox = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512)
 
     def deliver(self, event: DuplexEvent) -> None:
-        self._outbox.put_nowait(event)
+        self._outbox.put(event)
+
+    def output_guard(self, event: DuplexEvent):
+        return self._outbox.guard(event)
 
     async def submit(self, command: commands.DuplexCommand) -> None:
         if self.closed:
@@ -107,6 +111,10 @@ class FakeHandle:
     async def events(self):
         while True:
             event = await self._outbox.get()
+            if event is None:
+                return
+            if not self._outbox.is_valid(event):
+                continue
             yield event
             if event.is_terminal:
                 return

tests/entrypoints/openai/test_duplex_session_attachment.py

@@ -9,6 +9,8 @@ import json
 
 import pytest
 
+from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
+from vllm_omni.engine.duplex.events import AudioDelta, ResponseDone
 from vllm_omni.entrypoints.duplex.session_attachment import (
     DuplexEventJournal,
     DuplexJournalGapError,
@@ -22,6 +24,36 @@ from vllm_omni.entrypoints.duplex.session_attachment import (
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("journal", [False, True])
+async def test_audio_invalidated_while_waiting_for_send_is_not_sequenced_or_journaled(mocker, journal: bool):
+    registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096)
+    output = DuplexOutputBuffer(max_bytes=4096, max_events=8)
+    sent: list[dict[str, object]] = []
+    accepted = mocker.Mock()
+
+    await registry.create("s", send=mocker.AsyncMock(side_effect=sent.append), close=mocker.AsyncMock())
+    audio = AudioDelta(session_id="s", response_id="r", epoch=0, delta="AAAA")
+    output.put(audio)
+    assert await output.get() is audio
+    state = registry._sessions["s"]
+    async with state.outbound_lock:
+        delivery = asyncio.create_task(
+            registry.send_event(
+                "s", audio.to_realtime(), journal=journal, on_accepted=accepted, event_guard=lambda: output.guard(audio)
+            )
+        )
+        await asyncio.sleep(0)
+        assert not delivery.done()
+        output.invalidate("r", through_epoch=0)
+    assert await delivery is None
+    accepted.assert_not_called()
+    done = ResponseDone(response_id="r", response={"status": "cancelled"})
+    await registry.send_event("s", done.to_realtime())
+    assert [(entry["type"], entry["server_event_seq"]) for entry in sent] == [("response.done", 1)]
+    assert [entry.payload["event_id"] for entry in state.journal.replay_after(0)] == [done.event_id]
+
+
 class _Clock:
     def __init__(self) -> None:
         self.value = 0.0

tests/entrypoints/test_duplex_omni.py

@@ -15,6 +15,7 @@ from tests.entrypoints.test_omni_entrypoints import FakeAsyncOmniEngine
 from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
 from vllm_omni.engine.duplex import commands
 from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
+from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer
 from vllm_omni.engine.duplex.events import (
     AudioDelta,
     DuplexEvent,
@@ -41,6 +42,7 @@ class FakeDuplexEngine(FakeAsyncOmniEngine):
         self.duplex_session_config = DuplexSessionRuntimeConfig(max_sessions=2)
         self.duplex_capabilities = DuplexCapabilities(supports_session_resume=True)
         self.opened: list[tuple[str, DuplexSessionConfig]] = []
+        self.outputs: dict[str, DuplexOutputBuffer] = {}
         self.closed: list[tuple[str, str]] = []
         self.resumed: list[tuple[str, int]] = []
         #: The control id of every resume RPC, in order (a replay repeats one).
@@ -64,10 +66,11 @@ class FakeDuplexEngine(FakeAsyncOmniEngine):
     def _result(self, operation: str, session_id: str, **fields: Any) -> DuplexControlResultMessage:
         return DuplexControlResultMessage(control_id="c", operation=operation, session_id=session_id, ok=True, **fields)
 
-    async def open_session_async(self, session_id, session_config, *, timeout=None):
+    async def open_session_async(self, session_id, session_config, *, output_buffer, timeout=None):
         self.opened.append((session_id, session_config))
         if self.open_error is not None:
             raise self.open_error
+        self.outputs[session_id] = output_buffer
         return self._result(
             "open",
             session_id,
@@ -119,7 +122,11 @@ class FakeDuplexEngine(FakeAsyncOmniEngine):
     def emit(self, session_id: str, event: DuplexEvent) -> None:
         # The session manager binds the session identity to every event it emits.
         event = replace(event, session_id=session_id)
-        self.output_q.put(DuplexSessionEventMessage(session_id=session_id, event=event))
+        if isinstance(event, SessionClosed):
+            self.outputs.pop(session_id, None)
+            self.output_q.put(DuplexSessionEventMessage(session_id=session_id, event=event))
+        elif session_id in self.outputs:
+            self.outputs[session_id].put(event)
 
 
 def _make_omni(monkeypatch: pytest.MonkeyPatch) -> tuple[DuplexOmni, FakeDuplexEngine]:
@@ -152,6 +159,7 @@ async def test_open_session_allocates_the_id_and_ignores_client_ids(monkeypatch)
         assert handle.session_id != "mine"
         ((session_id, config),) = engine.opened
         assert session_id == handle.session_id
+        assert engine.outputs[session_id] is handle._outbox
         assert isinstance(config, DuplexSessionConfig)
         assert config.model == "dummy-model" and config.instructions == "hi"
         assert handle.capabilities is engine.duplex_capabilities
@@ -169,6 +177,32 @@ async def test_open_session_allocates_the_id_and_ignores_client_ids(monkeypatch)
         omni.shutdown()
 
 
+@pytest.mark.asyncio
+async def test_events_skips_audio_invalidated_after_dequeue(monkeypatch) -> None:
+    omni, engine = _make_omni(monkeypatch)
+    try:
+        handle = await omni.open_session()
+        output = engine.outputs[handle.session_id]
+        original_get = output.get
+
+        async def get_then_invalidate() -> DuplexEvent | None:
+            event = await original_get()
+            if isinstance(event, AudioDelta) and event.response_id == "cancelled":
+                output.invalidate("cancelled", through_epoch=0)
+            return event
+
+        monkeypatch.setattr(output, "get", get_then_invalidate)
+        engine.emit(handle.session_id, AudioDelta(response_id="cancelled", epoch=0, delta="old"))
+        engine.emit(handle.session_id, AudioDelta(response_id="current", epoch=1, delta="new"))
+
+        events = await _collect(handle, 1)
+
+        assert isinstance(events[0], AudioDelta)
+        assert events[0].response_id == "current" and events[0].delta == "new"
+    finally:
+        omni.shutdown()
+
+
 @pytest.mark.asyncio
 async def test_open_failure_drops_the_pending_handle_and_raises_the_engine_code(monkeypatch) -> None:
     omni, engine = _make_omni(monkeypatch)