From f781f7bb0878e83bd620377212d62e56a5a955bd Mon Sep 17 00:00:00 2001 From: bcsdhjew Date: Mon, 5 Oct 2026 05:06:04 +0800 Subject: [PATCH] [Core][Frontend] Bound duplex output delivery and reject cancelled audio (#7644) Signed-off-by: Nolen Liang Signed-off-by: bcsdhjew --- docs/design/fullduplex.md | 64 +++- docs/serving/realtime_duplex_api.md | 38 ++- .../duplex/test_aura_resumable_hooks.py | 2 + tests/engine/duplex/test_delivery.py | 123 ++++++++ tests/engine/duplex/test_session_manager.py | 290 ++++++++++++++---- tests/engine/duplex/test_session_runner.py | 90 +++++- tests/engine/test_duplex_omni_engine.py | 13 +- tests/engine/test_duplex_orchestrator.py | 31 +- .../entrypoints/duplex/test_duplex_serving.py | 12 +- .../openai/test_duplex_session_attachment.py | 32 ++ tests/entrypoints/test_duplex_omni.py | 38 ++- .../models/qwen3_omni/test_duplex_plugin.py | 17 +- .../test_realtime_codec_single_source.py | 9 +- vllm_omni/config/stage_config.py | 4 + vllm_omni/engine/duplex/delivery.py | 242 +++++++++++++++ vllm_omni/engine/duplex/messages.py | 8 +- vllm_omni/engine/duplex/realtime_events.py | 4 +- vllm_omni/engine/duplex/session/emitter.py | 58 +++- vllm_omni/engine/duplex/session/manager.py | 132 ++++++-- .../engine/duplex/session/model_channel.py | 33 +- vllm_omni/engine/duplex/session/runner.py | 13 + vllm_omni/engine/duplex_omni_engine.py | 9 +- vllm_omni/entrypoints/duplex/serving.py | 12 +- .../entrypoints/duplex/session_attachment.py | 14 +- vllm_omni/entrypoints/duplex_omni.py | 28 +- vllm_omni/protocol/duplex/errors.py | 1 + 26 files changed, 1170 insertions(+), 147 deletions(-) create mode 100644 tests/engine/duplex/test_delivery.py create mode 100644 vllm_omni/engine/duplex/delivery.py diff --git a/docs/design/fullduplex.md b/docs/design/fullduplex.md index da0584843..d7cd9bf54 100644 --- a/docs/design/fullduplex.md +++ b/docs/design/fullduplex.md @@ -220,6 +220,7 @@ vllm_omni/ │ │ DUPLEX_REALTIME_CAPABILITIES, duplex_response_format │ ├── events.py re-export shim (DuplexEvent = RealtimeEvent) + the runner's │ │ epoch-filter sets DOMAIN_TERMINAL_EVENTS / MODEL_OUTPUT_EVENTS +│ ├── delivery.py DuplexOutputBuffer (bounded output shared with the session handle) │ ├── realtime_events.py RealtimeProjectionState: internal event -> typed events │ ├── messages.py queue envelopes (Open/Close/Resume/Touch/Command/Result/Event), DuplexSessionError │ ├── config.py DuplexSessionConfig, DuplexCapabilities, ResponseCreateOptions @@ -284,7 +285,8 @@ command handle.submit(DuplexCommand) -> engine.submit_command_async -> DuplexSessionCommandMessage on the request queue [one-way] -> DuplexSessionManager.dispatch: unknown_session / input_backpressure checks, then runner mailbox output DuplexOrchestrator._intercept_stage_output -> runner.on_stage_output -> mailbox -> typed events - -> output_sink (DuplexSessionEventMessage) -> DuplexOmni._route_engine_message -> handle.events() + -> session's bounded output buffer -> handle.events(); closure and errors without a live session + use output_sink (DuplexSessionEventMessage) -> DuplexOmni._route_engine_message detach DuplexOmni.detach_session(expected_lease_generation) -> touch(DETACH): engine-owned disconnect grace, refused for a lease newer than the one the caller held; expiry -> SessionExpired resume DuplexOmni.resume_session(expected_lease_generation) -> lease CAS keyed by the control id (a replay @@ -303,7 +305,9 @@ reap DuplexSessionManager.reaper_loop: idle TTL / disconnect grace expiry, c ### Concurrency and ordering model of the runner -Everything below runs on the orchestrator asyncio loop; there is no lock. +The runner's session state is owned by the orchestrator asyncio loop and +needs no lock. The output buffer shared with the caller thread uses a lock +for queue updates and the final validity check before delivery. - **Single writer per session.** `DuplexEngineSession` is mutated only by its runner. Inputs reach the runner through one mailbox in this order: client @@ -332,13 +336,15 @@ Everything below runs on the orchestrator asyncio loop; there is no lock. on one loop in program order, nothing can be emitted for an old epoch after its terminal event; `runner.emit()` still drops a terminal carrying a stale epoch and stamps `epoch` on every event for clients that filter - defensively. This is a statement about order, not about latency: the engine - output queue and each `DuplexSessionHandle` outbox are unbounded and FIFO, - so a cancellation is delivered behind whatever audio was already emitted for - the response it cancels. A client that must stop quickly cancels its own - playback on `response.done` / `audio.cancelled` rather than waiting for the - stream to drain. Bounding those buffers and letting an accepted invalidation - skip undelivered media is left to the follow-up RFC. + defensively. In this experimental follow-up, public output goes directly to + one bounded buffer shared by the engine and `DuplexSessionHandle`. Accepted + cancellation removes queued `AudioDelta` events for that response through + the cancelled epoch. Text increments, audio/transcript completion markers, + other responses and normal completion retain FIFO order. The consumer also + checks an event already taken from the buffer: WebSocket delivery repeats + that check under the connection lock, before sequencing and journaling. + Already sequenced replay entries are unchanged; already delivered audio + still requires the client to stop its own playback. - **Backpressure before the mailbox.** `DuplexSessionManager.dispatch` checks `max_pending_input_bytes_per_session` and reserves a pending turn for `Commit` before the put; a rejected command is answered with @@ -349,6 +355,46 @@ Everything below runs on the orchestrator asyncio loop; there is no lock. `ResponseDone.stage_metrics` and `metadata.vllm_omni.stage_metrics` on the wire keep their meaning. +### Experimental output limits + +`max_pending_output_bytes_per_session` defaults to 2 MiB and +`max_pending_output_events_per_session` to 512. The byte limit counts compact +Realtime JSON, including base64 audio, not decoded PCM or Python object +overhead. Audio is already base64-encoded by the producer; counting its length +avoids serializing or decoding that large field again. Neither default has +been tuned from a performance study. + +Errors and response endings have an additional 64 KiB / eight-event reserve +and retain FIFO order. One final session-closure notification has an independent +slot, so an exhausted reserve cannot prevent closure or break delivery to +other sessions. If that reserve also fills, an error or response-ending +notification may be omitted; the final session-closure notification still +has its independent slot. Oversized close details (over 4 KiB with generated identity) +are replaced by a compact reason, preserving terminal type and identity. +Late writes after closure are ignored; a locally ended iterator is not reopened +by a late terminal, and a failed open closes the abandoned handle. Events other +than errors are not forwarded once their session's buffer has been removed. +One event held by the consumer is outside the queued budget. + +The overflow policy fails the active response without committing its partial +text to history and closes only the affected session. If overflow occurs while +a completion batch is being delivered, its not-yet-queued `ResponseDone` is +converted to failure with its identity and output preserved; any history entry +for that response is removed. An ending already accepted by the buffer is not +replaced or followed by a second ending. After overflow, newly +emitted ordinary events are suppressed, including text/audio completion, +content-part and output-item completion markers. Previously queued non-audio +events retain FIFO order; error and failed response-ending notifications may +use the reserve before the independent closure notification. This differs from +accepted cancellation, which removes only matching audio. See the +[public error and closure fields](../serving/realtime_duplex_api.md#output-limits-and-slow-consumers). +Response-only recovery is not implemented. These limits do not bound the runner's raw stage +output mailbox, replay journal, or client playback queues. The shared buffer +requires the current same-process, separate-thread engine arrangement; it is +not a cross-process delivery protocol. Direct handle iteration and WebSocket +delivery share buffer validity, but end-to-end client behavior still needs +independent real-model verification. + ## Orchestrator seams `OrchestratorBase` keeps the generic machinery (request handler skeleton, diff --git a/docs/serving/realtime_duplex_api.md b/docs/serving/realtime_duplex_api.md index 5095d1ca7..65c8fa64a 100644 --- a/docs/serving/realtime_duplex_api.md +++ b/docs/serving/realtime_duplex_api.md @@ -250,6 +250,40 @@ call; transport failures raise `DuplexConnectionError`; a session that ends while you are awaiting it raises `DuplexSessionClosedError`. All three derive from `DuplexClientError`. +### Output limits and slow consumers + +Deployment settings under `duplex_session` bound each session's pending +output: `max_pending_output_bytes_per_session` defaults to 2 MiB and +`max_pending_output_events_per_session` to 512 events. Bytes count compact +Realtime JSON, including base64 audio, not decoded audio. The limits apply +to both WebSocket delivery and direct Python handles, not to replay journals +or client playback queues. + +Exceeding a limit fails the active response without committing its partial +text to history, then closes only that session. If overflow occurs while a +completion is being queued, its not-yet-queued `response.done` becomes failed +with the same identity and output, and its response history entry is removed. +An ending already accepted by the buffer is not replaced or followed by a +second ending. Open a new session to continue; the closed session cannot be +resumed. The notifications use these fields: + +| Event | Field values | +| --- | --- | +| `error` | `error.type: "rate_limit_error"`, `error.code: "output_backpressure"` | +| `response.done` (active response or not-yet-queued ending) | `response.status: "failed"`, `response.status_details: {"type": "failed", "reason": "output_backpressure"}` | +| `session.closed` | `reason: "output_backpressure"` | + +Queued audio for the affected response is removed. Previously queued non-audio +events retain their order, but newly emitted ordinary events, including +audio/text, content-part and output-item completion markers, are suppressed. +Clients must therefore handle closure without waiting for every completion +marker. Errors and response endings have an extra 64 KiB / eight-event reserve; +if that also fills, either notification may be omitted. Closure has its own +independent slot. A closure notification exceeding 4 KiB is compacted to +`reason: "close_details_exceed_output_limit"` with empty details, preserving +the event type and identity. One event held by the consumer is outside the +queued budget; audio already sent still requires the client to stop playback. + ### Measure latency `EventCollector` accumulates events for assertions and metrics: @@ -315,6 +349,8 @@ single-consumer async iterator that ends after `session.closed` / `session.expired`; rejected commands come back as `ErrorEvent`s on it. Every event renders the wire JSON with `to_realtime()`, so anything written against the WebSocket protocol works unchanged on the typed stream. +If opening fails, the abandoned handle is closed. Late engine events do not +reopen a locally closed handle. `vllm_omni.clients.inline_duplex.InlineDuplexClient` wraps a `DuplexOmni` behind the `DuplexClient` API, so the same application code runs in-process @@ -573,7 +609,7 @@ formats (`pcm16`, `pcm_s16le`, `s16le`, `pcm_f32le`, `g711_ulaw`, | `session.created` | 2 | Session opened. OpenAI name; adds `attachment_generation`, `resume_token` and vLLM-Omni keys inside `session` (`id` is the server-allocated session id, `epoch`, `turn_id`, `playback`, `capabilities`, ...). | | `session.updated` | 2 | Echo of the effective session config after every `session.update` (same extended `session` object). | | `session.heartbeat_ack` | 3 | Reply to `session.heartbeat`. | -| `session.closed` | 3 | Last event on the socket, emitted once the engine released the session; `reason` ∈ `client_close`, `disconnect`, `timeout`, `transport_error`, `shutdown`, ... | +| `session.closed` | 3 | Last event on the socket, emitted once the engine released the session; `reason` ∈ `client_close`, `disconnect`, `timeout`, `transport_error`, `shutdown`, `output_backpressure`, `close_details_exceed_output_limit`, ... | | `session.resumed` | 3 | Resume accepted; carries the new `attachment_generation` and rotated `resume_token`; journaled events are replayed after it. | | `session.replaced` | 3 | Sent to the superseded socket when another socket resumes the session. | | `session.expired` | 3 | Engine lease reaped (`disconnect_grace_expired`, `idle_ttl_expired`); socket closes after it. | diff --git a/tests/engine/duplex/test_aura_resumable_hooks.py b/tests/engine/duplex/test_aura_resumable_hooks.py index e2ed43372..335ab97c6 100644 --- a/tests/engine/duplex/test_aura_resumable_hooks.py +++ b/tests/engine/duplex/test_aura_resumable_hooks.py @@ -370,6 +370,7 @@ def test_draining_empty_eos_completes_owning_response_before_shortcut() -> None: stage_port=port, model_state=SimpleNamespace(clear_continuation=lambda: None), services=SimpleNamespace(spawn=lambda *_a, **_k: None), + run=SimpleNamespace(closing=False), ), out, close_from_runtime=_noop, @@ -461,6 +462,7 @@ def test_draining_completion_drops_finished_response_books_only() -> None: stage_port=_Port(), model_state=SimpleNamespace(clear_continuation=lambda: None), services=SimpleNamespace(spawn=lambda *_a, **_k: None), + run=SimpleNamespace(closing=False), ), out, close_from_runtime=_noop, diff --git a/tests/engine/duplex/test_delivery.py b/tests/engine/duplex/test_delivery.py new file mode 100644 index 000000000..0dcdba42c --- /dev/null +++ b/tests/engine/duplex/test_delivery.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import asyncio +import json +from dataclasses import replace + +import pytest + +from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer, DuplexOutputOverflowError +from vllm_omni.engine.duplex.events import AudioDelta, ResponseDone, SessionClosed, TranscriptDelta +from vllm_omni.engine.duplex.realtime_events import RealtimeProjectionState, project_internal_event + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("event_type", "delta_field", "content_field"), + [ + ("response.text.delta", "delta", "text"), + ("response.output_audio.delta", "text", "transcript"), + ], +) +async def test_queued_response_creation_keeps_payload_and_byte_count(event_type, delta_field, content_field): + state = RealtimeProjectionState(session_id="s") + events = project_internal_event(state, {"type": "response.created", "response_id": "r", "modalities": ["text"]}) + payloads = [event.to_realtime() for event in events] + byte_count = sum( + len(json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode()) for payload in payloads + ) + output = DuplexOutputBuffer(max_bytes=byte_count, max_events=len(events)) + for event in events: + output.put(event) + assert output.pending_bytes == byte_count + + text = "Growing response. " * 512 + project_internal_event(state, {"type": event_type, "response_id": "r", delta_field: text}) + item = next(iter(state.conversation_items.values())) + assert item["content"][0][content_field] == text + assert output.pending_bytes == byte_count + delivered = [await output.get() for _ in events] + assert [event.to_realtime() for event in delivered] == payloads + assert output.pending_events == output.pending_bytes == 0 + + +@pytest.mark.asyncio +async def test_cancel_removes_only_matching_audio_and_invalidates_held_audio(): + output = DuplexOutputBuffer(max_bytes=4096, max_events=8) + held = AudioDelta(response_id="old", epoch=0, delta="AAAA") + text = TranscriptDelta(response_id="old", epoch=0, delta="kept") + other = replace(held, response_id="other") + newer = replace(held, epoch=1) + done = ResponseDone(response_id="old", response={"status": "cancelled"}) + for event in (held, replace(held), text, other, newer, done): + output.put(event) + assert await output.get() is held + assert output.is_valid(held) + output.invalidate("old", through_epoch=0) + assert output.pending_events == 4 + with output.guard(held) as valid: + assert not valid + assert [await output.get() for _ in range(4)] == [text, other, newer, done] + assert output.pending_events == output.pending_bytes == 0 + + +@pytest.mark.asyncio +async def test_full_reserve_cannot_block_closure_or_reorder_valid_audio(): + output = DuplexOutputBuffer(max_bytes=4096, max_events=1, reserve_events=1) + audio = AudioDelta(response_id="r", epoch=0, delta="AAAA") + done = ResponseDone(response_id="r", response={"status": "completed"}) + output.put(audio) + output.put(done) + with pytest.raises(DuplexOutputOverflowError): + output.put(done) + closed = SessionClosed(session_id="s", reason="client_close") + output.close(closed) + output.close(SessionClosed(reason="duplicate")) + assert not output.put(audio) + assert [await output.get() for _ in range(4)] == [audio, done, closed, None] + assert output.pending_events == output.pending_bytes == 0 + + +@pytest.mark.asyncio +async def test_oversized_close_still_ends_stream_without_reopening_it(): + output = DuplexOutputBuffer(max_bytes=128, max_events=1) + closed = SessionClosed(session_id="s", details={"large": "x" * 5000}) + output.close(closed) + terminal = await output.get() + assert isinstance(terminal, SessionClosed) + assert terminal.event_id == closed.event_id + assert terminal.session_id == "s" + assert terminal.reason == "close_details_exceed_output_limit" + assert terminal.details == {} + assert await output.get() is None + output.close(closed) + assert await output.get() is None + + +def test_byte_limit_rejects_audio_without_changing_pending_budget(): + output = DuplexOutputBuffer(max_bytes=128, max_events=8) + with pytest.raises(DuplexOutputOverflowError): + output.put(AudioDelta(delta="AAAA" * 64)) + assert output.pending_events == output.pending_bytes == 0 + + +@pytest.mark.asyncio +async def test_cancelled_waiter_can_be_replaced_and_woken_from_another_thread(): + output = DuplexOutputBuffer(max_bytes=4096, max_events=8) + abandoned = asyncio.create_task(output.get()) + await asyncio.sleep(0) + abandoned.cancel() + with pytest.raises(asyncio.CancelledError): + await abandoned + waiting = asyncio.create_task(output.get()) + await asyncio.sleep(0) + event = AudioDelta(response_id="r", epoch=0, delta="AAAA") + await asyncio.to_thread(output.put, event) + assert await asyncio.wait_for(waiting, timeout=2) is event + closing = asyncio.create_task(output.get()) + await asyncio.sleep(0) + await asyncio.to_thread(output.close) + assert await asyncio.wait_for(closing, timeout=2) is None diff --git a/tests/engine/duplex/test_session_manager.py b/tests/engine/duplex/test_session_manager.py index 5fe07d6bf..938fbba1f 100644 --- a/tests/engine/duplex/test_session_manager.py +++ b/tests/engine/duplex/test_session_manager.py @@ -25,13 +25,17 @@ from vllm_omni.engine.duplex.contracts import ( DuplexStageSubmissionResult, duplex_resource_request_belongs_to_session, ) +from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer from vllm_omni.engine.duplex.events import ( + AudioDelta, DuplexEvent, ErrorEvent, + ResponseDone, SessionClosed, SessionCreated, SessionExpired, SessionHeartbeatAck, + TranscriptDelta, ) from vllm_omni.engine.duplex.messages import ( CloseDuplexSessionMessage, @@ -297,6 +301,7 @@ class Harness: output_sink: asyncio.Queue = field(default_factory=asyncio.Queue) result_sink: asyncio.Queue = field(default_factory=asyncio.Queue) _control_seq: int = 0 + outputs: dict[str, DuplexOutputBuffer] = field(default_factory=dict) @classmethod def create( @@ -330,11 +335,20 @@ class Harness: await self.manager.shutdown() async def open(self, session_id: str, config: DuplexSessionConfig | None = None) -> DuplexControlResultMessage: + limits = self.manager.runtime_config + output_buffer = self.outputs.setdefault( + session_id, + DuplexOutputBuffer( + max_bytes=limits.max_pending_output_bytes_per_session, + max_events=limits.max_pending_output_events_per_session, + ), + ) await self.manager.handle( OpenDuplexSessionMessage( control_id=f"open-{session_id}", session_id=session_id, session_config=config or DuplexSessionConfig(model="fake-model"), + output_buffer=output_buffer, ) ) return await self.result() @@ -381,8 +395,14 @@ class Harness: assert isinstance(result, DuplexControlResultMessage) return result - def events(self, session_id: str | None = None) -> list[DuplexEvent]: + async def events(self, session_id: str | None = None) -> list[DuplexEvent]: events: list[DuplexEvent] = [] + for identity, output in self.outputs.items(): + while output.pending_events: + event = await output.get() + assert event is not None and event.session_id == identity + if output.is_valid(event) and (session_id is None or identity == session_id): + events.append(event) while not self.output_sink.empty(): message = self.output_sink.get_nowait() assert isinstance(message, DuplexSessionEventMessage) @@ -396,6 +416,22 @@ class Harness: assert session is not None return session + async def begin_output_response(self, session_id: str) -> str: + """Initialize real response, text and audio projection before filling output.""" + runner = self.manager.runners[session_id] + response_id = runner.session.begin_response() + runner.session.append_assistant_text("abcdefghij") + payloads: list[dict[str, object]] = [ + {"type": "response.created", "modalities": ["audio", "text"]}, + {"type": "response.text.delta", "delta": "abcdefghij"}, + {"type": "response.output_audio.delta", "audio": "AAAAAA==", "format": "pcm16"}, + ] + for payload in payloads: + runner.emit({**payload, "response_id": response_id, "epoch": runner.session.epoch}) + await self.events(session_id) + assert not runner.run.closing + return response_id + def capture_submissions(self, session_id: str) -> list[DuplexCommand]: """Replace the runner mailbox with a list so admitted commands can be inspected.""" runner = self.manager.runners[session_id] @@ -433,7 +469,7 @@ async def test_open_answers_with_capabilities_and_emits_session_created() -> Non assert result.public_session["voice"] == "test" assert result.error_code is None - events = harness.events("sid-open") + events = await harness.events("sid-open") assert isinstance(events[0], SessionCreated) assert events[0].epoch == 0 assert events[0].session["id"] == "sid-open" @@ -587,7 +623,7 @@ async def test_plugin_runtime_config_rejection_rolls_back_the_open(error: Except assert harness.manager.get("sid-runtime") is None assert harness.manager.active_count() == 0 assert harness.stage_port.ensure_calls == [] - assert harness.events() == [] + assert await harness.events() == [] # Nothing was registered under the rejected id and the slot is free again. plugin.runtime_config_error = None @@ -603,7 +639,7 @@ async def test_command_for_unknown_session_emits_unknown_session_error() -> None async with Harness.create() as harness: harness.command("sid-missing", Heartbeat(event_id="evt-hb")) - events = harness.events("sid-missing") + events = await harness.events("sid-missing") assert len(events) == 1 error = events[0] assert isinstance(error, ErrorEvent) @@ -618,11 +654,11 @@ async def test_command_for_closed_session_emits_unknown_session_error() -> None: async with Harness.create() as harness: await harness.open("sid-closed") assert (await harness.close("sid-closed")).ok is True - harness.events() + await harness.events() harness.command("sid-closed", Heartbeat(event_id="evt-late")) - events = harness.events("sid-closed") + events = await harness.events("sid-closed") assert [type(event) for event in events] == [ErrorEvent] assert events[0].code == "unknown_session" assert events[0].related_event_id == "evt-late" @@ -631,14 +667,14 @@ async def test_command_for_closed_session_emits_unknown_session_error() -> None: async def test_admitted_command_reaches_the_session_runner_in_order() -> None: async with Harness.create() as harness: await harness.open("sid-runner") - harness.events() + await harness.events() harness.clock.advance(5.0) harness.command("sid-runner", Heartbeat(event_id="evt-1")) harness.command("sid-runner", Heartbeat(event_id="evt-2")) await _settle() - events = harness.events("sid-runner") + events = await harness.events("sid-runner") assert [type(event) for event in events] == [SessionHeartbeatAck, SessionHeartbeatAck] assert all(event.epoch == 0 for event in events) assert harness.session("sid-runner").lease.last_activity == 5.0 @@ -647,7 +683,7 @@ async def test_admitted_command_reaches_the_session_runner_in_order() -> None: async def test_append_audio_is_backpressured_by_pending_input_bytes() -> None: async with Harness.create(max_pending_input_bytes_per_session=8) as harness: await harness.open("sid-bytes") - harness.events() + await harness.events() session = harness.session("sid-bytes") submitted = harness.capture_submissions("sid-bytes") assert session.reserve_input_bytes(6, limit=8) @@ -655,7 +691,7 @@ async def test_append_audio_is_backpressured_by_pending_input_bytes() -> None: harness.command("sid-bytes", AppendAudio(audio=b"pcm", event_id="evt-over")) harness.command("sid-bytes", AppendAudio(audio=b"pc", event_id="evt-fits")) - events = harness.events("sid-bytes") + events = await harness.events("sid-bytes") assert len(events) == 1 assert isinstance(events[0], ErrorEvent) assert events[0].code == "input_backpressure" @@ -670,7 +706,7 @@ async def test_append_audio_is_backpressured_by_pending_input_bytes() -> None: async def test_commit_holds_a_pending_turn_reservation_until_the_runner_dequeues_it() -> None: async with Harness.create(max_pending_turns_per_session=1) as harness: await harness.open("sid-turns") - harness.events() + await harness.events() session = harness.session("sid-turns") submitted = harness.capture_submissions("sid-turns") @@ -679,7 +715,7 @@ async def test_commit_holds_a_pending_turn_reservation_until_the_runner_dequeues assert [command.event_id for command in submitted] == ["evt-commit-1"] assert session.pending_input_turns == 1 - events = harness.events("sid-turns") + events = await harness.events("sid-turns") assert len(events) == 1 assert isinstance(events[0], ErrorEvent) assert events[0].code == "input_backpressure" @@ -689,7 +725,7 @@ async def test_commit_holds_a_pending_turn_reservation_until_the_runner_dequeues session.release_pending_turn() harness.command("sid-turns", Commit(event_id="evt-commit-3")) assert [command.event_id for command in submitted] == ["evt-commit-1", "evt-commit-3"] - assert harness.events("sid-turns") == [] + assert await harness.events("sid-turns") == [] def test_accepts_only_duplex_message_types() -> None: @@ -704,7 +740,14 @@ def test_accepts_only_duplex_message_types() -> None: ) try: config = DuplexSessionConfig() - assert manager.accepts(OpenDuplexSessionMessage(control_id="c", session_id="s", session_config=config)) + assert manager.accepts( + OpenDuplexSessionMessage( + control_id="c", + session_id="s", + session_config=config, + output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512), + ) + ) assert manager.accepts(CloseDuplexSessionMessage(control_id="c", session_id="s")) assert manager.accepts(ResumeDuplexSessionMessage(control_id="c", session_id="s", expected_lease_generation=0)) assert manager.accepts(TouchDuplexSessionMessage(control_id="c", session_id="s", activity="heartbeat")) @@ -729,7 +772,7 @@ def test_accepts_only_duplex_message_types() -> None: async def test_close_releases_stage_resources_and_is_idempotent() -> None: async with Harness.create() as harness: await harness.open("sid-close") - harness.events() + await harness.events() session = harness.session("sid-close") request_id = stage0_request_id("sid-close") session.bind_stage_request(0, request_id, fence=session.fence) @@ -750,7 +793,7 @@ async def test_close_releases_stage_resources_and_is_idempotent() -> None: assert harness.manager.runner_for_request_id(request_id) is None assert harness.stage_port.cleanup_calls == [([request_id], True), ([reserved_context.request_id], False)] assert harness.plugin.data_plane.closed_sessions == ["sid-close"] - closed = [event for event in harness.events("sid-close") if isinstance(event, SessionClosed)] + closed = [event for event in await harness.events("sid-close") if isinstance(event, SessionClosed)] assert len(closed) == 1 assert closed[0].reason == "client_close" assert closed[0].is_terminal is True @@ -760,7 +803,7 @@ async def test_close_releases_stage_resources_and_is_idempotent() -> None: assert again.lease_generation is None assert again.public_session is None assert harness.stage_port.cleanup_calls[2:] == [] - assert harness.events() == [] + assert await harness.events() == [] async def test_close_retains_the_admission_slot_until_stage_cleanup_succeeds() -> None: @@ -858,7 +901,7 @@ async def test_reaper_expires_idle_sessions_and_emits_session_expired() -> None: async with Harness.create(idle_ttl_s=2.0, disconnect_grace_s=1.0) as harness: await harness.open("sid-idle") await harness.open("sid-active") - harness.events() + await harness.events() harness.clock.advance(1.0) assert (await harness.touch("sid-active", DuplexLeaseActivity.HEARTBEAT.value)).ok is True harness.clock.advance(1.1) @@ -868,7 +911,7 @@ async def test_reaper_expires_idle_sessions_and_emits_session_expired() -> None: assert harness.manager.get("sid-idle") is None assert harness.manager.get("sid-active") is not None assert harness.manager.active_count() == 1 - expired = harness.events("sid-idle") + expired = await harness.events("sid-idle") assert len(expired) == 1 assert isinstance(expired[0], SessionExpired) assert expired[0].reason == "idle_ttl_expired" @@ -878,7 +921,7 @@ async def test_reaper_expires_idle_sessions_and_emits_session_expired() -> None: "session_id": "sid-idle", "reason": "idle_ttl_expired", } - assert harness.events("sid-active") == [] + assert await harness.events("sid-active") == [] assert harness.stage_port.cleanup_calls == [([stage0_request_id("sid-idle")], False)] assert harness.result_sink.empty() @@ -888,7 +931,7 @@ async def test_reaper_expires_idle_sessions_and_emits_session_expired() -> None: async def test_reaper_expires_detached_sessions_after_the_disconnect_grace() -> None: async with Harness.create(idle_ttl_s=300.0, disconnect_grace_s=1.0) as harness: await harness.open("sid-detached") - harness.events() + await harness.events() assert (await harness.touch("sid-detached", DuplexLeaseActivity.DETACH.value)).ok is True harness.clock.advance(0.5) assert await harness.manager.reap_expired() == 0 @@ -897,7 +940,7 @@ async def test_reaper_expires_detached_sessions_after_the_disconnect_grace() -> assert await harness.manager.reap_expired() == 1 assert harness.manager.get("sid-detached") is None - expired = harness.events("sid-detached") + expired = await harness.events("sid-detached") assert [type(event) for event in expired] == [SessionExpired] assert expired[0].reason == "disconnect_grace_expired" assert (await harness.resume("sid-detached", expected_lease_generation=0)).error_code == "unknown_session" @@ -906,7 +949,7 @@ async def test_reaper_expires_detached_sessions_after_the_disconnect_grace() -> async def test_session_closed_is_emitted_only_after_stage_cleanup_so_a_reopen_is_admitted() -> None: async with Harness.create(max_sessions=1) as harness: await harness.open("sid-closing") - harness.events() + await harness.events() session = harness.session("sid-closing") session.bind_stage_request(0, stage0_request_id("sid-closing"), fence=session.fence) harness.stage_port.cleanup_gate = asyncio.Event() @@ -915,11 +958,11 @@ async def test_session_closed_is_emitted_only_after_stage_cleanup_so_a_reopen_is await asyncio.wait_for(harness.stage_port.cleanup_started.wait(), timeout=1.0) # The runner is gone but the stage request is still being aborted: # no session.closed yet, so a client cannot race the admission slot. - assert harness.events("sid-closing") == [] + assert await harness.events("sid-closing") == [] harness.stage_port.cleanup_gate.set() assert (await close_task).ok is True - events = harness.events("sid-closing") + events = await harness.events("sid-closing") assert [type(event) for event in events] == [SessionClosed] assert events[0].reason == "client_close" assert (await harness.open("sid-replacement")).ok is True @@ -928,7 +971,7 @@ async def test_session_closed_is_emitted_only_after_stage_cleanup_so_a_reopen_is async def test_wire_close_command_frees_the_slot_before_session_closed() -> None: async with Harness.create(max_sessions=1) as harness: await harness.open("sid-wire-close") - harness.events() + await harness.events() session = harness.session("sid-wire-close") request_id = stage0_request_id("sid-wire-close") session.bind_stage_request(0, request_id, fence=session.fence) @@ -937,7 +980,7 @@ async def test_wire_close_command_frees_the_slot_before_session_closed() -> None harness.command("sid-wire-close", CloseSession(reason="client_close")) await asyncio.wait_for(harness.stage_port.cleanup_started.wait(), timeout=1.0) assert harness.manager.get("sid-wire-close") is None - assert harness.events("sid-wire-close") == [] + assert await harness.events("sid-wire-close") == [] harness.stage_port.cleanup_gate.set() for _ in range(50): @@ -945,7 +988,7 @@ async def test_wire_close_command_frees_the_slot_before_session_closed() -> None break await asyncio.sleep(0.01) assert harness.stage_port.cleanup_calls == [([request_id], True)] - events = harness.events("sid-wire-close") + events = await harness.events("sid-wire-close") assert [type(event) for event in events] == [SessionClosed] assert events[0].reason == "client_close" assert (await harness.open("sid-replacement")).ok is True @@ -954,7 +997,7 @@ async def test_wire_close_command_frees_the_slot_before_session_closed() -> None async def test_wire_close_stops_the_worker_and_refuses_later_control_ops() -> None: async with Harness.create(max_sessions=1) as harness: await harness.open("sid-wire") - harness.events() + await harness.events() runner = harness.manager.runners["sid-wire"] worker = runner._worker assert worker is not None @@ -975,7 +1018,7 @@ async def test_wire_close_stops_the_worker_and_refuses_later_control_ops() -> No break await asyncio.sleep(0.01) assert worker.done(), "the mailbox worker must exit after a wire close" - events = harness.events("sid-wire") + events = await harness.events("sid-wire") assert [type(event) for event in events] == [ErrorEvent, SessionClosed] or [ type(event) for event in events ] == [SessionClosed, ErrorEvent] @@ -1003,7 +1046,7 @@ async def test_concurrent_opens_cannot_exceed_max_sessions() -> None: async def test_runtime_close_without_stage_requests_frees_the_slot() -> None: async with Harness.create(max_sessions=1) as harness: await harness.open("sid-runtime") - harness.events() + await harness.events() runner = harness.manager.runners["sid-runtime"] runner.session.release_all_requests() @@ -1014,7 +1057,7 @@ async def test_runtime_close_without_stage_requests_frees_the_slot() -> None: await asyncio.sleep(0.01) assert harness.manager.get("sid-runtime") is None - assert [type(event) for event in harness.events("sid-runtime")] == [SessionClosed] + assert [type(event) for event in await harness.events("sid-runtime")] == [SessionClosed] assert (await harness.open("sid-next")).ok is True @@ -1028,7 +1071,7 @@ async def test_close_of_a_forgotten_session_is_idempotent() -> None: async def test_append_bytes_are_reserved_at_admission_until_the_runner_dequeues() -> None: async with Harness.create(max_sessions=1, max_pending_input_bytes_per_session=8) as harness: await harness.open("sid-bytes") - harness.events() + await harness.events() session = harness.session("sid-bytes") runner = harness.manager.runners["sid-bytes"] # Park the worker so admitted commands stay in the mailbox. @@ -1040,7 +1083,7 @@ async def test_append_bytes_are_reserved_at_admission_until_the_runner_dequeues( harness.command("sid-bytes", AppendAudio(audio=b"12345")) assert session.pending_input_bytes == 5 harness.command("sid-bytes", AppendAudio(audio=b"12345")) - errors = [event for event in harness.events("sid-bytes") if isinstance(event, ErrorEvent)] + errors = [event for event in await harness.events("sid-bytes") if isinstance(event, ErrorEvent)] assert [error.code for error in errors] == ["input_backpressure"] gate.set() @@ -1049,7 +1092,7 @@ async def test_append_admission_counts_audio_and_video_frame_bytes() -> None: """Manager reserves len(audio)+Σlen(frame); video bytes count toward the same limit.""" async with Harness.create(max_sessions=1, max_pending_input_bytes_per_session=20) as harness: await harness.open("sid-av") - harness.events() + await harness.events() session = harness.session("sid-av") runner = harness.manager.runners["sid-av"] gate = asyncio.Event() @@ -1072,7 +1115,7 @@ async def test_append_admission_counts_audio_and_video_frame_bytes() -> None: "sid-av", AppendAudio(audio=b"x" * 10, video_frames=("yyyyyyyyyy",), event_id="evt-over"), ) - errors = [event for event in harness.events("sid-av") if isinstance(event, ErrorEvent)] + errors = [event for event in await harness.events("sid-av") if isinstance(event, ErrorEvent)] assert [error.code for error in errors] == ["input_backpressure"] assert errors[0].related_event_id == "evt-over" assert session.pending_input_bytes == expected @@ -1082,7 +1125,7 @@ async def test_append_admission_counts_audio_and_video_frame_bytes() -> None: async def test_expired_session_retains_the_admission_slot_until_cleanup_succeeds() -> None: async with Harness.create(max_sessions=1, idle_ttl_s=1.0) as harness: await harness.open("sid-expired") - harness.events() + await harness.events() session = harness.session("sid-expired") request_id = stage0_request_id("sid-expired") session.bind_stage_request(0, request_id, fence=session.fence) @@ -1093,14 +1136,14 @@ async def test_expired_session_retains_the_admission_slot_until_cleanup_succeeds assert harness.manager.get("sid-expired") is None assert session.state == DuplexSessionState.CLOSED - assert [type(event) for event in harness.events("sid-expired")] == [SessionExpired] + assert [type(event) for event in await harness.events("sid-expired")] == [SessionExpired] assert (await harness.open("sid-replacement")).error_code == "resource_exhausted" assert await harness.manager.reap_expired() == 1 assert harness.stage_port.cleanup_calls == [([request_id], True)] * 2 assert session.resource_request_ids() == [] - assert harness.events("sid-expired") == [] + assert await harness.events("sid-expired") == [] assert (await harness.open("sid-replacement")).ok is True @@ -1131,13 +1174,13 @@ async def test_expired_cleanup_failure_does_not_block_other_sessions() -> None: async with Harness.create(stage_port=OneStuckStagePort(), idle_ttl_s=1.0) as harness: await harness.open("sid-stuck") await harness.open("sid-fine") - harness.events() + await harness.events() harness.clock.advance(2.0) assert await harness.manager.reap_expired() == 1 assert harness.manager.active_count() == 0 - assert {type(event) for event in harness.events()} == {SessionExpired} + assert {type(event) for event in await harness.events()} == {SessionExpired} assert (await harness.open("sid-fine")).ok is True assert (await harness.open("sid-stuck")).error_code == "session_exists" @@ -1150,7 +1193,7 @@ async def test_expired_cleanup_failure_does_not_block_other_sessions() -> None: async def test_close_sessions_for_request_ids_expires_the_owner_and_retries_cleanup() -> None: async with Harness.create(max_sessions=1) as harness: await harness.open("sid-request") - harness.events() + await harness.events() session = harness.session("sid-request") request_id = stage0_request_id("sid-request") session.bind_stage_request(0, request_id, fence=session.fence) @@ -1162,7 +1205,7 @@ async def test_close_sessions_for_request_ids_expires_the_owner_and_retries_clea assert harness.manager.get("sid-request") is None assert session.lease.terminal_reason == "request_cleanup" await _settle() - expired = harness.events("sid-request") + expired = await harness.events("sid-request") assert [type(event) for event in expired] == [SessionExpired] assert expired[0].reason == "request_cleanup" assert session.state == DuplexSessionState.CLOSED @@ -1256,7 +1299,12 @@ async def test_control_dispatch_is_ordered_per_session_without_blocking_other_se blocked_config = DuplexSessionConfig(model="fake-model", instructions="sid-blocked") manager.dispatch( - OpenDuplexSessionMessage(control_id="blocked-open", session_id="sid-blocked", session_config=blocked_config) + OpenDuplexSessionMessage( + control_id="blocked-open", + session_id="sid-blocked", + session_config=blocked_config, + output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512), + ) ) await asyncio.wait_for(plugin.runtime_config_started.wait(), timeout=1.0) manager.dispatch( @@ -1271,6 +1319,7 @@ async def test_control_dispatch_is_ordered_per_session_without_blocking_other_se control_id="independent-open", session_id="sid-independent", session_config=DuplexSessionConfig(model="fake-model"), + output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512), ) ) manager.dispatch( @@ -1338,16 +1387,18 @@ async def test_shutdown_closes_every_runner_and_stops_dispatch_tasks() -> None: await harness.open("sid-a") await harness.open("sid-b") sessions = [harness.session("sid-a"), harness.session("sid-b")] - harness.events() + await harness.events() await harness.manager.shutdown() assert harness.manager.active_count() == 0 assert harness.manager.runners == {} + assert harness.manager._outputs == {} + assert harness.manager._output_failed == set() assert all(session.state == DuplexSessionState.CLOSED for session in sessions) assert harness.manager.runner_for_request_id(stage0_request_id("sid-b")) is None # Shutdown is silent: no session.closed / session.expired is emitted. - assert harness.events() == [] + assert await harness.events() == [] assert harness.stage_port.cleanup_calls == [] @@ -1414,12 +1465,12 @@ async def test_expiry_emits_a_terminal_event_even_when_a_close_deferred_it() -> runner = harness.manager.runners["sid-deferred"] # What the runner does for a wire close: tear down, defer the terminal. await runner.close("client_close", emit_closed=False) - harness.events("sid-deferred") + await harness.events("sid-deferred") await runner.expire("request_cleanup", emit_expired=True) terminal = [ - event for event in harness.events("sid-deferred") if isinstance(event, SessionClosed | SessionExpired) + event for event in await harness.events("sid-deferred") if isinstance(event, SessionClosed | SessionExpired) ] assert len(terminal) == 1, terminal assert isinstance(terminal[0], SessionExpired) @@ -1466,7 +1517,7 @@ async def test_a_runtime_close_racing_a_manager_close_still_emits_one_terminal() """ async with Harness.create(max_sessions=1) as harness: await harness.open("sid-race") - harness.events() + await harness.events() session = harness.session("sid-race") session.bind_stage_request(0, stage0_request_id("sid-race"), fence=session.fence) runner = harness.manager.runners["sid-race"] @@ -1486,7 +1537,9 @@ async def test_a_runtime_close_racing_a_manager_close_still_emits_one_terminal() await harness.close("sid-race") - terminal = [event for event in harness.events("sid-race") if isinstance(event, SessionClosed | SessionExpired)] + terminal = [ + event for event in await harness.events("sid-race") if isinstance(event, SessionClosed | SessionExpired) + ] assert len(terminal) == 1, terminal @@ -1498,7 +1551,9 @@ async def test_a_terminal_event_is_emitted_once_even_if_expiry_runs_twice() -> N await runner.expire("idle_ttl_expired", emit_expired=True) await runner.expire("idle_ttl_expired", emit_expired=True) - terminal = [event for event in harness.events("sid-once") if isinstance(event, SessionClosed | SessionExpired)] + terminal = [ + event for event in await harness.events("sid-once") if isinstance(event, SessionClosed | SessionExpired) + ] assert len(terminal) == 1, terminal @@ -1590,12 +1645,12 @@ async def test_wire_close_of_an_idle_session_still_emits_session_closed() -> Non """ async with Harness.create(max_sessions=1) as harness: await harness.open("sid-idle-close") - harness.events() + await harness.events() harness.command("sid-idle-close", CloseSession(reason="client_close")) events: list[DuplexEvent] = [] for _ in range(200): - events.extend(harness.events("sid-idle-close")) + events.extend(await harness.events("sid-idle-close")) if events: break await asyncio.sleep(0.01) @@ -1725,11 +1780,13 @@ async def test_an_open_cancelled_after_admission_releases_the_runner_and_its_sta control_id="open-cancelled", session_id="sid-cancelled", session_config=DuplexSessionConfig(model="fake-model"), + output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512), ) ) ) await asyncio.wait_for(sink.entered.wait(), timeout=1.0) assert "sid-cancelled" in harness.manager.runners + assert "sid-cancelled" in harness.manager._outputs assert [context.request_id for context in harness.stage_port.ensure_calls] == [ stage0_request_id("sid-cancelled") ] @@ -1740,6 +1797,7 @@ async def test_an_open_cancelled_after_admission_releases_the_runner_and_its_sta harness.manager._result_sink = original_sink assert "sid-cancelled" not in harness.manager.runners + assert "sid-cancelled" not in harness.manager._outputs assert harness.manager.active_count() == 0 assert harness.stage_port.cleanup_calls == [([stage0_request_id("sid-cancelled")], False)] assert stage0_request_id("sid-cancelled") not in harness.manager._request_index @@ -1764,6 +1822,7 @@ async def test_an_open_cancelled_again_during_its_rollback_leaves_the_stage_clea control_id="open-cancelled-twice", session_id="sid-cancelled", session_config=DuplexSessionConfig(model="fake-model"), + output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512), ) ) ) @@ -1779,6 +1838,7 @@ async def test_an_open_cancelled_again_during_its_rollback_leaves_the_stage_clea open_task.cancel() await asyncio.wait_for(shutdown_started.wait(), timeout=1.0) + assert "sid-cancelled" not in harness.manager._outputs open_task.cancel() # lands inside the rollback, before the stage cleanup with pytest.raises(asyncio.CancelledError): await open_task @@ -1875,3 +1935,127 @@ async def test_a_resume_replayed_under_its_control_id_reports_the_generation_it_ other = await harness.resume("sid-replay", expected_lease_generation=0, control_id="rpc-b") assert other.ok is False and other.error_code == "session_resume_conflict" + + +async def test_output_overflow_closes_only_its_session_and_releases_capacity(mocker) -> None: + async with Harness.create(max_sessions=2, max_pending_output_events_per_session=8) as harness: + assert (await harness.open("slow")).ok + assert (await harness.open("normal")).ok + await harness.events() + assert not harness.manager.runners["slow"].run.closing + assert not harness.manager.runners["normal"].run.closing + slow = harness.session("slow") + normal = harness.session("normal") + slow_response = await harness.begin_output_response("slow") + normal_response = normal.begin_response() + request_id = stage0_request_id("slow") + slow.bind_request(request_id) + slow.mark_audio_sent(1000) + slow.acknowledge_playback(500, 500) + close_stream = mocker.spy(harness.plugin.data_plane, "close_stream") + + for _ in range(5): + harness.manager.runners["slow"].emit( + { + "type": "response.output_audio.delta", + "response_id": slow_response, + "audio": "AAAAAA==", + "format": "pcm16", + } + ) + harness.manager.emit(normal, [AudioDelta(response_id=normal_response, delta="CCCC")]) + await asyncio.wait_for(asyncio.gather(*tuple(harness.manager._dispatched_control_tasks)), timeout=2.0) + + events = await harness.events() + failed = [event for event in events if event.session_id == "slow"] + # Queued transcript deltas remain ordered; stale audio and the failed + # response's undelivered completion markers are discarded. + assert [type(event) for event in failed] == [TranscriptDelta] * 4 + [ErrorEvent, ResponseDone, SessionClosed] + assert failed[-3].code == "output_backpressure" + done = failed[-2] + assert isinstance(done, ResponseDone) + assert done.response_id == slow_response and done.status == "failed" + metadata = done.response["metadata"] + assert isinstance(metadata, dict) and metadata["committed"] is False + assert failed[-1].reason == "output_backpressure" + assert slow.history == () + assert slow.active_response_id is None + assert harness.stage_port.abort_calls == [[request_id]] + close_stream.assert_called_once_with(request_id) + delivered = [event for event in events if event.session_id == "normal"] + assert len(delivered) == 1 and isinstance(delivered[0], AudioDelta) + assert delivered[0].response_id == normal_response and delivered[0].delta == "CCCC" + assert harness.manager.get("slow") is None + assert harness.manager.get("normal") is normal + assert harness.manager.active_count() == 1 + assert (await harness.open("replacement")).ok + + +@pytest.mark.parametrize( + "capacity,end_before_emit,expected_status,closed", + [ + (8, True, "failed", True), + (8, False, "failed", True), + (14, True, "completed", True), + (64, True, "completed", False), + ], + ids=["already-ended-overflow", "active-overflow", "overflow-after-ending", "sufficient-capacity"], +) +async def test_completion_projection_retains_exactly_one_response_done( + capacity: int, end_before_emit: bool, expected_status: str, closed: bool +) -> None: + async with Harness.create(max_pending_output_events_per_session=capacity) as harness: + assert (await harness.open("completion")).ok + await harness.events() + response_id = await harness.begin_output_response("completion") + runner = harness.manager.runners["completion"] + for _ in range(4): + runner.emit( + { + "type": "response.output_audio.delta", + "response_id": response_id, + "audio": "AAAAAA==", + "format": "pcm16", + } + ) + assert harness.outputs["completion"].pending_events == 8 + runner.session.mark_audio_sent(1000) + runner.session.acknowledge_playback(500, 500) + if end_before_emit: + message = runner.session.end_response(commit_text=True) + assert message == {"role": "assistant", "content": "abcde"} + runner.session.register_history_item(f"item_{response_id}", message) + assert runner.session.history == (message,) + + runner.emit( + {"type": "response.done", "response_id": response_id, "status": "completed", "committed": end_before_emit} + ) + await asyncio.wait_for(asyncio.gather(*tuple(harness.manager._dispatched_control_tasks)), timeout=2.0) + + events = await harness.events() + endings = [event for event in events if isinstance(event, ResponseDone)] + assert len(endings) == 1 + assert endings[0].response_id == response_id and endings[0].status == expected_status + metadata = endings[0].response["metadata"] + assert isinstance(metadata, dict) and metadata["committed"] is (expected_status == "completed") + assert any(isinstance(event, SessionClosed) for event in events) is closed + if expected_status == "failed": + assert [type(event) for event in events] == [TranscriptDelta] * 4 + [ + ErrorEvent, + ResponseDone, + SessionClosed, + ] + assert endings[0].response["status_details"] == {"type": "failed", "reason": "output_backpressure"} + assert runner.session.history == () + else: + assert [type(event) for event in events[:8]] == [AudioDelta, TranscriptDelta] * 4 + assert runner.session.history == ({"role": "assistant", "content": "abcde"},) + if closed: + # The ending was accepted; overflowing on rate_limits.updated + # must not replace it, emit a second ending or undo its history. + assert [type(event) for event in events[-3:]] == [ResponseDone, ErrorEvent, SessionClosed] + assert events[-3] is endings[0] + else: + assert not any(isinstance(event, ErrorEvent) for event in events) + assert events[-2] is endings[0] + assert events[-1].type == "rate_limits.updated" diff --git a/tests/engine/duplex/test_session_runner.py b/tests/engine/duplex/test_session_runner.py index c51fc243c..af14c8ed9 100644 --- a/tests/engine/duplex/test_session_runner.py +++ b/tests/engine/duplex/test_session_runner.py @@ -6,7 +6,7 @@ The runner is driven exactly the way ``DuplexOrchestrator`` drives it: typed commands go through ``DuplexSessionManager.dispatch``, stage outputs are pushed with ``runner.on_stage_output`` and everything the session says is read back -from the manager's output sink as typed events. The stage port is a recording +from the session's output buffer as typed events. The stage port is a recording fake; the model plugin is the real MiniCPM-o 4.5 one so append planning and output projection are exercised end to end. """ @@ -38,7 +38,8 @@ from vllm_omni.engine.duplex.contracts import ( DuplexStageSubmissionResult, duplex_resource_request_id, ) -from vllm_omni.engine.duplex.events import DuplexEvent +from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer +from vllm_omni.engine.duplex.events import AudioDelta, DuplexEvent from vllm_omni.engine.duplex.messages import ( CloseDuplexSessionMessage, DuplexControlResultMessage, @@ -72,6 +73,8 @@ class RecordingStagePort(DuplexStagePort): self.submissions: list[DuplexStageSubmission] = [] self.cleanups: list[tuple[list[str], bool]] = [] self.aborts: list[list[str]] = [] + self.abort_started = asyncio.Event() + self.abort_gate: asyncio.Event | None = None self.fail_submit: Exception | None = None #: When set, ``submit`` parks on it after signalling ``submit_started``. self.submit_gate: asyncio.Event | None = None @@ -105,6 +108,9 @@ class RecordingStagePort(DuplexStagePort): async def abort_requests(self, request_ids: list[str]) -> None: self.aborts.append(list(request_ids)) + self.abort_started.set() + if self.abort_gate is not None: + await self.abort_gate.wait() def _fake_encode_audio(audio: object, sample_rate_hz: int, response_format: str, speed: float | None) -> str | None: @@ -120,6 +126,7 @@ class Harness: manager: DuplexSessionManager port: RecordingStagePort output: asyncio.Queue[Any] + output_buffer: DuplexOutputBuffer results: asyncio.Queue[Any] runner: DuplexSessionRunner events: list[DuplexEvent] = field(default_factory=list) @@ -139,6 +146,12 @@ class Harness: collected: list[DuplexEvent] = [] while True: drained = False + while self.output_buffer.pending_events: + event = await self.output_buffer.get() + assert event is not None + if self.output_buffer.is_valid(event): + collected.append(event) + drained = True while not self.output.empty(): message = self.output.get_nowait() if isinstance(message, DuplexSessionEventMessage): @@ -215,12 +228,17 @@ async def open_harness( port = RecordingStagePort(stage_count=stage_count) output: asyncio.Queue[Any] = asyncio.Queue() results: asyncio.Queue[Any] = asyncio.Queue() + limits = runtime_config or DuplexSessionRuntimeConfig() + output_buffer = DuplexOutputBuffer( + max_bytes=limits.max_pending_output_bytes_per_session, + max_events=limits.max_pending_output_events_per_session, + ) manager = DuplexSessionManager( plugin=plugin, stage_port=port, output_sink=output, result_sink=results, - runtime_config=runtime_config or DuplexSessionRuntimeConfig(), + runtime_config=limits, model_config=None, log_stats=log_stats, clock=clock, @@ -232,10 +250,21 @@ async def open_harness( instructions="You are a concise assistant.", extra_body=body, ) - await manager.handle(OpenDuplexSessionMessage(control_id="c-open", session_id=SESSION_ID, session_config=config)) + await manager.handle( + OpenDuplexSessionMessage( + control_id="c-open", session_id=SESSION_ID, session_config=config, output_buffer=output_buffer + ) + ) result = await asyncio.wait_for(results.get(), timeout=2.0) assert isinstance(result, DuplexControlResultMessage) and result.ok, result - harness = Harness(manager=manager, port=port, output=output, results=results, runner=manager.runners[SESSION_ID]) + harness = Harness( + manager=manager, + port=port, + output=output, + output_buffer=output_buffer, + results=results, + runner=manager.runners[SESSION_ID], + ) await harness.settle() return harness @@ -710,6 +739,8 @@ async def test_stale_epoch_output_is_dropped_after_barge_in() -> None: async def test_barge_in_aborts_draining_tts_as_well_as_the_active_request() -> None: h = await open_harness() + release_abort = asyncio.Event() + h.port.abort_gate = release_abort try: await h.run(append_audio()) request_id = h.stage0_request_id() @@ -730,7 +761,24 @@ async def test_barge_in_aborts_draining_tts_as_well_as_the_active_request() -> N h.session.request_resources[(1, "still-live")] = DuplexRequestResource( stage_id=1, request_id="still-live", fence=older, submitted=True ) - events = await h.run(commands.BargeIn()) + h.manager.emit( + h.session, + [ + AudioDelta(response_id="resp-old", delta="AAAA"), + AudioDelta(response_id="resp-old", delta="BBBB"), + AudioDelta(response_id=h.session.active_response_id, delta="CCCC"), + ], + ) + held = await h.output_buffer.get() + assert held is not None and h.output_buffer.is_valid(held) + h.submit(commands.BargeIn()) + await asyncio.wait_for(h.port.abort_started.wait(), timeout=2.0) + # Both queued audio and an event held by the consumer become stale + # before the stage abort can yield, including the older draining turn. + assert not h.output_buffer.is_valid(held) + assert h.output_buffer.pending_events == 0 + release_abort.set() + events = await h.settle() assert h.port.aborts == [[request_id, "duplex-drain-tts"]] assert not h.session.is_draining_request("duplex-drain-tts") done_ids = [ @@ -744,6 +792,36 @@ async def test_barge_in_aborts_draining_tts_as_well_as_the_active_request() -> N assert (0, request_id) not in h.session.request_resources assert (1, "still-live") in h.session.request_resources finally: + release_abort.set() + await close_harness(h) + + +@pytest.mark.asyncio +async def test_close_invalidates_draining_audio_before_stage_abort_finishes() -> None: + h = await open_harness() + release_abort = asyncio.Event() + h.port.abort_gate = release_abort + try: + h.session.bind_draining_request("duplex-drain-tts", "resp-old") + h.manager.emit( + h.session, + [AudioDelta(response_id="resp-old", delta="AAAA"), AudioDelta(response_id="resp-old", delta="BBBB")], + ) + held = await h.output_buffer.get() + assert held is not None and h.output_buffer.is_valid(held) + h.submit(commands.CloseSession()) + await asyncio.wait_for(h.port.abort_started.wait(), timeout=2.0) + # Session close suppresses audio.cancelled notifications; it must still + # invalidate the pending output without relying on those later events. + assert not h.output_buffer.is_valid(held) + assert h.output_buffer.pending_events == 0 + assert h.port.aborts == [["duplex-drain-tts"]] + release_abort.set() + events = await h.settle() + assert "session.closed" in types(events) + assert h.manager.active_count() == 0 + finally: + release_abort.set() await close_harness(h) diff --git a/tests/engine/test_duplex_omni_engine.py b/tests/engine/test_duplex_omni_engine.py index 06af9e7af..902972735 100644 --- a/tests/engine/test_duplex_omni_engine.py +++ b/tests/engine/test_duplex_omni_engine.py @@ -14,6 +14,7 @@ import pytest from vllm_omni.engine.duplex.commands import Heartbeat from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig +from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer from vllm_omni.engine.duplex.messages import ( CloseDuplexSessionMessage, DuplexControlResultMessage, @@ -72,14 +73,16 @@ async def test_open_session_runs_the_correlated_rpc_with_the_typed_config() -> N result = _ok("open", capabilities=DuplexCapabilities(), public_session={"id": "sid"}, lease_generation=0) engine = _engine(result) config = DuplexSessionConfig(model="m") + output_buffer = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512) - returned = await engine.open_session_async("sid", config, timeout=3.0) + returned = await engine.open_session_async("sid", config, output_buffer=output_buffer, timeout=3.0) assert returned is result ((key, message, timeout),) = engine.rpc_client.calls assert key == ("duplex", message.control_id) assert isinstance(message, OpenDuplexSessionMessage) assert message.session_id == "sid" and message.session_config is config + assert message.output_buffer is output_buffer assert timeout == 3.0 # A momentarily full request queue is backpressure, not a failed open: the # control RPC blocks rather than raising an untyped queue.Full at the caller. @@ -116,7 +119,9 @@ async def test_failed_control_result_maps_to_a_typed_session_error() -> None: ) ) with pytest.raises(DuplexSessionError) as excinfo: - await engine.open_session_async("sid", DuplexSessionConfig()) + await engine.open_session_async( + "sid", DuplexSessionConfig(), output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512) + ) assert excinfo.value.code == "resource_exhausted" assert excinfo.value.retryable is True assert excinfo.value.session_id == "sid" @@ -133,7 +138,9 @@ async def test_engine_error_and_timeout_map_to_session_errors() -> None: engine = _engine(TimeoutError("duplex open timed out for session sid")) with pytest.raises(DuplexSessionError) as excinfo: - await engine.open_session_async("sid", DuplexSessionConfig()) + await engine.open_session_async( + "sid", DuplexSessionConfig(), output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512) + ) assert excinfo.value.code == "timeout" and excinfo.value.retryable is True diff --git a/tests/engine/test_duplex_orchestrator.py b/tests/engine/test_duplex_orchestrator.py index 2b0c646e6..271b19f0e 100644 --- a/tests/engine/test_duplex_orchestrator.py +++ b/tests/engine/test_duplex_orchestrator.py @@ -24,6 +24,7 @@ from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig from vllm_omni.engine.duplex import commands from vllm_omni.engine.duplex.config import DuplexSessionConfig, DuplexSessionState from vllm_omni.engine.duplex.contracts import DuplexFence, duplex_resource_request_id +from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer from vllm_omni.engine.duplex.messages import ( CloseDuplexSessionMessage, DuplexControlResultMessage, @@ -78,10 +79,15 @@ def _build( return orchestrator, clients, rpc_q, output_q -def _open_message(extra_body: dict[str, object] | None = None) -> OpenDuplexSessionMessage: +def _open_message( + extra_body: dict[str, object] | None = None, + *, + output_buffer: DuplexOutputBuffer | None = None, +) -> OpenDuplexSessionMessage: return OpenDuplexSessionMessage( control_id=f"open-{SESSION_ID}", session_id=SESSION_ID, + output_buffer=output_buffer or DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512), session_config=DuplexSessionConfig( model="openbmb/MiniCPM-o-4_5", modalities=["text"], @@ -369,12 +375,13 @@ async def test_forwarded_stage_requests_are_bound_and_barge_in_aborts_them() -> @pytest.mark.asyncio async def test_session_owned_outputs_reach_the_runner_and_never_the_client_queue() -> None: orchestrator, _, rpc_q, output_q = _build(stages=2) - await _open(orchestrator, rpc_q) + output_buffer = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512) + await _open(orchestrator, rpc_q, output_buffer=output_buffer) request_id = _stage0_request_id() await _submit(orchestrator, _append_audio()) request_state = orchestrator.request_states[request_id] - while not output_q.empty(): - output_q.get_nowait() + while output_buffer.pending_events: + await output_buffer.get() consumed = await orchestrator._intercept_stage_output(1, 0, _tts_output(request_id), request_state, None, None) await _settle(orchestrator) @@ -382,8 +389,9 @@ async def test_session_owned_outputs_reach_the_runner_and_never_the_client_queue assert consumed is True session = orchestrator.session_manager.get(SESSION_ID) assert session is not None and session.active_response_id is not None - types = [message.event.type for message in [output_q.get_nowait() for _ in range(output_q.qsize())]] + types = [(await output_buffer.get()).type for _ in range(output_buffer.pending_events)] assert "response.created" in types and "response.output_audio.delta" in types + assert output_q.empty() orphan = DuplexOrchestratorRequestState( request_id="duplex-s.b3RoZXI.e.0.r.stage0", @@ -401,7 +409,8 @@ async def test_session_owned_outputs_reach_the_runner_and_never_the_client_queue @pytest.mark.asyncio async def test_forward_failure_closes_the_owning_session() -> None: orchestrator, clients, rpc_q, output_q = _build(stages=2) - await _open(orchestrator, rpc_q) + output_buffer = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512) + await _open(orchestrator, rpc_q, output_buffer=output_buffer) request_id = _stage0_request_id() await _submit(orchestrator, _append_audio()) request_state = orchestrator.request_states[request_id] @@ -417,7 +426,8 @@ async def test_forward_failure_closes_the_owning_session() -> None: assert SESSION_ID not in orchestrator.session_manager.runners assert session.state == DuplexSessionState.CLOSED assert orchestrator.session_manager.active_count() == 0 - types = [message.event.type for message in [output_q.get_nowait() for _ in range(output_q.qsize())]] + types = [(await output_buffer.get()).type for _ in range(output_buffer.pending_events)] + types.extend(output_q.get_nowait().event.type for _ in range(output_q.qsize())) # The stage failure is reported before the session expires, never after. assert types.index("error") < types.index("session.expired") assert types[-1] == "session.expired" @@ -428,7 +438,8 @@ async def test_aura_forward_failure_keeps_the_session() -> None: from vllm_omni.model_executor.models.aura_omni.duplex.plugin import AuraDuplexPlugin orchestrator, clients, rpc_q, output_q = _build(stages=2, plugin=AuraDuplexPlugin(_encode_audio)) - await _open(orchestrator, rpc_q) + output_buffer = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512) + await _open(orchestrator, rpc_q, output_buffer=output_buffer) request_id = next(iter(orchestrator.request_states)) await _submit(orchestrator, _append_audio()) request_state = orchestrator.request_states[request_id] @@ -442,7 +453,8 @@ async def test_aura_forward_failure_keeps_the_session() -> None: assert request_id not in orchestrator.request_states assert SESSION_ID in orchestrator.session_manager.runners assert session.state != DuplexSessionState.CLOSED - types = [message.event.type for message in [output_q.get_nowait() for _ in range(output_q.qsize())]] + types = [(await output_buffer.get()).type for _ in range(output_buffer.pending_events)] + types.extend(output_q.get_nowait().event.type for _ in range(output_q.qsize())) assert "error" in types assert "session.expired" not in types await orchestrator.session_manager.shutdown() @@ -604,6 +616,7 @@ async def test_turn_plugin_processes_multimodal_prompt_before_stage_submission(m control_id="open-qwen", session_id=SESSION_ID, session_config=DuplexSessionConfig(model="qwen", modalities=["text", "audio"]), + output_buffer=DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512), ) ) assert (await rpc_q.get()).ok diff --git a/tests/entrypoints/duplex/test_duplex_serving.py b/tests/entrypoints/duplex/test_duplex_serving.py index 70a0e9d48..06ae433e1 100644 --- a/tests/entrypoints/duplex/test_duplex_serving.py +++ b/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 diff --git a/tests/entrypoints/openai/test_duplex_session_attachment.py b/tests/entrypoints/openai/test_duplex_session_attachment.py index 10b600407..7e37719d9 100644 --- a/tests/entrypoints/openai/test_duplex_session_attachment.py +++ b/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 diff --git a/tests/entrypoints/test_duplex_omni.py b/tests/entrypoints/test_duplex_omni.py index ada212224..e075934ee 100644 --- a/tests/entrypoints/test_duplex_omni.py +++ b/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) diff --git a/tests/model_executor/models/qwen3_omni/test_duplex_plugin.py b/tests/model_executor/models/qwen3_omni/test_duplex_plugin.py index 908ae6cf6..fdd9d0674 100644 --- a/tests/model_executor/models/qwen3_omni/test_duplex_plugin.py +++ b/tests/model_executor/models/qwen3_omni/test_duplex_plugin.py @@ -27,6 +27,7 @@ from vllm_omni.engine.duplex.commands import ( UpdateSession, ) from vllm_omni.engine.duplex.config import DuplexSessionConfig +from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer from vllm_omni.engine.duplex.messages import OpenDuplexSessionMessage from vllm_omni.engine.duplex.plugin import DuplexRuntimeConfigError from vllm_omni.engine.duplex.session.manager import DuplexSessionManager @@ -77,6 +78,7 @@ async def open_qwen(): port = RecordingStagePort(stage_count=3) output: asyncio.Queue[Any] = asyncio.Queue() results: asyncio.Queue[Any] = asyncio.Queue() + output_buffer = DuplexOutputBuffer(max_bytes=2 * 1024 * 1024, max_events=512) manager = DuplexSessionManager( plugin=plugin, stage_port=port, @@ -86,10 +88,21 @@ async def open_qwen(): model_config=None, ) config = DuplexSessionConfig(model="qwen", modalities=["text", "audio"], overlap_policy="barge_in_on_speech") - await manager.handle(OpenDuplexSessionMessage(control_id="open", session_id=SESSION_ID, session_config=config)) + await manager.handle( + OpenDuplexSessionMessage( + control_id="open", session_id=SESSION_ID, session_config=config, output_buffer=output_buffer + ) + ) result = await results.get() assert result.ok, result - harness = Harness(manager, port, output, results, manager.runners[SESSION_ID]) + harness = Harness( + manager=manager, + port=port, + output=output, + output_buffer=output_buffer, + results=results, + runner=manager.runners[SESSION_ID], + ) await harness.settle() return harness diff --git a/tests/protocol/realtime/test_realtime_codec_single_source.py b/tests/protocol/realtime/test_realtime_codec_single_source.py index 554c99abe..a15600390 100644 --- a/tests/protocol/realtime/test_realtime_codec_single_source.py +++ b/tests/protocol/realtime/test_realtime_codec_single_source.py @@ -218,14 +218,17 @@ def test_the_error_code_vocabulary_is_tier3() -> None: assert not hasattr(realtime_errors, "REALTIME_ERROR_TYPES_BY_CODE") -def test_the_tier1_error_event_reports_only_openai_classes() -> None: +@pytest.mark.parametrize("code", ["resource_exhausted", "output_backpressure"]) +def test_the_tier1_error_event_reports_only_openai_classes(code: str) -> None: from vllm_omni.protocol.duplex import events as duplex_wire_events from vllm_omni.protocol.realtime import events as realtime_events # Tier 1 knows the envelope shape but not our code vocabulary. - assert realtime_events.ErrorEvent(code="resource_exhausted").error_type == "invalid_request_error" + assert realtime_events.ErrorEvent(code=code).error_type == "invalid_request_error" # Tier 2 resolves it through our table. - assert duplex_wire_events.ErrorEvent(code="resource_exhausted").error_type == "rate_limit_error" + error = duplex_wire_events.ErrorEvent(code=code) + assert error.error_type == "rate_limit_error" + assert error.to_realtime()["error"]["type"] == "rate_limit_error" @pytest.mark.parametrize("name", _TIER1_COMMANDS + _TIER2_COMMANDS + _TIER3_COMMANDS) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index eabea35c2..167664b37 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -536,6 +536,8 @@ class DuplexSessionRuntimeConfig: resume_replay_max_bytes_per_session: int = 8 * 1024 * 1024 max_pending_input_bytes_per_session: int = 16 * 1024 * 1024 max_pending_turns_per_session: int = 4 + max_pending_output_bytes_per_session: int = 2 * 1024 * 1024 + max_pending_output_events_per_session: int = 512 max_sessions: int = 1 # Unread by the plugin framework. It used to bound the per-session # completed-append table that made a retried append RPC submit once; the @@ -561,6 +563,8 @@ class DuplexSessionRuntimeConfig: "resume_replay_max_bytes_per_session": self.resume_replay_max_bytes_per_session, "max_pending_input_bytes_per_session": self.max_pending_input_bytes_per_session, "max_pending_turns_per_session": self.max_pending_turns_per_session, + "max_pending_output_bytes_per_session": self.max_pending_output_bytes_per_session, + "max_pending_output_events_per_session": self.max_pending_output_events_per_session, "max_sessions": self.max_sessions, "completed_append_cache_size": self.completed_append_cache_size, } diff --git a/vllm_omni/engine/duplex/delivery.py b/vllm_omni/engine/duplex/delivery.py new file mode 100644 index 000000000..b372a9615 --- /dev/null +++ b/vllm_omni/engine/duplex/delivery.py @@ -0,0 +1,242 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Bounded, single-consumer delivery across the engine and caller threads. + +Only queued, unsequenced audio can be removed. The producer still rejects +new stale model output; this buffer does not retain a growing response-id +tombstone table. It tracks just the most recently dequeued event until the +next ``get()`` so a consumer can recheck it immediately before delivery. + +Limits cover queued events, with a separate small termination reserve and +one final session-closure notification independent of that reserve. A +consumer may additionally hold one dequeued event, at most one queue budget +in size. Already journaled or delivered events are outside this buffer. +""" + +from __future__ import annotations + +import asyncio +import json +import threading +from collections import deque +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, replace + +from vllm_omni.engine.duplex.events import AudioDelta, DuplexEvent, ErrorEvent, ResponseDone, SessionClosed + +_MAX_TERMINAL_BYTES = 4096 + + +class DuplexOutputOverflowError(RuntimeError): + """One session exceeded its pending output budget; valid output was not dropped.""" + + +@dataclass(frozen=True, slots=True) +class _PendingEvent: + event: DuplexEvent + size: int + reserved: bool + + +class DuplexOutputBuffer: + """Nonblocking producer, one async consumer, and response-scoped audio removal. + + ``get()`` releases tracking of its previous result. The consumer must + finish delivering that result before requesting another. ``guard()`` + makes the final validity check and synchronous journal recording atomic + with invalidation; never await while holding that guard. + """ + + def __init__( + self, + *, + max_bytes: int, + max_events: int, + reserve_bytes: int = 64 * 1024, + reserve_events: int = 8, + ) -> None: + if min(max_bytes, max_events, reserve_bytes, reserve_events) <= 0: + raise ValueError("output buffer limits must be positive") + self._max_bytes = max_bytes + self._max_events = max_events + self._reserve_bytes = reserve_bytes + self._reserve_events = reserve_events + self._lock = threading.Lock() + self._pending: deque[_PendingEvent] = deque() + self._bytes = 0 + self._reserved_bytes = self._reserved_events = 0 + self._held: DuplexEvent | None = None + self._waiter: tuple[asyncio.AbstractEventLoop, asyncio.Future[None]] | None = None + self._closed = False + self._terminal: SessionClosed | None = None + + @property + def pending_bytes(self) -> int: + with self._lock: + return self._bytes + self._reserved_bytes + + @property + def pending_events(self) -> int: + with self._lock: + return len(self._pending) + + def put(self, event: DuplexEvent) -> bool: + """Append without blocking; overflow leaves the queue unchanged. + + Return false for late output after closure. A session terminal closes + through its own single slot, even when the ordinary reserve is full. + Error and response terminals can use the finite reserve. + They retain FIFO order and cannot overtake still-valid media. The + caller must stop the affected session after ordinary overflow instead + of repeatedly generating errors until the reserve also overflows. + """ + if isinstance(event, SessionClosed): + self.close(event) + return True + payload = event.to_realtime() + audio_size = 0 + if isinstance(event, AudioDelta): + # The producer has already encoded this as base64 ASCII. Count it + # directly without serializing the large string or decoding PCM. + audio_size = len(event.delta) + payload["delta"] = "" + size = audio_size + len(json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")) + with self._lock: + if self._closed: + return False + reserved = self._bytes + size > self._max_bytes or ( + len(self._pending) - self._reserved_events >= self._max_events + ) + if reserved and ( + not isinstance(event, ErrorEvent | ResponseDone) + or self._reserved_bytes + size > self._reserve_bytes + or self._reserved_events >= self._reserve_events + ): + raise DuplexOutputOverflowError("duplex session pending output limit exceeded") + self._pending.append(_PendingEvent(event=event, size=size, reserved=reserved)) + if reserved: + self._reserved_bytes += size + self._reserved_events += 1 + else: + self._bytes += size + waiter = self._waiter + self._waiter = None + self._notify(waiter) + return True + + @staticmethod + def _matches(event: DuplexEvent, response_id: str, through_epoch: int) -> bool: + return ( + isinstance(event, AudioDelta) + and event.response_id == response_id + and event.epoch is not None + and event.epoch <= through_epoch + ) + + def invalidate(self, response_id: str, through_epoch: int) -> None: + """Remove queued audio for one accepted cancellation. + + Non-audio events and other responses retain their exact order. An + already dequeued matching event becomes invalid too. The producer + must not enqueue further stale events after this call. + """ + with self._lock: + kept: deque[_PendingEvent] = deque() + for pending in self._pending: + if self._matches(pending.event, response_id, through_epoch): + self._release(pending) + else: + kept.append(pending) + self._pending = kept + if self._held is not None and self._matches(self._held, response_id, through_epoch): + self._held = None + + def _is_valid(self, event: DuplexEvent) -> bool: + return not isinstance(event, AudioDelta) or event is self._held + + def is_valid(self, event: DuplexEvent) -> bool: + """Check the current dequeued event; use ``guard`` for an atomic handoff.""" + with self._lock: + return self._is_valid(event) + + @contextmanager + def guard(self, event: DuplexEvent) -> Iterator[bool]: + """Hold validity stable during synchronous sequencing; never await here.""" + with self._lock: + yield self._is_valid(event) + + def _release(self, pending: _PendingEvent) -> None: + if pending.reserved: + self._reserved_bytes -= pending.size + self._reserved_events -= 1 + else: + self._bytes -= pending.size + + async def get(self) -> DuplexEvent | None: + """Return the next event, or ``None`` once a closed buffer is drained.""" + loop = asyncio.get_running_loop() + with self._lock: + self._held = None + while True: + with self._lock: + if self._pending: + pending = self._pending.popleft() + self._release(pending) + self._held = pending.event + return pending.event + if self._closed: + if self._terminal is not None: + terminal = self._terminal + self._terminal = None + return terminal + return None + if self._waiter is not None: + raise RuntimeError("duplex output buffer already has a waiting consumer") + future: asyncio.Future[None] = loop.create_future() + self._waiter = (loop, future) + try: + await future + finally: + with self._lock: + if self._waiter is not None and self._waiter[1] is future: + self._waiter = None + + def close(self, event: SessionClosed | None = None) -> None: + """Finish once, preserving FIFO output before an optional final notification. + + The single final slot cannot be consumed by errors or response endings. + Oversized close details are replaced by a compact notification retaining + the generated session/event identity and terminal type. Duplicate closure + and late output are harmless, including after a local + shutdown that already ended consumption without a session event. + """ + if event is not None: + size = len(json.dumps(event.to_realtime(), ensure_ascii=False, separators=(",", ":")).encode("utf-8")) + if size > _MAX_TERMINAL_BYTES: + event = replace(event, reason="close_details_exceed_output_limit", details={}) + with self._lock: + if self._closed: + return + self._closed = True + self._terminal = event + waiter = self._waiter + self._waiter = None + self._notify(waiter) + + @staticmethod + def _notify(waiter: tuple[asyncio.AbstractEventLoop, asyncio.Future[None]] | None) -> None: + if waiter is None: + return + loop, future = waiter + + def wake() -> None: + if not future.done(): + future.set_result(None) + + try: + loop.call_soon_threadsafe(wake) + except RuntimeError: + if not loop.is_closed(): + raise diff --git a/vllm_omni/engine/duplex/messages.py b/vllm_omni/engine/duplex/messages.py index 7c28b4638..a7a084c4a 100644 --- a/vllm_omni/engine/duplex/messages.py +++ b/vllm_omni/engine/duplex/messages.py @@ -5,8 +5,9 @@ Sessions live inside the engine (``DuplexSessionRunner`` on the orchestrator loop). The API layer only opens/closes/resumes/touches sessions through -correlated RPC and pushes ``DuplexCommand`` objects one-way; every session -output travels back as a ``DuplexSessionEventMessage``. +correlated RPC and pushes ``DuplexCommand`` objects one-way. Open binds a +bounded in-process output buffer; closure notifications and errors without +a live session travel back as ``DuplexSessionEventMessage``. """ from __future__ import annotations @@ -18,6 +19,7 @@ from vllm_omni.engine.messages import EngineQueueMessage if TYPE_CHECKING: from vllm_omni.engine.duplex.commands import DuplexCommand from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig + from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer from vllm_omni.engine.duplex.events import DuplexEvent @@ -44,6 +46,8 @@ class OpenDuplexSessionMessage(EngineQueueMessage, kw_only=True): session_id: str #: Already normalized by ``DuplexOmni``; the queue is in-process, so the object crosses as is. session_config: DuplexSessionConfig + #: In-process, bounded output shared with the single consumer; no media in the engine-wide queue. + output_buffer: DuplexOutputBuffer class CloseDuplexSessionMessage(EngineQueueMessage, kw_only=True): diff --git a/vllm_omni/engine/duplex/realtime_events.py b/vllm_omni/engine/duplex/realtime_events.py index 35eb0db5b..659ffd5b1 100644 --- a/vllm_omni/engine/duplex/realtime_events.py +++ b/vllm_omni/engine/duplex/realtime_events.py @@ -804,7 +804,9 @@ def _project(state: RealtimeProjectionState, event: dict[str, object]) -> list[D "status": "in_progress", "content": [], } - state.conversation_items[item_id] = item + # Later deltas grow the state item. Keep creation events independent + # so their queued payloads retain the size measured at admission. + state.conversation_items[item_id] = {**item, "content": []} events = [ _response_created_event(event), *_conversation_item_added_events(state, item), diff --git a/vllm_omni/engine/duplex/session/emitter.py b/vllm_omni/engine/duplex/session/emitter.py index 759bd6c27..6f36390b6 100644 --- a/vllm_omni/engine/duplex/session/emitter.py +++ b/vllm_omni/engine/duplex/session/emitter.py @@ -17,13 +17,15 @@ runner injects it as ``promote_deferred_overlap``. from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Mapping +from dataclasses import replace from vllm_omni.engine.duplex.config import DuplexSessionState from vllm_omni.engine.duplex.events import ( DOMAIN_TERMINAL_EVENTS, MODEL_OUTPUT_EVENTS, DuplexEvent, + ResponseDone, error_event, ) from vllm_omni.engine.duplex.realtime_events import ( @@ -84,8 +86,51 @@ class SessionEmitter: # ------------------------------------------------------------------ # def emit_events(self, events: list[DuplexEvent]) -> None: - for event in events: - self._ctx.manager.emit(self._ctx.session, event) + self._ctx.manager.emit(self._ctx.session, events) + + def fail_output(self, message: str, *, response_id: str | None, terminal: ResponseDone | None) -> list[DuplexEvent]: + """Fail undelivered output, returning recovery events without sending them again.""" + session = self._ctx.session + self._ctx.manager.invalidate_output(session.session_id, response_id, through_epoch=session.epoch) + if session.active_response_id is not None: + # Close still needs the request binding to abort the real engine work. + session.end_response(commit_text=False, preserve_request=True) + events: list[DuplexEvent] = [error_event("output_backpressure", message)] + if response_id is None: + return events + # Normal completion can commit just before its projected ending is + # delivered. An ending converted to failure must undo that commit. + session.delete_history_item(f"item_{response_id}") + payload = { + "type": "response.done", + "session_id": session.session_id, + "response_id": response_id, + "epoch": session.epoch, + "committed": False, + "status": "failed", + "status_details": {"type": "failed", "reason": "output_backpressure"}, + "playback": session.playback.as_dict(), + } + if terminal is None: + projected = project_internal_event(self.require_projector(), payload) + terminal = next((event for event in projected if isinstance(event, ResponseDone)), None) + else: + response = dict(terminal.response) + metadata = response.get("metadata") + response.update( + status="failed", + status_details=payload["status_details"], + metadata={**(metadata if isinstance(metadata, Mapping) else {}), **payload}, + ) + output = response.get("output") + if isinstance(output, list): + response["output"] = [ + {**item, "status": "failed"} if isinstance(item, Mapping) else item for item in output + ] + terminal = replace(terminal, response=response) + if terminal is not None: + events.append(terminal) + return events def emit_error( self, @@ -110,6 +155,13 @@ class SessionEmitter: accepted, deferred_overlap_payload = self._apply_outbound_session_event(payload) if not accepted: return + if payload.get("type") == "audio.cancelled": + response_id = payload.get("response_id") + cancelled_epoch = payload.get("cancelled_epoch") + if isinstance(response_id, str) and isinstance(cancelled_epoch, int): + self._ctx.manager.invalidate_output( + self._ctx.session.session_id, response_id, through_epoch=cancelled_epoch + ) self.emit_events(project_internal_event(self.require_projector(), payload)) if deferred_overlap_payload is not None and not self._ctx.run.closing: precreate_response = self._ctx.model_state.deferred_precreate_response diff --git a/vllm_omni/engine/duplex/session/manager.py b/vllm_omni/engine/duplex/session/manager.py index 38eeb7ab6..8e54341a3 100644 --- a/vllm_omni/engine/duplex/session/manager.py +++ b/vllm_omni/engine/duplex/session/manager.py @@ -11,8 +11,9 @@ state. Control operations (open / close / resume / touch) arrive as correlated RPC messages and answer through ``result_sink``; session commands arrive one-way as ``DuplexSessionCommandMessage`` and are pushed onto the runner's ordered -mailbox after backpressure admission; everything a session emits leaves -through ``output_sink`` as ``DuplexSessionEventMessage``. +mailbox after backpressure admission. Public output uses a bounded buffer +shared with that session's consumer. The engine-wide ``output_sink`` only +carries closure notifications and errors without a live session. """ from __future__ import annotations @@ -34,7 +35,16 @@ from vllm_omni.engine.duplex.contracts import ( duplex_resource_request_belongs_to_session, duplex_resource_request_id, ) -from vllm_omni.engine.duplex.events import DuplexEvent, ErrorEvent, SessionClosed, SessionExpired, error_event +from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer, DuplexOutputOverflowError +from vllm_omni.engine.duplex.events import ( + DuplexEvent, + ErrorEvent, + OutputAudioCleared, + ResponseDone, + SessionClosed, + SessionExpired, + error_event, +) from vllm_omni.engine.duplex.messages import ( CloseDuplexSessionMessage, DuplexControlResultMessage, @@ -128,6 +138,8 @@ class DuplexSessionManager: model_path=getattr(runtime_config, "server_vad_model_path", None) ) self.runners: dict[str, DuplexSessionRunner] = {} + self._outputs: dict[str, DuplexOutputBuffer] = {} + self._output_failed: set[str] = set() #: Sessions whose close/expiry began but whose stage cleanup has not finalized; #: they keep their admission slot until finalized. self._closing: dict[str, _PendingSessionCleanup] = {} @@ -270,20 +282,24 @@ class DuplexSessionManager: if runner is None: self._emit_raw( message.session_id, - self._error_event( - "unknown_session", - f"Unknown or closed duplex session: {message.session_id}", - command=command, - ), + [ + self._error_event( + "unknown_session", + f"Unknown or closed duplex session: {message.session_id}", + command=command, + ) + ], ) return session = runner.session if runner.closing: self.emit( session, - self._error_event( - "session_closed", f"Duplex session is closing: {session.session_id}", command=command - ), + [ + self._error_event( + "session_closed", f"Duplex session is closing: {session.session_id}", command=command + ) + ], ) return if isinstance(command, AppendAudio): @@ -293,7 +309,7 @@ class DuplexSessionManager: if modality_error is not None: self.emit( session, - self._error_event("invalid_input_modality", modality_error, command=command), + [self._error_event("invalid_input_modality", modality_error, command=command)], ) return limit = int(self.runtime_config.max_pending_input_bytes_per_session) @@ -301,22 +317,26 @@ class DuplexSessionManager: if not session.reserve_input_bytes(pending_bytes, limit=limit): self.emit( session, - self._error_event( - "input_backpressure", - "Duplex session has too many pending input bytes", - command=command, - ), + [ + self._error_event( + "input_backpressure", + "Duplex session has too many pending input bytes", + command=command, + ) + ], ) return elif isinstance(command, Commit): if not session.reserve_pending_turn(limit=int(self.runtime_config.max_pending_turns_per_session)): self.emit( session, - self._error_event( - "input_backpressure", - "Duplex session has too many pending input turns", - command=command, - ), + [ + self._error_event( + "input_backpressure", + "Duplex session has too many pending input turns", + command=command, + ) + ], ) return runner.submit(command) @@ -329,18 +349,72 @@ class DuplexSessionManager: # Emission # # ------------------------------------------------------------------ # - def emit(self, session: DuplexEngineSession, event: DuplexEvent) -> None: - """Bind the session identity to one typed event and push it to the engine output queue.""" - self._emit_raw(session.session_id, event, epoch=session.epoch) + def emit(self, session: DuplexEngineSession, events: list[DuplexEvent]) -> None: + """Bind identity and deliver through the session's bounded output.""" + self._emit_raw(session.session_id, events, epoch=session.epoch) + + def invalidate_output(self, session_id: str, response_id: str | None, *, through_epoch: int) -> None: + output = self._outputs.get(session_id) + if output is not None and response_id is not None: + output.invalidate(response_id, through_epoch=through_epoch) def _emit_raw( self, session_id: str, - event: DuplexEvent, + events: list[DuplexEvent], *, epoch: int | None = None, ) -> None: + ending_delivered = False + for index, event in enumerate(events): + try: + self._deliver(session_id, event, epoch=epoch) + except DuplexOutputOverflowError as exc: + if session_id in self._output_failed: + return + self._output_failed.add(session_id) + runner = self.runners.get(session_id) + if runner is None: + return + # Projection has already marked this batch's ending as emitted. + # Retain it instead of recursively projecting a second ending. + terminal = next((item for item in events[index:] if isinstance(item, ResponseDone)), None) + response_id = terminal.response_id if terminal is not None else event.response_id + if response_id is None and not ending_delivered: + response_id = runner.session.active_response_id + for recovery in runner.fail_output(str(exc), response_id=response_id, terminal=terminal): + try: + self._deliver(session_id, recovery, epoch=epoch) + except DuplexOutputOverflowError: + # A full control reserve may omit an error/response ending; + # session closure still has its independent terminal slot. + pass + return + ending_delivered |= isinstance(event, ResponseDone) + + def _deliver(self, session_id: str, event: DuplexEvent, *, epoch: int | None) -> None: event = replace(event, session_id=session_id, epoch=epoch) + output = self._outputs.get(session_id) + if output is not None and not isinstance(event, SessionClosed): + if isinstance(event, OutputAudioCleared) and epoch is not None: + self.invalidate_output(session_id, event.response_id, through_epoch=epoch) + if session_id in self._output_failed: + # The failure/close path has a finite reserve; stop adding ordinary output. + if not isinstance(event, ErrorEvent) and not ( + isinstance(event, ResponseDone) and event.status in {"failed", "cancelled"} + ): + return + output.put(event) + return + if isinstance(event, SessionClosed): + self._outputs.pop(session_id, None) + self._output_failed.discard(session_id) + elif not isinstance(event, ErrorEvent): + # A closed/rolled-back session must not move late audio back into + # the unbounded shared queue after its output buffer is removed. + return + # Only closure notifications and errors without a live session use the + # engine-wide queue. Audio never accumulates there before the bounded buffer. message = DuplexSessionEventMessage(session_id=session_id, event=event) put_nowait = getattr(self._output_sink, "put_nowait", None) if callable(put_nowait): @@ -514,6 +588,7 @@ class DuplexSessionManager: model_config=self.model_config, ) self.runners[session_id] = runner + self._outputs[session_id] = message.output_buffer # Reserve the Stage0 request resource atomically with admission. self.ensure_stage_request(session, stage_id=0) runner.start() @@ -566,6 +641,7 @@ class DuplexSessionManager: owned_runner = runner is not None and self.runners.get(session_id) is runner if owned_runner: self.runners.pop(session_id, None) + self._outputs.pop(session_id, None) pending: _PendingRequestCleanup | None = None key: tuple[str, int] | None = None if session is not None: @@ -686,7 +762,7 @@ class DuplexSessionManager: terminal: DuplexEvent = ( SessionExpired(reason=reason) if kind == "expired" else SessionClosed(reason=reason) ) - self.emit(session, terminal) + self.emit(session, [terminal]) async def _finalize_pending_cleanup( self, @@ -943,6 +1019,8 @@ class DuplexSessionManager: self._closing.clear() self._session_snapshots.clear() self._request_index.clear() + self._outputs.clear() + self._output_failed.clear() if self._owns_executor: self.executor.shutdown(wait=False, cancel_futures=True) diff --git a/vllm_omni/engine/duplex/session/model_channel.py b/vllm_omni/engine/duplex/session/model_channel.py index e2bccf9b7..15fca4f66 100644 --- a/vllm_omni/engine/duplex/session/model_channel.py +++ b/vllm_omni/engine/duplex/session/model_channel.py @@ -577,6 +577,8 @@ class ModelChannel: close_reason, emitted_response = await self._send_model_output_events( drain_result, expected_epoch=expected_epoch ) + if self._ctx.run.closing: + return if close_reason is not None: await self._close_from_runtime(close_reason) return @@ -607,7 +609,7 @@ class ModelChannel: expected_epoch: int | None = None, ) -> tuple[str | None, bool]: session = self._ctx.session - if expected_epoch is not None and session.epoch != expected_epoch: + if self._ctx.run.closing or (expected_epoch is not None and session.epoch != expected_epoch): return None, False close_reason: str | None = None emitted_response = False @@ -624,7 +626,7 @@ class ModelChannel: ) emitted_response = emitted_response or did_emit close_reason = close_reason or close_reason_for_result - if expected_epoch is not None and session.epoch != expected_epoch: + if self._ctx.run.closing or (expected_epoch is not None and session.epoch != expected_epoch): return None, emitted_response return close_reason, emitted_response @@ -636,7 +638,7 @@ class ModelChannel: str(model_result.get("error_code")), str(model_result.get("error") or "Duplex native data-plane error"), ) - if response_id is None: + if self._ctx.run.closing or response_id is None: return session.end_response(commit_text=False) self._out.emit( @@ -734,6 +736,8 @@ class ModelChannel: close_reason: str | None = None emitted_response = False self._end_active_response_before_future_model_turn(model_turn_id=model_turn_id) + if self._ctx.run.closing: + return close_reason, emitted_response if ( session.active_response_id is not None and model_turn_id is not None @@ -776,10 +780,14 @@ class ModelChannel: payload["response_id"] = response_id self._attach_runtime_metadata(payload, model_result) self._out.emit(payload) + if self._ctx.run.closing: + return close_reason, emitted_response if model_result.get("abort_data_plane_request") is True and isinstance(data_plane_request_id, str): # stage_port.abort_requests expects a list of ids; a bare str is # iterated as characters and never matches the prewarmed binding. await self._abort_request([data_plane_request_id], notify=False) + if self._ctx.run.closing: + return close_reason, emitted_response if response_id is not None: if not auto_response and self.response_continuations_remaining(response_id): self._ctx.services.spawn( @@ -822,7 +830,7 @@ class ModelChannel: data_plane = self._ctx.plugin.data_plane close_reason: str | None = None emitted_response = False - if expected_epoch is not None and session.epoch != expected_epoch: + if self._ctx.run.closing or (expected_epoch is not None and session.epoch != expected_epoch): return close_reason, emitted_response data_plane_request_id = model_result.get("data_plane_request_id") if isinstance(data_plane_request_id, str) and data_plane.is_terminal(data_plane_request_id): @@ -932,6 +940,8 @@ class ModelChannel: return close_reason, emitted_response if draining_response_id is None: self._end_active_response_before_future_model_turn(model_turn_id=model_turn_id) + if self._ctx.run.closing: + return close_reason, emitted_response if ( draining_response_id is None and session.active_response_id is not None @@ -955,6 +965,8 @@ class ModelChannel: if response_request_metrics: created_payload["response_request_metrics"] = response_request_metrics self._out.emit(created_payload) + if self._ctx.run.closing: + return close_reason, emitted_response stage_metrics = model_result.get("stage_metrics") response_stage_metrics = session.accumulate_response_stage_metrics( stage_metrics if isinstance(stage_metrics, Mapping) else None @@ -976,6 +988,8 @@ class ModelChannel: response_request_metrics=response_request_metrics, ) self._out.emit(speak_payload) + if self._ctx.run.closing: + return close_reason, emitted_response target_id = draining_response_id if draining_response_id not in (None, session.active_response_id) else None previous_sent_ms = session.playback_for_response(target_id).sent_ms text_chars_before_append = len(session.assistant_transcript(target_id)) @@ -1052,6 +1066,8 @@ class ModelChannel: response_request_metrics=response_request_metrics, ) self._out.emit(payload) + if self._ctx.run.closing: + return close_reason, emitted_response if ( not end_of_turn and model_result.get("stage_role") == "tts" @@ -1233,7 +1249,7 @@ class ModelChannel: response_id = session.active_response_id response_epoch = session.epoch response_turn_id = model_turn_id if model_turn_id is not None else session.active_response_turn_id - if expected_epoch is not None and response_epoch != expected_epoch: + if self._ctx.run.closing or (expected_epoch is not None and response_epoch != expected_epoch): return model_state.clear_continuation() payload: dict[str, object] = { @@ -1246,6 +1262,8 @@ class ModelChannel: if response_id is not None: payload["response_id"] = response_id self._out.emit(payload) + if self._ctx.run.closing: + return if response_id is None: if session.epoch == response_epoch and response_turn_id is not None: session.complete_model_turn(response_turn_id) @@ -1277,7 +1295,8 @@ class ModelChannel: ) -> bool: session = self._ctx.session stale_common_owner = ( - session.state == DuplexSessionState.CLOSED + self._ctx.run.closing + or session.state == DuplexSessionState.CLOSED or session.active_request_id != request_id or (expected_epoch is not None and session.epoch != expected_epoch) ) @@ -1372,6 +1391,6 @@ class ModelChannel: except Exception as exc: logger.exception("Failed to schedule duplex native response continuation: %s", exc) scheduled = False - if scheduled: + if scheduled and not self._ctx.run.closing: model_state.continuation_owner_id = owner_id model_state.continuation_units = count + 1 diff --git a/vllm_omni/engine/duplex/session/runner.py b/vllm_omni/engine/duplex/session/runner.py index ca9793938..81c89ac77 100644 --- a/vllm_omni/engine/duplex/session/runner.py +++ b/vllm_omni/engine/duplex/session/runner.py @@ -64,6 +64,7 @@ from vllm_omni.engine.duplex.contracts import ( from vllm_omni.engine.duplex.events import ( DuplexEvent, InputCleared, + ResponseDone, SessionExpired, SessionHeartbeatAck, ) @@ -347,6 +348,15 @@ class DuplexSessionRunner: """Whether ``session.closed`` / ``session.expired`` already left this runner.""" return self.run.closed_emitted + def fail_output(self, message: str, *, response_id: str | None, terminal: ResponseDone | None) -> list[DuplexEvent]: + """Experimental overflow policy: fail this response and close only its session.""" + if self.closing: + return [] + events = self.out.fail_output(message, response_id=response_id, terminal=terminal) + self._begin_close("output_backpressure") + self.manager.close_from_runner(self, "output_backpressure") + return events + def mark_closed_emitted(self) -> None: """Claim the session's one terminal event for the caller. @@ -1474,6 +1484,9 @@ class DuplexSessionRunner: if older_abort_ids: session.release_resources_for_request_ids(older_abort_ids) new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + self.manager.invalidate_output(session.session_id, old_response_id, through_epoch=old_epoch) + for response_id, _, _ in draining_cancels: + self.manager.invalidate_output(session.session_id, response_id, through_epoch=old_epoch) if old_request_id is not None: # Release projector/parser cursors so cancelled epochs do not # accumulate until the whole session closes. diff --git a/vllm_omni/engine/duplex_omni_engine.py b/vllm_omni/engine/duplex_omni_engine.py index c7c3a6328..98ea81fa0 100644 --- a/vllm_omni/engine/duplex_omni_engine.py +++ b/vllm_omni/engine/duplex_omni_engine.py @@ -37,6 +37,7 @@ from vllm_omni.engine.orchestrator import OrchestratorBase if TYPE_CHECKING: from vllm_omni.engine.duplex.commands import DuplexCommand + from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer logger = init_logger(__name__) @@ -151,6 +152,7 @@ class DuplexOmniEngine(AsyncOmniEngine): session_id: str, session_config: DuplexSessionConfig, *, + output_buffer: DuplexOutputBuffer, timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, ) -> DuplexControlResultMessage: control_id = uuid.uuid4().hex @@ -159,6 +161,7 @@ class DuplexOmniEngine(AsyncOmniEngine): control_id=control_id, session_id=session_id, session_config=session_config, + output_buffer=output_buffer, ), control_id=control_id, operation="open", @@ -171,10 +174,14 @@ class DuplexOmniEngine(AsyncOmniEngine): session_id: str, session_config: DuplexSessionConfig, *, + output_buffer: DuplexOutputBuffer, timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, ) -> DuplexControlResultMessage: loop = asyncio.get_running_loop() - return await loop.run_in_executor(None, lambda: self._open_session(session_id, session_config, timeout=timeout)) + return await loop.run_in_executor( + None, + lambda: self._open_session(session_id, session_config, output_buffer=output_buffer, timeout=timeout), + ) def _close_session( self, diff --git a/vllm_omni/entrypoints/duplex/serving.py b/vllm_omni/entrypoints/duplex/serving.py index 5e6f3da70..6751a9cf5 100644 --- a/vllm_omni/entrypoints/duplex/serving.py +++ b/vllm_omni/entrypoints/duplex/serving.py @@ -22,6 +22,7 @@ import json from collections.abc import Mapping from contextlib import suppress from dataclasses import dataclass, replace +from functools import partial from fastapi import WebSocket, WebSocketDisconnect from vllm.logger import init_logger @@ -478,7 +479,7 @@ class OmniDuplexSessionHandler: attachment_generation=credentials.attachment_generation, resume_token=credentials.resume_token, ) - await self._send_event(session_id, event) + await self._send_event(session_id, event, handle=handle) if isinstance(event, SessionClosed): close_reason = event.reason or event.type break @@ -523,19 +524,22 @@ class OmniDuplexSessionHandler: with suppress(Exception): await asyncio.wait_for(asyncio.shield(pump), _PUMP_DRAIN_TIMEOUT_S) - async def _send_event(self, session_id: str, event: DuplexEvent) -> None: + async def _send_event(self, session_id: str, event: DuplexEvent, *, handle: DuplexSessionHandle) -> None: payload = event.to_realtime() journal = not isinstance(event, _UNJOURNALED_EVENTS) and session_id not in self._resync_required_sessions + event_guard = partial(handle.output_guard, event) try: try: - await self._attachment_registry.send_event(session_id, payload, journal=journal) + await self._attachment_registry.send_event( + session_id, payload, journal=journal, event_guard=event_guard + ) except DuplexJournalOverflowError: first_overflow = session_id not in self._resync_required_sessions self._resync_required_sessions.add(session_id) if first_overflow: resync = SessionResyncRequired(session_id=session_id, reason="journal_overflow") await self._attachment_registry.send_event(session_id, resync.to_realtime(), journal=False) - await self._attachment_registry.send_event(session_id, payload, journal=False) + await self._attachment_registry.send_event(session_id, payload, journal=False, event_guard=event_guard) except KeyError: # Attachment already closed (takeover or teardown); the journal is gone. pass diff --git a/vllm_omni/entrypoints/duplex/session_attachment.py b/vllm_omni/entrypoints/duplex/session_attachment.py index b365995ec..e0c28f241 100644 --- a/vllm_omni/entrypoints/duplex/session_attachment.py +++ b/vllm_omni/entrypoints/duplex/session_attachment.py @@ -11,6 +11,7 @@ import secrets import time from collections import deque from collections.abc import Awaitable, Callable, Mapping +from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass, field from types import MappingProxyType @@ -291,6 +292,7 @@ class DuplexSessionAttachmentRegistry: *, journal: bool = True, on_accepted: Callable[[], None] | None = None, + event_guard: Callable[[], AbstractContextManager[bool]] | None = None, ) -> JournalEntry | None: """Sequence and dispatch one event to the current attachment. @@ -309,9 +311,15 @@ class DuplexSessionAttachmentRegistry: async with self._lock: if self._sessions.get(session_id) is not state: raise KeyError(f"unknown duplex attachment session: {session_id}") - entry = state.journal.record(payload) if journal else None - attachment = state.attachment - wire_payload = dict(entry.payload) if entry is not None else dict(payload) + # The producer may invalidate held audio while this coroutine + # waits for an attachment lock. Commit order before sequencing; + # release the cross-thread guard before awaiting network I/O. + with event_guard() if event_guard is not None else nullcontext(True) as valid: + if not valid: + return None + entry = state.journal.record(payload) if journal else None + attachment = state.attachment + wire_payload = dict(entry.payload) if entry is not None else dict(payload) if entry is not None and on_accepted is not None: on_accepted() if attachment is not None: diff --git a/vllm_omni/entrypoints/duplex_omni.py b/vllm_omni/entrypoints/duplex_omni.py index e5d8e5bf6..9baf6d45a 100644 --- a/vllm_omni/entrypoints/duplex_omni.py +++ b/vllm_omni/entrypoints/duplex_omni.py @@ -27,7 +27,7 @@ from __future__ import annotations import asyncio from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence -from contextlib import suppress +from contextlib import AbstractContextManager, suppress from typing import Any from uuid import uuid4 @@ -37,6 +37,7 @@ from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig from vllm_omni.engine.duplex import commands as duplex_commands from vllm_omni.engine.duplex.commands import DuplexCommand from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig, ResponseCreateOptions +from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer from vllm_omni.engine.duplex.events import DuplexEvent, SessionClosed from vllm_omni.engine.duplex.messages import ( DuplexControlResultMessage, @@ -66,7 +67,11 @@ class DuplexSessionHandle: self.capabilities: DuplexCapabilities = DuplexCapabilities() self.public_session: dict[str, object] = {} self.lease_generation: int = 0 - self._outbox: asyncio.Queue[DuplexEvent | None] = asyncio.Queue() + limits = omni.duplex_session_config + self._outbox = DuplexOutputBuffer( + max_bytes=limits.max_pending_output_bytes_per_session, + max_events=limits.max_pending_output_events_per_session, + ) self._closed = False self._close_reason: str | None = None self._closed_event = asyncio.Event() @@ -219,12 +224,10 @@ class DuplexSessionHandle: self._consumer_active = True try: while True: - if self._closed and self._outbox.empty(): - return event = await self._outbox.get() if event is None: - if self._closed and self._outbox.empty(): - return + return + if not self._outbox.is_valid(event): continue yield event if event.is_terminal: @@ -236,6 +239,10 @@ class DuplexSessionHandle: await self._closed_event.wait() return self._close_reason or "closed" + def output_guard(self, event: DuplexEvent) -> AbstractContextManager[bool]: + """Recheck held audio while committing its delivery; never await inside this guard.""" + return self._outbox.guard(event) + async def __aenter__(self) -> DuplexSessionHandle: return self @@ -256,7 +263,7 @@ class DuplexSessionHandle: self.lease_generation = int(result.lease_generation) def _deliver(self, event: DuplexEvent) -> None: - self._outbox.put_nowait(event) + self._outbox.put(event) if isinstance(event, SessionClosed): # SessionExpired is a SessionClosed self._mark_closed(event.reason or event.type) @@ -266,7 +273,7 @@ class DuplexSessionHandle: self._closed = True self._close_reason = reason self._closed_event.set() - self._outbox.put_nowait(None) + self._outbox.close() class DuplexOmni(AsyncOmni): @@ -356,8 +363,11 @@ class DuplexOmni(AsyncOmni): self._handles[session_id] = handle self._final_output_handler() try: - result = await self.engine.open_session_async(session_id, session_config, timeout=timeout) + result = await self.engine.open_session_async( + session_id, session_config, output_buffer=handle._outbox, timeout=timeout + ) except BaseException: + handle._mark_closed("open_failed") if self._handles.get(session_id) is handle: self._handles.pop(session_id, None) # Dropping the handle only forgets the id here. The engine may still diff --git a/vllm_omni/protocol/duplex/errors.py b/vllm_omni/protocol/duplex/errors.py index bc347d25e..6175758b9 100644 --- a/vllm_omni/protocol/duplex/errors.py +++ b/vllm_omni/protocol/duplex/errors.py @@ -45,6 +45,7 @@ REALTIME_ERROR_TYPES_BY_CODE: dict[str, str] = { "runtime_touch_failed": "server_error", "engine_error": "server_error", "input_backpressure": "rate_limit_error", + "output_backpressure": "rate_limit_error", "response_already_active": "invalid_request_error", "response_not_active": "invalid_request_error", "response_create_without_input": "invalid_request_error",