diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index 33def4de4..e29001d9a 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -38,9 +38,13 @@ x-minicpmo-4-5: &minicpmo_4_5 x-minicpmo-4-5-duplex: &minicpmo_4_5_duplex - vllm_omni/deploy/minicpmo_4_5_duplex.yaml - vllm_omni/engine/duplex/ + - vllm_omni/engine/duplex_omni_engine.py + - vllm_omni/engine/duplex_orchestrator.py + - vllm_omni/engine/omni_engine_base.py - vllm_omni/entrypoints/duplex/ + - vllm_omni/entrypoints/duplex_omni.py + - vllm_omni/entrypoints/async_omni_base.py - vllm_omni/clients/ - - vllm_omni/entrypoints/duplex_request_client.py - vllm_omni/outputs/duplex.py - vllm_omni/model_executor/duplex_sampling.py - vllm_omni/experimental/fullduplex/ @@ -249,7 +253,11 @@ source_file_dependencies: - tests/assets/minicpmo_4_5/ - tests/clients/ - tests/engine/duplex/ + - tests/engine/test_duplex_omni_engine.py + - tests/engine/test_duplex_orchestrator.py - tests/entrypoints/duplex/ + - tests/entrypoints/test_duplex_omni.py + - tests/entrypoints/openai_api/test_duplex_api_server.py - tests/model_executor/models/minicpmo_4_5/duplex/ omni_minicpmo_4_5_duplex_perf: diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index d600e8fbe..87d4c5362 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -66,7 +66,9 @@ steps: # 0.26 startup regression is understood and fixed. timeout_in_minutes: 90 commands: - - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model" + # Turn-based server VAD over /v1/realtime is deferred to RFC #7181 PR 5 + # (see the P0 thread on #7413); the test stays in the tree and runs again there. + - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model" --deselect tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py::TestQwen3OmniRealtimeWebSocket::test_server_vad_multi_turn_without_client_commit mirror_hardwares: h100_2 - label: "Entrypoint Test with L4" diff --git a/docs/api/README.md b/docs/api/README.md index b2d4ea1ba..db760506e 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -5,6 +5,9 @@ Main entry points for vLLM-Omni inference and serving. - [vllm_omni.entrypoints.async_omni.AsyncOmni][] +- [vllm_omni.entrypoints.async_omni_base.AsyncOmniBase][] +- [vllm_omni.entrypoints.duplex_omni.DuplexOmni][] +- [vllm_omni.entrypoints.duplex_omni.DuplexSessionHandle][] - [vllm_omni.engine.cfg_companion_tracker.CfgCompanionTracker][] - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][] - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][] @@ -17,12 +20,13 @@ Main entry points for vLLM-Omni inference and serving. ## Clients -Client libraries that talk to a running vLLM-Omni server (the full-duplex Realtime client and its per-model session presets). +Client libraries for full-duplex sessions: the Realtime WebSocket client, the in-process client over `DuplexOmni`, and the per-model session presets. - [vllm_omni.clients.duplex.AudioDelta][] - [vllm_omni.clients.duplex.AudioFormat][] - [vllm_omni.clients.duplex.ConnectionResumed][] - [vllm_omni.clients.duplex.DuplexClient][] +- [vllm_omni.clients.duplex.DuplexClientBase][] - [vllm_omni.clients.duplex.DuplexClientError][] - [vllm_omni.clients.duplex.DuplexConnectionError][] - [vllm_omni.clients.duplex.DuplexEvent][] @@ -44,6 +48,7 @@ Client libraries that talk to a running vLLM-Omni server (the full-duplex Realti - [vllm_omni.clients.duplex.TextDelta][] - [vllm_omni.clients.duplex.TranscriptDelta][] - [vllm_omni.clients.duplex.WebSocketTransport][] +- [vllm_omni.clients.inline_duplex.InlineDuplexClient][] ## Inputs diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 07beef2de..aa81d42c6 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -17,6 +17,7 @@ Common `PipelineConfig` fields include: | ------- | ------------- | | `model_type` | Pipeline identifier used during model and config resolution. | | `default_deploy_config_name` | Bundled deploy YAML loaded when the user does not pass `deploy_config`. | +| `duplex_plugin` | Dotted path of the model's `DuplexModelPlugin`. Set only for full-duplex models; it makes `vllm-omni serve` run the model through `DuplexOmni` (a server whose every surface runs on a duplex session) and the engine host a `DuplexOrchestrator` with the plugin loaded. | | `model_arch` | Default Hugging Face architecture for the pipeline. | | `hf_architectures` | Architecture names used to identify checkpoints whose `model_type` is shared. | | `hf_config_predicate` | Optional predicate used to select between pipelines with otherwise identical HF metadata. | @@ -52,7 +53,7 @@ The new deploy schema lives under `vllm_omni/deploy/` and is paired with a froze | ------- | ------ | ---------- | --------- | ------------- | | `base_config` | str (path) | optional | — | Overlay parent (relative or absolute). `stages:` / `platforms:` deep-merged by stage_id; other scalars overlay-wins. Intended for user-authored overlays; prod yamls stay flat. | | `async_chunk` | bool | optional | `true` | Enable chunked streaming between stages. Pin to `false` if the pipeline runs end-to-end. | -| `session_mode` | str | optional | `"turn"` | Session behavior. MiniCPM-o 4.5 deploy YAMLs set `"duplex"` so `/v1/realtime` and chat share the same profile. | +| `session_mode` | str | optional | `"turn"` | Session behavior. Duplex models (pipelines with a `duplex_plugin`, e.g. MiniCPM-o 4.5) set `"duplex"`; `vllm-omni serve` then serves it over `/v1/realtime?duplex=1` and `/v1/chat/completions`, both on duplex sessions. | | `active_stream_window` | int | optional | `0` | Number of active downstream stream slots; `0` preserves all-stream cycling. | | `duplex_session` | dict | optional | runtime defaults | Full-duplex session lifecycle, buffering, replay, and capacity limits. | | `connectors` | dict | optional | `null` | Named connector specs (`{name, extra}`). Referenced by each stage's `input_connectors` / `output_connectors`. See [Connector schema](#connector-schema). | diff --git a/docs/design/fullduplex.md b/docs/design/fullduplex.md index 7d3d5ec95..4646b85cd 100644 --- a/docs/design/fullduplex.md +++ b/docs/design/fullduplex.md @@ -1,1070 +1,373 @@ -# MiniCPM-o 4.5 Full-Duplex Runtime Architecture +# Unified Full-Duplex Framework -> Relocation note: after the H20 validation recorded below, the full-duplex -> implementation moved out of `vllm_omni/experimental/fullduplex/` into the -> stable tree (`vllm_omni/engine/duplex/`, `vllm_omni/entrypoints/duplex/`, -> `vllm_omni/model_executor/models/minicpmo_4_5/duplex/`, and companion -> modules). Module paths in this document reflect the new layout; the runtime -> evidence was produced on the pre-move tree, which is file-identical apart -> from module paths. +This document is the design record of the full-duplex serving framework: +the engine-resident session model, the typed command/event contract, the +single model plugin, and the serving/API surfaces built on top of them. It is +the owner document for `vllm_omni/engine/duplex/**`, +`vllm_omni/engine/duplex_omni_engine.py`, +`vllm_omni/engine/duplex_orchestrator.py`, +`vllm_omni/entrypoints/duplex_omni.py`, `vllm_omni/entrypoints/duplex/**`, +`vllm_omni/clients/**` and the per-model `duplex/` packages. + +The public wire contract lives in +[`docs/serving/realtime_duplex_api.md`](../serving/realtime_duplex_api.md); +this page describes how it is produced. ## Purpose -This document records the runtime architecture implemented by PR #3907, the -cleanup applied after published commit `e011d936`, and the remaining work that -must not be advertised as complete. Implemented contracts and future -architecture work are called out separately below. +A full-duplex model listens and speaks at the same time: audio streams into +the model while it produces audio, the model decides when to speak, and the +user can interrupt. Serving such a model needs a long-lived session with its +own identity (epoch, turn), ledgers (input, response, playback, history), a +lease (idle TTL, disconnect grace, resume) and a binding to the resumable +stage requests that carry the session's audio through the pipeline. -The checkpoint preserves the active runtime path already exercised on H20. It does -not introduce another reducer, controller, worker provider, or shadow runtime. +The framework gives that session exactly one owner, keeps the generic +turn-based stack free of duplex code, and exposes the same session to three +consumers: the OpenAI Realtime WebSocket route, the in-process Python API and +the model plugins. -## Review Snapshot - -- PR: `vllm-project/vllm-omni#3907` -- Published head before this refactor: `e011d936` -- Published base snapshot: `62589203` -- Published diff: 93 files, approximately `+29.2k/-0.6k` -- Local runtime and pytest validation: intentionally not used -- Required validation environment: an NVIDIA H20 CUDA host -- PR head reviewed for the current cleanup: - `f4f78fa555af4d65e46b003022d352ca2ea9d703` -- Validated tree: the uncommitted refactor synchronized file-for-file to the - isolated H20 worktree on 2026-07-20 - -The current tree has received fresh H20 focused and E2E validation against vLLM -0.25.0. The final affected matrix includes the pre-response continuation and -Stage1 CUDA Graph padding regressions. The checked-in MiniCPM profile admits -two concurrent sessions through an engine-owned limit; a third -session was rejected and capacity returned after one session closed. The two -accepted sessions independently completed their audio, transcript, response, -playback, and close lifecycles. This is isolation and admission evidence for the -tested two-session deployment, not a production fairness, arbitrary-capacity, -or failure-recovery claim. The validated paths include the MiniCPM duplex -runtime extension, typed resumable policy, extracted control plane and clients, request -preregistration, typed direct-output decision, engine-managed two-session -capacity in the checked-in deploy profile, separate public/runtime configuration -channels, private Session ledgers, ordered `session.update`, the MiniCPM runner -fast path, continuous -browser input, and the final Realtime input-lifecycle fixes. Exact evidence and -scope limits are recorded below. - -## Scope - -The checkpoint keeps these verified contracts: - -- MiniCPM Stage0 conversation KV continuity; -- Stage1 TTS and Token2Wav continuity; -- model-owned listen/speak decisions on the normal auto-response path; -- continuous browser PCM upload during assistant playback, without browser VAD - or browser-generated input commits; -- segment EOS and turn EOS as different boundaries; -- transcript/audio cursors scoped to a response and turn; -- playback acknowledgement and history commit; -- Stage0 request metadata routed per batch row rather than inherited from the - first request in the batch; -- Stage0 token timing snapshots and client-observed audio chunk cadence; -- scheduler data-plane append over a resumable request; -- stale epoch/turn/response fencing; -- `/v1/duplex` and OpenAI Realtime projection; -- optional OpenAI `server_vad` turn detection backed by a process-shared Silero - model and per-session VAD state; -- existing JoyVL behavior. - -The checkpoint does not claim: - -- scheduler-native KV append; -- an acoustic-onset-to-cancel latency target across arbitrary client chunk sizes; -- production multi-session admission, fairness, capacity, or failure recovery; -- bounded long-session KV; -- audio/video synchronization beyond one camera frame per model unit (camera - frames ride the audio appends as `video_frames`; the wire contract is in the - [Realtime Duplex API](../serving/realtime_duplex_api.md) guide). - -## Active Runtime Path - -The actual path is: +## Architecture ```text -WebSocket - -> OmniDuplexSessionHandler - -> DuplexSession + MiniCPMO45ServingSessionState - -> DuplexRequestClient - -> AsyncOmni thin open/append/signal/close proxies - -> DuplexControlClient over the engine-owned correlated RPC transport - -> DuplexControlPlane + DuplexSessionRuntimeManager - -> MiniCPM DuplexRuntimeExtension + engine session/stage bindings - -> StagePool - -> resumable scheduler request - -> MiniCPM Stage0 - -> Stage1 TTS / Token2Wav - -> output processor - -> Realtime protocol projection - -> WebSocket writer + Python users WebSocket clients + ┌──────────────────────┐ ┌────────────────────────┐ + │ InlineDuplexClient │ │ DuplexClient │ + │ (DuplexClientBase) │ │ (DuplexClientBase) │ + └──────────┬───────────┘ └───────────┬────────────┘ + │ DuplexSessionHandle │ /v1/realtime?duplex=1 (alias /v1/duplex) + │ ▼ + │ ┌──────────────────────────────────────┐ + │ │ OmniDuplexSessionHandler (thin) │ + │ │ websocket I/O, command_from_realtime│ + │ │ event.to_realtime(), attachment / │ + │ │ resume tokens / replay journal │ + │ └──────────────────┬───────────────────┘ + ▼ ▼ DuplexSessionHandle + ┌──────────────────────────────────────────────────────────────────────────────┐ + │ DuplexOmni(AsyncOmni) [entrypoints/duplex_omni.py, thin] │ + │ open_session / get_session / resume_session / detach_session / close_session│ + │ DuplexSessionHandle: submit(DuplexCommand) / events() -> DuplexEvent │ + └──────────────────────────────────────┬───────────────────────────────────────┘ + │ engine.open/close/resume/touch_session_async (RPC) + │ engine.submit_command_async (one-way) + ▼ + ┌──────────────────────────────────────────────────────────────────────────────┐ + │ DuplexOmniEngine(AsyncOmniEngine) [engine/duplex_omni_engine.py, thin] │ + │ loads the plugin, validates session_mode, builds queue messages │ + │ └─ DuplexOrchestrator(Orchestrator) [engine/duplex_orchestrator.py]│ + │ template seams + DuplexStagePort over the shared stage machinery │ + │ ├─ DuplexSessionManager admission, backpressure, lease reaper, │ + │ │ open/close/resume/touch, command dispatch │ + │ └─ DuplexSessionRunner ×N ONE session state (DuplexEngineSession), │ + │ ordered mailbox, append planning + stage │ + │ submit, output projection, response / │ + │ playback / history, VAD, typed events │ + │ plugin: DuplexModelPlugin (PipelineConfig.duplex_plugin, one/model) │ + └──────────────────────────────────────────────────────────────────────────────┘ ``` -The removed typed reducer/facade and worker-provider runtimes were not in this -path. This refactor continues to modify only the path above. +In each layer the duplex class extends the turn-based one, which extends +the shared base. The dependency runs duplex -> turn-based and never the +reverse, so the turn-based stack is unchanged and the duplex stack reuses +its request machinery instead of duplicating it: -## State Ownership +```text +OmniBase OmniEngineBase OrchestratorBase + ├─ Omni (sync, turn) └─ AsyncOmniEngine (turn requests) └─ Orchestrator (turn-based admission) + └─ AsyncOmniBase └─ DuplexOmniEngine (session msgs) └─ DuplexOrchestrator (hosts DuplexSessionManager, + └─ AsyncOmni (EngineClient) implements DuplexStagePort) + └─ DuplexOmni -The architecture allows several state objects only when their facts are -orthogonal. +AsyncOmni._create_engine -> AsyncOmniEngine AsyncOmniEngine._create_orchestrator -> Orchestrator +DuplexOmni._create_engine -> DuplexOmniEngine DuplexOmniEngine._create_orchestrator -> DuplexOrchestrator +``` -| Owner | Authoritative facts | +Each concrete class constructs its collaborator in an explicit factory +method the base calls at the right moment: `OmniBase.__init__` calls +`self._create_engine(...)`; `OmniEngineBase._bootstrap_orchestrator` calls +`self._create_orchestrator(...)` on the orchestrator thread once the stages +are initialized. `DuplexOmniEngine._create_orchestrator` passes the loaded +plugin and the session runtime config to `DuplexOrchestrator` directly. + +### Design rules + +1. **One session authority.** The whole duplex session (config, epoch/turn, + ledgers, lease, stage request resources, model state, projection state) is + one object, `DuplexEngineSession`, owned by one `DuplexSessionRunner` on + the orchestrator loop. Nothing above the engine keeps session state beyond + a handle (session id, event queue, capabilities, public session object). +2. **Typed contract.** Commands into a session are `DuplexCommand` + dataclasses (`engine/duplex/commands.py`); outputs are `DuplexEvent` + dataclasses (`engine/duplex/events.py`). `command_from_realtime()` and + `DuplexEvent.to_realtime()` derive the OpenAI Realtime JSON, so the wire + format is never hand-built, and the websocket handler and the inline client + share one conversion. +3. **Generic bases and the turn-based classes have zero duplex vocabulary.** + `OmniBase`, `AsyncOmniBase`, `AsyncOmni`, `OmniEngineBase`, + `AsyncOmniEngine`, `OrchestratorBase` and `Orchestrator` carry only + template seams; `tests/engine/test_duplex_import_boundary.py` checks that + importing the turn-based stack loads no duplex module. +4. **Every surface a duplex server exposes is backed by a session.** + `vllm-omni serve` constructs `DuplexOmni` when the pipeline declares + `duplex_plugin`; the server exposes `/v1/realtime?duplex=1` (alias + `/v1/duplex`), `POST /v1/chat/completions`, `/v1/models` and `/health`, and + every other turn-based route reports "not available". The chat route is + served by `DuplexChatCompletionsAdapter`, a Realtime client that runs one + short-lived session per request: speech becomes a committed turn, text is + seeded as the session's opening turn because a model-native model takes a + turn only when it hears speech. Whether a model can be reached with text is + `DuplexCapabilities.supports_chat_completions`; one that cannot is refused + at once rather than left to idle out. Turn-based use of the same model stays + available offline through `Omni` / `AsyncOmni`. +5. **Serving is transport only.** The websocket handler does socket I/O, + wire-envelope validation, command translation, event rendering and the + attachment/resume/replay bookkeeping; it holds no session state. +6. **No `typing.Protocol`** in the duplex surfaces: the plugin, data plane, + session state, PCM buffer, stage port and client transport seams are ABCs. + +## Package layout + +```text +vllm_omni/ +├── entrypoints/ +│ ├── omni_base.py OmniBase (+ abstract _create_engine) +│ ├── async_omni_base.py AsyncOmniBase (output pump, _route_engine_message seam) +│ ├── async_omni.py AsyncOmni (turn-based requests) +│ ├── duplex_omni.py DuplexOmni, DuplexSessionHandle (thin pipe) +│ ├── duplex/ SERVING (thin) +│ │ ├── serving.py OmniDuplexSessionHandler +│ │ ├── realtime_input.py RealtimeEnvelope (query rules, first message), parse_resume_request +│ │ ├── session_attachment.py DuplexSessionAttachmentRegistry (resume tokens, replay journal) +│ │ ├── audio_encoding.py encode_audio, injected into DuplexOmniEngine for the plugin's data plane +│ │ ├── chat_completions.py DuplexChatCompletionsAdapter (/v1/chat/completions on a session per request) +│ │ └── websocket.py websocket send/close/receive helpers +│ └── openai/api_server.py builds DuplexOmni for duplex models; session-backed app state +├── engine/ +│ ├── omni_engine_base.py OmniEngineBase (stage processes, orchestrator thread, queues, RPC) +│ ├── async_omni_engine.py AsyncOmniEngine (turn-based request building) +│ ├── duplex_omni_engine.py DuplexOmniEngine (session message surface, creates DuplexOrchestrator) +│ ├── orchestrator.py OrchestratorBase + Orchestrator +│ ├── duplex_orchestrator.py DuplexOrchestrator (+ DuplexOrchestratorRequestState; implements DuplexStagePort) +│ └── duplex/ +│ ├── commands.py DuplexCommand dataclasses, command_from_realtime +│ ├── realtime_commands.py Realtime client event -> DuplexCommand translation +│ ├── events.py DuplexEvent dataclasses (+ to_realtime), REALTIME_ERROR_TYPES_BY_CODE +│ ├── 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 +│ ├── contracts.py DuplexFence (session_id, epoch, turn_id), stage request records, DuplexStagePort +│ ├── plugin.py DuplexModelPlugin, DuplexModelSessionState, DuplexDataPlane, PcmAppendBuffer ABCs +│ ├── turn_detection.py server-side VAD turn detector used by the session +│ ├── audio.py / vad.py / intermediate.py +│ └── session/ one engine-resident session and everything that runs it +│ ├── engine_session.py DuplexEngineSession: ledgers, lease, fence, stage resources, append sequencing +│ ├── runner.py DuplexSessionRunner (per-session mailbox on the orchestrator loop) +│ ├── manager.py DuplexSessionManager (admission, backpressure, reaper, dispatch) +│ ├── context.py DuplexSessionContext / DuplexRunState: what the runner shares with its components +│ ├── emitter.py SessionEmitter: projection, epoch filter, domain effects of a terminal event +│ ├── model_channel.py ModelChannel: submit an append, project stage output, continue a turn +│ ├── control.py SessionControl: server VAD and the events that reconfigure it +│ ├── append_task.py AppendAttempt: one append in flight and the rollback its failure owes +│ ├── helpers.py pure reads and payload builders over a session +│ ├── lease.py DuplexLeaseState (idle TTL, disconnect grace, resume generation) +│ └── overlap_policy.py / commit_policy.py / playback_ledger.py +├── config/stage_config.py PipelineConfig.duplex_plugin; DuplexSessionRuntimeConfig +├── model_executor/models/minicpmo_4_5/duplex/plugin.py MiniCPMO45DuplexPlugin (+ data_plane, input, policy, ...) +└── clients/ + ├── duplex.py DuplexClientBase (ABC), DuplexClient (websocket), client-side events + ├── inline_duplex.py InlineDuplexClient (in-process, over DuplexOmni) + ├── minicpmo_4_5.py MiniCPM-o 4.5 session preset + └── personaplex.py PersonaPlex session preset +``` + +## Session identity + +- **The server allocates every session id.** `DuplexOmni.open_session` + generates `duplex-`; a `session_id` / `id` key inside the + Realtime session object is ignored (Realtime clients echo the session object + back). The allocated id reaches the client in `session.created` and is the + only handle for `session.resume`, `close` and every command. +- **Ids are never reused** within an engine's lifetime, so the id alone + identifies a session. There is no incarnation counter: stage request ids are + `duplex-s..e..r.`, model-side per-session state + is keyed by `session_id` alone, and a command, close, touch or resume for an + id the manager no longer holds is answered with `unknown_session`. +- **Staleness inside a live session is epoch-based.** `DuplexFence + (session_id, epoch, turn_id)` is an engine-internal identity: every stage + request is bound to the fence it was submitted under, cancels advance the + epoch, and tracked tasks re-check the epoch after each `await`. +- **Resume identity** is `(session_id, resume_token, expected_lease_generation)`: + the token belongs to the attachment registry, the lease generation is the + engine-side compare-and-swap. + +## Session lifecycle + +```text +open DuplexOmni.open_session(config) + -> engine.open_session_async(session_id, DuplexSessionConfig) [RPC] + -> DuplexSessionManager.open: admission (max_sessions, closing sessions keep their slot), + plugin.validate_client_extra_body / prepare_runtime_config, DuplexEngineSession, + Stage0 request id reserved (ensure_stage_request), DuplexSessionRunner.start() + -> first event: SessionCreated (announces the allocated id) +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() +detach DuplexOmni.detach_session -> touch(DETACH): engine-owned disconnect grace; expiry -> SessionExpired +resume DuplexOmni.resume_session(expected_lease_generation) -> lease CAS; the existing handle is re-entered +close DuplexOmni.close_session -> close RPC; the manager tears the runner down, then the stage cleanup + (abort submitted requests, release reserved ids), then SessionClosed, then the RPC result: seeing + the event means the admission slot is free (it is held until the cleanup succeeded) +reap DuplexSessionManager.reaper_loop: idle TTL / disconnect grace expiry, cleanup retries +``` + +### Concurrency and ordering model of the runner + +Everything below runs on the orchestrator asyncio loop; there is no lock. + +- **Single writer per session.** `DuplexEngineSession` is mutated only by its + runner. Inputs reach the runner through one mailbox in this order: client + commands (from `_request_handler`, in request-queue order), stage outputs + (`on_stage_output` enqueues an internal item; it does not mutate the session + inline), timers (silence continuation, bounded auto-response) and internal + items (deferred overlap promotion, VAD-synthesized commits). The mailbox + worker processes items one at a time. +- **Slow work is not awaited inline.** An `AppendAudio` handler reserves + bytes, then schedules a tracked append task on the per-session append tail: + the task awaits `_offload(decode/resample/VAD)` on a dedicated + `ThreadPoolExecutor` owned by the manager, then `plugin.plan_append`, then + `stage_port.submit`. The worker returns to the mailbox immediately, so a + later `CancelResponse` is not blocked behind an in-flight append. +- **Re-validation after every `await`.** A tracked task captures `(epoch, + turn_id)` when it starts and re-checks them before the stage submit and + before committing ledgers; a task that finds the epoch advanced rolls back + its PCM reservation and exits. +- **Cancel is atomic at the session.** `CancelResponse` / `BargeIn` advance + the epoch, cancel tracked append tasks, abort the owned stage bindings + through the stage port and emit the cancelled terminal events before + yielding. An append task already past its last check is harmless: its stage + request is bound to the old epoch's request id, which the abort covers, and + the base never emits session-owned outputs to a client. +- **Ordered event stream.** Because the epoch bump and every emission happen + 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. +- **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 + `ErrorEvent(code="input_backpressure")`. +- **Per-response stage metrics.** `_route_output` computes the per-segment + `StageMetrics` before calling `_intercept_stage_output`; the runner + accumulates the snapshot into the active response, so + `ResponseDone.stage_metrics` and `metadata.vllm_omni.stage_metrics` on the + wire keep their meaning. + +## Orchestrator seams + +`OrchestratorBase` keeps the generic machinery (request handler skeleton, +abort, collective RPC, output loops, stage error and dead-replica handling, +`_route_output`, stage forwarding, prewarm, cleanup, shutdown) and offers +these seams. `DuplexOrchestrator` extends `Orchestrator` rather than sitting +beside it: a seam it does not override falls through to the turn-based +implementation, and `_dispatch_message` hands anything that is not a session +message to `super()`: + +| Seam | `Orchestrator` | `DuplexOrchestrator` | +| --- | --- | --- | +| `_dispatch_message(msg) -> bool` | `add_request`, `streaming_update`, `add_companion_request`, `interaction` | open/command/close/resume/touch messages -> `DuplexSessionManager` | +| `_background_tasks()` | — | `[manager.reaper_loop()]` | +| `_shutdown_extensions()` | — | `await manager.shutdown()` | +| `_on_stage_submitted(stage_id, request_id, replica_id, req_state)` | — | bind the forwarded stage request to the owning session | +| `_intercept_stage_output(...) -> bool` | — | hand every session-owned output (Stage0 decision, Stage1 audio) plus its metric snapshot to `runner.on_stage_output`; `True` = do not forward | +| `_handle_forward_failure(...) -> bool` | — | fail the owning session's response and close the session, keep the loop alive | +| `_cleanup_request_ids(ids, abort=, release_owners=)` | ignores `release_owners` | brackets `super()` with `close_sessions_for_request_ids` / `defer_request_cleanups` / `finalize_closed_sessions` | +| `OrchestratorRequestState.session_owned` | never set | set by `ensure_request`: no synthetic terminal, no client emission, no terminal re-forward, no auto-cleanup on finish | + +`DuplexOrchestrator` also implements the `DuplexStagePort` ABC the manager and +runner use (`stage_count`, `sampling_defaults`, `ensure_request`, `submit`, +`cleanup`, `abort_requests`): the resumable Stage0 request is preregistered at +open with `DuplexOrchestratorRequestState(session_owned=True, session_id, +fence)`, submitted with `submit_initial` on the first unit and +`submit_update` afterwards, and the request's `streaming.bridge_states["duplex"]` +carries the session identity (`session_id`, `fence`, `epoch`, `turn_id`, +`model_turn_id`), the public session config and the server-owned runtime +config for the worker-side model hooks. + +## Model plugin + +`DuplexModelPlugin` (ABC, `engine/duplex/plugin.py`) is the one class a model +provides, selected by `PipelineConfig.duplex_plugin`; "is a duplex model" is +`duplex_plugin is not None`. It merges the engine policy and the session +policy that used to be two separately configured objects: + +| Half | Members | | --- | --- | -| `DuplexSession` | session state, epoch, turn, active/last response, pending input, overlap duration, playback cursor, history | -| `MiniCPMO45ServingSessionState` | MiniCPM PCM buffer, deferred overlap payload, data-plane task, continuation effect counter | -| `DuplexWebSocketActor` | inbound mailbox, outbound queue, writer, transport close state, effect task handles | -| `NativeRealtimeSessionProtocol` | wire-only response/item projection and input-buffer projection | -| `DuplexSessionRuntimeManager` | fence snapshot, append reservation, completed-operation cache, resource reservations, stage bindings | -| `DuplexControlPlane` | control-message validation, transactional open/append/update/close, extension invocation, typed segment/output decisions | - -Important distinctions: - -- Realtime response IDs are projection caches, not domain decision sources. -- Engine fences duplicate identity across a process boundary by design; they - are immutable validation snapshots. -- Stage bindings are engine resources, not serving response state. -- `continuation_owner_id` scopes an effect counter either to a visible response - or, before `response.created`, to an explicit model turn. A pre-response - continuation also revalidates request, epoch, turn state, and model-turn - identity immediately before append; it cannot survive turn completion. -- runtime open/close acknowledgements are handler-local effect bookkeeping, - not Session lifecycle. - -## Implemented Concurrency Contracts - -### Single inbound mailbox - -The Actor now uses one: - -```python -asyncio.Queue[event] -``` - -This replaces competing control/input/event `Queue.get()` tasks and the -deferred-event list while preserving WebSocket arrival order. A later close, -cancel, or clear must not overtake earlier append, playback ACK, or -session-update events. Slow runtime appends remain background tasks. Cancelling -their asyncio task does not stop a synchronous callable already running in the -default executor, so cancellation correctness is provided by the engine fence, -not by `Task.cancel()`. - -Properties covered by focused tests: - -- every enqueued event is delivered exactly once; -- terminal control preserves wire order while cancelling background work; -- clear and cancel never scan or delete later mailbox events; -- one writer owns all WebSocket sends. - -### Correlated RPC result router - -`AsyncOmniEngine` now has one consumer for `rpc_output_queue`. Results route by: - -```text -("duplex", control_id) -("collective", rpc_id) -``` - -The router replaces the global RPC lock. It supports out-of-order replies, -timeout unregister, late-result rejection, fatal-error broadcast, and close -unblocking. - -Additional lifecycle rules: - -- uncorrelated non-fatal errors are not broadcast to unrelated waiters; -- a latched fatal rejects new calls before they enter the request queue; -- `EngineDeadError` is broadcast directly to the RPC queue even though the - orchestrator consumes that exception during teardown; -- shutdown closes the router before joining the orchestrator; -- failure to enqueue shutdown closes the request queue and still attempts a - bounded orchestrator join; -- `StageRuntime` cleanup runs only after the orchestrator has stopped, avoiding - concurrent client teardown. - -The router removes wait-side head-of-line blocking. The orchestrator request -handler remains ordered, but executor submission order is not itself a session -ordering guarantee. - -### Extracted control ownership - -The generic engine mechanism and the duplex plugin are split by -responsibility: - -| Module | Responsibility | -| --- | --- | -| `engine.messages` | ordinary engine queue messages and collective RPC envelopes | -| `engine.rpc_result_router` | generic correlation routing through each result message's `rpc_correlation_key` | -| `engine.duplex.contracts` | immutable duplex DTOs, extension/stage protocols, typed decisions, request-ID codec | -| `engine.duplex.lease` | duplex lease activity, TTL/grace configuration, and expiry records | -| `engine.duplex.messages` | identity fence and duplex control message/result envelopes | -| `engine.duplex.runtime` | extension loading, validation, and compatibility exports | -| `engine.duplex.session` | session resources, fences, append reservations, idempotency cache, binding cleanup | -| `engine.duplex.control_plane` | control algorithms and the narrow `DuplexStagePort` used to submit and clean up stage requests | -| `engine.duplex.control_client` | typed control-message construction and calls through the engine-owned correlated RPC transport | -| `entrypoints.duplex_request_client` | deterministic request identity, client-state preregistration, acknowledgement validation, data-plane collection, rollback | - -`Orchestrator`, `AsyncOmniEngine`, and `AsyncOmni` retain explicit wiring and -compatibility proxies. They no longer own the corresponding duplex algorithms. -The ControlClient does not own waiters or consume the result queue; -`RpcResultRouter` remains the single correlation owner for both duplex and -collective RPC. The ControlPlane receives a narrow stage port rather than the -Orchestrator object. - -Control-plane enablement is independent from extension presence. -`PipelineConfig.duplex_control_enabled` determines whether an Orchestrator -constructs a ControlPlane. This preserves extension-free -`TURN_COMMIT_ONLY` deployments while allowing ordinary deployments to use a -single `None` fast-path check. The MiniCPM pipeline explicitly enables control -and separately selects its runtime extension. - -### Transactional open and append - -Opening a native session registers its engine state and reserves the Stage0 -request resource as one externally atomic operation. If capability, sampling, -extension, or request-state initialization fails, the ControlPlane removes the -session, reserved request resources, and any preregistered request state before -returning the error. Cleanup distinguishes a reserved resource from a submitted -request, so it does not decrement the running counter for unrelated work. -Closing a session that never appended also removes the preregistered resource. - -Expiry and request-triggered cleanup use a two-step close. The runtime manager -first marks the session closing but retains it in the admission set. Only a -successful stage cleanup finalizes and removes the session. A blocked or failed -cleanup therefore continues to consume its server-owned slot and is retried; -replacement admission cannot overlap request or KV resources that are still -owned by the closing session. - -Append uses prepare/submit/commit ordering. `prepare_append()` computes the -next sequence values without mutating the session. Prompt planning and stage -submission use that reservation, and only a successful submit commits -`input_seq`, `input_turn_seq`, and the accepted fence. Serving uses a matching -PCM reservation and consumes bytes only after the append acknowledgement; -failure or a closed predecessor rolls the bytes back. An `operation_id` caches -completed append results so a retry cannot submit the same physical append -twice. - -### Ordered append and configuration effects - -Native append effects form a per-session tail. Each append awaits its -predecessor before entering the engine, and `session.update` awaits that same -tail before changing either the serving aggregate or engine configuration. -Consequently, an append received before an update uses the old immutable -sampling snapshot, while the first accepted append after the update uses -the new generation. - -The engine increments `config_generation` when it accepts the replacement -configuration. At the next append boundary, the control plane rebuilds sampling -parameters from stage defaults through the model extension. It never mutates -the snapshot of a segment already in progress. - -### Irreversible cancellation fence - -Every cancel control carries the cancelled fence and the next session fence. -The `AsyncOmni` facade preserves both values across the serving, -entrypoint, engine-message, and orchestrator boundaries. The orchestrator then -atomically: - -1. advances the runtime state to the next fence; -2. releases and aborts bindings owned by the cancelled fence; -3. permanently rejects an append that arrives later with the old fence. - -Missing `next_fence` is rejected without releasing bindings; the engine never -advances independently from serving. A late cancel cannot move a session -backward when the runtime is already on a newer fence. Idle playback clear does -not send an engine cancel because no runtime resource is being cancelled. - -`model_native_duplex` also validates its engine-client contract before opening -the session. The open, append, signal, and close methods must exist and declare -the required `fence` arguments; signal must additionally accept `next_fence`. -An older or custom client that cannot preserve this boundary fails with -`runtime_contract_invalid` instead of silently running an unfenced native -session. The generic serving-session adapter retains its compatibility behavior. - -`DuplexFence` also carries a server-side session incarnation. Reusing a public -session ID after close increments that incarnation, and resource request IDs -include it after the first incarnation. A synchronous append from a closed -session therefore cannot be accepted by a newly opened session with the same -public ID. Stale close validates the full fence before removing the runtime -registry entry or releasing bindings. - -## Implemented Serving Ownership - -### Session-owned response and overlap state - -The Actor no longer stores: - -- a Session reference; -- lifecycle mirror state; -- active or last response ID; -- overlap duration; -- response-in-progress or playback business predicates; -- runtime open/close acknowledgements. - -`DuplexSession.begin_response()` owns both active and last response identity. -The Session owns overlap accumulation/reset and playback state. - -One `native_response_in_progress()` predicate now covers: - -- active response ID; -- active engine request ID; -- uncommitted ACK-only playback; -- active generic response task; -- response-bound append tasks; -- MiniCPM data-plane task; -- assistant-generating turn state. - -The write-only `MiniCPMO45ServingSessionState.response_emitted` field and a dead -data-plane task mirror were removed. - -### Response-scoped options - -`response.create` is represented by an immutable `ResponseCreateOptions` value. -The Session reserves it for one response, applies it to a response-local copy -when that response begins, and restores the base session configuration on -completion, listen, cancellation, or failure. Instructions, voice, modalities, -temperature, token limit, format, tools, and conversation metadata therefore do -not mutate the permanent session configuration or leak into a later response. - -### Explicit event effects and pure writer - -Business paths call `emit_event()`. It serializes domain effects and outbound -queue insertion with one session-local lock. Deferred append starts outside the -lock so recursive event emission cannot deadlock. - -The writer only: - -1. dequeues an event; -2. applies late stale filtering to streaming model output; -3. projects the event to the selected wire protocol; -4. sends JSON. - -The writer does not promote overlap, clear input, advance Session state, call -the engine, or commit history. - -### Single terminal acceptance point - -Terminal events are accepted or rejected before domain effects run: - -```text -response.done -response.listen -audio.cancelled -input.cancelled -session.closed -``` - -For a terminal carrying an epoch, that epoch must match the current Session. -Stale cancellation therefore cannot reset overlap belonging to a later turn. -Once a terminal is accepted and its effect is applied, the writer cannot revoke -it because close or epoch state changed while it waited in the output queue. - -Streaming audio/text deltas still receive late stale filtering at send time, -which prevents queued old audio from leaking after a cancel or epoch bump. - -## Runtime Type Boundary - -The stable kernel is deliberately limited to generic engine and scheduler -mechanisms: - -```text -vllm_omni.engine.messages # ordinary queue and collective RPC envelopes -vllm_omni.engine.rpc_result_router # generic correlation routing -vllm_omni.engine.async_engine_utils # ordinary engine lifecycle helpers -``` - -The full-duplex implementation lives in dedicated duplex modules: - -```text -vllm_omni.engine.duplex.runtime - # extension loading and validation -vllm_omni.engine.duplex.contracts - # immutable DTOs and narrow protocols -vllm_omni.engine.duplex.lease - # duplex lease primitives -vllm_omni.engine.duplex.messages - # duplex identity and queue envelopes -vllm_omni.engine.duplex.session - # session transaction implementation -vllm_omni.engine.duplex.control_plane - # engine-side control algorithms -vllm_omni.engine.duplex.control_client - # correlated control client -vllm_omni.entrypoints.duplex_request_client - # entrypoint request/output lifecycle -``` - -Stable engine and entrypoint modules import the duplex kernel eagerly at -module load (top-of-file imports). Duplex behavior remains opt-in: the -ControlPlane, ControlClient, and request client are *constructed* only on -duplex-enabled deployments, and the model runtime extension still loads only -through the dotted-string plugin path. The client-side package -(`vllm_omni/clients/`, home of the public `DuplexClient`) is never imported -by runtime code. - -`PipelineConfig.duplex_control_enabled` enables the generic mechanism, while -`PipelineConfig.duplex_runtime_extension` separately selects the model adapter. -The MiniCPM pipeline opts into both and installs -`MiniCPMO45DuplexRuntimeExtension`; ordinary pipelines construct neither a -ControlPlane nor a ControlClient. The model-neutral protocol exposes -sampling-parameter configuration, append planning, and a typed stage-output -decision. MiniCPM owns -stage-specific overrides, PCM-to-token budgeting, force-listen payload policy, -special-token interpretation, and `listen` response metadata. The generic -Orchestrator does not parse MiniCPM session keys, `listen_token_id`, or construct -MiniCPM metadata, and it does not import the adapter. - -Extension loading is fail-fast. Startup validates every required callable and -the number and type of returned stage sampling parameters. Invalid custom -extensions fail before a session can open. - -Direct listen/speak outcomes use `DuplexOutputDecision` on -`OmniRequestOutput`. That typed envelope is authoritative over inner completion -metadata for native projection, so a processed output cannot hide a listen -decision merely because it retains unrelated multimodal metadata. - -Generic multimodal chunk accumulation lives in -`vllm_omni.outputs.multimodal_accumulation`. It owns snapshot replacement, -DELTA draining, and audio-finality interpretation; `OmniRequestState` only -invokes that policy while assembling request output. Model output processors no -longer need to add chunk-retention rules to the main output orchestration file. - -The strict chunk-transfer schema still contains `hidden_states.tts` as the -generic AR-stage-to-TTS handoff and several MiniCPM streaming token IDs as -compatibility wire fields. The latter are not a preferred generic model -contract. Moving them behind a model-namespaced metadata envelope requires a -versioned producer/consumer migration because unknown msgspec fields are -rejected; silently deleting or renaming them in this checkpoint would break -existing Stage0-to-Stage1 transfer. - -### Public and runtime configuration channels - -Control messages carry two independent snapshots: - -- `session_config` contains the public Realtime/session values used by serving - and model prompt policy; -- `runtime_config` contains server-derived MiniCPM sampling, scheduler, context - reserve, prefix-budget, and resolved reference-audio values. - -The MiniCPM serving adapter derives `runtime_config` from deploy/model defaults -and validated public options. Clients cannot provide private runtime keys in -`session.create` or `session.update`; response-scoped `extra_body` drops those -keys rather than turning them into engine policy. The ControlPlane constructs -sampling parameters only from `runtime_config`, -while append planning receives both snapshots explicitly. Raw Realtime -`extra_body` is not used as the engine's sampling-policy channel. - -`session.update` prepares both candidates, validates the runtime candidate in -the engine, and replaces the serving snapshots only after the engine ACK. A -rejected candidate therefore leaves both the public configuration and runtime -sampling generation unchanged. - -`DeployConfig.duplex_session` is an immutable, server-owned -`DuplexSessionRuntimeConfig`. These limits are not read from client -`extra_body` or session payloads: - -| Field | Default | Owner | Lifecycle or resource effect | -| --- | ---: | --- | --- | -| `idle_ttl_s` | `300.0` | Engine lease manager | Expires an attached or detached session after inactivity; `null` disables only idle expiry. | -| `disconnect_grace_s` | `30.0` | Engine lease manager and API attachment registry | Bounds how long a detached session may retain engine resources before cleanup. | -| `reaper_interval_s` | `5.0` | Orchestrator | Sets the cadence for lease expiry and pending cleanup retries. | -| `resume_replay_ttl_s` | `60.0` | API attachment registry | Expires replay events retained for resume or takeover. | -| `resume_replay_max_bytes_per_session` | `8388608` | API attachment registry | Bounds each session's replay buffer. | -| `max_pending_input_bytes_per_session` | `16777216` | API session input ledger | Applies per-session byte backpressure before input is admitted. | -| `max_pending_turns_per_session` | `4` | API session input ledger | Bounds queued turns that have not completed runtime processing. | -| `max_sessions` | `1` | Engine session manager | Atomically enforces admission; closing sessions retain capacity until cleanup finalizes. | -| `completed_append_cache_size` | `256` | Engine session manager | Bounds completed append idempotency records per session. | - -Deploy profiles may lower or raise these values to match scheduler capacity, -but clients cannot override them. The engine remains authoritative for -admission and cleanup ownership even when the API performs an early check. - -### Client-side companion package - -The public client surface ships under `vllm_omni/clients/` and stays on the -consumer side of the boundary above: `vllm_omni.clients.duplex` provides the -model-agnostic `DuplexClient` (typed events, per-response demux via -`ResponseHandle`, `session.resume` with acknowledgement-based replay dedup) -and the `EventCollector` helper used for timing and request-metric summaries. -Model-specific session defaults live in thin per-model preset modules — -`vllm_omni.clients.minicpmo_4_5.create_duplex_session_config()` and -`vllm_omni.clients.personaplex.create_duplex_session_config()` — so the core -client never imports model packages and the model duplex plugin packages -never import the client (both directions are enforced by tests: -`tests/engine/test_duplex_import_boundary.py`, -`tests/entrypoints/openai_api/test_duplex_capability.py`, and the subprocess -purity guards in `tests/clients/test_model_session_configs.py`). -`examples/online_serving/barge_in_client.py` (with its flow companion -`barge_in_client_flow.md`) is the runnable reference for the client library -against a live duplex deployment. The wire vocabulary the client speaks — -every `/v1/realtime?duplex=1` event with its OpenAI-compatibility tier and a -JSON example, plus the native-to-Realtime name map — is catalogued in the -[Realtime Duplex API](../serving/realtime_duplex_api.md) serving guide; that -page derives from the normative contract table above and must be updated -with it. - -## Generic-Path Cleanup - -Model-specific `MINICPMO45_PROFILE_LOGS` probes were removed from: - -- the generic AR scheduler; -- orchestrator; -- AsyncOmni; -- the generic GPU AR runner. - -The cleanup removes the temporary serving and generic-runtime probes. The -generic AR runner builds `DuplexSamplingRow` values only when the model exposes -the optional `prepare_duplex_sampling()` hook, then invokes that hook exactly -once before the normal sampler. Ordinary models do not scan request rows. -MiniCPM owns force-listen logits, turn-ended state, and its row-local sampling -policy behind that hook; the runner no longer reads `_minicpmo45_*` attributes -or retries sampling after matching `TypeError` text. - -The scheduler no longer parses serving `session_config`, `extra_body`, or -`duplex_stage_sampling_params`. Stage-specific overrides are materialized as -ordinary `SamplingParams` by the model runtime extension before scheduler -admission. Existing vLLM stop-token, EOS, and max-token handling defines each -segment boundary; `resumable=True` parks the request for the next streaming -update. The scheduler does not inspect a duplex dictionary, Realtime -configuration, or MiniCPM sampling fields. - -The generic runner's unused `_duplex_force_listen_applied_segments` cleanup and -the obsolete `streaming_accumulated_keys` accumulation hook were removed. The -latter had no model producer after its upstream Qwen removal; streaming input -still performs the ordinary replacement merge needed by active models. - -Local `turn.signal` transitions such as `user_started` and -`assistant_started` no longer perform an engine RPC. Only cancellation and -`session.update` cross the engine boundary because those operations change -runtime resource identity or runtime configuration. - -During a stage transition, `StreamingInputState.source_token_decoder` exposes a -short-lived, model-neutral decoder sourced from the upstream output processor. -The MiniCPM Stage1 input processor consumes that capability for token-faithful -delta transcripts. The generic orchestrator no longer installs or names a -MiniCPM-specific bridge-state key. - -The stable API server imports the duplex serving package eagerly and inspects -the deployment `session_mode` before constructing the duplex handler. Ordinary -deployments do not create its registries and background state. -For an enabled deployment, the client still selects the Realtime duplex route -with `?duplex=1` (or an equivalent explicit true value). Model-name matching is -not used for routing or native-runtime activation. MiniCPM clients explicitly -set `extra_body.native_duplex=true`; the canonical way to do so is -the `vllm_omni.clients.minicpmo_4_5.create_duplex_session_config()` preset, -which the repository's barge-in example -(`examples/online_serving/barge_in_client.py`) uses in its session payload. - -`DuplexRequestClient` derives the data-plane request identity from the accepted -fence and preregisters a resumable `ClientRequestState` before the append crosses -the engine boundary. The acknowledgement must return the same identity or the -operation fails. Unregistered outputs, including names that merely use the old -`duplex-` prefix, are dropped. Timeout, cancel, and close remove the -preregistered state. - -## Why Scheduler Changes Remain - -Serving cannot preserve model KV after a segment ends. The scheduler must keep -one resumable request in a waiting state and accept a later update: - -```text -RUNNING - -> segment stop -WAITING_FOR_STREAMING_REQ (KV retained) - -> next append -RUNNING - -> session close -FINISHED (KV released) -``` - -Scheduler responsibilities are limited to resumable request state, runtime -context update, stop/boundary handling, and final release. Scheduler must not -own response IDs, playback, overlap, Realtime events, or model policy. - -## Serving and Compatibility Boundaries - -### Normative Realtime event contract - -The public contract in this section applies to `/v1/realtime?duplex=1`. Internal -duplex event names such as `response.output_audio.delta` are projection inputs, -not additional public aliases. Events for one attachment preserve mailbox -order; replay after resume preserves the original event order and `event_id`. - -Accepted compatibility aliases are normalized before mailbox admission: - -| Canonical input event | Accepted aliases | Contract | -| --- | --- | --- | -| `input.text.append` | `input_text.append`, `push_text` | Appends text to the open input item. | -| `input_audio_buffer.append` | `input.audio.append`, `push_chunk` | Appends audio; legacy aliases default to WAV when no format is supplied. | -| `turn.signal` | `signal_turn` | Sends an explicit turn signal; it does not imply response creation. | -| `playback.ack` | `audio.playback_ack` | Advances playback/history only for the identified response and attachment. | -| `session.close` | `close_session`, `close` | Closes the session after previously admitted mailbox events. | - -The Realtime projector emits the following response lifecycle. Cardinality is -scoped to one `response_id` unless stated otherwise. - -| Public event | Trigger and ordering | Cardinality | -| --- | --- | --- | -| `response.created` | First event for a visible assistant response; precedes its output-item, content-part, speak, text, and audio events. | Exactly once for every visible response. | -| `response.speak` | The model selected speak; emitted after `response.created` and no later than the first `response.output_audio.delta`. It carries decision metadata, not transcript text. | At most once. | -| `response.output_audio.delta` | Carries one ordered audio chunk. Every emitted audio delta is followed by a `response.output_audio_transcript.delta` for the same chunk. | Zero or more. | -| `response.output_audio_transcript.delta` | Append-only transcript contribution paired with an audio delta; it may be empty for a text-less audio unit and must not repeat or overlap earlier text. | One per audio delta. | -| `response.output_audio.done` | Closes the audio stream after its final delta and before response terminal events. It is omitted for a response with no audio or transcript. | At most once. | -| `response.output_audio_transcript.done` | Contains the exact concatenation of all non-overlapping transcript deltas for the response. | At most once, and only for a non-empty transcript. | -| `response.output_item.done` / `conversation.item.done` | Finalize the assistant item after content-part terminal events and before `response.done`. | At most once each. | -| `response.done` | Terminal event for a created response; follows audio, transcript, content-part, and item terminal events. | Exactly once for every created response that reaches a terminal state. | -| `rate_limits.updated` | Compatibility event emitted immediately after `response.done`; the current payload has an empty rate-limit list. | Once per emitted `response.done`. | -| `response.listen` | The model selected listen at a model-unit boundary. It may be emitted without `response.created`; clients must not infer a missing or pending response from it. | Zero or more per session, at most once per model-unit decision. | - -Session lifecycle events (`session.created`, `session.updated`, -`session.resumed`, `session.replaced`, `session.expired`, and -`session.resync_required`) are session-scoped rather than response-scoped. -`session.created` is emitted before held response output on a new attachment. -Cancellation and truncation use terminal status/details on the affected -response and never reuse that response ID for later model turns. - -Compatibility changes must update this table and its protocol contract tests in -the same change. The golden transcript test requires joined -`response.output_audio_transcript.delta` values to equal -`response.output_audio_transcript.done.transcript`; response tests separately enforce -at-most-once `response.speak` and terminal event cardinality. - -### Session internal ledgers - -`DuplexSession` remains the serving aggregate root and now composes -`InputBufferState`, `ResponseState`, `PlaybackLedger`, and -`ConversationHistory`. The ledgers are private. Read-only properties return -tuples, snapshots, or mapping proxies, and serving modules mutate them through -Session transition methods for request, response-turn, lifecycle, playback, -capability replacement, and history. The Realtime session-update parser still -normalizes individual fields before atomically retaining or rolling back the -aggregate configuration. Further work may narrow the read-only compatibility -surface, but it must not introduce another reducer or a second session state -machine. - -### Realtime input ownership - -The browser is a continuous PCM producer. While the session is open and the -microphone is unmuted, it sends `input_audio_buffer.append` every 200 ms, -including while assistant audio is being generated or played. It does not run -VAD and does not send `input_audio_buffer.commit`. Native Stage0 consumes the -stream in approximately one-second model units; MiniCPM listen, speak, and turn -EOS decisions advance the model conversation. - -Explicit Realtime clients may still use `input_audio_buffer.commit` to create a -conversation item. The translator validates that wire input before producing a -commit carrying `item_id`. Native auto-response may already have -streamed those PCM samples into the runtime, so the runner accepts a validated -commit even when no runtime-side chunk remains. A truly empty explicit buffer -continues to return `input_audio_buffer_empty`. - -The commit-to-response contract is: - -| Session state at `input_audio_buffer.commit` | Runtime action | Required event outcome | -| --- | --- | --- | -| Auto-response disabled, valid idle input | Commit and retain the input item. | Emit `input_audio_buffer.committed`; emit no response until an explicit `response.create`. | -| Auto-response enabled, valid idle speech | Submit the final native append for the current model turn. | Emit `input_audio_buffer.committed` before model output. Create a response only on the first visible text/audio output; a model listen decision may end with `response.listen` and no response. | -| Auto-response enabled, short overlap during an active response | Discard the short overlap as an acknowledgement. | Emit the committed/no-response acknowledgement; do not create or defer a response. | -| Auto-response enabled, meaningful overlap during an active response | Retain one deferred payload and promote it after the active response terminates. | Emit `input_audio_buffer.committed` with deferred metadata. A later response follows the same visible-output rule and never reuses the prior response ID. | -| Explicit silence/noise (`is_speech=false`) | Clear the physical input without a runtime append. | Emit the committed/no-response acknowledgement followed by `response.listen`. | -| Empty or invalid input | Reject without changing response ownership. | Emit one typed `error`; do not emit `response.created`. | -| Previous response cancelled or truncated and cleanup complete | Apply the corresponding idle rule above. | The prior response remains terminal; any later visible output receives a new response ID. | - -Append tasks, explicit commits, cancellation, and deferred promotion are -serialized by the session actor. Tests cover post-response commits, active -playback deferral, truncated-response ownership, and two-turn -response-required execution; the H20 validation record below is the golden -multi-turn runtime evidence. - -During auto-response overlap, `preserve_realtime_input` distinguishes "do not -append this silent chunk to the native buffer" from "clear the open Realtime -item". Silent overlap no longer discards earlier user PCM. - -The first chunk of one overlapping input item also reserves its target model -turn. A later Realtime commit uses that reserved identity even if response EOS -has already advanced the Session turn; clear, cancel, close, and successful -promotion release the reservation. - -### Physical input and model-turn identity - -A Realtime input item, an optional `input_audio_buffer.commit`, and a MiniCPM -model turn are intentionally different identities. The browser normally keeps -one input stream open without commits. Native Stage0 evaluates that stream in -approximately one-second model units. A sampled `<|turn_eos|>` closes the -current model turn and allows a later streamed unit to start another turn. -Explicitly committed input may also produce only model-listen decisions and no -spoken response. - -Therefore neither `model_turn_id` nor response cardinality is derived from the -number of physical inputs. `model_turn_id` advances at model EOS. The -`response-required` fixture may require exactly one response per requested -turn, but `model-policy` validation must accept additional model-owned responses -and cannot assert cross-turn transcript independence through physical-input -attribution. -Before closing, the model-policy demo drains every created response through -`response.done`, acknowledges every completed audio playback, and requires a -bounded quiet interval with no newly created response. This validates lifecycle -completion without suppressing or overriding the model's listen/speak policy. - -### Non-cancel signal semantics - -`turn.signal` remains a public compatibility surface for local Session and -Realtime transitions. Serving forwards only cancellation and `session.update` -to the engine. The engine rejects arbitrary non-cancel signal names instead of -returning a misleading supported acknowledgement. `session.update` remains a -real engine operation because it replaces the configuration consumed by later -append plans. - -## Remaining Architecture Work - -The following work is deliberately not hidden inside this checkpoint. - -### Serving composition - -The generic runner and runtime bridge now depend on the model-neutral -`ServingRuntimeAdapter` protocol and no longer import MiniCPM modules. The -`MiniCPMO45ServingRuntimeAdapter` owns MiniCPM serving state, PCM preparation, -data-plane projection, client/runtime configuration validation, prefix policy, -and capability projection. `PipelineConfig.duplex_serving_adapter` explicitly -selects its import path; `AsyncOmniEngine` and `AsyncOmni` carry that path to the -API server, which passes it into the generic handler. The handler has no model -default and fails startup when the duplex endpoint is enabled without an -adapter. Tests may inject an adapter instance directly. Import-boundary tests -verify that loading the generic runner, bridge, or handler does not load MiniCPM -modules. - -Behavior is still assembled through `DuplexSessionRunnerMixin`, -`NativeRuntimeBridgeMixin`, and `ChatFallbackProjectorMixin`, so explicit -adapter composition is only the model-policy boundary, not a completed removal -of host-method coupling. A later refactor may replace the remaining mixins with -explicit protocol and effect-runner dependencies while preserving one actor -mailbox and one `DuplexSession` transition owner. Splitting those components -must not create a second reducer or concurrent state machine. - -The four private Session ledgers are data partitions, not independent owners. -`DuplexSession` intentionally remains the aggregate root. Further encapsulation -may narrow compatibility properties, but serving code must continue to perform -state changes through Session transitions. - -### Plugin descriptor and typed payloads - -The pipeline currently selects the MiniCPM engine extension and Serving adapter -through two explicit fields, but there is not yet one versioned plugin descriptor -that binds those identities in an open handshake. Realtime client payloads do -not select either implementation. Adding a second native model should first -introduce that descriptor and reject serving/engine plugin mismatches explicitly. - -The stable boundary has typed fences, plans, policies, decisions, and control -messages, but `session_config`, `runtime_config`, and append payloads still use -generic mappings or objects at the extension boundary. Model-neutral typed -session and input-chunk contracts remain follow-up work, especially before -adding video. - -### Output modality policy - -The current Realtime duplex contract supports text and audio output. Unknown -output modalities fail closed with `unsupported_response_modality`; they are -never silently projected as text. This checkpoint deliberately does not add a -modality registry because there is only one native text/audio projection path -to register. - -A registry becomes justified when a second non-audio/non-text implementation -uses the same validation, encoding, lifecycle, and projection contract. At that -point the abstraction must be extracted from two working adapters and retain -the same fail-closed behavior. A speculative registry before that condition -would add selection state without removing any existing branch. - -### Performance measurement contract - -Stage0 token timing comes only from the engine's per-segment metric snapshot: -`vllm_ttft_ms`, `vllm_tpot_ms`, and `vllm_itls_ms`. A resumable Stage0 request -exports its segment-local `num_generation_tokens`; ordinary requests retain the -existing output-token fallback. Realtime response state clears its metric -accumulator at `response.created`, combines only model units owned by that -response, and ignores session-level events with no matching `response_id`. -Realtime projection carries the result under -`metadata.vllm_omni.stage_metrics`; clients do not estimate TPOT from transcript -deltas or subtract cumulative snapshots. - -Audio timing uses the client's monotonic receive clock and is partitioned by -`response_id`. `response_created_to_first_audio_ms` measures projection startup -after `response.created`; `commit_to_first_audio_ms` measures end-to-end output -from an explicit input commit. Neither field is called TTFP without naming its -origin. Reports also include inter-chunk and chunk-duration -count/mean/p50/p95/max plus maximum chunk gap. Engine and client measurements -have different clock origins and must not be merged into one timeline. - -### Stage1 CUDA Graph boundary - -Native Stage1 streaming advances Python-owned request generators and turn -state on every scheduler step. A full-model CUDA Graph would execute that -Python control flow only while capturing and then replay stale chunk and turn -metadata, even if its host-to-device tensor construction were made -capture-safe. The duplex deploy profile therefore keeps `enforce_eager=false` -but selects `cudagraph_mode=PIECEWISE` for Stage1. Python orchestration runs on -every step while fixed-topology compiled regions remain eligible for CUDA -Graph capture. This is a model execution constraint, not an E2E-only override. - -### Resource capabilities - -At the time of this snapshot the checked-in MiniCPM duplex deploy profile set -Stage0 and Stage1 `max_num_seqs` to two and configured the engine runtime -manager with `max_sessions=2`; the duplex settings have since been folded into -`vllm_omni/deploy/minicpmo_4_5.yaml`, which admits four sessions -(`duplex_session.max_sessions: 4`, matching its stage `max_num_seqs`). The -Realtime capability response advertises -multi-session and same-replica multi-session support with -`session_admission_mode="engine_managed"`. Client session fields cannot raise -this server-owned limit. Fresh E2E evidence covers two concurrent sessions in -both model-policy and response-required modes, rejection of a third session, -and admission of a replacement after one accepted session closes. It does not -establish behavior beyond that configured limit. - -Session leases, lease TTL, resumable attachment, orphan reaping, and bounded -Serving input/backpressure are implemented. A production admission controller, -KV-aware capacity budgeting, fairness, starvation metrics, and multi-session -worker-failure recovery remain follow-up work. Capability claims must stay -limited to the mechanisms and two-session execution shape validated here. - -## Validation Evidence - -All pytest and runtime evidence for this branch must run on the remote H20. - -### Current synchronized tree - -The current dirty tree was synchronized to the isolated H20 worktree at -`/home/admin/workspace/aop_lab/model_runner_v2/vllm-omni-worktrees/pr3907-boundary-cg-0721`. - -- Final affected matrix: 661 passed, one network-dependent config test - deselected, 20 warnings. It includes the complete MiniCPM native-duplex hook - file and all control/session/handler/Orchestrator/import-boundary suites. The - run is task `be221f67` (`/tmp/remote_gpu_logs/be221f67.log`). The deselected - `test_resolve_when_autodetect_resolves_none` attempted a Hugging Face model - lookup and hung in the isolated environment; the duplex deploy-config test - passed independently. -- Stage1 PIECEWISE CUDA Graph validation: the deploy-config assertion first - failed with no Stage1 compilation config, then passed in task `3ed6091d`. - Server task `95827027` started with `enforce_eager=False`, logged - `cudagraph_mode=PIECEWISE`, and captured batch sizes 1, 2, and 4 without the - CPU-to-CUDA capture error. The first concurrent E2E exposed graph padding as - `request_token_spans` covering three valid rows in a four-row capture. A - focused RED/GREEN now verifies both single-request and batched-request - padding, and the full MiniCPM hook file passed 47 tests in task `e89dba88`. - After the fix, two concurrent response-required sessions each completed two - spoken responses, playback acknowledgements, and history lifecycles with no - error or stale event. Artifacts are under - `/tmp/pr3907_piecewise_response_required_2x2_fixed_20260721`. -- An exploratory two-session, two-turn model-policy run still left one session - waiting after its second input commit while the other session completed both - turns. The same liveness symptom was reproduced with Stage1 eager mode before - the padding fix, and no Stage1 fatal occurred after the fix. This run is not - counted as PIECEWISE acceptance evidence; the remaining Stage0 resumable - reactivation path requires separate diagnosis. Artifacts are under - `/tmp/pr3907_piecewise_model_policy_2x2_fixed_20260721`. -- Latest broad affected matrix before the final pre-response continuation fix: - 608 passed, 19 warnings. This includes the - Stage0 row-routing regression, cleanup-held admission regressions, - segment-local token accounting, response-owned metric filtering, the generic - output-processor suite, direct E2E script execution, and continuous-response - overlap ownership. The run is task `e2146318` - (`/tmp/remote_gpu_logs/e2146318.log`). -- The pre-response continuation regression group passed 5 focused tests. The - response-required two-turn scenario then passed 5/5 repeated runs in task - two-session semantic isolation plus resume, takeover, and - admission passed in task `49225794`; model-unit overlap passed in task - `3888387d`. The final affected matrix above supersedes the earlier pending - rerun requirement. -- Serving plugin selection was first reproduced as three RED tests: the generic - handler silently selected MiniCPM and the pipeline had no explicit Serving - adapter field. The final pipeline/import/fail-fast group passed 9 tests in task - `4db6d959` (`/tmp/remote_gpu_logs/4db6d959.log`). The token-metric and client - audio-cadence producer-to-consumer group passed 8 tests in task `cf80f6b2` - (`/tmp/remote_gpu_logs/cf80f6b2.log`). -- Stable-import isolation was first reproduced as a RED that loaded twelve - `experimental.fullduplex` modules. After extracting the kernel and making - implementation loading conditional, the isolated import regression passed. - The focused ControlPlane/ControlClient/session/runtime/Orchestrator/AsyncOmni - suite then passed 139 tests in task `fad613ea` - (`/tmp/remote_gpu_logs/fad613ea.log`). - After the final contract-import cleanup, the import-boundary, runtime, and - lease suites passed 45 tests in task `bc138887` - (`/tmp/remote_gpu_logs/bc138887.log`). -- Concurrency and failure-path hardening now has direct regressions for - single-flight cancel/close cleanup, cleanup retry without dropping request - resources, failed-replica affinity cleanup and same-address reattachment, - append ordering across resume, takeover ownership of pending-turn - reservations, and a closed-transport send immediately after resumable - credentials are delivered. The final focused resume/takeover group passed 5 - tests in task `dc2ac96e` (`/tmp/remote_gpu_logs/dc2ac96e.log`). -- Stage1 batching regression: the focused test first failed because only the - first request in `runtime_additional_information` reached the talker. After - routing every request by `request_token_spans`, preserving per-row terminal - flags, and emitting request-keyed sparse audio, the complete MiniCPM duplex - hook file passed 42 tests. -- Stage0 batching regression: two requests in one batch now retain different - `duplex_prompt_token_ids` and `special_token_ids`. The model emits row-aligned - payload elements, so Stage0 no longer copies request zero's duplex metadata to - later sessions. -- Admission finalization regressions cover `max_sessions=1` with blocked and - failed expiry cleanup plus request-triggered cleanup. Every replacement is - rejected with `resource_exhausted` until stage resources are cleaned and the - old session is explicitly finalized. -- Resumable token accounting exports each model unit's native generation-token - count instead of recounting the request's cumulative token IDs. Realtime - response state resets and accumulates those segment snapshots, and the client - ignores unowned session-level metrics while summarizing a response. -- Engine-owned admission and two-session model-policy E2E: two random session - IDs were admitted and remained isolated, a third was rejected with - `resource_exhausted`, closing one accepted session returned capacity, and a - random replacement session was admitted. The capability projection reported - `configured_limit=2`, `advertised_multi_session=true`, and - `session_admission_mode=engine_managed`. In the same run one session selected - speak and the other selected listen without identity leakage. Detach/resume - advanced attachment generation from one to two and rotated its token; active - takeover rejected all four writes from the replaced connection. The run is - task `1d9dd1ed` (`/tmp/remote_gpu_logs/1d9dd1ed.log`), with artifacts under - `/tmp/pr3907_resume_takeover_native_metrics`. -- Concurrent two-session, two-turn response-required E2E: both sessions produced - two complete spoken responses, committed two playback acknowledgements and - history records, and emitted no error, cancel, truncate, or stale events. The - semantically distinct inputs produced isolated transcripts containing - "当时" and "一加一" respectively. Per-response Stage0 output counts were 20, - 38, 32, and 33 tokens rather than cumulative request totals. The run is task - `a16f1467` (`/tmp/remote_gpu_logs/a16f1467.log`), with artifacts under - `/tmp/pr3907_semantic_2turn_native_metrics`. -- Disconnect-grace E2E: a detached session expired after the configured - 30-second engine grace, and a later resume was rejected with - `session_resume_expired`. The run is task `9872d50c`, with artifacts under - `/tmp/pr3907_0719_disconnect_grace_expiry`. -- A separate black-box takeover client exercised the current backend with two - semantically distinct concurrent inputs. Active resume replayed 17 events - including five audio deltas and committed playback history; takeover rejected - all four stale-attachment sends, emitted `session.replaced`, and preserved - attachment/session isolation. Both transcripts contained only their expected - semantic token. The client task is `6773bd34` - (`/tmp/remote_gpu_logs/6773bd34.log`), with artifacts under - `/tmp/pr3907_gaohan_final_takeover`. The backend task was `5abcd1ee`. -- Across the final response-required, overlap, and resume/takeover runs, engine - Stage0 TTFT ranged from 31.17 to 310.18 ms, TPOT from 15.18 to 21.18 ms, and - reported ITL p95 from 16.35 to 32.83 ms. Explicit - `commit_to_first_audio_ms` ranged from 386.80 to 1033.25 ms, while the - client-monotonic maximum audio chunk gap ranged from 388.30 to 1316.37 ms. - These are response-scoped observations from a small, warm-state-sensitive - sample, not an SLO claim. The slowest response still does not satisfy a 200 ms - end-to-end first-output target. -- Post-merge focused handler/demo/Web/Stage0/runner/video suite: 324 passed, - 19 warnings. The run is task `f78cb7a8` - (`/tmp/remote_gpu_logs/f78cb7a8.log`). The Web contract explicitly rejects - browser VAD markers and browser-generated input commits while the merged - video path retains camera-frame admission. -- No-response listen and residual accounting: the focused handler regressions - passed after first reproducing both failures. A model-listen unit without a - Realtime response no longer defers the next committed input, and three - consecutive 1.2-second `pcm_f32le` inputs each returned pending input bytes - to zero. -- Model-unit overlap E2E: the second physical input streamed while the first - response remained active. Seven overlap decisions had - `defer_runtime_append=false`; the one-second model-unit boundary occurred - before the first `response.done`, and no barge-in, cancel, truncate, stale, or - error event occurred. The model-policy contract accepts a later unit ending - the continuous response instead of requiring a second physical-input-owned - response. The two response-owned Stage0 snapshots independently reported 20 - tokens / 152.8 ms TTFT and 19 tokens / 43.3 ms TTFT. The run is task - `a4977622` (`/tmp/remote_gpu_logs/a4977622.log`), with artifacts under - `/tmp/pr3907_overlap_native_metrics_scoped`. -- The split multi-session E2E runner is directly executable from the repository - even when another installed package owns the top-level `tests` name. Its - script-import and continuous-terminal regressions passed in task `f11b106c`. -- Pinned response-required E2E: the first 1400 ms of the pinned fixture produced - one complete spoken response, five audio/transcript deltas, and one playback - history commit, with symmetric created/audio-done/done lifecycle and no - duplicate `response.speak`. The client task is `a76dec52` - (`/tmp/remote_gpu_logs/a76dec52.log`). -- Continuous-input model-policy E2E: 36 browser-like 200 ms chunks produced - seven model-listen decisions with zero input commits and no - error/cancel/truncate/stale events. The client task is `fb256595`; artifacts - are in `/tmp/minicpmo_pr3907_continuous_model_policy_20260718`. -- Three-response continuous-input E2E: three consecutive 9469 ms streams, - composed of the original fixture plus four seconds of silence, produced three - spoken responses in one session. All three completed `response.done`, audio - transcript, playback ACK, and history commit; the run had 22 audio/transcript - deltas, zero input commits, and no error/cancel/truncate/stale events. This is - lifecycle evidence, not semantic cross-input independence, because the same - fixture was repeated. The client task is `2a5f29f0`; artifacts are in - `/tmp/minicpmo_pr3907_continuous_three_response_20260718`. -- The fresh validation backend was stopped after the E2E runs and GPU0/1 returned - to baseline. Audio-quality and ASR evaluation were not rerun for these - lifecycle-only changes. Other GPU workloads were not modified. - -### Regression lineage - -Earlier synchronized checkpoints passed broader affected matrices before the -extension, policy, and ledger consolidation. Those runs were useful while -developing the refactor but are superseded by the current file-identical H20 -matrix and E2E evidence above; they are intentionally omitted to avoid mixing -different trees in the acceptance argument. - -### Focused suites - -- RPC router and AsyncOmni shutdown/error routing; -- orchestrator fatal and duplex control paths; -- WebSocket Actor mailbox and writer; -- duplex protocol and Session ownership; -- Realtime handler, terminal ordering, overlap, playback ACK; -- scheduler resumable boundary; -- MiniCPM runner and Stage0/Stage1 lifecycle. - -### E2E contract - -Run two deliberately separate fixture contracts: - -- model-policy uses three distinct inputs, no transcript hints, 200 ms - real-time pacing, and continuous input with no browser commit. It allows - natural listen or speak decisions and assumes no 1:1 mapping between - physical inputs, model turns, and responses; -- response-required uses `minicpmo_pr3907_jiayan_16k.wav` pinned to SHA-256 - `2e5fd4eb3ee434ce107ee3a0591fa624a33f7683c7462f45fe651c443c9af941`, - sends its first 1400 ms, and requires one complete spoken response. - It is an output-lifecycle fixture, not evidence that arbitrary no-hint, - real-time input must select speak; -- late and old-response playback ACK handling in the affected test matrix; -- final model-policy response drain before `session.close`; -- close/disconnect with no post-terminal audio; -- transcript delta/final consistency; -- generated WAV validity and external ASR sanity. - -Distinct-input fixtures must also be semantically distinct. Different files or -waveform hashes containing the same greeting can legitimately produce the same -reply and are not sufficient evidence for cross-turn transcript independence. - -For each run record the exact SHA/worktree, command, server log, client log, -WAV path, and cleanup result. - -## Acceptance Criteria - -This checkpoint is ready to publish only when: - -- one mailbox passes exactly-once and wire-order tests; -- cancel advances the engine fence even when the old append continues in an - executor, and the next-epoch append remains accepted; -- request completion removes MiniCPM helper and model-owned force-listen state; -- RPC fatal/close wakes every waiter and no new work enters after terminal; -- Session is the only serving response/overlap authority; -- stale terminal events cannot mutate a newer epoch; -- an accepted terminal cannot be dropped by the writer; -- no duplicate `response.speak` or transcript delta appears; -- Stage0 KV continuity and Stage1 turn reset pass; -- the real H20 multi-turn audio E2E passes; -- ordinary non-duplex paths in the affected matrix still pass; -- no local path, proxy token, temporary profiling probe, or test-only switch is present. - -Passing this checkpoint supports the statement: - -> Single-session, model-owned MiniCPM-o 4.5 native duplex is reviewable on the -> validated H20 configuration. - -It does not support claims for an acoustic-onset barge-in latency target across -arbitrary client chunk sizes, multi-session production concurrency, bounded -long-session KV, or scheduler-native append. Camera-frame input -(`video_frames`) landed after this snapshot and is covered by its own E2E -test rather than by the evidence above. +| engine policy | `configure_sampling_params(runtime_config, defaults)`, `plan_append(...) -> DuplexAppendPlan` (the resumable Stage0 prompt for one unit), `decide_output(...) -> DuplexOutputDecision \| None` (e.g. the listen decision on a finished Stage0 segment) | +| session policy | `capabilities(max_sessions)`, `validate_client_extra_body`, `prepare_runtime_config(config, model_config)` (server-owned runtime keys, reference audio resolution), `runtime_config_for_update`, `runtime_config_for_function_output`, `create_session_state() -> DuplexModelSessionState`, `data_plane: DuplexDataPlane` (projects raw stage outputs into internal events), `data_plane_context(...)` | + +`DuplexOmniEngine._validate_deployment` loads the plugin before any stage +starts; `DuplexSessionManager.__init__` validates it against the stage +sampling defaults once the stage pools exist. Plugin hooks that may block +(`prepare_runtime_config` fetching `ref_audio`) are awaited in `open()` and +offloaded from the loop. + +The MiniCPM-o 4.5 plugin (`model_executor/models/minicpmo_4_5/duplex/plugin.py`) +is the only integration in this framework version. PersonaPlex and Nemotron +VoiceChat still carry their pre-framework duplex code (runtime extension plus +serving adapter) and are therefore not served over this framework yet: their +pipelines declare no `duplex_plugin`, so they run turn-based until the +follow-up PRs port them (RFC vllm-omni#7181, PR 2/3). + +## Serving + +`OmniDuplexSessionHandler.handle_realtime_session(websocket)`: + +1. accept; build the `RealtimeEnvelope` (autostart / default session from + query params, `event_id` correlation, error rendering through the typed + `ErrorEvent`); +2. handshake: first message `session.update` -> `omni.open_session(session)` + (the id inside the payload is ignored; the allocated id is announced in + `session.created` together with `resume_token` / `attachment_generation` + when the model supports resume); `session.resume(session_id, resume_token, + last_received_server_event_seq)` -> `attachment.authenticate_resume` -> + `omni.resume_session`; +3. reader loop: JSON -> `RealtimeEnvelope.translate` (`translate_realtime_command` + with the session's declared audio defaults) -> `handle.submit`; + envelope-level errors (invalid JSON, oversize frame, unknown type, + event acks) are answered locally; +4. writer pump (session-scoped, survives reconnects): `async for ev in + handle.events(): attachment.send_event(ev.to_realtime())`, journaling + for replay until the journal overflows (`session.resync_required`); +5. disconnect: a resumable session is detached (`attachment.detach` + + `omni.detach_session`, engine-owned grace); a superseded socket's + disconnect is a no-op; a non-resumable session is closed. Takeover sends + `session.replaced` to the old socket and closes it. + +## Validation + +- CPU suites: `tests/engine/duplex/**` (session, lease, manager, runner, + plugin, events, commands), `tests/engine/test_duplex_orchestrator.py`, + `tests/engine/test_duplex_omni_engine.py`, + `tests/entrypoints/test_duplex_omni.py`, + `tests/entrypoints/duplex/test_duplex_serving.py`, + `tests/entrypoints/openai/test_duplex_session_attachment.py`, + `tests/entrypoints/openai_api/test_duplex_api_server.py` (the duplex + server: pipeline probe, app state, routes, warmup gate), + `tests/entrypoints/duplex/test_chat_completions_adapter.py`, + `tests/clients/**`, `tests/engine/test_duplex_import_boundary.py`, + `tests/model_executor/models/minicpmo_4_5/duplex/**`, + `tests/worker/test_native_duplex_hooks.py`. +- GPU: `examples/online_serving/minicpmo/realtime_duplex_demo.py` (one turn + over the websocket), `examples/online_serving/barge_in_client.py` (cancel + during playback; `--inline` drives the same scenario through + `InlineDuplexClient`), `tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py` + (admission, resume, takeover, expiry). diff --git a/docs/design/module/engine_orchestration.md b/docs/design/module/engine_orchestration.md index d4dbf51a2..df22f3ce1 100644 --- a/docs/design/module/engine_orchestration.md +++ b/docs/design/module/engine_orchestration.md @@ -16,6 +16,7 @@ required_reviewers: - "@yenuo26" - "@NickCao" primary_code_paths: + - vllm_omni/engine/omni_engine_base.py - vllm_omni/engine/async_omni_engine.py - vllm_omni/engine/async_engine_utils.py - vllm_omni/engine/orchestrator.py @@ -28,8 +29,12 @@ related_code_paths: - vllm_omni/engine/stage_pool.py - vllm_omni/outputs/output_processor.py - vllm_omni/distributed/omni_coordinator/** - # Duplex kernel design is owned by ../fullduplex.md; listed here because the - # orchestrator hosts the duplex control plane. + # The duplex siblings (DuplexOmniEngine / DuplexOrchestrator) and the + # engine-resident session framework are owned by ../fullduplex.md; listed + # here because they fill the template seams of OmniEngineBase and + # OrchestratorBase. + - vllm_omni/engine/duplex_omni_engine.py + - vllm_omni/engine/duplex_orchestrator.py - vllm_omni/engine/duplex/** depends_on: - input_output_modality_contracts.md @@ -48,6 +53,8 @@ validation_paths: - tests/engine/test_cfg_companion_lifecycle.py - tests/engine/test_rpc_result_router.py - tests/engine/duplex/** + - tests/engine/test_duplex_orchestrator.py + - tests/engine/test_duplex_omni_engine.py - tests/engine/test_duplex_import_boundary.py upstream_refs: - vllm.v1.engine.EngineCoreRequest diff --git a/docs/design/module/entrypoints.md b/docs/design/module/entrypoints.md index a0dee90b9..db62f9cb6 100644 --- a/docs/design/module/entrypoints.md +++ b/docs/design/module/entrypoints.md @@ -19,7 +19,9 @@ required_reviewers: primary_code_paths: - vllm_omni/entrypoints/omni.py - vllm_omni/entrypoints/async_omni.py + - vllm_omni/entrypoints/async_omni_base.py - vllm_omni/entrypoints/omni_base.py + - vllm_omni/entrypoints/duplex_omni.py - vllm_omni/entrypoints/cli/** - vllm_omni/entrypoints/openai/** - vllm_omni/entrypoints/openpi/** @@ -27,14 +29,13 @@ primary_code_paths: - vllm_omni/entrypoints/stage_utils.py - vllm_omni/entrypoints/utils.py - vllm_omni/entrypoints/duplex/** - - vllm_omni/entrypoints/duplex_request_client.py - vllm_omni/clients/** primary_path_exceptions: - path: vllm_omni/entrypoints/openai/errors.py owner: error_contracts.md - path: vllm_omni/entrypoints/duplex/** owner: ../fullduplex.md - - path: vllm_omni/entrypoints/duplex_request_client.py + - path: vllm_omni/entrypoints/duplex_omni.py owner: ../fullduplex.md - path: vllm_omni/clients/** owner: ../fullduplex.md @@ -42,7 +43,9 @@ related_code_paths: - vllm_omni/errors.py - vllm_omni/inputs/** - vllm_omni/outputs/** + - vllm_omni/engine/omni_engine_base.py - vllm_omni/engine/async_omni_engine.py + - vllm_omni/engine/duplex_omni_engine.py - vllm_omni/config/** - vllm_omni/deploy/** depends_on: @@ -53,7 +56,9 @@ validation_paths: - tests/entrypoints/test_omni_entrypoints.py - tests/entrypoints/test_async_omni.py - tests/entrypoints/test_async_omni_pause_sleep_routing.py - - tests/entrypoints/test_async_omni_duplex.py + - tests/entrypoints/test_duplex_omni.py + - tests/entrypoints/duplex/** + - tests/entrypoints/openai_api/test_duplex_api_server.py - tests/entrypoints/test_serve.py - tests/entrypoints/test_stream_finish_reason.py - tests/entrypoints/openai/** diff --git a/docs/mkdocs/hooks/generate_api_readme.py b/docs/mkdocs/hooks/generate_api_readme.py index 3130e688f..d236c1129 100644 --- a/docs/mkdocs/hooks/generate_api_readme.py +++ b/docs/mkdocs/hooks/generate_api_readme.py @@ -25,8 +25,8 @@ CATEGORIES = { "clients": { "name": "Clients", "description": ( - "Client libraries that talk to a running vLLM-Omni server " - "(the full-duplex Realtime client and its per-model session presets)." + "Client libraries for full-duplex sessions: the Realtime WebSocket client, " + "the in-process client over `DuplexOmni`, and the per-model session presets." ), }, "inputs": { diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index f8db9eb92..42971ed97 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -107,7 +107,7 @@ th { | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — | | `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — | | `DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | ✅︎ | — | -| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — | +| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/OpenBMB/MiniCPM-o-4_5.md) | | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | ✅︎ | — | | `HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — | diff --git a/docs/serving/README.md b/docs/serving/README.md index 8167f4c9b..472ab2a49 100644 --- a/docs/serving/README.md +++ b/docs/serving/README.md @@ -145,15 +145,22 @@ interchangeably. | Incremental text input for speech synthesis | `WS /v1/audio/speech/stream` | Send text events and receive audio | [Streaming Text to Speech](speech_api.md#streaming-text-input-websocket) | | Live video understanding | `WS /v1/video/chat/stream` | Send video frames and receive text/audio | [Streaming Video Input](video_stream_api.md) | | Turn-based realtime audio | `WS /v1/realtime` | Stream one audio input and receive transcript/audio events | [Realtime Audio](realtime_api.md) | -| Continuous speech-to-speech interaction | `WS /v1/realtime?duplex=1` or `WS /v1/duplex` | Listen and speak concurrently with session control | [Full Duplex](full_duplex_api.md) | +| Continuous speech-to-speech interaction | `WS /v1/realtime?duplex=1` (alias `WS /v1/duplex`) | Listen and speak concurrently with session control | [Full Duplex](full_duplex_api.md) | | Generated video chunks | `WS /v1/realtime/video` | Start a diffusion request and receive fragmented MP4 | [Streaming Video Output](streaming_video_output_api.md) | | Robot policy inference | `WS /v1/realtime/robot/openpi` | Send MessagePack observations and receive action arrays | [OpenPI Robot Policy](openpi_api.md) | All six routes are model- or configuration-dependent. In particular, -`/v1/realtime` is not full duplex unless the client sets `duplex=1` and the -deployment explicitly enables duplex sessions. Clients should also verify the -duplex capability payload because the query-parameter form falls back to the -ordinary realtime handler when duplex is unavailable. +`/v1/realtime` is full duplex when the model is a duplex model (its pipeline +declares a `duplex_plugin` and the deploy configuration sets +`session_mode: duplex`); a stock Realtime client needs no vendor query +parameter, and `duplex=0` is the explicit opt-out. Such a server serves the +websocket route plus `POST /v1/chat/completions`, and no other turn-based +HTTP route. On a duplex server that chat route is not the turn-based path: +each request runs on a short-lived duplex session, so it holds one of +`duplex_session.max_sessions` for its lifetime and answers at the model's +real-time pace -- see [Full Duplex](full_duplex_api.md). Clients should +also verify the duplex capability payload because the query-parameter form +falls back to the ordinary realtime handler when duplex is unavailable. ## Related Endpoints diff --git a/docs/serving/full_duplex_api.md b/docs/serving/full_duplex_api.md index 4e8bf9748..92511e84f 100644 --- a/docs/serving/full_duplex_api.md +++ b/docs/serving/full_duplex_api.md @@ -15,40 +15,52 @@ capability gates, see the [Realtime Duplex API guide](realtime_duplex_api.md). | Endpoint | Protocol | Recommended use | | --- | --- | --- | -| `WS /v1/realtime?duplex=1` | OpenAI Realtime-style event projection | Applications and browser clients | -| `WS /v1/duplex` | Native vLLM-Omni duplex events | Runtime integration and low-level testing | - -Both endpoints use the same duplex engine and require the same model-side -adapter. Prefer the Realtime projection unless the native lifecycle events are -specifically required. +| `WS /v1/realtime?duplex=1` | OpenAI Realtime-style events (the normative contract) | Applications and browser clients | +| `WS /v1/duplex` | Alias of `/v1/realtime?duplex=1` (same protocol, same handler) | Clients that prefer a dedicated path | +| Python: `DuplexOmni` / `InlineDuplexClient` | Typed commands and events in-process | Embedding the model without a server | ## Enable Full Duplex -The route becomes usable only when the deployment configuration explicitly -sets: +A model is served full duplex when its registered pipeline declares a +`duplex_plugin` (the model's `DuplexModelPlugin`). That declaration alone +decides it: `vllm-omni serve` constructs `DuplexOmni` instead of `AsyncOmni`, +and every surface the server exposes is backed by a duplex session. It serves +`/v1/realtime?duplex=1` (and its alias `/v1/duplex`), +`POST /v1/chat/completions`, `/v1/models` and `/health`; every other +turn-based HTTP route (speech, batch, embeddings, video, ...) answers "not +available". Turn-based use of such a model stays available offline through the +Python API (`Omni` / `AsyncOmni`). + +`/v1/chat/completions` is the ordinary chat service running on the duplex +engine: a request is a turn-based generation on the same stages, served +alongside the live websocket sessions, with the request options the chat +service supports. It is wired only when the model's plugin declares +`DuplexCapabilities.supports_chat_completions`; on any other duplex model the +route reports "not available". A model that should not serve the route at all +lists it in the deploy configuration's `endpoint_restrictions`. + +The deploy configuration of such a model must agree: ```yaml session_mode: duplex ``` -The selected model pipeline must also provide a duplex serving adapter. -Model-native deployments configure the corresponding engine runtime extension -and control plane as part of their registered pipeline. +A duplex model started with a deploy configuration that does not set it fails +at startup rather than falling back to turn-based serving. !!! warning - `WS /v1/duplex` fails with `Duplex API is not available` when duplex is not - enabled. By contrast, `/v1/realtime?duplex=1` falls back to the ordinary - turn-based realtime handler when no duplex handler exists. Confirm that + On a deployment that is *not* duplex, `WS /v1/duplex` fails with + `Duplex API is not available` and `/v1/realtime?duplex=1` falls back to the + ordinary turn-based realtime handler. Confirm that `session.created.session.capabilities` is present before treating the connection as full duplex. -The current unified-runtime integrations are: - -- MiniCPM-o 4.5, using `vllm_omni/deploy/minicpmo_4_5.yaml` - (`session_mode: duplex`); -- PersonaPlex, whose default `vllm_omni/deploy/personaplex.yaml` enables duplex; -- Nemotron VoiceChat, via its registered duplex plugin package. +**MiniCPM-o 4.5** (`vllm_omni/deploy/minicpmo_4_5.yaml`) is the only model +served over this endpoint today. PersonaPlex and Nemotron VoiceChat still carry +their pre-framework duplex code: their pipelines declare no `duplex_plugin`, so +they run turn-based until the follow-up PRs port them to the plugin contract +(RFC [vllm-omni#7181](https://github.com/vllm-project/vllm-omni/issues/7181)). JoyVL is a separate HTTP interaction orchestrator and does not use these WebSocket endpoints. See [Standalone Experimental Servers](standalone_servers.md). @@ -101,22 +113,23 @@ deferred, treated as a short acknowledgement, or used to interrupt output. The `session.created` payload includes capability fields such as `supports_barge_in`, `supports_playback_ack`, `supports_multi_session`, `supports_session_resume`, and `chunk_period_ms`. Treat this payload as the -runtime contract. +runtime contract and branch on the flags, never on the model name: a model +that supports native overlapping speech may still advertise +`supports_barge_in=false` when destructive output interruption and model-state +rewind are not validated for it. Capacity and session-resume behavior also +depend on the selected deployment configuration. The per-model table lives in +the [Realtime Duplex API guide](realtime_duplex_api.md#capability-negotiation-by-model). -For example, PersonaPlex supports native overlapping speech but currently -advertises `supports_barge_in=false`; destructive output interruption and -model-state rewind have not been validated for that integration. Capacity and -session-resume behavior also depend on the selected deployment configuration. +## Python API -## Native Protocol +`vllm_omni.entrypoints.duplex_omni.DuplexOmni` runs the same engine-resident +sessions in-process: `open_session()` returns a `DuplexSessionHandle` whose +`submit()` takes typed `DuplexCommand` objects and whose `events()` yields +typed `DuplexEvent` objects (each with a `to_realtime()` wire rendering). +`vllm_omni.clients.inline_duplex.InlineDuplexClient` exposes that handle +behind the `DuplexClient` API. See the +[Realtime Duplex API guide](realtime_duplex_api.md#using-the-python-api). -`WS /v1/duplex` exposes lower-level lifecycle names including -`session.create`, `input_audio_buffer.append`, `turn.signal`, `playback.ack`, -and `session.close`. It returns native session, input, response, overlap, and -error events. This protocol is experimental and may evolve with the runtime; -applications should use the provided Realtime client where possible. - -See the [MiniCPM-o example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/minicpmo), -[PersonaPlex example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/personaplex), +See the [MiniCPM-o example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/minicpmo) and the [full-duplex runtime design](../design/fullduplex.md) for model-specific validation and architecture details. diff --git a/docs/serving/realtime_duplex_api.md b/docs/serving/realtime_duplex_api.md index 288c934e8..cc10d99c9 100644 --- a/docs/serving/realtime_duplex_api.md +++ b/docs/serving/realtime_duplex_api.md @@ -1,17 +1,18 @@ # Realtime Duplex API vLLM-Omni serves full-duplex speech models over a WebSocket endpoint, -`/v1/realtime`, that speaks the OpenAI Realtime vocabulary plus a +`/v1/realtime?duplex=1`, that speaks the OpenAI Realtime vocabulary plus a set of duplex extensions: the model decides when to listen and when to speak, the user can talk over the assistant, playback progress is reported back so history stays honest, and a dropped connection can resume the same session. This page covers how to run a duplex deployment, how to drive it from Python with `vllm_omni.clients.duplex.DuplexClient`, and the complete wire contract. -The handler is enabled by either a model-native `session_mode: duplex` -deployment or an explicit `duplex_session` configuration for turn-based -Server VAD. The runtime architecture is described in -[Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md). +The endpoint is served for models that ship a duplex plugin — currently +MiniCPM-o 4.5 only — and just for deploy configurations that declare +`session_mode: duplex`. PersonaPlex and Nemotron VoiceChat are ported to the +plugin contract in follow-up PRs and are not served over this endpoint yet. The runtime architecture +is described in [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md). ## Quick Start @@ -26,11 +27,38 @@ vllm-omni serve openbmb/MiniCPM-o-4_5 \ ``` `vllm_omni/deploy/minicpmo_4_5.yaml` declares `session_mode: duplex` and -`duplex_session.max_sessions: 4`, which mounts the WebSocket routes -`ws://:8099/v1/realtime` (this page) and `ws://:8099/v1/duplex` -(the native dialect for raw-protocol clients) next to the usual HTTP API. -PersonaPlex and Nemotron VoiceChat use `vllm_omni/deploy/personaplex.yaml` and -`vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml` in the same way. +`duplex_session.max_sessions: 4`. Because the MiniCPM-o 4.5 pipeline declares a +`duplex_plugin`, `vllm-omni serve` runs it through `DuplexOmni`: the server +mounts `ws://:8099/v1/realtime?duplex=1` (this page; `ws://:8099/v1/duplex` +is an alias of the same route), `POST /v1/chat/completions`, `/v1/models` and +`/health`, and nothing else. The remaining turn-based HTTP routes (speech, +batch, embeddings, video, ...) are not served by a duplex model. + +`/v1/chat/completions` is not a turn-based path on a duplex server: each +request opens a short-lived duplex session and reads the answer off it. Audio +content becomes an ordinary committed turn; text is seeded as the session's +opening turn, because a model-native model takes a turn only when it hears +speech. A model that cannot be seeded that way declares so, and a text request +to it is refused with 400. Two consequences are worth sizing for: + +- **Every request holds an admission slot** for its lifetime, so + `duplex_session.max_sessions` caps HTTP concurrency as well as websocket + sessions. Beyond it a request is refused with HTTP 503. +- **Latency follows the model's own clock.** A model-native session generates + per audio unit, so even a short answer costs what that many units cost -- + noticeably more than the same model would take turn-based. + +`n > 1`, `logprobs`, `tools`, and image or video content parts have no +representation in one duplex turn and are refused with HTTP 400 -- a Realtime +conversation item carries text and audio only. `temperature`, +`max_completion_tokens`, `modalities` and +`chat_template_kwargs.use_tts_template` apply; the other +`chat_template_kwargs` keys are logged as ignored. `extra_body` is passed to +the session unchanged, exactly as a websocket client's session `extra_body` +is, which is how a request asking for `modalities: ["text", "audio"]` supplies +the `ref_audio` MiniCPM-o requires for audio output. A model that should not +serve the route lists it in the deploy config's `endpoint_restrictions`, +exactly as a turn-based model does. ### Run the Example Client @@ -81,25 +109,24 @@ ships a preset: | Preset | Input / output audio | What it sets | | --- | --- | --- | -| `vllm_omni.clients.minicpmo_4_5.create_duplex_session_config(ref_audio=...)` | 16 kHz `pcm16` / 24 kHz `pcm16` | `extra_body.native_duplex=True` (opt in to the model-native lane), `force_listen_count=0`, `overlap_policy="listen_only"`, `playback_commit_policy="ack_only"`; `ref_audio` is the assistant voice clip | +| `vllm_omni.clients.minicpmo_4_5.create_duplex_session_config(ref_audio=...)` | 16 kHz `pcm16` / 24 kHz `pcm16` | `force_listen_count=0`, `overlap_policy="listen_only"`, `playback_commit_policy="ack_only"`; `ref_audio` is the assistant voice clip | | `vllm_omni.clients.personaplex.create_duplex_session_config(voice="NATF2.pt", persona="")` | 24 kHz `pcm_f32le` / 24 kHz `pcm16` | bundled `.pt` voice prompt and the persona as `instructions` | | `SessionConfig(...)` | 16 kHz `pcm16` / 24 kHz `pcm16` | model-neutral defaults; pass `extra_body`, `turn_detection`, `overlap_policy`, `playback_commit_policy`, `instructions`, `voice`, `temperature` yourself | Keyword overrides on a preset replace the corresponding `SessionConfig` field, e.g. `create_duplex_session_config(ref_audio=..., temperature=0.6)`. -Constructor options: `session_id` (names the session the client creates; it -does not resume or take over an existing session — entering the client always -performs the `session.update` handshake), `reconnect` (a `ReconnectPolicy`, -default five attempts with 0.25–4 s jittered backoff; `None` disables -auto-resume), `heartbeat_interval_s` (default 30 s; `None` disables the lease -heartbeat), `handshake_timeout_s`, and `connect` (a custom transport factory, -used by tests and benchmarks). Resume is supported only as automatic reconnect -within the same `DuplexClient` instance (see -[Reconnect and resume](#reconnect-and-resume)); re-attaching from a new client -or process requires the wire-level `session.resume` handshake -(`resume_token`, `incarnation`, `last_received_server_event_seq`), which this -client does not expose. +Constructor options: `reconnect` (a `ReconnectPolicy`, default five +attempts with 0.25–4 s jittered backoff; `None` disables auto-resume), +`heartbeat_interval_s` (default 30 s; `None` disables the lease heartbeat), +`handshake_timeout_s`, and `connect` (a custom transport factory, used by +tests and benchmarks). The session id is allocated by the server and becomes +`client.session_id` once `session.created` arrives; a client cannot choose +it. Resume is supported only as automatic reconnect within the same +`DuplexClient` instance (see [Reconnect and resume](#reconnect-and-resume)); +re-attaching from a new client or process requires the wire-level +`session.resume` handshake (`resume_token`, +`last_received_server_event_seq`), which this client does not expose. ### Stream audio and commit a turn @@ -195,7 +222,7 @@ the stamped `event_id`. On a transport drop the client reconnects with the `ReconnectPolicy` backoff and sends `session.resume` from the last acknowledged `server_event_seq`. Replayed events are deduplicated, consumers see one `ConnectionResumed` -event, and `client.incarnation` / `client.resume_token` are refreshed. The +event, and `client.resume_token` is refreshed. The server keeps a detached session alive for `disconnect_grace_s` (30 s by default); after that a resume fails with `session_resume_expired`. Resume is gated by `capabilities.supports_session_resume` (MiniCPM-o 4.5 and Nemotron @@ -238,9 +265,67 @@ definitions out of the dependency-free client. `audio_bytes(response_id)`, `coun `openai-realtime-duplex` benchmark backend (`vllm-omni bench serve`) is built on the same helpers. +## Using the Python API + +`vllm_omni.entrypoints.duplex_omni.DuplexOmni` is the in-process Python API +of a duplex model: sessions run inside the engine, and the class only opens, +resumes and closes them and pipes typed commands in and typed events out +through a `DuplexSessionHandle`. + +```python +from vllm_omni.engine.duplex.commands import AppendAudio, Commit +from vllm_omni.engine.duplex.events import AudioDelta, SessionClosed +from vllm_omni.entrypoints.duplex_omni import DuplexOmni + +omni = DuplexOmni(model="openbmb/MiniCPM-o-4_5", trust_remote_code=True) +async with await omni.open_session({"modalities": ["audio", "text"], "ref_audio": ref_audio_data_url}) as session: + async def consume(): + async for event in session.events(): # typed DuplexEvent objects + if isinstance(event, AudioDelta): + play(event.audio) + elif isinstance(event, ResponseDone): + break # one turn is enough here + task = asyncio.create_task(consume()) + await session.submit(AppendAudio(audio=pcm16_bytes, format="pcm16", sample_rate_hz=16000)) + await session.submit(Commit()) + await task +# leaving the block closes the session (the last event is SessionClosed) +``` + +`open_session` takes a `DuplexSessionConfig` or the same Realtime `session` +object the WebSocket handshake takes; the session id is always allocated by +`DuplexOmni`. `DuplexSessionHandle` offers `submit(command)` plus one +convenience wrapper per command (`append_audio`, `commit`, `create_response`, +`cancel_response`, `barge_in`, `ack_playback`, `update`, `create_item`, +`delete_item`, `truncate_item`, `heartbeat`, ...). `events()` is a +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. + +`vllm_omni.clients.inline_duplex.InlineDuplexClient` wraps a `DuplexOmni` +behind the `DuplexClient` API, so the same application code runs in-process +(no server) or over the network: + +```python +from vllm_omni.clients.inline_duplex import InlineDuplexClient +from vllm_omni.clients.minicpmo_4_5 import create_duplex_session_config + +config = create_duplex_session_config(ref_audio=audio_data_url(Path("reference_voice.wav"))) +async with InlineDuplexClient(omni, model="openbmb/MiniCPM-o-4_5", config=config) as client: + await client.stream_pcm(question) + await client.commit() + async for response in client.responses(): + ... +``` + +`examples/online_serving/barge_in_client.py --inline` drives the barge-in +scenario this way. The inline client has no transport, so it has no +reconnect, resume, heartbeat or event-acknowledgement machinery. + ## API Reference -This section is the complete wire contract of `/v1/realtime`: how +This section is the complete wire contract of `/v1/realtime?duplex=1`: how the dialect relates to the OpenAI Realtime protocol, which surfaces each model's capability flags gate, a catalogue of every client-to-server and server-to-client event, and a JSON example for each. `DuplexClient` speaks @@ -253,21 +338,26 @@ its protocol tests. ### Endpoint and transport -The endpoint is `ws(s):///v1/realtime`. Bare connections select the -session handler when the deployment declares either `session_mode: duplex` -or `duplex_session`. `?duplex=1` remains a compatibility alias, while -`?duplex=0` selects the legacy Realtime handler. Optional query parameters -are `model`, `session_id`, -`autostart` (`0` means resume-only), `resume`, and `native_duplex` -(`minicpmo45_native_duplex` is accepted as a deprecated alias and folded -into the canonical name). Every message is one JSON object per WebSocket +The endpoint is `ws(s):///v1/realtime`. On a deployment whose pipeline +declares a `duplex_plugin` and whose deploy configuration declares +`session_mode: duplex`, a bare connection is already a duplex session, so a +stock Realtime client needs no vendor query parameter. `?duplex=1` (`1`, +`true` or `on`) is the spelling this page uses and stays valid; +`?duplex=0` (`0`, `false`, `off`) explicitly opts out and selects the +turn-based Realtime handler, which a duplex server does not mount -- +that connection is answered with `Realtime API is not available`. +`ws(s):///v1/duplex` is an alias that always selects the duplex +handler. Optional query parameters are `model`, `autostart` (`0` means +resume-only) and `resume`. Every message is one JSON object per WebSocket text frame, discriminated by `type`. Inbound events are applied in arrival order through one per-session mailbox; outbound events preserve that order and carry a monotonically increasing `server_event_seq`, which is the replay cursor for `session.resume`. Client events may carry an `event_id`, echoed in `error.error.event_id` when the event is rejected. -Session identity is `(session_id, incarnation)`; `attachment_generation` +Session identity is the server-allocated `session_id` (`duplex-`, +announced in `session.created`; a `session_id` the client puts in the session +object is ignored, and ids are never reused); `attachment_generation` names the current socket; `epoch` is the interruption fence that advances on every cancel, clear, or barge-in; `turn_id` counts committed model turns; `response_id` (item id `item_`) names one assistant utterance. @@ -281,12 +371,13 @@ with no OpenAI counterpart. The event vocabulary is uniform, but several surfaces are gated by the `capabilities` object the server returns in `session.created`; a client must -branch on those flags rather than on the model name. The current model -plugins advertise: +branch on those flags rather than on the model name. MiniCPM-o 4.5 is the +only model on the plugin contract today; the other two columns record what +their integrations advertise once the follow-up PRs port them: | Capability | MiniCPM-o 4.5 | PersonaPlex | Nemotron VoiceChat | Gated surface | | --- | --- | --- | --- | --- | -| `implementation_level` | `model_native_duplex` (when `extra_body.native_duplex=true`) | `model_native_duplex` | `model_native_duplex` | model-owned `response.listen` / `response.speak`; the chat-fallback lane otherwise | +| `implementation_level` | `model_native_duplex` | `model_native_duplex` | `model_native_duplex` | model-owned `response.listen` / `response.speak` (constant: every duplex model is model-native) | | `chunk_period_ms` | 1000 | 80 | 80 | the model unit that `response.listen` decisions and camera frames align to | | `supports_session_resume` | yes | no | yes | `session.resume`, `session.resumed`, `session.replaced`, replay after a transport drop | | `supports_barge_in` | yes | no | no | `barge_in`, `turn.signal{event:"barge_in"}`, `overlap_policy=barge_in_on_speech`, `turn_detection.server_vad` | @@ -316,14 +407,8 @@ work unmodified. The overlap falls into the three tiers used by the event catalo | Object shapes | `realtime.session`, `realtime.item`, `realtime.response`; item content parts `input_text` / `input_audio` / `output_text` / `output_audio`; `function_call` and `function_call_output` items; `previous_item_id` chaining; `output_index` / `content_index` addressing | | Sequencing | `response.created → output_item.added → content_part.added → deltas → *.done → content_part.done → output_item.done → response.done → rate_limits.updated`; `speech_started → speech_stopped → committed` | -Audio/transcript output uses only the current (non-beta) OpenAI event names -(`response.output_audio.delta`, not the deprecated `response.audio.delta`), -matching `openai.types.realtime` in the official `openai` Python SDK and what -clients built against it (e.g. `livekit-plugins-openai`'s non-Azure code path) -expect; a client still on the old beta SDK types would not recognize these -events. Session fields and `conversation.item.*` events, in contrast, are -deliberately dual: the server emits and accepts **both** the beta and the GA -spellings at once (`output_modalities` alongside `modalities`, +The server deliberately emits and accepts **both** the beta and the GA +spellings at once (`response.output_audio.delta` alongside `output_modalities`, `conversation.item.added` alongside `conversation.item.created`) so either generation of OpenAI client parses the stream. @@ -333,7 +418,7 @@ A stock client ignores the extra keys; the extensions are additive. | Message | Extension | | --- | --- | -| `session.created` / `session.updated` | top-level `incarnation`, `attachment_generation`, `resume_token`; inside `session`: `state`, `turn_state`, `epoch`, `turn_id`, `active_request_id`, `active_response_id`, `active_response_turn_id`, `overlap_policy`, `overlap_*_ms/rms`, `playback_commit_policy`, `playback`, `capabilities`, `ref_audio`, `extra_body`, `idle_timeout_s`, `response_format` | +| `session.created` / `session.updated` | top-level `attachment_generation`, `resume_token`; inside `session`: `state`, `turn_state`, `epoch`, `turn_id`, `active_request_id`, `active_response_id`, `active_response_turn_id`, `overlap_policy`, `overlap_*_ms/rms`, `playback_commit_policy`, `playback`, `capabilities`, `ref_audio`, `extra_body`, `idle_timeout_s`, `response_format` | | `response.created` / `response.done` / `response.listen` | `response_id` at top level; the raw duplex event under `response.metadata` (`duplex_event`); `status_details.reason` uses vLLM reasons (`barge_in`, `client_cancelled`, `new_response`, …) | | `response.output_audio.delta` | `format`, `sample_rate_hz`, `metadata{session_id, epoch, model_speak, end_of_turn, audio_duration_ms, audio_text_marks, playback}` | | `response.speak` (inserted before the first delta) | not an OpenAI event, but rides the OpenAI response envelope (`response_id`, `item_id`, `output_index`, `content_index`) | @@ -349,14 +434,9 @@ Semantic divergences hidden behind shared names: before any commit; a commit may end in `response.listen` and no response, and the model may open a response with no commit at all. OpenAI: commit ⇒ item, `response.create` ⇒ exactly one response. -- With `turn_detection={"type":"server_vad"}`, omitted `interrupt_response` - defaults to `false` on turn-based models such as Qwen3-Omni, enabling - endpointing and automatic responses without interruption. Explicit `true` - is rejected there. Model-native duplex keeps the `true` default and rejects - `false`; its VAD implies `overlap_policy=barge_in_on_speech`. - This capability-dependent default differs from OpenAI's `true` default; - the effective configuration is returned in `session.created` / `session.updated`. - `semantic_vad` is unsupported. +- `turn_detection.interrupt_response=false` is rejected; + `create_response` is ignored; `semantic_vad` is unsupported; the VAD + runs per session (Silero) and implies `overlap_policy=barge_in_on_speech`. - `rate_limits.updated` is always an empty list (compatibility only). - Cancellation never reuses a `response_id`; truncation is driven by `playback.ack` as well as `conversation.item.truncate`. @@ -365,11 +445,11 @@ Semantic divergences hidden behind shared names: | Area | Messages / fields | | --- | --- | -| Session lifetime | `session.heartbeat` / `session.heartbeat_ack`, `session.event_ack`, `session.close` / `session.closed`, `session.resume` / `session.resumed` / `session.replaced` / `session.expired` / `session.resync_required`; `server_event_seq`, `resume_token`, `incarnation`, `attachment_generation`; query params `?duplex=1`, `autostart`, `resume`, `native_duplex` | +| Session lifetime | `session.heartbeat` / `session.heartbeat_ack`, `session.event_ack`, `session.close` / `session.closed`, `session.resume` / `session.resumed` / `session.replaced` / `session.expired` / `session.resync_required`; `server_event_seq`, `resume_token`, `attachment_generation`; query params `?duplex=1`, `autostart`, `resume` | | Turn-taking | `response.speak`, `response.listen`, `overlap.decision`, `overlap_policy`, `barge_in`, `turn.signal`, `input.cancel`, `input.text.append`, `epoch` / `turn_id` fencing, `force_listen` | | Playback truth | `playback.ack` / `playback.acknowledged`, `playback_commit_policy` (`ack_only` \| `commit_all_on_done`), `playback{generated_ms, sent_ms, played_ms, committed_ms}`, `history_committed`, `audio_text_marks` | -| Model negotiation | `capabilities{implementation_level, supports_input_append, supports_barge_in, supports_session_resume, chunk_period_ms, input_modes, …}`, `ref_audio`, `extra_body` (`auto_response`, `native_duplex` — deprecated alias `minicpmo45_native_duplex` is folded into it, `force_listen_count`, `duplex_initial_user_text`) | -| Diagnostics | `runtime.control`, `duplex.function_call.done`, the full error-code vocabulary (`REALTIME_ERROR_TYPES_BY_CODE` in `vllm_omni/entrypoints/duplex/realtime_state.py`) | +| Model negotiation | `capabilities{implementation_level, supports_input_append, supports_barge_in, supports_session_resume, chunk_period_ms, input_modes, …}`, `ref_audio`, `extra_body` (`auto_response`, `force_listen_count`, `duplex_initial_user_text`) | +| Diagnostics | `duplex.function_call.done`, the full error-code vocabulary (`REALTIME_ERROR_TYPES_BY_CODE` in `vllm_omni/engine/duplex/events.py`) | #### Consequences for clients @@ -433,11 +513,15 @@ remains an approximation, not word-level audio alignment. ### Event catalogue Every event type the Realtime route accepts or emits, with a worked example in the next section: 21 -client→server and 42 server→client. Aliases (`push_chunk`, -`input.audio.append`, `input_text.append`, `push_text`, `signal_turn`, -`audio.playback_ack`, `close_session`, `close`, `session_close`) share the -payload of their canonical event and are not listed separately; the alias -table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md). +client→server and 42 server→client. Aliases (`input_text.append`, +`push_text`, `signal_turn`, `audio.playback_ack`, `close_session`, `close`) +share the payload of their canonical event and are not listed +separately; the alias table is `translate_realtime_command` in +`vllm_omni/engine/duplex/realtime_commands.py`. The pre-Realtime WAV-append +aliases `push_chunk` and `input.audio.append` are **not** accepted on this +route: send `input_audio_buffer.append` with one of the supported input +formats (`pcm16`, `pcm_s16le`, `s16le`, `pcm_f32le`, `g711_ulaw`, +`g711_alaw`). #### Client to server @@ -447,18 +531,18 @@ table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o | `session.heartbeat` | 3 | Keepalive; refreshes the engine session lease. | | `session.event_ack` | 3 | Acknowledges received server events by `server_event_seq`; trims the resume replay journal. | | `session.close` | 3 | Graceful close; answered with `session.closed`. | -| `session.resume` | 3 | Re-attach a live session on a new socket (`resume_token`, `incarnation`, `last_received_server_event_seq`). | +| `session.resume` | 3 | Re-attach a live session on a new socket (`session_id`, `resume_token`, `last_received_server_event_seq`). | | `input_audio_buffer.append` | 2 | Append one audio chunk. OpenAI name; extended with `is_speech`, `video_frames`, `duration_ms`, `audio_end_ms`, per-event `format`/`sample_rate_hz`. | | `input_audio_buffer.commit` | 2 | Seal the buffered utterance into a user item. OpenAI name; extended with `final`, `response_create`, `is_speech:false`. | | `input.commit` | 3 | Alias of `input_audio_buffer.commit` accepted by the runner. | | `input_audio_buffer.clear` | 1 | Drop un-committed input audio. | | `input.cancel` | 3 | Cancel pending input; advances `epoch`. | -| `input.text.append` | 3 | Append text to the open input item (chat-fallback lane only). | +| `input.text.append` | 3 | Append text input; model-native sessions reject it with `native_text_append_unsupported`. | | `conversation.item.create` | 1 | Add a completed user item (text, audio, or `function_call_output`) to history. | | `conversation.item.retrieve` | 1 | Fetch a stored item by id. | | `conversation.item.delete` | 1 | Delete a stored item by id. | | `conversation.item.truncate` | 1 | Truncate an assistant item's audio/transcript at `audio_end_ms`. | -| `response.create` | 1 | Explicitly request a response with per-response overrides. | +| `response.create` | 1 | Explicitly request a response with per-response overrides. Requires committed audio input; otherwise `response_create_without_input`. | | `response.cancel` | 1 | Cancel the active (or named) response; advances `epoch`. | | `output_audio_buffer.clear` | 1 | Discard queued output audio; advances `epoch` (OpenAI: WebRTC-only, here also WebSocket). | | `barge_in` | 3 | Explicit hard interrupt; requires `capabilities.supports_barge_in`. | @@ -469,10 +553,10 @@ table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o | Event type | Tier | Description | | --- | --- | --- | -| `session.created` | 2 | Session opened. OpenAI name; adds `incarnation`, `attachment_generation`, `resume_token` and vLLM-Omni keys inside `session` (`epoch`, `turn_id`, `playback`, `capabilities`, ...). | +| `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; `reason` ∈ `session_close`, `timeout`, `disconnect`, ... | +| `session.closed` | 3 | Last event on the socket, emitted once the engine released the session; `reason` ∈ `client_close`, `disconnect`, `timeout`, `transport_error`, `shutdown`, ... | | `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. | @@ -508,8 +592,8 @@ table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o | `duplex.function_call.done` | 3 | Raw fallback when a model tool call is malformed. | | `output_audio_buffer.cleared` | 1 | Reply to `output_audio_buffer.clear`. | | `playback.acknowledged` | 3 | Reply to `playback.ack`; returns the playback ledger and `history_committed`. | -| `overlap.decision` | 3 | How an append that overlapped an active response was handled (`drop` \| `listen` \| `barge_in`). | -| `runtime.control` | 3 | Diagnostic: an engine control signal was partially unsupported. | +| `overlap.decision` | 3 | How an append that overlapped an active response was handled (`drop` \| `listen` \| `barge_in`); details under `metadata`. | +| `turn.event` | 3 | A turn-taking signal the session accepted (`event`, resulting `turn_state`). | | `error` | 2 | Error envelope. OpenAI shape and classes; vLLM-Omni error codes. | ### Message examples @@ -528,7 +612,6 @@ additionally carries `"server_event_seq": ` (omitted below for brevity). "event_id": "evt_0001", "session": { "model": "openbmb/MiniCPM-o-4_5", - "session_id": "sess_demo_01", "modalities": ["audio", "text"], "instructions": "You are a concise voice assistant.", "input_audio_format": "pcm16", @@ -541,7 +624,6 @@ additionally carries `"server_event_seq": ` (omitted below for brevity). "idle_timeout_s": 300, "extra_body": { "auto_response": true, - "native_duplex": true, "force_listen_count": 0 } } @@ -553,13 +635,12 @@ additionally carries `"server_event_seq": ` (omitted below for brevity). ```json { "type": "session.created", - "incarnation": 0, "attachment_generation": 1, "resume_token": "rt_3f9c…", "session": { "object": "realtime.session", "type": "realtime", - "id": "sess_demo_01", + "id": "duplex-5f2a…", "model": "openbmb/MiniCPM-o-4_5", "state": "open", "turn_state": "idle", @@ -612,7 +693,7 @@ additionally carries `"server_event_seq": ` (omitted below for brevity). `session.created` and every later `session.update`) ```json -{"type": "session.updated", "session": {"object": "realtime.session", "id": "sess_demo_01", "…": "…"}} +{"type": "session.updated", "session": {"object": "realtime.session", "id": "duplex-5f2a…", "…": "…"}} ``` `C→S session.heartbeat` / `S→C session.heartbeat_ack` @@ -622,7 +703,7 @@ additionally carries `"server_event_seq": ` (omitted below for brevity). ``` ```json -{"type": "session.heartbeat_ack", "session_id": "sess_demo_01"} +{"type": "session.heartbeat_ack", "session_id": "duplex-5f2a…"} ``` `C→S session.event_ack` @@ -638,7 +719,7 @@ additionally carries `"server_event_seq": ` (omitted below for brevity). ``` ```json -{"type": "session.closed", "session_id": "sess_demo_01", "reason": "session_close"} +{"type": "session.closed", "session_id": "duplex-5f2a…", "epoch": 3, "reason": "client_close", "event": {"type": "session.closed", "reason": "client_close"}} ``` `C→S session.resume` (new socket, `?duplex=1&autostart=0`) @@ -646,8 +727,7 @@ additionally carries `"server_event_seq": ` (omitted below for brevity). ```json { "type": "session.resume", - "session_id": "sess_demo_01", - "incarnation": 0, + "session_id": "duplex-5f2a…", "resume_token": "rt_3f9c…", "last_received_server_event_seq": 40 } @@ -659,8 +739,8 @@ events with `server_event_seq > 40` are replayed in order) ```json { "type": "session.resumed", - "session_id": "sess_demo_01", - "incarnation": 0, + "session_id": "duplex-5f2a…", + "session": {"id": "duplex-5f2a…", "state": "open", "epoch": 3, "turn_id": 2, "...": "the engine's current session object"}, "attachment_generation": 2, "resume_token": "rt_a71e…" } @@ -669,19 +749,19 @@ events with `server_event_seq > 40` are replayed in order) `S→C session.replaced` (sent to the superseded socket, which is then closed) ```json -{"type": "session.replaced", "session_id": "sess_demo_01", "attachment_generation": 1} +{"type": "session.replaced", "session_id": "duplex-5f2a…", "attachment_generation": 1} ``` `S→C session.expired` ```json -{"type": "session.expired", "session_id": "sess_demo_01", "incarnation": 0, "reason": "disconnect_grace_expired"} +{"type": "session.expired", "session_id": "duplex-5f2a…", "reason": "disconnect_grace_expired"} ``` `S→C session.resync_required` ```json -{"type": "session.resync_required", "session_id": "sess_demo_01", "reason": "journal_gap"} +{"type": "session.resync_required", "session_id": "duplex-5f2a…", "reason": "journal_gap"} ``` #### Input @@ -772,7 +852,7 @@ silence/noise declaration (no runtime append; answered with `response.listen`) "item_id": "item_9b2d…", "event": { "type": "input.committed", - "session_id": "sess_demo_01", + "session_id": "duplex-5f2a…", "turn_id": 1, "epoch": 0, "history_len": 1, @@ -802,7 +882,7 @@ deferred commit during an active response (`event.response_create_deferred`) "type": "input_audio_buffer.committed", "previous_item_id": "item_9b2d…", "item_id": "item_c044…", - "event": {"type": "input.committed", "session_id": "sess_demo_01", "turn_id": 2, "epoch": 0, "history_len": 3, "response_create_deferred": true, "message": {"…": "…"}} + "event": {"type": "input.committed", "session_id": "duplex-5f2a…", "turn_id": 2, "epoch": 0, "history_len": 3, "response_create_deferred": true, "message": {"…": "…"}} } ``` @@ -911,7 +991,7 @@ deferred commit during an active response (`event.response_create_deferred`) "status_details": null, "output": [], "modalities": ["audio", "text"], - "metadata": {"duplex_event": {"type": "response.created", "session_id": "sess_demo_01", "response_id": "resp_01", "epoch": 0, "turn_id": 1}} + "metadata": {"duplex_event": {"type": "response.created", "session_id": "duplex-5f2a…", "response_id": "resp_01", "epoch": 0, "turn_id": 1}} } } ``` @@ -947,7 +1027,7 @@ deferred commit during an active response (`event.response_create_deferred`) "item_id": "item_resp_01", "output_index": 0, "content_index": 0, - "metadata": {"session_id": "sess_demo_01", "epoch": 0, "model_speak": true} + "metadata": {"session_id": "duplex-5f2a…", "epoch": 0, "model_speak": true} } ``` @@ -964,7 +1044,7 @@ deferred commit during an active response (`event.response_create_deferred`) "format": "pcm16", "sample_rate_hz": 24000, "metadata": { - "session_id": "sess_demo_01", + "session_id": "duplex-5f2a…", "epoch": 0, "model_speak": true, "end_of_turn": false, @@ -1029,7 +1109,7 @@ deferred commit during an active response (`event.response_create_deferred`) "status": "completed", "status_details": null, "output": [{"id": "item_resp_01", "object": "realtime.item", "type": "message", "role": "assistant", "status": "completed", "content": [{"type": "output_audio", "transcript": "It is sunny, 24 degrees."}]}], - "metadata": {"type": "response.done", "session_id": "sess_demo_01", "response_id": "resp_01", "epoch": 0, "committed": true, "playback": {"generated_ms": 2400, "sent_ms": 2400, "played_ms": 2400, "committed_ms": 2400}} + "metadata": {"type": "response.done", "session_id": "duplex-5f2a…", "response_id": "resp_01", "epoch": 0, "committed": true, "playback": {"generated_ms": 2400, "sent_ms": 2400, "played_ms": 2400, "committed_ms": 2400}} } } ``` @@ -1055,9 +1135,9 @@ deferred commit during an active response (`event.response_create_deferred`) ```json { "type": "response.listen", - "session_id": "sess_demo_01", + "session_id": "duplex-5f2a…", "epoch": 0, - "response": {"object": "realtime.response", "status": "listening", "metadata": {"type": "response.listen", "session_id": "sess_demo_01", "epoch": 0, "reason": "model_listen"}} + "response": {"object": "realtime.response", "status": "listening", "metadata": {"type": "response.listen", "session_id": "duplex-5f2a…", "epoch": 0, "reason": "model_listen"}} } ``` @@ -1104,7 +1184,7 @@ function-call response (Nemotron VoiceChat tools) raw fallback when a model tool call is malformed (`call_id`/`name` missing) ```json -{"type": "duplex.function_call.done", "event": {"type": "function_call.done", "session_id": "sess_demo_01", "epoch": 0, "call_id": null, "name": null, "arguments": "{"}} +{"type": "duplex.function_call.done", "event": {"type": "function_call.done", "session_id": "duplex-5f2a…", "epoch": 0, "call_id": null, "name": null, "arguments": "{"}} ``` #### Playback @@ -1120,7 +1200,7 @@ raw fallback when a model tool call is malformed (`call_id`/`name` missing) "type": "playback.acknowledged", "event": { "type": "playback.acknowledged", - "session_id": "sess_demo_01", + "session_id": "duplex-5f2a…", "epoch": 0, "item_id": "item_resp_01", "played_ms": 1000, @@ -1145,20 +1225,17 @@ truncating ack after an interruption (what was actually heard) ```json { "type": "overlap.decision", - "session_id": "sess_demo_01", + "session_id": "duplex-5f2a…", "epoch": 0, "policy": "listen_only", "action": "listen", "reason": "long_overlap_speech", - "overlap_ms": 1400, - "buffer_audio": true, - "defer_runtime_append": true, - "force_listen": true + "metadata": {"overlap_ms": 1400, "buffer_audio": true, "defer_runtime_append": true, "force_listen": true} } ``` ```json -{"type": "overlap.decision", "session_id": "sess_demo_01", "epoch": 0, "policy": "barge_in_on_speech", "action": "barge_in", "reason": "server_vad_utterance_active"} +{"type": "overlap.decision", "session_id": "duplex-5f2a…", "epoch": 0, "policy": "barge_in_on_speech", "action": "barge_in", "reason": "server_vad_utterance_active"} ``` server-VAD session (`session.update` fragment) and the annotated append the @@ -1172,12 +1249,6 @@ translator forwards internally (shown for reference; clients never see it) {"type": "input_audio_buffer.append", "audio": "…", "format": "pcm16", "sample_rate_hz": 16000, "is_speech": true, "force_listen": true, "vad": {"backend": "silero", "is_speech": true, "speech_active": true, "speech_started": true, "speech_stopped": false, "speech_probability": 0.93}} ``` -`S→C runtime.control` (only when an engine signal was partially unsupported) - -```json -{"type": "runtime.control", "session_id": "sess_demo_01", "epoch": 1, "result": {"…": "redacted engine control result"}} -``` - #### Errors `S→C error` — Realtime shape @@ -1203,71 +1274,79 @@ translator forwards internally (shown for reference; clients never see it) ``` ```json -{"type": "error", "error": {"type": "rate_limit_error", "code": "duplex_session_busy", "message": "duplex session capacity reached (max_sessions=4)"}} +{"type": "error", "error": {"type": "rate_limit_error", "code": "resource_exhausted", "message": "duplex_session_capacity_exhausted: limit=4"}} ``` Handshake-stage errors (before the Realtime projector is bound) use the flat native shape: ```json -{"type": "error", "error": "Unknown or expired duplex session: sess_demo_01", "code": "session_resume_expired"} +{"type": "error", "error": "Unknown or expired duplex session: duplex-5f2a…", "code": "session_resume_expired"} ``` #### Coverage The examples above cover every event type the Realtime route emits or accepts (verified by diffing the `type` literals in -`vllm_omni/entrypoints/duplex/` against this document). Not shown, because -they never reach a `/v1/realtime?duplex=1` client unprojected: the native -`/v1/duplex` dialect (`session.create` / `open_session` / `session.config`, -`input.committed`, `input.cancelled`, `audio.cancelled`, -`response.output_audio.delta`, `response.text.delta`, `response.message`, -`function_call.done`) — their Realtime projections are the examples above -(see the name map at the end of this section) — and the runner-internal markers `__timeout__`, -`__disconnect__`, `__replaced_attachment__`. Input aliases (`push_chunk`, -`input.audio.append`, `input_text.append`, `push_text`, `signal_turn`, -`audio.playback_ack`, `close_session`, `close`, `session_close`) share the -payload of their canonical event (see the alias table in [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md)). +`vllm_omni/engine/duplex/events.py` and `vllm_omni/engine/duplex/commands.py` +against this document). Not shown, because they never reach a client +unprojected: the session-internal events the runner produces +(`input.committed`, `input.cancelled`, `audio.cancelled`, +`response.output_audio.delta`, `response.text.delta`, `function_call.done`) +— their Realtime projections are the examples above (see the name map at the +end of this section). Every public event is a typed `DuplexEvent` whose +`to_realtime()` is the wire object; every client event is parsed into a +typed `DuplexCommand` by `command_from_realtime()`. -#### Native to Realtime name map +#### Internal to Realtime name map -The native `/v1/duplex` dialect is what the session runner produces; the -Realtime projector renames and fans out these events before they reach a -`/v1/realtime?duplex=1` client. +The session runner emits these session-internal events; the Realtime +projector (`vllm_omni/engine/duplex/realtime_events.py`) renames and fans +them out into typed events before they reach a client. -| Native / internal event | Realtime projection | +| Internal event | Realtime projection | | --- | --- | -| `session.create` (`open_session`, `session.config`) | first `session.update` | | `input.committed` | `conversation.item.added` / `conversation.item.created`, `input_audio_buffer.committed`, `conversation.item.input_audio_transcription.completed`, `conversation.item.done` | | `input.cancelled` | `input_audio_buffer.cleared` | | `response.output_audio.delta` (with `audio_transcript`) | optional `response.speak`, then `response.output_audio.delta` and `response.output_audio_transcript.delta` | | `response.output_audio.done` / `response.output_text.done` | `response.output_audio.done` / `response.output_text.done` | | `response.text.delta` | `response.output_text.delta` | -| `response.message` (chat-fallback raw chunk) | passed through | | `audio.cancelled` | optional `output_audio_buffer.cleared`, then the cancelled terminal events and `response.done` with `status: "cancelled"` | | `function_call.done` | `response.created`, `response.output_item.added` (`function_call` item), `response.function_call_arguments.delta` / `.done`, `response.output_item.done`, `response.done` | | `playback.acknowledged` | `playback.acknowledged` (wrapped verbatim under `event`) | -| `runtime.control`, `session.resumed`, `session.replaced`, `session.expired`, `session.resync_required`, `session.heartbeat_ack` | passed through unchanged | +| `session.resumed`, `session.replaced`, `session.expired`, `session.resync_required`, `session.heartbeat_ack` | passed through unchanged | ## Known Limitations +- Only MiniCPM-o 4.5 is served over this endpoint today; PersonaPlex and + Nemotron VoiceChat arrive with the follow-up PRs that port them to the + plugin contract. - Several surfaces are capability-gated per model (see *Capability negotiation by model* above): PersonaPlex does not support session resume, barge-in, or audio truncation; Nemotron VoiceChat does not support barge-in or audio truncation; camera frames are consumed only by MiniCPM-o 4.5; tool calls are produced only by Nemotron VoiceChat. -- `turn_detection` supports `server_vad` or `null`, not `semantic_vad`. - Turn-based models support endpointing without interruption, including - `create_response=false` for explicit response creation. Model-native duplex - requires `interrupt_response=true` and `create_response=true`. +- A response is generated from committed **audio**, so the OpenAI text-prompt + shape does not drive one: `conversation.item.create` with an `input_text` + part adds the item to history, but a following `response.create` with no + committed audio is rejected with `response_create_without_input`, and + `input.text.append` is rejected with `native_text_append_unsupported`. + Text-to-speech has its own shape here — put the text in + `extra_body.duplex_initial_user_text` on `session.update`, then stream audio + units (silence is enough). The seeded turn is what the model answers. +- `turn_detection` supports only `server_vad` with `interrupt_response=true`; + `semantic_vad`, `interrupt_response=false`, and `create_response` are not + supported. - `input_audio_transcription` and `input_audio_noise_reduction` are accepted and echoed but no separate transcription or noise-reduction stage runs; transcripts come from the model. - `rate_limits.updated` is emitted for compatibility only and always carries an empty list. - Session capacity is bounded by `duplex_session.max_sessions` in the deploy - configuration; admission beyond it fails with `duplex_session_busy` or - `resource_exhausted`. -- The native `/v1/duplex` route speaks the internal dialect and does not - support `session.resume`; use `/v1/realtime?duplex=1` for the contract on - this page. + configuration; admission beyond it fails with `resource_exhausted` + (`rate_limit_error`, retryable). +- A duplex deployment serves the duplex session surfaces only: the websocket + route and `/v1/chat/completions`, the latter on a session per request (see + [Start the Server](#start-the-server)). The other turn-based HTTP routes + report "not available" on a MiniCPM-o 4.5 server; turn-based use of the + model stays available offline through `Omni` / `AsyncOmni`. diff --git a/examples/online_serving/barge_in_client.py b/examples/online_serving/barge_in_client.py index b144da9f4..b9d95c790 100644 --- a/examples/online_serving/barge_in_client.py +++ b/examples/online_serving/barge_in_client.py @@ -14,7 +14,10 @@ Scenario (see ``barge_in_client_flow.md`` for the diagram) The demo is model-neutral: it drives any duplex model through the public :class:`vllm_omni.clients.duplex.DuplexClient` over -``/v1/realtime?duplex=1``. Model-specific session shape comes from the +``/v1/realtime?duplex=1``, or — with ``--inline`` — through +:class:`vllm_omni.clients.inline_duplex.InlineDuplexClient` on an in-process +:class:`vllm_omni.entrypoints.duplex_omni.DuplexOmni` (no server needed; the +model loads in this process). Model-specific session shape comes from the per-model presets under ``vllm_omni.clients.`` selected with ``--preset`` (default ``minicpmo_4_5``, whose model decides listen/speak on its own — no VAD). @@ -31,6 +34,15 @@ then: --interrupt-wav follow_up_16k.wav \ --output-dir ./duplex_out +or, without a server, in-process (``--deploy-config`` is optional and defaults +to the model's bundled duplex profile): + + python barge_in_client.py --inline \ + --model /path/to/MiniCPM-o-4_5 \ + --ref-audio /path/to/reference_voice.wav \ + --question-wav question_16k.wav \ + --interrupt-wav follow_up_16k.wav + Input WAVs must be mono 16 kHz PCM16; when a preset's session takes a different capture format (the PersonaPlex preset streams 24 kHz float32), the audio is converted before streaming. ``--ref-audio`` is required by the @@ -55,6 +67,7 @@ if str(REPO_ROOT) not in sys.path: from vllm_omni.clients.duplex import ( # noqa: E402 AudioFormat, DuplexClient, + DuplexClientBase, EventCollector, SessionConfig, acknowledge_collected_playback, @@ -139,6 +152,39 @@ def _fold_responses(collector: EventCollector) -> dict[str, dict[str, object]]: return out +def _make_client(args: argparse.Namespace, config: SessionConfig) -> DuplexClientBase: + """Websocket client by default; ``--inline`` drives an in-process DuplexOmni.""" + if not args.inline: + return DuplexClient( + args.url, + model=args.model, + config=config, + reconnect=None, + heartbeat_interval_s=None, + handshake_timeout_s=args.timeout_s, + ) + # Imported lazily: the in-process path loads the model (and its engine + # dependencies) in this process, which the websocket path never needs. + from vllm_omni.clients.inline_duplex import InlineDuplexClient + from vllm_omni.entrypoints.duplex_omni import DuplexOmni + + # Models shipped with custom HF code (MiniCPM-o) need trust_remote_code to + # resolve their pipeline; harmless for the others. + omni_kwargs: dict[str, object] = {"model": args.model, "trust_remote_code": True} + if args.deploy_config: + omni_kwargs["deploy_config"] = args.deploy_config + omni = DuplexOmni(**omni_kwargs) + client = InlineDuplexClient( + omni, + model=args.model, + config=config, + handshake_timeout_s=args.timeout_s, + ) + # The client owns no engine; stop the in-process one when the run ends. + client.owned_omni = omni # type: ignore[attr-defined] + return client + + async def run(args: argparse.Namespace) -> int: config = _session_config(args.preset, ref_audio_path=args.ref_audio) question = _convert_to_session_format(read_pcm16_wav(Path(args.question_wav)), config.input_audio) @@ -148,105 +194,117 @@ async def run(args: argparse.Namespace) -> int: output_dir.mkdir(parents=True, exist_ok=True) collector = EventCollector() - client = DuplexClient( - args.url, - model=args.model, - config=config, - reconnect=None, - heartbeat_interval_s=None, - handshake_timeout_s=args.timeout_s, - ) + client = _make_client(args, config) # GREET: open the voice session (handshake happens on enter). - async with client: - consume_task = asyncio.create_task(collector.consume(client)) + try: + async with client: + consume_task = asyncio.create_task(collector.consume(client)) - # ASK: stream the question at realtime pace, then commit the turn. - # The model decides on its own to answer (no VAD). - print("[1] streaming question ...", file=sys.stderr) - await client.stream_pcm(question, chunk_ms=chunk_ms) - await client.commit(create_response=False) - await wait_for_condition( - lambda: collector.count("response.created") > 0, - timeout_s=args.timeout_s, - label="first response.created", - ) - print("[1] assistant is answering", file=sys.stderr) + # ASK: stream the question at realtime pace, then commit the turn. + # The model decides on its own to answer (no VAD). + print("[1] streaming question ...", file=sys.stderr) + await client.stream_pcm(question, chunk_ms=chunk_ms) + await client.commit(create_response=False) + await wait_for_condition( + lambda: collector.count("response.created") > 0, + timeout_s=args.timeout_s, + label="first response.created", + ) + print("[1] assistant is answering", file=sys.stderr) - # Let the answer play for a moment, as a listening user would. - await asyncio.sleep(args.listen_s) + # Let the answer play for a moment, as a listening user would. + await asyncio.sleep(args.listen_s) - # INTERRUPT: talk over the assistant while its audio still streams. - # The serving overlap policy decides: enough overlapped speech barges - # in (the first response is cancelled) or a short remark is deferred - # until the model yields the turn. - print("[2] interrupting mid-answer ...", file=sys.stderr) - ids_before_interrupt = set(collector.response_ids) - await client.stream_pcm(interrupt, chunk_ms=chunk_ms) - await client.commit(create_response=False) + # INTERRUPT: talk over the assistant while its audio still streams. + # The serving overlap policy decides: enough overlapped speech barges + # in (the first response is cancelled) or a short remark is deferred + # until the model yields the turn. + print("[2] interrupting mid-answer ...", file=sys.stderr) + ids_before_interrupt = set(collector.response_ids) + await client.stream_pcm(interrupt, chunk_ms=chunk_ms) + await client.commit(create_response=False) - # ANSWER AGAIN: wait for a response created *after* the interruption - # to finish. The first answer may still be streaming here; matching on - # any response.done would let its completion satisfy the wait and cut - # off the actual follow-up answer at close. - def follow_up_answer_done() -> bool: - return any( - event.get("type") == "response.done" - and collector.response_id(event) is not None - and collector.response_id(event) not in ids_before_interrupt - for event in collector.events + # ANSWER AGAIN: wait for a response created *after* the interruption + # to finish. The first answer may still be streaming here; matching on + # any response.done would let its completion satisfy the wait and cut + # off the actual follow-up answer at close. + def follow_up_answer_done() -> bool: + return any( + event.get("type") == "response.done" + and collector.response_id(event) is not None + and collector.response_id(event) not in ids_before_interrupt + for event in collector.events + ) + + await wait_for_condition( + follow_up_answer_done, + timeout_s=args.timeout_s, + label="response.done for the response answering the interruption", ) - await wait_for_condition( - follow_up_answer_done, - timeout_s=args.timeout_s, - label="response.done for the response answering the interruption", - ) + # HANG UP: report playback, close the session cleanly. + await acknowledge_collected_playback(client, collector) + await client.close(timeout_s=args.timeout_s) + try: + await asyncio.wait_for(consume_task, timeout=5.0) + except (asyncio.TimeoutError, asyncio.CancelledError): + consume_task.cancel() - # HANG UP: report playback, close the session cleanly. - await acknowledge_collected_playback(client, collector) - await client.close(timeout_s=args.timeout_s) - try: - await asyncio.wait_for(consume_task, timeout=5.0) - except (asyncio.TimeoutError, asyncio.CancelledError): - consume_task.cancel() - - # Report: one WAV per response; statuses show what the overlap did. - responses = _fold_responses(collector) - summary: list[dict[str, object]] = [] - for index, (response_id, item) in enumerate(responses.items(), start=1): - audio = item["audio"] - assert isinstance(audio, bytes) - sample_rate_hz = int(item["sample_rate_hz"]) # type: ignore[arg-type] - wav_path = output_dir / f"response_{index}_{item['status']}.wav" - if audio: - write_pcm16_wav(wav_path, audio, sample_rate_hz=sample_rate_hz) - summary.append( - { - "response_id": response_id, - "status": item["status"], - "audio_s": round(len(audio) / (sample_rate_hz * 2), 2), - "text": item["text"], - "wav": wav_path.name if audio else None, - } + # Report: one WAV per response; statuses show what the overlap did. + responses = _fold_responses(collector) + summary: list[dict[str, object]] = [] + for index, (response_id, item) in enumerate(responses.items(), start=1): + audio = item["audio"] + assert isinstance(audio, bytes) + sample_rate_hz = int(item["sample_rate_hz"]) # type: ignore[arg-type] + wav_path = output_dir / f"response_{index}_{item['status']}.wav" + if audio: + write_pcm16_wav(wav_path, audio, sample_rate_hz=sample_rate_hz) + summary.append( + { + "response_id": response_id, + "status": item["status"], + "audio_s": round(len(audio) / (sample_rate_hz * 2), 2), + "text": item["text"], + "wav": wav_path.name if audio else None, + } + ) + (output_dir / "summary.json").write_text(json.dumps(summary, indent=2, ensure_ascii=False)) + (output_dir / "events.jsonl").write_text( + "".join(json.dumps(event, ensure_ascii=False) + "\n" for event in collector.events), + encoding="utf-8", ) - (output_dir / "summary.json").write_text(json.dumps(summary, indent=2, ensure_ascii=False)) - (output_dir / "events.jsonl").write_text( - "".join(json.dumps(event, ensure_ascii=False) + "\n" for event in collector.events), - encoding="utf-8", - ) - print(json.dumps(summary, indent=2, ensure_ascii=False)) - cancelled = sum(1 for item in summary if item["status"] == "cancelled") - print( - f"\n{len(summary)} responses ({cancelled} interrupted); outputs in {output_dir}", - file=sys.stderr, - ) + print(json.dumps(summary, indent=2, ensure_ascii=False)) + cancelled = sum(1 for item in summary if item["status"] == "cancelled") + print( + f"\n{len(summary)} responses ({cancelled} interrupted); outputs in {output_dir}", + file=sys.stderr, + ) + finally: + owned_omni = getattr(client, "owned_omni", None) + if owned_omni is not None: + owned_omni.shutdown() return 0 def main() -> int: parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime") - parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5") + parser.add_argument( + "--model", + default="openbmb/MiniCPM-o-4_5", + help="served model name, or the model path with --inline", + ) + parser.add_argument( + "--inline", + action="store_true", + help="run the model in this process through DuplexOmni + InlineDuplexClient instead of connecting to --url", + ) + parser.add_argument( + "--deploy-config", + default=None, + help="deploy YAML for --inline (default: the model's bundled profile)", + ) parser.add_argument( "--preset", choices=["minicpmo_4_5", "personaplex", "none"], diff --git a/examples/online_serving/barge_in_client_flow.md b/examples/online_serving/barge_in_client_flow.md index d32a00e5a..9562d4c9c 100644 --- a/examples/online_serving/barge_in_client_flow.md +++ b/examples/online_serving/barge_in_client_flow.md @@ -1,8 +1,10 @@ # `barge_in_client.py` — what happens, at a glance -The demo plays one everyday conversation pattern against the full-duplex -server: ask something, get talked-over mid-answer, and see the assistant -cope. High-level flow below — for the exact wire events, read the +The demo plays one everyday conversation pattern against a full-duplex +model: ask something, get talked-over mid-answer, and see the assistant +cope. It runs either against a server over `/v1/realtime?duplex=1` or, with +`--inline`, against an in-process `DuplexOmni` — the conversation below is the +same in both modes. High-level flow — for the exact wire events, read the docstring in `barge_in_client.py`. ```mermaid @@ -48,7 +50,8 @@ answer is saved as `response_1_cancelled.wav`, a completed one as duration per answer. Companion: the runtime architecture lives in -[`docs/design/fullduplex.md`](../../docs/design/fullduplex.md); the client -API is `vllm_omni.clients.duplex.DuplexClient`, documented with the wire -protocol in +[`docs/design/fullduplex.md`](../../docs/design/fullduplex.md); the client API is +`vllm_omni.clients.duplex.DuplexClient` over the websocket and +`vllm_omni.clients.inline_duplex.InlineDuplexClient` in-process — both are +documented with the wire protocol in [`docs/serving/realtime_duplex_api.md`](../../docs/serving/realtime_duplex_api.md). diff --git a/examples/online_serving/minicpmo/barge_in_serve.sh b/examples/online_serving/minicpmo/barge_in_serve.sh index 6dade8c5c..1aabd3042 100644 --- a/examples/online_serving/minicpmo/barge_in_serve.sh +++ b/examples/online_serving/minicpmo/barge_in_serve.sh @@ -5,10 +5,10 @@ # Server side for barge_in_client.py: start MiniCPM-o 4.5 in native # full-duplex mode. # -# The duplex deploy profile registers the WS /v1/duplex endpoint (with -# /v1/realtime?duplex=1 projecting onto the same handler), enables the -# engine duplex control plane, and bounds the Thinker/Talker/Code2Wav -# stages to four live duplex sessions on one GPU. See +# MiniCPM-o 4.5 declares a duplex plugin, so vllm-omni serve runs it through +# DuplexOmni: the server is duplex-only (WS /v1/realtime?duplex=1, alias +# /v1/duplex, plus /v1/models and /health) and bounds the +# Thinker/Talker/Code2Wav stages to four live duplex sessions on one GPU. See # vllm_omni/deploy/minicpmo_4_5.yaml (session_mode: duplex) for the session # limits; every other duplex_session knob keeps its runtime default # (idle TTL 300s, disconnect grace 30s, 16MiB pending input per session). @@ -22,4 +22,4 @@ exec vllm-omni serve "$MODEL" \ --deploy-config vllm_omni/deploy/minicpmo_4_5.yaml \ --trust-remote-code \ --host 0.0.0.0 --port "$PORT" -# duplex endpoints: ws://:$PORT/v1/duplex and /v1/realtime?duplex=1 +# duplex endpoint: ws://:$PORT/v1/realtime?duplex=1 (alias /v1/duplex) diff --git a/examples/online_serving/minicpmo/realtime_duplex_demo.py b/examples/online_serving/minicpmo/realtime_duplex_demo.py index 9fc75621e..6dc3fd6e7 100644 --- a/examples/online_serving/minicpmo/realtime_duplex_demo.py +++ b/examples/online_serving/minicpmo/realtime_duplex_demo.py @@ -445,7 +445,6 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]: ref_audio=_ref_audio_data_url(args.ref_audio), temperature=args.temperature, ), - session_id=args.session_id, reconnect=None, heartbeat_interval_s=None, handshake_timeout_s=args.timeout_s, @@ -638,7 +637,6 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--url", default="ws://localhost:8099/v1/realtime?duplex=1") parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5") - parser.add_argument("--session-id") parser.add_argument( "--input-wav", help="Mono 16 kHz PCM16 clip to stream. Optional when --input-video carries the audio track.", diff --git a/examples/online_serving/minicpmo/realtime_web/app/static/app.js b/examples/online_serving/minicpmo/realtime_web/app/static/app.js index 1ea2ba093..4d2b78247 100644 --- a/examples/online_serving/minicpmo/realtime_web/app/static/app.js +++ b/examples/online_serving/minicpmo/realtime_web/app/static/app.js @@ -86,7 +86,6 @@ url.protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; url.searchParams.set('duplex', '1'); url.searchParams.set('model', config.model || 'openbmb/MiniCPM-o-4_5'); - url.searchParams.set('native_duplex', '1'); url.searchParams.set('autostart', '0'); return url.toString(); } @@ -458,7 +457,6 @@ settled = true; const extraBody = { auto_response: true, - native_duplex: true, }; const session = { modalities: ['audio', 'text'], diff --git a/recipes/OpenBMB/MiniCPM-o-4_5.md b/recipes/OpenBMB/MiniCPM-o-4_5.md index fb0934d4f..01580843f 100644 --- a/recipes/OpenBMB/MiniCPM-o-4_5.md +++ b/recipes/OpenBMB/MiniCPM-o-4_5.md @@ -1,6 +1,6 @@ # MiniCPM-o 4.5 -> Online serving and offline inference for omni multimodal chat +> Online full-duplex and chat-completions serving, and offline inference, for omni multimodal chat > (text / image / audio / video → text + 24 kHz speech) ## Summary @@ -9,8 +9,12 @@ - Model: [`openbmb/MiniCPM-o-4_5`](https://huggingface.co/openbmb/MiniCPM-o-4_5) - Task: Omni multimodal chat — accepts text / image / audio / video input; emits text and 24 kHz mono speech in the same response -- Mode: Online serving via the OpenAI-compatible `/v1/chat/completions` - API (plus Gradio demo), and offline inference via `Omni.generate` +- Mode: Online full-duplex serving over the OpenAI Realtime protocol + (`/v1/realtime?duplex=1`, alias `/v1/duplex`), the OpenAI-compatible + `/v1/chat/completions` API (plus Gradio demo) on the same server, and + offline turn-based inference via `Omni.generate`. Besides those the server + exposes only `/v1/models` and `/health`; the other turn-based HTTP routes + are not available on a duplex deployment. - Maintainer: [`@tc-mb`](https://github.com/tc-mb) (MiniCPM-V / MiniCPM-o team) ## When to use this recipe @@ -18,9 +22,10 @@ Use this recipe as a known-good starting point for serving `openbmb/MiniCPM-o-4_5` on vLLM-Omni. MiniCPM-o 4.5 is the omni member of the MiniCPM-o family — it runs a multimodal thinker, a streaming -MiniCPMTTS codec talker, and a separate batched Code2Wav stage so a single -`/v1/chat/completions` call can return text and 24 kHz speech in one -shot. The recommended batching deploy isolates the Thinker on GPU 0 and +MiniCPMTTS codec talker, and a separate batched Code2Wav stage, so a live +duplex session streams text and 24 kHz speech while it keeps listening, and a +single `/v1/chat/completions` call returns both in one shot. The +recommended batching deploy isolates the Thinker on GPU 0 and co-locates Talker and Code2Wav on GPU 1; 1-GPU, 3-GPU, and 8x4090 layouts are also provided. @@ -36,8 +41,11 @@ also provided. [`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml) - 8x RTX 4090 layout: [`vllm_omni/deploy/minicpmo_4_5_8x4090.yaml`](../../vllm_omni/deploy/minicpmo_4_5_8x4090.yaml) -- Online example + Gradio demo: +- Online examples (Realtime CLI demo, browser client, barge-in client, + chat-completions client, curl script, Gradio demo): [`examples/online_serving/minicpmo/`](../../examples/online_serving/minicpmo/) +- Duplex API and Python client docs: + [`docs/serving/realtime_duplex_api.md`](../../docs/serving/realtime_duplex_api.md) - Offline end-to-end example: [`examples/offline_inference/minicpmo/`](../../examples/offline_inference/minicpmo/) - Pipeline / talker source: @@ -179,6 +187,13 @@ as one aggregated chunk, so the table reports Stage 0 metrics. #### Verification +**Health and model listing**: + +```bash +curl http://localhost:8099/health +curl http://localhost:8099/v1/models +``` + **Quick smoke test (text-only output)**: ```bash @@ -191,8 +206,31 @@ curl http://localhost:8099/v1/chat/completions \ }' ``` -**Text + speech in one response** (the headline 4.5 feature). The model -bridge conditions the Talker from the generated assistant span, so the +**One duplex turn over the Realtime WebSocket** (text + speech stream back +while the session keeps listening): + +```bash +python examples/online_serving/minicpmo/realtime_duplex_demo.py \ + --url ws://localhost:8099/v1/realtime?duplex=1 \ + --model openbmb/MiniCPM-o-4_5 \ + --input-wav /path/to/input_16k_mono_pcm16.wav \ + --ref-audio /path/to/MiniCPM-o-Demo/assets/ref_audio/ref_minicpm_signature.wav \ + --output-dir /tmp/minicpmo_realtime_duplex_demo +``` + +The demo prints transcript deltas as they arrive and saves the streamed +24 kHz audio to WAV files. The same session shape is available from Python +through `vllm_omni.clients.duplex.DuplexClient` (over the server) or +`vllm_omni.clients.inline_duplex.InlineDuplexClient` (in-process, no server; +`examples/online_serving/barge_in_client.py --inline`). + +**Barge-in**: `examples/online_serving/barge_in_client.py` plays a question, +interrupts the answer with a second WAV and checks that the first response is +cancelled and the second one is answered. + +**Text + speech in one response over `/v1/chat/completions`** (the headline +4.5 feature, served by the ordinary chat service on the duplex engine). The +model bridge conditions the Talker from the generated assistant span, so the generic serving layer does not inject MiniCPM-specific template defaults. `use_tts_template=true` remains supported when explicitly requested: @@ -243,6 +281,12 @@ python examples/online_serving/minicpmo/gradio_demo.py \ Open `http://:7862` and try a text prompt with the **"Generate speech output (TTS)"** checkbox on / off. +**Turn-based generation** is also available offline +(`examples/offline_inference/minicpmo/`, `Omni.generate`). + +The Daily-Omni numbers above were measured on the `/v1/chat/completions` +route. + #### Notes - Memory budget: Thinker, Talker, and Code2Wav reserve 0.55, 0.15, and @@ -326,9 +370,13 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \ and defaults to deterministic seed 42. Stage-1 deploy sampling parameters control only vLLM's binary continue/stop token. -- **Output audio**: 24 kHz mono WAV inside the OpenAI-style - `message.audio.data` (base64). The Gradio demo's WAV player decodes - this automatically. +- **Output audio**: 24 kHz mono. On `/v1/chat/completions` it is base64 WAV + inside the OpenAI-style `message.audio.data` (the Gradio demo's WAV player + decodes this automatically); on a duplex session it streams as base64 + `response.audio.delta` events; offline `Omni.generate` returns WAV. +- **Response choices**: on `/v1/chat/completions` text and audio are separate + choices. SDK clients should select the choice whose `message.audio.data` is + populated rather than assuming `choices[0]` contains audio. - **Routing**: MiniCPM-o 4.5 and 2.6 both ship `architectures= ["MiniCPMO"]` in HF config; routing is disambiguated by @@ -340,6 +388,3 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \ - **Async chunking**: enabled in all deploy configs. Talker sends 25-code chunks with three-code left context to Code2Wav through `SharedMemoryConnector`; terminal chunks flush held lookahead state. -- **Response choices**: text and audio are separate choices. SDK clients - should select the choice whose `message.audio.data` is populated rather - than assuming `choices[0]` contains audio. diff --git a/recipes/README.md b/recipes/README.md index 33e29c086..078466e80 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -71,7 +71,7 @@ recipes/ | [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 | | [`cosmos3/Cosmos3-Edge.md`](./cosmos3/Cosmos3-Edge.md) | T2I / T2V / I2V + action (Physical-AI) generation on the Nemotron-based Edge transformer | 1x GPU (~8 GiB) | | [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 | -| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB | +| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online full-duplex serving and offline inference for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB | | [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB | | [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB | | [`OpenMOSS/MOSS-TTS-Nano-NPU.md`](./OpenMOSS/MOSS-TTS-Nano-NPU.md) | Online multilingual voice-cloned TTS | 1x Ascend 910B3 64GB | diff --git a/tests/benchmarks/patch/test_patch.py b/tests/benchmarks/patch/test_patch.py index 60e071ad5..d2552bbff 100644 --- a/tests/benchmarks/patch/test_patch.py +++ b/tests/benchmarks/patch/test_patch.py @@ -135,6 +135,7 @@ class MockResponse: async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch): class FakeRealtimeClient: last_instance = None + instances: list = [] def __init__(self, url): assert url == "ws://localhost:8000/v1/realtime?duplex=1" @@ -143,6 +144,9 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) self.sent = [] self.response_count = 0 self.ack_count = 0 + self.silence_seconds = [] + self.closed = 0 + type(self).instances.append(self) type(self).last_instance = self async def __aenter__(self): @@ -155,10 +159,22 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) assert model == "openbmb/MiniCPM-o-4_5" self.configure_kwargs = kwargs + async def stream_silence(self, *, seconds, chunk_ms=200, until=None): + # A model-native session speaks off its seeded context once audio + # units arrive; the silence itself carries no content. The probe + # stops the silence as soon as the turn settles. + self.silence_seconds.append(seconds) + self._emit_response() + assert until is not None and until() + return 1.0 + async def send(self, event): self.sent.append(event) if event["type"] != "response.create": return + self._emit_response() + + def _emit_response(self): self.response_count += 1 response_id = f"resp-{self.response_count}" now = time.monotonic() @@ -210,7 +226,7 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) self.ack_count += 1 async def close_session(self, **_kwargs): - return None + self.closed += 1 monkeypatch.setattr( "vllm_omni.benchmarks.patch.patch._RealtimeTTSProbe", @@ -239,30 +255,27 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) session=None, ) - client = FakeRealtimeClient.last_instance - assert client.configure_kwargs["native_duplex"] is False - assert client.configure_kwargs["extra_body"] == { - "ref_audio": "data:audio/wav;base64,AAAA", - "return_stage_metrics": True, - } - assert client.sent == [ - event - for index in range(4) - for event in ( - { - "type": "conversation.item.create", - "item": { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": f"text {index}"}], - }, - }, - {"type": "response.create"}, - ) - ] - assert client.ack_count == 4 + # One session per utterance: a model-native duplex session takes its text + # once, in the session context, so it cannot be re-seeded for a second one. + sessions = FakeRealtimeClient.instances + assert len(sessions) == 4 + for index, client in enumerate(sessions): + assert "native_duplex" not in client.configure_kwargs + assert client.configure_kwargs["extra_body"] == { + "ref_audio": "data:audio/wav;base64,AAAA", + "return_stage_metrics": True, + "duplex_initial_user_text": f"text {index}", + "force_listen_count": 0, + } + # The target text rides the session context, never a conversation item: + # a text-only response.create is rejected by a model-native session. + assert client.sent == [] + assert client.silence_seconds and all(seconds > 0 for seconds in client.silence_seconds) + assert client.ack_count == 1 + assert client.closed == 1 + assert output.success is True - assert output.generated_text == "turn 1 turn 2 turn 3 turn 4" + assert output.generated_text == "turn 1 turn 1 turn 1 turn 1" assert output.audio_duration == pytest.approx(0.4) assert output.ttft > 0 assert output.audio_ttfp > output.ttft @@ -281,11 +294,12 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) "session_id": session_id, "request_index": index, "utterance_id": f"utt-{index}", - "response_id": f"resp-{index + 1}", + # Response ids are per session, and each utterance now gets its own. + "response_id": "resp-1", "source": "client_monotonic_receive", "measurement_origin": { - "ttft": "conversation.item.create client send to first non-empty text delta", - "ttfp": "conversation.item.create client send to first audio packet", + "ttft": "first silence append client send to first non-empty text delta", + "ttfp": "first silence append client send to first audio packet", "rtf": "request-start-to-last-audio receive time divided by emitted audio duration", "tpot": "Stage-0 engine mean time per output token", }, diff --git a/tests/benchmarks/test_omniinteract.py b/tests/benchmarks/test_omniinteract.py index 73cedfb32..1a6acaca5 100644 --- a/tests/benchmarks/test_omniinteract.py +++ b/tests/benchmarks/test_omniinteract.py @@ -467,6 +467,27 @@ def test_tolerated_playback_ack_rejection_is_a_warning_not_a_failure(): oi._raise_if_session_terminated(collector, 0, warnings=warnings) +def test_our_own_close_is_expected_even_though_the_server_stamps_a_reason(): + """``session.close`` is answered with ``session.closed`` carrying ``client_close``. + + The guard exists to catch a session that ended for a reason we did not ask + for. A close we requested is not that, whether or not the server names it. + """ + collector = _collector(({"type": "session.closed", "reason": "client_close"}, 1.0)) + oi._raise_if_session_terminated(collector, 0, explicit_close_from=0) + + collector = _collector(({"type": "session.closed", "event": {"reason": "client_close"}}, 1.0)) + oi._raise_if_session_terminated(collector, 0, explicit_close_from=0) + + collector = _collector(({"type": "session.closed", "reason": "disconnect"}, 1.0)) + with pytest.raises(RuntimeError, match="Unexpected session.closed: disconnect"): + oi._raise_if_session_terminated(collector, 0, explicit_close_from=0) + + collector = _collector(({"type": "session.expired", "reason": "timeout"}, 1.0)) + with pytest.raises(RuntimeError, match="session.expired: timeout"): + oi._raise_if_session_terminated(collector, 0, explicit_close_from=0) + + @pytest.mark.parametrize( ("event", "match"), [ @@ -561,8 +582,8 @@ class _RealtimeClient(oi._RealtimeSession): instances: list[_RealtimeClient] = [] - def __init__(self, config: oi.OmniInteractBenchmarkConfig, session_id: str, reference_audio: str): - super().__init__(config, session_id, reference_audio) + def __init__(self, config: oi.OmniInteractBenchmarkConfig, reference_audio: str): + super().__init__(config, reference_audio) self.acks: list[tuple[str, int]] = [] self.instances.append(self) @@ -621,7 +642,7 @@ async def test_public_runner_executes_one_prepared_session(tmp_path: Path, monke assert "autostart=0" in _RealtimeClient.instances[-1].url session_config = _RealtimeClient.instances[-1].session_config assert session_config.extra_body["custom"] == "value" - assert session_config.extra_body["native_duplex"] is True + assert "native_duplex" not in session_config.extra_body assert session_config.ref_audio == "data:audio/wav;base64,ref" acks = _RealtimeClient.instances[-1].acks # Cumulative acks for the one response: an optional 0 ms checkpoint the diff --git a/tests/benchmarks/test_realtime_tts_probe.py b/tests/benchmarks/test_realtime_tts_probe.py new file mode 100644 index 000000000..c07a09b62 --- /dev/null +++ b/tests/benchmarks/test_realtime_tts_probe.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The Seed-TTS Realtime probe queues for a duplex session slot instead of failing.""" + +from __future__ import annotations + +import pytest + +from vllm_omni.benchmarks.patch import patch as bench_patch +from vllm_omni.clients import duplex as duplex_client + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _AdmissionGate: + """Stands in for ``DuplexClient``: refuses admission ``refusals`` times, then admits.""" + + instances: list[_AdmissionGate] = [] + refusals = 0 + refusal_code = "resource_exhausted" + + def __init__(self, url: str, **kwargs: object) -> None: + del url, kwargs + self.entered = False + type(self).instances.append(self) + + async def __aenter__(self) -> _AdmissionGate: + cls = type(self) + if cls.refusals > 0: + cls.refusals -= 1 + raise duplex_client.DuplexProtocolError("duplex_session_capacity_exhausted: limit=4", code=cls.refusal_code) + self.entered = True + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + del exc_type, exc, tb + + async def events(self): # the collector's consume loop ends at once + return + yield # pragma: no cover + + +@pytest.fixture +def gate(monkeypatch: pytest.MonkeyPatch) -> type[_AdmissionGate]: + _AdmissionGate.instances = [] + _AdmissionGate.refusals = 0 + _AdmissionGate.refusal_code = "resource_exhausted" + monkeypatch.setattr(duplex_client, "DuplexClient", _AdmissionGate) + return _AdmissionGate + + +async def _configure(probe: bench_patch._RealtimeTTSProbe) -> None: + try: + await probe.configure("test-model", auto_response=True) + finally: + await probe.__aexit__(None, None, None) + + +@pytest.mark.asyncio +async def test_configure_waits_for_a_free_session_slot(gate) -> None: + """Two refusals, then a slot: the request is served, not counted as failed.""" + gate.refusals = 2 + probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1") + await _configure(probe) + assert len(gate.instances) == 3 + assert gate.instances[-1].entered is True + assert probe._client is gate.instances[-1] + + +@pytest.mark.asyncio +async def test_configure_raises_other_rejections_at_once(gate) -> None: + gate.refusals = 1 + gate.refusal_code = "unsupported_turn_detection" + probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1") + with pytest.raises(duplex_client.DuplexProtocolError): + await _configure(probe) + assert len(gate.instances) == 1 + + +@pytest.mark.asyncio +async def test_configure_gives_up_when_the_slot_wait_budget_is_spent(gate, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(bench_patch._RealtimeTTSProbe, "_SESSION_SLOT_WAIT_S", 0.0) + gate.refusals = 5 + probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1") + with pytest.raises(duplex_client.DuplexProtocolError) as excinfo: + await _configure(probe) + assert excinfo.value.code == "resource_exhausted" + assert len(gate.instances) == 1 + + +class _SilenceSink: + """Stands in for a configured ``DuplexClient``: records every appended chunk.""" + + def __init__(self) -> None: + self.config = duplex_client.SessionConfig() + self.chunks: list[bytes] = [] + + async def append_audio(self, pcm: bytes, *, is_speech: bool | None = None, video_frames=None) -> None: + del video_frames + assert is_speech is False + self.chunks.append(pcm) + + +@pytest.mark.asyncio +async def test_stream_silence_stops_at_the_turn_end(monkeypatch: pytest.MonkeyPatch) -> None: + """Silence is fed until the turn settles, not for the whole cap.""" + probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1") + sink = _SilenceSink() + probe._client = sink + + async def no_wait(_: float) -> None: + return + + monkeypatch.setattr(bench_patch.asyncio, "sleep", no_wait) + streamed = await probe.stream_silence(seconds=12.0, chunk_ms=200, until=lambda: len(sink.chunks) >= 3) + assert len(sink.chunks) == 3 + assert streamed == pytest.approx(0.6) + assert all(len(chunk) == sink.config.input_audio.byte_count(200) for chunk in sink.chunks) + + +@pytest.mark.asyncio +async def test_stream_silence_runs_to_the_cap_when_the_turn_never_settles(monkeypatch: pytest.MonkeyPatch) -> None: + probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1") + sink = _SilenceSink() + probe._client = sink + + async def no_wait(_: float) -> None: + return + + monkeypatch.setattr(bench_patch.asyncio, "sleep", no_wait) + streamed = await probe.stream_silence(seconds=1.0, chunk_ms=200, until=lambda: False) + assert len(sink.chunks) == 5 + assert streamed == pytest.approx(1.0) + + +class _TurnEvents: + def __init__(self, response_ids: list[str], audio: dict[str, bytes], text: dict[str, str] | None = None) -> None: + self.response_ids = response_ids + self._audio = audio + self._text = text or {} + + def audio_bytes(self, response_id: str) -> bytes: + return self._audio.get(response_id, b"") + + def response_text(self, response_id: str) -> str: + return self._text.get(response_id, "") + + +def test_turn_response_is_the_first_one_with_audio_even_if_the_model_spoke_again(caplog) -> None: + events = _TurnEvents(["r0", "r1", "r2"], {"r1": b"\x01\x00", "r2": b"\x02\x00"}) + with caplog.at_level("WARNING", logger=bench_patch.logger.name): + assert bench_patch._seed_tts_turn_response_id(events, 0, 0) == "r1" + assert "spoke again" in caplog.text + + +def test_turn_response_ignores_responses_from_before_the_turn() -> None: + events = _TurnEvents(["r0", "r1"], {"r0": b"\x01\x00", "r1": b"\x02\x00"}) + assert bench_patch._seed_tts_turn_response_id(events, 1, 0) == "r1" + + +def test_turn_without_audio_is_an_error() -> None: + events = _TurnEvents(["r0"], {}) + with pytest.raises(RuntimeError, match="no audio response"): + bench_patch._seed_tts_turn_response_id(events, 0, 0) + + +def test_stall_report_says_the_model_never_answered() -> None: + events = _TurnEvents(["r0"], {"r0": b"\x01\x00"}) + report = bench_patch._seed_tts_turn_stall_report(events, 1, 3, 30.0) + assert "turn 3 never started a response" in report + assert "30.0s of silence" in report + + +def test_stall_report_lists_a_response_that_never_finished() -> None: + events = _TurnEvents(["r1"], {"r1": b"\x01\x00\x02\x00"}, {"r1": "hello"}) + report = bench_patch._seed_tts_turn_stall_report(events, 0, 0, 4.2) + assert "started 1 response(s) after 4.2s of silence" in report + assert "r1 (4 audio bytes, text 'hello')" in report diff --git a/tests/clients/test_duplex_client.py b/tests/clients/test_duplex_client.py index ab2a38951..43b7547cb 100644 --- a/tests/clients/test_duplex_client.py +++ b/tests/clients/test_duplex_client.py @@ -1,6 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Unit tests for the public duplex client (fake transport, no server).""" +"""Unit tests for the public duplex clients (fake transports, no server). + +``DuplexClient`` runs over a fake websocket; ``InlineDuplexClient`` over a +fake ``DuplexOmni`` whose handle records typed commands and replays typed +events. The demux / handshake tests that only need "feed a server event" are +parameterized over both clients. +""" from __future__ import annotations @@ -8,6 +14,8 @@ import asyncio import base64 import json import wave +from collections.abc import Callable +from typing import Any from urllib.parse import parse_qs, urlsplit import pytest @@ -16,8 +24,11 @@ from vllm_omni.clients.duplex import ( AudioFormat, ConnectionResumed, DuplexClient, + DuplexClientBase, + DuplexConnectionError, DuplexProtocolError, DuplexSessionClosedError, + ErrorEvent, EventCollector, ReconnectPolicy, SessionConfig, @@ -31,17 +42,26 @@ from vllm_omni.clients.duplex import ( summarize_session_request_metrics, write_pcm16_wav, ) +from vllm_omni.clients.inline_duplex import InlineDuplexClient +from vllm_omni.engine.duplex import commands as duplex_commands +from vllm_omni.engine.duplex import events as duplex_events +from vllm_omni.engine.duplex.messages import DuplexSessionError pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +SESSION_ID = "duplex-0123456789abcdef" +# session.created is not journaled by the server, so it carries no server_event_seq. SESSION_CREATED = { "type": "session.created", - "session": {"session_id": "sess-1"}, - "incarnation": 0, + "session": {"id": SESSION_ID, "model": "test-model"}, + "attachment_generation": 1, "resume_token": "tok-1", - "server_event_seq": 1, } -SESSION_CLOSED = {"type": "session.closed", "session_id": "sess-1", "server_event_seq": 99} +SESSION_CLOSED = {"type": "session.closed", "session_id": SESSION_ID, "reason": "client_close", "server_event_seq": 99} + + +# --------------------------------------------------------------------------- +# Fake websocket transport class FakeSocket: @@ -88,14 +108,139 @@ def make_client(*sockets: FakeSocket, **kwargs) -> tuple[DuplexClient, list[str] return client, calls +# --------------------------------------------------------------------------- +# Fake DuplexOmni for the inline client + + +class _WireEvent: + """A typed-event stand-in that renders a fixed wire payload.""" + + def __init__(self, raw: dict[str, object]) -> None: + self.raw = raw + + @property + def is_terminal(self) -> bool: + return self.raw.get("type") in {"session.closed", "session.expired"} + + def to_realtime(self) -> dict[str, object]: + return dict(self.raw) + + +class FakeInlineHandle: + def __init__(self, session_id: str) -> None: + self.session_id = session_id + self.closed = False + self.close_calls = 0 + self.commands: list[duplex_commands.DuplexCommand] = [] + self.consumer_active = False + self._outbox: asyncio.Queue = asyncio.Queue() + + def deliver(self, event: Any) -> None: + self._outbox.put_nowait(event) + + def feed(self, raw: dict[str, object]) -> None: + self.deliver(_WireEvent(raw)) + + async def submit(self, command: duplex_commands.DuplexCommand) -> None: + if self.closed: + raise DuplexSessionError("closed", code="session_closed", session_id=self.session_id) + self.commands.append(command) + + async def events(self): + self.consumer_active = True + try: + while True: + event = await self._outbox.get() + yield event + if event.is_terminal: + return + finally: + self.consumer_active = False + + async def close(self, *, reason: str = "client_close", timeout: float | None = None) -> None: + self.close_calls += 1 + if self.closed: + return + self.closed = True + self.deliver(duplex_events.SessionClosed(session_id=self.session_id, reason=reason)) + + +class FakeInlineOmni: + def __init__(self) -> None: + self.opened: list[dict[str, object]] = [] + self.handles: list[FakeInlineHandle] = [] + + async def open_session(self, config: dict[str, object]) -> FakeInlineHandle: + self.opened.append(dict(config)) + handle = FakeInlineHandle(SESSION_ID) + self.handles.append(handle) + handle.deliver( + duplex_events.SessionCreated( + session_id=SESSION_ID, + session={"id": SESSION_ID, "model": config.get("model")}, + ) + ) + return handle + + +def make_inline_client(**kwargs) -> tuple[InlineDuplexClient, FakeInlineOmni]: + omni = FakeInlineOmni() + kwargs.setdefault("handshake_timeout_s", 5.0) + return InlineDuplexClient(omni, model="test-model", **kwargs), omni + + +# --------------------------------------------------------------------------- +# One rig over both clients: ``feed`` a server event, drive ``client`` + + +class Rig: + def __init__(self, client: DuplexClientBase, feed: Callable[[dict[str, object]], None]) -> None: + self.client = client + self.feed = feed + + +def _websocket_rig(*, handshake: dict[str, object] | None = SESSION_CREATED) -> Rig: + sock = FakeSocket() + if handshake is not None: + sock.feed(handshake) + client, _ = make_client(sock) + return Rig(client, sock.feed) + + +def _inline_rig(*, handshake: dict[str, object] | None = SESSION_CREATED) -> Rig: + client, omni = make_inline_client() + original_open = omni.open_session + + async def open_session(config: dict[str, object]) -> FakeInlineHandle: + handle = await original_open(config) + if handshake is not SESSION_CREATED: + # Replace the default session.created with the requested first event. + handle._outbox = asyncio.Queue() + if handshake is not None: + handle.feed(handshake) + return handle + + omni.open_session = open_session # type: ignore[method-assign] + + def feed(raw: dict[str, object]) -> None: + omni.handles[-1].feed(raw) + + return Rig(client, feed) + + +@pytest.fixture(params=["websocket", "inline"]) +def make_rig(request: pytest.FixtureRequest) -> Callable[..., Rig]: + return _websocket_rig if request.param == "websocket" else _inline_rig + + # --------------------------------------------------------------------------- # SessionConfig def test_session_config_payload_defaults(): - payload = SessionConfig().to_session_payload(model="m", session_id="s") + payload = SessionConfig().to_session_payload(model="m") assert payload["model"] == "m" - assert payload["session_id"] == "s" + assert "session_id" not in payload assert payload["modalities"] == ["audio", "text"] assert payload["input_audio_format"] == "pcm16" assert payload["output_audio_format"] == "pcm16" @@ -137,19 +282,59 @@ def test_audio_format_math(): # --------------------------------------------------------------------------- -# Handshake and lifecycle +# Handshake and lifecycle (both clients) -async def test_handshake_adopts_session_and_acks(): +@pytest.mark.asyncio +async def test_handshake_adopts_server_allocated_session_id(make_rig): + rig = make_rig() + async with rig.client as client: + assert client.session_id == SESSION_ID + assert client.session_info == {"id": SESSION_ID, "model": "test-model"} + rig.feed(SESSION_CLOSED) + + +@pytest.mark.asyncio +async def test_handshake_error_raises_protocol_error(make_rig): + rig = make_rig(handshake={"type": "error", "error": {"code": "unsupported_audio_format", "message": "bad"}}) + with pytest.raises(DuplexProtocolError) as excinfo: + async with rig.client: + pass + assert excinfo.value.code == "unsupported_audio_format" + + +@pytest.mark.asyncio +async def test_session_expired_raises_from_event_stream(make_rig): + rig = make_rig() + async with rig.client as client: + rig.feed( + {"type": "session.expired", "session_id": SESSION_ID, "reason": "lease_expired", "server_event_seq": 2} + ) + with pytest.raises(DuplexSessionClosedError) as excinfo: + async for _ in client.events(): + pass + assert "lease_expired" in excinfo.value.reason + + +# --------------------------------------------------------------------------- +# Websocket transport specifics + + +@pytest.mark.asyncio +async def test_websocket_handshake_sends_session_update_and_acks_sequenced_events(): sock = FakeSocket() sock.feed(SESSION_CREATED) client, calls = make_client(sock) async with client: - assert client.session_id == "sess-1" + assert client.session_id == SESSION_ID assert client.resume_token == "tok-1" assert calls == ["ws://test-host:8099/v1/realtime?duplex=1&model=test-model&autostart=0"] assert sock.sent[0]["type"] == "session.update" assert sock.sent[0]["session"]["model"] == "test-model" + assert "session_id" not in sock.sent[0]["session"] + # session.created is unjournaled (no seq): nothing to ack yet. + assert _acks(sock) == [] + sock.feed({"type": "response.listen", "server_event_seq": 1}) await _drain(lambda: {"type": "session.event_ack", "server_event_seq": 1} in _acks(sock)) close_task = asyncio.create_task(client.close()) await _drain(lambda: "session.close" in sock.sent_types()) @@ -158,28 +343,7 @@ async def test_handshake_adopts_session_and_acks(): assert "session.close" in sock.sent_types() -async def test_handshake_error_raises_protocol_error(): - sock = FakeSocket() - sock.feed({"type": "error", "error": {"code": "unsupported_audio_format", "message": "bad"}}) - client, _ = make_client(sock) - with pytest.raises(DuplexProtocolError) as excinfo: - async with client: - pass - assert excinfo.value.code == "unsupported_audio_format" - - -async def test_session_expired_raises_from_event_stream(): - sock = FakeSocket() - sock.feed(SESSION_CREATED) - client, _ = make_client(sock) - async with client: - sock.feed({"type": "session.expired", "reason": "lease_expired", "server_event_seq": 2}) - with pytest.raises(DuplexSessionClosedError) as excinfo: - async for _ in client.events(): - pass - assert "lease_expired" in excinfo.value.reason - - +@pytest.mark.asyncio async def test_resync_required_drops_resume_credential(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -187,7 +351,7 @@ async def test_resync_required_drops_resume_credential(): async with client: assert client.resume_token == "tok-1" stream = client.events() - sock.feed({"type": "session.resync_required", "session_id": "sess-1", "server_event_seq": 2}) + sock.feed({"type": "session.resync_required", "session_id": SESSION_ID, "server_event_seq": 2}) async for event in stream: if event.type == "session.resync_required": break @@ -202,9 +366,10 @@ async def test_resync_required_drops_resume_credential(): # --------------------------------------------------------------------------- -# Input events +# Input events (websocket wire shape) +@pytest.mark.asyncio async def test_append_audio_tracks_cumulative_end_ms(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -224,6 +389,7 @@ async def test_append_audio_tracks_cumulative_end_ms(): sock.feed(SESSION_CLOSED) +@pytest.mark.asyncio async def test_append_audio_strips_video_frame_data_url_prefix(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -238,6 +404,7 @@ async def test_append_audio_strips_video_frame_data_url_prefix(): sock.feed(SESSION_CLOSED) +@pytest.mark.asyncio async def test_stream_pcm_chunking(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -249,6 +416,7 @@ async def test_stream_pcm_chunking(): sock.feed(SESSION_CLOSED) +@pytest.mark.asyncio async def test_interruption_primitives_send_documented_events(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -262,22 +430,21 @@ async def test_interruption_primitives_send_documented_events(): # --------------------------------------------------------------------------- -# Response demultiplexing +# Response demultiplexing (both clients) def _b64(data: bytes) -> str: return base64.b64encode(data).decode("ascii") -async def test_response_handle_flow(): - sock = FakeSocket() - sock.feed(SESSION_CREATED) +@pytest.mark.asyncio +async def test_response_handle_flow(make_rig): + rig = make_rig() chunk = b"\x00\x01" * 2400 # 100 ms of pcm16 @ 24 kHz - client, _ = make_client(sock) - async with client: - sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2}) - sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3}) - sock.feed( + async with rig.client as client: + rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2}) + rig.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3}) + rig.feed( { "type": "response.output_audio.delta", "response_id": "resp-1", @@ -286,7 +453,7 @@ async def test_response_handle_flow(): "server_event_seq": 4, } ) - sock.feed( + rig.feed( { "type": "response.output_audio_transcript.delta", "response_id": "resp-1", @@ -294,7 +461,7 @@ async def test_response_handle_flow(): "server_event_seq": 5, } ) - sock.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 6}) + rig.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 6}) async for response in client.responses(): chunks = [piece async for piece in response.audio()] @@ -304,49 +471,46 @@ async def test_response_handle_flow(): assert response.played_ms == pytest.approx(100.0) assert response.finished break - sock.feed(SESSION_CLOSED) + rig.feed(SESSION_CLOSED) -async def test_listen_decision_yields_finished_silent_handle(): - sock = FakeSocket() - sock.feed(SESSION_CREATED) - client, _ = make_client(sock) - async with client: - sock.feed({"type": "response.listen", "server_event_seq": 2}) +@pytest.mark.asyncio +async def test_listen_decision_yields_finished_silent_handle(make_rig): + rig = make_rig() + async with rig.client as client: + rig.feed({"type": "response.listen", "server_event_seq": 2}) async for response in client.responses(): assert response.decision == "listen" assert response.finished assert [piece async for piece in response.audio()] == [] break - sock.feed(SESSION_CLOSED) + rig.feed(SESSION_CLOSED) -async def test_listen_terminates_active_response(): - sock = FakeSocket() - sock.feed(SESSION_CREATED) - client, _ = make_client(sock) - async with client: - sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2}) - sock.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 3}) +@pytest.mark.asyncio +async def test_listen_terminates_active_response(make_rig): + rig = make_rig() + async with rig.client as client: + rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2}) + rig.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 3}) async for response in client.responses(): await response.wait(timeout_s=5.0) assert response.decision == "listen" break - sock.feed(SESSION_CLOSED) + rig.feed(SESSION_CLOSED) -async def test_id_less_listen_never_closes_an_in_flight_response(): +@pytest.mark.asyncio +async def test_id_less_listen_never_closes_an_in_flight_response(make_rig): # An id-less listen is a standalone decision beat (e.g. a silence-skip # while a response streams); it must surface as its own finished handle # and leave the in-flight response to its real terminal. - sock = FakeSocket() - sock.feed(SESSION_CREATED) - client, _ = make_client(sock) - async with client: - sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2}) - sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3}) - sock.feed({"type": "response.listen", "server_event_seq": 4}) - sock.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 5}) + rig = make_rig() + async with rig.client as client: + rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2}) + rig.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3}) + rig.feed({"type": "response.listen", "server_event_seq": 4}) + rig.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 5}) seen: list[tuple[str | None, str | None, bool]] = [] async for response in client.responses(): await response.wait(timeout_s=5.0) @@ -355,20 +519,19 @@ async def test_id_less_listen_never_closes_an_in_flight_response(): break assert ("resp-1", "speak", True) in seen assert (None, "listen", True) in seen - sock.feed(SESSION_CLOSED) + rig.feed(SESSION_CLOSED) -async def test_terminal_listen_keeps_spoken_decision(): +@pytest.mark.asyncio +async def test_terminal_listen_keeps_spoken_decision(make_rig): # A terminal listen after the response spoke is the model yielding the # turn; the handle must stay decision="speak" with its audio intact. - sock = FakeSocket() - sock.feed(SESSION_CREATED) + rig = make_rig() chunk = b"\x00\x01" * 2400 - client, _ = make_client(sock) - async with client: - sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2}) - sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3}) - sock.feed( + async with rig.client as client: + rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2}) + rig.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3}) + rig.feed( { "type": "response.output_audio.delta", "response_id": "resp-1", @@ -377,24 +540,23 @@ async def test_terminal_listen_keeps_spoken_decision(): "server_event_seq": 4, } ) - sock.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 5}) + rig.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 5}) async for response in client.responses(): chunks = [piece async for piece in response.audio()] assert chunks == [chunk] assert response.decision == "speak" assert response.finished break - sock.feed(SESSION_CLOSED) + rig.feed(SESSION_CLOSED) -async def test_error_event_surfaces_on_response_path(): +@pytest.mark.asyncio +async def test_error_event_surfaces_on_response_path(make_rig): # A rejected send produces no response; a consumer waiting in responses() # must see the rejection instead of waiting forever. - sock = FakeSocket() - sock.feed(SESSION_CREATED) - client, _ = make_client(sock) - async with client: - sock.feed( + rig = make_rig() + async with rig.client as client: + rig.feed( { "type": "error", "error": {"code": "input_audio_buffer_empty", "message": "empty commit"}, @@ -405,9 +567,10 @@ async def test_error_event_surfaces_on_response_path(): async for _ in client.responses(): pass assert excinfo.value.code == "input_audio_buffer_empty" - sock.feed(SESSION_CLOSED) + rig.feed(SESSION_CLOSED) +@pytest.mark.asyncio async def test_slow_audio_consumer_drops_oldest_instead_of_stalling(): from vllm_omni.clients.duplex import AudioDelta, ResponseHandle @@ -425,24 +588,24 @@ async def test_slow_audio_consumer_drops_oldest_instead_of_stalling(): # --------------------------------------------------------------------------- -# Resume +# Resume (websocket only) +@pytest.mark.asyncio async def test_resume_after_transport_drop(): first = FakeSocket() first.feed(SESSION_CREATED) second = FakeSocket() # Mirror the real resume activation payload # (entrypoints/duplex/serving.py, activation_payload_factory): it carries - # no nested "session" object, only the identity fields. + # the identity fields plus the (empty) public session object. second.feed( { "type": "session.resumed", - "session_id": "sess-1", - "incarnation": 0, - "attachment_generation": 1, + "session_id": SESSION_ID, + "session": {}, + "attachment_generation": 2, "resume_token": "tok-2", - "server_event_seq": 5, } ) client, calls = make_client( @@ -453,34 +616,39 @@ async def test_resume_after_transport_drop(): async with client: stream = client.events() - async def take_two(): + async def collect_until_resumed(): seen = [] async for event in stream: seen.append(event) - if len(seen) == 2: + if isinstance(event, SessionResumed): return seen - consumer = asyncio.create_task(take_two()) + consumer = asyncio.create_task(collect_until_resumed()) await asyncio.sleep(0) # let the consumer subscribe before the drop await asyncio.sleep(0) + first.feed({"type": "response.listen", "server_event_seq": 7}) + await _drain(lambda: {"type": "session.event_ack", "server_event_seq": 7} in _acks(first)) first.feed(RuntimeError("transport dropped")) seen = await asyncio.wait_for(consumer, timeout=5.0) - assert isinstance(seen[0], ConnectionResumed) - assert isinstance(seen[1], SessionResumed) + assert [type(event) for event in seen[-2:]] == [ConnectionResumed, SessionResumed] assert client.resume_token == "tok-2" + # The superseded transport is closed, not leaked. + assert first.closed # The activation payload has no session object; the info captured at # the original handshake must survive the resume. - assert client.session_info == {"session_id": "sess-1"} + assert client.session_info == {"id": SESSION_ID, "model": "test-model"} assert len(calls) == 2 resume = second.sent[0] - assert resume["type"] == "session.resume" - assert resume["session_id"] == "sess-1" - assert resume["incarnation"] == 0 - assert resume["resume_token"] == "tok-1" - assert resume["last_received_server_event_seq"] == 1 + assert resume == { + "type": "session.resume", + "session_id": SESSION_ID, + "resume_token": "tok-1", + "last_received_server_event_seq": 7, + } second.feed(SESSION_CLOSED) +@pytest.mark.asyncio async def test_no_reconnect_policy_surfaces_closed(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -493,6 +661,7 @@ async def test_no_reconnect_policy_surfaces_closed(): pass +@pytest.mark.asyncio async def test_heartbeat_survives_transient_send_failure(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -514,6 +683,7 @@ async def test_heartbeat_survives_transient_send_failure(): sock.feed(SESSION_CLOSED) +@pytest.mark.asyncio async def test_aexit_sends_session_close_on_error_path(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -534,6 +704,7 @@ def test_target_url_forces_autostart_off(): assert query["autostart"] == ["0"] +@pytest.mark.asyncio async def test_fatal_resume_error_gives_up(): first = FakeSocket() first.feed(SESSION_CREATED) @@ -553,6 +724,121 @@ async def test_fatal_resume_error_gives_up(): assert len(calls) == 2 # no retry after a fatal resume error +# --------------------------------------------------------------------------- +# Inline client over a fake DuplexOmni + + +@pytest.mark.asyncio +async def test_inline_handshake_opens_session_with_payload_and_adopts_id(): + client, omni = make_inline_client(config=SessionConfig(ref_audio="data:audio/wav;base64,AAA=")) + async with client: + assert omni.opened == [client.config.to_session_payload(model="test-model")] + assert omni.opened[0]["ref_audio"] == "data:audio/wav;base64,AAA=" + assert client.session_id == SESSION_ID + handle = omni.handles[0] + # The session was opened with the full session object already; the + # base handshake's session.update must not be re-submitted as a patch + # (a ref_audio patch would be rejected by the runner). + assert handle.commands == [] + await client.send({"type": "session.update", "session": {"instructions": "later"}}) + assert [type(command) for command in handle.commands] == [duplex_commands.UpdateSession] + assert handle.commands[0].patch == {"instructions": "later"} + + +@pytest.mark.asyncio +async def test_inline_client_submits_typed_commands_in_order(): + client, omni = make_inline_client() + async with client: + handle = omni.handles[0] + pcm = b"\x01\x00" * 1600 + await client.append_audio(pcm, is_speech=True) + await client.commit(create_response=True) + await client.cancel_response("resp-1") + await client.ack_playback(120.9, response_id="resp-1") + assert [type(command) for command in handle.commands] == [ + duplex_commands.AppendAudio, + duplex_commands.Commit, + duplex_commands.CancelResponse, + duplex_commands.AckPlayback, + ] + append, commit, cancel, ack = handle.commands + assert append.is_speech is True + assert append.audio_end_ms == 100 + assert append.event_id + assert commit.final is True and commit.create_response is True + assert cancel.response_id == "resp-1" + assert ack.played_ms == 120 and ack.response_id == "resp-1" + + +@pytest.mark.asyncio +async def test_inline_session_close_maps_to_handle_close(): + client, omni = make_inline_client() + async with client: + pass + handle = omni.handles[0] + assert handle.closed + assert handle.close_calls >= 1 + assert not any(isinstance(command, duplex_commands.CloseSession) for command in handle.commands) + # The typed session.closed reached the client as the expected end. + with pytest.raises(DuplexSessionClosedError): + await client.send({"type": "session.heartbeat"}) + + +@pytest.mark.asyncio +async def test_inline_client_forwards_heartbeats_and_drops_event_acks(): + client, omni = make_inline_client() + async with client: + handle = omni.handles[0] + await client.send({"type": "session.heartbeat"}) + await client.send({"type": "session.event_ack", "server_event_seq": 3}) + # A heartbeat is a lease touch the engine must see; the ack is a journal concern. + assert [type(command).__name__ for command in handle.commands] == ["Heartbeat"] + + +@pytest.mark.asyncio +async def test_inline_malformed_payload_surfaces_as_error_event(): + client, omni = make_inline_client() + async with client: + handle = omni.handles[0] + # Subscribe before sending: the translation error is dispatched inline. + first_event = asyncio.create_task(client.wait_for("error", timeout_s=5.0)) + await asyncio.sleep(0) + event_id = await client.send({"type": "playback.ack"}) # played_ms missing + event = await first_event + assert isinstance(event, ErrorEvent) + assert event.code == "bad_event" + assert event.related_event_id == event_id + assert handle.commands == [] + with pytest.raises(DuplexProtocolError) as excinfo: + async for _ in client.responses(): + pass + assert excinfo.value.code == "bad_event" + + +@pytest.mark.asyncio +async def test_inline_teardown_cancels_pump_and_closes_handle(): + client, omni = make_inline_client() + async with client: + handle = omni.handles[0] + await _drain(lambda: handle.consumer_active) + pump = client._pump_task + assert pump is not None and not pump.done() + assert pump.done() + assert not handle.consumer_active + assert handle.closed + + +@pytest.mark.asyncio +async def test_inline_send_after_handle_closed_raises_connection_error(): + client, omni = make_inline_client() + async with client: + handle = omni.handles[0] + handle.closed = True # the engine ended the session without an event yet + with pytest.raises(DuplexConnectionError, match="not open"): + await client.send({"type": "input_audio_buffer.clear"}) + handle.closed = False + + # --------------------------------------------------------------------------- # Collector @@ -600,8 +886,7 @@ def test_build_realtime_url_with_model_extra_query(): url = build_realtime_url( "ws://localhost:8099/v1/realtime?custom=1", "openbmb/MiniCPM-o-4_5", - session_id="session-a", - extra_query={"native_duplex": "1"}, + extra_query={"probe": "1"}, ) query = parse_qs(urlsplit(url).query) @@ -609,8 +894,7 @@ def test_build_realtime_url_with_model_extra_query(): "custom": ["1"], "duplex": ["1"], "model": ["openbmb/MiniCPM-o-4_5"], - "native_duplex": ["1"], - "session_id": ["session-a"], + "probe": ["1"], } @@ -619,12 +903,11 @@ def test_build_realtime_url_resume_only_when_autostart_disabled(): "ws://localhost:8099/v1/realtime?duplex=1", "openbmb/MiniCPM-o-4_5", autostart=False, - extra_query={"native_duplex": "1"}, ) query = parse_qs(urlsplit(url).query) assert query["autostart"] == ["0"] - assert query["native_duplex"] == ["1"] + assert "session_id" not in query def test_event_collector_partitions_audio_by_response(): @@ -989,6 +1272,7 @@ def test_duplex_unit_boundary_and_residual_math(): assert chunk_period_ms([{"type": "session.created", "session": {}}]) == 1000 +@pytest.mark.asyncio async def test_stream_pcm_sends_each_units_composite_beside_its_base_frame(): sock = FakeSocket() sock.feed(SESSION_CREATED) @@ -1011,16 +1295,16 @@ async def test_stream_pcm_sends_each_units_composite_beside_its_base_frame(): sock.feed(SESSION_CLOSED) -def test_build_realtime_url_native_duplex_flag_and_http_scheme(): - url = build_realtime_url("http://localhost:8099/v1/realtime", None, native_duplex=None) +def test_build_realtime_url_rewrites_http_scheme_and_rejects_others(): + url = build_realtime_url("http://localhost:8099/v1/realtime", None) parts = urlsplit(url) assert parts.scheme == "ws" assert parse_qs(parts.query) == {"duplex": ["1"]} - url = build_realtime_url("https://host/v1/realtime", "m", native_duplex=False) + url = build_realtime_url("https://host/v1/realtime", "m", autostart=True) parts = urlsplit(url) assert parts.scheme == "wss" - assert parse_qs(parts.query)["native_duplex"] == ["0"] + assert parse_qs(parts.query) == {"duplex": ["1"], "model": ["m"], "autostart": ["1"]} with pytest.raises(ValueError): build_realtime_url("ftp://host/v1/realtime", "m") diff --git a/tests/clients/test_model_session_configs.py b/tests/clients/test_model_session_configs.py index 00741d3c6..080fea879 100644 --- a/tests/clients/test_model_session_configs.py +++ b/tests/clients/test_model_session_configs.py @@ -18,13 +18,13 @@ from vllm_omni.clients.personaplex import ( pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def test_minicpmo45_session_config_matches_native_duplex_deployment(): +def test_minicpmo45_session_config_matches_duplex_deployment(): config = create_minicpmo45_session_config(ref_audio="data:audio/wav;base64,AAA=", temperature=0.0) payload = config.to_session_payload(model="openbmb/MiniCPM-o-4_5") assert payload["ref_audio"] == "data:audio/wav;base64,AAA=" assert payload["overlap_policy"] == "listen_only" assert payload["playback_commit_policy"] == "ack_only" - assert payload["extra_body"]["native_duplex"] is True + assert "native_duplex" not in payload["extra_body"] assert payload["extra_body"]["auto_response"] is True assert payload["temperature"] == 0.0 @@ -41,8 +41,9 @@ def test_personaplex_session_config_matches_deployment(): @pytest.mark.parametrize( "duplex_package", [ + # PersonaPlex and Nemotron VoiceChat rejoin this list with the follow-up + # PRs that port them to the duplex plugin framework. "vllm_omni.model_executor.models.minicpmo_4_5.duplex", - "vllm_omni.model_executor.models.personaplex.duplex", ], ) def test_model_duplex_packages_stay_clear_of_client_library(duplex_package: str): diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5.json b/tests/dfx/perf/tests/test_minicpmo_4_5.json index 7d7ca08ad..b384d61a0 100644 --- a/tests/dfx/perf/tests/test_minicpmo_4_5.json +++ b/tests/dfx/perf/tests/test_minicpmo_4_5.json @@ -26,8 +26,8 @@ { "dataset_name": "seed-tts", "dataset_path": "zhaochenyang20/seed-tts-eval", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", + "backend": "openai-realtime-tts", + "endpoint": "/v1/realtime", "num_prompts": [ 32, 64, @@ -42,73 +42,7 @@ "disable_shuffle": true, "trust_remote_code": true, "seed_tts_locale": "en", - "extra_body": { - "modalities": [ - "text", - "audio" - ], - "chat_template_kwargs": { - "enable_thinking": false, - "use_tts_template": true - } - }, - "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration", - "baseline": { - "H100": { - "request_throughput": [ - 0.9731, - 1.1986, - 1.686 - ], - "mean_ttft_ms": [ - 152.5665, - 270.2559, - 301.455 - ], - "mean_e2el_ms": [ - 1027.2953, - 3306.3735, - 4686.5366 - ], - "mean_audio_ttfp_ms": [ - 458.1553, - 2030.2335, - 2068.6233 - ], - "mean_audio_rtf": [ - 0.2357, - 0.791, - 1.0324 - ] - }, - "A3": { - "request_throughput": [ - 0.5383, - 0.6042, - 0.7547 - ], - "mean_ttft_ms": [ - 333.2633, - 533.8656, - 547.3388 - ], - "mean_e2el_ms": [ - 1857.2154, - 6600.5101, - 10450.6811 - ], - "mean_audio_ttfp_ms": [ - 986.4666, - 3411.0754, - 3352.7515 - ], - "mean_audio_rtf": [ - 0.4423, - 1.5734, - 2.3024 - ] - } - } + "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration" } ] }, @@ -138,8 +72,8 @@ { "dataset_name": "seed-tts", "dataset_path": "zhaochenyang20/seed-tts-eval", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", + "backend": "openai-realtime-tts", + "endpoint": "/v1/realtime", "num_prompts": [ 32, 64, @@ -154,16 +88,6 @@ "disable_shuffle": true, "trust_remote_code": true, "seed_tts_locale": "en", - "extra_body": { - "modalities": [ - "text", - "audio" - ], - "chat_template_kwargs": { - "enable_thinking": false, - "use_tts_template": true - } - }, "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration" } ] diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json b/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json index bb0599419..a145459ed 100644 --- a/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json +++ b/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json @@ -43,43 +43,7 @@ "extra_body": { "save_duplex_request_metrics": true }, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration", - "baseline": { - "H100": { - "request_throughput": [ - 0.3038 - ], - "mean_ttft_ms": [ - 139.997 - ], - "mean_e2el_ms": [ - 3269.8807 - ], - "mean_audio_ttfp_ms": [ - 456.073 - ], - "mean_audio_rtf": [ - 0.2254 - ] - }, - "A3": { - "request_throughput": [ - 0.1594 - ], - "mean_ttft_ms": [ - 274.164 - ], - "mean_e2el_ms": [ - 6250.7696 - ], - "mean_audio_ttfp_ms": [ - 994.422 - ], - "mean_audio_rtf": [ - 0.46163 - ] - } - } + "percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration" } ] } diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json b/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json index 26dc54aab..4fd78d681 100644 --- a/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json +++ b/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json @@ -26,8 +26,8 @@ { "dataset_name": "seed-tts", "dataset_path": "zhaochenyang20/seed-tts-eval", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", + "backend": "openai-realtime-tts", + "endpoint": "/v1/realtime", "num_prompts": [ 10, 40 @@ -40,41 +40,7 @@ "disable_shuffle": true, "trust_remote_code": true, "seed_tts_locale": "en", - "extra_body": { - "modalities": [ - "text", - "audio" - ], - "chat_template_kwargs": { - "enable_thinking": false, - "use_tts_template": true - } - }, - "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration", - "baseline": { - "H100": { - "request_throughput": [ - 0.9731, - 1.1986 - ], - "mean_ttft_ms": [ - 152.5665, - 270.2559 - ], - "mean_e2el_ms": [ - 1027.2953, - 3306.3735 - ], - "mean_audio_ttfp_ms": [ - 458.1553, - 2030.2335 - ], - "mean_audio_rtf": [ - 0.2357, - 0.791 - ] - } - } + "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration" } ] } diff --git a/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py b/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py index c833ae7a1..e0b0a6c53 100644 --- a/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py +++ b/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py @@ -39,7 +39,8 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import ( from vllm_omni.engine.async_omni_engine import StageRuntimeInfo from vllm_omni.engine.messages import ShutdownRequestMessage, StageSubmissionMessage from vllm_omni.engine.stage_init_utils import StageMetadata -from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.async_omni_base import AsyncEventResolver from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput diff --git a/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py b/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py index 2e5ceb8e2..7780f0475 100644 --- a/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py +++ b/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py @@ -42,7 +42,8 @@ from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState from vllm_omni.engine.stage_client import StagePoolClient from vllm_omni.engine.stage_init_utils import StageMetadata from vllm_omni.engine.stage_pool import StagePool -from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.async_omni_base import AsyncEventResolver from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniInteractionPrompt from vllm_omni.outputs import OmniRequestOutput diff --git a/tests/diffusion/test_diffusion_engine_metrics.py b/tests/diffusion/test_diffusion_engine_metrics.py index 1dad9face..ec222bcd3 100644 --- a/tests/diffusion/test_diffusion_engine_metrics.py +++ b/tests/diffusion/test_diffusion_engine_metrics.py @@ -90,6 +90,7 @@ def _get_function_source(source: str, class_name: str | None, func_name: str) -> Args: source: Full file source code. class_name: Enclosing class name, or None for module-level functions. + The method may be inherited from a base class defined in the same file. func_name: Function/method name. Returns: @@ -97,13 +98,24 @@ def _get_function_source(source: str, class_name: str | None, func_name: str) -> """ tree = ast.parse(source) if class_name is not None: - for node in ast.walk(tree): - if isinstance(node, ast.ClassDef) and node.name == class_name: - for item in node.body: - if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) and item.name == func_name: - result = ast.get_source_segment(source, item) - assert result is not None, f"{class_name}.{func_name} source not found" - return result + # Follow the bases: ``Orchestrator`` inherits its loops and helpers from + # ``OrchestratorBase`` in the same module, and the lookup stays anchored + # on the public class rather than on where the split happens to fall. + classes = {node.name: node for node in ast.walk(tree) if isinstance(node, ast.ClassDef)} + pending = [class_name] + seen: set[str] = set() + while pending: + name = pending.pop(0) + if name in seen or name not in classes: + continue + seen.add(name) + node = classes[name] + for item in node.body: + if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) and item.name == func_name: + result = ast.get_source_segment(source, item) + assert result is not None, f"{class_name}.{func_name} source not found" + return result + pending.extend(base.id for base in node.bases if isinstance(base, ast.Name)) else: for node in ast.walk(tree): if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name: diff --git a/tests/e2e/features/fullduplex/test_core_contracts.py b/tests/e2e/features/fullduplex/test_core_contracts.py index 12a87640f..45493745c 100644 --- a/tests/e2e/features/fullduplex/test_core_contracts.py +++ b/tests/e2e/features/fullduplex/test_core_contracts.py @@ -14,7 +14,7 @@ import dataclasses import pytest -from vllm_omni.engine.duplex.messages import DuplexFence +from vllm_omni.engine.duplex.contracts import DuplexFence from vllm_omni.experimental.fullduplex.core import protocol as ev from vllm_omni.experimental.fullduplex.core.adapter import ( DuplexAdapter, @@ -32,12 +32,15 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def test_fence_is_frozen_value_identity(): + # The fence carries (session_id, epoch, turn_id). response_seq was + # dropped with the incarnation concept: response identity is the session's + # active_response_id, not a field of the fence. f = DuplexFence(session_id="s1") - assert (f.epoch, f.turn_id, f.response_seq) == (0, 0, 0) + assert (f.epoch, f.turn_id) == (0, 0) with pytest.raises(dataclasses.FrozenInstanceError): f.epoch = 1 # type: ignore[misc] - assert DuplexFence("s1", 1, 2, 3) == DuplexFence("s1", 1, 2, 3) - assert DuplexFence("s1", 1, 2, 3) != DuplexFence("s1", 2, 2, 3) + assert DuplexFence("s1", 1, 2) == DuplexFence("s1", 1, 2) + assert DuplexFence("s1", 1, 2) != DuplexFence("s1", 2, 2) def test_session_identity_rules(): diff --git a/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py b/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py index f5bba11e4..0f9b4af51 100644 --- a/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py +++ b/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py @@ -71,15 +71,8 @@ from vllm_omni.experimental.fullduplex.video_stacking import ( # noqa: E402 from vllm_omni.metrics.definitions import compute_audio_rtf # noqa: E402 -def _url_with_model(url, model, *, autostart=None, session_id=None): - # This driver speaks the MiniCPM-o native duplex wire contract. - return build_realtime_url( - url, - model, - autostart=autostart, - session_id=session_id, - extra_query={"native_duplex": "1"}, - ) +def _url_with_model(url, model, *, autostart=None): + return build_realtime_url(url, model, autostart=autostart) _read_wav_pcm16 = read_pcm16_wav @@ -630,7 +623,6 @@ def _session_update_event(args: DemoArgs) -> dict[str, object]: "playback_commit_policy": "ack_only", "extra_body": { "auto_response": True, - "native_duplex": True, "force_listen_count": 0, }, } @@ -646,9 +638,6 @@ def _session_update_event(args: DemoArgs) -> dict[str, object]: "type": "session.update", "session": session_payload, } - session_id = getattr(args, "session_id", None) - if session_id: - session_payload["session_id"] = session_id return event @@ -1462,7 +1451,6 @@ async def run_demo(args: DemoArgs) -> dict[str, object]: args.url, args.model, autostart=False if getattr(args, "ref_audio", None) else None, - session_id=getattr(args, "session_id", None), ) state = DemoState() stop = asyncio.Event() @@ -1674,6 +1662,10 @@ async def run_demo(args: DemoArgs) -> dict[str, object]: ) unexpected_error_events = _unexpected_error_events(state) continuous_input_ok = not continuous_input or state.count("input_audio_buffer.committed") == 0 + # The server allocates every session id; label metrics with the one it announced. + created = state.events[0] if state.events else {} + created_session = created.get("session") if isinstance(created, dict) else None + server_session_id = created_session.get("id") if isinstance(created_session, dict) else None result = { "ok": terminal_activity_ok and state.count("session.closed") > 0 @@ -1711,8 +1703,8 @@ async def run_demo(args: DemoArgs) -> dict[str, object]: "input_transcription_ok": input_transcription_ok, "completed_response_ids": completed_response_ids, "response_timings": state.response_timing_summaries(), - "request_metrics": state.session_request_metrics(session_id=args.session_id), - "session_metrics": state.session_metric_summary(session_id=args.session_id), + "request_metrics": state.session_request_metrics(session_id=server_session_id), + "session_metrics": state.session_metric_summary(session_id=server_session_id), "lifecycle_counts_ok": lifecycle_counts_ok, "validation_mode": validation_mode, "scenario": scenario, @@ -1773,7 +1765,6 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--url", default="ws://localhost:8099/v1/realtime?duplex=1") parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5") - parser.add_argument("--session-id", help="Use an explicit public session ID, including for close/reopen tests.") parser.add_argument("--input-wav", required=True) parser.add_argument("--ref-audio", help="Optional WAV used as the MiniCPM-o voice prompt.") parser.add_argument( diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py index b5d04097b..19b4b8107 100644 --- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py @@ -17,7 +17,6 @@ import hashlib import io import json import sys -import uuid import wave from pathlib import Path @@ -266,8 +265,6 @@ async def _run_demo_process( str(args.chunk_ms), "--timeout-s", str(args.timeout_s), - "--session-id", - f"duplex-pair-{label}-{uuid.uuid4().hex}", ] if args.no_realtime_pacing: command.append("--no-realtime-pacing") diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py index 32e6e01da..c13f8ca7f 100644 --- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py @@ -21,7 +21,7 @@ from websockets.asyncio.client import ClientConnection from websockets.exceptions import ConnectionClosed SCRIPT_DIR = Path(__file__).resolve().parent -REPO_ROOT = SCRIPT_DIR.parents[3] +REPO_ROOT = SCRIPT_DIR.parents[2] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) @@ -42,6 +42,15 @@ def _url_with_model(*args, **kwargs) -> str: return _scenario_module()._url_with_model(*args, **kwargs) +def _created_session_id(created: dict[str, object]) -> str: + """The server-allocated session id announced by ``session.created``.""" + session = created.get("session") + session_id = session.get("id") if isinstance(session, dict) else None + if not isinstance(session_id, str) or not session_id: + raise RuntimeError("session.created omitted the server-allocated session id") + return session_id + + async def run_demo(args): return await _scenario_module().run_demo(args) @@ -217,7 +226,6 @@ async def _open_admission_session( args.url, args.model, autostart=False if getattr(args, "ref_audio", None) else None, - session_id=session_id, ) ws = await websockets.connect(url, max_size=64 * 1024 * 1024) await ws.send( @@ -225,10 +233,8 @@ async def _open_admission_session( { "type": "session.update", "session": { - "session_id": session_id, "model": args.model, "modalities": ["audio", "text"], - "extra_body": {"native_duplex": True}, **({"ref_audio": _ref_audio_data_url(args.ref_audio)} if getattr(args, "ref_audio", None) else {}), }, } @@ -255,12 +261,10 @@ async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str, for index in range(limit): accepted.append(await _open_admission_session(args, f"{prefix}-accepted-{index}")) - overflow_id = f"{prefix}-overflow" overflow_url = _url_with_model( args.url, args.model, autostart=False if getattr(args, "ref_audio", None) else None, - session_id=overflow_id, ) async with websockets.connect(overflow_url, max_size=64 * 1024 * 1024) as overflow: await overflow.send( @@ -268,10 +272,8 @@ async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str, { "type": "session.update", "session": { - "session_id": overflow_id, "model": args.model, "modalities": ["audio", "text"], - "extra_body": {"native_duplex": True}, **( {"ref_audio": _ref_audio_data_url(args.ref_audio)} if getattr(args, "ref_audio", None) @@ -326,7 +328,6 @@ async def _resume_probe( args.url, args.model, autostart=False if getattr(args, "ref_audio", None) else None, - session_id=session_id, ) async with websockets.connect(url, max_size=64 * 1024 * 1024) as first: await first.send( @@ -334,10 +335,8 @@ async def _resume_probe( { "type": "session.update", "session": { - "session_id": session_id, "model": args.model, "modalities": ["audio", "text"], - "extra_body": {"native_duplex": True}, **( {"ref_audio": _ref_audio_data_url(args.ref_audio)} if getattr(args, "ref_audio", None) @@ -348,10 +347,10 @@ async def _resume_probe( ) ) created, first_events = await _receive_until(first, "session.created", timeout_s=args.timeout_s) + session_id = _created_session_id(created) token = created.get("resume_token") - incarnation = created.get("incarnation") generation = created.get("attachment_generation") - if not isinstance(token, str) or not isinstance(incarnation, int): + if not isinstance(token, str): raise RuntimeError("session.created omitted resumable credentials") last_seq = max(_server_event_sequences(first_events), default=0) @@ -366,7 +365,6 @@ async def _resume_probe( { "type": "session.resume", "session_id": session_id, - "incarnation": incarnation, "resume_token": token, "last_received_server_event_seq": last_seq, } @@ -429,7 +427,6 @@ async def _takeover_probe( args.url, args.model, autostart=False if getattr(args, "ref_audio", None) else None, - session_id=session_id, ) first = await websockets.connect(url, max_size=64 * 1024 * 1024) second = None @@ -439,10 +436,8 @@ async def _takeover_probe( { "type": "session.update", "session": { - "session_id": session_id, "model": args.model, "modalities": ["audio", "text"], - "extra_body": {"native_duplex": True}, **( {"ref_audio": _ref_audio_data_url(args.ref_audio)} if getattr(args, "ref_audio", None) @@ -453,10 +448,10 @@ async def _takeover_probe( ) ) created, first_events = await _receive_until(first, "session.created", timeout_s=args.timeout_s) + session_id = _created_session_id(created) token = created.get("resume_token") - incarnation = created.get("incarnation") generation = created.get("attachment_generation") - if not isinstance(token, str) or not isinstance(incarnation, int) or not isinstance(generation, int): + if not isinstance(token, str) or not isinstance(generation, int): raise RuntimeError("session.created omitted takeover credentials") last_seq = max(_server_event_sequences(first_events), default=0) @@ -466,7 +461,6 @@ async def _takeover_probe( { "type": "session.resume", "session_id": session_id, - "incarnation": incarnation, "resume_token": token, "last_received_server_event_seq": last_seq, } diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py index 5e0cfff3a..2eee43fae 100644 --- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py @@ -18,7 +18,6 @@ import hashlib import io import json import sys -import uuid import wave from pathlib import Path @@ -412,8 +411,6 @@ async def run_soft_interrupt(args: argparse.Namespace) -> dict[str, object]: str(args.chunk_ms), "--timeout-s", str(args.timeout_s), - "--session-id", - f"duplex-soft-interrupt-{uuid.uuid4().hex}", ] command.extend(["--ref-audio", str(_canonical_path(args.ref_audio))]) temperature = getattr(args, "temperature", None) diff --git a/tests/e2e/online_serving/test_duplex_client_live.py b/tests/e2e/online_serving/test_duplex_client_live.py index d90a8a683..28044cdd1 100644 --- a/tests/e2e/online_serving/test_duplex_client_live.py +++ b/tests/e2e/online_serving/test_duplex_client_live.py @@ -12,7 +12,6 @@ incremental playback acks, a forced transport drop recovered by automatic from __future__ import annotations import asyncio -import uuid from pathlib import Path import pytest @@ -47,7 +46,6 @@ async def _run_live_session(*, url: str, model: str, ref_audio: Path, input_wav: url, model=model, config=config, - session_id=f"duplex-client-live-{uuid.uuid4().hex}", reconnect=ReconnectPolicy(max_attempts=5, backoff_s=(0.25, 1.0)), # Short interval so the periodic heartbeat path runs during the # session — including across the forced transport drop below. diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py b/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py index c04d3acec..59b40f5e5 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py @@ -6,8 +6,8 @@ from __future__ import annotations import asyncio +import base64 import json -import uuid from pathlib import Path import pytest @@ -90,25 +90,16 @@ async def _receive_protocol_events(ws, required_types: set[str], *, timeout_s: f async def _run_protocol_smoke(*, url: str, model: str, ref_audio: Path) -> list[dict[str, object]]: - session_id = f"duplex-ci-protocol-{uuid.uuid4().hex}" - websocket_url = build_realtime_url( - url, - model, - autostart=False, - session_id=session_id, - extra_query={"native_duplex": "1"}, - ) + websocket_url = build_realtime_url(url, model, autostart=False) async with websockets.connect(websocket_url, max_size=64 * 1024 * 1024) as ws: await ws.send( json.dumps( { "type": "session.update", "session": { - "session_id": session_id, "model": model, "modalities": ["audio", "text"], "ref_audio": _ref_audio_data_url(str(ref_audio)), - "extra_body": {"native_duplex": True}, }, } ) @@ -123,6 +114,170 @@ async def _run_protocol_smoke(*, url: str, model: str, ref_audio: Path) -> list[ return events +async def _run_text_only_response_create( + *, + url: str, + model: str, + ref_audio: Path, + text: str, + timeout_s: float = 120.0, +) -> dict[str, object]: + """Drive a text prompt with no audio and return the server's answer. + + A Realtime client asks for speech from text with a user message item + followed by ``response.create``. A model-native duplex session generates + from audio units, so this returns whatever the server answers rather than + assuming a response. + """ + websocket_url = build_realtime_url(url, model, autostart=False) + async with websockets.connect(websocket_url, max_size=64 * 1024 * 1024) as ws: + await ws.send( + json.dumps( + { + "type": "session.update", + "session": { + "model": model, + "modalities": ["audio", "text"], + "output_audio_format": "pcm16", + "turn_detection": None, + "temperature": 0.0, + "ref_audio": _ref_audio_data_url(str(ref_audio)), + "extra_body": {"auto_response": False}, + }, + } + ) + ) + await _receive_protocol_events(ws, {"session.created", "session.updated"}, timeout_s=60) + await ws.send( + json.dumps( + { + "type": "conversation.item.create", + "item": { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": text}], + }, + } + ) + ) + await ws.send(json.dumps({"type": "response.create"})) + + async def receive_outcome() -> dict[str, object]: + while True: + raw = await ws.recv() + if not isinstance(raw, str): + continue + event = json.loads(raw) + if isinstance(event, dict) and event.get("type") in {"error", "response.done"}: + return event + + return await asyncio.wait_for(receive_outcome(), timeout=timeout_s) + + +async def _run_seeded_text_to_audio( + *, + url: str, + model: str, + ref_audio: Path | None, + text: str, + modalities: tuple[str, ...] = ("audio", "text"), + silence_seconds: float = 12.0, + timeout_s: float = 180.0, +) -> dict[str, object]: + """Speak a seeded text: the duplex route's text-to-speech shape. + + A model-native session takes its text once, in the session context + (``duplex_initial_user_text``), and then generates per audio unit. Silence + carries no content of its own, so it only advances the clock and lets the + model answer the seeded turn. + """ + websocket_url = build_realtime_url(url, model, autostart=False) + audio_bytes = 0 + transcript: list[str] = [] + output_text: list[str] = [] + seen: list[str] = [] + session_payload: dict[str, object] = { + "model": model, + "modalities": list(modalities), + "input_audio_format": "pcm16", + "output_audio_format": "pcm16", + "turn_detection": None, + "temperature": 0.0, + "extra_body": { + "auto_response": True, + "force_listen_count": 0, + "duplex_initial_user_text": text, + }, + } + if ref_audio is not None: + session_payload["ref_audio"] = _ref_audio_data_url(str(ref_audio)) + async with websockets.connect(websocket_url, max_size=64 * 1024 * 1024) as ws: + await ws.send(json.dumps({"type": "session.update", "session": session_payload})) + + async def reader() -> None: + nonlocal audio_bytes + while True: + raw = await ws.recv() + if not isinstance(raw, str): + continue + event = json.loads(raw) + if not isinstance(event, dict): + continue + event_type = event.get("type") + if isinstance(event_type, str): + seen.append(event_type) + delta = event.get("delta") + if event_type == "response.output_audio.delta" and isinstance(delta, str): + audio_bytes += len(base64.b64decode(delta)) + elif event_type == "response.output_audio_transcript.delta" and isinstance(delta, str): + transcript.append(delta) + elif event_type == "response.output_text.delta" and isinstance(delta, str): + output_text.append(delta) + elif event_type == "error": + raise AssertionError(f"seeded text turn received an error: {event}") + + reader_task = asyncio.create_task(reader()) + silence = bytes(2 * 16_000 * 200 // 1000) + sent_ms = 0 + try: + while sent_ms < silence_seconds * 1000 and "response.done" not in seen: + sent_ms += 200 + await ws.send( + json.dumps( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(silence).decode("ascii"), + "input_audio_format": "pcm16", + "sample_rate_hz": 16_000, + "duration_ms": 200, + "audio_end_ms": sent_ms, + } + ) + ) + await asyncio.sleep(0.2) + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout_s + while loop.time() < deadline and "response.done" not in seen: + await asyncio.sleep(0.5) + if reader_task.done(): + await reader_task + # Close explicitly. Dropping the socket parks the session in its + # disconnect grace, where it keeps holding an admission slot and + # starves the tests that follow. + await ws.send(json.dumps({"type": "session.close"})) + deadline = loop.time() + 30 + while loop.time() < deadline and "session.closed" not in seen: + await asyncio.sleep(0.25) + finally: + reader_task.cancel() + return { + "audio_bytes": audio_bytes, + "transcript": "".join(transcript), + "output_text": "".join(output_text), + "event_types": seen, + } + + @pytest.mark.core_model @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) @@ -205,6 +360,215 @@ def test_duplex_single_session_video_input(omni_server, tmp_path: Path) -> None: _assert_session_metrics(result["session_metrics"], expected_count=2) +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +@pytest.mark.parametrize( + ("locale", "text"), + [ + ("en", "Please say exactly: the quick brown fox jumps over the lazy dog."), + ("zh", "请朗读：今天天气很好，我们一起去公园散步。"), + ], +) +def test_duplex_seeded_text_to_audio(omni_server, locale: str, text: str) -> None: + """text -> audio over the duplex route, in both locales. + + This is the case the deleted turn-based ``text -> audio`` tests covered. + It is also the shape the Seed-TTS Realtime backend uses, so it guards the + benchmark path as well as the modality. + """ + result = asyncio.run( + _run_seeded_text_to_audio( + url=realtime_url(omni_server), + model=omni_server.model, + ref_audio=resolve_ref_audio(), + text=text, + ) + ) + + assert "response.done" in result["event_types"], result["event_types"] + assert int(result["audio_bytes"]) > 0, f"{locale} seeded text produced no audio" + assert str(result["transcript"]).strip(), f"{locale} seeded text produced audio with no transcript" + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +def test_duplex_seeded_text_to_text_needs_no_reference_voice(omni_server) -> None: + """text -> text, the one duplex session that opens without a reference voice. + + ``ref_audio`` is required only when the session asks for audio output, so a + ``modalities: ["text"]`` session is the duplex equivalent of the deleted + turn-based ``text -> text`` case. The model is model-native and still + speaks its answer, so this asserts the text side and the absence of the + ``ref_audio_required`` rejection, not the absence of audio. + """ + result = asyncio.run( + _run_seeded_text_to_audio( + url=realtime_url(omni_server), + model=omni_server.model, + ref_audio=None, + text="What is the capital of France? Answer in one short sentence.", + modalities=("text",), + ) + ) + + assert "response.done" in result["event_types"], result["event_types"] + produced_text = str(result["output_text"]) or str(result["transcript"]) + assert produced_text.strip(), f"text-only session produced nothing: {result['event_types']}" + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +def test_duplex_seeded_text_to_long_audio_output(omni_server) -> None: + """A longer answer, so Code2Wav runs across many frames rather than one. + + Replaces the deleted ``text_to_audio_long_output`` case. The assertion is + on the amount of audio rather than its content: the point is that a long + answer streams to completion, not that the model says any given word. + """ + result = asyncio.run( + _run_seeded_text_to_audio( + url=realtime_url(omni_server), + model=omni_server.model, + ref_audio=resolve_ref_audio(), + text="What is the capital of China? Answer in about 40 words.", + silence_seconds=30.0, + timeout_s=240.0, + ) + ) + + assert "response.done" in result["event_types"], result["event_types"] + # 24 kHz mono pcm16: 2 s of speech is 96000 bytes, comfortably more than a + # single Code2Wav frame and well under a 40-word answer. + assert int(result["audio_bytes"]) > 96_000, f"expected a long answer, got {result['audio_bytes']} bytes" + assert str(result["transcript"]).strip() + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +def test_duplex_seeded_long_form_generation(omni_server) -> None: + """Long-form Chinese generation, replacing the deleted long-form case.""" + result = asyncio.run( + _run_seeded_text_to_audio( + url=realtime_url(omni_server), + model=omni_server.model, + ref_audio=resolve_ref_audio(), + text="帮我讲一个100字的故事。", + silence_seconds=30.0, + timeout_s=240.0, + ) + ) + + assert "response.done" in result["event_types"], result["event_types"] + assert int(result["audio_bytes"]) > 96_000, f"expected long-form audio, got {result['audio_bytes']} bytes" + assert str(result["transcript"]).strip() + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +def test_duplex_sequential_sessions_are_independent(omni_server) -> None: + """A second session must answer its own prompt, not replay the first one. + + Replaces the deleted ``sequential_requests_independent`` case. Sessions are + engine-resident now, so this is the check that one session's context does + not leak into the next one on the same stages. + """ + first = asyncio.run( + _run_seeded_text_to_audio( + url=realtime_url(omni_server), + model=omni_server.model, + ref_audio=resolve_ref_audio(), + text="What is the capital of France? Answer in one short sentence.", + ) + ) + second = asyncio.run( + _run_seeded_text_to_audio( + url=realtime_url(omni_server), + model=omni_server.model, + ref_audio=resolve_ref_audio(), + text="What is the capital of China? Answer in one short sentence.", + ) + ) + + for result in (first, second): + assert "response.done" in result["event_types"], result["event_types"] + assert int(result["audio_bytes"]) > 0 + assert str(result["transcript"]).strip() + + # Different prompts must not produce the same answer: that would mean the + # second session inherited the first one's context. + assert str(first["transcript"]).strip() != str(second["transcript"]).strip() + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +def test_duplex_text_only_response_create_is_rejected(omni_server) -> None: + """A duplex session generates from audio: a text-only prompt is refused, not ignored. + + Every session on this route is model-native, and such a model decides per + audio unit whether to speak: conversation items are context, not a turn. + Refusing is the point -- opening a response the model never fills costs the + caller the session's whole idle timeout to discover. A text prompt reaches + the model through the session's seeded opening turn + (``initial_user_text``) instead, which is what ``/v1/chat/completions`` + uses and what the Seed-TTS Realtime backend has to be ported to. + """ + outcome = asyncio.run( + _run_text_only_response_create( + url=realtime_url(omni_server), + model=omni_server.model, + ref_audio=resolve_ref_audio(), + text="Please say: the quick brown fox jumps over the lazy dog.", + ) + ) + + assert outcome.get("type") == "error", outcome + error = outcome.get("error") + assert isinstance(error, dict), outcome + assert error.get("code") == "text_only_turn_unsupported", outcome + assert error.get("type") == "invalid_request_error", outcome + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +def test_duplex_single_session_still_image_input(omni_server, tmp_path: Path) -> None: + """A spoken turn with one still image attached. + + The still-image case is a one-element frame list, which the client repeats + on every unit of the turn: unlike the camera track it never advances, so + this covers the image path the turn-based ``image -> text + audio`` request + used to cover. + """ + args = demo_args( + omni_server=omni_server, + input_wav=validated_input_wav(), + ref_audio=resolve_ref_audio(), + output_dir=tmp_path / "still_image", + ) + args.turns = 2 + args.turn_duration_ms = [args.first_turn_ms] * args.turns + args.video_frames_b64 = duplex_camera_frames(seconds=1, cache_dir=tmp_path / "still")[:1] + + result = asyncio.run(run_demo(args)) + + assert result["ok"] is True + assert result["video_frame_count"] == 1 + assert result["video_stacked_frame_count"] == 0 + _assert_positive_int(result["audio_delta_count"]) + assert result["done_count"] == 2 + assert result["error_count"] == 0 + assert result["all_audio_responses_have_transcript"] is True + assert result["transcript_delta_done_ok"] is True + _assert_request_metrics(result["request_metrics"], expected_count=2) + _assert_session_metrics(result["session_metrics"], expected_count=2) + + @pytest.mark.advanced_model @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py b/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py index 628040001..f931fb95d 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py @@ -30,7 +30,7 @@ from tests.e2e.online_serving.run_minicpmo_realtime_duplex_soft_interrupt import run_soft_interrupt, ) from tests.helpers.mark import hardware_test -from vllm_omni.entrypoints.duplex.server_vad import ( +from vllm_omni.engine.duplex.vad import ( SILERO_VAD_FILENAME, SILERO_VAD_REPO_ID, SILERO_VAD_REVISION, diff --git a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py index ecd885dd6..a23ce73e8 100644 --- a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py +++ b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py @@ -152,7 +152,6 @@ def test_realtime_duplex_demo_pair_launches_demo_processes_concurrently(tmp_path "parser.add_argument('--output-dir')", "parser.add_argument('--chunk-ms')", "parser.add_argument('--timeout-s')", - "parser.add_argument('--session-id')", "parser.add_argument('--ref-audio')", "parser.add_argument('--require-audio', action='store_true')", "parser.add_argument('--no-realtime-pacing', action='store_true')", @@ -1027,7 +1026,7 @@ def test_realtime_duplex_demo_resolves_distinct_turn_inputs(): demo._turn_input_paths(primary, ["second.wav"], turns=3) -def test_realtime_duplex_demo_explicitly_enables_native_runtime_before_connect(): +def test_realtime_duplex_demo_url_carries_no_model_opt_in_or_session_id(): demo = _load_demo_module() url = demo._url_with_model( @@ -1036,23 +1035,12 @@ def test_realtime_duplex_demo_explicitly_enables_native_runtime_before_connect() ) query = parse_qs(urlsplit(url).query) - assert query["native_duplex"] == ["1"] + assert query["duplex"] == ["1"] + assert "native_duplex" not in query + assert "session_id" not in query -def test_realtime_duplex_demo_explicit_session_id_reaches_autostart_query(): - demo = _load_demo_module() - - url = demo._url_with_model( - "ws://localhost:8099/v1/realtime?duplex=1", - "openbmb/MiniCPM-o-4_5", - session_id="reopen-e2e", - ) - - query = parse_qs(urlsplit(url).query) - assert query["session_id"] == ["reopen-e2e"] - - -def test_realtime_duplex_demo_session_update_uses_explicit_session_id(): +def test_realtime_duplex_demo_session_update_leaves_session_id_to_the_server(): demo = _load_demo_module() event = demo._session_update_event( @@ -1066,8 +1054,8 @@ def test_realtime_duplex_demo_session_update_uses_explicit_session_id(): assert event["type"] == "session.update" assert "session_id" not in event - assert event["session"]["session_id"] == "reopen-e2e" - assert event["session"]["extra_body"]["native_duplex"] is True + assert "session_id" not in event["session"] + assert "native_duplex" not in event["session"]["extra_body"] def test_realtime_duplex_demo_response_required_uses_deterministic_sampling(): diff --git a/tests/entrypoints/duplex/test_duplex_audio.py b/tests/engine/duplex/test_duplex_audio.py similarity index 95% rename from tests/entrypoints/duplex/test_duplex_audio.py rename to tests/engine/duplex/test_duplex_audio.py index d0be02072..ca358abb0 100644 --- a/tests/entrypoints/duplex/test_duplex_audio.py +++ b/tests/engine/duplex/test_duplex_audio.py @@ -6,7 +6,7 @@ import base64 import numpy as np import pytest -from vllm_omni.entrypoints.duplex.audio import convert_input_audio_with_rate +from vllm_omni.engine.duplex.audio import convert_input_audio_with_rate pytestmark = [pytest.mark.core_model, pytest.mark.cpu] diff --git a/tests/engine/duplex/test_duplex_commands.py b/tests/engine/duplex/test_duplex_commands.py new file mode 100644 index 000000000..ad45c8355 --- /dev/null +++ b/tests/engine/duplex/test_duplex_commands.py @@ -0,0 +1,453 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OpenAI Realtime client events -> typed ``DuplexCommand`` objects and their mailbox payloads.""" + +from __future__ import annotations + +import base64 + +import numpy as np +import pytest + +from vllm_omni.engine.duplex.commands import ( + REALTIME_COMMAND_TYPES, + AckPlayback, + AppendAudio, + AppendText, + BargeIn, + CancelInput, + CancelResponse, + ClearInput, + ClearOutputAudio, + CloseSession, + Commit, + CreateItem, + CreateResponse, + DeleteItem, + DuplexCommand, + DuplexCommandError, + Heartbeat, + SignalTurn, + TruncateItem, + UpdateSession, + command_from_realtime, +) +from vllm_omni.engine.duplex.realtime_commands import ( + RealtimeInputDefaults, + translate_realtime_command, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_LOUD_PCM16 = base64.b64encode(b"\x00\x10" * 8).decode("ascii") +_SILENT_PCM16 = base64.b64encode(b"\x00\x00" * 8).decode("ascii") +_JPEG_FRAME = base64.b64encode(b"\xff\xd8\xff\xe0" + b"\x00" * 12).decode("ascii") + +_MINIMAL_PAYLOADS: dict[str, tuple[dict[str, object], type[DuplexCommand]]] = { + "input_audio_buffer.append": ({"audio": _LOUD_PCM16}, AppendAudio), + "input_audio_buffer.commit": ({}, Commit), + "input_audio_buffer.clear": ({}, ClearInput), + "output_audio_buffer.clear": ({}, ClearOutputAudio), + "response.create": ({}, CreateResponse), + "response.cancel": ({}, CancelResponse), + "conversation.item.create": ({"item": {"type": "message", "role": "user", "content": []}}, CreateItem), + "conversation.item.delete": ({"item_id": "item_1"}, DeleteItem), + "conversation.item.truncate": ({"item_id": "item_1", "audio_end_ms": 100}, TruncateItem), + "session.update": ({"session": {"instructions": "hi"}}, UpdateSession), + "playback.ack": ({"played_ms": 10}, AckPlayback), + "session.heartbeat": ({}, Heartbeat), + "session.close": ({}, CloseSession), + "turn.signal": ({"event": "user_started"}, SignalTurn), + "input.text.append": ({"text": "hello"}, AppendText), + "input.cancel": ({}, CancelInput), + "barge_in": ({}, BargeIn), +} + + +def test_every_realtime_command_type_has_a_mapping_case(): + assert set(_MINIMAL_PAYLOADS) == set(REALTIME_COMMAND_TYPES) + + +@pytest.mark.parametrize("event_type", sorted(REALTIME_COMMAND_TYPES)) +def test_command_from_realtime_maps_each_type_to_its_dataclass_and_propagates_event_id(event_type: str): + body, expected_cls = _MINIMAL_PAYLOADS[event_type] + + command = command_from_realtime({"type": event_type, "event_id": f"evt-{event_type}", **body}) + + assert type(command) is expected_cls + assert command.event_id == f"evt-{event_type}" + payload = command.payload() + assert payload["type"] == expected_cls.type + assert payload["realtime_event_id"] == f"evt-{event_type}" + assert "event_id" not in payload + + anonymous = command_from_realtime({"type": event_type, **body}) + assert anonymous.event_id is None + assert "realtime_event_id" not in anonymous.payload() + + +@pytest.mark.parametrize( + ("alias", "expected_cls", "body"), + [ + ("push_text", AppendText, {"text": "hello"}), + ("input_text.append", AppendText, {"text": "hello"}), + ("signal_turn", SignalTurn, {"event": "user_started"}), + ("close_session", CloseSession, {}), + ("close", CloseSession, {}), + ("audio.playback_ack", AckPlayback, {"played_ms": 10}), + ("input.commit", Commit, {}), + ], +) +def test_pre_realtime_client_event_aliases_map_to_their_canonical_command( + alias: str, expected_cls: type[DuplexCommand], body: dict[str, object] +): + assert type(command_from_realtime({"type": alias, **body})) is expected_cls + + +def test_append_audio_decodes_and_converts_pcm16_to_16k_float32(): + command = command_from_realtime( + {"type": "input_audio_buffer.append", "audio": _LOUD_PCM16, "format": "pcm16", "sample_rate_hz": 16000} + ) + + assert isinstance(command, AppendAudio) + assert command.format == "pcm_f32le" + assert command.sample_rate_hz == 16000 + samples = np.frombuffer(command.audio, dtype=" 8 * 4 + + passthrough = translate_realtime_command( + {"type": "input_audio_buffer.append", "audio": base64.b64encode(np.zeros(4, dtype=" None: - self.calls = [] - - def execute(self, key, message, **kwargs): - self.calls.append((key, message, kwargs)) - return DuplexControlResultMessage( - control_id=message.control_id, - fence=message.fence, - operation="touch" if isinstance(message, TouchDuplexSessionMessage) else "resume", - session_id=message.session_id, - ok=True, - stage_results=[], - ) - - -class _UnsupportedTransport: - def execute(self, key, message, **kwargs): - del key, kwargs - return DuplexControlResultMessage( - control_id=message.control_id, - fence=message.fence, - operation="touch", - session_id=message.session_id, - ok=False, - stage_results=[{"result": {"supported": False}}], - unsupported_count=1, - error=DuplexControlError(code="invalid_capability", message="unsupported", retryable=False), - accepted_fence=DuplexFence(message.session_id, epoch=2), - lease_generation=7, - ) - - -def test_control_client_routes_touch_and_resume_by_control_id() -> None: - transport = _Transport() - control_ids = iter(("touch-id", "resume-id")) - client = DuplexControlClient(transport, control_id_factory=lambda: next(control_ids)) - fence = DuplexFence("sid-client") - - assert ( - client.touch( - fence.session_id, - fence=fence, - activity=DuplexLeaseActivity.PLAYBACK_ACK, - timeout=2.0, - )["ok"] - is True - ) - assert ( - client.resume( - fence.session_id, - fence=fence, - expected_lease_generation=7, - timeout=3.0, - )["ok"] - is True - ) - - touch_key, touch_message, touch_kwargs = transport.calls[0] - assert touch_key == ("duplex", "touch-id") - assert isinstance(touch_message, TouchDuplexSessionMessage) - assert touch_message.activity == DuplexLeaseActivity.PLAYBACK_ACK.value - assert touch_kwargs["timeout"] == 2.0 - - resume_key, resume_message, resume_kwargs = transport.calls[1] - assert resume_key == ("duplex", "resume-id") - assert isinstance(resume_message, ResumeDuplexSessionMessage) - assert resume_message.expected_lease_generation == 7 - assert resume_kwargs["timeout"] == 3.0 - - -def test_control_client_treats_ok_false_as_authoritative_failure() -> None: - client = DuplexControlClient(_UnsupportedTransport(), control_id_factory=lambda: "unsupported") - fence = DuplexFence("sid-unsupported") - - with pytest.raises(DuplexControlRequestError, match="duplex touch failed") as exc_info: - client.touch( - fence.session_id, - fence=fence, - activity=DuplexLeaseActivity.HEARTBEAT, - timeout=1.0, - ) - - assert exc_info.value.code == "invalid_capability" - assert exc_info.value.retryable is False - assert exc_info.value.accepted_fence == DuplexFence(fence.session_id, epoch=2) - assert exc_info.value.lease_generation == 7 diff --git a/tests/engine/duplex/test_duplex_control_plane.py b/tests/engine/duplex/test_duplex_control_plane.py deleted file mode 100644 index 14da4d21c..000000000 --- a/tests/engine/duplex/test_duplex_control_plane.py +++ /dev/null @@ -1,1270 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -import logging -from dataclasses import FrozenInstanceError - -import msgspec -import pytest - -from vllm_omni.engine.duplex.control_plane import ( - DuplexControlPlane, - DuplexOutputContext, - DuplexRequestIdentity, - DuplexStageRequestContext, - DuplexStageSubmission, - DuplexStageSubmissionResult, -) -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity, DuplexLeaseConfig -from vllm_omni.engine.duplex.messages import ( - AppendDuplexInputMessage, - CloseDuplexSessionMessage, - DuplexFence, - DuplexSessionLifecycleMessage, - OpenDuplexSessionMessage, - ResumeDuplexSessionMessage, - SignalDuplexTurnMessage, - TouchDuplexSessionMessage, -) -from vllm_omni.engine.duplex.runtime import ( - DuplexAppendPlan, - DuplexInputMode, - DuplexRuntimeCapabilities, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class _Extension: - def configure_sampling_params(self, *, runtime_config, defaults): - del runtime_config - return tuple(f"configured-{stage_id}" for stage_id, _ in enumerate(defaults)) - - def plan_append( - self, - *, - request_id, - fence, - session_config, - runtime_config, - seq, - turn_seq, - mode, - payload, - final, - sampling_params, - ): - del request_id, fence, session_config, runtime_config, seq, turn_seq, mode, payload, final - assert sampling_params == "configured-0" - return DuplexAppendPlan(prompt={"prompt_token_ids": [1, 2, 3]}) - - def decide_output(self, **kwargs): - del kwargs - return None - - -class _TypedStagePort: - stage_count = 2 - - def __init__(self) -> None: - self.ensure_calls: list[DuplexStageRequestContext] = [] - self.submit_calls: list[DuplexStageSubmission] = [] - self.cleanup_calls: list[tuple[list[str], bool]] = [] - - def sampling_defaults(self) -> tuple[object, ...]: - return ("default-0", "default-1") - - def ensure_request(self, context: DuplexStageRequestContext) -> None: - self.ensure_calls.append(context) - - async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: - self.submit_calls.append(submission) - return DuplexStageSubmissionResult( - request_id=submission.context.request_id, - stage_id=submission.context.stage_id, - replica_id=3, - ) - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - self.cleanup_calls.append((request_ids, abort)) - - -class _Clock: - def __init__(self) -> None: - self.value = 0.0 - - def __call__(self) -> float: - return self.value - - def advance(self, seconds: float) -> None: - self.value += seconds - - -@pytest.mark.asyncio -async def test_control_plane_uses_frozen_typed_stage_context_without_request_state() -> None: - stage_port = _TypedStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=_Extension(), - stage_port=stage_port, - result_sink=result_sink, - ) - fence = DuplexFence("typed-port") - - await plane.handle( - OpenDuplexSessionMessage( - control_id="open", - fence=fence, - session_id=fence.session_id, - capabilities={ - "input_modes": [DuplexInputMode.APPEND_AUDIO_CHUNK.value], - }, - session_config={"voice": "test"}, - runtime_config={"runtime": "test"}, - ) - ) - assert (await result_sink.get()).ok is True - - context = stage_port.ensure_calls[-1] - assert context.session_id == fence.session_id - assert context.fence == fence - assert context.stage_id == 0 - assert context.final_stage_id == 1 - assert context.sampling_params == ("configured-0", "configured-1") - assert context.session_config == {"voice": "test"} - assert context.runtime_config == {"runtime": "test"} - with pytest.raises(FrozenInstanceError): - context.stage_id = 1 # type: ignore[misc] - - await plane.handle( - AppendDuplexInputMessage( - control_id="append", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"audio": b"pcm"}, - final=True, - ) - ) - append_result = await result_sink.get() - assert append_result.ok is True - assert append_result.stage_results[0]["replica_id"] == 3 - assert append_result.stage_results[0]["result"]["response_stage_id"] == 1 - - submission = stage_port.submit_calls[-1] - assert submission.context == stage_port.ensure_calls[-1] - assert submission.prompt == {"prompt_token_ids": [1, 2, 3]} - assert submission.already_submitted is False - assert not hasattr(submission, "request_state") - - -def test_control_plane_accepts_only_typed_duplex_messages() -> None: - plane = DuplexControlPlane( - extension=None, - stage_port=_TypedStagePort(), - result_sink=asyncio.Queue(), - ) - - assert plane.accepts( - OpenDuplexSessionMessage( - control_id="open", - fence=DuplexFence("typed-message"), - session_id="typed-message", - capabilities=DuplexRuntimeCapabilities().__dict__, - ) - ) - assert plane.accepts(type("Lookalike", (), {"type": "open_duplex_session"})()) is False - - -@pytest.mark.parametrize( - "message", - [ - TouchDuplexSessionMessage( - control_id="touch-1", - fence=DuplexFence("sid-message"), - session_id="sid-message", - activity=DuplexLeaseActivity.HEARTBEAT.value, - ), - ResumeDuplexSessionMessage( - control_id="resume-1", - fence=DuplexFence("sid-message"), - session_id="sid-message", - expected_lease_generation=3, - ), - DuplexSessionLifecycleMessage( - fence=DuplexFence("sid-message"), - session_id="sid-message", - event="expired", - reason="idle_ttl_expired", - lease_generation=4, - submitted_request_ids=["req-submitted"], - reserved_request_ids=["req-reserved"], - ), - ], -) -def test_duplex_lease_messages_round_trip(message) -> None: - encoded = msgspec.json.encode(message) - decoded = msgspec.json.decode(encoded, type=type(message)) - - assert decoded == message - - -@pytest.mark.asyncio -async def test_control_plane_touches_append_signal_and_explicit_activity() -> None: - clock = _Clock() - stage_port = _TypedStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - lifecycle_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=_Extension(), - stage_port=stage_port, - result_sink=result_sink, - lifecycle_sink=lifecycle_sink, - lease_config=DuplexLeaseConfig(idle_ttl_s=30.0, disconnect_grace_s=5.0), - clock=clock, - ) - fence = DuplexFence("sid-touch") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open", - fence=fence, - session_id=fence.session_id, - capabilities={"input_modes": [DuplexInputMode.APPEND_AUDIO_CHUNK.value]}, - ) - ) - await result_sink.get() - session = plane.sessions.require(fence.session_id) - - clock.advance(1.0) - await plane.handle( - AppendDuplexInputMessage( - control_id="append", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"audio": b"pcm"}, - ) - ) - assert (await result_sink.get()).ok is True - assert session.lease.last_activity == 1.0 - - clock.advance(1.0) - await plane.handle( - SignalDuplexTurnMessage( - control_id="signal", - fence=fence, - session_id=fence.session_id, - event="session.update", - ) - ) - assert (await result_sink.get()).ok is True - assert session.lease.last_activity == 2.0 - - submit_count = len(stage_port.submit_calls) - clock.advance(1.0) - await plane.handle( - TouchDuplexSessionMessage( - control_id="heartbeat", - fence=fence, - session_id=fence.session_id, - activity=DuplexLeaseActivity.HEARTBEAT.value, - ) - ) - touch_result = await result_sink.get() - assert touch_result.ok is True - assert session.lease.last_activity == 3.0 - assert len(stage_port.submit_calls) == submit_count - assert lifecycle_sink.empty() - - -@pytest.mark.asyncio -async def test_control_plane_resume_requires_expected_lease_generation() -> None: - clock = _Clock() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=_TypedStagePort(), - result_sink=result_sink, - lifecycle_sink=asyncio.Queue(), - lease_config=DuplexLeaseConfig(), - clock=clock, - ) - fence = DuplexFence("sid-resume-control") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open", - fence=fence, - session_id=fence.session_id, - capabilities={}, - ) - ) - await result_sink.get() - await plane.handle( - TouchDuplexSessionMessage( - control_id="detach", - fence=fence, - session_id=fence.session_id, - activity=DuplexLeaseActivity.DETACH.value, - ) - ) - await result_sink.get() - - await plane.handle( - ResumeDuplexSessionMessage( - control_id="resume", - fence=fence, - session_id=fence.session_id, - expected_lease_generation=0, - ) - ) - result = await result_sink.get() - - assert result.ok is True - assert result.stage_results[0]["result"]["lease_generation"] == 1 - - await plane.handle( - ResumeDuplexSessionMessage( - control_id="stale-resume", - fence=fence, - session_id=fence.session_id, - expected_lease_generation=0, - ) - ) - stale_result = await result_sink.get() - assert stale_result.ok is False - assert stale_result.accepted_fence == fence - assert stale_result.lease_generation == 1 - - -@pytest.mark.asyncio -async def test_control_plane_reaps_one_session_through_cleanup_and_lifecycle_sink() -> None: - clock = _Clock() - stage_port = _TypedStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - lifecycle_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=stage_port, - result_sink=result_sink, - lifecycle_sink=lifecycle_sink, - lease_config=DuplexLeaseConfig(idle_ttl_s=2.0, disconnect_grace_s=1.0), - clock=clock, - ) - for session_id in ("sid-a", "sid-b"): - fence = DuplexFence(session_id) - await plane.handle( - OpenDuplexSessionMessage( - control_id=f"open-{session_id}", - fence=fence, - session_id=session_id, - capabilities={}, - ) - ) - await result_sink.get() - clock.advance(1.0) - plane.sessions.require("sid-b").touch(DuplexFence("sid-b"), DuplexLeaseActivity.HEARTBEAT) - clock.advance(1.1) - - expired_count = await plane.reap_expired() - - assert expired_count == 1 - assert plane.sessions.get("sid-a") is None - assert plane.sessions.get("sid-b") is not None - assert stage_port.cleanup_calls == [] - lifecycle = await lifecycle_sink.get() - assert isinstance(lifecycle, DuplexSessionLifecycleMessage) - assert lifecycle.session_id == "sid-a" - assert lifecycle.reason == "idle_ttl_expired" - assert result_sink.empty() - - -@pytest.mark.asyncio -async def test_control_plane_retries_expired_cleanup_before_publishing_lifecycle() -> None: - class _FailOnceStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.failures_remaining = 1 - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - if self.failures_remaining: - self.failures_remaining -= 1 - raise RuntimeError("transient cleanup failure") - await super().cleanup(request_ids, abort=abort) - - clock = _Clock() - stage_port = _FailOnceStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - lifecycle_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=stage_port, - result_sink=result_sink, - lifecycle_sink=lifecycle_sink, - lease_config=DuplexLeaseConfig(idle_ttl_s=1.0, disconnect_grace_s=1.0), - clock=clock, - ) - fence = DuplexFence("sid-retry-expiry") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-retry", - fence=fence, - session_id=fence.session_id, - capabilities={}, - ) - ) - await result_sink.get() - request_id = plane.stage_request_id(fence, stage_id=0) - plane.sessions.require(fence.session_id).reserve_stage_request(0, request_id, fence=fence) - clock.advance(2.0) - - assert await plane.reap_expired() == 0 - assert lifecycle_sink.empty() - - assert await plane.reap_expired() == 1 - assert (await lifecycle_sink.get()).session_id == fence.session_id - - -@pytest.mark.asyncio -async def test_expired_session_holds_admission_slot_while_cleanup_is_blocked() -> None: - class _BlockedCleanupStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.cleanup_started = asyncio.Event() - self.release_cleanup = asyncio.Event() - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - self.cleanup_started.set() - await self.release_cleanup.wait() - await super().cleanup(request_ids, abort=abort) - - clock = _Clock() - stage_port = _BlockedCleanupStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=stage_port, - result_sink=result_sink, - lease_config=DuplexLeaseConfig(idle_ttl_s=1.0, disconnect_grace_s=1.0), - clock=clock, - max_sessions=1, - ) - old_fence = DuplexFence("sid-expired-blocked") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-expired-blocked", - fence=old_fence, - session_id=old_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - old_session = plane.sessions.require(old_fence.session_id) - old_session.bind_stage_request(0, "req-expired-blocked", fence=old_fence) - clock.advance(2.0) - - reap_task = asyncio.create_task(plane.reap_expired()) - await asyncio.wait_for(stage_port.cleanup_started.wait(), timeout=1.0) - replacement_fence = DuplexFence("sid-replacement-blocked") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-replacement-blocked", - fence=replacement_fence, - session_id=replacement_fence.session_id, - capabilities={}, - ) - ) - - blocked_result = await result_sink.get() - assert blocked_result.ok is False - assert blocked_result.error is not None - assert blocked_result.error.code == "resource_exhausted" - assert plane.sessions.get(old_fence.session_id) is old_session - - stage_port.release_cleanup.set() - assert await reap_task == 1 - assert plane.sessions.get(old_fence.session_id) is None - - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-replacement-after-cleanup", - fence=replacement_fence, - session_id=replacement_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - - -@pytest.mark.asyncio -async def test_failed_expiry_cleanup_retains_admission_slot_for_retry() -> None: - class _FailingCleanupStagePort(_TypedStagePort): - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - del request_ids, abort - raise RuntimeError("cleanup failed") - - clock = _Clock() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=_FailingCleanupStagePort(), - result_sink=result_sink, - lease_config=DuplexLeaseConfig(idle_ttl_s=1.0, disconnect_grace_s=1.0), - clock=clock, - max_sessions=1, - ) - old_fence = DuplexFence("sid-expired-failed") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-expired-failed", - fence=old_fence, - session_id=old_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - old_session = plane.sessions.require(old_fence.session_id) - old_session.bind_stage_request(0, "req-expired-failed", fence=old_fence) - clock.advance(2.0) - - assert await plane.reap_expired() == 0 - replacement_fence = DuplexFence("sid-replacement-failed") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-replacement-failed", - fence=replacement_fence, - session_id=replacement_fence.session_id, - capabilities={}, - ) - ) - - blocked_result = await result_sink.get() - assert blocked_result.ok is False - assert blocked_result.error is not None - assert blocked_result.error.code == "resource_exhausted" - assert plane.sessions.get(old_fence.session_id) is old_session - - -@pytest.mark.asyncio -async def test_capacity_rejection_is_an_info_event(caplog: pytest.LogCaptureFixture) -> None: - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=_TypedStagePort(), - result_sink=result_sink, - max_sessions=1, - ) - first_fence = DuplexFence("sid-capacity-first") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-capacity-first", - fence=first_fence, - session_id=first_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - caplog.clear() - - rejected_fence = DuplexFence("sid-capacity-rejected") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-capacity-rejected", - fence=rejected_fence, - session_id=rejected_fence.session_id, - capabilities={}, - ) - ) - - result = await result_sink.get() - assert result.error.code == "resource_exhausted" - assert not [record for record in caplog.records if record.levelno >= logging.ERROR] - - -@pytest.mark.asyncio -async def test_failed_request_cleanup_retries_before_releasing_admission() -> None: - class _FailOnceCleanupStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.attempts = 0 - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - self.cleanup_calls.append((list(request_ids), abort)) - self.attempts += 1 - if self.attempts == 1: - raise RuntimeError("transient request cleanup failure") - - stage_port = _FailOnceCleanupStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=stage_port, - result_sink=result_sink, - max_sessions=1, - ) - old_fence = DuplexFence("sid-request-cleanup") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-request-cleanup", - fence=old_fence, - session_id=old_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - old_session = plane.sessions.require(old_fence.session_id) - old_session.bind_stage_request(0, "req-request-cleanup", fence=old_fence) - - closed = plane.close_sessions_for_request_ids( - ["req-request-cleanup"], - abort=True, - ) - assert closed == {old_fence.session_id: ["req-request-cleanup"]} - assert plane.sessions.get(old_fence.session_id) is old_session - with pytest.raises(RuntimeError, match="transient request cleanup failure"): - await stage_port.cleanup(closed[old_fence.session_id], abort=True) - - replacement_fence = DuplexFence("sid-request-cleanup-replacement") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-request-cleanup-replacement-blocked", - fence=replacement_fence, - session_id=replacement_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).error.code == "resource_exhausted" - - assert await plane.reap_expired() == 1 - assert plane.sessions.get(old_fence.session_id) is None - assert stage_port.cleanup_calls == [ - (["req-request-cleanup"], True), - (["req-request-cleanup"], True), - ] - - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-request-cleanup-replacement", - fence=replacement_fence, - session_id=replacement_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - - -@pytest.mark.asyncio -async def test_stale_request_cleanup_does_not_finalize_reopened_incarnation() -> None: - stage_port = _TypedStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=stage_port, - result_sink=result_sink, - max_sessions=1, - ) - old_fence = DuplexFence("sid-request-cleanup-reopen") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-request-cleanup-old", - fence=old_fence, - session_id=old_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - old_session = plane.sessions.require(old_fence.session_id) - old_session.bind_stage_request(0, "req-request-cleanup-old", fence=old_fence) - closed = plane.close_sessions_for_request_ids( - ["req-request-cleanup-old"], - abort=True, - cleanup_in_progress=True, - ) - assert closed == {old_fence.session_id: ["req-request-cleanup-old"]} - plane.defer_request_cleanups([old_fence.session_id]) - - await plane.handle( - CloseDuplexSessionMessage( - control_id="close-request-cleanup-old", - fence=old_fence, - session_id=old_fence.session_id, - reason="client_close", - ) - ) - assert (await result_sink.get()).ok is True - assert plane.sessions.get(old_fence.session_id) is None - - new_fence = DuplexFence(old_fence.session_id, incarnation=1) - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-request-cleanup-new", - fence=new_fence, - session_id=new_fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - new_session = plane.sessions.require(new_fence.session_id) - assert new_session.resume(new_fence, expected_lease_generation=0) == 1 - - assert await plane.reap_expired() == 1 - assert plane.sessions.get(new_fence.session_id) is new_session - - -@pytest.mark.asyncio -async def test_failed_open_cleanup_is_retried_before_releasing_admission() -> None: - class _FailOpenAndFirstCleanupStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.ensure_attempts = 0 - self.cleanup_attempts = 0 - - def ensure_request(self, context: DuplexStageRequestContext) -> None: - self.ensure_attempts += 1 - if self.ensure_attempts == 1: - raise RuntimeError("initial stage setup failed") - super().ensure_request(context) - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - self.cleanup_calls.append((list(request_ids), abort)) - self.cleanup_attempts += 1 - if self.cleanup_attempts == 1: - raise RuntimeError("transient open rollback failure") - - stage_port = _FailOpenAndFirstCleanupStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=_Extension(), - stage_port=stage_port, - result_sink=result_sink, - max_sessions=1, - ) - failed_fence = DuplexFence("sid-open-rollback") - - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-fails-during-stage-setup", - fence=failed_fence, - session_id=failed_fence.session_id, - capabilities={"input_modes": ["append_audio_chunk"]}, - ) - ) - - failed_result = await result_sink.get() - assert failed_result.ok is False - assert plane.sessions.get(failed_fence.session_id) is not None - - replacement_fence = DuplexFence("sid-open-rollback-replacement") - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-replacement-before-rollback", - fence=replacement_fence, - session_id=replacement_fence.session_id, - capabilities={"input_modes": ["append_audio_chunk"]}, - ) - ) - assert (await result_sink.get()).error.code == "resource_exhausted" - - assert await plane.reap_expired() == 0 - assert plane.sessions.get(failed_fence.session_id) is None - - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-replacement-after-rollback", - fence=replacement_fence, - session_id=replacement_fence.session_id, - capabilities={"input_modes": ["append_audio_chunk"]}, - ) - ) - assert (await result_sink.get()).ok is True - - -@pytest.mark.asyncio -async def test_expired_cleanup_failure_does_not_block_other_sessions() -> None: - class _OneStuckExpiryStagePort(_TypedStagePort): - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - if request_ids == [stuck_request_id]: - raise RuntimeError("stuck expiry cleanup") - await super().cleanup(request_ids, abort=abort) - - clock = _Clock() - stage_port = _OneStuckExpiryStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - lifecycle_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane( - extension=None, - stage_port=stage_port, - result_sink=result_sink, - lifecycle_sink=lifecycle_sink, - lease_config=DuplexLeaseConfig(idle_ttl_s=1.0, disconnect_grace_s=1.0), - clock=clock, - ) - stuck_fence = DuplexFence("sid-expiry-stuck") - ready_fence = DuplexFence("sid-expiry-ready") - stuck_request_id = plane.stage_request_id(stuck_fence, stage_id=0) - ready_request_id = plane.stage_request_id(ready_fence, stage_id=0) - for fence in (stuck_fence, ready_fence): - await plane.handle( - OpenDuplexSessionMessage( - control_id=f"open-{fence.session_id}", - fence=fence, - session_id=fence.session_id, - capabilities={}, - ) - ) - assert (await result_sink.get()).ok is True - plane.sessions.require(fence.session_id).reserve_stage_request( - 0, - plane.stage_request_id(fence, stage_id=0), - fence=fence, - ) - clock.advance(2.0) - - assert await plane.reap_expired() == 1 - - assert (await lifecycle_sink.get()).session_id == ready_fence.session_id - assert stage_port.cleanup_calls == [([ready_request_id], False)] - - -@pytest.mark.asyncio -async def test_cancel_retains_request_resources_until_abort_succeeds() -> None: - class _FailOnceStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.failures_remaining = 1 - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - if self.failures_remaining: - self.failures_remaining -= 1 - raise RuntimeError("transient abort failure") - await super().cleanup(request_ids, abort=abort) - - stage_port = _FailOnceStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=None, stage_port=stage_port, result_sink=result_sink) - cancelled_fence = DuplexFence("sid-cancel-retry") - next_fence = DuplexFence("sid-cancel-retry", epoch=1) - session = plane.sessions.open_session(cancelled_fence) - session.bind_stage_request(0, "req-cancel-retry", fence=cancelled_fence) - - def signal(control_id: str) -> SignalDuplexTurnMessage: - return SignalDuplexTurnMessage( - control_id=control_id, - fence=cancelled_fence, - next_fence=next_fence, - session_id=cancelled_fence.session_id, - event="input.cancel", - ) - - await plane.handle(signal("cancel-first")) - assert (await result_sink.get()).ok is False - assert session.fence == next_fence - assert session.resource_request_ids(cancelled_fence) == ["req-cancel-retry"] - - await plane.handle(signal("cancel-retry")) - assert (await result_sink.get()).ok is True - assert session.resource_request_ids(cancelled_fence) == [] - assert stage_port.cleanup_calls == [(["req-cancel-retry"], True)] - - -@pytest.mark.asyncio -async def test_close_retains_session_resources_until_cleanup_succeeds() -> None: - class _FailOnceStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.failures_remaining = 1 - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - if self.failures_remaining: - self.failures_remaining -= 1 - raise RuntimeError("transient close cleanup failure") - await super().cleanup(request_ids, abort=abort) - - stage_port = _FailOnceStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=None, stage_port=stage_port, result_sink=result_sink) - fence = DuplexFence("sid-close-retry") - session = plane.sessions.open_session(fence) - session.bind_stage_request(0, "req-close-retry", fence=fence) - - def close(control_id: str) -> CloseDuplexSessionMessage: - return CloseDuplexSessionMessage( - control_id=control_id, - fence=fence, - session_id=fence.session_id, - reason="client_close", - ) - - await plane.handle(close("close-first")) - assert (await result_sink.get()).ok is False - assert plane.sessions.get(fence.session_id) is session - assert session.resource_request_ids() == ["req-close-retry"] - - await plane.handle(close("close-retry")) - assert (await result_sink.get()).ok is True - assert plane.sessions.get(fence.session_id) is None - assert stage_port.cleanup_calls == [(["req-close-retry"], True)] - - -@pytest.mark.asyncio -async def test_control_cleanup_failure_does_not_block_other_sessions() -> None: - class _IndependentFailureStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.attempts: dict[str, int] = {} - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - request_id = request_ids[0] - self.attempts[request_id] = self.attempts.get(request_id, 0) + 1 - if request_id == "req-stuck" or self.attempts[request_id] == 1: - raise RuntimeError(f"cleanup failed for {request_id}") - await super().cleanup(request_ids, abort=abort) - - stage_port = _IndependentFailureStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=None, stage_port=stage_port, result_sink=result_sink) - - for session_id, request_id in (("sid-stuck", "req-stuck"), ("sid-ready", "req-ready")): - fence = DuplexFence(session_id) - session = plane.sessions.open_session(fence) - session.bind_stage_request(0, request_id, fence=fence) - await plane.handle( - CloseDuplexSessionMessage( - control_id=f"close-{session_id}", - fence=fence, - session_id=session_id, - reason="client_close", - ) - ) - assert (await result_sink.get()).ok is False - - await plane.reap_expired() - - assert plane.sessions.get("sid-stuck") is not None - assert plane.sessions.get("sid-ready") is None - assert stage_port.cleanup_calls == [(["req-ready"], True)] - - -@pytest.mark.asyncio -async def test_control_cleanup_is_single_flight_with_concurrent_reaper() -> None: - class _SlowCleanupStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.cleanup_started = asyncio.Event() - self.release_cleanup = asyncio.Event() - self.cleanup_attempts = 0 - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - self.cleanup_attempts += 1 - self.cleanup_started.set() - await self.release_cleanup.wait() - await super().cleanup(request_ids, abort=abort) - - stage_port = _SlowCleanupStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=None, stage_port=stage_port, result_sink=result_sink) - fence = DuplexFence("sid-concurrent-cleanup") - session = plane.sessions.open_session(fence) - session.bind_stage_request(0, "req-concurrent-cleanup", fence=fence) - - close_task = asyncio.create_task( - plane.handle( - CloseDuplexSessionMessage( - control_id="close-concurrent-cleanup", - fence=fence, - session_id=fence.session_id, - reason="client_close", - ) - ) - ) - await asyncio.wait_for(stage_port.cleanup_started.wait(), timeout=1) - - reaper_task = asyncio.create_task(plane.reap_expired()) - await asyncio.sleep(0) - await asyncio.sleep(0) - assert stage_port.cleanup_attempts == 1 - - stage_port.release_cleanup.set() - await asyncio.gather(close_task, reaper_task) - - assert (await result_sink.get()).ok is True - assert plane.sessions.get(fence.session_id) is None - assert stage_port.cleanup_calls == [(["req-concurrent-cleanup"], True)] - - -@pytest.mark.asyncio -async def test_unknown_input_mode_rejects_open_without_registering_session() -> None: - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=None, stage_port=_TypedStagePort(), result_sink=result_sink) - fence = DuplexFence("sid-unknown-mode") - - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-unknown-mode", - fence=fence, - session_id=fence.session_id, - capabilities={"input_modes": ["future_mode"]}, - ) - ) - - result = await result_sink.get() - assert result.ok is False - assert result.error is not None - assert result.error.code == "invalid_capability" - assert plane.sessions.get(fence.session_id) is None - - -@pytest.mark.asyncio -async def test_extension_free_turn_commit_does_not_allocate_or_submit_stage_request() -> None: - stage_port = _TypedStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=None, stage_port=stage_port, result_sink=result_sink) - fence = DuplexFence("sid-turn-commit") - - await plane.handle( - OpenDuplexSessionMessage( - control_id="open-turn-commit", - fence=fence, - session_id=fence.session_id, - capabilities={"input_modes": [DuplexInputMode.TURN_COMMIT_ONLY.value]}, - ) - ) - assert (await result_sink.get()).ok is True - - await plane.handle( - AppendDuplexInputMessage( - control_id="append-turn-commit", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.TURN_COMMIT_ONLY.value, - payload=None, - final=True, - ) - ) - result = await result_sink.get() - - assert result.ok is True - assert result.stage_results[0]["result"]["data_plane_append"] is False - assert stage_port.ensure_calls == [] - assert stage_port.submit_calls == [] - - -@pytest.mark.asyncio -async def test_session_update_replaces_both_configs_with_one_generation() -> None: - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=_Extension(), stage_port=_TypedStagePort(), result_sink=result_sink) - fence = DuplexFence("sid-atomic-config") - session = plane.sessions.open_session(fence) - - await plane.handle( - SignalDuplexTurnMessage( - control_id="update-config", - fence=fence, - session_id=fence.session_id, - event="session.update", - session_config={"voice": "new"}, - runtime_config={"temperature": 0.5}, - ) - ) - - assert (await result_sink.get()).ok is True - assert session.config_generation == 1 - assert session.session_config == {"voice": "new"} - assert session.runtime_config == {"temperature": 0.5} - - -@pytest.mark.asyncio -async def test_stage_submission_is_compensated_when_local_commit_fails() -> None: - stage_port = _TypedStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=_Extension(), stage_port=stage_port, result_sink=result_sink) - fence = DuplexFence("sid-submit-compensation") - session = plane.sessions.open_session( - fence, - capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}), - ) - request_id = plane.stage_request_id(fence, stage_id=0) - - def fail_commit(_reservation): - raise RuntimeError("local commit failed") - - session.commit_append = fail_commit # type: ignore[method-assign] - await plane.handle( - AppendDuplexInputMessage( - control_id="append-compensation", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"audio": b"pcm"}, - ) - ) - - assert (await result_sink.get()).ok is False - assert stage_port.cleanup_calls == [([request_id], True)] - - -@pytest.mark.asyncio -async def test_failed_submission_compensation_blocks_append_until_reaper_cleans_it() -> None: - class _FailingCleanupStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.cleanup_failures = 2 - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - await super().cleanup(request_ids, abort=abort) - if self.cleanup_failures > 0: - self.cleanup_failures -= 1 - raise RuntimeError("stage abort failed") - - stage_port = _FailingCleanupStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=_Extension(), stage_port=stage_port, result_sink=result_sink) - fence = DuplexFence("sid-durable-submit-compensation") - session = plane.sessions.open_session( - fence, - capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}), - ) - original_commit = session.commit_append - - def fail_commit(_reservation): - raise RuntimeError("local commit failed") - - session.commit_append = fail_commit # type: ignore[method-assign] - first = AppendDuplexInputMessage( - control_id="append-compensation-1", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"audio": b"pcm"}, - ) - await plane.handle(first) - assert (await result_sink.get()).ok is False - assert plane.pending_submission_cleanup_count == 1 - assert len(stage_port.submit_calls) == 1 - - await plane.handle( - AppendDuplexInputMessage( - control_id="append-compensation-2", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"audio": b"pcm"}, - ) - ) - assert (await result_sink.get()).ok is False - assert plane.pending_submission_cleanup_count == 1 - assert len(stage_port.submit_calls) == 1 - - session.commit_append = original_commit # type: ignore[method-assign] - stage_port.cleanup_failures = 0 - await plane.reap_expired() - - assert plane.pending_submission_cleanup_count == 0 - assert session.resource_request_ids() == [] - - -def test_stale_output_is_rejected_before_extension_decision() -> None: - class _CountingExtension(_Extension): - def __init__(self) -> None: - self.decision_calls = 0 - - def decide_output(self, **kwargs): - del kwargs - self.decision_calls += 1 - return None - - extension = _CountingExtension() - plane = DuplexControlPlane(extension=extension, stage_port=_TypedStagePort(), result_sink=asyncio.Queue()) - current = DuplexFence("sid-stale-output", epoch=1) - plane.sessions.open_session(current) - context = DuplexOutputContext( - identity=DuplexRequestIdentity( - session_id=current.session_id, - fence=DuplexFence(current.session_id, epoch=0), - ), - final_stage_id=1, - segment_finished=True, - ) - - assert plane.decide_output(0, object(), context) is None - assert extension.decision_calls == 0 - - -@pytest.mark.asyncio -async def test_control_dispatch_is_ordered_per_session_without_blocking_other_sessions() -> None: - class _BlockingStagePort(_TypedStagePort): - def __init__(self) -> None: - super().__init__() - self.started = asyncio.Event() - self.release = asyncio.Event() - - async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: - if submission.context.session_id == "sid-blocked": - self.started.set() - await self.release.wait() - return await super().submit(submission) - - stage_port = _BlockingStagePort() - result_sink: asyncio.Queue = asyncio.Queue() - plane = DuplexControlPlane(extension=_Extension(), stage_port=stage_port, result_sink=result_sink) - blocked = DuplexFence("sid-blocked") - independent = DuplexFence("sid-independent") - capabilities = DuplexRuntimeCapabilities(input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}) - plane.sessions.open_session(blocked, capabilities=capabilities) - plane.sessions.open_session(independent, capabilities=capabilities) - - plane.dispatch( - AppendDuplexInputMessage( - control_id="blocked-append", - fence=blocked, - session_id=blocked.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"audio": b"pcm"}, - operation_id="blocked-operation", - ) - ) - await asyncio.wait_for(stage_port.started.wait(), timeout=1) - plane.dispatch( - AppendDuplexInputMessage( - control_id="blocked-append-retry", - fence=blocked, - session_id=blocked.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"audio": b"pcm"}, - operation_id="blocked-operation", - ) - ) - plane.dispatch( - TouchDuplexSessionMessage( - control_id="blocked-touch", - fence=blocked, - session_id=blocked.session_id, - activity=DuplexLeaseActivity.HEARTBEAT.value, - ) - ) - plane.dispatch( - TouchDuplexSessionMessage( - control_id="independent-touch", - fence=independent, - session_id=independent.session_id, - activity=DuplexLeaseActivity.HEARTBEAT.value, - ) - ) - - independent_result = await asyncio.wait_for(result_sink.get(), timeout=1) - assert independent_result.control_id == "independent-touch" - - stage_port.release.set() - await plane.drain() - assert [result_sink.get_nowait().control_id for _ in range(3)] == [ - "blocked-append", - "blocked-append-retry", - "blocked-touch", - ] - assert [ - submission.context.session_id - for submission in stage_port.submit_calls - if submission.context.session_id == blocked.session_id - ] == [blocked.session_id] diff --git a/tests/engine/duplex/test_duplex_events.py b/tests/engine/duplex/test_duplex_events.py new file mode 100644 index 000000000..a4ebfbf5a --- /dev/null +++ b/tests/engine/duplex/test_duplex_events.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Typed duplex events: pure ``to_realtime()`` rendering and the stateful Realtime projection.""" + +from __future__ import annotations + +import base64 +import inspect + +import pytest + +from vllm_omni.engine.duplex import events as events_module +from vllm_omni.engine.duplex.events import ( + REALTIME_ERROR_TYPES_BY_CODE, + AudioDelta, + DuplexEvent, + DuplexRawEvent, + ErrorEvent, + InputCommitted, + ItemAdded, + ItemCreated, + ItemDone, + Listen, + OutputItemAdded, + OverlapDecision, + ResponseDone, + SessionClosed, + SessionCreated, + SessionExpired, + SessionHeartbeatAck, + SessionReplaced, + SessionResumed, + SessionResyncRequired, + TextDelta, + TranscriptDone, + TurnEvent, + error_event, +) +from vllm_omni.engine.duplex.realtime_events import ( + RealtimeProjectionState, + project_internal_event, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +#: Classes whose wire object carries ``session_id`` by contract. +_SESSION_ID_ON_WIRE = { + SessionClosed, + SessionExpired, + SessionHeartbeatAck, + Listen, + TurnEvent, + OverlapDecision, + SessionResumed, + SessionReplaced, + SessionResyncRequired, +} +#: Classes whose wire object carries ``epoch`` by contract. +_EPOCH_ON_WIRE = {Listen, TurnEvent, OverlapDecision} + + +def _public_event_classes() -> list[type[DuplexEvent]]: + classes = [ + getattr(events_module, name) + for name in events_module.__all__ + if inspect.isclass(getattr(events_module, name)) and issubclass(getattr(events_module, name), DuplexEvent) + ] + assert len(classes) > 40 + return [cls for cls in classes if cls is not DuplexEvent] + + +@pytest.mark.parametrize("event_cls", _public_event_classes(), ids=lambda cls: cls.__name__) +def test_to_realtime_renders_wire_type_event_id_and_documented_fields(event_cls: type[DuplexEvent]): + event = event_cls(session_id="duplex-hidden", epoch=7) + + wire = event.to_realtime() + + assert wire["type"] == event.type + assert wire["event_id"] == event.event_id + assert wire["event_id"].startswith("event_") + if event_cls is DuplexRawEvent: + assert wire["type"] == "duplex." + else: + assert wire["type"] == event_cls.wire_type + if event_cls in _SESSION_ID_ON_WIRE: + assert wire["session_id"] == "duplex-hidden" + else: + assert "session_id" not in wire + if event_cls in _EPOCH_ON_WIRE: + assert wire["epoch"] == 7 + else: + assert "epoch" not in wire + # Every value is plain JSON (mappings / sequences are copied into dict / list). + for value in wire.values(): + assert isinstance(value, str | int | float | bool | dict | list) or value is None + + +def test_optional_wire_fields_are_omitted_when_none(): + created = SessionCreated(session={"id": "sid"}) + assert created.optional_wire_fields == frozenset({"attachment_generation", "resume_token"}) + wire = created.to_realtime() + assert "attachment_generation" not in wire + assert "resume_token" not in wire + assert wire["session"] == {"id": "sid"} + + with_resume = SessionCreated(session={"id": "sid"}, attachment_generation=2, resume_token="tok").to_realtime() + assert with_resume["attachment_generation"] == 2 + assert with_resume["resume_token"] == "tok" + + delta = AudioDelta(response_id="resp_1", item_id="item_resp_1", delta="AAAA").to_realtime() + assert "sample_rate_hz" not in delta + assert "metadata" not in delta + assert delta["format"] == "pcm16" + assert delta["output_index"] == 0 + assert delta["content_index"] == 0 + + anonymous_listen = Listen(session_id="sid", epoch=0, details={"reason": "silence"}).to_realtime() + assert "response_id" not in anonymous_listen + assert "id" not in anonymous_listen["response"] + assert anonymous_listen["response"]["status"] == "listening" + assert anonymous_listen["response"]["metadata"] == {"reason": "silence"} + + bound_listen = Listen(session_id="sid", epoch=0, response_id="resp_1").to_realtime() + assert bound_listen["response_id"] == "resp_1" + assert bound_listen["response"]["id"] == "resp_1" + + +def test_error_event_maps_codes_to_openai_error_types_and_echoes_client_event_id(): + assert REALTIME_ERROR_TYPES_BY_CODE["resource_exhausted"] == "rate_limit_error" + assert REALTIME_ERROR_TYPES_BY_CODE["internal_error"] == "server_error" + assert REALTIME_ERROR_TYPES_BY_CODE["bad_event"] == "invalid_request_error" + + error = ErrorEvent( + code="resource_exhausted", + message="too many sessions", + related_event_id="evt_client", + param="session", + extra={"retryable": True}, + ) + + assert error.error_type == "rate_limit_error" + assert error.text == "too many sessions" + assert error.error == { + "type": "rate_limit_error", + "code": "resource_exhausted", + "message": "too many sessions", + "event_id": "evt_client", + "param": "session", + "retryable": True, + } + wire = error.to_realtime() + assert wire["type"] == "error" + assert wire["error"] == error.error + # The server event id and the echoed client event id are distinct. + assert wire["event_id"] != "evt_client" + assert ErrorEvent(code="never-seen-code").error_type == "invalid_request_error" + + +def test_error_event_helper_normalizes_optional_client_event_id_and_param(): + event = error_event("bad_event", "nope", event_id="evt_1", param="item", extra={"retryable": False}) + assert isinstance(event, ErrorEvent) + assert event.related_event_id == "evt_1" + assert event.param == "item" + assert event.error["retryable"] is False + + anonymous = error_event("bad_event", "nope", event_id=42, param="") + assert anonymous.related_event_id is None + assert anonymous.param is None + assert "event_id" not in anonymous.error + assert "param" not in anonymous.error + + +def test_audio_delta_decodes_base64_and_invalid_payloads_yield_none(): + raw = bytes(range(16)) + delta = AudioDelta(delta=base64.b64encode(raw).decode("ascii"), format="pcm16", sample_rate_hz=24000) + + assert delta.audio == raw + assert delta.to_realtime()["sample_rate_hz"] == 24000 + assert AudioDelta(delta="@@not base64@@").audio is None + + +def test_session_closed_and_expired_are_the_only_terminal_events(): + closed = SessionClosed(session_id="sid", reason="client_close", details={"type": "session.closed"}) + expired = SessionExpired(session_id="sid", reason="idle_timeout") + + assert closed.is_terminal is True + assert expired.is_terminal is True + assert closed.to_realtime() == { + "type": "session.closed", + "event_id": closed.event_id, + "session_id": "sid", + "reason": "client_close", + "event": {"type": "session.closed"}, + } + assert expired.to_realtime() == { + "type": "session.expired", + "event_id": expired.event_id, + "session_id": "sid", + "reason": "idle_timeout", + } + non_terminal = [cls for cls in _public_event_classes() if cls().is_terminal] + assert non_terminal == [SessionClosed, SessionExpired] + + +def test_generic_accessors_expose_response_item_text_and_status(): + assert TextDelta(response_id="resp_1", item_id="item_resp_1", delta="hi").text == "hi" + assert TranscriptDone(transcript="done").text == "done" + assert OutputItemAdded(item={"id": "item_1"}).item_id == "item_1" + assert ItemAdded(item={"id": 3}).item_id is None + assert ResponseDone(response={"status": "cancelled"}).status == "cancelled" + assert ResponseDone().status is None + assert TurnEvent(event="user_started", turn_state="user_speaking").response_id is None + assert TurnEvent().audio is None + + +def test_raw_event_prefixes_internal_type_and_wraps_details(): + raw = DuplexRawEvent(internal_type="model.custom", details={"type": "model.custom", "value": (1, 2)}) + + assert raw.type == "duplex.model.custom" + assert raw.to_realtime() == { + "type": "duplex.model.custom", + "event_id": raw.event_id, + "event": {"type": "model.custom", "value": [1, 2]}, + } + + +def test_turn_event_and_overlap_decision_carry_session_identity_on_the_wire(): + turn = TurnEvent(session_id="sid", epoch=2, event="barge_in", turn_state="barge_in").to_realtime() + assert turn == { + "type": "turn.event", + "event_id": turn["event_id"], + "session_id": "sid", + "event": "barge_in", + "turn_state": "barge_in", + "epoch": 2, + } + decision = OverlapDecision( + session_id="sid", epoch=2, policy="listen_only", action="ignore", reason="short", details={"ms": 300} + ).to_realtime() + assert decision["policy"] == "listen_only" + assert decision["action"] == "ignore" + assert decision["metadata"] == {"ms": 300} + assert decision["session_id"] == "sid" + + +# ---- RealtimeProjectionState / project_internal_event ---- + + +def _types(events: list[DuplexEvent]) -> list[str]: + return [event.type for event in events] + + +def test_projection_of_the_main_internal_event_sequence(): + state = RealtimeProjectionState(session_id="duplex-proj", model="test-model", initial_session_update=True) + pcm = base64.b64encode(b"\x00\x10" * 8).decode("ascii") + + created = project_internal_event(state, {"type": "session.created", "session": {"voice": "alloy"}}) + assert _types(created) == ["session.created", "session.updated"] + assert isinstance(created[0], SessionCreated) + session_payload = created[0].session + assert session_payload["id"] == "duplex-proj" + assert session_payload["model"] == "test-model" + assert session_payload["voice"] == "alloy" + assert session_payload["input_audio_format"] == "pcm16" + assert session_payload["output_audio_format"] == "pcm16" + assert state.initial_session_update is False + # The extra session.updated is a one-shot for the opening session.update. + assert _types(project_internal_event(state, {"type": "session.created", "session": {}})) == ["session.created"] + + response_created = project_internal_event( + state, {"type": "response.created", "response_id": "resp_1", "modalities": ["audio", "text"]} + ) + assert _types(response_created) == [ + "response.created", + "conversation.item.added", + "conversation.item.created", + "response.output_item.added", + "response.content_part.added", + ] + assert state.active_response_id == "resp_1" + assert response_created[0].response_id == "resp_1" + assert response_created[0].to_realtime()["response"]["status"] == "in_progress" + assert isinstance(response_created[1], ItemAdded) + assert isinstance(response_created[2], ItemCreated) + assert response_created[1].item_id == "item_resp_1" + assert response_created[3].item_id == "item_resp_1" + assert response_created[4].to_realtime()["part"] == {"type": "audio", "transcript": ""} + + audio_delta = project_internal_event( + state, + {"type": "response.output_audio.delta", "response_id": "resp_1", "audio": pcm, "text": "hi", "format": "pcm16"}, + ) + assert _types(audio_delta) == ["response.output_audio.delta", "response.output_audio_transcript.delta"] + assert isinstance(audio_delta[0], AudioDelta) + assert audio_delta[0].delta == pcm + assert audio_delta[0].format == "pcm16" + assert audio_delta[0].item_id == "item_resp_1" + assert audio_delta[1].text == "hi" + assert state.conversation_items["item_resp_1"]["content"][0]["transcript"] == "hi" + + done = project_internal_event(state, {"type": "response.done", "response_id": "resp_1"}) + assert _types(done) == [ + "response.output_audio.done", + "response.output_audio_transcript.done", + "response.content_part.done", + "response.output_item.done", + "conversation.item.done", + "response.done", + "rate_limits.updated", + ] + assert done[1].text == "hi" + assert isinstance(done[4], ItemDone) + assert done[4].item["status"] == "completed" + assert isinstance(done[5], ResponseDone) + assert done[5].status == "completed" + assert done[5].response["output"][0]["content"][0] == {"type": "output_audio", "transcript": "hi"} + assert state.active_response_id is None + # A duplicate terminal for the same response is idempotent. + assert project_internal_event(state, {"type": "response.done", "response_id": "resp_1"}) == [] + + committed = project_internal_event( + state, {"type": "input.committed", "message": {"role": "user", "content": "hello"}, "input_commit_seq": 1} + ) + assert _types(committed) == [ + "conversation.item.added", + "conversation.item.created", + "input_audio_buffer.committed", + "conversation.item.done", + ] + assert isinstance(committed[2], InputCommitted) + assert committed[2].previous_item_id == "item_resp_1" + assert committed[2].item_id == committed[0].item_id + assert committed[2].to_realtime()["event"]["input_commit_seq"] == 1 + assert committed[0].item["content"] == [{"type": "input_text", "text": "hello"}] + + project_internal_event(state, {"type": "response.created", "response_id": "resp_2"}) + project_internal_event(state, {"type": "response.output_audio.delta", "response_id": "resp_2", "audio": pcm}) + cancelled = project_internal_event( + state, {"type": "audio.cancelled", "response_id": "resp_2", "reason": "barge_in", "committed_ms": 0} + ) + assert _types(cancelled) == [ + "response.output_audio.done", + "response.content_part.done", + "response.output_item.done", + "conversation.item.done", + "response.done", + "rate_limits.updated", + ] + assert cancelled[4].status == "cancelled" + assert cancelled[4].response["status_details"] == {"type": "cancelled", "reason": "barge_in"} + assert state.active_response_id is None + + closed = project_internal_event(state, {"type": "session.closed", "reason": "client_close"}) + assert _types(closed) == ["session.closed"] + assert closed[0].is_terminal is True + assert closed[0].to_realtime()["reason"] == "client_close" + + +def test_projection_of_output_audio_buffer_clear_emits_cleared_before_terminals(): + state = RealtimeProjectionState(session_id="duplex-clear") + pcm = base64.b64encode(b"\x00\x10" * 8).decode("ascii") + project_internal_event(state, {"type": "response.created", "response_id": "resp_1"}) + project_internal_event(state, {"type": "response.output_audio.delta", "response_id": "resp_1", "audio": pcm}) + + cleared = project_internal_event( + state, {"type": "audio.cancelled", "reason": "output_audio_buffer_clear", "committed_ms": 250} + ) + + assert _types(cleared)[:2] == ["output_audio_buffer.cleared", "response.output_audio.done"] + assert cleared[0].response_id == "resp_1" + assert _types(cleared)[-2:] == ["response.done", "rate_limits.updated"] + assert state.item_truncation_cursors["item_resp_1"] == (0, 250) + + +def test_projection_leaves_error_to_typed_emit_sites(): + # Errors are constructed typed (``error_event``) at the emit site; the projector + # has no ``error`` branch and falls back to the raw wrapper for such a dict. + state = RealtimeProjectionState(session_id="duplex-err") + + projected = project_internal_event(state, {"type": "error", "code": "bad_event", "error": "nope"}) + + assert len(projected) == 1 + assert isinstance(projected[0], DuplexRawEvent) + assert projected[0].type == "duplex.error" + assert projected[0].to_realtime()["event"]["code"] == "bad_event" diff --git a/tests/engine/duplex/test_duplex_lease.py b/tests/engine/duplex/test_duplex_lease.py index 62fc8c70c..bd26b5740 100644 --- a/tests/engine/duplex/test_duplex_lease.py +++ b/tests/engine/duplex/test_duplex_lease.py @@ -1,19 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Lease semantics of the engine-resident duplex session (idle TTL, disconnect grace, resume CAS).""" + from __future__ import annotations from dataclasses import FrozenInstanceError import pytest -from vllm_omni.engine.duplex.lease import ( +from vllm_omni.engine.duplex.config import DuplexSessionConfig +from vllm_omni.engine.duplex.contracts import DuplexFence +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession, DuplexFenceMismatchError +from vllm_omni.engine.duplex.session.lease import ( DuplexLeaseActivity, DuplexLeaseConfig, + DuplexLeaseState, ) -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.runtime import DuplexInputMode -from vllm_omni.engine.duplex.session import DuplexSessionRuntimeManager pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -30,195 +33,148 @@ class FakeMonotonicClock: def _lease_config(*, idle_ttl_s: float | None = 300.0) -> DuplexLeaseConfig: - return DuplexLeaseConfig( - idle_ttl_s=idle_ttl_s, - disconnect_grace_s=30.0, + return DuplexLeaseConfig(idle_ttl_s=idle_ttl_s, disconnect_grace_s=30.0) + + +def _session(session_id: str, clock: FakeMonotonicClock, *, idle_ttl_s: float | None = 300.0) -> DuplexEngineSession: + return DuplexEngineSession( + session_id=session_id, + config=DuplexSessionConfig(), + lease=DuplexLeaseState(config=_lease_config(idle_ttl_s=idle_ttl_s), generation=0, last_activity=clock()), + _clock=clock, ) -def test_open_touch_detach_and_idle_expiry_use_monotonic_time() -> None: +def test_open_touch_detach_and_expiry_use_monotonic_time() -> None: clock = FakeMonotonicClock(100.0) - manager = DuplexSessionRuntimeManager(clock=clock) - fence = DuplexFence("sid-expiry") - session = manager.open_session(fence, lease_config=_lease_config()) + session = _session("sid-expiry", clock) assert session.lease.last_activity == 100.0 clock.advance(10.0) - session.touch(fence, DuplexLeaseActivity.HEARTBEAT) - session.detach(fence) + session.touch_lease(DuplexLeaseActivity.HEARTBEAT) + assert session.lease.last_activity == 110.0 + session.detach_lease() clock.advance(29.0) assert session.lease.disconnect_grace_expired(clock()) is False - assert manager.collect_expired() == [] + assert session.lease.idle_expired(clock()) is False - clock.advance(272.0) - expired = manager.collect_expired() + clock.advance(1.0) + assert session.lease.disconnect_grace_expired(clock()) is True - assert [item.session_id for item in expired] == ["sid-expiry"] - assert expired[0].reason == "disconnect_grace_expired" - assert manager.get("sid-expiry") is session - manager.finalize_close_session(session) - assert manager.get("sid-expiry") is None + clock.advance(271.0) + assert session.lease.idle_expired(clock()) is True -def test_detach_grace_and_resume_advance_token_independent_lease_generation() -> None: +def test_detach_grace_and_resume_advance_the_lease_generation() -> None: clock = FakeMonotonicClock(10.0) - manager = DuplexSessionRuntimeManager(clock=clock) - fence = DuplexFence("sid-resume") - session = manager.open_session(fence, lease_config=_lease_config()) + session = _session("sid-resume", clock) - session.detach(fence) + session.detach_lease() clock.advance(29.0) assert session.lease.disconnect_grace_expired(clock()) is False clock.advance(2.0) assert session.lease.disconnect_grace_expired(clock()) is True - generation = session.resume(fence, expected_lease_generation=0) + generation = session.resume_lease(expected_lease_generation=0) assert generation == 1 + assert session.lease_generation == 1 assert session.lease.detached_at is None assert session.lease.disconnect_grace_expired(clock()) is False with pytest.raises(ValueError, match="lease generation mismatch"): - session.resume(fence, expected_lease_generation=0) + session.resume_lease(expected_lease_generation=0) def test_active_operation_prevents_mid_transaction_expiry() -> None: clock = FakeMonotonicClock(0.0) - manager = DuplexSessionRuntimeManager(clock=clock) - fence = DuplexFence("sid-operation") - session = manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=5.0)) + session = _session("sid-operation", clock, idle_ttl_s=5.0) - session.begin_operation(fence, "append-1") + session.begin_lease_operation(session.fence, "append-1") clock.advance(10.0) - assert manager.collect_expired() == [] + assert session.lease.idle_expired(clock()) is False - session.end_operation(fence, "append-1") - assert manager.collect_expired() == [] + session.end_lease_operation("append-1") + assert session.lease.idle_expired(clock()) is False clock.advance(6.0) - assert [item.session_id for item in manager.collect_expired()] == ["sid-operation"] + assert session.lease.idle_expired(clock()) is True def test_sessions_have_independent_activity_deadlines_and_resources() -> None: clock = FakeMonotonicClock(0.0) - manager = DuplexSessionRuntimeManager(clock=clock) - fence_a = DuplexFence("sid-a") - fence_b = DuplexFence("sid-b") - session_a = manager.open_session(fence_a, lease_config=_lease_config(idle_ttl_s=10.0)) - session_b = manager.open_session(fence_b, lease_config=_lease_config(idle_ttl_s=10.0)) - session_a.reserve_stage_request(0, "req-a-reserved", fence=fence_a) - session_a.bind_stage_request(1, "req-a-submitted", fence=fence_a) - session_b.bind_stage_request(0, "req-b", fence=fence_b) + session_a = _session("sid-a", clock, idle_ttl_s=10.0) + session_b = _session("sid-b", clock, idle_ttl_s=10.0) + session_a.reserve_stage_request(0, "req-a-reserved", fence=session_a.fence) + session_a.bind_stage_request(1, "req-a-submitted", fence=session_a.fence) + session_b.bind_stage_request(0, "req-b", fence=session_b.fence) clock.advance(6.0) - session_b.touch(fence_b, DuplexLeaseActivity.MODEL_OUTPUT) + session_b.touch_lease(DuplexLeaseActivity.MODEL_OUTPUT) clock.advance(5.0) - expired = manager.collect_expired() - assert len(expired) == 1 - assert expired[0].session_id == "sid-a" - assert expired[0].reserved_request_ids == ("req-a-reserved",) - assert expired[0].submitted_request_ids == ("req-a-submitted",) - assert manager.get("sid-b") is session_b + assert session_a.lease.idle_expired(clock()) is True + assert session_b.lease.idle_expired(clock()) is False + assert session_a.resource_request_ids(submitted=False) == ["req-a-reserved"] + assert session_a.resource_request_ids(submitted=True) == ["req-a-submitted"] assert session_b.resource_request_ids() == ["req-b"] -def test_close_and_reaper_select_exactly_one_terminal_transition() -> None: +def test_close_is_a_single_terminal_transition() -> None: clock = FakeMonotonicClock(0.0) - manager = DuplexSessionRuntimeManager(clock=clock) - fence = DuplexFence("sid-race") - manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=1.0)) + session = _session("sid-race", clock, idle_ttl_s=1.0) clock.advance(2.0) + assert session.lease.idle_expired(clock()) is True - assert len(manager.collect_expired()) == 1 - assert manager.close_session(fence, reason="explicit_close") is None - assert manager.collect_expired() == [] - - fence_2 = DuplexFence("sid-race-2") - manager.open_session(fence_2, lease_config=_lease_config(idle_ttl_s=1.0)) - assert manager.close_session(fence_2, reason="explicit_close") is not None - clock.advance(2.0) - assert manager.collect_expired() == [] + assert session.begin_close(reason="explicit_close") is True + assert session.lease.terminal_reason == "explicit_close" + # A terminal lease is never reaped again and a second close is a no-op. + assert session.lease.idle_expired(clock()) is False + assert session.lease.disconnect_grace_expired(clock()) is False + assert session.begin_close(reason="idle_ttl_expired") is True + assert session.lease.terminal_reason == "explicit_close" + with pytest.raises(RuntimeError, match="terminal"): + session.touch_lease(DuplexLeaseActivity.HEARTBEAT) -def test_stale_fence_cannot_touch_detach_or_resume_lease() -> None: +def test_stale_fence_cannot_start_a_lease_operation() -> None: clock = FakeMonotonicClock(0.0) - manager = DuplexSessionRuntimeManager(clock=clock) - current = DuplexFence("sid-fence", epoch=1) + session = _session("sid-fence", clock) + session.epoch = 1 + session.sync_fence() stale = DuplexFence("sid-fence", epoch=0) - session = manager.open_session(current, lease_config=_lease_config()) - with pytest.raises(RuntimeError, match="fence mismatch"): - session.touch(stale, DuplexLeaseActivity.APPEND) - with pytest.raises(RuntimeError, match="fence mismatch"): - session.detach(stale) - with pytest.raises(RuntimeError, match="fence mismatch"): - session.resume(stale, expected_lease_generation=0) + with pytest.raises(DuplexFenceMismatchError, match="fence mismatch"): + session.begin_lease_operation(stale, "append-1") + with pytest.raises(DuplexFenceMismatchError, match="fence mismatch"): + session.begin_lease_operation(DuplexFence("other-session", epoch=1), "append-2") -def test_disabled_idle_expiry_never_collects_session() -> None: +def test_disabled_idle_expiry_never_expires_session() -> None: clock = FakeMonotonicClock(0.0) - manager = DuplexSessionRuntimeManager(clock=clock) - manager.open_session(DuplexFence("sid-no-expiry"), lease_config=_lease_config(idle_ttl_s=None)) + session = _session("sid-no-expiry", clock, idle_ttl_s=None) clock.advance(1_000_000.0) - assert manager.collect_expired() == [] + assert session.lease.expires_at is None + assert session.lease.idle_expired(clock()) is False def test_detached_session_expires_at_disconnect_grace_when_idle_ttl_is_disabled() -> None: clock = FakeMonotonicClock(0.0) - manager = DuplexSessionRuntimeManager(clock=clock) - fence = DuplexFence("sid-disconnect-grace") - session = manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=None)) + session = _session("sid-disconnect-grace", clock, idle_ttl_s=None) - session.detach(fence) + session.detach_lease() clock.advance(29.0) - assert manager.collect_expired() == [] + assert session.lease.disconnect_grace_expired(clock()) is False clock.advance(1.0) - expired = manager.collect_expired() - - assert [item.session_id for item in expired] == [fence.session_id] - assert expired[0].reason == "disconnect_grace_expired" + assert session.lease.disconnect_grace_expired(clock()) is True -def test_runtime_manager_enforces_server_owned_session_admission() -> None: - manager = DuplexSessionRuntimeManager(max_sessions=2) - manager.open_session(DuplexFence("sid-admission-a")) - manager.open_session(DuplexFence("sid-admission-b")) - - with pytest.raises(RuntimeError, match="duplex_session_capacity_exhausted"): - manager.open_session(DuplexFence("sid-admission-c")) - - -def test_completed_append_cache_is_bounded() -> None: - manager = DuplexSessionRuntimeManager(completed_append_limit=2) - fence = DuplexFence("sid-completed-cache") - session = manager.open_session(fence) - - for index in range(3): - session.record_completed_append( - f"operation-{index}", - fence=fence, - mode=DuplexInputMode.TURN_COMMIT_ONLY, - final=True, - stage_results=[{"index": index}], - ) - - assert list(session.completed_appends) == ["operation-1", "operation-2"] - - session.accept_fence(DuplexFence(fence.session_id, epoch=1)) - assert session.completed_appends == {} - - -def test_lease_config_and_expiry_record_are_immutable() -> None: +def test_lease_config_is_immutable_and_validated() -> None: config = _lease_config() with pytest.raises(FrozenInstanceError): config.idle_ttl_s = 1.0 # type: ignore[misc] - - clock = FakeMonotonicClock(0.0) - manager = DuplexSessionRuntimeManager(clock=clock) - manager.open_session(DuplexFence("sid-record"), lease_config=_lease_config(idle_ttl_s=1.0)) - clock.advance(2.0) - record = manager.collect_expired()[0] - with pytest.raises(FrozenInstanceError): - record.reason = "changed" # type: ignore[misc] + with pytest.raises(ValueError, match="idle_ttl_s"): + DuplexLeaseConfig(idle_ttl_s=0.0) + with pytest.raises(ValueError, match="disconnect_grace_s"): + DuplexLeaseConfig(disconnect_grace_s=0.0) diff --git a/tests/engine/duplex/test_duplex_plugin.py b/tests/engine/duplex/test_duplex_plugin.py new file mode 100644 index 000000000..f45b625e7 --- /dev/null +++ b/tests/engine/duplex/test_duplex_plugin.py @@ -0,0 +1,324 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""``DuplexModelPlugin`` loading/validation helpers and the ``PipelineConfig.duplex_plugin`` binding.""" + +from __future__ import annotations + +from dataclasses import fields +from pathlib import Path +from typing import Any + +import pytest +from vllm.sampling_params import SamplingParams + +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig, PipelineConfig +from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig +from vllm_omni.engine.duplex.contracts import DuplexAppendPlan +from vllm_omni.engine.duplex.plugin import ( + DuplexDataPlane, + DuplexModelPlugin, + DuplexModelSessionState, + DuplexRuntimeConfigError, + coerce_int, + load_duplex_plugin, + payload_turn_id, + reject_changed_runtime_value, + validate_duplex_plugin_sampling, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _encode_audio(audio: object, sample_rate_hz: int, fmt: str, speed: float | None) -> str | None: + del audio, sample_rate_hz, fmt, speed + return None + + +class FakeDataPlane(DuplexDataPlane): + def begin_request(self, request_id: str) -> None: + pass + + def is_terminal(self, request_id: str | None) -> bool: + return False + + def mark_terminal(self, request_id: str) -> None: + pass + + def close_stream(self, request_id: str) -> None: + pass + + def close_session(self, session_id: str, *, active_request_id: str | None = None) -> None: + pass + + def project(self, result, *, context=None): + del result, context + return [] + + +class FakePlugin(DuplexModelPlugin): + """Minimal plugin; referenced by dotted path from ``load_duplex_plugin`` tests.""" + + plugin_id = "fake-duplex" + + def __init__(self, encode_audio) -> None: + super().__init__(encode_audio) + self.data_plane = FakeDataPlane() + + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return tuple(defaults) + + def plan_append(self, **kwargs) -> DuplexAppendPlan: + del kwargs + return DuplexAppendPlan(prompt={"prompt_token_ids": [1]}) + + def decide_output(self, **kwargs): + del kwargs + return None + + def create_session_state(self) -> DuplexModelSessionState: + raise NotImplementedError + + def capabilities(self, *, max_sessions: int) -> DuplexCapabilities: + del max_sessions + return DuplexCapabilities() + + def validate_client_extra_body(self, extra_body: object) -> None: + pass + + async def prepare_runtime_config(self, config: DuplexSessionConfig, *, model_config: Any) -> dict[str, object]: + del config, model_config + return {} + + def runtime_config_for_update(self, config, current): + del config + return dict(current) + + def data_plane_context(self, **kwargs) -> object: + return dict(kwargs) + + +class PluginWithoutId(FakePlugin): + plugin_id = "" + + +class PluginWithoutDataPlane(FakePlugin): + def __init__(self, encode_audio) -> None: + DuplexModelPlugin.__init__(self, encode_audio) + + +class PluginWithWrongDataPlane(FakePlugin): + def __init__(self, encode_audio) -> None: + DuplexModelPlugin.__init__(self, encode_audio) + self.data_plane = object() # type: ignore[assignment] + + +class NotAPlugin: + plugin_id = "lookalike" + data_plane = FakeDataPlane() + + def __init__(self, encode_audio) -> None: + self.encode_audio = encode_audio + + +def _path(cls: type) -> str: + return f"{__name__}.{cls.__name__}" + + +# --------------------------------------------------------------------------- # +# load_duplex_plugin # +# --------------------------------------------------------------------------- # + + +def test_load_duplex_plugin_instantiates_the_class_at_the_dotted_path() -> None: + plugin = load_duplex_plugin(_path(FakePlugin), _encode_audio) + + assert type(plugin) is FakePlugin + assert plugin.plugin_id == "fake-duplex" + assert isinstance(plugin.data_plane, FakeDataPlane) + # Defaults of the optional class-level policy. + assert plugin.private_runtime_config_keys == frozenset() + assert plugin.silence_continuation_samples == 16000 + assert plugin.runtime_config_for_function_output(DuplexSessionConfig(), {}, {"name": "tool"}) is None + + +def test_load_duplex_plugin_rejects_a_non_plugin_class() -> None: + with pytest.raises(TypeError, match="not a DuplexModelPlugin"): + load_duplex_plugin(_path(NotAPlugin), _encode_audio) + + +def test_load_duplex_plugin_requires_a_plugin_id() -> None: + with pytest.raises(TypeError, match="plugin_id"): + load_duplex_plugin(_path(PluginWithoutId), _encode_audio) + + +@pytest.mark.parametrize("plugin_cls", [PluginWithoutDataPlane, PluginWithWrongDataPlane]) +def test_load_duplex_plugin_requires_a_data_plane(plugin_cls: type) -> None: + with pytest.raises(TypeError, match="DuplexDataPlane as data_plane"): + load_duplex_plugin(_path(plugin_cls), _encode_audio) + + +def test_load_duplex_plugin_rejects_paths_without_a_module() -> None: + with pytest.raises(ValueError, match="Invalid duplex plugin path"): + load_duplex_plugin("FakePlugin", _encode_audio) + + +def test_load_duplex_plugin_propagates_missing_attributes() -> None: + with pytest.raises(AttributeError): + load_duplex_plugin(f"{__name__}.NoSuchPlugin", _encode_audio) + + +# --------------------------------------------------------------------------- # +# validate_duplex_plugin_sampling # +# --------------------------------------------------------------------------- # + + +def test_validate_duplex_plugin_sampling_accepts_one_parameter_per_stage() -> None: + plugin = FakePlugin(_encode_audio) + + validate_duplex_plugin_sampling(plugin, sampling_defaults=(SamplingParams(), SamplingParams())) + validate_duplex_plugin_sampling(plugin, sampling_defaults=()) + + +def test_validate_duplex_plugin_sampling_rejects_stage_count_mismatch() -> None: + class WrongStageCountPlugin(FakePlugin): + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return defaults[:1] + + with pytest.raises(ValueError, match="one sampling parameter per stage"): + validate_duplex_plugin_sampling( + WrongStageCountPlugin(_encode_audio), + sampling_defaults=(SamplingParams(), SamplingParams()), + ) + + +def test_validate_duplex_plugin_sampling_rejects_type_mismatch() -> None: + class WrongSamplingTypePlugin(FakePlugin): + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return tuple(object() for _ in defaults) + + with pytest.raises(TypeError, match="sampling parameter type mismatch for stage 0"): + validate_duplex_plugin_sampling( + WrongSamplingTypePlugin(_encode_audio), + sampling_defaults=(SamplingParams(), SamplingParams()), + ) + + +def test_validate_duplex_plugin_sampling_skips_type_check_for_stages_without_defaults() -> None: + class ObjectSamplingPlugin(FakePlugin): + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return tuple(object() for _ in defaults) + + validate_duplex_plugin_sampling(ObjectSamplingPlugin(_encode_audio), sampling_defaults=(None, None)) + + +def test_validate_duplex_plugin_sampling_rejects_non_tuple_results() -> None: + class ListSamplingPlugin(FakePlugin): + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return list(defaults) + + with pytest.raises(TypeError, match="as a tuple"): + validate_duplex_plugin_sampling(ListSamplingPlugin(_encode_audio), sampling_defaults=(SamplingParams(),)) + + +# --------------------------------------------------------------------------- # +# Runtime-config helpers # +# --------------------------------------------------------------------------- # + + +def test_reject_changed_runtime_value_only_raises_on_change() -> None: + reject_changed_runtime_value("voice-a", "voice-a", message="voice is fixed", code="voice_changed") + + with pytest.raises(DuplexRuntimeConfigError, match="voice is fixed") as excinfo: + reject_changed_runtime_value("voice-b", "voice-a", message="voice is fixed", code="voice_changed") + + assert excinfo.value.code == "voice_changed" + assert isinstance(excinfo.value, ValueError) + + +def test_reject_changed_runtime_value_uses_the_requested_error_class() -> None: + class CustomRuntimeConfigError(DuplexRuntimeConfigError): + pass + + with pytest.raises(CustomRuntimeConfigError) as excinfo: + reject_changed_runtime_value(1, 2, message="changed", code="custom", error_cls=CustomRuntimeConfigError) + + assert excinfo.value.code == "custom" + + +def test_duplex_runtime_config_error_defaults_its_code() -> None: + assert DuplexRuntimeConfigError("bad").code == "invalid_duplex_runtime_config" + + +@pytest.mark.parametrize( + ("payload", "expected"), + [ + ({"duplex_turn_id": 3}, 3), + ({"duplex_turn_id": "7"}, 7), + ({"model_turn_id": 2}, 2), + ({"duplex_turn_id": 5, "model_turn_id": 9}, 5), + ({"duplex_turn_id": "not-a-number"}, None), + ({}, None), + ("turn", None), + (None, None), + ], +) +def test_payload_turn_id_prefers_duplex_turn_id_and_coerces(payload: object, expected: int | None) -> None: + assert payload_turn_id(payload) == expected + + +@pytest.mark.parametrize( + ("value", "expected"), + [(4, 4), ("11", 11), (2.9, 2), (True, 1), (None, None), ("x", None), (object(), None), ([1], None)], +) +def test_coerce_int(value: object, expected: int | None) -> None: + assert coerce_int(value) == expected + + +# --------------------------------------------------------------------------- # +# PipelineConfig binding # +# --------------------------------------------------------------------------- # + + +def test_pipeline_config_binds_one_duplex_plugin_path() -> None: + field_defaults = {f.name: f.default for f in fields(PipelineConfig)} + + assert "duplex_plugin" in field_defaults + assert field_defaults["duplex_plugin"] is None + + # The pre-framework fields survive only so that the pipelines of the models + # that are not ported yet still construct (RFC vllm-omni#7181 ports them). + # Nothing in the duplex framework may read them: a pipeline is a duplex + # model iff it declares ``duplex_plugin``. + legacy_fields = ("duplex_runtime_extension", "duplex_serving_adapter", "duplex_control_enabled") + for legacy_field in legacy_fields: + assert legacy_field in field_defaults + + framework_sources = [ + *(Path("vllm_omni/engine/duplex").glob("*.py")), + *(Path("vllm_omni/entrypoints/duplex").glob("*.py")), + Path("vllm_omni/engine/duplex_omni_engine.py"), + Path("vllm_omni/engine/duplex_orchestrator.py"), + Path("vllm_omni/entrypoints/duplex_omni.py"), + ] + readers = [ + f"{source}:{legacy_field}" + for source in framework_sources + for legacy_field in legacy_fields + if legacy_field in source.read_text(encoding="utf-8") + ] + assert readers == [] + + # Same rule for the runtime-config knob the framework stopped reading: the + # append-retry table it bounded went away with the correlated append RPC. + assert "completed_append_cache_size" in {f.name for f in fields(DuplexSessionRuntimeConfig)} + assert [ + str(source) + for source in framework_sources + if "completed_append_cache_size" in source.read_text(encoding="utf-8") + ] == [] diff --git a/tests/engine/duplex/test_duplex_runtime.py b/tests/engine/duplex/test_duplex_runtime.py deleted file mode 100644 index 475cf79d0..000000000 --- a/tests/engine/duplex/test_duplex_runtime.py +++ /dev/null @@ -1,563 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from dataclasses import FrozenInstanceError -from types import SimpleNamespace - -import pytest -from vllm.sampling_params import SamplingParams - -from vllm_omni.engine.duplex import runtime as duplex_runtime -from vllm_omni.engine.duplex.contracts import ( - duplex_data_plane_request_info, - duplex_resource_request_belongs_to_session, - duplex_resource_request_id, -) -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.runtime import ( - DuplexInputMode, - DuplexOutputAction, - DuplexOutputDecision, - DuplexRuntimeCapabilities, - DuplexSessionRuntimeManager, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.runtime import ( - MiniCPMO45DuplexRuntimeExtension, - _duplex_hd_slice_count, - build_duplex_data_plane_prompt, - duplex_scheduler_token_budget, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def test_duplex_fence_is_immutable(): - fence = DuplexFence("session") - - with pytest.raises(FrozenInstanceError): - fence.epoch = 1 # type: ignore[misc] - - assert not hasattr(fence, "__dict__") - - -def test_duplex_runtime_tracks_stage_bindings_and_barge_in_epoch(): - manager = DuplexSessionRuntimeManager() - session = manager.open_session( - DuplexFence("sid-1"), - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_TOKENS}, - ), - ) - session.bind_stage_request(stage_id=0, request_id="req-stage0", fence=session.fence) - session.bind_stage_request(stage_id=1, request_id="req-stage1", fence=session.fence) - - update = session.append_input(mode=DuplexInputMode.APPEND_TOKENS, fence=session.fence) - next_fence = DuplexFence("sid-1", epoch=1) - stale_request_ids = session.release_fence(session.fence) - session.accept_fence(next_fence) - - assert update.seq == 1 - assert session.fence == next_fence - assert stale_request_ids == ["req-stage0", "req-stage1"] - assert session.stage_bindings == {} - - -def test_duplex_runtime_tracks_same_request_id_for_each_pipeline_stage(): - manager = DuplexSessionRuntimeManager() - fence = DuplexFence("sid-shared-pipeline-request") - session = manager.open_session(fence) - - session.reserve_stage_request(0, "req-shared", fence=fence) - session.bind_stage_request(0, "req-shared", fence=fence) - session.bind_stage_request(1, "req-shared", fence=fence) - - assert session.stage_bindings[0].request_id == "req-shared" - assert session.stage_bindings[1].request_id == "req-shared" - assert session.resource_request_ids() == ["req-shared"] - assert session.input_seq == 0 - - -def test_duplex_runtime_extension_validation_rejects_missing_methods(): - class IncompleteExtension: - def configure_sampling_params(self, *, runtime_config, defaults): - del runtime_config - return defaults - - with pytest.raises(TypeError, match="plan_append"): - duplex_runtime.validate_duplex_runtime_extension(IncompleteExtension()) - - -def test_duplex_runtime_extension_validation_rejects_stage_count_mismatch(): - class WrongStageCountExtension(MiniCPMO45DuplexRuntimeExtension): - def configure_sampling_params(self, *, runtime_config, defaults): - del runtime_config - return defaults[:1] - - with pytest.raises(ValueError, match="one sampling parameter per stage"): - duplex_runtime.validate_duplex_runtime_extension( - WrongStageCountExtension(), - sampling_defaults=(SamplingParams(), SamplingParams()), - ) - - -def test_duplex_runtime_extension_validation_rejects_sampling_type_mismatch(): - class WrongSamplingTypeExtension(MiniCPMO45DuplexRuntimeExtension): - def configure_sampling_params(self, *, runtime_config, defaults): - del runtime_config - return tuple(object() for _ in defaults) - - with pytest.raises(TypeError, match="sampling parameter type"): - duplex_runtime.validate_duplex_runtime_extension( - WrongSamplingTypeExtension(), - sampling_defaults=(SamplingParams(), SamplingParams()), - ) - - -def test_duplex_output_decision_metadata_is_immutable(): - decision = DuplexOutputDecision( - action=DuplexOutputAction.DIRECT_RESPONSE, - metadata={"model_listen": True}, - ) - - with pytest.raises(TypeError): - decision.metadata["model_listen"] = False - - -def _decide_minicpmo_output( - output: object, - *, - segment_token_ids: tuple[int, ...] = (), - segment_output_metadata: dict | None = None, -): - return MiniCPMO45DuplexRuntimeExtension().decide_output( - stage_id=0, - final_stage_id=1, - segment_finished=True, - segment_token_ids=segment_token_ids, - segment_output_metadata=segment_output_metadata or {}, - output=output, - ) - - -def test_minicpmo_extension_owns_stage_sampling_overrides(): - defaults = ( - SamplingParams(max_tokens=4), - SamplingParams(max_tokens=8), - ) - - configured = MiniCPMO45DuplexRuntimeExtension().configure_sampling_params( - runtime_config={ - "duplex_stage_max_tokens": {"0": 20}, - "duplex_stage_sampling_params": {"1": {"stop_token_ids": [151645]}}, - }, - defaults=defaults, - ) - - assert configured[0].max_tokens == 20 - assert configured[1].stop_token_ids == [151645] - assert defaults[0].max_tokens == 4 - assert 151645 not in (defaults[1].stop_token_ids or []) - - -def test_minicpmo_output_decision_uses_raw_streaming_token_snapshot(): - decision = _decide_minicpmo_output( - SimpleNamespace(outputs=[SimpleNamespace()]), - segment_token_ids=(151705,), - segment_output_metadata={"special_token_ids": {"listen_token_id": 151705}}, - ) - - assert decision is not None - assert decision.action is DuplexOutputAction.DIRECT_RESPONSE - assert decision.metadata["duplex_native_decision"] == "listen" - assert decision.metadata["model_listen"] is True - - -@pytest.mark.parametrize("attr", ["token_ids", "cumulative_token_ids"]) -def test_minicpmo_output_decision_ignores_output_level_token_history(attr): - output = SimpleNamespace( - multimodal_output={"special_token_ids": {"listen_token_id": 151705}}, - outputs=[SimpleNamespace()], - **{attr: [42, 151705]}, - ) - - assert _decide_minicpmo_output(output) is None - - -@pytest.mark.parametrize("attr", ["token_ids", "cumulative_token_ids"]) -def test_minicpmo_output_decision_uses_completion_token_ids(attr): - output = SimpleNamespace( - multimodal_output={"special_token_ids": {"listen_token_id": 151705}}, - outputs=[SimpleNamespace(**{attr: [42, 151705]})], - ) - - assert _decide_minicpmo_output(output) is not None - - -def test_minicpmo_output_decision_uses_completion_stop_reason(): - output = SimpleNamespace( - multimodal_output={"special_token_ids": {"listen_token_id": 151705}}, - outputs=[SimpleNamespace(stop_reason=151705)], - ) - - assert _decide_minicpmo_output(output) is not None - - -def test_duplex_runtime_cancel_fence_rejects_late_append_and_accepts_next_epoch(): - manager = DuplexSessionRuntimeManager() - cancelled_fence = DuplexFence("sid-cancel-race") - next_fence = DuplexFence("sid-cancel-race", epoch=1) - session = manager.open_session( - cancelled_fence, - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - session.bind_stage_request( - stage_id=0, - request_id="req-cancelled", - fence=cancelled_fence, - ) - - stale_request_ids = session.cancel_fence(cancelled_fence, next_fence) - - assert stale_request_ids == ["req-cancelled"] - assert session.fence == next_fence - assert session.stage_bindings == {} - with pytest.raises(RuntimeError, match="fence mismatch"): - session.append_input( - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - fence=cancelled_fence, - ) - update = session.append_input( - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - fence=next_fence, - ) - assert update.seq == 1 - - -def test_duplex_runtime_stale_close_preserves_live_session_and_bindings(): - manager = DuplexSessionRuntimeManager() - current_fence = DuplexFence("sid-stale-close", epoch=1) - session = manager.open_session( - current_fence, - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - session.bind_stage_request(0, "req-live", fence=current_fence) - - with pytest.raises(RuntimeError, match="fence mismatch"): - manager.close_session(DuplexFence("sid-stale-close")) - - assert manager.get("sid-stale-close") is session - assert session.stage_request_ids() == ["req-live"] - - -def test_duplex_runtime_reopen_rejects_late_append_from_old_incarnation(): - manager = DuplexSessionRuntimeManager() - old_fence = DuplexFence("sid-reopen", incarnation=0) - old_session = manager.open_session( - old_fence, - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - manager.close_session(old_fence) - - new_fence = DuplexFence("sid-reopen", incarnation=1) - new_session = manager.open_session( - new_fence, - capabilities=old_session.capabilities, - ) - - with pytest.raises(RuntimeError, match="fence mismatch"): - new_session.append_input( - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - fence=old_fence, - ) - assert ( - new_session.append_input( - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - fence=new_fence, - ).seq - == 1 - ) - - -def test_duplex_prompt_expands_incarnation_metadata(): - fence = DuplexFence("sid-incarnation", incarnation=3) - - prompt = build_duplex_data_plane_prompt( - request_id=duplex_resource_request_id(fence, "stage0"), - fence=fence, - session_config={}, - runtime_config={}, - seq=1, - turn_seq=1, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - payload={"is_speech": True}, - final=False, - ) - - assert prompt["model_intermediate_buffer"]["duplex"]["incarnation"] == 3 - - -def test_duplex_runtime_tracks_turn_local_append_sequence(): - manager = DuplexSessionRuntimeManager() - session = manager.open_session( - DuplexFence("sid-turn-seq"), - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - - first = session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK, fence=session.fence) - second = session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK, fence=session.fence) - next_turn = DuplexFence("sid-turn-seq", turn_id=1, response_seq=1) - third = session.append_input( - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - fence=next_turn, - ) - fourth = session.append_input( - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - fence=next_turn, - ) - - assert [first.seq, second.seq, third.seq, fourth.seq] == [1, 2, 3, 4] - assert [first.turn_seq, second.turn_seq, third.turn_seq, fourth.turn_seq] == [1, 2, 1, 2] - assert [first.turn_id, second.turn_id, third.turn_id, fourth.turn_id] == [0, 0, 1, 1] - - -def test_duplex_runtime_rejects_unsupported_append_mode(): - manager = DuplexSessionRuntimeManager() - session = manager.open_session( - DuplexFence("sid-2"), - capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.TURN_COMMIT_ONLY}), - ) - - with pytest.raises(ValueError, match="not supported"): - session.append_input(mode=DuplexInputMode.APPEND_TOKENS, fence=session.fence) - - -def test_duplex_data_plane_request_info_extracts_structured_stage_result(): - request_id, response_stage_id = duplex_data_plane_request_info( - { - "stage_results": [ - {"result": {"supported": True}}, - { - "result": { - "data_plane_append": True, - "request_id": "duplex-sid-e0-stage0-s1", - "response_stage_id": 1, - } - }, - ] - } - ) - - assert request_id == "duplex-sid-e0-stage0-s1" - assert response_stage_id == 1 - - -def test_duplex_data_plane_request_info_rejects_missing_request_id(): - assert duplex_data_plane_request_info( - { - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": "", - "response_stage_id": 1, - } - } - ] - } - ) == (None, None) - - -def test_duplex_scheduler_token_budget_estimates_pcm_slots(): - assert ( - duplex_scheduler_token_budget( - { - "audio": "AAAAAA==", - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - == 16 - ) - - -def test_duplex_scheduler_token_budget_ignores_client_budget_fields(): - assert ( - duplex_scheduler_token_budget( - { - "audio": "AAAAAA==", - "format": "pcm_f32le", - "duplex_num_input_tokens": 999, - "num_input_tokens": 999, - } - ) - == 16 - ) - - -def test_resource_state_rejects_fence_regression_and_requires_explicit_fence(): - current = DuplexFence("sid", epoch=2, turn_id=3, response_seq=4) - manager = DuplexSessionRuntimeManager() - session = manager.open_session( - current, - capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}), - ) - - for stale in ( - DuplexFence("sid", epoch=1, turn_id=99, response_seq=99), - DuplexFence("sid", epoch=2, turn_id=2, response_seq=4), - DuplexFence("sid", epoch=2, turn_id=3, response_seq=3), - ): - with pytest.raises(RuntimeError, match="fence mismatch"): - session.accept_fence(stale) - assert session.fence == current - - with pytest.raises(TypeError, match="fence"): - session.bind_stage_request(0, "request") - with pytest.raises(TypeError, match="fence"): - session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK) - with pytest.raises(TypeError, match="DuplexFence"): - manager.open_session("legacy-session") - with pytest.raises(TypeError, match="DuplexFence"): - manager.close_session("legacy-session") - - -def test_resource_request_id_is_derived_from_fence_and_role(): - fence = DuplexFence("sid-with-dashes", epoch=7, turn_id=11, response_seq=13) - - assert duplex_resource_request_id(fence, "stage0") == "duplex-s.c2lkLXdpdGgtZGFzaGVz.i.0.e.7.r.stage0" - assert duplex_resource_request_id(fence, "stage1") == "duplex-s.c2lkLXdpdGgtZGFzaGVz.i.0.e.7.r.stage1" - - -def test_resource_request_id_codec_separates_session_id_from_incarnation(): - embedded_incarnation = duplex_resource_request_id( - DuplexFence("foo-i1", incarnation=0), - "stage0", - ) - actual_incarnation = duplex_resource_request_id( - DuplexFence("foo", incarnation=1), - "stage0", - ) - - assert embedded_incarnation != actual_incarnation - - -def test_resource_request_id_session_membership_uses_encoded_identity(): - request_id = duplex_resource_request_id( - DuplexFence("sid-with-dashes", incarnation=2, epoch=7), - "stage0", - ) - - assert duplex_resource_request_belongs_to_session(request_id, "sid-with-dashes") is True - assert duplex_resource_request_belongs_to_session(request_id, "sid") is False - assert duplex_resource_request_belongs_to_session("duplex-s.invalid.i.x.e.7.r.stage0", "sid") is False - - -def test_placeholder_budget_is_planned_inside_omni_engine_boundary(): - fence = DuplexFence("sid", turn_id=1, response_seq=1) - prompt = build_duplex_data_plane_prompt( - request_id=duplex_resource_request_id(fence, "stage0"), - fence=fence, - session_config={}, - runtime_config={}, - seq=2, - turn_seq=1, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - payload={ - "audio": "AAAAAA==", - "format": "pcm_f32le", - "duplex_num_input_tokens": 999, - "num_input_tokens": 999, - }, - final=False, - ) - - assert len(prompt["prompt_token_ids"]) == 16 - assert prompt["model_intermediate_buffer"]["duplex"]["fence"] == fence - assert prompt["model_intermediate_buffer"]["duplex"]["scheduler_token_budget"] == 16 - - -def _jpeg_b64(width: int, height: int) -> str: - import base64 - from io import BytesIO - - from PIL import Image - - buffer = BytesIO() - Image.new("RGB", (width, height), (32, 64, 96)).save(buffer, format="JPEG") - return base64.b64encode(buffer.getvalue()).decode("ascii") - - -def _pcm_unit_payload(frames: list[str]) -> dict[str, object]: - import base64 - - return { - "audio": base64.b64encode(b"\x00" * (16000 * 4)).decode("ascii"), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "video_frames": frames, - } - - -def test_duplex_scheduler_token_budget_one_frame_is_one_block(): - audio_only = duplex_scheduler_token_budget(_pcm_unit_payload([])) - one_frame = duplex_scheduler_token_budget(_pcm_unit_payload([_jpeg_b64(960, 540)])) - - assert one_frame - audio_only == 66 - - -@pytest.mark.parametrize( - ("base_size", "expected_blocks"), - [ - ((960, 540), 4), # official camera frame: source + 2 HD patches + composite - ((540, 960), 4), # portrait phone capture - ((640, 360), 4), - ((448, 448), 2), # fits one tile: the processor does not slice it - ((320, 240), 2), - ], -) -def test_duplex_scheduler_token_budget_matches_stage0_hd_slicing(base_size, expected_blocks): - """A stacked pair is processed with max_slice_nums=[2, 1]; the budget must - equal the blocks Stage0 builds for that base-frame size.""" - audio_only = duplex_scheduler_token_budget(_pcm_unit_payload([])) - pair = duplex_scheduler_token_budget(_pcm_unit_payload([_jpeg_b64(*base_size), _jpeg_b64(448, 448)])) - - assert pair - audio_only == expected_blocks * 66 - - -@pytest.mark.parametrize( - ("size", "expected_patches"), - [ - # multiple = min(ceil(w * h / 448**2), max_slice_nums); <= 1 is never sliced. - ((448, 448), 0), # exactly one tile - ((320, 240), 0), # smaller than one tile - ((447, 449), 0), # just under one tile - ((449, 449), 2), # just over: the processor cuts a 2-tile grid - ((640, 360), 2), - ((960, 540), 2), # official camera frame - ((540, 960), 2), # same frame rotated - ((1920, 1080), 2), # capped by max_slice_nums=2, not by area - ], -) -def test_duplex_hd_slice_count_matches_official_grid(size, expected_patches): - """Lock the port of ``MiniCPMVImageProcessor.get_sliced_grid``. - - Stage0 slices a stacked unit with ``max_slice_nums=[2, 1]``, so the patch - count for the base frame decides the scheduler budget. A frame that fits one - ``scale_resolution=448`` tile is not sliced at all. - """ - assert _duplex_hd_slice_count(size, 2) == expected_patches - - -def test_duplex_scheduler_token_budget_unreadable_base_frame_keeps_hd_fallback(): - audio_only = duplex_scheduler_token_budget(_pcm_unit_payload([])) - pair = duplex_scheduler_token_budget(_pcm_unit_payload(["not-an-image", _jpeg_b64(448, 448)])) - - assert pair - audio_only == 4 * 66 diff --git a/tests/engine/duplex/test_engine_session.py b/tests/engine/duplex/test_engine_session.py new file mode 100644 index 000000000..fccaa59e8 --- /dev/null +++ b/tests/engine/duplex/test_engine_session.py @@ -0,0 +1,612 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Session-state semantics of the engine-resident duplex session (``DuplexEngineSession``).""" + +from __future__ import annotations + +from dataclasses import FrozenInstanceError + +import pytest + +from vllm_omni.engine.duplex.config import ( + DuplexCapabilities, + DuplexConfigError, + DuplexOverlapPolicy, + DuplexSessionConfig, + DuplexTurnEventType, + DuplexTurnState, + ResponseCreateOptions, +) +from vllm_omni.engine.duplex.contracts import DuplexFence +from vllm_omni.engine.duplex.events import TurnEvent +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession, DuplexFenceMismatchError +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.capabilities import ( + minicpmo45_native_capabilities, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _session(session_id: str = "duplex-test", config: DuplexSessionConfig | None = None) -> DuplexEngineSession: + return DuplexEngineSession(session_id=session_id, config=config or DuplexSessionConfig(model="test-model")) + + +def test_commit_audio_input_does_not_advance_model_turn_identity(): + session = _session() + + first = session.commit_audio_input(transcript="first chunk") + second = session.commit_audio_input(transcript="second chunk") + + assert session.input_commit_seq == 2 + assert first.input_commit_seq == 1 + assert second.input_commit_seq == 2 + assert first.turn_id == second.turn_id == session.turn_id == 0 + assert first.message["transcript"] == "first chunk" + assert first.message["content"] == [ + {"type": "audio_url", "audio_url": {"url": "native-duplex:input-audio"}, "transcript": "first chunk"} + ] + assert len(session.history) == 2 + assert session.turn_state == DuplexTurnState.USER_COMMITTED + + +def test_commit_audio_input_accepts_explicit_turn_id(): + session = _session() + + committed = session.commit_audio_input(turn_id=4) + + assert committed.turn_id == 4 + assert session.turn_id == 0 + assert "transcript" not in committed.message + + +# ---- response options ---- + + +def test_response_options_apply_to_one_response_without_mutating_session_defaults(): + session = _session(config=DuplexSessionConfig(instructions="base", voice="base-voice", max_tokens=64)) + session.reserve_response_options( + ResponseCreateOptions(instructions="one response", voice="override-voice", max_tokens=8) + ) + + assert session.config.instructions == "base" + assert session.config.voice == "base-voice" + assert session.response_config is session.config + session.begin_response() + assert session.config.instructions == "base" + assert session.config.voice == "base-voice" + assert session.response_config.instructions == "one response" + assert session.response_config.voice == "override-voice" + assert session.response_config.max_tokens == 8 + + session.end_response() + assert session.response_config is session.config + assert session.config.instructions == "base" + assert session.config.voice == "base-voice" + assert session.config.max_tokens == 64 + + +def test_response_options_cannot_overwrite_an_unconsumed_reservation(): + session = _session(config=DuplexSessionConfig(instructions="base")) + session.reserve_response_options(ResponseCreateOptions(instructions="first")) + + with pytest.raises(RuntimeError, match="already reserved"): + session.reserve_response_options(ResponseCreateOptions(instructions="second")) + + session.begin_response() + assert session.response_config.instructions == "first" + + +def test_response_options_cannot_be_reserved_while_response_is_active(): + session = _session(config=DuplexSessionConfig(instructions="base")) + session.begin_response() + + with pytest.raises(RuntimeError, match="active"): + session.reserve_response_options(ResponseCreateOptions(instructions="too late")) + + +def test_discarded_response_options_do_not_apply(): + session = _session(config=DuplexSessionConfig(instructions="base")) + session.reserve_response_options(ResponseCreateOptions(instructions="dropped")) + session.discard_response_options() + + session.begin_response() + + assert session.response_config.instructions == "base" + + +# ---- response / overlap identity ---- + + +def test_session_owns_response_and_overlap_identity(): + session = _session() + + response_id = session.begin_response() + assert session.turn_state == DuplexTurnState.ASSISTANT_GENERATING + session.accumulate_overlap_speech(320) + session.accumulate_overlap_speech(180) + session.end_response() + + assert session.active_response_id is None + assert session.last_response_id == response_id + assert session.overlap_speech_ms == 500 + assert session.reset_overlap_speech() == 500 + assert session.overlap_speech_ms == 0 + assert session.turn_state == DuplexTurnState.IDLE + + +def test_session_composes_single_owner_ledgers_with_immutable_views(): + session = _session() + + session.bind_request("req-1") + response_id = session.begin_response(turn_id=3) + session.mark_audio_sent(duration_ms=240) + + assert session.active_request_id == "req-1" + assert session.active_response_id == response_id + assert session.active_response_turn_id == 3 + assert session.playback.sent_ms == 240 + assert session.turn_state == DuplexTurnState.ASSISTANT_PLAYING + with pytest.raises(FrozenInstanceError): + session.playback.sent_ms = 480 # type: ignore[misc] + assert session.playback.sent_ms == 240 + + +def test_barge_in_advances_epoch_and_drops_uncommitted_assistant_text(): + session = _session() + response_id = session.begin_response() + session.bind_request("chatcmpl-duplex-test") + session.append_assistant_text("unplayed answer") + + new_epoch = session.barge_in() + + assert response_id is not None + assert new_epoch == 1 + assert session.epoch == 1 + assert session.accepted_fence == DuplexFence(session.session_id, epoch=1, turn_id=0) + assert session.active_request_id is None + assert session.active_response_id is None + assert session.assistant_text_buffer == () + assert session.history == () + assert session.turn_state == DuplexTurnState.BARGE_IN + + +# ---- playback ledger ---- + + +def test_playback_ack_tracks_committed_cursor_separately(): + session = _session() + session.mark_audio_sent(duration_ms=10_000) + + session.acknowledge_playback(played_ms=2_000) + + assert session.playback.generated_ms == 10_000 + assert session.playback.sent_ms == 10_000 + assert session.playback.played_ms == 2_000 + assert session.playback.committed_ms == 2_000 + assert session.playback.as_dict() == { + "generated_ms": 10_000, + "sent_ms": 10_000, + "played_ms": 2_000, + "committed_ms": 2_000, + } + + +def test_playback_ack_is_scoped_per_response(): + session = _session() + first = session.begin_response() + session.mark_audio_sent(duration_ms=3_000) + session.end_response() + second = session.begin_response() + session.mark_audio_sent(duration_ms=1_000) + + session.acknowledge_playback(played_ms=2_500, response_id=first) + + assert session.playback_for_response(first).played_ms == 2_500 + assert session.playback_for_response(second).played_ms == 0 + assert session.playback.played_ms == 0 + + +def test_history_commit_uses_audio_text_alignment_marks(): + session = _session() + session.begin_response() + + session.append_assistant_text("hello ") + session.mark_audio_sent(duration_ms=1_000, text_chars=6) + session.append_assistant_text("world") + session.mark_audio_sent(duration_ms=2_000, text_chars=11) + assert [(mark.text_chars, mark.audio_end_ms) for mark in session.assistant_audio_text_marks] == [ + (6, 1_000), + (11, 2_000), + ] + session.acknowledge_playback(played_ms=1_200, committed_ms=1_200) + + committed = session.end_response(commit_text=True) + + assert committed == {"role": "assistant", "content": "hello w"} + assert session.history[-1] == committed + assert session.last_assistant_full_message == {"role": "assistant", "content": "hello world"} + + +def test_history_commit_with_ack_only_policy_defers_unacknowledged_text(): + session = _session(config=DuplexSessionConfig(playback_commit_policy="ack_only")) + response_id = session.begin_response() + session.append_assistant_text("never played") + session.mark_audio_sent(duration_ms=1_000, text_chars=12) + + committed = session.end_response(commit_text=True) + + assert committed is None + assert session.history == () + assert f"item_{response_id}" in session.pending_history_item_ids + + +def test_truncate_history_item_uses_response_alignment_marks(): + session = _session() + response_id = session.begin_response() + session.append_assistant_text("hello ") + session.mark_audio_sent(duration_ms=1_000, text_chars=6) + session.append_assistant_text("world") + session.mark_audio_sent(duration_ms=2_000, text_chars=11) + session.acknowledge_playback(played_ms=2_000) + committed = session.end_response(commit_text=True) + assert committed == {"role": "assistant", "content": "hello world"} + item_id = f"item_{response_id}" + session.register_history_item(item_id, committed) + + assert session.truncate_history_item(item_id, audio_end_ms=1_000) is True + + assert session.history[-1] == {"role": "assistant", "content": "hello"} + assert session.history_item_ids[item_id] == {"role": "assistant", "content": "hello"} + assert session.delete_history_item(item_id) is True + assert session.history == () + + +# ---- turn signals / model turns ---- + + +@pytest.mark.parametrize( + ("event_type", "turn_state"), + [ + (DuplexTurnEventType.USER_STARTED, DuplexTurnState.USER_SPEAKING), + (DuplexTurnEventType.USER_COMMITTED, DuplexTurnState.USER_COMMITTED), + (DuplexTurnEventType.ASSISTANT_STARTED, DuplexTurnState.ASSISTANT_GENERATING), + (DuplexTurnEventType.ASSISTANT_DONE, DuplexTurnState.IDLE), + (DuplexTurnEventType.BARGE_IN, DuplexTurnState.BARGE_IN), + ], +) +def test_signal_turn_transitions_and_returns_typed_turn_event(event_type, turn_state): + session = _session() + + event = session.signal_turn(event_type.value) + + assert isinstance(event, TurnEvent) + assert session.turn_state == turn_state + wire = event.to_realtime() + assert wire["type"] == "turn.event" + assert wire["event"] == event_type.value + assert wire["turn_state"] == turn_state.value + + +def test_signal_turn_playback_ack_updates_cursor_and_close_marks_session_closing(): + session = _session() + session.mark_audio_sent(duration_ms=5_000) + + ack = session.signal_turn("playback_ack", {"played_ms": 1_500, "committed_ms": 1_000}) + assert ack.event == "playback_ack" + assert session.playback.played_ms == 1_500 + assert session.playback.committed_ms == 1_000 + + close = session.signal_turn(DuplexTurnEventType.CLOSE.value) + assert close.event == "close" + assert session.state.value == "closing" + + +def test_complete_model_turn_advances_turn_id_and_fence(): + session = _session() + assert session.fence == DuplexFence(session.session_id, epoch=0, turn_id=0) + + session.complete_model_turn(0) + + assert session.turn_id == 1 + assert session.fence == DuplexFence(session.session_id, epoch=0, turn_id=1) + assert session.accepted_fence == session.fence + + # A stale terminal for an already-completed turn does not move identity backwards. + session.complete_model_turn(0) + assert session.turn_id == 1 + session.complete_model_turn(3) + assert session.turn_id == 4 + assert session.accepted_fence.turn_id == 4 + + +# ---- fence validation ---- + + +def test_accept_fence_rejects_stale_or_foreign_fences(): + session = _session("sid-fence") + session.complete_model_turn(0) + + with pytest.raises(DuplexFenceMismatchError): + session.accept_fence(DuplexFence("sid-fence", epoch=0, turn_id=0)) + with pytest.raises(DuplexFenceMismatchError): + session.accept_fence(DuplexFence("sid-other", epoch=0, turn_id=1)) + + session.accept_fence(DuplexFence("sid-fence", epoch=1, turn_id=0)) + assert session.accepted_fence == DuplexFence("sid-fence", epoch=1, turn_id=0) + + +def test_prepare_and_commit_append_sequence_chunks_per_turn_and_epoch(): + session = _session("sid-append") + fence = session.fence + + first = session.commit_append(session.prepare_append(fence)) + second = session.commit_append(session.prepare_append(fence)) + assert (first.seq, first.turn_seq, first.turn_id) == (1, 1, 0) + assert (second.seq, second.turn_seq, second.turn_id) == (2, 2, 0) + + session.complete_model_turn(0) + third = session.commit_append(session.prepare_append(session.fence)) + assert (third.seq, third.turn_seq, third.turn_id) == (3, 1, 1) + + # Barge-in advances the epoch (turn_id is kept) and restarts the append sequence. + session.barge_in() + assert session.fence == DuplexFence("sid-append", epoch=1, turn_id=1) + fourth = session.commit_append(session.prepare_append(session.fence)) + assert (fourth.seq, fourth.turn_seq, fourth.turn_id) == (1, 1, 1) + assert session.input_seq == 1 + + +def test_commit_append_rejects_stale_reservation(): + session = _session("sid-append-stale") + fence = session.fence + reservation = session.prepare_append(fence) + session.commit_append(session.prepare_append(fence)) + + with pytest.raises(RuntimeError, match="stale"): + session.commit_append(reservation) + + with pytest.raises(DuplexFenceMismatchError): + session.prepare_append(DuplexFence("sid-other", epoch=0, turn_id=0)) + + +def test_cancel_fence_releases_stage_requests_and_advances_identity(): + session = _session("sid-cancel") + cancelled = session.fence + session.reserve_stage_request(0, "req-a", fence=cancelled) + session.bind_stage_request(1, "req-b", fence=cancelled) + assert session.stage_request_submitted(1, "req-b") is True + assert session.stage_request_submitted(0, "req-a") is False + next_fence = DuplexFence("sid-cancel", epoch=1, turn_id=0) + + stale = session.cancel_fence(cancelled, next_fence) + + assert stale == ["req-a", "req-b"] + assert session.resource_request_ids() == [] + assert session.accepted_fence == next_fence + + with pytest.raises(DuplexFenceMismatchError): + session.cancel_fence(next_fence, DuplexFence("sid-cancel", epoch=1, turn_id=1)) + with pytest.raises(DuplexFenceMismatchError): + session.cancel_fence(DuplexFence("sid-other", epoch=1, turn_id=0), DuplexFence("sid-cancel", epoch=2)) + + +# ---- public view ---- + + +def test_as_public_dict_exposes_identity_capabilities_and_playback(): + session = _session("duplex-public", config=DuplexSessionConfig(model="m", voice="alloy")) + session.mark_audio_sent(duration_ms=100) + + payload = session.as_public_dict() + + assert payload["id"] == "duplex-public" + assert payload["model"] == "m" + assert payload["voice"] == "alloy" + assert payload["state"] == "open" + assert payload["turn_state"] == "assistant_playing" + assert payload["epoch"] == 0 + assert payload["turn_id"] == 0 + assert payload["capabilities"] == DuplexCapabilities().as_dict() + assert payload["playback"] == {"generated_ms": 100, "sent_ms": 100, "played_ms": 0, "committed_ms": 0} + + +# ---- DuplexSessionConfig.from_realtime ---- + + +def test_from_realtime_rejects_unsupported_audio_formats(): + with pytest.raises(DuplexConfigError) as excinfo: + DuplexSessionConfig.from_realtime({"input_audio_format": "mp3"}) + assert excinfo.value.code == "unsupported_audio_format" + + with pytest.raises(DuplexConfigError) as excinfo: + DuplexSessionConfig.from_realtime({"audio": {"output": {"format": {"type": "audio/opus"}}}}) + assert excinfo.value.code == "unsupported_audio_format" + + +def test_from_realtime_validates_turn_detection(): + with pytest.raises(DuplexConfigError) as excinfo: + DuplexSessionConfig.from_realtime({"turn_detection": {"type": "semantic_vad"}}) + assert excinfo.value.code == "unsupported_turn_detection" + assert excinfo.value.param == "turn_detection" + + with pytest.raises(DuplexConfigError) as excinfo: + DuplexSessionConfig.from_realtime({"overlap_policy": "barge_in_on_speech"}) + assert excinfo.value.code == "unsupported_turn_detection" + + server_vad = DuplexSessionConfig.from_realtime({"turn_detection": {"type": "server_vad"}}) + assert server_vad.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value + assert server_vad.extra_body["realtime_turn_detection"]["threshold"] == 0.5 + + model_owned = DuplexSessionConfig.from_realtime({"turn_detection": None}) + assert model_owned.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value + assert model_owned.extra_body["realtime_turn_detection"] is None + + +def test_from_realtime_maps_wire_fields_and_stores_realtime_keys_in_extra_body(): + config = DuplexSessionConfig.from_realtime( + { + "model": "openbmb/MiniCPM-o-4_5", + "instructions": "be brief", + "output_audio_format": "pcm16", + "max_response_output_tokens": "inf", + "audio": {"output": {"voice": "alloy", "speed": 1.25}}, + "tools": [{"type": "function", "name": "lookup"}], + "tool_choice": "auto", + "metadata": {"tenant": "t1"}, + "include": ["item.input_audio_transcription.logprobs"], + "input_audio_transcription": {"model": "whisper-1"}, + "extra_body": {"custom": True}, + } + ) + + assert config.model == "openbmb/MiniCPM-o-4_5" + assert config.instructions == "be brief" + assert config.voice == "alloy" + assert config.speed == 1.25 + assert config.max_tokens is None + assert config.response_format == "pcm" + assert config.modalities == ["text", "audio"] + assert config.extra_body["custom"] is True + assert config.extra_body["realtime_tools"] == [{"type": "function", "name": "lookup"}] + assert config.extra_body["realtime_tool_choice"] == "auto" + assert config.extra_body["realtime_metadata"] == {"tenant": "t1"} + assert config.extra_body["realtime_include"] == ["item.input_audio_transcription.logprobs"] + assert config.extra_body["realtime_input_audio_transcription"] == {"model": "whisper-1"} + assert config.extra_body["realtime_audio"] == {"output": {"voice": "alloy", "speed": 1.25}} + assert config.extra_body["realtime_output_audio_format"] == "pcm16" + assert "extra_body" not in config.extra_body["realtime_session_payload"] + assert config.extra_body["realtime_session_payload"]["model"] == "openbmb/MiniCPM-o-4_5" + + +def test_from_realtime_uses_served_model_when_payload_has_none(): + config = DuplexSessionConfig.from_realtime({}, model="served-model") + + assert config.model == "served-model" + assert config.response_format == "pcm" + assert config.idle_timeout_s == 300.0 + + +def test_from_realtime_ignores_client_chosen_session_ids(): + config = DuplexSessionConfig.from_realtime({"id": "client-chosen", "session_id": "also-client-chosen"}) + + assert not hasattr(config, "session_id") + assert "id" not in config.as_dict() + session = _session("duplex-engine-allocated", config=config) + assert session.as_public_dict()["id"] == "duplex-engine-allocated" + + +# ---- DuplexSessionConfig.apply_realtime_update ---- + + +def test_apply_realtime_update_rejects_model_change(): + config = DuplexSessionConfig(model="served-model") + + with pytest.raises(DuplexConfigError) as excinfo: + config.apply_realtime_update({"model": "other-model"}, session_id="sid") + + assert excinfo.value.code == "model_update_unsupported" + assert config.model == "served-model" + + +def test_apply_realtime_update_rejects_voice_change_after_audio_started(): + config = DuplexSessionConfig(model="m", voice="alloy") + + config.apply_realtime_update({"voice": "verse"}, audio_started=False) + assert config.voice == "verse" + with pytest.raises(DuplexConfigError) as excinfo: + config.apply_realtime_update({"audio": {"output": {"voice": "alloy"}}}, audio_started=True) + assert excinfo.value.code == "voice_update_after_audio_unsupported" + assert config.voice == "verse" + + +def test_apply_realtime_update_rejects_ref_audio_change(): + config = DuplexSessionConfig(model="m") + + with pytest.raises(DuplexConfigError) as excinfo: + config.apply_realtime_update({"ref_audio": "/tmp/voice.wav"}) + + assert excinfo.value.code == "ref_audio_update_unsupported" + assert config.ref_audio is None + + +def test_apply_realtime_update_patches_fields_and_realtime_keys(): + config = DuplexSessionConfig(model="m", instructions="old", extra_body={"realtime_tools": [{"name": "x"}]}) + + config.apply_realtime_update( + { + "instructions": "new", + "output_audio_format": "g711_ulaw", + "max_output_tokens": 32, + "tools": None, + "metadata": {"k": "v"}, + "overlap_policy": "not-a-policy", + } + ) + + assert config.instructions == "new" + assert config.response_format == "pcm" + assert config.max_tokens == 32 + assert "realtime_tools" not in config.extra_body + assert config.extra_body["realtime_metadata"] == {"k": "v"} + assert config.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value + assert config.extra_body["realtime_session_payload"]["instructions"] == "new" + + +# ---- overlap policy / capabilities ---- + + +def test_overlap_policy_defaults_and_invalid_values_to_listen_only(): + assert DuplexSessionConfig().overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value + assert DuplexSessionConfig._normalize_overlap_policy("auto") == DuplexOverlapPolicy.LISTEN_ONLY.value + assert DuplexSessionConfig._normalize_overlap_policy("not-a-policy") == DuplexOverlapPolicy.LISTEN_ONLY.value + assert ( + DuplexSessionConfig._normalize_overlap_policy(" Barge_In_On_Speech ") + == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value + ) + assert DuplexSessionConfig.from_event({"session": {"overlap_policy": "bogus"}}).overlap_policy == "listen_only" + + +def test_capabilities_as_dict_reports_model_native_duplex_constants(): + caps = DuplexCapabilities().as_dict() + + assert caps["implementation_level"] == "model_native_duplex" + assert caps["input_modes"] == ["append_audio_chunk"] + assert caps["supports_kv_lease"] is False + assert caps["supports_core_kv_lease"] is False + assert "client_event" in caps["signal_sources"] + assert set(caps) == set(minicpmo45_native_capabilities().as_dict()) + + +def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease(): + caps = minicpmo45_native_capabilities(max_sessions=2).as_dict() + + assert caps["implementation_level"] == "model_native_duplex" + assert caps["supports_input_append"] is True + assert caps["input_modes"] == ["append_audio_chunk"] + assert caps["adapter_patterns"] == ["scheduler_data_plane"] + assert caps["supports_model_internal_state"] is True + assert caps["requires_model_runner_kv"] is True + assert caps["requires_native_stage_role"] is True + assert caps["supports_kv_lease"] is False + assert caps["supports_core_kv_lease"] is False + assert caps["supports_stage_resumption"] is True + assert caps["supports_scheduler_native_append"] is False + assert caps["supports_core_resumable_request"] is True + assert caps["supports_stage_connector_handoff"] is True + assert caps["supports_audio_truncate"] is True + assert caps["supports_barge_in"] is True + assert caps["target_barge_in_latency_ms"] is None + assert caps["supports_multi_session"] is True + assert caps["supports_multi_session_same_replica"] is True + assert caps["supports_session_lease"] is True + assert caps["supports_session_resume"] is True + assert caps["session_admission_mode"] == "engine_managed" + assert caps["stage_handoff_transport"] == "scheduler_data_plane" + + +def test_minicpmo_native_capabilities_do_not_overclaim_single_session_deployment(): + caps = minicpmo45_native_capabilities(max_sessions=1).as_dict() + + assert caps["supports_multi_session"] is False + assert caps["supports_multi_session_same_replica"] is False diff --git a/tests/engine/duplex/test_overlap_policy.py b/tests/engine/duplex/test_overlap_policy.py new file mode 100644 index 000000000..096b7c7c9 --- /dev/null +++ b/tests/engine/duplex/test_overlap_policy.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The overlap decision, reached directly. + +These rules decide what happens when input arrives while the model is speaking. +They were unreachable except by driving a whole append through the runner, so +none of them had a test that named the rule it was checking. +""" + +from __future__ import annotations + +import base64 + +import numpy as np +import pytest + +from vllm_omni.engine.duplex.config import DuplexSessionConfig +from vllm_omni.engine.duplex.session import overlap_policy +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _session(**config: object) -> DuplexEngineSession: + values: dict[str, object] = {"model": "fake-model"} + values.update(config) + return DuplexEngineSession(session_id="sid", config=DuplexSessionConfig(**values)) + + +def _pcm16(rms: float, samples: int = 1600) -> str: + data = np.full(samples, int(rms * 32768), dtype=np.int16) + return base64.b64encode(data.tobytes()).decode("ascii") + + +# --------------------------------------------------------------------------- # +# Speech detection # +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("flag", [True, False]) +def test_an_explicit_is_speech_flag_wins_over_the_audio(flag: bool) -> None: + session = _session() + payload = {"format": "pcm16", "audio": _pcm16(0.9 if not flag else 0.0)} + assert overlap_policy.input_looks_like_speech(session, {"is_speech": flag}, payload) is flag + + +def test_vad_probability_decides_when_no_flag_is_given() -> None: + session = _session() + payload = {"format": "pcm16", "audio": _pcm16(0.0)} + assert overlap_policy.input_looks_like_speech(session, {"vad": {"speech_probability": 0.9}}, payload) is True + assert overlap_policy.input_looks_like_speech(session, {"vad": {"speech_probability": 0.1}}, payload) is False + + +def test_quiet_audio_below_the_configured_rms_is_not_speech() -> None: + session = _session(overlap_silence_rms=0.05) + loud = {"format": "pcm16", "audio": _pcm16(0.5)} + quiet = {"format": "pcm16", "audio": _pcm16(0.001)} + assert overlap_policy.input_looks_like_speech(session, {}, loud) is True + assert overlap_policy.input_looks_like_speech(session, {}, quiet) is False + + +def test_undecodable_audio_is_treated_as_speech() -> None: + """Fail toward listening: dropping real speech is worse than a spurious turn.""" + session = _session() + assert overlap_policy.input_looks_like_speech(session, {}, {"format": "pcm16", "audio": "!!not base64!!"}) is True + + +# --------------------------------------------------------------------------- # +# Barge-in # +# --------------------------------------------------------------------------- # + + +def test_an_explicit_client_overlap_action_requests_barge_in() -> None: + session = _session() + payload = {"format": "pcm16", "audio": _pcm16(0.5)} + decision = overlap_policy.decide(session, {"overlap_action": "barge_in"}, payload, auto_responds=True) + assert decision["action"] == "barge_in" + assert decision["reason"] == "client_overlap_action" + + +def test_barge_in_is_deferred_when_the_session_cannot_support_it() -> None: + session = _session() + session.capabilities = type(session.capabilities)(supports_barge_in=False) + payload = {"format": "pcm16", "audio": _pcm16(0.5)} + decision = overlap_policy.decide(session, {"force_barge_in": True}, payload, auto_responds=True) + assert decision["action"] != "barge_in" + + +# --------------------------------------------------------------------------- # +# Force-listen # +# --------------------------------------------------------------------------- # + + +def test_a_short_response_create_is_forced_to_listen() -> None: + """A sub-threshold ack should not take the turn from the model.""" + session = _session(overlap_short_ack_ms=700) + short = {"response_create": True, "duration_ms": 200} + long = {"response_create": True, "duration_ms": 1500} + assert overlap_policy.should_force_listen_for_short_commit(session, short, {}) is True + assert overlap_policy.should_force_listen_for_short_commit(session, long, {}) is False + + +def test_force_barge_in_overrides_the_short_ack_rule() -> None: + session = _session(overlap_short_ack_ms=700) + event = {"response_create": True, "duration_ms": 200, "force_barge_in": True} + assert overlap_policy.should_force_listen_for_short_commit(session, event, {}) is False + + +def test_auto_response_overlap_forces_listen_only_when_auto_responding() -> None: + event = {"force_listen": True} + assert overlap_policy.should_force_listen_for_auto_response_overlap(event, {}, auto_responds=True) is True + assert overlap_policy.should_force_listen_for_auto_response_overlap(event, {}, auto_responds=False) is False + + +# --------------------------------------------------------------------------- # +# Payload merging # +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("speech_first", [True, False], ids=["speech-then-silence", "silence-then-speech"]) +def test_merging_two_appends_keeps_speech_if_either_half_was_speech(speech_first: bool) -> None: + """Matching rates and decodable audio, so the halves really are concatenated. + + Without a ``sample_rate_hz`` on both, ``merge_audio_payloads`` falls back + to the newer chunk before merging, and the speech flag it returned was the + second chunk's own -- the test passed without exercising the merge. + """ + speech = { + "format": "pcm_f32le", + "audio": base64.b64encode(b"\x01" * 8).decode("ascii"), + "sample_rate_hz": 16000, + "is_speech": True, + } + silence = { + "format": "pcm_f32le", + "audio": base64.b64encode(b"\x00" * 8).decode("ascii"), + "sample_rate_hz": 16000, + "is_speech": False, + } + first, second = (speech, silence) if speech_first else (silence, speech) + merged = overlap_policy.merge_audio_payloads(first, second) + assert base64.b64decode(merged["audio"]) == base64.b64decode(first["audio"]) + base64.b64decode(second["audio"]) + assert merged["sample_rate_hz"] == 16000 + assert merged["is_speech"] is True + + +@pytest.mark.parametrize( + ("first", "second"), + [ + ({"format": "pcm_f32le", "audio": "", "sample_rate_hz": 16000}, {"format": "pcm16", "audio": ""}), + ( + {"format": "pcm_f32le", "audio": "", "sample_rate_hz": 16000}, + {"format": "pcm_f32le", "audio": "", "sample_rate_hz": 24000}, + ), + ( + {"format": "pcm_f32le", "audio": "!!not base64!!", "sample_rate_hz": 16000}, + {"format": "pcm_f32le", "audio": "", "sample_rate_hz": 16000}, + ), + ], + ids=["format-mismatch", "rate-mismatch", "undecodable"], +) +def test_unmergeable_payloads_fall_back_to_the_newer_one(first: dict, second: dict) -> None: + """Whatever cannot be concatenated keeps the newer chunk and drops the older. + + Worth pinning: the fallback discards ``first`` silently, so a bug that makes + two mergeable payloads look mismatched loses buffered audio rather than + raising. + """ + assert overlap_policy.merge_audio_payloads(first, second) == second + + +def test_mergeable_payloads_are_concatenated_in_order() -> None: + head = bytes([1, 2, 3, 4]) + tail = bytes([5, 6, 7, 8]) + first = {"format": "pcm_f32le", "audio": base64.b64encode(head).decode(), "sample_rate_hz": 16000} + second = {"format": "pcm_f32le", "audio": base64.b64encode(tail).decode(), "sample_rate_hz": 16000} + merged = overlap_policy.merge_audio_payloads(first, second) + assert base64.b64decode(merged["audio"]) == head + tail + assert merged["sample_rate_hz"] == 16000 + + +# --------------------------------------------------------------------------- # +# Server VAD + barge_in_on_speech # +# --------------------------------------------------------------------------- # + + +def test_server_vad_speech_started_reaches_a_barge_in_decision() -> None: + """The configuration this PR introduces must decide, not raise. + + Both call sites read the VAD verdict through the module-level + ``vad_speech_started`` helper while binding its result to a local of the + same name, which made the name local for the whole scope and raised + ``UnboundLocalError`` on the call itself. Every ``server_vad`` + + ``barge_in_on_speech`` append hit it, and the client saw ``internal_error`` + instead of a barge-in. The MiniCPM-o preset is ``listen_only``, so no demo + exercised this branch. + """ + session = _session(overlap_policy="barge_in_on_speech") + payload = {"format": "pcm16", "audio": _pcm16(0.5)} + event = {"vad": {"speech_started": True}} + + decision = overlap_policy.decide(session, event, payload, auto_responds=True) + + assert decision["action"] == "barge_in" + assert decision["reason"] == "server_vad_speech_started" + assert decision["cancel_reason"] == "turn_detected" + + +def test_server_vad_mid_utterance_keeps_listening_instead_of_barging_in() -> None: + """``speech_started`` false means the utterance is already running: not a new turn.""" + session = _session(overlap_policy="barge_in_on_speech") + payload = {"format": "pcm16", "audio": _pcm16(0.5)} + event = {"vad": {"speech_started": False}} + + decision = overlap_policy.decide(session, event, payload, auto_responds=True) + + assert decision["action"] != "barge_in" + + +def test_barge_in_on_speech_without_vad_still_barges_in_on_speech() -> None: + """No VAD verdict: the policy falls back to its own speech classification.""" + session = _session(overlap_policy="barge_in_on_speech") + payload = {"format": "pcm16", "audio": _pcm16(0.5)} + + decision = overlap_policy.decide(session, {}, payload, auto_responds=True) + + assert decision["action"] == "barge_in" + assert decision["reason"] == "barge_in_on_speech" diff --git a/tests/engine/duplex/test_session_manager.py b/tests/engine/duplex/test_session_manager.py new file mode 100644 index 000000000..0323f16d8 --- /dev/null +++ b/tests/engine/duplex/test_session_manager.py @@ -0,0 +1,1569 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""``DuplexSessionManager``: admission, control RPC, command backpressure, reaper and request cleanup.""" + +from __future__ import annotations + +import asyncio +import base64 +import logging +from dataclasses import FrozenInstanceError, dataclass, field +from typing import Any + +import pytest + +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.engine.duplex.commands import AppendAudio, CloseSession, Commit, DuplexCommand, Heartbeat +from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig, DuplexSessionState +from vllm_omni.engine.duplex.contracts import ( + DuplexAppendPlan, + DuplexFence, + DuplexStagePort, + DuplexStageRequestContext, + DuplexStageSubmission, + DuplexStageSubmissionResult, + duplex_resource_request_belongs_to_session, +) +from vllm_omni.engine.duplex.events import ( + DuplexEvent, + ErrorEvent, + SessionClosed, + SessionCreated, + SessionExpired, + SessionHeartbeatAck, +) +from vllm_omni.engine.duplex.messages import ( + CloseDuplexSessionMessage, + DuplexControlResultMessage, + DuplexSessionCommandMessage, + DuplexSessionEventMessage, + OpenDuplexSessionMessage, + ResumeDuplexSessionMessage, + TouchDuplexSessionMessage, +) +from vllm_omni.engine.duplex.plugin import ( + DuplexDataPlane, + DuplexModelPlugin, + DuplexModelSessionState, + DuplexRuntimeConfigError, + PcmAppendBuffer, + PcmAppendReservation, +) +from vllm_omni.engine.duplex.session import manager as session_manager_module +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession +from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity +from vllm_omni.engine.duplex.session.manager import DuplexSessionManager +from vllm_omni.engine.duplex.session.runner import _Internal + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +# --------------------------------------------------------------------------- # +# Fakes # +# --------------------------------------------------------------------------- # + + +class FakeClock: + def __init__(self, value: float = 0.0) -> None: + self.value = value + + def __call__(self) -> float: + return self.value + + def advance(self, seconds: float) -> None: + self.value += seconds + + +class FakePcmAppendBuffer(PcmAppendBuffer): + @property + def pending_byte_count(self) -> int: + return 0 + + def clear(self) -> None: + pass + + def clear_force_listen(self) -> None: + pass + + def has_pending(self) -> bool: + return False + + def has_reserved(self) -> bool: + return False + + def prepare_append(self, payload, *, operation_id, chunk_period_ms, allow_emit) -> PcmAppendReservation | None: + del payload, operation_id, chunk_period_ms, allow_emit + return None + + def prepare_commit(self, *, operation_id, chunk_period_ms) -> PcmAppendReservation: + del operation_id, chunk_period_ms + raise NotImplementedError("fake buffer never commits") + + def flush(self, *, chunk_period_ms) -> dict[str, object] | None: + del chunk_period_ms + return None + + +class FakeModelSessionState(DuplexModelSessionState): + def __init__(self) -> None: + self.audio_buffer = FakePcmAppendBuffer() + self.input_since_commit = False + self.speech_since_commit = False + self.context_locked = False + self.committed_audio_payload = None + self.committed_audio_operation_id = None + self.committed_audio_reserved_bytes = 0 + self.deferred_response_create = False + self.deferred_precreate_response = False + self.continuation_owner_id = None + self.continuation_units = 0 + self.pending_silence_task = None + self.pending_silence_owner_id = None + + def retain_committed_audio(self, payload, *, operation_id, reserved_bytes=0) -> None: + self.committed_audio_payload = payload + self.committed_audio_operation_id = operation_id + self.committed_audio_reserved_bytes = reserved_bytes + + def clear_committed_audio(self) -> int: + reserved = self.committed_audio_reserved_bytes + self.committed_audio_payload = None + self.committed_audio_operation_id = None + self.committed_audio_reserved_bytes = 0 + return reserved + + def clear_continuation(self) -> None: + self.continuation_owner_id = None + self.continuation_units = 0 + + +class FakeDataPlane(DuplexDataPlane): + def __init__(self) -> None: + self.closed_sessions: list[str] = [] + + def begin_request(self, request_id: str) -> None: + pass + + def is_terminal(self, request_id: str | None) -> bool: + return False + + def mark_terminal(self, request_id: str) -> None: + pass + + def close_stream(self, request_id: str) -> None: + pass + + def close_session(self, session_id: str, *, active_request_id: str | None = None) -> None: + self.closed_sessions.append(session_id) + + def project(self, result, *, context=None): + del result, context + return [] + + +class FakePlugin(DuplexModelPlugin): + plugin_id = "fake-duplex" + + def __init__(self, encode_audio=None) -> None: + super().__init__(encode_audio or (lambda *args: None)) + self.data_plane = FakeDataPlane() + self.runtime_config: dict[str, object] = {"runtime": "test"} + self.runtime_config_error: Exception | None = None + #: Session ids whose ``prepare_runtime_config`` waits for ``runtime_config_gate``. + self.blocked_session_ids: set[str] = set() + self.runtime_config_gate = asyncio.Event() + self.runtime_config_started = asyncio.Event() + + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return tuple(f"configured-{stage_id}" for stage_id, _ in enumerate(defaults)) + + def plan_append(self, **kwargs) -> DuplexAppendPlan: + del kwargs + return DuplexAppendPlan(prompt={"prompt_token_ids": [1, 2, 3]}) + + def decide_output(self, **kwargs): + del kwargs + return None + + def create_session_state(self) -> DuplexModelSessionState: + return FakeModelSessionState() + + def capabilities(self, *, max_sessions: int) -> DuplexCapabilities: + del max_sessions + return DuplexCapabilities(supports_input_append=True) + + def validate_client_extra_body(self, extra_body: object) -> None: + pass + + async def prepare_runtime_config(self, config: DuplexSessionConfig, *, model_config: Any) -> dict[str, object]: + del model_config + if self.runtime_config_error is not None: + raise self.runtime_config_error + if config.instructions in self.blocked_session_ids: + self.runtime_config_started.set() + await self.runtime_config_gate.wait() + return dict(self.runtime_config) + + def runtime_config_for_update(self, config, current): + del config + return dict(current) + + def data_plane_context(self, **kwargs) -> object: + return dict(kwargs) + + +class FakeStagePort(DuplexStagePort): + def __init__(self, stage_count: int = 2) -> None: + self._stage_count = stage_count + self.ensure_calls: list[DuplexStageRequestContext] = [] + self.submit_calls: list[DuplexStageSubmission] = [] + #: Every cleanup attempt, including the ones that were made to fail. + self.cleanup_calls: list[tuple[list[str], bool]] = [] + self.abort_calls: list[list[str]] = [] + self.cleanup_failures_remaining = 0 + self.cleanup_gate: asyncio.Event | None = None + self.cleanup_started = asyncio.Event() + + @property + def stage_count(self) -> int: + return self._stage_count + + def sampling_defaults(self) -> tuple[object, ...]: + return tuple(f"default-{stage_id}" for stage_id in range(self._stage_count)) + + def ensure_request(self, context: DuplexStageRequestContext) -> None: + self.ensure_calls.append(context) + + async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: + self.submit_calls.append(submission) + return DuplexStageSubmissionResult( + request_id=submission.context.request_id, + stage_id=submission.context.stage_id, + replica_id=3, + ) + + async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: + self.cleanup_calls.append((list(request_ids), abort)) + self.cleanup_started.set() + if self.cleanup_gate is not None: + await self.cleanup_gate.wait() + if self.cleanup_failures_remaining: + self.cleanup_failures_remaining -= 1 + raise RuntimeError("transient cleanup failure") + + async def abort_requests(self, request_ids: list[str]) -> None: + self.abort_calls.append(list(request_ids)) + + +# --------------------------------------------------------------------------- # +# Harness # +# --------------------------------------------------------------------------- # + + +def _runtime_config(**overrides: Any) -> DuplexSessionRuntimeConfig: + values: dict[str, Any] = { + "idle_ttl_s": 300.0, + "disconnect_grace_s": 30.0, + "max_sessions": 4, + "max_pending_input_bytes_per_session": 1024, + "max_pending_turns_per_session": 2, + } + values.update(overrides) + return DuplexSessionRuntimeConfig(**values) + + +def stage0_request_id(session_id: str, *, epoch: int = 0) -> str: + encoded = base64.urlsafe_b64encode(session_id.encode("utf-8")).decode("ascii").rstrip("=") + return f"duplex-s.{encoded}.e.{epoch}.r.stage0" + + +@dataclass +class Harness: + manager: DuplexSessionManager + plugin: FakePlugin + stage_port: FakeStagePort + clock: FakeClock + output_sink: asyncio.Queue = field(default_factory=asyncio.Queue) + result_sink: asyncio.Queue = field(default_factory=asyncio.Queue) + + @classmethod + def create( + cls, + *, + plugin: FakePlugin | None = None, + stage_port: FakeStagePort | None = None, + clock: FakeClock | None = None, + **runtime_overrides: Any, + ) -> Harness: + plugin = plugin or FakePlugin() + stage_port = stage_port or FakeStagePort() + clock = clock or FakeClock() + output_sink: asyncio.Queue = asyncio.Queue() + result_sink: asyncio.Queue = asyncio.Queue() + manager = DuplexSessionManager( + plugin=plugin, + stage_port=stage_port, + output_sink=output_sink, + result_sink=result_sink, + runtime_config=_runtime_config(**runtime_overrides), + model_config=None, + clock=clock, + ) + return cls(manager, plugin, stage_port, clock, output_sink, result_sink) + + async def __aenter__(self) -> Harness: + return self + + async def __aexit__(self, *exc_info: object) -> None: + await self.manager.shutdown() + + async def open(self, session_id: str, config: DuplexSessionConfig | None = None) -> DuplexControlResultMessage: + await self.manager.handle( + OpenDuplexSessionMessage( + control_id=f"open-{session_id}", + session_id=session_id, + session_config=config or DuplexSessionConfig(model="fake-model"), + ) + ) + return await self.result() + + async def close(self, session_id: str, *, reason: str = "client_close") -> DuplexControlResultMessage: + await self.manager.handle( + CloseDuplexSessionMessage(control_id=f"close-{session_id}", session_id=session_id, reason=reason) + ) + return await self.result() + + async def resume(self, session_id: str, *, expected_lease_generation: int) -> DuplexControlResultMessage: + await self.manager.handle( + ResumeDuplexSessionMessage( + control_id=f"resume-{session_id}-{expected_lease_generation}", + session_id=session_id, + expected_lease_generation=expected_lease_generation, + ) + ) + return await self.result() + + async def touch(self, session_id: str, activity: str) -> DuplexControlResultMessage: + await self.manager.handle( + TouchDuplexSessionMessage( + control_id=f"touch-{session_id}-{activity}", session_id=session_id, activity=activity + ) + ) + return await self.result() + + def command(self, session_id: str, command: DuplexCommand) -> None: + self.manager.dispatch(DuplexSessionCommandMessage(session_id=session_id, command=command)) + + async def result(self) -> DuplexControlResultMessage: + result = await asyncio.wait_for(self.result_sink.get(), timeout=2.0) + assert isinstance(result, DuplexControlResultMessage) + return result + + def events(self, session_id: str | None = None) -> list[DuplexEvent]: + events: list[DuplexEvent] = [] + while not self.output_sink.empty(): + message = self.output_sink.get_nowait() + assert isinstance(message, DuplexSessionEventMessage) + assert message.event.session_id == message.session_id + if session_id is None or message.session_id == session_id: + events.append(message.event) + return events + + def session(self, session_id: str) -> DuplexEngineSession: + session = self.manager.get(session_id) + assert session is not None + return session + + 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] + submitted: list[DuplexCommand] = [] + runner.submit = submitted.append # type: ignore[method-assign] + return submitted + + +async def _settle() -> None: + for _ in range(10): + await asyncio.sleep(0) + + +# --------------------------------------------------------------------------- # +# Open # +# --------------------------------------------------------------------------- # + + +async def test_open_answers_with_capabilities_and_emits_session_created() -> None: + async with Harness.create() as harness: + config = DuplexSessionConfig(model="fake-model", voice="test") + result = await harness.open("sid-open", config) + + assert result.ok is True + assert result.operation == "open" + assert result.control_id == "open-sid-open" + assert result.session_id == "sid-open" + assert result.lease_generation == 0 + assert result.capabilities == DuplexCapabilities(supports_input_append=True) + assert result.public_session is not None + assert result.public_session["id"] == "sid-open" + assert result.public_session["voice"] == "test" + assert result.error_code is None + + events = harness.events("sid-open") + assert isinstance(events[0], SessionCreated) + assert events[0].epoch == 0 + assert events[0].session["id"] == "sid-open" + wire = events[0].to_realtime() + assert wire["type"] == "session.created" + assert "incarnation" not in wire and "incarnation" not in wire["session"] + + session = harness.session("sid-open") + assert session.state == DuplexSessionState.OPEN + assert session.runtime_config == {"runtime": "test"} + assert isinstance(session.model_state, FakeModelSessionState) + assert harness.manager.active_count() == 1 + assert harness.manager.sessions() == {"sid-open": session} + + +async def test_open_reserves_the_fenced_stage0_request_with_a_frozen_context() -> None: + async with Harness.create() as harness: + await harness.open("sid-context", DuplexSessionConfig(model="fake-model", voice="test")) + + request_id = stage0_request_id("sid-context") + context = harness.stage_port.ensure_calls[-1] + assert isinstance(context, DuplexStageRequestContext) + assert context.request_id == request_id + assert context.session_id == "sid-context" + assert context.fence == DuplexFence("sid-context") + assert context.stage_id == 0 + assert context.final_stage_id == 1 + assert context.config_generation == 0 + assert context.sampling_params == ("configured-0", "configured-1") + assert context.stage_sampling_params == "configured-0" + assert context.session_config["voice"] == "test" + assert context.runtime_config == {"runtime": "test"} + with pytest.raises(FrozenInstanceError): + context.stage_id = 1 # type: ignore[misc] + + session = harness.session("sid-context") + assert session.resource_request_ids() == [request_id] + assert session.stage_request_submitted(0, request_id) is False + assert harness.manager.runner_for_request_id(request_id) is harness.manager.runners["sid-context"] + assert harness.stage_port.submit_calls == [] + + +def test_stage_request_id_is_derived_from_the_fence() -> None: + fence = DuplexFence("sid/with+special", epoch=2, turn_id=5) + + request_id = DuplexSessionManager.stage_request_id(fence, stage_id=0) + + assert request_id == stage0_request_id("sid/with+special", epoch=2) + assert duplex_resource_request_belongs_to_session(request_id, "sid/with+special") + assert not duplex_resource_request_belongs_to_session(request_id, "sid/other") + + +async def test_ensure_stage_request_ignores_stages_beyond_the_pipeline() -> None: + async with Harness.create() as harness: + await harness.open("sid-stages") + session = harness.session("sid-stages") + + assert harness.manager.ensure_stage_request(session, stage_id=2) is None + context = harness.manager.ensure_stage_request(session, stage_id=1) + + assert context is not None + assert context.request_id.endswith(".r.stage1") + assert harness.manager.runner_for_request_id(context.request_id) is harness.manager.runners["sid-stages"] + assert session.resource_request_ids() == [stage0_request_id("sid-stages"), context.request_id] + + +async def test_sampling_params_for_validates_the_plugin_policy() -> None: + class ShortSamplingPlugin(FakePlugin): + def configure_sampling_params(self, *, runtime_config, defaults): + if runtime_config: + return defaults[:1] + return tuple(defaults) + + class ListSamplingPlugin(FakePlugin): + def configure_sampling_params(self, *, runtime_config, defaults): + if runtime_config: + return list(defaults) + return tuple(defaults) + + async with Harness.create(plugin=ShortSamplingPlugin()) as harness: + result = await harness.open("sid-short") + assert result.ok is False + assert result.error_code == "invalid_argument" + assert "one sampling parameter per stage" in str(result.error_message) + assert harness.manager.get("sid-short") is None + assert harness.manager.active_count() == 0 + + async with Harness.create(plugin=ListSamplingPlugin()) as harness: + result = await harness.open("sid-list") + assert result.ok is False + assert result.error_code == "invalid_argument" + assert "as a tuple" in str(result.error_message) + + +async def test_duplicate_open_is_rejected_with_session_exists() -> None: + async with Harness.create() as harness: + assert (await harness.open("sid-dup")).ok is True + first_session = harness.session("sid-dup") + + result = await harness.open("sid-dup") + + assert result.ok is False + assert result.error_code == "session_exists" + assert result.error_retryable is False + assert harness.session("sid-dup") is first_session + assert harness.manager.active_count() == 1 + assert len(harness.stage_port.ensure_calls) == 1 + + +async def test_capacity_rejection_is_retryable_and_logged_at_info(caplog: pytest.LogCaptureFixture) -> None: + manager_logger = session_manager_module.logger + manager_logger.addHandler(caplog.handler) + try: + async with Harness.create(max_sessions=1) as harness: + assert (await harness.open("sid-first")).ok is True + caplog.clear() + + with caplog.at_level(logging.INFO, logger=manager_logger.name): + result = await harness.open("sid-rejected") + + assert result.ok is False + assert result.error_code == "resource_exhausted" + assert result.error_retryable is True + assert harness.manager.get("sid-rejected") is None + assert harness.manager.active_count() == 1 + rejections = [record for record in caplog.records if "rejected" in record.getMessage()] + assert rejections and all(record.levelno == logging.INFO for record in rejections) + assert not [record for record in caplog.records if record.levelno >= logging.ERROR] + finally: + manager_logger.removeHandler(caplog.handler) + + +@pytest.mark.parametrize( + ("error", "expected_code"), + [ + (ValueError("reference audio not found"), "invalid_duplex_runtime_config"), + (DuplexRuntimeConfigError("bad ref", code="unsupported_ref_audio_path"), "unsupported_ref_audio_path"), + (DuplexRuntimeConfigError("bad voice", code="invalid_voice"), "invalid_voice"), + ], +) +async def test_plugin_runtime_config_rejection_rolls_back_the_open(error: Exception, expected_code: str) -> None: + plugin = FakePlugin() + plugin.runtime_config_error = error + async with Harness.create(plugin=plugin) as harness: + result = await harness.open("sid-runtime") + + assert result.ok is False + assert result.error_code == expected_code + assert result.error_message == str(error) + assert result.capabilities is None + assert harness.manager.get("sid-runtime") is None + assert harness.manager.active_count() == 0 + assert harness.stage_port.ensure_calls == [] + assert harness.events() == [] + + # Nothing was registered under the rejected id and the slot is free again. + plugin.runtime_config_error = None + assert (await harness.open("sid-runtime")).ok is True + + +# --------------------------------------------------------------------------- # +# Command dispatch # +# --------------------------------------------------------------------------- # + + +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") + assert len(events) == 1 + error = events[0] + assert isinstance(error, ErrorEvent) + assert error.code == "unknown_session" + assert error.related_event_id == "evt-hb" + assert error.session_id == "sid-missing" + assert error.epoch is None + assert error.to_realtime()["error"]["code"] == "unknown_session" + + +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() + + harness.command("sid-closed", Heartbeat(event_id="evt-late")) + + events = 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" + + +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() + 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") + 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 + + +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() + session = harness.session("sid-bytes") + submitted = harness.capture_submissions("sid-bytes") + assert session.reserve_input_bytes(6, limit=8) + + 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") + assert len(events) == 1 + assert isinstance(events[0], ErrorEvent) + assert events[0].code == "input_backpressure" + assert events[0].related_event_id == "evt-over" + assert events[0].epoch == 0 + assert [command.event_id for command in submitted] == ["evt-fits"] + # Admission reserved the wire size of the admitted append; the runner + # released it on dequeue and re-reserved the decoded size (2 bytes). + assert session.pending_input_bytes == 8 + + +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() + session = harness.session("sid-turns") + submitted = harness.capture_submissions("sid-turns") + + harness.command("sid-turns", Commit(event_id="evt-commit-1")) + harness.command("sid-turns", Commit(event_id="evt-commit-2")) + + assert [command.event_id for command in submitted] == ["evt-commit-1"] + assert session.pending_input_turns == 1 + events = harness.events("sid-turns") + assert len(events) == 1 + assert isinstance(events[0], ErrorEvent) + assert events[0].code == "input_backpressure" + assert events[0].related_event_id == "evt-commit-2" + + # The runner releases the reservation when it dequeues the commit. + 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") == [] + + +def test_accepts_only_duplex_message_types() -> None: + manager = DuplexSessionManager( + plugin=FakePlugin(), + stage_port=FakeStagePort(), + output_sink=asyncio.Queue(), + result_sink=asyncio.Queue(), + runtime_config=_runtime_config(), + model_config=None, + clock=FakeClock(), + ) + try: + config = DuplexSessionConfig() + assert manager.accepts(OpenDuplexSessionMessage(control_id="c", session_id="s", session_config=config)) + 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")) + assert manager.accepts(DuplexSessionCommandMessage(session_id="s", command=Heartbeat())) + assert manager.accepts(type("Lookalike", (), {"type": "open_duplex_session", "session_id": "s"})()) is False + assert ( + manager.accepts(DuplexControlResultMessage(control_id="c", operation="open", session_id="s", ok=True)) + is False + ) + assert manager.accepts(Heartbeat()) is False + with pytest.raises(TypeError, match="Unsupported duplex control message"): + manager.dispatch(object()) + finally: + manager.executor.shutdown(wait=False) + + +# --------------------------------------------------------------------------- # +# Close / resume / touch # +# --------------------------------------------------------------------------- # + + +async def test_close_releases_stage_resources_and_is_idempotent() -> None: + async with Harness.create() as harness: + await harness.open("sid-close") + harness.events() + session = harness.session("sid-close") + request_id = stage0_request_id("sid-close") + session.bind_stage_request(0, request_id, fence=session.fence) + reserved_context = harness.manager.ensure_stage_request(session, stage_id=1) + assert reserved_context is not None + + result = await harness.close("sid-close", reason="client_close") + + assert result.ok is True + assert result.operation == "close" + assert result.lease_generation == 1 + assert result.public_session is not None and result.public_session["state"] == "closed" + assert session.state == DuplexSessionState.CLOSED + assert session.lease.terminal_reason == "client_close" + assert session.resource_request_ids() == [] + assert harness.manager.get("sid-close") is None + assert harness.manager.active_count() == 0 + 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)] + assert len(closed) == 1 + assert closed[0].reason == "client_close" + assert closed[0].is_terminal is True + + again = await harness.close("sid-close") + assert again.ok is True + assert again.lease_generation is None + assert again.public_session is None + assert harness.stage_port.cleanup_calls[2:] == [] + assert harness.events() == [] + + +async def test_close_retains_the_admission_slot_until_stage_cleanup_succeeds() -> None: + async with Harness.create(max_sessions=1) as harness: + await harness.open("sid-close-retry") + session = harness.session("sid-close-retry") + session.bind_stage_request(0, stage0_request_id("sid-close-retry"), fence=session.fence) + harness.stage_port.cleanup_failures_remaining = 1 + + result = await harness.close("sid-close-retry") + + # The session is closed from the client's point of view; only the + # stage cleanup is outstanding and the reaper retries it. + assert result.ok is True + assert harness.manager.get("sid-close-retry") is None + assert session.state == DuplexSessionState.CLOSED + assert (await harness.open("sid-replacement")).error_code == "resource_exhausted" + + assert await harness.manager.reap_expired() == 1 + assert harness.stage_port.cleanup_calls == [([stage0_request_id("sid-close-retry")], True)] * 2 + assert (await harness.open("sid-replacement")).ok is True + + +async def test_resume_requires_the_expected_lease_generation() -> None: + async with Harness.create() as harness: + await harness.open("sid-resume") + session = harness.session("sid-resume") + harness.clock.advance(3.0) + assert (await harness.touch("sid-resume", DuplexLeaseActivity.DETACH.value)).ok is True + assert session.lease.detached_at == 3.0 + + harness.clock.advance(1.0) + result = await harness.resume("sid-resume", expected_lease_generation=0) + + assert result.ok is True + assert result.operation == "resume" + assert result.lease_generation == 1 + assert result.public_session is not None and result.public_session["id"] == "sid-resume" + assert session.lease_generation == 1 + assert session.lease.detached_at is None + assert session.lease.last_activity == 4.0 + + stale = await harness.resume("sid-resume", expected_lease_generation=0) + assert stale.ok is False + assert stale.error_code == "session_resume_conflict" + assert stale.lease_generation == 1 + assert session.lease_generation == 1 + + unknown = await harness.resume("sid-unknown", expected_lease_generation=0) + assert unknown.ok is False + assert unknown.error_code == "unknown_session" + assert unknown.lease_generation is None + + +async def test_touch_records_activity_and_detach_marks_the_lease_detached() -> None: + async with Harness.create() as harness: + await harness.open("sid-touch") + session = harness.session("sid-touch") + assert session.lease.last_activity == 0.0 + + harness.clock.advance(2.0) + result = await harness.touch("sid-touch", DuplexLeaseActivity.HEARTBEAT.value) + assert result.ok is True + assert result.operation == "touch" + assert session.lease.last_activity == 2.0 + assert session.lease.detached_at is None + + harness.clock.advance(1.0) + assert (await harness.touch("sid-touch", DuplexLeaseActivity.DETACH.value)).ok is True + assert session.lease.last_activity == 3.0 + assert session.lease.detached_at == 3.0 + + harness.clock.advance(1.0) + assert (await harness.touch("sid-touch", DuplexLeaseActivity.ATTACH.value)).ok is True + assert session.lease.last_activity == 4.0 + # Only resume clears a detach. + assert session.lease.detached_at == 3.0 + + invalid = await harness.touch("sid-touch", "bogus") + assert invalid.ok is False + assert invalid.error_code == "invalid_argument" + + unknown = await harness.touch("sid-unknown", DuplexLeaseActivity.HEARTBEAT.value) + assert unknown.ok is False + assert unknown.error_code == "unknown_session" + assert harness.stage_port.submit_calls == [] + + +# --------------------------------------------------------------------------- # +# Reaper # +# --------------------------------------------------------------------------- # + + +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() + harness.clock.advance(1.0) + assert (await harness.touch("sid-active", DuplexLeaseActivity.HEARTBEAT.value)).ok is True + harness.clock.advance(1.1) + + assert await harness.manager.reap_expired() == 1 + + 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") + assert len(expired) == 1 + assert isinstance(expired[0], SessionExpired) + assert expired[0].reason == "idle_ttl_expired" + assert expired[0].to_realtime() == { + "type": "session.expired", + "event_id": expired[0].event_id, + "session_id": "sid-idle", + "reason": "idle_ttl_expired", + } + assert harness.events("sid-active") == [] + assert harness.stage_port.cleanup_calls == [([stage0_request_id("sid-idle")], False)] + assert harness.result_sink.empty() + + assert await harness.manager.reap_expired() == 0 + + +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() + assert (await harness.touch("sid-detached", DuplexLeaseActivity.DETACH.value)).ok is True + harness.clock.advance(0.5) + assert await harness.manager.reap_expired() == 0 + harness.clock.advance(0.5) + + assert await harness.manager.reap_expired() == 1 + + assert harness.manager.get("sid-detached") is None + expired = 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" + + +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() + 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() + + close_task = asyncio.create_task(harness.close("sid-closing")) + 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") == [] + + harness.stage_port.cleanup_gate.set() + assert (await close_task).ok is True + events = 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 + + +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() + session = harness.session("sid-wire-close") + request_id = stage0_request_id("sid-wire-close") + session.bind_stage_request(0, request_id, fence=session.fence) + harness.stage_port.cleanup_gate = asyncio.Event() + + 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") == [] + + harness.stage_port.cleanup_gate.set() + for _ in range(50): + if harness.manager._admission_count() == 0: + break + await asyncio.sleep(0.01) + assert harness.stage_port.cleanup_calls == [([request_id], True)] + events = 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 + + +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() + runner = harness.manager.runners["sid-wire"] + worker = runner._worker + assert worker is not None + harness.stage_port.cleanup_gate = asyncio.Event() + + harness.command("sid-wire", CloseSession(reason="client_close")) + await asyncio.wait_for(harness.stage_port.cleanup_started.wait(), timeout=1.0) + # Control operations issued while the close is in flight are refused, + # not silently accepted on a half-torn session. + resumed = await harness.resume("sid-wire", expected_lease_generation=0) + assert resumed.ok is False + assert resumed.error_code in {"session_closed", "unknown_session"} + harness.command("sid-wire", Heartbeat()) + + harness.stage_port.cleanup_gate.set() + for _ in range(50): + if worker.done() and harness.manager._admission_count() == 0: + break + await asyncio.sleep(0.01) + assert worker.done(), "the mailbox worker must exit after a wire close" + events = harness.events("sid-wire") + assert [type(event) for event in events] == [ErrorEvent, SessionClosed] or [ + type(event) for event in events + ] == [SessionClosed, ErrorEvent] + assert (await harness.open("sid-after")).ok is True + + +async def test_concurrent_opens_cannot_exceed_max_sessions() -> None: + async with Harness.create(max_sessions=1) as harness: + harness.plugin.blocked_session_ids.add("first") + first = asyncio.create_task( + harness.open("first", DuplexSessionConfig(model="fake-model", instructions="first")) + ) + await asyncio.wait_for(harness.plugin.runtime_config_started.wait(), timeout=1.0) + + # The first open holds the only slot while it awaits the plugin. + second = await harness.open("second", DuplexSessionConfig(model="fake-model", instructions="second")) + assert second.ok is False + assert second.error_code == "resource_exhausted" + + harness.plugin.runtime_config_gate.set() + assert (await first).ok is True + assert harness.manager.active_count() == 1 + + +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() + runner = harness.manager.runners["sid-runtime"] + runner.session.release_all_requests() + + await runner._close_from_runtime("model_failed") + for _ in range(50): + if harness.manager._admission_count() == 0: + break + 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 (await harness.open("sid-next")).ok is True + + +async def test_close_of_a_forgotten_session_is_idempotent() -> None: + async with Harness.create(max_sessions=1) as harness: + result = await harness.close("never-opened") + + assert result.ok is True + + +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() + session = harness.session("sid-bytes") + runner = harness.manager.runners["sid-bytes"] + # Park the worker so admitted commands stay in the mailbox. + gate = asyncio.Event() + runner._mailbox.put_nowait(_Internal("wait", {"gate": gate})) + runner._on_internal = lambda item: gate.wait() # type: ignore[method-assign] + await asyncio.sleep(0) + + 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)] + assert [error.code for error in errors] == ["input_backpressure"] + gate.set() + + +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() + session = harness.session("sid-expired") + request_id = stage0_request_id("sid-expired") + session.bind_stage_request(0, request_id, fence=session.fence) + harness.stage_port.cleanup_failures_remaining = 1 + harness.clock.advance(2.0) + + assert await harness.manager.reap_expired() == 0 + + 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 (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.open("sid-replacement")).ok is True + + +async def test_expired_session_holds_the_admission_slot_while_cleanup_is_blocked() -> None: + async with Harness.create(max_sessions=1, idle_ttl_s=1.0) as harness: + await harness.open("sid-blocked") + session = harness.session("sid-blocked") + session.bind_stage_request(0, stage0_request_id("sid-blocked"), fence=session.fence) + harness.stage_port.cleanup_gate = asyncio.Event() + harness.clock.advance(2.0) + + reap_task = asyncio.create_task(harness.manager.reap_expired()) + await asyncio.wait_for(harness.stage_port.cleanup_started.wait(), timeout=1.0) + assert (await harness.open("sid-replacement")).error_code == "resource_exhausted" + + harness.stage_port.cleanup_gate.set() + assert await reap_task == 1 + assert (await harness.open("sid-replacement")).ok is True + + +async def test_expired_cleanup_failure_does_not_block_other_sessions() -> None: + class OneStuckStagePort(FakeStagePort): + async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: + await super().cleanup(request_ids, abort=abort) + if request_ids == [stage0_request_id("sid-stuck")]: + raise RuntimeError("stuck expiry cleanup") + + 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() + 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 (await harness.open("sid-fine")).ok is True + assert (await harness.open("sid-stuck")).error_code == "session_exists" + + +# --------------------------------------------------------------------------- # +# Request-triggered cleanup (orchestrator error paths) # +# --------------------------------------------------------------------------- # + + +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() + session = harness.session("sid-request") + request_id = stage0_request_id("sid-request") + session.bind_stage_request(0, request_id, fence=session.fence) + harness.stage_port.cleanup_failures_remaining = 1 + + closed = harness.manager.close_sessions_for_request_ids([request_id, "req-unrelated"], abort=True) + + assert closed == {"sid-request": [request_id]} + assert harness.manager.get("sid-request") is None + assert session.lease.terminal_reason == "request_cleanup" + await _settle() + expired = harness.events("sid-request") + assert [type(event) for event in expired] == [SessionExpired] + assert expired[0].reason == "request_cleanup" + assert session.state == DuplexSessionState.CLOSED + # The stage resources are still owned by the pending cleanup. + assert session.resource_request_ids() == [request_id] + assert (await harness.open("sid-replacement")).error_code == "resource_exhausted" + + assert await harness.manager.reap_expired() == 0 + assert harness.stage_port.cleanup_calls == [([request_id], True)] + 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.manager.runner_for_request_id(request_id) is None + assert (await harness.open("sid-replacement")).ok is True + assert harness.manager.close_sessions_for_request_ids([request_id], abort=True) == {} + + +async def test_in_progress_request_cleanup_is_finalized_by_the_orchestrator_not_the_reaper() -> None: + async with Harness.create(max_sessions=1) as harness: + await harness.open("sid-in-progress") + session = harness.session("sid-in-progress") + request_id = stage0_request_id("sid-in-progress") + session.bind_stage_request(0, request_id, fence=session.fence) + + closed = harness.manager.close_sessions_for_request_ids([request_id], abort=True, cleanup_in_progress=True) + assert closed == {"sid-in-progress": [request_id]} + + # The orchestrator owns this cleanup: the reaper neither retries nor releases it. + assert await harness.manager.reap_expired() == 0 + assert harness.stage_port.cleanup_calls == [] + assert (await harness.open("sid-replacement")).error_code == "resource_exhausted" + + harness.manager.finalize_closed_sessions(["sid-in-progress"]) + + assert session.resource_request_ids() == [] + assert session.state == DuplexSessionState.CLOSED + assert harness.manager.runner_for_request_id(request_id) is None + assert await harness.manager.reap_expired() == 0 + assert harness.stage_port.cleanup_calls == [] + assert (await harness.open("sid-replacement")).ok is True + + +async def test_deferred_request_cleanup_is_retried_by_the_reaper() -> None: + async with Harness.create(max_sessions=1) as harness: + await harness.open("sid-deferred") + session = harness.session("sid-deferred") + request_id = stage0_request_id("sid-deferred") + session.bind_stage_request(0, request_id, fence=session.fence) + harness.manager.close_sessions_for_request_ids([request_id], abort=True, cleanup_in_progress=True) + assert await harness.manager.reap_expired() == 0 + assert harness.stage_port.cleanup_calls == [] + + # The orchestrator's own cleanup failed: hand the retry to the reaper. + harness.manager.defer_request_cleanups(["sid-deferred"]) + + assert await harness.manager.reap_expired() == 1 + assert harness.stage_port.cleanup_calls == [([request_id], True)] + assert session.resource_request_ids() == [] + assert (await harness.open("sid-replacement")).ok is True + + +async def test_request_cleanup_closes_the_owner_of_any_of_its_stage_requests() -> None: + async with Harness.create() as harness: + await harness.open("sid-merge") + session = harness.session("sid-merge") + stage0 = stage0_request_id("sid-merge") + session.bind_stage_request(0, stage0, fence=session.fence) + stage1 = harness.manager.ensure_stage_request(session, stage_id=1) + assert stage1 is not None + + closed = harness.manager.close_sessions_for_request_ids([stage1.request_id]) + + assert closed == {"sid-merge": [stage0, stage1.request_id]} + assert harness.manager.close_sessions_for_request_ids([stage0]) == {} + assert await harness.manager.reap_expired() == 1 + assert harness.stage_port.cleanup_calls == [([stage0, stage1.request_id], False)] + + +# --------------------------------------------------------------------------- # +# Dispatch ordering / shutdown # +# --------------------------------------------------------------------------- # + + +async def test_control_dispatch_is_ordered_per_session_without_blocking_other_sessions() -> None: + plugin = FakePlugin() + plugin.blocked_session_ids.add("sid-blocked") + async with Harness.create(plugin=plugin) as harness: + manager = harness.manager + blocked_config = DuplexSessionConfig(model="fake-model", instructions="sid-blocked") + + manager.dispatch( + OpenDuplexSessionMessage(control_id="blocked-open", session_id="sid-blocked", session_config=blocked_config) + ) + await asyncio.wait_for(plugin.runtime_config_started.wait(), timeout=1.0) + manager.dispatch( + TouchDuplexSessionMessage( + control_id="blocked-touch", + session_id="sid-blocked", + activity=DuplexLeaseActivity.HEARTBEAT.value, + ) + ) + manager.dispatch( + OpenDuplexSessionMessage( + control_id="independent-open", + session_id="sid-independent", + session_config=DuplexSessionConfig(model="fake-model"), + ) + ) + manager.dispatch( + TouchDuplexSessionMessage( + control_id="independent-touch", + session_id="sid-independent", + activity=DuplexLeaseActivity.HEARTBEAT.value, + ) + ) + + assert [(await harness.result()).control_id for _ in range(2)] == ["independent-open", "independent-touch"] + assert harness.result_sink.empty() + assert manager.get("sid-blocked") is None + + plugin.runtime_config_gate.set() + results = [await harness.result() for _ in range(2)] + assert [result.control_id for result in results] == ["blocked-open", "blocked-touch"] + assert all(result.ok for result in results) + assert manager.active_count() == 2 + + +async def test_shutdown_closes_every_runner_and_stops_dispatch_tasks() -> None: + harness = Harness.create() + await harness.open("sid-a") + await harness.open("sid-b") + sessions = [harness.session("sid-a"), harness.session("sid-b")] + harness.events() + + await harness.manager.shutdown() + + assert harness.manager.active_count() == 0 + assert harness.manager.runners == {} + 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 harness.stage_port.cleanup_calls == [] + + +def test_manager_rejects_a_plugin_whose_sampling_policy_mismatches_the_stages() -> None: + class WrongStageCountPlugin(FakePlugin): + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return defaults[:1] + + with pytest.raises(ValueError, match="one sampling parameter per stage"): + DuplexSessionManager( + plugin=WrongStageCountPlugin(), + stage_port=FakeStagePort(), + output_sink=asyncio.Queue(), + result_sink=asyncio.Queue(), + runtime_config=_runtime_config(), + model_config=None, + ) + + +# --------------------------------------------------------------------------- # +# Admission bookkeeping # +# --------------------------------------------------------------------------- # + + +async def test_duplicate_open_in_flight_does_not_release_the_first_opens_slot() -> None: + """A duplicate id arriving mid-flight must not free the slot the first open holds. + + The rejection runs the same ``finally`` as a successful open, so without a + guard it discards an ``_admitting`` entry it never added, and the next open + of a *different* id sails past ``max_sessions``. Server-allocated ids make + this a guard rather than a live path, which is exactly why a sequential + duplicate cannot exercise it. + """ + async with Harness.create(max_sessions=1) as harness: + harness.plugin.blocked_session_ids.add("dup") + first = asyncio.create_task(harness.open("dup", DuplexSessionConfig(model="fake-model", instructions="dup"))) + await asyncio.wait_for(harness.plugin.runtime_config_started.wait(), timeout=1.0) + + duplicate = await harness.open("dup", DuplexSessionConfig(model="fake-model", instructions="dup")) + assert duplicate.ok is False + assert duplicate.error_code == "session_exists" + + # The in-flight open still holds the only slot. + overflow = await harness.open("other", DuplexSessionConfig(model="fake-model", instructions="other")) + assert overflow.ok is False + assert overflow.error_code == "resource_exhausted" + + harness.plugin.runtime_config_gate.set() + assert (await first).ok is True + assert harness.manager.active_count() == 1 + + +async def test_expiry_emits_a_terminal_event_even_when_a_close_deferred_it() -> None: + """Every session ends with exactly one terminal event, whoever tears it down. + + A wire ``session.close`` defers its ``session.closed`` to the manager. If a + stage failure then takes the runner away from the manager (the orchestrator + cleanup path), nobody would emit the deferred event unless ``expire`` + upgrades it. + """ + async with Harness.create() as harness: + await harness.open("sid-deferred") + 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 runner.expire("request_cleanup", emit_expired=True) + + terminal = [ + event for event in harness.events("sid-deferred") if isinstance(event, SessionClosed | SessionExpired) + ] + assert len(terminal) == 1, terminal + assert isinstance(terminal[0], SessionExpired) + assert terminal[0].reason == "request_cleanup" + + +async def test_an_open_abandoned_by_its_caller_does_not_hold_the_admission_slot() -> None: + """A control-RPC timeout must not leave the engine holding a session. + + The caller's timeout unregisters its RPC waiter but cannot cancel the + manager's open, which is already past the capacity check. Without the + compensating close, that session lands, keeps its slot and its Stage0 + reservation, and with ``max_sessions=1`` the next open is refused until + idle expiry. + """ + async with Harness.create(max_sessions=1) as harness: + harness.plugin.blocked_session_ids.add("blocked") + config = DuplexSessionConfig(model="fake-model", instructions="blocked") + + open_task = asyncio.create_task(harness.open("sid-abandoned", config)) + await asyncio.wait_for(harness.plugin.runtime_config_started.wait(), timeout=1.0) + + # The caller gave up: DuplexOmni sends the compensating close while the + # open is still awaiting the plugin. + await harness.close("sid-abandoned", reason="open_abandoned") + + harness.plugin.runtime_config_gate.set() + await open_task + + assert "sid-abandoned" not in harness.manager.runners + # The slot came back, so a replacement is admitted. + assert (await harness.open("sid-replacement")).ok is True + + +async def test_a_runtime_close_racing_a_manager_close_still_emits_one_terminal() -> None: + """Two teardown drivers, one terminal event. + + The window is inside ``runner.close()``: ``cancel_append_tasks()`` is awaited + *before* the terminal is deferred, so an append task that dies there finds + the session merely CLOSING with neither ``closed_emitted`` nor + ``closed_deferred`` set, and emits ``session.closed`` itself. The manager + decided to emit before any of this and must re-check, or the session ends + with two terminal events. + """ + async with Harness.create(max_sessions=1) as harness: + await harness.open("sid-race") + 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"] + + async def dying_append() -> bool: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + # What AppendAttempt does when its submit fails mid-teardown. + await runner._close_from_runtime("runtime_append_task_failed") + raise + return True + + task = asyncio.ensure_future(dying_append()) + await asyncio.sleep(0) + runner.tasks.track_append_task(task, epoch=session.epoch, final=True, response_bound=True) + + await harness.close("sid-race") + + terminal = [event for event in harness.events("sid-race") if isinstance(event, SessionClosed | SessionExpired)] + assert len(terminal) == 1, terminal + + +async def test_a_terminal_event_is_emitted_once_even_if_expiry_runs_twice() -> None: + async with Harness.create() as harness: + await harness.open("sid-once") + runner = harness.manager.runners["sid-once"] + + 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)] + assert len(terminal) == 1, terminal + + +# --------------------------------------------------------------------------- # +# Resume and takeover lifecycle # +# --------------------------------------------------------------------------- # + + +async def test_resume_inside_the_grace_window_stops_the_reaper_expiring_the_session() -> None: + """Resuming clears the disconnect deadline, not just the lease generation. + + The reaper decides expiry from ``lease.detached_at``. A resume that bumped + the generation without clearing it would let the reaper kill a session a + client had already come back to. + """ + async with Harness.create(disconnect_grace_s=30.0) as harness: + await harness.open("sid-grace") + session = harness.session("sid-grace") + assert (await harness.touch("sid-grace", DuplexLeaseActivity.DETACH.value)).ok is True + harness.clock.advance(10.0) + + assert (await harness.resume("sid-grace", expected_lease_generation=0)).ok is True + assert session.lease.detached_at is None + + # Past the original deadline: the session survives because it is attached. + harness.clock.advance(25.0) + assert await harness.manager.reap_expired() == 0 + assert harness.manager.get("sid-grace") is not None + + +async def test_resume_after_the_grace_window_reports_unknown_session() -> None: + """Once the reaper has expired a detached session, resume has nothing to attach to.""" + async with Harness.create(disconnect_grace_s=5.0) as harness: + await harness.open("sid-late") + assert (await harness.touch("sid-late", DuplexLeaseActivity.DETACH.value)).ok is True + harness.clock.advance(6.0) + assert await harness.manager.reap_expired() == 1 + + late = await harness.resume("sid-late", expected_lease_generation=0) + + assert late.ok is False + assert late.error_code == "unknown_session" + + +async def test_takeover_invalidates_the_generation_the_replaced_client_held() -> None: + """Two clients cannot both hold the session: the first one's generation goes stale. + + This is the engine half of a websocket takeover — the transport picks the + winning socket, and the lease generation is what stops the loser from + resuming behind it. + """ + async with Harness.create() as harness: + await harness.open("sid-takeover") + session = harness.session("sid-takeover") + assert (await harness.touch("sid-takeover", DuplexLeaseActivity.DETACH.value)).ok is True + + winner = await harness.resume("sid-takeover", expected_lease_generation=0) + assert winner.ok is True + assert winner.lease_generation == 1 + + # The replaced client still believes it holds generation 0. + loser = await harness.resume("sid-takeover", expected_lease_generation=0) + assert loser.ok is False + assert loser.error_code == "session_resume_conflict" + assert session.lease_generation == 1 + + # The winner can still act on the session it took over. + assert (await harness.touch("sid-takeover", DuplexLeaseActivity.HEARTBEAT.value)).ok is True + + +async def test_resume_is_refused_once_a_close_has_begun() -> None: + async with Harness.create() as harness: + await harness.open("sid-closing") + await harness.close("sid-closing") + + resumed = await harness.resume("sid-closing", expected_lease_generation=0) + + assert resumed.ok is False + assert resumed.error_code == "unknown_session" + + +async def test_wire_close_of_an_idle_session_still_emits_session_closed() -> None: + """The protocol smoke shape: open, send nothing, close. + + A session that never bound a stage request has nothing for the reaper to + clean up, so the terminal event is the only thing the client ever gets + back for its ``session.close``. Regression guard for a close path that + only emitted after a cleanup it never performed. + """ + async with Harness.create(max_sessions=1) as harness: + await harness.open("sid-idle-close") + 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")) + if events: + break + await asyncio.sleep(0.01) + + assert [type(event) for event in events] == [SessionClosed] + assert events[0].reason == "client_close" + assert harness.manager.get("sid-idle-close") is None + assert (await harness.open("sid-idle-replacement")).ok is True + + +# --------------------------------------------------------------------------- # +# Reaper loop # +# --------------------------------------------------------------------------- # + + +async def test_reaper_loop_waits_between_ticks() -> None: + """The loop paces itself; it does not spin on the expiry check. + + Ported from the pre-framework ``Orchestrator._duplex_reaper_loop`` tests: + the loop moved to the session manager with the sessions it reaps. + """ + manager = object.__new__(DuplexSessionManager) + calls = 0 + + async def reap_expired(now: float | None = None) -> int: + nonlocal calls + calls += 1 + return 0 + + manager.reap_expired = reap_expired # type: ignore[method-assign] + manager.runtime_config = DuplexSessionRuntimeConfig(reaper_interval_s=0.01) + shutdown = asyncio.Event() + + task = asyncio.create_task(manager.reaper_loop(shutdown)) + await asyncio.sleep(0.035) + shutdown.set() + await asyncio.wait_for(task, timeout=5.0) + + assert 2 <= calls <= 5 + + +@pytest.mark.parametrize("first_cleanup_delay", [0.0, 0.05], ids=["immediate", "delayed"]) +async def test_reaper_loop_survives_one_cleanup_failure(first_cleanup_delay: float) -> None: + """A failed expiry sweep is retried on the next tick, not fatal to the loop.""" + recovered = asyncio.Event() + calls = 0 + + manager = object.__new__(DuplexSessionManager) + + async def reap_expired(now: float | None = None) -> int: + nonlocal calls + calls += 1 + if calls == 1: + await asyncio.sleep(first_cleanup_delay) + raise RuntimeError("transient cleanup failure") + recovered.set() + return 0 + + manager.reap_expired = reap_expired # type: ignore[method-assign] + manager.runtime_config = DuplexSessionRuntimeConfig(reaper_interval_s=0.01) + shutdown = asyncio.Event() + + task = asyncio.create_task(manager.reaper_loop(shutdown)) + try: + await asyncio.wait_for(recovered.wait(), timeout=5.0) + assert calls >= 2 + finally: + shutdown.set() + await asyncio.wait_for(task, timeout=5.0) diff --git a/tests/engine/duplex/test_session_runner.py b/tests/engine/duplex/test_session_runner.py new file mode 100644 index 000000000..51711f52a --- /dev/null +++ b/tests/engine/duplex/test_session_runner.py @@ -0,0 +1,1195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Engine-resident session runner scenarios (the old serving handler tests, re-homed). + +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 +fake; the model plugin is the real MiniCPM-o 4.5 one so append planning and +output projection are exercised end to end. +""" + +from __future__ import annotations + +import asyncio +import base64 +import struct +from collections.abc import Sequence +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +from vllm.sampling_params import SamplingParams + +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.engine.duplex import commands +from vllm_omni.engine.duplex.commands import DuplexCommand +from vllm_omni.engine.duplex.config import DuplexSessionConfig, DuplexSessionState +from vllm_omni.engine.duplex.contracts import ( + DuplexOutputContext, + DuplexRequestIdentity, + DuplexStagePort, + DuplexStageRequestContext, + DuplexStageSubmission, + DuplexStageSubmissionResult, + duplex_resource_request_id, +) +from vllm_omni.engine.duplex.events import DuplexEvent +from vllm_omni.engine.duplex.messages import ( + CloseDuplexSessionMessage, + DuplexControlResultMessage, + DuplexSessionCommandMessage, + DuplexSessionEventMessage, + OpenDuplexSessionMessage, +) +from vllm_omni.engine.duplex.session.manager import DuplexSessionManager +from vllm_omni.engine.duplex.session.runner import DuplexSessionRunner +from vllm_omni.metrics.stats import StageRequestStats, StageStats +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.plugin import MiniCPMO45DuplexPlugin + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +SESSION_ID = "duplex-test" +LISTEN_TOKEN_ID = 7 + + +# --------------------------------------------------------------------------- # +# Fakes # +# --------------------------------------------------------------------------- # + + +class RecordingStagePort(DuplexStagePort): + """Records what the runner asks of the orchestrator; never talks to a stage.""" + + def __init__(self, *, stage_count: int = 2) -> None: + self._stage_count = stage_count + self.ensured: list[DuplexStageRequestContext] = [] + self.submissions: list[DuplexStageSubmission] = [] + self.cleanups: list[tuple[list[str], bool]] = [] + self.aborts: list[list[str]] = [] + self.fail_submit: Exception | None = None + + @property + def stage_count(self) -> int: + return self._stage_count + + def sampling_defaults(self) -> tuple[object, ...]: + return tuple(SamplingParams(max_tokens=8) for _ in range(self._stage_count)) + + def ensure_request(self, context: DuplexStageRequestContext) -> None: + self.ensured.append(context) + + async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: + if self.fail_submit is not None: + raise self.fail_submit + self.submissions.append(submission) + return DuplexStageSubmissionResult( + request_id=submission.context.request_id, + stage_id=submission.context.stage_id, + replica_id=0, + ) + + async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: + self.cleanups.append((list(request_ids), abort)) + + async def abort_requests(self, request_ids: list[str]) -> None: + self.aborts.append(list(request_ids)) + + +def _fake_encode_audio(audio: object, sample_rate_hz: int, response_format: str, speed: float | None) -> str | None: + del sample_rate_hz, response_format, speed + if audio is None: + return None + samples = int(np.asarray(audio, dtype=np.float32).size) + return f"wav-{samples}" if samples > 0 else None + + +@dataclass +class Harness: + manager: DuplexSessionManager + port: RecordingStagePort + output: asyncio.Queue[Any] + results: asyncio.Queue[Any] + runner: DuplexSessionRunner + events: list[DuplexEvent] = field(default_factory=list) + + @property + def session(self): + return self.runner.session + + def submit(self, command: DuplexCommand) -> None: + self.manager.dispatch(DuplexSessionCommandMessage(session_id=SESSION_ID, command=command)) + + async def settle(self, *, idle_s: float = 0.05, timeout_s: float = 3.0) -> list[DuplexEvent]: + """Run the loop until the runner mailbox and append tasks are quiet; return new events.""" + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout_s + quiet_since: float | None = None + collected: list[DuplexEvent] = [] + while True: + drained = False + while not self.output.empty(): + message = self.output.get_nowait() + if isinstance(message, DuplexSessionEventMessage): + collected.append(message.event) + drained = True + busy = ( + drained + or not self.runner._mailbox.empty() + or any(not task.done() for task in self.runner.tasks.append_tasks) + or any(not task.done() for task in self.runner._background_tasks) + ) + now = loop.time() + if busy: + quiet_since = None + elif quiet_since is None: + quiet_since = now + elif now - quiet_since >= idle_s: + break + if now >= deadline: + break + await asyncio.sleep(0.005) + self.events.extend(collected) + return collected + + async def run(self, command: DuplexCommand) -> list[DuplexEvent]: + self.submit(command) + return await self.settle() + + def deliver( + self, + output: object, + *, + stage_id: int = 1, + segment_finished: bool = False, + segment_output_metadata: dict[str, Any] | None = None, + segment_token_ids: Sequence[int] = (), + epoch: int | None = None, + metrics: Any = None, + ) -> bool: + session = self.session + fence = session.fence if epoch is None else session.fence.__class__(SESSION_ID, epoch=epoch) + context = DuplexOutputContext( + identity=DuplexRequestIdentity(session_id=SESSION_ID, fence=fence), + final_stage_id=self.port.stage_count - 1, + segment_finished=segment_finished, + segment_token_ids=tuple(segment_token_ids), + segment_output_metadata=dict(segment_output_metadata or {}), + ) + return self.runner.on_stage_output(stage_id, output, metrics, request_id=output.request_id, context=context) + + async def deliver_and_settle(self, output: object, **kwargs: Any) -> list[DuplexEvent]: + self.deliver(output, **kwargs) + return await self.settle() + + def stage0_request_id(self, epoch: int | None = None) -> str: + fence = self.session.fence if epoch is None else self.session.fence.__class__(SESSION_ID, epoch=epoch) + return duplex_resource_request_id(fence, "stage0") + + +async def open_harness( + *, + auto_response: bool = True, + modalities: Sequence[str] = ("text",), + extra_body: dict[str, object] | None = None, + runtime_config: DuplexSessionRuntimeConfig | None = None, + stage_count: int = 2, + clock: Any = None, +) -> Harness: + plugin = MiniCPMO45DuplexPlugin(_fake_encode_audio) + port = RecordingStagePort(stage_count=stage_count) + output: asyncio.Queue[Any] = asyncio.Queue() + results: asyncio.Queue[Any] = asyncio.Queue() + manager = DuplexSessionManager( + plugin=plugin, + stage_port=port, + output_sink=output, + result_sink=results, + runtime_config=runtime_config or DuplexSessionRuntimeConfig(), + model_config=None, + clock=clock, + ) + body: dict[str, object] = {"auto_response": auto_response, **(extra_body or {})} + config = DuplexSessionConfig( + model="openbmb/MiniCPM-o-4_5", + modalities=list(modalities), + instructions="You are a concise assistant.", + extra_body=body, + ) + await manager.handle(OpenDuplexSessionMessage(control_id="c-open", session_id=SESSION_ID, session_config=config)) + 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]) + await harness.settle() + return harness + + +async def close_harness(harness: Harness) -> None: + await harness.manager.shutdown() + + +# --------------------------------------------------------------------------- # +# Payload / output builders # +# --------------------------------------------------------------------------- # + + +def pcm_f32(samples: int, *, value: float = 0.05) -> bytes: + return struct.pack(f"<{samples}f", *([value] * samples)) + + +def append_audio( + samples: int = 16000, + *, + value: float = 0.05, + is_speech: bool | None = True, + event_id: str | None = None, +) -> commands.AppendAudio: + return commands.AppendAudio( + audio=pcm_f32(samples, value=value), + format="pcm_f32le", + sample_rate_hz=16000, + is_speech=is_speech, + event_id=event_id, + ) + + +def stage_stats( + *, + stage_id: int, + request_id: str, + num_tokens_out: int, + itls_ms: list[float] | None = None, +) -> StageRequestStats: + """The per-request stats the orchestrator hands to the runner with an output.""" + return StageRequestStats( + batch_id=0, + batch_size=1, + num_tokens_in=7, + num_tokens_out=num_tokens_out, + stage_gen_time_ms=120.0, + rx_transfer_bytes=0, + rx_decode_time_ms=0.0, + rx_in_flight_time_ms=0.0, + stage_stats=StageStats(), + stage_id=stage_id, + request_id=request_id, + final_output_type="text", + vllm_itls_ms=list(itls_ms or []), + ) + + +def tts_output( + request_id: str, + *, + samples: int = 24000, + finished: bool = False, + text: str = "hello", + tts_is_last_chunk: bool = False, + turn_end: bool = False, + turn_id: int = 0, + epoch: int = 0, +) -> SimpleNamespace: + """A Stage1 (TTS) output the way the orchestrator hands it to the runner.""" + return SimpleNamespace( + request_id=request_id, + finished=finished, + outputs=[SimpleNamespace(text=text, token_ids=[], multimodal_output={})], + multimodal_output={ + "audio": np.zeros(samples, dtype=np.float32), + "sr": 24000, + "meta.tts_is_last_chunk": np.array([int(tts_is_last_chunk)], dtype=np.int32), + "meta.turn_end": np.array([int(turn_end)], dtype=np.int32), + "meta.duplex_turn_id": np.array([turn_id], dtype=np.int32), + "meta.duplex_epoch": np.array([epoch], dtype=np.int32), + }, + ) + + +def listen_output(request_id: str) -> SimpleNamespace: + """A finished Stage0 segment that ends with the listen token.""" + return SimpleNamespace( + request_id=request_id, + finished=True, + outputs=[ + SimpleNamespace( + text="", + token_ids=[11, 12, LISTEN_TOKEN_ID], + stop_reason=LISTEN_TOKEN_ID, + multimodal_output={}, + ) + ], + multimodal_output={"meta.listen_token_id": LISTEN_TOKEN_ID}, + ) + + +def types(events: Sequence[DuplexEvent]) -> list[str]: + return [event.type for event in events] + + +def find(events: Sequence[DuplexEvent], wire_type: str) -> DuplexEvent: + for event in events: + if event.type == wire_type: + return event + raise AssertionError(f"no {wire_type!r} in {types(events)}") + + +def b64(data: bytes) -> str: + return base64.b64encode(data).decode("ascii") + + +# --------------------------------------------------------------------------- # +# Session lifecycle # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_open_announces_the_session_and_reserves_the_stage0_request() -> None: + h = await open_harness() + try: + assert types(h.events) == ["session.created", "session.updated"] + created = h.events[0] + assert created.session_id == SESSION_ID + assert created.session["id"] == SESSION_ID + assert created.session["capabilities"]["supports_input_append"] is True + assert created.to_realtime()["type"] == "session.created" + assert "incarnation" not in created.to_realtime() + assert h.session.state == DuplexSessionState.OPEN + # Admission reserves the resumable Stage0 request atomically with the open. + assert [context.request_id for context in h.port.ensured] == [h.stage0_request_id(epoch=0)] + assert h.stage0_request_id(epoch=0) == "duplex-s.ZHVwbGV4LXRlc3Q.e.0.r.stage0" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_close_emits_session_closed_and_releases_stage_requests() -> None: + h = await open_harness() + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.manager.handle( + CloseDuplexSessionMessage(control_id="c-close", session_id=SESSION_ID, reason="client_close") + ) + result = await asyncio.wait_for(h.results.get(), timeout=2.0) + events = await h.settle() + + assert result.ok and result.operation == "close" + assert types(events) == ["session.closed"] + assert events[0].reason == "client_close" + assert events[0].is_terminal + assert h.session.state == DuplexSessionState.CLOSED + assert h.port.cleanups == [([request_id], True)] + assert SESSION_ID not in h.manager.runners + await close_harness(h) + + +@pytest.mark.asyncio +async def test_idle_expiry_emits_session_expired() -> None: + clock = {"now": 100.0} + h = await open_harness(runtime_config=DuplexSessionRuntimeConfig(idle_ttl_s=1.0), clock=lambda: clock["now"]) + clock["now"] += 2.0 + assert await h.manager.reap_expired() == 1 + events = await h.settle() + + assert types(events) == ["session.expired"] + assert events[0].reason == "idle_ttl_expired" + assert events[0].is_terminal + assert SESSION_ID not in h.manager.runners + assert not h.manager._closing + await close_harness(h) + + +@pytest.mark.asyncio +async def test_commands_for_a_closed_session_are_answered_with_unknown_session() -> None: + h = await open_harness() + await h.manager.handle(CloseDuplexSessionMessage(control_id="c-close", session_id=SESSION_ID, reason="bye")) + await h.results.get() + await h.settle() + + h.submit(commands.Heartbeat(event_id="evt-late")) + events = await h.settle() + + assert types(events) == ["error"] + assert events[0].code == "unknown_session" + assert events[0].related_event_id == "evt-late" + await close_harness(h) + + +# --------------------------------------------------------------------------- # +# Input path # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_append_plans_and_submits_one_model_unit_in_wire_order() -> None: + h = await open_harness() + try: + events = await h.run(append_audio()) + assert types(events) == ["input_audio_buffer.speech_started"] + assert len(h.port.submissions) == 1 + first = h.port.submissions[0] + assert first.context.request_id == h.stage0_request_id(epoch=0) + assert first.context.stage_id == 0 + assert first.already_submitted is False + duplex = first.prompt["model_intermediate_buffer"]["duplex"] + assert duplex["session_id"] == SESSION_ID + assert (duplex["epoch"], duplex["seq"], duplex["turn_id"], duplex["turn_seq"]) == (0, 1, 0, 1) + assert duplex["final"] is False + assert "incarnation" not in duplex + assert len(base64.b64decode(duplex["payload"]["audio"])) == 16000 * 4 + assert duplex["payload"]["is_speech"] is True + assert h.session.input_seq == 1 + assert h.session.stage_request_submitted(0, first.context.request_id) + + await h.run(append_audio()) + second = h.port.submissions[1] + assert second.already_submitted is True + assert second.prompt["model_intermediate_buffer"]["duplex"]["seq"] == 2 + assert h.session.input_seq == 2 + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_partial_units_are_buffered_until_a_whole_chunk_is_available() -> None: + h = await open_harness() + try: + await h.run(append_audio(samples=8000)) + assert h.port.submissions == [] + assert h.runner.model_state.audio_buffer.has_pending() + + await h.run(append_audio(samples=8000)) + assert len(h.port.submissions) == 1 + assert not h.runner.model_state.audio_buffer.has_pending() + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_input_backpressure_is_rejected_before_the_buffer_is_touched() -> None: + h = await open_harness(runtime_config=DuplexSessionRuntimeConfig(max_pending_input_bytes_per_session=1000)) + try: + events = await h.run(append_audio(event_id="evt-big")) + assert types(events) == ["error"] + assert events[0].code == "input_backpressure" + assert events[0].related_event_id == "evt-big" + assert h.port.submissions == [] + assert not h.runner.model_state.audio_buffer.has_pending() + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_text_append_is_not_supported_by_the_native_runtime() -> None: + h = await open_harness() + try: + events = await h.run(commands.AppendText(text="hello", event_id="evt-text")) + assert types(events) == ["error"] + assert events[0].code == "native_text_append_unsupported" + assert events[0].related_event_id == "evt-text" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_clear_input_drops_buffered_audio() -> None: + h = await open_harness() + try: + await h.run(append_audio(samples=8000)) + assert h.runner.model_state.audio_buffer.has_pending() + events = await h.run(commands.ClearInput()) + assert types(events) == ["input_audio_buffer.cleared"] + assert not h.runner.model_state.audio_buffer.has_pending() + assert h.session.pending_input_bytes == 0 + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_commit_without_audio_is_rejected() -> None: + h = await open_harness() + try: + events = await h.run(commands.Commit(event_id="evt-commit")) + assert types(events) == ["error"] + assert events[0].code == "input_audio_buffer_empty" + assert events[0].related_event_id == "evt-commit" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_commit_projects_the_user_item_and_does_not_resubmit_consumed_audio() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + events = await h.run(commands.Commit()) + assert types(events) == [ + "input_audio_buffer.speech_stopped", + "conversation.item.added", + "conversation.item.created", + "input_audio_buffer.committed", + "conversation.item.done", + ] + committed = find(events, "input_audio_buffer.committed") + assert committed.item_id == find(events, "conversation.item.added").item_id + # The unit already reached Stage0; a commit is only a turn marker here. + assert len(h.port.submissions) == 1 + assert h.session.input_commit_seq == 1 + assert h.session.history[-1]["role"] == "user" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_heartbeat_touches_the_lease_and_acks() -> None: + clock = {"now": 100.0} + h = await open_harness(clock=lambda: clock["now"]) + try: + clock["now"] = 150.0 + events = await h.run(commands.Heartbeat()) + assert types(events) == ["session.heartbeat_ack"] + assert h.session.lease.last_activity == 150.0 + finally: + await close_harness(h) + + +# --------------------------------------------------------------------------- # +# Model output path # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_stage1_audio_opens_a_response_and_streams_deltas() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + + events = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="he")) + assert types(events) == [ + "response.created", + "conversation.item.added", + "conversation.item.created", + "response.output_item.added", + "response.speak", + "response.content_part.added", + "response.output_audio.delta", + "response.output_audio_transcript.delta", + ] + response_id = find(events, "response.created").response_id + assert response_id == h.session.active_response_id + assert h.session.active_request_id == request_id + delta = find(events, "response.output_audio.delta") + assert (delta.delta, delta.format, delta.sample_rate_hz) == ("wav-24000", "wav", 24000) + assert delta.response_id == response_id + assert find(events, "response.output_audio_transcript.delta").delta == "he" + assert h.session.playback.sent_ms == 1000 + + # Cumulative Stage1 audio is sliced to the new samples only. + events = await h.deliver_and_settle(tts_output(request_id, samples=48000, text="hello")) + assert types(events) == ["response.output_audio.delta", "response.output_audio_transcript.delta"] + assert find(events, "response.output_audio.delta").delta == "wav-24000" + assert find(events, "response.output_audio_transcript.delta").delta == "llo" + assert h.session.playback.sent_ms == 2000 + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_turn_end_completes_the_response_and_advances_the_model_turn() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + + events = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello", turn_end=True)) + assert types(events) == [ + "response.output_audio.done", + "response.output_audio_transcript.done", + "response.content_part.done", + "response.output_item.done", + "conversation.item.done", + "response.done", + "rate_limits.updated", + ] + done = find(events, "response.done") + assert done.status == "completed" + assert find(events, "response.output_audio_transcript.done").transcript == "hello" + assert h.session.active_response_id is None + assert h.session.turn_id == 1 + # Ack-only playback: the assistant text enters history on playback ack. + assert h.session.history == () + + ack = await h.run(commands.AckPlayback(played_ms=1000, response_id=done.response_id)) + assert types(ack) == ["playback.acknowledged"] + assert ack[0].details["history_committed"] is True + assert ack[0].details["playback"]["committed_ms"] == 1000 + assert h.session.history == ({"role": "assistant", "content": "hello"},) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_playback_ack_for_an_unknown_response_is_rejected() -> None: + h = await open_harness() + try: + events = await h.run(commands.AckPlayback(played_ms=10, response_id="resp-unknown", event_id="evt-ack")) + assert types(events) == ["error"] + assert events[0].code == "playback_item_not_found" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_stale_epoch_output_is_dropped_after_barge_in() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="he")) + + events = await h.run(commands.BargeIn()) + assert h.session.epoch == 1 + assert h.session.active_response_id is None + assert h.port.aborts == [[request_id]] + assert h.port.cleanups == [([request_id], True)] + done = find(events, "response.done") + assert done.status == "cancelled" + assert types(events)[: types(events).index("response.done")] == [ + "response.output_audio.done", + "response.output_audio_transcript.done", + "response.content_part.done", + "response.output_item.done", + "conversation.item.done", + ] + + late = tts_output(request_id, samples=48000, text="hello", epoch=0) + assert h.deliver(late, epoch=0) is True # consumed, never forwarded + assert await h.settle() == [] + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_cancel_response_without_active_response_is_rejected() -> None: + h = await open_harness() + try: + events = await h.run(commands.CancelResponse(event_id="evt-cancel")) + assert types(events) == ["error"] + assert events[0].code == "response_not_active" + assert events[0].related_event_id == "evt-cancel" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_cancel_response_aborts_the_stage_request_and_reports_playback() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + await h.run(commands.AckPlayback(played_ms=400, response_id=h.session.active_response_id)) + + events = await h.run(commands.CancelResponse()) + done = find(events, "response.done") + assert done.status == "cancelled" + assert done.response["status_details"]["reason"] == "client_cancelled" + assert h.port.aborts == [[request_id]] + assert h.session.epoch == 1 + # Only the played prefix of the cancelled answer is kept in history. + assert h.session.history == ({"role": "assistant", "content": "he"},) + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_listen_decision_is_consumed_and_never_forwarded_to_tts() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + # An unfinished Stage0 segment is plain text for the TTS stage: forwarded. + text_output = SimpleNamespace( + request_id=request_id, + finished=False, + outputs=[SimpleNamespace(text="hi", token_ids=[11], multimodal_output={})], + multimodal_output={}, + ) + assert h.deliver(text_output, stage_id=0) is False + + consumed = h.deliver( + listen_output(request_id), + stage_id=0, + segment_finished=True, + segment_token_ids=[11, 12, LISTEN_TOKEN_ID], + segment_output_metadata={"meta.listen_token_id": LISTEN_TOKEN_ID}, + ) + events = await h.settle() + assert consumed is True + assert types(events) == ["response.listen"] + assert events[0].details["reason"] == "model_listen" + assert events[0].details["model_listen"] is True + assert events[0].to_realtime()["response"]["status"] == "listening" + assert h.session.active_response_id is None + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_listen_decision_on_a_resumable_request_closes_the_bounded_response() -> None: + """A live stage-0 request is resumable: its output never says ``finished``, + only its segment does. The listen that answers the last continuation unit + must still close the response, or the session sits forever with no + terminal event (the live-client E2E hang).""" + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + response_id = h.session.active_response_id + assert response_id is not None + # The continuation budget is spent: the unit being answered was the forced listen. + model_state = h.runner.model_state + model_state.continuation_owner_id = f"response:{response_id}" + model_state.continuation_units = h.runner.model._AUTO_RESPONSE_MAX_CONTINUATION_UNITS + + listen = listen_output(request_id) + listen.finished = False + events = await h.deliver_and_settle( + listen, + stage_id=0, + segment_finished=True, + segment_token_ids=[11, 12, LISTEN_TOKEN_ID], + segment_output_metadata={"meta.listen_token_id": LISTEN_TOKEN_ID}, + ) + + assert types(events)[0] == "response.listen" + assert events[0].details["model_listen"] is True + assert find(events, "response.done").response_id == response_id + assert h.session.active_response_id is None + assert model_state.continuation_units == 0 + assert len(h.port.submissions) == 1 # no further silence unit was scheduled + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_listen_decision_on_a_resumable_request_keeps_the_turn_going() -> None: + """Same unfinished listen, budget left: it must schedule the next unit. + + A non-terminal auto-response listen is answered with the next silence + unit, not with a ``response.listen`` event, so the submission is the + observable. Before the fix the projector saw ``finished=False`` and + yielded nothing: no silence unit, the response left open with nothing in + flight -- the same hang, one unit earlier. + """ + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + response_id = h.session.active_response_id + assert response_id is not None + submissions_before = len(h.port.submissions) + + listen = listen_output(request_id) + listen.finished = False + events = await h.deliver_and_settle( + listen, + stage_id=0, + segment_finished=True, + segment_token_ids=[11, 12, LISTEN_TOKEN_ID], + segment_output_metadata={"meta.listen_token_id": LISTEN_TOKEN_ID}, + ) + + assert not [event for event in events if event.type == "error"], types(events) + assert len(h.port.submissions) == submissions_before + 1, "the next silence unit was not scheduled" + assert h.runner.model_state.continuation_units == 1 + assert h.session.active_response_id == response_id + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_tts_segment_end_schedules_a_silence_continuation_unit() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + await h.deliver_and_settle( + tts_output(request_id, samples=48000, text="hello", tts_is_last_chunk=True, finished=True) + ) + + assert len(h.port.submissions) == 2 + silence = h.port.submissions[1] + assert silence.already_submitted is True + duplex = silence.prompt["model_intermediate_buffer"]["duplex"] + assert duplex["seq"] == 2 + assert duplex["payload"]["duplex_turn_id"] == 0 + assert set(base64.b64decode(duplex["payload"]["audio"])) == {0} + assert h.runner.model_state.continuation_units == 1 + # The response stays open across the segment boundary. + assert h.session.active_response_id is not None + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_stage_failure_fails_the_active_response() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hello")) + + h.runner.on_stage_failure(1, RuntimeError("boom")) + events = await h.settle() + assert types(events)[0] == "error" + assert events[0].code == "runtime_data_plane_stream_failed" + done = find(events, "response.done") + assert done.status == "failed" + assert done.response["status_details"]["reason"] == "runtime_data_plane_stream_failed" + assert h.session.active_response_id is None + finally: + await close_harness(h) + + +# --------------------------------------------------------------------------- # +# Turn mode (no auto response) # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_turn_mode_skips_silent_chunks() -> None: + h = await open_harness(auto_response=False) + try: + events = await h.run(append_audio(is_speech=False, value=0.0)) + assert types(events) == ["response.listen"] + assert events[0].details["reason"] == "silence_or_noise" + assert h.port.submissions == [] + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_turn_mode_commit_with_response_create_starts_one_response() -> None: + h = await open_harness(auto_response=False) + try: + events = await h.run(append_audio()) + assert types(events) == ["input_audio_buffer.speech_started"] + assert h.port.submissions == [] # buffered until the commit + + events = await h.run(commands.Commit(create_response=True)) + assert "input_audio_buffer.committed" in types(events) + assert "response.created" in types(events) + assert len(h.port.submissions) == 1 + final = h.port.submissions[0] + assert final.prompt["model_intermediate_buffer"]["duplex"]["final"] is True + request_id = final.context.request_id + response_id = h.session.active_response_id + assert response_id is not None + + events = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="sure")) + assert types(events)[-2:] == ["response.output_audio.delta", "response.output_audio_transcript.delta"] + events = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="sure", finished=True)) + done = find(events, "response.done") + assert done.response_id == response_id + assert done.status == "completed" + assert h.session.active_response_id is None + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_response_create_with_no_input_at_all_is_rejected() -> None: + """Nothing committed and no items: refuse rather than emit an empty turn.""" + h = await open_harness(auto_response=False) + try: + events = await h.run(commands.CreateResponse(event_id="evt-resp")) + assert types(events) == ["error"] + assert events[0].code == "response_create_without_input" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_an_item_only_response_create_is_refused_not_left_hanging() -> None: + """Items are context, not a turn, and the refusal must be immediate. + + A model-native model decides to speak from the audio it hears, so a turn + with no audio in it is one it never answers. Opening a response anyway + submits work that never completes, and the caller only finds out when the + session idles out -- observed as a 300-second HTTP 500 against a real + server. Refusing costs the same information and none of the wait. + """ + h = await open_harness(auto_response=False) + try: + await h.run( + commands.CreateItem( + item={ + "id": "item_prompt", + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "What is 2+2?"}], + } + ) + ) + events = await h.run(commands.CreateResponse(event_id="evt-resp")) + assert types(events) == ["error"], events + assert events[0].code == "text_only_turn_unsupported" + assert not h.port.submissions, "nothing may be submitted for a turn the model cannot answer" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_conversation_items_are_consumed_by_the_turn_they_start() -> None: + """The items belong to the turn that answered them, not to the next one. + + Without this, a session that once received an item would let every later + ``response.create`` start an empty turn off the same stale input. + """ + h = await open_harness(auto_response=False) + try: + await h.run( + commands.CreateItem( + item={ + "id": "item_hi", + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "hi"}], + } + ) + ) + assert h.session.unanswered_user_items() == 1 + # Consumed by the response.create that refused them: a later one must + # not see the same stale input and refuse for a second time. + await h.run(commands.CreateResponse(event_id="evt-first")) + assert h.session.unanswered_user_items() == 0 + events = await h.run(commands.CreateResponse(event_id="evt-second")) + assert events[0].code == "response_create_without_input" + finally: + await close_harness(h) + + +# --------------------------------------------------------------------------- # +# Session update and conversation items # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_session_update_rejects_instruction_changes_after_the_context_is_locked() -> None: + h = await open_harness() + try: + await h.run(append_audio()) + assert h.runner.model_state.context_locked + events = await h.run(commands.UpdateSession(patch={"instructions": "changed"}, event_id="evt-upd")) + assert types(events) == ["error"] + assert events[0].code == "instructions_update_unsupported" + assert events[0].related_event_id == "evt-upd" + assert h.session.config.instructions == "You are a concise assistant." + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_session_update_replaces_config_and_runtime_config_together() -> None: + h = await open_harness() + try: + generation = h.session.config_generation + events = await h.run(commands.UpdateSession(patch={"temperature": 0.3})) + assert types(events) == ["session.updated"] + assert events[0].session["temperature"] == 0.3 + assert h.session.config.temperature == 0.3 + assert h.session.runtime_config["duplex_stage_sampling_params"]["0"]["temperature"] == 0.3 + assert h.session.config_generation == generation + 2 + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_conversation_items_can_be_injected_and_deleted() -> None: + h = await open_harness() + try: + item = {"id": "item_hist", "type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]} + events = await h.run(commands.CreateItem(item=item)) + assert types(events) == ["conversation.item.added", "conversation.item.created", "conversation.item.done"] + assert h.session.history == ({"role": "user", "content": "hi"},) + + events = await h.run(commands.DeleteItem(item_id="item_hist")) + assert types(events) == ["conversation.item.deleted"] + assert h.session.history == () + + events = await h.run(commands.DeleteItem(item_id="item_missing", event_id="evt-del")) + assert types(events) == ["error"] + assert events[0].code == "item_not_found" + finally: + await close_harness(h) + + +def _stage_metrics_of(event: object) -> dict[str, dict[str, object]]: + """Per-stage engine metrics as the client reads them off one wire event.""" + payload = event.to_realtime() + metadata = payload.get("metadata") + assert isinstance(metadata, dict), payload + vllm_omni = metadata.get("vllm_omni") + assert isinstance(vllm_omni, dict), metadata + stage_metrics = vllm_omni.get("stage_metrics") + assert isinstance(stage_metrics, dict), vllm_omni + return stage_metrics + + +@pytest.mark.asyncio +async def test_stage0_metrics_reach_the_response_even_though_its_output_feeds_tts() -> None: + """A pass-through Stage0 output still has to report its tokens. + + Stage0 text with no decision is forwarded to the TTS stage rather than + consumed, and before sessions moved into the engine the orchestrator + published its metrics separately. Returning ``False`` here must not also + throw the metrics away, or every response reports no engine token count and + the benchmark cannot compute TPOT. + """ + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + text_output = SimpleNamespace( + request_id=request_id, + finished=False, + outputs=[SimpleNamespace(text="hi", token_ids=[11], multimodal_output={})], + multimodal_output={}, + ) + forwarded = h.deliver( + text_output, + stage_id=0, + metrics=stage_stats(stage_id=0, request_id=request_id, num_tokens_out=3, itls_ms=[9.0, 11.0]), + ) + assert forwarded is False + assert await h.settle() == [] + + events = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hi")) + stage_metrics = _stage_metrics_of(find(events, "response.output_audio.delta")) + assert stage_metrics["0"]["num_tokens_out"] == 3 + assert stage_metrics["0"]["vllm_itls_ms"] == [9.0, 11.0] + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_stage0_metrics_from_several_units_are_summed_into_one_response() -> None: + """Two pass-through segments before the first audio are one response's tokens.""" + h = await open_harness() + try: + await h.run(append_audio()) + request_id = h.stage0_request_id() + for tokens in (3, 4): + h.deliver( + SimpleNamespace( + request_id=request_id, + finished=False, + outputs=[SimpleNamespace(text="hi", token_ids=[11], multimodal_output={})], + multimodal_output={}, + ), + stage_id=0, + metrics=stage_stats(stage_id=0, request_id=request_id, num_tokens_out=tokens), + ) + assert await h.settle() == [] + + events = await h.deliver_and_settle(tts_output(request_id, samples=24000, text="hi")) + stage_metrics = _stage_metrics_of(find(events, "response.output_audio.delta")) + assert stage_metrics["0"]["num_tokens_out"] == 7 + finally: + await close_harness(h) + + +# --------------------------------------------------------------------------- # +# Server VAD # +# --------------------------------------------------------------------------- # + + +class _ScriptedDetector: + """A turn detector that answers each appended chunk with a pre-decided result.""" + + def __init__(self, results: list[object]) -> None: + self._results = list(results) + self.resets = 0 + + def process(self, base64_audio: str, *, fmt: str, sample_rate_hz: int | None, audio_end_ms: int | None = None): + del base64_audio, fmt, sample_rate_hz, audio_end_ms + return self._results.pop(0) + + def reset(self) -> None: + self.resets += 1 + + +def _speech_then_stop() -> _ScriptedDetector: + from vllm_omni.engine.duplex.turn_detection import TurnDetectionResult + + return _ScriptedDetector( + [ + TurnDetectionResult(is_speech=True, speech_active=True, speech_started=True, speech_probability=0.9), + TurnDetectionResult( + is_speech=False, + speech_active=False, + speech_stopped=True, + speech_probability=0.05, + should_commit=True, + create_response=True, + ), + ] + ) + + +def _final_submissions(h: Harness) -> list[object]: + return [s for s in h.port.submissions if s.prompt["model_intermediate_buffer"]["duplex"].get("final") is True] + + +@pytest.mark.asyncio +async def test_server_vad_speech_stopped_does_not_end_the_turn_of_an_auto_response_session() -> None: + """A model-native session keeps its turn open past the detector's stop. + + Committing at ``speech_stopped`` cut the utterance short: the model + listened on that early final unit, the trailing silence became a second, + near-empty turn, and the client's own commit answered nothing (the + server-VAD hard-interrupt E2E timed out waiting for the follow-up + response). The old translator synthesized that commit for turn-based + server VAD only. + """ + h = await open_harness() + try: + h.runner.control._detector = _speech_then_stop() + first = await h.run(append_audio()) + assert "input_audio_buffer.speech_started" in types(first) + # A partial unit, like the trailing silence of a real utterance: it is + # what a commit here would submit as the turn's final append. + second = await h.run(append_audio(samples=8000, is_speech=False, value=0.0)) + assert "input_audio_buffer.speech_stopped" in types(second) + assert "input_audio_buffer.committed" not in types(first) + types(second) + assert _final_submissions(h) == [], "the detector must not close a model-native turn" + + events = await h.run(commands.Commit()) + assert types(events).count("input_audio_buffer.committed") == 1 + assert len(_final_submissions(h)) == 1, "the client's commit closes the turn exactly once" + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_server_vad_speech_stopped_still_commits_a_turn_mode_session() -> None: + """Turn-based server VAD keeps the contract the old translator implemented.""" + h = await open_harness(auto_response=False) + try: + h.runner.control._detector = _speech_then_stop() + await h.run(append_audio()) + events = await h.run(append_audio(is_speech=False, value=0.0)) + assert "input_audio_buffer.speech_stopped" in types(events) + assert "input_audio_buffer.committed" in types(events) + assert len(_final_submissions(h)) == 1, "the detector's stop commits the turn and starts the response" + finally: + await close_harness(h) diff --git a/tests/engine/duplex/test_vad_backend.py b/tests/engine/duplex/test_vad_backend.py new file mode 100644 index 000000000..8a22259ac --- /dev/null +++ b/tests/engine/duplex/test_vad_backend.py @@ -0,0 +1,379 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The engine-side Silero VAD: backend selection and endpoint rules. + +The endpoint rules are the part worth pinning hardest. They were reimplemented +when turn detection moved into the engine, and drifted from upstream's +``ThresholdEndpointPolicy`` in two ways that only show up mid-utterance: a loud +frame used to cancel an in-progress silence timer, and ``audio_end_ms`` used to +exclude the trailing silence OpenAI says it should include. The parity cases +below hold upstream's own outputs, so neither can come back unnoticed. +""" + +from __future__ import annotations + +import base64 +import hashlib +import sys +import threading +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace + +import numpy as np +import pytest + +from vllm_omni.engine.duplex.turn_detection import validate_realtime_turn_detection +from vllm_omni.engine.duplex.vad import ( + SILERO_VAD_SHA256, + ServerVADUnavailableError, + SileroStreamingVAD, + SileroVADBackend, + SileroVADBackendProvider, + SileroVADConfig, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +FRAME = 512 +SAMPLE_RATE_HZ = 16_000 + + +def _drive(config: SileroVADConfig, probabilities: list[float]) -> list[tuple[str, int | None]]: + """Feed one probability per frame and collect the endpoint events.""" + scores = iter(probabilities) + vad = SileroStreamingVAD(config, frame_scorer=lambda _frame: next(scores)) + events: list[tuple[str, int | None]] = [] + for _ in probabilities: + result = vad.process(np.zeros(FRAME, dtype=np.float32)) + if result.speech_started: + events.append(("start", result.speech_start_ms)) + if result.speech_stopped: + events.append(("stop", result.speech_end_ms)) + return events + + +# --------------------------------------------------------------------------- # +# Parity with the serving-side policy this was ported from # +# --------------------------------------------------------------------------- # + +#: Probability sequences with the endpoint events upstream's +#: ``ThresholdEndpointPolicy`` produced for them, captured from +#: ``entrypoints/duplex/server_vad.py`` at `e2d2617f` before that module was +#: removed. Golden rather than computed because the oracle no longer ships: the +#: serving layer is transport only, so the VAD it used to carry is gone. Any +#: drift in the engine's rules still breaks these. +_PARITY_CASES = { + "one utterance": ( + [0.0, 0.9, 0.9, 0.9, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + {"threshold": 0.5, "prefix_padding_ms": 0, "silence_duration_ms": 100, "min_speech_duration_ms": 32}, + [("start", 32), ("stop", 256)], + ), + "prefix padding reaches back": ( + [0.0, 0.9, 0.9, 0.0, 0.0, 0.0, 0.0], + {"threshold": 0.5, "prefix_padding_ms": 300, "silence_duration_ms": 64, "min_speech_duration_ms": 32}, + [("start", 0), ("stop", 160)], + ), + "loud frames inside the silence timer": ( + [0.9, 0.9, 0.1, 0.45, 0.1, 0.1, 0.1], + {"threshold": 0.5, "prefix_padding_ms": 0, "silence_duration_ms": 96, "min_speech_duration_ms": 32}, + [("start", 0), ("stop", 160)], + ), + "min speech duration rejects a blip": ( + [0.9, 0.0, 0.9, 0.9, 0.9, 0.0, 0.0, 0.0, 0.0], + {"threshold": 0.5, "prefix_padding_ms": 0, "silence_duration_ms": 64, "min_speech_duration_ms": 96}, + [("start", 64), ("stop", 224)], + ), + "negative threshold floor at a low threshold": ( + [0.9, 0.9, 0.1, 0.1, 0.1, 0.1], + {"threshold": 0.16, "prefix_padding_ms": 0, "silence_duration_ms": 64, "min_speech_duration_ms": 32}, + [("start", 0)], + ), +} + + +@pytest.mark.parametrize(("probabilities", "options", "expected"), _PARITY_CASES.values(), ids=list(_PARITY_CASES)) +def test_endpoint_decisions_match_the_serving_side_policy( + probabilities: list[float], + options: dict[str, object], + expected: list[tuple[str, int | None]], +) -> None: + """The engine detector reproduces the policy server VAD shipped with upstream.""" + assert _drive(SileroVADConfig(**options), probabilities) == expected + + +# --------------------------------------------------------------------------- # +# The two rules that had drifted # +# --------------------------------------------------------------------------- # + + +def test_a_loud_frame_does_not_restart_the_silence_timer() -> None: + """Silero v6.2 hysteresis: once silence is running, a loud frame only delays it. + + The frame at 0.45 sits above the negative threshold (0.35) but below the + activation threshold. It must not cancel the pending endpoint, or a speaker + who trails off unevenly never gets their turn committed. + """ + config = SileroVADConfig(threshold=0.5, prefix_padding_ms=0, silence_duration_ms=96, min_speech_duration_ms=32) + assert _drive(config, [0.9, 0.9, 0.1, 0.45, 0.1, 0.1]) == [("start", 0), ("stop", 160)] + + +def test_audio_end_ms_includes_the_trailing_silence() -> None: + """OpenAI defines audio_end_ms as the end of the audio sent to the model. + + That includes the silence spent deciding the turn was over, so it is the + frame boundary the detector stopped on, not where speech last was. + """ + config = SileroVADConfig(threshold=0.5, prefix_padding_ms=0, silence_duration_ms=100, min_speech_duration_ms=32) + events = _drive(config, [0.0, 0.9, 0.9, 0.9, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + stop_ms = next(ms for kind, ms in events if kind == "stop") + # 8 frames consumed when the timer expires: 8 * 512 / 16 kHz = 256 ms. + assert stop_ms == 256 + + +def test_reset_keeps_the_session_clock_and_clamps_the_prefix() -> None: + """A barge-in resets the detector but not the session's timeline. + + Timestamps have to stay comparable with the rest of the session, and prefix + padding must not reach back into audio the reset discarded. + """ + config = SileroVADConfig(threshold=0.5, prefix_padding_ms=300, silence_duration_ms=64, min_speech_duration_ms=32) + scores = iter([0.0] * 10 + [0.9, 0.9, 0.9]) + vad = SileroStreamingVAD(config, frame_scorer=lambda _frame: next(scores)) + for _ in range(10): + vad.process(np.zeros(FRAME, dtype=np.float32)) + + stream_start_ms = round(10 * FRAME * 1000 / SAMPLE_RATE_HZ) + vad.reset() + result = vad.process(np.zeros(FRAME * 3, dtype=np.float32)) + + assert result.speech_started + # Without the clock the start would land at 0; without the clamp, 300 ms of + # prefix would reach back before the reset. + assert result.speech_start_ms == stream_start_ms + + +# --------------------------------------------------------------------------- # +# Input handling # +# --------------------------------------------------------------------------- # + + +def _pcm16_b64(samples: np.ndarray) -> str: + return base64.b64encode((samples * 32767).astype(" None: + config = SileroVADConfig(threshold=0.5) + tone = np.zeros(FRAME * 2, dtype=np.float32) + seen: list[int] = [] + vad = SileroStreamingVAD(config, frame_scorer=lambda frame: seen.append(frame.size) or 0.0) + vad.process_base64(_pcm16_b64(tone), fmt="pcm16", sample_rate_hz=SAMPLE_RATE_HZ) + assert seen == [FRAME, FRAME] + + +def test_an_odd_pcm16_byte_count_is_rejected() -> None: + vad = SileroStreamingVAD(SileroVADConfig(), frame_scorer=lambda _frame: 0.0) + with pytest.raises(ValueError, match="incomplete pcm16"): + vad.process_base64(base64.b64encode(b"\x00\x01\x02").decode(), fmt="pcm16", sample_rate_hz=SAMPLE_RATE_HZ) + + +def test_resampling_keeps_the_frame_grid_from_drifting() -> None: + """The 24 kHz remainder carries across chunks, so frame boundaries cannot drift. + + This is a weaker claim than upstream's ``StreamingAudioResampler``: the + interpolation is per chunk, so values near a chunk edge differ slightly. + What must hold is the frame *count* and alignment, which is what the + endpoint timestamps are derived from. + """ + source = np.sin(np.linspace(0, 40 * np.pi, 24_000, dtype=np.float32)) + + def frames_for(chunk_sizes: list[int]) -> int: + frames = 0 + + def count(_frame: np.ndarray) -> float: + nonlocal frames + frames += 1 + return 0.0 + + vad = SileroStreamingVAD(SileroVADConfig(), frame_scorer=count) + offset = 0 + for size in chunk_sizes: + vad.process_base64(_pcm16_b64(source[offset : offset + size]), fmt="pcm16", sample_rate_hz=24_000) + offset += size + assert offset == source.size + return frames + + whole = frames_for([source.size]) + split = frames_for([7, 13, 511, 512, 513, source.size - 1_556]) + assert whole == split == 16_000 // FRAME + + +def test_a_sample_rate_change_mid_stream_is_rejected() -> None: + vad = SileroStreamingVAD(SileroVADConfig(), frame_scorer=lambda _frame: 0.0) + chunk = _pcm16_b64(np.zeros(1_024, dtype=np.float32)) + vad.process_base64(chunk, fmt="pcm16", sample_rate_hz=24_000) + with pytest.raises(ValueError, match="sample rate cannot change"): + vad.process_base64(chunk, fmt="pcm16", sample_rate_hz=48_000) + + +# --------------------------------------------------------------------------- # +# Backend selection # +# --------------------------------------------------------------------------- # + + +def test_a_configured_model_path_that_does_not_exist_is_an_error() -> None: + """An operator who names an artifact gets told it is missing, not a fallback.""" + provider = SileroVADBackendProvider(model_path="/nonexistent/silero_vad.onnx") + with pytest.raises(ServerVADUnavailableError, match="does not exist"): + provider.get() + + +def test_a_configured_model_with_the_wrong_checksum_is_refused(tmp_path) -> None: + """A different Silero revision changes endpointing, so the pin is enforced.""" + impostor = tmp_path / "silero_vad.onnx" + impostor.write_bytes(b"not the pinned model") + provider = SileroVADBackendProvider(model_path=str(impostor)) + with pytest.raises(ServerVADUnavailableError, match="checksum mismatch"): + provider.get() + assert hashlib.sha256(impostor.read_bytes()).hexdigest() != SILERO_VAD_SHA256 + + +def test_the_onnx_backend_is_shared_by_every_session() -> None: + """One ONNX session for the whole engine: its state is passed in, not held.""" + provider = SileroVADBackendProvider() + try: + backend = provider.get() + except ServerVADUnavailableError as exc: + pytest.skip(f"no Silero backend available here: {exc}") + if not isinstance(backend, SileroVADBackend): + pytest.skip("no local ONNX artifact; the torch fallback is per-session by design") + + assert provider.get() is backend + # State is the caller's, so two sessions on one backend cannot interfere. + first, second = backend.new_state(), backend.new_state() + assert first is not second + probability, next_state = backend.infer(np.zeros(FRAME, dtype=np.float32), first) + assert 0.0 <= probability <= 1.0 + assert next_state is not first + + +def test_the_onnx_backend_holds_the_upstream_silero_v62_contract(monkeypatch, tmp_path) -> None: + """Pin the ONNX call shape without needing the real artifact or runtime. + + The context carry is the part worth pinning: Silero v6.2 is fed 64 samples + of the previous frame ahead of the current 512, and that tail travels in the + caller's state rather than inside the session -- which is exactly what lets + one session serve every duplex session at once. + """ + sessions: list[FakeInferenceSession] = [] + + class FakeSessionOptions: + inter_op_num_threads = 0 + intra_op_num_threads = 0 + + class FakeInferenceSession: + def __init__(self, path: str, *, providers: list[str], sess_options: object) -> None: + self.providers = providers + self.sess_options = sess_options + self.calls: list[dict[str, np.ndarray]] = [] + self.run_barrier: threading.Barrier | None = None + sessions.append(self) + + def get_inputs(self) -> list[SimpleNamespace]: + return [SimpleNamespace(name=name) for name in ("input", "state", "sr")] + + def run(self, _outputs: object, inputs: dict[str, np.ndarray]) -> list[np.ndarray]: + self.calls.append({name: np.array(value, copy=True) for name, value in inputs.items()}) + if self.run_barrier is not None: + self.run_barrier.wait(timeout=5) + return [ + np.asarray([[0.75]], dtype=np.float32), + np.full((2, 1, 128), len(self.calls), dtype=np.float32), + ] + + monkeypatch.setitem( + sys.modules, + "onnxruntime", + SimpleNamespace(SessionOptions=FakeSessionOptions, InferenceSession=FakeInferenceSession), + ) + + model_path = tmp_path / "silero_vad.onnx" + model_path.write_bytes(b"fake-model") + backend = SileroVADBackend(model_path) + session = sessions[0] + + # Single-threaded on purpose: many sessions share this one ORT session, so + # it must not fan out per call. + assert session.providers == ["CPUExecutionProvider"] + assert session.sess_options.inter_op_num_threads == 1 + assert session.sess_options.intra_op_num_threads == 1 + # Warm-up already ran: 64 context + 512 frame. + assert session.calls[0]["input"].shape == (1, 576) + assert session.calls[0]["state"].shape == (2, 1, 128) + assert session.calls[0]["sr"].item() == SAMPLE_RATE_HZ + + state = backend.new_state() + first_frame = np.arange(backend.frame_samples, dtype=np.float32) + probability, state = backend.infer(first_frame, state) + first_call = session.calls[1] + + assert probability == pytest.approx(0.75) + np.testing.assert_array_equal( + first_call["input"][:, : backend.context_samples], + np.zeros((1, backend.context_samples), dtype=np.float32), + ) + np.testing.assert_array_equal(first_call["input"][:, backend.context_samples :], first_frame[None, :]) + + second_frame = -first_frame + backend.infer(second_frame, state) + second_call = session.calls[2] + # The previous frame's tail becomes this frame's context. + np.testing.assert_array_equal( + second_call["input"][:, : backend.context_samples], + first_frame[None, -backend.context_samples :], + ) + np.testing.assert_array_equal(second_call["state"], np.full((2, 1, 128), 2, dtype=np.float32)) + + # Two sessions may enter the shared ORT session at the same time. + session.run_barrier = threading.Barrier(2) + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [ + executor.submit(backend.infer, first_frame, backend.new_state()), + executor.submit(backend.infer, second_frame, backend.new_state()), + ] + assert [future.result()[0] for future in futures] == pytest.approx([0.75, 0.75]) + + +# --------------------------------------------------------------------------- # +# Session payload validation # +# --------------------------------------------------------------------------- # + + +def test_an_unknown_server_vad_field_is_refused() -> None: + """A misspelled knob that silently does nothing is worse than a refusal. + + The endpointing would still look configured while running on defaults. + """ + error = validate_realtime_turn_detection({"turn_detection": {"type": "server_vad", "semantic_eagerness": "high"}}) + assert error is not None + assert "semantic_eagerness" in error + + +def test_the_documented_server_vad_fields_are_all_accepted() -> None: + assert ( + validate_realtime_turn_detection( + { + "turn_detection": { + "type": "server_vad", + "threshold": 0.6, + "prefix_padding_ms": 200, + "silence_duration_ms": 400, + "create_response": True, + "interrupt_response": True, + "min_speech_duration_ms": 120, + }, + "overlap_policy": "barge_in_on_speech", + } + ) + is None + ) diff --git a/tests/engine/test_async_omni_engine_abort_ack.py b/tests/engine/test_async_omni_engine_abort_ack.py index 86276ec66..43a29bf3a 100644 --- a/tests/engine/test_async_omni_engine_abort_ack.py +++ b/tests/engine/test_async_omni_engine_abort_ack.py @@ -54,7 +54,7 @@ def test_abort_async_waits_for_ack(mocker: MockerFixture): rpc_q: queue.Queue = queue.Queue() engine = _make_engine(request_q, rpc_q) mocker.patch( - "vllm_omni.engine.async_omni_engine.uuid.uuid4", + "vllm_omni.engine.omni_engine_base.uuid.uuid4", return_value=SimpleNamespace(hex="abort-rpc-1"), ) @@ -81,7 +81,7 @@ def test_abort_async_raises_on_orchestrator_error(mocker: MockerFixture): rpc_q: queue.Queue = queue.Queue() engine = _make_engine(request_q, rpc_q) mocker.patch( - "vllm_omni.engine.async_omni_engine.uuid.uuid4", + "vllm_omni.engine.omni_engine_base.uuid.uuid4", return_value=SimpleNamespace(hex="abort-rpc-err"), ) @@ -109,7 +109,7 @@ def test_abort_async_times_out_without_result(mocker: MockerFixture): rpc_q: queue.Queue = queue.Queue() engine = _make_engine(request_q, rpc_q) mocker.patch( - "vllm_omni.engine.async_omni_engine.uuid.uuid4", + "vllm_omni.engine.omni_engine_base.uuid.uuid4", return_value=SimpleNamespace(hex="abort-rpc-timeout"), ) @@ -139,7 +139,7 @@ def test_abort_async_preserves_request_queue_backpressure(mocker: MockerFixture) rpc_q: queue.Queue = queue.Queue() engine = _make_engine(request_q, rpc_q) mocker.patch( - "vllm_omni.engine.async_omni_engine.uuid.uuid4", + "vllm_omni.engine.omni_engine_base.uuid.uuid4", return_value=SimpleNamespace(hex="blocked-abort"), ) diff --git a/tests/engine/test_async_omni_engine_outputs.py b/tests/engine/test_async_omni_engine_outputs.py index 6b0d6a188..425a16fc1 100644 --- a/tests/engine/test_async_omni_engine_outputs.py +++ b/tests/engine/test_async_omni_engine_outputs.py @@ -7,7 +7,6 @@ Focuses on the critical behavior: when the orchestrator thread dies, subsequent attempts to collect output raise RuntimeError. """ -import asyncio import queue import threading from concurrent.futures import ThreadPoolExecutor @@ -18,16 +17,6 @@ from pytest_mock import MockerFixture from vllm_omni.engine.async_engine_utils import weak_shutdown_async_omni_engine from vllm_omni.engine.async_omni_engine import AsyncOmniEngine -from vllm_omni.engine.duplex.control_client import ( - DuplexControlClient, - DuplexControlRequestError, -) -from vllm_omni.engine.duplex.messages import ( - AppendDuplexInputMessage, - DuplexControlResultMessage, - DuplexFence, - SignalDuplexTurnMessage, -) from vllm_omni.engine.messages import ( CollectiveRPCResultMessage, ErrorMessage, @@ -166,9 +155,9 @@ def test_output_remains_on_shared_output_path(mocker: MockerFixture): raw_queue = queue.Queue() raw_queue.put_nowait( OutputMessage( - request_id="legacy-duplex-request", + request_id="shared-path-request", stage_id=0, - engine_outputs=OmniRequestOutput(request_id="legacy-duplex-request"), + engine_outputs=OmniRequestOutput(request_id="shared-path-request"), finished=False, ) ) @@ -176,293 +165,7 @@ def test_output_remains_on_shared_output_path(mocker: MockerFixture): output = engine.try_get_output(timeout=0.01) - assert output.request_id == "legacy-duplex-request" - - -def test_duplex_control_client_uses_engine_owned_correlated_transport(): - engine = object.__new__(AsyncOmniEngine) - engine._duplex_control_client = None - engine._correlated_rpc_client = CorrelatedRpcClient(queue.Queue(), queue.Queue()) - - try: - client = engine._get_duplex_control_client() - - assert isinstance(client, DuplexControlClient) - assert client._transport is engine._correlated_rpc_client - assert not hasattr(engine, "_submit_rpc_and_wait") - finally: - engine._correlated_rpc_client.close() - - -def test_duplex_control_api_requires_explicit_fence(): - engine = object.__new__(AsyncOmniEngine) - - with pytest.raises(TypeError, match="fence"): - engine.open_duplex_session("sid") - with pytest.raises(TypeError, match="fence"): - engine.append_duplex_input("sid", mode="append_audio_chunk", payload={}) - with pytest.raises(TypeError, match="fence"): - engine.signal_duplex_turn("sid", event="turn.end") - with pytest.raises(TypeError, match="fence"): - engine.close_duplex_session("sid") - - -def test_open_duplex_session_waits_for_control_ack(mocker: MockerFixture): - request_q = queue.Queue() - rpc_q = queue.Queue() - control_result = DuplexControlResultMessage( - control_id="ctrl-1", - fence=DuplexFence("sid"), - operation="open", - session_id="sid", - ok=False, - stage_results=[{"stage_id": 0, "replica_id": 0, "result": {"supported": False}}], - unsupported_count=1, - error_count=0, - ) - - engine = object.__new__(AsyncOmniEngine) - engine.request_queue = SimpleNamespace(sync_q=request_q) - engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) - engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) - engine._duplex_control_client = None - mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="ctrl-1")) - - with ThreadPoolExecutor(max_workers=1) as executor: - pending = executor.submit(engine.open_duplex_session, "sid", fence=DuplexFence("sid"), timeout=1) - msg = request_q.get(timeout=1) - assert msg.type == "open_duplex_session" - assert msg.control_id == "ctrl-1" - rpc_q.put(control_result) - with pytest.raises(DuplexControlRequestError) as exc_info: - pending.result(timeout=1) - - assert exc_info.value.result["unsupported_count"] == 1 - assert exc_info.value.result["stage_results"][0]["result"]["supported"] is False - engine._correlated_rpc_client.close() - - -def test_signal_duplex_turn_message_carries_next_fence(mocker: MockerFixture): - request_q = queue.Queue() - rpc_q = queue.Queue() - cancelled_fence = DuplexFence("sid-cancel", epoch=4, turn_id=2) - next_fence = DuplexFence("sid-cancel", epoch=5, turn_id=2) - control_result = DuplexControlResultMessage( - control_id="ctrl-signal", - fence=cancelled_fence, - operation="signal", - session_id="sid-cancel", - ok=True, - stage_results=[], - ) - - engine = object.__new__(AsyncOmniEngine) - engine.request_queue = SimpleNamespace(sync_q=request_q) - engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) - engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) - engine._duplex_control_client = None - mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="ctrl-signal")) - - with ThreadPoolExecutor(max_workers=1) as executor: - pending = executor.submit( - engine.signal_duplex_turn, - "sid-cancel", - event="input.cancel", - fence=cancelled_fence, - next_fence=next_fence, - session_config={"temperature": 0.0}, - timeout=1, - ) - msg = request_q.get(timeout=1) - assert isinstance(msg, SignalDuplexTurnMessage) - assert msg.fence == cancelled_fence - assert msg.next_fence == next_fence - assert msg.session_config == {"temperature": 0.0} - rpc_q.put(control_result) - assert pending.result(timeout=1)["ok"] is True - - engine._correlated_rpc_client.close() - - -@pytest.mark.asyncio -async def test_cancelled_async_append_does_not_stop_executor_callable(mocker: MockerFixture): - entered = threading.Event() - release = threading.Event() - completed = threading.Event() - engine = object.__new__(AsyncOmniEngine) - - def blocking_append(*args, **kwargs): - del args, kwargs - entered.set() - release.wait(timeout=1) - completed.set() - return {} - - mocker.patch.object(engine, "append_duplex_input", side_effect=blocking_append) - task = asyncio.create_task( - engine.append_duplex_input_async( - "sid-executor-cancel", - mode="append_audio_chunk", - payload=b"audio", - fence=DuplexFence("sid-executor-cancel"), - ) - ) - assert await asyncio.to_thread(entered.wait, 1) - - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - release.set() - - assert await asyncio.to_thread(completed.wait, 1) - - -def test_duplex_controls_route_out_of_order_without_global_lock(): - request_q = queue.Queue() - rpc_q = queue.Queue() - engine = object.__new__(AsyncOmniEngine) - engine.request_queue = SimpleNamespace(sync_q=request_q) - engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) - engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) - engine._duplex_control_client = None - first_fence = DuplexFence("sid", epoch=1) - second_fence = DuplexFence("sid", epoch=2) - first = AppendDuplexInputMessage( - control_id="first", - fence=first_fence, - session_id="sid", - expected_epoch=1, - mode="audio", - payload=b"first", - ) - second = AppendDuplexInputMessage( - control_id="second", - fence=second_fence, - session_id="sid", - expected_epoch=2, - mode="audio", - payload=b"second", - ) - - with ThreadPoolExecutor(max_workers=2) as executor: - client = engine._get_duplex_control_client() - first_result = executor.submit(client.execute, first, timeout=1) - second_result = executor.submit(client.execute, second, timeout=1) - submitted = {request_q.get(timeout=1).control_id, request_q.get(timeout=1).control_id} - assert submitted == {"first", "second"} - rpc_q.put( - DuplexControlResultMessage( - control_id="second", - fence=second_fence, - operation="append", - session_id="sid", - ok=True, - stage_results=[], - ) - ) - rpc_q.put( - DuplexControlResultMessage( - control_id="first", - fence=first_fence, - operation="append", - session_id="sid", - ok=True, - stage_results=[], - ) - ) - - assert first_result.result(timeout=1)["fence"] == first_fence - assert second_result.result(timeout=1)["fence"] == second_fence - - assert not hasattr(engine, "_rpc_lock") - engine._correlated_rpc_client.close() - - -def test_terminal_rpc_router_rejects_without_enqueuing_new_control(): - request_q = queue.Queue() - rpc_q = queue.Queue() - engine = object.__new__(AsyncOmniEngine) - engine.request_queue = SimpleNamespace(sync_q=request_q) - engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) - engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) - engine._duplex_control_client = None - pending_message = AppendDuplexInputMessage( - control_id="pending", - fence=DuplexFence("sid"), - session_id="sid", - mode="audio", - payload=b"audio", - ) - message = AppendDuplexInputMessage( - control_id="after-fatal", - fence=DuplexFence("sid"), - session_id="sid", - mode="audio", - payload=b"audio", - ) - - try: - with ThreadPoolExecutor(max_workers=1) as executor: - pending = executor.submit(engine._get_duplex_control_client().execute, pending_message, timeout=1) - assert request_q.get(timeout=1).control_id == "pending" - rpc_q.put(ErrorMessage(error="orchestrator failed", fatal=True)) - with pytest.raises(RuntimeError, match="orchestrator failed"): - pending.result(timeout=1) - - with pytest.raises(RuntimeError, match="orchestrator failed"): - engine._get_duplex_control_client().execute(message, timeout=1) - - assert request_q.empty() - finally: - engine._correlated_rpc_client.close() - - -def test_collective_and_duplex_results_share_router_without_swapping(mocker: MockerFixture): - request_q = queue.Queue() - rpc_q = queue.Queue() - engine = object.__new__(AsyncOmniEngine) - engine.request_queue = SimpleNamespace(sync_q=request_q) - engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) - engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) - engine._duplex_control_client = None - fence = DuplexFence("sid") - duplex = AppendDuplexInputMessage( - control_id="duplex-control", - fence=fence, - session_id="sid", - mode="audio", - payload=b"audio", - ) - mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="collective-rpc")) - - with ThreadPoolExecutor(max_workers=2) as executor: - duplex_result = executor.submit(engine._get_duplex_control_client().execute, duplex, timeout=1) - collective_result = executor.submit(engine.collective_rpc, "health", timeout=1) - submitted_types = {request_q.get(timeout=1).type, request_q.get(timeout=1).type} - assert submitted_types == {"append_duplex_input", "collective_rpc"} - rpc_q.put( - CollectiveRPCResultMessage( - rpc_id="collective-rpc", - method="health", - stage_ids=[0], - results=["healthy"], - ) - ) - rpc_q.put( - DuplexControlResultMessage( - control_id="duplex-control", - fence=fence, - operation="append", - session_id="sid", - ok=True, - stage_results=[], - ) - ) - - assert collective_result.result(timeout=1) == ["healthy"] - assert duplex_result.result(timeout=1)["fence"] == fence - - engine._correlated_rpc_client.close() + assert output.request_id == "shared-path-request" def test_collective_rpc_preserves_request_queue_backpressure(mocker: MockerFixture): @@ -483,8 +186,7 @@ def test_collective_rpc_preserves_request_queue_backpressure(mocker: MockerFixtu engine.request_queue = SimpleNamespace(sync_q=request_q) engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) - engine._duplex_control_client = None - mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="blocked-rpc")) + mocker.patch("vllm_omni.engine.omni_engine_base.uuid.uuid4", return_value=SimpleNamespace(hex="blocked-rpc")) with ThreadPoolExecutor(max_workers=1) as executor: pending = executor.submit(engine.collective_rpc, "health", timeout=1) @@ -506,36 +208,6 @@ def test_collective_rpc_preserves_request_queue_backpressure(mocker: MockerFixtu engine._correlated_rpc_client.close() -def test_open_duplex_session_raises_on_stage_control_error(mocker: MockerFixture): - request_q = queue.Queue() - rpc_q = queue.Queue() - control_result = DuplexControlResultMessage( - control_id="ctrl-error", - fence=DuplexFence("sid"), - operation="open", - session_id="sid", - ok=False, - stage_results=[{"stage_id": 0, "replica_id": 0, "result": {"supported": False, "error": "boom"}}], - unsupported_count=1, - error_count=1, - ) - - engine = object.__new__(AsyncOmniEngine) - engine.request_queue = SimpleNamespace(sync_q=request_q) - engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) - engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) - engine._duplex_control_client = None - mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="ctrl-error")) - - with ThreadPoolExecutor(max_workers=1) as executor: - pending = executor.submit(engine.open_duplex_session, "sid", fence=DuplexFence("sid"), timeout=1) - request_q.get(timeout=1) - rpc_q.put(control_result) - with pytest.raises(RuntimeError, match="duplex open failed"): - pending.result(timeout=1) - engine._correlated_rpc_client.close() - - def test_shutdown_does_not_race_runtime_cleanup_with_live_orchestrator(mocker: MockerFixture): engine = object.__new__(AsyncOmniEngine) engine._shutdown_called = False diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 27677bd54..f052f7e9d 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -12,7 +12,7 @@ import types import pytest from vllm_omni.diffusion.data import AttentionConfig -from vllm_omni.engine import async_omni_engine as async_omni_engine_module +from vllm_omni.engine import omni_engine_base as async_omni_engine_module from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.stage_init_utils import ( LogicalStageInitPlan, @@ -188,7 +188,7 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred(): def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod + import vllm_omni.engine.omni_engine_base as engine_mod engine = object.__new__(AsyncOmniEngine) engine.stage_configs = [types.SimpleNamespace()] diff --git a/tests/engine/test_diffusion_streaming_output_kwarg.py b/tests/engine/test_diffusion_streaming_output_kwarg.py index a6f3dff02..8bc00f587 100644 --- a/tests/engine/test_diffusion_streaming_output_kwarg.py +++ b/tests/engine/test_diffusion_streaming_output_kwarg.py @@ -33,7 +33,7 @@ def _resolve(monkeypatch, kwargs: dict[str, Any]) -> dict[str, Any]: return "config.yaml", [], None monkeypatch.setattr( - "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", + "vllm_omni.engine.omni_engine_base.load_and_resolve_stage_configs", _fake_load, ) engine = SimpleNamespace( diff --git a/tests/engine/test_duplex_import_boundary.py b/tests/engine/test_duplex_import_boundary.py index 013ab8661..1da513182 100644 --- a/tests/engine/test_duplex_import_boundary.py +++ b/tests/engine/test_duplex_import_boundary.py @@ -45,27 +45,62 @@ def test_isolated_import_ignores_conflicting_gpu_visibility(monkeypatch: pytest. _assert_isolated_import_succeeds("import vllm_omni.outputs") -def test_stable_engine_imports_load_duplex_kernel_eagerly() -> None: - # The duplex kernel is imported eagerly by the stable engine modules. - # Model-specific duplex adapters must still load only via the - # dotted-string plugin paths, and the client-side vllm_omni.clients - # package must never enter the stable runtime import graph. +def test_turn_based_stack_has_no_duplex_vocabulary() -> None: + # The generic bases and the turn-based siblings (AsyncOmni / AsyncOmniEngine / + # Orchestrator) carry no duplex code: importing them must not load the + # engine-side duplex package, the model plugins or the client library. _assert_isolated_import_succeeds(""" import sys import vllm_omni.engine.async_omni_engine +import vllm_omni.engine.omni_engine_base import vllm_omni.engine.orchestrator import vllm_omni.entrypoints.async_omni +import vllm_omni.entrypoints.async_omni_base + +forbidden_prefixes = ( + "vllm_omni.engine.duplex", + "vllm_omni.engine.duplex_omni_engine", + "vllm_omni.engine.duplex_orchestrator", + "vllm_omni.entrypoints.duplex", + "vllm_omni.entrypoints.duplex_omni", + "vllm_omni.clients", + "vllm_omni.model_executor.models.minicpmo_4_5.duplex", + "vllm_omni.model_executor.models.nemotron_voicechat.duplex", + "vllm_omni.model_executor.models.personaplex.duplex", +) +loaded = sorted( + name + for name in sys.modules + if any(name == prefix or name.startswith(prefix + ".") for prefix in forbidden_prefixes) +) +if loaded: + raise SystemExit("turn-based imports loaded duplex modules: " + ", ".join(loaded)) +""") + + +def test_duplex_stack_loads_the_framework_but_not_the_model_plugins() -> None: + # The duplex siblings import the engine-side duplex framework eagerly; + # model plugins still load only through PipelineConfig.duplex_plugin, and + # the client-side package never enters the server import graph. + _assert_isolated_import_succeeds(""" +import sys + +import vllm_omni.engine.duplex_omni_engine +import vllm_omni.engine.duplex_orchestrator +import vllm_omni.entrypoints.duplex_omni +import vllm_omni.entrypoints.duplex.serving expected_eager = ( - "vllm_omni.engine.duplex.contracts", - "vllm_omni.engine.duplex.control_plane", - "vllm_omni.entrypoints.duplex_request_client", - "vllm_omni.outputs.duplex", + "vllm_omni.engine.duplex.commands", + "vllm_omni.engine.duplex.events", + "vllm_omni.engine.duplex.session.engine_session", + "vllm_omni.engine.duplex.session.manager", + "vllm_omni.engine.duplex.plugin", ) missing = sorted(name for name in expected_eager if name not in sys.modules) if missing: - raise SystemExit("duplex kernel modules not imported eagerly: " + ", ".join(missing)) + raise SystemExit("duplex framework modules not imported eagerly: " + ", ".join(missing)) forbidden_prefixes = ( "vllm_omni.clients", @@ -79,7 +114,7 @@ loaded = sorted( if any(name == prefix or name.startswith(prefix + ".") for prefix in forbidden_prefixes) ) if loaded: - raise SystemExit("stable imports loaded plugin-only duplex modules: " + ", ".join(loaded)) + raise SystemExit("duplex imports loaded plugin-only or client modules: " + ", ".join(loaded)) """) @@ -190,16 +225,32 @@ def test_engine_duplex_uses_canonical_contract_module_names() -> None: core_dir = REPO_ROOT / "vllm_omni" / "experimental" / "fullduplex" / "core" for name in ( + "commands.py", "contracts.py", - "lease.py", + "events.py", "messages.py", - "session.py", - "control_plane.py", - "control_client.py", - "runtime.py", + "plugin.py", "intermediate.py", ): assert (engine_dir / name).is_file() + # one session and everything that runs it lives in the session subpackage + for name in ( + "__init__.py", + "engine_session.py", + "manager.py", + "runner.py", + "context.py", + "emitter.py", + "model_channel.py", + "control.py", + "append_task.py", + "helpers.py", + "lease.py", + ): + assert (engine_dir / "session" / name).is_file() + assert not (engine_dir / "session.py").exists() + for removed in ("control_plane.py", "control_client.py", "runtime.py"): + assert not (engine_dir / removed).exists() # the duplex_ prefix is dropped inside the duplex package assert not (engine_dir / "duplex_session.py").exists() assert not (engine_dir / "duplex_lease.py").exists() diff --git a/tests/engine/test_duplex_omni_engine.py b/tests/engine/test_duplex_omni_engine.py new file mode 100644 index 000000000..06af9e7af --- /dev/null +++ b/tests/engine/test_duplex_omni_engine.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""``DuplexOmniEngine``: the session message surface over the generic engine base.""" + +from __future__ import annotations + +import asyncio +import queue +from types import SimpleNamespace +from typing import Any + +import pytest + +from vllm_omni.engine.duplex.commands import Heartbeat +from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig +from vllm_omni.engine.duplex.messages import ( + CloseDuplexSessionMessage, + DuplexControlResultMessage, + DuplexSessionCommandMessage, + DuplexSessionError, + OpenDuplexSessionMessage, + ResumeDuplexSessionMessage, + TouchDuplexSessionMessage, +) +from vllm_omni.engine.duplex_omni_engine import DuplexOmniEngine + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FakeRpcClient: + def __init__(self, result: Any) -> None: + self.result = result + self.calls: list[tuple[tuple[str, str], Any, float | None]] = [] + #: Control RPCs must block rather than raise a raw queue.Full. + self.block_on_submit_flags: list[bool] = [] + + def execute(self, key, message, *, timeout=None, timeout_message=None, block_on_submit=False): + self.calls.append((key, message, timeout)) + self.block_on_submit_flags.append(block_on_submit) + if isinstance(self.result, BaseException): + raise self.result + return self.result + + +class _FakeSyncQueue: + def __init__(self, *, full: bool = False) -> None: + self.items: list[Any] = [] + self.full = full + + def put(self, item: Any, timeout: float | None = None) -> None: + if self.full: + raise queue.Full + self.items.append(item) + + +def _engine(result: Any, *, alive: bool = True, full: bool = False) -> DuplexOmniEngine: + engine = object.__new__(DuplexOmniEngine) + engine._correlated_rpc_client = _FakeRpcClient(result) + engine.request_queue = SimpleNamespace(sync_q=_FakeSyncQueue(full=full)) + engine.is_alive = lambda: alive + engine.plugin = None + return engine + + +def _ok(operation: str, **fields: Any) -> DuplexControlResultMessage: + return DuplexControlResultMessage(control_id="c", operation=operation, session_id="sid", ok=True, **fields) + + +@pytest.mark.asyncio +async def test_open_session_runs_the_correlated_rpc_with_the_typed_config() -> None: + result = _ok("open", capabilities=DuplexCapabilities(), public_session={"id": "sid"}, lease_generation=0) + engine = _engine(result) + config = DuplexSessionConfig(model="m") + + returned = await engine.open_session_async("sid", config, 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 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. + assert engine.rpc_client.block_on_submit_flags == [True] + + +@pytest.mark.asyncio +async def test_close_resume_and_touch_build_their_messages_without_incarnation() -> None: + engine = _engine(_ok("close")) + await engine.close_session_async("sid", reason="bye") + await engine.resume_session_async("sid", expected_lease_generation=4) + await engine.touch_session_async("sid", activity="detach") + + close, resume, touch = (call[1] for call in engine.rpc_client.calls) + assert isinstance(close, CloseDuplexSessionMessage) and close.reason == "bye" + assert isinstance(resume, ResumeDuplexSessionMessage) and resume.expected_lease_generation == 4 + assert isinstance(touch, TouchDuplexSessionMessage) and touch.activity == "detach" + for message in (close, resume, touch): + assert message.session_id == "sid" + assert not hasattr(message, "incarnation") + + +@pytest.mark.asyncio +async def test_failed_control_result_maps_to_a_typed_session_error() -> None: + engine = _engine( + DuplexControlResultMessage( + control_id="c", + operation="open", + session_id="sid", + ok=False, + error_code="resource_exhausted", + error_message="duplex_session_capacity_exhausted: limit=1", + error_retryable=True, + ) + ) + with pytest.raises(DuplexSessionError) as excinfo: + await engine.open_session_async("sid", DuplexSessionConfig()) + assert excinfo.value.code == "resource_exhausted" + assert excinfo.value.retryable is True + assert excinfo.value.session_id == "sid" + + +@pytest.mark.asyncio +async def test_engine_error_and_timeout_map_to_session_errors() -> None: + # The correlated RPC client raises RuntimeError for a closed router or a + # terminal orchestrator error; both become a typed engine_error. + engine = _engine(RuntimeError("orchestrator exploded")) + with pytest.raises(DuplexSessionError, match="orchestrator exploded") as excinfo: + await engine.close_session_async("sid") + assert excinfo.value.code == "engine_error" + + engine = _engine(TimeoutError("duplex open timed out for session sid")) + with pytest.raises(DuplexSessionError) as excinfo: + await engine.open_session_async("sid", DuplexSessionConfig()) + assert excinfo.value.code == "timeout" and excinfo.value.retryable is True + + +@pytest.mark.asyncio +async def test_submit_command_is_one_way_on_the_request_queue() -> None: + engine = _engine(None) + command = Heartbeat(event_id="evt-1") + + await engine.submit_command_async("sid", command) + + (message,) = engine.request_queue.sync_q.items + assert isinstance(message, DuplexSessionCommandMessage) + assert message.session_id == "sid" and message.command is command + assert engine.rpc_client.calls == [] + + +@pytest.mark.asyncio +async def test_submit_command_rejects_dead_engine_and_full_queue() -> None: + with pytest.raises(DuplexSessionError) as excinfo: + await _engine(None, alive=False).submit_command_async("sid", Heartbeat()) + assert excinfo.value.code == "engine_dead" + + with pytest.raises(DuplexSessionError) as excinfo: + await _engine(None, full=True).submit_command_async("sid", Heartbeat()) + assert excinfo.value.code == "engine_backpressure" and excinfo.value.retryable is True + + +def test_duplex_capabilities_come_from_the_loaded_plugin() -> None: + engine = object.__new__(DuplexOmniEngine) + engine.plugin = None + with pytest.raises(RuntimeError, match="not loaded"): + engine.duplex_capabilities + capabilities = DuplexCapabilities(supports_barge_in=False) + engine.plugin = SimpleNamespace(capabilities=lambda *, max_sessions: capabilities) + engine.duplex_session_config = SimpleNamespace(max_sessions=3) + assert engine.duplex_capabilities is capabilities + + +def test_validate_deployment_requires_a_plugin_and_duplex_session_mode() -> None: + engine = object.__new__(DuplexOmniEngine) + engine.model = "plain-model" + engine.pipeline_config = SimpleNamespace(duplex_plugin=None) + engine.deploy_config = SimpleNamespace(session_mode="duplex") + with pytest.raises(ValueError, match="not a duplex model"): + engine._validate_deployment() + + engine.pipeline_config = SimpleNamespace(duplex_plugin="pkg.mod.Plugin") + engine.deploy_config = SimpleNamespace(session_mode="turn") + with pytest.raises(ValueError, match="session_mode: duplex"): + engine._validate_deployment() + + engine.deploy_config = None + with pytest.raises(ValueError, match="none resolved"): + engine._validate_deployment() + + +@pytest.mark.asyncio +async def test_async_wrappers_run_the_blocking_calls_off_the_event_loop() -> None: + loop = asyncio.get_running_loop() + engine = _engine(_ok("touch")) + seen: list[bool] = [] + original = engine.rpc_client.execute + + def execute(*args, **kwargs): + # Called from the default executor thread, never on the event loop thread. + try: + seen.append(asyncio.get_running_loop() is loop) + except RuntimeError: + seen.append(False) + return original(*args, **kwargs) + + engine.rpc_client.execute = execute + await engine.touch_session_async("sid", activity="heartbeat") + assert seen == [False] diff --git a/tests/engine/test_duplex_orchestrator.py b/tests/engine/test_duplex_orchestrator.py new file mode 100644 index 000000000..5b8130c36 --- /dev/null +++ b/tests/engine/test_duplex_orchestrator.py @@ -0,0 +1,517 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""``DuplexOrchestrator``: the template seams and the stage port around ``OrchestratorBase``.""" + +from __future__ import annotations + +import asyncio +import struct +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +from vllm.v1.engine.exceptions import EngineDeadError + +from tests.engine.test_orchestrator import ( + FakeOutputProcessor, + FakeRunningCounter, + FakeStageClient, + _build_stage_pools, +) +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.messages import ( + CloseDuplexSessionMessage, + DuplexControlResultMessage, + DuplexSessionCommandMessage, + OpenDuplexSessionMessage, +) +from vllm_omni.engine.duplex_orchestrator import DuplexOrchestrator, DuplexOrchestratorRequestState +from vllm_omni.engine.messages import AbortRequestMessage, ShutdownRequestMessage +from vllm_omni.engine.orchestrator import Orchestrator +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.plugin import MiniCPMO45DuplexPlugin + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +SESSION_ID = "duplex-orch" + + +def _encode_audio(audio: object, sample_rate_hz: int, response_format: str, speed: float | None) -> str | None: + del sample_rate_hz, response_format, speed + samples = int(np.asarray(audio, dtype=np.float32).size) if audio is not None else 0 + return f"wav-{samples}" if samples > 0 else None + + +def _stage_configs(count: int) -> list[object]: + return [SimpleNamespace(model_config=SimpleNamespace(max_model_len=4096)) for _ in range(count)] + + +def _build( + *, + stages: int = 1, + running_counter: FakeRunningCounter | None = None, + runtime_config: DuplexSessionRuntimeConfig | None = None, +) -> tuple[DuplexOrchestrator, list[FakeStageClient], asyncio.Queue, asyncio.Queue]: + clients = [FakeStageClient(stage_type="llm", final_output=index == stages - 1) for index in range(stages)] + pools = _build_stage_pools( + [[client] for client in clients], + output_processors=[FakeOutputProcessor() for _ in clients], + stage_vllm_configs=_stage_configs(stages), + ) + rpc_q: asyncio.Queue = asyncio.Queue() + output_q: asyncio.Queue = asyncio.Queue() + orchestrator = DuplexOrchestrator( + request_async_queue=asyncio.Queue(), + output_async_queue=output_q, + rpc_async_queue=rpc_q, + stage_pools=pools, + running_counter=running_counter, + plugin=MiniCPMO45DuplexPlugin(_encode_audio), + duplex_session_config=runtime_config or DuplexSessionRuntimeConfig(reaper_interval_s=0.01), + model_config=None, + ) + return orchestrator, clients, rpc_q, output_q + + +def _open_message(extra_body: dict[str, object] | None = None) -> OpenDuplexSessionMessage: + return OpenDuplexSessionMessage( + control_id=f"open-{SESSION_ID}", + session_id=SESSION_ID, + session_config=DuplexSessionConfig( + model="openbmb/MiniCPM-o-4_5", + modalities=["text"], + instructions="public instructions", + extra_body={"auto_response": True, **(extra_body or {})}, + ), + ) + + +async def _open(orchestrator: DuplexOrchestrator, rpc_q: asyncio.Queue, **kwargs: Any) -> DuplexControlResultMessage: + await orchestrator._dispatch_message(_open_message(**kwargs)) + result = await asyncio.wait_for(rpc_q.get(), timeout=2.0) + assert isinstance(result, DuplexControlResultMessage) + return result + + +async def _close(orchestrator: DuplexOrchestrator, rpc_q: asyncio.Queue) -> DuplexControlResultMessage: + await orchestrator._dispatch_message( + CloseDuplexSessionMessage(control_id=f"close-{SESSION_ID}", session_id=SESSION_ID, reason="test") + ) + result = await asyncio.wait_for(rpc_q.get(), timeout=2.0) + await _settle(orchestrator) + return result + + +def _stage0_request_id(epoch: int = 0) -> str: + return duplex_resource_request_id(DuplexFence(SESSION_ID, epoch=epoch), "stage0") + + +def _append_audio(samples: int = 16000) -> commands.AppendAudio: + return commands.AppendAudio( + audio=struct.pack(f"<{samples}f", *([0.05] * samples)), + format="pcm_f32le", + sample_rate_hz=16000, + is_speech=True, + ) + + +async def _submit(orchestrator: DuplexOrchestrator, command: commands.DuplexCommand) -> None: + assert await orchestrator._dispatch_message(DuplexSessionCommandMessage(session_id=SESSION_ID, command=command)) + await _settle(orchestrator) + + +async def _settle(orchestrator: DuplexOrchestrator, *, timeout_s: float = 3.0) -> None: + """Let the session runner drain its mailbox and append tasks.""" + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout_s + quiet = 0 + while loop.time() < deadline: + runner = orchestrator.session_manager.runners.get(SESSION_ID) + busy = runner is not None and ( + not runner._mailbox.empty() + or any(not task.done() for task in runner.tasks.append_tasks) + or any(not task.done() for task in runner._background_tasks) + ) + busy = busy or any(not task.done() for task in orchestrator.session_manager._dispatched_control_tasks) + quiet = 0 if busy else quiet + 1 + if quiet >= 5: + return + await asyncio.sleep(0.01) + + +def _tts_output(request_id: str, *, samples: int = 24000, text: str = "hello") -> SimpleNamespace: + return SimpleNamespace( + request_id=request_id, + finished=False, + outputs=[SimpleNamespace(text=text, token_ids=[], multimodal_output={})], + multimodal_output={ + "audio": np.zeros(samples, dtype=np.float32), + "sr": 24000, + "meta.duplex_turn_id": np.array([0], dtype=np.int32), + "meta.duplex_epoch": np.array([0], dtype=np.int32), + }, + ) + + +# --------------------------------------------------------------------------- # +# Message dispatch and session lifecycle # +# --------------------------------------------------------------------------- # + + +def test_turn_based_orchestrator_has_no_session_manager() -> None: + orchestrator = Orchestrator( + request_async_queue=asyncio.Queue(), + output_async_queue=asyncio.Queue(), + rpc_async_queue=asyncio.Queue(), + stage_pools=[], + ) + assert not hasattr(orchestrator, "session_manager") + + +@pytest.mark.asyncio +async def test_duplex_messages_are_dispatched_and_generic_ones_are_not() -> None: + orchestrator, _, rpc_q, _ = _build() + assert await orchestrator._dispatch_message(AbortRequestMessage(request_ids=["x"])) is False + assert await orchestrator._dispatch_message(ShutdownRequestMessage()) is False + assert await orchestrator._dispatch_message(_open_message()) is True + assert (await asyncio.wait_for(rpc_q.get(), timeout=2.0)).ok is True + await orchestrator.session_manager.shutdown() + + +@pytest.mark.asyncio +async def test_open_preregisters_the_stage0_request_and_close_releases_it() -> None: + counter = FakeRunningCounter() + orchestrator, clients, rpc_q, _ = _build(running_counter=counter) + result = await _open(orchestrator, rpc_q) + assert result.ok is True + assert result.public_session["id"] == SESSION_ID + assert result.capabilities is not None and result.capabilities.supports_input_append + + request_id = _stage0_request_id() + request_state = orchestrator.request_states[request_id] + assert isinstance(request_state, DuplexOrchestratorRequestState) + assert request_state.session_owned is True + assert request_state.session_id == SESSION_ID + assert request_state.fence == DuplexFence(SESSION_ID) + assert request_state.streaming.enabled is True + # Preregistration reserves the id; nothing is running until an append submits. + assert counter.value == 0 + assert clients[0].add_request_calls == [] + assert orchestrator.session_manager.runner_for_request_id(request_id) is not None + + closed = await _close(orchestrator, rpc_q) + assert closed.ok is True + assert request_id not in orchestrator.request_states + assert orchestrator.session_manager.runner_for_request_id(request_id) is None + assert counter.value == 0 + assert orchestrator.session_manager.active_count() == 0 + + +@pytest.mark.asyncio +async def test_open_failure_rolls_back_session_and_reserved_request() -> None: + orchestrator, _, rpc_q, _ = _build() + # Server-owned runtime keys may not come from the client. + result = await _open(orchestrator, rpc_q, extra_body={"duplex_stage_max_tokens": {"0": 99}}) + + assert result.ok is False + assert result.error_code == "invalid_duplex_runtime_config" + assert orchestrator.request_states == {} + assert orchestrator.session_manager.active_count() == 0 + + +@pytest.mark.asyncio +async def test_open_fails_cleanly_when_a_stage_has_no_live_replica() -> None: + orchestrator, _, rpc_q, _ = _build(stages=2) + orchestrator.stage_pools[1].evict_replica(0) + + result = await _open(orchestrator, rpc_q) + + assert result.ok is False + assert "stage 1 has no live replica" in (result.error_message or "") + assert orchestrator.request_states == {} + assert orchestrator.session_manager.active_count() == 0 + + +@pytest.mark.asyncio +async def test_bridge_state_keeps_public_and_runtime_config_separate() -> None: + orchestrator, _, rpc_q, _ = _build() + await _open(orchestrator, rpc_q) + request_state = orchestrator.request_states[_stage0_request_id()] + bridge = request_state.streaming.bridge_states["duplex"] + + assert bridge["session_id"] == SESSION_ID + assert bridge["fence"] == DuplexFence(SESSION_ID) + assert (bridge["epoch"], bridge["turn_id"], bridge["model_turn_id"]) == (0, 0, 0) + assert "incarnation" not in bridge + assert bridge["session_config"]["instructions"] == "public instructions" + assert "duplex_stage_sampling_params" not in bridge["session_config"] + assert bridge["runtime_config"]["instructions"] == "public instructions" + assert bridge["runtime_config"]["duplex_stage_max_tokens"] == {"0": 20, "1": 8192} + # The plugin's sampling policy is applied to the preregistered request. + assert request_state.sampling_params_list[0].max_tokens == 20 + await orchestrator.session_manager.shutdown() + + +# --------------------------------------------------------------------------- # +# Stage port # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_append_submits_the_resumable_stage0_request_and_counts_it_running() -> None: + counter = FakeRunningCounter() + orchestrator, clients, rpc_q, _ = _build(running_counter=counter) + await _open(orchestrator, rpc_q) + request_id = _stage0_request_id() + + await _submit(orchestrator, _append_audio()) + assert len(clients[0].add_request_calls) == 1 + submitted = clients[0].add_request_calls[0][0] + assert submitted.request_id == request_id + assert submitted.resumable is True + assert submitted.sampling_params.max_tokens == 20 + request_state = orchestrator.request_states[request_id] + assert request_state.stage_fences[0] == DuplexFence(SESSION_ID) + assert 0 in request_state.stage_submit_ts + assert counter.value == 1 + session = orchestrator.session_manager.get(SESSION_ID) + assert session is not None and session.stage_request_submitted(0, request_id) + + # Later units update the same resumable request; the counter stays at one. + await _submit(orchestrator, _append_audio()) + assert len(clients[0].add_request_calls) == 2 + assert clients[0].add_request_calls[1][0].request_id == request_id + assert counter.value == 1 + + await _close(orchestrator, rpc_q) + assert clients[0].abort_calls == [[request_id]] + assert counter.value == 0 + + +@pytest.mark.asyncio +async def test_session_update_refreshes_the_next_append_sampling_params() -> None: + orchestrator, clients, rpc_q, _ = _build() + await _open(orchestrator, rpc_q) + request_state = orchestrator.request_states[_stage0_request_id()] + assert request_state.sampling_params_list[0].max_tokens == 20 + + await _submit(orchestrator, commands.UpdateSession(patch={"max_output_tokens": 7})) + await _submit(orchestrator, _append_audio()) + + assert request_state.sampling_params_list[0].max_tokens == 7 + assert clients[0].add_request_calls[-1][0].sampling_params.max_tokens == 7 + await orchestrator.session_manager.shutdown() + + +@pytest.mark.asyncio +async def test_forwarded_stage_requests_are_bound_and_barge_in_aborts_them() -> None: + orchestrator, clients, rpc_q, _ = _build(stages=2) + await _open(orchestrator, rpc_q) + request_id = _stage0_request_id() + await _submit(orchestrator, _append_audio()) + request_state = orchestrator.request_states[request_id] + session = orchestrator.session_manager.get(SESSION_ID) + assert session is not None + + # The base forwards Stage0 text to the TTS stage and reports the submission. + forwarded = SimpleNamespace(request_id=request_id, prompt_token_ids=[1, 2], resumable=True) + replica_id = await orchestrator.stage_pools[1].submit_initial(request_id, request_state, forwarded) + orchestrator._on_stage_submitted(1, request_id, replica_id, request_state) + assert request_state.stage_fences[1] == DuplexFence(SESSION_ID) + assert session.stage_request_submitted(1, request_id) + + await _submit(orchestrator, commands.BargeIn()) + assert session.epoch == 1 + assert clients[0].abort_calls == [[request_id]] + assert clients[1].abort_calls == [[request_id]] + assert request_id not in orchestrator.request_states + assert session.resource_request_ids() == [] + await orchestrator.session_manager.shutdown() + + +@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) + 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() + + consumed = await orchestrator._intercept_stage_output(1, 0, _tts_output(request_id), request_state, None, None) + await _settle(orchestrator) + + 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())]] + assert "response.created" in types and "response.output_audio.delta" in types + + orphan = DuplexOrchestratorRequestState( + request_id="duplex-s.b3RoZXI.e.0.r.stage0", + prompt=None, + sampling_params_list=[], + final_stage_id=1, + session_owned=True, + session_id="other", + fence=DuplexFence("other"), + ) + assert await orchestrator._intercept_stage_output(1, 0, _tts_output(orphan.request_id), orphan, None, None) + await orchestrator.session_manager.shutdown() + + +@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) + request_id = _stage0_request_id() + await _submit(orchestrator, _append_audio()) + request_state = orchestrator.request_states[request_id] + session = orchestrator.session_manager.get(SESSION_ID) + assert session is not None + + absorbed = await orchestrator._handle_forward_failure(request_id, 1, request_state, ValueError("bad thinker")) + await _settle(orchestrator) + + assert absorbed is True + assert request_id not in orchestrator.request_states + assert clients[0].abort_calls == [[request_id]] + 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())]] + # The stage failure is reported before the session expires, never after. + assert types.index("error") < types.index("session.expired") + assert types[-1] == "session.expired" + + +@pytest.mark.asyncio +async def test_request_cleanup_failure_is_deferred_and_retried_by_the_reaper(monkeypatch) -> None: + orchestrator, clients, rpc_q, _ = _build() + await _open(orchestrator, rpc_q) + request_id = _stage0_request_id() + await _submit(orchestrator, _append_audio()) + manager = orchestrator.session_manager + + async def failing_abort(request_ids: list[str]) -> None: + raise RuntimeError("stage abort failed") + + monkeypatch.setattr(clients[0], "abort_requests_async", failing_abort) + with pytest.raises(RuntimeError, match="stage abort failed"): + await orchestrator._cleanup_request_ids([request_id], abort=True, release_owners=True) + await _settle(orchestrator) + + assert SESSION_ID not in manager.runners + assert SESSION_ID in manager._closing # admission slot retained + assert manager._pending_request_cleanups and not manager._request_cleanups_in_progress + assert manager.active_count() == 0 and manager._admission_count() == 1 + + monkeypatch.undo() + assert await manager.reap_expired() >= 1 + assert manager._pending_request_cleanups == {} + assert manager._admission_count() == 0 + assert request_id not in orchestrator.request_states + + +@pytest.mark.asyncio +async def test_reaper_runs_as_a_background_task_and_survives_one_failure(monkeypatch) -> None: + orchestrator, _, _, _ = _build(runtime_config=DuplexSessionRuntimeConfig(reaper_interval_s=0.01)) + calls = {"count": 0} + + async def reap_expired(now: float | None = None) -> int: + calls["count"] += 1 + if calls["count"] == 1: + raise RuntimeError("transient cleanup failure") + return 0 + + monkeypatch.setattr(orchestrator.session_manager, "reap_expired", reap_expired) + tasks = orchestrator._background_tasks() + assert len(tasks) == 1 + task = asyncio.create_task(tasks[0]) + await asyncio.sleep(0.05) + orchestrator._shutdown_event.set() + await task + assert calls["count"] >= 2 + await orchestrator._shutdown_extensions() + + +# --------------------------------------------------------------------------- # +# The hierarchy the two surfaces depend on # +# --------------------------------------------------------------------------- # + + +def test_the_duplex_stack_extends_the_turn_based_one() -> None: + """One engine serves a session and an ordinary request, so duplex must extend turn-based. + + Pinned because it is the kind of relationship a later tidy-up turns back + into siblings over a shared base. If it does, a duplex server silently + stops serving /v1/chat/completions: the chat service needs a real + EngineClient with generate(), which only the turn-based side provides. + """ + from vllm_omni.engine.async_omni_engine import AsyncOmniEngine + from vllm_omni.engine.duplex_omni_engine import DuplexOmniEngine + from vllm_omni.engine.orchestrator import Orchestrator + from vllm_omni.entrypoints.async_omni import AsyncOmni + from vllm_omni.entrypoints.duplex_omni import DuplexOmni + + assert issubclass(DuplexOrchestrator, Orchestrator) + assert issubclass(DuplexOmniEngine, AsyncOmniEngine) + assert issubclass(DuplexOmni, AsyncOmni) + # What the chat route actually needs off the engine client. + assert hasattr(DuplexOmni, "generate") + + +@pytest.mark.asyncio +async def test_an_unrecognised_message_reaches_the_turn_based_handler() -> None: + """The duplex orchestrator consumes session messages and passes the rest down. + + Without the fall-through, an ordinary add_request on a duplex engine is + silently dropped and the chat request that sent it hangs forever. + """ + orchestrator, _clients, _rpc_q, _output_q = _build() + handled: list[str] = [] + + async def _turn_based(msg: object) -> None: + handled.append(getattr(msg, "type", "")) + + orchestrator._handle_add_request = _turn_based # type: ignore[method-assign] + msg = SimpleNamespace(type="add_request") + + assert await orchestrator._dispatch_message(msg) is True + assert handled == ["add_request"], "an add_request must reach Orchestrator, not be dropped" + + +@pytest.mark.asyncio +async def test_a_dead_replica_closes_the_sessions_it_was_serving() -> None: + """A session outlives its stage request, so replica death has to release the owner. + + ``_handle_dead_replica`` cleaned the stage request but not the runner that + owned it: the session stayed alive holding an admission slot, and the client + was told nothing — the request-scoped ``ErrorMessage`` it emits has no + frontend ``request_states`` entry to land in for a session-owned request, so + the frontend drops it. Heartbeats would then keep an unusable session + occupying capacity. + """ + orchestrator, clients, rpc_q, output_q = _build() + await _open(orchestrator, rpc_q) + request_id = _stage0_request_id() + await _submit(orchestrator, _append_audio()) + session = orchestrator.session_manager.get(SESSION_ID) + assert session is not None + assert orchestrator.session_manager.active_count() == 1 + + await orchestrator._handle_dead_replica(0, 0, EngineDeadError("stage 0 replica 0 died")) + await _settle(orchestrator) + + assert request_id not in orchestrator.request_states + assert SESSION_ID not in orchestrator.session_manager.runners + assert session.state == DuplexSessionState.CLOSED + assert orchestrator.session_manager.active_count() == 0, "the admission slot must come back" + messages = [output_q.get_nowait() for _ in range(output_q.qsize())] + types = [getattr(getattr(m, "event", None), "type", type(m).__name__) for m in messages] + assert types[-1] in {"session.expired", "session.closed"}, f"the client needs a terminal event, got {types}" diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index aa4f25ded..beddd2106 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -21,20 +21,6 @@ from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs, FinishReason from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats -from vllm_omni.engine.duplex.control_plane import DuplexControlPlane -from vllm_omni.engine.duplex.messages import ( - AppendDuplexInputMessage, - CloseDuplexSessionMessage, - DuplexFence, - OpenDuplexSessionMessage, - SignalDuplexTurnMessage, -) -from vllm_omni.engine.duplex.runtime import ( - DuplexInputMode, - DuplexRuntimeCapabilities, - DuplexSessionRuntimeState, - duplex_resource_request_id, -) from vllm_omni.engine.messages import ( AbortRequestMessage, AbortResultMessage, @@ -54,7 +40,6 @@ from vllm_omni.engine.orchestrator import ( ) from vllm_omni.engine.stage_pool import StagePool from vllm_omni.inputs.data import OmniDiffusionSamplingParams -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.runtime import MiniCPMO45DuplexRuntimeExtension from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -1118,778 +1103,6 @@ async def test_handle_abort_emits_error_result_on_failure() -> None: assert "cleanup exploded" in (result.error or "") -def _duplex_open_message( - session_id: str, - *, - incarnation: int = 0, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, -) -> OpenDuplexSessionMessage: - return OpenDuplexSessionMessage( - control_id=f"open-{session_id}", - fence=DuplexFence(session_id, incarnation=incarnation), - session_id=session_id, - capabilities={ - "input_modes": [DuplexInputMode.APPEND_AUDIO_CHUNK.value], - "implementation_level": "model_native_duplex", - }, - session_config=session_config or {}, - runtime_config=runtime_config or {}, - ) - - -async def _handle_duplex(orchestrator: Orchestrator, message: object) -> None: - await orchestrator._require_duplex_control_plane().handle(message) - - -def _duplex_request_state( - orchestrator: Orchestrator, - session: DuplexSessionRuntimeState, - *, - stage_id: int, -) -> OrchestratorRequestState | None: - context = orchestrator._require_duplex_control_plane().ensure_stage_request( - session, - stage_id=stage_id, - ) - if context is None: - return None - return orchestrator.request_states.get(context.request_id) - - -@pytest.mark.asyncio -async def test_duplex_control_plane_keeps_public_and_runtime_config_separate() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=_build_stage_pools([[stage0]]), - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - ) - open_message = _duplex_open_message( - "sid-config-channels", - incarnation=3, - session_config={ - "instructions": "public instructions", - "extra_body": {"duplex_stage_max_tokens": {"0": 99}}, - }, - runtime_config={ - "duplex_stage_max_tokens": {"0": 3}, - "duplex_stage_sampling_params": {"0": {"stop_token_ids": [151705]}}, - }, - ) - - await _handle_duplex(orchestrator, open_message) - - assert rpc_q.get_nowait().ok is True - session = orchestrator.duplex_sessions.require(open_message.session_id) - request_state = _duplex_request_state(orchestrator, session, stage_id=0) - assert request_state.sampling_params_list[0].max_tokens == 3 - assert request_state.sampling_params_list[0].stop_token_ids == [151705] - bridge = request_state.streaming.bridge_states["duplex"] - assert bridge["incarnation"] == open_message.fence.incarnation - assert bridge["session_config"] == open_message.session_config - assert bridge["runtime_config"] == open_message.runtime_config - - -@pytest.mark.asyncio -async def test_duplex_control_plane_preserves_turn_commit_without_model_extension() -> None: - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=[], - duplex_runtime_extension=None, - enable_duplex_control=True, - ) - assert isinstance(orchestrator.duplex_control_plane, DuplexControlPlane) - - message = OpenDuplexSessionMessage( - control_id="open-turn-commit", - fence=DuplexFence("sid-turn-commit"), - session_id="sid-turn-commit", - capabilities={"input_modes": [DuplexInputMode.TURN_COMMIT_ONLY.value]}, - session_config={}, - ) - await _handle_duplex(orchestrator, message) - - result = rpc_q.get_nowait() - assert result.ok is True - assert result.stage_results[0]["result"]["scheduler_request_context"] is False - - -def test_ordinary_orchestrator_bypasses_duplex_control_plane() -> None: - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=asyncio.Queue(), - stage_pools=[], - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - ) - - assert orchestrator.duplex_control_plane is None - - -@pytest.mark.asyncio -async def test_duplex_close_cleans_preregistered_request_without_append() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - rpc_q: asyncio.Queue = asyncio.Queue() - running_counter = FakeRunningCounter() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=_build_stage_pools([[stage0]]), - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - running_counter=running_counter, - ) - open_message = _duplex_open_message("sid-preregister-close") - - await _handle_duplex(orchestrator, open_message) - open_result = rpc_q.get_nowait() - request_id = open_result.stage_results[0]["result"]["request_id"] - assert request_id in orchestrator.request_states - - await _handle_duplex( - orchestrator, - CloseDuplexSessionMessage( - control_id="close-preregistered", - fence=open_message.fence, - session_id=open_message.session_id, - ), - ) - - assert rpc_q.get_nowait().ok is True - assert request_id not in orchestrator.request_states - assert orchestrator.duplex_sessions.get(open_message.session_id) is None - assert running_counter.value == 0 - - -@pytest.mark.asyncio -async def test_duplex_open_failure_rolls_back_session_and_reserved_request() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=_build_stage_pools([[stage0]]), - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - ) - open_message = _duplex_open_message( - "sid-open-rollback", - runtime_config={ - "duplex_stage_sampling_params": { - "0": {"stop_token_ids": 123}, - } - }, - ) - - await _handle_duplex(orchestrator, open_message) - - assert rpc_q.get_nowait().ok is False - assert orchestrator.duplex_sessions.get(open_message.session_id) is None - assert orchestrator.request_states == {} - - -@pytest.mark.asyncio -async def test_duplex_running_counter_tracks_only_submitted_request() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - rpc_q: asyncio.Queue = asyncio.Queue() - running_counter = FakeRunningCounter() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=_build_stage_pools([[stage0]]), - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - running_counter=running_counter, - ) - open_message = _duplex_open_message("sid-duplex-counter") - await _handle_duplex(orchestrator, open_message) - assert rpc_q.get_nowait().ok is True - assert running_counter.value == 0 - - await _handle_duplex( - orchestrator, - AppendDuplexInputMessage( - control_id="append-duplex-counter", - fence=open_message.fence, - session_id=open_message.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"is_speech": True}, - ), - ) - assert rpc_q.get_nowait().ok is True - assert running_counter.value == 1 - - await _handle_duplex( - orchestrator, - CloseDuplexSessionMessage( - control_id="close-duplex-counter", - fence=open_message.fence, - session_id=open_message.session_id, - ), - ) - assert rpc_q.get_nowait().ok is True - assert running_counter.value == 0 - - -@pytest.mark.asyncio -async def test_duplex_failed_append_does_not_advance_sequence_or_fence() -> None: - class FailingPlanExtension(MiniCPMO45DuplexRuntimeExtension): - def plan_append(self, **kwargs): - raise RuntimeError("planned append failure") - - stage0 = FakeStageClient(stage_type="llm", final_output=True) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=_build_stage_pools([[stage0]]), - duplex_runtime_extension=FailingPlanExtension(), - enable_duplex_control=True, - ) - fence = DuplexFence("sid-append-rollback", turn_id=1) - session = orchestrator.duplex_sessions.open_session( - DuplexFence("sid-append-rollback"), - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - - await _handle_duplex( - orchestrator, - AppendDuplexInputMessage( - control_id="append-rollback", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"is_speech": True}, - ), - ) - - result = rpc_q.get_nowait() - assert result.ok is False - assert session.fence == DuplexFence("sid-append-rollback") - assert session.input_seq == 0 - assert session.input_turn_seq == 0 - assert stage0.add_request_calls == [] - - -@pytest.mark.asyncio -async def test_duplex_duplicate_append_operation_is_submitted_once() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=_build_stage_pools([[stage0]]), - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - ) - fence = DuplexFence("sid-idempotent-append") - session = orchestrator.duplex_sessions.open_session( - fence, - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - - for control_id in ("append-first", "append-timeout-retry"): - await _handle_duplex( - orchestrator, - AppendDuplexInputMessage( - control_id=control_id, - operation_id="physical-input-1", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"is_speech": True}, - ), - ) - - first = rpc_q.get_nowait() - retry = rpc_q.get_nowait() - assert first.ok is True - assert retry.ok is True - assert first.stage_results == retry.stage_results - assert len(stage0.add_request_calls) == 1 - assert session.input_seq == 1 - - -@pytest.mark.asyncio -async def test_duplex_barge_in_aborts_bound_stage_requests_before_releasing_fence() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=False) - stage1 = FakeStageClient(stage_type="llm", final_output=True) - stage_pools = _build_stage_pools( - [[stage0], [stage1]], - output_processors=[FakeOutputProcessor(), FakeOutputProcessor()], - stage_vllm_configs=[ - SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), - SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), - ], - ) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=stage_pools, - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - ) - session = orchestrator.duplex_sessions.open_session( - DuplexFence("sid-stage-signal"), - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - session.bind_stage_request(0, "req-stage0", fence=session.fence) - session.bind_stage_request(1, "req-stage1", fence=session.fence) - assert stage_pools[0].select_replica_id("req-stage0") == 0 - assert stage_pools[1].select_replica_id("req-stage1") == 0 - session.append_input( - mode=DuplexInputMode.APPEND_AUDIO_CHUNK, - fence=session.fence, - ) - await _handle_duplex( - orchestrator, - SignalDuplexTurnMessage( - control_id="ctrl-signal", - fence=session.fence, - next_fence=DuplexFence("sid-stage-signal", epoch=1), - session_id="sid-stage-signal", - event="barge_in", - ), - ) - - assert stage0.abort_calls == [["req-stage0"]] - assert stage1.abort_calls == [["req-stage1"]] - assert stage0.collective_rpc_calls == [] - assert stage1.collective_rpc_calls == [] - assert session.stage_bindings == {} - assert session.fence == DuplexFence("sid-stage-signal", epoch=1) - result = rpc_q.get_nowait() - assert result.ok is True - - -@pytest.mark.asyncio -async def test_duplex_late_barge_in_releases_only_cancelled_fence_bindings() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=False) - stage1 = FakeStageClient(stage_type="llm", final_output=True) - stage_pools = _build_stage_pools( - [[stage0], [stage1]], - output_processors=[FakeOutputProcessor(), FakeOutputProcessor()], - stage_vllm_configs=[ - SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), - SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), - ], - ) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=stage_pools, - enable_duplex_control=True, - ) - cancelled_fence = DuplexFence("sid-late-stage-signal") - session = orchestrator.duplex_sessions.open_session( - cancelled_fence, - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - session.bind_stage_request(0, "req-stage0", fence=cancelled_fence) - session.bind_stage_request(1, "req-stage1", fence=cancelled_fence) - assert stage_pools[0].select_replica_id("req-stage0") == 0 - assert stage_pools[1].select_replica_id("req-stage1") == 0 - current_fence = DuplexFence("sid-late-stage-signal", epoch=1) - session.accept_fence(current_fence) - - await _handle_duplex( - orchestrator, - SignalDuplexTurnMessage( - control_id="ctrl-late-signal", - fence=cancelled_fence, - next_fence=current_fence, - session_id="sid-late-stage-signal", - event="barge_in", - ), - ) - - assert stage0.abort_calls == [["req-stage0"]] - assert stage1.abort_calls == [["req-stage1"]] - assert session.stage_bindings == {} - assert session.fence == current_fence - result = rpc_q.get_nowait() - assert result.ok is True - - -@pytest.mark.asyncio -async def test_duplex_cancel_without_next_fence_is_rejected_without_releasing_bindings() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - stage_pools = _build_stage_pools( - [[stage0]], - output_processors=[FakeOutputProcessor()], - stage_vllm_configs=[SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))], - ) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=stage_pools, - enable_duplex_control=True, - ) - fence = DuplexFence("sid-cancel-contract") - session = orchestrator.duplex_sessions.open_session(fence) - session.bind_stage_request(0, "req-live", fence=fence) - - await _handle_duplex( - orchestrator, - SignalDuplexTurnMessage( - control_id="cancel-without-next", - fence=fence, - session_id=fence.session_id, - event="input.cancel", - ), - ) - - result = rpc_q.get_nowait() - assert result.ok is False - assert result.error_count == 1 - assert "next_fence" in result.stage_results[0]["result"]["error"] - assert session.fence == fence - assert session.stage_request_ids() == ["req-live"] - assert stage0.abort_calls == [] - - -@pytest.mark.asyncio -async def test_duplex_session_update_replaces_runtime_config() -> None: - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=[], - enable_duplex_control=True, - ) - fence = DuplexFence("sid-update-config") - session = orchestrator.duplex_sessions.open_session( - fence, - session_config={"temperature": 0.7}, - ) - - await _handle_duplex( - orchestrator, - SignalDuplexTurnMessage( - control_id="update-config", - fence=fence, - session_id=fence.session_id, - event="session.update", - session_config={"temperature": 0.0, "instructions": "updated"}, - ), - ) - - result = rpc_q.get_nowait() - assert result.ok is True - assert session.session_config == {"temperature": 0.0, "instructions": "updated"} - - -@pytest.mark.asyncio -async def test_duplex_session_update_refreshes_next_append_sampling_params() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - stage_pools = _build_stage_pools( - [[stage0]], - output_processors=[FakeOutputProcessor()], - stage_vllm_configs=[SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))], - ) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=stage_pools, - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - ) - fence = DuplexFence("sid-update-policy") - session = orchestrator.duplex_sessions.open_session( - fence, - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - runtime_config={ - "duplex_stage_max_tokens": {"0": 2}, - "duplex_stage_sampling_params": {"0": {"stop_token_ids": [151645]}}, - }, - ) - req_state = _duplex_request_state(orchestrator, session, stage_id=0) - assert req_state is not None - assert req_state.sampling_params_list[0].max_tokens == 2 - - await _handle_duplex( - orchestrator, - SignalDuplexTurnMessage( - control_id="update-policy", - fence=fence, - session_id=fence.session_id, - event="session.update", - runtime_config={ - "duplex_stage_max_tokens": {"0": 7}, - "duplex_stage_sampling_params": {"0": {"stop_token_ids": [151705]}}, - }, - ), - ) - assert rpc_q.get_nowait().ok is True - - await _handle_duplex( - orchestrator, - AppendDuplexInputMessage( - control_id="append-updated-policy", - fence=fence, - session_id=fence.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"is_speech": True}, - ), - ) - - assert rpc_q.get_nowait().ok is True - assert req_state.sampling_params_list[0].max_tokens == 7 - assert req_state.sampling_params_list[0].stop_token_ids == [151705] - submitted_request = stage0.add_request_calls[0][0] - assert submitted_request.sampling_params.stop_token_ids == [151705] - assert submitted_request.sampling_params.max_tokens == 7 - - -@pytest.mark.asyncio -async def test_duplex_invalid_session_update_preserves_previous_config() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=_build_stage_pools([[stage0]]), - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - ) - fence = DuplexFence("sid-invalid-update") - original_runtime_config = { - "duplex_stage_sampling_params": {"0": {"stop_token_ids": [151645]}}, - } - session = orchestrator.duplex_sessions.open_session( - fence, - runtime_config=original_runtime_config, - ) - - await _handle_duplex( - orchestrator, - SignalDuplexTurnMessage( - control_id="invalid-update", - fence=fence, - session_id=fence.session_id, - event="session.update", - runtime_config={ - "duplex_stage_sampling_params": {"0": {"stop_token_ids": 123}}, - }, - ), - ) - - result = rpc_q.get_nowait() - assert result.ok is False - assert session.runtime_config == original_runtime_config - assert session.config_generation == 0 - - -@pytest.mark.asyncio -async def test_duplex_arbitrary_non_cancel_signal_is_rejected() -> None: - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=[], - enable_duplex_control=True, - ) - fence = DuplexFence("sid-unsupported-signal") - orchestrator.duplex_sessions.open_session(fence) - - await _handle_duplex( - orchestrator, - SignalDuplexTurnMessage( - control_id="unsupported-signal", - fence=fence, - session_id=fence.session_id, - event="turn.end", - ), - ) - - result = rpc_q.get_nowait() - assert result.ok is False - assert "unsupported duplex runtime signal" in result.stage_results[0]["result"]["error"] - - -@pytest.mark.asyncio -async def test_duplex_cancel_rejects_late_old_append_and_accepts_next_epoch() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - stage_pools = _build_stage_pools( - [[stage0]], - output_processors=[FakeOutputProcessor()], - stage_vllm_configs=[SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))], - ) - rpc_q: asyncio.Queue = asyncio.Queue() - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=rpc_q, - stage_pools=stage_pools, - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - ) - cancelled_fence = DuplexFence("sid-cancel-late-append") - next_fence = DuplexFence("sid-cancel-late-append", epoch=1) - session = orchestrator.duplex_sessions.open_session( - cancelled_fence, - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - ) - - await _handle_duplex( - orchestrator, - SignalDuplexTurnMessage( - control_id="cancel-old-fence", - fence=cancelled_fence, - next_fence=next_fence, - session_id=session.session_id, - event="input.cancel", - ), - ) - - cancel_result = rpc_q.get_nowait() - assert cancel_result.ok is True - assert session.fence == next_fence - - await _handle_duplex( - orchestrator, - AppendDuplexInputMessage( - control_id="late-old-append", - fence=cancelled_fence, - session_id=session.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"is_speech": True}, - ), - ) - - stale_result = rpc_q.get_nowait() - assert stale_result.ok is False - assert stage0.add_request_calls == [] - assert session.stage_bindings == {} - - await _handle_duplex( - orchestrator, - AppendDuplexInputMessage( - control_id="next-epoch-append", - fence=next_fence, - session_id=session.session_id, - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"is_speech": True}, - ), - ) - - next_result = rpc_q.get_nowait() - assert next_result.ok is True - assert len(stage0.add_request_calls) == 1 - assert session.stage_bindings[0].fence == next_fence - - -@pytest.mark.asyncio -async def test_duplex_append_updates_bridge_turn_id_on_long_lived_stage0_request() -> None: - stage0 = FakeStageClient(stage_type="llm", final_output=True) - stage_pools = _build_stage_pools( - [[stage0]], - output_processors=[FakeOutputProcessor()], - stage_vllm_configs=[SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))], - ) - orchestrator = Orchestrator( - request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), - rpc_async_queue=asyncio.Queue(), - stage_pools=stage_pools, - duplex_runtime_extension=MiniCPMO45DuplexRuntimeExtension(), - enable_duplex_control=True, - ) - session = orchestrator.duplex_sessions.open_session( - DuplexFence("sid-bridge-turn"), - capabilities=DuplexRuntimeCapabilities( - input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, - ), - session_config={ - "voice": "test", - }, - runtime_config={ - "duplex_stage_sampling_params": { - "0": {"stop_token_ids": [151645]}, - }, - }, - ) - - await _handle_duplex( - orchestrator, - AppendDuplexInputMessage( - control_id="append-1", - fence=DuplexFence("sid-bridge-turn"), - session_id="sid-bridge-turn", - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"is_speech": True}, - ), - ) - req_state1 = _duplex_request_state(orchestrator, session, stage_id=0) - assert req_state1 is not None - duplex_state1 = req_state1.streaming.bridge_states["duplex"] - assert duplex_state1["session_id"] == "sid-bridge-turn" - assert duplex_state1["epoch"] == 0 - assert duplex_state1["turn_id"] == 0 - assert duplex_state1["session_config"] == session.session_config - assert req_state1.sampling_params_list[0].stop_token_ids == [151645] - submitted_request = stage0.add_request_calls[0][0] - assert submitted_request.sampling_params.stop_token_ids == [151645] - assert submitted_request.sampling_params.max_tokens == 1 - - await _handle_duplex( - orchestrator, - AppendDuplexInputMessage( - control_id="append-2", - fence=DuplexFence("sid-bridge-turn", turn_id=1, response_seq=1), - session_id="sid-bridge-turn", - mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, - payload={"is_speech": True, "new_user_turn": True}, - ), - ) - req_state2 = _duplex_request_state(orchestrator, session, stage_id=0) - assert req_state2 is req_state1 - duplex_state2 = req_state2.streaming.bridge_states["duplex"] - assert duplex_state2["turn_id"] == 1 - assert duplex_state2["epoch"] == 0 - expected_request_id = duplex_resource_request_id(DuplexFence("sid-bridge-turn"), "stage0") - assert [call[0].request_id for call in stage0.add_request_calls] == [ - expected_request_id, - expected_request_id, - ] - - # --------------------------------------------------------------------------- # Multi-replica tests # --------------------------------------------------------------------------- @@ -2641,61 +1854,6 @@ async def test_stage_pool_process_llm_raw_outputs_mutates_iteration_stats() -> N assert iteration_stats.num_generation_tokens == 3 -@pytest.mark.asyncio -async def test_duplex_reaper_loop_waits_between_ticks(): - class _Plane: - def __init__(self) -> None: - self.calls = 0 - - async def reap_expired(self) -> int: - self.calls += 1 - return 0 - - orchestrator = object.__new__(Orchestrator) - orchestrator.duplex_control_plane = _Plane() - orchestrator._duplex_reaper_interval_s = 0.01 - orchestrator._shutdown_event = asyncio.Event() - - task = asyncio.create_task(orchestrator._duplex_reaper_loop()) - await asyncio.sleep(0.035) - orchestrator._shutdown_event.set() - await task - - assert 2 <= orchestrator.duplex_control_plane.calls <= 5 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("first_cleanup_delay", [0.0, 0.05], ids=["immediate", "delayed"]) -async def test_duplex_reaper_loop_survives_one_cleanup_failure(first_cleanup_delay: float) -> None: - recovered = asyncio.Event() - - class _Plane: - def __init__(self) -> None: - self.calls = 0 - - async def reap_expired(self) -> int: - self.calls += 1 - if self.calls == 1: - await asyncio.sleep(first_cleanup_delay) - raise RuntimeError("transient cleanup failure") - recovered.set() - return 0 - - orchestrator = object.__new__(Orchestrator) - orchestrator.duplex_control_plane = _Plane() - orchestrator._duplex_reaper_interval_s = 0.01 - orchestrator._shutdown_event = asyncio.Event() - - task = asyncio.create_task(orchestrator._duplex_reaper_loop()) - try: - # Wait for recovery itself, including when cleanup exceeds the old 35 ms window. - await asyncio.wait_for(recovered.wait(), timeout=5.0) - assert orchestrator.duplex_control_plane.calls >= 2 - finally: - orchestrator._shutdown_event.set() - await asyncio.wait_for(task, timeout=5.0) - - @pytest.mark.asyncio async def test_abort_retry_does_not_repeat_successful_stage_abort(): class _Pool: @@ -2733,48 +1891,3 @@ async def test_abort_retry_does_not_repeat_successful_stage_abort(): assert first.physical_abort_calls == 1 assert second.physical_abort_calls == 2 assert second.bound == set() - - -@pytest.mark.asyncio -async def test_request_cleanup_failure_is_deferred_to_control_plane(): - class _FailingPool: - async def abort_requests(self, request_ids: list[str]) -> None: - raise RuntimeError("stage abort failed") - - def release_bindings(self, request_ids: list[str]) -> None: - pass - - class _Plane: - def __init__(self) -> None: - self.deferred: list[str] = [] - self.finalized: list[str] = [] - - def close_sessions_for_request_ids(self, request_ids: list[str], **kwargs): - assert kwargs == {"abort": True, "cleanup_in_progress": True} - return {"sid-cleanup": ["req-a"]} - - def defer_request_cleanups(self, session_ids: list[str]) -> None: - self.deferred.extend(session_ids) - - def finalize_closed_sessions(self, session_ids: list[str]) -> None: - self.finalized.extend(session_ids) - - from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker - - orchestrator = object.__new__(Orchestrator) - orchestrator.stage_pools = [_FailingPool()] - orchestrator.duplex_control_plane = _Plane() - orchestrator._pd_kv_params = {} - orchestrator.request_states = {} - orchestrator._running_counter = None - orchestrator._cfg_tracker = CfgCompanionTracker() - - with pytest.raises(RuntimeError, match="stage abort failed"): - await orchestrator._cleanup_request_ids( - ["req-a"], - abort=True, - close_duplex_sessions=True, - ) - - assert orchestrator.duplex_control_plane.deferred == ["sid-cleanup"] - assert orchestrator.duplex_control_plane.finalized == [] diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index eaa246072..b8949786e 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -1054,7 +1054,7 @@ async def test_streaming_input_processor_client_error_does_not_forward_terminal_ final_stage_id=1, ) state.streaming.enabled = True - assert state.duplex_identity is None + assert state.session_owned is False orchestrator.request_states[state.request_id] = state try: await orchestrator._route_output(0, 0, _build_request_output(state.request_id), state, stage_metrics=None) @@ -1074,34 +1074,14 @@ async def test_streaming_input_processor_client_error_does_not_forward_terminal_ q.close() -@pytest.mark.asyncio -async def test_duplex_input_processor_failure_is_request_scoped(orchestrator_factory, monkeypatch) -> None: - class FailingStage(FakeStageClient): - def process_engine_inputs(self, *_args, **_kwargs): - raise ValueError("No latent or hidden_states found in thinker output") - - fixture = orchestrator_factory( - [FakeStageClient(stage_type="llm"), FailingStage(stage_type="llm", final_output=True)] - ) - state = OrchestratorRequestState( - request_id="bad", - prompt=SimpleNamespace(request_id="bad", prompt_token_ids=[1]), - sampling_params_list=[_sampling_params(), _sampling_params()], - final_stage_id=1, - duplex_identity=SimpleNamespace(), - ) - - async def no_cleanup(*_args, **_kwargs): - pass - - monkeypatch.setattr(fixture.orchestrator, "_cleanup_request_ids", no_cleanup) - try: - await fixture.orchestrator._forward_to_next_stage_unguarded("bad", 0, _build_request_output("raw"), state) - error = await _wait_for_error_message(fixture, request_id="bad") - assert error.fatal is False - finally: - fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) - fixture.thread.join(timeout=5) +# ``test_duplex_input_processor_failure_is_request_scoped`` lived here upstream, +# where one Orchestrator served both turn-based and duplex requests. Duplex +# request ownership now belongs to DuplexOrchestrator, and the base +# ``_handle_forward_failure`` deliberately absorbs nothing, so the equivalent +# assertion is +# ``tests/engine/test_duplex_orchestrator.py::test_forward_failure_closes_the_owning_session`` +# -- which also checks the stage abort, the session teardown and the +# error-before-session.expired ordering. @pytest.mark.asyncio diff --git a/tests/engine/test_orchestrator_segment_stage_keying.py b/tests/engine/test_orchestrator_segment_stage_keying.py index 5d3952e4b..9c5290245 100644 --- a/tests/engine/test_orchestrator_segment_stage_keying.py +++ b/tests/engine/test_orchestrator_segment_stage_keying.py @@ -2,10 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Streaming segment-boundary state must be keyed by stage. -``Orchestrator._orchestration_loop`` polls every stage into the same +``OrchestratorBase._orchestration_loop`` polls every stage into the same ``OrchestratorRequestState``. While the segment-boundary fields lived in a single flat slot per request, whichever stage polled most recently overwrote the others, and the per-stage consumers then read back a different stage's boundary. + +The duplex consumer of the per-stage segment (``DuplexOrchestrator``'s output +context) is covered in ``test_duplex_orchestrator.py``. """ from __future__ import annotations @@ -15,8 +18,6 @@ from dataclasses import dataclass, field import pytest from vllm.sampling_params import SamplingParams -from vllm_omni.engine.duplex.contracts import DuplexRequestIdentity -from vllm_omni.engine.duplex.messages import DuplexFence from vllm_omni.engine.orchestrator import ( Orchestrator, OrchestratorRequestState, @@ -54,22 +55,19 @@ class _RecordingPool: return _StageMetrics() -def _duplex_req_state(request_id: str = "req-duplex") -> OrchestratorRequestState: +def _streaming_req_state(request_id: str = "req-streaming") -> OrchestratorRequestState: req_state = OrchestratorRequestState( request_id=request_id, sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], final_stage_id=TALKER_STAGE, + session_owned=True, ) req_state.streaming.enabled = True - req_state.duplex_identity = DuplexRequestIdentity( - session_id="session-1", - fence=DuplexFence(session_id="session-1"), - ) return req_state -def test_duplex_output_context_reads_only_its_own_stage_segment() -> None: - req_state = _duplex_req_state() +def test_segment_accessor_reads_only_its_own_stage() -> None: + req_state = _streaming_req_state() req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState( finished=True, token_ids=[11, 22], @@ -78,26 +76,28 @@ def test_duplex_output_context_reads_only_its_own_stage_segment() -> None: # The talker polls after the thinker and is still mid-segment. req_state.streaming.segments[TALKER_STAGE] = StreamingSegmentState(finished=False) - thinker = Orchestrator._duplex_output_context(req_state, stage_id=THINKER_STAGE) - talker = Orchestrator._duplex_output_context(req_state, stage_id=TALKER_STAGE) - - assert thinker is not None - assert talker is not None + thinker = req_state.streaming.segment(THINKER_STAGE) + talker = req_state.streaming.segment(TALKER_STAGE) # These are the assertions that fail on the flat field: the talker is written # second, so the thinker loses its boundary to whichever stage polled last. - assert thinker.segment_finished is True - assert thinker.segment_token_ids == (11, 22) - assert thinker.segment_output_metadata == {"stage": "thinker"} + assert thinker.finished is True + assert thinker.token_ids == [11, 22] + assert thinker.output_metadata == {"stage": "thinker"} - assert talker.segment_finished is False - assert talker.segment_token_ids == () - assert talker.segment_output_metadata == {} + assert talker.finished is False + assert talker.token_ids == [] + assert talker.output_metadata == {} + + # A stage that never reported a boundary gets a fresh, unshared segment. + unreported = req_state.streaming.segment(2) + unreported.output_metadata["poisoned"] = True + assert req_state.streaming.segment(2).output_metadata == {} @pytest.mark.asyncio async def test_segment_boundary_on_one_stage_does_not_build_metrics_on_another() -> None: - req_state = _duplex_req_state() + req_state = _streaming_req_state() req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState(finished=True) talker_pool = _RecordingPool() diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 69df71231..842824c2d 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -7,23 +7,17 @@ import asyncio import queue from types import SimpleNamespace from typing import Any -from unittest.mock import ANY, AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock import janus import pytest from vllm.outputs import CompletionOutput, RequestOutput from vllm.sampling_params import SamplingParams -from vllm_omni.engine.duplex.contracts import ( - DuplexStageRequestContext, - DuplexStageSubmission, -) -from vllm_omni.engine.duplex.messages import DuplexFence from vllm_omni.engine.orchestrator import ( Orchestrator, OrchestratorRequestState, StreamingSegmentState, - _OrchestratorDuplexStagePort, ) from vllm_omni.engine.stage_pool import StagePool @@ -131,64 +125,6 @@ class FakePrewarmPool: return self.stage_client if self.get_bound_replica_id(request_id) is not None else None -def _duplex_stage_port_submission(): - stage_pools = [] - for stage_id in range(3): - pool = SimpleNamespace( - stage_client=SimpleNamespace(default_sampling_params=SamplingParams(max_tokens=1)), - stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), - submit_initial=AsyncMock(return_value=stage_id + 10), - submit_update=AsyncMock(return_value=stage_id + 20), - ) - stage_pools.append(pool) - request_states: dict[str, OrchestratorRequestState] = {} - prewarm = AsyncMock() - port = _OrchestratorDuplexStagePort( - stage_pools=stage_pools, - request_states=request_states, - running_counter=None, - cleanup_request_ids=AsyncMock(), - async_chunk=True, - prewarm_async_chunk_stages=prewarm, - ) - context = DuplexStageRequestContext( - request_id="req-duplex", - session_id="session-duplex", - fence=DuplexFence("session-duplex"), - stage_id=0, - final_stage_id=2, - config_generation=0, - sampling_params=tuple(SamplingParams(max_tokens=1) for _ in range(3)), - ) - port.ensure_request(context) - submission = DuplexStageSubmission( - context=context, - prompt={"prompt_token_ids": [1, 2]}, - already_submitted=False, - ) - return port, stage_pools, request_states, prewarm, submission - - -def test_duplex_bridge_state_catches_up_after_serving_turn_boundary() -> None: - port, _, request_states, _, submission = _duplex_stage_port_submission() - request_state = request_states[submission.context.request_id] - bridge_state = request_state.streaming.bridge_states["duplex"] - bridge_state["model_turn_id"] = 0 - - next_context = DuplexStageRequestContext( - request_id=submission.context.request_id, - session_id=submission.context.session_id, - fence=DuplexFence(submission.context.session_id, turn_id=1), - stage_id=submission.context.stage_id, - final_stage_id=submission.context.final_stage_id, - config_generation=submission.context.config_generation, - sampling_params=submission.context.sampling_params, - ) - port.ensure_request(next_context) - - assert bridge_state["model_turn_id"] == 1 - - def _request_output(request_id: str) -> RequestOutput: completion = CompletionOutput( index=0, @@ -279,13 +215,12 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N stage1.stage_client.get_payload_sender_info = MagicMock(return_value=payload_sender_info) orchestrator.stage_pools = [stage0, stage1, stage2] orchestrator._emit_tx_edge = lambda **_kwargs: None - orchestrator._record_duplex_stage_submission = MagicMock() + orchestrator._on_stage_submitted = MagicMock() req_state = OrchestratorRequestState( request_id="req-prewarm", prompt={"prompt_token_ids": [1, 2]}, sampling_params_list=[SamplingParams(max_tokens=1) for _ in range(3)], final_stage_id=2, - duplex_identity=SimpleNamespace(), ) prewarmed = await orchestrator._prewarm_async_chunk_stages( @@ -306,7 +241,7 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N assert stage2.submitted[0].resumable is True assert 1 not in req_state.stage_submit_ts assert 2 in req_state.stage_submit_ts - orchestrator._record_duplex_stage_submission.assert_called_once_with( + orchestrator._on_stage_submitted.assert_called_once_with( 2, "req-prewarm", 0, @@ -314,42 +249,6 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N ) -@pytest.mark.asyncio -async def test_duplex_prewarm_runs_after_first_stage0_submission() -> None: - port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission() - prewarm.return_value = True - - result = await port.submit(submission) - - assert result.stage_id == 0 - stage_pools[0].submit_initial.assert_awaited_once() - prewarm.assert_awaited_once_with("req-duplex", ANY, request_states["req-duplex"]) - - -@pytest.mark.asyncio -async def test_duplex_submit_bails_out_when_prewarm_failed_the_request() -> None: - """A failed prewarm has already aborted the request and popped its state. - - Returning a success result here would hand the control plane a replica for a - request that no longer exists, and the trailing bookkeeping would re-register - a running counter the cleanup just released -- a leak that never decrements. - """ - port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission() - counter = MagicMock() - port._running_counter = counter - prewarm.return_value = False - request_state = request_states["req-duplex"] - - with pytest.raises(RuntimeError, match="prewarm failed"): - await port.submit(submission) - - prewarm.assert_awaited_once() - assert request_state.duplex_stage_fences == {} - assert request_state.stage_submit_ts == {} - assert request_state.running_counter_registered is False - counter.increment.assert_not_called() - - @pytest.mark.asyncio async def test_async_route_forwards_to_outgoing_only_stage() -> None: orchestrator = object.__new__(Orchestrator) diff --git a/tests/engine/test_output_metadata_snapshots.py b/tests/engine/test_output_metadata_snapshots.py index 688f1252d..9670994c4 100644 --- a/tests/engine/test_output_metadata_snapshots.py +++ b/tests/engine/test_output_metadata_snapshots.py @@ -9,7 +9,7 @@ from vllm.sampling_params import RequestOutputKind from vllm.v1.engine import FinishReason from vllm_omni.outputs.mm_outputs import MultimodalPayload -from vllm_omni.outputs.output_modality import OutputModality, TensorAccumulationStrategy +from vllm_omni.outputs.output_modality import OutputModality from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -58,7 +58,7 @@ def test_rate_is_latest_snapshot_before_consolidation(key, representation): assert payload.get(key) is payload.to_dict()[key] payload = payload.merged_with(MultimodalPayload(tensors={"audio": torch.full((2,), 3)})) assert int(payload[key]) == 48000 # Missing metadata means retain, not clear. - payload.consolidate_tensors(TensorAccumulationStrategy.CONCAT_LAST) + payload.consolidate_tensors(OutputModality.AUDIO) payload.consolidate_metadata() torch.testing.assert_close(payload["audio"], torch.arange(4).repeat_interleave(2)) @@ -86,7 +86,7 @@ def test_self_merge_retains_snapshot_and_content(key): assert payload is not None payload = payload.merged_with(payload) assert int(payload[key]) == 24000 - payload.consolidate_tensors(TensorAccumulationStrategy.CONCAT_LAST) + payload.consolidate_tensors(OutputModality.AUDIO) torch.testing.assert_close(payload["audio"], torch.tensor([1.0, 2.0, 1.0, 2.0])) @@ -138,6 +138,6 @@ def test_non_snapshot_tensors_still_accumulate_and_empty_merge_keeps_identity(): assert first is not None and second is not None assert MultimodalPayload().merged_with(first) is first merged = first.merged_with(second) - merged.consolidate_tensors(TensorAccumulationStrategy.CONCAT_LAST) + merged.consolidate_tensors(OutputModality.AUDIO) torch.testing.assert_close(merged["latent"], torch.tensor([1.0, 3.0])) torch.testing.assert_close(merged["audio"], torch.tensor([2.0, 4.0])) diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 675a0650b..41d63af9e 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -389,7 +389,7 @@ class TestSingleStageModeDetection: return_value=None, ) mocker.patch( - "vllm_omni.engine.async_omni_engine.load_deploy_config", + "vllm_omni.engine.omni_engine_base.load_deploy_config", return_value=SimpleNamespace(duplex_session=DuplexSessionRuntimeConfig()), ) mocker.patch.object( @@ -446,11 +446,11 @@ class TestSingleStageModeDetection: def test_deploy_config_loads_duplex_runtime_config(self, mocker: MockerFixture): duplex_session = DuplexSessionRuntimeConfig(max_sessions=2) get_pipeline_config = mocker.patch( - "vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config", + "vllm_omni.engine.omni_engine_base.StageConfigFactory.get_pipeline_config", return_value=None, ) load_deploy_config = mocker.patch( - "vllm_omni.engine.async_omni_engine.load_deploy_config", + "vllm_omni.engine.omni_engine_base.load_deploy_config", return_value=SimpleNamespace(duplex_session=duplex_session), ) @@ -467,7 +467,8 @@ class TestSingleStageModeDetection: deploy_config_path="/fake/duplex.yaml", ) load_deploy_config.assert_called_once_with("/fake/duplex.yaml") - assert engine.duplex_session_config is duplex_session + # The turn-based engine only keeps the resolved deploy profile for introspection. + assert engine.deploy_config.duplex_session is duplex_session def test_auto_discovered_deploy_loads_duplex_runtime_config( self, @@ -476,11 +477,11 @@ class TestSingleStageModeDetection: deploy_path = "/resolved/qwen3_omni_moe.yaml" duplex_session = DuplexSessionRuntimeConfig(server_vad_model_path="/models/silero_vad.onnx") mocker.patch( - "vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config", + "vllm_omni.engine.omni_engine_base.StageConfigFactory.get_pipeline_config", return_value=None, ) load_deploy_config = mocker.patch( - "vllm_omni.engine.async_omni_engine.load_deploy_config", + "vllm_omni.engine.omni_engine_base.load_deploy_config", return_value=SimpleNamespace(duplex_session=duplex_session), ) @@ -557,7 +558,7 @@ class TestEndpointRestrictionsTrustRemoteCode: def _make_engine_no_thread(self, mocker: MockerFixture, **kwargs: Any) -> AsyncOmniEngine: mocker.patch( - "vllm_omni.engine.async_omni_engine.load_deploy_config", + "vllm_omni.engine.omni_engine_base.load_deploy_config", return_value=SimpleNamespace(duplex_session=DuplexSessionRuntimeConfig()), ) mocker.patch.object( diff --git a/tests/entrypoints/duplex/test_duplex_serving.py b/tests/entrypoints/duplex/test_duplex_serving.py new file mode 100644 index 000000000..dde616d94 --- /dev/null +++ b/tests/entrypoints/duplex/test_duplex_serving.py @@ -0,0 +1,609 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The thin websocket handler: handshake, command translation, event pump, resume/takeover.""" + +from __future__ import annotations + +import asyncio +import base64 +import json +from contextlib import suppress +from typing import Any + +import pytest +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.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 +from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler +from vllm_omni.entrypoints.duplex.websocket import MAX_EVENT_BYTES + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_DISCONNECT = object() + + +class FakeWebSocket: + def __init__(self, query: dict[str, str] | None = None) -> None: + self.query_params = dict(query or {}) + self.sent: list[dict[str, Any]] = [] + self.accepted = False + self.closed: list[tuple[int, str]] = [] + self._send_failure: str | None = None + self._inbound: asyncio.Queue[Any] = asyncio.Queue() + + async def accept(self) -> None: + self.accepted = True + + async def send_json(self, payload: dict[str, Any]) -> None: + if self.closed or self._send_failure is not None: + raise RuntimeError( + self._send_failure or "Unexpected ASGI message 'websocket.send', after sending 'websocket.close'." + ) + self.sent.append(json.loads(json.dumps(payload))) + + async def receive_text(self) -> str: + item = await self._inbound.get() + if item is _DISCONNECT: + raise WebSocketDisconnect(code=1000) + return item + + async def close(self, code: int = 1000, reason: str = "") -> None: + self.closed.append((code, reason)) + # A closed socket ends the reader exactly like a client disconnect. + self._inbound.put_nowait(_DISCONNECT) + + # ---- test helpers ---- + + def feed(self, payload: dict[str, Any] | str) -> None: + self._inbound.put_nowait(payload if isinstance(payload, str) else json.dumps(payload)) + + def disconnect(self) -> None: + self._inbound.put_nowait(_DISCONNECT) + + def break_sends(self) -> None: + """Kill the write half only: the reader stays parked, as it does in practice.""" + self._send_failure = "Unexpected ASGI message 'websocket.send', after sending 'websocket.close'." + + def types(self) -> list[str]: + return [payload["type"] for payload in self.sent] + + async def wait_for(self, wire_type: str, *, timeout_s: float = 2.0) -> dict[str, Any]: + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout_s + while loop.time() < deadline: + for payload in self.sent: + if payload["type"] == wire_type: + return payload + await asyncio.sleep(0.005) + raise AssertionError(f"no {wire_type!r} in {self.types()}") + + +class FakeHandle: + def __init__(self, session_id: str, capabilities: DuplexCapabilities, *, idle_timeout_s: float = 300) -> None: + self.session_id = session_id + self.capabilities = capabilities + self.public_session: dict[str, Any] = {"id": session_id, "idle_timeout_s": idle_timeout_s} + self.lease_generation = 0 + self.closed = False + self.close_reasons: list[str] = [] + self.commands: list[commands.DuplexCommand] = [] + self._outbox: asyncio.Queue[DuplexEvent] = asyncio.Queue() + + def deliver(self, event: DuplexEvent) -> None: + self._outbox.put_nowait(event) + + async def submit(self, command: commands.DuplexCommand) -> None: + if self.closed: + raise DuplexSessionError("closed", code="session_closed", session_id=self.session_id) + self.commands.append(command) + + async def events(self): + while True: + event = await self._outbox.get() + yield event + if event.is_terminal: + return + + async def close(self, *, reason: str = "client_close", timeout: float | None = None) -> None: + self.close_reasons.append(reason) + if self.closed: + return + self.closed = True + self.deliver(SessionClosed(session_id=self.session_id, reason=reason)) + + +class FakeOmni: + def __init__( + self, *, resumable: bool = True, replay_max_bytes: int = 64 * 1024, idle_timeout_s: float = 300 + ) -> None: + self.duplex_session_config = DuplexSessionRuntimeConfig( + resume_replay_ttl_s=60.0, resume_replay_max_bytes_per_session=replay_max_bytes + ) + self.capabilities = DuplexCapabilities(supports_session_resume=resumable) + self.idle_timeout_s = idle_timeout_s + self.opened: list[dict[str, Any]] = [] + self.handles: dict[str, FakeHandle] = {} + self.resumed: list[tuple[str, int]] = [] + self.detached: list[str] = [] + self.open_error: DuplexSessionError | None = None + + async def open_session(self, config: Any) -> FakeHandle: + self.opened.append(dict(config)) + if self.open_error is not None: + raise self.open_error + session_id = f"duplex-{len(self.handles) + 1:032x}" + handle = FakeHandle(session_id, self.capabilities, idle_timeout_s=self.idle_timeout_s) + self.handles[session_id] = handle + handle.deliver(SessionCreated(session_id=session_id, session={"id": session_id, "model": config.get("model")})) + return handle + + def get_session(self, session_id: str) -> FakeHandle | None: + return self.handles.get(session_id) + + async def resume_session(self, session_id: str, *, expected_lease_generation: int) -> FakeHandle: + self.resumed.append((session_id, expected_lease_generation)) + handle = self.handles[session_id] + handle.lease_generation = expected_lease_generation + 1 + return handle + + async def detach_session(self, session_id: str) -> None: + self.detached.append(session_id) + + +def _handler(omni: FakeOmni, **kwargs: Any) -> OmniDuplexSessionHandler: + kwargs.setdefault("config_timeout_s", 1.0) + kwargs.setdefault("idle_timeout_s", 5.0) + return OmniDuplexSessionHandler(duplex_omni=omni, **kwargs) + + +def _session_update(**session: Any) -> dict[str, Any]: + return {"type": "session.update", "session": {"model": "test-model", "modalities": ["audio", "text"], **session}} + + +async def _open( + handler: OmniDuplexSessionHandler, omni: FakeOmni, **session: Any +) -> tuple[FakeWebSocket, FakeHandle, asyncio.Task]: + ws = FakeWebSocket({"duplex": "1", "autostart": "0"}) + task = asyncio.create_task(handler.handle_realtime_session(ws)) + ws.feed(_session_update(**session)) + created = await ws.wait_for("session.created") + handle = omni.handles[created["session"]["id"]] + return ws, handle, task + + +# --------------------------------------------------------------------------- # +# Handshake and commands # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_session_update_opens_a_session_and_announces_the_server_id() -> None: + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni, session_id="mine", instructions="hi") + + assert ws.accepted + created = ws.sent[0] + assert created["type"] == "session.created" + assert created["session"]["id"] == handle.session_id != "mine" + assert created["attachment_generation"] == 1 + assert isinstance(created["resume_token"], str) and created["resume_token"] + assert "incarnation" not in created + assert "server_event_seq" not in created + # The whole session object went to open_session; the id inside it is ignored there. + assert omni.opened == [ + {"model": "test-model", "modalities": ["audio", "text"], "session_id": "mine", "instructions": "hi"} + ] + + ws.feed({"type": "input_audio_buffer.append", "audio": base64.b64encode(b"\x00\x01" * 8).decode("ascii")}) + ws.feed({"type": "input_audio_buffer.commit", "event_id": "evt-c"}) + await asyncio.sleep(0.05) + assert [type(command) for command in handle.commands] == [commands.AppendAudio, commands.Commit] + assert handle.commands[1].event_id == "evt-c" + + ws.disconnect() + await asyncio.wait_for(task, timeout=2.0) + # A resumable session is detached (engine-owned grace), not closed. + assert omni.detached == [handle.session_id] + assert handle.close_reasons == [] + + +@pytest.mark.asyncio +async def test_session_events_are_journaled_and_sent_in_order() -> None: + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + + handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="aGk=")) + delta = await ws.wait_for("response.output_audio.delta") + assert delta["delta"] == "aGk=" and delta["server_event_seq"] == 1 + + ws.feed({"type": "session.event_ack", "server_event_seq": 1}) + ws.feed({"type": "session.event_ack", "server_event_seq": 9, "event_id": "evt-ack"}) + error = await ws.wait_for("error") + assert error["error"]["code"] == "invalid_event_ack" and error["error"]["event_id"] == "evt-ack" + + await handle.close(reason="client_close") + closed = await ws.wait_for("session.closed") + assert closed["reason"] == "client_close" + await asyncio.wait_for(task, timeout=2.0) + assert ws.closed[0][1] == "client_close" + + +@pytest.mark.asyncio +async def test_envelope_errors_are_answered_locally_without_touching_the_session() -> None: + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + + ws.feed("not json") + ws.feed("x" * (MAX_EVENT_BYTES + 1)) + ws.feed({"type": "totally.unknown", "event_id": "evt-u"}) + ws.feed({"type": "session.resume", "event_id": "evt-r"}) + ws.feed({"type": "playback.ack", "event_id": "evt-p"}) # played_ms missing + await asyncio.sleep(0.1) + + codes = [payload["error"]["code"] for payload in ws.sent if payload["type"] == "error"] + assert codes == ["invalid_json", "event_too_large", "unknown_event", "unsupported_session_resume", "bad_event"] + errors = [payload for payload in ws.sent if payload["type"] == "error"] + assert errors[2]["error"]["event_id"] == "evt-u" + assert errors[4]["error"]["event_id"] == "evt-p" + assert handle.commands == [] + ws.disconnect() + await asyncio.wait_for(task, timeout=2.0) + + +@pytest.mark.asyncio +async def test_open_rejection_is_reported_and_the_socket_ends() -> None: + omni = FakeOmni() + omni.open_error = DuplexSessionError("no room", code="resource_exhausted", retryable=True) + handler = _handler(omni) + ws = FakeWebSocket({"duplex": "1"}) + task = asyncio.create_task(handler.handle_realtime_session(ws)) + ws.feed(_session_update()) + await asyncio.wait_for(task, timeout=2.0) + assert ws.types() == ["error"] + assert ws.sent[0]["error"]["code"] == "resource_exhausted" + + +@pytest.mark.asyncio +async def test_config_timeout_and_first_message_validation() -> None: + omni = FakeOmni() + handler = _handler(omni, config_timeout_s=0.05) + ws = FakeWebSocket({"duplex": "1"}) + await asyncio.wait_for(handler.handle_realtime_session(ws), timeout=2.0) + assert ws.sent[0]["error"]["code"] == "config_timeout" + + ws = FakeWebSocket({"duplex": "1"}) + task = asyncio.create_task(handler.handle_realtime_session(ws)) + ws.feed("{not json") + await asyncio.wait_for(task, timeout=2.0) + assert ws.sent[0]["error"]["code"] == "invalid_json" + assert omni.opened == [] + + +@pytest.mark.asyncio +async def test_idle_timeout_detaches_a_resumable_session_like_a_disconnect() -> None: + omni = FakeOmni(idle_timeout_s=0.05) + handler = _handler(omni, idle_timeout_s=0.05) + ws, handle, task = await _open(handler, omni) + + await asyncio.wait_for(task, timeout=2.0) + # Serving makes no session-lifetime decision: the engine lease decides + # whether a silent, detached session expires. + assert handle.close_reasons == [] + assert omni.detached == [handle.session_id] + + +@pytest.mark.asyncio +async def test_idle_timeout_closes_a_non_resumable_session() -> None: + omni = FakeOmni(idle_timeout_s=0.05, resumable=False) + handler = _handler(omni, idle_timeout_s=0.05) + ws, handle, task = await _open(handler, omni) + + await asyncio.wait_for(task, timeout=2.0) + assert handle.close_reasons == ["disconnect"] + assert omni.detached == [] + + +@pytest.mark.asyncio +async def test_transport_send_failure_detaches_the_session_instead_of_closing_it() -> None: + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + + # Only the write half dies, so the pump's send is the one and only report + # of the broken socket (the reader is still parked on receive). + ws.break_sends() + handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="aGk=")) + await asyncio.sleep(0.05) + + # The engine session is alive and resumable, so this is a disconnect of + # the current attachment, not a session close; the pump keeps journaling. + assert handle.close_reasons == [] + assert omni.detached == [handle.session_id] + assert not handler._pumps[handle.session_id].done() + + # The reader sees the same broken socket a moment later. The attachment is + # already gone, so this must not detach a second time: another detach would + # restart the engine's disconnect grace window. + ws.disconnect() + await asyncio.wait_for(task, timeout=2.0) + assert omni.detached == [handle.session_id] + assert handle.close_reasons == [] + + handle.deliver(SessionClosed(session_id=handle.session_id, reason="client_close")) + await asyncio.sleep(0.05) + assert handle.session_id not in handler._pumps + + +@pytest.mark.asyncio +async def test_non_resumable_session_is_closed_on_disconnect() -> None: + omni = FakeOmni(resumable=False) + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + assert "resume_token" not in ws.sent[0] + + ws.disconnect() + await asyncio.wait_for(task, timeout=2.0) + assert handle.close_reasons == ["disconnect"] + assert omni.detached == [] + + +@pytest.mark.asyncio +async def test_a_client_close_still_delivers_session_closed_before_the_socket_goes() -> None: + """The read loop must not outrun the pump's terminal event. + + ``DuplexSessionHandle._deliver`` queues ``SessionClosed`` and marks the + handle closed in one synchronous step, so ``handle.closed`` is already true + while the event is still sitting in the outbox. ``_read_loop`` loops on + exactly that flag: it returned, the endpoint returned, and the ASGI server + tore the socket down with ``session.closed`` unsent -- the client saw an + abrupt close (no close frame) instead of the terminal event it was waiting + for. + + The teardown is what makes this observable, so it is modelled here: a real + connection stops accepting writes the moment the endpoint returns, which is + why a pump that is merely *still scheduled* is already too late. + """ + omni = FakeOmni() + handler = _handler(omni) + ws = FakeWebSocket({"duplex": "1", "autostart": "0"}) + + async def serve() -> None: + try: + await handler.handle_realtime_session(ws) + finally: + # The ASGI server drops the connection when the endpoint returns; + # anything the pump writes after this never reaches the client. + ws.break_sends() + + task = asyncio.create_task(serve()) + ws.feed(_session_update()) + created = await ws.wait_for("session.created") + handle = omni.handles[created["session"]["id"]] + submit = handle.submit + + async def closing_submit(command: commands.DuplexCommand) -> None: + await submit(command) + if isinstance(command, commands.CloseSession): + # Same order as the real handle: queue the event, then flip the flag. + handle.deliver(SessionClosed(session_id=handle.session_id, reason="client_close")) + handle.closed = True + + handle.submit = closing_submit # type: ignore[method-assign] + + ws.feed({"type": "session.close"}) + await asyncio.wait_for(task, timeout=2.0) + + assert ws.types()[-1] == "session.closed", f"terminal event never reached the wire: {ws.types()}" + assert ws.closed and ws.closed[0][0] == 1000, "the socket must close with a normal close frame" + + +# --------------------------------------------------------------------------- # +# Resume and takeover # +# --------------------------------------------------------------------------- # + + +async def _resume( + handler: OmniDuplexSessionHandler, + session_id: str, + token: str, + *, + last_seq: int = 0, +) -> tuple[FakeWebSocket, asyncio.Task]: + ws = FakeWebSocket({"duplex": "1", "resume": "1"}) + task = asyncio.create_task(handler.handle_realtime_session(ws)) + ws.feed( + { + "type": "session.resume", + "session_id": session_id, + "resume_token": token, + "last_received_server_event_seq": last_seq, + } + ) + return ws, task + + +@pytest.mark.asyncio +async def test_resume_after_disconnect_replays_missed_events_and_rotates_the_token() -> None: + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + token = ws.sent[0]["resume_token"] + handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="one")) + await ws.wait_for("response.output_audio.delta") + ws.disconnect() + await asyncio.wait_for(task, timeout=2.0) + assert omni.detached == [handle.session_id] + # Emitted while detached: journaled for replay only. + handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="two")) + await asyncio.sleep(0.05) + + ws2, task2 = await _resume(handler, handle.session_id, token, last_seq=1) + resumed = await ws2.wait_for("session.resumed") + assert resumed["session_id"] == handle.session_id + assert resumed["attachment_generation"] == 2 + assert resumed["resume_token"] != token + assert "incarnation" not in resumed + replayed = await ws2.wait_for("response.output_audio.delta") + assert replayed["delta"] == "two" and replayed["server_event_seq"] == 2 + assert omni.resumed == [(handle.session_id, 0)] + + # The old token is revoked by the rotation. + ws3, task3 = await _resume(handler, handle.session_id, token) + await asyncio.wait_for(task3, timeout=2.0) + assert ws3.sent[0]["error"]["code"] == "invalid_resume_token" + + ws2.feed({"type": "session.close"}) + await asyncio.sleep(0.05) + ws2.disconnect() + await asyncio.wait_for(task2, timeout=2.0) + + +@pytest.mark.asyncio +async def test_resume_takes_over_a_live_attachment() -> None: + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + token = ws.sent[0]["resume_token"] + + ws2, task2 = await _resume(handler, handle.session_id, token) + await ws2.wait_for("session.resumed") + replaced = await ws.wait_for("session.replaced") + assert replaced["attachment_generation"] == 1 + assert ws.closed and ws.closed[0][1] == "session_replaced" + # The replaced socket's later input is ignored; the new one drives the session. + ws.feed({"type": "input_audio_buffer.clear"}) + ws2.feed({"type": "input_audio_buffer.clear", "event_id": "evt-new"}) + await asyncio.sleep(0.05) + assert [command.event_id for command in handle.commands] == ["evt-new"] + await asyncio.wait_for(task, timeout=2.0) + + ws2.disconnect() + await asyncio.wait_for(task2, timeout=2.0) + assert omni.detached == [handle.session_id] + + +@pytest.mark.asyncio +async def test_resume_validation_errors() -> None: + omni = FakeOmni() + handler = _handler(omni) + ws = FakeWebSocket({"duplex": "1", "resume": "1"}) + task = asyncio.create_task(handler.handle_realtime_session(ws)) + ws.feed({"type": "session.resume", "session_id": "duplex-x"}) + await asyncio.wait_for(task, timeout=2.0) + assert ws.sent[0]["error"]["code"] == "invalid_session_resume" + + ws, task = await _resume(handler, "duplex-unknown", "token") + await asyncio.wait_for(task, timeout=2.0) + assert ws.sent[0]["error"]["code"] == "session_resume_expired" + + +@pytest.mark.asyncio +async def test_journal_overflow_degrades_to_live_delivery_with_resync_required() -> None: + omni = FakeOmni(replay_max_bytes=256) + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + + handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="x" * 400)) + resync = await ws.wait_for("session.resync_required") + assert resync["reason"] == "journal_overflow" + delta = await ws.wait_for("response.output_audio.delta") + assert "server_event_seq" not in delta + assert ws.types().index("session.resync_required") < ws.types().index("response.output_audio.delta") + + ws.disconnect() + await asyncio.wait_for(task, timeout=2.0) + + +# --------------------------------------------------------------------------- # +# Envelope helpers # +# --------------------------------------------------------------------------- # + + +def test_realtime_envelope_query_rules() -> None: + envelope = RealtimeEnvelope.from_query_params({"model": "m"}) + assert envelope.initial_open_payload() == {"model": "m"} + assert envelope.initial_open_payload() is None # autostart happens once + + envelope = RealtimeEnvelope.from_query_params({"model": "m", "autostart": "0"}) + assert envelope.resume_only is True + assert envelope.initial_open_payload() is None + assert RealtimeEnvelope.from_query_params({"resume": "1"}).resume_only is True + assert RealtimeEnvelope.from_query_params({"session_id": "mine"}).default_session_payload() == {"model": None} + + +def test_realtime_envelope_first_message_classification() -> None: + envelope = RealtimeEnvelope.from_query_params({"model": "m", "autostart": "0"}) + resume = envelope.first_message({"type": "session.resume", "session_id": "s"}) + assert resume.kind == "resume" and resume.resume_payload["session_id"] == "s" + + envelope = RealtimeEnvelope.from_query_params({"model": "m", "autostart": "0"}) + opened = envelope.first_message({"type": "session.update", "session": {"model": "x", "instructions": "hi"}}) + assert opened.kind == "open" and opened.session_payload == {"model": "x", "instructions": "hi"} + assert opened.pending_command_payload is None + + envelope = RealtimeEnvelope.from_query_params({"model": "m", "autostart": "0"}) + autostarted = envelope.first_message({"type": "input_audio_buffer.commit"}) + assert autostarted.kind == "open" and autostarted.session_payload == {"model": "m"} + assert autostarted.pending_command_payload == {"type": "input_audio_buffer.commit"} + + +def test_parse_resume_request_requires_the_three_fields_only() -> None: + request = parse_resume_request({"session_id": "s", "resume_token": "t", "last_received_server_event_seq": 3}) + assert request is not None + assert (request.session_id, request.resume_token, request.last_received_server_event_seq) == ("s", "t", 3) + assert parse_resume_request({"session_id": "s", "resume_token": "t"}).last_received_server_event_seq == 0 + assert parse_resume_request({"session_id": "s", "resume_token": "t", "incarnation": 1}) is not None + assert parse_resume_request({"session_id": "s"}) is None + assert parse_resume_request({"session_id": "s", "resume_token": "t", "last_received_server_event_seq": -1}) is None + + +@pytest.mark.asyncio +async def test_a_resume_that_fails_to_activate_does_not_detach_the_live_attachment() -> None: + """A rejected resume must roll back only what it owns. + + Two reconnects can both authenticate and both complete the engine resume; + only one activates. The loser's activation raises, and the rollback used to + detach the *session*, starting the engine's disconnect grace for the winner + -- a grace ordinary heartbeats do not clear. + + The activation failure is injected rather than raced for: what is under test + is the rollback's precondition, not the window that produces it. + """ + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + token = ws.sent[0]["resume_token"] + registry = handler._attachment_registry + try: + assert await registry.has_attachment(handle.session_id), "the opener is attached" + detached_before = list(omni.detached) + + async def failing_resume(*args, **kwargs): + raise RuntimeError("transport activation lost the race") + + original_resume = registry.resume + registry.resume = failing_resume # type: ignore[method-assign] + try: + ws_lose, task_lose = await _resume(handler, handle.session_id, token) + error = await ws_lose.wait_for("error") + finally: + registry.resume = original_resume # type: ignore[method-assign] + + assert error["error"]["code"] == "session_resume_conflict" + assert omni.detached == detached_before, "the loser must not detach the winner" + assert await registry.has_attachment(handle.session_id), "the winner is still attached" + finally: + for pending in (task, locals().get("task_lose")): + if pending is not None: + pending.cancel() + with suppress(asyncio.CancelledError, Exception): + await pending diff --git a/tests/entrypoints/duplex/test_runtime_adapter_boundary.py b/tests/entrypoints/duplex/test_runtime_adapter_boundary.py deleted file mode 100644 index cb90e4bd5..000000000 --- a/tests/entrypoints/duplex/test_runtime_adapter_boundary.py +++ /dev/null @@ -1,145 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import subprocess -import sys -from pathlib import Path -from types import SimpleNamespace - -import pytest - -from vllm_omni.entrypoints.duplex.runtime_adapter import ( - validate_serving_runtime_adapter, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -@pytest.mark.parametrize( - "module_name", - [ - "vllm_omni.entrypoints.duplex.session_runner", - "vllm_omni.entrypoints.duplex.runtime_bridge", - "vllm_omni.entrypoints.duplex.serving", - ], -) -def test_generic_openai_runtime_import_does_not_load_model_adapters(module_name: str) -> None: - repo_root = Path(__file__).resolve().parents[3] - script = f""" -import sys - -import {module_name} - -model_duplex_prefixes = ( - "vllm_omni.model_executor.models.minicpmo_4_5.duplex", - "vllm_omni.model_executor.models.nemotron_voicechat.duplex", - "vllm_omni.model_executor.models.personaplex.duplex", -) -loaded = sorted( - name - for name in sys.modules - if any( - name == prefix or name.startswith(prefix + ".") - for prefix in model_duplex_prefixes - ) -) -if loaded: - raise SystemExit("generic OpenAI runtime loaded model duplex modules: " + ", ".join(loaded)) -""" - - result = subprocess.run( - [sys.executable, "-c", script], - cwd=repo_root, - capture_output=True, - text=True, - timeout=60, - check=False, - ) - - assert result.returncode == 0, result.stderr - - -def test_generic_handler_supports_adapterless_chat_fallback() -> None: - repo_root = Path(__file__).resolve().parents[3] - script = """ -from types import SimpleNamespace - -from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler - -chat_service = SimpleNamespace( - engine_client=SimpleNamespace(), - model_config=SimpleNamespace(model="test-model"), -) -handler = OmniDuplexSessionHandler(chat_service=chat_service) -if handler._serving_runtime_adapter is not None: - raise SystemExit("generic handler silently selected a model serving adapter") -""" - - result = subprocess.run( - [sys.executable, "-c", script], - cwd=repo_root, - capture_output=True, - text=True, - timeout=60, - check=False, - ) - - assert result.returncode == 0, result.stderr - - -def _valid_runtime_adapter() -> SimpleNamespace: - data_plane = SimpleNamespace( - begin_request=lambda request_id: None, - is_terminal=lambda request_id: False, - mark_terminal=lambda request_id: None, - close_stream=lambda request_id: None, - close_session=lambda session_id, **kwargs: None, - project=lambda result, **kwargs: (), - ) - return SimpleNamespace( - adapter_id="test-adapter", - session_states={}, - data_plane=data_plane, - clean_response_done_prefix="clean", - interrupted_tts_prefix="interrupted", - private_runtime_config_keys=frozenset({"private_key"}), - create_session_state=lambda: object(), - session_state=lambda session_id: object(), - remove_session_state=lambda session_id: None, - is_enabled=lambda config: True, - capabilities=lambda **kwargs: object(), - validate_client_extra_body=lambda extra_body: None, - prepare_runtime_config=lambda config, **kwargs: {}, - runtime_config_for_update=lambda config, current: {}, - data_plane_context=lambda **kwargs: object(), - ) - - -@pytest.mark.parametrize( - ("attribute", "invalid_value"), - [ - ("session_states", None), - ("clean_response_done_prefix", None), - ("interrupted_tts_prefix", None), - ("private_runtime_config_keys", {"private_key"}), - ], -) -def test_runtime_adapter_validator_rejects_invalid_protocol_attributes( - attribute: str, - invalid_value: object, -) -> None: - adapter = _valid_runtime_adapter() - setattr(adapter, attribute, invalid_value) - - with pytest.raises(TypeError, match=attribute): - validate_serving_runtime_adapter(adapter) - - -def test_runtime_adapter_validator_rejects_incomplete_data_plane() -> None: - adapter = _valid_runtime_adapter() - adapter.data_plane.close_session = None - - with pytest.raises(TypeError, match="data_plane.*close_session"): - validate_serving_runtime_adapter(adapter) diff --git a/tests/entrypoints/duplex/test_runtime_control_redaction.py b/tests/entrypoints/duplex/test_runtime_control_redaction.py deleted file mode 100644 index cfc5b68c5..000000000 --- a/tests/entrypoints/duplex/test_runtime_control_redaction.py +++ /dev/null @@ -1,115 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -"""Regression tests for duplex ``runtime_control`` payload redaction. - -The redacted result is emitted on the WebSocket (for example -``session.created.runtime_control``), so it must always survive -``json.dumps``. Control results carry typed values straight from the engine -- -``fence`` and ``accepted_fence`` are ``DuplexFence`` instances added by -``DuplexControlClient.execute`` -- which previously reached -``WebSocket.send_json`` unconverted and killed the session at handshake. -""" - -from __future__ import annotations - -import json -from dataclasses import dataclass -from enum import Enum - -import pytest - -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.entrypoints.duplex.runtime_bridge import NativeRuntimeBridgeMixin - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - -_redact = NativeRuntimeBridgeMixin._redact_runtime_control_result - - -def _control_result() -> dict[str, object]: - """The shape ``DuplexControlClient.execute`` returns for an open.""" - return { - "ok": True, - "operation": "open", - "session_id": "sess-1", - "fence": DuplexFence("sess-1", epoch=0, turn_id=0, response_seq=0, incarnation=0), - "accepted_fence": DuplexFence("sess-1", epoch=1, turn_id=2, response_seq=3, incarnation=4), - "stage_results": [{"stage_id": -1, "replica_id": -1, "result": {"supported": True}}], - } - - -def test_redacted_control_result_is_json_serializable() -> None: - """The whole payload must serialize the way send_json does.""" - redacted = _redact(_control_result()) - # Must not raise. - json.dumps({"type": "session.created", "runtime_control": redacted}) - - -def test_fences_are_converted_to_plain_dicts() -> None: - redacted = _redact(_control_result()) - assert redacted["fence"] == { - "session_id": "sess-1", - "epoch": 0, - "turn_id": 0, - "response_seq": 0, - "incarnation": 0, - } - # Field values are preserved, not just the type. - assert redacted["accepted_fence"]["epoch"] == 1 - assert redacted["accepted_fence"]["turn_id"] == 2 - assert redacted["accepted_fence"]["response_seq"] == 3 - assert redacted["accepted_fence"]["incarnation"] == 4 - - -def test_nested_dataclasses_are_converted() -> None: - """Fences can appear inside stage results, not only at the top level.""" - redacted = _redact( - { - "stage_results": [ - {"stage_id": 0, "result": {"fence": DuplexFence("sess-2", epoch=7)}}, - ] - } - ) - json.dumps(redacted) - assert redacted["stage_results"][0]["result"]["fence"]["epoch"] == 7 - - -def test_enums_are_reduced_to_values() -> None: - class _Mode(str, Enum): - APPEND = "append_audio_chunk" - - redacted = _redact({"mode": _Mode.APPEND}) - json.dumps(redacted) - assert redacted["mode"] == "append_audio_chunk" - - -def test_redaction_still_drops_heavy_payload_keys() -> None: - """The pre-existing contract must be preserved.""" - redacted = _redact( - { - "ok": True, - "stage_handoff": object(), - "tts_handoff": object(), - "omni_payload": object(), - "tts_hidden_states": object(), - "tts_token_ids": object(), - "traceback": "long trace", - } - ) - assert redacted == {"ok": True} - - -def test_plain_values_pass_through_unchanged() -> None: - payload = {"ok": True, "count": 3, "name": "x", "missing": None, "items": [1, "two"]} - assert _redact(payload) == payload - - -def test_dataclass_types_are_not_treated_as_instances() -> None: - """A class object is not a dataclass instance and must not be asdict'd.""" - - @dataclass - class _Shape: - x: int = 1 - - assert _redact(_Shape) is _Shape diff --git a/tests/entrypoints/duplex/test_websocket_actor.py b/tests/entrypoints/duplex/test_websocket_actor.py deleted file mode 100644 index 4e5d151b9..000000000 --- a/tests/entrypoints/duplex/test_websocket_actor.py +++ /dev/null @@ -1,145 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -import asyncio - -import pytest - -from vllm_omni.entrypoints.duplex.websocket import ( - DuplexWebSocketActor, - normalize_duplex_input_event, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class FakeWebSocket: - def __init__(self) -> None: - self.sent: list[dict[str, object]] = [] - - async def send_json(self, payload: dict[str, object]) -> None: - self.sent.append(dict(payload)) - - -@pytest.mark.parametrize( - ("event", "expected"), - [ - ({"type": "signal_turn", "event": "barge_in"}, {"type": "turn.signal", "event": "barge_in"}), - ({"type": "close_session"}, {"type": "session.close"}), - ({"type": "audio.playback_ack", "played_ms": 1}, {"type": "playback.ack", "played_ms": 1}), - ({"type": "input_text.append", "text": "a"}, {"type": "input.text.append", "text": "a"}), - ({"type": "push_text", "text": "b"}, {"type": "input.text.append", "text": "b"}), - ( - {"type": "input.audio.append", "audio": "wav"}, - {"type": "input_audio_buffer.append", "audio": "wav", "format": "wav"}, - ), - ( - {"type": "push_chunk", "audio": "wav", "format": "pcm_f32le"}, - {"type": "input_audio_buffer.append", "audio": "wav", "format": "pcm_f32le"}, - ), - ], -) -def test_input_aliases_normalize_once_at_mailbox_boundary(event, expected): - assert normalize_duplex_input_event(event) == expected - - -def test_actor_uses_one_fifo_mailbox_for_inbound_events(): - actor = DuplexWebSocketActor(FakeWebSocket()) - - assert isinstance(actor.mailbox, asyncio.Queue) - assert not isinstance(actor.mailbox, asyncio.PriorityQueue) - assert not hasattr(actor, "control_queue") - assert not hasattr(actor, "input_queue") - assert not hasattr(actor, "event_queue") - assert not hasattr(actor, "lifecycle_state") - assert not hasattr(actor, "last_response_id") - assert not hasattr(actor, "overlap_speech_ms") - assert not hasattr(actor, "transition") - assert not hasattr(actor, "output_generation_in_flight") - assert not hasattr(actor, "session") - assert not hasattr(actor, "runtime_opened") - assert not hasattr(actor, "runtime_closed") - assert not hasattr(actor, "drain_input_queue") - assert not hasattr(actor, "control_events_seen") - assert not hasattr(actor, "input_events_seen") - assert not hasattr(actor, "cancel_count") - - -@pytest.mark.asyncio -async def test_control_event_does_not_overtake_earlier_audio_input(): - actor = DuplexWebSocketActor(FakeWebSocket()) - await actor.enqueue_event({"type": "input_audio_buffer.append", "audio": "pcm"}) - await actor.enqueue_event({"type": "response.cancel"}) - - assert await actor.next_event() == {"type": "input_audio_buffer.append", "audio": "pcm"} - assert await actor.next_event() == {"type": "response.cancel"} - - -@pytest.mark.asyncio -@pytest.mark.parametrize("event_type", ["input.cancel", "session.close", "close_session"]) -async def test_terminal_control_preserves_wire_order(event_type: str): - actor = DuplexWebSocketActor(FakeWebSocket()) - await actor.enqueue_event({"type": "input_audio_buffer.append", "audio": "pcm"}) - await actor.enqueue_event({"type": event_type}) - - assert await actor.next_event() == {"type": "input_audio_buffer.append", "audio": "pcm"} - assert await actor.next_event() == {"type": event_type} - - -@pytest.mark.asyncio -async def test_mailbox_delivers_every_enqueued_event_exactly_once(): - actor = DuplexWebSocketActor(FakeWebSocket()) - events = [ - {"type": "input_audio_buffer.append", "audio": str(index)} - if index % 2 - else {"type": "response.created", "response_id": f"resp-{index}"} - for index in range(20) - ] - - await asyncio.gather(*(actor.enqueue_event(dict(event)) for event in events)) - received = [await actor.next_event() for _ in events] - - assert sorted(received, key=repr) == sorted(events, key=repr) - - -@pytest.mark.asyncio -async def test_writer_is_single_owner_of_websocket_send(): - websocket = FakeWebSocket() - actor = DuplexWebSocketActor(websocket) - writer = asyncio.create_task(actor.writer_loop()) - - await actor.send_json({"type": "one"}) - await actor.send_json({"type": "two"}) - await actor.close_writer() - await writer - - assert websocket.sent == [{"type": "one"}, {"type": "two"}] - - -@pytest.mark.asyncio -async def test_stale_fence_payload_is_dropped_before_websocket_send(): - websocket = FakeWebSocket() - actor = DuplexWebSocketActor(websocket, current_epoch=lambda: 2) - writer = asyncio.create_task(actor.writer_loop()) - - await actor.send_json({"type": "response.output_audio.delta", "epoch": 1}) - await actor.send_json({"type": "response.output_audio.delta", "epoch": 2}) - await actor.close_writer() - await writer - - assert websocket.sent == [{"type": "response.output_audio.delta", "epoch": 2}] - assert actor.stale_output_dropped == 1 - - -@pytest.mark.asyncio -async def test_writer_does_not_revoke_accepted_terminal_after_close_starts(): - websocket = FakeWebSocket() - actor = DuplexWebSocketActor(websocket, current_epoch=lambda: 1) - await actor.send_json({"type": "response.done", "epoch": 1, "response_id": "resp-1"}) - actor.closing = True - await actor.close_writer() - - await actor.writer_loop() - - assert websocket.sent == [{"type": "response.done", "epoch": 1, "response_id": "resp-1"}] - assert actor.stale_output_dropped == 0 diff --git a/tests/entrypoints/openai/test_duplex_protocol.py b/tests/entrypoints/openai/test_duplex_protocol.py deleted file mode 100644 index 77c8f2cb0..000000000 --- a/tests/entrypoints/openai/test_duplex_protocol.py +++ /dev/null @@ -1,601 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -import asyncio -import json - -import pytest - -from vllm_omni.entrypoints.duplex.protocol import ( - DuplexOverlapPolicy, - DuplexSession, - DuplexSessionConfig, - DuplexSessionRegistry, - DuplexTurnController, - DuplexTurnEventType, - DuplexTurnState, - ResponseCreateOptions, -) -from vllm_omni.entrypoints.duplex.realtime_session import ( - NativeRealtimeSessionProtocol, -) -from vllm_omni.entrypoints.duplex.realtime_state import RealtimeStateOwner -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.capabilities import ( - minicpmo45_native_capabilities, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def test_realtime_state_owner_uses_explicit_bindings_not_dynamic_attribute_proxy(): - assert "__getattr__" not in RealtimeStateOwner.__dict__ - assert "__setattr__" not in RealtimeStateOwner.__dict__ - - -class _ProtocolWebSocket: - def __init__(self, *events: dict[str, object]) -> None: - self._events = asyncio.Queue() - for event in events: - self._events.put_nowait(json.dumps(event)) - - async def receive_text(self) -> str: - return await self._events.get() - - -def test_duplex_session_commits_text_and_audio_as_one_turn(): - registry = DuplexSessionRegistry() - session = registry.create(DuplexSessionConfig(model="test-model")) - - session.append_text("hello") - session.append_audio("YWJj", fmt="wav", sample_rate_hz=16000) - committed = session.commit_user_input() - - assert committed is not None - assert committed.turn_id == 0 - assert committed.input_commit_seq == 1 - assert committed.epoch == 0 - assert len(session.history) == 1 - content = session.history[0]["content"] - assert isinstance(content, list) - assert content[0] == {"type": "text", "text": "hello"} - assert content[1]["type"] == "audio_url" - assert content[1]["audio_url"]["url"] == "data:audio/wav;base64,YWJj" - assert session.turn_state == DuplexTurnState.USER_COMMITTED - - -def test_response_options_apply_to_one_response_without_mutating_session_defaults(): - session = DuplexSession( - session_id="sid-response-options", - config=DuplexSessionConfig(instructions="base", voice="base-voice", max_tokens=64), - ) - session.reserve_response_options( - ResponseCreateOptions( - instructions="one response", - voice="override-voice", - max_tokens=8, - ) - ) - - assert session.config.instructions == "base" - assert session.config.voice == "base-voice" - session.begin_response() - assert session.config.instructions == "base" - assert session.config.voice == "base-voice" - assert session.response_config.instructions == "one response" - assert session.response_config.voice == "override-voice" - assert session.response_config.max_tokens == 8 - - session.end_response() - assert session.config.instructions == "base" - assert session.config.voice == "base-voice" - assert session.config.max_tokens == 64 - - -def test_response_options_cannot_overwrite_an_unconsumed_reservation(): - session = DuplexSession( - session_id="sid-response-options-pending", - config=DuplexSessionConfig(instructions="base"), - ) - session.reserve_response_options(ResponseCreateOptions(instructions="first")) - - with pytest.raises(RuntimeError, match="already reserved"): - session.reserve_response_options(ResponseCreateOptions(instructions="second")) - - session.begin_response() - assert session.response_config.instructions == "first" - - -def test_response_options_cannot_be_reserved_while_response_is_active(): - session = DuplexSession( - session_id="sid-response-options-active", - config=DuplexSessionConfig(instructions="base"), - ) - session.begin_response() - - with pytest.raises(RuntimeError, match="active"): - session.reserve_response_options(ResponseCreateOptions(instructions="too late")) - - -def test_realtime_model_name_does_not_implicitly_enable_native_duplex(): - protocol = NativeRealtimeSessionProtocol({}) - - event = protocol._session_create_from_realtime({"model": "openbmb/MiniCPM-o-4_5"}) - - assert "native_duplex" not in event["session"]["extra_body"] - assert "minicpmo45_native_duplex" not in event["session"]["extra_body"] - - -def test_realtime_explicit_native_duplex_flag_is_preserved(): - protocol = NativeRealtimeSessionProtocol({}) - - event = protocol._session_create_from_realtime( - { - "model": "openbmb/MiniCPM-o-4_5", - "extra_body": {"native_duplex": True}, - } - ) - - assert event["session"]["extra_body"]["native_duplex"] is True - - -def test_realtime_legacy_native_duplex_alias_is_normalized(): - # The deprecated model-prefixed spelling still opts in, but everything - # downstream (and every session echo) sees only the canonical key. - protocol = NativeRealtimeSessionProtocol({}) - - event = protocol._session_create_from_realtime( - { - "model": "openbmb/MiniCPM-o-4_5", - "extra_body": {"minicpmo45_native_duplex": True}, - } - ) - - assert event["session"]["extra_body"]["native_duplex"] is True - assert "minicpmo45_native_duplex" not in event["session"]["extra_body"] - - -def test_realtime_explicit_query_native_duplex_flag_is_available_before_autostart(): - protocol = NativeRealtimeSessionProtocol( - { - "model": "openbmb/MiniCPM-o-4_5", - "native_duplex": "1", - } - ) - - event = json.loads(asyncio.run(protocol.receive_internal_event_text(None))) - - assert event["type"] == "session.create" - assert event["session"]["extra_body"]["native_duplex"] is True - - -def test_realtime_legacy_query_native_duplex_alias_seeds_canonical_key(): - protocol = NativeRealtimeSessionProtocol( - { - "model": "openbmb/MiniCPM-o-4_5", - "minicpmo45_native_duplex": "1", - } - ) - - event = json.loads(asyncio.run(protocol.receive_internal_event_text(None))) - - assert event["type"] == "session.create" - assert event["session"]["extra_body"]["native_duplex"] is True - assert "minicpmo45_native_duplex" not in event["session"]["extra_body"] - - -@pytest.mark.asyncio -async def test_realtime_resume_heartbeat_and_event_ack_translate_without_session_autostart(): - protocol = NativeRealtimeSessionProtocol( - { - "model": "openbmb/MiniCPM-o-4_5", - "native_duplex": "1", - } - ) - - resume = await protocol._to_duplex_event( - { - "type": "session.resume", - "session_id": "sid-resume", - "incarnation": 3, - "resume_token": "opaque-token", - "last_received_server_event_seq": 17, - } - ) - heartbeat = await protocol._to_duplex_event( - { - "type": "session.heartbeat", - "event_id": "event-heartbeat", - } - ) - event_ack = await protocol._to_duplex_event( - { - "type": "session.event_ack", - "server_event_seq": 19, - } - ) - - assert resume == { - "type": "session.resume", - "session_id": "sid-resume", - "incarnation": 3, - "resume_token": "opaque-token", - "last_received_server_event_seq": 17, - } - assert heartbeat == { - "type": "session.heartbeat", - "event_id": "event-heartbeat", - } - assert event_ack == { - "type": "session.event_ack", - "server_event_seq": 19, - } - assert protocol._opened is False - - wire_protocol = NativeRealtimeSessionProtocol({}) - first_wire_event = json.loads( - await wire_protocol.receive_internal_event_text( - _ProtocolWebSocket( - { - "type": "session.resume", - "session_id": "sid-resume", - "incarnation": 3, - "resume_token": "opaque-token", - "last_received_server_event_seq": 17, - } - ) - ) - ) - - assert first_wire_event["type"] == "session.resume" - - -@pytest.mark.asyncio -async def test_realtime_resume_query_suppresses_model_autostart_on_same_url(): - protocol = NativeRealtimeSessionProtocol( - { - "model": "openbmb/MiniCPM-o-4_5", - "native_duplex": "1", - "resume": "1", - } - ) - resume = { - "type": "session.resume", - "session_id": "sid-resume", - "incarnation": 3, - "resume_token": "opaque-token", - "last_received_server_event_seq": 17, - } - - first_wire_event = json.loads(await protocol.receive_internal_event_text(_ProtocolWebSocket(resume))) - - assert first_wire_event == resume - assert protocol._autostarted_default_session is False - - -def test_realtime_projects_resume_lifecycle_events_without_duplex_prefix(): - protocol = NativeRealtimeSessionProtocol({}) - - created = protocol.encode_outbound_event( - { - "type": "session.created", - "session": {"id": "sid-resume"}, - "incarnation": 3, - "attachment_generation": 1, - "resume_token": "first-token", - } - )[0] - resumed = protocol.encode_outbound_event( - { - "type": "session.resumed", - "session_id": "sid-resume", - "incarnation": 3, - "attachment_generation": 2, - "resume_token": "second-token", - } - )[0] - heartbeat = protocol.encode_outbound_event( - { - "type": "session.heartbeat_ack", - "session_id": "sid-resume", - } - )[0] - replaced = protocol.encode_outbound_event( - { - "type": "session.replaced", - "session_id": "sid-resume", - "attachment_generation": 1, - } - )[0] - resync = protocol.encode_outbound_event( - { - "type": "session.resync_required", - "session_id": "sid-resume", - "reason": "journal_gap", - } - )[0] - - assert created["type"] == "session.created" - assert created["incarnation"] == 3 - assert created["attachment_generation"] == 1 - assert created["resume_token"] == "first-token" - assert resumed["type"] == "session.resumed" - assert resumed["resume_token"] == "second-token" - assert heartbeat["type"] == "session.heartbeat_ack" - assert replaced["type"] == "session.replaced" - assert resync == { - "type": "session.resync_required", - "session_id": "sid-resume", - "reason": "journal_gap", - "event_id": resync["event_id"], - } - - -def test_realtime_projects_listen_response_identity(): - protocol = NativeRealtimeSessionProtocol({}) - protocol.encode_outbound_event({"type": "session.created", "session": {"id": "sid-listen"}}) - - bound = protocol.encode_outbound_event( - { - "type": "response.listen", - "session_id": "sid-listen", - "epoch": 0, - "reason": "model_listen", - "model_listen": True, - "response_id": "resp_1", - } - )[0] - # A listen that terminates a precreated response keeps the response - # identity on the Realtime dialect; clients demultiplex on response_id. - assert bound["type"] == "response.listen" - assert bound["response_id"] == "resp_1" - assert bound["response"]["id"] == "resp_1" - assert bound["response"]["status"] == "listening" - - anonymous = protocol.encode_outbound_event( - { - "type": "response.listen", - "session_id": "sid-listen", - "epoch": 0, - "reason": "silence_or_noise", - } - )[0] - assert "response_id" not in anonymous - assert "id" not in anonymous["response"] - - -def test_duplex_session_registry_advances_incarnation_when_id_is_reused(): - registry = DuplexSessionRegistry() - - first = registry.create(session_id="reused-session") - registry.close(first.session_id) - second = registry.create(session_id="reused-session") - - assert first.incarnation == 0 - assert second.incarnation == 1 - - -def test_native_input_commit_does_not_advance_model_turn_identity(): - session = DuplexSession( - session_id="sid-model-owned-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - - first = session.commit_native_audio_input(transcript="first chunk") - second = session.commit_native_audio_input(transcript="second chunk") - - assert session.input_commit_seq == 2 - assert first.input_commit_seq == 1 - assert second.input_commit_seq == 2 - assert first.turn_id == second.turn_id == session.turn_id == 0 - - -def test_duplex_session_owns_response_and_overlap_identity(): - session = DuplexSession(session_id="sid-state-owner", config=DuplexSessionConfig()) - - response_id = session.begin_response() - session.accumulate_overlap_speech(320) - session.accumulate_overlap_speech(180) - session.end_response() - - assert session.active_response_id is None - assert session.last_response_id == response_id - assert session.overlap_speech_ms == 500 - assert session.reset_overlap_speech() == 500 - assert session.overlap_speech_ms == 0 - - -def test_duplex_session_composes_single_owner_ledgers(): - session = DuplexSession(session_id="sid-ledgers", config=DuplexSessionConfig()) - - session.append_text("hello") - session.bind_request("req-1") - response_id = session.begin_response(turn_id=3) - session.mark_audio_sent(duration_ms=240) - - assert not hasattr(session, "input") - assert not hasattr(session, "response") - assert not hasattr(session, "playback_ledger") - assert not hasattr(session, "conversation") - assert session.pending_text == ("hello",) - assert session.active_request_id == "req-1" - assert session.active_response_id == response_id - assert session.active_response_turn_id == 3 - assert session.playback.sent_ms == 240 - - try: - session.playback.sent_ms = 480 - except (AttributeError, TypeError): - pass - else: - raise AssertionError("playback view must be immutable") - - -def test_duplex_barge_in_advances_epoch_and_drops_uncommitted_assistant_text(): - registry = DuplexSessionRegistry() - session = registry.create() - response_id = session.begin_response() - session.bind_request("chatcmpl-duplex-test") - session.append_assistant_text("unplayed answer") - - new_epoch = session.barge_in() - - assert response_id is not None - assert new_epoch == 1 - assert session.epoch == 1 - assert session.active_request_id is None - assert session.active_response_id is None - assert session.assistant_text_buffer == () - assert session.history == () - assert session.turn_state == DuplexTurnState.BARGE_IN - - -def test_duplex_playback_ack_tracks_committed_cursor_separately(): - registry = DuplexSessionRegistry() - session = registry.create() - session.mark_audio_sent(duration_ms=10_000) - - session.acknowledge_playback(played_ms=2_000) - - assert session.playback.generated_ms == 10_000 - assert session.playback.sent_ms == 10_000 - assert session.playback.played_ms == 2_000 - assert session.playback.committed_ms == 2_000 - - -@pytest.mark.parametrize("sample_rate_hz,frame_count", [(22_050, 23), (24_000, 24)]) -def test_duplex_generated_audio_rounds_only_cumulative_frames(sample_rate_hz: int, frame_count: int): - session = DuplexSessionRegistry().create() - response_id = session.begin_response() - - durations = [ - session.record_generated_audio(response_id, frame_count=1, sample_rate_hz=sample_rate_hz) - for _ in range(frame_count) - ] - - assert durations[:-1] == [0] * (frame_count - 1) - assert durations[-1] == 1 - assert session.playback.generated_ms == 1 - assert session.playback.sent_ms == 0 - assert session.playback.played_ms == 0 - - -def test_duplex_generated_audio_rejects_rate_change_without_advancing(): - session = DuplexSessionRegistry().create() - response_id = session.begin_response() - session.record_generated_audio(response_id, frame_count=2205, sample_rate_hz=22_050) - - with pytest.raises(ValueError, match="sample rate changed"): - session.record_generated_audio(response_id, frame_count=2400, sample_rate_hz=24_000) - - assert session.playback.generated_ms == 100 - assert session.playback.sent_ms == 0 - - -def test_duplex_late_audio_accounting_cannot_change_new_response(): - session = DuplexSessionRegistry().create() - old_response = session.begin_response() - session.record_generated_audio(old_response, frame_count=2205, sample_rate_hz=22_050) - session.barge_in() - new_response = session.begin_response() - session.record_generated_audio(new_response, frame_count=2400, sample_rate_hz=48_000) - current_state = session.turn_state - - assert session.record_generated_audio(old_response, frame_count=2205, sample_rate_hz=22_050) is None - session.mark_audio_sent(100, response_id=old_response, text_chars=99) - - assert session.playback_for_response(old_response).sent_ms == 100 - assert session.playback.generated_ms == 50 - assert session.playback.sent_ms == 0 - assert session.assistant_audio_text_marks == () - assert session.turn_state == current_state - - -def test_duplex_history_commit_uses_audio_text_alignment_marks(): - registry = DuplexSessionRegistry() - session = registry.create() - session.begin_response() - - session.append_assistant_text("hello ") - session.mark_audio_sent(duration_ms=1_000, text_chars=6) - session.append_assistant_text("world") - session.mark_audio_sent(duration_ms=2_000, text_chars=11) - session.acknowledge_playback(played_ms=1_200, committed_ms=1_200) - - committed = session.end_response(commit_text=True) - - assert committed == {"role": "assistant", "content": "hello w"} - - -def test_turn_controller_accepts_external_signal_source(): - registry = DuplexSessionRegistry() - session = registry.create() - controller = DuplexTurnController() - - event = controller.signal(session, DuplexTurnEventType.USER_STARTED.value) - - assert event["type"] == "turn.event" - assert event["event"] == "user_started" - assert event["turn_state"] == "user_speaking" - assert session.turn_state == DuplexTurnState.USER_SPEAKING - - -def test_duplex_capabilities_do_not_claim_core_kv_or_input_append(): - registry = DuplexSessionRegistry() - session = registry.create() - caps = session.capabilities.as_dict() - - assert caps["implementation_level"] == "serving_session_adapter" - assert caps["supports_kv_lease"] is False - assert caps["supports_input_append"] is False - assert caps["supports_reencode_context"] is True - assert caps["adapter_patterns"] == ["chunk_group_append"] - assert "turn_commit_only" in caps["input_modes"] - assert "client_event" in caps["signal_sources"] - - -def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease(): - caps = minicpmo45_native_capabilities(max_sessions=2).as_dict() - - assert caps["implementation_level"] == "model_native_duplex" - assert caps["supports_input_append"] is True - assert caps["input_modes"] == ["append_audio_chunk"] - assert caps["adapter_patterns"] == ["scheduler_data_plane"] - assert caps["supports_model_internal_state"] is True - assert caps["requires_model_runner_kv"] is True - assert caps["requires_native_stage_role"] is True - assert caps["supports_kv_lease"] is False - assert caps["supports_core_kv_lease"] is False - assert caps["supports_stage_resumption"] is True - assert caps["supports_scheduler_native_append"] is False - assert caps["supports_core_resumable_request"] is True - assert caps["supports_stage_connector_handoff"] is True - assert caps["supports_audio_truncate"] is True - assert caps["supports_barge_in"] is True - assert caps["target_barge_in_latency_ms"] is None - assert caps["supports_multi_session"] is True - assert caps["supports_multi_session_same_replica"] is True - assert caps["supports_session_lease"] is True - assert caps["supports_session_resume"] is True - assert caps["session_admission_mode"] == "engine_managed" - assert caps["stage_handoff_transport"] == "scheduler_data_plane" - - -def test_duplex_overlap_policy_auto_falls_back_to_listen_only(): - from vllm_omni.entrypoints.duplex.protocol import ( - DuplexOverlapPolicy, - DuplexSessionConfig, - ) - - assert DuplexSessionConfig._normalize_overlap_policy("auto") == DuplexOverlapPolicy.LISTEN_ONLY.value - - -def test_minicpmo_native_capabilities_do_not_overclaim_single_session_deployment(): - caps = minicpmo45_native_capabilities(max_sessions=1).as_dict() - - assert caps["supports_multi_session"] is False - assert caps["supports_multi_session_same_replica"] is False - - -def test_duplex_overlap_policy_defaults_and_invalid_values_to_listen_only(): - assert DuplexSessionConfig().overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value - assert DuplexSessionConfig._normalize_overlap_policy("not-a-policy") == DuplexOverlapPolicy.LISTEN_ONLY.value diff --git a/tests/entrypoints/openai/test_duplex_session_attachment.py b/tests/entrypoints/openai/test_duplex_session_attachment.py index 13bd32761..9d7a208b5 100644 --- a/tests/entrypoints/openai/test_duplex_session_attachment.py +++ b/tests/entrypoints/openai/test_duplex_session_attachment.py @@ -133,7 +133,7 @@ async def test_registry_failed_send_is_accepted_only_when_journaled(mocker, jour assert accepted.call_args_list == expected_calls raise RuntimeError("transport lost") - await registry.create("sid-send-failure", incarnation=0, send=failing_send, close=mocker.AsyncMock()) + await registry.create("sid-send-failure", send=failing_send, close=mocker.AsyncMock()) with pytest.raises(RuntimeError, match="transport lost"): await registry.send_event( @@ -152,7 +152,7 @@ async def test_registry_detached_event_is_accepted_only_when_journaled(mocker, j registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096) send = mocker.AsyncMock() accepted = mocker.Mock() - created = await registry.create("sid-detached-acceptance", incarnation=0, send=send, close=mocker.AsyncMock()) + created = await registry.create("sid-detached-acceptance", send=send, close=mocker.AsyncMock()) await registry.detach("sid-detached-acceptance", attachment_generation=created.attachment_generation) entry = await registry.send_event( @@ -175,7 +175,7 @@ async def test_registry_unjournaled_event_is_accepted_after_successful_send(mock async def send(payload): accepted.assert_not_called() - await registry.create("sid-live-acceptance", incarnation=0, send=send, close=mocker.AsyncMock()) + await registry.create("sid-live-acceptance", send=send, close=mocker.AsyncMock()) entry = await registry.send_event( "sid-live-acceptance", @@ -193,7 +193,7 @@ async def test_registry_overflow_does_not_accept_or_send_event(mocker) -> None: registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=100) accepted = mocker.Mock() send = mocker.AsyncMock() - await registry.create("sid-overflow-acceptance", incarnation=0, send=send, close=mocker.AsyncMock()) + await registry.create("sid-overflow-acceptance", send=send, close=mocker.AsyncMock()) with pytest.raises(DuplexJournalOverflowError, match="byte limit"): await registry.send_event( @@ -231,13 +231,12 @@ async def test_registry_resume_rotates_token_replays_and_atomically_replaces_att async def close_b(reason): closes_b.append(reason) - created = await registry.create("sid", incarnation=2, send=send_a, close=close_a) + created = await registry.create("sid", send=send_a, close=close_a) await registry.send_event("sid", {"type": "event-1"}) await registry.send_event("sid", {"type": "event-2"}) resumed = await registry.resume( "sid", - incarnation=2, resume_token=created.resume_token.plaintext, last_received_server_event_seq=1, send=send_b, @@ -260,6 +259,34 @@ async def test_registry_resume_rotates_token_replays_and_atomically_replaces_att assert closes_a == [] and sends_b == [] and closes_b == [] +@pytest.mark.asyncio +async def test_registry_detach_without_a_generation_targets_the_current_attachment_once() -> None: + registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096) + sends: list[dict] = [] + + async def send(payload): + sends.append(payload) + + async def close(reason): + return None + + await registry.create("sid-current", send=send, close=close) + + # The outbound pump outlives any one connection and only knows the session: + # ``None`` means "whichever socket is attached right now". + assert await registry.detach("sid-current", attachment_generation=None) is True + # A second report of the same dead socket is not a new detach, so it cannot + # restart the engine's disconnect grace window. + assert await registry.detach("sid-current", attachment_generation=None) is False + assert await registry.detach("sid-current", attachment_generation=1) is False + assert await registry.is_current_attachment("sid-current", 1) is False + + # A detached session still journals for a later resume. + await registry.send_event("sid-current", {"type": "event-1"}) + assert sends == [] + assert await registry.detach("missing-session", attachment_generation=None) is False + + @pytest.mark.asyncio async def test_registry_resume_sends_activation_then_replay_before_new_live_events() -> None: registry = DuplexSessionAttachmentRegistry( @@ -282,13 +309,12 @@ async def test_registry_resume_sends_activation_then_replay_before_new_live_even activation_started.set() await release_activation.wait() - created = await registry.create("sid-order", incarnation=0, send=old_send, close=close) + created = await registry.create("sid-order", send=old_send, close=close) await registry.detach("sid-order", attachment_generation=1) await registry.send_event("sid-order", {"type": "replayed"}) resume_task = asyncio.create_task( registry.resume( "sid-order", - incarnation=0, resume_token=created.resume_token.plaintext, last_received_server_event_seq=0, send=new_send, @@ -322,13 +348,12 @@ async def test_registry_resume_delivery_failure_keeps_one_shot_old_token_recover async def close(reason): del reason - created = await registry.create("sid-recovery", incarnation=0, send=send, close=close) + created = await registry.create("sid-recovery", send=send, close=close) await registry.detach("sid-recovery", attachment_generation=1) with pytest.raises(RuntimeError, match="transport lost"): await registry.resume( "sid-recovery", - incarnation=0, resume_token=created.resume_token.plaintext, last_received_server_event_seq=0, send=failing_send, @@ -342,7 +367,6 @@ async def test_registry_resume_delivery_failure_keeps_one_shot_old_token_recover recovered = await registry.resume( "sid-recovery", - incarnation=0, resume_token=created.resume_token.plaintext, last_received_server_event_seq=0, send=send, @@ -353,7 +377,6 @@ async def test_registry_resume_delivery_failure_keeps_one_shot_old_token_recover with pytest.raises(InvalidResumeTokenError): await registry.resume( "sid-recovery", - incarnation=0, resume_token=created.resume_token.plaintext, last_received_server_event_seq=0, send=send, @@ -374,13 +397,12 @@ async def test_registry_concurrent_resume_allows_exactly_one_rotated_token_winne async def close(reason): del reason - created = await registry.create("sid-race", incarnation=0, send=send, close=close) + created = await registry.create("sid-race", send=send, close=close) async def attempt(): try: return await registry.resume( "sid-race", - incarnation=0, resume_token=created.resume_token.plaintext, last_received_server_event_seq=0, send=send, @@ -396,7 +418,7 @@ async def test_registry_concurrent_resume_allows_exactly_one_rotated_token_winne @pytest.mark.asyncio -async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> None: +async def test_registry_keeps_sessions_journals_and_tokens_isolated() -> None: registry = DuplexSessionAttachmentRegistry( replay_ttl_s=60.0, replay_max_bytes_per_session=4096, @@ -408,8 +430,8 @@ async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> N async def close(reason): del reason - created_a = await registry.create("sid-a", incarnation=1, send=send, close=close) - created_b = await registry.create("sid-b", incarnation=4, send=send, close=close) + created_a = await registry.create("sid-a", send=send, close=close) + created_b = await registry.create("sid-b", send=send, close=close) await registry.detach("sid-a", attachment_generation=1) await registry.detach("sid-b", attachment_generation=1) await registry.send_event("sid-a", {"type": "a-only"}) @@ -417,7 +439,6 @@ async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> N resumed_a = await registry.resume( "sid-a", - incarnation=1, resume_token=created_a.resume_token.plaintext, last_received_server_event_seq=0, send=send, @@ -425,7 +446,6 @@ async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> N ) resumed_b = await registry.resume( "sid-b", - incarnation=4, resume_token=created_b.resume_token.plaintext, last_received_server_event_seq=0, send=send, @@ -434,11 +454,10 @@ async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> N assert [entry.payload["type"] for entry in resumed_a.replay_entries] == ["a-only"] assert [entry.payload["type"] for entry in resumed_b.replay_entries] == ["b-only"] - with pytest.raises(ValueError, match="incarnation mismatch"): + with pytest.raises(InvalidResumeTokenError): await registry.resume( "sid-a", - incarnation=2, - resume_token=resumed_a.resume_token.plaintext, + resume_token=resumed_b.resume_token.plaintext, last_received_server_event_seq=1, send=send, close=close, @@ -458,7 +477,7 @@ async def test_registry_repr_never_contains_plaintext_tokens() -> None: async def close(reason): del reason - created = await registry.create("sid-repr", incarnation=0, send=send, close=close) + created = await registry.create("sid-repr", send=send, close=close) await registry.detach("sid-repr", attachment_generation=1) entry = await registry.send_event( "sid-repr", @@ -469,3 +488,67 @@ async def test_registry_repr_never_contains_plaintext_tokens() -> None: assert created.resume_token.plaintext not in repr(created) assert created.resume_token.plaintext not in repr(entry) assert created.resume_token.plaintext not in repr(registry) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel_on", ["activation", "replay"]) +async def test_registry_resume_cancelled_mid_delivery_rolls_back_like_a_failure(cancel_on: str) -> None: + """Cancellation must take the same rollback as an ordinary delivery failure. + + ``CancelledError`` derives from ``BaseException``, so an ``except + Exception`` around the activation and replay sends let a cancelled resume + keep its attachment as the current generation while the token it rotated + past was already gone — the session left attached to an unusable transport + with no way back. Both send points are covered because either can be the + one that gets cancelled. + """ + registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096) + + async def send(payload): + del payload + + async def close(reason): + del reason + + created = await registry.create("sid-cancel", send=send, close=close) + # One journaled event, so the replay loop has something to send. + await registry.send_event("sid-cancel", {"type": "response.output_text.delta", "delta": "a"}) + await registry.detach("sid-cancel", attachment_generation=1) + + sends = 0 + cancel_at = 1 if cancel_on == "activation" else 2 + + async def cancelling_send(payload): + del payload + nonlocal sends + sends += 1 + if sends == cancel_at: + raise asyncio.CancelledError + + with pytest.raises(asyncio.CancelledError): + await registry.resume( + "sid-cancel", + resume_token=created.resume_token.plaintext, + last_received_server_event_seq=0, + send=cancelling_send, + close=close, + activation_payload_factory=lambda token, generation: { + "type": "session.resumed", + "resume_token": token.plaintext, + "attachment_generation": generation, + }, + ) + + assert sends == cancel_at, "the send that was cancelled is the one under test" + # The cancelled attachment must not remain current... + assert not await registry.has_attachment("sid-cancel") + # ...and the token the caller presented has to work again, or the session + # is unreachable: resume() had already rotated past it. + recovered = await registry.resume( + "sid-cancel", + resume_token=created.resume_token.plaintext, + last_received_server_event_seq=0, + send=send, + close=close, + ) + assert recovered.attachment_generation == 3 diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py index 21725183a..37e2dc430 100644 --- a/tests/entrypoints/openai_api/test_api_server_guards.py +++ b/tests/entrypoints/openai_api/test_api_server_guards.py @@ -920,7 +920,6 @@ async def test_multistage_app_state_key_snapshot(monkeypatch) -> None: monkeypatch.setattr(api_server, "create_streaming_video_handler", lambda **_k: _marker("streaming_video")) monkeypatch.setattr(api_server, "OpenAIServingRealtime", _FakeCtor) monkeypatch.setattr(api_server, "OmniOpenAIServingVideo", _FakeCtor) - monkeypatch.setattr(api_server, "should_enable_duplex_endpoint", lambda *_a, **_k: False) state = State() await api_server.omni_init_app_state(engine, state, _minimal_args()) diff --git a/tests/entrypoints/openai_api/test_duplex_api_server.py b/tests/entrypoints/openai_api/test_duplex_api_server.py new file mode 100644 index 000000000..b528cbc2a --- /dev/null +++ b/tests/entrypoints/openai_api/test_duplex_api_server.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Guards for the duplex-only API server (unified full-duplex framework). + +A model whose pipeline declares ``duplex_plugin`` is served through +``DuplexOmni`` and nothing else: sessions run over ``/v1/realtime?duplex=1`` +(alias ``/v1/duplex``) and every turn-based route reports "not available". +These tests lock that shape, which ``test_api_server_guards.py`` covers for the +diffusion and multi-stage servers: + +* the pipeline probe that decides duplex vs turn-based at startup, +* the ``app.state`` snapshot of a duplex server (what is wired, what stays + ``None``), +* both duplex websocket routes reaching the session handler, and + ``/v1/realtime`` without the flag still reporting the turn-based Realtime API + as unavailable, +* the startup-warmup gate that holds real clients while the warmup session runs. + +CPU-only: no engine is started; ``DuplexOmni`` is subclassed with just the +accessors the app-state wiring reads. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from starlette.datastructures import State +from starlette.websockets import WebSocketDisconnect + +from vllm_omni.config.config_factory import StageConfigFactory +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.engine.duplex.config import DuplexCapabilities +from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler +from vllm_omni.entrypoints.duplex_omni import DuplexOmni +from vllm_omni.entrypoints.openai import api_server + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# ``app.state`` of a duplex-only server. Presence is not enough: a route whose +# handler is missing answers "not available", so "not wired" is written as None. +_DUPLEX_APP_STATE_KEYS = { + "engine_client", + "log_stats", + "args", + "sleeping_stages", + "stage_configs", + "vllm_config", + "diffusion_engine", + "openai_serving_models", + "serving_tokenization", + "serving_tokens", + "online_renderer", + "openai_serving_chat", + "openai_serving_chat_batch", + "openai_serving_completion", + "openai_serving_responses", + "openai_serving_embedding", + "openai_serving_pooling", + "openai_serving_classification", + "openai_serving_scores", + "openai_serving_transcription", + "openai_serving_translation", + "openai_serving_speech", + "openai_serving_audio_generate", + "openai_serving_video", + "openai_streaming_speech", + "openai_streaming_video", + "openai_streaming_video_output", + "openai_serving_realtime", + "openai_serving_realtime_robot", + "anthropic_serving_messages", + "openai_serving_duplex", + "enable_server_load_tracking", + "server_load_metrics", +} +#: Every turn-based service, plus the Realtime route that is not the duplex one. +#: ``openai_serving_chat`` is absent: a duplex engine also serves chat, through +#: the ordinary turn-based service, because DuplexOmni extends AsyncOmni. +#: ``online_renderer`` is absent for the same reason -- the chat service needs it. +_DUPLEX_MUST_BE_NONE = _DUPLEX_APP_STATE_KEYS - { + "engine_client", + "log_stats", + "args", + "sleeping_stages", + "stage_configs", + "vllm_config", + "openai_serving_models", + "openai_serving_duplex", + "openai_serving_chat", + "online_renderer", + "enable_server_load_tracking", + "server_load_metrics", +} +_DUPLEX_MUST_BE_WIRED = { + "engine_client", + "args", + "stage_configs", + "vllm_config", + "openai_serving_models", + "openai_serving_duplex", + "openai_serving_chat", +} + + +class _FakeDuplexOmni(DuplexOmni): + """A ``DuplexOmni`` without an engine: only what the app-state wiring reads.""" + + def __init__(self) -> None: + self.model = "demo-duplex-model" + self._stage_configs = [object(), object(), object()] + self._vllm_config = SimpleNamespace( + lora_config=None, + model_config=SimpleNamespace(), + parallel_config=SimpleNamespace(_api_process_rank=0), + ) + self._duplex_session_config = DuplexSessionRuntimeConfig() + self._capabilities = DuplexCapabilities(supports_chat_completions=True) + + async def get_vllm_config(self): + return self._vllm_config + + async def get_supported_tasks(self) -> tuple[str, ...]: + return ("generate",) + + async def get_tokenizer(self): + return SimpleNamespace(chat_template="{{ messages }}") + + @property + def stage_configs(self) -> list[object]: + return self._stage_configs + + @property + def duplex_session_config(self) -> DuplexSessionRuntimeConfig: + return self._duplex_session_config + + @property + def duplex_capabilities(self) -> DuplexCapabilities: + return self._capabilities + + # Properties on AsyncOmni, so the fake overrides rather than assigns them. + @property + def model_config(self) -> SimpleNamespace: + return SimpleNamespace() + + @property + def renderer(self) -> SimpleNamespace: + return SimpleNamespace() + + +class _FakeModels: + def __init__(self, *args, **kwargs) -> None: + self.base_model_paths = kwargs.get("base_model_paths") or [] + + +def _minimal_args(**overrides) -> SimpleNamespace: + args = SimpleNamespace( + model="demo-duplex-model", + served_model_name=None, + disable_log_stats=True, + enable_log_requests=False, + max_log_len=None, + enable_server_load_tracking=False, + trust_remote_code=True, + deploy_config=None, + # What the chat service reads when a duplex model also serves chat. + chat_template=None, + chat_template_content_format="auto", + trust_request_chat_template=False, + default_chat_template_kwargs=None, + response_role="assistant", + return_tokens_as_token_ids=False, + enable_auto_tool_choice=False, + exclude_tools_when_tool_choice_none=False, + tool_call_parser=None, + structured_outputs_config=SimpleNamespace(reasoning_parser=None), + enable_prompt_tokens_details=False, + enable_force_include_usage=False, + enable_log_outputs=False, + enable_log_deltas=False, + ) + for key, value in overrides.items(): + setattr(args, key, value) + return args + + +# --------------------------------------------------------------------------- # +# Startup: which entrypoint class serves the model # +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + ("pipeline_config", "expected"), + [ + (SimpleNamespace(duplex_plugin="pkg.mod.Plugin"), True), + (SimpleNamespace(duplex_plugin=None), False), + (SimpleNamespace(duplex_plugin=""), False), + (SimpleNamespace(), False), + (None, False), + ], +) +def test_duplex_model_probe_follows_the_pipeline_plugin(monkeypatch, pipeline_config, expected: bool) -> None: + """``duplex_plugin`` alone decides whether ``serve`` builds ``DuplexOmni``.""" + seen: dict[str, object] = {} + + def fake_get_pipeline_config(*, model, trust_remote_code, deploy_config_path=None, **_kwargs): + seen.update(model=model, trust_remote_code=trust_remote_code, deploy_config_path=deploy_config_path) + return pipeline_config + + monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", fake_get_pipeline_config) + + kwargs = {"trust_remote_code": True, "deploy_config": "deploy.yaml"} + assert api_server._is_duplex_model("demo-duplex-model", kwargs) is expected + assert seen == { + "model": "demo-duplex-model", + "trust_remote_code": True, + "deploy_config_path": "deploy.yaml", + } + + +def test_duplex_model_probe_propagates_resolution_errors(monkeypatch) -> None: + """A pipeline that cannot be resolved must fail startup, not fall back to turn-based.""" + + def fake_get_pipeline_config(**_kwargs): + raise ValueError("needs trust_remote_code=True") + + monkeypatch.setattr(StageConfigFactory, "get_pipeline_config", fake_get_pipeline_config) + + with pytest.raises(ValueError, match="trust_remote_code"): + api_server._is_duplex_model("demo-duplex-model", {}) + + +# --------------------------------------------------------------------------- # +# app.state of a duplex-only server # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_duplex_app_state_wires_only_the_session_surfaces(monkeypatch) -> None: + """Lock the duplex ``app.state``: the two session-backed surfaces are live, every turn-based service is None. + + Fails if a turn-based service is wired into a duplex server (its route + would answer instead of reporting "not available"), or if a key the routes + read disappears. ``openai_serving_chat`` is the one shared name: the route + is served, but by the adapter that runs it on a duplex session, never by + the turn-based chat service. + """ + monkeypatch.setattr(api_server, "OpenAIServingModels", _FakeModels) + monkeypatch.setattr(api_server, "OnlineRenderer", lambda **kwargs: SimpleNamespace(**kwargs)) + monkeypatch.setattr(api_server, "OmniOpenAIServingChat", lambda **kwargs: SimpleNamespace(**kwargs)) + engine = _FakeDuplexOmni() + state = State() + + await api_server.omni_init_app_state(engine, state, _minimal_args()) + + present = {key for key in _DUPLEX_APP_STATE_KEYS if hasattr(state, key)} + assert present == _DUPLEX_APP_STATE_KEYS + not_wired = sorted(key for key in _DUPLEX_MUST_BE_WIRED if getattr(state, key) is None) + assert not not_wired, f"duplex app.state keys registered but not wired: {not_wired}" + unexpectedly_set = sorted(key for key in _DUPLEX_MUST_BE_NONE if getattr(state, key) is not None) + assert not unexpectedly_set, f"turn-based services wired into a duplex server: {unexpectedly_set}" + assert isinstance(state.openai_serving_duplex, OmniDuplexSessionHandler) + assert state.openai_serving_chat is not None + assert state.engine_client is engine + assert state.vllm_config is engine._vllm_config + + +# --------------------------------------------------------------------------- # +# Routes # +# --------------------------------------------------------------------------- # + + +class _RecordingHandler: + """Stands in for ``OmniDuplexSessionHandler`` on the app state.""" + + def __init__(self) -> None: + self.queries: list[dict[str, str]] = [] + + async def handle_realtime_session(self, websocket) -> None: + await websocket.accept() + self.queries.append(dict(websocket.query_params)) + await websocket.send_json({"type": "session.created", "session": {"id": "duplex-test"}}) + await websocket.close() + + +def _duplex_app(handler: object) -> FastAPI: + app = FastAPI() + app.include_router(api_server.router) + app.state.openai_serving_duplex = handler + # A duplex server serves no turn-based Realtime route. + app.state.openai_serving_realtime = None + return app + + +@pytest.mark.parametrize("path", ["/v1/duplex", "/v1/realtime?duplex=1"]) +def test_both_duplex_routes_reach_the_session_handler(path: str) -> None: + """``/v1/duplex`` is an alias of ``/v1/realtime?duplex=1``: same handler, same protocol.""" + handler = _RecordingHandler() + + with TestClient(_duplex_app(handler)) as client: + with client.websocket_connect(path) as websocket: + assert websocket.receive_json()["type"] == "session.created" + with pytest.raises(WebSocketDisconnect): + websocket.receive_text() + + assert len(handler.queries) == 1 + + +@pytest.mark.parametrize("flag", ["1", "true", "on"]) +def test_realtime_route_switches_to_duplex_only_for_the_documented_flag_values(flag: str) -> None: + handler = _RecordingHandler() + + with TestClient(_duplex_app(handler)) as client: + with client.websocket_connect(f"/v1/realtime?duplex={flag}") as websocket: + assert websocket.receive_json()["type"] == "session.created" + with pytest.raises(WebSocketDisconnect): + websocket.receive_text() + + assert handler.queries == [{"duplex": flag}] + + +@pytest.mark.parametrize("query", ["", "?model=openbmb%2FMiniCPM-o-4_5"]) +def test_realtime_without_a_duplex_flag_reaches_the_session_handler(query: str) -> None: + """A bare connection to a duplex deployment is a duplex session. + + ``?duplex=1`` is the spelling the server advertises, but it is not a + requirement: a deployment that declares a duplex plugin answers plain + ``/v1/realtime`` too, so a stock Realtime client needs no vendor query + parameter. + """ + handler = _RecordingHandler() + + with TestClient(_duplex_app(handler)) as client: + with client.websocket_connect(f"/v1/realtime{query}") as websocket: + assert websocket.receive_json()["type"] == "session.created" + with pytest.raises(WebSocketDisconnect): + websocket.receive_text() + + assert len(handler.queries) == 1 + + +@pytest.mark.parametrize("query", ["?duplex=0", "?duplex=false", "?duplex=off"]) +def test_realtime_opted_out_of_duplex_falls_through_to_the_turn_based_route(query: str) -> None: + """An explicit opt-out selects the legacy handler, which a duplex server does not mount.""" + handler = _RecordingHandler() + + with TestClient(_duplex_app(handler)) as client: + with client.websocket_connect(f"/v1/realtime{query}") as websocket: + assert websocket.receive_json() == { + "type": "error", + "error": "Realtime API is not available", + "code": "unsupported", + } + with pytest.raises(WebSocketDisconnect): + websocket.receive_text() + + assert handler.queries == [] + + +# --------------------------------------------------------------------------- # +# Startup warmup gate # +# --------------------------------------------------------------------------- # + + +def _warmup_websocket(warmup_done: asyncio.Event | None, query: dict[str, str]) -> SimpleNamespace: + return SimpleNamespace( + app=SimpleNamespace(state=SimpleNamespace(duplex_warmup_done=warmup_done)), + query_params=query, + ) + + +@pytest.mark.asyncio +async def test_warmup_gate_holds_clients_until_the_warmup_session_finishes() -> None: + warmup_done = asyncio.Event() + websocket = _warmup_websocket(warmup_done, {}) + + waiting = asyncio.create_task(api_server._wait_for_duplex_warmup(websocket)) + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(asyncio.shield(waiting), timeout=0.05) + + warmup_done.set() + await asyncio.wait_for(waiting, timeout=1.0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("warmup_done", "query"), + [ + (None, {}), + (asyncio.Event(), {"vllm_omni_warmup": "1"}), + ], + ids=["no-warmup-configured", "the-warmup-connection-itself"], +) +async def test_warmup_gate_lets_the_warmup_connection_and_plain_servers_through(warmup_done, query) -> None: + websocket = _warmup_websocket(warmup_done, query) + + await asyncio.wait_for(api_server._wait_for_duplex_warmup(websocket), timeout=1.0) + + +@pytest.mark.asyncio +async def test_a_model_that_does_not_declare_chat_completions_does_not_get_the_route(monkeypatch) -> None: + """The decision stays the model's: no capability, no chat service, and the route says so.""" + monkeypatch.setattr(api_server, "OpenAIServingModels", _FakeModels) + engine = _FakeDuplexOmni() + engine._capabilities = DuplexCapabilities(supports_chat_completions=False) + state = State() + + await api_server.omni_init_app_state(engine, state, _minimal_args()) + + assert state.openai_serving_chat is None + assert isinstance(state.openai_serving_duplex, OmniDuplexSessionHandler) diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py deleted file mode 100644 index 0b64e7f14..000000000 --- a/tests/entrypoints/openai_api/test_duplex_handler.py +++ /dev/null @@ -1,9491 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -import base64 -import io -import json -import struct -import wave -from collections.abc import Callable -from contextlib import nullcontext -from types import SimpleNamespace -from typing import Any - -import numpy as np -import pytest -from starlette.websockets import WebSocketDisconnect - -# Import the OpenAI entrypoints package at collection time, before pytest -# attaches its caplog handlers to the root logger. vllm's api_server -# transitively imports model_hosting_container_standards, whose import-time -# configure_root_logger() raises the level of every handler then attached to -# the root logger to ERROR. The serving runtime under test reaches this chain -# lazily (runtime_bridge -> entrypoints.openai.protocol.audio), so without -# this eager import the side effect fires mid-test and silences caplog for -# the remainder of the session. Before the capability predicate moved to -# vllm_omni.entrypoints.duplex.capability, this file imported the package at -# module scope implicitly. -import vllm_omni.entrypoints.openai # noqa: F401 -from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig -from vllm_omni.engine.duplex.control_client import DuplexControlRequestError -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity -from vllm_omni.engine.duplex.messages import DuplexFence, DuplexSessionLifecycleMessage -from vllm_omni.engine.duplex.runtime import duplex_resource_request_id -from vllm_omni.entrypoints.duplex.protocol import ( - DuplexCapabilities, - DuplexOverlapPolicy, - DuplexPlaybackCommitPolicy, - DuplexSession, - DuplexSessionConfig, - ResponseCreateOptions, -) -from vllm_omni.entrypoints.duplex.realtime_session import NativeRealtimeSessionProtocol -from vllm_omni.entrypoints.duplex.runtime_adapter import ServingRuntimeConfigError -from vllm_omni.entrypoints.duplex.serving import ( - OmniDuplexSessionHandler, - should_enable_duplex_endpoint, -) -from vllm_omni.entrypoints.duplex.websocket import DuplexWebSocketActor -from vllm_omni.model_executor.models.minicpmo_4_5.duplex import ( - MiniCPMO45NativeDuplexServingAdapter, - MiniCPMO45PcmAppendBuffer, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.capabilities import ( - minicpmo45_native_capabilities, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.data_plane import ( - MiniCPMO45DataPlaneContext, - MiniCPMO45DataPlaneSession, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.runtime import ( - MiniCPMO45DuplexRuntimeExtension, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.session import ( - MiniCPMO45ServingSessionState, -) -from vllm_omni.model_executor.models.personaplex.duplex.serving_adapter import ( - PersonaPlexServingRuntimeAdapter, -) -from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.outputs.duplex import attach_duplex_output_decision - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -INVALID_VALIDATED_SAMPLE_RATES = ( - pytest.param(16_000.5, id="fractional"), - pytest.param(float("inf"), id="infinite"), - pytest.param(float("nan"), id="nan"), -) - - -def test_native_input_append_supports_explicit_session_opt_out(): - session = DuplexSession( - "sid-explicit-tts", - DuplexSessionConfig(model="test-model"), - capabilities=minicpmo45_native_capabilities(), - ) - - assert OmniDuplexSessionHandler._uses_native_input_append(session) is True - - session.config.extra_body["native_duplex"] = False - assert OmniDuplexSessionHandler._uses_native_input_append(session) is False - - -class _ModelConfig: - model = "test-model" - - -class FakeEngineClient: - output_modalities = ["text", "audio"] - - def __init__( - self, - *, - fail_open: bool = False, - fail_signal: bool = False, - fail_touch: bool = False, - fail_signal_events: set[str] | None = None, - fail_close: bool = False, - fail_abort: bool = False, - control_result: dict[str, object] | None = None, - open_result: dict[str, object] | None = None, - append_result: dict[str, object] | None = None, - append_results: list[dict[str, object]] | None = None, - collect_outputs: list[list[object]] | None = None, - collect_delay_s: float = 0.0, - signal_result: dict[str, object] | None = None, - close_result: dict[str, object] | None = None, - ) -> None: - self.fail_open = fail_open - self.fail_signal = fail_signal - self.fail_touch = fail_touch - self.fail_signal_events = set(fail_signal_events or ()) - self.fail_close = fail_close - self.fail_abort = fail_abort - self.control_result = control_result - self.open_result = open_result - self.append_result = append_result - self.append_results = list(append_results or []) - self.collect_outputs = list(collect_outputs or []) - self.collect_delay_s = collect_delay_s - self.signal_result = signal_result - self.close_result = close_result - self.opened: list[str] = [] - self.opened_fences: list[DuplexFence | None] = [] - self.appended: list[tuple[str, str, Any, bool]] = [] - self.append_operation_ids: list[str | None] = [] - self.appended_fences: list[DuplexFence | None] = [] - self.opened_configs: list[dict[str, Any]] = [] - self.opened_runtime_configs: list[dict[str, Any]] = [] - self.signals: list[tuple[str, str]] = [] - self.signal_fences: list[DuplexFence | None] = [] - self.signal_next_fences: list[DuplexFence | None] = [] - self.signal_session_configs: list[dict[str, Any] | None] = [] - self.signal_runtime_configs: list[dict[str, Any] | None] = [] - self.closed: list[tuple[str, str]] = [] - self.touched: list[tuple[str, DuplexLeaseActivity]] = [] - self.touch_fences: list[DuplexFence] = [] - self.resumed: list[tuple[str, int]] = [] - self.resume_fences: list[DuplexFence] = [] - self.aborted: list[str] = [] - self.abort_batches: list[list[str]] = [] - self.internal_abort_batches: list[list[str]] = [] - self.collected: list[tuple[str, int | None]] = [] - self.duplex_lifecycle_events: asyncio.Queue[DuplexSessionLifecycleMessage] = asyncio.Queue() - - async def open_duplex_session_async( - self, - session_id: str, - *, - session_mode: str = "duplex", - capabilities: dict[str, object] | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - timeout: float | None = None, - fence: DuplexFence | None = None, - ) -> Any: - del timeout - if self.fail_open: - raise RuntimeError("open failed") - self.opened.append(session_id) - self.opened_fences.append(fence) - self.opened_configs.append(dict(session_config or {})) - self.opened_runtime_configs.append(dict(runtime_config or {})) - return self.open_result if self.open_result is not None else self.control_result - - async def append_duplex_input_async( - self, - session_id: str, - *, - mode: str, - payload: object, - operation_id: str | None = None, - final: bool = False, - timeout: float | None = None, - collect_outputs: bool = True, - fence: DuplexFence | None = None, - ) -> Any: - del timeout, collect_outputs - self.appended.append((session_id, mode, payload, final)) - self.append_operation_ids.append(operation_id) - self.appended_fences.append(fence) - if self.append_results: - return self.append_results.pop(0) - return self.append_result if self.append_result is not None else self.control_result - - async def collect_duplex_data_plane_outputs_async( - self, - request_id: str, - *, - response_stage_id: int | None = None, - timeout: float | None = None, - ) -> list[object]: - del timeout - self.collected.append((request_id, response_stage_id)) - if self.collect_delay_s > 0: - await asyncio.sleep(self.collect_delay_s) - if not self.collect_outputs: - return [] - return self.collect_outputs.pop(0) - - async def signal_duplex_turn_async( - self, - session_id: str, - *, - event: str, - timeout: float | None = None, - fence: DuplexFence | None = None, - next_fence: DuplexFence | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - ) -> Any: - del timeout - if self.fail_signal or event in self.fail_signal_events: - raise RuntimeError("signal failed") - self.signals.append((session_id, event)) - self.signal_fences.append(fence) - self.signal_next_fences.append(next_fence) - self.signal_session_configs.append(session_config) - self.signal_runtime_configs.append(runtime_config) - return self.signal_result if self.signal_result is not None else self.control_result - - async def close_duplex_session_async( - self, - session_id: str, - *, - reason: str = "client_close", - timeout: float | None = None, - fence: DuplexFence | None = None, - ) -> Any: - del timeout - if self.fail_close: - raise RuntimeError("close failed") - self.closed.append((session_id, reason)) - return self.close_result if self.close_result is not None else self.control_result - - async def touch_duplex_session_async( - self, - session_id: str, - *, - fence: DuplexFence, - activity: DuplexLeaseActivity, - timeout: float | None = None, - ) -> dict[str, object]: - del timeout - if self.fail_touch: - raise RuntimeError("touch failed") - self.touched.append((session_id, activity)) - self.touch_fences.append(fence) - return { - "operation": "touch", - "session_id": session_id, - "ok": True, - "stage_results": [ - { - "stage_id": -1, - "replica_id": -1, - "result": {"supported": True, "lease_generation": 0}, - } - ], - } - - async def resume_duplex_session_async( - self, - session_id: str, - *, - fence: DuplexFence, - expected_lease_generation: int, - timeout: float | None = None, - ) -> dict[str, object]: - del timeout - self.resumed.append((session_id, expected_lease_generation)) - self.resume_fences.append(fence) - return { - "operation": "resume", - "session_id": session_id, - "ok": True, - "stage_results": [ - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": True, - "lease_generation": expected_lease_generation + 1, - }, - } - ], - } - - async def abort(self, request_ids: list[str]) -> None: - if self.fail_abort: - raise RuntimeError("abort failed") - assert isinstance(request_ids, list) - self.abort_batches.append(list(request_ids)) - self.aborted.extend(request_ids) - - async def _abort_internal_requests(self, request_ids: list[str]) -> None: - if self.fail_abort: - raise RuntimeError("abort failed") - assert isinstance(request_ids, list) - self.internal_abort_batches.append(list(request_ids)) - self.aborted.extend(request_ids) - - -class FakeChatService: - duplex_serving_adapter_path = ( - "vllm_omni.model_executor.models.minicpmo_4_5.duplex.serving_adapter.MiniCPMO45ServingRuntimeAdapter" - ) - - def __init__(self, engine_client: FakeEngineClient) -> None: - self.engine_client = engine_client - self.model_config = _ModelConfig() - self.seen_request_ids: list[str] = [] - - async def create_chat_completion(self, request, raw_request=None): - self.seen_request_ids.append(request.request_id) - - async def _gen(): - await asyncio.sleep(999) - yield "data: [DONE]\n\n" - - return _gen() - - def create_audio(self, audio_obj): - return SimpleNamespace(audio_data=f"wav-{int(audio_obj.audio_tensor.shape[0])}") - - -class TurnBasedFakeChatService(FakeChatService): - """Chat-fallback service without a model-native duplex adapter.""" - - duplex_serving_adapter_path = None - - def __init__(self, engine_client: FakeEngineClient) -> None: - super().__init__(engine_client) - self.requests: list[Any] = [] - - async def create_chat_completion(self, request, raw_request=None): - del raw_request - self.requests.append(request) - self.seen_request_ids.append(request.request_id) - - async def _gen(): - yield "data: [DONE]\n\n" - - return _gen() - - -class GateChatService(FakeChatService): - def __init__(self, engine_client: FakeEngineClient) -> None: - super().__init__(engine_client) - self.first_response_gate = asyncio.Event() - self.response_count = 0 - - async def create_chat_completion(self, request, raw_request=None): - del raw_request - self.seen_request_ids.append(request.request_id) - self.response_count += 1 - response_number = self.response_count - - async def _gen(): - if response_number == 1: - await self.first_response_gate.wait() - yield "data: [DONE]\n\n" - - return _gen() - - -def _test_data_plane() -> MiniCPMO45DataPlaneSession: - def encode_audio( - audio: object, - sample_rate_hz: int, - response_format: str, - speed: float | None, - ) -> str | None: - del sample_rate_hz, response_format, speed - if audio is None: - return None - if hasattr(audio, "numel"): - return f"wav-{int(audio.numel())}" - return f"wav-{int(np.asarray(audio).size)}" - - return MiniCPMO45DataPlaneSession(encode_audio) - - -def _data_plane_context(session: DuplexSession | None = None) -> MiniCPMO45DataPlaneContext: - if session is None: - return MiniCPMO45DataPlaneContext() - return MiniCPMO45DataPlaneContext( - epoch=session.epoch, - turn_id=session.turn_id, - active_response_turn_id=session.active_response_turn_id, - active_response_id=session.active_response_id, - auto_responds=bool(session.config.extra_body.get("auto_response", False)), - response_format=session.config.response_format, - speed=session.config.speed, - modalities=tuple(session.config.modalities), - ) - - -def _project_data_plane( - data_plane: MiniCPMO45DataPlaneSession, - result: object, - *, - session: DuplexSession | None = None, -) -> list[dict[str, Any]]: - return list(data_plane.project(result, context=_data_plane_context(session))) - - -class TimedWebSocket: - def __init__(self, *, on_send=None, receive_timeout_s: float = 1.0): - self._q: asyncio.Queue[str | BaseException] = asyncio.Queue() - self.sent: list[dict[str, Any]] = [] - self.accepted = False - self._on_send = on_send - self._receive_timeout_s = receive_timeout_s - self.query_params: dict[str, str] = {} - - async def accept(self): - self.accepted = True - - async def receive_text(self) -> str: - try: - item = await asyncio.wait_for(self._q.get(), timeout=self._receive_timeout_s) - except asyncio.TimeoutError as exc: - raise WebSocketDisconnect(code=1000) from exc - if isinstance(item, BaseException): - raise item - return item - - async def send_json(self, data: dict[str, Any]): - self.sent.append(data) - if self._on_send is not None: - self._on_send(self, data) - - def put(self, payload: dict[str, Any]) -> None: - self._q.put_nowait(json.dumps(payload)) - - async def close(self, code: int = 1000, reason: str | None = None) -> None: - del reason - self._q.put_nowait(WebSocketDisconnect(code=code)) - - def sent_types(self) -> list[str]: - return [m.get("type", "") for m in self.sent] - - -class FakeServerVADBackend: - frame_samples = 160 - - def __init__(self, probabilities: list[float]) -> None: - self.probabilities = probabilities - self.calls = 0 - - def new_state(self) -> int: - return 0 - - def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: - del frame - self.calls += 1 - index = int(state) - probability = self.probabilities[index] if index < len(self.probabilities) else 0.0 - return probability, index + 1 - - -class FakeServerVADProvider: - def __init__(self, backend: FakeServerVADBackend) -> None: - self.backend = backend - - def get(self) -> FakeServerVADBackend: - return self.backend - - -class FailingServerVADProvider: - def __init__(self) -> None: - self.calls = 0 - - def get(self) -> FakeServerVADBackend: - self.calls += 1 - raise RuntimeError("server VAD backend unavailable") - - -class FailOnceServerVADBackend(FakeServerVADBackend): - def __init__(self, probabilities: list[float]) -> None: - super().__init__(probabilities) - self._failed = False - - def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: - if not self._failed and int(state) == 2: - self._failed = True - raise RuntimeError("server VAD inference failed") - return super().infer(frame, state) - - -def _server_vad_turn_detection(*, create_response: bool = True) -> dict[str, Any]: - return { - "type": "server_vad", - "prefix_padding_ms": 10, - "silence_duration_ms": 20, - "create_response": create_response, - "interrupt_response": False, - } - - -def _server_vad_session_update( - turn_detection: dict[str, Any] | None, - *, - event_id: str | None = None, - sample_rate_hz: int = 16_000, -) -> dict[str, Any]: - event: dict[str, Any] = { - "type": "session.update", - "session": { - "model": "test-model", - "audio": { - "input": { - "format": {"type": "audio/pcm", "rate": sample_rate_hz}, - "turn_detection": turn_detection, - } - }, - }, - } - if event_id is not None: - event["event_id"] = event_id - return event - - -def _server_vad_audio_append( - frame_count: int = 5, - *, - event_id: str | None = None, - is_speech: bool | None = None, - sample_rate_hz: int = 16_000, -) -> dict[str, Any]: - sample_count = frame_count * 160 * sample_rate_hz // 16_000 - event: dict[str, Any] = { - "type": "input_audio_buffer.append", - "audio": base64.b64encode(np.zeros(sample_count, dtype=" None: - if data.get("type") == "response.done": - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=on_send) - if not at_creation: - ws.put(_server_vad_session_update(_server_vad_turn_detection())) - update = _server_vad_session_update(_server_vad_turn_detection(), event_id="bad-vad-audio-format") - update["session"]["audio"]["input"].update(input_update) - ws.put(update) - if not at_creation: - ws.put(_server_vad_audio_append()) - handler = OmniDuplexSessionHandler( - chat_service=TurnBasedFakeChatService(FakeEngineClient()), - server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([0.0, 0.9, 0.8, 0.0, 0.0])), - ) - - await handler.handle_realtime_session(ws) # type: ignore[arg-type] - - (error,) = [event["error"] for event in ws.sent if event.get("type") == "error"] - assert error["code"] == "unsupported_turn_detection" - assert error["event_id"] == "bad-vad-audio-format" - assert "mono PCM16" in error["message"] - assert ("session.created" in ws.sent_types()) is not at_creation - assert ("response.done" in ws.sent_types()) is not at_creation - - -@pytest.mark.asyncio -@pytest.mark.parametrize("sample_rate_hz", [16_000, 24_000]) -@pytest.mark.parametrize("create_response", [False, True], ids=["commit-only", "minimal-config"]) -@pytest.mark.parametrize("at_creation", [True, False], ids=["create", "update"]) -async def test_realtime_server_vad_commits_without_client_commit( - sample_rate_hz: int, create_response: bool, at_creation: bool -): - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - terminal_event = "response.done" if create_response else "input_audio_buffer.committed" - if data.get("type") == terminal_event: - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=on_send) - if not at_creation: - ws.put(_server_vad_session_update(None, sample_rate_hz=sample_rate_hz)) - ws.put( - _server_vad_session_update( - {"type": "server_vad"} if create_response else _server_vad_turn_detection(create_response=False), - sample_rate_hz=sample_rate_hz, - ) - ) - # Cover the default 500 ms silence timeout, plus streaming FIR right context. - ws.put(_server_vad_audio_append(54, sample_rate_hz=sample_rate_hz)) - engine = FakeEngineClient() - chat_service = TurnBasedFakeChatService(engine) - handler = OmniDuplexSessionHandler( - chat_service=chat_service, - server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([0.0, 0.9, 0.8, 0.0, 0.0])), - ) - - await handler.handle_realtime_session(ws) # type: ignore[arg-type] - - event_types = ws.sent_types() - assert "error" not in event_types - effective_session = [ - event["session"] for event in ws.sent if event["type"] in {"session.created", "session.updated"} - ][-1] - assert effective_session["turn_detection"]["interrupt_response"] is False - assert effective_session["audio"]["input"]["turn_detection"] == effective_session["turn_detection"] - assert effective_session["overlap_policy"] == "listen_only" - assert event_types.index("input_audio_buffer.speech_started") < event_types.index( - "input_audio_buffer.speech_stopped" - ) - assert event_types.index("input_audio_buffer.speech_stopped") < event_types.index("input_audio_buffer.committed") - assert event_types.index("input_audio_buffer.committed") < event_types.index("conversation.item.added") - assert event_types.index("conversation.item.added") < event_types.index("conversation.item.done") - assert "conversation.item.created" not in event_types - committed = next(event for event in ws.sent if event.get("type") == "input_audio_buffer.committed") - added = next(event for event in ws.sent if event.get("type") == "conversation.item.added") - done = next(event for event in ws.sent if event.get("type") == "conversation.item.done") - assert committed["item_id"] == added["item"]["id"] == done["item"]["id"] - assert event_types.count("response.created") == int(create_response) - assert len(chat_service.seen_request_ids) == int(create_response) - if create_response: - assert event_types.index("input_audio_buffer.committed") < event_types.index("response.created") - assert event_types.index("response.created") < event_types.index("response.done") - assert engine.opened == [] - assert engine.signals == [] - assert engine.closed == [] - - -@pytest.mark.asyncio -async def test_realtime_server_vad_create_response_false_accepts_response_create(): - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") == "input_audio_buffer.committed": - ws.put({"type": "response.create"}) - if data.get("type") == "response.done": - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=on_send) - ws.put(_server_vad_session_update(_server_vad_turn_detection(create_response=False))) - ws.put(_server_vad_audio_append()) - engine = FakeEngineClient() - chat_service = TurnBasedFakeChatService(engine) - handler = OmniDuplexSessionHandler( - chat_service=chat_service, - server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([0.0, 0.9, 0.8, 0.0, 0.0])), - ) - - await handler.handle_realtime_session(ws) # type: ignore[arg-type] - - event_types = ws.sent_types() - assert event_types.index("input_audio_buffer.committed") < event_types.index("response.created") - assert event_types.index("response.created") < event_types.index("response.done") - assert event_types.count("response.created") == 1 - assert len(chat_service.seen_request_ids) == 1 - assert engine.opened == engine.signals == engine.closed == [] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("sample_rate_hz", [16_000, 24_000]) -async def test_realtime_server_vad_complete_audio_item_bypasses_vad(sample_rate_hz: int): - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - item = data.get("item") - if ( - data.get("type") == "conversation.item.done" - and isinstance(item, dict) - and item.get("id") == "item-complete-audio" - ): - ws.put({"type": "response.create"}) - if data.get("type") == "response.done": - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=on_send) - ws.put(_server_vad_session_update(_server_vad_turn_detection(), sample_rate_hz=sample_rate_hz)) - samples = np.arange(sample_rate_hz // 10, dtype=np.float32) - pcm16 = np.rint(0.5 * np.sin(2 * np.pi * 1_000 * samples / sample_rate_hz) * 32_767).astype(" None: - if data.get("type") == "response.done": - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=on_send) - ws.put(_server_vad_session_update(_server_vad_turn_detection())) - first_append = _server_vad_audio_append(frame_count=2) - first_append["audio"] = base64.b64encode(np.full(320, 12_000, dtype=" None: - nonlocal committed_count, response_count - if data.get("type") == "input_audio_buffer.committed": - committed_count += 1 - if data.get("type") == "error" and data.get("error", {}).get("code") == "input_backpressure": - chat_service.first_response_gate.set() - if data.get("type") == "response.created": - response_count += 1 - if response_count == 2: - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=on_send) - ws.put(_server_vad_session_update(_server_vad_turn_detection())) - for _ in range(3): - ws.put(_server_vad_audio_append()) - handler = OmniDuplexSessionHandler( - chat_service=chat_service, - server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([0.0, 0.9, 0.8, 0.0, 0.0] * 3)), - ) - - await handler.handle_realtime_session(ws) # type: ignore[arg-type] - - assert committed_count == 2 - assert response_count == 2 - event_types = ws.sent_types() - committed_indices = [index for index, event in enumerate(event_types) if event == "input_audio_buffer.committed"] - response_created_indices = [index for index, event in enumerate(event_types) if event == "response.created"] - assert committed_indices[1] < event_types.index("response.done") < response_created_indices[1] - backpressure = [ - event - for event in ws.sent - if event.get("type") == "error" and event.get("error", {}).get("code") == "input_backpressure" - ] - assert len(backpressure) == 1 - - -@pytest.mark.asyncio -async def test_native_session_update_rejects_client_runtime_config(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-runtime-update-reject")) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": { - "extra_body": {"duplex_stage_max_tokens": {"0": 999}}, - }, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - error = next(message for message in ws.sent if message.get("type") == "error") - assert error["code"] == "invalid_duplex_runtime_config" - assert ("sid-runtime-update-reject", "session.update") not in engine.signals - - -@pytest.mark.asyncio -async def test_minicpmo_native_session_update_requires_ref_audio_before_enabling_audio(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-runtime-update-ref-required", modalities=["text"])) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": {"modalities": ["text", "audio"]}, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - error = next(message for message in ws.sent if message.get("type") == "error") - assert error["code"] == "ref_audio_required" - assert "session.updated" not in ws.sent_types() - assert ("sid-runtime-update-ref-required", "session.update") not in engine.signals - - -def test_personaplex_candidate_update_does_not_leak_minicpmo_ref_audio_check(): - handler = OmniDuplexSessionHandler( - chat_service=SimpleNamespace(engine_client=SimpleNamespace()), - served_model_name="test-model", - serving_runtime_adapter=PersonaPlexServingRuntimeAdapter(lambda *_: None), - ) - session = DuplexSession( - session_id="sid-personaplex-voice-change", - config=DuplexSessionConfig(modalities=["audio"], voice="NATF2.pt"), - capabilities=PersonaPlexServingRuntimeAdapter.capabilities(max_sessions=1), - ) - candidate_config = DuplexSessionConfig(modalities=["audio"], voice="OtherVoice.pt") - - error = handler._runtime_session_candidate_update_error(session, candidate_config) - - assert error is None - - -@pytest.mark.parametrize( - ("native", "turn_detection", "expected_code"), - [ - (False, {"type": "server_vad", "interrupt_response": False}, None), - (False, {"type": "server_vad"}, None), - (False, {"type": "server_vad", "interrupt_response": True}, "server_vad_requires_native_duplex"), - (True, {"type": "server_vad", "interrupt_response": False}, "unsupported_turn_detection"), - (True, {"type": "server_vad"}, None), - ], - ids=[ - "turn-based-endpointing", - "turn-based-default-endpointing", - "turn-based-explicit-interruption", - "native-non-interrupting", - "native-default-interruption", - ], -) -def test_server_vad_runtime_capability_matrix( - native: bool, - turn_detection: dict[str, object], - expected_code: str | None, -): - handler = OmniDuplexSessionHandler( - chat_service=TurnBasedFakeChatService(FakeEngineClient()), - ) - capabilities = minicpmo45_native_capabilities() if native else DuplexCapabilities() - session = DuplexSession( - "sid-server-vad-capability", - DuplexSessionConfig(), - capabilities=capabilities, - ) - candidate = DuplexSessionConfig.from_event({"session": {"turn_detection": turn_detection}}) - handler._resolve_server_vad_defaults(session, candidate, {"turn_detection": turn_detection}) - - error = handler._runtime_session_candidate_update_error(session, candidate) - - assert (error or {}).get("code") == expected_code - if expected_code is None: - assert candidate.server_vad is not None - assert candidate.server_vad.interrupt_response is native - assert candidate.overlap_policy == ("barge_in_on_speech" if native else "listen_only") - elif expected_code == "server_vad_requires_native_duplex": - assert "set interrupt_response=false" in error["error"] - - -@pytest.mark.asyncio -async def test_minicpmo_native_session_update_rejects_instructions_after_native_context_locked(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-instructions-locked", modalities=["text"])) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "duration_ms": 1000, - "is_speech": True, - } - ) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": {"instructions": "You are now a pirate."}, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - error = next(message for message in ws.sent if message.get("type") == "error") - assert error["code"] == "instructions_update_unsupported" - assert "session.updated" not in ws.sent_types() - assert ("sid-instructions-locked", "session.update") not in engine.signals - - -@pytest.mark.asyncio -async def test_minicpmo_native_session_update_allows_unchanged_instructions_after_lock(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-instructions-unchanged", modalities=["text"])) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "duration_ms": 1000, - "is_speech": True, - } - ) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": {"instructions": "You are a concise assistant.", "temperature": 0.3}, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert "session.updated" in ws.sent_types() - assert not any( - message.get("type") == "error" and message.get("code") == "instructions_update_unsupported" - for message in ws.sent - ) - - -@pytest.mark.asyncio -async def test_minicpmo_native_session_update_rejects_instructions_before_first_real_append(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-buffered-before-lock", modalities=["text"])) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(3200, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "duration_ms": 200, - "is_speech": True, - } - ) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": { - "instructions": ( - "You are now a swashbuckling pirate captain who speaks entirely in " - "nautical slang, references buried treasure constantly, and never " - "breaks character no matter what the user asks." - ) - }, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True, "response_create": True}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - error = next(message for message in ws.sent if message.get("type") == "error") - assert error["code"] == "instructions_update_unsupported" - assert "session.updated" not in ws.sent_types() - assert engine.appended, "the commit must still flush the buffered chunk despite the rejected update" - - -@pytest.mark.asyncio -async def test_minicpmo_native_session_update_rejects_native_duplex_mode_flip(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-duplex-flag-flip", modalities=["text"])) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": {"extra_body": {"native_duplex": False}}, - } - ) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": {"instructions": "You are now a pirate."}, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - errors = [message for message in ws.sent if message.get("type") == "error"] - assert [error["code"] for error in errors] == [ - "native_duplex_mode_update_unsupported", - "instructions_update_unsupported", - ] - assert "session.updated" not in ws.sent_types() - - -@pytest.mark.parametrize("sample_rate_hz", [22_050, 24_000]) -def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz(sample_rate_hz: int): - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - - payloads = protocol._from_duplex_event( - { - "type": "response.output_audio.delta", - "response_id": "resp-a", - "audio": "AAAA", - "format": "pcm", - "sample_rate_hz": sample_rate_hz, - "audio_duration_ms": 135, - } - ) - - audio_events = [ - payload for payload in payloads if payload["type"] in {"response.output_audio.delta", "response.audio.delta"} - ] - assert {payload["type"] for payload in audio_events} == {"response.output_audio.delta"} - assert {payload["format"] for payload in audio_events} == {"pcm16"} - assert {payload["sample_rate_hz"] for payload in audio_events} == {sample_rate_hz} - assert audio_events[0]["metadata"]["audio_duration_ms"] == 135 - - -def test_native_realtime_protocol_ignores_removed_legacy_event_switches(): - ws = TimedWebSocket() - ws.query_params = { - "legacy_audio_events": "true", - "vllm_omni_legacy_events": "true", - "output_audio_events": "true", - "vllm_omni_output_audio_events": "true", - } - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - - payloads = protocol._from_duplex_event( - { - "type": "response.output_audio.delta", - "response_id": "resp-canonical", - "audio": "AAAA", - "text": "hello", - "format": "pcm", - "sample_rate_hz": 24000, - "end_of_turn": True, - } - ) - - event_types = {payload["type"] for payload in payloads} - assert "response.output_audio.delta" in event_types - assert "response.output_audio.done" in event_types - assert "response.output_audio_transcript.delta" in event_types - assert "response.output_audio_transcript.done" in event_types - assert "response.audio.delta" not in event_types - assert "response.audio.done" not in event_types - assert "response.audio_transcript.delta" not in event_types - assert "response.audio_transcript.done" not in event_types - assert "response.output_item.created" not in event_types - assert "response.text.delta" not in event_types - - -def test_native_realtime_protocol_emits_terminal_audio_transcript_events(): - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - - payloads = protocol._from_duplex_event( - { - "type": "response.output_audio.delta", - "response_id": "resp-a", - "audio": "AAAA", - "text": "hello", - "format": "pcm", - "sample_rate_hz": 24000, - "end_of_turn": True, - } - ) - - by_type = {payload["type"]: payload for payload in payloads} - assert "response.output_audio.done" in by_type - assert "response.output_audio_transcript.delta" in by_type - assert "response.output_audio_transcript.done" in by_type - assert by_type["response.content_part.done"]["part"]["transcript"] == "hello" - assert by_type["response.output_item.done"]["item"]["object"] == "realtime.item" - assert by_type["response.done"]["response"]["output"][0]["object"] == "realtime.item" - - -def test_native_realtime_protocol_terminal_transcript_equals_joined_deltas(): - protocol = NativeRealtimeSessionProtocol(TimedWebSocket()) # type: ignore[arg-type] - response_id = "resp-transcript-contract" - events: list[dict[str, object]] = [] - - for text, end_of_turn in (("It's Canberra.", False), (" Next question.", True)): - events.extend( - protocol._from_duplex_event( - { - "type": "response.output_audio.delta", - "response_id": response_id, - "audio": "AAAA", - "text": text, - "format": "pcm", - "sample_rate_hz": 24000, - "end_of_turn": end_of_turn, - } - ) - ) - - transcript = "".join( - str(event["delta"]) for event in events if event["type"] == "response.output_audio_transcript.delta" - ) - done = [str(event["transcript"]) for event in events if event["type"] == "response.output_audio_transcript.done"] - - assert transcript == "It's Canberra. Next question." - assert done == [transcript] - - -def test_native_realtime_protocol_updates_in_progress_item_for_audio_truncate(): - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - - protocol._from_duplex_event({"type": "response.created", "response_id": "resp-truncate"}) - protocol._from_duplex_event( - { - "type": "response.output_audio.delta", - "response_id": "resp-truncate", - "audio": "AAAA", - "text": "hello", - "format": "pcm", - "sample_rate_hz": 24000, - "audio_duration_ms": 100, - "audio_text_marks": [{"text_chars": 5, "audio_end_ms": 100}], - } - ) - - item = protocol._conversation_items["item_resp-truncate"] - assert item["content"][0]["type"] == "output_audio" - assert item["content"][0]["transcript"] == "hello" - assert item["content"][0]["audio_duration_ms"] == 100 - - translated = asyncio.run( - protocol._to_duplex_event( - { - "type": "conversation.item.truncate", - "item_id": "item_resp-truncate", - "content_index": 0, - "audio_end_ms": 50, - } - ) - ) - - assert translated is None - assert protocol._conversation_items["item_resp-truncate"]["content"][0]["transcript"] == "he" - - -@pytest.mark.asyncio -async def test_native_realtime_protocol_rejects_truncate_for_user_item(): - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - protocol.bind_sender(ws.send_json) - protocol._conversation_items["item-user"] = { - "id": "item-user", - "object": "realtime.item", - "type": "message", - "role": "user", - "status": "completed", - "content": [{"type": "input_audio", "transcript": "hello"}], - } - - translated = await protocol._to_duplex_event( - { - "type": "conversation.item.truncate", - "item_id": "item-user", - "content_index": 0, - "audio_end_ms": 10, - } - ) - - assert translated is None - error = next(event for event in ws.sent if event.get("type") == "error") - assert error["error"]["code"] == "bad_event" - - -def test_native_realtime_protocol_response_done_emits_audio_done_before_terminal_events(): - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - protocol._from_duplex_event({"type": "response.created", "response_id": "resp-done"}) - protocol._from_duplex_event( - { - "type": "response.output_audio.delta", - "response_id": "resp-done", - "audio": "AAAA", - "format": "pcm", - "sample_rate_hz": 24000, - } - ) - - payloads = protocol._from_duplex_event({"type": "response.done", "response_id": "resp-done"}) - event_types = [payload["type"] for payload in payloads] - - assert event_types.index("response.output_audio.done") < event_types.index("response.content_part.done") - assert event_types.index("response.content_part.done") < event_types.index("response.output_item.done") - assert event_types.index("conversation.item.done") < event_types.index("response.done") - assert event_types.index("response.output_item.done") < event_types.index("response.done") - - -def test_native_realtime_protocol_audio_cancelled_uses_active_response_id(): - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - protocol._from_duplex_event({"type": "response.created", "response_id": "resp-cancel"}) - - payloads = protocol._from_duplex_event( - { - "type": "audio.cancelled", - "reason": "barge_in", - "committed_ms": 0, - } - ) - - done = [payload for payload in payloads if payload.get("type") == "response.done"] - assert len(done) == 1 - assert done[0]["response_id"] == "resp-cancel" - assert done[0]["response"]["id"] == "resp-cancel" - - -def test_native_realtime_protocol_audio_cancelled_does_not_reopen_completed_response(): - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - protocol._from_duplex_event({"type": "response.created", "response_id": "resp-complete"}) - protocol._from_duplex_event( - { - "type": "response.output_audio.delta", - "response_id": "resp-complete", - "audio": "AAAA", - "format": "pcm", - "sample_rate_hz": 24000, - } - ) - protocol._from_duplex_event({"type": "response.done", "response_id": "resp-complete"}) - - payloads = protocol._from_duplex_event( - { - "type": "audio.cancelled", - "reason": "session_close", - "committed_ms": 999999, - } - ) - - assert [payload.get("type") for payload in payloads] == [] - - -def test_native_audio_text_marks_are_normalized_to_session_cumulative_offsets(): - marks = OmniDuplexSessionHandler._normalize_native_audio_text_marks( - [{"text_chars": 3, "audio_end_ms": 200}], - audio_offset_ms=800, - text_offset_chars=5, - ) - - assert marks == [{"text_chars": 8, "audio_end_ms": 1000}] - - -def test_duplex_session_playback_commit_uses_multi_delta_audio_text_marks(): - session = DuplexSession( - session_id="sid-marks", - config=DuplexSessionConfig(playback_commit_policy=DuplexPlaybackCommitPolicy.ACK_ONLY.value), - ) - session.begin_response() - session.append_assistant_text("hello ") - session.mark_audio_sent(1000, text_chars=6) - session.append_assistant_text("world") - session.mark_audio_sent(2000, text_chars=11) - session.acknowledge_playback(played_ms=1500, committed_ms=1500) - - committed = session.end_response(commit_text=True) - - assert committed == {"role": "assistant", "content": "hello wo"} - assert session.history == (committed,) - - -def test_duplex_session_preserves_response_history_order_when_user_item_arrives_during_generation(): - session = DuplexSession(session_id="sid-history-order", config=DuplexSessionConfig()) - first_user = {"role": "user", "content": "first turn"} - next_user = { - "role": "user", - "content": [{"type": "audio_url", "audio_url": {"url": "data:audio/wav;base64,AAAA"}}], - } - session.append_history_message(first_user) - session.begin_response() - session.append_assistant_text("first answer") - session.append_history_message(next_user) - - assistant = session.end_response(commit_text=True) - - assert session.history == (first_user, assistant, next_user) - - -@pytest.mark.asyncio -async def test_duplex_session_actor_preserves_wire_order_before_control(): - ws = TimedWebSocket() - actor = DuplexWebSocketActor(ws) - - await actor.enqueue_event({"type": "input_audio_buffer.append", "audio": "AAAA"}) - await actor.enqueue_event({"type": "response.cancel"}) - - first = await actor.next_event() - second = await actor.next_event() - - assert first["type"] == "input_audio_buffer.append" - assert second["type"] == "response.cancel" - - -def _pcm_f32_b64(samples: int, *, value: float = 0.05) -> str: - return base64.b64encode(struct.pack(f"<{samples}f", *([value] * samples))).decode("ascii") - - -def _pcm16_b64(samples: int = 1, *, value: int = 1000) -> str: - return base64.b64encode(struct.pack(f"<{samples}h", *([value] * samples))).decode("ascii") - - -def _session_create(session_id: str = "duplex-test") -> dict[str, Any]: - return { - "type": "session.create", - "session_id": session_id, - "session": { - "model": "test-model", - "modalities": ["text", "audio"], - "idle_timeout_s": 1, - }, - } - - -def _native_session_create( - session_id: str = "duplex-native", - *, - modalities: list[str] | None = None, -) -> dict[str, Any]: - event = _session_create(session_id) - event["session"]["model"] = "openbmb/MiniCPM-o-4_5" - event["session"]["modalities"] = list(modalities or ["text"]) - event["session"]["instructions"] = "You are a concise assistant." - event["session"]["extra_body"] = {"native_duplex": True} - return event - - -def _native_realtime_session_update( - session_id: str, - *, - modalities: list[str] | None = None, -) -> dict[str, Any]: - return { - "type": "session.update", - "session": { - "session_id": session_id, - "model": "openbmb/MiniCPM-o-4_5", - "modalities": list(modalities or ["text"]), - "instructions": "You are a concise assistant.", - "idle_timeout_s": 1, - "extra_body": {"native_duplex": True}, - }, - } - - -def _native_server_vad_update(event_id: str, **session: object) -> dict[str, object]: - return { - "type": "session.update", - "event_id": event_id, - "session": {"turn_detection": {"type": "server_vad"}, **session}, - } - - -def test_native_stage0_request_id_matches_engine_after_same_session_id_reopen(): - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - first = handler._registry.create(session_id="sid-reused") - handler._registry.close(first.session_id) - reopened = handler._registry.create(session_id="sid-reused") - expected = duplex_resource_request_id( - DuplexFence( - reopened.session_id, - epoch=reopened.epoch, - incarnation=reopened.incarnation, - ), - "stage0", - ) - - assert reopened.incarnation == 1 - assert handler._native_stage0_request_id(reopened, reopened.epoch) == expected - - -def _native_audio_payload( - *, - samples: int = 16000, - value: float = 0.05, - is_speech: bool | None = True, - force_listen: bool | None = None, -) -> dict[str, Any]: - payload: dict[str, Any] = { - "type": "audio", - "audio": _pcm_f32_b64(samples, value=value), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - if is_speech is not None: - payload["is_speech"] = is_speech - if force_listen is not None: - payload["force_listen"] = force_listen - return payload - - -def _auto_response_context( - session_id: str, - *, - playback_active: bool = False, -) -> tuple[OmniDuplexSessionHandler, DuplexSession]: - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id=session_id, - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - if playback_active: - session.config.playback_commit_policy = DuplexPlaybackCommitPolicy.ACK_ONLY.value - session.begin_response() - session.mark_audio_sent(duration_ms=1000) - return handler, session - - -def _install_direct_silence_scheduler( - handler: OmniDuplexSessionHandler, - session: DuplexSession, -) -> None: - native = handler._runtime_session_state(session) - - async def _schedule(payload: object, **kwargs: Any) -> bool: - if handler._native_silence_continuation_is_stale( - session, - request_id=kwargs["request_id"], - response_id=kwargs["response_id"], - response_owned=kwargs["response_owned"], - expected_epoch=kwargs["expected_epoch"], - expected_incarnation=kwargs["expected_incarnation"], - expected_model_turn_id=kwargs["expected_model_turn_id"], - ): - return False - append_ok, _ = await handler._append_runtime_input( - session, - payload, - final=False, - send_json=kwargs["send_json"], - mode="append_audio_chunk", - expected_epoch=kwargs["expected_epoch"], - ) - return append_ok - - native.silence_continuation_scheduler = _schedule - - -def test_minicpmo_pcm_append_buffer_flush_preserves_accumulated_speech_marker(): - buffer = MiniCPMO45PcmAppendBuffer() - payload = _native_audio_payload(samples=8000) - - assert buffer.append(payload, chunk_period_ms=1000) is None - flushed = buffer.flush(chunk_period_ms=1000) - - assert flushed is not None - assert flushed["is_speech"] is True - - -def test_minicpmo_pcm_append_buffer_drops_serving_new_user_turn_marker(): - buffer = MiniCPMO45PcmAppendBuffer() - first = _native_audio_payload(samples=8000) - first.update( - new_user_turn=True, - new_user_turn_prefix_variant="clean_response_done", - force_speak=True, - ) - second = _native_audio_payload(samples=8000) - - assert buffer.append(first, chunk_period_ms=1000) is None - emitted = buffer.append(second, chunk_period_ms=1000) - - assert emitted is not None - assert emitted["is_speech"] is True - assert "new_user_turn" not in emitted - assert "new_user_turn_prefix_variant" not in emitted - assert "force_speak" not in emitted - - -def test_minicpmo_merge_native_audio_payloads_preserves_speech_marker(): - first = _native_audio_payload(samples=8000) - second = _native_audio_payload(samples=8000, value=0.0, is_speech=False) - - merged = OmniDuplexSessionHandler._merge_native_audio_payloads(first, second) - - assert merged["is_speech"] is True - - -def test_minicpmo_merge_drops_serving_new_user_turn_marker(): - first = _native_audio_payload(samples=8000) - first.update( - new_user_turn=True, - new_user_turn_prefix_variant="clean_response_done", - force_speak=True, - ) - second = _native_audio_payload(samples=8000, is_speech=False) - - merged = OmniDuplexSessionHandler._merge_native_audio_payloads(first, second) - - assert merged["is_speech"] is True - assert "new_user_turn" not in merged - assert "new_user_turn_prefix_variant" not in merged - assert "force_speak" not in merged - - -@pytest.mark.asyncio -async def test_minicpmo_clear_continuation_does_not_cancel_pending_silence_task(): - async def _pending() -> bool: - await asyncio.sleep(3600) - return True - - native = MiniCPMO45ServingSessionState() - task = asyncio.create_task(_pending()) - native.continuation_owner_id = "owner" - native.continuation_units = 1 - native.pending_silence_task = task - native.pending_silence_owner_id = "owner" - - native.clear_continuation() - - assert native.continuation_owner_id is None - assert native.continuation_units == 0 - assert native.pending_silence_task is None - assert native.pending_silence_owner_id is None - assert not task.cancelled() - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - - -def test_auto_response_playback_overlap_keeps_model_owned_listen_speak_decision(): - handler, session = _auto_response_context("sid-auto-silent-overlap", playback_active=True) - session.mark_audio_sent(duration_ms=1000) - silent_payload = _native_audio_payload(value=0.0, is_speech=None, force_listen=False) - speech_payload = _native_audio_payload(is_speech=None, force_listen=False) - - assert handler._should_force_listen_for_auto_response_overlap(session, {}, silent_payload) is False - assert handler._should_force_listen_for_auto_response_overlap(session, {}, speech_payload) is False - assert ( - handler._should_force_listen_for_auto_response_overlap( - session, - {"force_listen": True}, - speech_payload, - ) - is True - ) - assert ( - handler._should_force_listen_for_auto_response_overlap( - session, - {"force_barge_in": True}, - speech_payload, - ) - is False - ) - - session.acknowledge_playback(played_ms=1000, committed_ms=1000) - assert handler._should_force_listen_for_auto_response_overlap(session, {}, speech_payload) is False - - -def test_legacy_auto_overlap_policy_falls_back_to_listen_only(): - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-auto-policy-fallback", - config=DuplexSessionConfig( - overlap_policy=DuplexSessionConfig._normalize_overlap_policy("auto"), - ), - ) - session.capabilities = minicpmo45_native_capabilities() - - decision = handler._overlap_decision( - session, - {"duration_ms": session.config.overlap_barge_in_ms, "is_speech": True}, - _native_audio_payload(is_speech=True), - ) - - assert session.config.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value - assert decision["action"] == "listen" - assert decision["reason"] == "policy_listen_only" - - -def test_auto_response_overlap_silence_advances_model_unit_and_preserves_realtime_input(): - handler, session = _auto_response_context( - "sid-auto-silent-wire-buffer", - playback_active=True, - ) - decision = handler._overlap_decision( - session, - {"is_speech": False}, - _native_audio_payload(samples=3200, is_speech=False), - ) - - assert decision["action"] == "listen" - assert decision["buffer_audio"] is True - assert decision["defer_runtime_append"] is False - assert decision["force_listen"] is False - assert decision["preserve_realtime_input"] is True - - -def test_auto_response_playback_overlap_admits_model_units_and_tracks_speech(): - handler, session = _auto_response_context("sid-auto-deferred-overlap", playback_active=True) - payload = _native_audio_payload(samples=8640) - - decision = handler._overlap_decision( - session, - {"duration_ms": 540, "is_speech": True}, - payload, - ) - - assert decision["action"] == "listen" - assert decision["defer_runtime_append"] is False - assert decision["force_listen"] is False - assert decision["buffer_audio"] is True - assert session.overlap_speech_ms == 540 - - -@pytest.mark.asyncio -async def test_auto_response_overlap_exact_unit_commit_does_not_block_or_replay_next_unit(): - class ModelUnitEngine(FakeEngineClient): - def __init__(self) -> None: - super().__init__() - self.append_started = asyncio.Event() - self.active_response_ids_at_append: list[str | None] = [] - self.response_done_seen_at_append = False - self.session: DuplexSession | None = None - self.websocket: TimedWebSocket | None = None - - async def append_duplex_input_async(self, session_id: str, **kwargs): - assert self.session is not None - assert self.websocket is not None - self.active_response_ids_at_append.append(self.session.active_response_id) - self.response_done_seen_at_append = "response.done" in self.websocket.sent_types() - kwargs.pop("expected_epoch", None) - result = await super().append_duplex_input_async(session_id, **kwargs) - self.append_started.set() - return result - - engine = ModelUnitEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - engine.websocket = ws - create = _native_session_create("sid-auto-overlap-model-unit") - create["session"]["extra_body"]["auto_response"] = True - ws.put(create) - handler_task = asyncio.create_task(handler.handle_session(ws)) - - try: - for _ in range(100): - session = handler._registry.get("sid-auto-overlap-model-unit") - if session is not None and "session.created" in ws.sent_types(): - break - await asyncio.sleep(0.01) - else: - pytest.fail("native duplex session did not open") - - engine.session = session - session.config.playback_commit_policy = DuplexPlaybackCommitPolicy.ACK_ONLY.value - response_id = session.begin_response() - session.mark_audio_sent(duration_ms=2000) - request_id = handler._native_stage0_request_id(session, session.epoch) - session.bind_request(request_id) - expected_fence = DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ) - - chunk = { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(3200, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "duration_ms": 200, - "is_speech": True, - } - for _ in range(4): - ws.put(chunk) - await asyncio.sleep(0.05) - assert engine.appended == [] - - ws.put(chunk) - await asyncio.wait_for(engine.append_started.wait(), timeout=0.5) - - ws.put({"type": "input_audio_buffer.commit", "final": True}) - for _ in range(100): - if "input.committed" in ws.sent_types(): - break - await asyncio.sleep(0.01) - else: - pytest.fail("exact model unit input did not commit") - - engine.append_started.clear() - for _ in range(5): - ws.put(chunk) - await asyncio.wait_for(engine.append_started.wait(), timeout=0.5) - await asyncio.sleep(0.05) - assert "error" not in ws.sent_types() - finally: - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - assert len(engine.appended) == 2 - for appended_session_id, mode, payload, final in engine.appended: - assert appended_session_id == session.session_id - assert mode == "append_audio_chunk" - assert final is False - assert isinstance(payload, dict) - assert np.frombuffer(base64.b64decode(payload["audio"]), dtype=" None: - nonlocal overlap_sent - if data.get("type") != "response.output_audio.delta" or overlap_sent: - return - overlap_sent = True - chunk = { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(3200, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "duration_ms": 200, - "is_speech": True, - } - for _ in range(4): - ws.put(chunk) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - - ws = TimedWebSocket(on_send=commit_overlap_after_first_response, receive_timeout_s=3) - create = _native_realtime_session_update(session_id) - create["session"]["extra_body"]["auto_response"] = True - ws.put(create) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "duration_ms": 1000, - "is_speech": True, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - - handler_task = asyncio.create_task(handler.handle_realtime_session(ws)) - try: - for _ in range(300): - if len(engine.appended) >= 2: - break - await asyncio.sleep(0.01) - else: - pytest.fail( - f"committed overlap was not promoted after response.done: " - f"appends={len(engine.appended)}, events={ws.sent_types()}" - ) - - assert overlap_sent is True - assert len(engine.appended) == 2 - assert engine.appended[1][3] is True - assert ws.sent_types().count("response.created") == 1 - finally: - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=3) - - assert ws.sent_types().count("response.done") >= 1 - - -def test_auto_response_barge_in_does_not_mark_serving_new_user_turn_payload(): - handler, session = _auto_response_context("sid-auto-barge-new-turn", playback_active=True) - payload = _native_audio_payload() - - assert handler._should_force_listen_for_auto_response_overlap(session, {"force_barge_in": True}, payload) is False - assert "new_user_turn" not in payload - assert "new_user_turn_prefix_variant" not in payload - assert session.turn_id == 0 - assert payload.get("force_speak") is not True - - -def test_native_realtime_barge_in_payload_does_not_create_serving_turn_boundary(): - payload: dict[str, object] = { - "type": "audio", - "audio": _pcm_f32_b64(16000, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "is_speech": True, - } - session = DuplexSession(session_id="sid-realtime-barge", config=DuplexSessionConfig()) - session.turn_id = 2 - - assert "new_user_turn" not in payload - assert "new_user_turn_prefix_variant" not in payload - assert session.turn_id == 2 - - -def test_auto_response_force_barge_in_has_no_waiting_turn_variant(): - handler, session = _auto_response_context("sid-auto-barge-keeps-variant") - payload = _native_audio_payload() - - assert ( - handler._should_force_listen_for_auto_response_overlap( - session, - {"force_barge_in": True}, - payload, - ) - is False - ) - - -def test_duplex_endpoint_requires_explicit_realtime_configuration(): - assert should_enable_duplex_endpoint(None) is False - assert should_enable_duplex_endpoint([]) is False - assert should_enable_duplex_endpoint([SimpleNamespace(session_mode="turn")]) is False - assert should_enable_duplex_endpoint([{"session_mode": "turn"}]) is False - assert ( - should_enable_duplex_endpoint( - [ - SimpleNamespace(session_mode="turn"), - SimpleNamespace(session_mode="duplex"), - ] - ) - is True - ) - assert should_enable_duplex_endpoint([{"session_mode": "duplex"}]) is True - - -def test_duplex_endpoint_supports_top_level_session_mode(tmp_path): - config_path = tmp_path / "stage_config.yaml" - config_path.write_text( - """ -session_mode: duplex -stage_args: - - stage_id: 0 - engine_args: {} -""", - encoding="utf-8", - ) - - assert should_enable_duplex_endpoint([], config_path=str(config_path)) is True - - -def test_duplex_endpoint_supports_inherited_duplex_session(tmp_path): - base_path = tmp_path / "base.yaml" - base_path.write_text( - """ -session_mode: turn -duplex_session: {} -stages: [] -""", - encoding="utf-8", - ) - config_path = tmp_path / "deploy.yaml" - config_path.write_text("base_config: base.yaml\n", encoding="utf-8") - - assert should_enable_duplex_endpoint([], config_path=str(config_path)) is True - - -def test_duplex_handler_splits_data_plane_audio_list_into_deltas(): - import torch - - data_plane = _test_data_plane() - output = SimpleNamespace( - finished=True, - outputs=[ - SimpleNamespace( - text="hello", - multimodal_output={ - "audio": [ - torch.zeros(10, dtype=torch.float32), - torch.zeros(20, dtype=torch.float32), - ], - "sr": 24000, - }, - ) - ], - ) - - native_results = list(data_plane.project_output(output)) - - assert [result["audio_data"] for result in native_results] == ["wav-10", "wav-20"] - assert [result["text"] for result in native_results] == ["hello", ""] - assert [result["end_of_turn"] for result in native_results] == [False, True] - - -def test_duplex_listen_latent_does_not_poison_cumulative_audio_offset(): - import torch - - data_plane = _test_data_plane() - request_id = "duplex-duplex-sess-stage0" - - # A model-listen decision wraps the segment with a latent tensor that is - # NOT reply audio; it must not advance the cumulative audio offset. - listen_output = SimpleNamespace( - request_id=request_id, - finished=True, - outputs=[], - multimodal_output={ - "duplex_native_decision": "listen", - "model_listen": True, - "latent": torch.zeros(331776, dtype=torch.float32), - "meta": {"sr": 24000}, - }, - ) - listen_results = list(data_plane.project_output(listen_output)) - assert [result.get("is_listen") for result in listen_results] == [True] - - # The first speak unit carries cumulative stage-1 audio far smaller than - # the listen latent; it must still be delivered from sample 0. - speak_output = SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[ - SimpleNamespace( - text=" It was a very", - multimodal_output={}, - ) - ], - multimodal_output={ - "audio": torch.zeros(32768, dtype=torch.float32), - "sr": 24000, - }, - ) - speak_results = list(data_plane.project_output(speak_output)) - assert [result.get("audio_data") for result in speak_results] == ["wav-32768"] - assert speak_results[0]["text"] == " It was a very" - - -def test_direct_listen_decision_survives_inner_completion_metadata(): - inner_output = SimpleNamespace( - outputs=[ - SimpleNamespace( - multimodal_output={"special_token_ids": {"listen_token_id": 151705}}, - ) - ] - ) - decision = MiniCPMO45DuplexRuntimeExtension().decide_output( - stage_id=0, - final_stage_id=1, - segment_finished=True, - segment_token_ids=(151705,), - segment_output_metadata={"special_token_ids": {"listen_token_id": 151705}}, - output=inner_output, - ) - assert decision is not None - output = attach_duplex_output_decision( - OmniRequestOutput.from_stage_output( - inner_output, - request_id="duplex-direct-listen", - finished=True, - stage_id=0, - final_output_type=decision.final_output_type, - metrics={ - "stage_metrics": { - "0": { - "vllm_ttft_ms": 125.0, - "vllm_tpot_ms": 18.0, - "vllm_itl_ms": 17.5, - "vllm_itls_ms": [17.0, 18.0], - } - } - }, - ), - decision, - ) - - results = list(_test_data_plane().project_output(output)) - - assert len(results) == 1 - assert results[0]["is_listen"] is True - assert results[0]["model_listen"] is True - assert results[0]["listen_source"] == "model_listen" - assert results[0]["stage_metrics"] == output.metrics["stage_metrics"] - - -def test_duplex_segment_text_is_attached_once_across_streaming_batches(): - import torch - - class _ChatService: - model_config = _ModelConfig() - - def create_audio(self, audio_obj): - return SimpleNamespace(audio_data=f"wav-{int(audio_obj.audio_tensor.shape[0])}") - - data_plane = _test_data_plane() - request_id = "duplex-duplex-sess-stage0" - session = DuplexSession( - session_id="sid-auto-respond", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - - def _speak(total_samples: int, text: str, *, finished: bool): - output = SimpleNamespace( - request_id=request_id, - finished=finished, - outputs=[SimpleNamespace(text=text, multimodal_output={})], - multimodal_output={ - "audio": torch.zeros(total_samples, dtype=torch.float32), - "sr": 24000, - }, - ) - return _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - def _audio(results): - return [r for r in results if r.get("is_listen") is False] - - def _listen(results): - return [r for r in results if r.get("is_listen") is True] - - # One talker segment streams several cumulative-audio batches, each - # carrying the SAME segment text; only the first may attach it. - assert [r["text"] for r in _audio(_speak(100, " movie called", finished=False))] == [" movie called"] - assert [r["text"] for r in _audio(_speak(220, " movie called", finished=False))] == [""] - finished_results = _speak(300, " movie called", finished=True) - assert [r["text"] for r in _audio(finished_results)] == [""] - assert _listen(finished_results) == [] - - # Continuation units re-run the talker with the SAME handed text past a - # finished boundary (every engine segment ends finished=True); the text - # must stay suppressed or each continuation duplicates the transcript. - assert [r["text"] for r in _audio(_speak(400, " movie called", finished=False))] == [""] - assert [r["text"] for r in _audio(_speak(450, " movie called", finished=True))] == [""] - - # A genuinely new segment text is attached in full. - assert [r["text"] for r in _audio(_speak(500, " Titanic", finished=True))] == [" Titanic"] - - # Text growing within a segment is delivered as suffix deltas. - assert [r["text"] for r in _audio(_speak(600, " by", finished=False))] == [" by"] - assert [r["text"] for r in _audio(_speak(700, " by James", finished=True))] == [" James"] - - # A segment whose finished batch slices to an EMPTY audio delta (all - # samples already delivered) must not block the next segment's text. - assert [r["text"] for r in _audio(_speak(800, " James Cameron. The", finished=False))] == [" James Cameron. The"] - assert _speak(800, " James Cameron. The", finished=True) == [] - assert [r["text"] for r in _audio(_speak(900, " 997.", finished=False))] == [" 997."] - - -def test_duplex_data_plane_output_prefers_audio_segment_text_metadata(): - data_plane = _test_data_plane() - request_id = "duplex-sid-native-segment-text-e0-stage0" - segment_text = "你好，有什么想聊的吗？" - cumulative_text = segment_text * 2 - output = SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[SimpleNamespace(text=cumulative_text, multimodal_output={})], - multimodal_output={ - "audio": np.zeros(24000, dtype=np.float32), - "sr": 24000, - "meta.llm_output_text_utf8": np.frombuffer(segment_text.encode("utf-8"), dtype=np.uint8), - }, - ) - - results = list(data_plane.project_output(output)) - - assert len(results) == 1 - assert results[0]["audio_data"] == "wav-24000" - assert results[0]["text"] == segment_text - - next_output_without_segment_metadata = SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[SimpleNamespace(text=segment_text * 3, multimodal_output={})], - multimodal_output={ - "audio": np.zeros(30720, dtype=np.float32), - "sr": 24000, - }, - ) - next_results = list(data_plane.project_output(next_output_without_segment_metadata)) - - assert len(next_results) == 1 - assert next_results[0]["audio_data"] == "wav-6720" - assert next_results[0]["text"] == "" - - -@pytest.mark.parametrize( - ("request_id", "second_samples", "expected_second_audio"), - [ - ("duplex-sid-model-turn-audio-e0-stage0", 36000, "wav-12000"), - ("duplex-sid-model-turn-audio-restart-e0-stage0", 12000, "wav-12000"), - ], - ids=["cumulative-across-model-turns", "upstream-cumulative-audio-restarts"], -) -def test_duplex_audio_cursor_across_model_turns( - request_id: str, - second_samples: int, - expected_second_audio: str, -): - data_plane = _test_data_plane() - - def output(turn_id: int, samples: int): - return SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[SimpleNamespace(text="", multimodal_output={})], - multimodal_output={ - "audio": np.zeros(samples, dtype=np.float32), - "sr": 24000, - "meta.duplex_turn_id": np.array([turn_id], dtype=np.int32), - }, - ) - - first = list(data_plane.project_output(output(0, 24000))) - second = list(data_plane.project_output(output(1, second_samples))) - - assert first[0]["audio_data"] == "wav-24000" - assert second[0]["audio_data"] == expected_second_audio - - -def test_duplex_data_plane_accepts_model_turn_output_across_client_commits(): - data_plane = _test_data_plane() - session = DuplexSession( - session_id="sid-auto-stale-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.commit_native_audio_input() - session.commit_native_audio_input() - session.turn_id = 1 - session.epoch = 0 - request_id = "duplex-sid-auto-stale-turn-stage0" - output = SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[SimpleNamespace(text="old tail", multimodal_output={})], - multimodal_output={ - "audio": np.zeros(24000, dtype=np.float32), - "sr": 24000, - "meta.duplex_turn_id": np.array([1], dtype=np.int32), - "meta.llm_output_text_utf8": np.frombuffer("旧尾巴".encode(), dtype=np.uint8), - }, - ) - - results = list(data_plane.project_output(output, context=_data_plane_context(session))) - - assert len(results) == 1 - assert results[0]["audio_data"] == "wav-24000" - assert results[0]["model_turn_id"] == 1 - - -def test_duplex_data_plane_accepts_active_response_turn_identity(): - data_plane = _test_data_plane() - session = DuplexSession( - session_id="sid-active-response-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.turn_id = 1 - session.epoch = 0 - session.begin_response(turn_id=0) - output = SimpleNamespace( - request_id="duplex-sid-active-response-turn-stage0", - finished=True, - outputs=[SimpleNamespace(text="", multimodal_output={})], - multimodal_output={ - "audio": np.zeros(10, dtype=np.float32), - "sr": 24000, - "meta.duplex_turn_id": np.array([0], dtype=np.int32), - "meta.duplex_epoch": np.array([0], dtype=np.int32), - }, - ) - - results = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert len(results) == 1 - - -def test_duplex_data_plane_listen_preserves_model_turn_identity(): - data_plane = _test_data_plane() - session = DuplexSession( - session_id="sid-listen-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - output = SimpleNamespace( - request_id="duplex-sid-listen-turn-stage0", - finished=True, - outputs=[SimpleNamespace(text="", token_ids=[], multimodal_output={})], - multimodal_output={ - "duplex_native_decision": "listen", - "meta.duplex_turn_id": np.array([2], dtype=np.int32), - "meta.duplex_epoch": np.array([0], dtype=np.int32), - }, - ) - - results = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert len(results) == 1 - assert results[0]["is_listen"] is True - assert results[0]["model_turn_id"] == 2 - - -def test_duplex_data_plane_does_not_drop_future_model_turn_while_response_active(): - data_plane = _test_data_plane() - session = DuplexSession( - session_id="sid-future-response-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.turn_id = 0 - session.epoch = 0 - session.begin_response(turn_id=0) - output = SimpleNamespace( - request_id="duplex-sid-future-response-turn-stage0", - finished=False, - outputs=[SimpleNamespace(text="", multimodal_output={})], - multimodal_output={ - "audio": np.zeros(10, dtype=np.float32), - "sr": 24000, - "meta.duplex_turn_id": np.array([1], dtype=np.int32), - "meta.duplex_epoch": np.array([0], dtype=np.int32), - }, - ) - - results = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert len(results) == 1 - assert results[0]["model_turn_id"] == 1 - assert results[0]["audio_data"] == "wav-10" - - -def _duplex_tts_output( - *, - request_id: str, - samples: int, - finished: bool, - text: str = "hello", - tts_is_last_chunk: bool = False, - turn_end: bool = False, - turn_id: int = 0, - token_ids: list[int] | None = None, -): - return SimpleNamespace( - request_id=request_id, - finished=finished, - outputs=[SimpleNamespace(text=text, token_ids=list(token_ids or []), multimodal_output={})], - multimodal_output={ - "audio": np.zeros(samples, dtype=np.float32), - "sr": 24000, - "meta.tts_is_last_chunk": np.array([int(tts_is_last_chunk)], dtype=np.int32), - "meta.turn_end": np.array([int(turn_end)], dtype=np.int32), - "meta.duplex_turn_id": np.array([turn_id], dtype=np.int32), - "meta.duplex_epoch": np.array([0], dtype=np.int32), - }, - ) - - -@pytest.mark.parametrize( - ( - "finished", - "tts_is_last_chunk", - "turn_end", - "consume_audio_first", - "expected_end_of_turn", - "expected_audio", - ), - [ - (False, True, False, False, False, "wav-24000"), - (True, True, False, True, False, ""), - (True, True, True, True, True, ""), - ], - ids=[ - "last-audio-batch-does-not-end-response", - "scheduler-eos-only-ends-tts-segment", - "explicit-turn-end-metadata-ends-response", - ], -) -def test_duplex_auto_response_tts_boundary( - finished: bool, - tts_is_last_chunk: bool, - turn_end: bool, - consume_audio_first: bool, - expected_end_of_turn: bool, - expected_audio: str, -): - data_plane = _test_data_plane() - request_id = "duplex-sid-tts-boundary-e0-stage0" - session = DuplexSession( - session_id="sid-tts-boundary", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.bind_response_turn(0) - if consume_audio_first: - data_plane.slice_cumulative_audio(request_id, np.zeros(24000, dtype=np.float32)) - output = _duplex_tts_output( - request_id=request_id, - samples=24000, - finished=finished, - tts_is_last_chunk=tts_is_last_chunk, - turn_end=turn_end, - token_ids=[151645] if finished else None, - ) - - results = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert len(results) == 1 - boundary = results[0] - assert boundary["stage_role"] == "tts" - assert boundary["end_of_turn"] is expected_end_of_turn - assert boundary["abort_data_plane_request"] is True - assert boundary["audio_data"] == expected_audio - - -def test_duplex_auto_response_tts_scheduler_eos_fallback(): - data_plane = _test_data_plane() - request_id = "duplex-sid-tts-eos-no-profile-e0-stage0" - session = DuplexSession( - session_id="sid-tts-eos-no-profile", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.bind_response_turn(0) - data_plane.slice_cumulative_audio(request_id, np.zeros(24000, dtype=np.float32)) - output = _duplex_tts_output( - request_id=request_id, - samples=24000, - finished=True, - tts_is_last_chunk=False, - token_ids=[151645], - ) - - results = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert len(results) == 1 - assert results[0]["stage_role"] == "tts" - assert results[0]["abort_data_plane_request"] is True - - -def test_duplex_auto_response_rearms_tts_boundary_for_next_segment_in_same_turn(): - data_plane = _test_data_plane() - request_id = "duplex-sid-next-segment-same-turn-e0-stage0" - session = DuplexSession( - session_id="sid-next-segment-same-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.bind_response_turn(0) - data_plane.begin_request(request_id) - - first_segment = _duplex_tts_output( - request_id=request_id, - samples=24000, - finished=False, - tts_is_last_chunk=True, - turn_id=0, - ) - first_results = _project_data_plane( - data_plane, - {"data_plane_outputs": [first_segment]}, - session=session, - ) - - assert len(first_results) == 1 - assert first_results[0]["abort_data_plane_request"] is True - - data_plane.begin_request(request_id) - second_segment = _duplex_tts_output( - request_id=request_id, - samples=48000, - finished=False, - text="hello again", - tts_is_last_chunk=True, - turn_id=0, - ) - second_results = _project_data_plane( - data_plane, - {"data_plane_outputs": [second_segment]}, - session=session, - ) - - assert len(second_results) == 1 - assert second_results[0]["audio_data"] == "wav-24000" - assert second_results[0]["abort_data_plane_request"] is True - - -def test_duplex_turn_end_is_not_swallowed_by_finished_segment_fallback(): - data_plane = _test_data_plane() - request_id = "duplex-sid-turn-end-after-segment-e0-stage0" - session = DuplexSession( - session_id="sid-turn-end-after-segment", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.begin_response(turn_id=0) - session.bind_request(request_id) - _project_data_plane( - data_plane, - { - "data_plane_outputs": [ - _duplex_tts_output( - request_id=request_id, - samples=24000, - finished=False, - tts_is_last_chunk=True, - ) - ] - }, - session=session, - ) - output = _duplex_tts_output( - request_id=request_id, - samples=24000, - finished=True, - text="", - tts_is_last_chunk=False, - turn_end=True, - token_ids=[151645], - ) - - results = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert len(results) == 1 - assert results[0]["stage_role"] == "tts" - assert results[0]["end_of_turn"] is True - - duplicate = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert not any(result.get("end_of_turn") is True for result in duplicate) - - -def test_duplex_auto_response_discards_terminal_only_audio_before_response_creation(): - data_plane = _test_data_plane() - request_id = "duplex-sid-terminal-only-audio-e0-stage0" - session = DuplexSession( - session_id="sid-terminal-only-audio", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.bind_request(request_id) - first_audio = _duplex_tts_output( - request_id=request_id, - samples=3840, - finished=False, - text="", - ) - - assert _project_data_plane(data_plane, {"data_plane_outputs": [first_audio]}, session=session) == [] - - terminal = _duplex_tts_output( - request_id=request_id, - samples=3840, - finished=True, - text="", - tts_is_last_chunk=True, - turn_end=True, - token_ids=[151645], - ) - results = _project_data_plane(data_plane, {"data_plane_outputs": [terminal]}, session=session) - - assert len(results) == 1 - assert results[0]["audio_data"] == "" - assert results[0]["end_of_turn"] is True - assert session.active_request_id == request_id - assert not data_plane.is_terminal(request_id) - - -def test_duplex_auto_response_releases_buffered_audio_when_transcript_arrives(): - data_plane = _test_data_plane() - request_id = "duplex-sid-delayed-transcript-e0-stage0" - session = DuplexSession( - session_id="sid-delayed-transcript", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - audio_before_text = _duplex_tts_output( - request_id=request_id, - samples=100, - finished=False, - text="", - ) - - assert _project_data_plane(data_plane, {"data_plane_outputs": [audio_before_text]}, session=session) == [] - - audio_with_text = _duplex_tts_output( - request_id=request_id, - samples=200, - finished=False, - text="hello", - ) - results = _project_data_plane(data_plane, {"data_plane_outputs": [audio_with_text]}, session=session) - - assert [result["audio_data"] for result in results] == ["wav-100", "wav-100"] - assert [result["text"] for result in results] == ["", "hello"] - - -def test_duplex_auto_response_releases_buffered_audio_on_text_only_terminal(): - data_plane = _test_data_plane() - request_id = "duplex-sid-delayed-terminal-text-e0-stage0" - session = DuplexSession( - session_id="sid-delayed-terminal-text", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - audio_before_text = _duplex_tts_output( - request_id=request_id, - samples=100, - finished=False, - text="", - turn_id=1, - ) - - assert _project_data_plane(data_plane, {"data_plane_outputs": [audio_before_text]}, session=session) == [] - - textless_turn_terminal = _duplex_tts_output( - request_id=request_id, - samples=100, - finished=True, - text="", - tts_is_last_chunk=True, - turn_end=True, - turn_id=1, - token_ids=[151645], - ) - terminal_results = list( - data_plane.project( - {"data_plane_outputs": [textless_turn_terminal]}, - context=_data_plane_context(session), - ) - ) - - assert len(terminal_results) == 1 - assert terminal_results[0]["audio_data"] == "" - assert terminal_results[0]["end_of_turn"] is True - - text_only_terminal = SimpleNamespace( - request_id=request_id, - finished=True, - outputs=[SimpleNamespace(text="hello", token_ids=[151645], multimodal_output={})], - multimodal_output={ - "sr": 24000, - "meta.tts_is_last_chunk": np.array([1], dtype=np.int32), - "meta.turn_end": np.array([1], dtype=np.int32), - "meta.duplex_turn_id": np.array([2], dtype=np.int32), - "meta.duplex_epoch": np.array([0], dtype=np.int32), - }, - ) - results = _project_data_plane(data_plane, {"data_plane_outputs": [text_only_terminal]}, session=session) - - assert len(results) == 1 - assert results[0]["audio_data"] == "wav-100" - assert results[0]["text"] == "hello" - assert results[0]["end_of_turn"] is True - assert results[0]["abort_data_plane_request"] is True - assert not data_plane.has_pending_audio(request_id) - - -@pytest.mark.asyncio -async def test_native_append_propagates_current_turn_fence_to_engine(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession( - session_id="sid-fenced-append", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.turn_id = 2 - session.begin_response(turn_id=2) - ws = TimedWebSocket() - - await handler._append_runtime_input( - session, - {"duplex_turn_id": 2, "audio": "", "format": "pcm_f32le"}, - final=True, - send_json=ws.send_json, - mode="append_audio_chunk", - expected_epoch=0, - ) - - assert engine.appended_fences == [DuplexFence("sid-fenced-append", epoch=0, turn_id=2)] - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_tts_segment_boundary_appends_silence_unit(): - request_id = "duplex-sid-segment-boundary-e0-stage0" - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-segment-boundary", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.begin_response(turn_id=0) - session.turn_id = 1 - session.bind_request(request_id) - session.mark_audio_sent(100) - _install_direct_silence_scheduler(handler, session) - ws = TimedWebSocket() - - await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": request_id, - "text": "", - "audio_data": "", - "audio_format": "pcm16", - "end_of_turn": False, - "abort_data_plane_request": True, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - await asyncio.sleep(0.01) - - assert len(engine.appended) == 1 - _, mode, payload, final = engine.appended[0] - assert mode == "append_audio_chunk" - assert payload["format"] == "pcm_f32le" - assert payload["duplex_turn_id"] == 0 - assert "force_speak" not in payload - assert final is False - assert session.active_response_id is not None - assert "response.done" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_pre_response_tts_boundary_continues_model_turn(): - request_id = "duplex-sid-pre-response-boundary-e0-stage0" - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession( - session_id="sid-pre-response-boundary", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.bind_request(request_id) - _install_direct_silence_scheduler(handler, session) - ws = TimedWebSocket() - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": request_id, - "model_turn_id": session.turn_id, - "text": "", - "audio_data": "", - "audio_format": "pcm16", - "end_of_turn": False, - "abort_data_plane_request": True, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - await asyncio.sleep(0.01) - - assert close_reason is None - assert emitted is False - assert session.active_response_id is None - assert ws.sent == [] - assert len(engine.appended) == 1 - _, mode, payload, final = engine.appended[0] - assert mode == "append_audio_chunk" - assert payload["duplex_turn_id"] == session.turn_id - assert final is False - - -@pytest.mark.asyncio -async def test_minicpmo_pre_response_continuation_drops_after_model_turn_ends(): - request_id = "duplex-sid-stale-pre-response-boundary-e0-stage0" - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession( - session_id="sid-stale-pre-response-boundary", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.bind_request(request_id) - model_turn_id = session.turn_id - - native = handler._runtime_session_state(session) - - async def _stale_before_append(payload: object, **kwargs: Any) -> bool: - session.complete_model_turn(model_turn_id) - if handler._native_silence_continuation_is_stale( - session, - request_id=kwargs["request_id"], - response_id=kwargs["response_id"], - response_owned=kwargs["response_owned"], - expected_epoch=kwargs["expected_epoch"], - expected_incarnation=kwargs["expected_incarnation"], - expected_model_turn_id=kwargs["expected_model_turn_id"], - ): - return False - append_ok, _ = await handler._append_runtime_input( - session, - payload, - final=False, - send_json=kwargs["send_json"], - mode="append_audio_chunk", - expected_epoch=kwargs["expected_epoch"], - ) - return append_ok - - native.silence_continuation_scheduler = _stale_before_append - - await handler._maybe_continue_native_response( - TimedWebSocket().send_json, - session=session, - expected_epoch=session.epoch, - expected_model_turn_id=model_turn_id, - ) - await asyncio.sleep(0.01) - - assert engine.appended == [] - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_continuation_stops_at_large_safety_boundary(): - request_id = "duplex-sid-long-response-e0-stage0" - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession( - session_id="sid-long-response", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.begin_response(turn_id=0) - session.turn_id = 1 - session.bind_request(request_id) - session.mark_audio_sent(100) - _install_direct_silence_scheduler(handler, session) - ws = TimedWebSocket() - - for _ in range(handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS): - await handler._maybe_continue_native_response( - ws.send_json, - session=session, - expected_epoch=session.epoch, - ) - await asyncio.sleep(0.02) - - response_id = session.active_response_id - assert response_id is not None - assert len(engine.appended) == handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS - assert all(payload["duplex_turn_id"] == 0 for _, _, payload, _ in engine.appended) - assert all(not {"force_speak", "force_listen"} & payload.keys() for _, _, payload, _ in engine.appended[:-1]) - assert engine.appended[-1][2]["force_listen"] is True - assert handler._native_response_continuations_remaining(session, response_id) is False - - await handler._maybe_continue_native_response( - ws.send_json, - session=session, - expected_epoch=session.epoch, - ) - - assert len(engine.appended) == handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS - assert ws.sent_types()[-2:] == ["response.listen", "response.done"] - assert session.active_response_id is None - assert session.active_request_id == request_id - assert session.turn_id == 1 - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_boundary_listen_closes_response(): - request_id = "duplex-sid-boundary-listen-e0-stage0" - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession( - session_id="sid-boundary-listen", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - response_id = session.begin_response(turn_id=0) - session.bind_request(request_id) - native = handler._runtime_session_state(session) - native.continuation_owner_id = f"response:{response_id}" - native.continuation_units = handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS - _install_direct_silence_scheduler(handler, session) - ws = TimedWebSocket() - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "llm", - "is_listen": True, - "model_listen": True, - "data_plane_request_id": request_id, - "end_of_turn": False, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - - assert close_reason is None - assert emitted is True - assert engine.appended == [] - assert ws.sent_types() == ["response.listen", "response.done"] - # The listen terminates the precreated response: clients demultiplex the - # decision by its id, so the payload must carry the response identity. - listen_payload, done_payload = ws.sent[-2], ws.sent[-1] - assert listen_payload["response_id"] == response_id - assert done_payload["response_id"] == response_id - assert session.active_response_id is None - assert session.active_request_id == request_id - assert session.turn_id == 1 - assert native.continuation_owner_id is None - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_boundary_allows_next_model_turn_response(): - request_id = "duplex-sid-boundary-next-turn-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-boundary-next-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.begin_response(turn_id=0) - session.bind_request(request_id) - native = handler._runtime_session_state(session) - native.continuation_owner_id = f"response:{session.active_response_id}" - native.continuation_units = handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS - ws = TimedWebSocket() - - await handler._maybe_continue_native_response( - ws.send_json, - session=session, - expected_epoch=session.epoch, - ) - - assert session.active_response_id is None - assert session.turn_id == 1 - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": request_id, - "text": "next response", - "audio_data": "next-audio", - "audio_format": "pcm16", - "audio_duration_ms": 100, - "end_of_turn": False, - "model_turn_id": 1, - }, - session=session, - expected_epoch=session.epoch, - ) - - assert close_reason is None - assert emitted is True - assert ws.sent_types() == [ - "response.listen", - "response.done", - "response.created", - "response.speak", - "response.output_audio.delta", - ] - assert session.active_response_turn_id == 1 - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_boundary_does_not_close_new_epoch_response(): - request_id = "duplex-sid-boundary-race-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-boundary-race", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - old_response_id = session.begin_response(turn_id=0) - session.bind_request(request_id) - native = handler._runtime_session_state(session) - native.continuation_owner_id = f"response:{old_response_id}" - native.continuation_units = handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS - sent: list[dict[str, Any]] = [] - new_response_ids: list[str] = [] - - async def _barge_in_on_listen(payload: dict[str, Any]) -> None: - sent.append(payload) - if payload["type"] == "response.listen": - session.barge_in() - session.bind_request("duplex-sid-boundary-race-e1-stage0") - new_response_ids.append(session.begin_response(turn_id=1)) - - await handler._maybe_continue_native_response( - _barge_in_on_listen, - session=session, - expected_epoch=session.epoch, - ) - - assert [payload["type"] for payload in sent] == ["response.listen"] - assert len(new_response_ids) == 1 - assert session.active_response_id == new_response_ids[0] - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_pre_speak_listen_continues_same_response(): - request_id = "duplex-sid-pre-speak-listen-e0-stage0" - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession( - session_id="sid-pre-speak-listen", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - response_id = session.begin_response(turn_id=0) - session.turn_id = 1 - session.bind_request(request_id) - _install_direct_silence_scheduler(handler, session) - ws = TimedWebSocket() - - await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "llm", - "is_listen": True, - "model_listen": True, - "data_plane_request_id": request_id, - "end_of_turn": False, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - await asyncio.sleep(0.01) - - assert len(engine.appended) == 1 - assert session.active_response_id == response_id - assert session.active_request_id == request_id - assert "response.listen" not in ws.sent_types() - assert "response.done" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_listen_before_response_keeps_resumable_request(): - request_id = "duplex-sid-listen-before-response-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-listen-before-response", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.bind_request(request_id) - ws = TimedWebSocket() - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "llm", - "is_listen": True, - "model_listen": True, - "data_plane_request_id": request_id, - "end_of_turn": False, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - - assert close_reason is None - assert emitted is True - assert session.active_request_id == request_id - assert not handler._serving_runtime_adapter.data_plane.is_terminal(request_id) - assert ws.sent_types().count("response.listen") == 1 - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_listen_without_response_does_not_defer_commit(): - session_id = "sid-listen-before-response-commit" - request_id = duplex_resource_request_id(DuplexFence(session_id), "stage0") - inner_output = SimpleNamespace( - outputs=[ - SimpleNamespace( - multimodal_output={"special_token_ids": {"listen_token_id": 151705}}, - ) - ] - ) - decision = MiniCPMO45DuplexRuntimeExtension().decide_output( - stage_id=0, - final_stage_id=1, - segment_finished=True, - segment_token_ids=(151705,), - segment_output_metadata={"special_token_ids": {"listen_token_id": 151705}}, - output=inner_output, - ) - assert decision is not None - listen_output = attach_duplex_output_decision( - OmniRequestOutput.from_stage_output( - inner_output, - request_id=request_id, - finished=True, - stage_id=0, - final_output_type=decision.final_output_type, - ), - decision, - ) - append_result = { - "operation": "append", - "session_id": session_id, - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": { - "supported": True, - "implementation_level": "model_native_duplex", - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - }, - } - ], - "data_plane_outputs": [listen_output], - } - engine = FakeEngineClient(append_result=append_result) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - listen_count = 0 - - def append_residual_after_listen(ws: TimedWebSocket, data: dict[str, Any]) -> None: - nonlocal listen_count - if data.get("type") == "response.listen": - listen_count += 1 - if listen_count == 1: - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(3_200), - "format": "pcm_f32le", - "sample_rate_hz": 16_000, - "duration_ms": 200, - "is_speech": True, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - elif data.get("type") == "input_audio_buffer.committed": - asyncio.get_running_loop().call_later(0.05, ws.put, {"type": "session.close"}) - - ws = TimedWebSocket(on_send=append_residual_after_listen, receive_timeout_s=2) - create = _native_realtime_session_update(session_id) - create["session"]["extra_body"]["auto_response"] = True - ws.put(create) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16_000), - "format": "pcm_f32le", - "sample_rate_hz": 16_000, - "duration_ms": 1_000, - "is_speech": True, - } - ) - - await handler.handle_realtime_session(ws) - - committed = [event for event in ws.sent if event.get("type") == "input_audio_buffer.committed"] - assert len(engine.appended) == 2 - assert engine.appended[1][3] is True - assert len(committed) == 1 - assert committed[0].get("event", {}).get("overlap_deferred") is not True - assert committed[0].get("event", {}).get("response_create_deferred") is not True - assert "response.created" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_minicpmo_pcm_f32_residual_commit_releases_pending_input_bytes(): - session_id = "sid-pcm-residual-byte-accounting" - control_result = { - "operation": "append", - "session_id": session_id, - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [], - } - engine = FakeEngineClient(control_result=control_result) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket(receive_timeout_s=2) - create = _native_realtime_session_update(session_id) - create["session"]["extra_body"]["auto_response"] = True - ws.put(create) - handler_task = asyncio.create_task(handler.handle_realtime_session(ws)) - - try: - for _ in range(100): - session = handler._registry.get(session_id) - if session is not None and "session.created" in ws.sent_types(): - break - await asyncio.sleep(0.01) - else: - pytest.fail("native duplex session did not open") - - pending_bytes_after_turn: list[int] = [] - for turn_index in range(3): - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(19_200), - "format": "pcm_f32le", - "sample_rate_hz": 16_000, - "duration_ms": 1_200, - "is_speech": True, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - expected_appends = (turn_index + 1) * 2 - for _ in range(100): - if len(engine.appended) >= expected_appends: - await asyncio.sleep(0) - break - await asyncio.sleep(0.01) - else: - pytest.fail(f"turn {turn_index} residual append did not complete") - pending_bytes_after_turn.append(session.pending_input_bytes) - finally: - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - assert pending_bytes_after_turn == [0, 0, 0] - assert not any( - event.get("type") == "error" - and isinstance(event.get("error"), dict) - and event["error"].get("code") == "input_backpressure" - for event in ws.sent - ) - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_post_speak_listen_continues_same_response(): - request_id = "duplex-sid-post-speak-listen-e0-stage0" - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession( - session_id="sid-post-speak-listen", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - response_id = session.begin_response(turn_id=0) - session.turn_id = 1 - session.bind_request(request_id) - session.mark_audio_sent(100) - _install_direct_silence_scheduler(handler, session) - ws = TimedWebSocket() - - await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "llm", - "is_listen": True, - "model_listen": True, - "data_plane_request_id": request_id, - "end_of_turn": False, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - await asyncio.sleep(0.01) - - assert len(engine.appended) == 1 - assert session.active_response_id == response_id - assert session.active_request_id == request_id - assert "response.listen" not in ws.sent_types() - assert "response.done" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_turn_end_preserves_resumable_request(): - request_id = "duplex-sid-turn-end-preserve-request-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-turn-end-preserve-request", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.begin_response(turn_id=0) - session.bind_request(request_id) - ws = TimedWebSocket() - - await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": request_id, - "text": "done", - "audio_data": "audio", - "audio_format": "pcm16", - "audio_duration_ms": 100, - "end_of_turn": True, - "model_turn_id": 0, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - - assert session.active_response_id is None - assert session.active_request_id == request_id - assert "response.done" in ws.sent_types() - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_empty_turn_end_emits_model_listen(): - request_id = "duplex-sid-empty-turn-end-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-empty-turn-end", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.bind_request(request_id) - ws = TimedWebSocket() - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": request_id, - "text": "", - "audio_data": "", - "audio_format": "pcm16", - "end_of_turn": True, - "abort_data_plane_request": True, - "model_turn_id": 0, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - - assert close_reason is None - assert emitted is True - assert session.turn_id == 1 - assert session.active_request_id == request_id - assert session.active_response_id is None - assert ws.sent_types() == ["response.listen"] - assert ws.sent[0]["model_listen"] is True - assert ws.sent[0]["reason"] == "model_turn_completed_without_output" - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_drops_late_audio_from_completed_model_turn(): - request_id = "duplex-sid-late-completed-turn-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-late-completed-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.bind_request(request_id) - session.complete_model_turn(0) - ws = TimedWebSocket() - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": request_id, - "text": "late duplicate", - "audio_data": "late-audio", - "audio_format": "pcm16", - "audio_duration_ms": 100, - "end_of_turn": True, - "model_turn_id": 0, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - - assert close_reason is None - assert emitted is False - assert session.turn_id == 1 - assert session.active_response_id is None - assert "response.created" not in ws.sent_types() - assert "response.done" not in ws.sent_types() - - -def test_duplex_data_plane_drops_stale_epoch_audio_after_barge_in(): - class _ChatService: - model_config = _ModelConfig() - - def create_audio(self, audio_obj): - return SimpleNamespace(audio_data=f"wav-{int(audio_obj.audio_tensor.shape[0])}") - - data_plane = _test_data_plane() - session = DuplexSession( - session_id="sid-auto-stale-epoch", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.turn_id = 1 - session.epoch = 2 - output = SimpleNamespace( - request_id="duplex-sid-auto-stale-epoch-stage0", - finished=False, - outputs=[SimpleNamespace(text="old epoch", multimodal_output={})], - multimodal_output={ - "audio": np.zeros(24000, dtype=np.float32), - "sr": 24000, - "meta.duplex_epoch": np.array([1], dtype=np.int32), - "meta.duplex_turn_id": np.array([1], dtype=np.int32), - "meta.llm_output_text_utf8": np.frombuffer("旧轮".encode(), dtype=np.uint8), - }, - ) - - assert list(data_plane.project_output(output, context=_data_plane_context(session))) == [] - - -def test_duplex_data_plane_text_delta_preserves_repeated_suffix_growth(): - data_plane = _test_data_plane() - request_id = "duplex-sid-native-repeated-word-e0-stage0" - - assert data_plane.segment_text_delta(request_id, "好的") == "好的" - assert data_plane.segment_text_delta(request_id, "好的好的") == "好的" - - -def test_duplex_data_plane_text_delta_appends_distinct_non_prefix_segments(): - data_plane = _test_data_plane() - request_id = "duplex-sid-native-rewrite-e0-stage0" - - deltas = [ - data_plane.segment_text_delta(request_id, "It's Canberra."), - data_plane.segment_text_delta(request_id, " Next question."), - ] - - assert deltas == ["It's Canberra.", " Next question."] - assert "".join(deltas) == "It's Canberra. Next question." - - -def test_duplex_data_plane_close_session_removes_encoded_request_state(): - data_plane = _test_data_plane() - target_request_id = duplex_resource_request_id( - DuplexFence("sid-to-close", incarnation=1, epoch=2), - "stage0", - ) - other_request_id = duplex_resource_request_id( - DuplexFence("sid-to-keep", incarnation=1, epoch=2), - "stage0", - ) - data_plane.begin_request(target_request_id) - data_plane.begin_request(other_request_id) - - data_plane.close_session("sid-to-close") - - assert data_plane.has_request(target_request_id) is False - assert data_plane.has_request(other_request_id) is True - - -def test_duplex_auto_response_segment_complete_keeps_data_plane_request_open(): - import numpy as np - - class _ChatService: - model_config = _ModelConfig() - - def create_audio(self, audio_obj): - return SimpleNamespace(audio_data=f"wav-{int(audio_obj.audio_tensor.shape[0])}") - - data_plane = _test_data_plane() - request_id = "duplex-duplex-sess-stage0" - session = DuplexSession( - session_id="sid-auto-respond", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.begin_response() - session.mark_audio_sent(1280) - data_plane.slice_cumulative_audio(request_id, np.zeros(30720, dtype=np.float32)) - - output = SimpleNamespace( - request_id=request_id, - finished=True, - outputs=[SimpleNamespace(text="hello", multimodal_output={})], - multimodal_output={ - "audio": np.zeros(30720, dtype=np.float32), - "sr": 24000, - }, - ) - - native_results = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert len(native_results) == 1 - assert native_results[0]["is_listen"] is True - assert native_results[0]["reason"] == "auto_response_segment_complete" - assert native_results[0]["data_plane_request_id"] == request_id - assert "abort_data_plane_request" not in native_results[0] - - -def test_duplex_auto_response_prior_playback_does_not_abort_before_current_response(): - data_plane = _test_data_plane() - request_id = "duplex-duplex-sess-stage0" - session = DuplexSession( - session_id="sid-auto-respond-next-turn", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.mark_audio_sent(1280) - session.acknowledge_playback(played_ms=1280, committed_ms=1280) - - output = SimpleNamespace( - request_id=request_id, - finished=True, - outputs=[SimpleNamespace(text="", multimodal_output={})], - multimodal_output={}, - ) - - native_results = _project_data_plane(data_plane, {"data_plane_outputs": [output]}, session=session) - - assert native_results == [] - - -def test_duplex_auto_response_text_only_flush_does_not_consume_audio_transcript(): - import numpy as np - - data_plane = _test_data_plane() - request_id = "duplex-duplex-sess-stage0" - session = DuplexSession( - session_id="sid-auto-respond-delayed-audio", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - text = "你好呀，有什莫问的吗？" - - text_only_output = SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[SimpleNamespace(text=text, multimodal_output={})], - multimodal_output={"sr": 24000}, - ) - - assert _project_data_plane(data_plane, {"data_plane_outputs": [text_only_output]}, session=session) == [] - - audio_output = SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[SimpleNamespace(text=text, multimodal_output={})], - multimodal_output={ - "audio": np.zeros(24000, dtype=np.float32), - "sr": 24000, - }, - ) - - native_results = _project_data_plane(data_plane, {"data_plane_outputs": [audio_output]}, session=session) - - assert len(native_results) == 1 - assert native_results[0]["audio_data"].startswith("wav-") - assert native_results[0]["text"] == text - - -@pytest.mark.asyncio -async def test_duplex_auto_response_empty_terminal_does_not_create_empty_response(): - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - request_id = "duplex-sid-native-empty-terminal-e0-stage0" - session = DuplexSession( - session_id="sid-native-empty-terminal", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.mark_audio_sent(1280) - sent: list[dict[str, Any]] = [] - - async def send_json(payload: dict[str, Any]) -> None: - sent.append(payload) - - output = SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[SimpleNamespace(text="hello", multimodal_output={})], - multimodal_output={ - "sr": 24000, - "meta": {"turn_end": True, "duplex_turn_id": 0}, - }, - ) - - native_results = list( - handler._serving_runtime_adapter.data_plane.project( - {"data_plane_outputs": [output]}, - context=_data_plane_context(session), - ) - ) - assert len(native_results) == 1 - assert native_results[0]["end_of_turn"] is True - assert native_results[0]["model_turn_id"] == 0 - - for native_result in native_results: - await handler._send_one_native_duplex_event(send_json, native_result, session=session) - - assert "response.created" not in [m.get("type") for m in sent] - assert "response.done" not in [m.get("type") for m in sent] - assert session.turn_id == 1 - assert not handler._serving_runtime_adapter.data_plane.is_terminal(request_id) - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_session_close_cleans_auto_response_data_plane_state(): - session_id = "sid-native-cleanup" - request_id = f"duplex-{session_id}-e0-stage0" - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - handler._serving_runtime_adapter.data_plane.slice_cumulative_audio(request_id, np.zeros(24000, dtype=np.float32)) - handler._serving_runtime_adapter.data_plane.segment_text_delta(request_id, "hello") - handler._serving_runtime_adapter.data_plane.mark_terminal(request_id) - native = MiniCPMO45ServingSessionState() - handler._serving_runtime_adapter.session_states[session_id] = native - - event = _native_session_create(session_id) - event["session"]["extra_body"]["auto_response"] = True - ws = TimedWebSocket() - ws.put(event) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert not handler._serving_runtime_adapter.data_plane.has_request(request_id) - assert session_id not in handler._serving_runtime_adapter.session_states - - -@pytest.mark.asyncio -async def test_duplex_chat_audio_stream_uses_output_audio_delta_event(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - session = DuplexSession(session_id="sid-chat-audio", config=DuplexSessionConfig()) - response_id = session.begin_response() - sent: list[dict[str, Any]] = [] - - async def send_json(data: dict[str, Any], *, on_accepted: Callable[[], None] | None = None) -> None: - sent.append(data) - if on_accepted is not None: - on_accepted() - - await handler._emit_chat_payload( - session, - { - "modality": "audio", - "choices": [ - { - "delta": { - "content": "AAAA", - }, - "audio_metadata": { - "format": "wav", - "sample_rate_hz": 22_050, - "frame_count": 2205, - "channels": 1, - }, - } - ], - }, - session.epoch, - response_id, - send_json, - ) - - assert sent == [ - { - "type": "response.output_audio.delta", - "session_id": "sid-chat-audio", - "response_id": response_id, - "epoch": 0, - "audio": "AAAA", - "format": "wav", - "sample_rate_hz": 22_050, - "channels": 1, - "audio_duration_ms": 100, - } - ] - assert session.playback.generated_ms == 100 - assert session.playback.sent_ms == 100 - assert session.playback.played_ms == 0 - - -@pytest.mark.parametrize("response_format", ["pcm", "wav"]) -def test_duplex_chat_request_uses_response_audio_format(response_format: str): - handler = OmniDuplexSessionHandler( - chat_service=TurnBasedFakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - extra_audio = {"format": "flac", "voice": "alloy"} - session = DuplexSession( - session_id="sid-chat-audio-format", - config=DuplexSessionConfig(response_format=response_format, extra_body={"audio": extra_audio}), - ) - session.begin_response() - - request = handler._build_chat_request(session, "format-request") - - assert request.audio == {"format": response_format, "voice": "alloy"} - assert extra_audio == {"format": "flac", "voice": "alloy"} - - -def _chat_audio_completion_payload(*, frame_count: int = 2205, sample_rate_hz: int = 22_050) -> dict[str, object]: - return { - "modality": "audio", - "choices": [ - { - "delta": {"content": "AAAA"}, - "audio_metadata": { - "format": "pcm", - "sample_rate_hz": sample_rate_hz, - "frame_count": frame_count, - "channels": 1, - }, - } - ], - } - - -@pytest.mark.asyncio -async def test_duplex_chat_full_audio_message_preserves_actual_format_and_duration(): - handler = OmniDuplexSessionHandler( - chat_service=TurnBasedFakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-chat-full-audio", config=DuplexSessionConfig(response_format="wav")) - response_id = session.begin_response() - sent: list[dict[str, object]] = [] - - async def send_json(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: - sent.append(data) - if on_accepted is not None: - on_accepted() - - await handler._emit_chat_payload( - session, - { - "choices": [ - { - "message": {"role": "assistant", "content": "hello", "audio": {"data": "AAAA"}}, - "audio_metadata": { - "format": "pcm", - "sample_rate_hz": 22_050, - "frame_count": 2205, - "channels": 1, - }, - } - ], - }, - session.epoch, - response_id, - send_json, - ) - - assert [event["type"] for event in sent] == ["response.text.delta", "response.output_audio.delta"] - assert sent[1]["audio"] == "AAAA" - assert sent[1]["format"] == "pcm" - assert sent[1]["sample_rate_hz"] == 22_050 - assert sent[1]["audio_duration_ms"] == 100 - assert session.assistant_text_buffer == ("hello",) - assert session.playback.sent_ms == 100 - - -@pytest.mark.asyncio -async def test_duplex_chat_audio_does_not_round_each_fragment(): - handler = OmniDuplexSessionHandler( - chat_service=TurnBasedFakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-chat-sub-ms", config=DuplexSessionConfig()) - response_id = session.begin_response() - sent: list[dict[str, object]] = [] - - async def send_json(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: - sent.append(data) - if on_accepted is not None: - on_accepted() - - for _ in range(23): - await handler._emit_chat_payload( - session, - _chat_audio_completion_payload(frame_count=1), - session.epoch, - response_id, - send_json, - ) - - assert [event["audio_duration_ms"] for event in sent] == [0] * 22 + [1] - assert session.playback.generated_ms == 1 - assert session.playback.sent_ms == 1 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("send_error", [None, RuntimeError("output unavailable")], ids=["rejected", "failed"]) -async def test_duplex_chat_unaccepted_audio_is_generated_but_not_sent(send_error: RuntimeError | None): - handler = OmniDuplexSessionHandler( - chat_service=TurnBasedFakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-chat-unaccepted-audio", config=DuplexSessionConfig()) - response_id = session.begin_response() - - async def send_json(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: - if send_error is not None: - raise send_error - - with pytest.raises(RuntimeError, match="output unavailable") if send_error is not None else nullcontext(): - await handler._emit_chat_payload( - session, _chat_audio_completion_payload(), session.epoch, response_id, send_json - ) - - assert session.playback.generated_ms == 100 - assert session.playback.sent_ms == 0 - assert session.playback.played_ms == 0 - - -@pytest.mark.asyncio -async def test_duplex_chat_late_send_completion_only_updates_old_response(): - handler = OmniDuplexSessionHandler( - chat_service=TurnBasedFakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-chat-late-audio", config=DuplexSessionConfig()) - old_response = session.begin_response() - send_started = asyncio.Event() - finish_send = asyncio.Event() - - async def delayed_send(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: - send_started.set() - await finish_send.wait() - assert on_accepted is not None - on_accepted() - - emit_task = asyncio.create_task( - handler._emit_chat_payload(session, _chat_audio_completion_payload(), session.epoch, old_response, delayed_send) - ) - try: - await asyncio.wait_for(send_started.wait(), timeout=1) - session.barge_in() - new_response = session.begin_response() - current_state = session.turn_state - finish_send.set() - await asyncio.wait_for(emit_task, timeout=1) - - assert session.active_response_id == new_response - assert session.playback_for_response(old_response).sent_ms == 100 - assert session.playback.generated_ms == 0 - assert session.playback.sent_ms == 0 - assert session.turn_state == current_state - finally: - emit_task.cancel() - await asyncio.gather(emit_task, return_exceptions=True) - - -@pytest.mark.asyncio -async def test_duplex_chat_missing_audio_metadata_reports_error(): - class MissingMetadataChatService(TurnBasedFakeChatService): - async def create_chat_completion(self, request, raw_request=None): - async def generate(): - payload = {"modality": "audio", "choices": [{"delta": {"content": "AAAA"}}]} - yield f"data: {json.dumps(payload)}\n\n" - - return generate() - - handler = OmniDuplexSessionHandler( - chat_service=MissingMetadataChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-chat-missing-metadata", config=DuplexSessionConfig()) - sent: list[dict[str, object]] = [] - - async def send_json(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: - sent.append(data) - if on_accepted is not None: - on_accepted() - - await handler._run_response(session, send_json) - - assert [event["type"] for event in sent] == ["response.created", "error"] - assert sent[-1]["code"] == "response_error" - assert "AudioChunkMetadata" in sent[-1]["error"] - assert session.playback.generated_ms == 0 - assert session.playback.sent_ms == 0 - - -@pytest.mark.asyncio -async def test_duplex_chat_stage_metrics_use_latest_streaming_snapshot(): - class SnapshotChatService(FakeChatService): - async def create_chat_completion(self, request, raw_request=None): - del raw_request - self.seen_request_ids.append(request.request_id) - snapshots = [ - {"num_tokens_out": 1, "vllm_tpot_ms": 0.0, "vllm_itls_ms": []}, - {"num_tokens_out": 2, "vllm_tpot_ms": 10.0, "vllm_itls_ms": [10.0]}, - {"num_tokens_out": 3, "vllm_tpot_ms": 11.0, "vllm_itls_ms": [10.0, 12.0]}, - ] - - async def generate(): - for index, snapshot in enumerate(snapshots): - payload = { - "modality": "audio", - "choices": [ - { - "delta": {"content": f"audio-{index}"}, - "audio_metadata": { - "format": "pcm", - "sample_rate_hz": 24_000, - "frame_count": 2400, - "channels": 1, - }, - } - ], - "metrics": {"stage_metrics": {"0": snapshot}}, - } - yield f"data: {json.dumps(payload)}\n\n" - - return generate() - - handler = OmniDuplexSessionHandler( - chat_service=SnapshotChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-chat-snapshots", - config=DuplexSessionConfig(instructions="speak"), - ) - sent: list[dict[str, Any]] = [] - - async def send_json(data: dict[str, Any], *, on_accepted: Callable[[], None] | None = None) -> None: - sent.append(data) - if on_accepted is not None: - on_accepted() - - await handler._run_response(session, send_json) - - audio_deltas = [event for event in sent if event.get("type") == "response.output_audio.delta"] - done = next(event for event in sent if event.get("type") == "response.done") - assert [event["audio_duration_ms"] for event in audio_deltas] == [100, 200, 300] - assert done["playback"] == { - "generated_ms": 300, - "sent_ms": 300, - "played_ms": 0, - "committed_ms": 0, - } - assert all("vllm_omni" not in event for event in audio_deltas) - assert done["vllm_omni"]["stage_metrics"]["0"] == { - "num_tokens_out": 3, - "vllm_tpot_ms": 11.0, - "vllm_itls_ms": [10.0, 12.0], - } - - -@pytest.mark.asyncio -async def test_duplex_chat_rejects_unknown_output_modality(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-chat-unknown", config=DuplexSessionConfig()) - response_id = session.begin_response() - sent: list[dict[str, Any]] = [] - - async def send_json(data: dict[str, Any]) -> None: - sent.append(data) - - await handler._emit_chat_payload( - session, - { - "modality": "video", - "choices": [{"delta": {"content": "not-text"}}], - }, - session.epoch, - response_id, - send_json, - ) - - assert sent == [ - { - "type": "error", - "session_id": session.session_id, - "response_id": response_id, - "epoch": session.epoch, - "code": "unsupported_response_modality", - "error": "Unsupported chat response modality: video", - } - ] - assert session.assistant_text_buffer == () - - -@pytest.mark.asyncio -async def test_minicpmo_model_name_does_not_auto_enable_experimental_native_duplex(): - event = _session_create("sid-minicpmo-default-protocol") - event["session"]["model"] = "openbmb/MiniCPM-o-4_5" - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(event) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - created = next(m for m in ws.sent if m.get("type") == "session.created") - capabilities = created["session"]["capabilities"] - assert capabilities["implementation_level"] == "serving_session_adapter" - assert capabilities["supports_input_append"] is False - - -@pytest.mark.asyncio -async def test_duplex_handler_aborts_current_chat_request_id_on_barge_in(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") == "response.created": - ws.put({"type": "input.cancel", "reason": "test_barge_in"}) - - ws = TimedWebSocket(on_send=on_send) - ws.put(_session_create("sid-a")) - ws.put({"type": "input.text.append", "text": "hello"}) - ws.put({"type": "input.commit"}) - - await handler.handle_session(ws) - - assert ws.accepted - assert "response.created" in ws.sent_types() - assert "audio.cancelled" in ws.sent_types() - assert chat_service.seen_request_ids == ["duplex-sid-a-0-1"] - assert engine.aborted == ["chatcmpl-duplex-sid-a-0-1"] - assert engine.internal_abort_batches == [["chatcmpl-duplex-sid-a-0-1"]] - - -@pytest.mark.asyncio -async def test_duplex_handler_runtime_open_failure_is_reported_to_client(): - engine = FakeEngineClient(fail_open=True) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-open-fail")) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "runtime_open_failed" - assert engine.opened == [] - - -@pytest.mark.asyncio -async def test_duplex_handler_preserves_typed_admission_error(): - class AdmissionRejectingEngine(FakeEngineClient): - async def open_duplex_session_async(self, session_id: str, **kwargs): - del kwargs - raise DuplexControlRequestError( - { - "operation": "open", - "session_id": session_id, - "error": { - "code": "resource_exhausted", - "message": "duplex session capacity exhausted", - "retryable": True, - }, - "accepted_fence": None, - "lease_generation": None, - } - ) - - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(AdmissionRejectingEngine()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-admission-rejected")) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "resource_exhausted" - assert ws.sent[0]["retryable"] is True - - -@pytest.mark.asyncio -async def test_duplex_cancel_reports_playback_committed_cursor(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") == "response.created": - ws.put({"type": "playback.ack", "played_ms": 1200, "committed_ms": 1000}) - ws.put({"type": "input.cancel", "reason": "test_barge_in"}) - - ws = TimedWebSocket(on_send=on_send) - ws.put(_session_create("sid-playback")) - ws.put({"type": "input.text.append", "text": "hello"}) - ws.put({"type": "input.commit"}) - - await handler.handle_session(ws) - - ack = next(m for m in ws.sent if m.get("type") == "playback.acknowledged") - cancelled = next(m for m in ws.sent if m.get("type") == "audio.cancelled") - assert ack["playback"]["played_ms"] == 1200 - assert ack["playback"]["committed_ms"] == 1000 - assert cancelled["committed_ms"] == 1000 - assert cancelled["epoch"] == 1 - - -@pytest.mark.asyncio -async def test_playback_ack_response_id_selects_matching_pending_history_item(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-playback-response-id", - config=DuplexSessionConfig(), - ) - first_item_id = "item_resp-first" - second_item_id = "item_resp-second" - session.stage_pending_history_item( - first_item_id, - {"role": "assistant", "content": "first response"}, - ) - session.stage_pending_history_item( - second_item_id, - {"role": "assistant", "content": "second response"}, - ) - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": "resp-second", - "played_ms": 1000, - "committed_ms": 1000, - }, - ws.send_json, - ) - - ack = next(message for message in ws.sent if message.get("type") == "playback.acknowledged") - assert ack["item_id"] == second_item_id - assert ack["history_committed"] is True - assert first_item_id in session.pending_history_item_ids - assert second_item_id not in session.pending_history_item_ids - - -@pytest.mark.asyncio -async def test_playback_ack_recovers_missing_response_history_registration(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-playback-missing-registration", - config=DuplexSessionConfig(playback_commit_policy="ack_only"), - ) - response_id = session.begin_response() - session.append_assistant_text("response text") - session.mark_audio_sent(1000, text_chars=len("response text")) - session.end_response(commit_text=True, preserve_request=True) - assert not session.history_item_ids - assert f"item_{response_id}" in session.pending_history_item_ids - - ws = TimedWebSocket() - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": response_id, - "item_id": f"item_{response_id}", - "played_ms": 1000, - "committed_ms": 1000, - }, - ws.send_json, - ) - - ack = next(message for message in ws.sent if message.get("type") == "playback.acknowledged") - assert ack["history_committed"] is True - assert f"item_{response_id}" in session.history_item_ids - assert session.history[-1] == {"role": "assistant", "content": "response text"} - - -@pytest.mark.asyncio -async def test_playback_ack_rejects_response_after_followup_user_commit(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-playback-after-followup-user", - config=DuplexSessionConfig(playback_commit_policy="ack_only"), - ) - session.append_text("user A") - session.commit_user_input() - response_id = session.begin_response() - session.append_assistant_text("assistant A") - session.mark_audio_sent(1000, text_chars=len("assistant A")) - session.end_response(commit_text=False, preserve_request=True) - session.append_text("user B") - session.commit_user_input() - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": response_id, - "item_id": f"item_{response_id}", - "played_ms": 1000, - "committed_ms": 1000, - }, - ws.send_json, - ) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "playback_ack_too_late" - assert session.history == ( - {"role": "user", "content": "user A"}, - {"role": "user", "content": "user B"}, - ) - - -@pytest.mark.asyncio -async def test_precommit_zero_audio_ack_reserves_history_for_response_that_finishes_after_commit(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-zero-audio-checkpoint", - config=DuplexSessionConfig(playback_commit_policy="ack_only"), - ) - session.mark_user_input_activity() - response_id = session.begin_response() - item_id = f"item_{response_id}" - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": response_id, - "item_id": item_id, - "played_ms": 0, - "committed_ms": 0, - }, - ws.send_json, - ) - session.commit_native_audio_input(transcript="later user input") - session.append_assistant_text("hello world") - session.mark_audio_sent(1000, text_chars=len("hello world")) - committed_message = session.end_response(commit_text=True, preserve_request=True) - session.register_history_item(item_id, committed_message) - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": response_id, - "item_id": item_id, - "played_ms": 1000, - "committed_ms": 1000, - }, - ws.send_json, - ) - - assert not any(message.get("type") == "error" for message in ws.sent) - assert session.history == ( - {"role": "assistant", "content": "hello world"}, - { - "role": "user", - "content": [ - { - "type": "audio_url", - "audio_url": {"url": "native-duplex:input-audio"}, - "transcript": "later user input", - } - ], - "transcript": "later user input", - }, - ) - - -@pytest.mark.asyncio -async def test_precommit_partial_ack_extends_history_after_response_finishes(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-partial-checkpoint", - config=DuplexSessionConfig(playback_commit_policy="ack_only"), - ) - session.mark_user_input_activity() - response_id = session.begin_response() - item_id = f"item_{response_id}" - session.append_assistant_text("hello world") - session.mark_audio_sent(1000, text_chars=len("hello world")) - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": response_id, - "item_id": item_id, - "played_ms": 500, - "committed_ms": 500, - }, - ws.send_json, - ) - session.commit_native_audio_input(transcript="later user input") - committed_message = session.end_response(commit_text=True, preserve_request=True) - session.register_history_item(item_id, committed_message) - - assert session.history[0] == {"role": "assistant", "content": "hello"} - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": response_id, - "item_id": item_id, - "played_ms": 1000, - "committed_ms": 1000, - }, - ws.send_json, - ) - - assert not any(message.get("type") == "error" for message in ws.sent) - assert session.history[0] == {"role": "assistant", "content": "hello world"} - assert session.history[1]["role"] == "user" - - -@pytest.mark.asyncio -async def test_explicit_truncate_seals_snapshot_against_later_playback_ack(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-hard-playback-truncate", - config=DuplexSessionConfig(playback_commit_policy="ack_only"), - ) - response_id = session.begin_response() - item_id = f"item_{response_id}" - session.append_assistant_text("hello world") - session.mark_audio_sent(1000, text_chars=len("hello world")) - session.end_response(commit_text=False, preserve_request=True) - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": response_id, - "item_id": item_id, - "played_ms": 500, - "committed_ms": 500, - "truncate": True, - }, - ws.send_json, - ) - assert session.history == ({"role": "assistant", "content": "hello"},) - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": response_id, - "item_id": item_id, - "played_ms": 1000, - "committed_ms": 1000, - }, - ws.send_json, - ) - - assert not any(message.get("type") == "error" for message in ws.sent) - assert session.history == ({"role": "assistant", "content": "hello"},) - - -@pytest.mark.asyncio -async def test_cancel_does_not_append_old_assistant_after_followup_user(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-cancel-after-followup-user", - config=DuplexSessionConfig(playback_commit_policy="ack_only"), - ) - session.append_text("user A") - session.commit_user_input() - session.begin_response() - session.append_assistant_text("assistant A") - session.mark_audio_sent(1000, text_chars=len("assistant A")) - session.acknowledge_playback(500, 500) - session.append_text("user B") - session.commit_user_input() - ws = TimedWebSocket() - - cancelled = await handler._cancel_active_response(session, None, ws.send_json, reason="new_response") - - assert cancelled is True - assert session.history == ( - {"role": "user", "content": "user A"}, - {"role": "user", "content": "user B"}, - ) - - -@pytest.mark.asyncio -async def test_playback_ack_rejects_item_from_another_response(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-playback-mismatch", config=DuplexSessionConfig()) - first_response_id = session.begin_response() - session.append_assistant_text("first response") - session.mark_audio_sent(1000) - session.end_response(commit_text=False, preserve_request=True) - session.register_history_item(f"item_{first_response_id}", None) - second_response_id = session.begin_response() - session.mark_audio_sent(1000) - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": second_response_id, - "item_id": f"item_{first_response_id}", - "played_ms": 500, - "committed_ms": 500, - }, - ws.send_json, - ) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "playback_item_mismatch" - assert session.active_response_id == second_response_id - assert session.playback.as_dict() == { - "generated_ms": 1000, - "sent_ms": 1000, - "played_ms": 0, - "committed_ms": 0, - } - - -@pytest.mark.asyncio -async def test_playback_ack_rejects_user_history_item(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-playback-user-item", config=DuplexSessionConfig()) - user_message = {"role": "user", "content": "user input"} - session.append_history_message(user_message) - session.register_history_item("item_user", user_message) - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": "user", - "item_id": "item_user", - "played_ms": 500, - "committed_ms": 500, - }, - ws.send_json, - ) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "playback_item_not_found" - assert session.history == ({"role": "user", "content": "user input"},) - - -@pytest.mark.asyncio -async def test_late_playback_ack_recovers_response_local_history_after_next_response_begins(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession(session_id="sid-playback-response-local", config=DuplexSessionConfig()) - first_response_id = session.begin_response() - session.append_assistant_text("first response") - session.mark_audio_sent(1000, text_chars=len("first response")) - session.end_response(commit_text=False, preserve_request=True) - second_response_id = session.begin_response() - session.mark_audio_sent(2000, text_chars=len("second response")) - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": first_response_id, - "item_id": f"item_{first_response_id}", - "played_ms": 1000, - "committed_ms": 1000, - }, - ws.send_json, - ) - - ack = next(message for message in ws.sent if message.get("type") == "playback.acknowledged") - assert ack["history_committed"] is True - assert session.history[-1] == {"role": "assistant", "content": "first response"} - assert session.active_response_id == second_response_id - assert session.playback.as_dict() == { - "generated_ms": 2000, - "sent_ms": 2000, - "played_ms": 0, - "committed_ms": 0, - } - - -@pytest.mark.asyncio -async def test_late_playback_ack_does_not_advance_new_response_cursor(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-late-playback-ack", - config=DuplexSessionConfig(), - ) - first_response_id = session.begin_response() - session.append_assistant_text("first response") - session.mark_audio_sent(11480, text_chars=len("first response")) - session.end_response(commit_text=False, preserve_request=True) - session.register_history_item(f"item_{first_response_id}", None) - - second_response_id = session.begin_response() - session.mark_audio_sent(2200, text_chars=len("second response")) - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": first_response_id, - "item_id": f"item_{first_response_id}", - "played_ms": 11480, - "committed_ms": 11480, - }, - ws.send_json, - ) - - assert session.active_response_id == second_response_id - assert session.playback.as_dict() == { - "generated_ms": 2200, - "sent_ms": 2200, - "played_ms": 0, - "committed_ms": 0, - } - ack = next(message for message in ws.sent if message.get("type") == "playback.acknowledged") - assert ack["playback"] == { - "generated_ms": 11480, - "sent_ms": 11480, - "played_ms": 11480, - "committed_ms": 11480, - } - assert first_response_id not in session.response_playbacks - assert session.response_playbacks[second_response_id] == session.playback - - -@pytest.mark.asyncio -async def test_late_playback_ack_truncates_history_with_matching_response_cursor(): - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient()), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-late-playback-history", - config=DuplexSessionConfig(), - ) - first_response_id = session.begin_response() - session.append_assistant_text("abcdefghij") - session.mark_audio_sent(1000) - session.end_response(commit_text=False, preserve_request=True) - session.register_history_item(f"item_{first_response_id}", None) - - session.begin_response() - session.mark_audio_sent(10000) - ws = TimedWebSocket() - - await handler._handle_playback_ack( - session, - { - "type": "playback.ack", - "response_id": first_response_id, - "item_id": f"item_{first_response_id}", - "played_ms": 500, - "committed_ms": 500, - }, - ws.send_json, - ) - - assert session.history[-1] == {"role": "assistant", "content": "abcde"} - assert first_response_id in session.response_playbacks - - -@pytest.mark.asyncio -async def test_cancel_active_response_commits_only_played_assistant_history(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - session = DuplexSession( - session_id="sid-playback-partial", - config=DuplexSessionConfig(playback_commit_policy=DuplexPlaybackCommitPolicy.COMMIT_ALL_ON_DONE.value), - ) - response_id = session.begin_response() - session.append_assistant_text("hello world") - session.mark_audio_sent(1000, text_chars=len("hello world")) - session.acknowledge_playback(played_ms=500, committed_ms=500) - - cancelled = await handler._cancel_active_response(session, None, ws.send_json, reason="barge_in") - - assert cancelled is True - assert session.epoch == 1 - assert session.history == ({"role": "assistant", "content": "hello"},) - assert session.history_item_ids[f"item_{response_id}"] == session.history[0] - event = next(m for m in ws.sent if m.get("type") == "audio.cancelled") - assert event["committed_ms"] == 500 - - -@pytest.mark.asyncio -async def test_cancel_active_response_keeps_truncated_history_in_serving(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - session = DuplexSession( - session_id="sid-playback-runtime-rebuild", - config=DuplexSessionConfig(playback_commit_policy=DuplexPlaybackCommitPolicy.COMMIT_ALL_ON_DONE.value), - ) - response_id = session.begin_response() - session.append_assistant_text("hello world") - session.mark_audio_sent(1000, text_chars=len("hello world")) - session.acknowledge_playback(played_ms=500, committed_ms=500) - - cancelled = await handler._cancel_active_response( - session, - None, - ws.send_json, - reason="barge_in", - ) - - assert cancelled is True - assert session.history == ({"role": "assistant", "content": "hello"},) - assert session.history_item_ids[f"item_{response_id}"] == session.history[0] - assert engine.signals == [] - assert "error" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_cancel_active_native_data_plane_request_aborts_stage_request_id(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - session = DuplexSession(session_id="sid-native-abort", config=DuplexSessionConfig()) - session.bind_request("duplex-sid-native-abort-e0-stage0-s1") - - cancelled = await handler._cancel_active_response(session, None, ws.send_json, reason="barge_in") - - assert cancelled is True - assert engine.aborted == ["duplex-sid-native-abort-e0-stage0-s1"] - assert engine.internal_abort_batches == [["duplex-sid-native-abort-e0-stage0-s1"]] - assert "audio.cancelled" in ws.sent_types() - - -@pytest.mark.asyncio -async def test_cancel_chat_fallback_response_aborts_request_and_task(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler(chat_service=TurnBasedFakeChatService(engine)) - session = DuplexSession("sid-chat-cancel", DuplexSessionConfig()) - response_id = session.begin_response() - session.bind_request("chatcmpl-cancel") - ws = TimedWebSocket() - task = asyncio.create_task(asyncio.sleep(60)) - try: - cancelled = await handler._cancel_active_response(session, task, ws.send_json, reason="response.cancel") - - assert cancelled is True - assert task.cancelled() - assert engine.internal_abort_batches == [["chatcmpl-cancel"]] - assert "error" not in ws.sent_types() - (event,) = [event for event in ws.sent if event["type"] == "audio.cancelled"] - assert event["response_id"] == response_id - assert event["reason"] == "response.cancel" - finally: - task.cancel() - await asyncio.gather(task, return_exceptions=True) - - -@pytest.mark.asyncio -async def test_duplex_handler_explicit_close_closes_runtime_once_with_client_reason(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-close")) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["session.created", "session.closed"] - assert engine.closed == [("sid-close", "session_close")] - - -@pytest.mark.asyncio -async def test_realtime_disconnect_detaches_resumable_session_without_runtime_close(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket(receive_timeout_s=0.01) - ws.put(_native_realtime_session_update("sid-resumable-disconnect")) - - await handler.handle_realtime_session(ws) - - created = next(message for message in ws.sent if message.get("type") == "session.created") - assert created["session"]["id"] == "sid-resumable-disconnect" - assert isinstance(created["resume_token"], str) - assert created["resume_token"] - assert created["attachment_generation"] == 1 - assert engine.closed == [] - assert engine.touched == [ - ("sid-resumable-disconnect", DuplexLeaseActivity.DETACH), - ] - assert handler._registry.get("sid-resumable-disconnect") is not None - assert "sid-resumable-disconnect" in handler._serving_runtime_adapter.session_states - - -@pytest.mark.asyncio -async def test_realtime_initial_resume_token_delivery_failure_closes_unrecoverable_session(): - class FailingCreatedWebSocket(TimedWebSocket): - async def send_json(self, data: dict[str, Any]): - if data.get("type") == "session.created": - raise WebSocketDisconnect(code=1006) - await super().send_json(data) - - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = FailingCreatedWebSocket(receive_timeout_s=0.01) - ws.put(_native_realtime_session_update("sid-created-delivery-failure")) - - await handler.handle_realtime_session(ws) - - assert engine.closed == [("sid-created-delivery-failure", "disconnect")] - assert handler._registry.get("sid-created-delivery-failure") is None - assert "sid-created-delivery-failure" not in handler._serving_runtime_adapter.session_states - - -@pytest.mark.asyncio -async def test_realtime_send_after_created_disconnect_preserves_resumable_session(): - class ClosedAfterCreatedWebSocket(TimedWebSocket): - def __init__(self) -> None: - super().__init__(receive_timeout_s=0.01) - self.transport_closed = False - - async def send_json(self, data: dict[str, Any]): - if self.transport_closed: - raise RuntimeError( - "Unexpected ASGI message 'websocket.send', " - "after sending 'websocket.close' or response already completed." - ) - await super().send_json(data) - if data.get("type") == "session.created": - self.transport_closed = True - - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = ClosedAfterCreatedWebSocket() - ws.put(_native_realtime_session_update("sid-created-then-disconnected")) - - await handler.handle_realtime_session(ws) - - assert engine.closed == [] - assert engine.touched == [ - ("sid-created-then-disconnected", DuplexLeaseActivity.DETACH), - ] - assert handler._registry.get("sid-created-then-disconnected") is not None - assert "sid-created-then-disconnected" in handler._serving_runtime_adapter.session_states - - -@pytest.mark.asyncio -async def test_realtime_resume_rotates_token_replays_and_preserves_runtime_identity(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - first = TimedWebSocket(receive_timeout_s=0.01) - first.put(_native_realtime_session_update("sid-realtime-resume")) - first.put({"type": "session.heartbeat"}) - - await handler.handle_realtime_session(first) - - created = next(message for message in first.sent if message.get("type") == "session.created") - heartbeat = next(message for message in first.sent if message.get("type") == "session.heartbeat_ack") - token = created["resume_token"] - incarnation = created["incarnation"] - second = TimedWebSocket(receive_timeout_s=0.1) - second.query_params = { - "model": "openbmb/MiniCPM-o-4_5", - "native_duplex": "1", - "resume": "1", - } - second.put( - { - "type": "session.resume", - "session_id": "sid-realtime-resume", - "incarnation": incarnation, - "resume_token": token, - "last_received_server_event_seq": heartbeat["server_event_seq"] - 1, - } - ) - second.put({"type": "session.close"}) - - await handler.handle_realtime_session(second) - - resumed = next(message for message in second.sent if message.get("type") == "session.resumed") - replayed = [ - message - for message in second.sent - if message.get("type") == "session.heartbeat_ack" - and message.get("server_event_seq") == heartbeat["server_event_seq"] - ] - assert resumed["session_id"] == "sid-realtime-resume" - assert resumed["incarnation"] == incarnation - assert resumed["attachment_generation"] == 2 - assert resumed["resume_token"] != token - assert replayed == [heartbeat] - assert engine.opened == ["sid-realtime-resume"] - assert engine.resumed == [("sid-realtime-resume", 0)] - assert engine.resume_fences[0] == engine.opened_fences[0] - assert engine.closed == [("sid-realtime-resume", "session_close")] - - -@pytest.mark.asyncio -async def test_realtime_resume_preserves_append_tail_order_across_connections(): - class SlowFirstAppendEngine(FakeEngineClient): - def __init__(self) -> None: - super().__init__() - self.first_append_started = asyncio.Event() - self.second_append_started = asyncio.Event() - self.release_first_append = asyncio.Event() - self.append_start_order: list[int] = [] - - async def append_duplex_input_async(self, session_id: str, **kwargs): - call_index = len(self.append_start_order) + 1 - self.append_start_order.append(call_index) - if call_index == 1: - self.first_append_started.set() - await self.release_first_append.wait() - else: - self.second_append_started.set() - kwargs.pop("expected_epoch", None) - return await super().append_duplex_input_async(session_id, **kwargs) - - engine = SlowFirstAppendEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - first = TimedWebSocket(receive_timeout_s=1.0) - session_update = _native_realtime_session_update("sid-resume-append-order") - session_update["session"]["extra_body"]["auto_response"] = True - first.put(session_update) - first_task = asyncio.create_task(handler.handle_realtime_session(first)) - for _ in range(100): - created = next((message for message in first.sent if message.get("type") == "session.created"), None) - if created is not None: - break - await asyncio.sleep(0.01) - else: - pytest.fail("resumable session did not open") - assert created is not None - session = handler._registry.get("sid-resume-append-order") - assert session is not None - native_state = handler._runtime_session_state(session) - first_scheduler = native_state.silence_continuation_scheduler - assert first_scheduler is not None - first.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm16_b64(16000), - "sample_rate_hz": 16000, - "is_speech": True, - } - ) - await asyncio.wait_for(engine.first_append_started.wait(), timeout=1) - await first.close(code=1006) - await asyncio.wait_for(first_task, timeout=1) - - second = TimedWebSocket(receive_timeout_s=0.5) - second.query_params = { - "model": "openbmb/MiniCPM-o-4_5", - "native_duplex": "1", - "resume": "1", - } - second.put( - { - "type": "session.resume", - "session_id": "sid-resume-append-order", - "incarnation": created["incarnation"], - "resume_token": created["resume_token"], - "last_received_server_event_seq": created.get("server_event_seq", 0), - } - ) - second_task = asyncio.create_task(handler.handle_realtime_session(second)) - for _ in range(100): - if "session.resumed" in second.sent_types(): - break - await asyncio.sleep(0.01) - else: - pytest.fail("session did not resume") - assert native_state.silence_continuation_scheduler is not first_scheduler - second.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm16_b64(16000), - "sample_rate_hz": 16000, - "is_speech": True, - } - ) - try: - await asyncio.sleep(0.05) - assert not engine.second_append_started.is_set() - finally: - engine.release_first_append.set() - - await asyncio.wait_for(engine.second_append_started.wait(), timeout=1) - second.put({"type": "session.close"}) - await asyncio.wait_for(second_task, timeout=2) - assert engine.append_start_order == [1, 2] - - -@pytest.mark.asyncio -async def test_realtime_takeover_ignores_late_detach_from_replaced_attachment(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - first = TimedWebSocket(receive_timeout_s=1.0) - first.put(_native_realtime_session_update("sid-takeover")) - first_task = asyncio.create_task(handler.handle_realtime_session(first)) - for _ in range(100): - created = next((message for message in first.sent if message.get("type") == "session.created"), None) - if created is not None: - break - await asyncio.sleep(0.001) - else: - raise AssertionError("first attachment did not create the session") - - second = TimedWebSocket(receive_timeout_s=0.03) - second.put( - { - "type": "session.resume", - "session_id": "sid-takeover", - "incarnation": created["incarnation"], - "resume_token": created["resume_token"], - "last_received_server_event_seq": 0, - } - ) - await handler.handle_realtime_session(second) - await first_task - - detach_touches = [item for item in engine.touched if item[1] is DuplexLeaseActivity.DETACH] - assert detach_touches == [("sid-takeover", DuplexLeaseActivity.DETACH)] - assert "session.replaced" in first.sent_types() - assert engine.resumed == [("sid-takeover", 0)] - assert engine.closed == [] - - -@pytest.mark.asyncio -async def test_realtime_takeover_does_not_clear_new_attachment_pending_turns(): - class DelayedReplacementCloseWebSocket(TimedWebSocket): - def __init__(self) -> None: - super().__init__(receive_timeout_s=1.0) - self.replacement_close_started = asyncio.Event() - self.release_replacement_close = asyncio.Event() - - async def close(self, code: int = 1000, reason: str | None = None) -> None: - if not self.replacement_close_started.is_set(): - self.replacement_close_started.set() - await self.release_replacement_close.wait() - await super().close(code=code, reason=reason) - - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - first = DelayedReplacementCloseWebSocket() - first.put(_native_realtime_session_update("sid-takeover-pending-turn")) - first_task = asyncio.create_task(handler.handle_realtime_session(first)) - for _ in range(100): - created = next((message for message in first.sent if message.get("type") == "session.created"), None) - if created is not None: - break - await asyncio.sleep(0.001) - else: - pytest.fail("first attachment did not create the session") - assert created is not None - - second = TimedWebSocket(receive_timeout_s=1.0) - second.put( - { - "type": "session.resume", - "session_id": "sid-takeover-pending-turn", - "incarnation": created["incarnation"], - "resume_token": created["resume_token"], - "last_received_server_event_seq": 0, - } - ) - second_task = asyncio.create_task(handler.handle_realtime_session(second)) - await asyncio.wait_for(first.replacement_close_started.wait(), timeout=1) - - session = handler._registry.get("sid-takeover-pending-turn") - assert session is not None - assert session.reserve_pending_turn(limit=2) - first.release_replacement_close.set() - await asyncio.wait_for(first_task, timeout=1) - - assert session.pending_input_turns == 1 - session.release_pending_turn() - second.put({"type": "session.close"}) - await asyncio.wait_for(second_task, timeout=2) - - -@pytest.mark.asyncio -async def test_realtime_takeover_drops_input_queued_by_replaced_attachment(): - class DelayedReplacementCloseWebSocket(TimedWebSocket): - def __init__(self) -> None: - super().__init__(receive_timeout_s=1.0) - self.replacement_close_started = asyncio.Event() - self.release_replacement_close = asyncio.Event() - self.stale_heartbeat_read = asyncio.Event() - - async def receive_text(self) -> str: - payload = await super().receive_text() - if json.loads(payload).get("type") == "session.heartbeat": - self.stale_heartbeat_read.set() - return payload - - async def close(self, code: int = 1000, reason: str | None = None) -> None: - if not self.replacement_close_started.is_set(): - self.replacement_close_started.set() - await self.release_replacement_close.wait() - await super().close(code=code, reason=reason) - - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - first = DelayedReplacementCloseWebSocket() - first.put(_native_realtime_session_update("sid-takeover-stale-input")) - first_task = asyncio.create_task(handler.handle_realtime_session(first)) - for _ in range(100): - created = next((message for message in first.sent if message.get("type") == "session.created"), None) - if created is not None: - break - await asyncio.sleep(0.001) - else: - pytest.fail("first attachment did not create the session") - assert created is not None - - second = TimedWebSocket(receive_timeout_s=1.0) - second.put( - { - "type": "session.resume", - "session_id": "sid-takeover-stale-input", - "incarnation": created["incarnation"], - "resume_token": created["resume_token"], - "last_received_server_event_seq": 0, - } - ) - second_task = asyncio.create_task(handler.handle_realtime_session(second)) - await asyncio.wait_for(first.replacement_close_started.wait(), timeout=1) - - first.put({"type": "session.heartbeat"}) - await asyncio.wait_for(first.stale_heartbeat_read.wait(), timeout=1) - first.release_replacement_close.set() - await asyncio.wait_for(first_task, timeout=1) - - second.put({"type": "session.close"}) - await asyncio.wait_for(second_task, timeout=2) - assert all(activity is not DuplexLeaseActivity.HEARTBEAT for _, activity in engine.touched) - assert "session.heartbeat_ack" not in second.sent_types() - - -@pytest.mark.asyncio -async def test_realtime_resume_rejects_invalid_token_before_engine_control(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - first = TimedWebSocket(receive_timeout_s=0.01) - first.put(_native_realtime_session_update("sid-invalid-resume")) - await handler.handle_realtime_session(first) - - created = next(message for message in first.sent if message.get("type") == "session.created") - second = TimedWebSocket(receive_timeout_s=0.01) - second.put( - { - "type": "session.resume", - "session_id": "sid-invalid-resume", - "incarnation": created["incarnation"], - "resume_token": "invalid-token", - "last_received_server_event_seq": 0, - } - ) - - await handler.handle_realtime_session(second) - - error = next(message for message in second.sent if message.get("type") == "error") - assert error["error"]["code"] == "invalid_resume_token" - assert engine.resumed == [] - assert engine.closed == [] - - -@pytest.mark.asyncio -async def test_realtime_resume_recovers_when_rotated_token_delivery_is_lost(): - class FailingResumeWebSocket(TimedWebSocket): - async def send_json(self, data: dict[str, Any]): - if data.get("type") == "session.resumed": - raise RuntimeError("resume transport disappeared") - await super().send_json(data) - - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - first = TimedWebSocket(receive_timeout_s=0.01) - first.put(_native_realtime_session_update("sid-token-delivery")) - await handler.handle_realtime_session(first) - created = next(message for message in first.sent if message.get("type") == "session.created") - resume_event = { - "type": "session.resume", - "session_id": "sid-token-delivery", - "incarnation": created["incarnation"], - "resume_token": created["resume_token"], - "last_received_server_event_seq": 0, - } - - failed_transport = FailingResumeWebSocket(receive_timeout_s=0.01) - failed_transport.put(resume_event) - await handler.handle_realtime_session(failed_transport) - - recovered_transport = TimedWebSocket(receive_timeout_s=0.1) - recovered_transport.put(resume_event) - recovered_transport.put({"type": "session.close"}) - await handler.handle_realtime_session(recovered_transport) - - assert "session.resumed" in recovered_transport.sent_types() - assert engine.resumed == [ - ("sid-token-delivery", 0), - ("sid-token-delivery", 1), - ] - assert engine.closed == [("sid-token-delivery", "session_close")] - - -@pytest.mark.asyncio -async def test_realtime_disconnect_grace_cancels_only_orphan_response(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - duplex_session_config=DuplexSessionRuntimeConfig(disconnect_grace_s=0.02), - ) - ws = TimedWebSocket(receive_timeout_s=0.01) - ws.put(_native_realtime_session_update("sid-grace-response")) - - await handler.handle_realtime_session(ws) - session = handler._registry.get("sid-grace-response") - assert session is not None - session.begin_response() - session.bind_request("orphan-request") - orphan_task = asyncio.create_task(asyncio.sleep(10)) - handler._session_tasks[session.session_id].active_response_task = orphan_task - assert engine.closed == [] - - await asyncio.sleep(0.05) - - assert engine.closed == [] - assert engine.aborted == ["orphan-request"] - assert orphan_task.cancelled() - assert session.active_response_id is None - assert "sid-grace-response" in handler._serving_runtime_adapter.session_states - - -@pytest.mark.asyncio -async def test_realtime_heartbeat_control_failure_is_non_terminal(): - engine = FakeEngineClient(fail_touch=True) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket(receive_timeout_s=0.1) - ws.put(_native_realtime_session_update("sid-heartbeat-failure")) - ws.put({"type": "session.heartbeat"}) - ws.put({"type": "session.close"}) - - await handler.handle_realtime_session(ws) - - error = next(message for message in ws.sent if message.get("type") == "error") - assert error["error"]["code"] == "runtime_touch_failed" - assert engine.closed == [("sid-heartbeat-failure", "session_close")] - - -@pytest.mark.asyncio -async def test_realtime_idle_ttl_lifecycle_removes_detached_serving_projection(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket(receive_timeout_s=0.01) - ws.put(_native_realtime_session_update("sid-idle-expired")) - - await handler.handle_realtime_session(ws) - fence = engine.opened_fences[0] - assert fence is not None - await engine.duplex_lifecycle_events.put( - DuplexSessionLifecycleMessage( - fence=fence, - session_id="sid-idle-expired", - event="expired", - reason="idle_ttl_expired", - lease_generation=1, - submitted_request_ids=[], - reserved_request_ids=[], - ) - ) - await asyncio.sleep(0.05) - - assert handler._registry.get("sid-idle-expired") is None - assert "sid-idle-expired" not in handler._serving_runtime_adapter.session_states - assert "sid-idle-expired" not in handler._session_tasks - assert engine.closed == [] - - -@pytest.mark.asyncio -async def test_realtime_resume_reports_resync_required_after_journal_ttl_gap(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - duplex_session_config=DuplexSessionRuntimeConfig(resume_replay_ttl_s=0.01), - ) - first = TimedWebSocket(receive_timeout_s=0.01) - first.put(_native_realtime_session_update("sid-journal-gap")) - first.put({"type": "session.heartbeat"}) - await handler.handle_realtime_session(first) - created = next(message for message in first.sent if message.get("type") == "session.created") - await asyncio.sleep(0.02) - - second = TimedWebSocket(receive_timeout_s=0.01) - second.put( - { - "type": "session.resume", - "session_id": "sid-journal-gap", - "incarnation": created["incarnation"], - "resume_token": created["resume_token"], - "last_received_server_event_seq": 0, - } - ) - await handler.handle_realtime_session(second) - - assert "session.resync_required" in second.sent_types() - assert engine.resumed == [] - assert engine.closed == [] - - -@pytest.mark.asyncio -async def test_realtime_journal_overflow_degrades_to_live_resync_without_closing_runtime(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - duplex_session_config=DuplexSessionRuntimeConfig(resume_replay_max_bytes_per_session=128), - ) - ws = TimedWebSocket(receive_timeout_s=0.01) - ws.put(_native_realtime_session_update("sid-journal-overflow")) - - await handler.handle_realtime_session(ws) - - assert "session.created" in ws.sent_types() - assert "session.resync_required" in ws.sent_types() - assert engine.closed == [] - assert handler._registry.get("sid-journal-overflow") is not None - - -@pytest.mark.asyncio -async def test_realtime_input_backpressure_is_per_session_and_rejects_before_mutation(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - duplex_session_config=DuplexSessionRuntimeConfig( - max_pending_input_bytes_per_session=100, - ), - ) - first = TimedWebSocket(receive_timeout_s=0.1) - first.put(_native_realtime_session_update("sid-backpressure-a")) - for _ in range(2): - first.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm16_b64(20), - "format": "pcm16", - "sample_rate_hz": 16000, - } - ) - first.put({"type": "session.close"}) - - await handler.handle_realtime_session(first) - - backpressure = [ - event - for event in first.sent - if event.get("type") == "error" - and isinstance(event.get("error"), dict) - and event["error"].get("code") == "input_backpressure" - ] - assert len(backpressure) == 1 - - second = TimedWebSocket(receive_timeout_s=0.1) - second.put(_native_realtime_session_update("sid-backpressure-b")) - second.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm16_b64(20), - "format": "pcm16", - "sample_rate_hz": 16000, - } - ) - second.put({"type": "session.close"}) - - await handler.handle_realtime_session(second) - - assert not any( - event.get("type") == "error" - and isinstance(event.get("error"), dict) - and event["error"].get("code") == "input_backpressure" - for event in second.sent - ) - - -@pytest.mark.asyncio -async def test_duplex_handler_idle_timeout_close_does_not_emit_runtime_control(): - control_result = { - "operation": "close", - "session_id": "sid-disconnect", - "ok": True, - "unsupported_count": 1, - "error_count": 0, - "stage_results": [{"stage_id": 0, "replica_id": 0, "result": {"supported": False}}], - } - engine = FakeEngineClient(append_result=control_result) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=0.1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-disconnect")) - - await handler.handle_session(ws) - - assert engine.closed == [("sid-disconnect", "timeout")] - - -@pytest.mark.asyncio -async def test_idle_output_audio_clear_does_not_advance_runtime_fence(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_session_create("sid-idle-output-clear")) - ws.put({"type": "output_audio_buffer.clear"}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - cancelled = next(event for event in ws.sent if event.get("type") == "audio.cancelled") - assert cancelled["cancelled_epoch"] == cancelled["epoch"] == 0 - assert engine.signals == [] - - -@pytest.mark.asyncio -async def test_duplex_handler_runtime_close_failure_is_reported_without_closed_ack(): - engine = FakeEngineClient(fail_close=True) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-close-fail")) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["session.created", "error"] - assert ws.sent[-1]["code"] == "runtime_close_failed" - - -@pytest.mark.asyncio -async def test_duplex_handler_control_close_failure_is_reported_without_closed_ack(): - control_result = { - "operation": "close", - "session_id": "sid-control-close-fail", - "ok": False, - "unsupported_count": 0, - "error_count": 1, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": {"supported": False, "error": "stage close failed"}, - } - ], - } - engine = FakeEngineClient(close_result=control_result) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-control-close-fail")) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["session.created", "error"] - assert ws.sent[-1]["code"] == "runtime_close_failed" - assert ws.sent[-1]["runtime_control"]["error_count"] == 1 - - -@pytest.mark.asyncio -async def test_duplex_cancel_without_active_response_clears_pending_input_and_acks(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_session_create("sid-pending-cancel")) - ws.put({"type": "input.text.append", "text": "hello"}) - ws.put({"type": "input.cancel", "reason": "user_cancel"}) - ws.put({"type": "input.commit"}) - - await handler.handle_session(ws) - - cancelled = next(m for m in ws.sent if m.get("type") == "input.cancelled") - assert cancelled["cancelled"] == {"text_chunks": 1, "audio_chunks": 0} - assert cancelled["epoch"] == 1 - assert chat_service.seen_request_ids == [] - assert "response.created" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_turn_signal_input_cancel_uses_epoch_fence_transition(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-turn-signal-cancel")) - ws.put({"type": "input.text.append", "text": "discard me"}) - ws.put({"type": "turn.signal", "event": "input.cancel"}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert engine.signal_fences[-1] == DuplexFence("sid-turn-signal-cancel") - assert engine.signal_next_fences[-1] == DuplexFence("sid-turn-signal-cancel", epoch=1) - assert any(message.get("type") == "input.cancelled" for message in ws.sent) - - -@pytest.mark.asyncio -async def test_duplex_handler_local_turn_signal_does_not_round_trip_runtime(): - engine = FakeEngineClient(fail_signal=True) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_session_create("sid-signal-fail")) - ws.put({"type": "turn.signal", "event": "user_started"}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - turn_event = next(m for m in ws.sent if m.get("type") == "turn.event") - assert turn_event["event"] == "user_started" - assert not engine.signals - assert "runtime_signal_failed" not in {m.get("code") for m in ws.sent} - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_append_without_format_defaults_to_pcm16(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-default-pcm16")) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm16_b64(16000), - "sample_rate_hz": 16000, - "is_speech": True, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert len(engine.appended) == 1 - _, mode, payload, final = engine.appended[0] - assert mode == "append_audio_chunk" - assert final is False - assert isinstance(payload, dict) - assert payload["format"] == "pcm_f32le" - assert payload["sample_rate_hz"] == 16000 - samples = np.frombuffer(base64.b64decode(payload["audio"]), dtype=" None: - super().__init__() - self.append_started = asyncio.Event() - self.append_cancelled = asyncio.Event() - self.release_append = asyncio.Event() - - async def append_duplex_input_async( - self, - session_id: str, - *, - mode: str, - payload: object, - operation_id: str | None = None, - final: bool = False, - timeout: float | None = None, - collect_outputs: bool = True, - fence: DuplexFence | None = None, - ) -> Any: - self.append_started.set() - try: - await self.release_append.wait() - except asyncio.CancelledError: - self.append_cancelled.set() - raise - return await super().append_duplex_input_async( - session_id, - mode=mode, - payload=payload, - operation_id=operation_id, - final=final, - timeout=timeout, - collect_outputs=collect_outputs, - fence=fence, - ) - - rollback = mocker.spy(MiniCPMO45PcmAppendBuffer, "_rollback_reservation") - engine = SlowAppendEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-cancel-slow-append")) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm16_b64(16000), - "sample_rate_hz": 16000, - "is_speech": True, - } - ) - ws.put({"type": "input.cancel"}) - ws.put({"type": "session.close"}) - - handler_task = asyncio.create_task(handler.handle_session(ws)) - await asyncio.wait_for(engine.append_started.wait(), timeout=1) - try: - await asyncio.wait_for(engine.append_cancelled.wait(), timeout=0.2) - except TimeoutError: - pass - finally: - engine.release_append.set() - await asyncio.wait_for(handler_task, timeout=1) - - assert engine.append_cancelled.is_set() - rollback.assert_called_once() - assert engine.signal_fences[-1] == DuplexFence("sid-cancel-slow-append") - assert engine.signal_next_fences[-1] == DuplexFence("sid-cancel-slow-append", epoch=1) - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_append_control_error_does_not_emit_model_delta(): - control_result = { - "operation": "append", - "session_id": "sid-native-append-error", - "ok": False, - "unsupported_count": 0, - "error_count": 1, - "stage_results": [ - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": False, - "error": "duplex_stage_handoff_target_missing:tts", - }, - }, - ], - } - engine = FakeEngineClient(append_result=control_result) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-append-error")) - ws.put({"type": "input_audio_buffer.append", "audio": "AAAA", "format": "pcm_f32le"}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert "response.output_audio.delta" not in ws.sent_types() - error = next(m for m in ws.sent if m.get("type") == "error") - assert error["code"] == "runtime_append_failed" - assert error["runtime_control"]["error_count"] == 1 - - -@pytest.mark.asyncio -async def test_duplex_handler_session_update_control_failure_rolls_back_config(): - control_result = { - "operation": "signal", - "session_id": "sid-control-signal-fail", - "ok": False, - "unsupported_count": 0, - "error_count": 1, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": {"supported": False, "error": "stage signal failed"}, - } - ], - } - engine = FakeEngineClient(signal_result=control_result) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-control-signal-fail")) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": {"temperature": 0.2}, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert "session.updated" not in ws.sent_types() - error = next(m for m in ws.sent if m.get("type") == "error") - assert error["code"] == "runtime_signal_failed" - assert error["runtime_control"]["error_count"] == 1 - signaled_config = engine.signal_session_configs[-1] - assert signaled_config is not None - assert signaled_config["temperature"] == 0.2 - - -@pytest.mark.asyncio -async def test_native_session_update_commits_config_only_after_runtime_ack(): - class BlockingUpdateEngine(FakeEngineClient): - def __init__(self): - super().__init__() - self.update_started = asyncio.Event() - self.release_update = asyncio.Event() - - async def signal_duplex_turn_async(self, session_id: str, **kwargs): - if kwargs.get("event") == "session.update": - self.update_started.set() - await self.release_update.wait() - return await super().signal_duplex_turn_async(session_id, **kwargs) - - engine = BlockingUpdateEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([])), - ) - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - ws.put(_native_realtime_session_update("sid-update-two-phase")) - ws.put(_native_server_vad_update("update-two-phase")) - - handler_task = asyncio.create_task(handler.handle_session(ws, realtime_protocol=protocol)) - await asyncio.wait_for(engine.update_started.wait(), timeout=10) - live_session = handler._registry.get("sid-update-two-phase") - assert live_session is not None - assert live_session.config.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value - assert protocol._pending_turn_detection_update is not None - assert protocol._turn_detection is None - - protocol.encode_outbound_event({"type": "error", "code": "unrelated", "error": "unrelated"}) - assert protocol._pending_turn_detection_update is not None - - engine.release_update.set() - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - updated = next(message for message in reversed(ws.sent) if message.get("type") == "session.updated") - assert updated["session"]["overlap_policy"] == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value - assert protocol._turn_detection is not None - - -@pytest.mark.parametrize( - ("append_fails", "expected_code"), - [(True, "session_update_aborted"), (False, "instructions_update_unsupported")], -) -@pytest.mark.asyncio -async def test_native_server_vad_update_rejection_does_not_stall(append_fails, expected_code): - append_result = {"operation": "append", "ok": False, "error_count": 1} if append_fails else None - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(FakeEngineClient(append_result=append_result)), - config_timeout_s=0.1, - idle_timeout_s=1, - server_vad_backend_provider=FakeServerVADProvider(FakeServerVADBackend([])), - ) - ws = TimedWebSocket() - protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] - create = _native_realtime_session_update("sid-rejected-vad-update") - create["session"]["extra_body"]["auto_response"] = append_fails - ws.put(create) - if append_fails: - ws.put({"type": "input_audio_buffer.append", "audio": _pcm_f32_b64(16_000), "format": "pcm_f32le"}) - update = _native_server_vad_update("rejected-vad-update") - if not append_fails: - update["session"]["instructions"] = "You are now a pirate." - ws.put(update) - ws.put({"type": "session.close"}) - - await asyncio.wait_for(handler.handle_session(ws, realtime_protocol=protocol), timeout=10) - - errors = [message["error"] for message in ws.sent if isinstance(message.get("error"), dict)] - error = next(error for error in errors if error.get("code") == expected_code) - assert error["event_id"] == "rejected-vad-update" - assert protocol._pending_turn_detection_update is None - - -@pytest.mark.asyncio -async def test_native_session_update_waits_for_prior_append_effect(): - class OrderedUpdateEngine(FakeEngineClient): - def __init__(self): - super().__init__() - self.append_started = asyncio.Event() - self.release_append = asyncio.Event() - self.update_started = asyncio.Event() - self.effects: list[str] = [] - - async def append_duplex_input_async(self, session_id: str, **kwargs): - self.effects.append("append.started") - self.append_started.set() - await self.release_append.wait() - kwargs.pop("expected_epoch", None) - result = await super().append_duplex_input_async(session_id, **kwargs) - self.effects.append("append.done") - return result - - async def signal_duplex_turn_async(self, session_id: str, **kwargs): - if kwargs.get("event") == "session.update": - self.effects.append("session.update") - self.update_started.set() - return await super().signal_duplex_turn_async(session_id, **kwargs) - - engine = OrderedUpdateEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-ordered-update")) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000, value=0.05), - "format": "pcm_f32le", - } - ) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": {"temperature": 0.2}, - } - ) - - handler_task = asyncio.create_task(handler.handle_session(ws)) - await asyncio.wait_for(engine.append_started.wait(), timeout=1) - try: - await asyncio.sleep(0.05) - assert not engine.update_started.is_set() - finally: - engine.release_append.set() - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - assert engine.effects.index("append.done") < engine.effects.index("session.update") - - -@pytest.mark.asyncio -async def test_failed_native_append_prevents_queued_append_from_running(): - failed_result = { - "operation": "append", - "session_id": "sid-append-chain-failure", - "ok": False, - "unsupported_count": 0, - "error_count": 1, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": {"supported": False, "error": "append failed"}, - } - ], - } - - class BlockingFailureEngine(FakeEngineClient): - def __init__(self): - super().__init__(append_result=failed_result) - self.first_started = asyncio.Event() - self.release_first = asyncio.Event() - - async def append_duplex_input_async(self, session_id: str, **kwargs): - if not self.appended: - self.first_started.set() - await self.release_first.wait() - kwargs.pop("expected_epoch", None) - return await super().append_duplex_input_async(session_id, **kwargs) - - engine = BlockingFailureEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - create = _native_session_create("sid-append-chain-failure") - create["session"]["extra_body"]["auto_response"] = True - ws.put(create) - for _ in range(2): - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16_000, value=0.05), - "format": "pcm_f32le", - } - ) - - handler_task = asyncio.create_task(handler.handle_session(ws)) - await asyncio.wait_for(engine.first_started.wait(), timeout=1) - engine.release_first.set() - await asyncio.sleep(0.1) - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - assert len(engine.appended) == 1 - - -@pytest.mark.asyncio -async def test_failed_final_append_discards_retained_audio_without_retry(): - failed_result = { - "operation": "append", - "session_id": "sid-final-append-retry", - "ok": False, - "unsupported_count": 0, - "error_count": 1, - "stage_results": [], - } - engine = FakeEngineClient(append_result=failed_result) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - - def retry_after_runtime_error(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") == "error" and data.get("code") == "runtime_append_failed": - ws.put({"type": "response.create"}) - - ws = TimedWebSocket(on_send=retry_after_runtime_error) - ws.put(_native_session_create("sid-final-append-retry")) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(8_000, value=0.05), - "format": "pcm_f32le", - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - ws.put({"type": "response.create"}) - - handler_task = asyncio.create_task(handler.handle_session(ws)) - for _ in range(100): - state = handler._serving_runtime_adapter.session_states.get("sid-final-append-retry") - if len(engine.appended) >= 1 and state is not None and state.committed_audio_payload is None: - break - await asyncio.sleep(0.01) - state = handler._serving_runtime_adapter.session_states["sid-final-append-retry"] - assert state.committed_audio_payload is None - assert state.committed_audio_reserved_bytes == 0 - assert handler._registry.get("sid-final-append-retry").pending_input_bytes == 0 - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - assert len(engine.appended) == 1 - - -@pytest.mark.asyncio -async def test_duplex_handler_signal_unsupported_workers_with_data_plane_ack_is_not_fatal(): - control_result = { - "operation": "signal", - "session_id": "sid-control-signal-data-plane", - "ok": False, - "unsupported_count": 2, - "error_count": 0, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": {"supported": False, "reason": "worker_duplex_signal_not_implemented"}, - }, - { - "stage_id": 1, - "replica_id": 0, - "result": {"supported": False, "reason": "worker_duplex_signal_not_implemented"}, - }, - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": True, - "data_plane_signal": True, - "event": "barge_in", - "epoch": 1, - }, - }, - ], - } - engine = FakeEngineClient(signal_result=control_result) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-control-signal-data-plane")) - ws.put({"type": "turn.signal", "event": "barge_in"}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert "runtime_signal_failed" not in {m.get("code") for m in ws.sent} - assert "input.cancelled" in ws.sent_types() - runtime_control = next(m for m in ws.sent if m.get("type") == "runtime.control") - assert runtime_control["result"]["unsupported_count"] == 2 - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_open_unsupported_fails_session_create(): - control_result = { - "operation": "open", - "session_id": "sid-native-busy", - "ok": True, - "unsupported_count": 1, - "error_count": 0, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": {"supported": False, "reason": "native_duplex_session_busy"}, - } - ], - } - engine = FakeEngineClient(control_result=control_result) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-busy")) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "runtime_open_unsupported" - assert ws.sent[0]["runtime_control"]["unsupported_count"] == 1 - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_rejects_engine_without_fence_contract(): - class LegacyEngineClient(FakeEngineClient): - async def open_duplex_session_async( # type: ignore[override] - self, - session_id: str, - *, - session_mode: str = "duplex", - capabilities: dict[str, object] | None = None, - session_config: dict[str, object] | None = None, - timeout: float | None = None, - ) -> Any: - return await super().open_duplex_session_async( - session_id, - session_mode=session_mode, - capabilities=capabilities, - session_config=session_config, - timeout=timeout, - ) - - engine = LegacyEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-no-fence-contract")) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "runtime_contract_invalid" - assert "open_duplex_session_async" in ws.sent[0]["error"] - assert "fence" in ws.sent[0]["error"] - assert engine.opened == [] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_audio_append_does_not_retain_pending_audio(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-no-pending")) - short_chunk = base64.b64encode(b"\x00" * (800 * 4)).decode("ascii") - ws.put({"type": "input_audio_buffer.append", "audio": short_chunk, "format": "pcm_f32le"}) - ws.put({"type": "input.cancel", "reason": "barge_in"}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - cancelled = next(m for m in ws.sent if m.get("type") == "input.cancelled") - assert cancelled["cancelled"] == {"text_chunks": 0, "audio_chunks": 0} - assert cancelled["epoch"] == 1 - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_rejects_text_append_before_runtime_call(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-text")) - ws.put({"type": "input.text.append", "text": "hello"}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - error = next(m for m in ws.sent if m.get("type") == "error") - assert error["code"] == "native_text_append_unsupported" - assert engine.appended == [] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_separates_public_and_runtime_config(monkeypatch): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - - async def fake_resolve_ref_audio(_, *, model_config): - del model_config - return [0.25] * 1600, 16000 - - monkeypatch.setattr( - "vllm_omni.model_executor.models.minicpmo_4_5.duplex.MiniCPMO45NativeDuplexServingAdapter.resolve_ref_audio", - fake_resolve_ref_audio, - ) - event = _native_session_create("sid-native-ref-audio", modalities=["text", "audio"]) - event["session"]["ref_audio"] = "data:audio/wav;base64,AAAA" - ws = TimedWebSocket() - ws.put(event) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert ws.sent_types()[0] == "session.created" - opened_config = engine.opened_configs[0] - opened_runtime_config = engine.opened_runtime_configs[0] - extra_body = opened_config["extra_body"] - assert "ref_audio_data" not in extra_body - assert "duplex_stage_max_tokens" not in extra_body - assert opened_runtime_config["ref_audio_format"] == "pcm_f32le" - assert opened_runtime_config["ref_audio_sample_rate_hz"] == 16000 - assert opened_runtime_config["duplex_stage_max_tokens"] == {"0": 20, "1": 8192} - assert opened_runtime_config["duplex_stage_sampling_params"]["0"]["top_k"] == 20 - assert opened_runtime_config["duplex_stage_sampling_params"]["0"]["top_p"] == 0.8 - # The Talker YAML carries upstream's min_new_token=50 for chat TTS; a duplex - # chunk is 26 codec samples, so check_stop would never release the request. - assert opened_runtime_config["duplex_stage_sampling_params"]["1"] == {"min_tokens": 0} - assert base64.b64decode(opened_runtime_config["ref_audio_data"]) == struct.pack("<1600f", *([0.25] * 1600)) - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_rejects_client_runtime_config(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - event = _native_session_create("sid-native-private-config") - event["session"]["extra_body"]["duplex_stage_sampling_params"] = {"0": {"temperature": 0.0}} - ws = TimedWebSocket() - ws.put(event) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "invalid_duplex_runtime_config" - assert engine.opened == [] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_audio_output_requires_ref_audio(monkeypatch): - monkeypatch.setattr( - MiniCPMO45NativeDuplexServingAdapter, - "_load_native_tokenizer", - staticmethod(lambda model_config: None), - ) - - event = _native_session_create( - "sid-native-ref-required", - modalities=["text", "audio"], - ) - config = DuplexSessionConfig.from_event(event) - - with pytest.raises(ServingRuntimeConfigError) as exc_info: - await MiniCPMO45NativeDuplexServingAdapter.prepare_runtime_config( - config, - model_config=SimpleNamespace(model="openbmb/MiniCPM-o-4_5"), - ) - - assert exc_info.value.code == "ref_audio_required" - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_open_session_maps_missing_ref_audio_to_typed_error(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put( - _native_session_create( - "sid-native-ref-required-open", - modalities=["text", "audio"], - ) - ) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "ref_audio_required" - assert engine.opened == [] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_text_only_omits_ref_audio_when_client_does_not_provide_it(monkeypatch): - monkeypatch.setattr( - MiniCPMO45NativeDuplexServingAdapter, - "_load_native_tokenizer", - staticmethod(lambda model_config: None), - ) - - event = _native_session_create("sid-native-default-ref", modalities=["text"]) - config = DuplexSessionConfig.from_event(event) - - runtime_config = await MiniCPMO45NativeDuplexServingAdapter.prepare_runtime_config( - config, - model_config=SimpleNamespace(model="openbmb/MiniCPM-o-4_5"), - ) - - assert "ref_audio_data" not in runtime_config - assert "ref_audio_format" not in runtime_config - assert "ref_audio_sample_rate_hz" not in runtime_config - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_preserves_ref_audio_channels_until_normalize(monkeypatch): - class FakeMediaConnector: - def __init__(self, **kwargs): - self.kwargs = kwargs - - async def fetch_audio_async(self, ref_audio): - assert ref_audio == "data:audio/wav;base64,AAAA" - return np.tile(np.array([[0.25, -0.25], [0.5, -0.5]], dtype=np.float32), (800, 1)), 16000 - - monkeypatch.setattr( - "vllm_omni.model_executor.models.minicpmo_4_5.duplex.adapter.MediaConnector", - FakeMediaConnector, - ) - - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - event = _native_session_create("sid-native-ref-audio-stereo", modalities=["text", "audio"]) - event["session"]["ref_audio"] = "data:audio/wav;base64,AAAA" - ws = TimedWebSocket() - ws.put(event) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - runtime_config = engine.opened_runtime_configs[0] - assert base64.b64decode(runtime_config["ref_audio_data"]) == struct.pack("<1600f", *([0.0] * 1600)) - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_rejects_ref_audio_path(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - event = _native_session_create("sid-native-ref-path") - event["session"]["extra_body"] = { - "native_duplex": True, - "ref_audio_path": "/tmp/ref.wav", - } - ws = TimedWebSocket() - ws.put(event) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["error"] - assert ws.sent[0]["code"] == "unsupported_ref_audio_path" - assert engine.opened == [] - - -def test_minicpmo_native_duplex_explicit_barge_in_request_interrupts(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-native-explicit-barge-disabled", - config=DuplexSessionConfig(overlap_policy=DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value), - capabilities=minicpmo45_native_capabilities(), - ) - payload = { - "type": "audio", - "audio": _pcm_f32_b64(16000, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - "is_speech": True, - } - - decision = handler._overlap_decision( - session, - {"force_barge_in": True}, - payload, - ) - - assert decision == { - "action": "barge_in", - "reason": "client_force_barge_in", - "duration_ms": 1000, - "buffer_audio": True, - } - - -@pytest.mark.asyncio -async def test_native_server_vad_complete_audio_item_reaches_runtime(): - engine = FakeEngineClient() - backend = FakeServerVADBackend([0.9] * 10) - ws = TimedWebSocket() - create = _native_realtime_session_update("sid-native-complete-audio") - create["session"]["turn_detection"] = {"type": "server_vad"} - create["session"]["input_audio_format"] = "pcm16" - ws.put(create) - ws.put( - { - "type": "conversation.item.create", - "item": { - "id": "item-native-audio", - "type": "message", - "role": "user", - "content": [{"type": "input_audio", "audio": _pcm16_b64(2400)}], - }, - } - ) - ws.put({"type": "response.create"}) - ws.put({"type": "session.close"}) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - server_vad_backend_provider=FakeServerVADProvider(backend), - ) - - await handler.handle_realtime_session(ws) # type: ignore[arg-type] - - assert any(final for _, _, _, final in engine.appended) - assert "error" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_native_server_vad_uses_shared_pipeline_and_emits_boundaries(): - backend = FakeServerVADBackend([0.9] * 10 + [0.0] * 50) - engine = FakeEngineClient() - - def commit_after_speech_stops(ws: TimedWebSocket, event: dict[str, Any]) -> None: - if event.get("type") == "input_audio_buffer.speech_stopped": - ws.put({"type": "input_audio_buffer.commit", "final": True}) - ws.put({"type": "response.create"}) - ws.put({"type": "session.close"}) - - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - server_vad_backend_provider=FakeServerVADProvider(backend), - ) - ws = TimedWebSocket(on_send=commit_after_speech_stops) - create = _native_realtime_session_update("sid-native-shared-server-vad") - create["session"]["turn_detection"] = {"type": "server_vad"} - ws.put(create) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(60 * backend.frame_samples, value=0.0), - "format": "pcm_f32le", - "sample_rate_hz": 16_000, - } - ) - await asyncio.wait_for( - handler.handle_session(ws, realtime_protocol=NativeRealtimeSessionProtocol({})), - timeout=10, - ) - - boundary_events = [ - event - for event in ws.sent - if event.get("type") in {"input_audio_buffer.speech_started", "input_audio_buffer.speech_stopped"} - ] - assert [event["type"] for event in boundary_events] == [ - "input_audio_buffer.speech_started", - "input_audio_buffer.speech_stopped", - ] - assert boundary_events[0]["item_id"] == boundary_events[1]["item_id"] - created_session = next(event["session"] for event in ws.sent if event["type"] == "session.created") - assert created_session["audio"]["input"]["turn_detection"]["interrupt_response"] is True - assert created_session["overlap_policy"] == "barge_in_on_speech" - assert backend.calls == 60 - assert any(final for _, _, _, final in engine.appended) - assert not any(event.get("error", {}).get("code") == "server_vad_manual_commit_unsupported" for event in ws.sent) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "barge_in_event", - [ - {"type": "barge_in"}, - {"type": "turn.signal", "event": "barge_in"}, - ], -) -async def test_minicpmo_native_duplex_accepts_explicit_barge_in_control( - barge_in_event: dict[str, object], -): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-barge-control-disabled")) - ws.put(barge_in_event) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert not any(message.get("code") == "barge_in_unsupported" for message in ws.sent) - assert ("sid-native-barge-control-disabled", "barge_in") in engine.signals - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_explicit_non_speech_stays_listening_without_append(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-silence")) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": "AAAA", - "format": "pcm_f32le", - "is_speech": False, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert engine.appended == [] - assert "response.listen" in ws.sent_types() - assert "response.created" not in ws.sent_types() - assert "response.output_audio.delta" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_drains_data_plane_stream_until_done(): - def _stage_output(samples: int, *, finished: bool): - return SimpleNamespace( - request_id="duplex-sid-native-stream-e0-stage0-s1", - finished=finished, - outputs=[ - SimpleNamespace( - text="", - multimodal_output={ - "audio": np.zeros(samples, dtype=np.float32), - "sr": 24000, - }, - ) - ], - ) - - control_result = { - "operation": "append", - "session_id": "sid-native-stream", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": { - "supported": True, - "implementation_level": "model_native_duplex", - "data_plane_append": True, - "request_id": "duplex-sid-native-stream-e0-stage0-s1", - "response_stage_id": 1, - }, - } - ], - "data_plane_outputs": [_stage_output(10, finished=False)], - } - engine = FakeEngineClient( - append_result=control_result, - collect_outputs=[ - [_stage_output(20, finished=False)], - [_stage_output(30, finished=True)], - ], - ) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - - def close_on_second_created(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if ws.sent_types().count("response.created") >= 2: - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=close_on_second_created) - ws.put(_native_session_create("sid-native-stream")) - ws.put({"type": "input_audio_buffer.append", "audio": "AAAA", "format": "pcm_f32le"}) - - await handler.handle_session(ws) - - deltas = [m for m in ws.sent if m.get("type") == "response.output_audio.delta"] - assert [m["audio"] for m in deltas] == ["wav-10", "wav-10", "wav-10"] - assert [m["end_of_turn"] for m in deltas] == [False, False, True] - assert len([m for m in ws.sent if m.get("type") == "response.done"]) == 1 - assert engine.collected == [ - ("duplex-sid-native-stream-e0-stage0-s1", 1), - ("duplex-sid-native-stream-e0-stage0-s1", 1), - ] - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_restarts_drain_when_append_races_idle_exit( - monkeypatch: pytest.MonkeyPatch, -): - request_id = "duplex-sid-native-late-output-e0-stage0" - late_output = object() - engine = FakeEngineClient( - collect_outputs=[[], [], [], [late_output]], - collect_delay_s=0.05, - ) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - session = DuplexSession( - session_id="sid-native-late-output", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - session.bind_request(request_id) - native = handler._runtime_session_state(session) - projected_batches: list[object] = [] - projected = asyncio.Event() - - async def project_late_output(_send_json, result, **_kwargs): - projected_batches.append(result) - projected.set() - session.close() - return None, True - - monkeypatch.setattr(handler, "_send_native_duplex_events", project_late_output) - result = { - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - } - } - ] - } - - assert ( - await handler._start_native_data_plane_stream_task( - None, - result, - session=session, - expected_epoch=session.epoch, - ) - is True - ) - first_drain = native.data_plane_task - assert first_drain is not None - await asyncio.sleep(0.01) - - assert ( - await handler._start_native_data_plane_stream_task( - None, - result, - session=session, - expected_epoch=session.epoch, - ) - is False - ) - await asyncio.wait_for(projected.wait(), timeout=1) - await asyncio.wait_for(first_drain, timeout=1) - - assert len(engine.collected) == 4 - assert projected_batches == [{"data_plane_outputs": [late_output]}] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_ignores_outputs_after_data_plane_turn_done(): - request_id = "duplex-sid-native-post-done-e0-stage0-s1" - - def _stage_output(samples: int, *, turn_end: bool = False): - return SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[ - SimpleNamespace( - text="hello", - multimodal_output={ - "audio": np.zeros(samples, dtype=np.float32), - "sr": 24000, - "meta": { - "turn_end": turn_end, - "duplex_turn_id": 0, - "duplex_epoch": 0, - }, - }, - ) - ], - ) - - control_result = { - "operation": "append", - "session_id": "sid-native-post-done", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": { - "supported": True, - "implementation_level": "model_native_duplex", - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - }, - } - ], - "data_plane_outputs": [_stage_output(10)], - } - engine = FakeEngineClient( - append_result=control_result, - collect_outputs=[ - [_stage_output(20, turn_end=True)], - [_stage_output(30)], - ], - ) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - - def close_on_done(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") == "response.done": - ws.put({"type": "session.close"}) - - event = _native_session_create("sid-native-post-done") - event["session"]["extra_body"]["auto_response"] = True - ws = TimedWebSocket(on_send=close_on_done) - ws.put(event) - ws.put({"type": "input_audio_buffer.append", "audio": "AAAA", "format": "pcm_f32le"}) - - await handler.handle_session(ws) - - assert ws.sent_types().count("response.created") == 1 - assert ws.sent_types().count("response.done") == 1 - deltas = [m for m in ws.sent if m.get("type") == "response.output_audio.delta"] - assert [m["audio"] for m in deltas] == ["wav-10", "wav-10"] - assert engine.collected - assert all(collected == (request_id, 1) for collected in engine.collected) - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_accepts_next_append_after_response_done(): - def _stage_output(*, finished: bool): - return SimpleNamespace( - request_id="duplex-sid-native-next-turn-e0-stage0-s1", - finished=finished, - outputs=[ - SimpleNamespace( - text="", - multimodal_output={ - "audio": np.zeros(10, dtype=np.float32), - "sr": 24000, - }, - ) - ], - ) - - control_result = { - "operation": "append", - "session_id": "sid-native-next-turn", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": { - "supported": True, - "implementation_level": "model_native_duplex", - "data_plane_append": True, - "request_id": "duplex-sid-native-next-turn-e0-stage0-s1", - "response_stage_id": 1, - }, - } - ], - "data_plane_outputs": [_stage_output(finished=True)], - } - engine = FakeEngineClient(append_result=control_result) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - done_count = 0 - - def send_next_turn_or_close(ws: TimedWebSocket, data: dict[str, Any]) -> None: - nonlocal done_count - if data.get("type") != "response.done": - return - done_count += 1 - if done_count == 1: - ws.put({"type": "input_audio_buffer.append", "audio": "BBBB", "format": "pcm_f32le"}) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - ws.put({"type": "response.create"}) - else: - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=send_next_turn_or_close, receive_timeout_s=5) - create = _native_session_create("sid-native-next-turn") - create["session"]["idle_timeout_s"] = 5 - ws.put(create) - ws.put({"type": "input_audio_buffer.append", "audio": "AAAA", "format": "pcm_f32le"}) - - await handler.handle_session(ws) - - assert len(engine.appended) >= 2, [ - (message.get("type"), message.get("code"), message.get("reason")) for message in ws.sent - ] - assert len([m for m in ws.sent if m.get("type") == "response.done"]) == 2 - - -@pytest.mark.asyncio -async def test_minicpmo_native_auto_response_commit_finalizes_after_streamed_full_chunk(): - control_result = { - "operation": "append", - "session_id": "sid-native-auto-finalize", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [], - } - engine = FakeEngineClient(control_result=control_result) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - event = _native_session_create("sid-native-auto-finalize") - event["session"]["extra_body"]["auto_response"] = True - ws.put(event) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert len(engine.appended) == 1 - assert engine.appended[0][3] is False - assert "response.created" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_minicpmo_native_auto_response_commit_finalizes_with_active_data_plane_stream(): - request_id = "duplex-sid-native-auto-active-stream-stage0" - control_result = { - "operation": "append", - "session_id": "sid-native-auto-active-stream", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - } - } - ], - } - engine = FakeEngineClient( - control_result=control_result, - collect_outputs=[[], [], []], - collect_delay_s=0.05, - ) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - event = _native_session_create("sid-native-auto-active-stream") - event["session"]["extra_body"]["auto_response"] = True - ws.put(event) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert len(engine.appended) == 1 - assert engine.appended[0][3] is False - assert "response.created" not in ws.sent_types() - - -@pytest.mark.asyncio -async def test_minicpmo_native_auto_response_accepts_realtime_commit_after_stream_consumed(): - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket() - event = _native_session_create("sid-native-auto-consumed-stream") - event["session"]["extra_body"]["auto_response"] = True - ws.put(event) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - ws.put( - { - "type": "input_audio_buffer.commit", - "final": True, - "item_id": "item-consumed-stream", - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert not any(message.get("code") == "input_audio_buffer_empty" for message in ws.sent) - committed = next(message for message in ws.sent if message.get("type") == "input.committed") - assert committed["item_id"] == "item-consumed-stream" - assert committed["native_audio"] is True - - -@pytest.mark.asyncio -async def test_minicpmo_native_auto_response_post_response_commit_ignores_stale_progress_state(): - session_id = "sid-native-auto-post-response" - engine = FakeEngineClient() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") != "session.created": - return - session = handler._registry.get(session_id) - assert session is not None - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(8000), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - ws.put({"type": "session.close"}) - - ws = TimedWebSocket(on_send=on_send) - create = _native_session_create(session_id) - create["session"]["extra_body"]["auto_response"] = True - ws.put(create) - - await handler.handle_session(ws) - - assert len(engine.appended) == 1 - assert engine.appended[0][3] is True - assert "input.committed" in ws.sent_types() - assert not any(message.get("code") == "response_already_active" for message in ws.sent) - - -def test_realtime_overlap_commit_uses_current_generation_identity(): - handler, session = _auto_response_context("sid-realtime-reserved-input-turn") - session.turn_id = 1 - session.bind_response_turn(1) - session.turn_id = 3 - committed = handler._commit_native_audio_input(session) - - assert committed.turn_id == 3 - assert session.turn_id == 3 - - -@pytest.mark.asyncio -async def test_minicpmo_native_auto_response_keeps_request_bound_for_segment_continuation(): - request_id = "duplex-sid-native-auto-continuation-e0-stage0" - control_result = { - "operation": "append", - "session_id": "sid-native-auto-continuation", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - } - } - ], - } - terminal_segment = _duplex_tts_output( - request_id=request_id, - samples=0, - finished=True, - tts_is_last_chunk=True, - token_ids=[151645], - ) - engine = FakeEngineClient( - control_result=control_result, - collect_outputs=[[terminal_segment], [], []], - collect_delay_s=0.05, - ) - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") != "session.created": - return - session = handler._registry.get("sid-native-auto-continuation") - assert session is not None - session.capabilities.chunk_period_ms = 50 - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(800), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - ws.put({"type": "input_audio_buffer.commit", "final": True}) - - ws = TimedWebSocket(on_send=on_send) - event = _native_session_create("sid-native-auto-continuation") - event["session"]["extra_body"]["auto_response"] = True - ws.put(event) - - await handler.handle_session(ws) - - assert len(engine.appended) >= 3 - _, mode, payload, final = engine.appended[2] - assert mode == "append_audio_chunk" - assert payload["duplex_turn_id"] == 0 - assert final is False - - -@pytest.mark.asyncio -async def test_minicpmo_native_auto_response_real_input_preempts_unsent_silence(): - request_id = "duplex-sid-native-real-before-silence-e0-stage0" - - class SequencedAppendEngine(FakeEngineClient): - def __init__(self) -> None: - control_result = { - "operation": "append", - "session_id": "sid-native-real-before-silence", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - } - } - ], - } - terminal_segment = _duplex_tts_output( - request_id=request_id, - samples=0, - finished=True, - tts_is_last_chunk=True, - token_ids=[151645], - ) - super().__init__( - control_result=control_result, - collect_outputs=[[terminal_segment], [], []], - collect_delay_s=0.01, - ) - self.append_sequence: list[str] = [] - self.first_real_started = asyncio.Event() - self.second_real_started = asyncio.Event() - self.real_count = 0 - - async def append_duplex_input_async(self, session_id: str, **kwargs): - kwargs.pop("expected_epoch", None) - result = await super().append_duplex_input_async(session_id, **kwargs) - payload = kwargs.get("payload") - is_silence = ( - isinstance(payload, dict) - and payload.get("audio") == OmniDuplexSessionHandler._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO - ) - if is_silence: - self.append_sequence.append("silence") - else: - self.real_count += 1 - self.append_sequence.append(f"real-{self.real_count}") - if self.real_count == 1: - self.first_real_started.set() - elif self.real_count == 2: - self.second_real_started.set() - return result - - session_created = asyncio.Event() - engine = SequencedAppendEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") != "session.created": - return - session = handler._registry.get("sid-native-real-before-silence") - assert session is not None - session.capabilities.chunk_period_ms = 50 - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(800), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - session_created.set() - - ws = TimedWebSocket(on_send=on_send, receive_timeout_s=2.0) - event = _native_session_create("sid-native-real-before-silence") - event["session"]["extra_body"]["auto_response"] = True - ws.put(event) - handler_task = asyncio.create_task(handler.handle_session(ws)) - - await asyncio.wait_for(session_created.wait(), timeout=1) - await asyncio.wait_for(engine.first_real_started.wait(), timeout=1) - await asyncio.sleep(0.02) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(800, value=0.07), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - await asyncio.wait_for(engine.second_real_started.wait(), timeout=1) - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - assert engine.append_sequence[:2] == ["real-1", "real-2"] - - -@pytest.mark.asyncio -async def test_minicpmo_native_auto_response_real_input_waits_for_submitted_silence_tail(): - request_id = "duplex-sid-native-silence-tail-e0-stage0" - - class SubmittedSilenceEngine(FakeEngineClient): - def __init__(self) -> None: - control_result = { - "operation": "append", - "session_id": "sid-native-silence-tail", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - } - } - ], - } - terminal_segment = _duplex_tts_output( - request_id=request_id, - samples=0, - finished=True, - tts_is_last_chunk=True, - token_ids=[151645], - ) - super().__init__( - control_result=control_result, - collect_outputs=[[terminal_segment], [], []], - collect_delay_s=0.01, - ) - self.silence_append_started = asyncio.Event() - self.release_silence_append = asyncio.Event() - self.append_sequence: list[str] = [] - self.real_append_count = 0 - - async def append_duplex_input_async(self, session_id: str, **kwargs): - kwargs.pop("expected_epoch", None) - payload = kwargs.get("payload") - is_silence_continuation = ( - isinstance(payload, dict) - and payload.get("format") == "pcm_f32le" - and payload.get("audio") == OmniDuplexSessionHandler._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO - ) - if is_silence_continuation: - result = await super().append_duplex_input_async(session_id, **kwargs) - self.append_sequence.append("silence-started") - self.silence_append_started.set() - await self.release_silence_append.wait() - self.append_sequence.append("silence-returned") - return result - result = await super().append_duplex_input_async(session_id, **kwargs) - if ( - isinstance(payload, dict) - and payload.get("audio") != OmniDuplexSessionHandler._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO - ): - self.real_append_count += 1 - self.append_sequence.append(f"real-{self.real_append_count}") - return result - - engine = SubmittedSilenceEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - - def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") != "session.created": - return - session = handler._registry.get("sid-native-silence-tail") - assert session is not None - session.capabilities.chunk_period_ms = 50 - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(800, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - - ws = TimedWebSocket(on_send=on_send, receive_timeout_s=2.0) - event = _native_session_create("sid-native-silence-tail") - event["session"]["extra_body"]["auto_response"] = True - ws.put(event) - handler_task = asyncio.create_task(handler.handle_session(ws)) - - await asyncio.wait_for(engine.silence_append_started.wait(), timeout=1) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000, value=0.07), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - await asyncio.sleep(0.05) - - assert engine.real_append_count == 1 - assert engine.append_sequence == ["real-1", "silence-started"] - - engine.release_silence_append.set() - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - assert engine.real_append_count == 2 - assert engine.append_sequence == ["real-1", "silence-started", "silence-returned", "real-2"] - - -@pytest.mark.asyncio -async def test_minicpmo_native_auto_response_preserves_silence_continuations_across_append_tails(): - request_id = "duplex-sid-native-continuation-tails-e0-stage0" - - class SequencedAppendEngine(FakeEngineClient): - def __init__(self) -> None: - control_result = { - "operation": "append", - "session_id": "sid-native-continuation-tails", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - } - } - ], - } - super().__init__(control_result=control_result) - self.real_append_started = asyncio.Event() - self.release_real_append = asyncio.Event() - self.first_silence_started = asyncio.Event() - self.release_first_silence = asyncio.Event() - self.second_silence_started = asyncio.Event() - self.release_collect = asyncio.Event() - self.silence_count = 0 - - async def append_duplex_input_async(self, session_id: str, **kwargs): - kwargs.pop("expected_epoch", None) - payload = kwargs.get("payload") - assert isinstance(payload, dict) - result = await super().append_duplex_input_async(session_id, **kwargs) - if payload.get("audio") != OmniDuplexSessionHandler._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO: - self.real_append_started.set() - await self.release_real_append.wait() - return result - - self.silence_count += 1 - if self.silence_count == 1: - self.first_silence_started.set() - await self.release_first_silence.wait() - else: - self.second_silence_started.set() - return result - - async def collect_duplex_data_plane_outputs_async( - self, - request_id: str, - *, - response_stage_id: int | None = None, - timeout: float | None = None, - ) -> list[object]: - del timeout - self.collected.append((request_id, response_stage_id)) - await self.release_collect.wait() - return [] - - session_created = asyncio.Event() - - def on_send(_ws: TimedWebSocket, data: dict[str, Any]) -> None: - if data.get("type") == "session.created": - session_created.set() - - engine = SequencedAppendEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) - ws = TimedWebSocket(on_send=on_send, receive_timeout_s=2.0) - event = _native_session_create("sid-native-continuation-tails") - event["session"]["extra_body"]["auto_response"] = True - ws.put(event) - ws.put( - { - "type": "input_audio_buffer.append", - "audio": _pcm_f32_b64(16000, value=0.05), - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - handler_task = asyncio.create_task(handler.handle_session(ws)) - - await asyncio.wait_for(session_created.wait(), timeout=1) - await asyncio.wait_for(engine.real_append_started.wait(), timeout=1) - session = handler._registry.get("sid-native-continuation-tails") - assert session is not None - session.capabilities.chunk_period_ms = 50 - session.bind_request(request_id) - response_id = session.begin_response(turn_id=session.turn_id) - native = handler._runtime_session_state(session) - scheduler = native.silence_continuation_scheduler - assert scheduler is not None - - payload = handler._native_silence_unit_payload() - payload["duplex_turn_id"] = session.turn_id - kwargs = { - "request_id": request_id, - "owner_id": f"response:{response_id}", - "response_id": response_id, - "response_owned": True, - "expected_epoch": session.epoch, - "expected_incarnation": session.incarnation, - "expected_model_turn_id": None, - "send_json": ws.send_json, - } - assert await scheduler(payload, **kwargs) is True - - next_schedule = asyncio.create_task(scheduler(payload, **kwargs)) - await asyncio.sleep(0.05) - assert not next_schedule.done() - - engine.release_real_append.set() - await asyncio.wait_for(engine.first_silence_started.wait(), timeout=1) - assert not next_schedule.done() - - engine.release_first_silence.set() - assert await asyncio.wait_for(next_schedule, timeout=1) is True - await asyncio.wait_for(engine.second_silence_started.wait(), timeout=1) - ws.put({"type": "session.close"}) - await asyncio.wait_for(handler_task, timeout=2) - - assert engine.silence_count == 2 - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_uses_segment_text_metadata_for_transcript_cursor(): - request_id = "duplex-sid-native-repeat-text-e0-stage0" - - def _stage_output( - text: str, - samples: int, - *, - segment_text: str, - turn_id: int = 0, - turn_end: bool = False, - token_ids: list[int] | None = None, - ): - return SimpleNamespace( - request_id=request_id, - finished=False, - outputs=[ - SimpleNamespace( - text=text, - token_ids=list(token_ids or []), - multimodal_output={ - "audio": np.zeros(samples, dtype=np.float32), - "sr": 24000, - "meta.llm_output_text_utf8": np.frombuffer(segment_text.encode("utf-8"), dtype=np.uint8), - "meta": {"duplex_turn_id": turn_id, "turn_end": turn_end}, - }, - ) - ], - ) - - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-native-repeat-text", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - sent: list[dict[str, Any]] = [] - - async def send_json(payload: dict[str, Any]) -> None: - sent.append(payload) - - async def emit(output: object) -> None: - result = {"data_plane_outputs": [output]} - for native_result in handler._serving_runtime_adapter.data_plane.project( - result, - context=_data_plane_context(session), - ): - await handler._send_one_native_duplex_event(send_json, native_result, session=session) - - await emit(_stage_output("same reply", 24000, segment_text="same reply")) - await emit( - _stage_output("same replysame reply", 24000, segment_text="same reply", turn_end=True, token_ids=[151645]) - ) - - await emit(_stage_output("same replysame replysame reply", 48000, segment_text="same reply", turn_id=1)) - await emit( - _stage_output( - "same replysame replysame replysame reply", - 48000, - segment_text="same reply", - turn_id=1, - turn_end=True, - token_ids=[151645], - ) - ) - - created = [m for m in sent if m.get("type") == "response.created"] - done = [m for m in sent if m.get("type") == "response.done"] - deltas = [m for m in sent if m.get("type") == "response.output_audio.delta" and m.get("audio")] - protocol = NativeRealtimeSessionProtocol(TimedWebSocket()) # type: ignore[arg-type] - realtime_events = [event for payload in sent for event in protocol._from_duplex_event(payload)] - transcript_deltas = [ - m for m in realtime_events if m.get("type") == "response.output_audio_transcript.delta" and m.get("delta") - ] - assert len(created) == 2 - assert len(done) == 2 - assert len(deltas) == 2 - assert [m["text"] for m in deltas] == ["same reply", "same reply"] - assert [m["audio"] for m in deltas] == ["wav-24000", "wav-24000"] - assert [m["delta"] for m in transcript_deltas] == ["same reply", "same reply"] - assert [m["response_id"] for m in deltas] == [m["response_id"] for m in created] - assert [m["response_id"] for m in done] == [m["response_id"] for m in created] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_continuous_speak_reuses_active_response_until_turn_end(): - request_id = "duplex-sid-native-continuous-speak-e0-stage0" - - def _stage_output( - text: str, - samples: int, - *, - turn_end: bool = False, - token_ids: list[int] | None = None, - ): - return SimpleNamespace( - request_id=request_id, - finished=turn_end, - outputs=[ - SimpleNamespace( - text=text, - token_ids=list(token_ids or []), - multimodal_output={}, - ) - ], - multimodal_output={ - "audio": np.zeros(samples, dtype=np.float32), - "sr": 24000, - "meta.duplex_turn_id": np.array([0], dtype=np.int32), - "meta.duplex_epoch": np.array([0], dtype=np.int32), - "meta.turn_end": np.array([int(turn_end)], dtype=np.int32), - "meta.tts_is_last_chunk": np.array([int(turn_end)], dtype=np.int32), - }, - ) - - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-native-continuous-speak", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - sent: list[dict[str, Any]] = [] - - async def send_json(payload: dict[str, Any]) -> None: - sent.append(payload) - - async def emit(output: object) -> None: - result = {"data_plane_outputs": [output]} - for native_result in handler._serving_runtime_adapter.data_plane.project( - result, - context=_data_plane_context(session), - ): - await handler._send_one_native_duplex_event(send_json, native_result, session=session) - - await emit(_stage_output("hello", 24000)) - await emit(_stage_output("hello again", 48000)) - - created = [m for m in sent if m.get("type") == "response.created"] - deltas = [m for m in sent if m.get("type") == "response.output_audio.delta" and m.get("audio")] - assert len(created) == 1 - assert len(deltas) == 2 - assert {m["response_id"] for m in deltas} == {created[0]["response_id"]} - assert "response.done" not in [m.get("type") for m in sent] - - await emit(_stage_output("hello again", 48000, turn_end=True, token_ids=[151645])) - - done = [m for m in sent if m.get("type") == "response.done"] - assert len(done) == 1 - assert done[0]["response_id"] == created[0]["response_id"] - assert session.active_response_id is None - assert session.turn_id == 1 - - -def test_response_timing_binds_latest_request_start_for_model_turn(): - session = DuplexSession( - session_id="sid-response-request-timing", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.mark_model_turn_request_started(0, 10.0) - session.mark_model_turn_request_started(0, 11.0) - session.begin_response(turn_id=0) - - assert session.mark_response_first_outputs( - observed_at_s=11.2, - has_text=True, - has_audio=False, - )["ttft_ms"] == pytest.approx(200.0) - - session.mark_model_turn_request_started(0, 12.0) - assert session.mark_response_first_outputs( - observed_at_s=12.3, - has_text=False, - has_audio=True, - )["ttfp_ms"] == pytest.approx(1300.0) - - -def test_response_tpot_fallback_ignores_single_token_segment(): - session = DuplexSession( - session_id="sid-response-tpot", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.begin_response(turn_id=0) - - first_metrics = session.accumulate_response_stage_metrics({"0": {"num_tokens_out": 1, "vllm_tpot_ms": 900.0}}) - assert "vllm_tpot_ms" not in first_metrics["0"] - - second_metrics = session.accumulate_response_stage_metrics({"0": {"num_tokens_out": 2, "vllm_tpot_ms": 15.0}}) - assert second_metrics["0"]["num_tokens_out"] == 3 - assert second_metrics["0"]["vllm_tpot_ms"] == 15.0 - - -def test_response_tpot_keeps_token_weighted_value_when_itls_exist(): - session = DuplexSession( - session_id="sid-response-tpot-multi-token", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.begin_response(turn_id=0) - metrics = session.accumulate_response_stage_metrics( - { - "0": { - "num_tokens_out": 4, - "vllm_tpot_ms": 10.0, - "vllm_itls_ms": [30.0], - } - } - ) - assert metrics["0"]["vllm_itls_ms"] == [30.0] - assert metrics["0"]["vllm_itl_ms"] == 30.0 - assert metrics["0"]["vllm_tpot_ms"] == 10.0 - - -@pytest.mark.asyncio -async def test_continuous_response_metrics_accumulate_only_owned_model_units(monkeypatch): - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-response-metrics", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.mark_model_turn_request_started(0, 10.0) - monkeypatch.setattr( - "vllm_omni.entrypoints.duplex.runtime_bridge.time.monotonic", - lambda: 10.25, - ) - sent: list[dict[str, Any]] = [] - - async def send_json(payload: dict[str, Any]) -> None: - sent.append(payload) - - async def emit(*, text: str, audio: str, tokens: int, ttft_ms: float, itls: list[float], tpot_ms: float) -> None: - await handler._send_one_native_duplex_event( - send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": "duplex-sid-response-metrics-e0-stage0", - "text": text, - "audio_data": audio, - "audio_format": "pcm16", - "audio_duration_ms": 100, - "end_of_turn": False, - "model_turn_id": 0, - "stage_metrics": { - "0": { - "num_tokens_out": tokens, - "stage_gen_time_ms": 80.0, - "vllm_ttft_ms": ttft_ms, - "vllm_tpot_ms": tpot_ms, - "vllm_itls_ms": itls, - } - }, - }, - session=session, - ) - - await emit(text="first", audio="audio-a", tokens=3, ttft_ms=50.0, itls=[10.0, 11.0], tpot_ms=10.5) - await emit(text="second", audio="audio-b", tokens=4, ttft_ms=70.0, itls=[12.0, 13.0, 14.0], tpot_ms=13.0) - - deltas = [payload for payload in sent if payload.get("type") == "response.output_audio.delta"] - assert len(deltas) == 2 - first_metrics = deltas[0]["vllm_omni"]["stage_metrics"]["0"] - second_metrics = deltas[1]["vllm_omni"]["stage_metrics"]["0"] - assert first_metrics["num_tokens_out"] == 3 - assert first_metrics["vllm_tpot_ms"] == pytest.approx(10.5) - assert deltas[0]["vllm_omni"]["response_request_metrics"] == { - "source": "server_monotonic_request_start", - "measurement_origin": { - "ttft": "native model-turn request execution start to first non-empty text output", - "ttfp": "native model-turn request execution start to first audio output", - }, - "ttft_ms": 250.0, - "ttfp_ms": 250.0, - } - assert second_metrics["num_tokens_out"] == 7 - assert second_metrics["vllm_ttft_ms"] == 50.0 - assert second_metrics["vllm_itls_ms"] == [10.0, 11.0, 12.0, 13.0, 14.0] - assert second_metrics["vllm_tpot_ms"] == pytest.approx(12.0) - - first_response_id = deltas[0]["response_id"] - session.end_response() - await emit(text="third", audio="audio-c", tokens=3, ttft_ms=90.0, itls=[15.0, 16.0], tpot_ms=15.5) - - third_delta = [payload for payload in sent if payload.get("type") == "response.output_audio.delta"][-1] - third_metrics = third_delta["vllm_omni"]["stage_metrics"]["0"] - assert third_delta["response_id"] != first_response_id - assert third_metrics["num_tokens_out"] == 3 - assert third_metrics["vllm_ttft_ms"] == 90.0 - assert third_metrics["vllm_itls_ms"] == [15.0, 16.0] - assert third_metrics["vllm_tpot_ms"] == pytest.approx(15.5) - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_drops_old_turn_audio_while_new_response_is_active(): - request_id = "duplex-sid-native-old-turn-audio-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-native-old-turn-audio", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - response_id = session.begin_response(turn_id=1) - session.turn_id = 2 - session.bind_request(request_id) - ws = TimedWebSocket() - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": request_id, - "text": "stale", - "audio_data": "old-audio", - "audio_format": "pcm16", - "audio_duration_ms": 100, - "end_of_turn": False, - "model_turn_id": 0, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - - assert close_reason is None - assert emitted is False - assert session.active_response_id == response_id - assert ws.sent == [] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_future_turn_audio_starts_next_response(): - request_id = "duplex-sid-native-future-turn-audio-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-native-future-turn-audio", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - first_response_id = session.begin_response(turn_id=0) - session.bind_request(request_id) - ws = TimedWebSocket() - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "tts", - "is_listen": False, - "data_plane_request_id": request_id, - "text": "next", - "audio_data": "next-audio", - "audio_format": "pcm16", - "audio_duration_ms": 100, - "end_of_turn": False, - "model_turn_id": 1, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - - assert close_reason is None - assert emitted is True - sent_types = [event["type"] for event in ws.sent] - assert sent_types == [ - "response.done", - "response.created", - "response.speak", - "response.output_audio.delta", - ] - assert ws.sent[0]["response_id"] == first_response_id - assert ws.sent[1]["response_id"] != first_response_id - assert ws.sent[3]["response_id"] == ws.sent[1]["response_id"] - assert ws.sent[3]["vllm_omni"]["model_turn_id"] == 1 - assert session.active_response_turn_id == 1 - assert session.turn_id == 1 - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_drops_old_turn_listen_while_new_response_is_active(): - request_id = "duplex-sid-native-old-turn-listen-e0-stage0" - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) - session = DuplexSession( - session_id="sid-native-old-turn-listen", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - response_id = session.begin_response(turn_id=1) - session.turn_id = 2 - session.bind_request(request_id) - ws = TimedWebSocket() - - close_reason, emitted = await handler._send_one_native_duplex_event( - ws.send_json, - { - "supported": True, - "stage_role": "llm", - "is_listen": True, - "model_listen": True, - "data_plane_request_id": request_id, - "end_of_turn": False, - "model_turn_id": 0, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - }, - session=session, - expected_epoch=session.epoch, - ) - - assert close_reason is None - assert emitted is False - assert session.active_response_id == response_id - assert ws.sent == [] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_silence_continuation_requires_active_request(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - session = DuplexSession( - session_id="sid-native-stale-continuation", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.begin_response() - session.bind_request("duplex-sid-native-stale-continuation-stage0") - session.mark_audio_sent(100) - sent: list[dict[str, Any]] = [] - - async def send_json(data: dict[str, Any]) -> None: - sent.append(data) - - _install_direct_silence_scheduler(handler, session) - - await handler._maybe_continue_native_response(send_json, session=session, expected_epoch=session.epoch) - session.clear_request() - await asyncio.sleep(0.01) - - assert engine.appended == [] - - -@pytest.mark.asyncio -async def test_minicpmo_auto_response_segment_complete_continues_until_model_listen(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - session = DuplexSession( - session_id="sid-native-segment-complete", - config=DuplexSessionConfig(extra_body={"auto_response": True}), - ) - session.capabilities = minicpmo45_native_capabilities() - response_id = session.begin_response() - request_id = "duplex-sid-native-segment-complete-stage0" - session.bind_request(request_id) - session.mark_audio_sent(1000) - _install_direct_silence_scheduler(handler, session) - sent: list[dict[str, Any]] = [] - - async def send_json(data: dict[str, Any]) -> None: - sent.append(data) - - close_reason, emitted = await handler._send_one_native_duplex_event( - send_json, - { - "supported": True, - "stage_role": "llm", - "is_listen": True, - "model_listen": False, - "listen_source": "auto_response_segment_complete", - "reason": "auto_response_segment_complete", - "data_plane_request_id": request_id, - "end_of_turn": False, - "uses_model_runner_scheduler": True, - "runner_kv_backed": True, - "runtime_impl": "scheduler_data_plane", - "owned_runtime": False, - }, - session=session, - expected_epoch=session.epoch, - ) - await asyncio.sleep(0.01) - - assert close_reason is None - assert emitted is False - assert session.active_response_id == response_id - assert session.active_request_id == request_id - assert not handler._serving_runtime_adapter.data_plane.is_terminal(request_id) - assert sent == [] - assert len(engine.appended) == 1 - _, mode, payload, final = engine.appended[0] - assert mode == "append_audio_chunk" - assert payload["type"] == "audio" - assert payload.get("force_listen") is not True - assert final is False - - -@pytest.mark.parametrize( - ("steps", "expected_deltas"), - [ - ( - [("你好呀", 1), ("你好呀", 1), ("你好呀", 2)], - ["你好呀", "", "你好呀"], - ), - ( - [("你好呀", 1), ("你好呀", 2)], - ["你好呀", "你好呀"], - ), - ], - ids=["turn-scoped", "exact-repeated-turn-text"], -) -def test_minicpmo_data_plane_text_cursor( - steps: list[tuple[str, int]], - expected_deltas: list[str], -): - data_plane = _test_data_plane() - request_id = "duplex-sid-native-e0-stage0" - - deltas = [data_plane.segment_text_delta(request_id, text, turn_id=turn_id) for text, turn_id in steps] - - assert deltas == expected_deltas - - -def test_minicpmo_data_plane_text_cursor_does_not_clip_prior_turn_overlap(): - data_plane = _test_data_plane() - request_id = "duplex-sid-native-e0-stage0" - - assert data_plane.segment_text_delta(request_id, "你好呀。你有什么想聊的吗？", turn_id=1) - assert ( - data_plane.segment_text_delta(request_id, "你有什么想聊的吗？你好呀。", turn_id=2) - == "你有什么想聊的吗？你好呀。" - ) - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_cancel_interrupts_background_data_plane_stream(): - control_result = { - "operation": "append", - "session_id": "sid-native-cancel-stream", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [ - { - "stage_id": 0, - "replica_id": 0, - "result": { - "supported": True, - "implementation_level": "model_native_duplex", - "data_plane_append": True, - "request_id": "duplex-sid-native-cancel-stream-e0-stage0-s1", - "response_stage_id": 1, - }, - } - ], - } - engine = FakeEngineClient( - append_result=control_result, - collect_delay_s=0.2, - collect_outputs=[ - [ - SimpleNamespace( - request_id="duplex-sid-native-cancel-stream-e0-stage0-s1", - finished=True, - outputs=[ - SimpleNamespace( - text="", - multimodal_output={"audio": np.zeros(10, dtype=np.float32), "sr": 24000}, - ) - ], - ) - ] - ], - ) - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-cancel-stream")) - ws.put({"type": "input_audio_buffer.append", "audio": "AAAA", "format": "pcm_f32le"}) - ws.put({"type": "input.cancel"}) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert "audio.cancelled" in ws.sent_types() - assert "response.output_audio.delta" not in ws.sent_types() - assert engine.signals == [("sid-native-cancel-stream", "barge_in")] - assert engine.signal_fences == [DuplexFence("sid-native-cancel-stream")] - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_session_close_discards_partial_pcm_tail(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-short-chunks")) - chunk = _pcm_f32_b64(3200) - for _ in range(6): - ws.put( - { - "type": "input_audio_buffer.append", - "audio": chunk, - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - ) - ws.put({"type": "session.close"}) - - await handler.handle_session(ws) - - assert len(engine.appended) == 1 - _, mode, payload, final = engine.appended[0] - assert mode == "append_audio_chunk" - assert final is False - assert isinstance(payload, dict) - assert len(base64.b64decode(payload["audio"])) == 16000 * 4 - assert "duplex_num_input_tokens" not in payload - - -@pytest.mark.asyncio -async def test_minicpmo_native_duplex_idle_timeout_closes_runtime_with_timeout_reason(): - engine = FakeEngineClient() - chat_service = FakeChatService(engine) - handler = OmniDuplexSessionHandler(chat_service=chat_service, config_timeout_s=0.1, idle_timeout_s=0.1) - ws = TimedWebSocket() - ws.put(_native_session_create("sid-native-timeout")) - - await handler.handle_session(ws) - - assert ws.sent_types() == ["session.created", "session.closed"] - assert ws.sent[-1]["reason"] == "timeout" - assert engine.closed == [("sid-native-timeout", "timeout")] diff --git a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py index be83ea9db..d0476b56c 100644 --- a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py +++ b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py @@ -32,7 +32,7 @@ from tests.helpers.media import ( from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni.config.stage_config import load_deploy_config -from vllm_omni.entrypoints.duplex.server_vad import ( +from vllm_omni.engine.duplex.vad import ( SILERO_VAD_FILENAME, SILERO_VAD_REPO_ID, SILERO_VAD_REVISION, diff --git a/tests/entrypoints/openai_api/test_server_vad.py b/tests/entrypoints/openai_api/test_server_vad.py deleted file mode 100644 index 67476be73..000000000 --- a/tests/entrypoints/openai_api/test_server_vad.py +++ /dev/null @@ -1,429 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -import base64 -import io -import sys -import threading -import wave -import weakref -from concurrent.futures import ThreadPoolExecutor -from types import SimpleNamespace - -import numpy as np -import pytest - -from vllm_omni.entrypoints.duplex.protocol import ( - DuplexSession, - DuplexSessionConfig, -) -from vllm_omni.entrypoints.duplex.server_vad import ( - ServerVADConfig, - ServerVADFrame, - ServerVADPipeline, - SileroVADBackend, - SileroVADBackendProvider, - SpeechEndpointDecision, - ThresholdEndpointPolicy, - parse_session_turn_detection, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class SequenceDetector: - frame_samples = 160 - - def __init__(self, probabilities: list[float]) -> None: - self.probabilities = probabilities - - def new_state(self) -> int: - return 0 - - def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: - index = int(state) - probability = self.probabilities[index] if index < len(self.probabilities) else 0.0 - return probability, index + 1 - - -def test_server_vad_config_defaults_and_validation(): - config = ServerVADConfig.from_value({"type": "server_vad"}) - - assert config.type == "server_vad" - assert config.threshold == 0.5 - assert config.prefix_padding_ms == 300 - assert config.silence_duration_ms == 500 - assert config.create_response is True - assert config.interrupt_response is None - assert "interrupt_response" not in config.as_dict() - assert config.min_speech_duration_ms is None - assert config.as_dict()["type"] == config.type - - interrupting = ServerVADConfig.from_value( - { - "type": "server_vad", - "interrupt_response": True, - "min_speech_duration_ms": 96, - } - ) - assert interrupting.interrupt_response is True - assert interrupting.min_speech_duration_ms == 96 - for value in (None, 0, "true"): - with pytest.raises(ValueError, match="interrupt_response must be a boolean"): - ServerVADConfig.from_value({"type": "server_vad", "interrupt_response": value}) - for enabled in (False, True): - _, aliased = parse_session_turn_detection( - { - "turn_detection": {"type": "server_vad"}, - "audio": {"input": {"turn_detection": {"type": "server_vad", "interrupt_response": enabled}}}, - } - ) - assert aliased is not None and aliased.interrupt_response is enabled - with pytest.raises(ValueError, match="Unknown"): - ServerVADConfig.from_value({"type": "server_vad", "semantic_eagerness": "high"}) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("prefix_padding_ms", [0, 10, 30]) -async def test_server_vad_pipeline_handles_arbitrary_chunk_boundaries(prefix_padding_ms: int): - detector = SequenceDetector([0.0, 0.9, 0.8, 0.0, 0.0]) - pipeline = ServerVADPipeline( - detector, - ServerVADConfig(prefix_padding_ms=prefix_padding_ms, silence_duration_ms=20), - ) - samples = np.zeros(detector.frame_samples * 5, dtype=np.float32) - - batches = [ - await pipeline.push(samples[:73]), - await pipeline.push(samples[73:431]), - await pipeline.push(samples[431:]), - ] - frames = [frame for batch in batches for frame in batch.frames] - - assert len(frames) == 5 - assert [index for index, frame in enumerate(frames) if frame.decision.speech_started] == [1] - assert [index for index, frame in enumerate(frames) if frame.decision.speech_stopped] == [4] - assert frames[1].decision.audio_start_ms == max(0, 10 - prefix_padding_ms) - assert frames[4].decision.audio_end_ms == 50 - - pipeline.reset() - reset_batch = await pipeline.push(samples) - assert [index for index, frame in enumerate(reset_batch.frames) if frame.decision.speech_started] == [1] - assert [index for index, frame in enumerate(reset_batch.frames) if frame.decision.speech_stopped] == [4] - # Keep session-relative timestamps, but never reference prefix audio - # discarded by reset at 50 ms. - assert reset_batch.frames[1].decision.audio_start_ms == max(50, 60 - prefix_padding_ms) - assert reset_batch.frames[4].decision.audio_end_ms == 100 - - -@pytest.mark.asyncio -async def test_server_vad_retained_audio_does_not_keep_input_chunk(): - pipeline = ServerVADPipeline(SequenceDetector([]), ServerVADConfig()) - session = DuplexSession("server-vad-memory", DuplexSessionConfig()) - samples = np.zeros(160 * 100 + 1, dtype=np.float32) - source_ref = weakref.ref(samples) - batch = await pipeline.push(samples) - session.append_server_vad_frame( - batch.frames[-1].samples, - speech_started=False, - speech_stopped=False, - prefix_samples=80, - ) - session.append_server_vad_frame( - batch.frames[-1].samples, - speech_started=True, - speech_stopped=False, - prefix_samples=80, - ) - del batch, samples - # Allow the completed to_thread callback to release its result, too. - await asyncio.sleep(0) - - assert source_ref() is None - assert pipeline.scratch_bytes == 4 - assert session.server_vad_utterance_bytes == (80 + 160) * 4 - - -@pytest.mark.asyncio -async def test_server_vad_pipeline_pcm16_resampling_is_split_invariant(): - sample_index = np.arange(2_400, dtype=np.int32) - source = ((sample_index * 7919) % 65_536 - 32_768).astype(" tuple[np.ndarray, list[SpeechEndpointDecision]]: - pipeline = ServerVADPipeline( - SequenceDetector([0.0] * 10), - ServerVADConfig(), - ) - frames: list[ServerVADFrame] = [] - offset = 0 - for chunk_size in chunk_sizes: - chunk = continuous_source[offset : offset + chunk_size] - batch = await pipeline.push_pcm16(chunk.tobytes(), source_sample_rate_hz=24_000) - frames.extend(batch.frames) - offset += chunk_size - assert offset == continuous_source.size - return np.concatenate([frame.samples for frame in frames]), [frame.decision for frame in frames] - - whole_samples, whole_decisions = await process([continuous_source.size]) - split_samples, split_decisions = await process([7, 13, 511, 512, 513, continuous_source.size - 1_556]) - - np.testing.assert_array_equal(split_samples, whole_samples) - assert whole_samples.size == 1_600 - assert split_decisions == whole_decisions - - -@pytest.mark.asyncio -async def test_server_vad_pipeline_reset_clears_resampler_and_source_rate_lock(): - pipeline = ServerVADPipeline( - SequenceDetector([0.0]), - ServerVADConfig(), - ) - - batch = await pipeline.push_pcm16( - np.asarray([30_000, -30_000], dtype=" None: - self.providers = providers - self.sess_options = sess_options - self.calls: list[dict[str, np.ndarray]] = [] - self.run_barrier: threading.Barrier | None = None - sessions.append(self) - - def get_inputs(self) -> list[SimpleNamespace]: - return [SimpleNamespace(name=name) for name in ("input", "state", "sr")] - - def run(self, _outputs: object, inputs: dict[str, np.ndarray]) -> list[np.ndarray]: - self.calls.append({name: np.array(value, copy=True) for name, value in inputs.items()}) - if self.run_barrier is not None: - self.run_barrier.wait(timeout=5) - return [ - np.asarray([[0.75]], dtype=np.float32), - np.full((2, 1, 128), len(self.calls), dtype=np.float32), - ] - - sessions: list[FakeInferenceSession] = [] - monkeypatch.setitem( - sys.modules, - "onnxruntime", - SimpleNamespace( - SessionOptions=FakeSessionOptions, - InferenceSession=FakeInferenceSession, - ), - ) - - model_path = tmp_path / "silero_vad.onnx" - model_path.write_bytes(b"fake-model") - backend = SileroVADBackend(model_path) - session = sessions[0] - - assert session.providers == ["CPUExecutionProvider"] - assert session.sess_options.inter_op_num_threads == 1 - assert session.sess_options.intra_op_num_threads == 1 - assert session.calls[0]["input"].shape == (1, 576) - assert session.calls[0]["state"].shape == (2, 1, 128) - np.testing.assert_array_equal(session.calls[0]["state"], np.zeros((2, 1, 128), dtype=np.float32)) - assert session.calls[0]["sr"].shape == () - assert session.calls[0]["sr"].item() == 16_000 - - state = backend.new_state() - first_frame = np.arange(backend.frame_samples, dtype=np.float32) - probability, state = backend.infer(first_frame, state) - first_call = session.calls[1] - - assert probability == pytest.approx(0.75) - np.testing.assert_array_equal( - first_call["input"][:, : backend.context_samples], - np.zeros((1, backend.context_samples), dtype=np.float32), - ) - np.testing.assert_array_equal( - first_call["input"][:, backend.context_samples :], - first_frame[None, :], - ) - np.testing.assert_array_equal(first_call["state"], np.zeros((2, 1, 128), dtype=np.float32)) - - second_frame = -first_frame - backend.infer(second_frame, state) - second_call = session.calls[2] - np.testing.assert_array_equal( - second_call["input"][:, : backend.context_samples], - first_frame[None, -backend.context_samples :], - ) - np.testing.assert_array_equal( - second_call["input"][:, backend.context_samples :], - second_frame[None, :], - ) - np.testing.assert_array_equal( - second_call["state"], - np.full((2, 1, 128), 2, dtype=np.float32), - ) - - # Independent pipeline states may enter the shared ORT session concurrently. - session.run_barrier = threading.Barrier(2) - with ThreadPoolExecutor(max_workers=2) as executor: - futures = [ - executor.submit(backend.infer, first_frame, backend.new_state()), - executor.submit(backend.infer, second_frame, backend.new_state()), - ] - assert [future.result()[0] for future in futures] == pytest.approx([0.75, 0.75]) - - -def test_duplex_session_owns_server_vad_prefix_and_utterance_audio(): - session = DuplexSession("server-vad-session", DuplexSessionConfig()) - frame = np.zeros(160, dtype=np.float32) - - session.reserve_input_bytes(frame.nbytes * 3, limit=frame.nbytes * 4) - released = session.append_server_vad_frame( - frame, - speech_started=False, - speech_stopped=False, - prefix_samples=160, - ) - assert released == 0 - released = session.append_server_vad_frame( - frame, - speech_started=False, - speech_stopped=False, - prefix_samples=160, - ) - assert released == frame.nbytes - session.release_input_bytes(released) - session.append_server_vad_frame( - frame, - speech_started=True, - speech_stopped=False, - prefix_samples=160, - ) - - assert session.server_vad_utterance_bytes == frame.nbytes * 2 - assert session.stage_server_vad_audio_for_commit() is True - committed = session.commit_user_input() - assert committed is not None - assert committed.message["role"] == "user" - audio_url = committed.message["content"][0]["audio_url"]["url"] - assert audio_url.startswith("data:audio/wav;base64,") - wav_bytes = base64.b64decode(audio_url.partition(",")[2]) - with wave.open(io.BytesIO(wav_bytes), "rb") as wav_file: - assert wav_file.getnchannels() == 1 - assert wav_file.getsampwidth() == 2 - assert wav_file.getframerate() == 16_000 - assert wav_file.getnframes() == frame.size * 2 - - -def test_silero_provider_rejects_missing_or_invalid_local_artifact(tmp_path): - missing = SileroVADBackendProvider(model_path=str(tmp_path / "missing.onnx")) - with pytest.raises(RuntimeError, match="does not exist"): - missing.get() - - invalid_path = tmp_path / "silero_vad.onnx" - invalid_path.write_bytes(b"not-the-pinned-model") - invalid = SileroVADBackendProvider(model_path=str(invalid_path)) - with pytest.raises(RuntimeError, match="checksum mismatch"): - invalid.get() diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index dda9ab38c..3b378930a 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -674,7 +674,7 @@ def test_resolve_stage_configs_delegates_overrides_to_resolver(mocker): engine_args=SimpleNamespace(additional_config=additional_config), ) resolve_config = mocker.patch( - "vllm_omni.engine.async_omni_engine.resolve_omni_config", + "vllm_omni.engine.omni_engine_base.resolve_omni_config", return_value=OmniConfigResolution( config_path="dummy.yaml", stage_configs=(fake_diffusion_stage,), diff --git a/tests/entrypoints/test_async_omni_duplex.py b/tests/entrypoints/test_async_omni_duplex.py deleted file mode 100644 index 81e0d75cb..000000000 --- a/tests/entrypoints/test_async_omni_duplex.py +++ /dev/null @@ -1,601 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from types import SimpleNamespace - -import pytest - -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.runtime import ( - DuplexOutputAction, - DuplexOutputDecision, - duplex_resource_request_id, -) -from vllm_omni.engine.messages import OutputMessage -from vllm_omni.entrypoints.async_omni import AsyncOmni -from vllm_omni.entrypoints.client_request_state import ClientRequestState -from vllm_omni.entrypoints.duplex_request_client import DuplexRequestClient -from vllm_omni.metrics.stats import OrchestratorAggregator, StageRequestStats, StageStats -from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.outputs.duplex import attach_duplex_output_decision - - -def test_async_omni_uses_extracted_duplex_request_client(): - app = object.__new__(AsyncOmni) - app.engine = SimpleNamespace(num_stages=2) - app.request_states = {} - app._duplex_request_client = None - - client = app._get_duplex_request_client() - - assert isinstance(client, DuplexRequestClient) - assert not hasattr(client, "_async_omni") - - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -@pytest.mark.asyncio -async def test_async_omni_open_duplex_forwards_session_config_and_timeout(): - calls = [] - output_handler_starts = 0 - - async def open_duplex_session_async(session_id, **kwargs): - calls.append((session_id, kwargs)) - return {"ok": True} - - app = object.__new__(AsyncOmni) - app.engine = SimpleNamespace(open_duplex_session_async=open_duplex_session_async) - - def start_output_handler(): - nonlocal output_handler_starts - output_handler_starts += 1 - - app._final_output_handler = start_output_handler - - result = await app.open_duplex_session_async( - "sid", - fence=DuplexFence("sid"), - capabilities={"implementation_level": "model_native_duplex"}, - session_config={"instructions": "Be brief.", "idle_timeout_s": 30}, - timeout=7.5, - ) - - assert result == {"ok": True} - assert output_handler_starts == 1 - assert calls == [ - ( - "sid", - { - "session_mode": "duplex", - "capabilities": {"implementation_level": "model_native_duplex"}, - "session_config": {"instructions": "Be brief.", "idle_timeout_s": 30}, - "fence": DuplexFence("sid"), - "timeout": 7.5, - }, - ) - ] - - -@pytest.mark.asyncio -async def test_async_omni_duplex_runtime_controls_forward_timeout(): - calls = [] - cancelled_fence = DuplexFence("sid", epoch=2, turn_id=3) - next_fence = DuplexFence("sid", epoch=3, turn_id=3) - - async def append_duplex_input_async(session_id, **kwargs): - calls.append(("append", session_id, kwargs)) - return {"ok": True} - - async def signal_duplex_turn_async(session_id, **kwargs): - calls.append(("signal", session_id, kwargs)) - return {"ok": True} - - async def close_duplex_session_async(session_id, **kwargs): - calls.append(("close", session_id, kwargs)) - return {"ok": True} - - async def touch_duplex_session_async(session_id, **kwargs): - calls.append(("touch", session_id, kwargs)) - return {"ok": True} - - async def resume_duplex_session_async(session_id, **kwargs): - calls.append(("resume", session_id, kwargs)) - return {"ok": True} - - app = object.__new__(AsyncOmni) - app.request_states = {} - app._final_output_handler = lambda: None - app.engine = SimpleNamespace( - append_duplex_input_async=append_duplex_input_async, - signal_duplex_turn_async=signal_duplex_turn_async, - close_duplex_session_async=close_duplex_session_async, - touch_duplex_session_async=touch_duplex_session_async, - resume_duplex_session_async=resume_duplex_session_async, - ) - - await app.append_duplex_input_async( - "sid", - mode="append_audio_chunk", - payload={}, - fence=cancelled_fence, - timeout=12.5, - ) - await app.signal_duplex_turn_async( - "sid", - event="barge_in", - fence=cancelled_fence, - next_fence=next_fence, - timeout=13.5, - ) - await app.close_duplex_session_async("sid", reason="done", fence=next_fence, timeout=14.5) - await app.touch_duplex_session_async( - "sid", - fence=next_fence, - activity=DuplexLeaseActivity.HEARTBEAT, - timeout=15.5, - ) - await app.resume_duplex_session_async( - "sid", - fence=next_fence, - expected_lease_generation=3, - timeout=16.5, - ) - - assert calls == [ - ( - "append", - "sid", - { - "mode": "append_audio_chunk", - "payload": {}, - "final": False, - "expected_epoch": None, - "fence": cancelled_fence, - "timeout": 12.5, - }, - ), - ( - "signal", - "sid", - { - "event": "barge_in", - "fence": cancelled_fence, - "next_fence": next_fence, - "timeout": 13.5, - }, - ), - ( - "close", - "sid", - { - "reason": "done", - "fence": next_fence, - "timeout": 14.5, - }, - ), - ( - "touch", - "sid", - { - "fence": next_fence, - "activity": DuplexLeaseActivity.HEARTBEAT, - "timeout": 15.5, - }, - ), - ( - "resume", - "sid", - { - "fence": next_fence, - "expected_lease_generation": 3, - "timeout": 16.5, - }, - ), - ] - - -@pytest.mark.asyncio -async def test_async_omni_duplex_append_can_defer_data_plane_output_collection(): - app = object.__new__(AsyncOmni) - request_id = duplex_resource_request_id(DuplexFence("sid"), "stage0") - - async def append_duplex_input_async(session_id, **kwargs): - del session_id, kwargs - request_state = app.request_states.get(request_id) - assert request_state is not None - assert request_state.metrics is not None - return { - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": 1, - } - } - ] - } - - app.engine = SimpleNamespace(append_duplex_input_async=append_duplex_input_async) - app.request_states = {} - app._final_output_handler = lambda: None - - result = await app.append_duplex_input_async( - "sid", - mode="append_audio_chunk", - payload={}, - fence=DuplexFence("sid"), - collect_outputs=False, - ) - - assert "data_plane_outputs" not in result - - -@pytest.mark.asyncio -async def test_async_omni_duplex_append_forwards_fence(): - calls = [] - - async def append_duplex_input_async(session_id, **kwargs): - calls.append((session_id, kwargs)) - return {"ok": True} - - app = object.__new__(AsyncOmni) - app.request_states = {} - app._final_output_handler = lambda: None - app.engine = SimpleNamespace(append_duplex_input_async=append_duplex_input_async) - fence = DuplexFence("sid", epoch=1, turn_id=2) - - await app.append_duplex_input_async( - "sid", - mode="append_audio_chunk", - payload={}, - fence=fence, - collect_outputs=False, - ) - - assert calls[0][1]["fence"] is fence - - -@pytest.mark.asyncio -async def test_async_omni_duplex_collect_waits_for_response_stage(): - app = object.__new__(AsyncOmni) - req_state = ClientRequestState("duplex-sid-e0-stage0") - stage0_output = OmniRequestOutput( - request_id="duplex-sid-e0-stage0", - stage_id=0, - finished=False, - ) - stage1_output = OmniRequestOutput( - request_id="duplex-sid-e0-stage0", - stage_id=1, - final_output_type="audio", - finished=False, - ) - await req_state.queue.put( - OutputMessage( - request_id="duplex-sid-e0-stage0", - stage_id=0, - engine_outputs=stage0_output, - finished=False, - ) - ) - await req_state.queue.put( - OutputMessage( - request_id="duplex-sid-e0-stage0", - stage_id=1, - engine_outputs=stage1_output, - finished=False, - ) - ) - - outputs = await app._collect_duplex_data_plane_outputs( - "duplex-sid-e0-stage0", - req_state, - response_stage_id=1, - timeout=1.0, - ) - - assert outputs == [stage1_output] - - -@pytest.mark.asyncio -async def test_duplex_request_client_retains_output_route_at_segment_terminal(): - request_id = "duplex-sid-e0-stage0" - request_state = ClientRequestState(request_id) - request_states = {request_id: request_state} - output = OmniRequestOutput( - request_id=request_id, - stage_id=1, - final_output_type="audio", - finished=True, - ) - await request_state.queue.put( - OutputMessage( - request_id=request_id, - stage_id=1, - engine_outputs=output, - finished=True, - ) - ) - client = DuplexRequestClient( - SimpleNamespace(), - SimpleNamespace( - request_states=request_states, - num_stages=2, - log_stats=False, - ), - ) - - outputs = await client.collect_outputs( - request_id, - request_state, - response_stage_id=1, - timeout=1.0, - ) - - assert outputs == [output] - assert request_states[request_id] is request_state - - -@pytest.mark.asyncio -async def test_async_omni_duplex_direct_output_keeps_stage_metrics(): - app = object.__new__(AsyncOmni) - app._duplex_request_client = None - app._final_output_handler = lambda: None - request_id = "duplex-sid-e0-stage0" - req_state = ClientRequestState(request_id) - decision = DuplexOutputDecision( - action=DuplexOutputAction.DIRECT_RESPONSE, - metadata={ - "duplex_direct_response": True, - "duplex_native_decision": "speak", - }, - ) - direct_output = attach_duplex_output_decision( - OmniRequestOutput( - request_id=request_id, - stage_id=0, - final_output_type="text", - finished=True, - ), - decision, - ) - await req_state.queue.put( - OutputMessage( - request_id=request_id, - stage_id=0, - engine_outputs=direct_output, - metrics=StageRequestStats( - batch_id=1, - batch_size=1, - num_tokens_in=160, - num_tokens_out=8, - stage_gen_time_ms=80.0, - rx_transfer_bytes=0, - rx_decode_time_ms=0.0, - rx_in_flight_time_ms=0.0, - stage_stats=StageStats(total_token=8, total_gen_time_ms=80.0), - vllm_ttft_ms=12.5, - vllm_tpot_ms=9.0, - vllm_itl_ms=8.75, - vllm_itls_ms=[8.0, 9.5], - ), - finished=True, - ) - ) - - outputs = await app._collect_duplex_data_plane_outputs( - request_id, - req_state, - response_stage_id=1, - timeout=1.0, - ) - - assert outputs == [direct_output] - assert direct_output.metrics["stage_metrics"]["0"]["vllm_ttft_ms"] == 12.5 - assert direct_output.metrics["stage_metrics"]["0"]["vllm_tpot_ms"] == 9.0 - assert direct_output.metrics["stage_metrics"]["0"]["vllm_itls_ms"] == [8.0, 9.5] - - -@pytest.mark.asyncio -async def test_duplex_metrics_cursor_isolates_resumable_collect_windows(): - app = object.__new__(AsyncOmni) - app._duplex_request_client = None - app._final_output_handler = lambda: None - request_id = "duplex-sid-e0-stage0" - req_state = ClientRequestState(request_id) - req_state.metrics = OrchestratorAggregator( - num_stages=2, - log_stats=False, - wall_start_ts=0.0, - final_stage_id_for_e2e=1, - ) - - def stage_message(*, batch_id: int, tokens: int, ttft_ms: float, itls: list[float]) -> OutputMessage: - decision = DuplexOutputDecision( - action=DuplexOutputAction.DIRECT_RESPONSE, - metadata={"duplex_direct_response": True, "duplex_native_decision": "speak"}, - ) - output = attach_duplex_output_decision( - OmniRequestOutput( - request_id=request_id, - stage_id=0, - final_output_type="text", - finished=True, - ), - decision, - ) - return OutputMessage( - request_id=request_id, - stage_id=0, - engine_outputs=output, - metrics=StageRequestStats( - batch_id=batch_id, - batch_size=1, - num_tokens_in=160, - num_tokens_out=tokens, - stage_gen_time_ms=80.0, - rx_transfer_bytes=0, - rx_decode_time_ms=0.0, - rx_in_flight_time_ms=0.0, - stage_stats=StageStats(total_token=tokens, total_gen_time_ms=80.0), - vllm_ttft_ms=ttft_ms, - vllm_tpot_ms=9.0, - vllm_itl_ms=sum(itls) / len(itls), - vllm_itls_ms=itls, - ), - finished=True, - ) - - await req_state.queue.put(stage_message(batch_id=1, tokens=71, ttft_ms=62.5, itls=[10.0, 11.0])) - first = await app._collect_duplex_data_plane_outputs( - request_id, - req_state, - response_stage_id=1, - timeout=1.0, - ) - await req_state.queue.put(stage_message(batch_id=2, tokens=155, ttft_ms=81.0, itls=[12.0, 13.0, 14.0])) - second = await app._collect_duplex_data_plane_outputs( - request_id, - req_state, - response_stage_id=1, - timeout=1.0, - ) - - assert first[0].metrics["stage_metrics"]["0"]["num_tokens_out"] == 71 - second_stage0 = second[0].metrics["stage_metrics"]["0"] - assert second_stage0["num_tokens_out"] == 155 - assert second_stage0["vllm_ttft_ms"] == 81.0 - assert second_stage0["vllm_itls_ms"] == [12.0, 13.0, 14.0] - - -def test_async_omni_duplex_direct_output_prefers_outer_control_metadata(): - decision = DuplexOutputDecision( - action=DuplexOutputAction.DIRECT_RESPONSE, - metadata={ - "duplex_direct_response": True, - "duplex_native_decision": "listen", - }, - ) - output = attach_duplex_output_decision( - OmniRequestOutput( - request_id="duplex-sid-e0-stage0", - stage_id=0, - outputs=[], - _multimodal_output={"kind": "processed-payload"}, - _custom_output={"inner": "completion"}, - ), - decision, - ) - - assert AsyncOmni._is_direct_duplex_data_plane_response(output) - - -def test_async_omni_drops_unregistered_duplex_prefixed_output(): - app = object.__new__(AsyncOmni) - app.request_states = {} - output = OmniRequestOutput(request_id="duplex-unregistered-e0-stage0", stage_id=0) - - keep_processing, request_id, stage_id, request_state = app._handle_output_message( - OutputMessage( - request_id=output.request_id, - stage_id=0, - engine_outputs=output, - finished=False, - ) - ) - - assert keep_processing is True - assert request_id is stage_id is request_state is None - assert app.request_states == {} - - -@pytest.mark.asyncio -async def test_async_omni_duplex_collect_wraps_raw_response_stage_output(): - app = object.__new__(AsyncOmni) - app.log_stats = False - app._enable_ar_profiler = False - app.request_states = {} - app.prom_metrics = SimpleNamespace( - set_running=lambda value: None, - set_waiting=lambda value: None, - ) - app.engine = SimpleNamespace( - num_stages=2, - get_stage_metadata=lambda stage_id: SimpleNamespace( - final_output=stage_id == 1, - final_output_type="audio", - ), - ) - req_state = ClientRequestState("duplex-sid-e0-stage0") - raw_stage1_output = SimpleNamespace( - finished=False, - outputs=[], - stage_durations={}, - peak_memory_mb=0.0, - final_output_type="audio", - ) - await req_state.queue.put( - OutputMessage( - request_id="duplex-sid-e0-stage0", - stage_id=1, - engine_outputs=raw_stage1_output, - finished=False, - ) - ) - - outputs = await app._collect_duplex_data_plane_outputs( - "duplex-sid-e0-stage0", - req_state, - response_stage_id=1, - timeout=1.0, - ) - - assert len(outputs) == 1 - assert outputs[0].request_id == "duplex-sid-e0-stage0" - assert outputs[0].finished is False - assert outputs[0].stage_id == 1 - assert outputs[0].final_output_type == "audio" - assert outputs[0].peak_memory_mb == raw_stage1_output.peak_memory_mb - - -def test_async_omni_duplex_request_info_includes_response_stage(): - request_id, response_stage_id = AsyncOmni._duplex_data_plane_request_info( - { - "stage_results": [ - { - "result": { - "data_plane_append": True, - "request_id": "duplex-sid-e0-stage0", - "response_stage_id": 1, - } - } - ] - } - ) - - assert request_id == "duplex-sid-e0-stage0" - assert response_stage_id == 1 - - -@pytest.mark.asyncio -async def test_duplex_request_client_retains_output_route_when_close_fails(): - fence = DuplexFence("sid-close-route") - request_id = duplex_resource_request_id(fence, "stage0") - request_states = {request_id: ClientRequestState(request_id)} - - async def close_duplex_session_async(session_id, **kwargs): - del session_id, kwargs - raise RuntimeError("close failed") - - client = DuplexRequestClient( - SimpleNamespace(close_duplex_session_async=close_duplex_session_async), - SimpleNamespace(request_states=request_states), - ) - - with pytest.raises(RuntimeError, match="close failed"): - await client.close(fence.session_id, reason="test", fence=fence, timeout=1.0) - - assert request_id in request_states diff --git a/tests/entrypoints/test_disable_log_stats_routing.py b/tests/entrypoints/test_disable_log_stats_routing.py index f9dd2da4a..623cc9af4 100644 --- a/tests/entrypoints/test_disable_log_stats_routing.py +++ b/tests/entrypoints/test_disable_log_stats_routing.py @@ -53,7 +53,7 @@ def test_async_engine_consumes_disable_log_stats_before_resolution(monkeypatch: ) monkeypatch.setattr( - "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", + "vllm_omni.engine.omni_engine_base.load_and_resolve_stage_configs", fake_resolve, ) engine = SimpleNamespace(_apply_strategy_lb_policy=lambda *args, **kwargs: None) diff --git a/tests/entrypoints/test_duplex_fence_propagation.py b/tests/entrypoints/test_duplex_fence_propagation.py deleted file mode 100644 index 550723ec8..000000000 --- a/tests/entrypoints/test_duplex_fence_propagation.py +++ /dev/null @@ -1,126 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -import pytest - -from tests.entrypoints.openai_api.test_duplex_handler import FakeChatService, TimedWebSocket -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.entrypoints.async_omni import AsyncOmni -from vllm_omni.entrypoints.duplex.protocol import DuplexSession, DuplexSessionConfig -from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class FenceRecordingEngine: - output_modalities = ["text", "audio"] - - def __init__(self) -> None: - self.calls: list[tuple[str, DuplexFence | None]] = [] - self.signal_next_fences: list[DuplexFence | None] = [] - - async def open_duplex_session_async(self, session_id, *, fence=None, **kwargs): - del session_id, kwargs - self.calls.append(("open", fence)) - return {"ok": True} - - async def append_duplex_input_async(self, session_id, *, fence=None, **kwargs): - del session_id, kwargs - self.calls.append(("append", fence)) - return {"ok": True} - - async def signal_duplex_turn_async( - self, - session_id, - *, - event, - fence=None, - next_fence=None, - timeout=None, - ): - del session_id, event, timeout - self.calls.append(("signal", fence)) - self.signal_next_fences.append(next_fence) - return {"ok": True} - - async def close_duplex_session_async(self, session_id, *, fence=None, **kwargs): - del session_id, kwargs - self.calls.append(("close", fence)) - return {"ok": True} - - -@pytest.mark.asyncio -async def test_async_omni_forwards_fence_for_all_runtime_operations(): - engine = FenceRecordingEngine() - app = object.__new__(AsyncOmni) - app.engine = engine - app.request_states = {} - app._final_output_handler = lambda: None - fence = DuplexFence("sid", epoch=1, turn_id=2) - - await app.open_duplex_session_async("sid", fence=fence) - await app.append_duplex_input_async( - "sid", mode="append_audio_chunk", payload={}, fence=fence, collect_outputs=False - ) - await app.signal_duplex_turn_async("sid", event="turn.end", fence=fence) - await app.close_duplex_session_async("sid", reason="done", fence=fence) - - assert engine.calls == [(operation, fence) for operation in ("open", "append", "signal", "close")] - - -@pytest.mark.asyncio -async def test_openai_handler_passes_current_session_fence_for_runtime_controls(): - engine = FenceRecordingEngine() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession( - session_id="sid", - config=DuplexSessionConfig(extra_body={"native_duplex": True}), - epoch=1, - turn_id=2, - ) - ws = TimedWebSocket() - fence = DuplexFence("sid", epoch=1, turn_id=2) - - await handler._open_runtime_session(session, ws.send_json) - await handler._signal_runtime_session(session, "turn.end", send_json=ws.send_json) - await handler._close_runtime_session(session, reason="done", send_json=ws.send_json) - - assert engine.calls == [(operation, fence) for operation in ("open", "signal", "close")] - - -@pytest.mark.asyncio -async def test_runtime_signal_uses_fence_captured_before_session_turn_advances(): - engine = FenceRecordingEngine() - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) - session = DuplexSession(session_id="sid", config=DuplexSessionConfig(extra_body={"native_duplex": True})) - captured_fence = DuplexFence("sid", epoch=0, turn_id=0) - session.turn_id = 1 - - await handler._signal_runtime_session( - session, - "input.commit", - fence=captured_fence, - ) - - assert engine.calls == [("signal", captured_fence)] - - -@pytest.mark.asyncio -async def test_openai_handler_forwards_cancel_fence_through_async_omni_facade(): - engine = FenceRecordingEngine() - app = object.__new__(AsyncOmni) - app.engine = engine - handler = OmniDuplexSessionHandler(chat_service=FakeChatService(app)) - session = DuplexSession(session_id="sid-cancel", config=DuplexSessionConfig(extra_body={"native_duplex": True})) - cancelled_fence = DuplexFence("sid-cancel", epoch=2, turn_id=3) - next_fence = DuplexFence("sid-cancel", epoch=3, turn_id=3) - - assert await handler._signal_runtime_session( - session, - "input.cancel", - fence=cancelled_fence, - next_fence=next_fence, - ) - - assert engine.calls == [("signal", cancelled_fence)] - assert engine.signal_next_fences == [next_fence] diff --git a/tests/entrypoints/test_duplex_omni.py b/tests/entrypoints/test_duplex_omni.py new file mode 100644 index 000000000..f33b3ea4d --- /dev/null +++ b/tests/entrypoints/test_duplex_omni.py @@ -0,0 +1,326 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""``DuplexOmni`` / ``DuplexSessionHandle``: the thin Python API over a fake engine.""" + +from __future__ import annotations + +import asyncio +from dataclasses import replace +from typing import Any + +import pytest + +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.events import ( + AudioDelta, + DuplexEvent, + SessionClosed, + SessionCreated, + SessionExpired, +) +from vllm_omni.engine.duplex.messages import ( + DuplexControlResultMessage, + DuplexSessionError, + DuplexSessionEventMessage, +) +from vllm_omni.engine.messages import ErrorMessage +from vllm_omni.entrypoints.duplex_omni import DuplexOmni, DuplexSessionHandle + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class FakeDuplexEngine(FakeAsyncOmniEngine): + """The engine surface ``DuplexOmni`` uses; every call is recorded.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.duplex_session_config = DuplexSessionRuntimeConfig(max_sessions=2) + self.duplex_capabilities = DuplexCapabilities(supports_session_resume=True) + self.opened: list[tuple[str, DuplexSessionConfig]] = [] + self.closed: list[tuple[str, str]] = [] + self.resumed: list[tuple[str, int]] = [] + self.touched: list[tuple[str, str]] = [] + self.commands: list[tuple[str, commands.DuplexCommand]] = [] + self.open_error: DuplexSessionError | None = None + self.close_error: DuplexSessionError | None = None + + 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): + self.opened.append((session_id, session_config)) + if self.open_error is not None: + raise self.open_error + return self._result( + "open", + session_id, + capabilities=self.duplex_capabilities, + public_session={"id": session_id, "idle_timeout_s": 7}, + lease_generation=0, + ) + + async def close_session_async(self, session_id, *, reason="client_close", timeout=None): + self.closed.append((session_id, reason)) + if self.close_error is not None: + raise self.close_error + # The runner emits session.closed before answering the RPC. + self.emit(session_id, SessionClosed(reason=reason)) + return self._result("close", session_id) + + async def resume_session_async(self, session_id, *, expected_lease_generation, timeout=None): + self.resumed.append((session_id, expected_lease_generation)) + return self._result("resume", session_id, lease_generation=expected_lease_generation + 1) + + async def touch_session_async(self, session_id, *, activity, timeout=None): + self.touched.append((session_id, activity)) + return self._result("touch", session_id) + + async def submit_command_async(self, session_id, command): + self.commands.append((session_id, command)) + + 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)) + + +def _make_omni(monkeypatch: pytest.MonkeyPatch) -> tuple[DuplexOmni, FakeDuplexEngine]: + engine = FakeDuplexEngine() + monkeypatch.setattr("vllm_omni.entrypoints.duplex_omni.DuplexOmniEngine", lambda *args, **kwargs: engine) + monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model) + return DuplexOmni("dummy-model"), engine + + +async def _collect(handle: DuplexSessionHandle, count: int, *, timeout_s: float = 2.0) -> list[DuplexEvent]: + events: list[DuplexEvent] = [] + + async def consume() -> None: + async for event in handle.events(): + events.append(event) + if len(events) >= count: + break + + await asyncio.wait_for(consume(), timeout=timeout_s) + return events + + +@pytest.mark.asyncio +async def test_open_session_allocates_the_id_and_ignores_client_ids(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session({"session_id": "mine", "id": "mine-too", "instructions": "hi"}) + + assert handle.session_id.startswith("duplex-") and len(handle.session_id) == len("duplex-") + 32 + assert handle.session_id != "mine" + ((session_id, config),) = engine.opened + assert session_id == handle.session_id + assert isinstance(config, DuplexSessionConfig) + assert config.model == "dummy-model" and config.instructions == "hi" + assert handle.capabilities is engine.duplex_capabilities + assert handle.public_session["id"] == handle.session_id + assert omni.sessions == {handle.session_id: handle} + assert omni.get_session(handle.session_id) is handle + assert omni.active_session_count() == 1 + + second = await omni.open_session(DuplexSessionConfig(instructions="typed")) + assert second.session_id != handle.session_id + assert engine.opened[1][1].model == "dummy-model" + assert omni.duplex_session_config.max_sessions == 2 + assert omni.duplex_capabilities.supports_session_resume is True + 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) + engine.open_error = DuplexSessionError("full", code="resource_exhausted", retryable=True) + try: + with pytest.raises(DuplexSessionError) as excinfo: + await omni.open_session() + assert excinfo.value.code == "resource_exhausted" + assert omni.sessions == {} + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_events_are_routed_to_their_handle_and_end_on_session_closed(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session() + engine.emit("duplex-unknown", AudioDelta(delta="zzz")) # dropped + engine.emit(handle.session_id, SessionCreated(session={"id": handle.session_id})) + engine.emit(handle.session_id, AudioDelta(response_id="r1", delta="abc")) + engine.emit(handle.session_id, SessionClosed(reason="client_close")) + + events = await _collect(handle, 10) + + assert [event.type for event in events] == ["session.created", "response.output_audio.delta", "session.closed"] + assert events[0].session_id == handle.session_id + assert handle.closed and handle.close_reason == "client_close" + assert await handle.wait_closed() == "client_close" + assert handle.session_id not in omni.sessions + with pytest.raises(DuplexSessionError) as excinfo: + await handle.submit(commands.Heartbeat()) + assert excinfo.value.code == "session_closed" + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_events_is_single_consumer_but_may_be_reentered(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session() + engine.emit(handle.session_id, AudioDelta(delta="one")) + first = handle.events() + assert (await first.__anext__()).type == "response.output_audio.delta" + with pytest.raises(RuntimeError, match="active events"): + await handle.events().__anext__() + await first.aclose() + + engine.emit(handle.session_id, SessionExpired(reason="idle_ttl_expired")) + events = await _collect(handle, 1) + assert events[0].type == "session.expired" and events[0].is_terminal + assert handle.closed and handle.close_reason == "idle_ttl_expired" + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_close_session_waits_for_the_typed_session_closed_event(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session() + consumer = asyncio.create_task(_collect(handle, 1)) + await handle.close(reason="done") + + assert engine.closed == [(handle.session_id, "done")] + assert handle.closed and handle.close_reason == "done" + assert (await consumer)[0].type == "session.closed" + # Idempotent: a second close is a no-op. + await omni.close_session(handle.session_id) + assert len(engine.closed) == 1 + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_close_failure_marks_the_handle_closed_and_propagates(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + engine.close_error = DuplexSessionError("boom", code="engine_error") + try: + handle = await omni.open_session() + with pytest.raises(DuplexSessionError, match="boom"): + await omni.close_session(handle.session_id) + assert handle.closed + # A failed close must not leave a dead handle shadowing the id. + assert omni.sessions == {} + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_engine_death_ends_every_session_handle(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session() + engine.output_q.put(ErrorMessage(error="orchestrator died", fatal=True)) + + await asyncio.wait_for(handle.wait_closed(), timeout=2.0) + assert handle.closed + assert "engine_dead" in (handle.close_reason or "") + assert omni.sessions == {} + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_resume_detach_and_touch_forward_to_the_engine(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session() + resumed = await omni.resume_session(handle.session_id, expected_lease_generation=0) + assert resumed is handle and handle.lease_generation == 1 + assert engine.resumed == [(handle.session_id, 0)] + + await omni.detach_session(handle.session_id) + await omni.touch_session(handle.session_id, activity="heartbeat") + assert engine.touched == [(handle.session_id, "detach"), (handle.session_id, "heartbeat")] + + for call in ( + omni.resume_session("duplex-nope", expected_lease_generation=0), + omni.detach_session("duplex-nope"), + ): + with pytest.raises(DuplexSessionError) as excinfo: + await call + assert excinfo.value.code == "unknown_session" + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_handle_wrappers_build_typed_commands_in_caller_order(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session() + await handle.append_audio(b"\x00\x01", sample_rate_hz=16000, is_speech=True) + await handle.commit(create_response=True) + await handle.create_response({"instructions": "x"}) + await handle.cancel_response("resp-1") + await handle.barge_in() + await handle.ack_playback(120.7, response_id="resp-1", committed_ms=100) + await handle.update({"temperature": 0.2}) + await handle.create_item({"id": "item_1", "type": "message"}, previous_item_id="root") + await handle.delete_item("item_1") + await handle.truncate_item("item_1", audio_end_ms=50) + await handle.signal_turn("user_started", {"x": 1}) + await handle.heartbeat() + + submitted = [command for session_id, command in engine.commands if session_id == handle.session_id] + assert [type(command) for command in submitted] == [ + commands.AppendAudio, + commands.Commit, + commands.CreateResponse, + commands.CancelResponse, + commands.BargeIn, + commands.AckPlayback, + commands.UpdateSession, + commands.CreateItem, + commands.DeleteItem, + commands.TruncateItem, + commands.SignalTurn, + commands.Heartbeat, + ] + append, commit, create, cancel, _, ack, update, item, delete, truncate, signal, _ = submitted + assert append.audio == b"\x00\x01" and append.sample_rate_hz == 16000 and append.is_speech is True + assert commit.create_response is True and commit.final is True + assert create.options == {"instructions": "x"} + assert cancel.response_id == "resp-1" + assert ack.played_ms == 120 and ack.committed_ms == 100 and ack.response_id == "resp-1" + assert update.patch == {"temperature": 0.2} + assert item.previous_item_id == "root" and item.item["id"] == "item_1" + assert delete.item_id == "item_1" + assert truncate.audio_end_ms == 50 + assert signal.event == "user_started" and signal.signal_payload == {"x": 1} + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_shutdown_closes_every_handle(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + first = await omni.open_session() + second = await omni.open_session() + + omni.shutdown() + + assert first.closed and second.closed + assert first.close_reason == "shutdown" + assert omni.sessions == {} + assert engine.shutdown_called diff --git a/tests/entrypoints/test_omni_base_profiler.py b/tests/entrypoints/test_omni_base_profiler.py index 827ab67b6..d343e2f03 100644 --- a/tests/entrypoints/test_omni_base_profiler.py +++ b/tests/entrypoints/test_omni_base_profiler.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Unit tests for OmniBase and AsyncOmni profiler methods.""" from types import SimpleNamespace @@ -28,10 +31,14 @@ class TestOmniBaseProfiler: @pytest.fixture def omni_base_instance(self, mock_engine, mocker: MockerFixture): """Create an OmniBase instance with mocked dependencies.""" - mocker.patch("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", return_value=mock_engine) + from vllm_omni.entrypoints.omni_base import OmniBase + + # ``OmniBase`` builds its engine through the ``_create_engine`` seam that + # ``AsyncOmni`` / ``DuplexOmni`` implement; there is no module-level + # engine class to patch any more. + mocker.patch.object(OmniBase, "_create_engine", return_value=mock_engine) mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", side_effect=lambda x: x) mocker.patch("vllm_omni.entrypoints.omni_base.weakref.finalize") - from vllm_omni.entrypoints.omni_base import OmniBase instance = OmniBase(model="test-model") return instance @@ -48,10 +55,14 @@ class TestOmniBaseProfiler: def test_removed_diffusion_batch_size_fails_loudly(self, mock_engine, mocker: MockerFixture): """Legacy diffusion_batch_size callers receive an actionable error.""" - mocker.patch("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", return_value=mock_engine) + from vllm_omni.entrypoints.omni_base import OmniBase + + # ``OmniBase`` builds its engine through the ``_create_engine`` seam that + # ``AsyncOmni`` / ``DuplexOmni`` implement; there is no module-level + # engine class to patch any more. + mocker.patch.object(OmniBase, "_create_engine", return_value=mock_engine) mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", side_effect=lambda x: x) mocker.patch("vllm_omni.entrypoints.omni_base.weakref.finalize") - from vllm_omni.entrypoints.omni_base import OmniBase with pytest.raises(TypeError, match="diffusion_batch_size.*max_num_seqs"): OmniBase(model="test-model", diffusion_batch_size=8) diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 01264eaea..32d7d9988 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -118,8 +118,6 @@ class FakeAsyncOmniEngine: self.output_processors = [SimpleNamespace(tokenizer=None) for _ in range(self.num_stages)] self.input_processor = None self.endpoint_restrictions = () - self.duplex_session_config = None - self.duplex_serving_adapter_path = None self.output_q: queue.Queue[Any] = queue.Queue() self.submitted: list[dict[str, Any]] = [] @@ -183,7 +181,8 @@ class FakeAsyncOmniEngine: def _patch_engine(monkeypatch: pytest.MonkeyPatch, engine: FakeAsyncOmniEngine) -> None: - monkeypatch.setattr("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", lambda *args, **kwargs: engine) + monkeypatch.setattr("vllm_omni.entrypoints.omni.AsyncOmniEngine", lambda *args, **kwargs: engine) + monkeypatch.setattr("vllm_omni.entrypoints.async_omni.AsyncOmniEngine", lambda *args, **kwargs: engine) monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model) # Don't add random UUIDs to requests calling .generate since we usually # just want to check for present requests anyway, and would need to just diff --git a/tests/examples/test_minicpmo_realtime_web_static.py b/tests/examples/test_minicpmo_realtime_web_static.py index 926ae36ef..2027d5aa3 100644 --- a/tests/examples/test_minicpmo_realtime_web_static.py +++ b/tests/examples/test_minicpmo_realtime_web_static.py @@ -36,7 +36,7 @@ def test_client_uses_proxy_relative_realtime_url_and_model_policy_session(): assert "new URL(config.realtimePath, window.location.href)" in source assert "url.protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'" in source assert "url.searchParams.set('autostart', '0')" in source - assert "url.searchParams.set('native_duplex', '1')" in source + assert "native_duplex" not in source assert "auto_response: true" in source assert "input_audio_buffer.append" in source assert "input_audio_buffer.commit" not in source diff --git a/tests/model_executor/models/minicpmo_4_5/duplex/test_commit_policy.py b/tests/model_executor/models/minicpmo_4_5/duplex/test_commit_policy.py index b933f7173..d7f3db633 100644 --- a/tests/model_executor/models/minicpmo_4_5/duplex/test_commit_policy.py +++ b/tests/model_executor/models/minicpmo_4_5/duplex/test_commit_policy.py @@ -5,7 +5,7 @@ from __future__ import annotations import pytest -from vllm_omni.entrypoints.duplex.commit_policy import ( +from vllm_omni.engine.duplex.session.commit_policy import ( CommitAction, CommitSnapshot, decide_commit_action, @@ -24,7 +24,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] speech_since_commit=True, active_response_id=None, overlap_speech_ms=0, - native_response_in_progress=False, + response_in_progress=False, playback_active=False, ), CommitAction.START_AUTO_RESPONSE, @@ -36,7 +36,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] speech_since_commit=True, active_response_id=None, overlap_speech_ms=800, - native_response_in_progress=True, + response_in_progress=True, playback_active=True, ), CommitAction.START_AUTO_RESPONSE, @@ -48,7 +48,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] speech_since_commit=False, active_response_id=None, overlap_speech_ms=0, - native_response_in_progress=True, + response_in_progress=True, playback_active=False, ), CommitAction.COMMIT_ONLY, @@ -60,7 +60,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] speech_since_commit=True, active_response_id=None, overlap_speech_ms=800, - native_response_in_progress=True, + response_in_progress=True, playback_active=False, ), CommitAction.DEFER_ACTIVE_RESPONSE, @@ -72,7 +72,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] speech_since_commit=True, active_response_id=None, overlap_speech_ms=800, - native_response_in_progress=False, + response_in_progress=False, playback_active=True, ), CommitAction.DEFER_ACTIVE_RESPONSE, @@ -84,7 +84,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] speech_since_commit=True, active_response_id=None, overlap_speech_ms=0, - native_response_in_progress=False, + response_in_progress=False, playback_active=False, ), CommitAction.COMMIT_ONLY, diff --git a/tests/model_executor/models/minicpmo_4_5/duplex/test_data_plane_audio_cursor.py b/tests/model_executor/models/minicpmo_4_5/duplex/test_data_plane_audio_cursor.py new file mode 100644 index 000000000..b2be91366 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/duplex/test_data_plane_audio_cursor.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The cumulative Stage1 audio cursor. + +Stage1 hands the data plane a snapshot of every sample it has generated for the +request so far, and only the new tail may reach the client. The cursor that +makes that possible is the one piece of per-request state whose loss is silent +and expensive: rewinding it re-sends the whole session as a single delta. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.data_plane import ( + MiniCPMO45DataPlaneSession, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +REQUEST_ID = "duplex-s.c2Vzc2lvbg==.e.0.r.stage0" + + +def _session() -> MiniCPMO45DataPlaneSession: + return MiniCPMO45DataPlaneSession(lambda *args, **kwargs: "") + + +def _snapshot(num_samples: int) -> np.ndarray: + return np.arange(num_samples, dtype=np.float32) + + +def _sliced_len(session: MiniCPMO45DataPlaneSession, num_samples: int) -> int: + sliced = session.slice_cumulative_audio(REQUEST_ID, _snapshot(num_samples)) + if sliced is None: + return 0 + return int(np.asarray(sliced).reshape(-1).size) + + +def test_a_growing_snapshot_yields_only_the_new_tail() -> None: + session = _session() + assert _sliced_len(session, 24000) == 24000 + assert _sliced_len(session, 48000) == 24000 + assert _sliced_len(session, 48000) == 0 # unchanged snapshot emits nothing + + +def test_a_placeholder_unit_does_not_rewind_the_cursor() -> None: + """The failure behind a 300s benchmark case producing 1373s of audio. + + A continuation unit arrives as a one-sample snapshot. Taking that as "the + producer restarted" drops the cursor to 1, so the next real snapshot is + sliced from the start and the client is handed the entire session again. + """ + session = _session() + assert _sliced_len(session, 637440) == 637440 + assert _sliced_len(session, 1) == 1 + assert _sliced_len(session, 1) == 1 + assert _sliced_len(session, 1) == 1 + + # 649920 - 637440, not 649920 - 1. + assert _sliced_len(session, 649920) == 12480 + + +def test_a_genuinely_restarted_producer_is_still_picked_up() -> None: + """Two consecutive below-cursor snapshots that grow are a new buffer.""" + session = _session() + assert _sliced_len(session, 48000) == 48000 + assert _sliced_len(session, 1000) == 1000 # head of the new buffer + assert _sliced_len(session, 3000) == 2000 # its tail, not 3000 - 48000 + assert _sliced_len(session, 5000) == 2000 + + +def test_close_stream_releases_the_cursor_for_the_next_stream() -> None: + session = _session() + assert _sliced_len(session, 48000) == 48000 + session.close_stream(REQUEST_ID) + assert _sliced_len(session, 24000) == 24000 diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index e12010ce4..6c22b818b 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -55,20 +55,18 @@ class TestRegistryDeclaration: assert pipeline.model_type == _PIPELINE_KEY assert pipeline.model_arch == "MiniCPMO45OmniForConditionalGeneration" - def test_native_duplex_control_is_explicit_without_a_fixed_session_cap(self) -> None: + def test_duplex_plugin_is_declared_without_a_fixed_session_cap(self) -> None: pipeline = OMNI_PIPELINES[_PIPELINE_KEY] - assert pipeline.duplex_control_enabled is True - assert pipeline.duplex_serving_adapter == ( - "vllm_omni.model_executor.models.minicpmo_4_5.duplex.serving_adapter.MiniCPMO45ServingRuntimeAdapter" + assert pipeline.duplex_plugin == ( + "vllm_omni.model_executor.models.minicpmo_4_5.duplex.plugin.MiniCPMO45DuplexPlugin" ) assert not hasattr(pipeline, "max_native_duplex_sessions") - def test_ordinary_pipeline_defaults_to_no_duplex_control(self) -> None: + def test_ordinary_pipeline_declares_no_duplex_plugin(self) -> None: pipeline = PipelineConfig( model_type="ordinary", stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),) ) - assert pipeline.duplex_control_enabled is False - assert pipeline.duplex_serving_adapter is None + assert pipeline.duplex_plugin is None assert not hasattr(pipeline, "max_native_duplex_sessions") diff --git a/tests/model_executor/models/nemotron_voicechat/duplex/conftest.py b/tests/model_executor/models/nemotron_voicechat/duplex/conftest.py new file mode 100644 index 000000000..134731bf6 --- /dev/null +++ b/tests/model_executor/models/nemotron_voicechat/duplex/conftest.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Nemotron VoiceChat duplex is not on the duplex plugin framework yet. + +RFC vllm-omni#7181 splits the port: this PR lands the unified framework with +MiniCPM-o 4.5, and the follow-up PR moves Nemotron VoiceChat onto +``DuplexModelPlugin``. Until then the modules under test still import the +pre-framework duplex runtime / serving adapter that this PR removed, so +collecting them would fail at import time and abort the whole pytest session. +The tests themselves are unchanged and come back with the port. + +Listed per file rather than globbed, so the coverage loss is visible and shrinks +file by file as the port lands. ``test_input.py`` imports nothing that was +removed and keeps running. +""" + +#: Both import ``DuplexFence`` from ``engine.duplex.messages``, which no longer +#: exports it; ``test_data_plane`` also imports ``engine.duplex.runtime``. +collect_ignore = ["test_data_plane.py", "test_runtime_contract.py"] diff --git a/tests/model_executor/models/nemotron_voicechat/duplex/test_data_plane.py b/tests/model_executor/models/nemotron_voicechat/duplex/test_data_plane.py index 4bdc835d9..04e57c397 100644 --- a/tests/model_executor/models/nemotron_voicechat/duplex/test_data_plane.py +++ b/tests/model_executor/models/nemotron_voicechat/duplex/test_data_plane.py @@ -7,9 +7,9 @@ from types import SimpleNamespace import numpy as np import pytest +from vllm_omni.engine.duplex.runtime import duplex_resource_request_id from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.runtime import duplex_resource_request_id from vllm_omni.model_executor.models.nemotron_voicechat.duplex.data_plane import ( NemotronVoiceChatDataPlaneContext, NemotronVoiceChatDataPlaneSession, diff --git a/tests/model_executor/models/personaplex/duplex/conftest.py b/tests/model_executor/models/personaplex/duplex/conftest.py new file mode 100644 index 000000000..d49e8c617 --- /dev/null +++ b/tests/model_executor/models/personaplex/duplex/conftest.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""PersonaPlex duplex is not on the duplex plugin framework yet. + +RFC vllm-omni#7181 splits the port: this PR lands the unified framework with +MiniCPM-o 4.5, and the follow-up PR moves PersonaPlex onto ``DuplexModelPlugin``. +Until then the modules under test still import the pre-framework duplex +runtime / serving adapter that this PR removed, so collecting them would fail +at import time and abort the whole pytest session. The tests themselves are +unchanged and come back with the port. + +Listed per file rather than globbed, so the coverage loss is visible and shrinks +file by file as the port lands. +""" + +#: ``test_unified_runtime`` imports ``engine.duplex.runtime`` and +#: ``entrypoints.duplex.runtime_adapter`` directly. ``test_stage0_runtime`` fails +#: transitively: ``personaplex.duplex.runtime_extension`` imports ``DuplexFence`` +#: from ``engine.duplex.messages``, which no longer exports it. +collect_ignore = ["test_stage0_runtime.py", "test_unified_runtime.py"] diff --git a/tests/model_executor/models/personaplex/duplex/test_unified_runtime.py b/tests/model_executor/models/personaplex/duplex/test_unified_runtime.py index 1f03e894b..e4f9ad9b1 100644 --- a/tests/model_executor/models/personaplex/duplex/test_unified_runtime.py +++ b/tests/model_executor/models/personaplex/duplex/test_unified_runtime.py @@ -8,6 +8,12 @@ from types import SimpleNamespace import numpy as np import pytest from vllm.sampling_params import SamplingParams +from vllm_omni.engine.duplex.runtime import ( + DuplexInputMode, +) +from vllm_omni.entrypoints.duplex.runtime_adapter import ( + ServingRuntimeConfigError, +) from tests.e2e.online_serving import personaplex_realtime_duplex as e2e_driver from vllm_omni.config.stage_config import ( @@ -15,12 +21,6 @@ from vllm_omni.config.stage_config import ( merge_pipeline_deploy, ) from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.runtime import ( - DuplexInputMode, -) -from vllm_omni.entrypoints.duplex.runtime_adapter import ( - ServingRuntimeConfigError, -) from vllm_omni.model_executor.models.personaplex.duplex.config import DEFAULT_PERSONA from vllm_omni.model_executor.models.personaplex.duplex.data_plane import ( PersonaPlexDataPlaneContext, diff --git a/tests/worker/test_native_duplex_hooks.py b/tests/worker/test_native_duplex_hooks.py index c47802fbc..388e334cc 100644 --- a/tests/worker/test_native_duplex_hooks.py +++ b/tests/worker/test_native_duplex_hooks.py @@ -23,7 +23,7 @@ def _minicpmo_duplex_policy_case( MiniCPMO45OmniForConditionalGeneration, ) - session_key = ("sid-policy", 1) + session_key = "sid-policy" model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) model.model_stage = "llm" model._minicpmo45_native_duplex_token_ids_cache = { @@ -35,8 +35,7 @@ def _minicpmo_duplex_policy_case( row = DuplexSamplingRow( row_idx=0, request_id="req-policy", - session_id=session_key[0], - incarnation=session_key[1], + session_id=session_key, seq=3, payload=payload, max_tokens=20, @@ -63,13 +62,12 @@ def test_minicpmo_model_hook_owns_duplex_sampling_rows_and_force_listen(): "tts_bos_token_id": 8, "turn_eos_token_id": 9, } - model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={("sid-hook", 2): state}) + model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={"sid-hook": state}) logits = torch.zeros((1, 16), dtype=torch.float32) row = DuplexSamplingRow( row_idx=0, request_id="req-hook", session_id="sid-hook", - incarnation=2, seq=3, payload={"force_listen": True, "is_speech": True}, max_tokens=20, @@ -78,7 +76,7 @@ def test_minicpmo_model_hook_owns_duplex_sampling_rows_and_force_listen(): model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,)) assert model._minicpmo45_active_duplex_rows == [0] - assert model._minicpmo45_duplex_row_sessions == {0: ("sid-hook", 2)} + assert model._minicpmo45_duplex_row_sessions == {0: "sid-hook"} assert model._minicpmo45_duplex_row_payloads == {0: row.payload} assert model._minicpmo45_duplex_row_max_tokens == {0: 20} assert logits[0, listen_id].item() == 0.0 @@ -152,7 +150,7 @@ def test_minicpmo_model_hook_old_response_output_does_not_clear_pending_speech_c pending_speech_context=True, ) model, row = _minicpmo_duplex_policy_case(state, {"is_speech": True}) - model._minicpmo45_duplex_row_sessions = {0: (row.session_id, row.incarnation)} + model._minicpmo45_duplex_row_sessions = {0: row.session_id} model._minicpmo45_duplex_row_payloads = {0: row.payload} model._record_minicpmo45_duplex_terminator( @@ -173,7 +171,7 @@ def test_minicpmo_model_hook_new_response_output_clears_pending_speech_context() pending_speech_context=True, ) model, row = _minicpmo_duplex_policy_case(state, {"is_speech": False}) - model._minicpmo45_duplex_row_sessions = {0: (row.session_id, row.incarnation)} + model._minicpmo45_duplex_row_sessions = {0: row.session_id} model._minicpmo45_duplex_row_payloads = {0: row.payload} model._record_minicpmo45_duplex_terminator( @@ -195,7 +193,7 @@ def test_minicpmo_model_hook_empty_speak_envelope_preserves_pending_speech_conte pending_speech_response_open=False, ) model, row = _minicpmo_duplex_policy_case(state, {"is_speech": False}) - model._minicpmo45_duplex_row_sessions = {0: (row.session_id, row.incarnation)} + model._minicpmo45_duplex_row_sessions = {0: row.session_id} model._minicpmo45_duplex_row_payloads = {0: row.payload} token_ids = { "listen_token_id": 7, @@ -236,7 +234,7 @@ def test_minicpmo_model_hook_second_new_response_step_is_not_forced_to_listen(): original_logits = logits.clone() model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,)) - model._minicpmo45_duplex_row_sessions = {0: (row.session_id, row.incarnation)} + model._minicpmo45_duplex_row_sessions = {0: row.session_id} model._minicpmo45_duplex_row_payloads = {0: row.payload} model._record_minicpmo45_duplex_terminator( 0, @@ -269,7 +267,7 @@ def test_minicpmo_model_hook_same_speech_row_second_step_does_not_rearm_pending_ logits[0, 10] = 20.0 model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,)) - model._minicpmo45_duplex_row_sessions = {0: (row.session_id, row.incarnation)} + model._minicpmo45_duplex_row_sessions = {0: row.session_id} model._minicpmo45_duplex_row_payloads = {0: row.payload} model._record_minicpmo45_duplex_terminator( 0, @@ -339,7 +337,6 @@ def test_generic_ar_runner_builds_typed_duplex_sampling_rows(): "duplex": { "data_plane": True, "session_id": "sid-runner-hook", - "incarnation": 4, "seq": 4, "payload": {"is_speech": True}, } @@ -359,7 +356,6 @@ def test_generic_ar_runner_builds_typed_duplex_sampling_rows(): assert rows[0].row_idx == 0 assert rows[0].request_id == "req-duplex" assert rows[0].session_id == "sid-runner-hook" - assert rows[0].incarnation == 4 assert rows[0].seq == 4 assert rows[0].payload == {"is_speech": True} assert rows[0].max_tokens == 32 @@ -458,13 +454,13 @@ def test_generic_ar_runner_has_no_minicpmo_sampler_state_or_typeerror_probe(): assert 'if "duplex_rows" not in str(exc)' not in source -def test_minicpmo_model_cleans_incarnation_state_when_request_finishes(): +def test_minicpmo_model_cleans_session_state_when_request_finishes(): from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( MiniCPMO45OmniForConditionalGeneration, ) request_id = "duplex-sid-cleanup-i3-e0-stage0" - session_key = ("sid-cleanup", 3) + session_key = "sid-cleanup" model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) model.model = SimpleNamespace() model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: object()}) @@ -1731,8 +1727,8 @@ def test_minicpmo_stage0_native_sampler_penalizes_text_from_prior_chunk(): model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) model.model_stage = "llm" model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) - session_key = ("sid-cross-chunk-repetition", 0) - state = _MiniCPMO45Stage0SessionState(session_id=session_key[0]) + session_key = "sid-cross-chunk-repetition" + state = _MiniCPMO45Stage0SessionState(session_id=session_key) state.generated_tokens = [198] * 8 model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state}) model._minicpmo45_duplex_row_sessions = {0: session_key} @@ -1791,8 +1787,8 @@ def test_minicpmo_stage0_native_sampler_matches_official_negative_logit_penalty( model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) model.model_stage = "llm" model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) - session_key = ("sid-negative-logit-repetition", 0) - state = _MiniCPMO45Stage0SessionState(session_id=session_key[0]) + session_key = "sid-negative-logit-repetition" + state = _MiniCPMO45Stage0SessionState(session_id=session_key) repeated = 198 alternative = 1234 state.generated_tokens = [repeated] @@ -1827,8 +1823,8 @@ def test_minicpmo_stage0_records_bounded_model_policy_history(): MiniCPMO45OmniForConditionalGeneration, ) - session_key = ("sid-text-history", 0) - state = _MiniCPMO45Stage0SessionState(session_id=session_key[0]) + session_key = "sid-text-history" + state = _MiniCPMO45Stage0SessionState(session_id=session_key) model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state}) model._minicpmo45_duplex_row_sessions = {0: session_key} diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index c5a96148e..959007c30 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -79,16 +79,15 @@ CHECK_IMPORTS = { "vllm_omni/entrypoints/openai/video_stream_base.py", "vllm_omni/entrypoints/openai/video_stream_session.py", "vllm_omni/experimental/fullduplex/client.py", + "vllm_omni/engine/duplex/audio.py", + "vllm_omni/engine/duplex/commands.py", "vllm_omni/engine/duplex/contracts.py", + "vllm_omni/engine/duplex/events.py", + "vllm_omni/engine/duplex/realtime_commands.py", "vllm_omni/experimental/fullduplex/joyvl/memory/memory.py", - "vllm_omni/model_executor/models/minicpmo_4_5/duplex/adapter.py", "vllm_omni/model_executor/models/minicpmo_4_5/duplex/input.py", - "vllm_omni/model_executor/models/minicpmo_4_5/duplex/runtime.py", + "vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py", "vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py", - "vllm_omni/entrypoints/duplex/audio.py", - "vllm_omni/entrypoints/duplex/realtime_input.py", - "vllm_omni/entrypoints/duplex/runtime_bridge.py", - "vllm_omni/entrypoints/duplex/serving.py", "vllm_omni/model_executor/models/personaplex/duplex/input.py", "vllm_omni/model_executor/models/personaplex/duplex/runtime_extension.py", "vllm_omni/model_executor/models/personaplex/duplex/stage0.py", diff --git a/vllm_omni/benchmarks/omniinteract.py b/vllm_omni/benchmarks/omniinteract.py index 83e224f34..e566ac12b 100644 --- a/vllm_omni/benchmarks/omniinteract.py +++ b/vllm_omni/benchmarks/omniinteract.py @@ -417,6 +417,13 @@ def _has_post_commit_decision( # refusal is harmless — record it instead of failing the case. _TOLERATED_ERROR_CODES = frozenset({"playback_ack_too_late"}) +# ``session.close`` carries no reason of its own, so the server stamps the +# terminal event with the reason it inferred for it. Anything else on a +# ``session.closed`` after our own close request means the session ended for a +# reason we did not ask for (a disconnect, a lease timeout, a shutdown), which +# is what this guard is looking for. +_CLIENT_CLOSE_REASONS = frozenset({"client_close"}) + def _error_code(event: dict[str, object]) -> str | None: error = event.get("error") @@ -452,7 +459,7 @@ def _raise_if_session_terminated( event_type == "session.closed" and explicit_close_from is not None and index >= explicit_close_from - and reason is None + and (reason is None or reason in _CLIENT_CLOSE_REASONS) ) if expected: continue @@ -826,7 +833,7 @@ def write_batch_artifacts( ) -def _websocket_url(config: OmniInteractBenchmarkConfig, session_id: str) -> str: +def _websocket_url(config: OmniInteractBenchmarkConfig) -> str: endpoint = ( config.endpoint if urlsplit(config.endpoint).scheme @@ -834,13 +841,7 @@ def _websocket_url(config: OmniInteractBenchmarkConfig, session_id: str) -> str: ) from vllm_omni.clients.duplex import build_realtime_url - return build_realtime_url( - endpoint, - config.model, - autostart=False, - native_duplex=True, - session_id=session_id, - ) + return build_realtime_url(endpoint, config.model, autostart=False) class _RealtimeSession: @@ -855,7 +856,7 @@ class _RealtimeSession: _MAX_FRAME_BYTES = 64 * 1024 * 1024 - def __init__(self, config: OmniInteractBenchmarkConfig, session_id: str, reference_audio: str) -> None: + def __init__(self, config: OmniInteractBenchmarkConfig, reference_audio: str) -> None: from vllm_omni.clients.duplex import DuplexClient, EventCollector, SessionConfig self.session_config = SessionConfig( @@ -864,12 +865,11 @@ class _RealtimeSession: playback_commit_policy="ack_only", idle_timeout_s=float(config.timeout_s), extra_body={ - "native_duplex": True, "force_listen_count": 0, **(config.extra_body or {}), }, ) - self.url = _websocket_url(config, session_id) + self.url = _websocket_url(config) headers = dict(config.extra_headers or {}) async def connect(url: str) -> WebSocketTransport: @@ -886,7 +886,6 @@ class _RealtimeSession: self.url, model=config.model, config=self.session_config, - session_id=session_id, reconnect=None, heartbeat_interval_s=None, handshake_timeout_s=min(config.timeout_s, 20.0), @@ -1013,7 +1012,7 @@ async def run_omniinteract_case( frames = prepared_input.video_frames if not any(frames): raise ValueError(f"No video frames were decoded from {case.video_path}") - async with _RealtimeSession(config, session_id, reference_audio) as client: + async with _RealtimeSession(config, reference_audio) as client: session_from = 0 # the collector holds only this session's events pcm = _ensure_final_commit_tail(pcm, client.events.events) playback = _Playback() diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index a230b9c2e..bee910c90 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -14,7 +14,7 @@ import time import traceback import uuid import wave -from collections.abc import Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass, replace from datetime import datetime from pathlib import Path @@ -2153,10 +2153,68 @@ async def async_request_openai_audio_speech( return output +#: Silence budget per Seed-TTS turn: a model-native duplex session generates +#: per audio unit, and the target text rides the session context, so the +#: silence only advances the clock. It stops at the turn's response.done -- +#: a native model that keeps hearing silence after its turn may decide to +#: speak again, and the benchmark measures one response per utterance -- so +#: the budget is only spent on a turn the model is slow to take: it may +#: choose to listen on a few units first, and a turn that has not settled +#: when the budget runs out is reported with what the model did. +_SEED_TTS_SILENCE_SECONDS = 30.0 +#: A native model normally answers the seeded text within this much silence; +#: a turn that needs more is logged so a slow-to-speak model shows in the run. +_SEED_TTS_PROMPT_RESPONSE_S = 12.0 +#: MiniCPM-o emits 24 kHz mono; used to report audio_frames after the session closed. +_SEED_TTS_OUTPUT_SAMPLE_RATE_HZ = 24_000 + + +def _seed_tts_turn_stall_report(events: object, response_offset: int, request_index: int, silence_s: float) -> str: + """Explain a Seed-TTS turn that never settled: what the model did with the silence.""" + response_ids = list(getattr(events, "response_ids")[response_offset:]) + if not response_ids: + return ( + f"Seed-TTS Realtime TTS turn {request_index} never started a response: the model listened " + f"through {silence_s:.1f}s of silence and the wait that followed" + ) + audio_bytes = getattr(events, "audio_bytes") + response_text = getattr(events, "response_text") + started = ", ".join( + f"{response_id} ({len(audio_bytes(response_id))} audio bytes, text {response_text(response_id)!r})" + for response_id in response_ids + ) + return ( + f"Seed-TTS Realtime TTS turn {request_index} started {len(response_ids)} response(s) after " + f"{silence_s:.1f}s of silence but none reached response.done: {started}" + ) + + +def _seed_tts_turn_response_id(events: object, response_offset: int, request_index: int) -> str: + """The response id the Seed-TTS turn is measured on: the first one with audio. + + A model-native session answers the seeded text once, but nothing in the + protocol stops it from speaking again on silence it hears afterwards, so + a later audio response is the model's own and not a failed turn. + """ + response_ids = getattr(events, "response_ids") + audio_bytes = getattr(events, "audio_bytes") + audio_response_ids = [response_id for response_id in response_ids[response_offset:] if audio_bytes(response_id)] + if not audio_response_ids: + raise RuntimeError(f"Seed-TTS Realtime TTS turn {request_index} produced no audio response") + if len(audio_response_ids) > 1: + logger.warning( + "Seed-TTS Realtime TTS turn %d: model spoke again after its response (%d audio responses); " + "measuring the first", + request_index, + len(audio_response_ids), + ) + return audio_response_ids[0] + + def _realtime_websocket_url(api_url: str) -> str: from vllm_omni.clients.duplex import build_realtime_url - return build_realtime_url(api_url, None, native_duplex=None) + return build_realtime_url(api_url, None) def _nonnegative_number(value: object) -> bool: @@ -2321,6 +2379,13 @@ class _RealtimeTTSProbe: request time. """ + #: How long ``configure`` waits for a free duplex session. Every Seed-TTS + #: utterance is one session and the deploy config admits ``max_sessions`` + #: of them, so a benchmark run above that concurrency queues for a slot + #: rather than counting the server's (retryable) refusal as a failed + #: request. + _SESSION_SLOT_WAIT_S = 120.0 + def __init__(self, url: str) -> None: from vllm_omni.clients.duplex import EventCollector @@ -2347,16 +2412,13 @@ class _RealtimeTTSProbe: *, output_audio_format: str = "pcm16", instructions: str | None = None, - native_duplex: bool = False, auto_response: bool = False, extra_body: dict[str, object] | None = None, - session_id: str | None = None, timeout_s: float = 120.0, ) -> None: - from vllm_omni.clients.duplex import AudioFormat, DuplexClient, SessionConfig + from vllm_omni.clients.duplex import AudioFormat, DuplexClient, DuplexProtocolError, SessionConfig session_extra_body: dict[str, object] = dict(extra_body or {}) - session_extra_body["native_duplex"] = bool(native_duplex) config = SessionConfig( output_audio=AudioFormat(output_audio_format, 24_000), instructions=instructions, @@ -2365,16 +2427,33 @@ class _RealtimeTTSProbe: playback_commit_policy="ack_only", extra_body=session_extra_body, ) - self._client = DuplexClient( - self._url, - model=model, - config=config, - session_id=session_id, - reconnect=None, - heartbeat_interval_s=None, - handshake_timeout_s=timeout_s, - ) - await self._client.__aenter__() + deadline = time.monotonic() + self._SESSION_SLOT_WAIT_S + delay_s = 0.25 + waited = False + while True: + # A failed handshake closes the socket on the client's side, so + # a refused attempt leaves nothing behind to clean up. + client = DuplexClient( + self._url, + model=model, + config=config, + reconnect=None, + heartbeat_interval_s=None, + handshake_timeout_s=timeout_s, + ) + try: + await client.__aenter__() + except DuplexProtocolError as exc: + if exc.code != "resource_exhausted" or time.monotonic() >= deadline: + raise + if not waited: + logger.info("Seed-TTS Realtime TTS: no free duplex session (%s); waiting for a slot", exc) + waited = True + await asyncio.sleep(delay_s) + delay_s = min(delay_s * 2.0, 2.0) + continue + self._client = client + break self._consume_task = asyncio.create_task(self.events.consume(self._client)) async def send(self, event: dict[str, object]) -> None: @@ -2387,6 +2466,32 @@ class _RealtimeTTSProbe: assert self._client is not None await acknowledge_collected_playback(self._client, self.events) + async def stream_silence( + self, + *, + seconds: float, + chunk_ms: int = 200, + until: Callable[[], bool] | None = None, + ) -> float: + """Append silent PCM16 units so a model-native session has units to speak on. + + A duplex model generates per audio unit. The target text rides the + session context (``duplex_initial_user_text``), so the audio only has + to advance the clock; silence keeps it from adding content of its own. + Streams in real time for at most ``seconds``, stopping as soon as + ``until`` holds, and returns the seconds actually appended. + """ + assert self._client is not None + input_format = self._client.config.input_audio + chunk = bytes(max(input_format.byte_count(chunk_ms), input_format.bytes_per_sample)) + chunk_s = input_format.duration_ms(len(chunk)) / 1000.0 + streamed_s = 0.0 + while streamed_s < seconds and not (until is not None and until()): + await self._client.append_audio(chunk, is_speech=False) + streamed_s += chunk_s + await asyncio.sleep(chunk_s) + return streamed_s + async def close_session(self, *, timeout_s: float = 20.0) -> None: assert self._client is not None await self._client.close(timeout_s=timeout_s) @@ -2423,69 +2528,70 @@ async def async_request_openai_realtime_duplex( if not turn_prompts: turn_prompts = [("", request_func_input.prompt)] session_id = f"seed-tts-{request_func_input.request_id or uuid.uuid4().hex}" + silence_seconds = float(getattr(request_func_input, "seed_tts_silence_seconds", 0.0) or _SEED_TTS_SILENCE_SECONDS) + turn_metrics: list[dict[str, object]] = [] + turn_timings: list[dict[str, object]] = [] + turn_pcm_bytes: list[bytes] = [] + turn_transcripts: list[str] = [] + measurement_origin = { + "ttft": "first silence append client send to first non-empty text delta", + "ttfp": "first silence append client send to first audio packet", + "rtf": "request-start-to-last-audio receive time divided by emitted audio duration", + } try: - async with _RealtimeTTSProbe(_realtime_websocket_url(request_func_input.api_url)) as client: - await client.configure( - request_func_input.model_name or request_func_input.model, - output_audio_format="pcm16", - instructions=getattr( - request_func_input, - "seed_tts_system_prompt", - SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT, - ), - native_duplex=False, - auto_response=False, - extra_body=speech_extra, - session_id=session_id, - timeout_s=120.0, - ) - turn_metrics: list[dict[str, object]] = [] - turn_timings: list[dict[str, object]] = [] - turn_pcm_bytes: list[bytes] = [] - turn_transcripts: list[str] = [] - measurement_origin = { - "ttft": "conversation.item.create client send to first non-empty text delta", - "ttfp": "conversation.item.create client send to first audio packet", - "rtf": "request-start-to-last-audio receive time divided by emitted audio duration", - } - for request_index, (utterance_id, target_text) in enumerate(turn_prompts): + # One session per utterance. A model-native duplex session takes its + # text once, in the session context (``duplex_initial_user_text``), so + # a session cannot be re-seeded for a second target text. + for request_index, (utterance_id, target_text) in enumerate(turn_prompts): + async with _RealtimeTTSProbe(_realtime_websocket_url(request_func_input.api_url)) as client: + await client.configure( + request_func_input.model_name or request_func_input.model, + output_audio_format="pcm16", + instructions=getattr( + request_func_input, + "seed_tts_system_prompt", + SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT, + ), + auto_response=True, + extra_body={ + **speech_extra, + "duplex_initial_user_text": target_text, + "force_listen_count": 0, + }, + timeout_s=120.0, + ) response_offset = len(client.events.response_ids) done_before = client.events.count("response.done") errors_before = len(client.events.errors()) - turn_started_at_s = time.monotonic() - await client.send( - { - "type": "conversation.item.create", - "item": { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": target_text}], - }, - } - ) - await client.send({"type": "response.create"}) - await wait_for_condition( - lambda: ( + + def turn_settled() -> bool: + return ( client.events.count("response.done") > done_before or len(client.events.errors()) > errors_before - ), - timeout_s=180.0, - label=f"Seed-TTS Realtime TTS turn {request_index} response.done", - ) + ) + + turn_started_at_s = time.monotonic() + silence_s = await client.stream_silence(seconds=silence_seconds, until=turn_settled) + if silence_s > _SEED_TTS_PROMPT_RESPONSE_S: + logger.warning( + "Seed-TTS Realtime TTS turn %d: the model took %.1fs of silence to settle its response", + request_index, + silence_s, + ) + try: + await wait_for_condition( + turn_settled, + timeout_s=180.0, + label=f"Seed-TTS Realtime TTS turn {request_index} response.done", + ) + except TimeoutError as exc: + raise RuntimeError( + _seed_tts_turn_stall_report(client.events, response_offset, request_index, silence_s) + ) from exc errors = client.events.errors() if len(errors) > errors_before: raise RuntimeError(f"Seed-TTS Realtime TTS server error: {errors[-1]}") - new_audio_response_ids = [ - response_id - for response_id in client.events.response_ids[response_offset:] - if client.events.audio_bytes(response_id) - ] - if len(new_audio_response_ids) != 1: - raise RuntimeError( - f"Seed-TTS Realtime TTS turn {request_index} expected one audio response, " - f"got {len(new_audio_response_ids)}" - ) - response_id = new_audio_response_ids[0] + response_id = _seed_tts_turn_response_id(client.events, response_offset, request_index) timing = client.events.timing_summary( after_s=turn_started_at_s, input_committed_at_s=turn_started_at_s, @@ -2518,45 +2624,43 @@ async def async_request_openai_realtime_duplex( ) turn_transcripts.append(client.events.response_text(response_id)) await client.acknowledge_playback() - request_finished_at = time.perf_counter() - session_metrics = summarize_session_request_metrics( - turn_metrics, - session_id=session_id, - ) - await client.close_session(timeout_s=30.0) + await client.close_session(timeout_s=30.0) + request_finished_at = time.perf_counter() + session_metrics = summarize_session_request_metrics( + turn_metrics, + session_id=session_id, + ) - output.generated_text = " ".join(filter(None, turn_transcripts)) - from vllm_omni.clients.duplex import metric_mean + output.generated_text = " ".join(filter(None, turn_transcripts)) + from vllm_omni.clients.duplex import metric_mean - output.ttft = (metric_mean(session_metrics.get("ttft_ms")) or 0.0) / 1000.0 - output.audio_ttfp = (metric_mean(session_metrics.get("ttfp_ms")) or 0.0) / 1000.0 - output.audio_rtf = metric_mean(session_metrics.get("rtf")) or 0.0 - output.audio_duration = ( - sum(float(metric.get("audio_duration_ms") or 0.0) for metric in turn_metrics) / 1000.0 + output.ttft = (metric_mean(session_metrics.get("ttft_ms")) or 0.0) / 1000.0 + output.audio_ttfp = (metric_mean(session_metrics.get("ttfp_ms")) or 0.0) / 1000.0 + output.audio_rtf = metric_mean(session_metrics.get("rtf")) or 0.0 + output.audio_duration = sum(float(metric.get("audio_duration_ms") or 0.0) for metric in turn_metrics) / 1000.0 + output.audio_frames = int(output.audio_duration * _SEED_TTS_OUTPUT_SAMPLE_RATE_HZ) + output.latency = request_finished_at - output.start_time + output.tts_turn_pcm_bytes = turn_pcm_bytes + output.tts_output_pcm_bytes = b"".join(turn_pcm_bytes) + if bool((request_func_input.extra_body or {}).get("save_duplex_request_metrics")): + output.duplex_request_metrics = turn_metrics + output.duplex_session_metrics = session_metrics + output.output_tokens = sum( + int(stage0.get("output_token_count") or 0) + for timing in turn_timings + if isinstance((stage0 := timing.get("stage0_tokens")), dict) + ) + token_timing_measured = _apply_stage0_token_timings( + output, + [timing.get("stage0_tokens") for timing in turn_timings], + expected_output_tokens=output.output_tokens, + ) + if not token_timing_measured and output.output_tokens > 1: + logger.warning( + "Realtime TTS session %s omitted complete engine token timing; standard TPOT/ITL are unavailable", + session_id, ) - output.audio_frames = int(output.audio_duration * client.events.output_sample_rate_hz) - output.latency = request_finished_at - output.start_time - output.tts_turn_pcm_bytes = turn_pcm_bytes - output.tts_output_pcm_bytes = b"".join(turn_pcm_bytes) - if bool((request_func_input.extra_body or {}).get("save_duplex_request_metrics")): - output.duplex_request_metrics = turn_metrics - output.duplex_session_metrics = session_metrics - output.output_tokens = sum( - int(stage0.get("output_token_count") or 0) - for timing in turn_timings - if isinstance((stage0 := timing.get("stage0_tokens")), dict) - ) - token_timing_measured = _apply_stage0_token_timings( - output, - [timing.get("stage0_tokens") for timing in turn_timings], - expected_output_tokens=output.output_tokens, - ) - if not token_timing_measured and output.output_tokens > 1: - logger.warning( - "Realtime TTS session %s omitted complete engine token timing; standard TPOT/ITL are unavailable", - session_id, - ) - output.success = True + output.success = True except Exception: output.success = False output.error = traceback.format_exc() diff --git a/vllm_omni/clients/__init__.py b/vllm_omni/clients/__init__.py index 924582bf4..35f904778 100644 --- a/vllm_omni/clients/__init__.py +++ b/vllm_omni/clients/__init__.py @@ -2,8 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Client-side libraries for vLLM-Omni serving APIs. -Modules under this package speak to a running vLLM-Omni server over the -network. They are for applications and tests; server runtime code +``vllm_omni.clients.duplex`` speaks to a running vLLM-Omni server over the +network; ``vllm_omni.clients.inline_duplex`` drives an in-process +``DuplexOmni`` handed in by the caller with the same client API. Modules here +import no server runtime code at module load, and server runtime code (``vllm_omni.engine``, ``vllm_omni.entrypoints``, ``vllm_omni.model_executor``) must never import from here. """ diff --git a/vllm_omni/clients/duplex.py b/vllm_omni/clients/duplex.py index a1d9fc1df..52266b64e 100644 --- a/vllm_omni/clients/duplex.py +++ b/vllm_omni/clients/duplex.py @@ -2,9 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Public asynchronous client for the vLLM-Omni full-duplex Realtime API. -Connects to ``/v1/realtime?duplex=1`` (the normative duplex contract) and -provides typed events, response demultiplexing, incremental playback -acking, and transparent session resume on transport drops. +:class:`DuplexClientBase` holds the transport-agnostic session client (typed +events, response demultiplexing, incremental playback acking, the session +handshake). :class:`DuplexClient` connects it to ``/v1/realtime?duplex=1`` +(the normative duplex contract) over WebSocket with transparent session +resume on transport drops; :class:`vllm_omni.clients.inline_duplex.InlineDuplexClient` +drives an in-process :class:`~vllm_omni.entrypoints.duplex_omni.DuplexOmni` +with exactly the same usage. Example:: @@ -35,10 +39,10 @@ import math import random import time import wave +from abc import ABC, abstractmethod from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence from dataclasses import dataclass, field from pathlib import Path -from typing import Protocol from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit from uuid import uuid4 @@ -53,6 +57,7 @@ __all__ = [ "AudioDelta", "ConnectionResumed", "DuplexClient", + "DuplexClientBase", "DuplexClientError", "DuplexConnectionError", "DuplexEvent", @@ -229,7 +234,7 @@ class SessionConfig: idle_timeout_s: float | None = None extra_body: dict[str, object] = field(default_factory=dict) - def to_session_payload(self, *, model: str, session_id: str | None = None) -> dict[str, object]: + def to_session_payload(self, *, model: str) -> dict[str, object]: payload: dict[str, object] = { "model": model, "modalities": list(self.modalities), @@ -246,8 +251,6 @@ class SessionConfig: }, "turn_detection": self.turn_detection, } - if session_id: - payload["session_id"] = session_id if self.instructions is not None: payload["instructions"] = self.instructions if self.ref_audio is not None: @@ -368,11 +371,6 @@ class SessionCreated(DuplexEvent): value = self.raw.get("resume_token") return value if isinstance(value, str) else None - @property - def incarnation(self) -> int: - value = self.raw.get("incarnation") - return value if isinstance(value, int) else 0 - class SessionResumed(SessionCreated): pass @@ -598,79 +596,84 @@ class ResponseHandle: # Client -class WebSocketTransport(Protocol): +class WebSocketTransport(ABC): """Minimal WebSocket surface the client needs (satisfied by ``websockets``).""" + @abstractmethod async def send(self, data: str) -> None: ... + @abstractmethod async def recv(self) -> str | bytes: ... + @abstractmethod async def close(self) -> None: ... ConnectFn = Callable[[str], Awaitable["WebSocketTransport"]] -class DuplexClient: - """Async client for one duplex session over ``/v1/realtime?duplex=1``. +class DuplexClientBase(ABC): + """Transport-agnostic duplex session client. - ``session_id`` only names the session this client creates: entering the - client always performs the ``session.update`` handshake. Resuming an - existing session is supported only as automatic reconnect within the same - client instance (see ``ReconnectPolicy``); taking over a session from a - new client requires the wire-level ``session.resume`` handshake - (``resume_token``, ``incarnation``, ``last_received_server_event_seq``), - which this client does not expose. + Everything a caller touches (typed events, response demux, input helpers, + playback acks, the ``session.update`` handshake) lives here. A subclass + only supplies the transport: :meth:`_open` establishes it and must deliver + the server's ``session.created`` payload through :meth:`_dispatch`, + :meth:`_send_command` delivers one client event, :meth:`_teardown` + releases the transport. """ def __init__( self, - url: str, *, model: str, config: SessionConfig | None = None, - session_id: str | None = None, - reconnect: ReconnectPolicy | None = ReconnectPolicy(), - heartbeat_interval_s: float | None = 30.0, handshake_timeout_s: float = 30.0, - connect: ConnectFn | None = None, ) -> None: - self.url = url self.model = model self.config = config or SessionConfig() - self.session_id = session_id + #: Allocated by the server; ``None`` until ``session.created`` arrives. + self.session_id: str | None = None self.session_info: dict[str, object] = {} - self.incarnation = 0 - self.resume_token: str | None = None - self._reconnect = reconnect - self._heartbeat_interval_s = heartbeat_interval_s self._handshake_timeout_s = handshake_timeout_s - self._connect_fn: ConnectFn = connect or self._default_connect - self._ws: WebSocketTransport | None = None - self._reader_task: asyncio.Task[None] | None = None - self._heartbeat_task: asyncio.Task[None] | None = None self._subscribers: list[asyncio.Queue[DuplexEvent | _ClosedMarker]] = [] self._response_queue: asyncio.Queue[ResponseHandle | _ClosedMarker | _ErrorMarker] = asyncio.Queue( maxsize=_MAX_BUFFERED_EVENTS ) self._responses: dict[str, ResponseHandle] = {} - self._last_server_event_seq: int | None = None self._input_audio_end_ms = 0.0 self._closing = False self._closed = asyncio.Event() self._closed_marker: _ClosedMarker | None = None + # -- transport hooks ------------------------------------------------------ + + @abstractmethod + async def _open(self) -> None: + """Establish the transport and start delivering server events via ``_dispatch``. + + Called once from :meth:`__aenter__`; the base then sends the + ``session.update`` handshake and waits for ``session.created``. + """ + + @abstractmethod + async def _send_command(self, payload: dict[str, object]) -> None: + """Deliver one client event (already stamped with ``event_id``).""" + + @abstractmethod + async def _teardown(self) -> None: + """Release the transport; must be idempotent.""" + # -- lifecycle ---------------------------------------------------------- - async def __aenter__(self) -> DuplexClient: - self._ws = await self._connect_fn(self._target_url()) + async def __aenter__(self) -> DuplexClientBase: handshake_queue = self._add_subscriber() try: - self._reader_task = asyncio.create_task(self._read_loop(), name="duplex-client-reader") + await self._open() await self.send( { "type": "session.update", - "session": self.config.to_session_payload(model=self.model, session_id=self.session_id), + "session": self.config.to_session_payload(model=self.model), } ) created = await self._wait_on_queue( @@ -679,14 +682,12 @@ class DuplexClient: timeout_s=self._handshake_timeout_s, ) except BaseException: - await self._teardown() + await self._teardown_quietly() raise finally: self._remove_subscriber(handshake_queue) assert isinstance(created, SessionCreated) self._adopt_session(created) - if self._heartbeat_interval_s is not None: - self._heartbeat_task = asyncio.create_task(self._heartbeat_loop(), name="duplex-client-heartbeat") return self async def __aexit__(self, exc_type, exc, traceback) -> None: @@ -706,7 +707,7 @@ class DuplexClient: except DuplexClientError: pass finally: - await self._teardown() + await self._teardown_quietly() async def close(self, *, timeout_s: float = 20.0) -> None: """Send ``session.close`` and wait for the server to confirm.""" @@ -723,6 +724,12 @@ class DuplexClient: except asyncio.TimeoutError: self._finalize("close timed out", expected=True) + async def _teardown_quietly(self) -> None: + self._closing = True + if not self._closed.is_set(): + self._finalize("closed", expected=True) + await self._teardown() + # -- input ---------------------------------------------------------------- async def append_audio( @@ -925,53 +932,25 @@ class DuplexClient: payload = dict(event) payload.setdefault("event_id", f"evt_{uuid4().hex}") try: - await self._ws.send(json.dumps(payload)) + await self._send_command(payload) except asyncio.CancelledError: raise + except DuplexClientError: + raise except Exception as exc: raise DuplexConnectionError(f"send failed: {exc}") from exc return payload["event_id"] # -- internals --------------------------------------------------------------- - def _target_url(self) -> str: - parts = urlsplit(self.url) - path = parts.path if parts.path not in ("", "/") else "/v1/realtime" - query = dict(parse_qsl(parts.query, keep_blank_values=True)) - query.setdefault("duplex", "1") - query.setdefault("model", self.model) - # The client always opens with an explicit session.update (and resumes - # with an explicit session.resume). Autostart would race that handshake: - # with a model query param, the server creates a bare default session - # before reading the first client event, silently dropping ref_audio - # and the model-specific extra_body. Force it off even when the caller's - # URL carries autostart=1 — this client cannot operate over it. - query["autostart"] = "0" - return urlunsplit((parts.scheme or "ws", parts.netloc, path, urlencode(query), parts.fragment)) - - async def _default_connect(self, url: str) -> WebSocketTransport: - try: - import websockets - except ImportError as exc: # pragma: no cover - websockets is a pinned dep - raise DuplexConnectionError("The duplex client requires the 'websockets' package") from exc - try: - return await websockets.connect(url, max_size=_MAX_FRAME_BYTES) - except asyncio.CancelledError: - raise - except Exception as exc: - raise DuplexConnectionError(f"connect to {url} failed: {exc}") from exc - def _adopt_session(self, event: SessionCreated) -> None: - # A resume activation carries no nested session object; keep the - # session_info captured at the original handshake in that case. + # session.created and session.resumed both carry the engine's current + # session object; keep the last one seen. if event.session: self.session_info = event.session session_id = event.session.get("session_id") or event.session.get("id") or event.session_id if isinstance(session_id, str) and session_id: self.session_id = session_id - self.incarnation = event.incarnation - if event.resume_token: - self.resume_token = event.resume_token def _announce(self, handle: ResponseHandle) -> None: """Deliver a handle to responses() exactly once, decision known.""" @@ -1013,36 +992,18 @@ class DuplexClient: if item.type in types: return item - async def _read_loop(self) -> None: - while True: - try: - while True: - raw = await self._ws.recv() - if isinstance(raw, (bytes, bytearray)): - continue - try: - data = json.loads(raw) - except json.JSONDecodeError: - continue - if isinstance(data, dict): - await self._dispatch(data) - except asyncio.CancelledError: - raise - except Exception as exc: - if self._closing or self._closed.is_set(): - self._finalize("closed", expected=True) - return - if await self._attempt_resume(): - continue - self._finalize(f"connection lost: {exc}", expected=False) - return + def _accept_event(self, data: dict[str, object]) -> bool: + """Transport hook: return False to drop ``data`` before dispatch (duplicates).""" + del data + return True + + async def _after_dispatch(self, event: DuplexEvent, data: dict[str, object]) -> None: + """Transport hook run after an event was delivered to subscribers.""" + del event, data async def _dispatch(self, data: dict[str, object]) -> None: - seq = data.get("server_event_seq") - if isinstance(seq, int): - if self._last_server_event_seq is not None and seq <= self._last_server_event_seq: - return # duplicate delivered by resume replay - self._last_server_event_seq = seq + if not self._accept_event(data): + return event = wrap_event(data) if isinstance(event, SessionResumed): @@ -1109,14 +1070,130 @@ class DuplexClient: elif isinstance(event, SessionExpired): reason = data.get("reason") self._finalize(f"expired: {reason}" if reason else "expired", expected=False) - elif event.type == "session.resync_required": + + await self._after_dispatch(event, data) + + def _finalize(self, reason: str, *, expected: bool) -> None: + if self._closed.is_set(): + return + marker = _ClosedMarker(reason=reason, expected=expected) + self._closed_marker = marker + self._closed.set() + for handle in self._responses.values(): + handle._finish(None) + self._responses.clear() + for queue in (*self._subscribers, self._response_queue): + _put_drop_oldest(queue, marker) + + +class DuplexClient(DuplexClientBase): + """Async client for one duplex session over ``/v1/realtime?duplex=1``.""" + + def __init__( + self, + url: str, + *, + model: str, + config: SessionConfig | None = None, + reconnect: ReconnectPolicy | None = ReconnectPolicy(), + heartbeat_interval_s: float | None = 30.0, + handshake_timeout_s: float = 30.0, + connect: ConnectFn | None = None, + ) -> None: + super().__init__(model=model, config=config, handshake_timeout_s=handshake_timeout_s) + self.url = url + self.resume_token: str | None = None + self._reconnect = reconnect + self._heartbeat_interval_s = heartbeat_interval_s + self._connect_fn: ConnectFn = connect or self._default_connect + self._ws: WebSocketTransport | None = None + self._reader_task: asyncio.Task[None] | None = None + self._heartbeat_task: asyncio.Task[None] | None = None + self._last_server_event_seq: int | None = None + + # -- transport hooks ------------------------------------------------------ + + async def _open(self) -> None: + self._ws = await self._connect_fn(self._target_url()) + self._reader_task = asyncio.create_task(self._read_loop(), name="duplex-client-reader") + + async def __aenter__(self) -> DuplexClient: + await super().__aenter__() + if self._heartbeat_interval_s is not None: + self._heartbeat_task = asyncio.create_task(self._heartbeat_loop(), name="duplex-client-heartbeat") + return self + + async def _send_command(self, payload: dict[str, object]) -> None: + if self._ws is None: + raise DuplexConnectionError("send failed: transport is not connected") + await self._ws.send(json.dumps(payload)) + + async def _teardown(self) -> None: + for task in (self._heartbeat_task, self._reader_task): + if task is not None and not task.done(): + task.cancel() + try: + await task + except (asyncio.CancelledError, Exception): # noqa: BLE001 + pass + if self._ws is not None: + try: + await self._ws.close() + except Exception: # noqa: BLE001 + pass + + # -- internals --------------------------------------------------------------- + + def _target_url(self) -> str: + parts = urlsplit(self.url) + path = parts.path if parts.path not in ("", "/") else "/v1/realtime" + query = dict(parse_qsl(parts.query, keep_blank_values=True)) + query.setdefault("duplex", "1") + query.setdefault("model", self.model) + # The client always opens with an explicit session.update (and resumes + # with an explicit session.resume). Autostart would race that handshake: + # with a model query param, the server creates a bare default session + # before reading the first client event, silently dropping ref_audio + # and the model-specific extra_body. Force it off even when the caller's + # URL carries autostart=1 — this client cannot operate over it. + query["autostart"] = "0" + return urlunsplit((parts.scheme or "ws", parts.netloc, path, urlencode(query), parts.fragment)) + + async def _default_connect(self, url: str) -> WebSocketTransport: + try: + import websockets + except ImportError as exc: # pragma: no cover - websockets is a pinned dep + raise DuplexConnectionError("The duplex client requires the 'websockets' package") from exc + try: + return await websockets.connect(url, max_size=_MAX_FRAME_BYTES) + except asyncio.CancelledError: + raise + except Exception as exc: + raise DuplexConnectionError(f"connect to {url} failed: {exc}") from exc + + def _adopt_session(self, event: SessionCreated) -> None: + super()._adopt_session(event) + if event.resume_token: + self.resume_token = event.resume_token + + def _accept_event(self, data: dict[str, object]) -> bool: + seq = data.get("server_event_seq") + if isinstance(seq, int): + if self._last_server_event_seq is not None and seq <= self._last_server_event_seq: + return False # duplicate delivered by resume replay + self._last_server_event_seq = seq + return True + + async def _after_dispatch(self, event: DuplexEvent, data: dict[str, object]) -> None: + if event.type == "session.resync_required": # The server stopped journaling this session's events, so any # retained resume credential will be refused on the next # session.resume. Drop it now: a later transport failure then # finalizes immediately instead of burning reconnect attempts. self.resume_token = None - if isinstance(seq, int) and not self._closed.is_set(): + seq = data.get("server_event_seq") + if isinstance(seq, int) and not self._closed.is_set() and self._ws is not None: try: await self._ws.send(json.dumps({"type": "session.event_ack", "server_event_seq": seq})) except asyncio.CancelledError: @@ -1124,6 +1201,30 @@ class DuplexClient: except Exception: pass # the reader will observe the transport failure itself + async def _read_loop(self) -> None: + while True: + try: + while True: + raw = await self._ws.recv() + if isinstance(raw, (bytes, bytearray)): + continue + try: + data = json.loads(raw) + except json.JSONDecodeError: + continue + if isinstance(data, dict): + await self._dispatch(data) + except asyncio.CancelledError: + raise + except Exception as exc: + if self._closing or self._closed.is_set(): + self._finalize("closed", expected=True) + return + if await self._attempt_resume(): + continue + self._finalize(f"connection lost: {exc}", expected=False) + return + async def _attempt_resume(self) -> bool: policy = self._reconnect if policy is None or not self.resume_token or not self.session_id: @@ -1142,7 +1243,6 @@ class DuplexClient: { "type": "session.resume", "session_id": self.session_id, - "incarnation": self.incarnation, "resume_token": self.resume_token, "last_received_server_event_seq": self._last_server_event_seq or 0, } @@ -1162,7 +1262,10 @@ class DuplexClient: continue event_type = data.get("type") if event_type == "session.resumed": + previous = self._ws self._ws = ws + if previous is not None and previous is not ws: + await self._close_quietly(previous) await self._dispatch({"type": "connection.resumed", "attempt": attempt}) await self._dispatch(data) for entry in pending: @@ -1212,35 +1315,6 @@ class DuplexClient: except DuplexClientError: return - def _finalize(self, reason: str, *, expected: bool) -> None: - if self._closed.is_set(): - return - marker = _ClosedMarker(reason=reason, expected=expected) - self._closed_marker = marker - self._closed.set() - for handle in self._responses.values(): - handle._finish(None) - self._responses.clear() - for queue in (*self._subscribers, self._response_queue): - _put_drop_oldest(queue, marker) - - async def _teardown(self) -> None: - self._closing = True - if not self._closed.is_set(): - self._finalize("closed", expected=True) - for task in (self._heartbeat_task, self._reader_task): - if task is not None and not task.done(): - task.cancel() - try: - await task - except (asyncio.CancelledError, Exception): # noqa: BLE001 - pass - if self._ws is not None: - try: - await self._ws.close() - except Exception: # noqa: BLE001 - pass - # --------------------------------------------------------------------------- # Test/benchmark collector @@ -1368,7 +1442,7 @@ class EventCollector: self.response_ids: list[str] = [] self.output_sample_rate_hz = 24_000 - async def consume(self, client: DuplexClient) -> None: + async def consume(self, client: DuplexClientBase) -> None: """Subscribe to ``client`` and collect until the session ends.""" try: async for event in client.events(): @@ -1795,15 +1869,12 @@ def build_realtime_url( model: str | None, *, autostart: bool | None = None, - native_duplex: bool | None = None, - session_id: str | None = None, extra_query: dict[str, str] | None = None, ) -> str: """Add explicit duplex query parameters to a Realtime URL. :class:`DuplexClient` builds its own URL; this helper is for drivers that - speak the wire protocol directly. ``native_duplex`` sets the per-session - model-native opt-in query flag; other model-specific query flags ride in + speak the wire protocol directly. Model-specific query flags ride in ``extra_query``. ``http(s)`` URLs are rewritten to ``ws(s)``. """ parts = urlsplit(url) @@ -1815,14 +1886,10 @@ def build_realtime_url( query.setdefault("duplex", "1") if model: query.setdefault("model", model) - if native_duplex is not None: - query["native_duplex"] = "1" if native_duplex else "0" for key, value in (extra_query or {}).items(): query.setdefault(key, value) if autostart is not None: query.setdefault("autostart", "1" if autostart else "0") - if session_id: - query.setdefault("session_id", session_id) return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment)) @@ -1880,7 +1947,7 @@ def summarize_session_request_metrics( return summary -async def acknowledge_collected_playback(client: DuplexClient, collector: EventCollector) -> None: +async def acknowledge_collected_playback(client: DuplexClientBase, collector: EventCollector) -> None: """Ack playback of every collected response's audio (probe shorthand). A response that has produced no audio yet is still acked (0 ms) unless it diff --git a/vllm_omni/clients/inline_duplex.py b/vllm_omni/clients/inline_duplex.py new file mode 100644 index 000000000..79c09e014 --- /dev/null +++ b/vllm_omni/clients/inline_duplex.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""In-process duplex client: the :class:`DuplexClient` usage over a ``DuplexOmni``. + +Example:: + + from vllm_omni.clients.inline_duplex import InlineDuplexClient + from vllm_omni.clients.minicpmo_4_5 import create_duplex_session_config + from vllm_omni.entrypoints.duplex_omni import DuplexOmni + + omni = DuplexOmni(model="openbmb/MiniCPM-o-4_5", trust_remote_code=True) + cfg = create_duplex_session_config(ref_audio=audio_data_url(wav_bytes)) + async with InlineDuplexClient(omni, model="openbmb/MiniCPM-o-4_5", config=cfg) as client: + await client.stream_pcm(pcm16) + await client.commit() + async for response in client.responses(): + async for chunk in response.audio(): + play(chunk) + break + +There is no transport to drop, so this client has no reconnect, resume, +heartbeat or event-acknowledgement machinery. The ``DuplexOmni`` instance is +passed in by the caller; this module never imports server runtime code at +import time. +""" + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING + +from vllm_omni.clients.duplex import DuplexClientBase, DuplexConnectionError, SessionConfig + +if TYPE_CHECKING: + from vllm_omni.engine.duplex.realtime_commands import RealtimeInputDefaults + from vllm_omni.entrypoints.duplex_omni import DuplexOmni, DuplexSessionHandle + +__all__ = ["InlineDuplexClient"] + +#: Acknowledgements are a websocket-journal concern; heartbeats are lease touches and go through. +_TRANSPORT_ONLY_EVENTS = frozenset({"session.event_ack"}) + + +class InlineDuplexClient(DuplexClientBase): + """Async client for one duplex session on an in-process :class:`DuplexOmni`.""" + + def __init__( + self, + omni: DuplexOmni, + *, + model: str, + config: SessionConfig | None = None, + handshake_timeout_s: float = 30.0, + ) -> None: + super().__init__(model=model, config=config, handshake_timeout_s=handshake_timeout_s) + self.omni = omni + self._handle: DuplexSessionHandle | None = None + self._pump_task: asyncio.Task[None] | None = None + self._handshake_pending = False + + # -- transport hooks ------------------------------------------------------ + + async def _open(self) -> None: + # The handle's first event is ``session.created``; the pump delivers + # it through ``_dispatch`` where the base handshake is waiting. + self._handle = await self.omni.open_session(self.config.to_session_payload(model=self.model)) + # The base class follows up with the ``session.update`` handshake; the + # open above already carried that session object, so it is not resent. + self._handshake_pending = True + self._pump_task = asyncio.create_task(self._pump_events(), name="inline-duplex-client-pump") + + def _input_defaults(self) -> RealtimeInputDefaults: + """Audio format / sample-rate / VAD defaults declared by this session's object.""" + from vllm_omni.engine.duplex.realtime_commands import RealtimeInputDefaults + + session_payload = self.session_info or self.config.to_session_payload(model=self.model) + return RealtimeInputDefaults().with_session_payload(session_payload) + + async def _send_command(self, payload: dict[str, object]) -> None: + handle = self._handle + if handle is None or handle.closed: + raise DuplexConnectionError("send failed: session is not open") + event_type = payload.get("type") + if event_type == "session.update" and self._handshake_pending: + self._handshake_pending = False + return + if event_type == "session.close": + await handle.close() + return + if event_type in _TRANSPORT_ONLY_EVENTS: + return + # Lazy import keeps ``vllm_omni.clients`` free of runtime imports at + # module load; the caller already holds a live DuplexOmni. + from vllm_omni.engine.duplex.commands import DuplexCommandError, command_from_realtime + from vllm_omni.engine.duplex.events import error_event + + try: + command = command_from_realtime(payload, defaults=self._input_defaults()) + except DuplexCommandError as exc: + await self._dispatch( + error_event(exc.code, str(exc), event_id=exc.event_id or payload.get("event_id")).to_realtime() + ) + return + await handle.submit(command) + + async def _teardown(self) -> None: + pump = self._pump_task + if pump is not None and not pump.done(): + pump.cancel() + try: + await pump + except (asyncio.CancelledError, Exception): # noqa: BLE001 + pass + handle = self._handle + if handle is not None and not handle.closed: + try: + await handle.close() + except Exception: # noqa: BLE001 + pass + + # -- internals --------------------------------------------------------------- + + async def _pump_events(self) -> None: + handle = self._handle + assert handle is not None + reason = "closed" + try: + async for event in handle.events(): + await self._dispatch(event.to_realtime()) + except asyncio.CancelledError: + raise + except Exception as exc: # noqa: BLE001 + reason = f"session failed: {exc}" + self._finalize(reason, expected=False) + return + self._finalize(reason, expected=True) diff --git a/vllm_omni/clients/minicpmo_4_5.py b/vllm_omni/clients/minicpmo_4_5.py index 4b25346d4..b43346324 100644 --- a/vllm_omni/clients/minicpmo_4_5.py +++ b/vllm_omni/clients/minicpmo_4_5.py @@ -13,7 +13,7 @@ __all__ = ["create_duplex_session_config"] def create_duplex_session_config(*, ref_audio: str | None = None, **overrides: object) -> SessionConfig: """Session preset matching the MiniCPM-o 4.5 native duplex deployment.""" - extra_body: dict[str, object] = {"native_duplex": True, "force_listen_count": 0} + extra_body: dict[str, object] = {"force_listen_count": 0} extra_body.update(overrides.pop("extra_body", {})) config = SessionConfig( ref_audio=ref_audio, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 9c9186c16..c96f514e8 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -309,14 +309,17 @@ class PipelineConfig: diffusers_class_name: str | None = None diffusers_class_aliases: tuple[str, ...] = () endpoint_restrictions: tuple[EndpointRestriction, ...] = () - # Optional model-owned duplex planner loaded by the stable engine runtime. + # Dotted path of the model's ``DuplexModelPlugin``. A pipeline is a duplex + # model iff this is set: ``vllm-omni serve`` then always runs it through + # ``DuplexOmni`` (every served surface runs on a duplex session) and the engine hosts one + # ``DuplexOrchestrator`` with the plugin loaded. + duplex_plugin: str | None = None + # Legacy duplex wiring of the models that are not ported to the plugin + # framework yet (PersonaPlex, Nemotron VoiceChat). Nothing reads them: a + # pipeline that only declares these is served turn-based. Each field goes + # away with the follow-up PR that ports its model to ``duplex_plugin``. duplex_runtime_extension: str | None = None - # Optional model-owned Serving adapter loaded only when the Realtime duplex - # endpoint is enabled. Generic OpenAI modules must not select a model. duplex_serving_adapter: str | None = None - # Explicitly enable the stable duplex control mechanism. This is separate - # from the optional model extension because turn-commit-only deployments - # do not require a model planner. duplex_control_enabled: bool = False # Bundled deploy defaults for this concrete pipeline topology. The file is # loaded from vllm_omni/deploy; None uses DeployConfig defaults. @@ -510,6 +513,12 @@ class DuplexSessionRuntimeConfig: max_pending_input_bytes_per_session: int = 16 * 1024 * 1024 max_pending_turns_per_session: int = 4 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 + # framework carries appends as one-way commands on the runner's ordered + # mailbox, so there is no client-visible retry to deduplicate. The field + # stays so the deploy configs of the models that are not ported yet still + # parse, and goes away with the PR that ports the last of them. completed_append_cache_size: int = 256 server_vad_model_path: str | None = None # Startup warmup: run this many silent 80 ms-style frames through a diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 4e9fd5a77..ca772afbd 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -1,8 +1,8 @@ # MiniCPM-o 4.5 three-stage deploy on one large-memory GPU. pipeline: minicpmo_4_5 async_chunk: true -# Native full-duplex shares this profile with /v1/chat/completions. Every other -# duplex_session knob keeps its runtime default. +# session_mode: duplex makes vllm-omni serve run this model duplex-only +# (/v1/realtime?duplex=1); every other duplex_session knob keeps its runtime default. session_mode: duplex # Keep the active window aligned with the session and stage capacities below. # A smaller window serializes first-packet generation for concurrent chat turns. diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml index 9c3818813..9600e5959 100644 --- a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml @@ -4,8 +4,8 @@ # isolated from codec/vocoder contention while preserving stage-2 batching. pipeline: minicpmo_4_5 async_chunk: true -# Native full-duplex shares this profile with /v1/chat/completions. Every other -# duplex_session knob keeps its runtime default. +# session_mode: duplex makes vllm-omni serve run this model duplex-only +# (/v1/realtime?duplex=1); every other duplex_session knob keeps its runtime default. session_mode: duplex # Keep the active window aligned with the session and stage capacities below. # A smaller window serializes first-packet generation for concurrent chat turns. diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml index 5eddbb6f4..aae531218 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml @@ -4,8 +4,8 @@ # Code2Wav run independently on GPUs 0, 1, and 2. pipeline: minicpmo_4_5 async_chunk: true -# Native full-duplex shares this profile with /v1/chat/completions. Every other -# duplex_session knob keeps its runtime default. +# session_mode: duplex makes vllm-omni serve run this model duplex-only +# (/v1/realtime?duplex=1); every other duplex_session knob keeps its runtime default. session_mode: duplex # Keep the active window aligned with the session and stage capacities below. # A smaller window serializes first-packet generation for concurrent chat turns. diff --git a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml index eacc069d1..00b1090f8 100644 --- a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml @@ -15,8 +15,9 @@ # so 4096 is the empirically-verified safe ceiling. pipeline: minicpmo_4_5 async_chunk: true -# Native full-duplex shares this profile with /v1/chat/completions. The default -# duplex_session caps sessions at 1, matching Stage 0/1 max_num_seqs below. +# session_mode: duplex makes vllm-omni serve run this model duplex-only +# (/v1/realtime?duplex=1). The default duplex_session caps sessions at 1, +# matching Stage 0/1 max_num_seqs below. session_mode: duplex active_stream_window: 1 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 6850e6b29..3469d8fee 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1,498 +1,45 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -""" -Async Omni Engine for vLLM-Omni multi-stage runtime. - -AsyncOmniEngine in the caller's thread is a thin proxy that communicates -with the Orchestrator (running in a background thread) via janus queues. -""" +"""AsyncOmniEngine: the turn-based engine (add_request / streaming updates / CFG companions / interaction).""" from __future__ import annotations -import asyncio -import concurrent.futures import copy -import queue import shutil -import threading import time -import uuid -import weakref from collections.abc import Mapping, Sequence from typing import Any, Literal, cast -import janus -from vllm import envs as vllm_envs from vllm.inputs import PromptType from vllm.logger import init_logger from vllm.v1.engine import EngineCoreRequest -from vllm.v1.engine.input_processor import InputProcessor -from vllm_omni.config.config_factory import StageConfigFactory -from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config -from vllm_omni.config.stage_config import ( - DuplexSessionRuntimeConfig, - PipelineConfig, - load_deploy_config, -) from vllm_omni.data_entry_keys import REQUEST_ARTIFACT_DIRS_KEY, TRANSFORM_OWNED_META_KEYS from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.async_engine_utils import ( - SHUTDOWN_ENQUEUE_TIMEOUT_S, - SHUTDOWN_JOIN_TIMEOUT_S, apply_omni_final_stage_metadata, - enqueue_orchestrator_shutdown, inject_global_id, - is_abort_transport_shutdown, - is_janus_sync_queue_shutdown, - shutdown_runtime_after_orchestrator, upgrade_to_omni_request, - weak_shutdown_async_omni_engine, -) -from vllm_omni.engine.duplex.control_client import DuplexControlClient -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.runtime import ( - load_duplex_runtime_extension, - validate_duplex_runtime_extension, ) from vllm_omni.engine.messages import ( - AbortRequestMessage, - AbortResultMessage, AddCompanionRequestMessage, - CollectiveRPCRequestMessage, - CollectiveRPCResultMessage, - EngineQueueMessage, - ErrorMessage, InteractionMessage, - OutputMessage, StageSubmissionMessage, ) -from vllm_omni.engine.orchestrator import Orchestrator -from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient +from vllm_omni.engine.omni_engine_base import OmniEngineBase, StageRuntimeInfo +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorBase from vllm_omni.engine.serialization import deserialize_additional_information -from vllm_omni.engine.stage_client import StageClient -from vllm_omni.engine.stage_init_utils import build_stage0_input_processor -from vllm_omni.engine.stage_pool import StagePool -from vllm_omni.engine.stage_runtime import ( - StageRuntimeInfo, - create_stage_runtime, -) -from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin -from vllm_omni.entrypoints.utils import parse_stage_overrides from vllm_omni.inputs.data import OmniInteractionPrompt, OmniSamplingParams -from vllm_omni.metrics.prometheus import OmniRequestCounter logger = init_logger(__name__) -_STARTUP_POLL_INTERVAL_S = 1.0 -_REQUEST_QUEUE_MAXSIZE = 256 -_ConfigResolutionResult = OmniConfigResolution | tuple[str | None, list[Any], str | None] +class AsyncOmniEngine(OmniEngineBase): + """Turn-based engine used by ``Omni`` / ``AsyncOmni``.""" -def load_and_resolve_stage_configs( - model: str, - kwargs: dict[str, Any], - *, - trust_remote_code: bool | None, - deploy_config_path: str | None, - stage_overrides: Mapping[str, Mapping[str, Any]] | None, - strategy_config_path: str | None, -) -> OmniConfigResolution: - """Compatibility seam delegating to the single config resolver.""" - return resolve_omni_config( - model, - trust_remote_code=trust_remote_code, - cli_overrides=kwargs, - deploy_config_path=deploy_config_path, - stage_overrides=stage_overrides, - strategy_config_path=strategy_config_path, - ) - - -class AsyncOmniEngine: - """Thin proxy that launches an Orchestrator in a background thread. - - All stage clients, input/output processors, and stage-to-stage transfer - logic live inside the Orchestrator coroutine (running in its own thread - with a dedicated asyncio event loop). This class communicates with it - via janus queues (sync side for callers, async side for orchestrator). - - Args: - model: Model name or path - init_timeout: Total timeout waiting for orchestrator startup (seconds). - stage_init_timeout: Timeout for stage initialization (seconds) - **kwargs: Additional arguments - """ - - # Class-level defaults so tests that bypass __init__ via object.__new__ - # don't AttributeError when stage-init / forward paths touch these attrs. - _log_stats: bool = False - _coordinator_runtime: Any = None - _transfer_emitter: Any = None - _prom_metrics: Any = None - _enable_orch_monitor: bool = False - # Lazily created by get_output_blocking_async(). - _output_drain_executor: concurrent.futures.ThreadPoolExecutor | None = None - - def __init__( - self, - model: str, - stage_init_timeout: int = 300, - init_timeout: int = 600, - single_stage_mode: bool = False, - transfer_emitter: Any = None, - prom_metrics: Any = None, - log_stats: bool = False, - tokenizer: str | None = None, - trust_remote_code: bool | None = None, - **kwargs: Any, - ) -> None: - self.model = model - self.tokenizer = tokenizer - # Cached by get_diffusion_od_config(). - self._diffusion_od_config_view: Any = None - startup_timeout = int(init_timeout) - # Forwarded into Orchestrator so its _forward_to_next_stage path can - # emit per-edge transfer_tx_s / transfer_size_bytes histograms. - # Optional: when None, Orchestrator silently skips TX emit (existing - # RX path still works via OrchestratorAggregator). - self._transfer_emitter = transfer_emitter - self._prom_metrics = prom_metrics - # Drives upstream EngineCore + scheduler stats production. When False - # the engine skips SchedulerStats / IterationStats; the per-(stage, - # replica) vllm:* wrap stays registered but reads zero. Respects the - # --log-stats CLI flag set by the user via OmniBase. - self._log_stats = log_stats - self._enable_orch_monitor = bool(kwargs.pop("enable_orch_monitor", False)) - - logger.info(f"[AsyncOmniEngine] Initializing with model {model}") - - # ------------------------------------------------------------------ # - # Single-stage mode detection # - # ------------------------------------------------------------------ # - # Single-stage mode is enabled when the caller explicitly passes # - # single_stage_mode=True, or when a stage_id is provided in the args. # - _stage_id_kwarg = kwargs.get("stage_id") - if isinstance(_stage_id_kwarg, int) and not single_stage_mode: - single_stage_mode = True - - self.single_stage_mode: bool = single_stage_mode - self._single_stage_id_filter: int | None = ( - int(_stage_id_kwarg) if single_stage_mode and isinstance(_stage_id_kwarg, int) else None - ) - self._omni_master_address: str | None = kwargs.get("omni_master_address") - self._omni_master_port: int | None = kwargs.get("omni_master_port") - - # New omni-coordinator flags. Consumed only in single_stage_mode. - # ``omni_dp_size_local`` is process-local: each invocation (head and - # every headless) launches that many replicas for its own stage. - self._omni_dp_size_local: int = int(kwargs.get("omni_dp_size_local") or 1) - if self._omni_dp_size_local < 1: - raise ValueError(f"--omni-dp-size-local must be >= 1, got {self._omni_dp_size_local}") - self._omni_lb_policy: str = str(kwargs.get("omni_lb_policy") or "random") - self._omni_heartbeat_timeout: float = float(kwargs.get("omni_heartbeat_timeout") or 30.0) - if self._omni_heartbeat_timeout <= 0: - raise ValueError(f"--omni-heartbeat-timeout must be > 0, got {self._omni_heartbeat_timeout}") - # Concurrent same-device stage init (admission + SH/EX phase locks). - # Sourced from the parallel_stage_init orchestrator/CLI arg (config, - # not an env var); default False preserves serial init. - self._parallel_stage_init: bool = bool(kwargs.get("parallel_stage_init") or False) - - if single_stage_mode: - logger.info( - "[AsyncOmniEngine] Single-stage mode enabled (stage_id_filter=%s, master=%s:%s)", - self._single_stage_id_filter, - self._omni_master_address, - self._omni_master_port, - ) - - # Keep the historical tuple return from _resolve_stage_configs while - # retaining the richer resolver result for pipeline-wide settings. - # Overrides of that private seam fall back to the factory below. - deploy_config_path = kwargs.get("deploy_config") - self._config_resolution: OmniConfigResolution | None = None - self.config_path, self.stage_configs = self._resolve_stage_configs( - model, - kwargs, - trust_remote_code=trust_remote_code, - ) - if self._config_resolution is None: - pipeline_config = StageConfigFactory.get_pipeline_config( - model=model, - trust_remote_code=bool(trust_remote_code), - deploy_config_path=deploy_config_path, - ) - self._set_pipeline_runtime_config(pipeline_config, self.config_path) - else: - self._set_pipeline_runtime_config( - self._config_resolution.pipeline_config, - self._config_resolution.config_path, - ) - - self.num_stages = len(self.stage_configs) - stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None - self.async_chunk = bool(getattr(stage0_args, "async_chunk", False)) - self.stage_pools: list[StagePool] = [] - self.stage_clients: list[StageClient] = [] # logical-stage view for external readers - self.input_processor: InputProcessor | None = None - self.prompt_transform_func: Any | None = None - self.prompt_expand_func: Any | None = None - self.supported_tasks: tuple[str, ...] = ("generate",) - self.default_sampling_params_list: list[OmniSamplingParams] = [] - self.stage_metadata: list[StageRuntimeInfo] = [] - # Janus queues are constructed eagerly here (not deferred to the - # orchestrator thread) so the master server's ROUTER thread always - # sees a non-None ``self.request_queue`` when on_register fires. - # ``async_q`` lazily binds to whatever event loop first awaits on - # it (the orchestrator loop), so cross-thread use stays correct. - self.request_queue: janus.Queue[EngineQueueMessage] = janus.Queue(maxsize=_REQUEST_QUEUE_MAXSIZE) - self.output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() - self.rpc_output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() - self._shutdown_called = False - self._weak_finalizer: weakref.finalize | None = None - self._correlated_rpc_client: CorrelatedRpcClient | None = None - self._duplex_control_client: DuplexControlClient | None = None - self._running_counter = OmniRequestCounter() - self._engines_waiting_counter = OmniRequestCounter() - - logger.info(f"[AsyncOmniEngine] Launching Orchestrator thread with {self.num_stages} stages") - - # Launch orchestrator background thread - startup_future: concurrent.futures.Future = concurrent.futures.Future() - - self.orchestrator_thread = threading.Thread( - target=self._bootstrap_orchestrator, - args=( - stage_init_timeout, - startup_future, - ), - daemon=True, - name="orchestrator", - ) - self.orchestrator_thread.start() - self._wait_for_orchestrator_init(startup_future, startup_timeout) - self._correlated_rpc_client = CorrelatedRpcClient( - self.request_queue.sync_q, - self.rpc_output_queue.sync_q, - ) - - # Stage runtime fields are assigned directly on self by the bootstrap thread. - self._weak_finalizer = weakref.finalize( - self, - weak_shutdown_async_omni_engine, - self.orchestrator_thread, - self.request_queue, - self.output_queue, - self.rpc_output_queue, - self._correlated_rpc_client, - ) - - logger.info(f"[AsyncOmniEngine] Orchestrator ready with {self.num_stages} stages") - - def get_diffusion_od_config(self) -> Any: - """Expose the diffusion ``model_class_name`` to client-side model-extras. - - The worker holds the full config; here we just resolve the pipeline class - name from the model config (cached). ``model_class_name`` may be ``None``. - """ - if self._diffusion_od_config_view is None: - from types import SimpleNamespace - - from vllm_omni.diffusion.data import resolve_model_class_name - from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata - - model_class_name = resolve_model_class_name(self.model) - metadata = get_diffusion_model_metadata(model_class_name) - self._diffusion_od_config_view = SimpleNamespace( - model_class_name=model_class_name, - supports_multimodal_inputs=metadata.supports_multimodal_inputs, - max_multimodal_image_inputs=metadata.max_multimodal_image_inputs, - supports_mixed_reference_inputs=metadata.supports_mixed_reference_inputs, - ) - return self._diffusion_od_config_view - - def _initialize_stages(self, stage_init_timeout: int) -> None: - """Initialize stage clients/processors via StageRuntime and assign to self.""" - self._runtime = create_stage_runtime( - stage_configs=self.stage_configs, - model=self.model, - config_path=self.config_path, - single_stage_mode=self.single_stage_mode, - stage_init_timeout=stage_init_timeout, - async_chunk=self.async_chunk, - tokenizer=self.tokenizer, - parallel_stage_init=self._parallel_stage_init, - single_stage_id_filter=self._single_stage_id_filter, - omni_master_address=self._omni_master_address, - omni_master_port=self._omni_master_port, - omni_dp_size_local=self._omni_dp_size_local, - omni_heartbeat_timeout=self._omni_heartbeat_timeout, - omni_lb_policy=self._omni_lb_policy, - request_queue=self.request_queue, - log_stats=self._log_stats, - ) - self._runtime.initialize() - - self.num_stages = len(self.stage_configs) - self.stage_pools = self._runtime.stage_pools - self.stage_clients = [ - cast(StageClient, pool.stage_client) for pool in self.stage_pools if pool.stage_client is not None - ] - self.stage_vllm_configs = [pool.stage_vllm_config for pool in self.stage_pools] - self.output_processors = [pool.output_processor for pool in self.stage_pools] - self.input_processor = ( - build_stage0_input_processor(self.stage_vllm_configs[0]) - if self.stage_vllm_configs and self.stage_vllm_configs[0] is not None - else None - ) - self.prompt_transform_func = ( - getattr(self.stage_clients[0], "prompt_transform_func", None) if self.stage_clients else None - ) - self.prompt_expand_func = next( - ( - getattr(client, "prompt_expand_func", None) - for client in self.stage_clients - if getattr(client, "prompt_expand_func", None) is not None - ), - None, - ) - self.default_sampling_params_list = [client.default_sampling_params for client in self.stage_clients] - self.stage_metadata = [ - StageRuntimeInfo( - final_output=client.final_output, - final_output_type=client.final_output_type, - stage_type=client.stage_type, - model_stage=getattr(client, "model_stage", None), - ) - for client in self.stage_clients - ] - supported_tasks: set[str] = set() - if any(getattr(client, "is_comprehension", False) for client in self.stage_clients): - supported_tasks.add("generate") - if any(meta.final_output_type == "audio" for meta in self.stage_metadata): - supported_tasks.add("speech") - self.supported_tasks = tuple(supported_tasks) if supported_tasks else ("generate",) - - def _bootstrap_orchestrator( - self, - stage_init_timeout: int, - startup_future: concurrent.futures.Future, - ) -> None: - """Create loop, initialize stages, then run Orchestrator.""" - - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - - async def _run_orchestrator() -> None: - self._initialize_stages(stage_init_timeout) - - pd_config = self._detect_pd_config() - - membership_controller = self._runtime.create_membership_controller() - duplex_runtime_extension = None - if self._duplex_control_enabled: - duplex_runtime_extension = load_duplex_runtime_extension( - getattr(self, "_duplex_runtime_extension_path", None) - ) - if duplex_runtime_extension is not None: - stage_clients = [] - for pool in self.stage_pools: - assert pool.stage_client is not None - stage_clients.append(pool.stage_client) - validate_duplex_runtime_extension( - duplex_runtime_extension, - sampling_defaults=tuple(client.default_sampling_params for client in stage_clients), - ) - - orchestrator = Orchestrator( - request_async_queue=self.request_queue.async_q, - output_async_queue=self.output_queue.async_q, - rpc_async_queue=self.rpc_output_queue.async_q, - stage_pools=self.stage_pools, - async_chunk=self.async_chunk, - pd_config=pd_config, - membership_controller=membership_controller, - running_counter=self._running_counter, - engines_waiting_counter=self._engines_waiting_counter, - transfer_emitter=self._transfer_emitter, - prom_metrics=self._prom_metrics, - log_stats=self._log_stats, - enable_orch_monitor=self._enable_orch_monitor, - duplex_runtime_extension=duplex_runtime_extension, - enable_duplex_control=self._duplex_control_enabled, - duplex_session_config=self.duplex_session_config, - ) - if not startup_future.done(): - startup_future.set_result(asyncio.get_running_loop()) - await orchestrator.run() - - try: - loop.run_until_complete(_run_orchestrator()) - except Exception as e: - if not startup_future.done(): - wrapped = RuntimeError(f"Orchestrator initialization failed: {e}") - wrapped.__cause__ = e - startup_future.set_exception(wrapped) - logger.exception("[AsyncOmniEngine] Orchestrator thread crashed") - error_text = str(e) or "Orchestrator thread crashed" - try: - error_msg = ErrorMessage(error=error_text, fatal=True) - if self.output_queue is not None: - self.output_queue.sync_q.put_nowait(error_msg) - if self.rpc_output_queue is not None: - self.rpc_output_queue.sync_q.put_nowait(error_msg) - except Exception: - pass - raise - finally: - try: - pending = [task for task in asyncio.all_tasks(loop) if not task.done()] - for task in pending: - task.cancel() - if pending: - loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) - loop.run_until_complete(loop.shutdown_asyncgens()) - if hasattr(loop, "shutdown_default_executor"): - loop.run_until_complete(loop.shutdown_default_executor()) - except Exception: - logger.exception("[AsyncOmniEngine] Failed during orchestrator loop cleanup") - finally: - asyncio.set_event_loop(None) - loop.close() - - def _wait_for_orchestrator_init(self, startup_future: concurrent.futures.Future, startup_timeout: int) -> None: - """ - Wait for orchestrator startup future to return ready. Raises exception on any failures to the init process. - """ - deadline = time.monotonic() + startup_timeout - while True: - remaining = deadline - time.monotonic() - if remaining <= 0: - logger.warning( - "[AsyncOmniEngine] Orchestrator startup timed out after %ss. " - "Multi-stage deployments that initialize stages sequentially on one device " - "or load checkpoints from slow storage may need larger --init-timeout and " - "--stage-init-timeout values.", - startup_timeout, - ) - self._try_shutdown("[AsyncOmniEngine] Failed to cleanup after orchestrator startup timeout") - raise TimeoutError(f"Orchestrator did not become ready within {startup_timeout}s") - try: - startup_future.result( - timeout=min(remaining, _STARTUP_POLL_INTERVAL_S), - ) - break - except concurrent.futures.TimeoutError: - if not self.orchestrator_thread.is_alive(): - self._try_shutdown("[AsyncOmniEngine] Failed to cleanup after orchestrator startup failure") - if startup_future.done(): - startup_future.result() # re-raises the real exception - raise RuntimeError("Orchestrator thread died during startup") - except Exception: - self._try_shutdown("[AsyncOmniEngine] Failed to cleanup after orchestrator startup failure") - raise + def _create_orchestrator(self, **orchestrator_kwargs: Any) -> OrchestratorBase: + return Orchestrator(**orchestrator_kwargs) # ---- request helpers ---- @@ -918,158 +465,6 @@ class AsyncOmniEngine: len(companions), ) - def _detect_pd_config(self) -> dict[str, Any] | None: - """Detect PD (Prefill-Decode) disaggregation config from stage_configs. - Returns a dict with 'pd_pair' and 'bootstrap_addr', or None. - """ - pd_pair = PDDisaggregationMixin.detect_pd_separation_from_stage_configs(self.stage_configs) - if pd_pair is None: - return None - prefill_idx, decode_idx = pd_pair - - # Extract bootstrap address from prefill stage engine_args - bootstrap_addr: str | None = None - try: - prefill_cfg = self.stage_configs[prefill_idx] - ea = getattr(prefill_cfg, "engine_args", None) - kv_cfg = getattr(ea, "kv_transfer_config", None) if ea is not None else None - if kv_cfg is not None: - port = vllm_envs.VLLM_MOONCAKE_BOOTSTRAP_PORT - kv_ip = getattr(kv_cfg, "kv_ip", None) or "127.0.0.1" - bootstrap_addr = f"http://{kv_ip}:{port}" - except Exception as exc: - logger.warning("[AsyncOmniEngine] Could not extract PD bootstrap address: %s", exc) - - logger.info( - "[AsyncOmniEngine] PD disaggregation detected: prefill=stage-%d, decode=stage-%d, bootstrap=%s", - prefill_idx, - decode_idx, - bootstrap_addr, - ) - prefill_engine_id: str | None = None - try: - prefill_client = self.stage_clients[prefill_idx] - kv_cfg = getattr(getattr(prefill_client, "vllm_config", None), "kv_transfer_config", None) - prefill_engine_id = getattr(kv_cfg, "engine_id", None) - except Exception as exc: - logger.warning("[AsyncOmniEngine] Could not extract prefill engine_id: %s", exc) - - return { - "pd_pair": (prefill_idx, decode_idx), - "bootstrap_addr": bootstrap_addr, - "prefill_engine_id": prefill_engine_id, - } - - def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None: - """Apply a strategy-derived ``omni_lb_policy`` to the engine. - - Precedence: an explicit ``--omni-lb-policy`` always wins. ``"random"`` is - the engine default and is treated as "unset" (indistinguishable from no - flag), so a strategy value overrides it. If the user explicitly passed a - non-default policy that conflicts with the strategy-derived one, raise so - the mismatch is not silently ignored. - """ - if not derived: - return - explicit = kwargs.get("omni_lb_policy") - user_set = explicit is not None and str(explicit) != "random" - if user_set: - if str(explicit) != str(derived): - raise ValueError( - f"Conflicting load-balancer policy: --omni-lb-policy={explicit!r} was given " - f"but the composable-parallel strategy derived omni_lb_policy={derived!r}. " - "Drop --omni-lb-policy to use the strategy value, or make them match." - ) - return - if self._omni_lb_policy != str(derived): - logger.info( - "[composable_parallel] applying strategy-derived omni_lb_policy=%r (was %r).", - derived, - self._omni_lb_policy, - ) - self._omni_lb_policy = str(derived) - - @staticmethod - def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list[dict[str, Any]]: - """Compatibility seam for the factory-owned diffusion fallback.""" - return StageConfigFactory.create_default_diffusion(kwargs) - - def _set_pipeline_runtime_config( - self, - pipeline_config: PipelineConfig | None, - config_path: str | None, - ) -> None: - """Initialize engine-wide settings resolved from pipeline metadata.""" - self.endpoint_restrictions = pipeline_config.endpoint_restrictions if pipeline_config is not None else () - self._duplex_runtime_extension_path = ( - pipeline_config.duplex_runtime_extension if pipeline_config is not None else None - ) - self.duplex_serving_adapter_path = ( - pipeline_config.duplex_serving_adapter if pipeline_config is not None else None - ) - self._duplex_control_enabled = bool(pipeline_config and pipeline_config.duplex_control_enabled) - self.duplex_session_config = DuplexSessionRuntimeConfig() - if config_path is not None: - self.duplex_session_config = load_deploy_config(config_path).duplex_session - - def _resolve_stage_configs( - self, - model: str, - kwargs: dict[str, Any], - *, - trust_remote_code: bool | None, - ) -> tuple[str, list[Any]]: - """Resolve stage configs and inject defaults shared by orchestrator/headless.""" - - for legacy_arg in ("stage_configs_path", "stage_configs"): - if legacy_arg in kwargs: - raise ValueError(f"`{legacy_arg}` is no longer supported; use `deploy_config` instead.") - - # log_stats is captured by __init__; its CLI-only negative alias must - # not cross into per-stage structured config ownership validation. - kwargs.pop("disable_log_stats", None) - deploy_config_path = kwargs.pop("deploy_config", None) - strategy_config_path = kwargs.pop("strategy_config", None) - # CLI callers arrive pre-parsed; offline Python callers may use the - # JSON-string form documented in recipes. - stage_overrides = parse_stage_overrides(kwargs.pop("stage_overrides", None)) - - # ``diffusion_streaming_output`` is the public AsyncOmni/serve kwarg; - # stage configs know the field as ``streaming_output``. Mirror only a - # truthy value so the CLI's False default does not override deploy YAML. - if kwargs.get("diffusion_streaming_output") and kwargs.get("streaming_output") is None: - kwargs["streaming_output"] = True - - resolution = cast( - _ConfigResolutionResult, - load_and_resolve_stage_configs( - model, - kwargs, - trust_remote_code=trust_remote_code, - deploy_config_path=deploy_config_path, - stage_overrides=stage_overrides, - strategy_config_path=strategy_config_path, - ), - ) - if isinstance(resolution, OmniConfigResolution): - self._config_resolution = resolution - config_path = resolution.config_path - stage_configs = list(resolution.stage_configs) - strategy_lb_policy = resolution.omni_lb_policy - else: - # Compatibility for overrides of the historical tuple-returning - # seam. Production always receives OmniConfigResolution above. - config_path, stage_configs, strategy_lb_policy = resolution - - # A strategy.yaml may derive a pipeline-wide load-balancer policy. It is - # an orchestrator-level knob (read once at construction), so apply it here - # rather than as a per-stage config field. - self._apply_strategy_lb_policy(strategy_lb_policy, kwargs) - - return cast(str, config_path), stage_configs - - # ==================== Public API ==================== - def add_request( self, request_id: str, @@ -1241,390 +636,6 @@ class AsyncOmniEngine: resumable=resumable, ) - def open_duplex_session( - self, - session_id: str, - *, - session_mode: str = "duplex", - capabilities: dict[str, object] | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - fence: DuplexFence, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Open an engine-level duplex session.""" - return self._get_duplex_control_client().open( - session_id, - session_mode=session_mode, - capabilities=capabilities, - session_config=session_config, - runtime_config=runtime_config, - fence=fence, - timeout=timeout, - ) - - async def open_duplex_session_async( - self, - session_id: str, - *, - session_mode: str = "duplex", - capabilities: dict[str, object] | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - fence: DuplexFence, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Async wrapper for opening an engine-level duplex session.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - None, - lambda: self.open_duplex_session( - session_id, - session_mode=session_mode, - capabilities=capabilities, - session_config=session_config, - runtime_config=runtime_config, - fence=fence, - timeout=timeout, - ), - ) - - def append_duplex_input( - self, - session_id: str, - *, - mode: str, - payload: object, - operation_id: str | None = None, - final: bool = False, - expected_epoch: int | None = None, - fence: DuplexFence, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Append input to an engine-level duplex session.""" - return self._get_duplex_control_client().append( - session_id, - mode=mode, - payload=payload, - operation_id=operation_id, - final=final, - expected_epoch=expected_epoch, - fence=fence, - timeout=timeout, - ) - - async def append_duplex_input_async( - self, - session_id: str, - *, - mode: str, - payload: object, - operation_id: str | None = None, - final: bool = False, - expected_epoch: int | None = None, - fence: DuplexFence, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Async wrapper for appending duplex input.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - None, - lambda: self.append_duplex_input( - session_id, - mode=mode, - payload=payload, - operation_id=operation_id, - final=final, - expected_epoch=expected_epoch, - fence=fence, - timeout=timeout, - ), - ) - - def signal_duplex_turn( - self, - session_id: str, - *, - event: str, - fence: DuplexFence, - next_fence: DuplexFence | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Signal an engine-level duplex turn.""" - return self._get_duplex_control_client().signal( - session_id, - event=event, - fence=fence, - next_fence=next_fence, - session_config=session_config, - runtime_config=runtime_config, - timeout=timeout, - ) - - async def signal_duplex_turn_async( - self, - session_id: str, - *, - event: str, - fence: DuplexFence, - next_fence: DuplexFence | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Async wrapper for signaling a duplex turn.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - None, - lambda: self.signal_duplex_turn( - session_id, - event=event, - fence=fence, - next_fence=next_fence, - session_config=session_config, - runtime_config=runtime_config, - timeout=timeout, - ), - ) - - def close_duplex_session( - self, - session_id: str, - *, - reason: str = "client_close", - fence: DuplexFence, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Close an engine-level duplex session.""" - return self._get_duplex_control_client().close( - session_id, - reason=reason, - fence=fence, - timeout=timeout, - ) - - def touch_duplex_session( - self, - session_id: str, - *, - fence: DuplexFence, - activity: DuplexLeaseActivity, - timeout: float | None = 10.0, - ) -> dict[str, object]: - return self._get_duplex_control_client().touch( - session_id, - fence=fence, - activity=activity, - timeout=timeout, - ) - - async def touch_duplex_session_async( - self, - session_id: str, - *, - fence: DuplexFence, - activity: DuplexLeaseActivity, - timeout: float | None = 10.0, - ) -> dict[str, object]: - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - None, - lambda: self.touch_duplex_session( - session_id, - fence=fence, - activity=activity, - timeout=timeout, - ), - ) - - def resume_duplex_session( - self, - session_id: str, - *, - fence: DuplexFence, - expected_lease_generation: int, - timeout: float | None = 10.0, - ) -> dict[str, object]: - return self._get_duplex_control_client().resume( - session_id, - fence=fence, - expected_lease_generation=expected_lease_generation, - timeout=timeout, - ) - - async def resume_duplex_session_async( - self, - session_id: str, - *, - fence: DuplexFence, - expected_lease_generation: int, - timeout: float | None = 10.0, - ) -> dict[str, object]: - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - None, - lambda: self.resume_duplex_session( - session_id, - fence=fence, - expected_lease_generation=expected_lease_generation, - timeout=timeout, - ), - ) - - def _get_duplex_control_client(self) -> DuplexControlClient: - client = getattr(self, "_duplex_control_client", None) - if client is None: - transport = getattr(self, "_correlated_rpc_client", None) - if transport is None: - raise RuntimeError("correlated RPC client is not initialized") - client = DuplexControlClient( - transport, - control_id_factory=lambda: uuid.uuid4().hex, - ) - self._duplex_control_client = client - return client - - async def close_duplex_session_async( - self, - session_id: str, - *, - reason: str = "client_close", - fence: DuplexFence, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Async wrapper for closing an engine-level duplex session.""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - None, - lambda: self.close_duplex_session(session_id, reason=reason, fence=fence, timeout=timeout), - ) - - def try_get_output(self, timeout: float = 0.001) -> EngineQueueMessage | None: - """Read one output message from the Orchestrator output queue.""" - try: - return self.output_queue.sync_q.get(timeout=timeout) - except queue.Empty: - if not self.is_alive(): - raise RuntimeError("Orchestrator died unexpectedly. See logs above.") - return None - - async def try_get_output_async(self) -> EngineQueueMessage | None: - """Async read from the Orchestrator output queue.""" - try: - return self.output_queue.sync_q.get_nowait() - except queue.Empty: - if not self.is_alive(): - raise RuntimeError("Orchestrator died unexpectedly. See logs above.") - return None - - async def get_output_blocking_async(self, timeout: float = 1.0) -> EngineQueueMessage | None: - """Blocking-wait read from the Orchestrator output queue. - - Waits up to ``timeout`` seconds in a dedicated drain thread for the - next message (condition-variable wakeup instead of a poll cadence); - returns ``None`` on timeout so the caller keeps its liveness check, - mirroring ``try_get_output_async``'s contract. Used by the serving - final-output drain when ``VLLM_OMNI_EVENT_DRIVEN_ORCH`` is on. - """ - executor = self._output_drain_executor - if executor is None: - executor = concurrent.futures.ThreadPoolExecutor( - max_workers=1, - thread_name_prefix="omni-output-drain", - ) - self._output_drain_executor = executor - - sync_q = self.output_queue.sync_q - - def _drain_get() -> EngineQueueMessage | None: - # Exceptions are swallowed to a None sentinel: the queue may be - # closed mid-shutdown, and an exception left on an executor future - # after task cancellation would warn as never-retrieved. - try: - return sync_q.get(timeout=timeout) - except queue.Empty: - return None - except Exception: - return None - - loop = asyncio.get_running_loop() - msg = await loop.run_in_executor(executor, _drain_get) - if msg is None and not self.is_alive(): - raise RuntimeError("Orchestrator died unexpectedly. See logs above.") - return msg - - def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo: - """Get cached metadata for a stage.""" - return self.stage_metadata[stage_id] - - def abort(self, request_ids: list[str]) -> None: - """Fire-and-forget abort: enqueue and return without waiting. - - Prefer :meth:`abort_async` when the caller needs acknowledgment that - stage aborts, binding release, and orchestrator request cleanup finished. - """ - if not request_ids or getattr(self, "_shutdown_called", False): - return - if self.request_queue is None: - raise RuntimeError("request_queue is not initialized") - try: - self.request_queue.sync_q.put(AbortRequestMessage(request_ids=request_ids)) - except Exception as exc: - if getattr(self, "_shutdown_called", False) and is_janus_sync_queue_shutdown(exc): - return - raise - - async def abort_async( - self, - request_ids: list[str], - timeout: float | None = None, - ) -> list[OutputMessage]: - """Abort requests and wait for orchestrator acknowledgment. - - Unlike :meth:`abort`, this generates an ``rpc_id``, correlates the - :class:`AbortResultMessage` via :class:`CorrelatedRpcClient`, and - raises if the orchestrator reports failure or times out. - - Returns: - Final-stage AR abort ``OutputMessage`` list carrying partial - tokens generated before abort (empty for diffusion / no OP state). - """ - if not request_ids or getattr(self, "_shutdown_called", False): - return [] - if self.request_queue is None: - raise RuntimeError("request_queue is not initialized") - transport = self._correlated_rpc_client - if transport is None: - raise RuntimeError("correlated RPC client is not initialized") - - rpc_id = uuid.uuid4().hex - msg = AbortRequestMessage(request_ids=request_ids, rpc_id=rpc_id) - - def _wait() -> AbortResultMessage: - result_msg = transport.execute( - ("abort", rpc_id), - msg, - timeout=timeout, - timeout_message=f"abort timed out after {timeout} seconds", - block_on_submit=True, - ) - if not isinstance(result_msg, AbortResultMessage): - raise RuntimeError(f"unexpected abort result type: {type(result_msg).__name__}") - return result_msg - - loop = asyncio.get_running_loop() - try: - result_msg = await loop.run_in_executor(None, _wait) - except Exception as exc: - if getattr(self, "_shutdown_called", False) and is_abort_transport_shutdown(exc): - return [] - raise - if not result_msg.success: - raise RuntimeError(result_msg.error or "abort failed") - return list(result_msg.abort_outputs or []) - def submit_interaction( self, request_id: str, @@ -1649,166 +660,6 @@ class AsyncOmniEngine: """Async interaction API.""" self.submit_interaction(request_id, interaction) - def collective_rpc( - self, - method: str, - timeout: float | None = None, - args: tuple[Any, ...] = (), - kwargs: dict[str, Any] | None = None, - stage_ids: list[int] | None = None, - ) -> list[Any]: - """Send a control RPC to the Orchestrator and wait for aggregated results. - This uses a dedicated RPC output queue so control-plane messages do not - race with the normal request output polling loop. - """ - rpc_id = uuid.uuid4().hex - msg = CollectiveRPCRequestMessage( - rpc_id=rpc_id, - method=method, - timeout=timeout, - args=tuple(args), - kwargs=kwargs or {}, - stage_ids=stage_ids, - ) - - transport = self._correlated_rpc_client - if transport is None: - raise RuntimeError("correlated RPC client is not initialized") - result_msg = transport.execute( - ("collective", rpc_id), - msg, - timeout=timeout, - timeout_message=f"collective_rpc timed out after {timeout} seconds", - block_on_submit=True, - ) - if not isinstance(result_msg, CollectiveRPCResultMessage): - raise RuntimeError(f"unexpected collective RPC result type: {type(result_msg).__name__}") - return list(result_msg.results) - - async def collective_rpc_async( - self, - method: str, - timeout: float | None = None, - args: tuple[Any, ...] = (), - kwargs: dict[str, Any] | None = None, - stage_ids: list[int] | None = None, - ) -> list[Any]: - """Async wrapper around collective_rpc().""" - loop = asyncio.get_running_loop() - return await loop.run_in_executor( - None, - lambda: self.collective_rpc( - method=method, - timeout=timeout, - args=args, - kwargs=kwargs, - stage_ids=stage_ids, - ), - ) - - def is_alive(self) -> bool: - """Whether the orchestrator thread is alive.""" - return bool(self.orchestrator_thread.is_alive()) - - def shutdown(self) -> None: - """Send shutdown message and wait for the Orchestrator thread to exit.""" - if getattr(self, "_shutdown_called", False): - return - self._shutdown_called = True - finalizer = getattr(self, "_weak_finalizer", None) - if finalizer is not None and finalizer.alive: - finalizer.detach() - - logger.info("[AsyncOmniEngine] Shutting down Orchestrator") - request_queue_closed = False - shutdown_enqueued = enqueue_orchestrator_shutdown( - self.request_queue, - timeout=SHUTDOWN_ENQUEUE_TIMEOUT_S, - ) - if self.request_queue is not None and not shutdown_enqueued: - logger.error( - "[AsyncOmniEngine] Failed to enqueue orchestrator shutdown; " - "closing the request queue to wake the request handler" - ) - try: - self.request_queue.close() - request_queue_closed = True - except Exception: - logger.exception("[AsyncOmniEngine] Failed to close the request queue") - - if self._correlated_rpc_client is not None: - try: - self._correlated_rpc_client.close() - except Exception: - logger.exception("[AsyncOmniEngine] Failed to close correlated RPC client") - - orchestrator_stopped = False - try: - if self.is_alive(): - self.orchestrator_thread.join(timeout=SHUTDOWN_JOIN_TIMEOUT_S) - orchestrator_stopped = not self.is_alive() - if not orchestrator_stopped: - logger.error( - "[AsyncOmniEngine] Orchestrator did not stop within %.1f seconds; continuing cleanup", - SHUTDOWN_JOIN_TIMEOUT_S, - ) - except Exception: - logger.exception("[AsyncOmniEngine] Failed to join Orchestrator thread") - - for q in (self.request_queue, self.output_queue, self.rpc_output_queue): - try: - if not (q is self.request_queue and request_queue_closed): - q.close() - except Exception: - pass - - if self._output_drain_executor is not None: - # Any in-flight blocking get bails out within its ≤1 s timeout - # (or immediately via the queue close above), so don't wait. - self._output_drain_executor.shutdown(wait=False) - self._output_drain_executor = None - - if hasattr(self, "_runtime") and self._runtime is not None and orchestrator_stopped: - try: - self._runtime.shutdown() - except Exception: - logger.exception("[AsyncOmniEngine] Failed to shutdown StageRuntime") - elif hasattr(self, "_runtime") and self._runtime is not None: - logger.warning("[AsyncOmniEngine] Deferring StageRuntime shutdown until the Orchestrator exits") - threading.Thread( - target=shutdown_runtime_after_orchestrator, - args=(self.orchestrator_thread, self._runtime), - daemon=True, - name="omni-stage-runtime-shutdown", - ).start() - - # ── Release CuMem allocator memory pool ────────────────────────────── - # When enable_sleep_mode is in use, the CuMem (CUDA Virtual Memory - # Management) allocator holds model weights in a singleton memory pool - # that lives in the parent process. Killing the engine-core subprocess - # does NOT release this pool — the weights stay resident on the GPU - # and can cause CUDA OOM for subsequent engine instances (especially - # large models like BAGEL-7B-MoT whose weights alone consume ~134 GiB). - # - # CuMemAllocator.sleep() is NOT idempotent — calling it on already- - # slept entries causes CUDA_ERROR_INVALID_VALUE at cumem_allocator - # cuMemRelease (double-free of the memory handle). Use release_pools() - # instead, which is the designed cleanup path: it drops MemPool refs - # and lets the destructor/free path handle asleep entries correctly - # (returns a null handle so the C extension skips unmap/release). - try: - from vllm.device_allocator.cumem import CuMemAllocator, cumem_available - - if cumem_available: - allocator = CuMemAllocator.get_instance() - allocator.release_pools() - logger.debug("[AsyncOmniEngine] Released CuMem memory pool during shutdown") - except Exception: - pass - - def _try_shutdown(self, *args, **kwargs) -> None: - try: - self.shutdown() - except Exception: - logger.exception(*args, **kwargs) +# ``StageRuntimeInfo`` moved to ``omni_engine_base``; keep the historical import path. +__all__ = ["AsyncOmniEngine", "StageRuntimeInfo"] diff --git a/vllm_omni/engine/duplex/__init__.py b/vllm_omni/engine/duplex/__init__.py index bd5ae0b3f..533e31e23 100644 --- a/vllm_omni/engine/duplex/__init__.py +++ b/vllm_omni/engine/duplex/__init__.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Engine-side helpers for the full-duplex runtime.""" +"""Engine-resident full-duplex sessions: typed commands/events, the merged session +state, the per-session runner, the session manager and the model plugin contract.""" diff --git a/vllm_omni/entrypoints/duplex/audio.py b/vllm_omni/engine/duplex/audio.py similarity index 96% rename from vllm_omni/entrypoints/duplex/audio.py rename to vllm_omni/engine/duplex/audio.py index 747293a1a..5f163826a 100644 --- a/vllm_omni/entrypoints/duplex/audio.py +++ b/vllm_omni/engine/duplex/audio.py @@ -11,6 +11,8 @@ import wave import numpy as np +from vllm_omni.utils.audio_resample import StreamingAudioResampler + try: from audioop import alaw2lin, lin2alaw, lin2ulaw, ulaw2lin except ImportError: # pragma: no cover - audioop is removed in newer Python. @@ -39,14 +41,10 @@ def validate_input_sample_rate_hz(sample_rate_hz: object) -> int: def resample_pcm16_mono(raw: bytes, *, source_rate_hz: int, target_rate_hz: int) -> bytes: if source_rate_hz <= 0 or target_rate_hz <= 0 or source_rate_hz == target_rate_hz: return raw - # Import lazily so the generic Duplex audio module does not pull in the - # OpenAI API server while the Duplex stack itself is being imported. - from vllm_omni.entrypoints.openai.audio_utils_mixin import StreamingAudioResampler - samples = np.frombuffer(raw, dtype=" StreamingAudioResampler._max_polyphase_factor: + if reduced_factor > StreamingAudioResampler.max_polyphase_factor: # The FIR size scales with the reduced ratio. Keep unusual client # rates bounded while common audio rates use the high-quality path. if samples.size <= 1: diff --git a/vllm_omni/engine/duplex/commands.py b/vllm_omni/engine/duplex/commands.py new file mode 100644 index 000000000..90d48d050 --- /dev/null +++ b/vllm_omni/engine/duplex/commands.py @@ -0,0 +1,324 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Typed commands accepted by a duplex session. + +``DuplexSessionHandle.submit()`` takes one of these; the websocket handler and +``InlineDuplexClient`` build them with :func:`command_from_realtime` from the +OpenAI Realtime client event vocabulary. Every command can also render the +session-internal mailbox payload (``payload()``), which is the dictionary +vocabulary the session runner bodies were written against. +""" + +from __future__ import annotations + +import base64 +from collections.abc import Mapping +from dataclasses import dataclass, field, fields +from typing import TYPE_CHECKING, ClassVar + +if TYPE_CHECKING: + from vllm_omni.engine.duplex.realtime_commands import RealtimeInputDefaults + + +class DuplexCommandError(ValueError): + """A client payload could not be turned into a command.""" + + def __init__(self, message: str, *, code: str = "bad_event", event_id: str | None = None) -> None: + super().__init__(message) + self.code = code + self.event_id = event_id + + +@dataclass(frozen=True, slots=True, kw_only=True) +class DuplexCommand: + #: Mailbox event type this command renders to (see ``payload()``). + type: ClassVar[str] = "" + #: Client correlation id (OpenAI ``event_id``), echoed on error events. + event_id: str | None = None + + def payload(self) -> dict[str, object]: + """Render the session-internal mailbox dictionary.""" + data: dict[str, object] = {"type": self.type} + for f in fields(self): + if f.name == "event_id": + if self.event_id is not None: + data["realtime_event_id"] = self.event_id + continue + value = getattr(self, f.name) + if value is None: + continue + if isinstance(value, tuple): + value = list(value) + elif isinstance(value, Mapping): + value = dict(value) + data[f.name] = value + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AppendAudio(DuplexCommand): + type: ClassVar[str] = "input_audio_buffer.append" + #: Raw audio bytes in ``format`` at ``sample_rate_hz`` (base64 only on the wire). + audio: bytes + format: str = "pcm16" + sample_rate_hz: int | None = None + is_speech: bool | None = None + video_frames: tuple[str, ...] = () + duration_ms: int | None = None + audio_end_ms: int | None = None + #: Model-neutral hints carried through from the wire (rms, vad, transcript hints ...). + hints: Mapping[str, object] = field(default_factory=dict) + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + hints = data.pop("hints", None) + if isinstance(hints, Mapping): + # Hints are raw wire values; the typed fields went through + # ``build_append_audio``'s normalization, so they have to win. An + # unset typed field is simply absent here (``payload`` skips None), + # which is what lets a hint still carry it. Merging the other way + # round let a client's ``"is_speech": 0`` override the computed + # ``bool | None`` and silently miss the runner's silent-commit path. + merged: dict[str, object] = dict(hints) + merged.update(data) + data = merged + data["audio"] = base64.b64encode(self.audio).decode("ascii") + if not data.get("video_frames"): + data.pop("video_frames", None) + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AppendText(DuplexCommand): + type: ClassVar[str] = "input.text.append" + text: str + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Commit(DuplexCommand): + type: ClassVar[str] = "input_audio_buffer.commit" + final: bool = True + #: ``None`` means "no explicit request": the runner decides on commit (auto-response + #: sessions answer on their own); ``True`` / ``False`` force it. + create_response: bool | None = None + is_speech: bool | None = None + #: Realtime conversation item created for this commit (wire correlation only). + realtime_item_id: str | None = None + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + if self.create_response is not None: + data["response_create"] = data.pop("create_response") + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CreateResponse(DuplexCommand): + type: ClassVar[str] = "response.create" + #: Raw Realtime ``response`` object (instructions, voice, modalities, ...). + options: Mapping[str, object] = field(default_factory=dict) + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + options = data.pop("options", None) + if isinstance(options, Mapping): + data["response"] = dict(options) + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ClearInput(DuplexCommand): + type: ClassVar[str] = "input_audio_buffer.clear" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CancelInput(DuplexCommand): + type: ClassVar[str] = "input.cancel" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CancelResponse(DuplexCommand): + type: ClassVar[str] = "response.cancel" + response_id: str | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class BargeIn(DuplexCommand): + type: ClassVar[str] = "barge_in" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ClearOutputAudio(DuplexCommand): + type: ClassVar[str] = "output_audio_buffer.clear" + #: Explicit response to clear; ``None`` targets the active/last response. + response_id: str | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SignalTurn(DuplexCommand): + """Generic ``turn.signal`` (local turn transitions such as ``user_started``).""" + + type: ClassVar[str] = "turn.signal" + event: str + signal_payload: Mapping[str, object] = field(default_factory=dict) + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + signal_payload = data.pop("signal_payload", None) + if isinstance(signal_payload, Mapping) and signal_payload: + data["payload"] = dict(signal_payload) + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class UpdateSession(DuplexCommand): + """``session.update``: a Realtime ``session`` object patch.""" + + type: ClassVar[str] = "turn.signal" + patch: Mapping[str, object] = field(default_factory=dict) + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + data["event"] = "session.update" + data["payload"] = dict(data.pop("patch", {}) or {}) + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AckPlayback(DuplexCommand): + type: ClassVar[str] = "playback.ack" + played_ms: int + committed_ms: int | None = None + response_id: str | None = None + item_id: str | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Heartbeat(DuplexCommand): + type: ClassVar[str] = "session.heartbeat" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CreateItem(DuplexCommand): + """``conversation.item.create`` (history injection or function-call output).""" + + type: ClassVar[str] = "turn.signal" + item: Mapping[str, object] + previous_item_id: str | None = None + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + data["event"] = "conversation.item.create" + payload: dict[str, object] = {"item": dict(data.pop("item"))} + previous = data.pop("previous_item_id", None) + if previous is not None: + payload["previous_item_id"] = previous + data["payload"] = payload + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class DeleteItem(DuplexCommand): + type: ClassVar[str] = "turn.signal" + item_id: str + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + data["event"] = "conversation.item.delete" + data["payload"] = {"item_id": data.pop("item_id")} + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TruncateItem(DuplexCommand): + type: ClassVar[str] = "turn.signal" + item_id: str + audio_end_ms: int + content_index: int = 0 + + def payload(self) -> dict[str, object]: + data = DuplexCommand.payload(self) + data["event"] = "conversation.item.truncate" + data["payload"] = { + "item_id": data.pop("item_id"), + "audio_end_ms": data.pop("audio_end_ms"), + "content_index": data.pop("content_index", 0), + } + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class CloseSession(DuplexCommand): + """Graceful close requested through the command stream (``session.close``).""" + + type: ClassVar[str] = "session.close" + reason: str = "client_close" + + +#: Realtime client event types that map onto commands (wire vocabulary). +REALTIME_COMMAND_TYPES: frozenset[str] = frozenset( + { + "input_audio_buffer.append", + "input_audio_buffer.commit", + "input_audio_buffer.clear", + "output_audio_buffer.clear", + "response.create", + "response.cancel", + "conversation.item.create", + "conversation.item.delete", + "conversation.item.truncate", + "session.update", + "playback.ack", + "session.heartbeat", + "session.close", + "turn.signal", + "input.text.append", + "input.cancel", + "barge_in", + } +) + + +def command_from_realtime( + payload: Mapping[str, object], + *, + defaults: RealtimeInputDefaults | None = None, +) -> DuplexCommand: + """Translate one OpenAI Realtime client event into a command. + + Envelope concerns (``event_id`` acknowledgements, resume bookkeeping) are + handled by the transport; this function only validates and maps the + payload shape. ``defaults`` carries the audio format / sample rate / VAD + defaults the session declared, so every transport decodes an append the + same way. Raises :class:`DuplexCommandError` for malformed or unsupported + payloads. + """ + from vllm_omni.engine.duplex.realtime_commands import translate_realtime_command + + return translate_realtime_command(payload, defaults=defaults) + + +__all__ = [ + "REALTIME_COMMAND_TYPES", + "AckPlayback", + "AppendAudio", + "AppendText", + "BargeIn", + "CancelInput", + "CancelResponse", + "ClearInput", + "ClearOutputAudio", + "CloseSession", + "Commit", + "CreateItem", + "CreateResponse", + "DeleteItem", + "DuplexCommand", + "DuplexCommandError", + "Heartbeat", + "SignalTurn", + "TruncateItem", + "UpdateSession", + "command_from_realtime", +] diff --git a/vllm_omni/engine/duplex/config.py b/vllm_omni/engine/duplex/config.py new file mode 100644 index 000000000..aab0fb1c5 --- /dev/null +++ b/vllm_omni/engine/duplex/config.py @@ -0,0 +1,815 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Session configuration, capabilities and value types of the duplex session model. + +These classes are the typed half of the public contract: ``DuplexOmni`` +accepts a ``DuplexSessionConfig`` (or the equivalent OpenAI Realtime +``session`` object via :meth:`DuplexSessionConfig.from_realtime`) and reports +``DuplexCapabilities``. They used to live in ``entrypoints/duplex/protocol.py``. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, cast + + +class DuplexConfigError(ValueError): + """A Realtime session/response object was rejected while mapping it onto the duplex config.""" + + def __init__(self, message: str, *, code: str = "bad_event", param: str | None = None) -> None: + super().__init__(message) + self.code = code + self.param = param + + +class DuplexOverlapPolicy(str, Enum): + LISTEN_ONLY = "listen_only" + BARGE_IN_ON_SPEECH = "barge_in_on_speech" + + +class DuplexPlaybackCommitPolicy(str, Enum): + COMMIT_ALL_ON_DONE = "commit_all_on_done" + ACK_ONLY = "ack_only" + + +class DuplexSessionState(str, Enum): + OPEN = "open" + CLOSING = "closing" + CLOSED = "closed" + + +class DuplexTurnState(str, Enum): + IDLE = "idle" + USER_SPEAKING = "user_speaking" + USER_COMMITTED = "user_committed" + ASSISTANT_GENERATING = "assistant_generating" + ASSISTANT_PLAYING = "assistant_playing" + BARGE_IN = "barge_in" + + +class DuplexTurnEventType(str, Enum): + USER_STARTED = "user_started" + USER_COMMITTED = "user_committed" + ASSISTANT_STARTED = "assistant_started" + ASSISTANT_DONE = "assistant_done" + BARGE_IN = "barge_in" + PLAYBACK_ACK = "playback_ack" + TIMEOUT = "timeout" + CLOSE = "close" + + +@dataclass +class DuplexCapabilities: + """Runtime/model capabilities exposed by the duplex serving protocol. + + These are intentionally explicit so the serving layer does not assume all + duplex models support the same input append, rollback, or turn policy. + ``supports_core_kv_lease`` is reserved for scheduler-owned KV lifecycle; + model-owned decoder/TTS state must use ``supports_model_internal_state``. + ``supports_core_resumable_request`` means the scheduler can resume the + same request id across streaming updates, but it is not a KV lease by + itself. Realtime support means this endpoint can speak the native Realtime + event schema for the supported audio duplex paths while keeping model- or + scheduler-specific limits explicit in the capability payload. + """ + + supports_session_adapter: bool = True + supports_model_native_turn_policy: bool = False + supports_external_turn_signal: bool = True + supports_client_commit: bool = True + supports_barge_in: bool = True + supports_playback_ack: bool = True + supports_input_append: bool = False + supports_replace_latest_chunk: bool = True + supports_reencode_context: bool = True + supports_rollback_to_checkpoint: bool = False + supports_turn_commit_only: bool = True + supports_kv_lease: bool = False + supports_core_kv_lease: bool = False + supports_model_internal_state: bool = False + supports_stage_resumption: bool = False + supports_scheduler_native_append: bool = False + supports_core_resumable_request: bool = False + supports_stage_connector_handoff: bool = False + supports_independent_io_streams: bool = False + supports_realtime_endpoint: bool = False + supports_multi_session: bool = False + supports_multi_session_same_replica: bool = False + supports_session_lease: bool = False + supports_session_resume: bool = False + session_admission_mode: str = "serving_managed" + supports_audio_truncate: bool = False + #: Whether this model can serve ``/v1/chat/completions`` on a duplex + #: session. That needs one thing the protocol cannot assume: the model + #: honours ``DuplexSessionConfig.initial_user_text``, so a turn can be + #: seeded from text. A model-native model decides to speak from the audio + #: it hears, and silence is its signal *not* to take a turn, so without + #: seeding a text prompt has no representation and the turn never + #: completes. ``text_turn_priming_units`` is how many silence units such a + #: seeded turn must be given to generate on -- it still speaks per unit. + supports_chat_completions: bool = False + text_turn_priming_units: int = 0 + requires_model_runner_kv: bool = False + requires_native_stage_role: bool = False + adapter_patterns: list[str] = field(default_factory=lambda: ["chunk_group_append"]) + signal_sources: list[str] = field(default_factory=lambda: ["client_event", "server_policy", "model_native"]) + stage_handoff_transport: str | None = None + chunk_period_ms: int | None = 1000 + target_barge_in_latency_ms: int | None = 1000 + + def as_dict(self) -> dict[str, object]: + return { + "supports_session_adapter": self.supports_session_adapter, + "supports_model_native_turn_policy": self.supports_model_native_turn_policy, + "supports_external_turn_signal": self.supports_external_turn_signal, + "supports_client_commit": self.supports_client_commit, + "supports_barge_in": self.supports_barge_in, + "supports_playback_ack": self.supports_playback_ack, + "supports_input_append": self.supports_input_append, + "supports_replace_latest_chunk": self.supports_replace_latest_chunk, + "supports_reencode_context": self.supports_reencode_context, + "supports_rollback_to_checkpoint": self.supports_rollback_to_checkpoint, + "supports_turn_commit_only": self.supports_turn_commit_only, + "supports_kv_lease": self.supports_kv_lease, + "supports_core_kv_lease": self.supports_core_kv_lease, + "supports_model_internal_state": self.supports_model_internal_state, + "supports_stage_resumption": self.supports_stage_resumption, + "supports_scheduler_native_append": self.supports_scheduler_native_append, + "supports_core_resumable_request": self.supports_core_resumable_request, + "supports_stage_connector_handoff": self.supports_stage_connector_handoff, + "supports_independent_io_streams": self.supports_independent_io_streams, + "supports_realtime_endpoint": self.supports_realtime_endpoint, + "supports_multi_session": self.supports_multi_session, + "supports_multi_session_same_replica": self.supports_multi_session_same_replica, + "supports_session_lease": self.supports_session_lease, + "supports_session_resume": self.supports_session_resume, + "session_admission_mode": self.session_admission_mode, + "supports_audio_truncate": self.supports_audio_truncate, + "requires_model_runner_kv": self.requires_model_runner_kv, + "requires_native_stage_role": self.requires_native_stage_role, + # Every duplex model is model-native and appends audio chunks (§7 D4); + # kept on the wire as constants for clients that still read them. + "implementation_level": "model_native_duplex", + "adapter_patterns": self.adapter_patterns, + "input_modes": ["append_audio_chunk"], + "signal_sources": self.signal_sources, + "stage_handoff_transport": self.stage_handoff_transport, + "chunk_period_ms": self.chunk_period_ms, + "target_barge_in_latency_ms": self.target_barge_in_latency_ms, + } + + +@dataclass +class DuplexPlaybackCursor: + generated_ms: int = 0 + sent_ms: int = 0 + played_ms: int = 0 + committed_ms: int = 0 + + def acknowledge(self, played_ms: int, committed_ms: int | None = None) -> None: + self.played_ms = max(self.played_ms, max(0, int(played_ms))) + if committed_ms is None: + committed_ms = self.played_ms + self.committed_ms = max(self.committed_ms, max(0, int(committed_ms))) + + def truncate_committed(self, committed_ms: int) -> None: + self.committed_ms = max(0, min(max(self.sent_ms, self.generated_ms), int(committed_ms))) + + def as_dict(self) -> dict[str, int]: + return { + "generated_ms": self.generated_ms, + "sent_ms": self.sent_ms, + "played_ms": self.played_ms, + "committed_ms": self.committed_ms, + } + + def snapshot(self) -> DuplexPlaybackView: + return DuplexPlaybackView(**self.as_dict()) + + +@dataclass(frozen=True, slots=True) +class DuplexPlaybackView: + generated_ms: int = 0 + sent_ms: int = 0 + played_ms: int = 0 + committed_ms: int = 0 + + def as_dict(self) -> dict[str, int]: + return { + "generated_ms": self.generated_ms, + "sent_ms": self.sent_ms, + "played_ms": self.played_ms, + "committed_ms": self.committed_ms, + } + + +@dataclass +class DuplexAudioChunk: + data: str + format: str = "wav" + sample_rate_hz: int | None = None + + +@dataclass +class DuplexSessionConfig: + model: str | None = None + modalities: list[str] = field(default_factory=lambda: ["text", "audio"]) + instructions: str | None = None + #: Text the session treats as the user's opening turn. A caller that cannot + #: reach the model by speech (an HTTP request, a batch job) puts its prompt + #: here; ``supports_text_only_turn`` says whether the model can answer it. + initial_user_text: str | None = None + voice: str | None = None + ref_audio: str | None = None + response_format: str = "wav" + temperature: float | None = None + max_tokens: int | None = None + speed: float | None = None + use_tts_template: bool = True + idle_timeout_s: float = 300.0 + overlap_policy: str = DuplexOverlapPolicy.LISTEN_ONLY.value + overlap_short_ack_ms: int = 700 + overlap_barge_in_ms: int = 1200 + overlap_silence_rms: float = 0.003 + playback_commit_policy: str = DuplexPlaybackCommitPolicy.COMMIT_ALL_ON_DONE.value + extra_body: dict[str, object] = field(default_factory=dict) + + def as_dict(self) -> dict[str, object]: + return { + "model": self.model, + "modalities": list(self.modalities), + "instructions": self.instructions, + "initial_user_text": self.initial_user_text, + "voice": self.voice, + "ref_audio": self.ref_audio, + "response_format": self.response_format, + "temperature": self.temperature, + "max_tokens": self.max_tokens, + "speed": self.speed, + "use_tts_template": self.use_tts_template, + "idle_timeout_s": self.idle_timeout_s, + "overlap_policy": self.overlap_policy, + "overlap_short_ack_ms": self.overlap_short_ack_ms, + "overlap_barge_in_ms": self.overlap_barge_in_ms, + "overlap_silence_rms": self.overlap_silence_rms, + "playback_commit_policy": self.playback_commit_policy, + "extra_body": dict(self.extra_body), + } + + def normalized(self) -> DuplexSessionConfig: + """Apply the clamps / defaults ``from_event`` applies to a wire object (typed configs skip them).""" + return DuplexSessionConfig.from_event({"session": self.as_dict()}) + + @classmethod + def from_event(cls, event: dict[str, object]) -> DuplexSessionConfig: + payload = event.get("session") + if isinstance(payload, dict): + source = payload + else: + source = event + + config = cls() + if isinstance(source.get("model"), str): + config.model = source["model"] + if isinstance(source.get("instructions"), str): + config.instructions = source["instructions"] + if isinstance(source.get("initial_user_text"), str): + config.initial_user_text = source["initial_user_text"] + if isinstance(source.get("voice"), str): + config.voice = source["voice"] + if isinstance(source.get("ref_audio"), str): + config.ref_audio = source["ref_audio"] + if isinstance(source.get("response_format"), str): + config.response_format = source["response_format"] + if isinstance(source.get("use_tts_template"), bool): + config.use_tts_template = bool(source["use_tts_template"]) + if isinstance(source.get("temperature"), int | float): + config.temperature = float(source["temperature"]) + if isinstance(source.get("max_tokens"), int): + config.max_tokens = int(source["max_tokens"]) + if isinstance(source.get("speed"), int | float): + config.speed = float(source["speed"]) + if isinstance(source.get("idle_timeout_s"), int | float): + config.idle_timeout_s = float(source["idle_timeout_s"]) + if isinstance(source.get("overlap_policy"), str): + config.overlap_policy = cls._normalize_overlap_policy(source["overlap_policy"]) + if isinstance(source.get("overlap_short_ack_ms"), int | float): + config.overlap_short_ack_ms = max(0, int(source["overlap_short_ack_ms"])) + if isinstance(source.get("overlap_barge_in_ms"), int | float): + config.overlap_barge_in_ms = max(0, int(source["overlap_barge_in_ms"])) + if isinstance(source.get("overlap_silence_rms"), int | float): + config.overlap_silence_rms = max(0.0, float(source["overlap_silence_rms"])) + if isinstance(source.get("playback_commit_policy"), str): + config.playback_commit_policy = cls._normalize_playback_commit_policy(source["playback_commit_policy"]) + if isinstance(source.get("modalities"), list) and all(isinstance(x, str) for x in source["modalities"]): + config.modalities = list(source["modalities"]) + if isinstance(source.get("extra_body"), dict): + config.extra_body = dict(source["extra_body"]) + extra = config.extra_body + if isinstance(extra.get("overlap_policy"), str): + config.overlap_policy = cls._normalize_overlap_policy(extra["overlap_policy"]) + if isinstance(extra.get("overlap_short_ack_ms"), int | float): + config.overlap_short_ack_ms = max(0, int(extra["overlap_short_ack_ms"])) + if isinstance(extra.get("overlap_barge_in_ms"), int | float): + config.overlap_barge_in_ms = max(0, int(extra["overlap_barge_in_ms"])) + if isinstance(extra.get("overlap_silence_rms"), int | float): + config.overlap_silence_rms = max(0.0, float(extra["overlap_silence_rms"])) + if isinstance(extra.get("playback_commit_policy"), str): + config.playback_commit_policy = cls._normalize_playback_commit_policy(extra["playback_commit_policy"]) + return config + + @classmethod + def from_realtime( + cls, + session_payload: Mapping[str, object], + *, + model: str | None = None, + ) -> DuplexSessionConfig: + """Build a config from an OpenAI Realtime ``session`` object (session.update / open_session). + + Applies the wire defaults, validates audio formats and turn detection + (raising :class:`DuplexConfigError` with the same codes the Realtime + translator used), fills ``turn_detection`` defaults / derives + ``overlap_policy``, and stores the Realtime-only fields under the + ``realtime_*`` keys of ``extra_body`` exactly like the old + ``_session_create_from_realtime``. + """ + from vllm_omni.engine.duplex.realtime_commands import ( + RealtimeInputDefaults, + duplex_response_format, + input_audio_transcription_config, + json_safe_realtime_payload, + realtime_max_output_tokens, + realtime_overlap_fields, + validate_realtime_session_audio_formats, + ) + from vllm_omni.engine.duplex.turn_detection import ( + normalize_turn_detection_session_payload, + validate_realtime_turn_detection, + ) + + payload: dict[str, object] = dict(session_payload) + format_error = validate_realtime_session_audio_formats(payload) + if format_error is not None: + raise DuplexConfigError(format_error, code="unsupported_audio_format") + turn_detection_error = validate_realtime_turn_detection(payload) + if turn_detection_error is not None: + raise DuplexConfigError(turn_detection_error, code="unsupported_turn_detection", param="turn_detection") + normalize_turn_detection_session_payload(payload) + defaults = RealtimeInputDefaults().with_session_payload(payload) + payload.update(realtime_overlap_fields(payload)) + + audio_config = payload.get("audio") + audio_input = audio_config.get("input") if isinstance(audio_config, dict) else None + audio_output = audio_config.get("output") if isinstance(audio_config, dict) else None + extra_body_payload = payload.get("extra_body") + extra_body: dict[str, object] = dict(extra_body_payload) if isinstance(extra_body_payload, dict) else {} + extra_body["realtime_session_payload"] = json_safe_realtime_payload(payload) + if isinstance(payload.get("tools"), list): + extra_body["realtime_tools"] = payload["tools"] + if isinstance(payload.get("tool_choice"), str | dict): + extra_body["realtime_tool_choice"] = payload["tool_choice"] + metadata = payload.get("metadata") + if isinstance(metadata, dict): + extra_body["realtime_metadata"] = dict(metadata) + include = payload.get("include") + if isinstance(include, list): + extra_body["realtime_include"] = list(include) + prompt = payload.get("prompt") + if isinstance(prompt, dict): + extra_body["realtime_prompt"] = dict(prompt) + transcription = input_audio_transcription_config(payload) + if isinstance(transcription, dict): + extra_body["realtime_input_audio_transcription"] = dict(transcription) + noise_reduction = payload.get("input_audio_noise_reduction") + if isinstance(noise_reduction, dict): + extra_body["realtime_input_audio_noise_reduction"] = dict(noise_reduction) + audio_noise_reduction = audio_input.get("noise_reduction") if isinstance(audio_input, dict) else None + if isinstance(audio_noise_reduction, dict): + extra_body["realtime_input_audio_noise_reduction"] = dict(audio_noise_reduction) + if isinstance(audio_config, dict): + extra_body["realtime_audio"] = dict(audio_config) + if isinstance(payload.get("tracing"), str | dict): + extra_body["realtime_tracing"] = payload["tracing"] + if isinstance(payload.get("turn_detection"), dict): + extra_body["realtime_turn_detection"] = dict(cast("Mapping[str, object]", payload["turn_detection"])) + elif "turn_detection" in payload and payload.get("turn_detection") is None: + extra_body["realtime_turn_detection"] = None + extra_body.setdefault("realtime_output_audio_format", defaults.output_audio_format) + voice = payload.get("voice") + if not isinstance(voice, str) and isinstance(audio_output, dict): + voice = audio_output.get("voice") + speed = payload.get("speed") + if not isinstance(speed, int | float) and isinstance(audio_output, dict): + speed = audio_output.get("speed") + resolved_model = payload.get("model") if isinstance(payload.get("model"), str) else model + session: dict[str, object] = { + "model": resolved_model, + "modalities": payload.get("modalities") or payload.get("output_modalities") or ["text", "audio"], + "instructions": payload.get("instructions"), + # Not a Realtime field, but a session one: a client that cannot + # reach the model by speech seeds its opening turn here. Without + # this a websocket client could only use the older + # ``extra_body.duplex_initial_user_text`` spelling. + "initial_user_text": payload.get("initial_user_text"), + "voice": voice, + "ref_audio": payload.get("ref_audio"), + "response_format": duplex_response_format(defaults.output_audio_format), + "temperature": payload.get("temperature"), + "max_tokens": realtime_max_output_tokens( + payload.get("max_response_output_tokens") + or payload.get("max_output_tokens") + or payload.get("max_tokens") + ), + "speed": speed, + "idle_timeout_s": payload.get("idle_timeout_s") or 300.0, + **realtime_overlap_fields(payload), + "extra_body": extra_body, + } + return cls.from_event({"session": session}) + + def apply_realtime_update( + self, + patch: Mapping[str, object], + *, + session_id: str | None = None, + audio_started: bool = False, + ) -> None: + """Apply a Realtime ``session.update`` patch in place. + + Raises :class:`DuplexConfigError` (``code`` in ``model_update_unsupported``, + ``voice_update_after_audio_unsupported``, ``ref_audio_update_unsupported``) + when the patch changes something a live session cannot change. + ``audio_started`` is ``playback.generated_ms > 0 or playback.sent_ms > 0``. + """ + from vllm_omni.engine.duplex.realtime_commands import ( + REALTIME_OUTPUT_AUDIO_FORMATS, + duplex_response_format, + input_audio_transcription_config, + json_safe_realtime_payload, + parse_realtime_audio_format, + realtime_max_output_tokens, + ) + + # Any: the patch is read with ``isinstance(payload.get(...))`` guards and consumed + # through ``payload[...]``; typing the values ``object`` would need a cast per read. + payload: dict[str, Any] = dict(patch) + model = payload.get("model") + audio_config = payload.get("audio") + audio_input = audio_config.get("input") if isinstance(audio_config, dict) else None + audio_output = audio_config.get("output") if isinstance(audio_config, dict) else None + voice = payload.get("voice") + if not isinstance(voice, str) and isinstance(audio_output, dict): + voice = audio_output.get("voice") + if isinstance(model, str) and self.model is not None and model != self.model: + raise DuplexConfigError( + "session.update cannot change model for an open realtime duplex session", + code="model_update_unsupported", + ) + if isinstance(model, str) and self.model is None: + self.model = model + if isinstance(voice, str) and audio_started: + raise DuplexConfigError( + "session.update cannot change voice after audio output has started", + code="voice_update_after_audio_unsupported", + ) + if isinstance(payload.get("ref_audio"), str): + raise DuplexConfigError( + "session.update cannot change ref_audio after the session is open", code="ref_audio_update_unsupported" + ) + if isinstance(payload.get("instructions"), str): + self.instructions = str(payload["instructions"]) + elif "instructions" in payload and payload.get("instructions") is None: + self.instructions = None + if isinstance(voice, str): + self.voice = str(voice) + elif "voice" in payload and payload.get("voice") is None: + self.voice = None + response_format = payload.get("output_audio_format") or payload.get("response_format") + if response_format is None and isinstance(audio_output, dict): + response_format = audio_output.get("format") + response_format, _ = parse_realtime_audio_format(response_format) + if isinstance(response_format, str) and response_format.lower() in REALTIME_OUTPUT_AUDIO_FORMATS: + self.response_format = duplex_response_format(response_format) + if isinstance(payload.get("temperature"), int | float): + self.temperature = float(payload["temperature"]) + speed = payload.get("speed") + if not isinstance(speed, int | float) and isinstance(audio_output, dict): + speed = audio_output.get("speed") + if isinstance(speed, int | float): + self.speed = float(speed) + max_tokens = ( + payload.get("max_response_output_tokens") + if "max_response_output_tokens" in payload + else payload.get("max_output_tokens") + if "max_output_tokens" in payload + else payload.get("max_tokens") + ) + if "max_response_output_tokens" in payload or "max_output_tokens" in payload or "max_tokens" in payload: + self.max_tokens = realtime_max_output_tokens(max_tokens) + if isinstance(payload.get("overlap_policy"), str): + self.overlap_policy = self._normalize_overlap_policy(str(payload["overlap_policy"])) + if isinstance(payload.get("overlap_short_ack_ms"), int | float): + self.overlap_short_ack_ms = max(0, int(payload["overlap_short_ack_ms"])) + if isinstance(payload.get("overlap_barge_in_ms"), int | float): + self.overlap_barge_in_ms = max(0, int(payload["overlap_barge_in_ms"])) + if isinstance(payload.get("overlap_silence_rms"), int | float): + self.overlap_silence_rms = max(0.0, float(payload["overlap_silence_rms"])) + if isinstance(payload.get("playback_commit_policy"), str): + self.playback_commit_policy = self._normalize_playback_commit_policy(str(payload["playback_commit_policy"])) + modalities = payload.get("modalities") or payload.get("output_modalities") + if isinstance(modalities, list) and all(isinstance(item, str) for item in modalities): + self.modalities = list(modalities) + if isinstance(payload.get("extra_body"), dict): + self.extra_body.update(payload["extra_body"]) + extra = payload["extra_body"] + if isinstance(extra.get("overlap_policy"), str): + self.overlap_policy = self._normalize_overlap_policy(str(extra["overlap_policy"])) + if isinstance(extra.get("playback_commit_policy"), str): + self.playback_commit_policy = self._normalize_playback_commit_policy( + str(extra["playback_commit_policy"]) + ) + if isinstance(payload.get("tools"), list): + self.extra_body["realtime_tools"] = payload["tools"] + elif "tools" in payload and payload.get("tools") is None: + self.extra_body.pop("realtime_tools", None) + if isinstance(payload.get("tool_choice"), str | dict): + self.extra_body["realtime_tool_choice"] = payload["tool_choice"] + elif "tool_choice" in payload and payload.get("tool_choice") is None: + self.extra_body.pop("realtime_tool_choice", None) + if isinstance(payload.get("metadata"), dict): + self.extra_body["realtime_metadata"] = dict(payload["metadata"]) + elif "metadata" in payload and payload.get("metadata") is None: + self.extra_body.pop("realtime_metadata", None) + if isinstance(payload.get("include"), list): + self.extra_body["realtime_include"] = list(payload["include"]) + elif "include" in payload and payload.get("include") is None: + self.extra_body.pop("realtime_include", None) + if isinstance(payload.get("prompt"), dict): + self.extra_body["realtime_prompt"] = dict(payload["prompt"]) + elif "prompt" in payload and payload.get("prompt") is None: + self.extra_body.pop("realtime_prompt", None) + transcription = input_audio_transcription_config(payload) + if isinstance(transcription, dict): + self.extra_body["realtime_input_audio_transcription"] = dict(transcription) + elif "input_audio_transcription" in payload and payload.get("input_audio_transcription") is None: + self.extra_body.pop("realtime_input_audio_transcription", None) + if isinstance(payload.get("input_audio_noise_reduction"), dict): + self.extra_body["realtime_input_audio_noise_reduction"] = dict(payload["input_audio_noise_reduction"]) + elif "input_audio_noise_reduction" in payload and payload.get("input_audio_noise_reduction") is None: + self.extra_body.pop("realtime_input_audio_noise_reduction", None) + if isinstance(audio_input, dict) and isinstance(audio_input.get("noise_reduction"), dict): + self.extra_body["realtime_input_audio_noise_reduction"] = dict(audio_input["noise_reduction"]) + elif isinstance(audio_input, dict) and audio_input.get("noise_reduction") is None: + self.extra_body.pop("realtime_input_audio_noise_reduction", None) + if isinstance(payload.get("audio"), dict): + self.extra_body["realtime_audio"] = dict(payload["audio"]) + elif "audio" in payload and payload.get("audio") is None: + self.extra_body.pop("realtime_audio", None) + if isinstance(payload.get("tracing"), str | dict): + self.extra_body["realtime_tracing"] = payload["tracing"] + elif "tracing" in payload and payload.get("tracing") is None: + self.extra_body.pop("realtime_tracing", None) + if isinstance(payload.get("turn_detection"), dict): + self.extra_body["realtime_turn_detection"] = dict(payload["turn_detection"]) + elif "turn_detection" in payload and payload.get("turn_detection") is None: + self.extra_body["realtime_turn_detection"] = None + self.extra_body["realtime_session_payload"] = json_safe_realtime_payload(payload) + return None + + @staticmethod + def _normalize_overlap_policy(value: str) -> str: + normalized = value.strip().lower() + if normalized in {policy.value for policy in DuplexOverlapPolicy}: + return normalized + return DuplexOverlapPolicy.LISTEN_ONLY.value + + @staticmethod + def _normalize_playback_commit_policy(value: str) -> str: + normalized = value.strip().lower() + if normalized in {policy.value for policy in DuplexPlaybackCommitPolicy}: + return normalized + return DuplexPlaybackCommitPolicy.COMMIT_ALL_ON_DONE.value + + +@dataclass(frozen=True) +class ResponseCreateOptions: + instructions: str | None = None + voice: str | None = None + response_format: str | None = None + temperature: float | None = None + max_tokens: int | None = None + speed: float | None = None + modalities: tuple[str, ...] | None = None + extra_body: Mapping[str, object] = field(default_factory=dict) + + def apply_to(self, config: DuplexSessionConfig) -> None: + for field_name in ("instructions", "voice", "response_format", "temperature", "max_tokens", "speed"): + value = getattr(self, field_name) + if value is not None: + setattr(config, field_name, value) + if self.modalities is not None: + config.modalities = list(self.modalities) + config.extra_body.update(self.extra_body) + + @classmethod + def from_realtime( + cls, + response_payload: Mapping[str, object], + *, + private_runtime_config_keys: frozenset[str] = frozenset(), + ) -> ResponseCreateOptions: + """Parse an OpenAI Realtime ``response`` object into response-scoped options. + + Raises :class:`DuplexConfigError` (``code="unsupported_native_response_options"``) + for options a model-native duplex session cannot apply per response. + Private runtime keys in ``extra_body`` are dropped. + """ + from vllm_omni.engine.duplex.realtime_commands import ( + REALTIME_OUTPUT_AUDIO_FORMATS, + duplex_response_format, + parse_realtime_audio_format, + realtime_max_output_tokens, + ) + + payload: dict[str, object] = dict(response_payload) + audio_config = payload.get("audio") + audio_output = audio_config.get("output") if isinstance(audio_config, dict) else None + nested_voice = audio_output.get("voice") if isinstance(audio_output, dict) else None + unsupported = ( + payload.get("instructions") is not None + or payload.get("voice") is not None + or nested_voice is not None + or payload.get("temperature") is not None + or any( + payload.get(field_name) is not None + for field_name in ("max_response_output_tokens", "max_output_tokens", "max_tokens") + ) + or payload.get("tools") is not None + or payload.get("tool_choice") is not None + ) + if unsupported: + raise DuplexConfigError( + "response.create options are not supported by a model-native duplex session", + code="unsupported_native_response_options", + ) + instructions = str(payload["instructions"]) if isinstance(payload.get("instructions"), str) else None + voice = payload.get("voice") + if not isinstance(voice, str) and isinstance(audio_output, dict): + voice = audio_output.get("voice") + voice = str(voice) if isinstance(voice, str) else None + response_format = payload.get("output_audio_format") or payload.get("response_format") + if response_format is None and isinstance(audio_output, dict): + response_format = audio_output.get("format") + response_format, _ = parse_realtime_audio_format(response_format) + if isinstance(response_format, str) and response_format.lower() in REALTIME_OUTPUT_AUDIO_FORMATS: + response_format = duplex_response_format(response_format) + else: + response_format = None + temperature = ( + float(cast("int | float", payload["temperature"])) + if isinstance(payload.get("temperature"), int | float) + else None + ) + speed = payload.get("speed") + if not isinstance(speed, int | float) and isinstance(audio_output, dict): + speed = audio_output.get("speed") + speed = float(speed) if isinstance(speed, int | float) else None + max_tokens = ( + payload.get("max_response_output_tokens") + if "max_response_output_tokens" in payload + else payload.get("max_output_tokens") + if "max_output_tokens" in payload + else payload.get("max_tokens") + ) + if "max_response_output_tokens" in payload or "max_output_tokens" in payload or "max_tokens" in payload: + max_tokens = realtime_max_output_tokens(max_tokens) + else: + max_tokens = None + modalities = payload.get("modalities") or payload.get("output_modalities") + if isinstance(modalities, list) and all(isinstance(item, str) for item in modalities): + modalities = tuple(modalities) + else: + modalities = None + response_extra: dict[str, object] = {} + conversation = payload.get("conversation") + if isinstance(conversation, str): + response_extra["realtime_response_conversation"] = conversation + metadata = payload.get("metadata") + if isinstance(metadata, dict): + response_extra["realtime_response_metadata"] = dict(metadata) + prompt = payload.get("prompt") + if isinstance(prompt, dict): + response_extra["realtime_response_prompt"] = dict(prompt) + if isinstance(payload.get("tools"), list): + response_extra["realtime_response_tools"] = payload["tools"] + if isinstance(payload.get("tool_choice"), str | dict): + response_extra["realtime_response_tool_choice"] = payload["tool_choice"] + extra_body = payload.get("extra_body") + if isinstance(extra_body, dict): + response_extra.update( + (key, value) for key, value in extra_body.items() if key not in private_runtime_config_keys + ) + return cls( + instructions=instructions, + voice=voice, + response_format=response_format, + temperature=temperature, + max_tokens=max_tokens, + speed=speed, + modalities=modalities, + extra_body=response_extra, + ) + + +@dataclass +class DuplexCommittedInput: + message: dict[str, object] + turn_id: int + epoch: int + input_commit_seq: int + + +@dataclass +class DuplexAssistantAudioTextMark: + text_chars: int + audio_end_ms: int + + +def realtime_item_to_history_message(item: object) -> dict[str, object] | None: + """Convert a Realtime conversation item into a chat history message (None when empty).""" + if not isinstance(item, dict): + return None + role = item.get("role") + if role not in {"system", "user", "assistant"}: + return None + content = item.get("content") + if isinstance(content, str): + text = content.strip() + return {"role": role, "content": text} if text else None + if not isinstance(content, list): + return None + text_chunks: list[str] = [] + audio_chunks: list[dict[str, object]] = [] + for part in content: + if not isinstance(part, dict): + continue + part_type = part.get("type") + if part_type in {"input_text", "text", "output_text"} and isinstance(part.get("text"), str): + text_chunks.append(str(part["text"])) + elif part_type in {"input_audio", "audio"}: + audio = part.get("audio") or part.get("data") + fmt = part.get("format") if isinstance(part.get("format"), str) else "wav" + if isinstance(audio, str) and audio: + audio_chunks.append({"type": "audio_url", "audio_url": {"url": f"data:audio/{fmt};base64,{audio}"}}) + elif part_type in {"audio_transcript", "transcript"} and isinstance(part.get("text"), str): + text_chunks.append(str(part["text"])) + text = "".join(text_chunks).strip() + if audio_chunks: + content_items: list[dict[str, object]] = [] + if text: + content_items.append({"type": "text", "text": text}) + content_items.extend(audio_chunks) + return {"role": role, "content": content_items} + if text: + return {"role": role, "content": text} + return None + + +def realtime_max_output_tokens(value: object) -> int | None: + """Normalize Realtime max output tokens (``"inf"`` -> ``None``).""" + from vllm_omni.engine.duplex.realtime_commands import realtime_max_output_tokens as _impl + + return _impl(value) + + +def input_audio_transcription_config(session_payload: Mapping[str, object]) -> dict[str, object] | None: + """Return the ``input_audio_transcription`` object of a Realtime session payload.""" + from vllm_omni.engine.duplex.realtime_commands import input_audio_transcription_config as _impl + + return _impl(session_payload) + + +__all__ = [ + "DuplexAssistantAudioTextMark", + "DuplexAudioChunk", + "DuplexCapabilities", + "DuplexCommittedInput", + "DuplexConfigError", + "DuplexOverlapPolicy", + "DuplexPlaybackCommitPolicy", + "DuplexPlaybackCursor", + "DuplexPlaybackView", + "DuplexSessionConfig", + "DuplexSessionState", + "DuplexTurnEventType", + "DuplexTurnState", + "ResponseCreateOptions", + "input_audio_transcription_config", + "realtime_item_to_history_message", + "realtime_max_output_tokens", +] diff --git a/vllm_omni/engine/duplex/contracts.py b/vllm_omni/engine/duplex/contracts.py index 6af4fdea5..ed14f5fb8 100644 --- a/vllm_omni/engine/duplex/contracts.py +++ b/vllm_omni/engine/duplex/contracts.py @@ -1,102 +1,50 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Model-neutral contracts for the experimental duplex engine plugin. +"""Model-neutral value types shared by the duplex engine components. -This module contains only immutable data transfer objects and narrow protocols. -Duplex control algorithms, session implementations, model policy, and Realtime -serving remain in sibling experimental modules. +Immutable DTOs plus the ``DuplexStagePort`` base class that ``DuplexOrchestrator`` +implements for the session manager/runner. """ from __future__ import annotations import base64 -from collections.abc import Iterable, Mapping +from abc import ABC, abstractmethod +from collections.abc import Mapping from dataclasses import dataclass, field from enum import Enum from types import MappingProxyType -from typing import Any, Protocol - -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.messages import EngineQueueMessage -class SessionMode(str, Enum): - TURN = "turn" - DUPLEX = "duplex" +@dataclass(frozen=True, slots=True) +class DuplexFence: + """Engine-internal session identity used for stage request ids and stale filtering.""" - -class DuplexInputMode(str, Enum): - APPEND_TOKENS = "append_tokens" - APPEND_AUDIO_CHUNK = "append_audio_chunk" - REPLACE_LATEST_CHUNK = "replace_latest_chunk" - REENCODE_CONTEXT = "reencode_context" - ROLLBACK_TO_CHECKPOINT = "rollback_to_checkpoint" - TURN_COMMIT_ONLY = "turn_commit_only" + session_id: str + epoch: int = 0 + turn_id: int = 0 class DuplexOutputAction(str, Enum): DIRECT_RESPONSE = "direct_response" -@dataclass -class DuplexRuntimeCapabilities: - input_modes: set[DuplexInputMode] = field(default_factory=lambda: {DuplexInputMode.TURN_COMMIT_ONLY}) - implementation_level: str = "serving_session_adapter" - - @dataclass(frozen=True) class DuplexAppendPlan: - prompt: dict[str, Any] + prompt: dict[str, object] @dataclass(frozen=True) class DuplexOutputDecision: action: DuplexOutputAction - metadata: Mapping[str, Any] = field(default_factory=dict) + metadata: Mapping[str, object] = field(default_factory=dict) final_output_type: str = "text" def __post_init__(self) -> None: object.__setattr__(self, "metadata", MappingProxyType(dict(self.metadata))) -class DuplexRuntimeExtension(Protocol): - """Pure model policy invoked by the experimental duplex control plane.""" - - def configure_sampling_params( - self, - *, - runtime_config: dict[str, Any], - defaults: tuple[object, ...], - ) -> tuple[object, ...]: ... - - def plan_append( - self, - *, - request_id: str, - fence: DuplexFence, - session_config: dict[str, Any], - runtime_config: dict[str, Any], - seq: int, - turn_seq: int, - mode: DuplexInputMode, - payload: object, - final: bool, - sampling_params: object, - ) -> DuplexAppendPlan: ... - - def decide_output( - self, - *, - stage_id: int, - final_stage_id: int, - segment_finished: bool, - segment_token_ids: tuple[int, ...], - segment_output_metadata: dict[str, Any], - output: object, - ) -> DuplexOutputDecision | None: ... - - @dataclass(frozen=True) class DuplexRequestIdentity: session_id: str @@ -112,8 +60,8 @@ class DuplexStageRequestContext: final_stage_id: int config_generation: int sampling_params: tuple[object, ...] - session_config: Mapping[str, Any] = field(default_factory=dict) - runtime_config: Mapping[str, Any] = field(default_factory=dict) + session_config: Mapping[str, object] = field(default_factory=dict) + runtime_config: Mapping[str, object] = field(default_factory=dict) def __post_init__(self) -> None: object.__setattr__(self, "sampling_params", tuple(self.sampling_params)) @@ -128,7 +76,7 @@ class DuplexStageRequestContext: @dataclass(frozen=True) class DuplexStageSubmission: context: DuplexStageRequestContext - prompt: Mapping[str, Any] + prompt: Mapping[str, object] already_submitted: bool def __post_init__(self) -> None: @@ -148,7 +96,7 @@ class DuplexOutputContext: final_stage_id: int segment_finished: bool segment_token_ids: tuple[int, ...] = () - segment_output_metadata: Mapping[str, Any] = field(default_factory=dict) + segment_output_metadata: Mapping[str, object] = field(default_factory=dict) def __post_init__(self) -> None: object.__setattr__(self, "segment_token_ids", tuple(self.segment_token_ids)) @@ -159,53 +107,27 @@ class DuplexOutputContext: ) -class DuplexStagePort(Protocol): +class DuplexStagePort(ABC): + """Narrow stage-management surface the session runner/manager use (implemented by DuplexOrchestrator).""" + @property + @abstractmethod def stage_count(self) -> int: ... + @abstractmethod def sampling_defaults(self) -> tuple[object, ...]: ... + @abstractmethod def ensure_request(self, context: DuplexStageRequestContext) -> None: ... + @abstractmethod async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: ... + @abstractmethod async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: ... - -class DuplexControlPlanePort(Protocol): - @property - def sessions(self) -> object: ... - - def accepts(self, message: object) -> bool: ... - - def dispatch(self, message: object) -> None: ... - - async def shutdown(self) -> None: ... - - def close_sessions_for_request_ids(self, request_ids: list[str]) -> dict[str, list[str]]: ... - - def finalize_closed_sessions(self, session_ids: Iterable[str]) -> None: ... - - def session_for_identity(self, identity: DuplexRequestIdentity | None) -> object | None: ... - - def decide_output( - self, - stage_id: int, - output: object, - context: DuplexOutputContext | None, - ) -> DuplexOutputDecision | None: ... - - -class CorrelatedRpcTransport(Protocol): - def execute( - self, - key: tuple[str, str], - message: EngineQueueMessage, - *, - timeout: float | None, - timeout_message: str, - block_on_submit: bool = False, - ) -> EngineQueueMessage: ... + @abstractmethod + async def abort_requests(self, request_ids: list[str]) -> None: ... def duplex_data_plane_request_info(result: dict[str, object]) -> tuple[str | None, int | None]: @@ -232,20 +154,19 @@ def duplex_resource_request_id(fence: DuplexFence, role: str) -> str: ): raise ValueError(f"invalid duplex resource role: {role!r}") encoded_session_id = base64.urlsafe_b64encode(fence.session_id.encode("utf-8")).decode("ascii").rstrip("=") - return f"duplex-s.{encoded_session_id}.i.{fence.incarnation}.e.{fence.epoch}.r.{role}" + return f"duplex-s.{encoded_session_id}.e.{fence.epoch}.r.{role}" def duplex_resource_request_belongs_to_session(request_id: str, session_id: str) -> bool: """Return whether a current-format resource request belongs to a session.""" parts = request_id.split(".") - if len(parts) != 8 or parts[0] != "duplex-s" or parts[2] != "i" or parts[4] != "e" or parts[6] != "r": + if len(parts) != 6 or parts[0] != "duplex-s" or parts[2] != "e" or parts[4] != "r": return False try: int(parts[3]) - int(parts[5]) except ValueError: return False - role = parts[7] + role = parts[5] if not role or any( character not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-" for character in role ): @@ -255,21 +176,16 @@ def duplex_resource_request_belongs_to_session(request_id: str, session_id: str) __all__ = [ - "CorrelatedRpcTransport", + "DuplexFence", "DuplexAppendPlan", - "DuplexControlPlanePort", - "DuplexInputMode", "DuplexOutputAction", "DuplexOutputContext", "DuplexOutputDecision", "DuplexRequestIdentity", - "DuplexRuntimeCapabilities", - "DuplexRuntimeExtension", "DuplexStagePort", "DuplexStageRequestContext", "DuplexStageSubmission", "DuplexStageSubmissionResult", - "SessionMode", "duplex_data_plane_request_info", "duplex_resource_request_belongs_to_session", "duplex_resource_request_id", diff --git a/vllm_omni/engine/duplex/control_client.py b/vllm_omni/engine/duplex/control_client.py deleted file mode 100644 index 05ec46e49..000000000 --- a/vllm_omni/engine/duplex/control_client.py +++ /dev/null @@ -1,228 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import uuid -from collections.abc import Callable - -from vllm_omni.engine.duplex.contracts import CorrelatedRpcTransport -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity -from vllm_omni.engine.duplex.messages import ( - AppendDuplexInputMessage, - CloseDuplexSessionMessage, - DuplexControlResultMessage, - DuplexFence, - OpenDuplexSessionMessage, - ResumeDuplexSessionMessage, - SignalDuplexTurnMessage, - TouchDuplexSessionMessage, -) - -DuplexControlMessage = ( - OpenDuplexSessionMessage - | AppendDuplexInputMessage - | SignalDuplexTurnMessage - | CloseDuplexSessionMessage - | TouchDuplexSessionMessage - | ResumeDuplexSessionMessage -) - - -class DuplexControlRequestError(RuntimeError): - """Typed client-side failure returned by the duplex control plane.""" - - def __init__(self, result: dict[str, object]) -> None: - error = result.get("error") - error_data = error if isinstance(error, dict) else {} - self.result = result - self.code = str(error_data.get("code") or "internal_error") - self.retryable = bool(error_data.get("retryable", False)) - accepted_fence = result.get("accepted_fence") - self.accepted_fence = accepted_fence if isinstance(accepted_fence, DuplexFence) else None - lease_generation = result.get("lease_generation") - self.lease_generation = lease_generation if isinstance(lease_generation, int) else None - super().__init__(f"duplex {result.get('operation')} failed: {result}") - - -class DuplexControlClient: - """Builds duplex commands on top of an engine-owned RPC transport.""" - - def __init__( - self, - transport: CorrelatedRpcTransport, - *, - control_id_factory: Callable[[], str] | None = None, - ) -> None: - self._transport = transport - self._control_id_factory = control_id_factory or (lambda: uuid.uuid4().hex) - - def execute(self, message: DuplexControlMessage, *, timeout: float | None) -> dict[str, object]: - result_message = self._transport.execute( - ("duplex", message.control_id), - message, - timeout=timeout, - timeout_message=f"duplex control timed out after {timeout} seconds", - ) - if not isinstance(result_message, DuplexControlResultMessage): - raise RuntimeError(f"unexpected duplex control result type: {type(result_message).__name__}") - if result_message.fence != message.fence: - raise RuntimeError( - f"duplex control fence mismatch: expected {message.fence!r}, got {result_message.fence!r}" - ) - result = { - "fence": result_message.fence, - "operation": result_message.operation, - "session_id": result_message.session_id, - "ok": result_message.ok, - "stage_results": list(result_message.stage_results), - "unsupported_count": result_message.unsupported_count, - "error_count": result_message.error_count, - "accepted_fence": result_message.accepted_fence, - "lease_generation": result_message.lease_generation, - "error": ( - { - "code": result_message.error.code, - "message": result_message.error.message, - "retryable": result_message.error.retryable, - } - if result_message.error is not None - else None - ), - } - if not result_message.ok: - raise DuplexControlRequestError(result) - return result - - def open( - self, - session_id: str, - *, - session_mode: str, - capabilities: dict[str, object] | None, - session_config: dict[str, object] | None, - runtime_config: dict[str, object] | None, - fence: DuplexFence, - timeout: float | None, - ) -> dict[str, object]: - return self.execute( - OpenDuplexSessionMessage( - control_id=self._control_id_factory(), - fence=fence, - session_id=session_id, - session_mode=session_mode, - capabilities=dict(capabilities or {}), - session_config=dict(session_config or {}), - runtime_config=dict(runtime_config or {}), - ), - timeout=timeout, - ) - - def append( - self, - session_id: str, - *, - mode: str, - payload: object, - operation_id: str | None, - final: bool, - expected_epoch: int | None, - fence: DuplexFence, - timeout: float | None, - ) -> dict[str, object]: - if expected_epoch is not None and expected_epoch != fence.epoch: - raise ValueError("expected_epoch must match fence.epoch") - return self.execute( - AppendDuplexInputMessage( - control_id=self._control_id_factory(), - operation_id=operation_id, - fence=fence, - session_id=session_id, - expected_epoch=expected_epoch, - mode=mode, - payload=payload, - final=final, - ), - timeout=timeout, - ) - - def signal( - self, - session_id: str, - *, - event: str, - fence: DuplexFence, - next_fence: DuplexFence | None, - session_config: dict[str, object] | None, - runtime_config: dict[str, object] | None, - timeout: float | None, - ) -> dict[str, object]: - return self.execute( - SignalDuplexTurnMessage( - control_id=self._control_id_factory(), - fence=fence, - session_id=session_id, - event=event, - next_fence=next_fence, - session_config=dict(session_config) if session_config is not None else None, - runtime_config=dict(runtime_config) if runtime_config is not None else None, - ), - timeout=timeout, - ) - - def close( - self, - session_id: str, - *, - reason: str, - fence: DuplexFence, - timeout: float | None, - ) -> dict[str, object]: - return self.execute( - CloseDuplexSessionMessage( - control_id=self._control_id_factory(), - fence=fence, - session_id=session_id, - reason=reason, - ), - timeout=timeout, - ) - - def touch( - self, - session_id: str, - *, - fence: DuplexFence, - activity: DuplexLeaseActivity, - timeout: float | None, - ) -> dict[str, object]: - return self.execute( - TouchDuplexSessionMessage( - control_id=self._control_id_factory(), - fence=fence, - session_id=session_id, - activity=activity.value, - ), - timeout=timeout, - ) - - def resume( - self, - session_id: str, - *, - fence: DuplexFence, - expected_lease_generation: int, - timeout: float | None, - ) -> dict[str, object]: - return self.execute( - ResumeDuplexSessionMessage( - control_id=self._control_id_factory(), - fence=fence, - session_id=session_id, - expected_lease_generation=expected_lease_generation, - ), - timeout=timeout, - ) - - -__all__ = ["CorrelatedRpcTransport", "DuplexControlClient", "DuplexControlMessage"] diff --git a/vllm_omni/engine/duplex/control_plane.py b/vllm_omni/engine/duplex/control_plane.py deleted file mode 100644 index b425542b8..000000000 --- a/vllm_omni/engine/duplex/control_plane.py +++ /dev/null @@ -1,1063 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -"""Orchestrator-side control plane for experimental duplex sessions. - -``Orchestrator`` creates this component only when duplex control is enabled -and supplies a narrow :class:`DuplexStagePort` implementation. Queue messages -are routed through :meth:`DuplexControlPlane.handle`; stage outputs are passed -to :meth:`DuplexControlPlane.decide_output`; request cleanup is reported back -through the control plane so session leases and stage bindings stay coherent. - -The module owns duplex session/control algorithms. It deliberately does not -own stage pools, request queues, or OpenAI Realtime protocol state. -""" - -from __future__ import annotations - -import asyncio -from collections.abc import Callable, Iterable -from dataclasses import dataclass -from typing import Any, Protocol - -from vllm.logger import init_logger - -from vllm_omni.engine.duplex.contracts import ( - DuplexAppendPlan, - DuplexInputMode, - DuplexOutputContext, - DuplexOutputDecision, - DuplexRequestIdentity, - DuplexRuntimeCapabilities, - DuplexRuntimeExtension, - DuplexStagePort, - DuplexStageRequestContext, - DuplexStageSubmission, - DuplexStageSubmissionResult, - SessionMode, - duplex_resource_request_id, -) -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity, DuplexLeaseConfig -from vllm_omni.engine.duplex.messages import ( - AppendDuplexInputMessage, - CloseDuplexSessionMessage, - DuplexControlError, - DuplexControlResultMessage, - DuplexFence, - DuplexSessionLifecycleMessage, - OpenDuplexSessionMessage, - ResumeDuplexSessionMessage, - SignalDuplexTurnMessage, - TouchDuplexSessionMessage, -) -from vllm_omni.engine.duplex.session import ( - DuplexAppendReservation, - DuplexFenceMismatchError, - DuplexSessionExpiry, - DuplexSessionRuntimeManager, - DuplexSessionRuntimeState, -) - -logger = init_logger(__name__) - - -@dataclass(frozen=True) -class _PendingControlCleanup: - kind: str - session_id: str - fence: DuplexFence - submitted_request_ids: tuple[str, ...] - reserved_request_ids: tuple[str, ...] = () - - -@dataclass(frozen=True) -class _PendingSubmissionCleanup: - session_id: str - request_ids: tuple[str, ...] - - -@dataclass(frozen=True) -class _PendingRequestCleanup: - session_id: str - fence: DuplexFence - lease_generation: int - request_ids: tuple[str, ...] - abort: bool - - -class DuplexResultSink(Protocol): - async def put(self, message: DuplexControlResultMessage) -> None: ... - - -class DuplexLifecycleSink(Protocol): - async def put(self, message: DuplexSessionLifecycleMessage) -> None: ... - - -class DuplexControlPlane: - _MESSAGE_TYPES = ( - OpenDuplexSessionMessage, - AppendDuplexInputMessage, - SignalDuplexTurnMessage, - CloseDuplexSessionMessage, - TouchDuplexSessionMessage, - ResumeDuplexSessionMessage, - ) - - def __init__( - self, - *, - extension: DuplexRuntimeExtension | None, - stage_port: DuplexStagePort, - result_sink: DuplexResultSink, - lifecycle_sink: DuplexLifecycleSink | None = None, - lease_config: DuplexLeaseConfig | None = None, - clock: Callable[[], float] | None = None, - max_sessions: int | None = None, - completed_append_limit: int = 256, - ) -> None: - self._extension = extension - self._stage_port = stage_port - self._result_sink = result_sink - self._lifecycle_sink = lifecycle_sink - self._lease_config = lease_config or DuplexLeaseConfig() - self._sessions = DuplexSessionRuntimeManager( - clock=clock, - max_sessions=max_sessions, - completed_append_limit=completed_append_limit, - ) - self._pending_expirations: dict[tuple[str, int], DuplexSessionExpiry] = {} - self._pending_control_cleanups: dict[tuple[str, str, int, int, int, int], _PendingControlCleanup] = {} - self._control_cleanup_tasks: dict[ - tuple[str, str, int, int, int, int], - asyncio.Task[None], - ] = {} - self._pending_submission_cleanups: dict[str, _PendingSubmissionCleanup] = {} - self._submission_cleanup_tasks: dict[str, asyncio.Task[None]] = {} - self._pending_request_cleanups: dict[tuple[str, int, int], _PendingRequestCleanup] = {} - self._request_cleanup_tasks: dict[tuple[str, int, int], asyncio.Task[None]] = {} - self._request_cleanups_in_progress: set[tuple[str, int, int]] = set() - self._session_control_tails: dict[str, asyncio.Task[None]] = {} - self._dispatched_control_tasks: set[asyncio.Task[None]] = set() - - @property - def sessions(self) -> DuplexSessionRuntimeManager: - return self._sessions - - @property - def pending_submission_cleanup_count(self) -> int: - return len(self._pending_submission_cleanups) - - def accepts(self, message: object) -> bool: - return isinstance(message, self._MESSAGE_TYPES) - - async def handle(self, message: object) -> None: - if isinstance(message, OpenDuplexSessionMessage): - await self.handle_open(message) - elif isinstance(message, AppendDuplexInputMessage): - await self.handle_append(message) - elif isinstance(message, SignalDuplexTurnMessage): - await self.handle_signal(message) - elif isinstance(message, CloseDuplexSessionMessage): - await self.handle_close(message) - elif isinstance(message, TouchDuplexSessionMessage): - await self.handle_touch(message) - elif isinstance(message, ResumeDuplexSessionMessage): - await self.handle_resume(message) - else: - raise TypeError(f"Unsupported duplex control message: {type(message).__name__}") - - def dispatch(self, message: object) -> None: - """Schedule one control command without blocking unrelated sessions.""" - if not self.accepts(message): - raise TypeError(f"Unsupported duplex control message: {type(message).__name__}") - session_id = message.session_id - predecessor = self._session_control_tails.get(session_id) - - async def run_ordered() -> None: - if predecessor is not None: - await predecessor - await self.handle(message) - - task = asyncio.create_task( - run_ordered(), - name=f"duplex-control-{session_id}-{message.control_id}", - ) - self._session_control_tails[session_id] = task - self._dispatched_control_tasks.add(task) - - def discard(completed: asyncio.Task[None]) -> None: - self._dispatched_control_tasks.discard(completed) - if self._session_control_tails.get(session_id) is completed: - self._session_control_tails.pop(session_id, None) - if not completed.cancelled(): - completed.exception() - - task.add_done_callback(discard) - - async def drain(self) -> None: - while self._dispatched_control_tasks: - await asyncio.gather(*tuple(self._dispatched_control_tasks)) - - async def shutdown(self) -> None: - tasks = tuple( - self._dispatched_control_tasks - | set(self._control_cleanup_tasks.values()) - | set(self._submission_cleanup_tasks.values()) - | set(self._request_cleanup_tasks.values()) - ) - for task in tasks: - task.cancel() - if tasks: - await asyncio.gather(*tasks, return_exceptions=True) - - @staticmethod - def coerce_capabilities(raw: dict[str, object]) -> DuplexRuntimeCapabilities: - input_modes: set[DuplexInputMode] = set() - values = raw.get("input_modes") - if values is not None and not isinstance(values, list): - raise TypeError("duplex input_modes capability must be a list") - if isinstance(values, list): - for value in values: - try: - input_modes.add(DuplexInputMode(value)) - except (TypeError, ValueError) as exc: - raise ValueError(f"unknown duplex input mode: {value!r}") from exc - return DuplexRuntimeCapabilities( - input_modes=input_modes or {DuplexInputMode.TURN_COMMIT_ONLY}, - implementation_level=str(raw.get("implementation_level") or "serving_session_adapter"), - ) - - def sampling_params_for_config(self, runtime_config: dict[str, Any]) -> list[object]: - defaults = self._stage_port.sampling_defaults() - if self._extension is None: - return list(defaults) - configured = self._extension.configure_sampling_params( - runtime_config=runtime_config, - defaults=defaults, - ) - if not isinstance(configured, tuple): - raise TypeError("duplex runtime extension must return sampling parameters as a tuple") - if len(configured) != len(defaults): - raise ValueError("duplex runtime extension must return one sampling parameter per stage") - return list(configured) - - @staticmethod - def stage_request_id(fence: DuplexFence, *, stage_id: int) -> str: - return duplex_resource_request_id(fence, f"stage{stage_id}") - - def ensure_stage_request( - self, - session: DuplexSessionRuntimeState, - *, - stage_id: int, - fence: DuplexFence | None = None, - ) -> DuplexStageRequestContext | None: - if stage_id >= self._stage_port.stage_count: - return None - effective_fence = fence or session.fence - request_id = self.stage_request_id(effective_fence, stage_id=stage_id) - session.reserve_stage_request(stage_id, request_id, fence=effective_fence) - context = DuplexStageRequestContext( - request_id=request_id, - session_id=session.session_id, - fence=effective_fence, - stage_id=stage_id, - final_stage_id=self._stage_port.stage_count - 1, - config_generation=session.config_generation, - sampling_params=tuple(self.sampling_params_for_config(session.runtime_config)), - session_config=session.session_config, - runtime_config=session.runtime_config, - ) - self._stage_port.ensure_request(context) - return context - - async def handle_open(self, message: OpenDuplexSessionMessage) -> None: - session: DuplexSessionRuntimeState | None = None - try: - session_mode = SessionMode(message.session_mode) - capabilities = self.coerce_capabilities(message.capabilities) - native_append_modes = capabilities.input_modes - {DuplexInputMode.TURN_COMMIT_ONLY} - if native_append_modes and self._extension is None: - raise RuntimeError("duplex_runtime_extension_not_configured") - session = self.sessions.open_session( - message.fence, - capabilities=capabilities, - session_config=message.session_config, - runtime_config=message.runtime_config, - lease_config=self._lease_config, - ) - request_context = self.ensure_stage_request(session, stage_id=0) if self._extension is not None else None - await self.put_result( - message.control_id, - fence=message.fence, - operation="open", - session_id=message.session_id, - stage_results=[ - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": True, - "implementation_level": capabilities.implementation_level, - "data_plane_session": True, - "session_mode": session_mode.value, - "scheduler_request_context": request_context is not None, - "request_id": request_context.request_id if request_context is not None else None, - }, - } - ], - ) - except Exception as exc: - if self._control_error(exc).code == "resource_exhausted": - logger.info("open_duplex_session rejected: %s", exc) - else: - logger.exception("open_duplex_session failed: %s", exc) - if session is not None and self.sessions.get(message.session_id) is session: - reserved_request_ids = tuple(session.resource_request_ids()) - self.sessions.begin_close_session(session.fence, reason="open_rollback") - cleanup_key = self._cleanup_key("open_rollback", session.fence) - pending = _PendingControlCleanup( - kind="open_rollback", - session_id=session.session_id, - fence=session.fence, - reserved_request_ids=reserved_request_ids, - submitted_request_ids=(), - ) - self._pending_control_cleanups[cleanup_key] = pending - try: - await self._complete_control_cleanup(cleanup_key, pending) - except Exception as cleanup_exc: - logger.warning( - "duplex open rollback remains pending for session %s: %s", - session.session_id, - cleanup_exc, - ) - await self.put_result( - message.control_id, - fence=message.fence, - operation="open", - session_id=message.session_id, - stage_results=[], - error=exc, - ) - - async def handle_append(self, message: AppendDuplexInputMessage) -> None: - session: DuplexSessionRuntimeState | None = None - lease_operation_id = f"append:{message.control_id}" - operation_started = False - try: - session = self.sessions.require(message.session_id) - await self._complete_pending_submission_cleanup(message.session_id) - if message.expected_epoch is not None and message.expected_epoch != message.fence.epoch: - raise ValueError("expected_epoch must match fence.epoch") - mode = DuplexInputMode(message.mode) - if message.operation_id is not None: - completed = session.completed_append( - message.operation_id, - fence=message.fence, - mode=mode, - final=message.final, - ) - if completed is not None: - session.touch(message.fence, DuplexLeaseActivity.APPEND) - await self.put_result( - message.control_id, - fence=message.fence, - operation="append", - session_id=message.session_id, - stage_results=completed, - ) - return - session.begin_operation(message.fence, lease_operation_id) - operation_started = True - reservation = session.prepare_append(mode=mode, fence=message.fence) - stage_results = await self.append_via_data_plane( - message, - session=session, - reservation=reservation, - mode=mode, - ) - if message.operation_id is not None: - session.record_completed_append( - message.operation_id, - fence=message.fence, - mode=mode, - final=message.final, - stage_results=stage_results, - ) - await self.put_result( - message.control_id, - fence=message.fence, - operation="append", - session_id=message.session_id, - stage_results=stage_results, - ) - except Exception as exc: - logger.exception("append_duplex_input failed: %s", exc) - await self.put_result( - message.control_id, - fence=message.fence, - operation="append", - session_id=message.session_id, - stage_results=[], - error=exc, - ) - finally: - if session is not None and operation_started and lease_operation_id in session.lease.active_operations: - session.end_operation(session.fence, lease_operation_id) - - async def append_via_data_plane( - self, - message: AppendDuplexInputMessage, - *, - session: DuplexSessionRuntimeState, - reservation: DuplexAppendReservation, - mode: DuplexInputMode, - ) -> list[dict[str, object]]: - if self._extension is None and mode is DuplexInputMode.TURN_COMMIT_ONLY: - update = session.commit_append(reservation) - return [ - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": True, - "data_plane_append": False, - "seq": update.seq, - "turn_id": update.turn_id, - "turn_seq": update.turn_seq, - "mode": mode.value, - }, - } - ] - if self._stage_port.stage_count == 0: - return [ - { - "stage_id": -1, - "replica_id": -1, - "result": {"supported": False, "error": "duplex_data_plane_has_no_stage"}, - } - ] - if self._extension is None: - raise RuntimeError("duplex_runtime_extension_not_configured") - - stage_id = 0 - request_id = self.stage_request_id(message.fence, stage_id=stage_id) - existing_binding = session.stage_bindings.get(stage_id) - already_submitted = existing_binding is not None and existing_binding.request_id == request_id - request_context = self.ensure_stage_request( - session, - stage_id=stage_id, - fence=message.fence, - ) - if request_context is None: - raise RuntimeError("duplex_data_plane_has_no_stage") - append_plan = self._extension.plan_append( - request_id=request_id, - fence=message.fence, - session_config=dict(request_context.session_config), - runtime_config=dict(request_context.runtime_config), - seq=reservation.update.seq, - turn_seq=reservation.update.turn_seq, - mode=mode, - payload=message.payload, - final=message.final, - sampling_params=request_context.stage_sampling_params, - ) - if not isinstance(append_plan, DuplexAppendPlan): - raise TypeError("duplex runtime extension plan_append() must return DuplexAppendPlan") - submission = DuplexStageSubmission( - context=request_context, - prompt=append_plan.prompt, - already_submitted=already_submitted, - ) - submission_result = await self._stage_port.submit(submission) - try: - if submission_result.request_id != request_id or submission_result.stage_id != stage_id: - raise RuntimeError("duplex stage adapter returned a mismatched submission result") - update = session.commit_append(reservation) - session.bind_stage_request(stage_id, request_id, fence=message.fence) - except BaseException: - self._pending_submission_cleanups[session.session_id] = _PendingSubmissionCleanup( - session_id=session.session_id, - request_ids=(request_id,), - ) - try: - await self._complete_pending_submission_cleanup(session.session_id) - except Exception as cleanup_exc: - logger.warning( - "duplex append compensation remains pending for session %s: %s", - session.session_id, - cleanup_exc, - ) - raise - return [ - { - "stage_id": stage_id, - "replica_id": submission_result.replica_id, - "result": { - "supported": True, - "implementation_level": session.capabilities.implementation_level, - "data_plane_append": True, - "request_id": request_id, - "response_stage_id": request_context.final_stage_id, - "seq": update.seq, - "turn_id": update.turn_id, - "response_seq": message.fence.response_seq, - "turn_seq": update.turn_seq, - "mode": mode.value, - "resumable": True, - }, - } - ] - - async def handle_signal(self, message: SignalDuplexTurnMessage) -> None: - try: - cancel_events = {"barge_in", "input.cancel", "response.cancel"} - if message.event not in {*cancel_events, "session.update"}: - raise ValueError(f"unsupported duplex runtime signal: {message.event}") - session = self.sessions.require(message.session_id) - await self._complete_pending_submission_cleanup(message.session_id) - effective_next_fence = message.next_fence - if message.event in cancel_events: - if effective_next_fence is None: - raise ValueError(f"{message.event} requires next_fence") - cleanup_key = self._cleanup_key("cancel", message.fence) - pending = self._pending_control_cleanups.get(cleanup_key) - if pending is None: - stale_request_ids = session.prepare_cancel_fence(message.fence, effective_next_fence) - pending = _PendingControlCleanup( - kind="cancel", - session_id=message.session_id, - fence=message.fence, - submitted_request_ids=tuple(stale_request_ids), - ) - self._pending_control_cleanups[cleanup_key] = pending - await self._complete_control_cleanup(cleanup_key, pending) - else: - session.accept_fence(message.fence) - if message.runtime_config is not None: - self.sampling_params_for_config(message.runtime_config) - session.replace_configs( - session_config=message.session_config, - runtime_config=message.runtime_config, - ) - session.touch(session.fence, DuplexLeaseActivity.SIGNAL) - await self.put_result( - message.control_id, - fence=message.fence, - operation="signal", - session_id=message.session_id, - stage_results=[ - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": True, - "data_plane_signal": True, - "event": message.event, - "fence": message.fence, - "next_fence": effective_next_fence, - }, - } - ], - ) - except Exception as exc: - logger.exception("signal_duplex_turn failed: %s", exc) - await self.put_result( - message.control_id, - fence=message.fence, - operation="signal", - session_id=message.session_id, - stage_results=[], - error=exc, - ) - - async def handle_close(self, message: CloseDuplexSessionMessage) -> None: - try: - await self._complete_pending_submission_cleanup(message.session_id) - session = self.sessions.get(message.session_id) - if session is None: - await self.put_result( - message.control_id, - fence=message.fence, - operation="close", - session_id=message.session_id, - stage_results=[], - ) - return - cleanup_key = self._cleanup_key("close", message.fence) - pending = self._pending_control_cleanups.get(cleanup_key) - if pending is None: - submitted = tuple(session.resource_request_ids(submitted=True)) - reserved = tuple(session.resource_request_ids(submitted=False)) - self.sessions.begin_close_session(message.fence, reason=message.reason) - pending = _PendingControlCleanup( - kind="close", - session_id=message.session_id, - fence=message.fence, - submitted_request_ids=submitted, - reserved_request_ids=reserved, - ) - self._pending_control_cleanups[cleanup_key] = pending - await self._complete_control_cleanup(cleanup_key, pending) - await self.put_result( - message.control_id, - fence=message.fence, - operation="close", - session_id=message.session_id, - stage_results=[ - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": True, - "data_plane_close": True, - "reason": message.reason, - }, - } - ], - ) - except Exception as exc: - logger.exception("close_duplex_session failed: %s", exc) - await self.put_result( - message.control_id, - fence=message.fence, - operation="close", - session_id=message.session_id, - stage_results=[], - error=exc, - ) - - async def handle_touch(self, message: TouchDuplexSessionMessage) -> None: - try: - session = self.sessions.require(message.session_id) - activity = DuplexLeaseActivity(message.activity) - if activity is DuplexLeaseActivity.DETACH: - session.detach(message.fence) - else: - session.touch(message.fence, activity) - await self.put_result( - message.control_id, - fence=message.fence, - operation="touch", - session_id=message.session_id, - stage_results=[ - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": True, - "activity": activity.value, - "lease_generation": session.lease.generation, - }, - } - ], - ) - except Exception as exc: - logger.exception("touch_duplex_session failed: %s", exc) - await self.put_result( - message.control_id, - fence=message.fence, - operation="touch", - session_id=message.session_id, - stage_results=[], - error=exc, - ) - - async def handle_resume(self, message: ResumeDuplexSessionMessage) -> None: - try: - session = self.sessions.require(message.session_id) - generation = session.resume( - message.fence, - expected_lease_generation=message.expected_lease_generation, - ) - await self.put_result( - message.control_id, - fence=message.fence, - operation="resume", - session_id=message.session_id, - stage_results=[ - { - "stage_id": -1, - "replica_id": -1, - "result": { - "supported": True, - "lease_generation": generation, - }, - } - ], - ) - except Exception as exc: - logger.exception("resume_duplex_session failed: %s", exc) - await self.put_result( - message.control_id, - fence=message.fence, - operation="resume", - session_id=message.session_id, - stage_results=[], - error=exc, - ) - - async def reap_expired(self, now: float | None = None) -> int: - completed = 0 - for session_id in list(self._pending_submission_cleanups): - try: - await self._complete_pending_submission_cleanup(session_id) - except Exception as exc: - logger.warning( - "duplex append compensation remains pending for session %s: %s", - session_id, - exc, - ) - for key, pending in list(self._pending_control_cleanups.items()): - try: - await self._complete_control_cleanup(key, pending) - except Exception as exc: - logger.warning( - "duplex %s cleanup remains pending for session %s: %s", - pending.kind, - pending.session_id, - exc, - ) - for key, pending in list(self._pending_request_cleanups.items()): - if key in self._request_cleanups_in_progress: - continue - try: - await self._complete_request_cleanup(key, pending) - except Exception as exc: - logger.warning( - "duplex request cleanup remains pending for session %s: %s", - pending.session_id, - exc, - ) - continue - completed += 1 - for item in self.sessions.collect_expired( - now, - excluded_session_ids=set(self._pending_submission_cleanups), - ): - self._pending_expirations.setdefault( - (item.session_id, item.lease_generation), - item, - ) - for key, item in list(self._pending_expirations.items()): - try: - if item.submitted_request_ids: - await self._stage_port.cleanup(list(item.submitted_request_ids), abort=True) - if item.reserved_request_ids: - await self._stage_port.cleanup(list(item.reserved_request_ids)) - if self._lifecycle_sink is not None: - await self._lifecycle_sink.put( - DuplexSessionLifecycleMessage( - fence=item.fence, - session_id=item.session_id, - event="expired", - reason=item.reason, - lease_generation=item.lease_generation, - submitted_request_ids=list(item.submitted_request_ids), - reserved_request_ids=list(item.reserved_request_ids), - ) - ) - except Exception as exc: - logger.warning( - "duplex expiry cleanup remains pending for session %s: %s", - item.session_id, - exc, - ) - continue - session = self.sessions.get(item.session_id) - if session is not None and session.lease.generation == item.lease_generation: - self.sessions.finalize_close_session(session) - self._pending_expirations.pop(key, None) - completed += 1 - return completed - - async def _complete_pending_submission_cleanup(self, session_id: str) -> None: - pending = self._pending_submission_cleanups.get(session_id) - if pending is None: - return - task = self._submission_cleanup_tasks.get(session_id) - if task is not None and task.done(): - self._submission_cleanup_tasks.pop(session_id, None) - task = None - if task is None: - task = asyncio.create_task( - self._run_submission_cleanup(pending), - name=f"duplex-append-compensation-{session_id}", - ) - self._submission_cleanup_tasks[session_id] = task - - def discard(completed: asyncio.Task[None]) -> None: - if self._submission_cleanup_tasks.get(session_id) is completed: - self._submission_cleanup_tasks.pop(session_id, None) - - task.add_done_callback(discard) - await asyncio.shield(task) - - async def _run_submission_cleanup(self, pending: _PendingSubmissionCleanup) -> None: - await self._stage_port.cleanup(list(pending.request_ids), abort=True) - session = self.sessions.get(pending.session_id) - if session is not None: - session.release_request_ids(list(pending.request_ids)) - if self._pending_submission_cleanups.get(pending.session_id) is pending: - self._pending_submission_cleanups.pop(pending.session_id, None) - - async def _complete_request_cleanup( - self, - key: tuple[str, int, int], - pending: _PendingRequestCleanup, - ) -> None: - task = self._request_cleanup_tasks.get(key) - if task is not None and task.done(): - self._request_cleanup_tasks.pop(key, None) - task = None - if task is None: - task = asyncio.create_task( - self._run_request_cleanup(key, pending), - name=f"duplex-request-cleanup-{pending.session_id}", - ) - self._request_cleanup_tasks[key] = task - - def discard(completed: asyncio.Task[None]) -> None: - if self._request_cleanup_tasks.get(key) is completed: - self._request_cleanup_tasks.pop(key, None) - - task.add_done_callback(discard) - await asyncio.shield(task) - - async def _run_request_cleanup( - self, - key: tuple[str, int, int], - pending: _PendingRequestCleanup, - ) -> None: - await self._stage_port.cleanup(list(pending.request_ids), abort=pending.abort) - session = self.sessions.get(pending.session_id) - if ( - session is not None - and session.fence.incarnation == pending.fence.incarnation - and session.lease.generation == pending.lease_generation - ): - self.sessions.finalize_close_session(session) - if self._pending_request_cleanups.get(key) is pending: - self._pending_request_cleanups.pop(key, None) - self._request_cleanups_in_progress.discard(key) - - @staticmethod - def _cleanup_key(kind: str, fence: DuplexFence) -> tuple[str, str, int, int, int, int]: - return (kind, fence.session_id, fence.incarnation, fence.epoch, fence.turn_id, fence.response_seq) - - async def _complete_control_cleanup( - self, - key: tuple[str, str, int, int, int, int], - pending: _PendingControlCleanup, - ) -> None: - task = self._control_cleanup_tasks.get(key) - if task is not None and task.done(): - self._control_cleanup_tasks.pop(key, None) - task = None - if task is None: - task = asyncio.create_task( - self._run_control_cleanup(key, pending), - name=f"duplex-{pending.kind}-cleanup-{pending.session_id}", - ) - self._control_cleanup_tasks[key] = task - - def discard(completed: asyncio.Task[None]) -> None: - if self._control_cleanup_tasks.get(key) is completed: - self._control_cleanup_tasks.pop(key, None) - - task.add_done_callback(discard) - await asyncio.shield(task) - - async def _run_control_cleanup( - self, - key: tuple[str, str, int, int, int, int], - pending: _PendingControlCleanup, - ) -> None: - if pending.submitted_request_ids: - await self._stage_port.cleanup(list(pending.submitted_request_ids), abort=True) - if pending.reserved_request_ids: - await self._stage_port.cleanup(list(pending.reserved_request_ids)) - session = self.sessions.get(pending.session_id) - if session is not None: - if pending.kind == "cancel": - session.release_fence(pending.fence) - elif pending.kind in {"close", "open_rollback"}: - self.sessions.finalize_close_session(session) - self._pending_control_cleanups.pop(key, None) - - @classmethod - def _iter_result_dicts(cls, result: object): - if isinstance(result, dict): - yield result - elif isinstance(result, list | tuple): - for item in result: - yield from cls._iter_result_dicts(item) - - @classmethod - def _result_counts(cls, stage_results: list[dict[str, object]]) -> tuple[int, int]: - unsupported_count = 0 - error_count = 0 - for item in stage_results: - for result in cls._iter_result_dicts(item.get("result")): - if result.get("supported") is False: - unsupported_count += 1 - if result.get("error"): - error_count += 1 - return unsupported_count, error_count - - async def put_result( - self, - control_id: str, - *, - fence: DuplexFence, - operation: str, - session_id: str, - stage_results: list[dict[str, object]], - error: BaseException | str | None = None, - ) -> None: - control_error = self._control_error(error) if error is not None else None - if error is not None: - stage_results = [ - { - "stage_id": -1, - "replica_id": -1, - "result": {"supported": False, "error": control_error.message}, - } - ] - unsupported_count, error_count = self._result_counts(stage_results) - session = self.sessions.get(session_id) - await self._result_sink.put( - DuplexControlResultMessage( - control_id=control_id, - fence=fence, - operation=operation, - session_id=session_id, - ok=error_count == 0 and unsupported_count == 0, - stage_results=stage_results, - unsupported_count=unsupported_count, - error_count=error_count, - error=control_error, - accepted_fence=session.fence if session is not None else None, - lease_generation=session.lease.generation if session is not None else None, - ) - ) - - @staticmethod - def _control_error(error: BaseException | str) -> DuplexControlError: - message = str(error) - if isinstance(error, DuplexFenceMismatchError): - code = "stale_fence" - retryable = False - elif "unknown duplex input mode" in message: - code = "invalid_capability" - retryable = False - elif "duplex_session_capacity_exhausted" in message: - code = "resource_exhausted" - retryable = True - elif isinstance(error, KeyError): - code = "not_found" - retryable = False - elif isinstance(error, (TypeError, ValueError)): - code = "invalid_argument" - retryable = False - elif isinstance(error, TimeoutError): - code = "timeout" - retryable = True - else: - code = "failed_precondition" - retryable = False - return DuplexControlError(code=code, message=message, retryable=retryable) - - def decide_output( - self, - stage_id: int, - output: object, - context: DuplexOutputContext | None, - ) -> DuplexOutputDecision | None: - if context is None or self._extension is None: - return None - session = self._sessions.get(context.identity.session_id) - if session is None: - return None - try: - session.touch(context.identity.fence, DuplexLeaseActivity.MODEL_OUTPUT) - except RuntimeError: - return None - decision = self._extension.decide_output( - stage_id=stage_id, - final_stage_id=context.final_stage_id, - segment_finished=context.segment_finished, - segment_token_ids=context.segment_token_ids, - segment_output_metadata=dict(context.segment_output_metadata), - output=output, - ) - if decision is not None and not isinstance(decision, DuplexOutputDecision): - raise TypeError("duplex runtime extension decide_output() must return DuplexOutputDecision or None") - return decision - - def session_for_identity(self, identity: DuplexRequestIdentity | None) -> DuplexSessionRuntimeState | None: - if identity is None: - return None - return self._sessions.get(identity.session_id) - - def close_sessions_for_request_ids( - self, - request_ids: list[str], - *, - abort: bool = False, - cleanup_in_progress: bool = False, - ) -> dict[str, list[str]]: - closed = self._sessions.close_sessions_for_request_ids(request_ids) - for session_id, stale_request_ids in closed.items(): - session = self._sessions.get(session_id) - if session is None: - continue - key = (session_id, session.fence.incarnation, session.lease.generation) - existing = self._pending_request_cleanups.get(key) - merged_request_ids = tuple( - dict.fromkeys( - [ - *(existing.request_ids if existing is not None else ()), - *stale_request_ids, - ] - ) - ) - self._pending_request_cleanups[key] = _PendingRequestCleanup( - session_id=session_id, - fence=session.fence, - lease_generation=session.lease.generation, - request_ids=merged_request_ids, - abort=abort or (existing.abort if existing is not None else False), - ) - if cleanup_in_progress: - self._request_cleanups_in_progress.add(key) - return closed - - def defer_request_cleanups(self, session_ids: Iterable[str]) -> None: - session_id_set = set(session_ids) - active_keys = {key for key in self._request_cleanups_in_progress if key[0] in session_id_set} - self._request_cleanups_in_progress.difference_update(active_keys) - - def finalize_closed_sessions(self, session_ids: Iterable[str]) -> None: - session_id_set = set(session_ids) - for key in list(self._pending_request_cleanups): - if key[0] not in session_id_set: - continue - self._pending_request_cleanups.pop(key, None) - self._request_cleanups_in_progress.discard(key) - self._sessions.finalize_closed_sessions(session_ids) - - -__all__ = [ - "DuplexControlPlane", - "DuplexOutputContext", - "DuplexRequestIdentity", - "DuplexResultSink", - "DuplexLifecycleSink", - "DuplexStagePort", - "DuplexStageRequestContext", - "DuplexStageSubmission", - "DuplexStageSubmissionResult", -] diff --git a/vllm_omni/engine/duplex/events.py b/vllm_omni/engine/duplex/events.py new file mode 100644 index 000000000..9c131a893 --- /dev/null +++ b/vllm_omni/engine/duplex/events.py @@ -0,0 +1,756 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Typed events emitted by a duplex session. + +Every public event is a frozen dataclass with explicit fields. The Realtime +wire JSON is *derived* from those fields by :meth:`DuplexEvent.to_realtime`, +which is pure: it never consults session state. The stateful part (response / +item ids, content-part bookkeeping) lives in +``vllm_omni.engine.duplex.realtime_events`` on the session runner and is +consumed when the events are *constructed*. + +``session_id`` / ``epoch`` are bound by +``DuplexSessionManager.emit`` when an event leaves the runner; producers +(projector, serving handshake helpers) build events with the defaults. +""" + +from __future__ import annotations + +import base64 +from collections.abc import Mapping +from dataclasses import dataclass, field, fields +from typing import ClassVar, cast +from uuid import uuid4 + + +def new_event_id() -> str: + return f"event_{uuid4().hex}" + + +#: OpenAI Realtime ``error.type`` for each internal error code. +REALTIME_ERROR_TYPES_BY_CODE: dict[str, str] = { + "bad_event": "invalid_request_error", + "bad_audio": "invalid_request_error", + "config_timeout": "invalid_request_error", + "invalid_json": "invalid_request_error", + "event_too_large": "invalid_request_error", + "unknown_event": "invalid_request_error", + "internal_error": "server_error", + "runtime_append_failed": "server_error", + "runtime_append_task_failed": "server_error", + "runtime_signal_failed": "server_error", + "runtime_abort_failed": "server_error", + "runtime_data_plane_stream_failed": "server_error", + "runtime_data_plane_text_without_audio": "server_error", + "resource_exhausted": "rate_limit_error", + "session_exists": "invalid_request_error", + "session_closed": "invalid_request_error", + "unknown_session": "invalid_request_error", + "invalid_duplex_runtime_config": "invalid_request_error", + "instructions_update_unsupported": "invalid_request_error", + "persona_update_unsupported": "invalid_request_error", + "voice_update_unsupported": "invalid_request_error", + "unsupported_nemotron_duplex_mode": "invalid_request_error", + "unsupported_native_response_options": "invalid_request_error", + "runtime_touch_failed": "server_error", + "engine_error": "server_error", + "input_backpressure": "rate_limit_error", + "response_already_active": "invalid_request_error", + "response_not_active": "invalid_request_error", + "response_create_without_input": "invalid_request_error", + "text_only_turn_unsupported": "invalid_request_error", + "commit_aborted": "server_error", + "input_audio_buffer_empty": "invalid_request_error", + "missing_item_id": "invalid_request_error", + "item_not_found": "invalid_request_error", + "playback_item_mismatch": "invalid_request_error", + "playback_item_not_found": "invalid_request_error", + "playback_ack_too_late": "invalid_request_error", + "unsupported_audio_format": "invalid_request_error", + "unsupported_turn_detection": "invalid_request_error", + "unsupported_ref_audio_path": "invalid_request_error", + "ref_audio_required": "invalid_request_error", + "model_update_unsupported": "invalid_request_error", + "voice_update_after_audio_unsupported": "invalid_request_error", + "ref_audio_update_unsupported": "invalid_request_error", + "native_text_append_unsupported": "invalid_request_error", + "invalid_video_frames": "invalid_request_error", + "invalid_function_call_output": "invalid_request_error", + "server_vad_unavailable": "server_error", +} + +_IDENTITY_FIELDS = frozenset({"session_id", "epoch", "event_id"}) + + +def _wire_value(value: object) -> object: + if isinstance(value, Mapping): + return {str(k): _wire_value(v) for k, v in value.items()} + if isinstance(value, tuple | list): + return [_wire_value(v) for v in value] + return value + + +@dataclass(frozen=True, slots=True, kw_only=True) +class DuplexEvent: + """Base of every public session event.""" + + #: Realtime ``type`` this class renders to. + wire_type: ClassVar[str] = "" + #: Fields that are omitted from the wire object when ``None``. + optional_wire_fields: ClassVar[frozenset[str]] = frozenset() + + session_id: str = "" + epoch: int | None = None + #: Server event id (OpenAI ``event_id``); generated at construction. + event_id: str = field(default_factory=new_event_id) + + # ---- identity ---- + + @property + def type(self) -> str: + return self.wire_type + + @property + def is_terminal(self) -> bool: + return False + + # ---- generic accessors (None when the class has no such field) ---- + + @property + def response_id(self) -> str | None: + return None + + @property + def item_id(self) -> str | None: + return None + + @property + def text(self) -> str | None: + return None + + @property + def audio(self) -> bytes | None: + return None + + # ---- wire rendering (pure) ---- + + def _wire_fields(self) -> dict[str, object]: + data: dict[str, object] = {} + for f in fields(self): + if f.name in _IDENTITY_FIELDS: + continue + value = getattr(self, f.name) + if value is None and f.name in self.optional_wire_fields: + continue + data[f.name] = _wire_value(value) + return data + + def to_realtime(self) -> dict[str, object]: + """The Realtime wire JSON object for this event (derived from the fields).""" + return {"type": self.wire_type, "event_id": self.event_id, **self._wire_fields()} + + def __repr__(self) -> str: # pragma: no cover - debugging aid + return f"{type(self).__name__}(type={self.wire_type!r}, session_id={self.session_id!r}, epoch={self.epoch!r})" + + +# ---- session lifecycle ---- + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionCreated(DuplexEvent): + wire_type = "session.created" + optional_wire_fields = frozenset({"attachment_generation", "resume_token"}) + + session: Mapping[str, object] = field(default_factory=dict) + #: Transport resume credentials (set by the websocket handler when resume is supported). + attachment_generation: int | None = None + resume_token: str | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionUpdated(DuplexEvent): + wire_type = "session.updated" + + session: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionResumed(SessionCreated): + wire_type = "session.resumed" + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, **DuplexEvent._wire_fields(self)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionReplaced(DuplexEvent): + """A newer connection took the session over (sent to the replaced socket).""" + + wire_type = "session.replaced" + + attachment_generation: int = 0 + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, **DuplexEvent._wire_fields(self)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionResyncRequired(DuplexEvent): + """The replay journal cannot bridge the gap; the client must start over.""" + + wire_type = "session.resync_required" + + reason: str = "journal_gap" + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, **DuplexEvent._wire_fields(self)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionClosed(DuplexEvent): + wire_type = "session.closed" + + reason: str = "closed" + #: The session-internal close event (kept for clients that inspect it). + details: Mapping[str, object] = field(default_factory=dict) + + @property + def is_terminal(self) -> bool: + return True + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, "reason": self.reason, "event": _wire_value(self.details)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionExpired(SessionClosed): + wire_type = "session.expired" + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, "reason": self.reason} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SessionHeartbeatAck(DuplexEvent): + wire_type = "session.heartbeat_ack" + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TurnEvent(DuplexEvent): + """Turn-state transition (``turn.event``).""" + + wire_type = "turn.event" + + event: str = "" + turn_state: str = "" + + def _wire_fields(self) -> dict[str, object]: + return {"session_id": self.session_id, "event": self.event, "turn_state": self.turn_state, "epoch": self.epoch} + + +# ---- response lifecycle ---- + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _ResponseEvent(DuplexEvent): + """Events addressed to one response (and usually one output item / content part).""" + + response_id: str | None = None # type: ignore[assignment] + item_id: str | None = None # type: ignore[assignment] + output_index: int = 0 + content_index: int = 0 + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ResponseCreated(DuplexEvent): + wire_type = "response.created" + + response_id: str | None = None # type: ignore[assignment] + response: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Listen(DuplexEvent): + """The model decided to keep listening (no spoken response for this turn).""" + + wire_type = "response.listen" + optional_wire_fields = frozenset({"response_id"}) + + response_id: str | None = None # type: ignore[assignment] + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + response: dict[str, object] = { + "object": "realtime.response", + "status": "listening", + "metadata": _wire_value(self.details), + } + data: dict[str, object] = {"session_id": self.session_id, "epoch": self.epoch, "response": response} + if self.response_id: + response["id"] = self.response_id + data["response_id"] = self.response_id + return data + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Speak(_ResponseEvent): + wire_type = "response.speak" + + metadata: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class OverlapDecision(DuplexEvent): + wire_type = "overlap.decision" + + policy: str | None = None + action: str | None = None + reason: str | None = None + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + return { + "session_id": self.session_id, + "epoch": self.epoch, + "policy": self.policy, + "action": self.action, + "reason": self.reason, + "metadata": _wire_value(self.details), + } + + +@dataclass(frozen=True, slots=True, kw_only=True) +class OutputItemAdded(DuplexEvent): + wire_type = "response.output_item.added" + + response_id: str | None = None # type: ignore[assignment] + output_index: int = 0 + item: Mapping[str, object] = field(default_factory=dict) + + @property + def item_id(self) -> str | None: + value = self.item.get("id") + return value if isinstance(value, str) else None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class OutputItemDone(OutputItemAdded): + wire_type = "response.output_item.done" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ContentPartAdded(_ResponseEvent): + wire_type = "response.content_part.added" + + part: Mapping[str, object] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ContentPartDone(ContentPartAdded): + wire_type = "response.content_part.done" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AudioDelta(_ResponseEvent): + wire_type = "response.output_audio.delta" + optional_wire_fields = frozenset({"sample_rate_hz", "metadata"}) + + #: Base64 audio in ``format``. + delta: str = "" + format: str = "pcm16" + sample_rate_hz: int | None = None + metadata: Mapping[str, object] | None = None + + @property + def audio(self) -> bytes | None: + try: + return base64.b64decode(self.delta) + except (ValueError, TypeError): + return None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class AudioDone(_ResponseEvent): + wire_type = "response.output_audio.done" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TranscriptDelta(_ResponseEvent): + wire_type = "response.output_audio_transcript.delta" + + delta: str = "" + + @property + def text(self) -> str | None: + return self.delta + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TranscriptDone(_ResponseEvent): + wire_type = "response.output_audio_transcript.done" + + transcript: str = "" + + @property + def text(self) -> str | None: + return self.transcript + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TextDelta(_ResponseEvent): + wire_type = "response.output_text.delta" + + delta: str = "" + + @property + def text(self) -> str | None: + return self.delta + + +@dataclass(frozen=True, slots=True, kw_only=True) +class TextDone(_ResponseEvent): + wire_type = "response.output_text.done" + + text: str = "" # type: ignore[assignment] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class FunctionCallArgumentsDelta(DuplexEvent): + wire_type = "response.function_call_arguments.delta" + + response_id: str | None = None # type: ignore[assignment] + item_id: str | None = None # type: ignore[assignment] + output_index: int = 0 + call_id: str = "" + delta: str = "" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class FunctionCallArgumentsDone(DuplexEvent): + wire_type = "response.function_call_arguments.done" + + response_id: str | None = None # type: ignore[assignment] + item_id: str | None = None # type: ignore[assignment] + output_index: int = 0 + call_id: str = "" + arguments: str = "" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ResponseDone(DuplexEvent): + wire_type = "response.done" + + response_id: str | None = None # type: ignore[assignment] + response: Mapping[str, object] = field(default_factory=dict) + + @property + def status(self) -> str | None: + value = self.response.get("status") + return value if isinstance(value, str) else None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class RateLimitsUpdated(DuplexEvent): + wire_type = "rate_limits.updated" + + rate_limits: tuple[Mapping[str, object], ...] = () + + +# ---- input buffer / conversation ---- + + +@dataclass(frozen=True, slots=True, kw_only=True) +class InputCommitted(DuplexEvent): + wire_type = "input_audio_buffer.committed" + + previous_item_id: str | None = None + item_id: str | None = None # type: ignore[assignment] + #: The session-internal ``input.committed`` event. + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + return { + "previous_item_id": self.previous_item_id, + "item_id": self.item_id, + "event": _wire_value(self.details), + } + + +@dataclass(frozen=True, slots=True, kw_only=True) +class InputCleared(DuplexEvent): + wire_type = "input_audio_buffer.cleared" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SpeechStarted(DuplexEvent): + wire_type = "input_audio_buffer.speech_started" + + audio_start_ms: int = 0 + item_id: str | None = None # type: ignore[assignment] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SpeechStopped(DuplexEvent): + wire_type = "input_audio_buffer.speech_stopped" + + audio_end_ms: int = 0 + item_id: str | None = None # type: ignore[assignment] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _ItemEvent(DuplexEvent): + previous_item_id: str | None = None + item: Mapping[str, object] = field(default_factory=dict) + + @property + def item_id(self) -> str | None: + value = self.item.get("id") + return value if isinstance(value, str) else None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemAdded(_ItemEvent): + wire_type = "conversation.item.added" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemCreated(_ItemEvent): + wire_type = "conversation.item.created" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemDone(_ItemEvent): + wire_type = "conversation.item.done" + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemRetrieved(DuplexEvent): + wire_type = "conversation.item.retrieved" + + item: Mapping[str, object] = field(default_factory=dict) + + @property + def item_id(self) -> str | None: + value = self.item.get("id") + return value if isinstance(value, str) else None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemDeleted(DuplexEvent): + wire_type = "conversation.item.deleted" + + item_id: str | None = None # type: ignore[assignment] + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + return {"item_id": self.item_id, "event": _wire_value(self.details)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ItemTruncated(DuplexEvent): + wire_type = "conversation.item.truncated" + + item_id: str | None = None # type: ignore[assignment] + content_index: int = 0 + audio_end_ms: int = 0 + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + return { + "item_id": self.item_id, + "content_index": self.content_index, + "audio_end_ms": self.audio_end_ms, + "event": _wire_value(self.details), + } + + +@dataclass(frozen=True, slots=True, kw_only=True) +class InputTranscriptionCompleted(DuplexEvent): + wire_type = "conversation.item.input_audio_transcription.completed" + + item_id: str | None = None # type: ignore[assignment] + content_index: int = 0 + transcript: str = "" + + @property + def text(self) -> str | None: + return self.transcript + + +# ---- playback / control ---- + + +@dataclass(frozen=True, slots=True, kw_only=True) +class OutputAudioCleared(DuplexEvent): + wire_type = "output_audio_buffer.cleared" + + response_id: str | None = None # type: ignore[assignment] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class PlaybackAcknowledged(DuplexEvent): + wire_type = "playback.acknowledged" + + #: The session-internal ``playback.acknowledged`` event (cursor, committed_ms ...). + details: Mapping[str, object] = field(default_factory=dict) + + def _wire_fields(self) -> dict[str, object]: + return {"event": _wire_value(self.details)} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class ErrorEvent(DuplexEvent): + wire_type = "error" + + code: str = "internal_error" + message: str = "" + #: Client ``event_id`` this error answers (OpenAI ``error.event_id``). + related_event_id: str | None = None + param: str | None = None + #: Extra keys merged into the wire ``error`` object (``retryable`` ...). + extra: Mapping[str, object] = field(default_factory=dict) + + @property + def error_type(self) -> str: + return REALTIME_ERROR_TYPES_BY_CODE.get(self.code, "invalid_request_error") + + @property + def error(self) -> dict[str, object]: + error: dict[str, object] = {"type": self.error_type, "code": self.code, "message": self.message} + if self.related_event_id: + error["event_id"] = self.related_event_id + if self.param: + error["param"] = self.param + error.update(cast("Mapping[str, object]", _wire_value(self.extra))) + return error + + @property + def text(self) -> str | None: + return self.message + + def _wire_fields(self) -> dict[str, object]: + return {"error": self.error} + + +@dataclass(frozen=True, slots=True, kw_only=True) +class DuplexRawEvent(DuplexEvent): + """A session-internal event with no dedicated Realtime type (``duplex.``).""" + + wire_type = "duplex.raw" + + internal_type: str = "" + details: Mapping[str, object] = field(default_factory=dict) + + @property + def type(self) -> str: + return f"duplex.{self.internal_type}" + + def to_realtime(self) -> dict[str, object]: + return {"type": self.type, "event_id": self.event_id, "event": _wire_value(self.details)} + + +def error_event( + code: str, + message: str, + *, + event_id: object | None = None, + param: object | None = None, + extra: Mapping[str, object] | None = None, +) -> ErrorEvent: + """Build the error event for an internal code; ``event_id`` is the *client* event id.""" + return ErrorEvent( + code=code, + message=message, + related_event_id=event_id if isinstance(event_id, str) and event_id else None, + param=param if isinstance(param, str) and param else None, + extra=dict(extra or {}), + ) + + +#: Internal events that terminate a response/session and must never be dropped as stale. +DOMAIN_TERMINAL_EVENTS = frozenset( + { + "response.done", + "response.listen", + "audio.cancelled", + "input.cancelled", + "session.closed", + } +) + +#: Internal streaming model-output events subject to epoch/close stale filtering. +MODEL_OUTPUT_EVENTS = frozenset( + { + "response.created", + "response.listen", + "response.speak", + "response.output_item.added", + "response.content_part.added", + "response.output_audio.delta", + "response.output_audio.done", + "response.output_text.delta", + "response.output_text.done", + "response.text.delta", + "response.text.done", + "response.message", + "response.output_item.done", + "response.content_part.done", + "response.done", + } +) + + +__all__ = [ + "DOMAIN_TERMINAL_EVENTS", + "MODEL_OUTPUT_EVENTS", + "REALTIME_ERROR_TYPES_BY_CODE", + "AudioDelta", + "AudioDone", + "ContentPartAdded", + "ContentPartDone", + "DuplexEvent", + "DuplexRawEvent", + "ErrorEvent", + "FunctionCallArgumentsDelta", + "FunctionCallArgumentsDone", + "InputCleared", + "InputCommitted", + "InputTranscriptionCompleted", + "ItemAdded", + "ItemCreated", + "ItemDeleted", + "ItemDone", + "ItemRetrieved", + "ItemTruncated", + "Listen", + "OutputAudioCleared", + "OutputItemAdded", + "OutputItemDone", + "OverlapDecision", + "PlaybackAcknowledged", + "RateLimitsUpdated", + "ResponseCreated", + "ResponseDone", + "SessionClosed", + "SessionCreated", + "SessionExpired", + "SessionHeartbeatAck", + "SessionReplaced", + "SessionResumed", + "SessionResyncRequired", + "SessionUpdated", + "Speak", + "SpeechStarted", + "SpeechStopped", + "TextDelta", + "TextDone", + "TranscriptDelta", + "TranscriptDone", + "TurnEvent", + "error_event", + "new_event_id", +] diff --git a/vllm_omni/engine/duplex/messages.py b/vllm_omni/engine/duplex/messages.py index 2c49d573f..ad264e0c5 100644 --- a/vllm_omni/engine/duplex/messages.py +++ b/vllm_omni/engine/duplex/messages.py @@ -1,130 +1,114 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Queue contracts owned by the experimental duplex control plane.""" +"""Queue envelopes exchanged between ``DuplexOmniEngine`` and ``DuplexOrchestrator``. + +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``. +""" from __future__ import annotations -from dataclasses import dataclass -from typing import Literal +from typing import TYPE_CHECKING, Literal 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.events import DuplexEvent -@dataclass(frozen=True, slots=True) -class DuplexFence: - """Identity fence carried by experimental duplex control messages.""" - session_id: str - epoch: int = 0 - turn_id: int = 0 - response_seq: int = 0 - incarnation: int = 0 +class DuplexSessionError(RuntimeError): + """A session control operation was rejected by the engine.""" + + def __init__( + self, + message: str, + *, + code: str = "internal_error", + retryable: bool = False, + session_id: str | None = None, + ) -> None: + super().__init__(message) + self.code = code + self.retryable = retryable + self.session_id = session_id class OpenDuplexSessionMessage(EngineQueueMessage, kw_only=True): type: Literal["open_duplex_session"] = "open_duplex_session" control_id: str - fence: DuplexFence session_id: str - session_mode: str = "duplex" - capabilities: dict[str, object] - session_config: dict[str, object] | None = None - runtime_config: dict[str, object] | None = None - - -class AppendDuplexInputMessage(EngineQueueMessage, kw_only=True): - type: Literal["append_duplex_input"] = "append_duplex_input" - control_id: str - operation_id: str | None = None - fence: DuplexFence - session_id: str - expected_epoch: int | None = None - mode: str - payload: object - final: bool = False - - -class SignalDuplexTurnMessage(EngineQueueMessage, kw_only=True): - type: Literal["signal_duplex_turn"] = "signal_duplex_turn" - control_id: str - fence: DuplexFence - session_id: str - event: str - next_fence: DuplexFence | None = None - session_config: dict[str, object] | None = None - runtime_config: dict[str, object] | None = None + #: Already normalized by ``DuplexOmni``; the queue is in-process, so the object crosses as is. + session_config: DuplexSessionConfig class CloseDuplexSessionMessage(EngineQueueMessage, kw_only=True): type: Literal["close_duplex_session"] = "close_duplex_session" control_id: str - fence: DuplexFence session_id: str reason: str = "client_close" -class TouchDuplexSessionMessage(EngineQueueMessage, kw_only=True): - type: Literal["touch_duplex_session"] = "touch_duplex_session" - control_id: str - fence: DuplexFence - session_id: str - activity: str - - class ResumeDuplexSessionMessage(EngineQueueMessage, kw_only=True): type: Literal["resume_duplex_session"] = "resume_duplex_session" control_id: str - fence: DuplexFence session_id: str expected_lease_generation: int -class DuplexSessionLifecycleMessage(EngineQueueMessage, kw_only=True): - type: Literal["duplex_session_lifecycle"] = "duplex_session_lifecycle" - fence: DuplexFence +class TouchDuplexSessionMessage(EngineQueueMessage, kw_only=True): + type: Literal["touch_duplex_session"] = "touch_duplex_session" + control_id: str session_id: str - event: str - reason: str - lease_generation: int - submitted_request_ids: list[str] - reserved_request_ids: list[str] + activity: str -class DuplexControlError(EngineQueueMessage, kw_only=True): - code: str - message: str - retryable: bool = False +class DuplexSessionCommandMessage(EngineQueueMessage, kw_only=True): + """One-way session command; rejections come back as ``ErrorEvent``.""" + + type: Literal["duplex_session_command"] = "duplex_session_command" + session_id: str + command: DuplexCommand class DuplexControlResultMessage(EngineQueueMessage, kw_only=True): type: Literal["duplex_control_result"] = "duplex_control_result" control_id: str - fence: DuplexFence operation: str session_id: str ok: bool - stage_results: list[dict[str, object]] - unsupported_count: int = 0 - error_count: int = 0 - error: DuplexControlError | None = None - accepted_fence: DuplexFence | None = None lease_generation: int | None = None + capabilities: DuplexCapabilities | None = None + #: The Realtime ``session`` object of the session (wire payload, already a dict). + public_session: dict[str, object] | None = None + error_code: str | None = None + error_message: str | None = None + error_retryable: bool = False @property def rpc_correlation_key(self) -> tuple[str, str]: return ("duplex", self.control_id) +class DuplexSessionEventMessage(EngineQueueMessage, kw_only=True): + """A typed session event on the engine output queue.""" + + type: Literal["duplex_session_event"] = "duplex_session_event" + session_id: str + event: DuplexEvent + + __all__ = [ - "AppendDuplexInputMessage", "CloseDuplexSessionMessage", - "DuplexControlError", "DuplexControlResultMessage", - "DuplexFence", - "DuplexSessionLifecycleMessage", + "DuplexSessionCommandMessage", + "DuplexSessionError", + "DuplexSessionEventMessage", "OpenDuplexSessionMessage", "ResumeDuplexSessionMessage", - "SignalDuplexTurnMessage", "TouchDuplexSessionMessage", ] diff --git a/vllm_omni/engine/duplex/plugin.py b/vllm_omni/engine/duplex/plugin.py new file mode 100644 index 000000000..130d5204d --- /dev/null +++ b/vllm_omni/engine/duplex/plugin.py @@ -0,0 +1,329 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Model plugin contract for full-duplex models. + +One ``DuplexModelPlugin`` subclass per model binds what used to be two dotted +paths (the engine ``DuplexRuntimeExtension`` and the serving +``ServingRuntimeAdapter``). Everything runs engine-side now, so the plugin is +loaded once by ``DuplexOmniEngine`` and handed to ``DuplexOrchestrator`` / +``DuplexSessionManager``. +""" + +from __future__ import annotations + +import asyncio +from abc import ABC, abstractmethod +from collections.abc import Callable, Iterable, Mapping +from importlib import import_module +from typing import TYPE_CHECKING + +from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig +from vllm_omni.engine.duplex.contracts import ( + DuplexAppendPlan, + DuplexFence, + DuplexOutputDecision, +) + +if TYPE_CHECKING: + from vllm.config import ModelConfig + + +class DuplexRuntimeConfigError(ValueError): + """A model plugin rejected client-visible runtime configuration.""" + + def __init__(self, message: str, *, code: str = "invalid_duplex_runtime_config") -> None: + super().__init__(message) + self.code = code + + +def reject_changed_runtime_value( + new_value: object, + current_value: object, + *, + message: str, + code: str, + error_cls: type[DuplexRuntimeConfigError] = DuplexRuntimeConfigError, +) -> None: + if new_value != current_value: + raise error_cls(message, code=code) + + +class PcmAppendReservation(ABC): + operation_id: str + payload: dict[str, object] | None + + @property + @abstractmethod + def active(self) -> bool: ... + + @property + @abstractmethod + def byte_count(self) -> int: ... + + @abstractmethod + def commit(self) -> None: ... + + @abstractmethod + def rollback(self) -> None: ... + + +class PcmAppendBuffer(ABC): + @property + @abstractmethod + def pending_byte_count(self) -> int: ... + + @abstractmethod + def clear(self) -> None: ... + + @abstractmethod + def clear_force_listen(self) -> None: ... + + @abstractmethod + def has_pending(self) -> bool: ... + + @abstractmethod + def has_reserved(self) -> bool: ... + + @abstractmethod + def prepare_append( + self, + payload: dict[str, object], + *, + operation_id: str, + chunk_period_ms: int, + allow_emit: bool, + ) -> PcmAppendReservation | None: ... + + @abstractmethod + def prepare_commit( + self, + *, + operation_id: str, + chunk_period_ms: int, + ) -> PcmAppendReservation: ... + + @abstractmethod + def flush(self, *, chunk_period_ms: int) -> dict[str, object] | None: ... + + +class DuplexModelSessionState(ABC): + """Model-owned per-session state; owned by the session runner (one per session).""" + + audio_buffer: PcmAppendBuffer + input_since_commit: bool + speech_since_commit: bool + context_locked: bool + committed_audio_payload: dict[str, object] | None + committed_audio_operation_id: str | None + committed_audio_reserved_bytes: int + deferred_response_create: bool + deferred_precreate_response: bool + continuation_owner_id: str | None + continuation_units: int + pending_silence_task: asyncio.Task[bool] | None + pending_silence_owner_id: str | None + + @abstractmethod + def retain_committed_audio( + self, + payload: dict[str, object], + *, + operation_id: str | None, + reserved_bytes: int = 0, + ) -> None: ... + + @abstractmethod + def clear_committed_audio(self) -> int: ... + + @abstractmethod + def clear_continuation(self) -> None: ... + + +class DuplexDataPlane(ABC): + """Projects raw stage outputs of one model into internal duplex events.""" + + @abstractmethod + def begin_request(self, request_id: str) -> None: ... + + @abstractmethod + def is_terminal(self, request_id: str | None) -> bool: ... + + @abstractmethod + def mark_terminal(self, request_id: str) -> None: ... + + @abstractmethod + def close_stream(self, request_id: str) -> None: ... + + @abstractmethod + def close_session(self, session_id: str, *, active_request_id: str | None = None) -> None: ... + + @abstractmethod + def project(self, result: object, *, context: object | None = None) -> Iterable[dict[str, object]]: ... + + +EncodeAudio = Callable[[object, int, str, float | None], str | None] + + +class DuplexModelPlugin(ABC): + """Everything vLLM-Omni needs to know about one full-duplex model. + + Engine policy (sampling params, append planning, output decisions) and + session policy (capabilities, runtime configuration, per-session state, + data-plane projection) live on the same object so a mismatch between the + two halves is impossible by construction. + """ + + plugin_id: str = "" + private_runtime_config_keys: frozenset[str] = frozenset() + #: Samples per silence unit the runner appends to keep a model turn going. + silence_continuation_samples: int = 16000 + data_plane: DuplexDataPlane + + def __init__(self, encode_audio: EncodeAudio) -> None: + # Constructor-only: concrete plugins hand the encoder to their data plane. + del encode_audio + + # ---- engine policy (was DuplexRuntimeExtension) ---- + + @abstractmethod + def configure_sampling_params( + self, + *, + runtime_config: dict[str, object], + defaults: tuple[object, ...], + ) -> tuple[object, ...]: ... + + @abstractmethod + def plan_append( + self, + *, + request_id: str, + fence: DuplexFence, + session_config: dict[str, object], + runtime_config: dict[str, object], + seq: int, + turn_seq: int, + payload: object, + final: bool, + sampling_params: object, + ) -> DuplexAppendPlan: ... + + @abstractmethod + def decide_output( + self, + *, + stage_id: int, + final_stage_id: int, + segment_finished: bool, + segment_token_ids: tuple[int, ...], + segment_output_metadata: dict[str, object], + output: object, + ) -> DuplexOutputDecision | None: ... + + # ---- session policy (was ServingRuntimeAdapter) ---- + + @abstractmethod + def create_session_state(self) -> DuplexModelSessionState: ... + + @abstractmethod + def capabilities(self, *, max_sessions: int) -> DuplexCapabilities: ... + + @abstractmethod + def validate_client_extra_body(self, extra_body: object) -> None: ... + + @abstractmethod + async def prepare_runtime_config( + self, config: DuplexSessionConfig, *, model_config: ModelConfig | None + ) -> dict[str, object]: ... + + @abstractmethod + def runtime_config_for_update( + self, + config: DuplexSessionConfig, + current: Mapping[str, object], + ) -> dict[str, object]: ... + + @abstractmethod + def data_plane_context( + self, + *, + epoch: int, + turn_id: int, + active_response_turn_id: int | None, + active_response_id: str | None, + auto_responds: bool, + response_format: str, + speed: float | None, + modalities: tuple[str, ...], + ) -> object: ... + + # Optional hook: build the runtime config patch for a function-call output + # item. Plugins without tools keep the default (no change). + def runtime_config_for_function_output( + self, + config: DuplexSessionConfig, + current: Mapping[str, object], + item: Mapping[str, object], + ) -> dict[str, object] | None: + del config, current, item + return None + + +def load_duplex_plugin(path: str, encode_audio: EncodeAudio) -> DuplexModelPlugin: + module_name, separator, attribute_name = path.rpartition(".") + if not separator: + raise ValueError(f"Invalid duplex plugin path: {path!r}") + plugin_type = getattr(import_module(module_name), attribute_name) + plugin = plugin_type(encode_audio) + if not isinstance(plugin, DuplexModelPlugin): + raise TypeError(f"{path!r} is not a DuplexModelPlugin") + if not plugin.plugin_id: + raise TypeError("Duplex plugin must declare plugin_id") + if not isinstance(getattr(plugin, "data_plane", None), DuplexDataPlane): + raise TypeError("Duplex plugin must declare a DuplexDataPlane as data_plane") + return plugin + + +def validate_duplex_plugin_sampling(plugin: DuplexModelPlugin, *, sampling_defaults: tuple[object, ...]) -> None: + """Fail fast when the plugin cannot produce one sampling parameter per stage.""" + configured = plugin.configure_sampling_params(runtime_config={}, defaults=sampling_defaults) + if not isinstance(configured, tuple): + raise TypeError("Duplex plugin must return sampling parameters as a tuple") + if len(configured) != len(sampling_defaults): + raise ValueError("Duplex plugin must return one sampling parameter per stage") + for stage_id, (value, default) in enumerate(zip(configured, sampling_defaults, strict=True)): + if default is not None and not isinstance(value, type(default)): + raise TypeError( + "Duplex plugin sampling parameter type mismatch " + f"for stage {stage_id}: expected {type(default).__name__}, got {type(value).__name__}" + ) + + +def payload_turn_id(payload: object) -> int | None: + if not isinstance(payload, Mapping): + return None + return coerce_int(payload.get("duplex_turn_id", payload.get("model_turn_id"))) + + +def coerce_int(value: object) -> int | None: + try: + return int(value) if value is not None else None + except (TypeError, ValueError): + return None + + +__all__ = [ + "DuplexDataPlane", + "DuplexModelPlugin", + "DuplexModelSessionState", + "DuplexRuntimeConfigError", + "EncodeAudio", + "PcmAppendBuffer", + "PcmAppendReservation", + "coerce_int", + "load_duplex_plugin", + "payload_turn_id", + "reject_changed_runtime_value", + "validate_duplex_plugin_sampling", +] diff --git a/vllm_omni/engine/duplex/realtime_commands.py b/vllm_omni/engine/duplex/realtime_commands.py new file mode 100644 index 000000000..370364017 --- /dev/null +++ b/vllm_omni/engine/duplex/realtime_commands.py @@ -0,0 +1,887 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OpenAI Realtime client events -> :class:`DuplexCommand` (stateless). + +This is the mapping half of the former ``entrypoints/duplex/realtime_input.py`` +translator. Everything that needed per-session state in the old translator +(input-buffer emptiness for commits, response-id fallbacks for cancels, +conversation-item lookups, VAD) is resolved by the session runner through the +helpers on :class:`~vllm_omni.engine.duplex.realtime_events.RealtimeProjectionState`; +the commands produced here carry the raw client intent only. +""" + +from __future__ import annotations + +import base64 +import binascii +from collections.abc import Mapping +from dataclasses import dataclass, replace +from typing import Any, cast +from uuid import uuid4 + +import numpy as np + +from vllm_omni.engine.duplex.audio import convert_input_audio_with_rate +from vllm_omni.engine.duplex.commands import ( + AckPlayback, + AppendAudio, + AppendText, + BargeIn, + CancelInput, + CancelResponse, + ClearInput, + ClearOutputAudio, + CloseSession, + Commit, + CreateItem, + CreateResponse, + DeleteItem, + DuplexCommand, + DuplexCommandError, + Heartbeat, + SignalTurn, + TruncateItem, + UpdateSession, +) + +REALTIME_INPUT_AUDIO_FORMATS = { + "pcm16", + "pcm_s16le", + "s16le", + "pcm_f32le", + "g711_ulaw", + "g711_alaw", +} +REALTIME_OUTPUT_AUDIO_FORMATS = { + "pcm16", + "pcm_s16le", + "s16le", + "wav", + "pcm", + "g711_ulaw", + "g711_alaw", +} + +#: Wire hint keys copied verbatim from a client append onto the internal payload. +REALTIME_INPUT_HINT_KEYS = ( + "duration_ms", + "audio_duration_ms", + "audio_start_ms", + "audio_end_ms", + "is_speech", + "speech", + "speech_probability", + "vad", + "overlap_action", + "overlap", + "force_barge_in", + "force_listen", + "text", + "transcript", +) + + +@dataclass(frozen=True, slots=True) +class RealtimeInputDefaults: + """Session-level wire defaults an append may omit (derived from the session object).""" + + input_audio_format: str = "pcm16" + input_sample_rate_hz: int = 16000 + output_audio_format: str = "pcm16" + output_sample_rate_hz: int | None = None + overlap_silence_rms: float = 0.003 + + def with_session_payload(self, session_payload: Mapping[str, object]) -> RealtimeInputDefaults: + """Return defaults updated from a Realtime ``session`` object (session.update).""" + values = apply_realtime_session_defaults(self, session_payload) + return values + + +# ---- audio format helpers ---- + + +def parse_realtime_audio_format(raw_format: object) -> tuple[object, int | None]: + def normalize_format(fmt: str) -> str: + normalized = fmt.lower() + if normalized in {"audio/pcm", "pcm"}: + return "pcm16" + if normalized in {"audio/wav", "wav"}: + return "wav" + if normalized in {"audio/pcm16", "pcm16", "pcm_s16le", "s16le"}: + return "pcm16" + if normalized in {"audio/pcm_f32le", "pcm_f32le", "f32le"}: + return "pcm_f32le" + if normalized in {"audio/g711_ulaw", "g711_ulaw", "g711-ulaw", "ulaw", "mulaw"}: + return "g711_ulaw" + if normalized in {"audio/g711_alaw", "g711_alaw", "g711-alaw", "alaw"}: + return "g711_alaw" + return fmt + + if isinstance(raw_format, str): + return normalize_format(raw_format), None + if not isinstance(raw_format, dict): + return raw_format, None + rate = raw_format.get("rate") or raw_format.get("sample_rate_hz") or raw_format.get("sample_rate") + sample_rate_hz = int(rate) if isinstance(rate, int | float) and rate > 0 else None + fmt = raw_format.get("type") or raw_format.get("format") + if not isinstance(fmt, str): + return raw_format, sample_rate_hz + return normalize_format(fmt), sample_rate_hz + + +def duplex_response_format(realtime_format: str) -> str: + normalized = realtime_format.lower() + if normalized in {"pcm16", "pcm_s16le", "s16le"}: + return "pcm" + if normalized in {"g711_ulaw", "g711_alaw"}: + return "pcm" + if normalized in {"wav", "pcm"}: + return normalized + return "wav" + + +def realtime_output_format(duplex_format: object) -> str: + if isinstance(duplex_format, str) and duplex_format.lower() in {"g711_ulaw", "g711_alaw"}: + return duplex_format.lower() + if isinstance(duplex_format, str) and duplex_format.lower() == "pcm": + return "pcm16" + return str(duplex_format or "wav") + + +def is_supported_realtime_input_format(fmt: object) -> bool: + return isinstance(fmt, str) and fmt.lower() in REALTIME_INPUT_AUDIO_FORMATS + + +def realtime_audio_format_object(fmt: object, *, sample_rate_hz: int | None = None) -> dict[str, object]: + if isinstance(fmt, str) and fmt.lower() in {"pcm16", "pcm_s16le", "s16le", "pcm"}: + payload: dict[str, object] = {"type": "audio/pcm"} + elif isinstance(fmt, str) and fmt.lower() == "pcm_f32le": + payload = {"type": "audio/pcm_f32le"} + elif isinstance(fmt, str) and fmt.lower() == "g711_ulaw": + payload = {"type": "audio/g711_ulaw"} + elif isinstance(fmt, str) and fmt.lower() == "g711_alaw": + payload = {"type": "audio/g711_alaw"} + else: + payload = {"type": "audio/wav"} + if sample_rate_hz is not None: + payload["rate"] = int(sample_rate_hz) + return payload + + +def validate_realtime_session_audio_formats(session_payload: Mapping[str, object]) -> str | None: + audio_config = session_payload.get("audio") + input_format: object = session_payload.get("input_audio_format") + if input_format is None and isinstance(audio_config, dict): + audio_input = audio_config.get("input") + if isinstance(audio_input, dict): + input_format = audio_input.get("format") + parsed_input, _ = parse_realtime_audio_format(input_format) + if input_format is not None and not ( + isinstance(parsed_input, str) and parsed_input.lower() in REALTIME_INPUT_AUDIO_FORMATS + ): + return f"Unsupported input_audio_format: {input_format}" + + output_format: object = session_payload.get("output_audio_format") or session_payload.get("response_format") + if output_format is None and isinstance(audio_config, dict): + audio_output = audio_config.get("output") + if isinstance(audio_output, dict): + output_format = audio_output.get("format") + parsed_output, _ = parse_realtime_audio_format(output_format) + if output_format is not None and not ( + isinstance(parsed_output, str) and parsed_output.lower() in REALTIME_OUTPUT_AUDIO_FORMATS + ): + return f"Unsupported output_audio_format: {output_format}" + return None + + +def validate_realtime_response_audio_formats(response_payload: Mapping[str, object]) -> str | None: + output_format: object = response_payload.get("output_audio_format") or response_payload.get("response_format") + audio_config = response_payload.get("audio") + if output_format is None and isinstance(audio_config, dict): + audio_output = audio_config.get("output") + if isinstance(audio_output, dict): + output_format = audio_output.get("format") + parsed_output, _ = parse_realtime_audio_format(output_format) + if output_format is not None and not ( + isinstance(parsed_output, str) and parsed_output.lower() in REALTIME_OUTPUT_AUDIO_FORMATS + ): + return f"Unsupported output_audio_format: {output_format}" + return None + + +def validate_conversation_item_audio_formats(item: object) -> str | None: + if not isinstance(item, dict): + return None + content = item.get("content") + if not isinstance(content, list): + return None + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") not in {"input_audio", "audio"}: + continue + raw_format = part.get("format") + parsed_format, _ = parse_realtime_audio_format(raw_format) + if raw_format is not None and not ( + isinstance(parsed_format, str) and parsed_format.lower() in REALTIME_INPUT_AUDIO_FORMATS + ): + return f"Unsupported input_audio format in conversation.item.create: {raw_format}" + return None + + +def apply_realtime_session_defaults( + defaults: RealtimeInputDefaults, + session_payload: Mapping[str, object], +) -> RealtimeInputDefaults: + """Derive the wire defaults a Realtime ``session`` object declares.""" + input_format: object = session_payload.get("input_audio_format") + audio_config = session_payload.get("audio") + if input_format is None and isinstance(audio_config, dict): + audio_input = audio_config.get("input") + if isinstance(audio_input, dict): + input_format = audio_input.get("format") + input_format, input_rate = parse_realtime_audio_format(input_format) + # Any: heterogeneous field values splatted into ``dataclasses.replace`` below. + updates: dict[str, Any] = {} + if isinstance(input_format, str) and input_format.lower() in REALTIME_INPUT_AUDIO_FORMATS: + updates["input_audio_format"] = input_format + output_format: object = session_payload.get("output_audio_format") or session_payload.get("response_format") + output_rate_raw: object | None = None + if output_format is None and isinstance(audio_config, dict): + audio_output = audio_config.get("output") + if isinstance(audio_output, dict): + output_format = audio_output.get("format") + output_rate_raw = audio_output.get("sample_rate_hz") or audio_output.get("sample_rate") + output_format, output_rate = parse_realtime_audio_format(output_format) + if output_rate is None and isinstance(output_rate_raw, int | float) and output_rate_raw > 0: + output_rate = int(output_rate_raw) + if isinstance(output_format, str) and output_format.lower() in REALTIME_OUTPUT_AUDIO_FORMATS: + updates["output_audio_format"] = realtime_output_format(output_format) + sample_rate = session_payload.get("sample_rate_hz") or session_payload.get("sample_rate") + if sample_rate is None and isinstance(audio_config, dict): + audio_input = audio_config.get("input") + if isinstance(audio_input, dict): + sample_rate = audio_input.get("sample_rate_hz") or audio_input.get("sample_rate") + if sample_rate is None: + sample_rate = input_rate + if isinstance(sample_rate, int | float) and sample_rate > 0: + updates["input_sample_rate_hz"] = int(sample_rate) + if isinstance(output_rate, int | float) and output_rate > 0: + updates["output_sample_rate_hz"] = int(output_rate) + overlap_fields = realtime_overlap_fields(session_payload) + overlap_silence_rms = overlap_fields.get("overlap_silence_rms") + if isinstance(overlap_silence_rms, int | float): + updates["overlap_silence_rms"] = max(0.0, float(overlap_silence_rms)) + return replace(defaults, **updates) if updates else defaults + + +def realtime_overlap_fields(session_payload: Mapping[str, object]) -> dict[str, object]: + fields: dict[str, object] = {} + if isinstance(session_payload.get("overlap_policy"), str): + fields["overlap_policy"] = session_payload["overlap_policy"] + for key in ("overlap_short_ack_ms", "overlap_barge_in_ms", "overlap_silence_rms"): + value = session_payload.get(key) + if isinstance(value, int | float): + fields[key] = value + if isinstance(session_payload.get("playback_commit_policy"), str): + fields["playback_commit_policy"] = session_payload["playback_commit_policy"] + return fields + + +def json_safe_realtime_payload(payload: Mapping[str, object]) -> dict[str, object]: + clean: dict[str, object] = {} + for key, value in payload.items(): + if key == "extra_body": + continue + if isinstance(value, str | int | float | bool) or value is None: + clean[key] = value + elif isinstance(value, dict): + clean[key] = json_safe_realtime_payload(value) + elif isinstance(value, list): + clean[key] = [ + (json_safe_realtime_payload(item) if isinstance(item, dict) else item) + for item in value + if isinstance(item, str | int | float | bool | dict) or item is None + ] + return clean + + +def input_audio_transcription_config(session_payload: Mapping[str, object]) -> dict[str, object] | None: + transcription = session_payload.get("input_audio_transcription") + if isinstance(transcription, dict): + return transcription + audio_config = session_payload.get("audio") + if not isinstance(audio_config, dict): + return None + audio_input = audio_config.get("input") + if not isinstance(audio_input, dict): + return None + transcription = audio_input.get("transcription") + return transcription if isinstance(transcription, dict) else None + + +def realtime_max_output_tokens(value: object) -> int | None: + """Normalize Realtime max output tokens (``"inf"`` -> ``None``).""" + if value is None: + return None + if isinstance(value, str) and value.strip().lower() in {"inf", "infinity", "unlimited"}: + return None + if isinstance(value, int) and value > 0: + return int(value) + return None + + +# ---- speech hints ---- + + +def input_explicitly_non_speech(event: Mapping[str, object]) -> bool: + for key in ("is_speech", "speech"): + value = event.get(key) + if isinstance(value, bool): + return not value + vad = event.get("vad") + if isinstance(vad, dict): + value = vad.get("is_speech") + if isinstance(value, bool): + return not value + probability = vad.get("speech_probability", vad.get("probability")) + if isinstance(probability, int | float): + return float(probability) < 0.5 + probability = event.get("speech_probability") + return isinstance(probability, int | float) and float(probability) < 0.5 + + +def input_looks_like_speech( + event: Mapping[str, object], + *, + audio: object, + fmt: object, + overlap_silence_rms: float, +) -> bool: + if input_explicitly_non_speech(event): + return False + for key in ("is_speech", "speech"): + value = event.get(key) + if isinstance(value, bool): + return value + vad = event.get("vad") + if isinstance(vad, dict): + probability = vad.get("speech_probability", vad.get("probability")) + if isinstance(probability, int | float): + return float(probability) >= 0.5 + probability = event.get("speech_probability") + if isinstance(probability, int | float): + return float(probability) >= 0.5 + if fmt != "pcm_f32le": + return True + if isinstance(audio, bytes | bytearray): + raw = bytes(audio) + elif isinstance(audio, str): + try: + raw = base64.b64decode(audio, validate=True) + except (binascii.Error, ValueError): + return True + else: + return True + if len(raw) < 4 or len(raw) % 4 != 0: + return True + samples = np.frombuffer(raw, dtype=np.float32) + if samples.size == 0: + return False + rms = float(np.sqrt(np.mean(np.square(samples.astype(np.float32))))) + threshold = event.get("overlap_silence_rms") + if not isinstance(threshold, int | float): + vad = event.get("vad") + if isinstance(vad, dict): + threshold = vad.get("silence_rms") + silence_rms = float(threshold) if isinstance(threshold, int | float) else overlap_silence_rms + return rms >= max(0.0, silence_rms) + + +def copy_realtime_input_hints(source: Mapping[str, object], target: dict[str, object]) -> None: + for key in REALTIME_INPUT_HINT_KEYS: + if key in source: + target[key] = source[key] + + +def validate_realtime_video_frames(video_frames: object, max_slice_nums: object) -> str | None: + """Validate omni-duplex camera frames on input_audio_buffer.append. + + Wire contract matches the official MiniCPM-o duplex loop: one base + base64 JPEG per ~1 s audio chunk, optionally followed by that unit's + stacked composite tiling the sub-frames captured inside it (at most 2 + images either way). A caller-supplied ``max_slice_nums`` is rejected rather + than silently ignored: slicing is Stage 0's decision here, and Stage 0 + already applies the official HD suggestion for a stacked unit + (``max_slice_nums=[2, 1]``). The wire simply does not let the client choose + it. + """ + if max_slice_nums not in (None, 1): + return "max_slice_nums is not selectable on the wire; Stage0 slices stacked units itself" + if not isinstance(video_frames, list): + return "video_frames must be a list of base64-encoded images" + frames = [frame for frame in video_frames if frame is not None] + if len(frames) > 2: + return "video_frames carries more than 2 frames for one append; send ~1 frame per 1 s chunk" + for frame in frames: + if not isinstance(frame, str) or not frame: + return "video_frames entries must be non-empty base64 strings" + if len(frame) > 4_000_000: + return "video_frames entry exceeds 4MB base64; reduce capture resolution or JPEG quality" + try: + header = base64.b64decode(frame[:64] + "=" * (-len(frame[:64]) % 4)) + except (binascii.Error, ValueError): + return "video_frames entries must be valid base64" + if not (header.startswith(b"\xff\xd8") or header.startswith(b"\x89PNG")): + return "video_frames entries must be JPEG or PNG images" + return None + + +# ---- conversation items ---- + + +def normalize_conversation_item(item: Mapping[str, object]) -> dict[str, object]: + normalized = dict(item) + normalized.setdefault("id", f"item_{uuid4().hex}") + normalized.setdefault("object", "realtime.item") + normalized.setdefault("type", "message") + normalized.setdefault("status", "completed") + if "role" not in normalized and normalized.get("type") == "message": + normalized["role"] = "user" + if not isinstance(normalized.get("content"), list): + normalized["content"] = [] + return normalized + + +def text_chars_for_audio_ms_from_marks( + audio_end_ms: int, + text_len: int, + marks: list[object], + *, + final_ms: object | None = None, +) -> int: + if text_len <= 0: + return 0 + clean_marks: list[tuple[int, int]] = [] + for mark in marks: + if not isinstance(mark, dict): + continue + raw_text_chars = mark.get("text_chars") + raw_audio_end_ms = mark.get("audio_end_ms", mark.get("audio_ms")) + if not isinstance(raw_text_chars, int | float) or not isinstance(raw_audio_end_ms, int | float): + continue + clean_marks.append((max(0, int(raw_audio_end_ms)), min(text_len, max(0, int(raw_text_chars))))) + if not clean_marks: + return 0 if audio_end_ms <= 0 else text_len + clean_marks.sort(key=lambda item: item[0]) + audio_end_ms = max(0, int(audio_end_ms)) + if audio_end_ms <= 0: + return 0 + previous_ms = 0 + previous_chars = 0 + for mark_ms, mark_chars in clean_marks: + mark_ms = max(previous_ms, mark_ms) + mark_chars = max(previous_chars, min(text_len, mark_chars)) + if audio_end_ms <= mark_ms: + if mark_ms <= previous_ms: + return mark_chars + ratio = (audio_end_ms - previous_ms) / max(1, mark_ms - previous_ms) + return int(previous_chars + (mark_chars - previous_chars) * max(0.0, min(1.0, ratio))) + previous_ms = mark_ms + previous_chars = mark_chars + if isinstance(final_ms, int | float) and int(final_ms) > previous_ms: + if audio_end_ms >= int(final_ms): + return text_len + ratio = (audio_end_ms - previous_ms) / max(1, int(final_ms) - previous_ms) + return int(previous_chars + (text_len - previous_chars) * max(0.0, min(1.0, ratio))) + return text_len if audio_end_ms >= previous_ms else previous_chars + + +def truncate_realtime_item_content(item: dict[str, object], *, content_index: int, audio_end_ms: int) -> None: + content = item.get("content") + if not isinstance(content, list) or not content: + return + index = max(0, int(content_index)) + if index >= len(content): + return + part = content[index] + if not isinstance(part, dict): + return + transcript = part.get("transcript") + if not isinstance(transcript, str) or not transcript: + return + marks = part.get("audio_text_marks") + if isinstance(marks, list): + keep_chars = text_chars_for_audio_ms_from_marks( + audio_end_ms, + len(transcript), + marks, + final_ms=part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms"), + ) + part["transcript"] = transcript[:keep_chars].rstrip() + return + duration_ms = part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms") + if isinstance(duration_ms, int | float) and duration_ms > 0: + keep_chars = int(len(transcript) * max(0.0, min(1.0, int(audio_end_ms) / float(duration_ms)))) + part["transcript"] = transcript[:keep_chars].rstrip() + elif audio_end_ms <= 0: + part["transcript"] = "" + + +def validate_realtime_item_truncate(item: Mapping[str, object], *, content_index: int, audio_end_ms: int) -> str | None: + if item.get("type") != "message" or item.get("role") != "assistant": + return "conversation.item.truncate only supports assistant message items" + if audio_end_ms < 0: + return "conversation.item.truncate requires non-negative audio_end_ms" + content = item.get("content") + if not isinstance(content, list) or not content: + return None + index = max(0, int(content_index)) + if index >= len(content): + return f"conversation.item.truncate content_index out of range: {content_index}" + part = content[index] + if not isinstance(part, dict): + return "conversation.item.truncate target content part is invalid" + if part.get("type") not in {"audio", "output_audio"}: + return "conversation.item.truncate target content part must be audio" + duration_ms = part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms") + if isinstance(duration_ms, int | float) and int(duration_ms) >= 0 and audio_end_ms > int(duration_ms): + return "conversation.item.truncate audio_end_ms exceeds item audio duration" + return None + + +def input_transcript_from_item(item: Mapping[str, object]) -> str: + content = item.get("content") + if not isinstance(content, list): + return "" + parts: list[str] = [] + for part in content: + if not isinstance(part, dict): + continue + for key in ("transcript", "text"): + value = part.get(key) + if isinstance(value, str) and value: + if part.get("type") in {"input_audio", "audio", "audio_transcript", "transcript"}: + parts.append(value) + break + return "".join(parts).strip() + + +# ---- append / item audio conversion ---- + + +def build_append_audio( + event: Mapping[str, object], + *, + defaults: RealtimeInputDefaults, + hints_source: Mapping[str, object] | None = None, +) -> AppendAudio: + """Validate, convert (to 16 kHz ``pcm_f32le`` base64) and pack one audio append. + + Speech classification without VAD (``is_speech`` on the command) uses the + explicit wire hints and the RMS fallback exactly like the old translator. + """ + event_id = cast("str", event.get("event_id")) if isinstance(event.get("event_id"), str) else None + audio = event.get("audio") or event.get("delta") + fmt, format_rate = parse_realtime_audio_format( + event.get("format") or event.get("input_audio_format") or defaults.input_audio_format + ) + sample_rate_hz = ( + event.get("sample_rate_hz") or event.get("sample_rate") or format_rate or defaults.input_sample_rate_hz + ) + if not is_supported_realtime_input_format(fmt): + raise DuplexCommandError( + f"Unsupported input_audio_format: {fmt}", + code="unsupported_audio_format", + event_id=event_id, + ) + try: + audio, fmt, sample_rate_hz = convert_input_audio_with_rate( + audio, + fmt, + sample_rate_hz=sample_rate_hz if isinstance(sample_rate_hz, int | float) else None, + ) + except ValueError as exc: + raise DuplexCommandError(str(exc), code="bad_event", event_id=event_id) from exc + hints: dict[str, object] = {} + if hints_source is not None: + copy_realtime_input_hints(hints_source, hints) + copy_realtime_input_hints(event, hints) + looks_like_speech = input_looks_like_speech( + {**(hints_source or {}), **event}, + audio=audio, + fmt=fmt, + overlap_silence_rms=defaults.overlap_silence_rms, + ) + video_frames: tuple[str, ...] = () + raw_frames = event.get("video_frames") + if raw_frames is not None: + frames_error = validate_realtime_video_frames(raw_frames, event.get("max_slice_nums")) + if frames_error is not None: + raise DuplexCommandError(frames_error, code="invalid_video_frames", event_id=event_id) + if isinstance(raw_frames, list): + video_frames = tuple(frame for frame in raw_frames if isinstance(frame, str) and frame) + duration_ms = hints.get("duration_ms", hints.get("audio_duration_ms")) + audio_end_ms = hints.get("audio_end_ms") + try: + audio_bytes = base64.b64decode(audio, validate=True) if isinstance(audio, str) and audio else b"" + except (binascii.Error, ValueError) as exc: + raise DuplexCommandError("input audio is not valid base64", code="bad_audio", event_id=event_id) from exc + return AppendAudio( + event_id=event_id, + audio=audio_bytes, + format=str(fmt), + sample_rate_hz=int(sample_rate_hz) if isinstance(sample_rate_hz, int | float) else None, + is_speech=looks_like_speech, + video_frames=video_frames, + duration_ms=int(duration_ms) if isinstance(duration_ms, int | float) else None, + audio_end_ms=int(audio_end_ms) if isinstance(audio_end_ms, int | float) else None, + hints=hints, + ) + + +# ---- translation ---- + + +def translate_realtime_command( + payload: Mapping[str, object], + *, + defaults: RealtimeInputDefaults | None = None, +) -> DuplexCommand: + """Map one OpenAI Realtime client event onto a :class:`DuplexCommand`. + + Raises :class:`DuplexCommandError` for malformed or unsupported payloads. + ``session.resume`` and ``session.event_ack`` are transport concerns and are + rejected with ``code="unknown_event"``. + """ + defaults = defaults or RealtimeInputDefaults() + event_type = payload.get("type") + event_id = cast("str", payload.get("event_id")) if isinstance(payload.get("event_id"), str) else None + if not isinstance(event_type, str): + raise DuplexCommandError("Duplex event missing string type", code="bad_event", event_id=event_id) + + if event_type == "session.update": + session = payload.get("session") + session_payload: Mapping[str, object] = session if isinstance(session, dict) else payload + format_error = validate_realtime_session_audio_formats(session_payload) + if format_error is not None: + raise DuplexCommandError(format_error, code="unsupported_audio_format", event_id=event_id) + from vllm_omni.engine.duplex.turn_detection import validate_realtime_turn_detection + + turn_detection_error = validate_realtime_turn_detection(session_payload) + if turn_detection_error is not None: + raise DuplexCommandError(turn_detection_error, code="unsupported_turn_detection", event_id=event_id) + return UpdateSession(event_id=event_id, patch=dict(session_payload)) + + if event_type == "conversation.item.create": + item = payload.get("item") + format_error = validate_conversation_item_audio_formats(item) + if format_error is not None: + raise DuplexCommandError(format_error, code="unsupported_audio_format", event_id=event_id) + if not isinstance(item, dict): + raise DuplexCommandError("conversation.item.create requires item", code="bad_event", event_id=event_id) + previous_item_id = payload.get("previous_item_id") + return CreateItem( + event_id=event_id, + item=normalize_conversation_item(item), + previous_item_id=previous_item_id if isinstance(previous_item_id, str) else None, + ) + + if event_type == "conversation.item.delete": + item_id = payload.get("item_id") + if not isinstance(item_id, str) or not item_id: + raise DuplexCommandError( + "conversation.item.delete requires item_id", code="missing_item_id", event_id=event_id + ) + return DeleteItem(event_id=event_id, item_id=item_id) + + if event_type == "conversation.item.truncate": + item_id = payload.get("item_id") + audio_end_ms = payload.get("audio_end_ms") + content_index = payload.get("content_index", 0) + if not isinstance(item_id, str) or not item_id: + raise DuplexCommandError( + "conversation.item.truncate requires item_id", code="missing_item_id", event_id=event_id + ) + if not isinstance(audio_end_ms, int | float): + raise DuplexCommandError( + "conversation.item.truncate requires numeric audio_end_ms", code="bad_event", event_id=event_id + ) + return TruncateItem( + event_id=event_id, + item_id=item_id, + audio_end_ms=int(audio_end_ms), + content_index=int(content_index) if isinstance(content_index, int | float) else 0, + ) + + if event_type == "input_audio_buffer.append": + return build_append_audio(dict(payload), defaults=defaults) + + if event_type in {"input_audio_buffer.commit", "input.commit"}: + final = payload.get("final", True) + create_response = payload.get("response_create", payload.get("create_response")) + is_speech = payload.get("is_speech") + return Commit( + event_id=event_id, + final=bool(final) if isinstance(final, bool) else True, + create_response=( + bool(create_response) + if isinstance(create_response, bool) + else (True if event_type == "input.commit" and create_response is None else None) + ), + is_speech=is_speech if isinstance(is_speech, bool) else None, + ) + + if event_type == "input_audio_buffer.clear": + return ClearInput(event_id=event_id) + + if event_type == "output_audio_buffer.clear": + response_id = payload.get("response_id") + return ClearOutputAudio( + event_id=event_id, + response_id=response_id if isinstance(response_id, str) and response_id else None, + ) + + if event_type == "response.cancel": + response_id = payload.get("response_id") + return CancelResponse( + event_id=event_id, + response_id=response_id if isinstance(response_id, str) and response_id else None, + ) + + if event_type == "response.create": + response_payload = payload.get("response") + if isinstance(response_payload, dict): + format_error = validate_realtime_response_audio_formats(response_payload) + if format_error is not None: + raise DuplexCommandError(format_error, code="unsupported_audio_format", event_id=event_id) + return CreateResponse( + event_id=event_id, + options=dict(response_payload) if isinstance(response_payload, dict) else {}, + ) + + if event_type in {"playback.ack", "audio.playback_ack"}: + played_ms = payload.get("played_ms") + if not isinstance(played_ms, int | float): + raise DuplexCommandError("playback.ack requires numeric played_ms", code="bad_event", event_id=event_id) + committed_ms = payload.get("committed_ms") + response_id = payload.get("response_id") + item_id = payload.get("item_id") + return AckPlayback( + event_id=event_id, + played_ms=int(played_ms), + committed_ms=int(committed_ms) if isinstance(committed_ms, int | float) else None, + response_id=response_id if isinstance(response_id, str) and response_id else None, + item_id=item_id if isinstance(item_id, str) and item_id else None, + ) + + if event_type == "session.heartbeat": + return Heartbeat(event_id=event_id) + + if event_type == "conversation.item.retrieve": + # The projected conversation items live engine-side; the runner answers. + return SignalTurn(event_id=event_id, event="conversation.item.retrieve", signal_payload=dict(payload)) + + if event_type in {"session.close", "close", "close_session"}: + reason = payload.get("reason") + return CloseSession(event_id=event_id, reason=reason if isinstance(reason, str) and reason else "client_close") + + if event_type in {"input.text.append", "input_text.append", "push_text"}: + text = payload.get("text") + if not isinstance(text, str): + raise DuplexCommandError("input.text.append requires text", code="bad_event", event_id=event_id) + return AppendText(event_id=event_id, text=text) + + if event_type == "input.cancel": + return CancelInput(event_id=event_id) + + if event_type == "barge_in": + return BargeIn(event_id=event_id) + + if event_type in {"turn.signal", "signal_turn"}: + signal_event = payload.get("event") + if not isinstance(signal_event, str) or not signal_event: + raise DuplexCommandError("turn.signal requires event", code="bad_event", event_id=event_id) + signal_payload = payload.get("payload") + signal_payload = dict(signal_payload) if isinstance(signal_payload, dict) else {} + if signal_event == "input.cancel": + return CancelInput(event_id=event_id) + if signal_event == "barge_in": + return BargeIn(event_id=event_id) + if signal_event == "response.cancel": + response_id = signal_payload.get("response_id", payload.get("response_id")) + return CancelResponse( + event_id=event_id, + response_id=response_id if isinstance(response_id, str) and response_id else None, + ) + if signal_event == "session.update": + return UpdateSession(event_id=event_id, patch=signal_payload) + if signal_event == "conversation.item.create": + item = signal_payload.get("item") + if not isinstance(item, dict): + raise DuplexCommandError("conversation.item.create requires item", code="bad_event", event_id=event_id) + previous_item_id = signal_payload.get("previous_item_id") + return CreateItem( + event_id=event_id, + item=normalize_conversation_item(item), + previous_item_id=previous_item_id if isinstance(previous_item_id, str) else None, + ) + if signal_event == "conversation.item.delete": + item_id = signal_payload.get("item_id") + if not isinstance(item_id, str) or not item_id: + raise DuplexCommandError( + "conversation.item.delete requires item_id", code="missing_item_id", event_id=event_id + ) + return DeleteItem(event_id=event_id, item_id=item_id) + if signal_event == "conversation.item.truncate": + item_id = signal_payload.get("item_id") + audio_end_ms = signal_payload.get("audio_end_ms") + content_index = signal_payload.get("content_index", 0) + if not isinstance(item_id, str) or not item_id or not isinstance(audio_end_ms, int | float): + raise DuplexCommandError( + "conversation.item.truncate requires item_id and numeric audio_end_ms", + code="bad_event", + event_id=event_id, + ) + return TruncateItem( + event_id=event_id, + item_id=item_id, + audio_end_ms=int(audio_end_ms), + content_index=int(content_index) if isinstance(content_index, int | float) else 0, + ) + return SignalTurn(event_id=event_id, event=signal_event, signal_payload=signal_payload) + + raise DuplexCommandError(f"Unknown duplex event type: {event_type}", code="unknown_event", event_id=event_id) + + +__all__ = [ + "REALTIME_INPUT_AUDIO_FORMATS", + "REALTIME_INPUT_HINT_KEYS", + "REALTIME_OUTPUT_AUDIO_FORMATS", + "RealtimeInputDefaults", + "apply_realtime_session_defaults", + "build_append_audio", + "copy_realtime_input_hints", + "duplex_response_format", + "input_audio_transcription_config", + "input_explicitly_non_speech", + "input_looks_like_speech", + "input_transcript_from_item", + "is_supported_realtime_input_format", + "json_safe_realtime_payload", + "normalize_conversation_item", + "parse_realtime_audio_format", + "realtime_audio_format_object", + "realtime_max_output_tokens", + "realtime_output_format", + "realtime_overlap_fields", + "text_chars_for_audio_ms_from_marks", + "translate_realtime_command", + "truncate_realtime_item_content", + "validate_conversation_item_audio_formats", + "validate_realtime_item_truncate", + "validate_realtime_response_audio_formats", + "validate_realtime_session_audio_formats", + "validate_realtime_video_frames", +] diff --git a/vllm_omni/engine/duplex/realtime_events.py b/vllm_omni/engine/duplex/realtime_events.py new file mode 100644 index 000000000..91359f297 --- /dev/null +++ b/vllm_omni/engine/duplex/realtime_events.py @@ -0,0 +1,1526 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Session-internal duplex events -> typed :class:`DuplexEvent` objects (stateful projection). + +This is the former ``entrypoints/duplex/realtime_output.py`` projector plus the +projection-relevant fields of ``realtime_state.py``, now owned by the session +runner through :class:`RealtimeProjectionState`. The projection consumes the +state (response / item ids, content-part bookkeeping) when it *constructs* +events; rendering an event to wire JSON (``event.to_realtime()``) is pure and +lives on the event classes in ``vllm_omni.engine.duplex.events``. + +Besides the output projection (:func:`project_internal_event`) the state also +carries the input-side bookkeeping the old input translator kept (input buffer +flags, conversation items, response-id fallbacks); the ``resolve_*`` / +``note_*`` helpers give the runner the same behaviour for the corresponding +commands. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, cast +from uuid import uuid4 + +from vllm_omni.engine.duplex.audio import convert_output_audio +from vllm_omni.engine.duplex.commands import ( + AppendAudio, + CancelResponse, + ClearOutputAudio, + Commit, + CreateItem, + DeleteItem, + TruncateItem, +) +from vllm_omni.engine.duplex.events import ( + AudioDelta, + AudioDone, + ContentPartAdded, + ContentPartDone, + DuplexEvent, + DuplexRawEvent, + FunctionCallArgumentsDelta, + FunctionCallArgumentsDone, + InputCleared, + InputCommitted, + InputTranscriptionCompleted, + ItemAdded, + ItemCreated, + ItemDeleted, + ItemDone, + ItemRetrieved, + ItemTruncated, + Listen, + OutputAudioCleared, + OutputItemAdded, + OutputItemDone, + RateLimitsUpdated, + ResponseCreated, + ResponseDone, + SessionClosed, + SessionCreated, + SessionReplaced, + SessionResumed, + SessionResyncRequired, + SessionUpdated, + Speak, + SpeechStarted, + SpeechStopped, + TextDelta, + TextDone, + TranscriptDelta, + TranscriptDone, + error_event, +) +from vllm_omni.engine.duplex.realtime_commands import ( + RealtimeInputDefaults, + apply_realtime_session_defaults, + build_append_audio, + copy_realtime_input_hints, + input_looks_like_speech, + input_transcript_from_item, + parse_realtime_audio_format, + realtime_audio_format_object, + realtime_output_format, + truncate_realtime_item_content, + validate_realtime_item_truncate, +) + +if TYPE_CHECKING: + # Runtime import pulls in the VAD backend; the projection only needs the type. + from vllm_omni.engine.duplex.turn_detection import TurnDetectionResult + + +def _str_or_none(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + +def _int_or(value: object, default: int = 0) -> int: + return int(value) if isinstance(value, int | float) else default + + +@dataclass(slots=True) +class _ResponseProjection: + item_id: str + transcript_parts: list[str] = field(default_factory=list) + text_parts: list[str] = field(default_factory=list) + audio_duration_ms: int | None = None + audio_text_marks: list[dict[str, int]] = field(default_factory=list) + audio_delta_emitted: bool = False + audio_done_emitted: bool = False + audio_part_added: bool = False + audio_part_done: bool = False + text_part_added: bool = False + text_part_done: bool = False + output_text_done: bool = False + output_item_done: bool = False + conversation_item_done: bool = False + speak_emitted: bool = False + done_emitted: bool = False + + @property + def transcript(self) -> str: + return "".join(self.transcript_parts) + + @property + def text(self) -> str: + return "".join(self.text_parts) + + +@dataclass(slots=True) +class RealtimeProjectionState: + """Per-session Realtime projection state (lives on the engine-side runner).""" + + session_id: str + model: str | None = None + default_payload: Mapping[str, object] | None = None + #: Wire defaults (input/output formats, rates, silence rms) declared by the session object. + defaults: RealtimeInputDefaults = field(default_factory=RealtimeInputDefaults) + #: The first ``session.created`` also emits ``session.updated`` (client sent session.update to open). + initial_session_update: bool = False + # ---- response / item projection ---- + response_states: dict[str | int, _ResponseProjection] = field(default_factory=dict) + item_truncation_cursors: dict[str, tuple[int, int]] = field(default_factory=dict) + active_response_id: str | None = None + last_response_id: str | None = None + conversation_items: dict[str, dict[str, object]] = field(default_factory=dict) + last_conversation_item_id: str | None = None + pending_commit_item_ids: list[str] = field(default_factory=list) + # ---- input buffer projection (was on the input translator) ---- + input_speech_started: bool = False + active_input_item_id: str | None = None + input_audio_buffer_has_audio: bool = False + input_audio_buffer_had_non_speech: bool = False + input_audio_buffer_transcript_parts: list[str] = field(default_factory=list) + + def __post_init__(self) -> None: + if self.default_payload is not None: + self.apply_session_defaults(self.default_payload) + + # ---- session defaults ---- + + def apply_session_defaults(self, session_payload: Mapping[str, object]) -> None: + self.defaults = apply_realtime_session_defaults(self.defaults, session_payload) + + @property + def input_audio_format(self) -> str: + return self.defaults.input_audio_format + + @property + def input_sample_rate_hz(self) -> int: + return self.defaults.input_sample_rate_hz + + @property + def output_audio_format(self) -> str: + return self.defaults.output_audio_format + + @property + def output_sample_rate_hz(self) -> int | None: + return self.defaults.output_sample_rate_hz + + @property + def overlap_silence_rms(self) -> float: + return self.defaults.overlap_silence_rms + + +# ---- response projection helpers ---- + + +def _response_state( + state: RealtimeProjectionState, + response_id: object, + *, + event: Mapping[str, object] | None = None, + create: bool = True, +) -> _ResponseProjection | None: + if not (isinstance(response_id, str) and response_id): + # An event without a response id has no projection to attach to; a + # key derived from the transient event object would alias later + # anonymous events (and leak an entry per terminal). + return None + key = response_id + item_id = f"item_{response_id}" + projection = state.response_states.get(key) + if projection is None and create: + projection = _ResponseProjection(item_id=item_id) + state.response_states[key] = projection + return projection + + +def response_is_done(state: RealtimeProjectionState, response_id: object) -> bool: + projection = _response_state(state, response_id, create=False) + return projection is not None and projection.done_emitted + + +def _response_item_id(state: RealtimeProjectionState, response_id: object) -> str: + projection = _response_state(state, response_id) + if projection is not None: + return projection.item_id + return f"item_{uuid4().hex}" + + +def _response_content_part(*, transcript: str = "") -> dict[str, object]: + return {"type": "audio", "transcript": transcript} + + +def _response_text_content_part(*, text: str = "") -> dict[str, object]: + return {"type": "text", "text": text} + + +def _response_item_content_part( + *, + transcript: str = "", + audio_duration_ms: int | None = None, + audio_text_marks: list[dict[str, int]] | None = None, +) -> dict[str, object]: + part: dict[str, object] = {"type": "output_audio", "transcript": transcript} + if audio_duration_ms is not None: + part["audio_duration_ms"] = int(audio_duration_ms) + if audio_text_marks: + part["audio_text_marks"] = [dict(mark) for mark in audio_text_marks] + return part + + +def _response_item_text_content_part(*, text: str = "") -> dict[str, object]: + return {"type": "output_text", "text": text} + + +def _previous_item_id(state: RealtimeProjectionState, item_id: str) -> str | None: + previous: str | None = None + for known_id in state.conversation_items: + if known_id == item_id: + return previous + previous = known_id + if state.last_conversation_item_id == item_id: + return previous + return state.last_conversation_item_id + + +def _conversation_item_added_events(state: RealtimeProjectionState, item: dict[str, object]) -> list[DuplexEvent]: + item_id = item.get("id") + explicit_previous_item_id = item.pop("_previous_item_id", None) + previous_item_id = ( + explicit_previous_item_id if isinstance(explicit_previous_item_id, str) else state.last_conversation_item_id + ) + if isinstance(item_id, str) and item_id: + state.last_conversation_item_id = item_id + return [ + ItemAdded(previous_item_id=previous_item_id, item=item), + ItemCreated(previous_item_id=previous_item_id, item=item), + ] + + +def _conversation_item_done_event(state: RealtimeProjectionState, item: dict[str, object]) -> ItemDone: + item_id = item.get("id") + previous_item_id = _previous_item_id(state, item_id) if isinstance(item_id, str) else None + if isinstance(item_id, str) and item_id: + state.conversation_items[item_id] = item + state.last_conversation_item_id = item_id + return ItemDone(previous_item_id=previous_item_id, item=item) + + +def _remove_conversation_item(state: RealtimeProjectionState, item_id: str) -> bool: + removed = state.conversation_items.pop(item_id, None) is not None + if state.last_conversation_item_id == item_id: + remaining_ids = list(state.conversation_items) + state.last_conversation_item_id = remaining_ids[-1] if remaining_ids else None + state.item_truncation_cursors.pop(item_id, None) + return removed + + +def _apply_pending_item_truncation(state: RealtimeProjectionState, item: dict[str, object]) -> None: + item_id = item.get("id") + if not isinstance(item_id, str) or not item_id: + return + cursor = state.item_truncation_cursors.get(item_id) + if cursor is None: + return + content_index, audio_end_ms = cursor + truncate_realtime_item_content(item, content_index=content_index, audio_end_ms=audio_end_ms) + + +def _response_done_output_item( + state: RealtimeProjectionState, response_id: object, *, status: str +) -> dict[str, object]: + item_id = _response_item_id(state, response_id) + projection = _response_state(state, response_id) + transcript = projection.transcript if projection is not None else "" + text = projection.text if projection is not None else "" + audio_duration_ms = projection.audio_duration_ms if projection is not None else None + audio_text_marks = projection.audio_text_marks if projection is not None else None + content: list[dict[str, object]] = [] + if (projection is not None and projection.audio_part_added) or transcript or audio_duration_ms is not None: + content.append( + _response_item_content_part( + transcript=transcript, + audio_duration_ms=audio_duration_ms, + audio_text_marks=audio_text_marks, + ) + ) + if text: + content.append(_response_item_text_content_part(text=text)) + item = { + "id": item_id, + "object": "realtime.item", + "type": "message", + "role": "assistant", + "status": status, + "content": content, + } + _apply_pending_item_truncation(state, item) + return item + + +def _refresh_in_progress_response_item(state: RealtimeProjectionState, response_id: object) -> None: + if not isinstance(response_id, str) or not response_id: + return + item_id = _response_item_id(state, response_id) + item = state.conversation_items.get(item_id) + if not isinstance(item, dict): + return + content = item.get("content") + if not isinstance(content, list): + content = [] + item["content"] = content + projection = _response_state(state, response_id) + if projection is None: + return + transcript = projection.transcript + audio_duration_ms = projection.audio_duration_ms + audio_text_marks = projection.audio_text_marks + has_audio = projection.audio_part_added or bool(transcript) or audio_duration_ms is not None + if has_audio: + audio_part = _response_item_content_part( + transcript=transcript, + audio_duration_ms=audio_duration_ms, + audio_text_marks=audio_text_marks, + ) + if content and isinstance(content[0], dict) and content[0].get("type") in {"audio", "output_audio"}: + content[0] = audio_part + else: + content.insert(0, audio_part) + text = projection.text + if text: + text_index = ( + 1 if content and isinstance(content[0], dict) and content[0].get("type") in {"audio", "output_audio"} else 0 + ) + text_part = _response_item_text_content_part(text=text) + if ( + len(content) > text_index + and isinstance(content[text_index], dict) + and content[text_index].get("type") in {"text", "output_text"} + ): + content[text_index] = text_part + else: + content.insert(text_index, text_part) + _apply_pending_item_truncation(state, item) + + +def _append_response_transcript(state: RealtimeProjectionState, response_id: object, text: str) -> None: + projection = _response_state(state, response_id) + if projection is None or not text: + return + projection.transcript_parts.append(text) + + +def _ensure_response_text_part_added(state: RealtimeProjectionState, response_id: object) -> list[DuplexEvent]: + projection = _response_state(state, response_id) + if projection is None or projection.text_part_added: + return [] + projection.text_part_added = True + return [ + ContentPartAdded( + response_id=_str_or_none(response_id), + item_id=_response_item_id(state, response_id), + content_index=1 if projection.audio_part_added else 0, + part=_response_text_content_part(), + ) + ] + + +def _ensure_response_audio_part_added(state: RealtimeProjectionState, response_id: object) -> list[DuplexEvent]: + projection = _response_state(state, response_id) + if projection is None or projection.audio_part_added: + return [] + projection.audio_part_added = True + return [ + ContentPartAdded( + response_id=_str_or_none(response_id), + item_id=_response_item_id(state, response_id), + part=_response_content_part(), + ) + ] + + +def _remember_response_audio_metadata( + state: RealtimeProjectionState, response_id: object, event: Mapping[str, object] +) -> None: + projection = _response_state(state, response_id) + if projection is None: + return + duration = event.get("audio_duration_ms") + playback = event.get("playback") + if not isinstance(duration, int | float) and isinstance(playback, dict): + duration = playback.get("sent_ms") or playback.get("generated_ms") + if isinstance(duration, int | float): + projection.audio_duration_ms = max(projection.audio_duration_ms or 0, int(duration)) + marks = event.get("audio_text_marks") + if not isinstance(marks, list): + return + clean_marks: list[dict[str, int]] = [] + for mark in marks: + if not isinstance(mark, dict): + continue + text_chars = mark.get("text_chars") + audio_end_ms = mark.get("audio_end_ms", mark.get("audio_ms")) + if not isinstance(text_chars, int | float) or not isinstance(audio_end_ms, int | float): + continue + clean_marks.append({"text_chars": max(0, int(text_chars)), "audio_end_ms": max(0, int(audio_end_ms))}) + if clean_marks: + merged = list(projection.audio_text_marks) + merged.extend(clean_marks) + deduped: dict[tuple[int, int], dict[str, int]] = {} + for mark in merged: + deduped[(int(mark["audio_end_ms"]), int(mark["text_chars"]))] = mark + projection.audio_text_marks = sorted( + deduped.values(), key=lambda mark: (mark["audio_end_ms"], mark["text_chars"]) + ) + + +def _response_created_event(event: Mapping[str, object]) -> ResponseCreated: + response_id = event.get("response_id") + metadata = event.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + metadata = {**metadata, "duplex_event": dict(event)} + return ResponseCreated( + response_id=_str_or_none(response_id), + response={ + "id": response_id, + "object": "realtime.response", + "status": "in_progress", + "status_details": None, + "output": [], + "modalities": event.get("modalities") or ["audio", "text"], + "metadata": metadata, + }, + ) + + +def _response_speak_metadata(event: Mapping[str, object]) -> dict[str, object]: + return {key: event[key] for key in ("session_id", "epoch", "model_speak") if key in event} + + +def _realtime_audio_delta_events( + state: RealtimeProjectionState, + event: Mapping[str, object], + response_id: object, + audio: str, +) -> list[DuplexEvent]: + item_id = _response_item_id(state, response_id) + fmt, format_rate = parse_realtime_audio_format(event.get("format", "wav")) + source_fmt = realtime_output_format(fmt) + source_sample_rate_hz = event.get("sample_rate_hz") or format_rate + target_sample_rate_hz = ( + state.output_sample_rate_hz if state.output_audio_format in {"g711_ulaw", "g711_alaw"} else None + ) + audio, fmt, converted_sample_rate_hz = convert_output_audio( + audio, + source_fmt=source_fmt, + target_fmt=state.output_audio_format, + source_sample_rate_hz=(int(source_sample_rate_hz) if isinstance(source_sample_rate_hz, int | float) else None), + target_sample_rate_hz=target_sample_rate_hz, + ) + sample_rate_hz = ( + converted_sample_rate_hz + if fmt in {"g711_ulaw", "g711_alaw"} + else event.get("sample_rate_hz") or format_rate or state.output_sample_rate_hz + ) + metadata: dict[str, object] = {} + for key in ("session_id", "epoch", "model_speak", "end_of_turn", "playback", "vllm_omni"): + if key in event: + metadata[key] = event[key] + duration_ms = event.get("audio_duration_ms") + if isinstance(duration_ms, int | float): + metadata["audio_duration_ms"] = int(duration_ms) + marks = event.get("audio_text_marks") + if isinstance(marks, list): + metadata["audio_text_marks"] = marks + projection = _response_state(state, response_id) + if projection is not None: + projection.audio_delta_emitted = True + _remember_response_audio_metadata(state, response_id, event) + events: list[DuplexEvent] = [] + if projection is not None and not projection.speak_emitted and metadata.get("model_speak") is True: + projection.speak_emitted = True + events.append( + Speak(response_id=_str_or_none(response_id), item_id=item_id, metadata=_response_speak_metadata(event)) + ) + events.append( + AudioDelta( + response_id=_str_or_none(response_id), + item_id=item_id, + delta=str(audio), + format=str(fmt), + sample_rate_hz=int(sample_rate_hz) if isinstance(sample_rate_hz, int | float) else None, + metadata=metadata or None, + ) + ) + return events + + +def _realtime_audio_done_events( + state: RealtimeProjectionState, event: Mapping[str, object], response_id: object +) -> list[DuplexEvent]: + item_id = _response_item_id(state, response_id) + projection = _response_state(state, response_id, event=event) + transcript = projection.transcript if projection is not None else "" + events: list[DuplexEvent] = [] + if ( + isinstance(response_id, str) + and projection is not None + and not projection.audio_delta_emitted + and not transcript + ): + return events + if projection is None or not projection.audio_done_emitted: + if projection is not None: + projection.audio_done_emitted = True + events.append(AudioDone(response_id=_str_or_none(response_id), item_id=item_id)) + if transcript: + events.append(TranscriptDone(response_id=_str_or_none(response_id), item_id=item_id, transcript=transcript)) + return events + + +def _realtime_response_done_event( + state: RealtimeProjectionState, + event: Mapping[str, object], + *, + projection: _ResponseProjection | None = None, + status: str = "completed", + status_details: dict[str, object] | None = None, +) -> ResponseDone | None: + response_id = event.get("response_id") + if projection is None: + projection = _response_state(state, response_id, event=event) + if projection is not None: + if projection.done_emitted: + return None + projection.done_emitted = True + return ResponseDone( + response_id=_str_or_none(response_id), + response={ + "id": response_id, + "object": "realtime.response", + "status": status, + "status_details": status_details, + "output": [_response_done_output_item(state, response_id, status=status)], + "metadata": dict(event), + }, + ) + + +def _realtime_response_terminal_events( + state: RealtimeProjectionState, + event: Mapping[str, object], + response_id: object, + *, + status: str = "completed", + status_details: dict[str, object] | None = None, +) -> list[DuplexEvent]: + item_id = _response_item_id(state, response_id) + projection = _response_state(state, response_id, event=event) + transcript = projection.transcript if projection is not None else "" + rid = _str_or_none(response_id) + events: list[DuplexEvent] = [] + if projection is not None and projection.audio_part_added and not projection.audio_part_done: + projection.audio_part_done = True + events.append( + ContentPartDone(response_id=rid, item_id=item_id, part=_response_content_part(transcript=transcript)) + ) + if projection is not None and projection.text_parts and not projection.output_text_done: + projection.output_text_done = True + events.append( + TextDone( + response_id=rid, + item_id=item_id, + content_index=1 if projection.audio_part_added else 0, + text=projection.text, + ) + ) + if projection is not None and projection.text_parts and not projection.text_part_done: + projection.text_part_done = True + events.append( + ContentPartDone( + response_id=rid, + item_id=item_id, + content_index=1 if projection.audio_part_added else 0, + part=_response_text_content_part(text=projection.text), + ) + ) + if projection is None or not projection.output_item_done: + if projection is not None: + projection.output_item_done = True + item = _response_done_output_item(state, response_id, status=status) + events.append(OutputItemDone(response_id=rid, item=item)) + if projection is None or not projection.conversation_item_done: + if projection is not None: + projection.conversation_item_done = True + events.append(_conversation_item_done_event(state, item)) + done_event = _realtime_response_done_event( + state, + {**event, "response_id": response_id}, + projection=projection, + status=status, + status_details=status_details, + ) + if done_event is not None: + events.append(done_event) + # vLLM-Omni does not expose a Realtime-specific quota budget; emit the + # terminal event with an empty list so clients that sequence on it work. + events.append(RateLimitsUpdated()) + if isinstance(response_id, str) and response_id == state.active_response_id: + state.active_response_id = None + return events + + +def _user_item_content_from_duplex_message(message: object) -> list[dict[str, object]]: + if not isinstance(message, dict): + return [{"type": "input_audio", "transcript": ""}] + message_transcript = message.get("transcript") + content = message.get("content") + if isinstance(content, str): + return [{"type": "input_text", "text": content}] + if not isinstance(content, list): + return [ + {"type": "input_audio", "transcript": message_transcript if isinstance(message_transcript, str) else ""} + ] + parts: list[dict[str, object]] = [] + for part in content: + if not isinstance(part, dict): + continue + part_type = part.get("type") + if part_type in {"text", "input_text"} and isinstance(part.get("text"), str): + parts.append({"type": "input_text", "text": part["text"]}) + continue + if part_type == "audio_url": + transcript = ( + part.get("transcript") + if isinstance(part.get("transcript"), str) + else message_transcript + if isinstance(message_transcript, str) + else "" + ) + parts.append({"type": "input_audio", "transcript": transcript}) + continue + if part_type in {"audio", "input_audio"}: + transcript = part.get("transcript") if isinstance(part.get("transcript"), str) else "" + parts.append({"type": "input_audio", "transcript": transcript}) + return parts or [{"type": "input_audio", "transcript": ""}] + + +def _input_audio_transcription_completed_event( + item_id: str, item: Mapping[str, object] +) -> InputTranscriptionCompleted | None: + content = item.get("content") + if not isinstance(content, list): + return None + for index, part in enumerate(content): + if not isinstance(part, dict) or part.get("type") != "input_audio": + continue + transcript = part.get("transcript") + if isinstance(transcript, str) and transcript: + return InputTranscriptionCompleted(item_id=item_id, content_index=index, transcript=transcript) + return None + + +def _pop_pending_commit_item_id(state: RealtimeProjectionState) -> str: + if state.pending_commit_item_ids: + return state.pending_commit_item_ids.pop(0) + return f"item_{uuid4().hex}" + + +def realtime_session_payload(state: RealtimeProjectionState, session: object) -> dict[str, object]: + """Fill the Realtime ``session`` object defaults for session.created/updated.""" + payload = dict(session) if isinstance(session, Mapping) else {} + payload.setdefault("object", "realtime.session") + payload.setdefault("type", "realtime") + payload.setdefault("id", payload.get("id") or state.session_id) + payload.setdefault("model", payload.get("model") or state.model) + payload.setdefault("input_audio_format", state.input_audio_format) + payload.setdefault("output_audio_format", state.output_audio_format) + payload.setdefault("modalities", payload.get("modalities") or ["text", "audio"]) + payload.setdefault("output_modalities", payload.get("output_modalities") or payload.get("modalities")) + payload.setdefault( + "audio", + { + "input": { + "format": realtime_audio_format_object( + state.input_audio_format, sample_rate_hz=state.input_sample_rate_hz + ), + "sample_rate_hz": state.input_sample_rate_hz, + }, + "output": { + "format": realtime_audio_format_object( + state.output_audio_format, sample_rate_hz=state.output_sample_rate_hz + ), + }, + }, + ) + payload.setdefault("turn_detection", payload.get("turn_detection")) + payload.setdefault("input_audio_transcription", payload.get("input_audio_transcription")) + payload.setdefault("tracing", payload.get("tracing")) + return payload + + +# ---- main projection ---- + + +def project_internal_event(state: RealtimeProjectionState, event: Mapping[str, object]) -> list[DuplexEvent]: + """Project one session-internal event onto 0..n typed public events.""" + return _project(state, dict(event)) + + +def _project(state: RealtimeProjectionState, event: dict[str, object]) -> list[DuplexEvent]: + event_type = event.get("type") + if event_type == "session.created": + session = realtime_session_payload(state, event.get("session")) + events: list[DuplexEvent] = [ + SessionCreated( + session=session, + attachment_generation=( + cast("int", event["attachment_generation"]) + if isinstance(event.get("attachment_generation"), int) + else None + ), + resume_token=_str_or_none(event.get("resume_token")), + ) + ] + if state.initial_session_update: + events.append(SessionUpdated(session=session)) + state.initial_session_update = False + return events + if event_type == "session.updated": + return [SessionUpdated(session=realtime_session_payload(state, event.get("session")))] + if event_type == "session.resumed": + return [ + SessionResumed( + session=realtime_session_payload(state, event.get("session")), + attachment_generation=( + cast("int", event["attachment_generation"]) + if isinstance(event.get("attachment_generation"), int) + else None + ), + resume_token=_str_or_none(event.get("resume_token")), + ) + ] + if event_type == "session.replaced": + return [SessionReplaced(attachment_generation=_int_or(event.get("attachment_generation")))] + if event_type == "session.resync_required": + return [SessionResyncRequired(reason=str(event.get("reason") or "journal_gap"))] + if event_type == "session.closed": + return [SessionClosed(reason=str(event.get("reason") or "closed"), details=event)] + if event_type == "response.created": + response_id = event.get("response_id") + if isinstance(response_id, str) and response_id: + state.active_response_id = response_id + state.last_response_id = response_id + item_id = _response_item_id(state, response_id) + modalities = event.get("modalities") + has_audio_modality = not isinstance(modalities, list) or "audio" in modalities + item = { + "id": item_id, + "object": "realtime.item", + "type": "message", + "role": "assistant", + "status": "in_progress", + "content": [], + } + state.conversation_items[item_id] = item + events = [ + _response_created_event(event), + *_conversation_item_added_events(state, item), + OutputItemAdded(response_id=_str_or_none(response_id), item=item), + ] + if has_audio_modality: + events.extend(_ensure_response_audio_part_added(state, response_id)) + return events + if event_type == "response.listen": + return [Listen(response_id=_str_or_none(event.get("response_id")), details=event)] + if event_type == "response.speak": + response_id = event.get("response_id") + projection = _response_state(state, response_id) + if projection is not None: + if projection.speak_emitted: + return [] + projection.speak_emitted = True + return [ + Speak( + response_id=_str_or_none(response_id), + item_id=_response_item_id(state, response_id), + metadata=_response_speak_metadata(event), + ) + ] + if event_type == "response.output_audio.delta": + response_id = event.get("response_id") + audio = event.get("audio", "") + events = [] + if isinstance(audio, str) and audio: + events.extend(_ensure_response_audio_part_added(state, response_id)) + events.extend(_realtime_audio_delta_events(state, event, response_id, audio)) + _refresh_in_progress_response_item(state, response_id) + text = event.get("text") + has_text = isinstance(text, str) and bool(text) + if has_text: + _append_response_transcript(state, response_id, cast("str", text)) + _refresh_in_progress_response_item(state, response_id) + # Keep the audio.delta + transcript.delta pair invariant even for + # text-less units so clients that treat the pair as unit-complete work. + if has_text or (isinstance(audio, str) and bool(audio)): + events.append( + TranscriptDelta( + response_id=_str_or_none(response_id), + item_id=_response_item_id(state, response_id), + delta=cast("str", text) if has_text else "", + ) + ) + if event.get("end_of_turn") is True: + events.extend(_realtime_audio_done_events(state, event, response_id)) + events.extend( + _realtime_response_terminal_events( + state, + event, + response_id, + status="completed", + status_details={"type": "completed", "reason": event.get("finish_reason") or "stop"}, + ) + ) + return events + if event_type == "response.text.delta": + response_id = event.get("response_id") + text = event.get("delta", "") + projection = _response_state(state, response_id) + if projection is not None and isinstance(text, str) and text: + projection.text_parts.append(text) + _refresh_in_progress_response_item(state, response_id) + events = _ensure_response_text_part_added(state, response_id) + events.append( + TextDelta( + response_id=_str_or_none(response_id), + item_id=_response_item_id(state, response_id), + content_index=1 if projection is not None and projection.audio_part_added else 0, + delta=str(text) if isinstance(text, str) else "", + ) + ) + return events + if event_type == "response.done": + response_id = event.get("response_id") + status = cast("str", event.get("status")) if isinstance(event.get("status"), str) else "completed" + status_details = ( + cast("dict[str, object]", event.get("status_details")) + if isinstance(event.get("status_details"), dict) + else None + ) + return [ + *_realtime_audio_done_events(state, event, response_id), + *_realtime_response_terminal_events( + state, event, response_id, status=status, status_details=status_details + ), + ] + if event_type == "input.committed": + event_item_id = event.get("realtime_item_id") + item_id = ( + event_item_id if isinstance(event_item_id, str) and event_item_id else _pop_pending_commit_item_id(state) + ) + item = state.conversation_items.get(item_id) + events = [] + if item is None: + message = event.get("message") + no_response = event.get("no_response") is True + is_speech = event.get("is_speech") + item = { + "id": item_id, + "object": "realtime.item", + "type": "message", + "role": "user", + "status": "completed", + "content": ( + [{"type": "input_audio", "transcript": "", "is_speech": False}] + if no_response and is_speech is False + else _user_item_content_from_duplex_message(message) + ), + } + state.conversation_items[item_id] = item + events.extend(_conversation_item_added_events(state, item)) + item["status"] = "completed" + events.append( + InputCommitted(previous_item_id=_previous_item_id(state, item_id), item_id=item_id, details=event) + ) + transcription_event = _input_audio_transcription_completed_event(item_id, item) + if transcription_event is not None: + events.append(transcription_event) + events.append(_conversation_item_done_event(state, item)) + return events + if event_type == "input.cancelled": + return [InputCleared()] + if event_type == "audio.cancelled": + response_id = event.get("response_id") + events = [] + if event.get("reason") == "output_audio_buffer_clear": + if not isinstance(response_id, str) or not response_id: + response_id = state.active_response_id or state.last_response_id + events.append(OutputAudioCleared(response_id=_str_or_none(response_id))) + if not isinstance(response_id, str) or not response_id: + return events + elif not isinstance(response_id, str) or not response_id: + response_id = state.active_response_id + if not isinstance(response_id, str) or not response_id: + return events + if response_is_done(state, response_id): + return events + committed_ms = event.get("committed_ms") + if isinstance(committed_ms, int | float): + item_id = _response_item_id(state, response_id) + committed_audio_ms = max(0, int(committed_ms)) + state.item_truncation_cursors[item_id] = (0, committed_audio_ms) + item = state.conversation_items.get(item_id) + if item is not None: + truncate_realtime_item_content(item, content_index=0, audio_end_ms=committed_audio_ms) + events.extend(_realtime_audio_done_events(state, event, response_id)) + events.extend( + _realtime_response_terminal_events( + state, + event, + response_id, + status="cancelled", + status_details={"type": "cancelled", "reason": event.get("reason") or "client_cancelled"}, + ) + ) + if isinstance(response_id, str) and response_id == state.active_response_id: + state.active_response_id = None + return events + if event_type == "conversation.item.created": + item = event.get("item") + if isinstance(item, dict) and isinstance(item.get("id"), str): + item_id = str(item["id"]) + already_known = item_id in state.conversation_items + state.conversation_items[item_id] = item + if already_known: + if item.get("status") == "completed": + return [_conversation_item_done_event(state, item)] + return [] + events = _conversation_item_added_events(state, item) + if item.get("status") == "completed": + events.append(_conversation_item_done_event(state, item)) + return events + return [ItemCreated(item=item if isinstance(item, Mapping) else {})] + if event_type == "conversation.item.deleted": + item_id = event.get("item_id") + if isinstance(item_id, str): + _remove_conversation_item(state, item_id) + return [ItemDeleted(item_id=_str_or_none(item_id), details=event)] + if event_type == "conversation.item.truncated": + item_id = event.get("item_id") + audio_end_ms = event.get("audio_end_ms") + content_index = event.get("content_index", 0) + if isinstance(item_id, str): + item = state.conversation_items.get(item_id) + if item is not None: + truncate_realtime_item_content( + item, + content_index=_int_or(content_index), + audio_end_ms=_int_or(audio_end_ms), + ) + return [ + ItemTruncated( + item_id=_str_or_none(item_id), + content_index=_int_or(content_index), + audio_end_ms=_int_or(audio_end_ms), + details=event, + ) + ] + if event_type == "conversation.item.retrieved": + item = event.get("item") + return [ItemRetrieved(item=item if isinstance(item, Mapping) else {})] + if event_type == "response.output_item.done": + response_id = event.get("response_id") + return _realtime_response_terminal_events( + state, + event, + response_id, + status="completed", + status_details={"type": "completed", "reason": event.get("finish_reason") or "stop"}, + ) + if event_type == "function_call.done": + return _function_call_done_events(state, event) + if event_type == "input_audio_buffer.speech_started": + return [ + SpeechStarted( + audio_start_ms=_int_or(event.get("audio_start_ms")), item_id=_str_or_none(event.get("item_id")) + ) + ] + if event_type == "input_audio_buffer.speech_stopped": + return [ + SpeechStopped(audio_end_ms=_int_or(event.get("audio_end_ms")), item_id=_str_or_none(event.get("item_id"))) + ] + return [DuplexRawEvent(internal_type=str(event_type), details=event)] + + +def _function_call_done_events(state: RealtimeProjectionState, event: dict[str, object]) -> list[DuplexEvent]: + call_id = event.get("call_id") + name = event.get("name") + arguments = event.get("arguments", "") + if not isinstance(call_id, str) or not call_id or not isinstance(name, str) or not name: + return [DuplexRawEvent(internal_type="function_call.done", details=event)] + if not isinstance(arguments, str): + arguments = str(arguments) + response_id = f"resp_{uuid4().hex}" + item_id = f"item_{uuid4().hex}" + item = { + "id": item_id, + "object": "realtime.item", + "type": "function_call", + "status": "completed", + "name": name, + "call_id": call_id, + "arguments": arguments, + } + state.conversation_items[item_id] = item + response = { + "id": response_id, + "object": "realtime.response", + "status": "completed", + "status_details": {"type": "completed", "reason": "completed"}, + "output": [dict(item)], + "metadata": {}, + } + return [ + ResponseCreated(response_id=response_id, response={**response, "status": "in_progress", "output": []}), + *_conversation_item_added_events(state, item), + OutputItemAdded(response_id=response_id, item=dict(item)), + FunctionCallArgumentsDelta(response_id=response_id, item_id=item_id, call_id=call_id, delta=arguments), + FunctionCallArgumentsDone(response_id=response_id, item_id=item_id, call_id=call_id, arguments=arguments), + OutputItemDone(response_id=response_id, item=dict(item)), + _conversation_item_done_event(state, item), + ResponseDone(response_id=response_id, response=response), + ] + + +# ---- input-side helpers for the runner (were on the input translator) ---- + + +def register_user_item( + state: RealtimeProjectionState, + item: Mapping[str, object], + *, + previous_item_id: str | None = None, +) -> list[DuplexEvent]: + """Record a client-created user message item and return its ack events. + + Replaces the old translator's ``_send_realtime_input_ack``: call it when + handling a ``CreateItem`` whose item role is ``user`` (before the internal + ``conversation.item.created`` is emitted, which then projects to ``done``). + """ + normalized = dict(item) + if isinstance(previous_item_id, str): + normalized["_previous_item_id"] = previous_item_id + if normalized.get("role") != "user": + return [] + state.conversation_items[str(normalized["id"])] = normalized + return _conversation_item_added_events(state, normalized) + + +def retrieve_item_events(state: RealtimeProjectionState, payload: Mapping[str, object]) -> list[DuplexEvent]: + """Answer ``conversation.item.retrieve`` from the projected conversation items.""" + event_id = payload.get("event_id") + item_id = payload.get("item_id") + if not isinstance(item_id, str) or not item_id: + return [error_event("missing_item_id", "conversation.item.retrieve requires item_id", event_id=event_id)] + item = state.conversation_items.get(item_id) + if item is None: + return [error_event("item_not_found", f"Conversation item not found: {item_id}", event_id=event_id)] + return [ItemRetrieved(item=item)] + + +def emit_input_speech_started(state: RealtimeProjectionState, audio_start_ms: object = 0) -> list[DuplexEvent]: + if state.input_speech_started: + return [] + state.input_speech_started = True + if state.active_input_item_id is None: + state.active_input_item_id = f"item_{uuid4().hex}" + return [SpeechStarted(audio_start_ms=_int_or(audio_start_ms), item_id=state.active_input_item_id)] + + +def emit_input_speech_stopped( + state: RealtimeProjectionState, *, item_id: str, audio_end_ms: object = 0 +) -> list[DuplexEvent]: + if not state.input_speech_started: + return [] + state.input_speech_started = False + return [SpeechStopped(audio_end_ms=_int_or(audio_end_ms), item_id=item_id)] + + +def _remember_input_transcript_hint(state: RealtimeProjectionState, event: Mapping[str, object]) -> None: + transcript = event.get("transcript") + if not isinstance(transcript, str): + transcript = event.get("text") if isinstance(event.get("text"), str) else None + if not isinstance(transcript, str): + hints = event.get("hints") + if isinstance(hints, dict): + transcript = hints.get("transcript") + if not isinstance(transcript, str): + transcript = hints.get("text") if isinstance(hints.get("text"), str) else None + if isinstance(transcript, str) and transcript: + parts = state.input_audio_buffer_transcript_parts + if parts and parts[-1] == transcript: + return + parts.append(transcript) + + +def _consume_input_transcript_hint(state: RealtimeProjectionState) -> str: + transcript = "".join(state.input_audio_buffer_transcript_parts).strip() + state.input_audio_buffer_transcript_parts.clear() + return transcript + + +def note_input_append( + state: RealtimeProjectionState, + payload: dict[str, object], + *, + vad_result: TurnDetectionResult | None = None, +) -> list[DuplexEvent]: + """Update the input-buffer projection for one appended chunk; returns typed events. + + ``payload`` is the internal ``input_audio_buffer.append`` dictionary + (``AppendAudio.payload()``), optionally after ``apply_turn_detection_result``; + ``vad_result`` is the ``TurnDetectionResult`` when server VAD is active. The + returned events are ``input_audio_buffer.speech_started`` / + ``speech_stopped`` exactly as the old translator produced them. + """ + audio = payload.get("audio") + looks_like_speech = bool(vad_result.is_speech) if vad_result is not None else payload.get("is_speech") is not False + has_audio = isinstance(audio, str) and bool(audio) + state.input_audio_buffer_has_audio = state.input_audio_buffer_has_audio or (looks_like_speech and has_audio) + state.input_audio_buffer_had_non_speech = state.input_audio_buffer_had_non_speech or ( + not looks_like_speech and has_audio + ) + events: list[DuplexEvent] = [] + stop_ms: object = payload.get("audio_end_ms", payload.get("audio_ms", 0)) + if vad_result is not None: + if vad_result.speech_stopped and vad_result.audio_end_ms is not None: + stop_ms = vad_result.audio_end_ms + if vad_result.speech_stopped and vad_result.speech_active: + events.extend( + emit_input_speech_stopped( + state, + item_id=state.active_input_item_id or f"item_{uuid4().hex}", + audio_end_ms=stop_ms, + ) + ) + if vad_result.speech_started: + start_ms = ( + vad_result.audio_start_ms if vad_result.audio_start_ms is not None else payload.get("audio_start_ms", 0) + ) + events.extend(emit_input_speech_started(state, start_ms)) + if vad_result.speech_stopped and not vad_result.speech_active: + events.extend( + emit_input_speech_stopped( + state, + item_id=state.active_input_item_id or f"item_{uuid4().hex}", + audio_end_ms=stop_ms, + ) + ) + elif looks_like_speech: + events.extend(emit_input_speech_started(state, payload.get("audio_start_ms", 0))) + if looks_like_speech: + _remember_input_transcript_hint(state, payload) + return events + + +def discard_pending_input_audio(state: RealtimeProjectionState, audio_end_ms: int | None = None) -> list[DuplexEvent]: + """Drop Realtime input-buffer state that was consumed as overlap (returns typed events).""" + events: list[DuplexEvent] = [] + if state.input_speech_started and state.active_input_item_id is not None: + events.append(SpeechStopped(audio_end_ms=max(0, int(audio_end_ms or 0)), item_id=state.active_input_item_id)) + state.input_speech_started = False + state.active_input_item_id = None + state.input_audio_buffer_has_audio = False + state.input_audio_buffer_had_non_speech = False + state.input_audio_buffer_transcript_parts.clear() + return events + + +def clear_input_buffer(state: RealtimeProjectionState) -> None: + """Input-buffer projection reset for ``input_audio_buffer.clear``.""" + state.input_speech_started = False + state.active_input_item_id = None + state.input_audio_buffer_has_audio = False + state.input_audio_buffer_had_non_speech = False + state.input_audio_buffer_transcript_parts.clear() + + +@dataclass(frozen=True, slots=True) +class ResolvedCommit: + """What a ``Commit`` means given the projected input buffer.""" + + #: Internal ``input_audio_buffer.commit`` payload to run, or None when rejected. + payload: dict[str, object] | None + #: Events to emit before running (speech_stopped) or instead (error). + events: list[DuplexEvent] + #: True when the VAD detector should be reset. + reset_vad: bool = False + + +def resolve_commit(state: RealtimeProjectionState, command: Commit) -> ResolvedCommit: + """Apply the old translator's commit rules (empty buffer, non-speech commit, item id).""" + if command.realtime_item_id is None and not state.input_audio_buffer_has_audio: + if state.input_audio_buffer_had_non_speech: + state.input_audio_buffer_had_non_speech = False + state.active_input_item_id = None + payload: dict[str, object] = { + "type": "input_audio_buffer.commit", + "final": command.final, + "response_create": False, + "is_speech": False, + } + if command.event_id is not None: + payload["realtime_event_id"] = command.event_id + return ResolvedCommit(payload=payload, events=[], reset_vad=True) + return ResolvedCommit( + payload=None, + events=[ + error_event( + "input_audio_buffer_empty", + "input_audio_buffer.commit requires a non-empty input audio buffer", + event_id=command.event_id, + ) + ], + ) + item_id = command.realtime_item_id or state.active_input_item_id or f"item_{uuid4().hex}" + state.pending_commit_item_ids.append(item_id) + events = emit_input_speech_stopped(state, item_id=item_id, audio_end_ms=0) + transcript = _consume_input_transcript_hint(state) + state.active_input_item_id = None + state.input_audio_buffer_has_audio = False + state.input_audio_buffer_had_non_speech = False + payload = { + "type": "input_audio_buffer.commit", + "final": command.final, + "realtime_item_id": item_id, + "response_create": bool(command.create_response) if command.create_response is not None else False, + } + if command.is_speech is not None: + payload["is_speech"] = command.is_speech + if transcript: + payload["transcript"] = transcript + if command.event_id is not None: + payload["realtime_event_id"] = command.event_id + return ResolvedCommit(payload=payload, events=events, reset_vad=True) + + +@dataclass(frozen=True, slots=True) +class ResolvedControl: + """Result of resolving a cancel/clear/delete/truncate against the projection.""" + + #: Internal payload(s) to run in order (empty when nothing to do). + payloads: list[dict[str, object]] + #: Events to emit instead (already-done response ack, errors). + events: list[DuplexEvent] + + +def resolve_cancel_response(state: RealtimeProjectionState, command: CancelResponse) -> ResolvedControl: + payload: dict[str, object] = {"type": "response.cancel", "reason": "response.cancel"} + response_id = command.response_id or state.active_response_id or state.last_response_id + if response_is_done(state, response_id): + return ResolvedControl( + payloads=[], + events=[ + error_event( + "response_not_active", f"Response is already complete: {response_id}", event_id=command.event_id + ) + ], + ) + if isinstance(response_id, str) and response_id: + payload["response_id"] = response_id + if command.event_id is not None: + payload["realtime_event_id"] = command.event_id + return ResolvedControl(payloads=[payload], events=[]) + + +def resolve_clear_output_audio(state: RealtimeProjectionState, command: ClearOutputAudio) -> ResolvedControl: + payload: dict[str, object] = {"type": "output_audio_buffer.clear", "reason": "output_audio_buffer.clear"} + response_id = command.response_id or state.active_response_id or state.last_response_id + if response_is_done(state, response_id): + return ResolvedControl(payloads=[], events=[OutputAudioCleared(response_id=_str_or_none(response_id))]) + if isinstance(response_id, str) and response_id: + payload["response_id"] = response_id + if command.event_id is not None: + payload["realtime_event_id"] = command.event_id + return ResolvedControl(payloads=[payload], events=[]) + + +def resolve_delete_item(state: RealtimeProjectionState, command: DeleteItem) -> ResolvedControl: + if command.item_id not in state.conversation_items: + return ResolvedControl( + payloads=[], + events=[ + error_event( + "item_not_found", f"Conversation item not found: {command.item_id}", event_id=command.event_id + ) + ], + ) + _remove_conversation_item(state, command.item_id) + return ResolvedControl(payloads=[command.payload()], events=[]) + + +def resolve_truncate_item(state: RealtimeProjectionState, command: TruncateItem) -> ResolvedControl: + item = state.conversation_items.get(command.item_id) + if not isinstance(item, dict): + return ResolvedControl( + payloads=[], + events=[ + error_event( + "item_not_found", f"Conversation item not found: {command.item_id}", event_id=command.event_id + ) + ], + ) + truncate_error = validate_realtime_item_truncate( + item, content_index=command.content_index, audio_end_ms=command.audio_end_ms + ) + if truncate_error is not None: + return ResolvedControl( + payloads=[], events=[error_event("bad_event", truncate_error, event_id=command.event_id)] + ) + state.item_truncation_cursors[command.item_id] = (command.content_index, command.audio_end_ms) + truncate_realtime_item_content(item, content_index=command.content_index, audio_end_ms=command.audio_end_ms) + ack_payload: dict[str, object] = { + "type": "playback.ack", + "item_id": command.item_id, + "committed_ms": int(command.audio_end_ms), + "played_ms": int(command.audio_end_ms), + "truncate": True, + } + return ResolvedControl(payloads=[command.payload(), ack_payload], events=[]) + + +def _duplicate_function_call_output(state: RealtimeProjectionState, call_id: object) -> tuple[bool, bool]: + matching_call = any( + known.get("type") == "function_call" and known.get("call_id") == call_id + for known in state.conversation_items.values() + ) + duplicate_output = any( + known.get("type") == "function_call_output" and known.get("call_id") == call_id + for known in state.conversation_items.values() + ) + return matching_call, duplicate_output + + +def resolve_create_item(state: RealtimeProjectionState, command: CreateItem) -> ResolvedControl: + """Expand ``conversation.item.create`` the way the old translator did. + + Assistant/system/function items and text-only user items become one + ``turn.signal conversation.item.create`` payload. User items carrying audio + become one ``input_audio_buffer.append`` per audio part followed by an + ``input_audio_buffer.commit`` (no response). The user-item ack events are + returned in ``events`` and must be emitted *before* running the payloads. + """ + item = dict(command.item) + if isinstance(command.previous_item_id, str): + item["_previous_item_id"] = command.previous_item_id + item_id = str(item["id"]) + item_type = item.get("type") + role = item.get("role") + if item_type == "function_call_output": + call_id = item.get("call_id") + output = item.get("output") + matching_call, duplicate_output = _duplicate_function_call_output(state, call_id) + if not isinstance(call_id, str) or not call_id or not matching_call: + return ResolvedControl( + payloads=[], + events=[ + error_event( + "invalid_function_call_output", + "function_call_output requires the call_id of a completed function call", + event_id=command.event_id, + param="item.call_id", + ) + ], + ) + if not isinstance(output, str) or duplicate_output: + return ResolvedControl( + payloads=[], + events=[ + error_event( + "invalid_function_call_output", + ( + "function_call_output requires a string output" + if not isinstance(output, str) + else f"function_call_output already exists for call_id {call_id}" + ), + event_id=command.event_id, + param="item.output", + ) + ], + ) + ack_events = register_user_item(state, item) if item_type == "message" and role == "user" else [] + signal_payload = { + "type": "turn.signal", + "event": "conversation.item.create", + "payload": {"item": item}, + } + if command.event_id is not None: + signal_payload["realtime_event_id"] = command.event_id + if item_type != "message" or role in {"assistant", "system"}: + return ResolvedControl(payloads=[signal_payload], events=ack_events) + content = item.get("content") + if not isinstance(content, list): + return ResolvedControl(payloads=[], events=ack_events) + text_chunks: list[str] = [] + audio_payloads: list[dict[str, object]] = [] + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") in {"input_text", "text"} and isinstance(part.get("text"), str): + text_chunks.append(str(part["text"])) + if part.get("type") in {"input_audio", "audio"}: + audio = part.get("audio") or part.get("data") + if not isinstance(audio, str) or not audio: + continue + try: + append = build_append_audio( + { + "audio": audio, + "format": part.get("format"), + "sample_rate_hz": part.get("sample_rate_hz") or part.get("sample_rate"), + "event_id": command.event_id, + }, + defaults=state.defaults, + hints_source={**item, **part}, + ) + except Exception as exc: # DuplexCommandError from validation/conversion + code = getattr(exc, "code", "bad_event") + if code == "unsupported_audio_format": + continue + return ResolvedControl( + payloads=[], + events=[error_event(code, str(exc), event_id=command.event_id, param="sample_rate_hz")], + ) + speech_hints = {**item, **part} + if not input_looks_like_speech( + speech_hints, audio=append.audio, fmt=append.format, overlap_silence_rms=state.overlap_silence_rms + ): + continue + state.input_speech_started = True + payload = append.payload() + copy_realtime_input_hints(part, payload) + copy_realtime_input_hints(item, payload) + audio_payloads.append(payload) + if audio_payloads: + state.input_audio_buffer_has_audio = True + transcript = input_transcript_from_item(item) + commit_payload: dict[str, object] = { + "type": "input_audio_buffer.commit", + "final": True, + "realtime_item_id": item_id, + "response_create": False, + } + if transcript: + commit_payload["transcript"] = transcript + return ResolvedControl(payloads=[*audio_payloads, commit_payload], events=ack_events) + if not text_chunks: + return ResolvedControl(payloads=[], events=ack_events) + text_item = dict(item) + text_item["status"] = "completed" + signal_payload["payload"] = {"item": text_item} + return ResolvedControl(payloads=[signal_payload], events=ack_events) + + +def append_audio_payload(command: AppendAudio) -> dict[str, object]: + """Internal ``input_audio_buffer.append`` payload for a command (convenience for the runner).""" + return command.payload() + + +__all__ = [ + "RealtimeProjectionState", + "ResolvedCommit", + "ResolvedControl", + "append_audio_payload", + "clear_input_buffer", + "discard_pending_input_audio", + "emit_input_speech_started", + "emit_input_speech_stopped", + "note_input_append", + "project_internal_event", + "realtime_session_payload", + "register_user_item", + "resolve_cancel_response", + "resolve_clear_output_audio", + "resolve_commit", + "resolve_create_item", + "resolve_delete_item", + "resolve_truncate_item", + "response_is_done", + "retrieve_item_events", +] diff --git a/vllm_omni/engine/duplex/runtime.py b/vllm_omni/engine/duplex/runtime.py deleted file mode 100644 index dede0e0a6..000000000 --- a/vllm_omni/engine/duplex/runtime.py +++ /dev/null @@ -1,108 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -"""Experimental duplex runtime loading and compatibility exports.""" - -from __future__ import annotations - -from importlib import import_module - -from vllm_omni.engine.duplex.contracts import ( - DuplexAppendPlan, - DuplexInputMode, - DuplexOutputAction, - DuplexOutputDecision, - DuplexRuntimeCapabilities, - DuplexRuntimeExtension, - SessionMode, - duplex_data_plane_request_info, - duplex_resource_request_belongs_to_session, - duplex_resource_request_id, -) - - -def load_duplex_runtime_extension(path: str | None) -> DuplexRuntimeExtension | None: - if not path: - return None - module_name, separator, attribute_name = path.rpartition(".") - if not separator: - raise ValueError(f"Invalid duplex runtime extension path: {path!r}") - extension_type = getattr(import_module(module_name), attribute_name) - return validate_duplex_runtime_extension(extension_type()) - - -def validate_duplex_runtime_extension( - extension: object, - *, - sampling_defaults: tuple[object, ...] | None = None, -) -> DuplexRuntimeExtension: - required_methods = ( - "configure_sampling_params", - "plan_append", - "decide_output", - ) - missing = [name for name in required_methods if not callable(getattr(extension, name, None))] - if missing: - raise TypeError(f"Duplex runtime extension is missing callable method(s): {', '.join(missing)}") - typed_extension = extension # type: ignore[assignment] - if sampling_defaults is not None: - configured = typed_extension.configure_sampling_params( - runtime_config={}, - defaults=sampling_defaults, - ) - if not isinstance(configured, tuple): - raise TypeError("Duplex runtime extension must return sampling parameters as a tuple") - if len(configured) != len(sampling_defaults): - raise ValueError("Duplex runtime extension must return one sampling parameter per stage") - for stage_id, (value, default) in enumerate(zip(configured, sampling_defaults, strict=True)): - if default is not None and not isinstance(value, type(default)): - raise TypeError( - "Duplex runtime extension sampling parameter type mismatch " - f"for stage {stage_id}: expected {type(default).__name__}, got {type(value).__name__}" - ) - return typed_extension - - -from vllm_omni.engine.duplex.lease import ( # noqa: E402, F401 - DuplexLeaseActivity, - DuplexLeaseConfig, - DuplexLeaseState, - DuplexSessionExpiry, -) -from vllm_omni.engine.duplex.session import ( # noqa: E402, F401 - DuplexAppendReservation, - DuplexCompletedAppend, - DuplexFenceMismatchError, - DuplexInputAppend, - DuplexRequestResource, - DuplexSessionRuntimeManager, - DuplexSessionRuntimeState, - DuplexStageBinding, -) - -__all__ = [ - "DuplexAppendPlan", - "DuplexAppendReservation", - "DuplexCompletedAppend", - "DuplexFenceMismatchError", - "DuplexInputAppend", - "DuplexInputMode", - "DuplexLeaseActivity", - "DuplexLeaseConfig", - "DuplexLeaseState", - "DuplexOutputAction", - "DuplexOutputDecision", - "DuplexRequestResource", - "DuplexRuntimeCapabilities", - "DuplexRuntimeExtension", - "DuplexSessionExpiry", - "DuplexSessionRuntimeManager", - "DuplexSessionRuntimeState", - "DuplexStageBinding", - "SessionMode", - "duplex_data_plane_request_info", - "duplex_resource_request_belongs_to_session", - "duplex_resource_request_id", - "load_duplex_runtime_extension", - "validate_duplex_runtime_extension", -] diff --git a/vllm_omni/engine/duplex/session.py b/vllm_omni/engine/duplex/session.py deleted file mode 100644 index b7ad51504..000000000 --- a/vllm_omni/engine/duplex/session.py +++ /dev/null @@ -1,534 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import time -from collections import OrderedDict -from collections.abc import Callable, Iterable -from dataclasses import dataclass, field -from typing import Any - -from vllm_omni.engine.duplex.contracts import ( - DuplexInputMode, - DuplexRuntimeCapabilities, -) -from vllm_omni.engine.duplex.lease import ( - DuplexLeaseActivity, - DuplexLeaseConfig, - DuplexLeaseState, - DuplexSessionExpiry, -) -from vllm_omni.engine.duplex.messages import DuplexFence - - -def _default_capabilities() -> DuplexRuntimeCapabilities: - return DuplexRuntimeCapabilities() - - -class DuplexFenceMismatchError(RuntimeError): - def __init__(self, expected: DuplexFence, actual: DuplexFence) -> None: - super().__init__(f"duplex fence mismatch: expected {expected!r}, got {actual!r}") - self.expected = expected - self.actual = actual - - -@dataclass -class DuplexStageBinding: - request_id: str - fence: DuplexFence - - -@dataclass -class DuplexRequestResource: - stage_id: int - request_id: str - fence: DuplexFence - submitted: bool = False - - -@dataclass -class DuplexInputAppend: - seq: int - turn_seq: int - turn_id: int - - -@dataclass(frozen=True) -class DuplexAppendReservation: - fence: DuplexFence - mode: DuplexInputMode - base_fence: DuplexFence - base_input_seq: int - base_input_turn_seq: int - base_append_turn_key: tuple[int, int, int] | None - update: DuplexInputAppend - - -@dataclass(frozen=True) -class DuplexCompletedAppend: - fence: DuplexFence - mode: DuplexInputMode - final: bool - stage_results: tuple[dict[str, object], ...] - - -@dataclass -class DuplexSessionRuntimeState: - """Engine resource handles associated with a core-owned identity fence.""" - - fence: DuplexFence - lease: DuplexLeaseState - _clock: Callable[[], float] = field(repr=False) - capabilities: DuplexRuntimeCapabilities = field(default_factory=_default_capabilities) - session_config: dict[str, Any] = field(default_factory=dict) - runtime_config: dict[str, Any] = field(default_factory=dict) - config_generation: int = 0 - stage_bindings: dict[int, DuplexStageBinding] = field(default_factory=dict) - request_resources: dict[tuple[int, str], DuplexRequestResource] = field(default_factory=dict) - input_seq: int = 0 - input_turn_seq: int = 0 - _append_turn_key: tuple[int, int, int] | None = None - completed_append_limit: int = 256 - completed_appends: OrderedDict[str, DuplexCompletedAppend] = field(default_factory=OrderedDict) - - def __post_init__(self) -> None: - if self.completed_append_limit <= 0: - raise ValueError("completed_append_limit must be positive") - - @property - def session_id(self) -> str: - return self.fence.session_id - - @property - def epoch(self) -> int: - return self.fence.epoch - - @property - def turn_id(self) -> int: - return self.fence.turn_id - - def _validate_fence(self, fence: DuplexFence) -> None: - if fence.session_id != self.session_id or fence.incarnation != self.fence.incarnation: - raise DuplexFenceMismatchError(self.fence, fence) - current = self.fence - if fence.epoch < current.epoch or ( - fence.epoch == current.epoch - and (fence.turn_id < current.turn_id or fence.response_seq < current.response_seq) - ): - raise DuplexFenceMismatchError(current, fence) - - def accept_fence(self, fence: DuplexFence) -> None: - self._validate_fence(fence) - if fence.epoch != self.fence.epoch: - self.input_seq = 0 - self.input_turn_seq = 0 - self._append_turn_key = None - self.completed_appends.clear() - self.fence = fence - - def touch(self, fence: DuplexFence, activity: DuplexLeaseActivity) -> None: - self._validate_fence(fence) - self.lease.touch(self._clock(), activity) - - def detach(self, fence: DuplexFence) -> None: - self._validate_fence(fence) - self.lease.detach(self._clock()) - - def resume(self, fence: DuplexFence, *, expected_lease_generation: int) -> int: - self._validate_fence(fence) - return self.lease.resume( - self._clock(), - expected_generation=expected_lease_generation, - ) - - def begin_operation(self, fence: DuplexFence, operation_id: str) -> None: - self._validate_fence(fence) - self.lease.begin_operation(self._clock(), operation_id) - - def end_operation(self, fence: DuplexFence, operation_id: str) -> None: - self._validate_fence(fence) - self.lease.end_operation(self._clock(), operation_id) - - def replace_session_config(self, session_config: dict[str, Any]) -> None: - self.session_config = dict(session_config) - self.config_generation += 1 - - def replace_runtime_config(self, runtime_config: dict[str, Any]) -> None: - self.runtime_config = dict(runtime_config) - self.config_generation += 1 - - def replace_configs( - self, - *, - session_config: dict[str, Any] | None = None, - runtime_config: dict[str, Any] | None = None, - ) -> None: - """Atomically publish one validated configuration generation.""" - if session_config is None and runtime_config is None: - return - next_session_config = self.session_config if session_config is None else dict(session_config) - next_runtime_config = self.runtime_config if runtime_config is None else dict(runtime_config) - self.session_config = next_session_config - self.runtime_config = next_runtime_config - self.config_generation += 1 - - def reserve_stage_request(self, stage_id: int, request_id: str, *, fence: DuplexFence) -> None: - self._validate_fence(fence) - resource_key = (stage_id, request_id) - existing = self.request_resources.get(resource_key) - if existing is not None: - if ( - existing.fence.session_id != fence.session_id - or existing.fence.incarnation != fence.incarnation - or existing.fence.epoch != fence.epoch - ): - raise ValueError(f"Duplex request resource already reserved with different identity: {request_id}") - return - self.request_resources[resource_key] = DuplexRequestResource( - stage_id=stage_id, - request_id=request_id, - fence=fence, - ) - - def bind_stage_request(self, stage_id: int, request_id: str, *, fence: DuplexFence) -> None: - self.reserve_stage_request(stage_id, request_id, fence=fence) - self.accept_fence(fence) - resource = self.request_resources[(stage_id, request_id)] - resource.fence = fence - resource.submitted = True - self.stage_bindings[stage_id] = DuplexStageBinding(request_id=request_id, fence=fence) - - def stage_request_ids(self, fence: DuplexFence | None = None) -> list[str]: - return [ - binding.request_id for binding in self.stage_bindings.values() if fence is None or binding.fence == fence - ] - - def resource_request_ids( - self, - fence: DuplexFence | None = None, - *, - submitted: bool | None = None, - ) -> list[str]: - return list( - dict.fromkeys( - resource.request_id - for resource in self.request_resources.values() - if (fence is None or resource.fence == fence) and (submitted is None or resource.submitted is submitted) - ) - ) - - def release_request_ids(self, request_ids: list[str]) -> None: - released = set(request_ids) - if not released: - return - self.stage_bindings = { - stage_id: binding for stage_id, binding in self.stage_bindings.items() if binding.request_id not in released - } - self.request_resources = { - resource_key: resource - for resource_key, resource in self.request_resources.items() - if resource.request_id not in released - } - - def prepare_append(self, *, mode: DuplexInputMode, fence: DuplexFence) -> DuplexAppendReservation: - if mode not in self.capabilities.input_modes: - raise ValueError(f"Duplex input mode {mode.value!r} is not supported by session {self.session_id}") - self._validate_fence(fence) - input_seq = 0 if fence.epoch != self.fence.epoch else self.input_seq - input_turn_seq = 0 if fence.epoch != self.fence.epoch else self.input_turn_seq - append_turn_key = None if fence.epoch != self.fence.epoch else self._append_turn_key - turn_key = (fence.epoch, fence.turn_id, fence.response_seq) - turn_seq = input_turn_seq + 1 if turn_key == append_turn_key else 1 - return DuplexAppendReservation( - fence=fence, - mode=mode, - base_fence=self.fence, - base_input_seq=self.input_seq, - base_input_turn_seq=self.input_turn_seq, - base_append_turn_key=self._append_turn_key, - update=DuplexInputAppend( - seq=input_seq + 1, - turn_seq=turn_seq, - turn_id=fence.turn_id, - ), - ) - - def commit_append(self, reservation: DuplexAppendReservation) -> DuplexInputAppend: - if ( - self.fence != reservation.base_fence - or self.input_seq != reservation.base_input_seq - or self.input_turn_seq != reservation.base_input_turn_seq - or self._append_turn_key != reservation.base_append_turn_key - ): - raise RuntimeError("duplex append reservation is stale") - self.accept_fence(reservation.fence) - self.input_seq = reservation.update.seq - self.input_turn_seq = reservation.update.turn_seq - self._append_turn_key = ( - reservation.fence.epoch, - reservation.fence.turn_id, - reservation.fence.response_seq, - ) - return reservation.update - - def completed_append( - self, - operation_id: str, - *, - fence: DuplexFence, - mode: DuplexInputMode, - final: bool, - ) -> list[dict[str, object]] | None: - self._validate_fence(fence) - completed = self.completed_appends.get(operation_id) - if completed is None: - return None - if completed.fence != fence or completed.mode != mode or completed.final != final: - raise ValueError(f"duplex append operation {operation_id!r} was reused with different metadata") - return [dict(result) for result in completed.stage_results] - - def record_completed_append( - self, - operation_id: str, - *, - fence: DuplexFence, - mode: DuplexInputMode, - final: bool, - stage_results: list[dict[str, object]], - ) -> None: - self.completed_appends[operation_id] = DuplexCompletedAppend( - fence=fence, - mode=mode, - final=final, - stage_results=tuple(dict(result) for result in stage_results), - ) - self.completed_appends.move_to_end(operation_id) - while len(self.completed_appends) > self.completed_append_limit: - self.completed_appends.popitem(last=False) - - def append_input(self, *, mode: DuplexInputMode, fence: DuplexFence) -> DuplexInputAppend: - return self.commit_append(self.prepare_append(mode=mode, fence=fence)) - - def release_fence(self, fence: DuplexFence) -> list[str]: - stale = self.resource_request_ids(fence) - self.stage_bindings = { - stage_id: binding for stage_id, binding in self.stage_bindings.items() if binding.fence != fence - } - self.request_resources = { - resource_key: resource - for resource_key, resource in self.request_resources.items() - if resource.fence != fence - } - return stale - - def cancel_fence(self, cancelled_fence: DuplexFence, next_fence: DuplexFence) -> list[str]: - stale = self.prepare_cancel_fence(cancelled_fence, next_fence) - self.release_fence(cancelled_fence) - return stale - - def prepare_cancel_fence(self, cancelled_fence: DuplexFence, next_fence: DuplexFence) -> list[str]: - """Advance the cancellation fence without dropping cleanup records.""" - if cancelled_fence.session_id != self.session_id or cancelled_fence.incarnation != self.fence.incarnation: - raise DuplexFenceMismatchError(self.fence, cancelled_fence) - if ( - next_fence.session_id != self.session_id - or next_fence.incarnation != self.fence.incarnation - or next_fence.epoch <= cancelled_fence.epoch - ): - raise DuplexFenceMismatchError(cancelled_fence, next_fence) - current_key = (self.fence.epoch, self.fence.turn_id, self.fence.response_seq) - cancelled_key = (cancelled_fence.epoch, cancelled_fence.turn_id, cancelled_fence.response_seq) - next_key = (next_fence.epoch, next_fence.turn_id, next_fence.response_seq) - if cancelled_key > current_key: - raise DuplexFenceMismatchError(self.fence, cancelled_fence) - if next_key > current_key: - self.accept_fence(next_fence) - return self.resource_request_ids(cancelled_fence) - - def begin_close(self, fence: DuplexFence, *, reason: str) -> bool: - """Make close irreversible while retaining resources for cleanup retry.""" - self.accept_fence(fence) - if self.lease.terminal_reason is not None: - return True - return self.lease.mark_terminal(reason) - - def finalize_close(self) -> list[str]: - return self.close() - - def close(self, fence: DuplexFence | None = None) -> list[str]: - if fence is not None: - self.accept_fence(fence) - stale = self.resource_request_ids() - self.stage_bindings.clear() - self.request_resources.clear() - return stale - - def terminate(self, fence: DuplexFence, *, reason: str) -> bool: - if not self.begin_close(fence, reason=reason): - return False - self.finalize_close() - return True - - -class DuplexSessionRuntimeManager: - def __init__( - self, - *, - clock: Callable[[], float] | None = None, - max_sessions: int | None = None, - completed_append_limit: int = 256, - ) -> None: - if max_sessions is not None and max_sessions <= 0: - raise ValueError("max_sessions must be positive or null") - if completed_append_limit <= 0: - raise ValueError("completed_append_limit must be positive") - self._sessions: dict[str, DuplexSessionRuntimeState] = {} - self._clock = clock or time.monotonic - self._max_sessions = max_sessions - self._completed_append_limit = completed_append_limit - - def open_session( - self, - fence: DuplexFence, - *, - capabilities: DuplexRuntimeCapabilities | None = None, - session_config: dict[str, Any] | None = None, - runtime_config: dict[str, Any] | None = None, - lease_config: DuplexLeaseConfig | None = None, - ) -> DuplexSessionRuntimeState: - if not isinstance(fence, DuplexFence): - raise TypeError("open_session requires DuplexFence") - if fence.session_id in self._sessions: - raise ValueError(f"Duplex session already exists: {fence.session_id}") - if self._max_sessions is not None and len(self._sessions) >= self._max_sessions: - raise RuntimeError(f"duplex_session_capacity_exhausted: limit={self._max_sessions}") - session = DuplexSessionRuntimeState( - fence=fence, - lease=DuplexLeaseState( - config=lease_config or DuplexLeaseConfig(), - generation=0, - last_activity=self._clock(), - ), - _clock=self._clock, - capabilities=capabilities or _default_capabilities(), - session_config=dict(session_config or {}), - runtime_config=dict(runtime_config or {}), - completed_append_limit=self._completed_append_limit, - ) - self._sessions[fence.session_id] = session - return session - - def get(self, session_id: str) -> DuplexSessionRuntimeState | None: - return self._sessions.get(session_id) - - def require(self, session_id: str) -> DuplexSessionRuntimeState: - session = self.get(session_id) - if session is None: - raise KeyError(f"Unknown duplex session: {session_id}") - return session - - def close_session( - self, - fence: DuplexFence, - *, - reason: str = "explicit_close", - ) -> DuplexSessionRuntimeState | None: - if not isinstance(fence, DuplexFence): - raise TypeError("close_session requires DuplexFence") - session = self._sessions.get(fence.session_id) - if session is not None: - if session.lease.terminal_reason is not None: - return None - if not session.terminate(fence, reason=reason): - return None - if self._sessions.get(fence.session_id) is session: - self._sessions.pop(fence.session_id) - return session - - def begin_close_session( - self, - fence: DuplexFence, - *, - reason: str = "explicit_close", - ) -> DuplexSessionRuntimeState | None: - if not isinstance(fence, DuplexFence): - raise TypeError("begin_close_session requires DuplexFence") - session = self._sessions.get(fence.session_id) - if session is not None: - session.begin_close(fence, reason=reason) - return session - - def finalize_close_session(self, session: DuplexSessionRuntimeState) -> None: - session.finalize_close() - if self._sessions.get(session.session_id) is session: - self._sessions.pop(session.session_id, None) - - def collect_expired( - self, - now: float | None = None, - *, - excluded_session_ids: set[str] | None = None, - ) -> list[DuplexSessionExpiry]: - effective_now = self._clock() if now is None else now - excluded = excluded_session_ids or set() - expired: list[DuplexSessionExpiry] = [] - for session_id, session in list(self._sessions.items()): - if session_id in excluded: - continue - if session.lease.disconnect_grace_expired(effective_now): - reason = "disconnect_grace_expired" - elif session.lease.idle_expired(effective_now): - reason = "idle_ttl_expired" - else: - continue - submitted = tuple(session.resource_request_ids(submitted=True)) - reserved = tuple(session.resource_request_ids(submitted=False)) - if not session.begin_close(session.fence, reason=reason): - continue - expired.append( - DuplexSessionExpiry( - session_id=session_id, - fence=session.fence, - lease_generation=session.lease.generation, - reason=reason, - submitted_request_ids=submitted, - reserved_request_ids=reserved, - ) - ) - return expired - - def close_sessions_for_request_ids(self, request_ids: list[str]) -> dict[str, list[str]]: - request_id_set = set(request_ids) - closed: dict[str, list[str]] = {} - for session_id, session in list(self._sessions.items()): - stale = session.resource_request_ids() - if request_id_set.isdisjoint(stale): - continue - if not session.begin_close(session.fence, reason="request_cleanup"): - continue - closed[session_id] = stale - return closed - - def finalize_closed_sessions(self, session_ids: Iterable[str]) -> None: - for session_id in session_ids: - session = self._sessions.get(session_id) - if session is None or session.lease.terminal_reason is None: - continue - self.finalize_close_session(session) - - -__all__ = [ - "DuplexAppendReservation", - "DuplexCompletedAppend", - "DuplexFenceMismatchError", - "DuplexInputAppend", - "DuplexLeaseActivity", - "DuplexLeaseConfig", - "DuplexLeaseState", - "DuplexRequestResource", - "DuplexSessionExpiry", - "DuplexSessionRuntimeManager", - "DuplexSessionRuntimeState", - "DuplexStageBinding", -] diff --git a/vllm_omni/engine/duplex/session/__init__.py b/vllm_omni/engine/duplex/session/__init__.py new file mode 100644 index 000000000..a062ebf0e --- /dev/null +++ b/vllm_omni/engine/duplex/session/__init__.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""One engine-resident duplex session, and everything that runs it. + +``DuplexEngineSession`` (``engine_session``) is the state -- the input, +response, playback and conversation ledgers, the lease and the identity fence. +``DuplexSessionRunner`` (``runner``) owns one of those on the orchestrator loop +and is the only writer of it; ``DuplexSessionManager`` (``manager``) owns the +runners and the admission slots. + +The rest of the package is what the runner was split into, and each module +names the thing it owns rather than a layer: + +* ``context`` -- the state a runner and its components share, and the only + services a component may ask the runner for. +* ``emitter`` -- everything the session sends: the Realtime projection, the + epoch filter, the domain effects a terminal event applies. +* ``model_channel`` -- the model side: submitting an append, projecting the + stage output that comes back, continuing a turn with silence. +* ``control`` -- server VAD and the events that reconfigure the session. +* ``append_task`` -- one append in flight, and the rollback its failure owes. +* ``helpers`` -- the small pure reads and payload builders over a session. +* ``overlap_policy`` / ``commit_policy`` / ``playback_ledger`` -- the decisions + a session makes about speech that arrives while the model is speaking, what a + commit should do, and what the client has actually played. +* ``lease`` -- the idle/activity lease that decides when a session expires. + +Imports here are the package's public surface; a module of this package is fair +game for anything inside it, but code outside ``engine.duplex`` should reach for +these three names. +""" + +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession +from vllm_omni.engine.duplex.session.manager import DuplexSessionManager +from vllm_omni.engine.duplex.session.runner import DuplexSessionRunner + +__all__ = ["DuplexEngineSession", "DuplexSessionManager", "DuplexSessionRunner"] diff --git a/vllm_omni/engine/duplex/session/append_task.py b/vllm_omni/engine/duplex/session/append_task.py new file mode 100644 index 000000000..724dc7f82 --- /dev/null +++ b/vllm_omni/engine/duplex/session/append_task.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""One append in flight, and what its failure owes the session. + +An append is not a call: it is a task queued behind the previous append, which +may be cancelled, may find the session closed before it starts, and on any of +those paths has to give back what it reserved --- the PCM bytes it took from the +input budget, the committed audio it was going to consume, and the response it +precreated. + +That compensation is the reason this is an object. The five steps of an append +all need the same eleven values, so as closures over ``_start_append`` they read +as one function with five entry points and no way to test any of them. Named +fields make the captured state explicit and the compensation a method. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from dataclasses import dataclass + +from vllm.logger import init_logger + +from vllm_omni.engine.duplex.config import DuplexSessionState, DuplexTurnEventType +from vllm_omni.engine.duplex.plugin import PcmAppendReservation +from vllm_omni.engine.duplex.session import helpers +from vllm_omni.engine.duplex.session.context import DuplexSessionContext +from vllm_omni.engine.duplex.session.emitter import SessionEmitter +from vllm_omni.engine.duplex.session.model_channel import ModelChannel + +logger = init_logger(__name__) + + +@dataclass(slots=True) +class AppendAttempt: + """One queued append to the model, with the rollback each failure path owes.""" + + ctx: DuplexSessionContext + out: SessionEmitter + model: ModelChannel + #: Marks the session closing before the close coroutine gets to run. + fail_session: Callable[[str], None] + + payload: dict[str, object] + epoch: int + request_id: str + final: bool + pcm_reservation: PcmAppendReservation | None + operation_id: str | None + #: Committed audio this append consumes; released once it is safely submitted. + retained_committed_payload: dict[str, object] | None + #: Response reserved before submission, to be failed if submission does not happen. + precreated_response_id: str | None + #: Last chance to call the append off, checked once the predecessor is done. + before_append: Callable[[], bool] | None + + # ------------------------------------------------------------------ # + # Compensation # + # ------------------------------------------------------------------ # + + def discard_retained_audio(self) -> None: + """Release the committed audio this append was going to consume. + + Only if it is still the session's: a later commit may have replaced it, + and that one belongs to the append that will carry it. + """ + model_state = self.ctx.model_state + if self.retained_committed_payload is not None and ( + model_state.committed_audio_payload is self.retained_committed_payload + ): + self.ctx.session.release_input_bytes(model_state.clear_committed_audio()) + + def abandon(self) -> None: + """Give back everything this append reserved but never used.""" + if self.pcm_reservation is not None: + self.pcm_reservation.rollback() + self.discard_retained_audio() + + def release_on_failure(self, done: asyncio.Task[bool]) -> None: + """Done-callback: a cancelled or failed append must not hold the audio.""" + if done.cancelled(): + self.discard_retained_audio() + return + try: + append_ok = done.result() + except Exception: + append_ok = False + if not append_ok: + self.discard_retained_audio() + + def clear_pending_silence(self, done: asyncio.Task[bool]) -> None: + """Done-callback of a silence continuation: it is no longer pending.""" + model_state = self.ctx.model_state + if model_state.pending_silence_task is done: + model_state.pending_silence_task = None + model_state.pending_silence_owner_id = None + + # ------------------------------------------------------------------ # + # Running # + # ------------------------------------------------------------------ # + + async def run_in_wire_order(self, predecessor: asyncio.Task[bool] | None) -> bool: + """Wait for the previous append, then submit this one if it still applies.""" + if predecessor is not None: + try: + predecessor_ok = await predecessor + except asyncio.CancelledError: + current = asyncio.current_task() + if current is not None and current.cancelling(): + raise + predecessor_ok = False + except Exception: + predecessor_ok = False + if not predecessor_ok: + self.abandon() + return False + run = self.ctx.run + if run.closing or run.runtime_closed or self.ctx.session.state != DuplexSessionState.OPEN: + self.abandon() + return False + if self.before_append is not None and not self.before_append(): + self.abandon() + # Called off, not failed: the chain behind it still runs. + return True + if self.pcm_reservation is not None and not self.pcm_reservation.active: + self.discard_retained_audio() + return False + return await self._submit() + + async def _submit(self) -> bool: + session = self.ctx.session + model_state = self.ctx.model_state + try: + append_ok, emitted_response = await self.model.append_runtime_input( + self.payload, + operation_id=( + self.pcm_reservation.operation_id if self.pcm_reservation is not None else self.operation_id + ), + final=self.final, + expected_epoch=self.epoch, + ) + if append_ok: + model_state.context_locked = True + if self.pcm_reservation is not None: + self.pcm_reservation.commit() + session.release_input_bytes(self.pcm_reservation.byte_count) + self.discard_retained_audio() + else: + self.abandon() + if not append_ok and self._precreated_response_still_active(): + self._fail_precreated_response() + if not append_ok and session.state == DuplexSessionState.CLOSED: + self.ctx.run.runtime_closed = True + return False + if not emitted_response and session.epoch == self.epoch: + if session.active_request_id == helpers.stage0_request_id(session, self.epoch): + session.clear_request(self.request_id) + if self.final: + self.out.emit_events([session.signal_turn(DuplexTurnEventType.USER_STARTED.value)]) + return append_ok + except asyncio.CancelledError: + if self.pcm_reservation is not None: + self.pcm_reservation.rollback() + raise + except Exception as exc: + self.abandon() + logger.exception("Native duplex append task failed: %s", exc) + self.model.send_runtime_error("runtime_append_task_failed", exc) + if session.state != DuplexSessionState.CLOSED: + self.fail_session("runtime_append_task_failed") + return False + + def _precreated_response_still_active(self) -> bool: + return ( + self.precreated_response_id is not None + and self.ctx.session.active_response_id == self.precreated_response_id + ) + + def _fail_precreated_response(self) -> None: + session = self.ctx.session + session.end_response(commit_text=False) + self.out.emit( + { + "type": "response.done", + "session_id": session.session_id, + "response_id": self.precreated_response_id, + "epoch": session.epoch, + "committed": False, + "status": "failed", + "status_details": {"type": "failed", "reason": "runtime_append_failed"}, + "playback": session.playback.as_dict(), + } + ) diff --git a/vllm_omni/entrypoints/duplex/commit_policy.py b/vllm_omni/engine/duplex/session/commit_policy.py similarity index 91% rename from vllm_omni/entrypoints/duplex/commit_policy.py rename to vllm_omni/engine/duplex/session/commit_policy.py index 868b98c98..0c07a6fa3 100644 --- a/vllm_omni/entrypoints/duplex/commit_policy.py +++ b/vllm_omni/engine/duplex/session/commit_policy.py @@ -21,7 +21,7 @@ class CommitSnapshot: speech_since_commit: bool active_response_id: str | None overlap_speech_ms: int - native_response_in_progress: bool + response_in_progress: bool playback_active: bool @@ -31,7 +31,7 @@ def decide_commit_action(snapshot: CommitSnapshot) -> CommitAction: if snapshot.speech_since_commit: return CommitAction.START_AUTO_RESPONSE return CommitAction.COMMIT_ONLY - if snapshot.native_response_in_progress and snapshot.overlap_speech_ms > 0: + if snapshot.response_in_progress and snapshot.overlap_speech_ms > 0: return CommitAction.DEFER_ACTIVE_RESPONSE if snapshot.playback_active and snapshot.overlap_speech_ms > 0: return CommitAction.DEFER_ACTIVE_RESPONSE diff --git a/vllm_omni/engine/duplex/session/context.py b/vllm_omni/engine/duplex/session/context.py new file mode 100644 index 000000000..1e7b861bc --- /dev/null +++ b/vllm_omni/engine/duplex/session/context.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Shared state for :class:`DuplexSessionRunner` and its components. + +The runner is a state machine that several components read and write. Splitting +it into modules only helps if the shared part stops being implicit: a component +holding a back-reference to the runner is a mixin with extra steps, and one +holding its own copy of a flag diverges from the others. + +So the mutable part is named. ``DuplexRunState`` is the small set of flags that +more than one component touches -- ``grep "\\.run\\."`` finds every mutation -- +and ``DuplexSessionContext`` is the read-mostly collaborator bundle everything +is constructed with. Anything a component needs from the runner that is neither +of those is infrastructure, and goes through :class:`RunnerServices`. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Protocol, TypeVar + +if TYPE_CHECKING: + from vllm.outputs import RequestOutput + + from vllm_omni.engine.duplex.contracts import ( + DuplexOutputContext, + DuplexOutputDecision, + DuplexStagePort, + ) + from vllm_omni.engine.duplex.plugin import DuplexModelPlugin, DuplexModelSessionState + from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession + from vllm_omni.engine.duplex.session.manager import DuplexSessionManager + from vllm_omni.metrics.stats import StageRequestStats + +_OffloadT = TypeVar("_OffloadT") + + +@dataclass(slots=True) +class DuplexRunState: + """Mutable per-run flags shared by the runner and its components. + + Every field here is written by one component and read by another; that is + the criterion for being in this object rather than private to a component. + """ + + #: An irreversible close has begun: commands and control ops are refused. + closing: bool = False + #: ``session.closed`` / ``session.expired`` already left the runner. + closed_emitted: bool = False + #: The terminal event is deferred to the manager, which emits it after the + #: stage cleanup so it also means "the admission slot is free again". + closed_deferred: bool = False + close_reason: str | None = None + #: The model runtime reported a close; further data-plane work is pointless. + runtime_closed: bool = False + #: Request id of the resumable data-plane stream currently bound to the session. + stream_request_id: str | None = None + + +class RunnerServices(Protocol): + """The only things a component may ask the runner for. + + Deliberately two methods. A component that needs more than task scheduling + from the runner is reaching for orchestration that belongs in the runner. + """ + + def spawn(self, coro: Awaitable[None], *, name: str) -> None: + """Run ``coro`` as a tracked background task on the session's loop.""" + ... + + async def offload(self, fn: Callable[..., _OffloadT], *args: object, **kwargs: object) -> _OffloadT: + """Run a blocking call off the orchestrator loop.""" + ... + + +@dataclass(slots=True) +class DuplexSessionContext: + """Collaborators one session's components are built with. + + Read-mostly: the session and the model state are mutated through their own + APIs, not by rebinding these fields. The mutable flags live in ``run``. + """ + + session: DuplexEngineSession + model_state: DuplexModelSessionState + plugin: DuplexModelPlugin + stage_port: DuplexStagePort + manager: DuplexSessionManager + tasks: DuplexSessionTasks + run: DuplexRunState + services: RunnerServices + + +# --------------------------------------------------------------------------- # +# Task bookkeeping # +# --------------------------------------------------------------------------- # + + +@dataclass(frozen=True, slots=True) +class DuplexAppendTaskMeta: + epoch: int + final: bool + response_bound: bool + + +@dataclass +class DuplexSessionTasks: + """Tracked task handles of one session (append tail, active response, pending silence).""" + + append_tasks: dict[asyncio.Task[bool], DuplexAppendTaskMeta] = field(default_factory=dict) + append_tail: asyncio.Task[bool] | None = None + active_response_task: asyncio.Task[None] | None = None + + def track_append_task( + self, + task: asyncio.Task[bool], + *, + epoch: int, + final: bool, + response_bound: bool, + ) -> None: + self.append_tasks[task] = DuplexAppendTaskMeta(epoch, final, response_bound) + task.add_done_callback(self.append_tasks.pop) + + def has_response_bound_append_tasks(self) -> bool: + return any(meta.response_bound for meta in self.append_tasks.values()) + + async def cancel_append_tasks(self, timeout_s: float = 0.25, *, response_bound_only: bool = False) -> bool: + tasks = [task for task, meta in self.append_tasks.items() if not response_bound_only or meta.response_bound] + if not tasks: + return False + cancelled_tail = self.append_tail if self.append_tail in tasks else None + for task in tasks: + task.cancel() + try: + await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=timeout_s) + except TimeoutError: + pass + if cancelled_tail is not None and self.append_tail is cancelled_tail: + self.append_tail = None + return True + + +# --------------------------------------------------------------------------- # +# Stage output (the mailbox item the orchestrator hands the runner) # +# --------------------------------------------------------------------------- # + + +@dataclass(frozen=True, slots=True) +class StageOutput: + """One stage result, queued for the session that owns the request.""" + + stage_id: int + output: RequestOutput + metrics: StageRequestStats | None + request_id: str + context: DuplexOutputContext + decision: DuplexOutputDecision | None diff --git a/vllm_omni/engine/duplex/session/control.py b/vllm_omni/engine/duplex/session/control.py new file mode 100644 index 000000000..2d46154fe --- /dev/null +++ b/vllm_omni/engine/duplex/session/control.py @@ -0,0 +1,328 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Control of one duplex session: server VAD, and the events that reconfigure it. + +Server turn detection and ``session.update`` are the same concern seen from two +sides. The detector is per-session mutable state, and ``session.update`` is the +event that replaces it mid-flight -- along with the session config, the runtime +config and the sampling policy, each of which has to be validated against a +candidate before the live session adopts it. The rest of ``turn.signal`` +(conversation items, local turn markers) travels the same path and is small, so +it lives here rather than in a module of its own. + +The one thing this does not own is ordering against the data plane: a +``session.update`` must not overtake an append still in flight, so the runner +supplies ``wait_for_append_tail``. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from copy import deepcopy + +from vllm.logger import init_logger + +from vllm_omni.engine.duplex.config import DuplexConfigError, realtime_item_to_history_message +from vllm_omni.engine.duplex.plugin import DuplexRuntimeConfigError +from vllm_omni.engine.duplex.session import helpers +from vllm_omni.engine.duplex.session.context import DuplexSessionContext +from vllm_omni.engine.duplex.session.emitter import SessionEmitter +from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity +from vllm_omni.engine.duplex.session.model_channel import ModelChannel +from vllm_omni.engine.duplex.turn_detection import ( + PendingTurnDetectionUpdate, + ServerTurnDetector, + ServerVADUnavailableError, + SileroVADBackendProvider, + TurnDetectionConfig, + TurnDetectionResult, + apply_turn_detection_result, +) + +logger = init_logger(__name__) + + +class SessionControl: + """Server VAD and the control events that reconfigure one session.""" + + def __init__( + self, + ctx: DuplexSessionContext, + out: SessionEmitter, + model: ModelChannel, + *, + wait_for_append_tail: Callable[[], Awaitable[bool]], + ) -> None: + self._ctx = ctx + self._out = out + self._model = model + self._wait_for_append_tail = wait_for_append_tail + self._config: TurnDetectionConfig | None = None + self._detector: ServerTurnDetector | None = None + + def init_turn_detection(self) -> None: + turn_detection = self._ctx.session.config.extra_body.get("realtime_turn_detection") + if not isinstance(turn_detection, dict) or turn_detection.get("type") != "server_vad": + self._config = None + self._detector = None + return + try: + config = TurnDetectionConfig.from_realtime(turn_detection) + self._config = config + self._detector = config.build_detector(self._vad_backend_provider()) + except Exception as exc: + logger.warning("Duplex session %s: turn detection disabled: %s", self._ctx.session.session_id, exc) + self._config = None + self._detector = None + + def _vad_backend_provider(self) -> SileroVADBackendProvider | None: + """The engine-wide Silero backend, so one model serves every session. + + Held by the manager rather than built here: the ONNX backend is shared, + and ``duplex_session.server_vad_model_path`` is a deploy-level setting. + """ + return getattr(self._ctx.manager, "vad_backend_provider", None) + + def reset_vad(self) -> None: + """Drop the detector's speech state (a barge-in or a clear starts a new turn).""" + if self._detector is not None: + self._detector.reset() + + async def run_turn_detection(self, event: dict[str, object]) -> TurnDetectionResult | None: + detector = self._detector + if detector is None: + return None + audio = event.get("audio") + if not isinstance(audio, str) or not audio: + return None + fmt = event.get("format") if isinstance(event.get("format"), str) else "pcm_f32le" + sample_rate_hz = event.get("sample_rate_hz") + try: + result = await self._ctx.services.offload( + detector.process, + audio, + fmt=fmt, + sample_rate_hz=sample_rate_hz if isinstance(sample_rate_hz, int) else None, + audio_end_ms=event.get("audio_end_ms") if isinstance(event.get("audio_end_ms"), int) else None, + ) + except ServerVADUnavailableError as exc: + self._out.emit_error("server_vad_unavailable", str(exc)) + self._detector = None + return None + except ValueError as exc: + self._out.emit_error("bad_audio", str(exc)) + return None + apply_turn_detection_result(event, result) + return result + + # ------------------------------------------------------------------ # + # Control events # + # ------------------------------------------------------------------ # + + async def on_turn_signal(self, event: dict[str, object]) -> None: + session = self._ctx.session + turn_event = event.get("event") + realtime_event_id = event.get("realtime_event_id") + if not isinstance(turn_event, str): + self._out.emit_error("bad_event", "turn.signal requires event") + return + if turn_event == "barge_in" and not session.capabilities.supports_barge_in: + self._out.emit_events([helpers.barge_in_unsupported_error()]) + return + if turn_event == "session.update": + payload = event.get("payload") + if not isinstance(payload, dict): + self._out.emit_error( + "bad_event", "session.update requires a session payload", event_id=realtime_event_id + ) + return + await self.on_session_update( + payload, + realtime_event_id=realtime_event_id if isinstance(realtime_event_id, str) else None, + ) + return + if turn_event == "conversation.item.create": + await self._on_conversation_item_create(event) + return + if turn_event == "conversation.item.delete": + payload = event.get("payload") + item_id = payload.get("item_id") if isinstance(payload, dict) else None + deleted = session.delete_history_item(item_id) if isinstance(item_id, str) else False + self._out.emit( + { + "type": "conversation.item.deleted", + "session_id": session.session_id, + "item_id": item_id, + "deleted": deleted, + } + ) + return + if turn_event == "conversation.item.truncate": + payload = event.get("payload") + item_id = payload.get("item_id") if isinstance(payload, dict) else None + audio_end_ms = payload.get("audio_end_ms") if isinstance(payload, dict) else None + truncated = ( + session.truncate_history_item( + item_id, + audio_end_ms=int(audio_end_ms) if isinstance(audio_end_ms, int | float) else 0, + hard=True, + ) + if isinstance(item_id, str) + else False + ) + self._out.emit( + { + "type": "conversation.item.truncated", + "session_id": session.session_id, + "item_id": item_id, + "content_index": (payload.get("content_index", 0) if isinstance(payload, dict) else 0), + "audio_end_ms": audio_end_ms, + "truncated": truncated, + } + ) + return + self._out.emit_events([session.signal_turn(turn_event, event)]) + + async def on_session_update(self, payload: dict[str, object], *, realtime_event_id: str | None) -> None: + session = self._ctx.session + model_state = self._ctx.model_state + pending_turn_detection: PendingTurnDetectionUpdate | None = None + + def reject_update() -> None: + if pending_turn_detection is not None: + pending_turn_detection.reject() + + try: + pending_turn_detection = PendingTurnDetectionUpdate.prepare( + payload, backend_provider=self._vad_backend_provider() + ) + except Exception as exc: + self._out.emit_error("unsupported_turn_detection", str(exc), event_id=realtime_event_id) + return + if not await self._wait_for_append_tail(): + self._out.emit_error( + "session_update_aborted", + "session.update was not applied because the preceding append failed", + event_id=realtime_event_id, + ) + reject_update() + return + try: + self._ctx.plugin.validate_client_extra_body(payload.get("extra_body")) + except DuplexRuntimeConfigError as exc: + self._out.emit_error(exc.code, str(exc), event_id=realtime_event_id) + reject_update() + return + candidate_config = deepcopy(session.config) + audio_started = session.playback.generated_ms > 0 or session.playback.sent_ms > 0 + try: + candidate_config.apply_realtime_update( + payload, + session_id=session.session_id, + audio_started=audio_started, + ) + except DuplexConfigError as exc: + self._out.emit_error(exc.code, str(exc) or "session.update was rejected", event_id=realtime_event_id) + reject_update() + return + if candidate_config.instructions != session.config.instructions and model_state.context_locked: + self._out.emit_error( + "instructions_update_unsupported", + "session.update cannot change instructions after the native duplex context is initialized", + event_id=realtime_event_id, + ) + reject_update() + return + requests_audio = any(str(modality).lower() == "audio" for modality in candidate_config.modalities) + if ( + requests_audio + and "ref_audio_data" not in session.runtime_config + and getattr(self._ctx.plugin, "requires_ref_audio", False) + ): + self._out.emit_error( + "ref_audio_required", "Native duplex audio output requires ref_audio", event_id=realtime_event_id + ) + reject_update() + return + try: + candidate_runtime_config = self._ctx.plugin.runtime_config_for_update( + candidate_config, + dict(session.runtime_config), + ) + # Validate the sampling policy for the candidate before adopting it. + self._ctx.plugin.configure_sampling_params( + runtime_config=dict(candidate_runtime_config), + defaults=tuple(self._ctx.stage_port.sampling_defaults()), + ) + except (DuplexRuntimeConfigError, DuplexConfigError) as exc: + self._out.emit_error(exc.code, str(exc), event_id=realtime_event_id) + reject_update() + return + except Exception as exc: + self._out.emit_error("runtime_signal_failed", str(exc), event_id=realtime_event_id) + reject_update() + return + session.replace_config(candidate_config) + session.replace_runtime_config(candidate_runtime_config) + try: + session.touch_lease(DuplexLeaseActivity.SIGNAL) + except Exception: + pass + if pending_turn_detection is not None: + self._config, self._detector = pending_turn_detection.commit(self._detector) + projector = self._out.require_projector() + projector.apply_session_defaults(payload) + self._out.emit({"type": "session.updated", "session": session.as_public_dict()}) + + async def _on_conversation_item_create(self, event: dict[str, object]) -> None: + session = self._ctx.session + payload = event.get("payload") + item = payload.get("item") if isinstance(payload, dict) else None + item_type = item.get("type") if isinstance(item, dict) else None + if item_type == "function_call_output" and isinstance(item, dict): + if not await self._wait_for_append_tail(): + return + try: + candidate_runtime_config = self._ctx.plugin.runtime_config_for_function_output( + session.config, + dict(session.runtime_config), + item, + ) + except DuplexRuntimeConfigError as exc: + self._out.emit_error(exc.code, str(exc)) + return + if candidate_runtime_config is not None: + session.replace_runtime_config(candidate_runtime_config) + self._out.emit( + { + "type": "conversation.item.created", + "session_id": session.session_id, + "item": item, + "created": True, + } + ) + self._ctx.services.spawn( + self._model.maybe_continue_response( + expected_epoch=session.epoch, + expected_model_turn_id=session.turn_id, + ), + name="duplex-continue", + ) + return + message = realtime_item_to_history_message(item) + item_id = item.get("id") if isinstance(item, dict) else None + if message is not None: + session.append_history_message(message) + session.register_history_item(item_id if isinstance(item_id, str) else None, message) + if message.get("role") == "user": + # A later response.create may answer this without any audio. + session.notify_new_user_item() + self._out.emit( + { + "type": "conversation.item.created", + "session_id": session.session_id, + "item": item, + "created": message is not None, + } + ) diff --git a/vllm_omni/engine/duplex/session/emitter.py b/vllm_omni/engine/duplex/session/emitter.py new file mode 100644 index 000000000..508bc7ef1 --- /dev/null +++ b/vllm_omni/engine/duplex/session/emitter.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Everything one duplex session sends to its client. + +Emission is the hub of the session runner: near enough every path through the +runner ends in an event, which is why this is a collaborator rather than a +module of functions. It owns three things the rest of the runner should not +have to think about --- the Realtime projection state, the epoch filter that +drops output from a superseded turn, and the domain transitions an outbound +terminal event applies to the session before it is projected. + +The one thing it cannot own is what happens *after* a terminal promotes +deferred overlap audio: that re-enters the mailbox of the runner, so the +runner injects it as ``promote_deferred_overlap``. +""" + +from __future__ import annotations + +from collections.abc import Callable + +from vllm_omni.engine.duplex.config import DuplexSessionState +from vllm_omni.engine.duplex.events import ( + DOMAIN_TERMINAL_EVENTS, + MODEL_OUTPUT_EVENTS, + DuplexEvent, + error_event, +) +from vllm_omni.engine.duplex.realtime_events import ( + RealtimeProjectionState, + discard_pending_input_audio, + project_internal_event, +) +from vllm_omni.engine.duplex.session import overlap_policy +from vllm_omni.engine.duplex.session.context import DuplexSessionContext + + +class SessionEmitter: + """Outbound events for one session.""" + + def __init__( + self, + ctx: DuplexSessionContext, + *, + promote_deferred_overlap: Callable[[dict[str, object], bool], None], + ) -> None: + self._ctx = ctx + self._promote_deferred_overlap = promote_deferred_overlap + self._projector: RealtimeProjectionState | None = ctx.session.projector + + # ------------------------------------------------------------------ # + # Projection state # + # ------------------------------------------------------------------ # + + @property + def projector(self) -> RealtimeProjectionState | None: + """The Realtime projection state, or ``None`` before anything needs it.""" + return self._projector + + @projector.setter + def projector(self, value: RealtimeProjectionState | None) -> None: + self._projector = value + self._ctx.session.projector = value + + def require_projector(self) -> RealtimeProjectionState: + """The projection state, created on first use. + + A session torn down before ``start()`` never builds one, so the lazy + path here is reachable rather than defensive. + """ + projector = self._projector + if projector is None: + session = self._ctx.session + projector = RealtimeProjectionState( + session_id=session.session_id, + model=session.config.model, + initial_session_update=True, + ) + self.projector = projector + return projector + + # ------------------------------------------------------------------ # + # Sending # + # ------------------------------------------------------------------ # + + def emit_events(self, events: list[DuplexEvent]) -> None: + for event in events: + self._ctx.manager.emit(self._ctx.session, event) + + def emit_error( + self, + code: str, + message: str, + *, + event_id: object | None = None, + retryable: bool | None = None, + ) -> None: + """Send one typed ``error`` event (``event_id`` is the client event it answers).""" + extra = {} if retryable is None else {"retryable": retryable} + self.emit_events([error_event(code, message, event_id=event_id, extra=extra)]) + + def emit(self, payload: dict[str, object]) -> None: + """Apply the domain effects of an internal event, then project it to typed events and send them. + + Only events with domain effects (response / cancel / close terminals) or + Realtime projection state (response items, content parts) still travel as + internal dictionaries; stateless events are constructed typed at the + emit site. + """ + accepted, deferred_overlap_payload = self._apply_outbound_session_event(payload) + if not accepted: + return + 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 + self._ctx.model_state.deferred_precreate_response = False + self._promote_deferred_overlap(deferred_overlap_payload, precreate_response) + + # ------------------------------------------------------------------ # + # Filtering and domain effects # + # ------------------------------------------------------------------ # + + def auto_responds(self) -> bool: + """Whether the session answers committed input without a ``response.create``.""" + extra = getattr(self._ctx.session.config, "extra_body", None) + if not isinstance(extra, dict): + return False + return extra.get("auto_response") is True or extra.get("full_duplex") is True + + def is_stale_model_output(self, payload: dict[str, object]) -> bool: + """Whether ``payload`` belongs to a turn the session has already moved past.""" + session = self._ctx.session + event_type = payload.get("type") + if event_type in DOMAIN_TERMINAL_EVENTS: + return False + if event_type not in MODEL_OUTPUT_EVENTS: + return False + if self._ctx.run.closing and event_type != "response.listen": + return True + if session.state == DuplexSessionState.CLOSED and event_type != "response.listen": + return True + epoch = payload.get("epoch") + return isinstance(epoch, int) and epoch != session.epoch + + def _apply_outbound_session_event(self, payload: dict[str, object]) -> tuple[bool, dict[str, object] | None]: + """Apply domain transitions before an event is projected (moved from serving).""" + session = self._ctx.session + model_state = self._ctx.model_state + payload_type = payload.get("type") + is_terminal = payload_type in DOMAIN_TERMINAL_EVENTS + if is_terminal: + payload_epoch = payload.get("epoch") + if isinstance(payload_epoch, int) and payload_epoch != session.epoch: + return False, None + if payload_type in {"response.done", "response.listen"} and ( + self._ctx.run.closing or session.state == DuplexSessionState.CLOSED + ): + return False, None + elif self.is_stale_model_output(payload): + return False, None + + if payload_type == "session.closed": + self._ctx.run.close_reason = self._ctx.run.close_reason or str(payload.get("reason") or "closed") + session.mark_closing() + + if not is_terminal: + return True, None + + terminal_status = payload.get("status") + terminal_status_details = payload.get("status_details") + if terminal_status is None and isinstance(terminal_status_details, dict): + terminal_status = terminal_status_details.get("type") + response_terminal = payload_type == "response.done" or ( + payload_type == "response.listen" and session.active_response_id is not None + ) + can_promote_overlap = response_terminal and terminal_status not in {"cancelled", "failed"} + deferred_overlap_payload: dict[str, object] | None = None + continuous_input_crosses_terminal = ( + can_promote_overlap + and self.auto_responds() + and model_state.input_since_commit + and not model_state.deferred_response_create + ) + if continuous_input_crosses_terminal: + session.reset_overlap_speech() + return True, None + realtime_input_still_open = ( + can_promote_overlap and model_state.input_since_commit and not model_state.deferred_response_create + ) + if realtime_input_still_open: + session.reset_overlap_speech() + return True, None + if can_promote_overlap and session.overlap_speech_ms > 0: + has_deferred_overlap = ( + model_state.audio_buffer.has_pending() or model_state.committed_audio_payload is not None + ) + should_promote_overlap = ( + session.state == DuplexSessionState.OPEN + and has_deferred_overlap + and session.overlap_speech_ms > session.config.overlap_short_ack_ms + ) + if should_promote_overlap: + flushed_reserved_bytes = model_state.audio_buffer.pending_byte_count + deferred_overlap_payload = model_state.audio_buffer.flush( + chunk_period_ms=session.capabilities.chunk_period_ms or 1000 + ) + if model_state.committed_audio_payload is not None: + if deferred_overlap_payload is not None: + deferred_overlap_payload = overlap_policy.merge_audio_payloads( + model_state.committed_audio_payload, + deferred_overlap_payload, + ) + else: + deferred_overlap_payload = model_state.committed_audio_payload + if self.auto_responds() and deferred_overlap_payload is not None: + deferred_overlap_payload = dict(deferred_overlap_payload) + deferred_overlap_payload["force_listen"] = False + if deferred_overlap_payload is not None: + model_state.retain_committed_audio( + deferred_overlap_payload, + operation_id=model_state.committed_audio_operation_id, + reserved_bytes=flushed_reserved_bytes, + ) + model_state.input_since_commit = deferred_overlap_payload is not None + # Realtime path: defer the promoted response to the next terminal. + model_state.deferred_response_create = True + model_state.deferred_precreate_response = False + deferred_overlap_payload = None + else: + had_pending_overlap_audio = model_state.audio_buffer.has_pending() + model_state.audio_buffer.clear() + model_state.input_since_commit = False + model_state.speech_since_commit = False + if had_pending_overlap_audio and self._projector is not None: + self.emit_events(discard_pending_input_audio(self._projector, session.overlap_speech_ms)) + if payload_type in {"audio.cancelled", "input.cancelled", "session.closed"}: + session.release_input_bytes(model_state.clear_committed_audio()) + + session.reset_overlap_speech() + if ( + can_promote_overlap + and model_state.deferred_response_create + and model_state.committed_audio_payload is not None + ): + deferred_overlap_payload = model_state.committed_audio_payload + model_state.deferred_response_create = False + model_state.input_since_commit = False + model_state.speech_since_commit = False + return True, deferred_overlap_payload diff --git a/vllm_omni/entrypoints/duplex/protocol.py b/vllm_omni/engine/duplex/session/engine_session.py similarity index 56% rename from vllm_omni/entrypoints/duplex/protocol.py rename to vllm_omni/engine/duplex/session/engine_session.py index 042180838..bb5c38515 100644 --- a/vllm_omni/entrypoints/duplex/protocol.py +++ b/vllm_omni/engine/duplex/session/engine_session.py @@ -1,411 +1,64 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The one duplex session state, owned by the engine-side ``DuplexSessionRunner``. + +``DuplexEngineSession`` is the single session object: the input / response / +playback / conversation ledgers, the lease, the identity fence, the stage +request resources, the append sequencing, the model plugin's per-session state +(``model_state``) and the Realtime projection state (``projector``). There is +no cross-boundary fence protocol; ``session.fence`` is derived from the +session's own epoch/turn and ``accepted_fence`` is the monotonic high-water +mark of fences accepted for stage requests (``sync_fence()`` publishes the +current identity there). +""" + from __future__ import annotations import copy -from collections.abc import Mapping +import time +from collections.abc import Callable, Mapping from dataclasses import dataclass, field -from enum import Enum from types import MappingProxyType +from typing import TYPE_CHECKING from uuid import uuid4 -import numpy as np - -from vllm_omni.entrypoints.duplex.audio import ( - encode_float32_mono_wav_base64, +from vllm_omni.engine.duplex.config import ( + DuplexAssistantAudioTextMark, + DuplexCapabilities, + DuplexCommittedInput, + DuplexPlaybackCommitPolicy, + DuplexPlaybackCursor, + DuplexPlaybackView, + DuplexSessionConfig, + DuplexSessionState, + DuplexTurnEventType, + DuplexTurnState, + ResponseCreateOptions, ) -from vllm_omni.entrypoints.duplex.server_vad import ( - ServerVADConfig, - parse_session_turn_detection, +from vllm_omni.engine.duplex.contracts import DuplexFence +from vllm_omni.engine.duplex.events import TurnEvent +from vllm_omni.engine.duplex.session.lease import ( + DuplexLeaseActivity, + DuplexLeaseConfig, + DuplexLeaseState, ) - -class DuplexOverlapPolicy(str, Enum): - LISTEN_ONLY = "listen_only" - BARGE_IN_ON_SPEECH = "barge_in_on_speech" +if TYPE_CHECKING: + from vllm_omni.engine.duplex.plugin import DuplexModelSessionState + from vllm_omni.engine.duplex.realtime_events import RealtimeProjectionState -class DuplexPlaybackCommitPolicy(str, Enum): - COMMIT_ALL_ON_DONE = "commit_all_on_done" - ACK_ONLY = "ack_only" - - -class DuplexSessionState(str, Enum): - OPEN = "open" - CLOSING = "closing" - CLOSED = "closed" - - -# Canonical extra_body key for the per-session opt-in to a model-native -# duplex runtime, plus the deprecated model-prefixed spelling it replaced. -# Client payloads may still use the legacy key; it is folded into the -# canonical one at ingestion so everything downstream (and every echo in -# session.created/session.updated) sees only NATIVE_DUPLEX_KEY. -NATIVE_DUPLEX_KEY = "native_duplex" -LEGACY_NATIVE_DUPLEX_KEY = "minicpmo45_native_duplex" - - -def normalize_native_duplex_key(extra_body: dict[str, object]) -> dict[str, object]: - """Fold the deprecated alias into the canonical key (canonical wins).""" - if LEGACY_NATIVE_DUPLEX_KEY in extra_body: - legacy = extra_body.pop(LEGACY_NATIVE_DUPLEX_KEY) - extra_body.setdefault(NATIVE_DUPLEX_KEY, legacy) - return extra_body - - -def native_duplex_opt_in(extra_body: Mapping[str, object]) -> object: - """Read the opt-in flag, accepting the deprecated alias for raw dicts.""" - if NATIVE_DUPLEX_KEY in extra_body: - return extra_body[NATIVE_DUPLEX_KEY] - return extra_body.get(LEGACY_NATIVE_DUPLEX_KEY) - - -class DuplexTurnState(str, Enum): - IDLE = "idle" - USER_SPEAKING = "user_speaking" - USER_COMMITTED = "user_committed" - ASSISTANT_GENERATING = "assistant_generating" - ASSISTANT_PLAYING = "assistant_playing" - BARGE_IN = "barge_in" - - -class DuplexTurnEventType(str, Enum): - USER_STARTED = "user_started" - USER_COMMITTED = "user_committed" - ASSISTANT_STARTED = "assistant_started" - ASSISTANT_DONE = "assistant_done" - BARGE_IN = "barge_in" - PLAYBACK_ACK = "playback_ack" - TIMEOUT = "timeout" - CLOSE = "close" - - -@dataclass -class DuplexCapabilities: - """Runtime/model capabilities exposed by the duplex serving protocol. - - These are intentionally explicit so the serving layer does not assume all - duplex models support the same input append, rollback, or turn policy. - ``supports_core_kv_lease`` is reserved for scheduler-owned KV lifecycle; - model-owned decoder/TTS state must use ``supports_model_internal_state``. - ``supports_core_resumable_request`` means the scheduler can resume the - same request id across streaming updates, but it is not a KV lease by - itself. Realtime support means this endpoint can speak the native Realtime - event schema for the supported audio duplex paths while keeping model- or - scheduler-specific limits explicit in the capability payload. - """ - - supports_session_adapter: bool = True - supports_model_native_turn_policy: bool = False - supports_external_turn_signal: bool = True - supports_client_commit: bool = True - supports_barge_in: bool = True - supports_playback_ack: bool = True - supports_input_append: bool = False - supports_replace_latest_chunk: bool = True - supports_reencode_context: bool = True - supports_rollback_to_checkpoint: bool = False - supports_turn_commit_only: bool = True - supports_kv_lease: bool = False - supports_core_kv_lease: bool = False - supports_model_internal_state: bool = False - supports_stage_resumption: bool = False - supports_scheduler_native_append: bool = False - supports_core_resumable_request: bool = False - supports_stage_connector_handoff: bool = False - supports_independent_io_streams: bool = False - supports_realtime_endpoint: bool = False - supports_multi_session: bool = False - supports_multi_session_same_replica: bool = False - supports_session_lease: bool = False - supports_session_resume: bool = False - session_admission_mode: str = "serving_managed" - supports_audio_truncate: bool = False - requires_model_runner_kv: bool = False - requires_native_stage_role: bool = False - implementation_level: str = "serving_session_adapter" - adapter_patterns: list[str] = field(default_factory=lambda: ["chunk_group_append"]) - input_modes: list[str] = field(default_factory=lambda: ["turn_commit_only", "reencode_context"]) - signal_sources: list[str] = field(default_factory=lambda: ["client_event", "server_policy", "model_native"]) - stage_handoff_transport: str | None = None - chunk_period_ms: int | None = 1000 - target_barge_in_latency_ms: int | None = 1000 - - def as_dict(self) -> dict[str, object]: - return { - "supports_session_adapter": self.supports_session_adapter, - "supports_model_native_turn_policy": self.supports_model_native_turn_policy, - "supports_external_turn_signal": self.supports_external_turn_signal, - "supports_client_commit": self.supports_client_commit, - "supports_barge_in": self.supports_barge_in, - "supports_playback_ack": self.supports_playback_ack, - "supports_input_append": self.supports_input_append, - "supports_replace_latest_chunk": self.supports_replace_latest_chunk, - "supports_reencode_context": self.supports_reencode_context, - "supports_rollback_to_checkpoint": self.supports_rollback_to_checkpoint, - "supports_turn_commit_only": self.supports_turn_commit_only, - "supports_kv_lease": self.supports_kv_lease, - "supports_core_kv_lease": self.supports_core_kv_lease, - "supports_model_internal_state": self.supports_model_internal_state, - "supports_stage_resumption": self.supports_stage_resumption, - "supports_scheduler_native_append": self.supports_scheduler_native_append, - "supports_core_resumable_request": self.supports_core_resumable_request, - "supports_stage_connector_handoff": self.supports_stage_connector_handoff, - "supports_independent_io_streams": self.supports_independent_io_streams, - "supports_realtime_endpoint": self.supports_realtime_endpoint, - "supports_multi_session": self.supports_multi_session, - "supports_multi_session_same_replica": self.supports_multi_session_same_replica, - "supports_session_lease": self.supports_session_lease, - "supports_session_resume": self.supports_session_resume, - "session_admission_mode": self.session_admission_mode, - "supports_audio_truncate": self.supports_audio_truncate, - "requires_model_runner_kv": self.requires_model_runner_kv, - "requires_native_stage_role": self.requires_native_stage_role, - "implementation_level": self.implementation_level, - "adapter_patterns": self.adapter_patterns, - "input_modes": self.input_modes, - "signal_sources": self.signal_sources, - "stage_handoff_transport": self.stage_handoff_transport, - "chunk_period_ms": self.chunk_period_ms, - "target_barge_in_latency_ms": self.target_barge_in_latency_ms, - } - - -@dataclass -class DuplexPlaybackCursor: - generated_ms: int = 0 - sent_ms: int = 0 - played_ms: int = 0 - committed_ms: int = 0 - generated_frames: int = 0 - sample_rate_hz: int | None = None - - def acknowledge(self, played_ms: int, committed_ms: int | None = None) -> None: - self.played_ms = max(self.played_ms, max(0, int(played_ms))) - if committed_ms is None: - committed_ms = self.played_ms - self.committed_ms = max(self.committed_ms, max(0, int(committed_ms))) - - def truncate_committed(self, committed_ms: int) -> None: - self.committed_ms = max(0, min(max(self.sent_ms, self.generated_ms), int(committed_ms))) - - def as_dict(self) -> dict[str, int]: - return { - "generated_ms": self.generated_ms, - "sent_ms": self.sent_ms, - "played_ms": self.played_ms, - "committed_ms": self.committed_ms, - } - - def snapshot(self) -> DuplexPlaybackView: - return DuplexPlaybackView(**self.as_dict()) - - -@dataclass(frozen=True, slots=True) -class DuplexPlaybackView: - generated_ms: int = 0 - sent_ms: int = 0 - played_ms: int = 0 - committed_ms: int = 0 - - def as_dict(self) -> dict[str, int]: - return { - "generated_ms": self.generated_ms, - "sent_ms": self.sent_ms, - "played_ms": self.played_ms, - "committed_ms": self.committed_ms, - } - - -@dataclass -class DuplexAudioChunk: - data: str - format: str = "wav" - sample_rate_hz: int | None = None - - -@dataclass -class DuplexSessionConfig: - model: str | None = None - modalities: list[str] = field(default_factory=lambda: ["text", "audio"]) - instructions: str | None = None - voice: str | None = None - ref_audio: str | None = None - response_format: str = "wav" - temperature: float | None = None - max_tokens: int | None = None - speed: float | None = None - use_tts_template: bool = True - idle_timeout_s: float = 300.0 - overlap_policy: str = DuplexOverlapPolicy.LISTEN_ONLY.value - overlap_short_ack_ms: int = 700 - overlap_barge_in_ms: int = 1200 - overlap_silence_rms: float = 0.003 - playback_commit_policy: str = DuplexPlaybackCommitPolicy.COMMIT_ALL_ON_DONE.value - turn_detection_configured: bool = False - server_vad: ServerVADConfig | None = None - extra_body: dict[str, object] = field(default_factory=dict) - - def as_dict(self) -> dict[str, object]: - payload: dict[str, object] = { - "model": self.model, - "modalities": list(self.modalities), - "instructions": self.instructions, - "voice": self.voice, - "ref_audio": self.ref_audio, - "response_format": self.response_format, - "temperature": self.temperature, - "max_tokens": self.max_tokens, - "speed": self.speed, - "use_tts_template": self.use_tts_template, - "idle_timeout_s": self.idle_timeout_s, - "overlap_policy": self.overlap_policy, - "overlap_short_ack_ms": self.overlap_short_ack_ms, - "overlap_barge_in_ms": self.overlap_barge_in_ms, - "overlap_silence_rms": self.overlap_silence_rms, - "playback_commit_policy": self.playback_commit_policy, - "extra_body": dict(self.extra_body), - } - if self.turn_detection_configured: - payload["turn_detection"] = self.server_vad.as_dict() if self.server_vad is not None else None - return payload - - @classmethod - def from_event(cls, event: dict[str, object]) -> DuplexSessionConfig: - payload = event.get("session") - if isinstance(payload, dict): - source = payload - else: - source = event - - config = cls() - if isinstance(source.get("model"), str): - config.model = source["model"] - if isinstance(source.get("instructions"), str): - config.instructions = source["instructions"] - if isinstance(source.get("voice"), str): - config.voice = source["voice"] - if isinstance(source.get("ref_audio"), str): - config.ref_audio = source["ref_audio"] - if isinstance(source.get("response_format"), str): - config.response_format = source["response_format"] - if isinstance(source.get("use_tts_template"), bool): - config.use_tts_template = bool(source["use_tts_template"]) - if isinstance(source.get("temperature"), int | float): - config.temperature = float(source["temperature"]) - if isinstance(source.get("max_tokens"), int): - config.max_tokens = int(source["max_tokens"]) - if isinstance(source.get("speed"), int | float): - config.speed = float(source["speed"]) - if isinstance(source.get("idle_timeout_s"), int | float): - config.idle_timeout_s = float(source["idle_timeout_s"]) - if isinstance(source.get("overlap_policy"), str): - config.overlap_policy = cls._normalize_overlap_policy(source["overlap_policy"]) - if isinstance(source.get("overlap_short_ack_ms"), int | float): - config.overlap_short_ack_ms = max(0, int(source["overlap_short_ack_ms"])) - if isinstance(source.get("overlap_barge_in_ms"), int | float): - config.overlap_barge_in_ms = max(0, int(source["overlap_barge_in_ms"])) - if isinstance(source.get("overlap_silence_rms"), int | float): - config.overlap_silence_rms = max(0.0, float(source["overlap_silence_rms"])) - if isinstance(source.get("playback_commit_policy"), str): - config.playback_commit_policy = cls._normalize_playback_commit_policy(source["playback_commit_policy"]) - if isinstance(source.get("modalities"), list) and all(isinstance(x, str) for x in source["modalities"]): - config.modalities = list(source["modalities"]) - turn_detection_configured, server_vad = parse_session_turn_detection(source) - if isinstance(source.get("extra_body"), dict): - config.extra_body = normalize_native_duplex_key(dict(source["extra_body"])) - extra = config.extra_body - raw_realtime_session = extra.get("realtime_session_payload") - if not turn_detection_configured and isinstance(raw_realtime_session, dict): - turn_detection_configured, server_vad = parse_session_turn_detection(raw_realtime_session) - if isinstance(extra.get("overlap_policy"), str): - config.overlap_policy = cls._normalize_overlap_policy(extra["overlap_policy"]) - if isinstance(extra.get("overlap_short_ack_ms"), int | float): - config.overlap_short_ack_ms = max(0, int(extra["overlap_short_ack_ms"])) - if isinstance(extra.get("overlap_barge_in_ms"), int | float): - config.overlap_barge_in_ms = max(0, int(extra["overlap_barge_in_ms"])) - if isinstance(extra.get("overlap_silence_rms"), int | float): - config.overlap_silence_rms = max(0.0, float(extra["overlap_silence_rms"])) - if isinstance(extra.get("playback_commit_policy"), str): - config.playback_commit_policy = cls._normalize_playback_commit_policy(extra["playback_commit_policy"]) - if turn_detection_configured: - config.turn_detection_configured = True - config.server_vad = server_vad - return config - - @staticmethod - def _normalize_overlap_policy(value: str) -> str: - normalized = value.strip().lower() - if normalized in {policy.value for policy in DuplexOverlapPolicy}: - return normalized - return DuplexOverlapPolicy.LISTEN_ONLY.value - - @staticmethod - def _normalize_playback_commit_policy(value: str) -> str: - normalized = value.strip().lower() - if normalized in {policy.value for policy in DuplexPlaybackCommitPolicy}: - return normalized - return DuplexPlaybackCommitPolicy.COMMIT_ALL_ON_DONE.value - - -@dataclass(frozen=True) -class ResponseCreateOptions: - instructions: str | None = None - voice: str | None = None - response_format: str | None = None - temperature: float | None = None - max_tokens: int | None = None - speed: float | None = None - modalities: tuple[str, ...] | None = None - extra_body: Mapping[str, object] = field(default_factory=dict) - - def apply_to(self, config: DuplexSessionConfig) -> None: - for field_name in ("instructions", "voice", "response_format", "temperature", "max_tokens", "speed"): - value = getattr(self, field_name) - if value is not None: - setattr(config, field_name, value) - if self.modalities is not None: - config.modalities = list(self.modalities) - config.extra_body.update(self.extra_body) - - -@dataclass -class DuplexCommittedInput: - message: dict[str, object] - turn_id: int - epoch: int - input_commit_seq: int - - -@dataclass -class PendingServerVADTurn: - committed: DuplexCommittedInput - create_response: bool - item_id: str - - -@dataclass -class DuplexAssistantAudioTextMark: - text_chars: int - audio_end_ms: int +def _default_lease() -> DuplexLeaseState: + return DuplexLeaseState(config=DuplexLeaseConfig(), generation=0, last_activity=time.monotonic()) @dataclass class InputBufferState: commit_seq: int = 0 - pending_text: list[str] = field(default_factory=list) - pending_audio: list[DuplexAudioChunk] = field(default_factory=list) overlap_speech_ms: int = 0 reserved_input_bytes: int = 0 pending_turns: int = 0 - turn_detection_config_locked: bool = False - vad_prefix_audio: list[np.ndarray] = field(default_factory=list) - vad_prefix_samples: int = 0 - vad_utterance_audio: list[np.ndarray] = field(default_factory=list) - vad_speech_active: bool = False - active_server_vad_item_id: str | None = None - pending_server_vad_turn: PendingServerVADTurn | None = None @dataclass @@ -413,10 +66,6 @@ class ResponseState: active_request_id: str | None = None active_response_id: str | None = None active_response_turn_id: int | None = None - request_started_at_s_by_turn: dict[int, float] = field(default_factory=dict) - active_response_request_started_at_s: float | None = None - active_response_ttft_ms: float | None = None - active_response_ttfp_ms: float | None = None active_response_input_commit_seq: int | None = None active_response_awaits_input_commit: bool = False last_response_id: str | None = None @@ -439,6 +88,10 @@ class PlaybackLedger: @dataclass class ConversationHistory: messages: list[dict[str, object]] = field(default_factory=list) + #: User items added by ``conversation.item.create`` that no response has + #: answered yet. A turn can start from these alone, the way the Realtime + #: protocol allows, without any audio having been committed. + unanswered_user_items: int = 0 item_ids: dict[str, dict[str, object]] = field(default_factory=dict) history_item_placeholders: dict[str, dict[str, object]] = field(default_factory=dict) item_audio_text_marks: dict[str, list[DuplexAssistantAudioTextMark]] = field(default_factory=dict) @@ -454,21 +107,264 @@ class ConversationHistory: ) +class DuplexFenceMismatchError(RuntimeError): + def __init__(self, expected: DuplexFence, actual: DuplexFence) -> None: + super().__init__(f"duplex fence mismatch: expected {expected!r}, got {actual!r}") + self.expected = expected + self.actual = actual + + @dataclass -class DuplexSession: +class DuplexRequestResource: + stage_id: int + request_id: str + fence: DuplexFence + submitted: bool = False + + +@dataclass +class DuplexInputAppend: + seq: int + turn_seq: int + turn_id: int + + +@dataclass(frozen=True) +class DuplexAppendReservation: + fence: DuplexFence + base_fence: DuplexFence + base_input_seq: int + base_input_turn_seq: int + base_append_turn_key: tuple[int, int, int] | None + update: DuplexInputAppend + + +@dataclass +class DuplexEngineSession: + """The one session state (see module docstring). + + Owned and mutated only by its ``DuplexSessionRunner`` on the orchestrator loop. + """ + session_id: str config: DuplexSessionConfig capabilities: DuplexCapabilities = field(default_factory=DuplexCapabilities) - incarnation: int = 0 state: DuplexSessionState = DuplexSessionState.OPEN turn_state: DuplexTurnState = DuplexTurnState.IDLE epoch: int = 0 turn_id: int = 0 + lease: DuplexLeaseState = field(default_factory=_default_lease, repr=False) + _clock: Callable[[], float] = field(default=time.monotonic, repr=False) + _runtime_config: dict[str, object] = field(default_factory=dict, repr=False) + #: Bumped on every published session / runtime config change. + config_generation: int = 0 + #: Stage request ids reserved or submitted for this session, keyed by ``(stage_id, request_id)``. + request_resources: dict[tuple[int, str], DuplexRequestResource] = field(default_factory=dict, repr=False) + #: Highest fence accepted for a stage request (monotonic; reset of the append + #: sequence happens when its epoch advances). + accepted_fence: DuplexFence = field(default=None, repr=False) # type: ignore[assignment] + input_seq: int = 0 + input_turn_seq: int = 0 + _append_turn_key: tuple[int, int, int] | None = field(default=None, repr=False) _input: InputBufferState = field(default_factory=InputBufferState, repr=False) _response: ResponseState = field(default_factory=ResponseState, repr=False) + #: Stage snapshots that arrived while no response was open. A stage whose + #: output feeds the next stage rather than the client reports its metrics + #: before the response those tokens end up in exists, so they are held here + #: and folded into the first response that opens after them. + _pending_stage_metrics: list[dict[str, dict[str, object]]] = field(default_factory=list, repr=False) _playback: PlaybackLedger = field(default_factory=PlaybackLedger, repr=False) _conversation: ConversationHistory = field(default_factory=ConversationHistory, repr=False) - _runtime_config: dict[str, object] = field(default_factory=dict, repr=False) + model_state: DuplexModelSessionState | None = field(default=None, repr=False) + projector: RealtimeProjectionState | None = field(default=None, repr=False) + created_monotonic: float = field(default_factory=time.monotonic) + + def __post_init__(self) -> None: + if self.accepted_fence is None: + self.accepted_fence = self.fence + else: + self.accept_fence(self.fence) + if self.config.initial_user_text: + # Seeded text is user input like any other: it waits for a response. + self.notify_new_user_item() + + # ---- identity / fence ---- + + @property + def fence(self) -> DuplexFence: + return DuplexFence(self.session_id, epoch=self.epoch, turn_id=self.turn_id) + + def sync_fence(self) -> DuplexFence: + """Publish the current (epoch, turn_id) identity as the accepted fence.""" + fence = self.fence + self.accept_fence(fence) + return fence + + def _validate_fence(self, fence: DuplexFence) -> None: + current = self.accepted_fence + if fence.session_id != self.session_id: + raise DuplexFenceMismatchError(current, fence) + if fence.epoch < current.epoch or (fence.epoch == current.epoch and fence.turn_id < current.turn_id): + raise DuplexFenceMismatchError(current, fence) + + def accept_fence(self, fence: DuplexFence) -> None: + self._validate_fence(fence) + if fence.epoch != self.accepted_fence.epoch: + self.input_seq = 0 + self.input_turn_seq = 0 + self._append_turn_key = None + self.accepted_fence = fence + + # ---- lease ---- + + @property + def lease_generation(self) -> int: + return self.lease.generation + + def touch_lease(self, activity: DuplexLeaseActivity) -> None: + self._validate_fence(self.fence) + self.lease.touch(self._clock(), activity) + + def detach_lease(self) -> None: + self._validate_fence(self.fence) + self.lease.detach(self._clock()) + + def resume_lease(self, *, expected_lease_generation: int) -> int: + self._validate_fence(self.fence) + return self.lease.resume(self._clock(), expected_generation=expected_lease_generation) + + def begin_lease_operation(self, fence: DuplexFence, operation_id: str) -> None: + self._validate_fence(fence) + self.lease.begin_operation(self._clock(), operation_id) + + def end_lease_operation(self, operation_id: str) -> None: + self._validate_fence(self.fence) + self.lease.end_operation(self._clock(), operation_id) + + # ---- stage request resources ---- + + def reserve_stage_request(self, stage_id: int, request_id: str, *, fence: DuplexFence) -> None: + self._validate_fence(fence) + resource_key = (stage_id, request_id) + existing = self.request_resources.get(resource_key) + if existing is not None: + if existing.fence.session_id != fence.session_id or existing.fence.epoch != fence.epoch: + raise ValueError(f"Duplex request resource already reserved with different identity: {request_id}") + return + self.request_resources[resource_key] = DuplexRequestResource( + stage_id=stage_id, + request_id=request_id, + fence=fence, + ) + + def bind_stage_request(self, stage_id: int, request_id: str, *, fence: DuplexFence) -> None: + self.reserve_stage_request(stage_id, request_id, fence=fence) + self.accept_fence(fence) + resource = self.request_resources[(stage_id, request_id)] + resource.fence = fence + resource.submitted = True + + def stage_request_submitted(self, stage_id: int, request_id: str) -> bool: + resource = self.request_resources.get((stage_id, request_id)) + return resource is not None and resource.submitted + + def resource_request_ids( + self, + fence: DuplexFence | None = None, + *, + submitted: bool | None = None, + ) -> list[str]: + return list( + dict.fromkeys( + resource.request_id + for resource in self.request_resources.values() + if (fence is None or resource.fence == fence) and (submitted is None or resource.submitted is submitted) + ) + ) + + def release_all_requests(self) -> list[str]: + """Drop every reserved/submitted stage request; returns their ids for stage cleanup.""" + stale = self.resource_request_ids() + self.request_resources.clear() + return stale + + # ---- append sequencing ---- + + def prepare_append(self, fence: DuplexFence) -> DuplexAppendReservation: + self._validate_fence(fence) + same_epoch = fence.epoch == self.accepted_fence.epoch + input_seq = self.input_seq if same_epoch else 0 + input_turn_seq = self.input_turn_seq if same_epoch else 0 + append_turn_key = self._append_turn_key if same_epoch else None + turn_key = (fence.epoch, fence.turn_id, 0) + turn_seq = input_turn_seq + 1 if turn_key == append_turn_key else 1 + return DuplexAppendReservation( + fence=fence, + base_fence=self.accepted_fence, + base_input_seq=self.input_seq, + base_input_turn_seq=self.input_turn_seq, + base_append_turn_key=self._append_turn_key, + update=DuplexInputAppend( + seq=input_seq + 1, + turn_seq=turn_seq, + turn_id=fence.turn_id, + ), + ) + + def commit_append(self, reservation: DuplexAppendReservation) -> DuplexInputAppend: + if ( + self.accepted_fence != reservation.base_fence + or self.input_seq != reservation.base_input_seq + or self.input_turn_seq != reservation.base_input_turn_seq + or self._append_turn_key != reservation.base_append_turn_key + ): + raise RuntimeError("duplex append reservation is stale") + self.accept_fence(reservation.fence) + self.input_seq = reservation.update.seq + self.input_turn_seq = reservation.update.turn_seq + self._append_turn_key = (reservation.fence.epoch, reservation.fence.turn_id, 0) + return reservation.update + + # ---- cancel / close of stage resources ---- + + def release_fence(self, fence: DuplexFence) -> list[str]: + stale = self.resource_request_ids(fence) + self.request_resources = { + resource_key: resource + for resource_key, resource in self.request_resources.items() + if resource.fence != fence + } + return stale + + def cancel_fence(self, cancelled_fence: DuplexFence, next_fence: DuplexFence) -> list[str]: + stale = self.prepare_cancel_fence(cancelled_fence, next_fence) + self.release_fence(cancelled_fence) + return stale + + def prepare_cancel_fence(self, cancelled_fence: DuplexFence, next_fence: DuplexFence) -> list[str]: + """Advance the cancellation fence without dropping cleanup records.""" + current = self.accepted_fence + if cancelled_fence.session_id != self.session_id: + raise DuplexFenceMismatchError(current, cancelled_fence) + if next_fence.session_id != self.session_id or next_fence.epoch <= cancelled_fence.epoch: + raise DuplexFenceMismatchError(cancelled_fence, next_fence) + current_key = (current.epoch, current.turn_id) + cancelled_key = (cancelled_fence.epoch, cancelled_fence.turn_id) + next_key = (next_fence.epoch, next_fence.turn_id) + if cancelled_key > current_key: + raise DuplexFenceMismatchError(current, cancelled_fence) + if next_key > current_key: + self.accept_fence(next_fence) + return self.resource_request_ids(cancelled_fence) + + def begin_close(self, *, reason: str) -> bool: + """Make close irreversible while retaining stage resources for cleanup retry.""" + self.accept_fence(self.fence) + if self.lease.terminal_reason is not None: + return True + return self.lease.mark_terminal(reason) + + # ---- configuration ---- @property def response_config(self) -> DuplexSessionConfig: @@ -486,19 +382,12 @@ class DuplexSession: def replace_runtime_config(self, runtime_config: Mapping[str, object]) -> None: self._runtime_config = dict(runtime_config) + self.config_generation += 1 @property def input_commit_seq(self) -> int: return self._input.commit_seq - @property - def pending_text(self) -> tuple[str, ...]: - return tuple(self._input.pending_text) - - @property - def pending_audio(self) -> tuple[DuplexAudioChunk, ...]: - return tuple(self._input.pending_audio) - @property def history(self) -> tuple[dict[str, object], ...]: placeholders = {id(message) for message in self._conversation.history_item_placeholders.values()} @@ -620,10 +509,6 @@ class DuplexSession: def playback(self) -> DuplexPlaybackView: return self._playback.current.snapshot() - @property - def response_playbacks(self) -> Mapping[str, DuplexPlaybackView]: - return MappingProxyType({key: value.snapshot() for key, value in self._playback.by_response.items()}) - @property def history_item_ids(self) -> Mapping[str, dict[str, object]]: return MappingProxyType({key: dict(value) for key, value in self._conversation.item_ids.items()}) @@ -637,10 +522,11 @@ class DuplexSession: return MappingProxyType(dict(self._conversation.pending_truncations_ms)) def replace_config(self, config: DuplexSessionConfig) -> None: + previous_seed = self.config.initial_user_text self.config = config - - def replace_capabilities(self, capabilities: DuplexCapabilities) -> None: - self.capabilities = capabilities + self.config_generation += 1 + if config.initial_user_text and config.initial_user_text != previous_seed: + self.notify_new_user_item() def transition_turn(self, state: DuplexTurnState) -> None: self.turn_state = state @@ -660,45 +546,6 @@ class DuplexSession: def bind_response_turn(self, turn_id: int | None) -> None: self._response.active_response_turn_id = turn_id - def mark_model_turn_request_started(self, turn_id: int, started_at_s: float) -> None: - """Record the latest native request start that can own one model turn.""" - turn_id = int(turn_id) - started_at_s = float(started_at_s) - if self.active_response_turn_id == turn_id: - if self._response.active_response_request_started_at_s is None: - self._response.active_response_request_started_at_s = started_at_s - return - self._response.request_started_at_s_by_turn[turn_id] = started_at_s - - def mark_response_first_outputs( - self, - *, - observed_at_s: float, - has_text: bool, - has_audio: bool, - ) -> dict[str, object]: - """Return server-monotonic TTF metrics observed for the active response.""" - started_at_s = self._response.active_response_request_started_at_s - if started_at_s is None: - return {} - elapsed_ms = max(0.0, (float(observed_at_s) - started_at_s) * 1000.0) - if has_text and self._response.active_response_ttft_ms is None: - self._response.active_response_ttft_ms = elapsed_ms - if has_audio and self._response.active_response_ttfp_ms is None: - self._response.active_response_ttfp_ms = elapsed_ms - metrics: dict[str, object] = { - "source": "server_monotonic_request_start", - "measurement_origin": { - "ttft": "native model-turn request execution start to first non-empty text output", - "ttfp": "native model-turn request execution start to first audio output", - }, - } - if self._response.active_response_ttft_ms is not None: - metrics["ttft_ms"] = self._response.active_response_ttft_ms - if self._response.active_response_ttfp_ms is not None: - metrics["ttfp_ms"] = self._response.active_response_ttfp_ms - return metrics - def active_response_accepts_model_turn(self, turn_id: int | None) -> bool: if self._response.active_response_id is None: return False @@ -708,19 +555,19 @@ class DuplexSession: return active_turn_id is None or int(turn_id) == active_turn_id def append_history_message(self, message: dict[str, object]) -> None: - response_id = self.active_response_id - if response_id is not None: - self.reserve_history_item(f"item_{response_id}") self._conversation.messages.append(message) - def stage_pending_history_item(self, item_id: str, message: dict[str, object]) -> None: - self._conversation.pending_item_ids[item_id] = dict(message) + def notify_new_user_item(self) -> None: + """Record a user item that a later ``response.create`` may answer.""" + self._conversation.unanswered_user_items += 1 - def append_text(self, text: str) -> None: - if not text: - return - self._input.pending_text.append(text) - self.turn_state = DuplexTurnState.USER_SPEAKING + def unanswered_user_items(self) -> int: + """How many user items are waiting for a response.""" + return self._conversation.unanswered_user_items + + def reset_unanswered_user_items(self) -> None: + """A turn has started: the pending items are now its input.""" + self._conversation.unanswered_user_items = 0 @property def pending_input_bytes(self) -> int: @@ -737,13 +584,6 @@ class DuplexSession: self._input.reserved_input_bytes += size return True - @property - def turn_detection_config_locked(self) -> bool: - return self._input.turn_detection_config_locked - - def lock_turn_detection_config(self) -> None: - self._input.turn_detection_config_locked = True - def release_input_bytes(self, size: int) -> None: self._input.reserved_input_bytes = max(0, self._input.reserved_input_bytes - max(0, int(size))) @@ -759,200 +599,18 @@ class DuplexSession: def release_pending_turn(self) -> None: self._input.pending_turns = max(0, self._input.pending_turns - 1) - def clear_pending_turn_reservations(self) -> None: - self._input.pending_turns = 0 - - def append_audio(self, data: str, *, fmt: str = "wav", sample_rate_hz: int | None = None) -> None: - if not data: - return - self._input.pending_audio.append(DuplexAudioChunk(data=data, format=fmt, sample_rate_hz=sample_rate_hz)) - self.turn_state = DuplexTurnState.USER_SPEAKING - - def append_server_vad_frame( - self, - samples: np.ndarray, - *, - speech_started: bool, - speech_stopped: bool, - prefix_samples: int, - ) -> int: - """Store one detector frame while keeping commit-eligible audio session-owned. - - Returns the number of raw float32 bytes released when the rolling - pre-speech prefix is trimmed. - """ - # Session-owned audio must not retain the batch backing a detector frame. - frame = np.array(samples, dtype=np.float32, copy=True).reshape(-1) - released_samples = 0 - if not self._input.vad_speech_active: - if speech_started: - self._input.vad_utterance_audio = [*self._input.vad_prefix_audio, frame] - self._input.vad_prefix_audio = [] - self._input.vad_prefix_samples = 0 - self._input.vad_speech_active = True - else: - self._input.vad_prefix_audio.append(frame) - self._input.vad_prefix_samples += frame.size - while self._input.vad_prefix_samples > prefix_samples and self._input.vad_prefix_audio: - overflow = self._input.vad_prefix_samples - prefix_samples - first = self._input.vad_prefix_audio[0] - if first.size <= overflow: - self._input.vad_prefix_audio.pop(0) - self._input.vad_prefix_samples -= first.size - released_samples += first.size - else: - self._input.vad_prefix_audio[0] = first[overflow:].copy() - self._input.vad_prefix_samples -= overflow - released_samples += overflow - else: - self._input.vad_utterance_audio.append(frame) - - if speech_stopped: - self._input.vad_speech_active = False - if self._input.vad_speech_active or speech_started: - self.turn_state = DuplexTurnState.USER_SPEAKING - return released_samples * np.dtype(np.float32).itemsize - - def stage_server_vad_audio_for_commit(self) -> bool: - if not self._input.vad_utterance_audio: - return False - samples = np.concatenate(self._input.vad_utterance_audio) - encoded = encode_float32_mono_wav_base64(samples, sample_rate_hz=16_000) - self._input.pending_audio.append(DuplexAudioChunk(data=encoded, format="wav", sample_rate_hz=16_000)) - self._input.vad_utterance_audio.clear() - self._input.vad_speech_active = False - return True - - def clear_server_vad_audio(self) -> None: - self._input.vad_prefix_audio.clear() - self._input.vad_prefix_samples = 0 - self._input.vad_utterance_audio.clear() - self._input.vad_speech_active = False - self._input.active_server_vad_item_id = None - - @property - def active_server_vad_item_id(self) -> str | None: - return self._input.active_server_vad_item_id - - @property - def server_vad_utterance_bytes(self) -> int: - return sum(frame.nbytes for frame in self._input.vad_utterance_audio) - - def discard_uncommitted_server_vad_utterance(self) -> int: - released = self.server_vad_utterance_bytes - self._input.vad_utterance_audio.clear() - self._input.vad_speech_active = False - self._input.active_server_vad_item_id = None - self._input.pending_text.clear() - return released - - def begin_server_vad_speech(self, item_id: str) -> None: - self._input.active_server_vad_item_id = item_id - - def finish_server_vad_speech(self) -> str | None: - item_id = self._input.active_server_vad_item_id - self._input.active_server_vad_item_id = None - return item_id - def mark_user_input_activity(self) -> None: self.turn_state = DuplexTurnState.USER_SPEAKING - def mark_assistant_generating(self) -> None: - self.turn_state = DuplexTurnState.ASSISTANT_GENERATING - def cancel_pending_input(self) -> dict[str, int]: - cancelled = { - "text_chunks": len(self._input.pending_text), - "audio_chunks": len(self._input.pending_audio), - } - self._input.pending_text.clear() - self._input.pending_audio.clear() + """Drop every pending reservation; the PCM buffer itself is model state cleared by the runner.""" + cancelled = {"text_chunks": 0, "audio_chunks": 0} self._input.reserved_input_bytes = 0 self._input.pending_turns = 0 - self.clear_server_vad_audio() self.turn_state = DuplexTurnState.IDLE return cancelled - def commit_user_input( - self, - *, - append_history: bool = True, - retain_input_bytes: int = 0, - ) -> DuplexCommittedInput | None: - text = "".join(self._input.pending_text).strip() - audio_chunks = list(self._input.pending_audio) - if not text and not audio_chunks: - return None - - content: str | list[dict[str, object]] - if audio_chunks: - content_items: list[dict[str, object]] = [] - if text: - content_items.append({"type": "text", "text": text}) - for chunk in audio_chunks: - content_items.append( - { - "type": "audio_url", - "audio_url": { - "url": f"data:audio/{chunk.format};base64,{chunk.data}", - }, - } - ) - content = content_items - else: - content = text - - self._bind_active_response_to_input_commit(self._input.commit_seq + 1) - self._input.commit_seq += 1 - message = {"role": "user", "content": content} - if append_history: - self._conversation.messages.append(message) - self._input.pending_text.clear() - self._input.pending_audio.clear() - self._input.reserved_input_bytes = max(0, int(retain_input_bytes)) - self.turn_state = DuplexTurnState.USER_COMMITTED - return DuplexCommittedInput( - message=message, - turn_id=self.turn_id, - epoch=self.epoch, - input_commit_seq=self._input.commit_seq, - ) - - @property - def pending_server_vad_turn(self) -> PendingServerVADTurn | None: - return self._input.pending_server_vad_turn - - def stage_server_vad_turn( - self, - committed: DuplexCommittedInput, - *, - create_response: bool, - item_id: str, - ) -> None: - if self._input.pending_server_vad_turn is not None: - raise RuntimeError("A Server VAD turn is already pending") - response_id = self.active_response_id - if response_id is not None: - # A pending user turn advances input_commit_seq before the active - # response completes. Reserve its current history position first. - self.reserve_history_item(f"item_{response_id}") - self._input.pending_server_vad_turn = PendingServerVADTurn( - committed=committed, - create_response=create_response, - item_id=item_id, - ) - - def activate_server_vad_turn(self) -> PendingServerVADTurn | None: - pending = self._input.pending_server_vad_turn - if pending is None: - return None - self._input.pending_server_vad_turn = None - self._conversation.messages.append(pending.committed.message) - self.register_history_item(pending.item_id, pending.committed.message) - self.turn_state = DuplexTurnState.USER_COMMITTED - return pending - - def commit_native_audio_input( + def commit_audio_input( self, *, transcript: str | None = None, @@ -970,8 +628,6 @@ class DuplexSession: if transcript: message["transcript"] = transcript self._conversation.messages.append(message) - self._input.pending_text.clear() - self._input.pending_audio.clear() self.turn_state = DuplexTurnState.USER_COMMITTED return DuplexCommittedInput( message=message, @@ -983,13 +639,9 @@ class DuplexSession: def complete_model_turn(self, turn_id: int) -> None: """Advance the model-owned output identity after its terminal signal.""" completed_turn_id = int(turn_id) - self._response.request_started_at_s_by_turn = { - pending_turn_id: started_at_s - for pending_turn_id, started_at_s in self._response.request_started_at_s_by_turn.items() - if pending_turn_id > completed_turn_id - } if completed_turn_id >= self.turn_id: self.turn_id = completed_turn_id + 1 + self.sync_fence() def _bind_active_response_to_input_commit(self, input_commit_seq: int) -> None: if self.active_response_id is not None and self._response.active_response_awaits_input_commit: @@ -1024,12 +676,6 @@ class DuplexSession: response_id = f"resp-{self.session_id}-{self.epoch}-{uuid4().hex[:8]}" self._response.active_response_id = response_id self._response.active_response_turn_id = self.turn_id if turn_id is None else int(turn_id) - self._response.active_response_request_started_at_s = self._response.request_started_at_s_by_turn.pop( - self._response.active_response_turn_id, - None, - ) - self._response.active_response_ttft_ms = None - self._response.active_response_ttfp_ms = None self._response.active_response_input_commit_seq = self.input_commit_seq self._response.active_response_awaits_input_commit = self.turn_state == DuplexTurnState.USER_SPEAKING self._response.last_response_id = response_id @@ -1048,11 +694,39 @@ class DuplexSession: self._response.stage_metric_tpot_weighted_ms.clear() self._response.stage_metric_tpot_weight.clear() + def stash_stage_metrics(self, stage_metrics: Mapping[object, object] | None) -> None: + """Hold a stage snapshot until a response exists to attribute it to. + + A stage that hands its output to the next stage instead of the client + (stage 0 feeding the TTS stage) reports its token metrics before the + response carrying those tokens is created. Dropping them would leave + every response without an engine-side token count; holding them keeps + the count whole, at the cost of attributing a turn the model never + spoke to the next response it does speak. + """ + if not isinstance(stage_metrics, Mapping): + return + snapshot = { + str(stage_id): dict(values) for stage_id, values in stage_metrics.items() if isinstance(values, Mapping) + } + if not snapshot: + return + if self.active_response_id is not None: + self.accumulate_response_stage_metrics(snapshot) + return + self._pending_stage_metrics.append(snapshot) + def accumulate_response_stage_metrics( self, stage_metrics: Mapping[object, object] | None, ) -> dict[str, dict[str, object]]: - if self.active_response_id is None or not isinstance(stage_metrics, Mapping): + if self.active_response_id is None: + return copy.deepcopy(self._response.stage_metrics) + if self._pending_stage_metrics: + pending, self._pending_stage_metrics = self._pending_stage_metrics, [] + for held in pending: + self.accumulate_response_stage_metrics(held) + if not isinstance(stage_metrics, Mapping): return copy.deepcopy(self._response.stage_metrics) additive_fields = ( @@ -1111,22 +785,17 @@ class DuplexSession: tpot_ms = raw_values.get("vllm_tpot_ms") token_count = raw_values.get("num_tokens_out") - if isinstance(tpot_ms, int | float) and not isinstance(tpot_ms, bool) and tpot_ms > 0: - weight = ( - max(int(token_count) - 1, 0) - if isinstance(token_count, int | float) and not isinstance(token_count, bool) - else 1 + if isinstance(tpot_ms, int | float) and tpot_ms > 0: + weight = max(int(token_count) - 1, 1) if isinstance(token_count, int | float) else 1 + self._response.stage_metric_tpot_weighted_ms[stage_id] = ( + self._response.stage_metric_tpot_weighted_ms.get(stage_id, 0.0) + float(tpot_ms) * weight + ) + self._response.stage_metric_tpot_weight[stage_id] = ( + self._response.stage_metric_tpot_weight.get(stage_id, 0) + weight + ) + current["vllm_tpot_ms"] = self._response.stage_metric_tpot_weighted_ms[stage_id] / float( + self._response.stage_metric_tpot_weight[stage_id] ) - if weight > 0: - self._response.stage_metric_tpot_weighted_ms[stage_id] = ( - self._response.stage_metric_tpot_weighted_ms.get(stage_id, 0.0) + float(tpot_ms) * weight - ) - self._response.stage_metric_tpot_weight[stage_id] = ( - self._response.stage_metric_tpot_weight.get(stage_id, 0) + weight - ) - current["vllm_tpot_ms"] = self._response.stage_metric_tpot_weighted_ms[stage_id] / float( - self._response.stage_metric_tpot_weight[stage_id] - ) for name, value in raw_values.items(): if name not in handled_fields: @@ -1164,54 +833,25 @@ class DuplexSession: if text: self._response.assistant_text_buffer.append(text) - def record_generated_audio(self, response_id: str, *, frame_count: int, sample_rate_hz: int) -> int | None: - """Accumulate source frames, rounding only the cumulative duration. - - A late chunk cannot change a different response's playback cursor. - This records encoded audio, not transport delivery or client playback. - """ - if response_id != self.active_response_id: - return None - if type(frame_count) is not int or frame_count < 0: - raise ValueError("Audio frame_count must be a non-negative integer") - if type(sample_rate_hz) is not int or sample_rate_hz <= 0: - raise ValueError("Audio sample_rate_hz must be a positive integer") - playback = self._playback.current - if playback.sample_rate_hz is not None and playback.sample_rate_hz != sample_rate_hz: - raise ValueError("Audio sample rate changed within a response") - playback.sample_rate_hz = sample_rate_hz - playback.generated_frames += frame_count - playback.generated_ms = playback.generated_frames * 1000 // sample_rate_hz - return playback.generated_ms - def mark_audio_sent( self, duration_ms: int | None = None, *, - response_id: str | None = None, text_chars: int | None = None, audio_text_marks: list[dict[str, object]] | None = None, ) -> None: - """Record cumulative audio accepted by the output path. - - ``sent_ms`` includes queued/journaled output; only client acknowledgements - establish ``played_ms``. It is not confirmation of socket delivery. - """ - playback = self._playback.current if response_id is None else self._playback.by_response.get(response_id) - if playback is None: - return - is_current = response_id is None or response_id == self.active_response_id + playback = self._playback.current if duration_ms is not None: playback.generated_ms = max(playback.generated_ms, duration_ms) playback.sent_ms = max(playback.sent_ms, duration_ms) - if is_current and text_chars is not None and text_chars >= 0: + if text_chars is not None and text_chars >= 0: self._response.assistant_audio_text_marks.append( DuplexAssistantAudioTextMark( text_chars=int(text_chars), audio_end_ms=max(0, int(duration_ms)), ) ) - if is_current and audio_text_marks: + if audio_text_marks: for raw_mark in audio_text_marks: if not isinstance(raw_mark, dict): continue @@ -1225,8 +865,7 @@ class DuplexSession: audio_end_ms=max(0, int(raw_audio_end_ms)), ) ) - if is_current: - self.turn_state = DuplexTurnState.ASSISTANT_PLAYING + self.turn_state = DuplexTurnState.ASSISTANT_PLAYING def _playback_cursor_for_response(self, response_id: str | None = None) -> DuplexPlaybackCursor: if response_id is None: @@ -1284,14 +923,7 @@ class DuplexSession: response_input_commit_seq = self._response.active_response_input_commit_seq if response_input_commit_seq is None: response_input_commit_seq = self.input_commit_seq - response_item_id = f"item_{response_id}" if response_id is not None else None - response_history_is_reserved = response_item_id is not None and ( - response_item_id in self._conversation.item_ids - or response_item_id in self._conversation.history_item_placeholders - ) - response_history_is_late = ( - self.input_commit_seq > response_input_commit_seq and not response_history_is_reserved - ) + response_history_is_late = self.input_commit_seq > response_input_commit_seq assistant_text = "".join(self._response.assistant_text_buffer).strip() message = None if assistant_text: @@ -1312,8 +944,9 @@ class DuplexSession: committed_text = "" if commit_text and committed_text: message = {"role": "assistant", "content": committed_text} - if response_item_id is not None and response_history_is_reserved: - self._store_history_item_message(response_item_id, message) + item_id = f"item_{response_id}" if response_id is not None else None + if item_id is not None and item_id in self._conversation.history_item_placeholders: + self._store_history_item_message(item_id, message) else: self._conversation.messages.append(message) effective_playback_policy = playback_commit_policy or self.config.playback_commit_policy @@ -1324,16 +957,13 @@ class DuplexSession: and effective_playback_policy == DuplexPlaybackCommitPolicy.ACK_ONLY.value ): self.register_history_item(f"item_{response_id}", None) - elif response_id is not None and (not assistant_text or not commit_text): + elif response_id is not None and not assistant_text: self._discard_history_item_placeholder(f"item_{response_id}") self._response.assistant_text_buffer.clear() if not preserve_request: self._response.active_request_id = None self._response.active_response_id = None self._response.active_response_turn_id = None - self._response.active_response_request_started_at_s = None - self._response.active_response_ttft_ms = None - self._response.active_response_ttfp_ms = None self._response.active_response_input_commit_seq = None self._response.active_response_awaits_input_commit = False self._clear_response_metrics() @@ -1386,11 +1016,6 @@ class DuplexSession: ) def delete_history_item(self, item_id: str) -> bool: - removed_pending_turn = False - pending_server_vad_turn = self._input.pending_server_vad_turn - if pending_server_vad_turn is not None and pending_server_vad_turn.item_id == item_id: - self._input.pending_server_vad_turn = None - removed_pending_turn = True response_id = item_id.removeprefix("item_") if item_id.startswith("item_") else None message = self._conversation.item_ids.pop(item_id, None) self._conversation.item_audio_text_marks.pop(item_id, None) @@ -1403,7 +1028,7 @@ class DuplexSession: if response_id is not None: self._conversation.assistant_response_snapshots.pop(response_id, None) if message is None: - return removed_pending_turn or pending is not None or removed_placeholder + return pending is not None or removed_placeholder try: self._conversation.messages.remove(message) except ValueError: @@ -1644,17 +1269,12 @@ class DuplexSession: def barge_in(self) -> int: self.epoch += 1 - self._input.pending_text.clear() - self._input.pending_audio.clear() + self.sync_fence() self._response.assistant_text_buffer.clear() self._response.assistant_audio_text_marks.clear() self._response.active_request_id = None self._response.active_response_id = None self._response.active_response_turn_id = None - self._response.request_started_at_s_by_turn.clear() - self._response.active_response_request_started_at_s = None - self._response.active_response_ttft_ms = None - self._response.active_response_ttfp_ms = None self._response.active_response_input_commit_seq = None self._response.active_response_awaits_input_commit = False self._clear_response_metrics() @@ -1670,17 +1290,35 @@ class DuplexSession: self.state = DuplexSessionState.CLOSED self.turn_state = DuplexTurnState.IDLE self._response.active_response_turn_id = None - self._response.request_started_at_s_by_turn.clear() - self._response.active_response_request_started_at_s = None - self._response.active_response_ttft_ms = None - self._response.active_response_ttfp_ms = None self._response.active_response_input_commit_seq = None self._response.active_response_awaits_input_commit = False - self.clear_server_vad_audio() - self._input.pending_server_vad_turn = None self._clear_response_metrics() self._restore_response_config() + def signal_turn(self, event_type: str, payload: Mapping[str, object] | None = None) -> TurnEvent: + """Apply one external turn signal and return the typed ``turn.event``.""" + payload = payload or {} + if event_type == DuplexTurnEventType.USER_STARTED.value: + self.transition_turn(DuplexTurnState.USER_SPEAKING) + elif event_type == DuplexTurnEventType.USER_COMMITTED.value: + self.transition_turn(DuplexTurnState.USER_COMMITTED) + elif event_type == DuplexTurnEventType.ASSISTANT_STARTED.value: + self.transition_turn(DuplexTurnState.ASSISTANT_GENERATING) + elif event_type == DuplexTurnEventType.ASSISTANT_DONE.value: + self.transition_turn(DuplexTurnState.IDLE) + elif event_type == DuplexTurnEventType.PLAYBACK_ACK.value: + played_ms = int(payload.get("played_ms", 0) or 0) + committed_ms = payload.get("committed_ms") + self.acknowledge_playback( + played_ms, + int(committed_ms) if isinstance(committed_ms, int | float) else None, + ) + elif event_type == DuplexTurnEventType.BARGE_IN.value: + self.transition_turn(DuplexTurnState.BARGE_IN) + elif event_type in {DuplexTurnEventType.CLOSE.value, DuplexTurnEventType.TIMEOUT.value}: + self.transition_session(DuplexSessionState.CLOSING) + return TurnEvent(event=event_type, turn_state=self.turn_state.value) + def as_public_dict(self) -> dict[str, object]: payload: dict[str, object] = { "id": self.session_id, @@ -1708,8 +1346,6 @@ class DuplexSession: "playback": self.playback.as_dict(), "capabilities": self.capabilities.as_dict(), } - if self.config.turn_detection_configured: - payload["turn_detection"] = self.config.server_vad.as_dict() if self.config.server_vad is not None else None if isinstance(self.config.extra_body.get("realtime_tools"), list): payload["tools"] = self.config.extra_body["realtime_tools"] if isinstance(self.config.extra_body.get("realtime_tool_choice"), str | dict): @@ -1738,73 +1374,14 @@ class DuplexSession: return payload -class DuplexTurnController: - """Interaction controller that accepts signals from multiple sources.""" - - def signal( - self, - session: DuplexSession, - event_type: str, - payload: dict[str, object] | None = None, - ) -> dict[str, object]: - payload = payload or {} - if event_type == DuplexTurnEventType.USER_STARTED.value: - session.transition_turn(DuplexTurnState.USER_SPEAKING) - elif event_type == DuplexTurnEventType.USER_COMMITTED.value: - session.transition_turn(DuplexTurnState.USER_COMMITTED) - elif event_type == DuplexTurnEventType.ASSISTANT_STARTED.value: - session.transition_turn(DuplexTurnState.ASSISTANT_GENERATING) - elif event_type == DuplexTurnEventType.ASSISTANT_DONE.value: - session.transition_turn(DuplexTurnState.IDLE) - elif event_type == DuplexTurnEventType.PLAYBACK_ACK.value: - played_ms = int(payload.get("played_ms", 0) or 0) - committed_ms = payload.get("committed_ms") - session.acknowledge_playback( - played_ms, - int(committed_ms) if isinstance(committed_ms, int | float) else None, - ) - elif event_type == DuplexTurnEventType.BARGE_IN.value: - session.transition_turn(DuplexTurnState.BARGE_IN) - elif event_type in {DuplexTurnEventType.CLOSE.value, DuplexTurnEventType.TIMEOUT.value}: - session.transition_session(DuplexSessionState.CLOSING) - return { - "type": "turn.event", - "session_id": session.session_id, - "event": event_type, - "turn_state": session.turn_state.value, - "epoch": session.epoch, - } - - -class DuplexSessionRegistry: - def __init__(self, capabilities: DuplexCapabilities | None = None) -> None: - self._capabilities = capabilities or DuplexCapabilities() - self._sessions: dict[str, DuplexSession] = {} - self._next_incarnation_by_session_id: dict[str, int] = {} - - def create(self, config: DuplexSessionConfig | None = None, session_id: str | None = None) -> DuplexSession: - sid = session_id or f"duplex-{uuid4().hex}" - if sid in self._sessions: - raise ValueError(f"Duplex session already exists: {sid}") - incarnation = self._next_incarnation_by_session_id.get(sid, 0) - self._next_incarnation_by_session_id[sid] = incarnation + 1 - session = DuplexSession( - session_id=sid, - config=config or DuplexSessionConfig(), - capabilities=self._capabilities, - incarnation=incarnation, - ) - self._sessions[sid] = session - return session - - def active_count(self) -> int: - return len(self._sessions) - - def get(self, session_id: str) -> DuplexSession | None: - return self._sessions.get(session_id) - - def close(self, session_id: str) -> DuplexSession | None: - session = self._sessions.pop(session_id, None) - if session is not None: - session.close() - return session +__all__ = [ + "ConversationHistory", + "DuplexAppendReservation", + "DuplexEngineSession", + "DuplexFenceMismatchError", + "DuplexInputAppend", + "DuplexRequestResource", + "InputBufferState", + "PlaybackLedger", + "ResponseState", +] diff --git a/vllm_omni/engine/duplex/session/helpers.py b/vllm_omni/engine/duplex/session/helpers.py new file mode 100644 index 000000000..f8b705e6d --- /dev/null +++ b/vllm_omni/engine/duplex/session/helpers.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Small reads and payload builders over one duplex session. + +None of these hold state or decide anything: they answer a question about the +session, or shape one of the wire payloads its events carry. They were methods +on the runner only because that is where the session reference lived, which +made every component that needed one a reason to reach back into the runner. + +Functions, not a component, because there is nothing here to own. +""" + +from __future__ import annotations + +import binascii +from typing import TYPE_CHECKING + +import pybase64 as base64 + +from vllm_omni.engine.duplex.config import DuplexPlaybackCommitPolicy +from vllm_omni.engine.duplex.contracts import DuplexFence, duplex_resource_request_id +from vllm_omni.engine.duplex.events import ErrorEvent, OverlapDecision, error_event + +if TYPE_CHECKING: + from collections.abc import Mapping + + from vllm_omni.engine.duplex.session.context import DuplexSessionTasks + from vllm_omni.engine.duplex.session.engine_session import DuplexCommittedInput, DuplexEngineSession + + +# --------------------------------------------------------------------------- # +# Questions about the session # +# --------------------------------------------------------------------------- # + + +def stage0_request_id(session: DuplexEngineSession, epoch: int) -> str: + return duplex_resource_request_id(DuplexFence(session.session_id, epoch=epoch), "stage0") + + +def response_in_progress(session: DuplexEngineSession, tasks: DuplexSessionTasks) -> bool: + """Whether the model still owns the turn. + + Broader than ``active_response_id``: audio already sent but not yet played + back, and an append still in flight, both mean the turn is not free. + """ + if session.active_response_id is not None: + return True + if assistant_playback_active(session): + return True + active_task = tasks.active_response_task + if active_task is not None and not active_task.done(): + return True + return tasks.has_response_bound_append_tasks() + + +def assistant_playback_active(session: DuplexEngineSession) -> bool: + """Whether the client is still playing audio the session has sent.""" + return ( + session.config.playback_commit_policy == DuplexPlaybackCommitPolicy.ACK_ONLY.value + and session.playback.sent_ms > session.playback.committed_ms + ) + + +def audio_payload_size_bytes(payload: Mapping[str, object]) -> int: + audio = payload.get("audio") or payload.get("data") + if not isinstance(audio, str): + return 0 + try: + return len(base64.b64decode(audio, validate=True)) + except (ValueError, binascii.Error): + return 0 + + +# --------------------------------------------------------------------------- # +# Turn transitions # +# --------------------------------------------------------------------------- # + + +def advance_barge_in_epoch(session: DuplexEngineSession) -> tuple[int, dict[str, int]]: + """Start a new epoch for a barge-in, returning it with the playback it cut off.""" + old_playback = session.playback.as_dict() + new_epoch = session.barge_in() + session.clear_playback_cursor() + return new_epoch, old_playback + + +def commit_played_response_history( + session: DuplexEngineSession, + response_id: str | None, + committed_ms: int, +) -> None: + """Truncate an interrupted response in history to what the client actually heard.""" + if not response_id or committed_ms < 0: + return + session.truncate_history_item( + f"item_{response_id}", + audio_end_ms=committed_ms, + playback=session.playback_for_response(response_id), + ) + + +def commit_audio_input( + session: DuplexEngineSession, + *, + realtime_item_id: object | None = None, + transcript: object | None = None, + turn_id: int | None = None, +) -> DuplexCommittedInput: + clean_transcript = transcript.strip() if isinstance(transcript, str) else None + committed = session.commit_audio_input( + transcript=clean_transcript or None, + turn_id=turn_id, + ) + if isinstance(realtime_item_id, str) and realtime_item_id: + session.register_history_item(realtime_item_id, committed.message) + return committed + + +# --------------------------------------------------------------------------- # +# Payloads # +# --------------------------------------------------------------------------- # + + +def barge_in_unsupported_error() -> ErrorEvent: + return error_event("barge_in_unsupported", "Barge-in is not supported by this duplex model") + + +def overlap_decision_event(session: DuplexEngineSession, decision: dict[str, object]) -> OverlapDecision: + details: dict[str, object] = { + "type": "overlap.decision", + "session_id": session.session_id, + "epoch": session.epoch, + "policy": session.config.overlap_policy, + **decision, + } + action = decision.get("action") + reason = decision.get("reason") + return OverlapDecision( + policy=session.config.overlap_policy, + action=action if isinstance(action, str) else None, + reason=reason if isinstance(reason, str) else None, + details=details, + ) + + +def input_committed_payload( + session: DuplexEngineSession, + committed: DuplexCommittedInput, + *, + realtime_item_id: object | None = None, +) -> dict[str, object]: + payload: dict[str, object] = { + "type": "input.committed", + "session_id": session.session_id, + "turn_id": committed.turn_id, + "epoch": committed.epoch, + "history_len": len(session.history), + "message": committed.message, + } + if isinstance(realtime_item_id, str) and realtime_item_id: + payload["realtime_item_id"] = realtime_item_id + return payload + + +def audio_committed_payload( + session: DuplexEngineSession, + *, + committed: DuplexCommittedInput | None = None, + realtime_item_id: object | None = None, + transcript: object | None = None, +) -> dict[str, object]: + message = committed.message if committed is not None else None + if not isinstance(message, dict): + input_audio_part: dict[str, object] = { + "type": "audio_url", + "audio_url": {"url": "native-duplex:input-audio"}, + } + if isinstance(transcript, str) and transcript: + input_audio_part["transcript"] = transcript + message = {"role": "user", "content": [input_audio_part]} + payload: dict[str, object] = { + "type": "input.committed", + "session_id": session.session_id, + "turn_id": committed.turn_id if committed is not None else session.turn_id, + "epoch": committed.epoch if committed is not None else session.epoch, + "history_len": len(session.history), + "native_audio": True, + "message": message, + } + if isinstance(transcript, str) and transcript: + payload["transcript"] = transcript + if isinstance(realtime_item_id, str) and realtime_item_id: + payload["realtime_item_id"] = realtime_item_id + return payload diff --git a/vllm_omni/engine/duplex/lease.py b/vllm_omni/engine/duplex/session/lease.py similarity index 92% rename from vllm_omni/engine/duplex/lease.py rename to vllm_omni/engine/duplex/session/lease.py index c38faacd1..aea204039 100644 --- a/vllm_omni/engine/duplex/lease.py +++ b/vllm_omni/engine/duplex/session/lease.py @@ -8,8 +8,6 @@ from __future__ import annotations from dataclasses import dataclass, field from enum import Enum -from vllm_omni.engine.duplex.messages import DuplexFence - class DuplexLeaseActivity(str, Enum): APPEND = "append" @@ -111,19 +109,8 @@ class DuplexLeaseState: return True -@dataclass(frozen=True) -class DuplexSessionExpiry: - session_id: str - fence: DuplexFence - lease_generation: int - reason: str - submitted_request_ids: tuple[str, ...] - reserved_request_ids: tuple[str, ...] - - __all__ = [ "DuplexLeaseActivity", "DuplexLeaseConfig", "DuplexLeaseState", - "DuplexSessionExpiry", ] diff --git a/vllm_omni/engine/duplex/session/manager.py b/vllm_omni/engine/duplex/session/manager.py new file mode 100644 index 000000000..0b7a7a07e --- /dev/null +++ b/vllm_omni/engine/duplex/session/manager.py @@ -0,0 +1,872 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Engine-side owner of duplex sessions: admission, lease reaping, command dispatch. + +``DuplexOrchestrator`` hosts exactly one manager. The whole session +(``DuplexEngineSession``) lives here, owned by one ``DuplexSessionRunner`` per +session on the orchestrator loop; nothing outside the engine holds session +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``. +""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +import time +from collections.abc import Awaitable, Callable, Iterable +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING + +from vllm.logger import init_logger + +from vllm_omni.engine.duplex.commands import AppendAudio, Commit, DuplexCommand +from vllm_omni.engine.duplex.contracts import ( + DuplexFence, + DuplexStagePort, + DuplexStageRequestContext, + 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.messages import ( + CloseDuplexSessionMessage, + DuplexControlResultMessage, + DuplexSessionCommandMessage, + DuplexSessionError, + DuplexSessionEventMessage, + OpenDuplexSessionMessage, + ResumeDuplexSessionMessage, + TouchDuplexSessionMessage, +) +from vllm_omni.engine.duplex.plugin import DuplexModelPlugin, DuplexRuntimeConfigError, validate_duplex_plugin_sampling +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession, DuplexFenceMismatchError +from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity, DuplexLeaseConfig, DuplexLeaseState +from vllm_omni.engine.duplex.turn_detection import SileroVADBackendProvider + +if TYPE_CHECKING: + import janus + from vllm.config import ModelConfig + + from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig + from vllm_omni.engine.duplex.session.runner import DuplexSessionRunner + from vllm_omni.engine.messages import EngineQueueMessage + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class _PendingRequestCleanup: + session_id: str + lease_generation: int + request_ids: tuple[str, ...] + abort: bool + + +@dataclass(frozen=True) +class _PendingSessionCleanup: + """A close/expiry whose stage cleanup has not succeeded yet (retried by the reaper).""" + + kind: str + session_id: str + lease_generation: int + reason: str + submitted_request_ids: tuple[str, ...] + reserved_request_ids: tuple[str, ...] + + +class DuplexSessionManager: + _CONTROL_TYPES = ( + OpenDuplexSessionMessage, + CloseDuplexSessionMessage, + ResumeDuplexSessionMessage, + TouchDuplexSessionMessage, + ) + _MESSAGE_TYPES = (*_CONTROL_TYPES, DuplexSessionCommandMessage) + + def __init__( + self, + *, + plugin: DuplexModelPlugin, + stage_port: DuplexStagePort, + output_sink: janus.AsyncQueue[EngineQueueMessage], + result_sink: janus.AsyncQueue[EngineQueueMessage], + runtime_config: DuplexSessionRuntimeConfig, + model_config: ModelConfig | None, + clock: Callable[[], float] | None = None, + executor: concurrent.futures.ThreadPoolExecutor | None = None, + ) -> None: + validate_duplex_plugin_sampling(plugin, sampling_defaults=tuple(stage_port.sampling_defaults())) + self.plugin = plugin + self.stage_port = stage_port + self.model_config = model_config + self.runtime_config = runtime_config + self._output_sink = output_sink + self._result_sink = result_sink + self._clock = clock or time.monotonic + self.executor = executor or concurrent.futures.ThreadPoolExecutor( + max_workers=4, + thread_name_prefix="duplex-session", + ) + self._owns_executor = executor is None + self._lease_config = DuplexLeaseConfig( + idle_ttl_s=runtime_config.idle_ttl_s, + disconnect_grace_s=runtime_config.disconnect_grace_s, + ) + #: One Silero backend for the whole engine: the ONNX graph is stateless + #: per call, so every session's detector can share it. Resolved lazily, + #: on the first session that actually turns server VAD on. + self.vad_backend_provider = SileroVADBackendProvider( + model_path=getattr(runtime_config, "server_vad_model_path", None) + ) + self.runners: dict[str, DuplexSessionRunner] = {} + #: Sessions whose close/expiry began but whose stage cleanup has not finalized; + #: they keep their admission slot until finalized. + self._closing: dict[str, _PendingSessionCleanup] = {} + #: Session ids whose ``open`` passed the capacity check and is still awaiting the plugin. + self._admitting: set[str] = set() + #: Opens whose caller gave up (control RPC timeout / cancellation) before + #: admission finished. The open completes anyway -- it is already past + #: the capacity check -- so it must tear itself down on arrival instead + #: of holding a slot and a stage resource until idle expiry. + self._abandoned_opens: set[str] = set() + self._request_index: dict[str, str] = {} + self._session_control_tails: dict[str, asyncio.Task[None]] = {} + self._dispatched_control_tasks: set[asyncio.Task[None]] = set() + self._pending_request_cleanups: dict[tuple[str, int], _PendingRequestCleanup] = {} + self._request_cleanup_tasks: dict[tuple[str, int], asyncio.Task[None]] = {} + self._request_cleanups_in_progress: set[tuple[str, int]] = set() + #: Sessions retained for cleanup retry after their runner is gone. + self._session_snapshots: dict[str, DuplexEngineSession] = {} + + # ------------------------------------------------------------------ # + # Views # + # ------------------------------------------------------------------ # + + def active_count(self) -> int: + return len(self.runners) + + def get(self, session_id: str) -> DuplexEngineSession | None: + runner = self.runners.get(session_id) + return runner.session if runner is not None else None + + def sessions(self) -> dict[str, DuplexEngineSession]: + return {session_id: runner.session for session_id, runner in self.runners.items()} + + def runner_for_request_id(self, request_id: str) -> DuplexSessionRunner | None: + session_id = self._request_index.get(request_id) + if session_id is None: + for candidate_id in self.runners: + if duplex_resource_request_belongs_to_session(request_id, candidate_id): + session_id = candidate_id + break + if session_id is None: + return None + return self.runners.get(session_id) + + def register_request(self, request_id: str, session_id: str) -> None: + self._request_index[request_id] = session_id + + def unregister_request(self, request_id: str) -> None: + self._request_index.pop(request_id, None) + + def _unregister_session_requests(self, session_id: str) -> None: + for request_id in [rid for rid, sid in self._request_index.items() if sid == session_id]: + self._request_index.pop(request_id, None) + + # ------------------------------------------------------------------ # + # Message routing # + # ------------------------------------------------------------------ # + + def accepts(self, message: object) -> bool: + return isinstance(message, self._MESSAGE_TYPES) + + def dispatch(self, message: object) -> None: + """Route one engine request-queue message without blocking the request handler.""" + if isinstance(message, DuplexSessionCommandMessage): + self._dispatch_command(message) + return + if not isinstance(message, self._CONTROL_TYPES): + raise TypeError(f"Unsupported duplex control message: {type(message).__name__}") + session_id = message.session_id + + async def handle_message() -> None: + await self.handle(message) + + self._run_control(session_id, f"duplex-control-{session_id}-{message.control_id}", handle_message) + + def _run_control( + self, + session_id: str, + name: str, + operation: Callable[[], Awaitable[None]], + ) -> asyncio.Task[None]: + """Run one control operation after the session's previous one (per-session order).""" + predecessor = self._session_control_tails.get(session_id) + + async def run_ordered() -> None: + if predecessor is not None: + try: + await predecessor + except BaseException: + pass + await operation() + + task = asyncio.create_task(run_ordered(), name=name) + self._session_control_tails[session_id] = task + self._track_task(task) + + def discard(completed: asyncio.Task[None]) -> None: + if self._session_control_tails.get(session_id) is completed: + self._session_control_tails.pop(session_id, None) + + task.add_done_callback(discard) + return task + + def _track_task(self, task: asyncio.Task[None]) -> None: + """Keep a background task alive until it finishes and log an escaped exception.""" + self._dispatched_control_tasks.add(task) + + def done(completed: asyncio.Task[None]) -> None: + self._dispatched_control_tasks.discard(completed) + if not completed.cancelled() and completed.exception() is not None: + logger.error("duplex control task %s failed: %r", completed.get_name(), completed.exception()) + + task.add_done_callback(done) + + async def handle(self, message: object) -> None: + if isinstance(message, OpenDuplexSessionMessage): + await self.open(message) + elif isinstance(message, CloseDuplexSessionMessage): + await self.close(message) + elif isinstance(message, ResumeDuplexSessionMessage): + await self.resume(message) + elif isinstance(message, TouchDuplexSessionMessage): + await self.touch(message) + elif isinstance(message, DuplexSessionCommandMessage): + self._dispatch_command(message) + else: + raise TypeError(f"Unsupported duplex control message: {type(message).__name__}") + + def _dispatch_command(self, message: DuplexSessionCommandMessage) -> None: + """Admit one command: identity check, backpressure, then the runner's ordered mailbox. + + Backpressure contract shared with the runner: a ``Commit`` that passes + admission holds one pending-turn reservation and an ``AppendAudio`` + holds its wire-size byte reservation until the runner dequeues the + command (the runner then re-reserves the decoded size around its PCM + reservation), so the mailbox never holds more bytes than the session + limit even while the worker is busy. + """ + runner = self.runners.get(message.session_id) + command = message.command + 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, + ), + ) + 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 + ), + ) + return + if isinstance(command, AppendAudio): + limit = int(self.runtime_config.max_pending_input_bytes_per_session) + if not session.reserve_input_bytes(len(command.audio), limit=limit): + self.emit( + session, + 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, + ), + ) + return + runner.submit(command) + + @staticmethod + def _error_event(code: str, message: str, *, command: DuplexCommand | None = None) -> ErrorEvent: + return error_event(code, message, event_id=command.event_id if command is not None else None) + + # ------------------------------------------------------------------ # + # 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_raw( + self, + session_id: str, + event: DuplexEvent, + *, + epoch: int | None = None, + ) -> None: + event = replace(event, session_id=session_id, epoch=epoch) + message = DuplexSessionEventMessage(session_id=session_id, event=event) + put_nowait = getattr(self._output_sink, "put_nowait", None) + if callable(put_nowait): + put_nowait(message) + else: # pragma: no cover - sink without put_nowait + asyncio.ensure_future(self._output_sink.put(message)) + + async def _put_result( + self, + message: ( + OpenDuplexSessionMessage + | CloseDuplexSessionMessage + | ResumeDuplexSessionMessage + | TouchDuplexSessionMessage + ), + *, + operation: str, + ok: bool, + session: DuplexEngineSession | None = None, + error: BaseException | None = None, + ) -> None: + error_code, error_message, error_retryable = ( + self._control_error(error) if error is not None else (None, None, False) + ) + result = DuplexControlResultMessage( + control_id=message.control_id, + operation=operation, + session_id=message.session_id, + ok=ok, + lease_generation=session.lease_generation if session is not None else None, + capabilities=session.capabilities if session is not None and ok else None, + public_session=session.as_public_dict() if session is not None and ok else None, + error_code=error_code, + error_message=error_message, + error_retryable=error_retryable, + ) + await self._result_sink.put(result) + + @staticmethod + def _control_error(error: BaseException) -> tuple[str, str, bool]: + """Map an exception to ``(code, message, retryable)`` for a control result.""" + message = str(error) + retryable = False + if isinstance(error, DuplexSessionError): + code = error.code + retryable = error.retryable + elif isinstance(error, DuplexRuntimeConfigError): + code = error.code + elif isinstance(error, DuplexFenceMismatchError): + code = "stale_fence" + elif isinstance(error, KeyError): + code = "not_found" + elif isinstance(error, TypeError | ValueError): + code = "invalid_argument" + elif isinstance(error, TimeoutError): + code = "timeout" + retryable = True + else: + code = "failed_precondition" + return code, message, retryable + + # ------------------------------------------------------------------ # + # Sampling / stage request helpers shared with the runner # + # ------------------------------------------------------------------ # + + def sampling_params_for(self, session: DuplexEngineSession) -> tuple[object, ...]: + defaults = tuple(self.stage_port.sampling_defaults()) + configured = self.plugin.configure_sampling_params( + runtime_config=dict(session.runtime_config), + defaults=defaults, + ) + if not isinstance(configured, tuple): + raise TypeError("duplex plugin must return sampling parameters as a tuple") + if len(configured) != len(defaults): + raise ValueError("duplex plugin must return one sampling parameter per stage") + return configured + + @staticmethod + def stage_request_id(fence: DuplexFence, *, stage_id: int) -> str: + return duplex_resource_request_id(fence, f"stage{stage_id}") + + def ensure_stage_request( + self, + session: DuplexEngineSession, + *, + stage_id: int, + fence: DuplexFence | None = None, + ) -> DuplexStageRequestContext | None: + """Reserve the session's stage request id and register it with the stage port.""" + if stage_id >= self.stage_port.stage_count: + return None + effective_fence = fence or session.fence + request_id = self.stage_request_id(effective_fence, stage_id=stage_id) + session.reserve_stage_request(stage_id, request_id, fence=effective_fence) + context = DuplexStageRequestContext( + request_id=request_id, + session_id=session.session_id, + fence=effective_fence, + stage_id=stage_id, + final_stage_id=self.stage_port.stage_count - 1, + config_generation=session.config_generation, + sampling_params=self.sampling_params_for(session), + session_config=session.config.as_dict(), + runtime_config=session.runtime_config, + ) + self.stage_port.ensure_request(context) + self.register_request(request_id, session.session_id) + return context + + # ------------------------------------------------------------------ # + # Control operations # + # ------------------------------------------------------------------ # + + def _admission_count(self) -> int: + return len(set(self.runners) | set(self._closing) | self._admitting) + + async def open(self, message: OpenDuplexSessionMessage) -> None: + from vllm_omni.engine.duplex.session.runner import DuplexSessionRunner + + session_id = message.session_id + session: DuplexEngineSession | None = None + runner: DuplexSessionRunner | None = None + # Only the open that put the id in ``_admitting`` may take it out again: + # a duplicate open must not release the slot the first one is holding. + holds_admission_slot = False + try: + if session_id in self.runners or session_id in self._closing or session_id in self._admitting: + raise DuplexSessionError(f"Duplex session already exists: {session_id}", code="session_exists") + max_sessions = int(self.runtime_config.max_sessions) + if self._admission_count() >= max_sessions: + raise DuplexSessionError( + f"duplex_session_capacity_exhausted: limit={max_sessions}", + code="resource_exhausted", + retryable=True, + ) + # Hold the slot across the plugin await: opens of different + # sessions run concurrently and must not all pass the check above. + self._admitting.add(session_id) + holds_admission_slot = True + config = message.session_config + self.plugin.validate_client_extra_body(config.extra_body) + try: + runtime_config = await self.plugin.prepare_runtime_config(config, model_config=self.model_config) + except DuplexRuntimeConfigError: + raise + except ValueError as exc: + raise DuplexRuntimeConfigError(str(exc)) from exc + capabilities = self.plugin.capabilities(max_sessions=max_sessions) + session = DuplexEngineSession( + session_id=session_id, + config=config, + capabilities=capabilities, + lease=DuplexLeaseState(config=self._lease_config, generation=0, last_activity=self._clock()), + _clock=self._clock, + _runtime_config=dict(runtime_config), + ) + # Validates the plugin's sampling policy for this runtime config before admission. + self.sampling_params_for(session) + session.model_state = self.plugin.create_session_state() + runner = DuplexSessionRunner( + session=session, + plugin=self.plugin, + stage_port=self.stage_port, + manager=self, + model_config=self.model_config, + ) + self.runners[session_id] = runner + # Reserve the Stage0 request resource atomically with admission. + self.ensure_stage_request(session, stage_id=0) + runner.start() + await self._put_result(message, operation="open", ok=True, session=session) + if session_id in self._abandoned_opens: + # The caller timed out while we were admitting; hand the slot + # and the stage resource straight back. + self._abandoned_opens.discard(session_id) + await self._close_runner(runner, kind="close", reason="open_abandoned") + except Exception as exc: + error_code, _, _ = self._control_error(exc) + if error_code in {"resource_exhausted", "session_exists"}: + logger.info("open_duplex_session rejected: %s", exc) + else: + logger.exception("open_duplex_session failed: %s", exc) + if runner is not None and self.runners.get(session_id) is runner: + self.runners.pop(session_id, None) + try: + await runner.shutdown() + except Exception: + logger.exception("duplex open rollback: runner shutdown failed for %s", session_id) + if session is not None: + reserved = session.release_all_requests() + if reserved: + try: + await self.stage_port.cleanup(list(reserved)) + except Exception: + logger.warning("duplex open rollback: request cleanup pending for %s", session_id) + self._unregister_session_requests(session_id) + await self._put_result(message, operation="open", ok=False, error=exc) + finally: + if holds_admission_slot: + self._admitting.discard(session_id) + self._abandoned_opens.discard(session_id) + + def _require_runner(self, session_id: str) -> DuplexSessionRunner: + runner = self.runners.get(session_id) + if runner is None: + raise DuplexSessionError(f"Unknown duplex session: {session_id}", code="unknown_session") + if runner.closing: + raise DuplexSessionError(f"Duplex session is closing: {session_id}", code="session_closed") + return runner + + async def close(self, message: CloseDuplexSessionMessage) -> None: + session_id = message.session_id + runner = self.runners.get(session_id) + if runner is None: + if session_id in self._admitting: + # The open is still awaiting the plugin. It cannot be cancelled + # from here, so leave a note it checks before returning. + self._abandoned_opens.add(session_id) + # Already closing, or closed and forgotten: ids are never reused + # (D16), so a repeated close is idempotent rather than an error. + await self._put_result(message, operation="close", ok=True) + return + session = runner.session + try: + await self._close_runner(runner, kind="close", reason=message.reason) + except Exception as exc: + # The session is closed either way; only the stage cleanup is + # outstanding and the reaper retries it while the slot stays held. + logger.warning("close_duplex_session: stage cleanup for %s remains pending: %s", session_id, exc) + await self._put_result(message, operation="close", ok=True, session=session) + + def close_from_runner(self, runner: DuplexSessionRunner, reason: str) -> None: + """Finish a close the runner started from its command stream (``session.close``). + + Runs on the orchestrator loop as a tracked task (the runner's worker + cannot await its own teardown): stage requests are aborted, the + admission slot is released and ``session.closed`` is emitted last. + """ + session_id = runner.session.session_id + if self.runners.get(session_id) is not runner: + return + + async def close_runner() -> None: + if self.runners.get(session_id) is runner: + await self._close_runner(runner, kind="close", reason=reason) + + # Chained on the session's control tail so a control RPC issued after + # the wire close observes the closed session, never a half-torn one. + self._run_control(session_id, f"duplex-close-{session_id}", close_runner) + + async def _close_runner(self, runner: DuplexSessionRunner, *, kind: str, reason: str) -> None: + """Begin an irreversible close, tear down the runner, then release stage resources. + + The session keeps its admission slot (``_closing``) until stage cleanup + succeeded; a failed cleanup is retried by the reaper. For an explicit + close the ``session.closed`` event is emitted only after the cleanup + attempt, so a client that sees it can open a replacement session at + once (the runner used to emit it before the stage requests were + aborted, which let a prompt reopen hit ``resource_exhausted``). + """ + session = runner.session + submitted = tuple(session.resource_request_ids(submitted=True)) + reserved = tuple(session.resource_request_ids(submitted=False)) + session.begin_close(reason=reason) + pending = _PendingSessionCleanup( + kind=kind, + session_id=session.session_id, + lease_generation=session.lease_generation, + reason=reason, + submitted_request_ids=submitted, + reserved_request_ids=reserved, + ) + self._closing[session.session_id] = pending + if self.runners.get(session.session_id) is runner: + self.runners.pop(session.session_id, None) + emit_after_cleanup = kind in {"close", "expired"} and not runner.closed_emitted + try: + if kind == "expired": + await runner.expire(reason, emit_expired=not emit_after_cleanup) + else: + await runner.close(reason, emit_closed=not emit_after_cleanup) + except Exception: + logger.exception("duplex runner %s failed for session %s", kind, session.session_id) + session.close() + try: + await self._finalize_pending_cleanup(pending, session) + finally: + # ``emit_after_cleanup`` was decided before the teardown awaits; a + # runtime close racing them may have emitted the terminal already. + if emit_after_cleanup and not runner.closed_emitted: + runner.mark_closed_emitted() + terminal: DuplexEvent = ( + SessionExpired(reason=reason) if kind == "expired" else SessionClosed(reason=reason) + ) + self.emit(session, terminal) + + async def _finalize_pending_cleanup( + self, + pending: _PendingSessionCleanup, + session: DuplexEngineSession | None, + ) -> None: + if pending.submitted_request_ids: + await self.stage_port.cleanup(list(pending.submitted_request_ids), abort=True) + if pending.reserved_request_ids: + await self.stage_port.cleanup(list(pending.reserved_request_ids)) + if session is not None: + session.release_all_requests() + self._unregister_session_requests(pending.session_id) + if self._closing.get(pending.session_id) is pending: + self._closing.pop(pending.session_id, None) + self._session_snapshots.pop(pending.session_id, None) + + async def resume(self, message: ResumeDuplexSessionMessage) -> None: + session: DuplexEngineSession | None = None + try: + runner = self._require_runner(message.session_id) + session = runner.session + try: + session.resume_lease(expected_lease_generation=message.expected_lease_generation) + except ValueError as exc: + raise DuplexSessionError(str(exc), code="session_resume_conflict") from exc + await self._put_result(message, operation="resume", ok=True, session=session) + except Exception as exc: + logger.exception("resume_duplex_session failed: %s", exc) + await self._put_result(message, operation="resume", ok=False, session=session, error=exc) + + async def touch(self, message: TouchDuplexSessionMessage) -> None: + session: DuplexEngineSession | None = None + try: + runner = self._require_runner(message.session_id) + session = runner.session + activity = DuplexLeaseActivity(message.activity) + if activity is DuplexLeaseActivity.DETACH: + session.detach_lease() + else: + session.touch_lease(activity) + await self._put_result(message, operation="touch", ok=True, session=session) + except DuplexSessionError as exc: + # A lease touch racing the session's own close is routine: the + # caller learns the session is gone from the result, not a fault. + logger.debug("touch_duplex_session skipped: %s", exc) + await self._put_result(message, operation="touch", ok=False, session=session, error=exc) + except Exception as exc: + logger.exception("touch_duplex_session failed: %s", exc) + await self._put_result(message, operation="touch", ok=False, session=session, error=exc) + + # ------------------------------------------------------------------ # + # Lease expiry / reaper # + # ------------------------------------------------------------------ # + + async def reaper_loop(self, shutdown_event: asyncio.Event) -> None: + interval = float(self.runtime_config.reaper_interval_s) + while not shutdown_event.is_set(): + try: + await asyncio.wait_for(shutdown_event.wait(), timeout=interval) + except TimeoutError: + try: + await self.reap_expired() + except Exception: + logger.exception("[DuplexSessionManager] expiry cleanup failed; retrying on next tick") + + async def reap_expired(self, now: float | None = None) -> int: + completed = 0 + effective_now = self._clock() if now is None else now + # Retry request cleanups recorded by the orchestrator error paths. + for key, request_cleanup in list(self._pending_request_cleanups.items()): + if key in self._request_cleanups_in_progress: + continue + try: + await self._complete_request_cleanup(key, request_cleanup) + except Exception as exc: + logger.warning( + "duplex request cleanup remains pending for session %s: %s", + request_cleanup.session_id, + exc, + ) + continue + completed += 1 + # Retry closes/expiries whose stage cleanup failed. + for session_id, session_cleanup in list(self._closing.items()): + if session_id in self.runners: + continue + if session_cleanup.kind == "request_cleanup" and any( + key[0] == session_id for key in self._pending_request_cleanups + ): + # The stage cleanup is owned by the request-cleanup path (orchestrator + # or the retry above); ``finalize_closed_sessions`` releases the slot. + continue + try: + await self._finalize_pending_cleanup(session_cleanup, self._session_snapshots.get(session_id)) + except Exception as exc: + logger.warning( + "duplex %s cleanup remains pending for session %s: %s", + session_cleanup.kind, + session_id, + exc, + ) + continue + completed += 1 + # Expire leases. + for session_id, runner in list(self.runners.items()): + lease = runner.session.lease + if lease.disconnect_grace_expired(effective_now): + reason = "disconnect_grace_expired" + elif lease.idle_expired(effective_now): + reason = "idle_ttl_expired" + else: + continue + self._session_snapshots[session_id] = runner.session + try: + await self._close_runner(runner, kind="expired", reason=reason) + except Exception as exc: + logger.warning("duplex expiry cleanup remains pending for session %s: %s", session_id, exc) + continue + completed += 1 + return completed + + # ------------------------------------------------------------------ # + # Request-triggered cleanup (orchestrator error paths) # + # ------------------------------------------------------------------ # + + def close_sessions_for_request_ids( + self, + request_ids: list[str], + *, + abort: bool = False, + cleanup_in_progress: bool = False, + ) -> dict[str, list[str]]: + request_id_set = set(request_ids) + closed: dict[str, list[str]] = {} + for session_id, runner in list(self.runners.items()): + session = runner.session + stale = session.resource_request_ids() + if request_id_set.isdisjoint(stale): + continue + if not session.begin_close(reason="request_cleanup"): + continue + closed[session_id] = stale + key = (session_id, session.lease_generation) + existing = self._pending_request_cleanups.get(key) + merged = tuple(dict.fromkeys([*(existing.request_ids if existing is not None else ()), *stale])) + self._pending_request_cleanups[key] = _PendingRequestCleanup( + session_id=session_id, + lease_generation=session.lease_generation, + request_ids=merged, + abort=abort or (existing.abort if existing is not None else False), + ) + if cleanup_in_progress: + self._request_cleanups_in_progress.add(key) + # Tear the runner down in the background; the orchestrator owns the stage cleanup. + self._session_snapshots[session_id] = session + self._closing[session_id] = _PendingSessionCleanup( + kind="request_cleanup", + session_id=session_id, + lease_generation=session.lease_generation, + reason="request_cleanup", + submitted_request_ids=(), + reserved_request_ids=(), + ) + self.runners.pop(session_id, None) + task = asyncio.create_task(runner.expire("request_cleanup"), name=f"duplex-request-cleanup-{session_id}") + self._track_task(task) + return closed + + def defer_request_cleanups(self, session_ids: Iterable[str]) -> None: + session_id_set = set(session_ids) + active_keys = {key for key in self._request_cleanups_in_progress if key[0] in session_id_set} + self._request_cleanups_in_progress.difference_update(active_keys) + + def finalize_closed_sessions(self, session_ids: Iterable[str]) -> None: + session_id_set = set(session_ids) + for key in list(self._pending_request_cleanups): + if key[0] not in session_id_set: + continue + self._pending_request_cleanups.pop(key, None) + self._request_cleanups_in_progress.discard(key) + for session_id in session_id_set: + session = self._session_snapshots.pop(session_id, None) + if session is not None and session.lease.terminal_reason is not None: + session.release_all_requests() + session.close() + self._closing.pop(session_id, None) + self._unregister_session_requests(session_id) + + async def _complete_request_cleanup( + self, + key: tuple[str, int], + pending: _PendingRequestCleanup, + ) -> None: + task = self._request_cleanup_tasks.get(key) + if task is not None and task.done(): + self._request_cleanup_tasks.pop(key, None) + task = None + if task is None: + task = asyncio.create_task( + self._run_request_cleanup(key, pending), + name=f"duplex-request-cleanup-{pending.session_id}", + ) + self._request_cleanup_tasks[key] = task + + def discard(completed: asyncio.Task[None]) -> None: + if self._request_cleanup_tasks.get(key) is completed: + self._request_cleanup_tasks.pop(key, None) + + task.add_done_callback(discard) + await asyncio.shield(task) + + async def _run_request_cleanup(self, key: tuple[str, int], pending: _PendingRequestCleanup) -> None: + await self.stage_port.cleanup(list(pending.request_ids), abort=pending.abort) + if self._pending_request_cleanups.get(key) is pending: + self._pending_request_cleanups.pop(key, None) + self._request_cleanups_in_progress.discard(key) + self.finalize_closed_sessions([pending.session_id]) + + # ------------------------------------------------------------------ # + # Shutdown # + # ------------------------------------------------------------------ # + + async def shutdown(self) -> None: + runners = list(self.runners.values()) + self.runners.clear() + for runner in runners: + try: + await runner.shutdown() + except Exception: + logger.exception("duplex runner shutdown failed for %s", runner.session.session_id) + tasks = tuple(self._dispatched_control_tasks | set(self._request_cleanup_tasks.values())) + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + self._session_control_tails.clear() + self._closing.clear() + self._session_snapshots.clear() + self._request_index.clear() + if self._owns_executor: + self.executor.shutdown(wait=False, cancel_futures=True) + + +__all__ = ["DuplexSessionManager"] diff --git a/vllm_omni/engine/duplex/session/model_channel.py b/vllm_omni/engine/duplex/session/model_channel.py new file mode 100644 index 000000000..322db7ede --- /dev/null +++ b/vllm_omni/engine/duplex/session/model_channel.py @@ -0,0 +1,1100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The model side of one duplex session. + +Everything between the session and the model runtime: planning an append with +the plugin and submitting it to the stage port, turning the stage output that +comes back into session events, and -- when the model chooses to keep listening +-- offering it another unit of silence so the turn can continue. + +Those three read as separate concerns but they call each other in a cycle. An +append submits and immediately projects its own first outputs; a projected +output can schedule a continuation; a continuation is another append. Splitting +them into three modules would mean three mutual imports and no fewer edges, so +they are one component whose seam with the runner is narrow instead. + +That seam is three callbacks, for the three things this component triggers but +does not own: closing the session, scheduling the continuation append as a +tracked task, and aborting a stage request in the background. +""" + +from __future__ import annotations + +import asyncio +import uuid +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import replace +from typing import TYPE_CHECKING, Protocol + +import pybase64 as base64 +from vllm.logger import init_logger + +from vllm_omni.engine.duplex.config import DuplexSessionState +from vllm_omni.engine.duplex.contracts import ( + DuplexAppendPlan, + DuplexFence, + DuplexOutputContext, + DuplexOutputDecision, + DuplexStageSubmission, + duplex_data_plane_request_info, +) +from vllm_omni.engine.duplex.plugin import ( + DuplexRuntimeConfigError, + coerce_int, + payload_turn_id, +) +from vllm_omni.engine.duplex.session.context import DuplexSessionContext, StageOutput +from vllm_omni.engine.duplex.session.emitter import SessionEmitter +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession, DuplexFenceMismatchError +from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity +from vllm_omni.metrics.stats import OrchestratorAggregator, StageRequestStats +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.outputs.duplex import attach_duplex_output_decision + +if TYPE_CHECKING: + from vllm.outputs import RequestOutput + +logger = init_logger(__name__) + + +class SilenceContinuationScheduler(Protocol): + """Runs one continuation append as a tracked task (see ``_schedule_silence_continuation``).""" + + async def __call__( + self, + payload: dict[str, object], + *, + request_id: str, + owner_id: str, + response_id: str | None, + response_owned: bool, + expected_epoch: int | None, + expected_model_turn_id: int | None, + ) -> bool: ... + + +class ModelChannel: + """Appends out to the model, events back from it, for one session.""" + + # One MiniCPM model unit (1 s at 16 kHz) is the compatibility default. + _SILENCE_UNIT_PAYLOAD_AUDIO = base64.b64encode(bytes(16000 * 4)).decode("ascii") + _RESPONSE_MAX_CONTINUATION_UNITS = 8 + _AUTO_RESPONSE_MAX_CONTINUATION_UNITS = 64 + + def __init__( + self, + ctx: DuplexSessionContext, + out: SessionEmitter, + *, + close_from_runtime: Callable[[str], Awaitable[None]], + schedule_silence_continuation: SilenceContinuationScheduler, + abort_request: Callable[..., Awaitable[None]], + ) -> None: + self._ctx = ctx + self._out = out + self._close_from_runtime = close_from_runtime + self._schedule_silence_continuation = schedule_silence_continuation + self._abort_request = abort_request + + @staticmethod + def should_commit_response_to_history(session: DuplexEngineSession, response_id: str | None) -> bool: + if response_id is not None and response_id != session.active_response_id: + return True + mode = session.response_config.extra_body.get("realtime_response_conversation") + return not isinstance(mode, str) or mode.strip().lower() != "none" + + def response_created_payload( + self, + response_id: str, + *, + epoch: int, + request_id: str | None = None, + ) -> dict[str, object]: + session = self._ctx.session + response_config = session.response_config + payload: dict[str, object] = { + "type": "response.created", + "session_id": session.session_id, + "response_id": response_id, + "epoch": epoch, + "modalities": list(response_config.modalities), + } + if request_id is not None: + payload["request_id"] = request_id + metadata = response_config.extra_body.get("realtime_response_metadata") + if not isinstance(metadata, dict): + metadata = response_config.extra_body.get("realtime_metadata") + if isinstance(metadata, dict): + payload["metadata"] = dict(metadata) + conversation = response_config.extra_body.get("realtime_response_conversation") + if isinstance(conversation, str): + payload["conversation"] = conversation + prompt = response_config.extra_body.get("realtime_response_prompt") + if isinstance(prompt, dict): + payload["prompt"] = dict(prompt) + return payload + + # ------------------------------------------------------------------ # + # Submitting an append # + # ------------------------------------------------------------------ # + + async def append_runtime_input( + self, + payload: object, + *, + operation_id: str | None = None, + final: bool, + expected_epoch: int | None = None, + ) -> tuple[bool, bool]: + session = self._ctx.session + if not session.capabilities.supports_input_append: + return True, False + if expected_epoch is not None and session.epoch != expected_epoch: + return True, False + try: + result = await self._append_via_data_plane( + payload, + final=final, + operation_id=operation_id, + expected_epoch=expected_epoch, + ) + except asyncio.CancelledError: + raise + except Exception as exc: + logger.exception("Failed to append duplex runtime input: %s", exc) + self.send_runtime_error("runtime_append_failed", exc) + return False, False + if result is None: + return True, False + if expected_epoch is not None and session.epoch != expected_epoch: + return True, False + request_id, _ = duplex_data_plane_request_info(result) + if request_id is not None: + self._ctx.plugin.data_plane.begin_request(request_id) + close_reason, emitted_response = await self._send_model_output_events(result, expected_epoch=expected_epoch) + emitted_response = emitted_response or request_id is not None + if close_reason is None and self._start_data_plane_stream(result): + emitted_response = True + if close_reason is not None: + await self._close_from_runtime(close_reason) + return False, emitted_response + return True, emitted_response + + async def _append_via_data_plane( + self, + payload: object, + *, + final: bool, + operation_id: str | None, + expected_epoch: int | None, + ) -> dict[str, object] | None: + """Submit one planned append to the stage port and bind the resulting stage request.""" + session = self._ctx.session + if self._ctx.stage_port.stage_count == 0: + raise RuntimeError("duplex_data_plane_has_no_stage") + payload_turn = payload_turn_id(payload) + fence = DuplexFence( + session.session_id, + epoch=session.epoch, + turn_id=( + payload_turn + if payload_turn is not None + else ( + session.active_response_turn_id if session.active_response_turn_id is not None else session.turn_id + ) + ), + ) + lease_operation_id = f"append:{operation_id or uuid.uuid4().hex}" + operation_started = False + stage_id = 0 + request_id = self._ctx.manager.stage_request_id(fence, stage_id=stage_id) + try: + session.begin_lease_operation(fence, lease_operation_id) + operation_started = True + reservation = session.prepare_append(fence) + already_submitted = session.stage_request_submitted(stage_id, request_id) + request_context = self._ctx.manager.ensure_stage_request(session, stage_id=stage_id, fence=fence) + if request_context is None: + raise RuntimeError("duplex_data_plane_has_no_stage") + append_plan = self._ctx.plugin.plan_append( + request_id=request_id, + fence=fence, + session_config=dict(request_context.session_config), + runtime_config=dict(request_context.runtime_config), + seq=reservation.update.seq, + turn_seq=reservation.update.turn_seq, + payload=payload, + final=final, + sampling_params=request_context.stage_sampling_params, + ) + if not isinstance(append_plan, DuplexAppendPlan): + raise TypeError("duplex plugin plan_append() must return DuplexAppendPlan") + submission = DuplexStageSubmission( + context=request_context, + prompt=append_plan.prompt, + already_submitted=already_submitted, + ) + submission_result = await self._ctx.stage_port.submit(submission) + try: + if submission_result.request_id != request_id or submission_result.stage_id != stage_id: + raise RuntimeError("duplex stage adapter returned a mismatched submission result") + if expected_epoch is not None and session.epoch != expected_epoch: + raise DuplexFenceMismatchError(session.fence, fence) + update = session.commit_append(reservation) + session.bind_stage_request(stage_id, request_id, fence=fence) + except BaseException: + try: + await self._ctx.stage_port.cleanup([request_id]) + except Exception as cleanup_exc: + logger.warning( + "duplex append compensation remains pending for session %s: %s", + session.session_id, + cleanup_exc, + ) + raise + session.touch_lease(DuplexLeaseActivity.APPEND) + return { + "ok": True, + "operation": "append", + "session_id": session.session_id, + "stage_results": [ + { + "stage_id": stage_id, + "replica_id": submission_result.replica_id, + "result": { + "supported": True, + "data_plane_append": True, + "request_id": request_id, + "response_stage_id": request_context.final_stage_id, + "seq": update.seq, + "turn_id": update.turn_id, + "turn_seq": update.turn_seq, + "resumable": True, + }, + } + ], + } + finally: + if operation_started and lease_operation_id in session.lease.active_operations: + try: + session.end_lease_operation(lease_operation_id) + except Exception: + session.lease.active_operations.discard(lease_operation_id) + + def _start_data_plane_stream(self, result: object) -> bool: + """Bind the resumable stage request as the session's active data-plane stream.""" + request_id, _ = duplex_data_plane_request_info(result) if isinstance(result, dict) else (None, None) + if request_id is None or self._data_plane_outputs_finished(result): + return False + session = self._ctx.session + session.bind_request(request_id) + if self._ctx.run.stream_request_id == request_id: + return False + self._ctx.run.stream_request_id = request_id + return True + + def cancel_data_plane_stream(self) -> bool: + had_stream = self._ctx.run.stream_request_id is not None + self._ctx.run.stream_request_id = None + return had_stream + + async def signal_cancel_fence(self, cancelled_fence: DuplexFence) -> bool: + """In-process equivalent of the old ``barge_in`` control signal.""" + session = self._ctx.session + try: + next_fence = session.sync_fence() + if next_fence.epoch > cancelled_fence.epoch: + stale_request_ids = session.cancel_fence(cancelled_fence, next_fence) + else: + stale_request_ids = [] + if stale_request_ids: + await self._ctx.stage_port.cleanup(list(stale_request_ids), abort=True) + session.touch_lease(DuplexLeaseActivity.SIGNAL) + except asyncio.CancelledError: + raise + except Exception as exc: + logger.exception("Failed to signal duplex runtime session: %s", exc) + self.send_runtime_error("runtime_signal_failed", exc) + return False + return True + + def send_runtime_error(self, code: str, exc: BaseException) -> None: + if isinstance(exc, DuplexRuntimeConfigError): + code = exc.code + self._out.emit_error(code, str(exc), retryable=False) + + @staticmethod + def _data_plane_outputs_finished(result: object) -> bool: + if not isinstance(result, dict): + return False + outputs = result.get("data_plane_outputs") + if not isinstance(outputs, list) or not outputs: + return False + return bool(getattr(outputs[-1], "finished", False)) + + def _runtime_data_plane_context(self) -> object: + session = self._ctx.session + response_config = session.response_config + return self._ctx.plugin.data_plane_context( + epoch=session.epoch, + turn_id=session.turn_id, + active_response_turn_id=session.active_response_turn_id, + active_response_id=session.active_response_id, + auto_responds=self._out.auto_responds(), + response_format=response_config.response_format, + speed=response_config.speed, + modalities=tuple(response_config.modalities), + ) + + # ------------------------------------------------------------------ # + # Stage outputs # + # ------------------------------------------------------------------ # + + def decide_output( + self, stage_id: int, output: RequestOutput, context: DuplexOutputContext + ) -> DuplexOutputDecision | None: + """Pure plugin decision (no session mutation: this runs inline on the orchestrator loop).""" + decision = self._ctx.plugin.decide_output( + stage_id=stage_id, + final_stage_id=context.final_stage_id, + segment_finished=context.segment_finished, + segment_token_ids=context.segment_token_ids, + segment_output_metadata=dict(context.segment_output_metadata), + output=output, + ) + if decision is not None and not isinstance(decision, DuplexOutputDecision): + raise TypeError("duplex plugin decide_output() must return DuplexOutputDecision or None") + return decision + + @staticmethod + def stage_metrics_snapshot(stage_id: int, metrics: object, output: object) -> dict[str, dict[str, object]] | None: + if not isinstance(metrics, StageRequestStats): + return None + event = metrics + if event.stage_id is None: + event = replace(event, stage_id=stage_id) + if event.final_output_type is None: + final_output_type = getattr(output, "final_output_type", None) + if isinstance(final_output_type, str): + event = replace(event, final_output_type=final_output_type) + try: + merged = OrchestratorAggregator._merge_stage_metric_event(None, event) + except Exception: + return None + return {str(stage_id): merged} + + def _build_stage_output(self, item: StageOutput) -> OmniRequestOutput: + output = item.output + finished = bool(getattr(output, "finished", False)) or item.context.segment_finished + if item.decision is not None: + engine_output = OmniRequestOutput.from_stage_output( + output, + request_id=item.request_id, + stage_id=item.stage_id, + final_output_type=item.decision.final_output_type, + ) + # Set after construction, not as a keyword: ``from_stage_output`` + # copies ``finished`` from the source last, and a resumable stage + # request never finishes -- only its segment does. The decision + # closes that segment, and the projector recognises a listen only + # on a finished output. The orchestrator queue used to re-assert + # ``finished`` on every direct-response message; the mailbox hands + # the raw stage output over, so it has to happen here. + engine_output.finished = True + engine_output = attach_duplex_output_decision(engine_output, item.decision) + elif isinstance(output, OmniRequestOutput): + engine_output = output + if finished: + engine_output.finished = True + else: + final_output_type = getattr(output, "final_output_type", None) + engine_output = OmniRequestOutput.from_stage_output( + output, + request_id=item.request_id, + finished=finished, + stage_id=item.stage_id, + final_output_type=final_output_type if isinstance(final_output_type, str) else "audio", + ) + snapshot = self.stage_metrics_snapshot(item.stage_id, item.metrics, output) + if snapshot is not None: + existing = engine_output.metrics if isinstance(engine_output.metrics, dict) else {} + existing_stage_metrics = existing.get("stage_metrics") + merged_stage_metrics = ( + {**existing_stage_metrics, **snapshot} if isinstance(existing_stage_metrics, dict) else snapshot + ) + engine_output.metrics = {**existing, "stage_metrics": merged_stage_metrics} + return engine_output + + async def on_stage_output_item(self, item: StageOutput) -> None: + session = self._ctx.session + if session.state == DuplexSessionState.CLOSED or self._ctx.run.closing: + return + expected_epoch = item.context.identity.fence.epoch + if session.epoch != expected_epoch: + return + if session.lease.terminal_reason is not None: + # Lease already closed/expired: the output belongs to a dead session. + return + session.touch_lease(DuplexLeaseActivity.MODEL_OUTPUT) + if self._ctx.plugin.data_plane.is_terminal(item.request_id): + return + if self._out.auto_responds(): + active_request_id = session.active_request_id + if active_request_id is not None and active_request_id != item.request_id: + return + engine_output = self._build_stage_output(item) + drain_result = {"data_plane_outputs": [engine_output]} + close_reason, emitted_response = await self._send_model_output_events( + drain_result, expected_epoch=expected_epoch + ) + if close_reason is not None: + await self._close_from_runtime(close_reason) + return + finished = self._data_plane_outputs_finished(drain_result) + if finished and emitted_response and not self._out.auto_responds(): + # A finished, emitted response releases the per-request projector + # cursor on its way out and offers the model another + # silence unit. + self._ctx.plugin.data_plane.close_stream(item.request_id) + if self._ctx.run.stream_request_id == item.request_id: + self._ctx.run.stream_request_id = None + self._ctx.services.spawn( + self.maybe_continue_response(expected_epoch=expected_epoch), + name="duplex-continue", + ) + + async def _send_model_output_events( + self, + result: object, + *, + expected_epoch: int | None = None, + ) -> tuple[str | None, bool]: + session = self._ctx.session + if expected_epoch is not None and session.epoch != expected_epoch: + return None, False + close_reason: str | None = None + emitted_response = False + request_id, _ = duplex_data_plane_request_info(result) if isinstance(result, dict) else (None, None) + if self._ctx.plugin.data_plane.is_terminal(request_id): + return None, False + if request_id is not None and session.active_request_id is None: + session.bind_request(request_id) + context = self._runtime_data_plane_context() + for model_result in self._ctx.plugin.data_plane.project(result, context=context): + close_reason_for_result, did_emit = await self._send_one_model_output_event( + model_result, + expected_epoch=expected_epoch, + ) + 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: + return None, emitted_response + return close_reason, emitted_response + + def _fail_response_from_model_error(self, model_result: dict[str, object]) -> None: + """Report a data-plane error and fail the response it interrupted.""" + session = self._ctx.session + response_id = session.active_response_id + self._out.emit_error( + str(model_result.get("error_code")), + str(model_result.get("error") or "Duplex native data-plane error"), + ) + if response_id is None: + return + session.end_response(commit_text=False) + self._out.emit( + { + "type": "response.done", + "session_id": session.session_id, + "response_id": response_id, + "epoch": session.epoch, + "committed": False, + "status": "failed", + "status_details": {"type": "failed", "reason": model_result.get("error_code")}, + "playback": session.playback.as_dict(), + } + ) + + def _complete_model_turn_without_output( + self, + model_result: dict[str, object], + *, + model_turn_id: int | None, + data_plane_request_id: object, + ) -> bool: + """The model ended its turn with nothing to say. + + Returns whether an event was emitted: an auto-response session reports + the decision as a listen, a turn-based one simply releases the request. + """ + session = self._ctx.session + data_plane = self._ctx.plugin.data_plane + auto_response = self._out.auto_responds() + if isinstance(data_plane_request_id, str): + if not auto_response and data_plane_request_id == session.active_request_id: + session.clear_request() + if not auto_response: + data_plane.mark_terminal(data_plane_request_id) + if model_turn_id is not None: + session.complete_model_turn(model_turn_id) + if not auto_response: + return False + self._ctx.model_state.clear_continuation() + payload = { + "type": "response.listen", + "session_id": session.session_id, + "epoch": session.epoch, + "reason": "model_turn_completed_without_output", + "model_listen": True, + } + self._attach_runtime_metadata(payload, model_result) + self._out.emit(payload) + return True + + async def _on_model_listen( + self, + model_result: dict[str, object], + *, + model_turn_id: int | None, + data_plane_request_id: object, + expected_epoch: int | None, + ) -> tuple[str | None, bool]: + """The model chose to keep listening rather than speak. + + Either it continues the current response with another unit, or the + response ends here: an auto-response session that has continuations + left schedules one, and anything else closes the turn out. + """ + session = self._ctx.session + model_state = self._ctx.model_state + data_plane = self._ctx.plugin.data_plane + auto_response = self._out.auto_responds() + close_reason: str | None = None + emitted_response = False + self._end_active_response_before_future_model_turn(model_turn_id=model_turn_id) + if ( + session.active_response_id is not None + and model_turn_id is not None + and not session.active_response_accepts_model_turn(model_turn_id) + ): + return close_reason, emitted_response + active_response_id = session.active_response_id + auto_continuations_remaining = active_response_id is None or self.response_continuations_remaining( + active_response_id + ) + non_terminal_auto_listen = ( + auto_response + and active_response_id is not None + and session.active_request_id is not None + and model_result.get("end_of_turn") is not True + and auto_continuations_remaining + ) + if non_terminal_auto_listen: + self._ctx.services.spawn( + self.maybe_continue_response(expected_epoch=expected_epoch), name="duplex-continue" + ) + return close_reason, emitted_response + if not auto_response and data_plane_request_id == session.active_request_id: + session.clear_request() + model_listen = model_result.get("model_listen") + if not isinstance(model_listen, bool): + model_listen = model_result.get("reason") in {None, "", "model_listen"} + response_id = session.active_response_id + if isinstance(data_plane_request_id, str) and not auto_response: + data_plane.mark_terminal(data_plane_request_id) + emitted_response = True + payload = { + "type": "response.listen", + "session_id": session.session_id, + "epoch": session.epoch, + "reason": model_result.get("reason") or "model_listen", + "model_listen": model_listen, + } + if response_id is not None: + payload["response_id"] = response_id + self._attach_runtime_metadata(payload, model_result) + self._out.emit(payload) + if model_result.get("abort_data_plane_request") is True and isinstance(data_plane_request_id, str): + await self._abort_request(data_plane_request_id, notify=False) + if response_id is not None: + if not auto_response and self.response_continuations_remaining(response_id): + self._ctx.services.spawn( + self.maybe_continue_response(expected_epoch=expected_epoch), name="duplex-continue" + ) + return close_reason, emitted_response + if auto_response: + model_state.clear_continuation() + if not auto_continuations_remaining: + completed_turn_id = model_turn_id + if completed_turn_id is None: + completed_turn_id = session.active_response_turn_id + if completed_turn_id is not None: + session.complete_model_turn(completed_turn_id) + session.end_response(commit_text=False, preserve_request=auto_response) + self._out.emit( + { + "type": "response.done", + "session_id": session.session_id, + "response_id": response_id, + "epoch": session.epoch, + "committed": False, + "playback": session.playback.as_dict(), + } + ) + return close_reason, emitted_response + + async def _send_one_model_output_event( + self, + model_result: dict[str, object], + *, + expected_epoch: int | None = None, + ) -> tuple[str | None, bool]: + session = self._ctx.session + 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: + 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): + return close_reason, emitted_response + auto_response = self._out.auto_responds() + active_request_matches = session.active_request_id == data_plane_request_id or ( + auto_response and session.active_request_id is None + ) + if isinstance(data_plane_request_id, str) and not active_request_matches: + return close_reason, emitted_response + if isinstance(model_result.get("error_code"), str): + self._fail_response_from_model_error(model_result) + return close_reason, True + if model_result.get("function_call") is True: + self._out.emit( + { + "type": "function_call.done", + "session_id": session.session_id, + "epoch": session.epoch, + "call_id": model_result.get("call_id"), + "name": model_result.get("name"), + "arguments": model_result.get("arguments", ""), + } + ) + return close_reason, True + if model_result.get("requires_stage_handoff") is True or model_result.get("requires_tts_stage") is True: + # Reserve the protocol response only when Stage1 emits text/audio. + return close_reason, False + is_listen = model_result.get("is_listen") + model_turn_id = coerce_int(model_result.get("model_turn_id")) + if model_result.get("is_buffering") is True or model_result.get("prefill_success") is False: + if model_result.get("data_plane_request_id") == session.active_request_id: + session.clear_request() + payload = { + "type": "response.listen", + "session_id": session.session_id, + "epoch": session.epoch, + "reason": model_result.get("reason") or "buffering", + "model_listen": False, + "buffering": True, + } + self._attach_runtime_metadata(payload, model_result) + self._out.emit(payload) + return close_reason, emitted_response + if is_listen is True: + return await self._on_model_listen( + model_result, + model_turn_id=model_turn_id, + data_plane_request_id=data_plane_request_id, + expected_epoch=expected_epoch, + ) + + text = model_result.get("text") + audio = model_result.get("audio_data", model_result.get("audio")) + end_of_turn = bool(model_result.get("end_of_turn", False)) + has_text = isinstance(text, str) and bool(text) + has_audio = isinstance(audio, str) and bool(audio) + if not has_text and not has_audio and not end_of_turn: + tts_segment_ended = ( + model_result.get("stage_role") == "tts" and model_result.get("abort_data_plane_request") is True + ) + if ( + tts_segment_ended + and auto_response + and (model_turn_id is not None or session.active_response_id is not None) + ): + self._ctx.services.spawn( + self.maybe_continue_response(expected_epoch=expected_epoch, expected_model_turn_id=model_turn_id), + name="duplex-continue", + ) + return close_reason, emitted_response + if end_of_turn and not has_text and not has_audio and session.active_response_id is None: + emitted_response = self._complete_model_turn_without_output( + model_result, + model_turn_id=model_turn_id, + data_plane_request_id=data_plane_request_id, + ) + return close_reason, emitted_response + if session.active_response_id is None and model_turn_id is not None and model_turn_id < session.turn_id: + # Late audio of a completed model turn must not reserve a second response. + return close_reason, emitted_response + self._end_active_response_before_future_model_turn(model_turn_id=model_turn_id) + if ( + session.active_response_id is not None + and model_turn_id is not None + and not session.active_response_accepts_model_turn(model_turn_id) + ): + return close_reason, emitted_response + emitted_response = True + response_created = False + response_id = session.active_response_id + if response_id is None: + response_id = session.begin_response(turn_id=model_turn_id) + response_created = True + self._out.emit(self.response_created_payload(response_id, epoch=session.epoch)) + response_stage_metrics = session.accumulate_response_stage_metrics( + model_result.get("stage_metrics") if isinstance(model_result.get("stage_metrics"), Mapping) else None + ) + if response_created: + speak_payload = { + "type": "response.speak", + "session_id": session.session_id, + "response_id": response_id, + "epoch": session.epoch, + "text": text if isinstance(text, str) else "", + "end_of_turn": end_of_turn, + "model_speak": True, + } + self._attach_runtime_metadata(speak_payload, model_result, stage_metrics=response_stage_metrics) + self._out.emit(speak_payload) + previous_sent_ms = session.playback.sent_ms + text_chars_before_append = len("".join(session.assistant_text_buffer)) + if isinstance(text, str): + session.append_assistant_text(text) + duration_ms = model_result.get("audio_duration_ms") + text_chars = len("".join(session.assistant_text_buffer)) + mark_duration_ms = None + mark_text_chars: int | None = text_chars + if model_result.get("audio_text_mark") is False: + mark_text_chars = None + if isinstance(duration_ms, int | float): + mark_duration_ms = int(duration_ms) + if model_result.get("audio_duration_is_cumulative") is not True: + mark_duration_ms += session.playback.sent_ms + audio_text_marks = model_result.get("audio_text_marks") + audio_text_marks = self._normalize_audio_text_marks( + audio_text_marks if isinstance(audio_text_marks, list) else None, + audio_offset_ms=( + 0 + if model_result.get("audio_text_marks_are_cumulative") is True + or model_result.get("audio_duration_is_cumulative") is True + else previous_sent_ms + ), + text_offset_chars=( + 0 if model_result.get("audio_text_marks_are_cumulative") is True else text_chars_before_append + ), + ) + session.mark_audio_sent( + mark_duration_ms, + text_chars=mark_text_chars if mark_duration_ms is not None else None, + audio_text_marks=audio_text_marks, + ) + payload = { + "type": "response.output_audio.delta", + "session_id": session.session_id, + "response_id": response_id, + "epoch": session.epoch, + "text": text if isinstance(text, str) else "", + "audio": audio if isinstance(audio, str) else "", + "format": ( + model_result.get("audio_format") + if isinstance(model_result.get("audio_format"), str) + else session.response_config.response_format + ), + "end_of_turn": end_of_turn, + "model_speak": True, + } + if mark_duration_ms is not None: + payload["audio_duration_ms"] = mark_duration_ms + if audio_text_marks: + payload["audio_text_marks"] = audio_text_marks + elif mark_duration_ms is not None and mark_text_chars is not None: + payload["audio_text_marks"] = [ + {"text_chars": max(0, int(mark_text_chars)), "audio_end_ms": max(0, int(mark_duration_ms))} + ] + payload["playback"] = session.playback.as_dict() + sample_rate_hz = model_result.get("sample_rate_hz") or model_result.get("audio_sample_rate_hz") + if isinstance(sample_rate_hz, int | float) and int(sample_rate_hz) > 0: + payload["sample_rate_hz"] = int(sample_rate_hz) + self._attach_runtime_metadata(payload, model_result, stage_metrics=response_stage_metrics) + self._out.emit(payload) + if ( + not end_of_turn + and model_result.get("stage_role") == "tts" + and model_result.get("abort_data_plane_request") is True + and auto_response + ): + self._ctx.services.spawn( + self.maybe_continue_response(expected_epoch=expected_epoch), name="duplex-continue" + ) + if end_of_turn: + data_plane_request_id = model_result.get("data_plane_request_id") + if isinstance(data_plane_request_id, str) and not auto_response: + data_plane.close_stream(data_plane_request_id) + if isinstance(data_plane_request_id, str): + if not auto_response and data_plane_request_id == session.active_request_id: + session.clear_request() + if not auto_response: + data_plane.mark_terminal(data_plane_request_id) + should_commit = self.should_commit_response_to_history(session, response_id) + committed_message = session.end_response(commit_text=should_commit, preserve_request=auto_response) + model_turn_id = coerce_int(model_result.get("model_turn_id")) + if model_turn_id is not None: + session.complete_model_turn(model_turn_id) + if should_commit: + session.register_history_item(f"item_{response_id}", committed_message) + self._out.emit( + { + "type": "response.done", + "session_id": session.session_id, + "response_id": response_id, + "epoch": session.epoch, + "committed": committed_message is not None, + "playback": session.playback.as_dict(), + } + ) + return close_reason, emitted_response + + def _end_active_response_before_future_model_turn(self, *, model_turn_id: int | None) -> None: + session = self._ctx.session + if not self._out.auto_responds(): + return + response_id = session.active_response_id + active_turn_id = session.active_response_turn_id + if response_id is None or model_turn_id is None or active_turn_id is None: + return + if int(model_turn_id) <= int(active_turn_id): + return + session.complete_model_turn(int(model_turn_id) - 1) + should_commit = self.should_commit_response_to_history(session, response_id) + committed_message = session.end_response(commit_text=should_commit, preserve_request=True) + if should_commit: + session.register_history_item(f"item_{response_id}", committed_message) + self._out.emit( + { + "type": "response.done", + "session_id": session.session_id, + "response_id": response_id, + "epoch": session.epoch, + "committed": committed_message is not None, + "playback": session.playback.as_dict(), + } + ) + + @staticmethod + def _normalize_audio_text_marks( + audio_text_marks: list[object] | None, + *, + audio_offset_ms: int, + text_offset_chars: int, + ) -> list[dict[str, int]] | None: + if not audio_text_marks: + return None + normalized: list[dict[str, int]] = [] + for raw_mark in audio_text_marks: + if not isinstance(raw_mark, dict): + continue + raw_text_chars = raw_mark.get("text_chars") + raw_audio_end_ms = raw_mark.get("audio_end_ms", raw_mark.get("audio_ms")) + if not isinstance(raw_text_chars, int | float) or not isinstance(raw_audio_end_ms, int | float): + continue + normalized.append( + { + "text_chars": max(0, int(raw_text_chars) + int(text_offset_chars)), + "audio_end_ms": max(0, int(raw_audio_end_ms) + int(audio_offset_ms)), + } + ) + return normalized or None + + @staticmethod + def _attach_runtime_metadata( + payload: dict[str, object], + model_result: dict[str, object], + *, + stage_metrics: Mapping[str, object] | None = None, + ) -> None: + metadata: dict[str, object] = {} + runtime_impl = model_result.get("runtime_impl") + if isinstance(runtime_impl, str) and runtime_impl: + metadata["runtime_impl"] = runtime_impl + owned_runtime = model_result.get("owned_runtime") + if isinstance(owned_runtime, bool): + metadata["owned_runtime"] = owned_runtime + model_turn_id = coerce_int(model_result.get("model_turn_id")) + if model_turn_id is not None: + metadata["model_turn_id"] = model_turn_id + for name in ("uses_model_runner_scheduler", "runner_kv_backed"): + value = model_result.get(name) + if isinstance(value, bool): + metadata[name] = value + effective_stage_metrics = stage_metrics if stage_metrics is not None else model_result.get("stage_metrics") + if isinstance(effective_stage_metrics, Mapping): + metadata["stage_metrics"] = { + str(stage_id): dict(values) + for stage_id, values in effective_stage_metrics.items() + if isinstance(values, Mapping) + } + if metadata: + payload["vllm_omni"] = metadata + + # ------------------------------------------------------------------ # + # Silence continuation # + # ------------------------------------------------------------------ # + + def silence_unit_payload(self) -> dict[str, object]: + samples = int(self._ctx.plugin.silence_continuation_samples) + audio = ( + self._SILENCE_UNIT_PAYLOAD_AUDIO + if samples == 16000 + else base64.b64encode(bytes(samples * 4)).decode("ascii") + ) + return {"type": "audio", "audio": audio, "format": "pcm_f32le", "sample_rate_hz": 16000} + + def response_continuations_remaining(self, response_id: str) -> bool: + model_state = self._ctx.model_state + owner_id = f"response:{response_id}" + count = model_state.continuation_units if model_state.continuation_owner_id == owner_id else 0 + limit = ( + self._AUTO_RESPONSE_MAX_CONTINUATION_UNITS + if self._out.auto_responds() + else self._RESPONSE_MAX_CONTINUATION_UNITS + ) + return count < limit + + async def _finish_bounded_auto_response( + self, + *, + expected_epoch: int | None, + model_turn_id: int | None = None, + ) -> None: + session = self._ctx.session + model_state = self._ctx.model_state + 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: + return + model_state.clear_continuation() + payload: dict[str, object] = { + "type": "response.listen", + "session_id": session.session_id, + "epoch": response_epoch, + "reason": "continuation_limit", + "model_listen": False, + } + if response_id is not None: + payload["response_id"] = response_id + self._out.emit(payload) + if response_id is None: + if session.epoch == response_epoch and response_turn_id is not None: + session.complete_model_turn(response_turn_id) + return + if session.epoch != response_epoch or session.active_response_id != response_id: + return + if response_turn_id is not None: + session.complete_model_turn(response_turn_id) + session.end_response(commit_text=False, preserve_request=True) + self._out.emit( + { + "type": "response.done", + "session_id": session.session_id, + "response_id": response_id, + "epoch": response_epoch, + "committed": False, + "playback": session.playback.as_dict(), + } + ) + + def silence_continuation_is_stale( + self, + *, + request_id: str, + response_id: str | None, + response_owned: bool, + expected_epoch: int | None, + expected_model_turn_id: int | None, + ) -> bool: + session = self._ctx.session + stale_common_owner = ( + session.state == DuplexSessionState.CLOSED + or session.active_request_id != request_id + or (expected_epoch is not None and session.epoch != expected_epoch) + ) + stale_response_owner = response_owned and session.active_response_id != response_id + stale_model_turn_owner = not response_owned and ( + session.active_response_id is not None or session.turn_id != expected_model_turn_id + ) + return stale_common_owner or stale_response_owner or stale_model_turn_owner + + async def maybe_continue_response( + self, + *, + expected_epoch: int | None, + expected_model_turn_id: int | None = None, + ) -> None: + session = self._ctx.session + model_state = self._ctx.model_state + response_id = session.active_response_id + if session.state == DuplexSessionState.CLOSED or self._ctx.run.closing: + model_state.clear_continuation() + return + request_id = session.active_request_id + if request_id is None: + model_state.clear_continuation() + return + if expected_epoch is not None and session.epoch != expected_epoch: + return + auto_response = self._out.auto_responds() + response_owned = response_id is not None + if response_owned: + owner_id = f"response:{response_id}" + payload_turn_id_value = ( + session.active_response_turn_id if session.active_response_turn_id is not None else session.turn_id + ) + else: + if not auto_response or expected_model_turn_id is None or session.turn_id != expected_model_turn_id: + model_state.clear_continuation() + return + owner_id = f"model-turn:{expected_model_turn_id}" + payload_turn_id_value = expected_model_turn_id + count = model_state.continuation_units if model_state.continuation_owner_id == owner_id else 0 + continuation_limit = ( + self._AUTO_RESPONSE_MAX_CONTINUATION_UNITS if auto_response else self._RESPONSE_MAX_CONTINUATION_UNITS + ) + if count >= continuation_limit: + if auto_response: + await self._finish_bounded_auto_response( + expected_epoch=expected_epoch, + model_turn_id=payload_turn_id_value, + ) + return + payload = self.silence_unit_payload() + if auto_response and count + 1 == continuation_limit: + payload["force_listen"] = True + payload["duplex_turn_id"] = payload_turn_id_value + try: + scheduled = await self._schedule_silence_continuation( + payload, + request_id=request_id, + owner_id=owner_id, + response_id=response_id, + response_owned=response_owned, + expected_epoch=expected_epoch, + expected_model_turn_id=expected_model_turn_id, + ) + except asyncio.CancelledError: + raise + except Exception as exc: + logger.exception("Failed to schedule duplex native response continuation: %s", exc) + scheduled = False + if scheduled: + model_state.continuation_owner_id = owner_id + model_state.continuation_units = count + 1 diff --git a/vllm_omni/engine/duplex/session/overlap_policy.py b/vllm_omni/engine/duplex/session/overlap_policy.py new file mode 100644 index 000000000..0aad04368 --- /dev/null +++ b/vllm_omni/engine/duplex/session/overlap_policy.py @@ -0,0 +1,388 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""How one duplex session reacts to input that arrives while it is speaking. + +Short acknowledgement, barge-in, or keep listening: the rules read the session +config and the incoming append, and nothing else. Extracted from +``DuplexSessionRunner`` because they were a decision function wearing a +method's clothes -- the only runner state they ever touched was ``session`` -- +and because reaching them previously meant driving a whole append through the +runner. +""" + +from __future__ import annotations + +import binascii +from collections.abc import Mapping + +import numpy as np +import pybase64 as base64 + +from vllm_omni.engine.duplex.config import DuplexOverlapPolicy +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession + + +def decide( + session: DuplexEngineSession, + event: dict[str, object], + payload: dict[str, object], + *, + auto_responds: bool, +) -> dict[str, object]: + duration_ms = input_audio_duration_ms(event, payload) + is_speech = input_looks_like_speech(session, event, payload) + if not session.capabilities.supports_barge_in and event_requests_barge_in(event): + return defer_unsupported_barge_in(session, duration_ms=duration_ms, is_speech=is_speech) + explicit = event.get("overlap_action") or event.get("overlap") + if isinstance(explicit, str): + normalized = explicit.strip().lower() + if normalized in {"barge_in", "interrupt", "cancel"}: + return { + "action": "barge_in", + "reason": "client_overlap_action", + "duration_ms": duration_ms, + "buffer_audio": True, + } + if normalized in {"listen", "continue", "continue_output", "ack"}: + session.reset_overlap_speech() + return { + "action": "listen", + "reason": "client_overlap_action", + "duration_ms": duration_ms, + "buffer_audio": ( + normalized == "listen" and is_speech and duration_ms > session.config.overlap_short_ack_ms + ), + "defer_runtime_append": True, + } + if normalized in {"drop", "ignore", "silence"}: + session.reset_overlap_speech() + return { + "action": "drop", + "reason": "client_overlap_action", + "duration_ms": duration_ms, + "buffer_audio": False, + } + + if bool(event.get("force_barge_in", False)): + return { + "action": "barge_in", + "reason": "client_force_barge_in", + "duration_ms": duration_ms, + "buffer_audio": True, + } + if auto_responds: + if is_speech: + session.accumulate_overlap_speech(duration_ms) + speech_started = vad_speech_started(event, payload) + if ( + session.capabilities.supports_barge_in + and is_speech + and session.config.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value + and speech_started is not False + ): + return { + "action": "barge_in", + "reason": ("server_vad_speech_started" if speech_started is True else "barge_in_on_speech"), + "cancel_reason": "turn_detected" if speech_started is True else "barge_in", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + } + return { + "action": "listen", + "reason": "auto_response_continuous", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + "defer_runtime_append": False, + "force_listen": event.get("force_listen") is True or payload.get("force_listen") is True, + "preserve_realtime_input": True, + } + if bool(event.get("force_listen", False)): + session.reset_overlap_speech() + return { + "action": "listen", + "reason": "client_force_listen", + "duration_ms": duration_ms, + "buffer_audio": is_speech, + "defer_runtime_append": True, + } + + policy = session.config.overlap_policy + if not is_speech: + if session.overlap_speech_ms <= 0: + session.reset_overlap_speech() + return { + "action": "drop", + "reason": "silence_or_noise", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": False, + } + + if is_short_ack_transcript_hint(event, payload): + session.reset_overlap_speech() + return { + "action": "listen", + "reason": "short_ack_transcript", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": False, + "defer_runtime_append": True, + } + + if policy == DuplexOverlapPolicy.LISTEN_ONLY.value: + session.accumulate_overlap_speech(duration_ms) + return { + "action": "listen", + "reason": "policy_listen_only", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + "defer_runtime_append": True, + } + + if policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value and not session.capabilities.supports_barge_in: + return defer_unsupported_barge_in(session, duration_ms=duration_ms, is_speech=True) + + session.accumulate_overlap_speech(duration_ms) + if policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value: + speech_started = vad_speech_started(event, payload) + if speech_started is False: + return { + "action": "listen", + "reason": "server_vad_utterance_active", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + "defer_runtime_append": False, + "force_listen": True, + "preserve_realtime_input": True, + } + return { + "action": "barge_in", + "reason": ("server_vad_speech_started" if speech_started is True else "policy_barge_in_on_speech"), + "cancel_reason": "turn_detected" if speech_started is True else "barge_in", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + } + + if ( + duration_ms <= session.config.overlap_short_ack_ms + and session.overlap_speech_ms <= session.config.overlap_short_ack_ms + ): + return { + "action": "listen", + "reason": "short_ack", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + "defer_runtime_append": True, + } + if session.overlap_speech_ms >= session.config.overlap_barge_in_ms: + if not session.capabilities.supports_barge_in: + return { + "action": "listen", + "reason": "barge_in_unsupported", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + "defer_runtime_append": True, + } + return { + "action": "barge_in", + "reason": "long_overlap_speech", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + } + return { + "action": "listen", + "reason": "accumulating_overlap_speech", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + "defer_runtime_append": True, + } + + +def vad_speech_started(event: Mapping[str, object], payload: Mapping[str, object]) -> bool | None: + for source in (event, payload): + vad = source.get("vad") + if isinstance(vad, Mapping) and isinstance(vad.get("speech_started"), bool): + return bool(vad["speech_started"]) + return None + + +def event_requests_barge_in(event: Mapping[str, object]) -> bool: + if event.get("force_barge_in") is True: + return True + explicit = event.get("overlap_action") or event.get("overlap") + return isinstance(explicit, str) and explicit.strip().lower() in {"barge_in", "interrupt", "cancel"} + + +def defer_unsupported_barge_in( + session: DuplexEngineSession, + *, + duration_ms: int, + is_speech: bool, +) -> dict[str, object]: + if is_speech: + session.accumulate_overlap_speech(duration_ms) + return { + "action": "listen", + "reason": "barge_in_unsupported", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": is_speech, + "defer_runtime_append": True, + } + + +def is_short_ack_transcript_hint(event: dict[str, object], payload: dict[str, object]) -> bool: + raw_text = event.get("transcript") or event.get("text") or payload.get("transcript") or payload.get("text") + if not isinstance(raw_text, str): + return False + normalized = raw_text.strip().lower() + if not normalized: + return False + compact = "".join(ch for ch in normalized if ch.isalnum() or "一" <= ch <= "鿿") + if compact in { + "嗯", + "嗯嗯", + "对", + "对的", + "好", + "好的", + "继续", + "继续说", + "可以", + "是的", + "yes", + "yeah", + "yep", + "ok", + "okay", + "continue", + "goon", + "right", + }: + return True + return normalized in {"go on", "keep going", "please continue"} + + +def input_audio_duration_ms(event: dict[str, object], payload: dict[str, object]) -> int: + for key in ("duration_ms", "audio_duration_ms"): + value = event.get(key) + if isinstance(value, int | float): + return max(0, int(value)) + fmt = payload.get("format") + sample_rate_hz = payload.get("sample_rate_hz") + audio = payload.get("audio") + if fmt == "pcm_f32le" and isinstance(sample_rate_hz, int) and sample_rate_hz > 0 and isinstance(audio, str): + try: + raw = base64.b64decode(audio, validate=True) + except (binascii.Error, ValueError): + return 0 + return int((len(raw) // 4) * 1000 / sample_rate_hz) + return 0 + + +def merge_audio_payloads(first: dict[str, object], second: dict[str, object]) -> dict[str, object]: + if first.get("format") != "pcm_f32le" or second.get("format") != "pcm_f32le": + return second + first_rate = first.get("sample_rate_hz") + second_rate = second.get("sample_rate_hz") + if not isinstance(first_rate, int) or not isinstance(second_rate, int) or first_rate != second_rate: + return second + first_audio = first.get("audio") + second_audio = second.get("audio") + if not isinstance(first_audio, str) or not isinstance(second_audio, str): + return second + try: + first_raw = base64.b64decode(first_audio, validate=True) + second_raw = base64.b64decode(second_audio, validate=True) + except (binascii.Error, ValueError): + return second + merged = dict(second) + merged["audio"] = base64.b64encode(first_raw + second_raw).decode("ascii") + merged["sample_rate_hz"] = first_rate + merged_frames = [ + frame + for source in (first.get("video_frames"), second.get("video_frames")) + if isinstance(source, list) + for frame in source + if isinstance(frame, str) and frame + ] + if merged_frames: + merged["video_frames"] = merged_frames + else: + merged.pop("video_frames", None) + merged["force_listen"] = bool(first.get("force_listen", False)) or bool(second.get("force_listen", False)) + merged.pop("force_speak", None) + merged["is_speech"] = bool(first.get("is_speech", False)) or bool(second.get("is_speech", False)) + return merged + + +def should_force_listen_for_short_commit( + session: DuplexEngineSession, event: dict[str, object], payload: dict[str, object] +) -> bool: + if event.get("force_listen") is True or payload.get("force_listen") is True: + return True + if event.get("force_barge_in") is True: + return False + if event.get("response_create") is not True: + return False + duration_ms = input_audio_duration_ms(event, payload) + return 0 < duration_ms <= session.config.overlap_short_ack_ms + + +def should_force_listen_for_auto_response_overlap( + event: dict[str, object], payload: dict[str, object], *, auto_responds: bool +) -> bool: + if not auto_responds: + return False + if event.get("force_barge_in") is True: + return False + return event.get("force_listen") is True or payload.get("force_listen") is True + + +def input_looks_like_speech(session: DuplexEngineSession, event: dict[str, object], payload: dict[str, object]) -> bool: + for key in ("is_speech", "speech"): + value = event.get(key) + if isinstance(value, bool): + return value + vad = event.get("vad") + if isinstance(vad, dict): + value = vad.get("is_speech") + if isinstance(value, bool): + return value + probability = vad.get("speech_probability", vad.get("probability")) + if isinstance(probability, int | float): + return float(probability) >= 0.5 + probability = event.get("speech_probability") + if isinstance(probability, int | float): + return float(probability) >= 0.5 + + fmt = payload.get("format") + audio = payload.get("audio") + if fmt in {"pcm_f32le", "pcm16"} and isinstance(audio, str): + try: + raw = base64.b64decode(audio, validate=True) + except (binascii.Error, ValueError): + return True + if fmt == "pcm_f32le": + if len(raw) < 4 or len(raw) % 4 != 0: + return True + samples = np.frombuffer(raw, dtype=np.float32) + else: + if len(raw) < 2 or len(raw) % 2 != 0: + return True + samples = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 + if samples.size == 0: + return False + rms = float(np.sqrt(np.mean(np.square(samples.astype(np.float32))))) + return rms >= session.config.overlap_silence_rms + return True diff --git a/vllm_omni/engine/duplex/session/playback_ledger.py b/vllm_omni/engine/duplex/session/playback_ledger.py new file mode 100644 index 000000000..c6843f64b --- /dev/null +++ b/vllm_omni/engine/duplex/session/playback_ledger.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Applying one ``playback.ack`` to a duplex session. + +The client reports how much of a response it has actually played; the session +uses that to commit or truncate the assistant item in conversation history, and +to release the response's playback state once it is fully played. + +Extracted from ``DuplexSessionRunner`` because the only runner state it touched +was ``session``. It returns the events it produces rather than emitting them, +so the rules can be exercised without a runner. +""" + +from __future__ import annotations + +from vllm_omni.engine.duplex.events import DuplexEvent, PlaybackAcknowledged, error_event +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession +from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity + + +def apply_playback_ack(session: DuplexEngineSession, event: dict[str, object]) -> list[DuplexEvent]: + """Apply one ``playback.ack`` and return the events it produces.""" + played_ms = event.get("played_ms", event.get("audio_ms", 0)) + committed_ms = event.get("committed_ms") + if not isinstance(played_ms, int | float): + return [error_event("bad_event", "playback.ack requires played_ms")] + try: + session.touch_lease(DuplexLeaseActivity.PLAYBACK_ACK) + except Exception: + pass + committed_cursor = int(committed_ms) if isinstance(committed_ms, int | float) else int(played_ms) + item_id = event.get("item_id") + response_id = event.get("response_id") + response_id = response_id if isinstance(response_id, str) and response_id else None + if not isinstance(item_id, str) or not item_id: + item_id = f"item_{response_id}" if response_id is not None else None + elif response_id is None and item_id.startswith("item_"): + response_id = item_id.removeprefix("item_") + if response_id is None and item_id is None and len(session.pending_history_item_ids) == 1: + item_id = next(iter(session.pending_history_item_ids)) + if item_id.startswith("item_"): + response_id = item_id.removeprefix("item_") + if response_id is None and item_id is None and session.active_response_id is not None: + response_id = session.active_response_id + item_id = f"item_{response_id}" + if response_id is not None: + expected_item_id = f"item_{response_id}" + if item_id is None: + item_id = expected_item_id + elif item_id != expected_item_id: + return [error_event("playback_item_mismatch", "playback.ack item_id must match item_.")] + if not session.has_assistant_response_item(response_id, item_id): + return [ + error_event("playback_item_not_found", f"No assistant response item is registered for {response_id}.") + ] + if session.playback_ack_is_too_late(response_id, item_id): + return [ + error_event("playback_ack_too_late", "playback.ack arrived after a later user input was committed.") + ] + session.reserve_history_item(item_id) + elif item_id is not None: + return [error_event("playback_item_not_found", "playback.ack requires a response-owned assistant item.")] + hard_truncate = event.get("truncate") is True + if hard_truncate: + playback = session.acknowledge_playback(int(played_ms), committed_cursor, response_id=response_id) + playback = session.truncate_playback_commit(committed_cursor, response_id=response_id) + else: + playback = session.acknowledge_playback(int(played_ms), committed_cursor, response_id=response_id) + committed_history = False + if isinstance(item_id, str) and item_id: + expected_item_id = f"item_{response_id}" if response_id is not None else None + if ( + expected_item_id == item_id + and item_id not in session.history_item_ids + and item_id not in session.pending_history_item_ids + ): + session.register_history_item(item_id, None) + committed_history = session.truncate_history_item( + item_id, + audio_end_ms=committed_cursor, + playback=playback, + hard=hard_truncate, + ) + elif session.pending_history_item_ids: + pending_ids = list(session.pending_history_item_ids) + if len(pending_ids) == 1: + item_id = pending_ids[0] + committed_history = session.truncate_history_item( + item_id, + audio_end_ms=committed_cursor, + playback=playback, + hard=hard_truncate, + ) + elif session.active_response_id is not None: + item_id = f"item_{session.active_response_id}" + committed_history = session.truncate_history_item( + item_id, + audio_end_ms=committed_cursor, + playback=playback, + hard=hard_truncate, + ) + elif session.last_assistant_full_message is not None: + if item_id is None and session.history_item_ids: + assistant_item_ids = [ + known_item_id + for known_item_id, message in session.history_item_ids.items() + if message.get("role") == "assistant" + ] + if len(assistant_item_ids) == 1: + item_id = assistant_item_ids[0] + if isinstance(item_id, str) and item_id: + committed_history = session.truncate_history_item( + item_id, + audio_end_ms=committed_cursor, + playback=playback, + hard=hard_truncate, + ) + events: list[DuplexEvent] = [ + PlaybackAcknowledged( + details={ + "type": "playback.acknowledged", + "session_id": session.session_id, + "epoch": session.epoch, + "item_id": item_id, + "played_ms": int(played_ms), + "committed_ms": committed_cursor, + "truncate": event.get("truncate") is True, + "playback": playback.as_dict(), + "history_committed": committed_history, + } + ) + ] + if committed_history and committed_cursor >= max(playback.sent_ms, playback.generated_ms): + session.release_response_playback(response_id) + session.release_response_history_snapshot(response_id) + return events diff --git a/vllm_omni/engine/duplex/session/runner.py b/vllm_omni/engine/duplex/session/runner.py new file mode 100644 index 000000000..7915500b7 --- /dev/null +++ b/vllm_omni/engine/duplex/session/runner.py @@ -0,0 +1,1772 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Engine-resident session runner: one ordered mailbox and one session state per duplex session. + +``DuplexSessionRunner`` owns the whole lifecycle of one session on the +orchestrator loop (the session is never touched from another thread): + +* every mutation of ``DuplexEngineSession`` happens on this loop, through the + mailbox worker (commands, stage outputs, internal items) or through tracked + tasks that re-validate ``(epoch, turn_id)`` after each ``await``; +* appends are planned with the model plugin and submitted to the stage port + in wire order on the per-session append tail (no RPC hop); +* stage outputs are pushed in by ``DuplexOrchestrator._intercept_stage_output`` + instead of being polled through request queues; +* everything the session says leaves as typed events through + ``DuplexSessionManager.emit`` after terminal-acceptance / stale-epoch + filtering, so a cancelled epoch can never speak again. +""" + +from __future__ import annotations + +import asyncio +import uuid +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, TypeVar + +from vllm.logger import init_logger + +from vllm_omni.engine.duplex.audio import convert_input_audio_with_rate +from vllm_omni.engine.duplex.commands import ( + AckPlayback, + AppendAudio, + AppendText, + BargeIn, + CancelInput, + CancelResponse, + ClearInput, + ClearOutputAudio, + CloseSession, + Commit, + CreateItem, + CreateResponse, + DeleteItem, + DuplexCommand, + Heartbeat, + SignalTurn, + TruncateItem, + UpdateSession, +) +from vllm_omni.engine.duplex.config import ( + DuplexConfigError, + DuplexOverlapPolicy, + DuplexPlaybackCommitPolicy, + DuplexSessionState, + ResponseCreateOptions, +) +from vllm_omni.engine.duplex.contracts import ( + DuplexOutputContext, + DuplexOutputDecision, + DuplexStagePort, +) +from vllm_omni.engine.duplex.events import ( + DuplexEvent, + InputCleared, + SessionExpired, + SessionHeartbeatAck, +) +from vllm_omni.engine.duplex.plugin import ( + DuplexModelPlugin, + DuplexModelSessionState, + PcmAppendReservation, + payload_turn_id, +) +from vllm_omni.engine.duplex.realtime_events import ( + RealtimeProjectionState, + discard_pending_input_audio, + note_input_append, + resolve_cancel_response, + resolve_clear_output_audio, + resolve_commit, + resolve_create_item, + resolve_delete_item, + resolve_truncate_item, + retrieve_item_events, +) +from vllm_omni.engine.duplex.session import helpers, overlap_policy, playback_ledger +from vllm_omni.engine.duplex.session.append_task import AppendAttempt +from vllm_omni.engine.duplex.session.commit_policy import CommitAction, CommitSnapshot, decide_commit_action +from vllm_omni.engine.duplex.session.context import ( + DuplexRunState, + DuplexSessionContext, + DuplexSessionTasks, + StageOutput, +) +from vllm_omni.engine.duplex.session.control import SessionControl +from vllm_omni.engine.duplex.session.emitter import SessionEmitter +from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession +from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity +from vllm_omni.engine.duplex.session.model_channel import ModelChannel +from vllm_omni.engine.duplex.turn_detection import ( + TurnDetectionResult, +) +from vllm_omni.metrics.stats import StageRequestStats + +if TYPE_CHECKING: + from vllm.config import ModelConfig + from vllm.outputs import RequestOutput + + from vllm_omni.engine.duplex.session.manager import DuplexSessionManager + +logger = init_logger(__name__) + +_OffloadT = TypeVar("_OffloadT") + + +@dataclass(frozen=True, slots=True) +class _Internal: + """Runner-internal mailbox item (its payload runs through the dict-based handlers).""" + + kind: str + payload: dict[str, object] = field(default_factory=dict) + + +_CANCEL_EVENTS = frozenset({"input.cancel", "response.cancel", "barge_in", "output_audio_buffer.clear"}) + + +class DuplexSessionRunner: + """Owns one ``DuplexEngineSession`` on the orchestrator loop (see module docstring).""" + + def __init__( + self, + *, + session: DuplexEngineSession, + plugin: DuplexModelPlugin, + stage_port: DuplexStagePort, + manager: DuplexSessionManager, + model_config: ModelConfig | None, + ) -> None: + self.session = session + self.plugin = plugin + self.stage_port = stage_port + self.manager = manager + self.model_config = model_config + if session.model_state is None: + session.model_state = plugin.create_session_state() + self.model_state: DuplexModelSessionState = session.model_state + self.tasks = DuplexSessionTasks() + self._mailbox: asyncio.Queue[DuplexCommand | StageOutput | _Internal] = asyncio.Queue() + self._worker: asyncio.Task[None] | None = None + self._worker_stopped = False + self._loop: asyncio.AbstractEventLoop | None = None + self._background_tasks: set[asyncio.Task[None]] = set() + #: Flags more than one component reads and writes (see session_context). + self.run = DuplexRunState() + self.ctx = DuplexSessionContext( + session=session, + model_state=self.model_state, + plugin=plugin, + stage_port=stage_port, + manager=manager, + tasks=self.tasks, + run=self.run, + services=self, + ) + self.out = SessionEmitter(self.ctx, promote_deferred_overlap=self._promote_deferred_overlap_later) + self.model = ModelChannel( + self.ctx, + self.out, + close_from_runtime=self._close_from_runtime, + schedule_silence_continuation=self._schedule_silence_continuation, + abort_request=self._abort_request_background, + ) + self.control = SessionControl( + self.ctx, + self.out, + self.model, + wait_for_append_tail=self._wait_for_append_tail, + ) + + # ------------------------------------------------------------------ # + # Public interface # + # ------------------------------------------------------------------ # + + def start(self) -> None: + self._loop = asyncio.get_running_loop() + session = self.session + if self.out.projector is None: + default_payload = session.config.extra_body.get("realtime_session_payload") + self.out.projector = RealtimeProjectionState( + session_id=session.session_id, + model=session.config.model, + default_payload=default_payload if isinstance(default_payload, Mapping) else None, + initial_session_update=True, + ) + # Every client speaks the Realtime protocol now; the old runner forced + # the ACK-only playback ledger for that path. + session.config.playback_commit_policy = DuplexPlaybackCommitPolicy.ACK_ONLY.value + self.control.init_turn_detection() + self._worker = self._loop.create_task(self._run(), name=f"duplex-session-{session.session_id}") + self.emit({"type": "session.created", "session": session.as_public_dict()}) + + def submit(self, command: DuplexCommand) -> None: + self._mailbox.put_nowait(command) + + def on_stage_output( + self, + stage_id: int, + output: RequestOutput, + metrics: StageRequestStats | None, + *, + request_id: str, + context: DuplexOutputContext, + ) -> bool: + """Accept one stage output (orchestrator loop); return True when it must not be forwarded.""" + decision: DuplexOutputDecision | None = None + if stage_id < context.final_stage_id: + decision = self.model.decide_output(stage_id, output, context) + consume = decision is not None or stage_id >= context.final_stage_id + if not consume: + # Stage0 text without a direct decision feeds the TTS stage as before. + # Its metrics still have to reach the client: before sessions moved + # into the engine the orchestrator published them as a standalone + # ``StageMetricsMessage``, a path session-owned requests no longer + # take. Hand them to the session instead of dropping them, on the + # mailbox so they stay ordered with this session's other work. + snapshot = self.model.stage_metrics_snapshot(stage_id, metrics, output) + if snapshot is not None and not self.run.closing and self.session.state != DuplexSessionState.CLOSED: + self._mailbox.put_nowait(_Internal("stage_metrics", {"stage_metrics": snapshot})) + return False + if self.run.closing or self.session.state == DuplexSessionState.CLOSED: + return True + self._mailbox.put_nowait( + StageOutput( + stage_id=stage_id, + output=output, + metrics=metrics, + request_id=request_id, + context=context, + decision=decision, + ) + ) + return True + + def on_stage_failure(self, stage_id: int, exc: BaseException) -> None: + """A stage rejected this session's request: fail the active response now. + + Runs synchronously on the loop (no mailbox hop): the orchestrator + expires the session right after this call, so a queued item could be + cancelled with the worker and the client would only see + ``session.expired``. Emitting here keeps the order + ``error`` -> failed ``response.done`` -> ``session.expired``. + """ + session = self.session + if session.state == DuplexSessionState.CLOSED: + return + self._emit_error( + "runtime_data_plane_stream_failed", + f"Stage-{stage_id} input processor failed: {type(exc).__name__}: {exc}", + ) + response_id = session.active_response_id + if response_id is not None: + session.end_response(commit_text=False) + self.emit( + { + "type": "response.done", + "session_id": session.session_id, + "response_id": response_id, + "epoch": session.epoch, + "committed": False, + "status": "failed", + "status_details": {"type": "failed", "reason": "runtime_data_plane_stream_failed"}, + "playback": session.playback.as_dict(), + } + ) + + @property + def closed_emitted(self) -> bool: + """Whether ``session.closed`` / ``session.expired`` already left this runner.""" + return self.run.closed_emitted + + def mark_closed_emitted(self) -> None: + """Claim the session's one terminal event for the caller. + + The manager emits the deferred terminal itself, after the stage cleanup; + recording it here stops a late runtime close emitting a second one. + """ + self.run.closed_emitted = True + self.run.closed_deferred = False + + @property + def closing(self) -> bool: + """Whether an irreversible close has begun (commands and control ops are refused).""" + return self.run.closing or self.session.state != DuplexSessionState.OPEN + + async def close(self, reason: str, *, emit_closed: bool = True) -> None: + """Graceful close: cancel work, release the data plane, emit ``session.closed``. + + With ``emit_closed=False`` the manager emits ``session.closed`` itself + once the stage resources are released, so the event also means "the + admission slot is free again". + """ + session = self.session + if session.state == DuplexSessionState.CLOSED and (self.run.closed_emitted or self.run.closed_deferred): + await self._stop_worker() + return + self._begin_close(reason) + self.model_state.audio_buffer.clear() + session.release_all_input_bytes() + self.model_state.input_since_commit = False + self.model_state.speech_since_commit = False + self.model_state.clear_committed_audio() + await self.tasks.cancel_append_tasks() + self.model.cancel_data_plane_stream() + await self._cancel_active_response(self.tasks.active_response_task, reason=reason, notify=False) + self.tasks.active_response_task = None + self._cleanup_duplex_session_state() + if not self.run.closed_emitted and not self.run.closed_deferred: + self.run.close_reason = self.run.close_reason or reason + if emit_closed: + self.run.closed_emitted = True + self.emit({"type": "session.closed", "session_id": session.session_id, "reason": reason}) + else: + self.run.closed_deferred = True + session.close() + await self._stop_worker() + + async def expire(self, reason: str, *, emit_expired: bool = True) -> None: + """Lease expiry / runtime cleanup: emit ``session.expired`` and tear down. + + With ``emit_expired=False`` the manager emits the event after the stage + cleanup (see ``close``). + """ + session = self.session + self._begin_close(reason) + if not self.run.closed_emitted: + if emit_expired: + # Also when the event was deferred: the manager only emits a + # deferred terminal for the teardown it drives itself, and an + # expiry that arrives first (a stage failure racing a wire + # ``session.close``) takes the runner away from it. Emitting + # here keeps "every session ends with one terminal event". + self.run.closed_emitted = True + self._emit_events([SessionExpired(reason=reason)]) + else: + self.run.closed_deferred = True + await self.tasks.cancel_append_tasks() + self.model.cancel_data_plane_stream() + active_response_task = self.tasks.active_response_task + if active_response_task is not None and not active_response_task.done(): + active_response_task.cancel() + await asyncio.gather(active_response_task, return_exceptions=True) + self.tasks.active_response_task = None + if session.active_response_id is not None: + session.end_response(commit_text=False) + self._cleanup_duplex_session_state() + session.close() + await self._stop_worker() + + async def shutdown(self) -> None: + self.run.closing = True + self.run.closed_emitted = True + await self.tasks.cancel_append_tasks() + self.model.cancel_data_plane_stream() + active_response_task = self.tasks.active_response_task + if active_response_task is not None and not active_response_task.done(): + active_response_task.cancel() + await asyncio.gather(active_response_task, return_exceptions=True) + self.session.close() + await self._stop_worker() + + # ------------------------------------------------------------------ # + # Worker # + # ------------------------------------------------------------------ # + + async def _run(self) -> None: + while not self._worker_stopped: + item = await self._mailbox.get() + try: + await self._handle_item(item) + except asyncio.CancelledError: + raise + except Exception as exc: + logger.exception("Duplex session %s failed handling %r: %s", self.session.session_id, item, exc) + self._emit_error("internal_error", str(exc)) + + async def _stop_worker(self) -> None: + worker = self._worker + self._worker = None + # Set even when the worker itself is stopping (a wire ``session.close`` + # runs inside it): its loop exits after the current item instead of + # parking on the mailbox forever. + self._worker_stopped = True + for task in list(self._background_tasks): + task.cancel() + if self._background_tasks: + await asyncio.gather(*self._background_tasks, return_exceptions=True) + self._background_tasks.clear() + if worker is not None and worker is not asyncio.current_task() and not worker.done(): + worker.cancel() + await asyncio.gather(worker, return_exceptions=True) + + def spawn(self, coro: Awaitable[None], *, name: str) -> asyncio.Task[None]: + task = asyncio.ensure_future(coro) + task.set_name(name) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) + return task + + async def offload(self, fn: Callable[..., _OffloadT], *args: object, **kwargs: object) -> _OffloadT: + loop = self._loop or asyncio.get_running_loop() + if kwargs: + return await loop.run_in_executor(self.manager.executor, lambda: fn(*args, **kwargs)) + return await loop.run_in_executor(self.manager.executor, fn, *args) + + async def _handle_item(self, item: DuplexCommand | StageOutput | _Internal) -> None: + session = self.session + if isinstance(item, StageOutput): + await self.model.on_stage_output_item(item) + return + if isinstance(item, _Internal): + await self._on_internal(item) + return + if self.run.closing or session.state == DuplexSessionState.CLOSED: + if isinstance(item, Commit): + session.release_pending_turn() + elif isinstance(item, AppendAudio): + session.release_input_bytes(len(item.audio)) + return + await self._on_command(item) + + async def _on_internal(self, item: _Internal) -> None: + if item.kind == "stage_metrics": + stage_metrics = item.payload.get("stage_metrics") + if isinstance(stage_metrics, Mapping): + self.session.stash_stage_metrics(stage_metrics) + return + if item.kind == "promote_deferred_overlap": + if self.run.closing or self.session.state != DuplexSessionState.OPEN: + return + payload = item.payload.get("payload") + if not isinstance(payload, dict): + return + await self._start_append( + payload, + final=True, + precreate_response=bool(item.payload.get("precreate_response", False)), + operation_id=self.model_state.committed_audio_operation_id, + retained_committed_payload=(payload if self.model_state.committed_audio_payload is payload else None), + ) + return + if item.kind == "commit": + await self._on_commit(dict(item.payload)) + return + if item.kind == "run_payload": + await self._run_internal_payload(dict(item.payload)) + return + logger.warning("Unknown duplex runner internal item: %s", item.kind) + + async def _run_internal_payload(self, payload: dict[str, object]) -> None: + """Run one internal event dictionary through the matching handler.""" + event_type = payload.get("type") + if event_type == "input_audio_buffer.append": + await self._on_append_audio(payload) + elif event_type in {"input_audio_buffer.commit", "input.commit", "response.create"}: + await self._on_commit(payload) + elif event_type == "playback.ack": + self._emit_events(playback_ledger.apply_playback_ack(self.session, payload)) + elif event_type in _CANCEL_EVENTS: + await self._on_cancel(payload) + elif event_type == "turn.signal": + await self.control.on_turn_signal(payload) + else: + self._emit_error("unknown_event", f"Unknown duplex event: {event_type}") + + # ------------------------------------------------------------------ # + # Commands # + # ------------------------------------------------------------------ # + + async def _on_command(self, command: DuplexCommand) -> None: + session = self.session + projector = self._require_projector() + if isinstance(command, AppendAudio): + # The manager reserved the wire size at admission; the handler + # re-reserves the decoded size around its PCM reservation. + session.release_input_bytes(len(command.audio)) + await self._on_append_audio(command.payload()) + elif isinstance(command, AppendText): + session.mark_user_input_activity() + self._emit_error( + "native_text_append_unsupported", + "The selected native duplex runtime accepts audio append only", + event_id=command.event_id, + ) + elif isinstance(command, Commit): + session.release_pending_turn() + resolved = resolve_commit(projector, command) + self._emit_events(resolved.events) + if resolved.reset_vad: + self.control.reset_vad() + if resolved.payload is not None: + await self._on_commit(resolved.payload) + elif isinstance(command, CreateResponse): + await self._on_commit(command.payload()) + elif isinstance(command, ClearInput): + self._on_clear_input() + elif isinstance(command, CancelResponse): + resolved = resolve_cancel_response(projector, command) + self._emit_events(resolved.events) + for payload in resolved.payloads: + await self._on_cancel(payload) + elif isinstance(command, ClearOutputAudio): + resolved = resolve_clear_output_audio(projector, command) + self._emit_events(resolved.events) + for payload in resolved.payloads: + await self._on_cancel(payload) + elif isinstance(command, CancelInput | BargeIn): + await self._on_cancel(command.payload()) + elif isinstance(command, SignalTurn): + if command.event == "conversation.item.retrieve": + self._emit_events( + retrieve_item_events( + projector, + { + **dict(command.signal_payload), + **({"event_id": command.event_id} if command.event_id else {}), + }, + ) + ) + return + payload = command.payload() + if command.event in _CANCEL_EVENTS: + normalized = dict(payload.get("payload") or {}) + normalized.update(payload) + normalized["type"] = command.event + await self._on_cancel(normalized) + else: + await self.control.on_turn_signal(payload) + elif isinstance(command, UpdateSession): + await self.control.on_session_update(dict(command.patch), realtime_event_id=command.event_id) + elif isinstance(command, AckPlayback): + self._emit_events(playback_ledger.apply_playback_ack(self.session, command.payload())) + elif isinstance(command, Heartbeat): + self._on_heartbeat(command) + elif isinstance(command, CreateItem): + resolved = resolve_create_item(projector, command) + self._emit_events(resolved.events) + for payload in resolved.payloads: + await self._run_internal_payload(payload) + elif isinstance(command, DeleteItem): + resolved = resolve_delete_item(projector, command) + self._emit_events(resolved.events) + for payload in resolved.payloads: + await self.control.on_turn_signal(payload) + elif isinstance(command, TruncateItem): + resolved = resolve_truncate_item(projector, command) + self._emit_events(resolved.events) + for payload in resolved.payloads: + await self._run_internal_payload(payload) + elif isinstance(command, CloseSession): + await self._on_close_command(command.reason) + else: + self._emit_error("unknown_event", f"Unknown duplex command: {type(command).__name__}") + + async def _on_close_command(self, reason: str) -> None: + # A close requested through the ordered command stream (the client's + # ``session.close``) tears the runner down here; the manager then + # aborts the stage requests, frees the admission slot and emits + # ``session.closed`` exactly like a close through the control RPC. + # Mark the session closing before the first await so a concurrent + # resume / touch / command is refused instead of racing the teardown. + self._begin_close(reason) + await self.close(reason, emit_closed=False) + self.manager.close_from_runner(self, reason) + + def _on_heartbeat(self, command: Heartbeat) -> None: + try: + self.session.touch_lease(DuplexLeaseActivity.HEARTBEAT) + except Exception as exc: + self._emit_error("runtime_touch_failed", str(exc), event_id=command.event_id) + return + self._emit_events([SessionHeartbeatAck()]) + + def _on_clear_input(self) -> None: + session = self.session + model_state = self.model_state + model_state.audio_buffer.clear() + session.release_all_input_bytes() + model_state.input_since_commit = False + model_state.speech_since_commit = False + model_state.clear_committed_audio() + session.cancel_pending_input() + projector = self.out.projector + if projector is not None: + from vllm_omni.engine.duplex.realtime_events import clear_input_buffer + + clear_input_buffer(projector) + self._emit_events([InputCleared()]) + + # ------------------------------------------------------------------ # + # Emission (delegated to SessionEmitter) # + # ------------------------------------------------------------------ # + + def emit(self, payload: dict[str, object]) -> None: + """Apply an internal event to the session, project it, and send it.""" + self.out.emit(payload) + + def _emit_events(self, events: list[DuplexEvent]) -> None: + self.out.emit_events(events) + + def _emit_error( + self, + code: str, + message: str, + *, + event_id: object | None = None, + retryable: bool | None = None, + ) -> None: + self.out.emit_error(code, message, event_id=event_id, retryable=retryable) + + def _require_projector(self) -> RealtimeProjectionState: + return self.out.require_projector() + + def _is_stale_model_output(self, payload: dict[str, object]) -> bool: + return self.out.is_stale_model_output(payload) + + def _promote_deferred_overlap_later(self, payload: dict[str, object], precreate_response: bool) -> None: + """Re-enter the mailbox so the promoted turn is handled in command order.""" + self._mailbox.put_nowait( + _Internal( + "promote_deferred_overlap", + {"payload": payload, "precreate_response": precreate_response}, + ) + ) + + # ------------------------------------------------------------------ # + # Session lifecycle helpers # + # ------------------------------------------------------------------ # + + def _begin_close(self, reason: str) -> None: + self.run.closing = True + self.run.close_reason = self.run.close_reason or reason + self.session.mark_closing() + + def _session_auto_responds(self) -> bool: + return self.out.auto_responds() + + def _cleanup_duplex_session_state(self) -> None: + session = self.session + self.plugin.data_plane.close_session(session.session_id, active_request_id=session.active_request_id) + self.run.stream_request_id = None + + # ------------------------------------------------------------------ # + # Append path (was the audio-append branch + start_native_append) # + # ------------------------------------------------------------------ # + + async def _barge_in_for_overlap(self, event: dict[str, object], decision: dict[str, object]) -> bool: + """Interrupt the model so the user's overlapping speech becomes the turn. + + Cancels the active response and the data-plane stream, advances the + barge-in epoch so stale output is dropped, and reports what was cut. + Returns whether the append should continue: losing the fence race + means this append belongs to a superseded epoch and is abandoned. + """ + session = self.session + model_state = self.model_state + event["force_barge_in"] = True + cancelled_fence = session.fence + playback_was_active = helpers.assistant_playback_active(self.session) + model_state.audio_buffer.clear_force_listen() + session.reset_overlap_speech() + model_state.input_since_commit = False + model_state.speech_since_commit = False + await self.tasks.cancel_append_tasks() + had_stream = self.run.stream_request_id is not None + cancel_reason = str(decision.get("cancel_reason") or "barge_in") + cancelled = await self._cancel_active_response( + self.tasks.active_response_task, + reason=cancel_reason, + ) + had_stream = self.model.cancel_data_plane_stream() or had_stream + if not cancelled and had_stream: + old_epoch = session.epoch + old_response_id = session.active_response_id + committed_ms = session.playback.committed_ms + helpers.commit_played_response_history(session, old_response_id, committed_ms) + new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + self.emit( + { + "type": "audio.cancelled", + "session_id": session.session_id, + "response_id": old_response_id, + "reason": cancel_reason, + "cancelled_epoch": old_epoch, + "epoch": new_epoch, + "committed_ms": committed_ms, + "playback": old_playback, + } + ) + cancelled = True + if not cancelled and playback_was_active: + old_epoch = session.epoch + committed_ms = session.playback.committed_ms + helpers.commit_played_response_history(session, session.last_response_id, committed_ms) + new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + self.emit( + { + "type": "audio.cancelled", + "session_id": session.session_id, + "response_id": session.last_response_id, + "reason": cancel_reason, + "cancelled_epoch": old_epoch, + "epoch": new_epoch, + "committed_ms": committed_ms, + "playback": old_playback, + } + ) + cancelled = True + if session.epoch > cancelled_fence.epoch: + if not await self.model.signal_cancel_fence(cancelled_fence): + return False + self.tasks.active_response_task = None + return True + + async def _on_append_audio(self, event: dict[str, object]) -> None: + session = self.session + model_state = self.model_state + session.mark_user_input_activity() + audio = event.get("audio") or event.get("data") + if not isinstance(audio, str): + self._emit_error("bad_event", "input_audio_buffer.append requires audio") + return + if not session.capabilities.supports_barge_in and overlap_policy.event_requests_barge_in(event): + self._emit_events([helpers.barge_in_unsupported_error()]) + event = dict(event) + event.pop("force_barge_in", None) + for key in ("overlap_action", "overlap"): + value = event.get(key) + if isinstance(value, str) and value.strip().lower() in {"barge_in", "interrupt", "cancel"}: + event.pop(key, None) + fmt = event.get("format") if isinstance(event.get("format"), str) else "pcm16" + sr_raw = event.get("sample_rate_hz") or event.get("sample_rate") + sample_rate_hz = sr_raw if isinstance(sr_raw, int | float) else 16000 + try: + audio, fmt, sample_rate_hz = await self.offload( + convert_input_audio_with_rate, + audio, + fmt, + sample_rate_hz=sample_rate_hz, + ) + except ValueError as exc: + self._emit_error("bad_event", str(exc)) + return + if isinstance(fmt, str) and fmt.lower() in {"pcm16", "pcm_s16le", "s16le"}: + self._emit_error("bad_audio", "input_audio_buffer.append pcm16 audio could not be decoded") + return + event["audio"] = audio + event["format"] = fmt + event["sample_rate_hz"] = sample_rate_hz + vad_result = await self.control.run_turn_detection(event) + projector = self._require_projector() + self._emit_events(note_input_append(projector, event, vad_result=vad_result)) + if self.run.closing or session.state != DuplexSessionState.OPEN: + return + + force_listen = bool(event.get("force_listen", False)) + payload: dict[str, object] = { + "type": "audio", + "audio": audio, + "format": fmt, + "sample_rate_hz": sample_rate_hz, + "force_listen": force_listen, + } + video_frames = event.get("video_frames") + if isinstance(video_frames, list): + frames = [frame for frame in video_frames if isinstance(frame, str) and frame] + if frames: + payload["video_frames"] = frames + payload["is_speech"] = overlap_policy.input_looks_like_speech(self.session, event, payload) + auto_responds = self._session_auto_responds() + defer_append = False + buffer_overlap_audio = True + self._mark_pending_silence_superseded() + overlap_active = helpers.response_in_progress(self.session, self.tasks) and ( + not auto_responds + or ( + session.capabilities.supports_barge_in + and ( + session.config.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value + or overlap_policy.event_requests_barge_in(event) + ) + ) + ) + if overlap_active: + decision = overlap_policy.decide(self.session, event, payload, auto_responds=auto_responds) + self._emit_events([helpers.overlap_decision_event(self.session, decision)]) + action = decision.get("action") + if action == "drop": + duration_ms = overlap_policy.input_audio_duration_ms(event, payload) + self._emit_events(discard_pending_input_audio(projector, duration_ms)) + self._maybe_schedule_vad_commit(vad_result) + return + if action == "listen": + buffer_overlap_audio = bool(decision.get("buffer_audio", True)) + defer_append = bool(decision.get("defer_runtime_append", True)) + if not buffer_overlap_audio and decision.get("preserve_realtime_input") is not True: + self._emit_events( + discard_pending_input_audio(projector, overlap_policy.input_audio_duration_ms(event, payload)) + ) + if decision.get("force_listen", True) is True: + payload["force_listen"] = True + else: + if not await self._barge_in_for_overlap(event, decision): + return + buffer_overlap_audio = True + defer_append = False + elif not auto_responds and not overlap_policy.input_looks_like_speech(self.session, event, payload): + # Turn-mode only: skip silent chunks so they don't open a response. + self.emit( + { + "type": "response.listen", + "session_id": session.session_id, + "epoch": session.epoch, + "reason": "silence_or_noise", + } + ) + self._maybe_schedule_vad_commit(vad_result) + return + if overlap_policy.should_force_listen_for_auto_response_overlap(event, payload, auto_responds=auto_responds): + payload["force_listen"] = True + if not buffer_overlap_audio: + self._maybe_schedule_vad_commit(vad_result) + return + session.mark_user_input_activity() + model_state.input_since_commit = True + model_state.speech_since_commit = model_state.speech_since_commit or overlap_policy.input_looks_like_speech( + self.session, event, payload + ) + raw_audio_bytes = helpers.audio_payload_size_bytes(payload) + try: + if not session.reserve_input_bytes( + raw_audio_bytes, + limit=int(self.manager.runtime_config.max_pending_input_bytes_per_session), + ): + self._emit_error("input_backpressure", "Duplex session pending input exceeds server limit") + return + # Full-duplex: emit each ~chunk_period of audio so the model runs + # per-chunk generation without an explicit response.create. + allow_emit = not defer_append and self._session_auto_responds() + pcm_reservation = model_state.audio_buffer.prepare_append( + payload, + operation_id=uuid.uuid4().hex, + chunk_period_ms=session.capabilities.chunk_period_ms or 1000, + allow_emit=allow_emit, + ) + except ValueError as exc: + session.release_input_bytes(raw_audio_bytes) + self._emit_error("bad_event", str(exc)) + return + if pcm_reservation is None: + self._maybe_schedule_vad_commit(vad_result) + return + if pcm_reservation.byte_count == 0: + session.release_input_bytes(raw_audio_bytes) + payload = pcm_reservation.payload + await self._start_append(payload, final=False, pcm_reservation=pcm_reservation) + self._maybe_schedule_vad_commit(vad_result) + + def _maybe_schedule_vad_commit(self, vad_result: TurnDetectionResult | None) -> None: + """Server VAD ended the user turn: run the same commit the old translator synthesized.""" + if vad_result is None or not vad_result.should_commit: + return + if self._session_auto_responds(): + # A model-native session decides its own turns; server VAD is there + # to hear the user (speech_started / speech_stopped, barge-in), not + # to end the turn. Committing at the detector's stop cuts the + # utterance short: the model listens on that early final unit, the + # trailing silence lands in a second, near-empty turn, and the + # client's own commit then has nothing left to answer. The old + # translator synthesized this commit for turn-based server VAD only. + return + command = Commit(final=True, create_response=vad_result.create_response) + resolved = resolve_commit(self._require_projector(), command) + self._emit_events(resolved.events) + if resolved.reset_vad: + self.control.reset_vad() + if resolved.payload is not None: + self._mailbox.put_nowait(_Internal("commit", resolved.payload)) + + def _clear_completed_pending_silence(self) -> None: + task = self.model_state.pending_silence_task + if task is not None and task.done(): + self.model_state.pending_silence_task = None + self.model_state.pending_silence_owner_id = None + + def _mark_pending_silence_superseded(self) -> None: + task = self.model_state.pending_silence_task + if task is None: + return + if task.done(): + self.model_state.pending_silence_task = None + # Do not cancel the task here: a silence append may already have + # reached the stage; before_append skips silence that has not started. + self.model_state.pending_silence_owner_id = None + + def _real_input_waiting(self) -> bool: + self._clear_completed_pending_silence() + return self.model_state.audio_buffer.has_pending() or self.model_state.audio_buffer.has_reserved() + + async def _start_append( + self, + payload: dict[str, object], + *, + final: bool, + precreate_response: bool = False, + pcm_reservation: PcmAppendReservation | None = None, + operation_id: str | None = None, + retained_committed_payload: dict[str, object] | None = None, + silence_continuation: bool = False, + before_append: Callable[[], bool] | None = None, + ) -> asyncio.Task[bool]: + session = self.session + model_state = self.model_state + if not silence_continuation: + self._mark_pending_silence_superseded() + append_epoch = session.epoch + append_turn_id = payload_turn_id(payload) + if append_turn_id is None: + append_turn_id = session.turn_id + request_id = helpers.stage0_request_id(self.session, append_epoch) + if final or precreate_response: + session.bind_request(request_id) + if precreate_response: + session.bind_response_turn(append_turn_id) + if precreate_response and session.active_response_id is None: + response_id = session.begin_response(turn_id=append_turn_id) + self.emit(self.model.response_created_payload(response_id, epoch=append_epoch)) + attempt = AppendAttempt( + ctx=self.ctx, + out=self.out, + model=self.model, + fail_session=self._fail_session_from_append, + payload=payload, + epoch=append_epoch, + request_id=request_id, + final=final, + pcm_reservation=pcm_reservation, + operation_id=operation_id, + retained_committed_payload=retained_committed_payload, + precreated_response_id=session.active_response_id if precreate_response else None, + before_append=before_append, + ) + + predecessor = self.tasks.append_tail + if predecessor is not None and predecessor.done(): + try: + predecessor_ok = predecessor.result() + except (asyncio.CancelledError, Exception): + predecessor_ok = False + if not predecessor_ok: + # Appends queued behind a failed predecessor stop; a later + # command is an explicit retry and starts a new chain. + predecessor = None + task = asyncio.create_task(attempt.run_in_wire_order(predecessor)) + task.add_done_callback(attempt.release_on_failure) + self.tasks.append_tail = task + self.tasks.track_append_task( + task, + epoch=append_epoch, + final=final, + response_bound=final or precreate_response, + ) + if silence_continuation: + model_state.pending_silence_task = task + task.add_done_callback(attempt.clear_pending_silence) + # Let this wire-order effect start before the next mailbox item can cancel it. + await asyncio.sleep(0) + return task + + def _fail_session_from_append(self, reason: str) -> None: + """An append task died: mark the session closing now, close it out next.""" + self._begin_close(reason) + self.spawn(self._close_from_runtime(reason), name="duplex-runtime-close") + + async def _wait_for_append_tail(self) -> bool: + predecessor = self.tasks.append_tail + if predecessor is None: + return True + try: + return await predecessor + except asyncio.CancelledError: + current = asyncio.current_task() + if current is not None and current.cancelling(): + raise + return False + except Exception: + return False + + async def _schedule_silence_continuation( + self, + payload: object, + *, + request_id: str, + owner_id: str, + response_id: str | None, + response_owned: bool, + expected_epoch: int | None, + expected_model_turn_id: int | None, + ) -> bool: + session = self.session + model_state = self.model_state + self._clear_completed_pending_silence() + pending_silence = model_state.pending_silence_task + if pending_silence is not None and not pending_silence.done(): + if pending_silence is asyncio.current_task(): + return False + try: + if not await pending_silence: + return False + except asyncio.CancelledError: + current = asyncio.current_task() + if current is not None and current.cancelling(): + raise + return False + except Exception: + return False + self._clear_completed_pending_silence() + pending_silence = model_state.pending_silence_task + if pending_silence is not None and not pending_silence.done(): + return False + append_tail = self.tasks.append_tail + if (append_tail is None or append_tail.done()) and self._real_input_waiting(): + return False + continuation_delay_s = max(0.0, float(session.capabilities.chunk_period_ms or 1000) / 1000.0) + if continuation_delay_s > 0: + await asyncio.sleep(continuation_delay_s) + if ( + self.tasks.append_tail is not append_tail + or ((append_tail is None or append_tail.done()) and self._real_input_waiting()) + or self.model.silence_continuation_is_stale( + request_id=request_id, + response_id=response_id, + response_owned=response_owned, + expected_epoch=expected_epoch, + expected_model_turn_id=expected_model_turn_id, + ) + ): + return False + + def _still_valid() -> bool: + return not self._real_input_waiting() and not self.model.silence_continuation_is_stale( + request_id=request_id, + response_id=response_id, + response_owned=response_owned, + expected_epoch=expected_epoch, + expected_model_turn_id=expected_model_turn_id, + ) + + model_state.pending_silence_owner_id = owner_id + task = await self._start_append( + dict(payload) if isinstance(payload, dict) else {}, + final=False, + silence_continuation=True, + before_append=_still_valid, + ) + return task is not None + + # ------------------------------------------------------------------ # + # Runtime-initiated close # + # ------------------------------------------------------------------ # + + async def _close_from_runtime(self, reason: str) -> None: + """A runtime failure closed the session from inside the engine.""" + session = self.session + if session.state == DuplexSessionState.CLOSED: + return + self.run.runtime_closed = True + self._begin_close(reason) + self._cleanup_duplex_session_state() + # ``closed_deferred`` means a manager-driven close already promised the + # terminal after its stage cleanup. Emitting here too would give the + # session two terminal events: this path can run concurrently, when an + # append task abandoned by that close fails and calls back in. + if not self.run.closed_emitted and not self.run.closed_deferred: + self.run.closed_emitted = True + self.emit({"type": "session.closed", "session_id": session.session_id, "reason": reason}) + session.close() + # The manager aborts the stage requests (if any) and frees the + # admission slot; a session without stage requests must not stay + # registered until its idle TTL. + self.manager.close_from_runner(self, reason) + + # ------------------------------------------------------------------ # + # Cancel / barge-in # + # ------------------------------------------------------------------ # + + async def _on_cancel(self, event: dict[str, object]) -> None: + session = self.session + model_state = self.model_state + event_type = str(event.get("type")) + if event_type == "barge_in" and not session.capabilities.supports_barge_in: + self._emit_events([helpers.barge_in_unsupported_error()]) + return + cancel_reason = ( + "output_audio_buffer_clear" + if event_type == "output_audio_buffer.clear" + else "client_cancelled" + if event_type == "response.cancel" + else "barge_in" + ) + cancelled_fence = session.fence + if event_type == "response.cancel": + requested_response_id = event.get("response_id") + has_active_response_work = helpers.response_in_progress(self.session, self.tasks) + if ( + isinstance(requested_response_id, str) + and session.active_response_id is not None + and requested_response_id != session.active_response_id + ): + self._emit_error( + "response_not_active", + f"Response is not active: {requested_response_id}", + event_id=event.get("realtime_event_id"), + ) + return + if not has_active_response_work: + if isinstance(requested_response_id, str): + return + self._emit_error( + "response_not_active", + "response.cancel requires an active response", + event_id=event.get("realtime_event_id"), + ) + return + had_unbuffered_append = model_state.input_since_commit and not model_state.audio_buffer.has_pending() + playback_was_active = helpers.assistant_playback_active(self.session) + if event_type in {"input.cancel", "barge_in"}: + model_state.audio_buffer.clear() + session.release_all_input_bytes() + model_state.input_since_commit = False + model_state.speech_since_commit = False + model_state.clear_committed_audio() + had_append = await self.tasks.cancel_append_tasks( + response_bound_only=event_type in {"response.cancel", "output_audio_buffer.clear"}, + ) + if event_type == "response.cancel": + session.release_input_bytes(model_state.clear_committed_audio()) + had_stream = self.run.stream_request_id is not None + cancelled = await self._cancel_active_response(self.tasks.active_response_task, reason=cancel_reason) + had_stream = self.model.cancel_data_plane_stream() or had_stream + if not cancelled and (had_stream or had_append or had_unbuffered_append): + old_epoch = session.epoch + old_response_id = session.active_response_id + committed_ms = session.playback.committed_ms + helpers.commit_played_response_history(session, old_response_id, committed_ms) + new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + self.emit( + { + "type": "audio.cancelled", + "session_id": session.session_id, + "response_id": old_response_id, + "reason": cancel_reason, + "cancelled_epoch": old_epoch, + "epoch": new_epoch, + "committed_ms": committed_ms, + "playback": old_playback, + } + ) + cancelled = True + if not cancelled and playback_was_active: + old_epoch = session.epoch + committed_ms = session.playback.committed_ms + helpers.commit_played_response_history(session, session.last_response_id, committed_ms) + new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + self.emit( + { + "type": "audio.cancelled", + "session_id": session.session_id, + "response_id": session.last_response_id, + "reason": cancel_reason, + "cancelled_epoch": old_epoch, + "epoch": new_epoch, + "committed_ms": committed_ms, + "playback": old_playback, + } + ) + cancelled = True + if not cancelled and event_type == "response.cancel": + old_epoch = session.epoch + old_response_id = session.active_response_id + committed_ms = session.playback.committed_ms + helpers.commit_played_response_history(session, old_response_id, committed_ms) + new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + self.emit( + { + "type": "audio.cancelled", + "session_id": session.session_id, + "response_id": old_response_id, + "reason": cancel_reason, + "cancelled_epoch": old_epoch, + "epoch": new_epoch, + "committed_ms": committed_ms, + "playback": old_playback, + } + ) + cancelled = True + if not cancelled and event_type == "output_audio_buffer.clear": + old_playback = session.playback.as_dict() + committed_ms = session.playback.committed_ms + session.clear_playback_cursor() + self.emit( + { + "type": "audio.cancelled", + "session_id": session.session_id, + "response_id": session.active_response_id, + "reason": cancel_reason, + "cancelled_epoch": session.epoch, + "epoch": session.epoch, + "committed_ms": committed_ms, + "playback": old_playback, + } + ) + self.tasks.active_response_task = None + return + if not cancelled: + self._cancel_pending_input(reason="barge_in") + if not await self.model.signal_cancel_fence(cancelled_fence): + return + self.tasks.active_response_task = None + + async def _cancel_active_response( + self, + active_task: asyncio.Task[None] | None, + *, + reason: str, + notify: bool = True, + ) -> bool: + session = self.session + has_running_task = active_task is not None and not active_task.done() + if not has_running_task and session.active_request_id is None and session.active_response_id is None: + return False + + old_epoch = session.epoch + old_request_id = session.active_request_id + old_response_id = session.active_response_id + committed_ms = session.playback.committed_ms + committed_message = session.end_response( + commit_text=self.model.should_commit_response_to_history(session, old_response_id), + playback_commit_policy=DuplexPlaybackCommitPolicy.ACK_ONLY.value, + ) + if old_response_id is not None: + item_id = f"item_{old_response_id}" + if committed_message is not None: + session.register_history_item(item_id, committed_message) + elif committed_ms > 0 and not session.playback_ack_is_too_late(old_response_id, item_id): + session.truncate_history_item(item_id, audio_end_ms=committed_ms) + # The epoch bump is the atomic part: from here on every model output + # and append of the old epoch is dropped by the stale-epoch filter in + # ``emit`` / the append tail, whatever the awaits below interleave with. + new_epoch, old_playback = helpers.advance_barge_in_epoch(session) + if old_request_id is not None: + # Release projector/parser cursors so cancelled epochs do not + # accumulate until the whole session closes. + self.plugin.data_plane.close_stream(old_request_id) + await self._abort_request_background(old_request_id, notify=notify) + if has_running_task and active_task is not None: + active_task.cancel() + try: + await asyncio.wait_for(asyncio.gather(active_task, return_exceptions=True), timeout=0.25) + except TimeoutError: + pass + if notify: + self.emit( + { + "type": "audio.cancelled", + "session_id": session.session_id, + "response_id": old_response_id, + "reason": reason, + "cancelled_epoch": old_epoch, + "epoch": new_epoch, + "committed_ms": committed_ms, + "playback": old_playback, + } + ) + return True + + async def _abort_request_background(self, request_id: str, *, notify: bool) -> None: + try: + await self.stage_port.abort_requests([request_id]) + except asyncio.CancelledError: + raise + except Exception as exc: + logger.exception("Failed to abort duplex request %s: %s", request_id, exc) + if notify and self.session.state != DuplexSessionState.CLOSED: + self.model.send_runtime_error("runtime_abort_failed", exc) + + def _cancel_pending_input(self, *, reason: str) -> None: + session = self.session + cancelled = session.cancel_pending_input() + helpers.advance_barge_in_epoch(session) + self.emit( + { + "type": "input.cancelled", + "session_id": session.session_id, + "reason": reason, + "epoch": session.epoch, + "cancelled": cancelled, + } + ) + + # ------------------------------------------------------------------ # + # Commit / response.create # + # ------------------------------------------------------------------ # + + def _apply_response_create_options(self, response_payload: dict[str, object]) -> str | None: + session = self.session + try: + options = ResponseCreateOptions.from_realtime( + response_payload, + private_runtime_config_keys=self.plugin.private_runtime_config_keys, + ) + except DuplexConfigError as exc: + return exc.code + try: + session.reserve_response_options(options) + except Exception: + return "response_already_active" + return None + + async def _flush_and_submit_committed_turn( + self, + event: dict[str, object], + *, + event_type: str, + realtime_item_id: object, + should_create_response: bool, + ) -> bool: + """Flush the buffered turn and submit it, with no response yet running. + + Returns whether the turn was flushed. An empty buffer falls through to + the bare acknowledgement at the end of ``_on_commit``. + """ + session = self.session + model_state = self.model_state + commit_reservation = ( + model_state.audio_buffer.prepare_commit( + operation_id=uuid.uuid4().hex, + chunk_period_ms=session.capabilities.chunk_period_ms or 1000, + ) + if event_type in {"input.commit", "input_audio_buffer.commit"} + else None + ) + flushed_buffer_reserved_bytes = model_state.audio_buffer.pending_byte_count if commit_reservation is None else 0 + flushed = ( + commit_reservation.payload + if commit_reservation is not None + else model_state.audio_buffer.flush(chunk_period_ms=session.capabilities.chunk_period_ms or 1000) + ) + if model_state.committed_audio_payload is not None: + if flushed is not None: + flushed = overlap_policy.merge_audio_payloads(model_state.committed_audio_payload, flushed) + else: + flushed = model_state.committed_audio_payload + if commit_reservation is not None: + commit_reservation.commit() + if flushed is None: + return False + if overlap_policy.should_force_listen_for_short_commit(self.session, event, flushed): + flushed = dict(flushed) + flushed["force_listen"] = True + model_state.input_since_commit = False + committed = helpers.commit_audio_input( + session, + realtime_item_id=realtime_item_id, + transcript=event.get("transcript"), + ) + self.emit( + helpers.audio_committed_payload( + session, + committed=committed, + realtime_item_id=realtime_item_id, + transcript=event.get("transcript"), + ) + ) + operation_id = commit_reservation.operation_id if commit_reservation is not None else uuid.uuid4().hex + reserved_bytes = ( + commit_reservation.byte_count if commit_reservation is not None else flushed_buffer_reserved_bytes + ) + model_state.retain_committed_audio(flushed, operation_id=operation_id, reserved_bytes=reserved_bytes) + if should_create_response: + await self._start_append( + flushed, + final=True, + precreate_response=True, + operation_id=model_state.committed_audio_operation_id, + retained_committed_payload=flushed, + ) + else: + model_state.deferred_response_create = False + return True + + async def _start_response_from_committed_audio(self) -> None: + """Answer a client ``response.create``. + + A turn starts from committed input. Audio is the usual kind, but the + Realtime protocol also lets a client build the turn out of + ``conversation.item.create`` and ask for a response with no audio at + all; those items are input too. A model-native session still generates + per audio unit, so an item-only turn is primed with the model's own + silence unit -- the one continuations already use -- rather than with + anything this layer invents. + + Refused only when nothing is waiting, or a response is already running, + rather than silently producing an empty turn. + """ + session = self.session + model_state = self.model_state + if ( + helpers.response_in_progress(self.session, self.tasks) + or self.tasks.append_tasks + or self.run.stream_request_id is not None + ): + if session.active_response_id is None and ( + session.active_request_id is not None + or self.tasks.append_tasks + or self.run.stream_request_id is not None + ): + return + self._emit_error( + "response_already_active", "response.create cannot start while another response is active." + ) + session.discard_response_options() + return + if model_state.committed_audio_payload is not None: + committed_payload = model_state.committed_audio_payload + operation_id = model_state.committed_audio_operation_id + if operation_id is None: + operation_id = uuid.uuid4().hex + model_state.committed_audio_operation_id = operation_id + session.reset_unanswered_user_items() + await self._start_append( + committed_payload, + final=True, + precreate_response=True, + operation_id=operation_id, + retained_committed_payload=committed_payload, + ) + return + if session.unanswered_user_items(): + # Conversation items are context, not a turn. A model-native model + # decides to speak from the audio it hears, and there is no audio + # here: opening a response anyway produces one the model never + # fills, which the caller only discovers when the session idles + # out. Text reaches such a model through the seeded opening turn + # (``DuplexSessionConfig.initial_user_text``) instead. + self._emit_error( + "text_only_turn_unsupported", + "This duplex model answers speech input only. Seed the session with " + "initial_user_text to ask it in text." + if session.capabilities.supports_chat_completions + else "This duplex model answers speech input only.", + ) + session.reset_unanswered_user_items() + session.discard_response_options() + return + self._emit_error( + "response_create_without_input", + "Duplex response.create requires committed audio to answer.", + ) + session.discard_response_options() + + def _discard_short_overlap_ack(self) -> None: + """Drop a sub-threshold interjection made while the model is speaking. + + Too short to be a turn, so it neither interrupts the response nor + becomes one: the buffered audio is discarded and the turn stays with + the model. + """ + session = self.session + model_state = self.model_state + model_state.audio_buffer.clear() + session.release_all_input_bytes() + model_state.input_since_commit = False + model_state.speech_since_commit = False + model_state.clear_committed_audio() + self._emit_events(discard_pending_input_audio(self._require_projector(), session.overlap_speech_ms)) + self.emit( + { + "type": "input.committed", + "session_id": session.session_id, + "turn_id": session.turn_id, + "epoch": session.epoch, + "empty": True, + "is_speech": False, + "overlap_ack": True, + "no_response": True, + } + ) + session.reset_overlap_speech() + session.discard_response_options() + + def _defer_commit_behind_active_response( + self, + event: dict[str, object], + *, + realtime_item_id: object, + should_create_response: bool, + precreate_response_requested: bool, + ) -> bool: + """Retain a commit that arrived while a response is still running. + + The audio is kept so the turn is not lost, and the response it should + start is remembered; ``_maybe_promote_deferred_overlap`` replays it once + the active response ends. Returns whether the commit was deferred -- + an empty buffer has nothing to retain and falls through. + """ + session = self.session + model_state = self.model_state + commit_reservation = model_state.audio_buffer.prepare_commit( + operation_id=uuid.uuid4().hex, + chunk_period_ms=session.capabilities.chunk_period_ms or 1000, + ) + deferred_payload = commit_reservation.payload + if deferred_payload is None: + commit_reservation.commit() + return False + if model_state.committed_audio_payload is not None: + deferred_payload = overlap_policy.merge_audio_payloads( + model_state.committed_audio_payload, + deferred_payload, + ) + model_state.retain_committed_audio( + deferred_payload, + operation_id=commit_reservation.operation_id, + reserved_bytes=commit_reservation.byte_count, + ) + commit_reservation.commit() + model_state.deferred_response_create = should_create_response + model_state.deferred_precreate_response = precreate_response_requested + model_state.input_since_commit = False + model_state.speech_since_commit = False + committed = helpers.commit_audio_input( + session, + realtime_item_id=realtime_item_id, + transcript=event.get("transcript"), + ) + committed_payload = helpers.audio_committed_payload( + session, + committed=committed, + realtime_item_id=realtime_item_id, + transcript=event.get("transcript"), + ) + committed_payload["overlap_deferred"] = True + committed_payload["response_create_deferred"] = should_create_response + self.emit(committed_payload) + return True + + async def _commit_and_start_auto_response(self, event: dict[str, object], *, realtime_item_id: object) -> None: + """Commit the buffered turn and submit it, letting the model answer it.""" + session = self.session + model_state = self.model_state + commit_reservation = model_state.audio_buffer.prepare_commit( + operation_id=uuid.uuid4().hex, + chunk_period_ms=session.capabilities.chunk_period_ms or 1000, + ) + final_payload = commit_reservation.payload + if model_state.committed_audio_payload is not None: + if final_payload is not None: + final_payload = overlap_policy.merge_audio_payloads(model_state.committed_audio_payload, final_payload) + else: + final_payload = model_state.committed_audio_payload + commit_reservation.commit() + if final_payload is not None: + model_state.retain_committed_audio( + final_payload, + operation_id=commit_reservation.operation_id, + reserved_bytes=commit_reservation.byte_count, + ) + model_state.deferred_response_create = False + model_state.input_since_commit = False + model_state.speech_since_commit = False + data_plane_turn_id = session.turn_id + committed = helpers.commit_audio_input( + session, + realtime_item_id=realtime_item_id, + transcript=event.get("transcript"), + turn_id=data_plane_turn_id, + ) + self.emit( + helpers.audio_committed_payload( + session, + committed=committed, + realtime_item_id=realtime_item_id, + transcript=event.get("transcript"), + ) + ) + if final_payload is not None: + await self._start_append( + {**final_payload, "duplex_turn_id": data_plane_turn_id}, + final=True, + precreate_response=False, + operation_id=commit_reservation.operation_id, + retained_committed_payload=final_payload, + ) + + def _commit_silent_input(self) -> None: + """Commit a turn the client itself marked as silence: drop it and keep listening.""" + session = self.session + model_state = self.model_state + model_state.input_since_commit = False + model_state.speech_since_commit = False + model_state.audio_buffer.clear() + session.release_all_input_bytes() + model_state.clear_committed_audio() + self.emit( + { + "type": "input.committed", + "session_id": session.session_id, + "turn_id": session.turn_id, + "epoch": session.epoch, + "empty": True, + "is_speech": False, + "no_response": True, + } + ) + self.emit( + { + "type": "response.listen", + "session_id": session.session_id, + "epoch": session.epoch, + "reason": "silence_or_noise", + } + ) + + async def _on_commit(self, event: dict[str, object]) -> None: + session = self.session + model_state = self.model_state + event_type = str(event.get("type")) + realtime_item_id = event.get("realtime_item_id") + realtime_validated_audio_commit = ( + event_type == "input_audio_buffer.commit" and isinstance(realtime_item_id, str) and bool(realtime_item_id) + ) + if event_type in {"input.commit", "input_audio_buffer.commit"} and not await self._wait_for_append_tail(): + # resolve_commit has already applied this commit's side effects + # (speech_stopped emitted, input flags cleared), so returning + # silently leaves the client with those effects, no + # input_audio_buffer.committed and no error. The session.update + # path says so explicitly; this one now does too. + self._emit_error( + "commit_aborted", + "the commit was not applied because the preceding append failed", + event_id=event.get("realtime_event_id"), + ) + return + if event_type == "input_audio_buffer.commit" and event.get("is_speech") is False: + self._commit_silent_input() + return + should_create_response = ( + event_type == "response.create" + or bool(event.get("response_create", event_type == "input.commit")) + or (event_type == "input_audio_buffer.commit" and self._session_auto_responds()) + ) + precreate_response_requested = event_type == "response.create" or bool( + event.get("response_create", event_type == "input.commit") + ) + if event_type == "response.create": + response_payload = event.get("response") + if isinstance(response_payload, dict): + response_options_error = self._apply_response_create_options(response_payload) + if response_options_error is not None: + error_message = ( + "The selected native duplex runtime does not support generation " + "overrides for instructions, voice, temperature, max tokens, tools, or " + "tool_choice." + if response_options_error == "unsupported_native_response_options" + else "response.create cannot reserve options while another response is active." + ) + self._emit_error(response_options_error, error_message, event_id=event.get("realtime_event_id")) + return + if event_type == "input_audio_buffer.commit": + has_pending_audio = ( + model_state.input_since_commit + or model_state.audio_buffer.has_pending() + or model_state.committed_audio_payload is not None + or realtime_validated_audio_commit + ) + if not has_pending_audio and not self.tasks.append_tasks and self.run.stream_request_id is None: + self._emit_error( + "input_audio_buffer_empty", "input_audio_buffer.commit requires a non-empty input audio buffer." + ) + return + commit_action = decide_commit_action( + CommitSnapshot( + auto_responds=self._session_auto_responds(), + speech_since_commit=model_state.speech_since_commit, + active_response_id=session.active_response_id, + overlap_speech_ms=session.overlap_speech_ms, + response_in_progress=helpers.response_in_progress(self.session, self.tasks), + playback_active=helpers.assistant_playback_active(self.session), + ) + ) + if commit_action is CommitAction.DEFER_ACTIVE_RESPONSE: + if session.overlap_speech_ms <= session.config.overlap_short_ack_ms: + self._discard_short_overlap_ack() + return + + if self._defer_commit_behind_active_response( + event, + realtime_item_id=realtime_item_id, + should_create_response=should_create_response, + precreate_response_requested=precreate_response_requested, + ): + return + if commit_action is CommitAction.START_AUTO_RESPONSE: + await self._commit_and_start_auto_response(event, realtime_item_id=realtime_item_id) + return + if event_type == "response.create": + await self._start_response_from_committed_audio() + return + if not helpers.response_in_progress(self.session, self.tasks) and await self._flush_and_submit_committed_turn( + event, + event_type=event_type, + realtime_item_id=realtime_item_id, + should_create_response=should_create_response, + ): + return + # Nothing flushed (or a response is still in progress): acknowledge the + # commit without starting a new response. + had_uncommitted_audio = ( + model_state.input_since_commit + or model_state.audio_buffer.has_pending() + or model_state.committed_audio_payload is not None + or realtime_validated_audio_commit + ) + committed = None + if event_type in {"input_audio_buffer.commit", "input.commit"}: + model_state.input_since_commit = False + model_state.speech_since_commit = False + if event_type != "response.create": + committed = ( + helpers.commit_audio_input( + session, + realtime_item_id=realtime_item_id, + transcript=event.get("transcript"), + ) + if had_uncommitted_audio + else None + ) + self.emit( + helpers.audio_committed_payload( + session, + committed=committed, + realtime_item_id=realtime_item_id, + transcript=event.get("transcript"), + ) + ) + return + if committed is not None: + if isinstance(realtime_item_id, str): + session.register_history_item(realtime_item_id, committed.message) + self.emit(helpers.input_committed_payload(session, committed, realtime_item_id=realtime_item_id)) + # NOTE(refactor): the old generic chat-completion fallback response + # (`_run_response`) is gone; a native session with a response already in + # progress only acknowledges the commit here. + + # ------------------------------------------------------------------ # + # playback.ack (moved from OmniDuplexSessionHandler) # + # ------------------------------------------------------------------ # + + +__all__ = ["DuplexSessionRunner"] diff --git a/vllm_omni/engine/duplex/turn_detection.py b/vllm_omni/engine/duplex/turn_detection.py new file mode 100644 index 000000000..6c5c84f8a --- /dev/null +++ b/vllm_omni/engine/duplex/turn_detection.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Server-side turn detection (``turn_detection: server_vad``) for the session runner. + +Moved from the Realtime input translator: the runner owns one +``ServerTurnDetector`` per session, feeds each appended chunk to it and emits +``input_audio_buffer.speech_started`` / ``speech_stopped`` from the result. +The two-phase ``session.update`` commit (prepare / commit / reject) that used +to live on the translator is expressed by :class:`PendingTurnDetectionUpdate`. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass + +import numpy as np + +from vllm_omni.engine.duplex.vad import ( + SILERO_VAD_MIN_THRESHOLD, + ServerVADUnavailableError, + SileroStreamingVAD, + SileroVADBackendProvider, + SileroVADConfig, + StreamingVADResult, +) + +__all__ = [ + "PendingTurnDetectionUpdate", + "SileroVADBackendProvider", + "ServerTurnDetector", + "ServerVADUnavailableError", + "TurnDetectionConfig", + "TurnDetectionResult", + "apply_turn_detection_result", + "configured_realtime_turn_detection", + "normalize_turn_detection_session_payload", + "validate_realtime_turn_detection", +] + + +def configured_realtime_turn_detection(session_payload: Mapping[str, object]) -> tuple[str | None, object]: + """Return ``(field_path, value)`` for the turn detection setting present in the payload.""" + field = "turn_detection" if "turn_detection" in session_payload else None + selected = session_payload.get("turn_detection") + audio_config = session_payload.get("audio") + if isinstance(audio_config, dict): + audio_input = audio_config.get("input") + if isinstance(audio_input, dict) and "turn_detection" in audio_input: + field, selected = "audio.input.turn_detection", audio_input["turn_detection"] + return field, selected + + +def _valid_vad_number(value: object, *, minimum: float, maximum: float | None = None, strict: bool = False) -> bool: + return ( + isinstance(value, int | float) + and not isinstance(value, bool) + and np.isfinite(float(value)) + and (float(value) > minimum if strict else float(value) >= minimum) + and (maximum is None or float(value) <= maximum) + ) + + +#: Every field a ``server_vad`` object may carry. Unknown keys are refused +#: rather than ignored: a misspelled tuning knob that silently does nothing is +#: worse than a rejected session, because the endpointing still *looks* applied. +_SERVER_VAD_FIELDS = frozenset( + { + "type", + "threshold", + "prefix_padding_ms", + "silence_duration_ms", + "create_response", + "interrupt_response", + "min_speech_duration_ms", + } +) + + +def validate_realtime_turn_detection(session_payload: Mapping[str, object]) -> str | None: + """Validate the ``turn_detection`` object of a Realtime session payload (None when valid).""" + field, turn_detection = configured_realtime_turn_detection(session_payload) + if field is None: + if session_payload.get("overlap_policy") == "barge_in_on_speech": + return ( + "overlap_policy='barge_in_on_speech' requires turn_detection.type='server_vad' on the Realtime endpoint" + ) + return None + if turn_detection is not None and not isinstance(turn_detection, dict): + return f"{field} must be null or an object with type='server_vad'" + if isinstance(turn_detection, dict): + if turn_detection.get("type") != "server_vad": + return f"{field}.type must be 'server_vad'" + unknown = sorted(set(turn_detection) - _SERVER_VAD_FIELDS) + if unknown: + return f"Unknown {field} field(s): {', '.join(unknown)}" + threshold = turn_detection.get("threshold", 0.5) + if not _valid_vad_number(threshold, minimum=SILERO_VAD_MIN_THRESHOLD, strict=True, maximum=1): + return f"{field}.threshold must be greater than {SILERO_VAD_MIN_THRESHOLD} and at most 1" + for name in ("prefix_padding_ms", "silence_duration_ms", "min_speech_duration_ms"): + value = turn_detection.get(name) + if value is not None and not _valid_vad_number(value, minimum=0): + return f"{field}.{name} must be a non-negative number" + if turn_detection.get("interrupt_response", True) is not True: + return ( + f"{field}.interrupt_response=false is unsupported; use turn_detection=null for model-owned listen/speak" + ) + desired_policy = "barge_in_on_speech" if turn_detection is not None else "listen_only" + overlap_policy = session_payload.get("overlap_policy") + if isinstance(overlap_policy, str) and overlap_policy != desired_policy: + return f"overlap_policy={overlap_policy!r} conflicts with turn_detection; expected {desired_policy!r}" + return None + + +@dataclass(frozen=True, slots=True) +class TurnDetectionConfig: + """Normalized ``server_vad`` configuration (all fields filled with defaults).""" + + threshold: float = 0.5 + prefix_padding_ms: int = 300 + silence_duration_ms: int = 500 + min_speech_duration_ms: int = 96 + create_response: bool = True + interrupt_response: bool = True + + @classmethod + def from_realtime(cls, turn_detection: Mapping[str, object]) -> TurnDetectionConfig: + return cls( + threshold=float(turn_detection.get("threshold", 0.5)), + prefix_padding_ms=int(turn_detection.get("prefix_padding_ms", 300)), + silence_duration_ms=int(turn_detection.get("silence_duration_ms", 500)), + min_speech_duration_ms=int(turn_detection.get("min_speech_duration_ms", 96)), + create_response=bool(turn_detection.get("create_response", True)), + interrupt_response=bool(turn_detection.get("interrupt_response", True)), + ) + + @property + def overlap_policy(self) -> str: + return "barge_in_on_speech" + + def build_detector(self, backend_provider: SileroVADBackendProvider | None = None) -> ServerTurnDetector: + return ServerTurnDetector(self, backend_provider=backend_provider) + + +def normalize_turn_detection_session_payload( + session_payload: dict[str, object], +) -> tuple[bool, TurnDetectionConfig | None]: + """Fill ``turn_detection`` defaults and derive ``overlap_policy`` in place. + + Returns ``(configured, config)`` where ``configured`` is False when the + payload does not mention turn detection at all (nothing changes), and + ``config`` is ``None`` for ``turn_detection: null`` (model-owned turns, + ``overlap_policy`` forced to ``listen_only``). + """ + field, configured = configured_realtime_turn_detection(session_payload) + if field is None: + return False, None + if configured is None: + session_payload["overlap_policy"] = "listen_only" + return True, None + assert isinstance(configured, dict) + merged = { + "type": "server_vad", + "interrupt_response": True, + "threshold": 0.5, + "prefix_padding_ms": 300, + "silence_duration_ms": 500, + "min_speech_duration_ms": 96, + **configured, + } + config = TurnDetectionConfig.from_realtime(merged) + session_payload["turn_detection"] = dict(merged) + session_payload["overlap_policy"] = config.overlap_policy + return True, config + + +@dataclass(frozen=True, slots=True) +class TurnDetectionResult: + is_speech: bool + speech_active: bool + speech_started: bool = False + speech_stopped: bool = False + speech_probability: float = 0.0 + audio_start_ms: int | None = None + audio_end_ms: int | None = None + #: A ``speech_stopped`` under ``server_vad`` ends the user turn. + should_commit: bool = False + #: Whether the auto-commit should also request a response (``create_response``). + create_response: bool = True + + def as_payload(self) -> dict[str, object]: + return { + "backend": "silero", + "is_speech": self.is_speech, + "speech_active": self.speech_active, + "speech_started": self.speech_started, + "speech_stopped": self.speech_stopped, + "speech_probability": self.speech_probability, + } + + +class ServerTurnDetector: + """One Silero-backed streaming detector per session.""" + + def __init__( + self, + config: TurnDetectionConfig, + *, + backend_provider: SileroVADBackendProvider | None = None, + ) -> None: + self.config = config + self._vad = SileroStreamingVAD( + SileroVADConfig( + threshold=float(config.threshold), + prefix_padding_ms=int(config.prefix_padding_ms), + silence_duration_ms=int(config.silence_duration_ms), + min_speech_duration_ms=max(32, int(config.min_speech_duration_ms)), + ), + backend_provider=backend_provider, + ) + + @property + def vad(self) -> SileroStreamingVAD: + return self._vad + + def reset(self) -> None: + self._vad.reset() + + def process( + self, + base64_audio: str, + *, + fmt: str, + sample_rate_hz: int | None, + audio_end_ms: int | None = None, + ) -> TurnDetectionResult: + """Score one appended chunk. Raises ``ServerVADUnavailableError`` / ``ValueError`` like the VAD.""" + result: StreamingVADResult = self._vad.process_base64(base64_audio, fmt=fmt, sample_rate_hz=sample_rate_hz) + stop_ms = result.speech_end_ms if result.speech_end_ms is not None else audio_end_ms + return TurnDetectionResult( + is_speech=result.is_speech, + speech_active=result.speech_active, + speech_started=result.speech_started, + speech_stopped=result.speech_stopped, + speech_probability=result.speech_probability, + audio_start_ms=result.speech_start_ms, + audio_end_ms=stop_ms, + should_commit=bool(result.speech_stopped and not result.speech_active), + create_response=self.config.create_response, + ) + + +def apply_turn_detection_result(payload: dict[str, object], result: TurnDetectionResult) -> None: + """Merge a detector result onto an internal ``input_audio_buffer.append`` payload. + + Mirrors what the old translator attached: ``is_speech`` from the detector, + the ``vad`` hint block, and ``force_listen`` while speech is active. + """ + payload["is_speech"] = result.is_speech + payload["vad"] = result.as_payload() + if result.speech_active: + payload["force_listen"] = True + if result.audio_start_ms is not None and result.speech_started: + payload["audio_start_ms"] = result.audio_start_ms + if result.audio_end_ms is not None and result.speech_stopped: + payload["audio_end_ms"] = result.audio_end_ms + + +@dataclass +class PendingTurnDetectionUpdate: + """A ``session.update`` turn-detection change awaiting the engine ACK (two-phase).""" + + config: TurnDetectionConfig | None + detector: ServerTurnDetector | None + + @classmethod + def prepare( + cls, + session_payload: dict[str, object], + *, + backend_provider: SileroVADBackendProvider | None = None, + ) -> PendingTurnDetectionUpdate | None: + """Normalize the payload in place and stage the new detector; None when not configured.""" + configured, config = normalize_turn_detection_session_payload(session_payload) + if not configured: + return None + detector = config.build_detector(backend_provider) if config is not None else None + return cls(config=config, detector=detector) + + def commit( + self, current: ServerTurnDetector | None + ) -> tuple[TurnDetectionConfig | None, ServerTurnDetector | None]: + """Adopt the staged detector; the replaced one is reset. Returns the new (config, detector).""" + if current is not None: + current.reset() + return self.config, self.detector + + def reject(self) -> None: + if self.detector is not None: + self.detector.reset() diff --git a/vllm_omni/engine/duplex/vad.py b/vllm_omni/engine/duplex/vad.py new file mode 100644 index 000000000..ad3c26820 --- /dev/null +++ b/vllm_omni/engine/duplex/vad.py @@ -0,0 +1,488 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Silero server VAD for the engine-resident duplex session. + +The detector keeps the split upstream's serving-side ``server_vad`` module +used, because the split is what makes one model serve many sessions: a +**backend** scores one 512-sample frame and is otherwise stateless, and the +per-stream state -- the partial frame, the model state, the endpoint counters -- +belongs to :class:`SileroStreamingVAD`, one per session. + +Two backends. :class:`SileroVADBackend` runs the pinned Silero v6.2 ONNX graph on +CPU and is shared process-wide, which is why its state is passed in and out +rather than held on the session object. :class:`TorchSileroBackend` loads the +same model through the ``silero-vad`` package for environments that have no +local ONNX artifact; torch keeps its state inside the module, so that one is per +session. + +Endpointing follows Silero v6.2's streaming hysteresis: activation uses the +configured threshold, while a turn can only *end* on frames below +``max(threshold - 0.15, 0.01)``. Once a silence candidate exists, louder frames +keep the elapsed-silence clock running but cannot themselves close the turn. +""" + +from __future__ import annotations + +import binascii +import hashlib +import threading +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Protocol + +import numpy as np +import pybase64 as base64 +from vllm.logger import init_logger + +logger = init_logger(__name__) + +SILERO_VAD_REPO_ID = "istupakov/silero-vad-onnx" +SILERO_VAD_REVISION = "8b14476858ef240c50b3884bb38cc67290c1cc70" +SILERO_VAD_FILENAME = "silero_vad.onnx" +SILERO_VAD_SHA256 = "1a153a22f4509e292a94e67d6f9b85e8deb25b4988682b7e174c65279d8788e3" +SILERO_VAD_MIN_THRESHOLD = 0.15 + +_SAMPLE_RATE_HZ = 16_000 +_FRAME_SAMPLES = 512 + + +class ServerVADUnavailableError(RuntimeError): + pass + + +@dataclass(frozen=True, slots=True) +class SileroVADConfig: + threshold: float = 0.5 + prefix_padding_ms: int = 300 + silence_duration_ms: int = 500 + min_speech_duration_ms: int = 96 + + +@dataclass(frozen=True, slots=True) +class StreamingVADResult: + is_speech: bool + speech_active: bool + speech_started: bool = False + speech_stopped: bool = False + speech_probability: float = 0.0 + speech_start_ms: int | None = None + speech_end_ms: int | None = None + + +# --------------------------------------------------------------------------- # +# Backends # +# --------------------------------------------------------------------------- # + + +class SpeechDetectorBackend(Protocol): + """Scores one frame. Per-stream state is created here but owned by the caller.""" + + frame_samples: int + + def new_state(self) -> object: ... + + def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: ... + + +@dataclass(frozen=True, slots=True) +class _SileroVADState: + model_state: np.ndarray + context: np.ndarray + + +class SileroVADBackend: + """Shared ONNX Runtime Silero v6.2 detector running on CPU.""" + + sample_rate_hz = _SAMPLE_RATE_HZ + frame_samples = _FRAME_SAMPLES + context_samples = 64 + model_state_shape = (2, 1, 128) + + def __init__(self, model_path: str | Path) -> None: + try: + import onnxruntime as ort + except ImportError as exc: # pragma: no cover - platform packaging supplies ORT. + raise ServerVADUnavailableError("server_vad requires ONNX Runtime") from exc + + session_options = ort.SessionOptions() + session_options.inter_op_num_threads = 1 + session_options.intra_op_num_threads = 1 + + self.model_path = Path(model_path) + self._session = ort.InferenceSession( + str(self.model_path), + providers=["CPUExecutionProvider"], + sess_options=session_options, + ) + input_names = {item.name for item in self._session.get_inputs()} + if not {"input", "sr", "state"} <= input_names: + raise ServerVADUnavailableError(f"Unsupported Silero ONNX input contract: {sorted(input_names)}") + self._warm_up() + + def _warm_up(self) -> None: + self.infer(np.zeros(self.frame_samples, dtype=np.float32), self.new_state()) + + def new_state(self) -> _SileroVADState: + return _SileroVADState( + model_state=np.zeros(self.model_state_shape, dtype=np.float32), + context=np.zeros((1, self.context_samples), dtype=np.float32), + ) + + def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: + if not isinstance(state, _SileroVADState): + raise TypeError("Silero detector state must be created by SileroVADBackend.new_state()") + model_state = np.ascontiguousarray(state.model_state, dtype=np.float32) + context = np.ascontiguousarray(state.context, dtype=np.float32) + if model_state.shape != self.model_state_shape: + raise ValueError(f"Silero model state must have shape {self.model_state_shape}, got {model_state.shape}") + expected_context_shape = (1, self.context_samples) + if context.shape != expected_context_shape: + raise ValueError(f"Silero context must have shape {expected_context_shape}, got {context.shape}") + + audio = np.ascontiguousarray(frame, dtype=np.float32).reshape(1, -1) + if audio.shape[1] != self.frame_samples: + raise ValueError( + f"Silero detector frame must contain exactly {self.frame_samples} samples, got {audio.shape[1]}" + ) + model_input = np.concatenate((context, audio), axis=1) + # ONNX Runtime permits concurrent Run calls on one CPU session. Stream + # state is explicit, so independent sessions can safely share this one. + output = self._session.run( + None, + { + "input": model_input, + "state": model_state, + "sr": np.asarray(self.sample_rate_hz, dtype=np.int64), + }, + ) + if len(output) < 2: + raise ServerVADUnavailableError("Silero ONNX model did not return probability and model state") + probability = float(np.asarray(output[0]).reshape(-1)[0]) + next_state = _SileroVADState( + model_state=np.ascontiguousarray(output[1], dtype=np.float32), + context=np.ascontiguousarray(model_input[:, -self.context_samples :], dtype=np.float32), + ) + return probability, next_state + + +class TorchSileroBackend: + """The same model through the ``silero-vad`` package, for hosts with no ONNX artifact. + + Torch keeps the recurrent state inside the module, so unlike the ONNX + backend this one cannot be shared: each session gets its own instance and + ``new_state`` only resets it. + """ + + sample_rate_hz = _SAMPLE_RATE_HZ + frame_samples = _FRAME_SAMPLES + + def __init__(self) -> None: + try: + import torch + from silero_vad import load_silero_vad + except ImportError as exc: + raise ServerVADUnavailableError( + "server_vad needs either a local Silero ONNX artifact (see " + "duplex_session.server_vad_model_path) or the optional 'silero-vad' " + "package; install vllm-omni[server-vad]" + ) from exc + self._torch = torch + self._model = load_silero_vad(onnx=False) + + def new_state(self) -> object: + reset_states = getattr(self._model, "reset_states", None) + if callable(reset_states): + reset_states() + return None + + def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: + tensor = self._torch.from_numpy(np.ascontiguousarray(frame, dtype=np.float32)) + return float(self._model(tensor, self.sample_rate_hz).item()), state + + +class SileroVADBackendProvider: + """Resolve, verify and load one detector backend per engine process. + + ONNX is preferred and shared; the torch package is the fallback and is built + per call because its state is internal. + """ + + def __init__(self, *, model_path: str | None = None) -> None: + self.model_path = model_path + self._backend: SileroVADBackend | None = None + self._lock = threading.Lock() + + def get(self) -> SpeechDetectorBackend: + """The shared ONNX backend, or a per-session torch one when it cannot be built. + + An explicitly configured ``server_vad_model_path`` never falls back: the + operator named that artifact, so a missing file or a missing ONNX Runtime + is their error, not something to paper over with a different model. + """ + if self._backend is not None: + return self._backend + with self._lock: + if self._backend is not None: + return self._backend + path = self._resolve_local_artifact() + if path is None: + # No pinned artifact: the torch package keeps state internally, + # so this instance belongs to one session only. + return TorchSileroBackend() + self._verify_checksum(path) + try: + self._backend = SileroVADBackend(path) + except ServerVADUnavailableError: + if self.model_path: + raise + logger.warning( + "Silero ONNX artifact %s found but ONNX Runtime is unavailable; " + "falling back to the torch 'silero-vad' package (one model per session)", + path, + ) + return TorchSileroBackend() + return self._backend + + def _resolve_local_artifact(self) -> Path | None: + if self.model_path: + path = Path(self.model_path).expanduser() + if path.is_file(): + return path + # Explicitly configured: a missing file is an error, not a fallback. + raise ServerVADUnavailableError(f"Configured Silero VAD model does not exist: {path}") + try: + from vllm.transformers_utils.repo_utils import try_get_local_file + + cached = try_get_local_file( + model=SILERO_VAD_REPO_ID, + file_name=SILERO_VAD_FILENAME, + revision=SILERO_VAD_REVISION, + ) + except Exception: + return None + return cached if isinstance(cached, Path) and cached.is_file() else None + + def _verify_checksum(self, path: Path) -> None: + digest = hashlib.sha256(path.read_bytes()).hexdigest() + if digest != SILERO_VAD_SHA256: + raise ServerVADUnavailableError( + f"Silero VAD checksum mismatch for {path}: expected {SILERO_VAD_SHA256}, got {digest}" + ) + + +# --------------------------------------------------------------------------- # +# Per-session streaming detector # +# --------------------------------------------------------------------------- # + + +class SileroStreamingVAD: + _SAMPLE_RATE_HZ = _SAMPLE_RATE_HZ + _WINDOW_SAMPLES = _FRAME_SAMPLES + + def __init__( + self, + config: SileroVADConfig, + *, + backend: SpeechDetectorBackend | None = None, + backend_provider: SileroVADBackendProvider | None = None, + frame_scorer: Callable[[np.ndarray], float] | None = None, + ) -> None: + self.config = config + self._explicit_backend = backend + self._provider = backend_provider + self._frame_scorer = frame_scorer + self._backend: SpeechDetectorBackend | None = backend + self._detector_state: object = None + self._clear_stream_state() + + def _clear_stream_state(self) -> None: + self._pending = np.empty(0, dtype=np.float32) + self._speech_active = False + self._candidate_samples = 0 + self._candidate_start_sample: int | None = None + self._silence_samples = 0 + self._processed_samples = 0 + self._stream_start_ms = 0 + self._resample_rate_hz: int | None = None + self._resample_remainder = 0 + + @property + def scratch_bytes(self) -> int: + """Bytes this detector holds between chunks (the partial frame only).""" + return int(self._pending.nbytes) + + @property + def speech_active(self) -> bool: + return self._speech_active + + def reset(self) -> None: + """Drop stream state while preserving the session-relative audio clock. + + The clock survives so ``speech_start_ms`` after a barge-in still refers + to the session's timeline; ``_stream_start_ms`` then stops the prefix + padding reaching back into audio that was discarded. + """ + stream_start_ms = round(self._processed_samples * 1000 / self._SAMPLE_RATE_HZ) + processed = self._processed_samples + self._clear_stream_state() + self._processed_samples = processed + self._stream_start_ms = stream_start_ms + if self._backend is not None: + self._detector_state = self._backend.new_state() + + # ------------------------------------------------------------------ # + # Input # + # ------------------------------------------------------------------ # + + def process_base64( + self, + audio: object, + *, + fmt: object, + sample_rate_hz: object, + ) -> StreamingVADResult: + if not isinstance(audio, str): + raise ValueError("Silero server VAD requires base64 audio") + try: + raw = base64.b64decode(audio, validate=True) + except (binascii.Error, ValueError) as exc: + raise ValueError("Silero server VAD received invalid base64 audio") from exc + rate = int(sample_rate_hz) if isinstance(sample_rate_hz, int | float) else self._SAMPLE_RATE_HZ + if rate <= 0: + raise ValueError("Silero server VAD requires a positive sample rate") + + if fmt == "pcm_f32le": + if len(raw) % 4: + raise ValueError("Silero server VAD received an incomplete pcm_f32le frame") + samples = np.frombuffer(raw, dtype=" np.ndarray: + """Resample to 16 kHz, carrying the fractional sample-count remainder. + + A fallback path: the session runner already normalises input to 16 kHz + through ``convert_input_audio_with_rate`` before the detector sees it, + so this runs only for callers that feed the VAD directly. + + Carrying the remainder keeps the *number* of output samples exact, so + frame boundaries cannot drift over a long stream. The interpolation + itself is per chunk and has no filter state, so sample values near a + chunk edge differ slightly from resampling the same audio in one go -- + enough for the frame grid to stay put, not enough to call it a + streaming resampler. + """ + if rate == self._SAMPLE_RATE_HZ: + self._resample_rate_hz = None + self._resample_remainder = 0 + return samples + if self._resample_rate_hz is not None and rate != self._resample_rate_hz: + raise ValueError("server_vad input sample rate cannot change within a continuous audio stream") + self._resample_rate_hz = rate + numerator = samples.size * self._SAMPLE_RATE_HZ + self._resample_remainder + target_size, self._resample_remainder = divmod(numerator, rate) + if target_size == 0: + return np.empty(0, dtype=np.float32) + if samples.size == 1: + return np.full(target_size, samples[0], dtype=np.float32) + source_x = np.linspace(0.0, 1.0, num=samples.size, endpoint=True) + target_x = np.linspace(0.0, 1.0, num=target_size, endpoint=True) + return np.interp(target_x, source_x, samples).astype(np.float32) + + # ------------------------------------------------------------------ # + # Endpointing # + # ------------------------------------------------------------------ # + + def process(self, samples: np.ndarray) -> StreamingVADResult: + samples = np.asarray(samples, dtype=np.float32).reshape(-1) + if samples.size: + self._pending = np.concatenate((self._pending, samples)) + + started = False + stopped = False + start_ms: int | None = None + end_ms: int | None = None + max_probability = 0.0 + contained_speech = self._speech_active + min_speech_samples = max(1, round(self.config.min_speech_duration_ms * self._SAMPLE_RATE_HZ / 1000)) + min_silence_samples = max(1, round(self.config.silence_duration_ms * self._SAMPLE_RATE_HZ / 1000)) + negative_threshold = max(self.config.threshold - SILERO_VAD_MIN_THRESHOLD, 0.01) + + pending = self._pending + offset = 0 + while pending.size - offset >= self._WINDOW_SAMPLES: + frame = np.ascontiguousarray(pending[offset : offset + self._WINDOW_SAMPLES], dtype=np.float32) + offset += self._WINDOW_SAMPLES + frame_start = self._processed_samples + self._processed_samples += self._WINDOW_SAMPLES + probability = min(1.0, max(0.0, self._score_frame(frame))) + max_probability = max(max_probability, probability) + + if self._speech_active: + contained_speech = True + below_negative_threshold = probability < negative_threshold + if not below_negative_threshold and self._silence_samples == 0: + continue + # A silence candidate is running: louder frames keep its clock + # moving but cannot themselves close the turn. + self._silence_samples += self._WINDOW_SAMPLES + if not below_negative_threshold or self._silence_samples < min_silence_samples: + continue + self._speech_active = False + self._silence_samples = 0 + stopped = True + # OpenAI defines audio_end_ms as the end of the audio sent to + # the model, trailing silence included. + end_ms = max(0, round(self._processed_samples * 1000 / self._SAMPLE_RATE_HZ)) + continue + + if probability >= self.config.threshold: + if self._candidate_samples == 0: + self._candidate_start_sample = frame_start + self._candidate_samples += self._WINDOW_SAMPLES + if self._candidate_samples >= min_speech_samples: + candidate_start = self._candidate_start_sample or 0 + detected_start_ms = round(candidate_start * 1000 / self._SAMPLE_RATE_HZ) + self._speech_active = True + self._candidate_samples = self._silence_samples = 0 + self._candidate_start_sample = None + contained_speech = True + started = True + # Prefix audio from before the last reset is gone. + start_ms = max( + self._stream_start_ms, + max(0, detected_start_ms - self.config.prefix_padding_ms), + ) + else: + self._candidate_samples = 0 + self._candidate_start_sample = None + + # Copy the residual once, after the frames are scored, so a large input + # chunk is neither pinned by a view nor re-copied for every frame. + self._pending = pending[offset:].copy() + + return StreamingVADResult( + contained_speech, self._speech_active, started, stopped, max_probability, start_ms, end_ms + ) + + def _score_frame(self, frame: np.ndarray) -> float: + if self._frame_scorer is not None: + return float(self._frame_scorer(frame)) + backend = self._require_backend() + probability, self._detector_state = backend.infer(frame, self._detector_state) + return float(probability) + + def _require_backend(self) -> SpeechDetectorBackend: + if self._backend is None: + provider = self._provider or SileroVADBackendProvider() + self._backend = provider.get() + self._detector_state = self._backend.new_state() + return self._backend diff --git a/vllm_omni/engine/duplex_omni_engine.py b/vllm_omni/engine/duplex_omni_engine.py new file mode 100644 index 000000000..56ca3e352 --- /dev/null +++ b/vllm_omni/engine/duplex_omni_engine.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""DuplexOmniEngine: the duplex engine sibling of ``AsyncOmniEngine``. + +It is deliberately small. The engine base owns the stage processes, the +orchestrator thread and the correlated RPC transport; this class only knows +how duplex session messages enter the engine (open/close/resume/touch through +correlated RPC, session commands one-way) and which orchestrator to build +(``DuplexOrchestrator`` with the model plugin loaded). +""" + +from __future__ import annotations + +import asyncio +import queue +import uuid +from typing import TYPE_CHECKING, Any + +from vllm.logger import init_logger + +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig +from vllm_omni.engine.duplex.messages import ( + CloseDuplexSessionMessage, + DuplexControlResultMessage, + DuplexSessionCommandMessage, + DuplexSessionError, + OpenDuplexSessionMessage, + ResumeDuplexSessionMessage, + TouchDuplexSessionMessage, +) +from vllm_omni.engine.duplex.plugin import DuplexModelPlugin, EncodeAudio, load_duplex_plugin +from vllm_omni.engine.messages import EngineQueueMessage +from vllm_omni.engine.orchestrator import OrchestratorBase + +if TYPE_CHECKING: + from vllm_omni.engine.duplex.commands import DuplexCommand + +logger = init_logger(__name__) + +_DEFAULT_CONTROL_TIMEOUT_S = 10.0 +_COMMAND_PUT_TIMEOUT_S = 30.0 + + +class DuplexOmniEngine(AsyncOmniEngine): + """Engine for full-duplex models; sessions run inside ``DuplexOrchestrator``.""" + + plugin: DuplexModelPlugin | None = None + #: Set by ``_validate_deployment`` from the deploy config before any stage starts. + duplex_session_config: DuplexSessionRuntimeConfig + + # Any: **kwargs is forwarded verbatim to ``OmniEngineBase.__init__``. + def __init__(self, model: str, *, duplex_audio_encoder: EncodeAudio, **kwargs: Any) -> None: + # Injected by ``DuplexOmni`` (entrypoint layer): the engine package + # knows how to hand model audio to the plugin, not how to encode it. + self._audio_encoder = duplex_audio_encoder + super().__init__(model, **kwargs) + + # ---- orchestrator construction (orchestrator thread) ---- + + def _validate_deployment(self) -> None: + """Load the plugin and check the scheduler contract before any stage process starts.""" + pipeline_config = self.pipeline_config + plugin_path = getattr(pipeline_config, "duplex_plugin", None) if pipeline_config is not None else None + if pipeline_config is None: + raise ValueError( + f"{self.model!r}: the model pipeline could not be resolved (models with custom HF code " + "such as MiniCPM-o need trust_remote_code=True)" + ) + if not plugin_path: + raise ValueError(f"{self.model!r} is not a duplex model: the pipeline declares no duplex_plugin") + deploy_config = self.deploy_config + if deploy_config is None: + raise ValueError("DuplexOmniEngine requires a deploy config with session_mode: duplex (none resolved)") + if getattr(deploy_config, "session_mode", "turn") != "duplex": + raise ValueError( + "DuplexOmniEngine requires a deploy config with session_mode: duplex " + f"(got {getattr(deploy_config, 'session_mode', None)!r})" + ) + self.duplex_session_config = deploy_config.duplex_session + self.plugin = load_duplex_plugin(plugin_path, self._audio_encoder) + + # Any: signature of the ``OmniEngineBase._create_orchestrator`` seam it overrides. + def _create_orchestrator(self, **orchestrator_kwargs: Any) -> OrchestratorBase: + from vllm_omni.engine.duplex_orchestrator import DuplexOrchestrator + + assert self.plugin is not None, "_validate_deployment() must run before the orchestrator is created" + stage0_vllm_config = self.stage_pools[0].stage_vllm_config if self.stage_pools else None + model_config = getattr(stage0_vllm_config, "model_config", None) + return DuplexOrchestrator( + plugin=self.plugin, + duplex_session_config=self.duplex_session_config, + model_config=model_config, + **orchestrator_kwargs, + ) + + # ---- deployment facts ---- + + @property + def duplex_capabilities(self) -> DuplexCapabilities: + if self.plugin is None: + raise RuntimeError("duplex plugin is not loaded") + return self.plugin.capabilities(max_sessions=self.duplex_session_config.max_sessions) + + # ---- session message surface ---- + + def _execute_control( + self, + message: EngineQueueMessage, + *, + control_id: str, + operation: str, + session_id: str, + timeout: float | None, + ) -> DuplexControlResultMessage: + try: + result = self.rpc_client.execute( + ("duplex", control_id), + message, + timeout=timeout, + timeout_message=f"duplex {operation} timed out for session {session_id}", + # Block rather than surface a raw queue.Full: the handlers below + # promise a typed DuplexSessionError, and a momentarily full + # request queue is backpressure, not a failed control op. + block_on_submit=True, + ) + except TimeoutError as exc: + raise DuplexSessionError(str(exc), code="timeout", retryable=True, session_id=session_id) from exc + except RuntimeError as exc: + # The RPC router is closed or the orchestrator reported a terminal error. + raise DuplexSessionError(str(exc), code="engine_error", session_id=session_id) from exc + if not isinstance(result, DuplexControlResultMessage): + raise DuplexSessionError( + f"unexpected duplex control result: {type(result).__name__}", + code="internal_error", + session_id=session_id, + ) + if not result.ok: + raise DuplexSessionError( + result.error_message or f"duplex {operation} failed", + code=result.error_code or "internal_error", + retryable=result.error_retryable, + session_id=session_id, + ) + return result + + def _open_session( + self, + session_id: str, + session_config: DuplexSessionConfig, + *, + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexControlResultMessage: + control_id = uuid.uuid4().hex + return self._execute_control( + OpenDuplexSessionMessage( + control_id=control_id, + session_id=session_id, + session_config=session_config, + ), + control_id=control_id, + operation="open", + session_id=session_id, + timeout=timeout, + ) + + async def open_session_async( + self, + session_id: str, + session_config: DuplexSessionConfig, + *, + 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)) + + def _close_session( + self, + session_id: str, + *, + reason: str = "client_close", + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexControlResultMessage: + control_id = uuid.uuid4().hex + return self._execute_control( + CloseDuplexSessionMessage( + control_id=control_id, + session_id=session_id, + reason=reason, + ), + control_id=control_id, + operation="close", + session_id=session_id, + timeout=timeout, + ) + + async def close_session_async( + self, + session_id: str, + *, + reason: str = "client_close", + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexControlResultMessage: + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + None, + lambda: self._close_session(session_id, reason=reason, timeout=timeout), + ) + + def _resume_session( + self, + session_id: str, + *, + expected_lease_generation: int, + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexControlResultMessage: + control_id = uuid.uuid4().hex + return self._execute_control( + ResumeDuplexSessionMessage( + control_id=control_id, + session_id=session_id, + expected_lease_generation=expected_lease_generation, + ), + control_id=control_id, + operation="resume", + session_id=session_id, + timeout=timeout, + ) + + async def resume_session_async( + self, + session_id: str, + *, + expected_lease_generation: int, + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexControlResultMessage: + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + None, + lambda: self._resume_session( + session_id, + expected_lease_generation=expected_lease_generation, + timeout=timeout, + ), + ) + + def _touch_session( + self, + session_id: str, + *, + activity: str, + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexControlResultMessage: + control_id = uuid.uuid4().hex + return self._execute_control( + TouchDuplexSessionMessage( + control_id=control_id, + session_id=session_id, + activity=activity, + ), + control_id=control_id, + operation="touch", + session_id=session_id, + timeout=timeout, + ) + + async def touch_session_async( + self, + session_id: str, + *, + activity: str, + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexControlResultMessage: + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + None, + lambda: self._touch_session(session_id, activity=activity, timeout=timeout), + ) + + def _submit_command(self, session_id: str, command: DuplexCommand) -> None: + """One-way: enqueue a session command in caller order (blocks only on queue backpressure).""" + if not self.is_alive(): + raise DuplexSessionError("engine is not alive", code="engine_dead", session_id=session_id) + message = DuplexSessionCommandMessage(session_id=session_id, command=command) + try: + self.request_queue.sync_q.put(message, timeout=_COMMAND_PUT_TIMEOUT_S) + except queue.Full as exc: + raise DuplexSessionError( + "engine request queue is full", code="engine_backpressure", retryable=True, session_id=session_id + ) from exc + + async def submit_command_async(self, session_id: str, command: DuplexCommand) -> None: + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, lambda: self._submit_command(session_id, command)) + + +__all__ = ["DuplexOmniEngine"] diff --git a/vllm_omni/engine/duplex_orchestrator.py b/vllm_omni/engine/duplex_orchestrator.py new file mode 100644 index 000000000..09a767fb1 --- /dev/null +++ b/vllm_omni/engine/duplex_orchestrator.py @@ -0,0 +1,352 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""DuplexOrchestrator: the duplex sibling of ``Orchestrator``. + +It hosts the ``DuplexSessionManager`` (admission, leases, session runners) on +the orchestrator loop and implements the ``DuplexStagePort`` the runners use +to submit resumable Stage0 requests. Generic stage forwarding, prewarm and +cleanup stay in ``OrchestratorBase``; this class only fills the template seams +and applies the session-owned policy. +""" + +from __future__ import annotations + +import time as _time +from collections.abc import Coroutine +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +from vllm.logger import init_logger + +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.engine.duplex.contracts import ( + DuplexFence, + DuplexOutputContext, + DuplexRequestIdentity, + DuplexStagePort, + DuplexStageRequestContext, + DuplexStageSubmission, + DuplexStageSubmissionResult, +) +from vllm_omni.engine.duplex.plugin import DuplexModelPlugin +from vllm_omni.engine.duplex.session.engine_session import DuplexFenceMismatchError +from vllm_omni.engine.duplex.session.manager import DuplexSessionManager +from vllm_omni.engine.messages import EngineQueueMessage, OutputMessage +from vllm_omni.engine.orchestrator import ( + Orchestrator, + OrchestratorRequestState, + build_engine_core_request_from_tokens, +) + +if TYPE_CHECKING: + from vllm.config import ModelConfig + from vllm.outputs import RequestOutput + + from vllm_omni.metrics.stats import StageRequestStats + +logger = init_logger(__name__) + + +@dataclass +class DuplexOrchestratorRequestState(OrchestratorRequestState): + """Request bookkeeping for a session-owned (resumable duplex) stage request.""" + + session_id: str = "" + fence: DuplexFence | None = None + stage_fences: dict[int, DuplexFence] = field(default_factory=dict) + config_generation: int = -1 + + +class DuplexOrchestrator(Orchestrator, DuplexStagePort): + """Stage management for a duplex deployment; owns one ``DuplexSessionManager``. + + Extends the turn-based orchestrator rather than sitting beside it, so one + engine serves both a duplex session and an ordinary request. The + dependency runs duplex -> turn-based and never the reverse, which is what + keeps ``Orchestrator`` free of duplex code. + """ + + def __init__( + self, + # Any: ``*args`` / ``**kwargs`` are forwarded verbatim to ``OrchestratorBase.__init__``. + *args: Any, + plugin: DuplexModelPlugin, + duplex_session_config: DuplexSessionRuntimeConfig, + model_config: ModelConfig | None = None, + **kwargs: Any, + ) -> None: + super().__init__(*args, **kwargs) + self.plugin = plugin + self.duplex_session_config = duplex_session_config + self.session_manager = DuplexSessionManager( + plugin=plugin, + stage_port=self, + output_sink=self.output_async_queue, + result_sink=self.rpc_async_queue, + runtime_config=self.duplex_session_config, + model_config=model_config, + ) + + # ------------------------------------------------------------------ # + # Template seams # + # ------------------------------------------------------------------ # + + async def _dispatch_message(self, msg: EngineQueueMessage) -> bool: + if self.session_manager.accepts(msg): + self.session_manager.dispatch(msg) + return True + # Not a session message: it is an ordinary turn-based request. + return await super()._dispatch_message(msg) + + # Any: ``Coroutine``'s send/yield parameters, as in the ``OrchestratorBase`` seam. + def _background_tasks(self) -> list[Coroutine[Any, Any, None]]: + return [self.session_manager.reaper_loop(self._shutdown_event)] + + async def _shutdown_extensions(self) -> None: + await self.session_manager.shutdown() + + def _on_stage_submitted( + self, + stage_id: int, + request_id: str, + replica_id: int, + req_state: OrchestratorRequestState, + ) -> None: + del replica_id + if not isinstance(req_state, DuplexOrchestratorRequestState) or req_state.fence is None: + return + runner = self.session_manager.runner_for_request_id(request_id) + if runner is None: + return + fence = req_state.fence + req_state.stage_fences[stage_id] = fence + try: + runner.session.bind_stage_request(stage_id, request_id, fence=fence) + except DuplexFenceMismatchError: + # The session already advanced past this request's epoch (cancel + # raced the submit); the cancel path aborts the stale request id. + logger.debug("[DuplexOrchestrator] stale stage binding ignored for %s stage-%s", request_id, stage_id) + return + req_state.stage_submit_ts[stage_id] = _time.time() + self._register_running_request(req_state) + self.session_manager.register_request(request_id, req_state.session_id) + + async def _intercept_stage_output( + self, + stage_id: int, + replica_id: int, + output: RequestOutput, + req_state: OrchestratorRequestState, + stage_metrics: StageRequestStats | None, + submit_ts: float | None, + ) -> bool: + del replica_id, submit_ts + if not isinstance(req_state, DuplexOrchestratorRequestState) or req_state.fence is None: + return False + request_id = output.request_id + runner = self.session_manager.runner_for_request_id(request_id) + if runner is None: + # Session gone: nothing may forward or reach a client. + return True + segment = req_state.streaming.segment(stage_id) + context = DuplexOutputContext( + identity=DuplexRequestIdentity( + session_id=req_state.session_id, + fence=req_state.stage_fences.get(stage_id, req_state.fence), + ), + final_stage_id=req_state.final_stage_id, + segment_finished=req_state.streaming.enabled and segment.finished, + segment_token_ids=tuple(segment.token_ids), + segment_output_metadata=segment.output_metadata, + ) + return runner.on_stage_output( + stage_id, + output, + stage_metrics, + request_id=request_id, + context=context, + ) + + async def _handle_forward_failure( + self, + req_id: str, + next_stage_id: int, + req_state: OrchestratorRequestState, + exc: BaseException, + ) -> bool: + if not req_state.session_owned: + return False + runner = self.session_manager.runner_for_request_id(req_id) + if runner is not None: + runner.on_stage_failure(next_stage_id, exc) + await self._cleanup_request_ids( + [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + abort=True, + release_owners=True, + ) + return True + + async def _cleanup_request_ids( + self, + request_ids: list[str], + *, + abort: bool = False, + release_owners: bool = False, + ) -> list[OutputMessage]: + if not request_ids: + return [] + cleanup_ids = list(dict.fromkeys(request_ids)) + closing_session_ids: list[str] = [] + if release_owners: + closed_sessions = self.session_manager.close_sessions_for_request_ids( + cleanup_ids, + abort=abort, + cleanup_in_progress=True, + ) + closing_session_ids.extend(closed_sessions) + for session_id, stale_request_ids in closed_sessions.items(): + logger.info( + "[DuplexOrchestrator] closed duplex session %s while cleaning failed request ids %s", + session_id, + stale_request_ids, + ) + cleanup_ids.extend(stale_request_ids) + cleanup_ids = list(dict.fromkeys(cleanup_ids)) + try: + outputs = await super()._cleanup_request_ids(cleanup_ids, abort=abort) + except BaseException: + if closing_session_ids: + self.session_manager.defer_request_cleanups(closing_session_ids) + raise + if closing_session_ids: + self.session_manager.finalize_closed_sessions(closing_session_ids) + for request_id in cleanup_ids: + self.session_manager.unregister_request(request_id) + return outputs + + # ------------------------------------------------------------------ # + # DuplexStagePort # + # ------------------------------------------------------------------ # + + @property + def stage_count(self) -> int: + return len(self.stage_pools) + + def sampling_defaults(self) -> tuple[object, ...]: + defaults = [] + for pool in self.stage_pools: + client = pool.stage_client + if client is None: + # Every replica of this stage was evicted (engine core died); a + # session cannot be admitted until the deployment is restarted. + raise RuntimeError(f"stage {pool.stage_id} has no live replica") + defaults.append(client.default_sampling_params) + return tuple(defaults) + + @staticmethod + def _sync_bridge_state( + request_state: OrchestratorRequestState, + context: DuplexStageRequestContext, + ) -> None: + duplex_state = request_state.streaming.bridge_states.setdefault("duplex", {}) + if not isinstance(duplex_state, dict): + duplex_state = {} + request_state.streaming.bridge_states["duplex"] = duplex_state + previous_epoch = duplex_state.get("epoch") + current_model_turn_id = duplex_state.get("model_turn_id") + if ( + not isinstance(current_model_turn_id, int) + or previous_epoch != context.fence.epoch + or current_model_turn_id < context.fence.turn_id + ): + # A safety boundary can close a response before the model emits + # its normal turn_eos. Catch up the engine-owned identity when the + # next fenced append starts. + duplex_state["model_turn_id"] = context.fence.turn_id + duplex_state.update( + { + "session_id": context.session_id, + "fence": context.fence, + "epoch": context.fence.epoch, + "turn_id": context.fence.turn_id, + "session_config": dict(context.session_config), + "runtime_config": dict(context.runtime_config), + } + ) + + def ensure_request(self, context: DuplexStageRequestContext) -> None: + request_state = self.request_states.get(context.request_id) + if request_state is None: + request_state = DuplexOrchestratorRequestState( + request_id=context.request_id, + prompt=None, + sampling_params_list=list(context.sampling_params), + final_stage_id=context.final_stage_id, + session_owned=True, + session_id=context.session_id, + fence=context.fence, + config_generation=context.config_generation, + ) + request_state.streaming.enabled = True + self.request_states[context.request_id] = request_state + elif isinstance(request_state, DuplexOrchestratorRequestState): + if request_state.config_generation != context.config_generation: + request_state.sampling_params_list = list(context.sampling_params) + request_state.config_generation = context.config_generation + request_state.session_id = context.session_id + request_state.fence = context.fence + else: + raise RuntimeError(f"request {context.request_id} is not session-owned") + self._sync_bridge_state(request_state, context) + self.session_manager.register_request(context.request_id, context.session_id) + + async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: + context = submission.context + request_state = self.request_states.get(context.request_id) + if not isinstance(request_state, DuplexOrchestratorRequestState): + raise RuntimeError(f"duplex request was not preregistered: {context.request_id}") + request = build_engine_core_request_from_tokens( + request_id=context.request_id, + prompt=dict(submission.prompt), + params=context.stage_sampling_params, + model_config=self.stage_pools[context.stage_id].stage_vllm_config.model_config, + resumable=True, + ) + request.external_req_id = request.request_id + pool = self.stage_pools[context.stage_id] + if submission.already_submitted: + replica_id = await pool.submit_update(context.request_id, request_state, request) + else: + replica_id = await pool.submit_initial(context.request_id, request_state, request, prompt_text=None) + if self.async_chunk and context.stage_id == 0: + prewarmed = await self._prewarm_async_chunk_stages( + context.request_id, + request, + request_state, + ) + if not prewarmed: + # The prewarm already failed the request, aborted it and + # popped its state; the runner turns this into an append + # error instead of writing onto an orphaned object. + raise RuntimeError( + f"async-chunk prewarm failed for duplex request {context.request_id}; the request was aborted" + ) + request_state.stage_fences[context.stage_id] = context.fence + request_state.stage_submit_ts[context.stage_id] = _time.time() + self._register_running_request(request_state) + return DuplexStageSubmissionResult( + request_id=context.request_id, + stage_id=context.stage_id, + replica_id=replica_id, + ) + + async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: + await self._cleanup_request_ids(request_ids, abort=abort) + + async def abort_requests(self, request_ids: list[str]) -> None: + if request_ids: + await self._abort_request_ids(list(dict.fromkeys(request_ids))) + + +__all__ = ["DuplexOrchestrator", "DuplexOrchestratorRequestState"] diff --git a/vllm_omni/engine/omni_engine_base.py b/vllm_omni/engine/omni_engine_base.py new file mode 100644 index 000000000..87ef8d5e6 --- /dev/null +++ b/vllm_omni/engine/omni_engine_base.py @@ -0,0 +1,1229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OmniEngineBase: process/thread ownership and transport shared by every Omni engine.""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +import json +import queue +import threading +import time +import uuid +import weakref +from collections.abc import Mapping +from dataclasses import asdict +from pathlib import Path +from typing import Any, cast + +import janus +import torch +from vllm import envs as vllm_envs +from vllm.logger import init_logger +from vllm.v1.engine.input_processor import InputProcessor + +from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remote_code_override +from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config +from vllm_omni.config.stage_config import ( + _DEPLOY_DIR, + DuplexSessionRuntimeConfig, + PipelineConfig, + load_deploy_config, +) +from vllm_omni.diffusion.data import ( + DiffusionParallelConfig, + parse_attention_config, +) +from vllm_omni.diffusion.io_support import get_diffusion_output_type +from vllm_omni.engine.async_engine_utils import ( + SHUTDOWN_ENQUEUE_TIMEOUT_S, + SHUTDOWN_JOIN_TIMEOUT_S, + enqueue_orchestrator_shutdown, + is_abort_transport_shutdown, + is_janus_sync_queue_shutdown, + shutdown_runtime_after_orchestrator, + weak_shutdown_async_omni_engine, +) +from vllm_omni.engine.messages import ( + AbortRequestMessage, + AbortResultMessage, + CollectiveRPCRequestMessage, + CollectiveRPCResultMessage, + EngineQueueMessage, + ErrorMessage, + OutputMessage, +) +from vllm_omni.engine.orchestrator import OrchestratorBase +from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient +from vllm_omni.engine.stage_client import StageClient +from vllm_omni.engine.stage_init_utils import build_stage0_input_processor +from vllm_omni.engine.stage_pool import StagePool +from vllm_omni.engine.stage_runtime import ( + StageRuntimeInfo, + create_stage_runtime, +) +from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin +from vllm_omni.entrypoints.utils import parse_stage_overrides +from vllm_omni.inputs.data import OmniSamplingParams +from vllm_omni.metrics.prometheus import OmniRequestCounter + +logger = init_logger(__name__) + +_STARTUP_POLL_INTERVAL_S = 1.0 +_REQUEST_QUEUE_MAXSIZE = 256 +_ConfigResolutionResult = OmniConfigResolution | tuple[str | None, list[Any], str | None] + + +def load_and_resolve_stage_configs( + model: str, + kwargs: dict[str, Any], + *, + trust_remote_code: bool | None, + deploy_config_path: str | None, + stage_overrides: Mapping[str, Mapping[str, Any]] | None, + strategy_config_path: str | None, +) -> OmniConfigResolution: + """Compatibility seam delegating to the single config resolver.""" + return resolve_omni_config( + model, + trust_remote_code=trust_remote_code, + cli_overrides=kwargs, + deploy_config_path=deploy_config_path, + stage_overrides=stage_overrides, + strategy_config_path=strategy_config_path, + ) + + +class OmniEngineBase: + """Generic engine: launches an orchestrator in a background thread. + + Shared by ``AsyncOmniEngine`` (turn-based requests) and ``DuplexOmniEngine`` + (duplex sessions). Subclasses construct their orchestrator in + ``_create_orchestrator``. + + All stage clients, input/output processors, and stage-to-stage transfer + logic live inside the Orchestrator coroutine (running in its own thread + with a dedicated asyncio event loop). This class communicates with it + via janus queues (sync side for callers, async side for orchestrator). + + Args: + model: Model name or path + init_timeout: Total timeout waiting for orchestrator startup (seconds). + stage_init_timeout: Timeout for stage initialization (seconds) + **kwargs: Additional arguments + """ + + # Class-level defaults so tests that bypass __init__ via object.__new__ + # don't AttributeError when stage-init / forward paths touch these attrs. + _log_stats: bool = False + _coordinator_runtime: Any = None + _transfer_emitter: Any = None + _prom_metrics: Any = None + _enable_orch_monitor: bool = False + # Lazily created by get_output_blocking_async(). + _output_drain_executor: concurrent.futures.ThreadPoolExecutor | None = None + + def __init__( + self, + model: str, + stage_init_timeout: int = 300, + init_timeout: int = 600, + single_stage_mode: bool = False, + transfer_emitter: Any = None, + prom_metrics: Any = None, + log_stats: bool = False, + tokenizer: str | None = None, + trust_remote_code: bool | None = None, + **kwargs: Any, + ) -> None: + self.model = model + self.tokenizer = tokenizer + # Cached by get_diffusion_od_config(). + self._diffusion_od_config_view: Any = None + startup_timeout = int(init_timeout) + # Forwarded into Orchestrator so its _forward_to_next_stage path can + # emit per-edge transfer_tx_s / transfer_size_bytes histograms. + # Optional: when None, Orchestrator silently skips TX emit (existing + # RX path still works via OrchestratorAggregator). + self._transfer_emitter = transfer_emitter + self._prom_metrics = prom_metrics + # Drives upstream EngineCore + scheduler stats production. When False + # the engine skips SchedulerStats / IterationStats; the per-(stage, + # replica) vllm:* wrap stays registered but reads zero. Respects the + # --log-stats CLI flag set by the user via OmniBase. + self._log_stats = log_stats + self._enable_orch_monitor = bool(kwargs.pop("enable_orch_monitor", False)) + + logger.info(f"[OmniEngine] Initializing with model {model}") + + # ------------------------------------------------------------------ # + # Single-stage mode detection # + # ------------------------------------------------------------------ # + # Single-stage mode is enabled when the caller explicitly passes # + # single_stage_mode=True, or when a stage_id is provided in the args. # + _stage_id_kwarg = kwargs.get("stage_id") + if isinstance(_stage_id_kwarg, int) and not single_stage_mode: + single_stage_mode = True + + self.single_stage_mode: bool = single_stage_mode + self._single_stage_id_filter: int | None = ( + int(_stage_id_kwarg) if single_stage_mode and isinstance(_stage_id_kwarg, int) else None + ) + self._omni_master_address: str | None = kwargs.get("omni_master_address") + self._omni_master_port: int | None = kwargs.get("omni_master_port") + + # New omni-coordinator flags. Consumed only in single_stage_mode. + # ``omni_dp_size_local`` is process-local: each invocation (head and + # every headless) launches that many replicas for its own stage. + self._omni_dp_size_local: int = int(kwargs.get("omni_dp_size_local") or 1) + if self._omni_dp_size_local < 1: + raise ValueError(f"--omni-dp-size-local must be >= 1, got {self._omni_dp_size_local}") + self._omni_lb_policy: str = str(kwargs.get("omni_lb_policy") or "random") + self._omni_heartbeat_timeout: float = float(kwargs.get("omni_heartbeat_timeout") or 30.0) + if self._omni_heartbeat_timeout <= 0: + raise ValueError(f"--omni-heartbeat-timeout must be > 0, got {self._omni_heartbeat_timeout}") + # Concurrent same-device stage init (admission + SH/EX phase locks). + # Sourced from the parallel_stage_init orchestrator/CLI arg (config, + # not an env var); default False preserves serial init. + self._parallel_stage_init: bool = bool(kwargs.get("parallel_stage_init") or False) + + if single_stage_mode: + logger.info( + "[OmniEngine] Single-stage mode enabled (stage_id_filter=%s, master=%s:%s)", + self._single_stage_id_filter, + self._omni_master_address, + self._omni_master_port, + ) + + # Keep the historical tuple return from _resolve_stage_configs while + # retaining the richer resolver result for pipeline-wide settings. + # Overrides of that private seam fall back to the factory below. + deploy_config_path = kwargs.get("deploy_config") + # ``trust_remote_code`` is tri-state (bool | None): ``None`` means "not + # specified" so stage-config resolution can defer to the deploy yaml's + # per-stage value (see ``with_trust_remote_code_override``). The + # restriction path below loads the top-level HF config via vLLM's + # ``get_config``, which needs a real bool, so collapse ``None`` to the + # default ``False`` here (#5495). + pipeline_config = StageConfigFactory.get_pipeline_config( + model=model, + trust_remote_code=bool(trust_remote_code), + deploy_config_path=deploy_config_path, + ) + self.pipeline_config = pipeline_config + self.endpoint_restrictions = pipeline_config.endpoint_restrictions if pipeline_config is not None else () + # The resolved deploy profile (explicit --deploy-config or the pipeline + # default). Duplex engines read ``deploy_config.duplex_session`` and + # ``session_mode`` from it; the turn-based engine only keeps it for + # introspection. + self.deploy_config = None + #: Path ``self.deploy_config`` was parsed from, so the duplex runtime + #: config below can reuse it instead of re-reading the same yaml. + self._deploy_config_source: str | None = None + deploy_config_source: str | Path | None = deploy_config_path + if deploy_config_source is None and pipeline_config is not None: + default_name = getattr(pipeline_config, "default_deploy_config_name", None) + if default_name: + deploy_config_source = _DEPLOY_DIR / default_name + elif isinstance(deploy_config_source, str) and not Path(deploy_config_source).exists(): + candidate = _DEPLOY_DIR / deploy_config_source + if candidate.exists(): + deploy_config_source = candidate + if deploy_config_source is not None: + try: + self.deploy_config = load_deploy_config(deploy_config_source) + self._deploy_config_source = str(deploy_config_source) + except Exception: + if deploy_config_path is not None: + # The caller named this file: failing to load it is their error. + raise + logger.warning( + "[OmniEngine] Could not load the pipeline's default deploy config %s", + deploy_config_source, + exc_info=True, + ) + + # Subclasses check deployment invariants here, before any stage + # process is launched (a failure is cheap at this point). + self._validate_deployment() + + # Tri-state: None means "not specified" — the deploy yaml's per-stage + # trust_remote_code stays in effect. An explicit True/False here is a + # global override (precedence: caller > deploy yaml > default False); + # the merge rule lives in with_trust_remote_code_override. + kwargs = with_trust_remote_code_override(kwargs, trust_remote_code) + self._config_resolution: OmniConfigResolution | None = None + self.config_path, self.stage_configs = self._resolve_stage_configs( + model, + kwargs, + trust_remote_code=trust_remote_code, + ) + if self._config_resolution is None: + # Same model, trust_remote_code and deploy path as the resolution + # above, so reuse it rather than paying for the HF config again. + self._set_pipeline_runtime_config(self.pipeline_config, self.config_path) + else: + self._set_pipeline_runtime_config( + self._config_resolution.pipeline_config, + self._config_resolution.config_path, + ) + + self.num_stages = len(self.stage_configs) + stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None + self.async_chunk = bool(getattr(stage0_args, "async_chunk", False)) + self.stage_pools: list[StagePool] = [] + self.stage_clients: list[StageClient] = [] # logical-stage view for external readers + self.input_processor: InputProcessor | None = None + self.prompt_transform_func: Any | None = None + self.prompt_expand_func: Any | None = None + self.supported_tasks: tuple[str, ...] = ("generate",) + self.default_sampling_params_list: list[OmniSamplingParams] = [] + self.stage_metadata: list[StageRuntimeInfo] = [] + # Janus queues are constructed eagerly here (not deferred to the + # orchestrator thread) so the master server's ROUTER thread always + # sees a non-None ``self.request_queue`` when on_register fires. + # ``async_q`` lazily binds to whatever event loop first awaits on + # it (the orchestrator loop), so cross-thread use stays correct. + self.request_queue: janus.Queue[EngineQueueMessage] = janus.Queue(maxsize=_REQUEST_QUEUE_MAXSIZE) + self.output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() + self.rpc_output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() + self._shutdown_called = False + self._weak_finalizer: weakref.finalize | None = None + self._correlated_rpc_client: CorrelatedRpcClient | None = None + self._running_counter = OmniRequestCounter() + self._engines_waiting_counter = OmniRequestCounter() + + logger.info(f"[OmniEngine] Launching Orchestrator thread with {self.num_stages} stages") + + # Launch orchestrator background thread + startup_future: concurrent.futures.Future = concurrent.futures.Future() + + self.orchestrator_thread = threading.Thread( + target=self._bootstrap_orchestrator, + args=( + stage_init_timeout, + startup_future, + ), + daemon=True, + name="orchestrator", + ) + self.orchestrator_thread.start() + self._wait_for_orchestrator_init(startup_future, startup_timeout) + self._correlated_rpc_client = CorrelatedRpcClient( + self.request_queue.sync_q, + self.rpc_output_queue.sync_q, + ) + + # Stage runtime fields are assigned directly on self by the bootstrap thread. + self._weak_finalizer = weakref.finalize( + self, + weak_shutdown_async_omni_engine, + self.orchestrator_thread, + self.request_queue, + self.output_queue, + self.rpc_output_queue, + self._correlated_rpc_client, + ) + + logger.info(f"[OmniEngine] Orchestrator ready with {self.num_stages} stages") + + def get_diffusion_od_config(self) -> Any: + """Expose the diffusion ``model_class_name`` to client-side model-extras. + + The worker holds the full config; here we just resolve the pipeline class + name from the model config (cached). ``model_class_name`` may be ``None``. + """ + if self._diffusion_od_config_view is None: + from types import SimpleNamespace + + from vllm_omni.diffusion.data import resolve_model_class_name + from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata + + model_class_name = resolve_model_class_name(self.model) + metadata = get_diffusion_model_metadata(model_class_name) + self._diffusion_od_config_view = SimpleNamespace( + model_class_name=model_class_name, + supports_multimodal_inputs=metadata.supports_multimodal_inputs, + max_multimodal_image_inputs=metadata.max_multimodal_image_inputs, + supports_mixed_reference_inputs=metadata.supports_mixed_reference_inputs, + ) + return self._diffusion_od_config_view + + def _initialize_stages(self, stage_init_timeout: int) -> None: + """Initialize stage clients/processors via StageRuntime and assign to self.""" + self._runtime = create_stage_runtime( + stage_configs=self.stage_configs, + model=self.model, + config_path=self.config_path, + single_stage_mode=self.single_stage_mode, + stage_init_timeout=stage_init_timeout, + async_chunk=self.async_chunk, + tokenizer=self.tokenizer, + parallel_stage_init=self._parallel_stage_init, + single_stage_id_filter=self._single_stage_id_filter, + omni_master_address=self._omni_master_address, + omni_master_port=self._omni_master_port, + omni_dp_size_local=self._omni_dp_size_local, + omni_heartbeat_timeout=self._omni_heartbeat_timeout, + omni_lb_policy=self._omni_lb_policy, + request_queue=self.request_queue, + log_stats=self._log_stats, + ) + self._runtime.initialize() + + self.num_stages = len(self.stage_configs) + self.stage_pools = self._runtime.stage_pools + self.stage_clients = [ + cast(StageClient, pool.stage_client) for pool in self.stage_pools if pool.stage_client is not None + ] + self.stage_vllm_configs = [pool.stage_vllm_config for pool in self.stage_pools] + self.output_processors = [pool.output_processor for pool in self.stage_pools] + self.input_processor = ( + build_stage0_input_processor(self.stage_vllm_configs[0]) + if self.stage_vllm_configs and self.stage_vllm_configs[0] is not None + else None + ) + self.prompt_transform_func = ( + getattr(self.stage_clients[0], "prompt_transform_func", None) if self.stage_clients else None + ) + self.prompt_expand_func = next( + ( + getattr(client, "prompt_expand_func", None) + for client in self.stage_clients + if getattr(client, "prompt_expand_func", None) is not None + ), + None, + ) + self.default_sampling_params_list = [client.default_sampling_params for client in self.stage_clients] + self.stage_metadata = [ + StageRuntimeInfo( + final_output=client.final_output, + final_output_type=client.final_output_type, + stage_type=client.stage_type, + model_stage=getattr(client, "model_stage", None), + ) + for client in self.stage_clients + ] + supported_tasks: set[str] = set() + if any(getattr(client, "is_comprehension", False) for client in self.stage_clients): + supported_tasks.add("generate") + if any(meta.final_output_type == "audio" for meta in self.stage_metadata): + supported_tasks.add("speech") + self.supported_tasks = tuple(supported_tasks) if supported_tasks else ("generate",) + + def _bootstrap_orchestrator( + self, + stage_init_timeout: int, + startup_future: concurrent.futures.Future, + ) -> None: + """Create loop, initialize stages, then run Orchestrator.""" + + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + async def _run_orchestrator() -> None: + self._initialize_stages(stage_init_timeout) + + pd_config = self._detect_pd_config() + + membership_controller = self._runtime.create_membership_controller() + + orchestrator = self._create_orchestrator( + request_async_queue=self.request_queue.async_q, + output_async_queue=self.output_queue.async_q, + rpc_async_queue=self.rpc_output_queue.async_q, + stage_pools=self.stage_pools, + async_chunk=self.async_chunk, + pd_config=pd_config, + membership_controller=membership_controller, + running_counter=self._running_counter, + engines_waiting_counter=self._engines_waiting_counter, + transfer_emitter=self._transfer_emitter, + prom_metrics=self._prom_metrics, + log_stats=self._log_stats, + enable_orch_monitor=self._enable_orch_monitor, + ) + if not startup_future.done(): + startup_future.set_result(asyncio.get_running_loop()) + await orchestrator.run() + + try: + loop.run_until_complete(_run_orchestrator()) + except Exception as e: + if not startup_future.done(): + wrapped = RuntimeError(f"Orchestrator initialization failed: {e}") + wrapped.__cause__ = e + startup_future.set_exception(wrapped) + logger.exception("[OmniEngine] Orchestrator thread crashed") + error_text = str(e) or "Orchestrator thread crashed" + try: + error_msg = ErrorMessage(error=error_text, fatal=True) + if self.output_queue is not None: + self.output_queue.sync_q.put_nowait(error_msg) + if self.rpc_output_queue is not None: + self.rpc_output_queue.sync_q.put_nowait(error_msg) + except Exception: + pass + raise + finally: + try: + pending = [task for task in asyncio.all_tasks(loop) if not task.done()] + for task in pending: + task.cancel() + if pending: + loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + loop.run_until_complete(loop.shutdown_asyncgens()) + if hasattr(loop, "shutdown_default_executor"): + loop.run_until_complete(loop.shutdown_default_executor()) + except Exception: + logger.exception("[OmniEngine] Failed during orchestrator loop cleanup") + finally: + asyncio.set_event_loop(None) + loop.close() + + def _validate_deployment(self) -> None: + """Seam: check ``pipeline_config`` / ``deploy_config`` before stages start (default: nothing).""" + + def _create_orchestrator(self, **orchestrator_kwargs: Any) -> OrchestratorBase: + """Construct this engine's orchestrator (runs on the orchestrator thread after stage init).""" + raise NotImplementedError + + def _wait_for_orchestrator_init(self, startup_future: concurrent.futures.Future, startup_timeout: int) -> None: + """ + Wait for orchestrator startup future to return ready. Raises exception on any failures to the init process. + """ + deadline = time.monotonic() + startup_timeout + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + logger.warning( + "[OmniEngine] Orchestrator startup timed out after %ss. " + "Multi-stage deployments that initialize stages sequentially on one device " + "or load checkpoints from slow storage may need larger --init-timeout and " + "--stage-init-timeout values.", + startup_timeout, + ) + self._try_shutdown("[OmniEngine] Failed to cleanup after orchestrator startup timeout") + raise TimeoutError(f"Orchestrator did not become ready within {startup_timeout}s") + try: + startup_future.result( + timeout=min(remaining, _STARTUP_POLL_INTERVAL_S), + ) + break + except concurrent.futures.TimeoutError: + if not self.orchestrator_thread.is_alive(): + self._try_shutdown("[OmniEngine] Failed to cleanup after orchestrator startup failure") + if startup_future.done(): + startup_future.result() # re-raises the real exception + raise RuntimeError("Orchestrator thread died during startup") + except Exception: + self._try_shutdown("[OmniEngine] Failed to cleanup after orchestrator startup failure") + raise + + @staticmethod + def _get_default_cache_config(cache_backend: str | None) -> dict[str, Any] | None: + if cache_backend == "cache_dit": + return { + "Fn_compute_blocks": 1, + "Bn_compute_blocks": 0, + "max_warmup_steps": 4, + "residual_diff_threshold": 0.24, + "max_continuous_cached_steps": 3, + "enable_taylorseer": False, + "taylorseer_order": 1, + "scm_steps_mask_policy": None, + "scm_steps_policy": "dynamic", + } + if cache_backend == "tea_cache": + return { + "rel_l1_thresh": 0.2, + } + if cache_backend == "mag_cache": + return { + "mag_threshold": 0.24, + "mag_max_skip_steps": 5, + "mag_retention_ratio": 0.1, + } + if cache_backend == "step_cache": + return { + "step_cache_dit_enabled": True, + "velocity_sim_thresholds": [0.95, 0.93], + "velocity_skip_countdowns": [4, 2], + "step_cache_dit_min_history": 2, + "step_cache_dit_max_history": 2, + } + return None + + @staticmethod + def _normalize_cache_config(cache_backend: str | None, cache_config: Any | None) -> Any | None: + if isinstance(cache_config, str): + try: + cache_config = json.loads(cache_config) + except json.JSONDecodeError: + logger.warning("Invalid cache_config JSON, using defaults.") + cache_config = None + if cache_config is None and cache_backend not in (None, "", "none"): + cache_config = OmniEngineBase._get_default_cache_config(cache_backend) + return cache_config + + def _detect_pd_config(self) -> dict[str, Any] | None: + """Detect PD (Prefill-Decode) disaggregation config from stage_configs. + Returns a dict with 'pd_pair' and 'bootstrap_addr', or None. + """ + pd_pair = PDDisaggregationMixin.detect_pd_separation_from_stage_configs(self.stage_configs) + if pd_pair is None: + return None + prefill_idx, decode_idx = pd_pair + + # Extract bootstrap address from prefill stage engine_args + bootstrap_addr: str | None = None + try: + prefill_cfg = self.stage_configs[prefill_idx] + ea = getattr(prefill_cfg, "engine_args", None) + kv_cfg = getattr(ea, "kv_transfer_config", None) if ea is not None else None + if kv_cfg is not None: + port = vllm_envs.VLLM_MOONCAKE_BOOTSTRAP_PORT + kv_ip = getattr(kv_cfg, "kv_ip", None) or "127.0.0.1" + bootstrap_addr = f"http://{kv_ip}:{port}" + except Exception as exc: + logger.warning("[OmniEngine] Could not extract PD bootstrap address: %s", exc) + + logger.info( + "[OmniEngine] PD disaggregation detected: prefill=stage-%d, decode=stage-%d, bootstrap=%s", + prefill_idx, + decode_idx, + bootstrap_addr, + ) + prefill_engine_id: str | None = None + try: + prefill_client = self.stage_clients[prefill_idx] + kv_cfg = getattr(getattr(prefill_client, "vllm_config", None), "kv_transfer_config", None) + prefill_engine_id = getattr(kv_cfg, "engine_id", None) + except Exception as exc: + logger.warning("[OmniEngine] Could not extract prefill engine_id: %s", exc) + + return { + "pd_pair": (prefill_idx, decode_idx), + "bootstrap_addr": bootstrap_addr, + "prefill_engine_id": prefill_engine_id, + } + + @staticmethod + def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list: + """Create a default single-stage diffusion config from kwargs.""" + # We temporally create a default config for diffusion stage. + # In the future, we should merge the default config with the user-provided config. + normalized_kwargs = dict(kwargs) + default_sampling_params = normalized_kwargs.get("default_sampling_params") + if isinstance(default_sampling_params, str): + try: + default_sampling_params = json.loads(default_sampling_params) + except json.JSONDecodeError: + logger.warning("Invalid default_sampling_params JSON, ignoring stage defaults.") + default_sampling_params = None + if not isinstance(default_sampling_params, dict): + default_sampling_params = None + stage_default_sampling_params = default_sampling_params.get("0", {}) if default_sampling_params else {} + if normalized_kwargs.get("dtype") is None: + normalized_kwargs["dtype"] = "auto" + + # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error. + if "dtype" in normalized_kwargs and not isinstance(normalized_kwargs["dtype"], str): + if not isinstance(normalized_kwargs["dtype"], torch.dtype): + raise TypeError( + f"Provided dtype must be a string or torch.dtype, got {type(normalized_kwargs['dtype']).__name__}" + ) + normalized_kwargs["dtype"] = str(normalized_kwargs["dtype"]).removeprefix("torch.") + + cache_backend = normalized_kwargs.get("cache_backend", "none") + cache_config = OmniEngineBase._normalize_cache_config( + cache_backend, + normalized_kwargs.get("cache_config", None), + ) + + parallel_config = normalized_kwargs.get("parallel_config") + if isinstance(parallel_config, dict): + parallel_config = DiffusionParallelConfig.from_dict(parallel_config) + if parallel_config is None: + ulysses_degree = normalized_kwargs.get("ulysses_degree") or 1 + ring_degree = normalized_kwargs.get("ring_degree") or 1 + allgather_degree = normalized_kwargs.get("allgather_degree") or 1 + ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict" + ulysses_a2a_permute = bool(normalized_kwargs.get("ulysses_a2a_permute", False)) + sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size") + pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 + data_parallel_size = normalized_kwargs.get("data_parallel_size") + tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1 + cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1 + pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 + vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1 + vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "tile" + text_encoder_tp_size = normalized_kwargs.get("text_encoder_tp_size") or 1 + enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False + use_hsdp = normalized_kwargs.get("use_hsdp", False) + hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size", -1) + hsdp_replicate_size = normalized_kwargs.get("hsdp_replicate_size", 1) + if sequence_parallel_size is None: + sequence_parallel_size = allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree + + parallel_config = DiffusionParallelConfig( + pipeline_parallel_size=pipeline_parallel_size, + data_parallel_size=data_parallel_size, + tensor_parallel_size=tensor_parallel_size, + enable_expert_parallel=enable_expert_parallel, + sequence_parallel_size=sequence_parallel_size, + ulysses_degree=ulysses_degree, + ring_degree=ring_degree, + allgather_degree=allgather_degree, + ulysses_mode=ulysses_mode, + ulysses_a2a_permute=ulysses_a2a_permute, + cfg_parallel_size=cfg_parallel_size, + vae_patch_parallel_size=vae_patch_parallel_size, + vae_parallel_mode=vae_parallel_mode, + text_encoder_tp_size=text_encoder_tp_size, + use_hsdp=use_hsdp, + hsdp_shard_size=hsdp_shard_size, + hsdp_replicate_size=hsdp_replicate_size, + ) + + num_gpus = normalized_kwargs.get("num_gpus") + if num_gpus is not None: + num_gpus = int(num_gpus) + parallel_config.resolve_data_parallel_size(num_gpus) + + num_devices = max(1, int(parallel_config.world_size)) + devices = ",".join(str(i) for i in range(num_devices)) + model_class_name = kwargs.get("model_class_name", None) + final_output_type = get_diffusion_output_type(model_class_name) + + attention_config = None + if ( + kwargs.get("diffusion_attention_config") is not None + or kwargs.get("diffusion_attention_backend") is not None + or kwargs.get("fastvideo_vsa_topk") is not None + ): + attention_config = parse_attention_config( + kwargs.get("diffusion_attention_config"), + attention_backend=kwargs.get("diffusion_attention_backend"), + fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"), + ) + + extras = dict(kwargs.get("extras") or {}) + for key, default in ( + ("auxiliary_text_encoder", None), + ("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"), + ): + top_level_value = kwargs.get(key) + if top_level_value is not None: + extras[key] = top_level_value + else: + extras.setdefault(key, default) + + stage_engine_args = { + "max_num_seqs": kwargs.get("max_num_seqs") or 1, + "parallel_config": parallel_config, + # Default-stage construction bypasses the structured projection. + # Runner selection remains owned by the selected engine/platform. + "engine_backend": kwargs.get("engine_backend", "default"), + "model_class_name": kwargs.get("model_class_name", None), + "task_type": kwargs.get("task_type", None), + "model_config": kwargs.get("model_config", None), + "additional_config": kwargs.get("additional_config", None), + "step_execution": kwargs.get("step_execution", False), + "request_batch_max_wait_ms": kwargs.get("request_batch_max_wait_ms", 0.0), + "vae_use_slicing": kwargs.get("vae_use_slicing", False), + "vae_use_tiling": kwargs.get("vae_use_tiling", False), + "cache_backend": cache_backend, + "cache_config": cache_config, + "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False), + "diffusion_offload_config": kwargs.get("diffusion_offload_config", None), + "enable_cpu_offload": kwargs.get("enable_cpu_offload", False), + "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False), + "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False), + "dlo_use_allgather": kwargs.get("dlo_use_allgather", True), + "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0), + "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"), + "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"), + "dlo_host_registration_limit_gib": kwargs.get("dlo_host_registration_limit_gib", 0.0), + "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"), + "diffusion_compile_granularity": ( + "regional" + if kwargs.get("diffusion_compile_granularity") is None + else kwargs["diffusion_compile_granularity"] + ), + "diffusion_compile_dynamic": ( + True if kwargs.get("diffusion_compile_dynamic") is None else kwargs["diffusion_compile_dynamic"] + ), + "fa_deterministic": bool(kwargs.get("fa_deterministic", False)), + "boundary_ratio": kwargs.get("boundary_ratio", None), + "flow_shift": kwargs.get("flow_shift", None), + "diffusion_load_format": kwargs.get("diffusion_load_format", "default"), + "lora_path": kwargs.get("lora_path", None), + "lora_scale": kwargs.get("lora_scale", 1.0), + "lora_backend": kwargs.get("lora_backend", "peft"), + "custom_pipeline_args": kwargs.get("custom_pipeline_args", None), + "worker_extension_cls": kwargs.get("worker_extension_cls", None), + "trust_remote_code": (False if kwargs.get("trust_remote_code") is None else kwargs["trust_remote_code"]), + "distributed_executor_backend": kwargs.get("distributed_executor_backend"), + "enable_sleep_mode": kwargs.get("enable_sleep_mode", False), + "enable_prompt_embed_cache": kwargs.get("enable_prompt_embed_cache", False), + "prompt_embed_cache_size": kwargs.get("prompt_embed_cache_size", 32), + "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), + "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4), + "quantization": kwargs.get("quantization", None), + "quantization_config": kwargs.get("quantization_config", None), + "diffusion_kv_cache_dtype": kwargs.get("diffusion_kv_cache_dtype", None), + "diffusion_kv_cache_skip_steps": kwargs.get("diffusion_kv_cache_skip_steps", None), + "diffusion_kv_cache_skip_layers": kwargs.get("diffusion_kv_cache_skip_layers", None), + **({"diffusion_attention_config": attention_config} if attention_config is not None else {}), + "force_cutlass_fp8": bool(kwargs.get("force_cutlass_fp8", False)), + "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False), + "streaming_output": kwargs.get("diffusion_streaming_output", False), + "enable_ar_profiler": kwargs.get("enable_ar_profiler", False), + "extras": extras, + **( + { + "profiler_config": asdict(kwargs["profiler_config"]) + if hasattr(kwargs["profiler_config"], "__dataclass_fields__") + else kwargs["profiler_config"] + } + if kwargs.get("profiler_config") is not None + else {} + ), + } + if num_gpus is not None: + stage_engine_args["num_gpus"] = num_gpus + # Only set dtype if it was already explicitly passed and normalized + if "dtype" in normalized_kwargs: + stage_engine_args["dtype"] = normalized_kwargs["dtype"] + + # New split fields for diffusers adapter kwargs. + if kwargs.get("diffusers_load_kwargs") is not None: + stage_engine_args["diffusers_load_kwargs"] = kwargs["diffusers_load_kwargs"] + if kwargs.get("diffusers_call_kwargs") is not None: + stage_engine_args["diffusers_call_kwargs"] = kwargs["diffusers_call_kwargs"] + + default_stage_cfg = [ + { + "stage_id": 0, + "stage_type": "diffusion", + "runtime": { + "process": True, + "devices": devices, + }, + "engine_args": stage_engine_args, + "engine_input_source": [], + "default_sampling_params": stage_default_sampling_params, + "final_output": True, + "final_output_type": final_output_type, + } + ] + default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion" + return default_stage_cfg + + def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None: + """Apply a strategy-derived ``omni_lb_policy`` to the engine. + + Precedence: an explicit ``--omni-lb-policy`` always wins. ``"random"`` is + the engine default and is treated as "unset" (indistinguishable from no + flag), so a strategy value overrides it. If the user explicitly passed a + non-default policy that conflicts with the strategy-derived one, raise so + the mismatch is not silently ignored. + """ + if not derived: + return + explicit = kwargs.get("omni_lb_policy") + user_set = explicit is not None and str(explicit) != "random" + if user_set: + if str(explicit) != str(derived): + raise ValueError( + f"Conflicting load-balancer policy: --omni-lb-policy={explicit!r} was given " + f"but the composable-parallel strategy derived omni_lb_policy={derived!r}. " + "Drop --omni-lb-policy to use the strategy value, or make them match." + ) + return + if self._omni_lb_policy != str(derived): + logger.info( + "[composable_parallel] applying strategy-derived omni_lb_policy=%r (was %r).", + derived, + self._omni_lb_policy, + ) + self._omni_lb_policy = str(derived) + + @staticmethod + def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list[dict[str, Any]]: + """Compatibility seam for the factory-owned diffusion fallback.""" + return StageConfigFactory.create_default_diffusion(kwargs) + + def _set_pipeline_runtime_config( + self, + pipeline_config: PipelineConfig | None, + config_path: str | None, + ) -> None: + """Initialize engine-wide settings resolved from pipeline metadata.""" + self.endpoint_restrictions = pipeline_config.endpoint_restrictions if pipeline_config is not None else () + # No duplex_runtime_extension / duplex_serving_adapter / duplex_control + # here: the pre-framework wiring they named is gone, and design rule 3 + # in docs/design/fullduplex.md keeps duplex vocabulary out of the base. + self.duplex_session_config = DuplexSessionRuntimeConfig() + if config_path is not None: + if self.deploy_config is not None and self._deploy_config_source == str(config_path): + # Already parsed in __init__; the same yaml twice is pure cost. + self.duplex_session_config = self.deploy_config.duplex_session + else: + self.duplex_session_config = load_deploy_config(config_path).duplex_session + + def _resolve_stage_configs( + self, + model: str, + kwargs: dict[str, Any], + *, + trust_remote_code: bool | None, + ) -> tuple[str, list[Any]]: + """Resolve stage configs and inject defaults shared by orchestrator/headless.""" + + for legacy_arg in ("stage_configs_path", "stage_configs"): + if legacy_arg in kwargs: + raise ValueError(f"`{legacy_arg}` is no longer supported; use `deploy_config` instead.") + + # log_stats is captured by __init__; its CLI-only negative alias must + # not cross into per-stage structured config ownership validation. + kwargs.pop("disable_log_stats", None) + deploy_config_path = kwargs.pop("deploy_config", None) + strategy_config_path = kwargs.pop("strategy_config", None) + # CLI callers arrive pre-parsed; offline Python callers may use the + # JSON-string form documented in recipes. + stage_overrides = parse_stage_overrides(kwargs.pop("stage_overrides", None)) + + # ``diffusion_streaming_output`` is the public AsyncOmni/serve kwarg; + # stage configs know the field as ``streaming_output``. Mirror only a + # truthy value so the CLI's False default does not override deploy YAML. + if kwargs.get("diffusion_streaming_output") and kwargs.get("streaming_output") is None: + kwargs["streaming_output"] = True + + resolution = cast( + _ConfigResolutionResult, + load_and_resolve_stage_configs( + model, + kwargs, + trust_remote_code=trust_remote_code, + deploy_config_path=deploy_config_path, + stage_overrides=stage_overrides, + strategy_config_path=strategy_config_path, + ), + ) + if isinstance(resolution, OmniConfigResolution): + self._config_resolution = resolution + config_path = resolution.config_path + stage_configs = list(resolution.stage_configs) + strategy_lb_policy = resolution.omni_lb_policy + else: + # Compatibility for overrides of the historical tuple-returning + # seam. Production always receives OmniConfigResolution above. + config_path, stage_configs, strategy_lb_policy = resolution + + # A strategy.yaml may derive a pipeline-wide load-balancer policy. It is + # an orchestrator-level knob (read once at construction), so apply it here + # rather than as a per-stage config field. + self._apply_strategy_lb_policy(strategy_lb_policy, kwargs) + + return cast(str, config_path), stage_configs + + # ==================== Public API ==================== + + def try_get_output(self, timeout: float = 0.001) -> EngineQueueMessage | None: + """Read one output message from the Orchestrator output queue.""" + try: + return self.output_queue.sync_q.get(timeout=timeout) + except queue.Empty: + if not self.is_alive(): + raise RuntimeError("Orchestrator died unexpectedly. See logs above.") + return None + + async def try_get_output_async(self) -> EngineQueueMessage | None: + """Async read from the Orchestrator output queue.""" + try: + return self.output_queue.sync_q.get_nowait() + except queue.Empty: + if not self.is_alive(): + raise RuntimeError("Orchestrator died unexpectedly. See logs above.") + return None + + async def get_output_blocking_async(self, timeout: float = 1.0) -> EngineQueueMessage | None: + """Blocking-wait read from the Orchestrator output queue. + + Waits up to ``timeout`` seconds in a dedicated drain thread for the + next message (condition-variable wakeup instead of a poll cadence); + returns ``None`` on timeout so the caller keeps its liveness check, + mirroring ``try_get_output_async``'s contract. Used by the serving + final-output drain when ``VLLM_OMNI_EVENT_DRIVEN_ORCH`` is on. + """ + executor = self._output_drain_executor + if executor is None: + executor = concurrent.futures.ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="omni-output-drain", + ) + self._output_drain_executor = executor + + sync_q = self.output_queue.sync_q + + def _drain_get() -> EngineQueueMessage | None: + # Exceptions are swallowed to a None sentinel: the queue may be + # closed mid-shutdown, and an exception left on an executor future + # after task cancellation would warn as never-retrieved. + try: + return sync_q.get(timeout=timeout) + except queue.Empty: + return None + except Exception: + return None + + loop = asyncio.get_running_loop() + msg = await loop.run_in_executor(executor, _drain_get) + if msg is None and not self.is_alive(): + raise RuntimeError("Orchestrator died unexpectedly. See logs above.") + return msg + + def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo: + """Get cached metadata for a stage.""" + return self.stage_metadata[stage_id] + + def abort(self, request_ids: list[str]) -> None: + """Fire-and-forget abort: enqueue and return without waiting. + + Prefer :meth:`abort_async` when the caller needs acknowledgment that + stage aborts, binding release, and orchestrator request cleanup finished. + """ + if not request_ids or getattr(self, "_shutdown_called", False): + return + if self.request_queue is None: + raise RuntimeError("request_queue is not initialized") + try: + self.request_queue.sync_q.put(AbortRequestMessage(request_ids=request_ids)) + except Exception as exc: + if getattr(self, "_shutdown_called", False) and is_janus_sync_queue_shutdown(exc): + return + raise + + async def abort_async( + self, + request_ids: list[str], + timeout: float | None = None, + ) -> list[OutputMessage]: + """Abort requests and wait for orchestrator acknowledgment. + + Unlike :meth:`abort`, this generates an ``rpc_id``, correlates the + :class:`AbortResultMessage` via :class:`CorrelatedRpcClient`, and + raises if the orchestrator reports failure or times out. + + Returns: + Final-stage AR abort ``OutputMessage`` list carrying partial + tokens generated before abort (empty for diffusion / no OP state). + """ + if not request_ids or getattr(self, "_shutdown_called", False): + return [] + if self.request_queue is None: + raise RuntimeError("request_queue is not initialized") + transport = self._correlated_rpc_client + if transport is None: + raise RuntimeError("correlated RPC client is not initialized") + + rpc_id = uuid.uuid4().hex + msg = AbortRequestMessage(request_ids=request_ids, rpc_id=rpc_id) + + def _wait() -> AbortResultMessage: + result_msg = transport.execute( + ("abort", rpc_id), + msg, + timeout=timeout, + timeout_message=f"abort timed out after {timeout} seconds", + block_on_submit=True, + ) + if not isinstance(result_msg, AbortResultMessage): + raise RuntimeError(f"unexpected abort result type: {type(result_msg).__name__}") + return result_msg + + loop = asyncio.get_running_loop() + try: + result_msg = await loop.run_in_executor(None, _wait) + except Exception as exc: + if getattr(self, "_shutdown_called", False) and is_abort_transport_shutdown(exc): + return [] + raise + if not result_msg.success: + raise RuntimeError(result_msg.error or "abort failed") + return list(result_msg.abort_outputs or []) + + def collective_rpc( + self, + method: str, + timeout: float | None = None, + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + stage_ids: list[int] | None = None, + ) -> list[Any]: + """Send a control RPC to the Orchestrator and wait for aggregated results. + + This uses a dedicated RPC output queue so control-plane messages do not + race with the normal request output polling loop. + """ + rpc_id = uuid.uuid4().hex + msg = CollectiveRPCRequestMessage( + rpc_id=rpc_id, + method=method, + timeout=timeout, + args=tuple(args), + kwargs=kwargs or {}, + stage_ids=stage_ids, + ) + + transport = self._correlated_rpc_client + if transport is None: + raise RuntimeError("correlated RPC client is not initialized") + result_msg = transport.execute( + ("collective", rpc_id), + msg, + timeout=timeout, + timeout_message=f"collective_rpc timed out after {timeout} seconds", + block_on_submit=True, + ) + if not isinstance(result_msg, CollectiveRPCResultMessage): + raise RuntimeError(f"unexpected collective RPC result type: {type(result_msg).__name__}") + return list(result_msg.results) + + async def collective_rpc_async( + self, + method: str, + timeout: float | None = None, + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + stage_ids: list[int] | None = None, + ) -> list[Any]: + """Async wrapper around collective_rpc().""" + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + None, + lambda: self.collective_rpc( + method=method, + timeout=timeout, + args=args, + kwargs=kwargs, + stage_ids=stage_ids, + ), + ) + + @property + def rpc_client(self) -> CorrelatedRpcClient: + """The correlated RPC transport (available once the orchestrator is up).""" + if self._correlated_rpc_client is None: + raise RuntimeError("engine RPC transport is not initialized") + return self._correlated_rpc_client + + def is_alive(self) -> bool: + """Whether the orchestrator thread is alive.""" + return bool(self.orchestrator_thread.is_alive()) + + def shutdown(self) -> None: + """Send shutdown message and wait for the Orchestrator thread to exit.""" + if getattr(self, "_shutdown_called", False): + return + self._shutdown_called = True + finalizer = getattr(self, "_weak_finalizer", None) + if finalizer is not None and finalizer.alive: + finalizer.detach() + + logger.info("[OmniEngine] Shutting down Orchestrator") + request_queue_closed = False + shutdown_enqueued = enqueue_orchestrator_shutdown( + self.request_queue, + timeout=SHUTDOWN_ENQUEUE_TIMEOUT_S, + ) + if self.request_queue is not None and not shutdown_enqueued: + logger.error( + "[OmniEngine] Failed to enqueue orchestrator shutdown; " + "closing the request queue to wake the request handler" + ) + try: + self.request_queue.close() + request_queue_closed = True + except Exception: + logger.exception("[OmniEngine] Failed to close the request queue") + + if self._correlated_rpc_client is not None: + try: + self._correlated_rpc_client.close() + except Exception: + logger.exception("[OmniEngine] Failed to close correlated RPC client") + + orchestrator_stopped = False + try: + if self.is_alive(): + self.orchestrator_thread.join(timeout=SHUTDOWN_JOIN_TIMEOUT_S) + orchestrator_stopped = not self.is_alive() + if not orchestrator_stopped: + logger.error( + "[OmniEngine] Orchestrator did not stop within %.1f seconds; continuing cleanup", + SHUTDOWN_JOIN_TIMEOUT_S, + ) + except Exception: + logger.exception("[OmniEngine] Failed to join Orchestrator thread") + + for q in (self.request_queue, self.output_queue, self.rpc_output_queue): + try: + if not (q is self.request_queue and request_queue_closed): + q.close() + except Exception: + pass + + if self._output_drain_executor is not None: + # Any in-flight blocking get bails out within its ≤1 s timeout + # (or immediately via the queue close above), so don't wait. + self._output_drain_executor.shutdown(wait=False) + self._output_drain_executor = None + + if hasattr(self, "_runtime") and self._runtime is not None and orchestrator_stopped: + try: + self._runtime.shutdown() + except Exception: + logger.exception("[OmniEngine] Failed to shutdown StageRuntime") + elif hasattr(self, "_runtime") and self._runtime is not None: + logger.warning("[OmniEngine] Deferring StageRuntime shutdown until the Orchestrator exits") + threading.Thread( + target=shutdown_runtime_after_orchestrator, + args=(self.orchestrator_thread, self._runtime), + daemon=True, + name="omni-stage-runtime-shutdown", + ).start() + + # ── Release CuMem allocator memory pool ────────────────────────────── + # When enable_sleep_mode is in use, the CuMem (CUDA Virtual Memory + # Management) allocator holds model weights in a singleton memory pool + # that lives in the parent process. Killing the engine-core subprocess + # does NOT release this pool — the weights stay resident on the GPU + # and can cause CUDA OOM for subsequent engine instances (especially + # large models like BAGEL-7B-MoT whose weights alone consume ~134 GiB). + # + # CuMemAllocator.sleep() is NOT idempotent — calling it on already- + # slept entries causes CUDA_ERROR_INVALID_VALUE at cumem_allocator + # cuMemRelease (double-free of the memory handle). Use release_pools() + # instead, which is the designed cleanup path: it drops MemPool refs + # and lets the destructor/free path handle asleep entries correctly + # (returns a null handle so the C extension skips unmap/release). + try: + from vllm.device_allocator.cumem import CuMemAllocator, cumem_available + + if cumem_available: + allocator = CuMemAllocator.get_instance() + allocator.release_pools() + logger.debug("[OmniEngine] Released CuMem memory pool during shutdown") + except Exception: + pass + + def _try_shutdown(self, *args, **kwargs) -> None: + try: + self.shutdown() + except Exception: + logger.exception(*args, **kwargs) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 75620bec9..aced2c1b3 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -18,7 +18,7 @@ import asyncio import os import shutil import time as _time -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Coroutine from dataclasses import dataclass, field from http import HTTPStatus from typing import Any @@ -34,28 +34,10 @@ from vllm.v1.engine import EngineCoreOutputs from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats -from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig from vllm_omni.diffusion.data import is_diffusion_request_started_output from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker -from vllm_omni.engine.duplex.contracts import ( - DuplexControlPlanePort, - DuplexOutputContext, - DuplexOutputDecision, - DuplexRequestIdentity, - DuplexRuntimeExtension, - DuplexStageRequestContext, - DuplexStageSubmission, - DuplexStageSubmissionResult, -) -from vllm_omni.engine.duplex.control_plane import DuplexControlPlane -from vllm_omni.engine.duplex.lease import DuplexLeaseConfig -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.session import ( - DuplexSessionRuntimeManager, - DuplexSessionRuntimeState, -) from vllm_omni.engine.membership_controller import MembershipController from vllm_omni.engine.messages import ( AbortRequestMessage, @@ -82,7 +64,6 @@ from vllm_omni.metrics.prometheus import OmniRequestCounter from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.outputs.duplex import attach_duplex_output_decision logger = init_logger(__name__) @@ -224,9 +205,10 @@ class OrchestratorRequestState: # Per-request pipeline timing accumulator (milliseconds) pipeline_timings: dict[str, float] = field(default_factory=dict) - duplex_identity: DuplexRequestIdentity | None = None - duplex_stage_fences: dict[int, DuplexFence] = field(default_factory=dict) - duplex_config_generation: int = -1 + # A session-owned request's lifetime belongs to a session (duplex), not to + # its own ``finished`` flag: no synthetic terminal, no client emission of + # stage-0 segment ends, no terminal re-forward, no cleanup on finish. + session_owned: bool = False running_counter_registered: bool = False request_artifact_dirs: set[str] = field(default_factory=set) @@ -279,143 +261,14 @@ class StreamingInputState: return segment if segment is not None else StreamingSegmentState() -class _OrchestratorDuplexStagePort: - """Adapts generic stage pools to the model-neutral duplex control plane.""" +class OrchestratorBase: + """Stage-management loop shared by the turn-based and duplex orchestrators. - def __init__( - self, - *, - stage_pools: list[StagePool], - request_states: dict[str, OrchestratorRequestState], - running_counter: OmniRequestCounter | None, - cleanup_request_ids: Callable[..., Any], - async_chunk: bool, - prewarm_async_chunk_stages: Callable[ - [str, Any, OrchestratorRequestState], - Awaitable[bool], - ], - ) -> None: - self._stage_pools = stage_pools - self._request_states = request_states - self._running_counter = running_counter - self._cleanup_request_ids = cleanup_request_ids - self._async_chunk = async_chunk - self._prewarm_async_chunk_stages = prewarm_async_chunk_stages - - @property - def stage_count(self) -> int: - return len(self._stage_pools) - - def sampling_defaults(self) -> tuple[object, ...]: - return tuple(pool.stage_client.default_sampling_params for pool in self._stage_pools) - - @staticmethod - def _sync_bridge_state( - request_state: OrchestratorRequestState, - context: DuplexStageRequestContext, - ) -> None: - duplex_state = request_state.streaming.bridge_states.setdefault("duplex", {}) - if not isinstance(duplex_state, dict): - duplex_state = {} - request_state.streaming.bridge_states["duplex"] = duplex_state - previous_epoch = duplex_state.get("epoch") - current_model_turn_id = duplex_state.get("model_turn_id") - if ( - not isinstance(current_model_turn_id, int) - or previous_epoch != context.fence.epoch - or current_model_turn_id < context.fence.turn_id - ): - # A serving-side safety boundary can close a Realtime response - # before the model emits its normal turn_eos. Catch up the - # engine-owned identity when the next fenced append starts. - duplex_state["model_turn_id"] = context.fence.turn_id - duplex_state.update( - { - "session_id": context.session_id, - "fence": context.fence, - "incarnation": context.fence.incarnation, - "epoch": context.fence.epoch, - "turn_id": context.fence.turn_id, - "response_seq": context.fence.response_seq, - "session_config": dict(context.session_config), - "runtime_config": dict(context.runtime_config), - } - ) - - def ensure_request(self, context: DuplexStageRequestContext) -> None: - request_state = self._request_states.get(context.request_id) - if request_state is None: - request_state = OrchestratorRequestState( - request_id=context.request_id, - prompt=None, - sampling_params_list=list(context.sampling_params), - final_stage_id=context.final_stage_id, - duplex_config_generation=context.config_generation, - ) - request_state.streaming.enabled = True - self._request_states[context.request_id] = request_state - elif request_state.duplex_config_generation != context.config_generation: - request_state.sampling_params_list = list(context.sampling_params) - request_state.duplex_config_generation = context.config_generation - request_state.duplex_identity = DuplexRequestIdentity( - session_id=context.session_id, - fence=context.fence, - ) - self._sync_bridge_state(request_state, context) - - async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: - context = submission.context - request_state = self._request_states.get(context.request_id) - if request_state is None: - raise RuntimeError(f"duplex request was not preregistered: {context.request_id}") - request = build_engine_core_request_from_tokens( - request_id=context.request_id, - prompt=dict(submission.prompt), - params=context.stage_sampling_params, - model_config=self._stage_pools[context.stage_id].stage_vllm_config.model_config, - resumable=True, - ) - request.external_req_id = request.request_id - pool = self._stage_pools[context.stage_id] - if submission.already_submitted: - replica_id = await pool.submit_update(context.request_id, request_state, request) - else: - replica_id = await pool.submit_initial(context.request_id, request_state, request, prompt_text=None) - if self._async_chunk and context.stage_id == 0: - prewarmed = await self._prewarm_async_chunk_stages( - context.request_id, - request, - request_state, - ) - if not prewarmed: - # The prewarm already failed the request, aborted it, and - # popped its state. Falling through would write fences and - # a submit timestamp onto that orphaned object and - # re-increment the running counter the cleanup just - # released, then hand the control plane a success result - # for a request that no longer exists. Raise instead: - # ``DuplexControlPlane.handle_append`` turns this into an - # error result for the append. - raise RuntimeError( - f"async-chunk prewarm failed for duplex request {context.request_id}; the request was aborted" - ) - request_state.duplex_stage_fences[context.stage_id] = context.fence - request_state.stage_submit_ts[context.stage_id] = _time.time() - if not request_state.running_counter_registered and self._running_counter is not None: - self._running_counter.increment() - request_state.running_counter_registered = True - return DuplexStageSubmissionResult( - request_id=context.request_id, - stage_id=context.stage_id, - replica_id=replica_id, - ) - - async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: - await self._cleanup_request_ids(request_ids, abort=abort) - - -class Orchestrator: - """Runs inside a background thread's asyncio event loop.""" + Runs inside a background thread's asyncio event loop. Subclasses fill the + template seams (``_dispatch_message``, ``_background_tasks``, + ``_shutdown_extensions``, ``_on_stage_submitted``, ``_intercept_stage_output``, + ``_handle_forward_failure``, ``_cleanup_request_ids``). + """ # Class-level defaults so tests that bypass __init__ via object.__new__ # don't AttributeError when transfer / counter emit paths access them. @@ -424,7 +277,6 @@ class Orchestrator: _transfer_emitter: Any = None _prom_metrics: Any = None _stat_logger: OmniPrometheusStatLogger | None = None - duplex_control_plane: DuplexControlPlanePort | None = None def __init__( self, @@ -442,9 +294,6 @@ class Orchestrator: prom_metrics: Any = None, log_stats: bool = False, enable_orch_monitor: bool = False, - duplex_runtime_extension: DuplexRuntimeExtension | None = None, - enable_duplex_control: bool = False, - duplex_session_config: DuplexSessionRuntimeConfig | None = None, ) -> None: self.request_async_queue = request_async_queue self.output_async_queue = output_async_queue @@ -485,32 +334,6 @@ class Orchestrator: self._cfg_tracker = CfgCompanionTracker() self._stage_input_processors: dict[int, Any] = {} - self.duplex_control_plane: DuplexControlPlanePort | None = None - self._duplex_reaper_interval_s = 1.0 - if enable_duplex_control: - runtime_session_config = duplex_session_config or DuplexSessionRuntimeConfig() - self._duplex_reaper_interval_s = runtime_session_config.reaper_interval_s - - self.duplex_control_plane = DuplexControlPlane( - extension=duplex_runtime_extension, - stage_port=_OrchestratorDuplexStagePort( - stage_pools=self.stage_pools, - request_states=self.request_states, - running_counter=self._running_counter, - cleanup_request_ids=self._cleanup_request_ids, - async_chunk=self.async_chunk, - prewarm_async_chunk_stages=self._prewarm_async_chunk_stages, - ), - result_sink=self.rpc_async_queue, - lifecycle_sink=self.output_async_queue, - lease_config=DuplexLeaseConfig( - idle_ttl_s=runtime_session_config.idle_ttl_s, - disconnect_grace_s=runtime_session_config.disconnect_grace_s, - ), - max_sessions=runtime_session_config.max_sessions, - completed_append_limit=runtime_session_config.completed_append_cache_size, - ) - self._shutdown_event = asyncio.Event() self._stages_shutdown = False self._event_driven_orch = _event_driven_orch_enabled() @@ -588,17 +411,6 @@ class Orchestrator: logger.exception("[Orchestrator] OmniPrometheusStatLogger init failed; metrics wrap disabled") self._stat_logger = None - @property - def duplex_sessions(self) -> DuplexSessionRuntimeManager: - if self.duplex_control_plane is None: - raise RuntimeError("duplex control plane is disabled") - return self.duplex_control_plane.sessions - - def _require_duplex_control_plane(self) -> DuplexControlPlanePort: - if self.duplex_control_plane is None: - raise RuntimeError("duplex control plane is disabled") - return self.duplex_control_plane - async def run(self) -> None: """Main entry point for the Orchestrator event loop.""" logger.info("[Orchestrator] Starting event loop") @@ -620,8 +432,8 @@ class Orchestrator: membership_watcher = self._membership.start() tasks = [request_task, output_task] - if self.duplex_control_plane is not None: - tasks.append(asyncio.create_task(self._duplex_reaper_loop(), name="orchestrator-duplex-reaper")) + for index, coro in enumerate(self._background_tasks()): + tasks.append(asyncio.create_task(coro, name=f"orchestrator-background-{index}")) if membership_watcher is not None: tasks.append(membership_watcher) @@ -655,8 +467,7 @@ class Orchestrator: await asyncio.gather(*tasks, return_exceptions=True) except Exception: pass - if self.duplex_control_plane is not None: - await self.duplex_control_plane.shutdown() + await self._shutdown_extensions() if self._membership is not None: await self._membership.drain_tasks(timeout=10.0) @@ -680,18 +491,10 @@ class Orchestrator: msg = await self.request_async_queue.get() msg_type = msg.type - if msg_type == "add_request": - await self._handle_add_request(msg) - elif msg_type == "streaming_update": - await self._handle_streaming_update(msg) - elif msg_type == "add_companion_request": - await self._handle_add_companion(msg) - elif self.duplex_control_plane is not None and self.duplex_control_plane.accepts(msg): - self.duplex_control_plane.dispatch(msg) - elif msg_type == "abort": + if await self._dispatch_message(msg): + continue + if msg_type == "abort": await self._handle_abort(msg) - elif msg_type == "interaction": - await self._handle_interaction(msg) elif msg_type == "collective_rpc": await self._handle_collective_rpc(msg) elif isinstance(msg, RegisterRemoteReplicaMessage): @@ -717,183 +520,58 @@ class Orchestrator: else: logger.warning("[Orchestrator] Unknown message type: %s", msg_type) - async def _duplex_reaper_loop(self) -> None: - while not self._shutdown_event.is_set(): - try: - await asyncio.wait_for( - self._shutdown_event.wait(), - timeout=self._duplex_reaper_interval_s, - ) - except TimeoutError: - plane = self.duplex_control_plane - if plane is not None: - try: - await plane.reap_expired() - except Exception: - logger.exception("[Orchestrator] Duplex expiry cleanup failed; retrying on next tick") + # ---- Template seams filled by the turn-based / duplex subclasses ---- - async def _handle_add_request(self, msg: StageSubmissionMessage) -> None: - """Handle an add_request message from the main thread.""" - stage_id = 0 - request_id = msg.request_id - prompt = msg.prompt - original_prompt = msg.original_prompt - sampling_params_list = msg.sampling_params_list - if not sampling_params_list: - raise ValueError(f"Missing sampling params for stage 0. Got {len(sampling_params_list)} stage params.") - final_stage_id = msg.final_stage_id - final_output_stage_ids = set(msg.final_output_stage_ids or [final_stage_id]) + async def _dispatch_message(self, msg: EngineQueueMessage) -> bool: + """Handle a subclass-specific request-queue message; return True when consumed.""" + del msg + return False - if not self.stage_pools[stage_id].live_replica_ids(): - # Stage 0 lost all replicas between the HTTP-layer errored check and - # dispatch. Runs before request state / running counter registration, - # so the helper's cleanup is a no-op here. - await self._fail_request_dead_stage(request_id, stage_id) - return + def _background_tasks(self) -> list[Coroutine[Any, Any, None]]: + """Extra coroutines gathered with the request/output handlers in ``run()``.""" + return [] - logger.debug( - "[Orchestrator] _handle_add_request: stage=%s req=%s " - "prompt_type=%s original_prompt_type=%s final_stage=%s " - "num_sampling_params=%d", - stage_id, - request_id, - type(prompt).__name__, - type(original_prompt).__name__, - final_stage_id, - len(sampling_params_list), - ) + async def _shutdown_extensions(self) -> None: + """Runs in ``run()``'s ``finally`` after task cancellation, before membership drain.""" - req_state = OrchestratorRequestState( - request_id=request_id, - prompt=original_prompt, - sampling_params_list=sampling_params_list, - final_stage_id=final_stage_id, - final_output_stage_ids=final_output_stage_ids, - request_timestamp=float(msg.request_timestamp or _time.time()), - mm_features=getattr(prompt, "mm_features", None), - request_artifact_dirs=set(msg.request_artifact_dirs or ()), - ) - self.request_states[request_id] = req_state - self._register_running_request(req_state) - req_state.streaming.enabled = bool(getattr(prompt, "resumable", False)) - req_state.stage_submit_ts[stage_id] = _time.time() - enqueue_ts = msg.enqueue_ts - if enqueue_ts > 0: - req_state.pipeline_timings["queue_wait_ms"] = (_time.perf_counter() - enqueue_ts) * 1000.0 - preprocess_ms = msg.preprocess_ms - if preprocess_ms > 0: - req_state.pipeline_timings["preprocess_ms"] = preprocess_ms - if not await self._dispatch_or_fail_request( - lambda: self.stage_pools[stage_id].submit_initial( - request_id, - req_state, - prompt, - prompt_text=msg.output_prompt_text, - ), - req_id=request_id, - stage_id=stage_id, - operation="add_request", - ): - return + def _on_stage_submitted( + self, + stage_id: int, + request_id: str, + replica_id: int, + req_state: OrchestratorRequestState, + ) -> None: + """Called after every successful ``submit_initial`` / ``submit_update``.""" + del stage_id, request_id, replica_id, req_state - if self.async_chunk and stage_id == 0 and final_stage_id > 0: - await self._prewarm_async_chunk_stages(request_id, prompt, req_state) + async def _intercept_stage_output( + self, + stage_id: int, + replica_id: int, + output: Any, + req_state: OrchestratorRequestState, + stage_metrics: Any, + submit_ts: float | None, + ) -> bool: + """Offer a processed output to the subclass before forwarding. - async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None: - """Handle a streaming_update message for an existing request.""" - stage_id = 0 - request_id = msg.request_id - request = msg.prompt - final_stage_id = msg.final_stage_id - req_state = self.request_states.get(request_id) - if req_state is None: - # Streaming updates always follow the first-chunk add_request - # through the same ordered submission queue, so an unknown id - # here can only mean the request already finished or was - # aborted (e.g. the client disconnected). Re-adding it would - # resurrect a headless session that keeps cycling through the - # stages with nobody consuming its outputs (issue #4271). - logger.warning( - "[Orchestrator] streaming_update for unknown req=%s; dropping (request finished or aborted)", - request_id, - ) - return + Return True when the output must not be forwarded to the next stage + (the subclass consumed it). Client emission for session-owned requests + is already suppressed by ``OrchestratorRequestState.session_owned``. + """ + del stage_id, replica_id, output, req_state, stage_metrics, submit_ts + return False - if msg.sampling_params_list: - req_state.sampling_params_list = msg.sampling_params_list - - req_state.streaming.enabled = True - req_state.stage_submit_ts[stage_id] = _time.time() - if not await self._dispatch_or_fail_request( - lambda: self.stage_pools[stage_id].submit_update( - request_id, - req_state, - request, - prompt_text=msg.output_prompt_text, - ), - req_id=request_id, - stage_id=stage_id, - operation="streaming_update", - ): - return - - if self.async_chunk and stage_id == 0 and final_stage_id > 0: - await self._prewarm_async_chunk_stages(request_id, request, req_state) - - async def _handle_add_companion(self, msg: AddCompanionRequestMessage) -> None: - """Handle an add_companion_request message: submit companion to stage 0.""" - companion_id = msg.companion_id - parent_id = msg.parent_id - role = msg.role - companion_prompt = msg.prompt - sampling_params_list = msg.sampling_params_list - - parent_state = self.request_states.get(parent_id) - if parent_state is None: - logger.info( - "[Orchestrator] Dropping CFG companion %s (role=%s): parent %s is no longer active", - companion_id, - role, - parent_id, - ) - return - - self._cfg_tracker.register_companion(parent_id, role, companion_id) - - companion_state = OrchestratorRequestState( - request_id=companion_id, - prompt=companion_prompt, - sampling_params_list=sampling_params_list, - final_stage_id=0, - final_output_stage_ids={0}, - request_timestamp=parent_state.request_timestamp, - ) - self.request_states[companion_id] = companion_state - companion_state.stage_submit_ts[0] = _time.time() - # A dead-stage failure is attributed to the parent request; its cleanup - # also releases this companion. - if not await self._dispatch_or_fail_request( - lambda: self.stage_pools[0].submit_initial( - companion_id, - companion_state, - companion_prompt, - prompt_text=msg.companion_prompt_text, - affinity_request_id=parent_id, - ), - req_id=parent_id, - stage_id=0, - operation=f"companion {companion_id}", - dispatch_req_id=companion_id, - ): - return - - logger.info( - "[Orchestrator] CFG companion submitted: %s (role=%s, parent=%s, stage-0 replica-%s)", - companion_id, - role, - parent_id, - self.stage_pools[0].get_bound_replica_id(companion_id), - ) + async def _handle_forward_failure( + self, + req_id: str, + next_stage_id: int, + req_state: OrchestratorRequestState, + exc: BaseException, + ) -> bool: + """Return True when a next-stage input-processing failure was absorbed (else it re-raises).""" + del req_id, next_stage_id, req_state, exc + return False async def _handle_abort(self, msg: AbortRequestMessage) -> None: """Handle an abort message from the main thread. @@ -935,44 +613,6 @@ class Orchestrator: for output_msg in abort_outputs: await self.output_async_queue.put(output_msg) - async def _handle_interaction(self, msg: InteractionMessage) -> None: - """Handle a midway interaction for an active streaming diffusion request.""" - stage_id = 0 - request_id = msg.request_id - event_id = msg.interaction.get("event_id") - req_state = self.request_states.get(request_id) - if req_state is None: - logger.info("[Orchestrator] Dropping interaction for inactive req %s", request_id) - await self.output_async_queue.put( - ErrorMessage( - error=f"No active request for interaction: {request_id}", - fatal=False, - request_id=request_id, - event_id=event_id, - stage_id=stage_id, - ) - ) - return - - try: - await self.stage_pools[stage_id].submit_interaction(request_id, msg.interaction) - except Exception as exc: - logger.info( - "[Orchestrator] Failed interaction for req %s: %s", - request_id, - exc, - exc_info=True, - ) - await self.output_async_queue.put( - ErrorMessage( - error=f"Failed interaction for request {request_id}: {exc}", - fatal=False, - request_id=request_id, - event_id=event_id, - stage_id=stage_id, - ) - ) - async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage]: """Forward abort requests to all stage pools. @@ -1539,7 +1179,7 @@ class Orchestrator: await self._cleanup_request_ids( [parent_id, *self._cfg_tracker.cleanup_parent(parent_id)], abort=True, - close_duplex_sessions=True, + release_owners=True, ) def _update_stage_replica_waiting(self, stage_id: int, replica_id: int, n_waiting: int) -> None: @@ -1615,6 +1255,13 @@ class Orchestrator: await self._cleanup_request_ids( [req_id, *self._cfg_tracker.cleanup_parent(req_id)], abort=True, + # A session-owned request has an owner that outlives the stage + # request: without this the runner and its admission slot + # survive the replica that was serving them, and the client is + # told nothing -- the request-scoped ErrorMessage has no + # frontend request_states entry to land in. No-op for + # turn-based requests, which have no owner. + release_owners=True, ) async def _fail_request_dead_stage(self, req_id: str, stage_id: int) -> None: @@ -1642,6 +1289,7 @@ class Orchestrator: await self._cleanup_request_ids( [req_id, *self._cfg_tracker.cleanup_parent(req_id)], abort=True, + release_owners=True, ) async def _fail_request_client_error( @@ -1652,7 +1300,7 @@ class Orchestrator: *, status_code: int = HTTPStatus.BAD_REQUEST.value, error_type: str = DEFAULT_CLIENT_ERROR_TYPE, - close_duplex_sessions: bool = False, + release_owners: bool = False, ) -> None: """Fail one request with a non-fatal client error (400 by default). @@ -1672,7 +1320,7 @@ class Orchestrator: await self._cleanup_request_ids( [req_id, *self._cfg_tracker.cleanup_parent(req_id)], abort=True, - close_duplex_sessions=close_duplex_sessions, + release_owners=release_owners, ) async def _dispatch_or_fail_request( @@ -1767,7 +1415,7 @@ class Orchestrator: request_ids: list[str], *, abort: bool = False, - close_duplex_sessions: bool = False, + release_owners: bool = False, ) -> list[OutputMessage]: """Release pool bindings and logical request state for the given ids. @@ -1778,10 +1426,14 @@ class Orchestrator: loss, membership unregister) funnels through here, so tracker state cannot outlive its requests. + ``release_owners`` asks the subclass to also release any session that + owns these requests (error paths pass True); the base ignores it. + Returns: Final-stage AR abort ``OutputMessage`` list when ``abort=True``; otherwise an empty list. """ + del release_owners if not request_ids: return [] @@ -1808,42 +1460,18 @@ class Orchestrator: if cid not in batch: batch.add(cid) cleanup_ids.append(cid) - closing_session_ids: list[str] = [] - if close_duplex_sessions and self.duplex_control_plane is not None: - closed_sessions = self.duplex_control_plane.close_sessions_for_request_ids( - cleanup_ids, - abort=abort, - cleanup_in_progress=True, - ) - closing_session_ids.extend(closed_sessions) - for session_id, stale_request_ids in closed_sessions.items(): - logger.info( - "[Orchestrator] closed duplex session %s while cleaning failed request ids %s", - session_id, - stale_request_ids, - ) - cleanup_ids.extend(stale_request_ids) - cleanup_ids = list(dict.fromkeys(cleanup_ids)) - abort_outputs: list[OutputMessage] = [] - try: - if abort: - abort_outputs = await self._abort_request_ids(cleanup_ids) - self._release_request_bindings(cleanup_ids) - for request_id in cleanup_ids: - self._pd_kv_params.pop(request_id, None) - req_state = self.request_states.pop(request_id, None) - if req_state is not None: - cleanup_request_artifact_dirs(getattr(req_state, "request_artifact_dirs", ())) - if req_state is not None and req_state.running_counter_registered and self._running_counter is not None: - self._running_counter.decrement() - req_state.running_counter_registered = False - except BaseException: - if closing_session_ids and self.duplex_control_plane is not None: - self.duplex_control_plane.defer_request_cleanups(closing_session_ids) - raise - if closing_session_ids and self.duplex_control_plane is not None: - self.duplex_control_plane.finalize_closed_sessions(closing_session_ids) + if abort: + abort_outputs = await self._abort_request_ids(cleanup_ids) + self._release_request_bindings(cleanup_ids) + for request_id in cleanup_ids: + self._pd_kv_params.pop(request_id, None) + req_state = self.request_states.pop(request_id, None) + if req_state is not None: + cleanup_request_artifact_dirs(getattr(req_state, "request_artifact_dirs", ())) + if req_state is not None and req_state.running_counter_registered and self._running_counter is not None: + self._running_counter.decrement() + req_state.running_counter_registered = False return abort_outputs async def _apply_raw_terminal_stage_finish( @@ -1888,7 +1516,7 @@ class Orchestrator: for request_id in request_ids: req_state = self.request_states.get(request_id) - if req_state is None or self._is_duplex_session_request(req_state): + if req_state is None or req_state.session_owned: continue final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} @@ -1937,28 +1565,6 @@ class Orchestrator: params.cfg_kv_request_ids = companion_request_ids return params - def _duplex_session_for_req_state(self, req_state: OrchestratorRequestState) -> DuplexSessionRuntimeState | None: - if self.duplex_control_plane is None: - return None - return self.duplex_control_plane.session_for_identity(req_state.duplex_identity) - - def _record_duplex_stage_submission( - self, - stage_id: int, - request_id: str, - replica_id: int, - req_state: OrchestratorRequestState, - ) -> None: - del replica_id - identity = req_state.duplex_identity - session = self._duplex_session_for_req_state(req_state) - if identity is None or session is None: - return - req_state.duplex_stage_fences[stage_id] = identity.fence - session.bind_stage_request(stage_id, request_id, fence=identity.fence) - req_state.stage_submit_ts[stage_id] = _time.time() - self._register_running_request(req_state) - def _register_running_request(self, req_state: OrchestratorRequestState) -> None: if req_state.running_counter_registered or self._running_counter is None: return @@ -1991,17 +1597,11 @@ class Orchestrator: req_state.finished_final_output_stage_ids.add(stage_id) final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} request_finished = final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids) - # Duplex stage-0 segment boundaries are not client-visible outputs: - # direct decisions are emitted by the model runtime extension below, - # while spoken content flows through the next stage. Forwarding - # the raw stage-0 output as well injects one cumulative-text, - # no-audio message per unit that every downstream consumer must - # filter out again (the official implementation returns exactly one - # result per audio chunk). - is_duplex_stage0_segment = ( - stage_id == 0 and self._is_duplex_session_request(req_state) and req_state.streaming.segment(0).finished - ) - if self.stage_pools[stage_id].final_output and not is_duplex_stage0_segment: + if req_state.session_owned: + # Session-owned outputs are delivered to their session through + # ``_intercept_stage_output`` below, never to a request queue. + pass + elif self.stage_pools[stage_id].final_output: await self.output_async_queue.put( OutputMessage( request_id=req_id, @@ -2009,13 +1609,7 @@ class Orchestrator: replica_id=replica_id, engine_outputs=output, metrics=stage_metrics, - finished=( - request_finished - or ( - self._is_duplex_session_request(req_state) - and req_state.streaming.segment(stage_id).finished - ) - ), + finished=request_finished, stage_submit_ts=submit_ts, ) ) @@ -2036,16 +1630,7 @@ class Orchestrator: self._pd_kv_params[req_id] = kv_params if isinstance(kv_params, dict) else dict(kv_params) req_state.pd_prefill_multimodal_output = getattr(output, "multimodal_output", None) - duplex_output_decision = self._duplex_output_decision(stage_id, output, req_state) - if duplex_output_decision is not None: - await self._emit_duplex_direct_output( - stage_id, - req_id, - output, - duplex_output_decision, - stage_metrics, - submit_ts, - ) + if await self._intercept_stage_output(stage_id, replica_id, output, req_state, stage_metrics, submit_ts): return if ( @@ -2082,7 +1667,7 @@ class Orchestrator: and self.request_states.get(req_id) is req_state and finished and not final_only_finished - and not self._is_duplex_session_request(req_state) + and not req_state.session_owned ): # For streaming sessions, send the terminal (resumable=False) update only on a finish await self._forward_to_next_stage( @@ -2095,7 +1680,7 @@ class Orchestrator: is_final_update=True, ) - if request_finished and not self._is_duplex_session_request(req_state): + if request_finished and not req_state.session_owned: await self._cleanup_request_ids([req_id, *self._cfg_tracker.cleanup_parent(req_id)]) def _next_stage_already_submitted(self, stage_id: int, req_state: OrchestratorRequestState) -> bool: @@ -2189,92 +1774,6 @@ class Orchestrator: request.payload_sender_info = payload_sender_info return request - @staticmethod - def _duplex_output_context( - req_state: OrchestratorRequestState, - *, - # Required: segment state is stage-keyed, so a caller that omitted this - # would silently read "no boundary" rather than the intended stage's. - stage_id: int | None, - ) -> DuplexOutputContext | None: - identity = req_state.duplex_identity - if identity is None: - return None - fence = req_state.duplex_stage_fences.get(stage_id, identity.fence) if stage_id is not None else identity.fence - segment = req_state.streaming.segment(stage_id) - return DuplexOutputContext( - identity=DuplexRequestIdentity( - session_id=identity.session_id, - fence=fence, - ), - final_stage_id=req_state.final_stage_id, - segment_finished=req_state.streaming.enabled and segment.finished, - segment_token_ids=tuple(segment.token_ids), - segment_output_metadata=segment.output_metadata, - ) - - @staticmethod - def _is_duplex_session_request(req_state: OrchestratorRequestState) -> bool: - return req_state.duplex_identity is not None - - @staticmethod - def _duplex_fence_for_req_state( - req_state: OrchestratorRequestState, - *, - stage_id: int | None = None, - ) -> DuplexFence | None: - context = Orchestrator._duplex_output_context(req_state, stage_id=stage_id) - return context.identity.fence if context is not None else None - - def _duplex_output_decision( - self, - stage_id: int, - output: Any, - req_state: OrchestratorRequestState, - ) -> DuplexOutputDecision | None: - if self.duplex_control_plane is None: - return None - context = self._duplex_output_context(req_state, stage_id=stage_id) - decision = self.duplex_control_plane.decide_output( - stage_id, - output, - context, - ) - return decision - - async def _emit_duplex_direct_output( - self, - stage_id: int, - req_id: str, - output: Any, - decision: DuplexOutputDecision, - stage_metrics: Any, - submit_ts: float | None, - ) -> None: - action = getattr(decision.action, "value", decision.action) - if action != "direct_response": - raise ValueError(f"Unsupported duplex output action: {action}") - engine_output = attach_duplex_output_decision( - OmniRequestOutput.from_stage_output( - output, - request_id=req_id, - finished=True, - stage_id=stage_id, - final_output_type=decision.final_output_type, - ), - decision, - ) - await self.output_async_queue.put( - OutputMessage( - request_id=req_id, - stage_id=stage_id, - engine_outputs=engine_output, - metrics=stage_metrics, - finished=True, - stage_submit_ts=submit_ts, - ) - ) - @staticmethod def _completion_multimodal_output(output: Any, completion: Any) -> dict[str, Any]: mm_output = getattr(output, "multimodal_output", None) @@ -2653,7 +2152,7 @@ class Orchestrator: }, params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params), ) - self._record_duplex_stage_submission( + self._on_stage_submitted( next_logical, req_id, replica_id, @@ -2706,7 +2205,7 @@ class Orchestrator: replica_id = await next_pool.submit_update(req_id, req_state, request) else: replica_id = await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) - self._record_duplex_stage_submission( + self._on_stage_submitted( next_logical, req_id, replica_id, @@ -2751,7 +2250,7 @@ class Orchestrator: str(exc), status_code=exc.status_code, error_type=exc.error_type, - close_duplex_sessions=self._is_duplex_session_request(req_state), + release_owners=req_state.session_owned, ) return except Exception as exc: @@ -2760,21 +2259,8 @@ class Orchestrator: req_id, next_logical, ) - if not self._is_duplex_session_request(req_state): + if not await self._handle_forward_failure(req_id, next_logical, req_state, exc): raise - await self.output_async_queue.put( - ErrorMessage( - request_id=req_id, - stage_id=next_logical, - error=f"Stage-{next_logical} input processor failed: {type(exc).__name__}: {exc}", - error_type=type(exc).__name__, - ) - ) - await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], - abort=True, - close_duplex_sessions=True, - ) return finally: req_state.streaming.source_token_decoder = previous_decoder @@ -2842,7 +2328,7 @@ class Orchestrator: replica_id = await next_pool.submit_update(req_id, req_state, request) else: replica_id = await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) - self._record_duplex_stage_submission( + self._on_stage_submitted( next_logical, req_id, replica_id, @@ -2900,7 +2386,7 @@ class Orchestrator: 0, "async_chunk requires prompt_token_ids on the stage-0 request; " "an embeds-only prompt cannot prewarm downstream stages", - close_duplex_sessions=True, + release_owners=True, ) return False @@ -2996,7 +2482,7 @@ class Orchestrator: # The guard swallows submit_initial's return value; the pool binding # it recorded carries the same replica. bound_replica_id = next_pool.get_bound_replica_id(request_id) - self._record_duplex_stage_submission( + self._on_stage_submitted( next_stage_id, request_id, bound_replica_id if bound_replica_id is not None else 0, @@ -3081,3 +2567,222 @@ class Orchestrator: for pool in self.stage_pools: for replica_id in pool.live_replica_ids(): pool.shutdown_replica(replica_id) + + +class Orchestrator(OrchestratorBase): + """Turn-based orchestrator: admits ``add_request`` / streaming / companion / interaction messages.""" + + async def _dispatch_message(self, msg: EngineQueueMessage) -> bool: + msg_type = msg.type + if msg_type == "add_request": + await self._handle_add_request(msg) + elif msg_type == "streaming_update": + await self._handle_streaming_update(msg) + elif msg_type == "add_companion_request": + await self._handle_add_companion(msg) + elif msg_type == "interaction": + await self._handle_interaction(msg) + else: + return False + return True + + async def _handle_add_request(self, msg: StageSubmissionMessage) -> None: + """Handle an add_request message from the main thread.""" + stage_id = 0 + request_id = msg.request_id + prompt = msg.prompt + original_prompt = msg.original_prompt + sampling_params_list = msg.sampling_params_list + if not sampling_params_list: + raise ValueError(f"Missing sampling params for stage 0. Got {len(sampling_params_list)} stage params.") + final_stage_id = msg.final_stage_id + final_output_stage_ids = set(msg.final_output_stage_ids or [final_stage_id]) + + if not self.stage_pools[stage_id].live_replica_ids(): + # Stage 0 lost all replicas between the HTTP-layer errored check and + # dispatch. Runs before request state / running counter registration, + # so the helper's cleanup is a no-op here. + await self._fail_request_dead_stage(request_id, stage_id) + return + + logger.debug( + "[Orchestrator] _handle_add_request: stage=%s req=%s " + "prompt_type=%s original_prompt_type=%s final_stage=%s " + "num_sampling_params=%d", + stage_id, + request_id, + type(prompt).__name__, + type(original_prompt).__name__, + final_stage_id, + len(sampling_params_list), + ) + + req_state = OrchestratorRequestState( + request_id=request_id, + prompt=original_prompt, + sampling_params_list=sampling_params_list, + final_stage_id=final_stage_id, + final_output_stage_ids=final_output_stage_ids, + request_timestamp=float(msg.request_timestamp or _time.time()), + mm_features=getattr(prompt, "mm_features", None), + request_artifact_dirs=set(msg.request_artifact_dirs or ()), + ) + self.request_states[request_id] = req_state + self._register_running_request(req_state) + req_state.streaming.enabled = bool(getattr(prompt, "resumable", False)) + req_state.stage_submit_ts[stage_id] = _time.time() + enqueue_ts = msg.enqueue_ts + if enqueue_ts > 0: + req_state.pipeline_timings["queue_wait_ms"] = (_time.perf_counter() - enqueue_ts) * 1000.0 + preprocess_ms = msg.preprocess_ms + if preprocess_ms > 0: + req_state.pipeline_timings["preprocess_ms"] = preprocess_ms + if not await self._dispatch_or_fail_request( + lambda: self.stage_pools[stage_id].submit_initial( + request_id, + req_state, + prompt, + prompt_text=msg.output_prompt_text, + ), + req_id=request_id, + stage_id=stage_id, + operation="add_request", + ): + return + + if self.async_chunk and stage_id == 0 and final_stage_id > 0: + await self._prewarm_async_chunk_stages(request_id, prompt, req_state) + + async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None: + """Handle a streaming_update message for an existing request.""" + stage_id = 0 + request_id = msg.request_id + request = msg.prompt + final_stage_id = msg.final_stage_id + req_state = self.request_states.get(request_id) + if req_state is None: + # Streaming updates always follow the first-chunk add_request + # through the same ordered submission queue, so an unknown id + # here can only mean the request already finished or was + # aborted (e.g. the client disconnected). Re-adding it would + # resurrect a headless session that keeps cycling through the + # stages with nobody consuming its outputs (issue #4271). + logger.warning( + "[Orchestrator] streaming_update for unknown req=%s; dropping (request finished or aborted)", + request_id, + ) + return + + if msg.sampling_params_list: + req_state.sampling_params_list = msg.sampling_params_list + + req_state.streaming.enabled = True + req_state.stage_submit_ts[stage_id] = _time.time() + if not await self._dispatch_or_fail_request( + lambda: self.stage_pools[stage_id].submit_update( + request_id, + req_state, + request, + prompt_text=msg.output_prompt_text, + ), + req_id=request_id, + stage_id=stage_id, + operation="streaming_update", + ): + return + + if self.async_chunk and stage_id == 0 and final_stage_id > 0: + await self._prewarm_async_chunk_stages(request_id, request, req_state) + + async def _handle_add_companion(self, msg: AddCompanionRequestMessage) -> None: + """Handle an add_companion_request message: submit companion to stage 0.""" + companion_id = msg.companion_id + parent_id = msg.parent_id + role = msg.role + companion_prompt = msg.prompt + sampling_params_list = msg.sampling_params_list + + parent_state = self.request_states.get(parent_id) + if parent_state is None: + logger.info( + "[Orchestrator] Dropping CFG companion %s (role=%s): parent %s is no longer active", + companion_id, + role, + parent_id, + ) + return + + self._cfg_tracker.register_companion(parent_id, role, companion_id) + + companion_state = OrchestratorRequestState( + request_id=companion_id, + prompt=companion_prompt, + sampling_params_list=sampling_params_list, + final_stage_id=0, + final_output_stage_ids={0}, + request_timestamp=parent_state.request_timestamp, + ) + self.request_states[companion_id] = companion_state + companion_state.stage_submit_ts[0] = _time.time() + # A dead-stage failure is attributed to the parent request; its cleanup + # also releases this companion. + if not await self._dispatch_or_fail_request( + lambda: self.stage_pools[0].submit_initial( + companion_id, + companion_state, + companion_prompt, + prompt_text=msg.companion_prompt_text, + affinity_request_id=parent_id, + ), + req_id=parent_id, + stage_id=0, + operation=f"companion {companion_id}", + dispatch_req_id=companion_id, + ): + return + + logger.info( + "[Orchestrator] CFG companion submitted: %s (role=%s, parent=%s, stage-0 replica-%s)", + companion_id, + role, + parent_id, + self.stage_pools[0].get_bound_replica_id(companion_id), + ) + + async def _handle_interaction(self, msg: InteractionMessage) -> None: + """Handle a midway interaction for an active streaming diffusion request.""" + stage_id = 0 + request_id = msg.request_id + event_id = msg.interaction.get("event_id") + req_state = self.request_states.get(request_id) + if req_state is None: + logger.info("[Orchestrator] Dropping interaction for inactive req %s", request_id) + await self.output_async_queue.put( + ErrorMessage( + error=f"No active request for interaction: {request_id}", + fatal=False, + request_id=request_id, + event_id=event_id, + stage_id=stage_id, + ) + ) + return + + try: + await self.stage_pools[stage_id].submit_interaction(request_id, msg.interaction) + except Exception as exc: + logger.info( + "[Orchestrator] Failed interaction for req %s: %s", + request_id, + exc, + exc_info=True, + ) + await self.output_async_queue.put( + ErrorMessage( + error=f"Failed interaction for request {request_id}: {exc}", + fatal=False, + request_id=request_id, + event_id=event_id, + stage_id=stage_id, + ) + ) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 993f0ca1e..6114bbea2 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -11,7 +11,6 @@ StageEngineCoreClient instances) instead of OmniStage with worker processes. from __future__ import annotations import asyncio -import os import time import uuid from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence @@ -21,31 +20,17 @@ from vllm import TokensPrompt from vllm.engine.protocol import EngineClient, StreamingInput from vllm.logger import init_logger from vllm.lora.request import LoRARequest -from vllm.outputs import CompletionOutput, PoolingRequestOutput -from vllm.plugins.io_processors import get_io_processor +from vllm.outputs import PoolingRequestOutput from vllm.pooling_params import PoolingParams from vllm.renderers.inputs.preprocess import extract_prompt_components from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.tasks import SupportedTask -from vllm.utils import random_uuid -from vllm.v1.engine.exceptions import EngineDeadError from vllm_omni.diffusion.data import CuMemTag, OmniACK, OmniSleepTask, OmniWakeTask -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity -from vllm_omni.engine.duplex.messages import ( - DuplexFence, - DuplexSessionLifecycleMessage, -) -from vllm_omni.engine.messages import ErrorMessage, OutputMessage +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.engine.messages import ErrorMessage +from vllm_omni.entrypoints.async_omni_base import ABORT_TIMEOUT_S, AsyncOmniBase from vllm_omni.entrypoints.client_request_state import ClientRequestState -from vllm_omni.entrypoints.duplex_request_client import ( - DuplexRequestClient, - DuplexRequestOutputPort, -) -from vllm_omni.entrypoints.omni_base import ( - OmniBase, - OmniEngineDeadError, -) from vllm_omni.errors import client_error_metadata from vllm_omni.inputs.data import OmniSamplingParams from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics @@ -59,74 +44,9 @@ if TYPE_CHECKING: from vllm_omni.inputs.data import OmniInteractionPrompt, OmniPromptType logger = init_logger(__name__) -_FINAL_OUTPUT_IDLE_SLEEP_S = 0.001 -# Blocking-wait interval for the event-driven final-output drain -# (VLLM_OMNI_EVENT_DRIVEN_ORCH=1): a message wakes the drain immediately via -# the janus queue's condition variable; this timeout only bounds how often the -# orchestrator liveness check runs while the pipeline is idle. -_FINAL_OUTPUT_BLOCKING_WAIT_S = 1.0 -# Shared DELETE / generate() cleanup abort bound. Env is the documented knob. -ABORT_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_ABORT_TIMEOUT", 2.0)) -class AsyncEventResolver: - """ - A generic signal aggregator designed for synchronized handshakes in - distributed or multi-stage environments. Supports waiting for a specified - number (expected_count) of worker signals in both inline and multiprocess modes. - """ - - def __init__(self, orchestrator=None): - self._pending_tasks: dict[str, dict] = {} - self.orchestrator = orchestrator - self._lock = asyncio.Lock() - - def watch_task(self, task_id: str, expected_count: int = 1) -> asyncio.Future: - loop = asyncio.get_running_loop() - fut = loop.create_future() - self._pending_tasks[task_id] = { - "future": fut, - "expected_count": expected_count, - "received": [], - "start_time": time.time(), - } - return fut - - async def resolve(self, ack: OmniACK): - tid = getattr(ack, "task_id", None) - - if tid is None and isinstance(ack, dict): - tid = ack.get("task_id") - - async with self._lock: - task_info = self._pending_tasks.get(tid) - if task_info is None: - logger.warning(f"Received stray ACK for task_id {tid}. Task might have timed out.") - return - - task_info["received"].append(ack) - current_count = len(task_info["received"]) - expected = task_info["expected_count"] - - orchestrator = self.orchestrator - if orchestrator and hasattr(orchestrator, "metrics") and orchestrator.metrics: - freed = getattr(ack, "freed_bytes", 0) - if freed == 0 and isinstance(ack, dict): - freed = ack.get("freed_bytes", 0) - orchestrator.metrics.record_vram_reclaimed(freed) - - logger.info(f"[Resolver] Task {tid} progress: {current_count}/{expected} ACKs received.") - - if current_count >= expected: - self._pending_tasks.pop(tid) - fut = task_info["future"] - if not fut.done(): - elapsed = time.time() - task_info["start_time"] - logger.info(f"[Resolver] Task {tid} completed successfully in {elapsed:.2f}s.") - fut.set_result(task_info["received"]) - - -class AsyncOmni(EngineClient, OmniBase): +class AsyncOmni(AsyncOmniBase, EngineClient): """Asynchronous unified entry point for multi-stage pipelines using AsyncOmniEngine. This is the refactored version that uses AsyncOmniEngine instead of @@ -155,8 +75,13 @@ class AsyncOmni(EngineClient, OmniBase): ... print(output) """ - def __init__(self, *args: Any, model: str = "", **kwargs: Any) -> None: - OmniBase.__init__(self, model=model, **kwargs) + def _create_engine(self, **engine_kwargs: Any) -> AsyncOmniEngine: + return AsyncOmniEngine(**engine_kwargs) + + def __init__(self, model: str = "", *args: Any, **kwargs: Any) -> None: + tts_max_instructions_length = kwargs.get("tts_max_instructions_length", None) + AsyncOmniBase.__init__(self, *args, model=model, **kwargs) + self.tts_max_instructions_length = tts_max_instructions_length self._pause_cond: asyncio.Condition = asyncio.Condition() self._paused: bool = False # In-flight EngineCore submits (non-streaming add_request, or each @@ -173,294 +98,6 @@ class AsyncOmni(EngineClient, OmniBase): self._sleeping_tags: set[str] = set() self._stage_sleeping_tags: dict[int, set[str]] = {} self._level2_sleeping: bool = False - self._duplex_request_client: DuplexRequestClient | None = None - self.duplex_lifecycle_events: asyncio.Queue[DuplexSessionLifecycleMessage] = asyncio.Queue() - self.final_output_task: asyncio.Task | None = None - self.event_resolver = AsyncEventResolver(orchestrator=self) - self.config_path = self.engine.config_path - self.tts_max_instructions_length = kwargs.get("tts_max_instructions_length", None) - self.input_processor = self.engine.input_processor - self.endpoint_restrictions = self.engine.endpoint_restrictions - self.duplex_session_config = self.engine.duplex_session_config - self.duplex_serving_adapter_path = self.engine.duplex_serving_adapter_path - - stage_index = self._get_comprehension_stage_index() - if stage_index is None: - self.io_processor = None - else: - vllm_config = self.engine.stage_vllm_configs[stage_index] - io_processor_plugin = vllm_config.model_config.io_processor_plugin - renderer = self.renderer - if renderer is None: - from vllm.renderers import renderer_from_config - - renderer = renderer_from_config(vllm_config) - self.io_processor = get_io_processor(vllm_config, renderer, io_processor_plugin) - - def _resolve_transfer_replica(self, stage_id: int, request_id: str) -> int | None: - """Look up the sticky-routed replica for (stage_id, request_id). - - Used as the ``replica_resolver`` callback by ``OrchestratorAggregator`` - to label transfer_* metrics without plumbing replica ids through - ``TransferEdgeStats`` / ``StageRequestStats`` / connector adapters. - Returns None when stage_id is out of range or the request hasn't been - bound to a replica yet — the metric emit then defensive-skips. - """ - pools = getattr(self.engine, "stage_pools", None) - if pools is None or not (0 <= stage_id < len(pools)): - return None - return pools[stage_id].get_bound_replica_id(request_id) - - def _get_comprehension_stage_index(self) -> int | None: - fallback_idx: int | None = None - for idx, stage_client in enumerate(self.engine.stage_clients): - stage_vllm_config = self.engine.stage_vllm_configs[idx] - if stage_vllm_config is None: - continue - if fallback_idx is None: - fallback_idx = idx - if stage_client.is_comprehension: - return idx - return fallback_idx - - @property - def renderer(self): - """Return the renderer from the engine input processor when available.""" - if self.input_processor is None: - return None - return self.input_processor.renderer - - @property - def vllm_config(self): - """Return the vLLM config for the comprehension stage when present.""" - stage_index = self._get_comprehension_stage_index() - if stage_index is None: - return None - return self.engine.stage_vllm_configs[stage_index] - - async def get_vllm_config(self) -> Any: - """Compatibility helper for call sites expecting async vllm config access.""" - return self.vllm_config - - def get_diffusion_od_config(self) -> Any | None: - """Return the diffusion-stage config when the pipeline has one.""" - saw_diffusion_stage = False - for stage_client in self.engine.stage_clients: - if getattr(stage_client, "stage_type", None) != "diffusion": - continue - - saw_diffusion_stage = True - - od_config = getattr(stage_client, "od_config", None) - if od_config is not None: - return od_config - - inner_engine = getattr(stage_client, "_engine", None) - od_config = getattr(inner_engine, "od_config", None) - if od_config is not None: - return od_config - - # Out-of-process diffusion clients don't carry od_config (it lives in the - # worker); fall back to the engine's model_class_name resolution. - if saw_diffusion_stage: - return self.engine.get_diffusion_od_config() - - return None - - @property - def model_config(self): - """Return the model config for the comprehension stage when present.""" - vllm_config = self.vllm_config - if vllm_config is None: - return None - return vllm_config.model_config - - @staticmethod - def _get_unique_request_id(external_request_id: str): - """Get a random new request ID for this request; at the server level, - this is usually set by the calling entrypoint, but in direct calls, we - need to set it explicitly since we do not allow empty IDs. - - NOTE: in the upstream vLLM, this is done in the InputProcessor's - `assign_request_id`. - """ - uuid = random_uuid() - prefix = "" if not external_request_id else f"{external_request_id}-" - return f"{prefix}{uuid:.8}" - - async def open_duplex_session_async( - self, - session_id: str, - *, - session_mode: str = "duplex", - capabilities: dict[str, object] | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - fence: DuplexFence, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Open an engine-level duplex session when the backend supports it.""" - return await self._get_duplex_request_client().open( - session_id, - session_mode=session_mode, - capabilities=capabilities, - session_config=session_config, - runtime_config=runtime_config, - fence=fence, - timeout=timeout, - ) - - async def append_duplex_input_async( - self, - session_id: str, - *, - mode: str, - payload: object, - operation_id: str | None = None, - final: bool = False, - expected_epoch: int | None = None, - fence: DuplexFence, - timeout: float | None = 10.0, - collect_outputs: bool = True, - ) -> dict[str, object]: - """Append input to an engine-level duplex session.""" - return await self._get_duplex_request_client().append( - session_id, - mode=mode, - payload=payload, - operation_id=operation_id, - final=final, - expected_epoch=expected_epoch, - fence=fence, - timeout=timeout, - collect_outputs=collect_outputs, - ) - - async def collect_duplex_data_plane_outputs_async( - self, - request_id: str, - *, - response_stage_id: int | None = None, - timeout: float | None = 10.0, - ) -> list[OmniRequestOutput]: - """Collect the next duplex data-plane output batch for a live request.""" - return await self._get_duplex_request_client().collect_registered_outputs( - request_id, - response_stage_id=response_stage_id, - timeout=timeout, - ) - - async def signal_duplex_turn_async( - self, - session_id: str, - *, - event: str, - fence: DuplexFence, - next_fence: DuplexFence | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Send a turn/control signal to an engine-level duplex session.""" - return await self._get_duplex_request_client().signal( - session_id, - event=event, - fence=fence, - next_fence=next_fence, - session_config=session_config, - runtime_config=runtime_config, - timeout=timeout, - ) - - async def close_duplex_session_async( - self, - session_id: str, - *, - reason: str = "client_close", - fence: DuplexFence, - timeout: float | None = 10.0, - ) -> dict[str, object]: - """Close an engine-level duplex session.""" - return await self._get_duplex_request_client().close( - session_id, - reason=reason, - fence=fence, - timeout=timeout, - ) - - async def touch_duplex_session_async( - self, - session_id: str, - *, - fence: DuplexFence, - activity: DuplexLeaseActivity, - timeout: float | None = 10.0, - ) -> dict[str, object]: - return await self._get_duplex_request_client().touch( - session_id, - fence=fence, - activity=activity, - timeout=timeout, - ) - - async def resume_duplex_session_async( - self, - session_id: str, - *, - fence: DuplexFence, - expected_lease_generation: int, - timeout: float | None = 10.0, - ) -> dict[str, object]: - return await self._get_duplex_request_client().resume( - session_id, - fence=fence, - expected_lease_generation=expected_lease_generation, - timeout=timeout, - ) - - def _get_duplex_request_client(self) -> DuplexRequestClient: - client = getattr(self, "_duplex_request_client", None) - if client is None: - engine = getattr(self, "engine", None) - client = DuplexRequestClient( - engine, - DuplexRequestOutputPort( - request_states=getattr(self, "request_states", {}), - num_stages=getattr(engine, "num_stages", 1), - log_stats=getattr(self, "log_stats", False), - start_output_handler=self._final_output_handler, - process_single_result=self._process_single_result, - ), - ) - self._duplex_request_client = client - return client - - @staticmethod - def _duplex_data_plane_request_info(result: dict[str, object]) -> tuple[str | None, int | None]: - return DuplexRequestClient.request_info(result) - - async def _collect_duplex_data_plane_outputs( - self, - request_id: str, - req_state: ClientRequestState, - *, - response_stage_id: int | None, - timeout: float | None, - ) -> list[OmniRequestOutput]: - return await self._get_duplex_request_client().collect_outputs( - request_id, - req_state, - response_stage_id=response_stage_id, - timeout=timeout, - ) - - @classmethod - def _is_direct_duplex_data_plane_response(cls, output: object) -> bool: - return DuplexRequestClient.is_direct_response(output) - - @classmethod - def _duplex_multimodal_output(cls, output: object) -> dict[str, object]: - return DuplexRequestClient.multimodal_output(output) # ==================== Generate Method ==================== @@ -857,198 +494,6 @@ class AsyncOmni(EngineClient, OmniBase): """ raise NotImplementedError("AsyncOmni.encode is not implemented.") - # ==================== Processing Methods ==================== - - async def _process_orchestrator_results( - self, - request_id: str, - metrics: OrchestratorMetrics, - final_stage_id_for_e2e: int, - req_start_ts: dict[str, float], - wall_start_ts: float, - ) -> AsyncGenerator[OmniRequestOutput, None]: - """Read results from the Orchestrator (via the request's asyncio.Queue) - and yield OmniRequestOutput objects. - - The Orchestrator handles all stage-to-stage transfers. This method - only processes final outputs that arrive on the per-request queue. - """ - req_state = self.request_states.get(request_id) - if req_state is None: - return - - while True: - result = await req_state.queue.get() - - if isinstance(result, ErrorMessage): - logger.error( - "[AsyncOmni] Orchestrator error for req=%s stage-%s: %s", - request_id, - result.stage_id, - result.error, - ) - if result.fatal: - raise OmniEngineDeadError( - result.error, - error_stage_id=result.stage_id, - ) - self._raise_nonfatal_error_message(result) - - if not isinstance(result, OutputMessage): - logger.warning("[AsyncOmni] Dropping unexpected per-request message %r", result) - continue - - stage_id = result.stage_id - - self._check_engine_output_error(result, request_id, stage_id) - - # Process the result (constructs OmniRequestOutput) - output_to_yield = self._process_single_result( - result, - stage_id, - metrics, - req_start_ts, - wall_start_ts, - final_stage_id_for_e2e, - ) - - if output_to_yield: - # Set the external request ID back to the user yielded input - output_to_yield.request_id = req_state.external_request_id or output_to_yield.request_id - logger.debug( - "[AsyncOmni] req=%s stage-%s yielding final_output_type=%s", - request_id, - stage_id, - getattr(output_to_yield, "final_output_type", None), - ) - yield output_to_yield - - # The Orchestrator sets "finished" when the final stage is done - if result.finished: - break - - # ==================== Output Handler ==================== - - def _final_output_handler(self) -> None: - """Start the final output handler if not already running. - - This handler reads messages from the Orchestrator output queue and - routes them to per-request asyncio.Queues. - """ - if self.final_output_task is not None: - return - - engine = self.engine - - # Event-driven drain (VLLM_OMNI_EVENT_DRIVEN_ORCH=1): block on the - # queue's condition variable in a dedicated thread instead of the - # get_nowait + 1 ms sleep cadence. Same flag as the orchestrator-side - # event-driven loop (vllm_omni/engine/orchestrator.py). - from vllm_omni.engine.orchestrator import _event_driven_orch_enabled - - event_driven_drain = _event_driven_orch_enabled() and hasattr(engine, "get_output_blocking_async") - - async def _final_output_loop(): - """Background coroutine that dispatches final outputs to request queues.""" - try: - while True: - if event_driven_drain: - msg = await engine.get_output_blocking_async(timeout=_FINAL_OUTPUT_BLOCKING_WAIT_S) - if msg is None: - # Timed out with the orchestrator alive; loop for - # the periodic liveness check. - continue - else: - msg = await engine.try_get_output_async() - if msg is None: - await asyncio.sleep(_FINAL_OUTPUT_IDLE_SLEEP_S) - continue - - if isinstance(msg, dict) and msg.get("type") == "ack": - ack_data = msg.get("ack") - tid = getattr(ack_data, "task_id", "unknown") - logger.info(f"[{self._name}] Intercepted wrapped ACK for task {tid}") - await self.event_resolver.resolve(ack_data) - continue - if isinstance(msg, OmniACK): - logger.info(f"[{self._name}] Intercepted raw ACK object: {msg.task_id}") - await self.event_resolver.resolve(msg) - continue - if hasattr(msg, "task_id"): - tid = getattr(msg, "task_id") - logger.info(f"[{self._name}] Intercepted task-ID object: {tid}") - await self.event_resolver.resolve(msg) - continue - - if getattr(msg, "type", None) == "duplex_session_lifecycle": - await self.duplex_lifecycle_events.put(msg) - continue - - if isinstance(msg, ErrorMessage): - # Route request-scoped errors to that request's queue and - # keep the loop alive. A request whose stage replica died - # and was evicted gets a fatal error delivered here; only - # that request fails (its consumer raises), the server - # stays up for other stages/requests (#4285). A fatal - # error without a request_id is a genuine engine-wide - # death and falls through to the except handler below. - if msg.request_id is not None: - req_state = self.request_states.get(msg.request_id) - if req_state is not None: - await req_state.queue.put(msg) - else: - logger.warning( - "[%s] dropping error for unknown req %s", - self._name, - msg.request_id, - ) - continue - if not msg.fatal: - continue - - should_continue, _, stage_id, req_state = self._handle_output_message(msg) - if should_continue: - continue - - req_state.stage_id = stage_id - - # Route to the per-request queue - await req_state.queue.put(msg) - - except asyncio.CancelledError: - raise - except OmniEngineDeadError as e: - logger.error("[AsyncOmni] Engine dead: %s", e) - for req_state in list(self.request_states.values()): - error_msg = ErrorMessage( - error=str(e), - fatal=True, - request_id=req_state.request_id, - stage_id=e.error_stage_id, - ) - await req_state.queue.put(error_msg) - except EngineDeadError as e: - logger.error("[AsyncOmni] Engine dead: %s", e) - for req_state in list(self.request_states.values()): - error_msg = ErrorMessage( - error=str(e), - fatal=True, - request_id=req_state.request_id, - ) - await req_state.queue.put(error_msg) - except Exception as e: - logger.exception("[AsyncOmni] final_output_loop failed.") - for req_state in list(self.request_states.values()): - error_msg = ErrorMessage( - request_id=req_state.request_id, - error=str(e), - ) - await req_state.queue.put(error_msg) - self.final_output_task = None - - self.final_output_task = asyncio.create_task(_final_output_loop()) - logger.debug("[AsyncOmni] Final output handler started") - # ==================== Control Methods ==================== async def collective_rpc( @@ -1172,102 +617,6 @@ class AsyncOmni(EngineClient, OmniBase): if self.log_stats: logger.info("[AsyncOmni] Queued interaction for request %s", request_id) - async def _abort_internal_requests( - self, - request_id: str | Iterable[str], - *, - timeout: float = ABORT_TIMEOUT_S, - ): - """Abort request(s) via the Orchestrator given internal request IDs, - which take the format -. - """ - request_ids = [request_id] if isinstance(request_id, str) else list(request_id) - # Request IDs are already internal, so we just need to get the matching states. - internal_req_ids = [rid for rid in request_ids if rid in self.request_states] - try: - # Unbind generate() if abort_async blocks in the executor. - await asyncio.wait_for(self._abort(internal_req_ids, timeout=timeout), timeout=timeout) - except TimeoutError: - logger.warning( - "[AsyncOmni] Timed out aborting %s after %.1fs; " - "engine abort is best-effort until the current batch drains", - ",".join(internal_req_ids), - timeout, - ) - except Exception: - logger.exception("[AsyncOmni] Cleanup abort failed for %s", ",".join(internal_req_ids)) - - async def _abort(self, request_ids: list[str], *, timeout: float | None = None) -> None: - """Abort request IDs via the engine and enqueue terminal abort outputs. - - Waits for orchestrator abort acknowledgment, enqueues any AR terminal - abort outputs (partial tokens) into each request's asyncio queue, then - cancels the input pump. Frontend ``request_states`` stay registered so - ``generate()`` can consume the terminal message in - ``_process_orchestrator_results`` and run normal cleanup. - - When ``abort_async`` returns no output for an active request (OP not - registered yet, unbound replica, or orchestrator id drop), enqueue a - synthetic finished abort so ``generate()`` cannot hang on ``queue.get``. - """ - abort_outputs = await self.engine.abort_async(request_ids, timeout=timeout) or [] - delivered: set[str] = set() - for output_msg in abort_outputs: - req_id = getattr(output_msg, "request_id", None) - if req_id is None: - continue - state = self.request_states.get(req_id) - if state is None: - logger.debug("[AsyncOmni] Dropping abort output for unknown req %s", req_id) - continue - await state.queue.put(output_msg) - delivered.add(req_id) - for rid in request_ids: - state = self.request_states.get(rid) - if state is not None and rid not in delivered: - queue = getattr(state, "queue", None) - if queue is not None: - await state.queue.put(self._synthetic_abort_output_message(rid)) - delivered.add(rid) - for rid in request_ids: - self._record_request_failure_once(rid, reason="client_abort") - state = self.request_states.get(rid) - input_stream_task = getattr(state, "input_stream_task", None) - if input_stream_task is not None and not input_stream_task.done(): - input_stream_task.cancel() - if self.log_stats: - logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids)) - - @staticmethod - def _synthetic_abort_output_message(request_id: str) -> OutputMessage: - """Terminal abort OutputMessage used when the engine returned none.""" - engine_output = OmniRequestOutput( - request_id=request_id, - finished=True, - stage_id=0, - final_output_type="text", - outputs=[ - CompletionOutput( - index=0, - text="", - token_ids=[], - cumulative_logprob=None, - logprobs=None, - finish_reason="abort", - stop_reason=None, - ) - ], - ) - return OutputMessage( - request_id=request_id, - stage_id=0, - replica_id=None, - engine_outputs=engine_output, - metrics=None, - finished=True, - stage_submit_ts=None, - ) - def _split_stage_ids_by_type(self, stage_ids: list[int] | None = None) -> tuple[list[int], list[int]]: """Split stage ids into AR/LLM (EngineCore) vs diffusion (worker RPC).""" n_stages = len(self.engine.stage_clients) @@ -1697,42 +1046,6 @@ class AsyncOmni(EngineClient, OmniBase): # ==================== Properties ==================== - @property - def is_running(self) -> bool: - """Check if the engine is running.""" - orchestrator_alive = self.engine.is_alive() - task_alive = self.final_output_task is not None and not self.final_output_task.done() - return orchestrator_alive and task_alive - - @property - def errored(self) -> bool: - """Whether the engine is in a process-fatal error state. - - Delegates to ``OmniBase.errored``, which is true only when the - orchestrator thread is dead; per-stage liveness is reported via - ``check_health`` instead. Redeclared here to satisfy the - ``EngineClient`` abstract-property requirement (Python's ABC - mechanism does not resolve abstract methods from sibling MRO - entries). - """ - return OmniBase.errored.fget(self) # type: ignore[union-attr] - - @property - def _name(self) -> str: - return "AsyncOrchestrator" - - @property - def is_stopped(self) -> bool: - """EngineClient abstract property implementation.""" - return self.errored - - @property - def dead_error(self) -> BaseException: - """EngineClient abstract property implementation.""" - return OmniEngineDeadError() - - # ==================== EngineClient Interface ==================== - async def get_input_preprocessor(self) -> InputProcessor: """Get input preprocessor.""" return self.input_processor @@ -1793,16 +1106,3 @@ class AsyncOmni(EngineClient, OmniBase): async def get_supported_tasks(self) -> tuple[SupportedTask, ...]: """Return the task set exposed by the orchestrator-backed engine.""" return tuple(self.engine.supported_tasks) - - async def check_health(self) -> None: - """Check engine health by verifying the Orchestrator process is alive.""" - OmniBase.check_health(self) - - # ==================== Shutdown ==================== - - def shutdown(self, timeout: float | None = None) -> None: - """Shutdown the engine.""" - if self.final_output_task is not None: - self.final_output_task.cancel() - self.final_output_task = None - OmniBase.shutdown(self) diff --git a/vllm_omni/entrypoints/async_omni_base.py b/vllm_omni/entrypoints/async_omni_base.py new file mode 100644 index 000000000..e7c86de97 --- /dev/null +++ b/vllm_omni/entrypoints/async_omni_base.py @@ -0,0 +1,557 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""AsyncOmniBase: the asyncio-side foundation shared by ``AsyncOmni`` and ``DuplexOmni``. + +Owns the engine output pump (one consumer of the engine output queue that +routes messages to per-request queues, ACK resolvers, or a subclass hook), +engine-dead fan-out, health/shutdown, and the config accessors. Turn-based +request handling lives in ``AsyncOmni``; duplex sessions in ``DuplexOmni``. +""" + +from __future__ import annotations + +import asyncio +import os +import time +from collections.abc import AsyncGenerator, Iterable +from typing import Any + +from vllm.logger import init_logger +from vllm.outputs import CompletionOutput +from vllm.plugins.io_processors import get_io_processor +from vllm.utils import random_uuid +from vllm.v1.engine.exceptions import EngineDeadError + +from vllm_omni.diffusion.data import OmniACK +from vllm_omni.engine.messages import ErrorMessage, OutputMessage +from vllm_omni.entrypoints.omni_base import ( + OmniBase, + OmniEngineDeadError, +) +from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics +from vllm_omni.outputs import OmniRequestOutput + +logger = init_logger(__name__) +_FINAL_OUTPUT_IDLE_SLEEP_S = 0.001 +# Blocking-wait interval for the event-driven final-output drain +# (VLLM_OMNI_EVENT_DRIVEN_ORCH=1): a message wakes the drain immediately via +# the janus queue's condition variable; this timeout only bounds how often the +# orchestrator liveness check runs while the pipeline is idle. +_FINAL_OUTPUT_BLOCKING_WAIT_S = 1.0 +# Shared DELETE / generate() cleanup abort bound. Env is the documented knob. +ABORT_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_ABORT_TIMEOUT", 2.0)) + + +class AsyncEventResolver: + """ + A generic signal aggregator designed for synchronized handshakes in + distributed or multi-stage environments. Supports waiting for a specified + number (expected_count) of worker signals in both inline and multiprocess modes. + """ + + def __init__(self, orchestrator=None): + self._pending_tasks: dict[str, dict] = {} + self.orchestrator = orchestrator + self._lock = asyncio.Lock() + + def watch_task(self, task_id: str, expected_count: int = 1) -> asyncio.Future: + loop = asyncio.get_running_loop() + fut = loop.create_future() + self._pending_tasks[task_id] = { + "future": fut, + "expected_count": expected_count, + "received": [], + "start_time": time.time(), + } + return fut + + async def resolve(self, ack: OmniACK): + tid = getattr(ack, "task_id", None) + + if tid is None and isinstance(ack, dict): + tid = ack.get("task_id") + + async with self._lock: + task_info = self._pending_tasks.get(tid) + if task_info is None: + logger.warning(f"Received stray ACK for task_id {tid}. Task might have timed out.") + return + + task_info["received"].append(ack) + current_count = len(task_info["received"]) + expected = task_info["expected_count"] + + orchestrator = self.orchestrator + if orchestrator and hasattr(orchestrator, "metrics") and orchestrator.metrics: + freed = getattr(ack, "freed_bytes", 0) + if freed == 0 and isinstance(ack, dict): + freed = ack.get("freed_bytes", 0) + orchestrator.metrics.record_vram_reclaimed(freed) + + logger.info(f"[Resolver] Task {tid} progress: {current_count}/{expected} ACKs received.") + + if current_count >= expected: + self._pending_tasks.pop(tid) + fut = task_info["future"] + if not fut.done(): + elapsed = time.time() - task_info["start_time"] + logger.info(f"[Resolver] Task {tid} completed successfully in {elapsed:.2f}s.") + fut.set_result(task_info["received"]) + + +class AsyncOmniBase(OmniBase): + """Shared asyncio foundation of the async entrypoints (see module docstring).""" + + def __init__(self, *args: Any, model: str = "", **kwargs: Any) -> None: + OmniBase.__init__(self, model=model, **kwargs) + self.final_output_task: asyncio.Task | None = None + self.event_resolver = AsyncEventResolver(orchestrator=self) + self.config_path = self.engine.config_path + self.input_processor = self.engine.input_processor + self.endpoint_restrictions = self.engine.endpoint_restrictions + + stage_index = self._get_comprehension_stage_index() + if stage_index is None: + self.io_processor = None + else: + vllm_config = self.engine.stage_vllm_configs[stage_index] + io_processor_plugin = vllm_config.model_config.io_processor_plugin + renderer = self.renderer + if renderer is None: + from vllm.renderers import renderer_from_config + + renderer = renderer_from_config(vllm_config) + self.io_processor = get_io_processor(vllm_config, renderer, io_processor_plugin) + + # ==================== Subclass seam ==================== + + def _route_engine_message(self, msg: Any) -> bool: + """Consume an engine output message the subclass owns (return True when consumed).""" + del msg + return False + + def _on_engine_dead(self, error: str) -> None: + """Seam: the output pump is ending because the engine died (default: nothing). + + ``DuplexOmni`` ends every session handle here; turn-based requests are + already failed through their per-request queues. + """ + del error + + def _resolve_transfer_replica(self, stage_id: int, request_id: str) -> int | None: + """Look up the sticky-routed replica for (stage_id, request_id). + + Used as the ``replica_resolver`` callback by ``OrchestratorAggregator`` + to label transfer_* metrics without plumbing replica ids through + ``TransferEdgeStats`` / ``StageRequestStats`` / connector adapters. + Returns None when stage_id is out of range or the request hasn't been + bound to a replica yet — the metric emit then defensive-skips. + """ + pools = getattr(self.engine, "stage_pools", None) + if pools is None or not (0 <= stage_id < len(pools)): + return None + return pools[stage_id].get_bound_replica_id(request_id) + + def _get_comprehension_stage_index(self) -> int | None: + fallback_idx: int | None = None + for idx, stage_client in enumerate(self.engine.stage_clients): + stage_vllm_config = self.engine.stage_vllm_configs[idx] + if stage_vllm_config is None: + continue + if fallback_idx is None: + fallback_idx = idx + if stage_client.is_comprehension: + return idx + return fallback_idx + + @property + def renderer(self): + """Return the renderer from the engine input processor when available.""" + if self.input_processor is None: + return None + return self.input_processor.renderer + + @property + def vllm_config(self): + """Return the vLLM config for the comprehension stage when present.""" + stage_index = self._get_comprehension_stage_index() + if stage_index is None: + return None + return self.engine.stage_vllm_configs[stage_index] + + async def get_vllm_config(self) -> Any: + """Compatibility helper for call sites expecting async vllm config access.""" + return self.vllm_config + + def get_diffusion_od_config(self) -> Any | None: + """Return the diffusion-stage config when the pipeline has one.""" + saw_diffusion_stage = False + for stage_client in self.engine.stage_clients: + if getattr(stage_client, "stage_type", None) != "diffusion": + continue + + saw_diffusion_stage = True + + od_config = getattr(stage_client, "od_config", None) + if od_config is not None: + return od_config + + inner_engine = getattr(stage_client, "_engine", None) + od_config = getattr(inner_engine, "od_config", None) + if od_config is not None: + return od_config + + # Out-of-process diffusion clients don't carry od_config (it lives in the + # worker); fall back to the engine's model_class_name resolution. + if saw_diffusion_stage: + return self.engine.get_diffusion_od_config() + + return None + + @property + def model_config(self): + """Return the model config for the comprehension stage when present.""" + vllm_config = self.vllm_config + if vllm_config is None: + return None + return vllm_config.model_config + + @staticmethod + def _get_unique_request_id(external_request_id: str): + """Get a random new request ID for this request; at the server level, + this is usually set by the calling entrypoint, but in direct calls, we + need to set it explicitly since we do not allow empty IDs. + + NOTE: in the upstream vLLM, this is done in the InputProcessor's + `assign_request_id`. + """ + uuid = random_uuid() + prefix = "" if not external_request_id else f"{external_request_id}-" + return f"{prefix}{uuid:.8}" + + # ==================== Processing Methods ==================== + + async def _process_orchestrator_results( + self, + request_id: str, + metrics: OrchestratorMetrics, + final_stage_id_for_e2e: int, + req_start_ts: dict[str, float], + wall_start_ts: float, + ) -> AsyncGenerator[OmniRequestOutput, None]: + """Read results from the Orchestrator (via the request's asyncio.Queue) + and yield OmniRequestOutput objects. + + The Orchestrator handles all stage-to-stage transfers. This method + only processes final outputs that arrive on the per-request queue. + """ + req_state = self.request_states.get(request_id) + if req_state is None: + return + + while True: + result = await req_state.queue.get() + + if isinstance(result, ErrorMessage): + logger.error( + "[AsyncOmni] Orchestrator error for req=%s stage-%s: %s", + request_id, + result.stage_id, + result.error, + ) + if result.fatal: + raise OmniEngineDeadError( + result.error, + error_stage_id=result.stage_id, + ) + self._raise_nonfatal_error_message(result) + + if not isinstance(result, OutputMessage): + logger.warning("[AsyncOmni] Dropping unexpected per-request message %r", result) + continue + + stage_id = result.stage_id + + self._check_engine_output_error(result, request_id, stage_id) + + # Process the result (constructs OmniRequestOutput) + output_to_yield = self._process_single_result( + result, + stage_id, + metrics, + req_start_ts, + wall_start_ts, + final_stage_id_for_e2e, + ) + + if output_to_yield: + # Set the external request ID back to the user yielded input + output_to_yield.request_id = req_state.external_request_id or output_to_yield.request_id + logger.debug( + "[AsyncOmni] req=%s stage-%s yielding final_output_type=%s", + request_id, + stage_id, + getattr(output_to_yield, "final_output_type", None), + ) + yield output_to_yield + + # The Orchestrator sets "finished" when the final stage is done + if result.finished: + break + + # ==================== Output Handler ==================== + + def _final_output_handler(self) -> None: + """Start the final output handler if not already running. + + This handler reads messages from the Orchestrator output queue and + routes them to per-request asyncio.Queues. + """ + if self.final_output_task is not None: + return + + engine = self.engine + + # Event-driven drain (VLLM_OMNI_EVENT_DRIVEN_ORCH=1): block on the + # queue's condition variable in a dedicated thread instead of the + # get_nowait + 1 ms sleep cadence. Same flag as the orchestrator-side + # event-driven loop (vllm_omni/engine/orchestrator.py). + from vllm_omni.engine.orchestrator import _event_driven_orch_enabled + + event_driven_drain = _event_driven_orch_enabled() and hasattr(engine, "get_output_blocking_async") + + async def _final_output_loop(): + """Background coroutine that dispatches final outputs to request queues.""" + try: + while True: + if event_driven_drain: + msg = await engine.get_output_blocking_async(timeout=_FINAL_OUTPUT_BLOCKING_WAIT_S) + if msg is None: + # Timed out with the orchestrator alive; loop for + # the periodic liveness check. + continue + else: + msg = await engine.try_get_output_async() + if msg is None: + await asyncio.sleep(_FINAL_OUTPUT_IDLE_SLEEP_S) + continue + + if self._route_engine_message(msg): + continue + + if isinstance(msg, dict) and msg.get("type") == "ack": + ack_data = msg.get("ack") + tid = getattr(ack_data, "task_id", "unknown") + logger.info(f"[{self._name}] Intercepted wrapped ACK for task {tid}") + await self.event_resolver.resolve(ack_data) + continue + if isinstance(msg, OmniACK): + logger.info(f"[{self._name}] Intercepted raw ACK object: {msg.task_id}") + await self.event_resolver.resolve(msg) + continue + if hasattr(msg, "task_id"): + tid = getattr(msg, "task_id") + logger.info(f"[{self._name}] Intercepted task-ID object: {tid}") + await self.event_resolver.resolve(msg) + continue + + if isinstance(msg, ErrorMessage): + # Route request-scoped errors to that request's queue and + # keep the loop alive. A request whose stage replica died + # and was evicted gets a fatal error delivered here; only + # that request fails (its consumer raises), the server + # stays up for other stages/requests (#4285). A fatal + # error without a request_id is a genuine engine-wide + # death and falls through to the except handler below. + if msg.request_id is not None: + req_state = self.request_states.get(msg.request_id) + if req_state is not None: + await req_state.queue.put(msg) + else: + logger.warning( + "[%s] dropping error for unknown req %s", + self._name, + msg.request_id, + ) + continue + if not msg.fatal: + continue + + should_continue, _, stage_id, req_state = self._handle_output_message(msg) + if should_continue: + continue + + req_state.stage_id = stage_id + + # Route to the per-request queue + await req_state.queue.put(msg) + + except asyncio.CancelledError: + raise + except OmniEngineDeadError as e: + logger.error("[%s] Engine dead: %s", self._name, e) + for req_state in list(self.request_states.values()): + error_msg = ErrorMessage( + error=str(e), + fatal=True, + request_id=req_state.request_id, + stage_id=e.error_stage_id, + ) + await req_state.queue.put(error_msg) + self._on_engine_dead(str(e)) + except EngineDeadError as e: + logger.error("[%s] Engine dead: %s", self._name, e) + for req_state in list(self.request_states.values()): + error_msg = ErrorMessage( + error=str(e), + fatal=True, + request_id=req_state.request_id, + ) + await req_state.queue.put(error_msg) + self._on_engine_dead(str(e)) + except Exception as e: + logger.exception("[%s] final_output_loop failed.", self._name) + for req_state in list(self.request_states.values()): + error_msg = ErrorMessage( + request_id=req_state.request_id, + error=str(e), + ) + await req_state.queue.put(error_msg) + self.final_output_task = None + self._on_engine_dead(str(e)) + + self.final_output_task = asyncio.create_task(_final_output_loop()) + logger.debug("[AsyncOmni] Final output handler started") + + async def _abort_internal_requests( + self, + request_id: str | Iterable[str], + *, + timeout: float = ABORT_TIMEOUT_S, + ): + """Abort request(s) via the Orchestrator given internal request IDs, + which take the format -. + """ + request_ids = [request_id] if isinstance(request_id, str) else list(request_id) + # Request IDs are already internal, so we just need to get the matching states. + internal_req_ids = [rid for rid in request_ids if rid in self.request_states] + try: + # Unbind generate() if abort_async blocks in the executor. + await asyncio.wait_for(self._abort(internal_req_ids, timeout=timeout), timeout=timeout) + except TimeoutError: + logger.warning( + "[AsyncOmni] Timed out aborting %s after %.1fs; " + "engine abort is best-effort until the current batch drains", + ",".join(internal_req_ids), + timeout, + ) + except Exception: + logger.exception("[AsyncOmni] Cleanup abort failed for %s", ",".join(internal_req_ids)) + + async def _abort(self, request_ids: list[str], *, timeout: float | None = None) -> None: + """Abort request IDs via the engine and enqueue terminal abort outputs. + + Waits for orchestrator abort acknowledgment, enqueues any AR terminal + abort outputs (partial tokens) into each request's asyncio queue, then + cancels the input pump. Frontend ``request_states`` stay registered so + ``generate()`` can consume the terminal message in + ``_process_orchestrator_results`` and run normal cleanup. + + When ``abort_async`` returns no output for an active request (OP not + registered yet, unbound replica, or orchestrator id drop), enqueue a + synthetic finished abort so ``generate()`` cannot hang on ``queue.get``. + """ + abort_outputs = await self.engine.abort_async(request_ids, timeout=timeout) or [] + delivered: set[str] = set() + for output_msg in abort_outputs: + req_id = getattr(output_msg, "request_id", None) + if req_id is None: + continue + state = self.request_states.get(req_id) + if state is None: + logger.debug("[AsyncOmni] Dropping abort output for unknown req %s", req_id) + continue + await state.queue.put(output_msg) + delivered.add(req_id) + for rid in request_ids: + state = self.request_states.get(rid) + if state is not None and rid not in delivered: + queue = getattr(state, "queue", None) + if queue is not None: + await state.queue.put(self._synthetic_abort_output_message(rid)) + delivered.add(rid) + for rid in request_ids: + self._record_request_failure_once(rid, reason="client_abort") + state = self.request_states.get(rid) + input_stream_task = getattr(state, "input_stream_task", None) + if input_stream_task is not None and not input_stream_task.done(): + input_stream_task.cancel() + if self.log_stats: + logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids)) + + @staticmethod + def _synthetic_abort_output_message(request_id: str) -> OutputMessage: + """Terminal abort OutputMessage used when the engine returned none.""" + engine_output = OmniRequestOutput( + request_id=request_id, + finished=True, + stage_id=0, + final_output_type="text", + outputs=[ + CompletionOutput( + index=0, + text="", + token_ids=[], + cumulative_logprob=None, + logprobs=None, + finish_reason="abort", + stop_reason=None, + ) + ], + ) + return OutputMessage( + request_id=request_id, + stage_id=0, + replica_id=None, + engine_outputs=engine_output, + metrics=None, + finished=True, + stage_submit_ts=None, + ) + + @property + def is_running(self) -> bool: + """Check if the engine is running.""" + orchestrator_alive = self.engine.is_alive() + task_alive = self.final_output_task is not None and not self.final_output_task.done() + return orchestrator_alive and task_alive + + @property + def _name(self) -> str: + return type(self).__name__ + + @property + def is_stopped(self) -> bool: + """EngineClient abstract property implementation.""" + return self.errored + + @property + def dead_error(self) -> BaseException: + """EngineClient abstract property implementation.""" + return OmniEngineDeadError() + + # ==================== EngineClient Interface ==================== + + async def check_health(self) -> None: + """Check engine health by verifying the Orchestrator process is alive.""" + OmniBase.check_health(self) + + # ==================== Shutdown ==================== + + def shutdown(self, timeout: float | None = None) -> None: + """Shutdown the engine.""" + if self.final_output_task is not None: + self.final_output_task.cancel() + self.final_output_task = None + OmniBase.shutdown(self) diff --git a/vllm_omni/entrypoints/cli/benchmark/cli_args.py b/vllm_omni/entrypoints/cli/benchmark/cli_args.py index 21f528954..626951298 100644 --- a/vllm_omni/entrypoints/cli/benchmark/cli_args.py +++ b/vllm_omni/entrypoints/cli/benchmark/cli_args.py @@ -206,8 +206,10 @@ def add_seed_tts_cli_args(parser: argparse.ArgumentParser) -> None: "--seed-tts-turns-per-session", type=int, default=1, - help="Group this many Seed-TTS target texts into one Realtime session. " - "The first row's reference audio and transcript are reused for every turn.", + help="Group this many Seed-TTS target texts into one benchmark request. " + "The first row's reference audio and transcript are reused for every turn. " + "A model-native duplex session takes its target text once, in the session " + "context, so each turn runs in its own session.", ) group.add_argument( "--seed-tts-root", diff --git a/vllm_omni/entrypoints/duplex/__init__.py b/vllm_omni/entrypoints/duplex/__init__.py index 736e98a6a..327eab07f 100644 --- a/vllm_omni/entrypoints/duplex/__init__.py +++ b/vllm_omni/entrypoints/duplex/__init__.py @@ -1,6 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -from .websocket import DuplexWebSocketActor +"""Thin websocket serving for engine-resident duplex sessions. -__all__ = ["DuplexWebSocketActor"] +The session logic lives in ``vllm_omni.engine.duplex`` and is reached through +``vllm_omni.entrypoints.duplex_omni.DuplexOmni``; this package only speaks the +OpenAI Realtime wire protocol over ``/v1/realtime?duplex=1``. +""" + +from .serving import OmniDuplexSessionHandler + +__all__ = ["OmniDuplexSessionHandler"] diff --git a/vllm_omni/entrypoints/duplex/audio_encoding.py b/vllm_omni/entrypoints/duplex/audio_encoding.py new file mode 100644 index 000000000..f790b029b --- /dev/null +++ b/vllm_omni/entrypoints/duplex/audio_encoding.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Audio encoding for duplex model output (entrypoint layer, injected into ``DuplexOmniEngine``).""" + +from __future__ import annotations + +import numpy as np +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def encode_audio( + audio_data: object, + sample_rate_hz: int, + response_format: str, + speed: float | None, +) -> str | None: + """Encode a model audio tensor/array into base64 in ``response_format``. + + Moved from the serving runtime bridge; uses the shared ``AudioMixin`` + encoder directly instead of going through the chat service. + """ + if audio_data is None: + return None + try: + import torch + + from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin + from vllm_omni.entrypoints.openai.protocol.audio import CreateAudio + + if isinstance(audio_data, torch.Tensor): + audio_tensor = audio_data.detach().cpu().float().numpy() + else: + audio_tensor = np.asarray(audio_data, dtype=np.float32) + if audio_tensor.ndim > 1: + audio_tensor = audio_tensor.reshape(-1) + audio_response = AudioMixin().create_audio( + CreateAudio( + audio_tensor=audio_tensor, + sample_rate=sample_rate_hz, + response_format=response_format, + speed=float(speed) if isinstance(speed, int | float) and speed > 0 else 1.0, + stream_format="audio", + base64_encode=True, + ) + ) + return str(audio_response.audio_data) + except Exception: + logger.exception("Failed to encode duplex data-plane audio output") + return None + + +__all__ = ["encode_audio"] diff --git a/vllm_omni/entrypoints/duplex/capability.py b/vllm_omni/entrypoints/duplex/capability.py deleted file mode 100644 index 50126886b..000000000 --- a/vllm_omni/entrypoints/duplex/capability.py +++ /dev/null @@ -1,34 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from vllm.logger import init_logger - -logger = init_logger(__name__) - - -def should_enable_duplex_endpoint( - stage_configs: list | None, - *, - config_path: str | None = None, -) -> bool: - """Enable the realtime session handler for explicitly configured deployments.""" - if stage_configs: - for stage in stage_configs: - session_mode = ( - stage.get("session_mode") if isinstance(stage, dict) else getattr(stage, "session_mode", None) - ) - if session_mode == "duplex": - return True - if config_path: - try: - from vllm_omni.config.stage_config import resolve_deploy_yaml - - raw_config = resolve_deploy_yaml(config_path) - if raw_config.get("session_mode") == "duplex" or isinstance(raw_config.get("duplex_session"), dict): - return True - except Exception as exc: - logger.warning("Failed to inspect realtime session configuration from %s: %s", config_path, exc) - return False - - -__all__ = ["should_enable_duplex_endpoint"] diff --git a/vllm_omni/entrypoints/duplex/chat_fallback.py b/vllm_omni/entrypoints/duplex/chat_fallback.py deleted file mode 100644 index bd682c4f3..000000000 --- a/vllm_omni/entrypoints/duplex/chat_fallback.py +++ /dev/null @@ -1,302 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -import json -from collections.abc import AsyncGenerator -from functools import partial -from typing import Any - -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest -from vllm.entrypoints.serve.engine.protocol import ErrorResponse -from vllm.logger import init_logger - -from vllm_omni.entrypoints.duplex.protocol import DuplexSession - -logger = init_logger(__name__) - - -class ChatFallbackProjectorMixin: - """Project generic chat completion output into duplex response events.""" - - async def _run_response(self, session: DuplexSession, send_json) -> None: - response_id = session.begin_response() - epoch = session.epoch - request_id = f"duplex-{session.session_id}-{epoch}-{session.input_commit_seq}" - session.bind_request(f"chatcmpl-{request_id}") - await send_json( - self._response_created_payload( - session, - response_id, - epoch=epoch, - request_id=session.active_request_id, - ) - ) - - try: - request = self._build_chat_request(session, request_id) - result = await self._chat_service.create_chat_completion(request, raw_request=None) - if isinstance(result, ErrorResponse): - error = result.error - await send_json( - { - "type": "error", - "error": error.message if error else "Chat request failed", - "code": error.type if error else "chat_error", - } - ) - session.end_response(commit_text=False) - return - if hasattr(result, "__aiter__"): - await self._drain_streaming_response(session, result, epoch, response_id, send_json) - else: - await self._emit_full_response(session, result, epoch, response_id, send_json) - if session.epoch == epoch: - final_stage_metrics = session.accumulate_response_stage_metrics(None) - should_commit = self._should_commit_response_to_history(session, response_id) - committed_message = session.end_response(commit_text=should_commit) - if should_commit: - session.register_history_item(f"item_{response_id}", committed_message) - done_payload: dict[str, object] = { - "type": "response.done", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "committed": committed_message is not None, - "playback": session.playback.as_dict(), - } - if final_stage_metrics: - done_payload["vllm_omni"] = {"stage_metrics": final_stage_metrics} - await send_json(done_payload) - except asyncio.CancelledError: - raise - except Exception as exc: - logger.exception("Duplex response failed: %s", exc) - session.end_response(commit_text=False) - await send_json( - { - "type": "error", - "session_id": session.session_id, - "response_id": response_id, - "error": str(exc), - "code": "response_error", - } - ) - - def _build_chat_request(self, session: DuplexSession, request_id: str) -> ChatCompletionRequest: - response_config = session.response_config - messages: list[dict[str, object]] = [] - if response_config.instructions: - messages.append({"role": "system", "content": response_config.instructions}) - messages.extend(session.history) - - kwargs: dict[str, Any] = { - "model": response_config.model or self._chat_service.model_config.model, - "messages": messages, - "stream": True, - } - if response_config.temperature is not None: - kwargs["temperature"] = response_config.temperature - if response_config.max_tokens is not None: - kwargs["max_tokens"] = response_config.max_tokens - model_extra = dict(response_config.extra_body) - model_extra.pop("native_duplex", None) - model_extra.pop("minicpmo45_native_duplex", None) - tools = model_extra.pop("realtime_response_tools", model_extra.pop("realtime_tools", None)) - tool_choice = model_extra.pop( - "realtime_response_tool_choice", - model_extra.pop("realtime_tool_choice", None), - ) - for protocol_key in ( - "realtime_response_conversation", - "realtime_response_metadata", - "realtime_response_prompt", - ): - model_extra.pop(protocol_key, None) - kwargs.update(model_extra) - if "audio" in response_config.modalities: - audio_options = kwargs.get("audio") - audio_options = dict(audio_options) if isinstance(audio_options, dict) else {} - # Chat defaults to WAV. Request the session's internal output format - # explicitly rather than relabeling the returned bytes as PCM. - audio_options["format"] = response_config.response_format - kwargs["audio"] = audio_options - if isinstance(tools, list): - kwargs["tools"] = tools - if isinstance(tool_choice, str | dict): - kwargs["tool_choice"] = tool_choice - - request = ChatCompletionRequest(**kwargs) - object.__setattr__(request, "modalities", response_config.modalities) - object.__setattr__(request, "request_id", request_id) - object.__setattr__( - request, - "chat_template_kwargs", - {"use_tts_template": response_config.use_tts_template}, - ) - return request - - async def _drain_streaming_response( - self, - session: DuplexSession, - result: AsyncGenerator[str, None], - epoch: int, - response_id: str, - send_json, - ) -> None: - async for raw_chunk in result: - if session.epoch != epoch: - return - for payload in self._parse_sse_payloads(raw_chunk): - if payload == "[DONE]": - continue - if isinstance(payload, dict): - await self._emit_chat_payload(session, payload, epoch, response_id, send_json) - - async def _emit_full_response( - self, - session: DuplexSession, - result: Any, - epoch: int, - response_id: str, - send_json, - ) -> None: - if hasattr(result, "model_dump"): - payload = result.model_dump(mode="json", exclude_unset=True) - else: - payload = {"response": str(result)} - await self._emit_chat_payload(session, payload, epoch, response_id, send_json) - - def _parse_sse_payloads(self, raw_chunk: str) -> list[dict[str, object] | str]: - payloads: list[dict[str, object] | str] = [] - for line in raw_chunk.splitlines(): - line = line.strip() - if not line.startswith("data:"): - continue - data = line[5:].strip() - if not data: - continue - if data == "[DONE]": - payloads.append(data) - continue - try: - parsed = json.loads(data) - except json.JSONDecodeError: - logger.debug("Skipping non-JSON duplex stream payload: %s", data) - continue - if isinstance(parsed, dict): - payloads.append(parsed) - return payloads - - async def _emit_chat_payload( - self, - session: DuplexSession, - payload: dict[str, object], - epoch: int, - response_id: str, - send_json, - ) -> None: - if session.epoch != epoch or session.active_response_id != response_id: - return - metrics = payload.get("metrics") - stage_metrics = metrics.get("stage_metrics") if isinstance(metrics, dict) else None - if isinstance(stage_metrics, dict): - session.replace_response_stage_metric_snapshots(stage_metrics) - - modality = payload.get("modality") - if modality not in {None, "text", "audio"}: - await send_json( - { - "type": "error", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "code": "unsupported_response_modality", - "error": f"Unsupported chat response modality: {modality}", - } - ) - return - choices = payload.get("choices") - if not isinstance(choices, list): - await send_json( - { - "type": "response.message", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "payload": payload, - } - ) - return - - for choice in choices: - if not isinstance(choice, dict): - continue - delta = choice.get("delta") - message = choice.get("message") - content = None - if isinstance(delta, dict): - content = delta.get("content") - elif isinstance(message, dict): - content = message.get("content") - - if isinstance(content, str) and content: - if modality != "audio": - session.append_assistant_text(content) - await send_json( - { - "type": "response.text.delta", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "delta": content, - } - ) - - audio_content = content if modality == "audio" else None - if isinstance(message, dict): - message_audio = message.get("audio") - if isinstance(message_audio, dict): - audio_content = message_audio.get("data") - if isinstance(audio_content, str) and audio_content: - # The OpenAI package imports duplex serving during initialization. - # Defer this import until chat fallback has finished loading. - from vllm_omni.entrypoints.openai.protocol.audio import AudioChunkMetadata - - audio_metadata = AudioChunkMetadata.model_validate(choice.get("audio_metadata")) - duration_ms = session.record_generated_audio( - response_id, - frame_count=audio_metadata.frame_count, - sample_rate_hz=audio_metadata.sample_rate_hz, - ) - if duration_ms is None: - return - await send_json( - { - "type": "response.output_audio.delta", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "audio": audio_content, - "format": audio_metadata.format, - "sample_rate_hz": audio_metadata.sample_rate_hz, - "channels": audio_metadata.channels, - "audio_duration_ms": duration_ms, - }, - on_accepted=partial(session.mark_audio_sent, duration_ms, response_id=response_id), - ) - - finish_reason = choice.get("finish_reason") - if finish_reason is not None and modality != "audio": - await send_json( - { - "type": "response.output_item.done", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "finish_reason": finish_reason, - "modality": modality, - } - ) diff --git a/vllm_omni/entrypoints/duplex/realtime_input.py b/vllm_omni/entrypoints/duplex/realtime_input.py index 386b4cfc3..be700ef62 100644 --- a/vllm_omni/entrypoints/duplex/realtime_input.py +++ b/vllm_omni/entrypoints/duplex/realtime_input.py @@ -1,1584 +1,179 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Thin OpenAI Realtime wire envelope for one websocket connection. + +Everything that needed session state now lives engine-side +(``vllm_omni.engine.duplex.realtime_commands`` for the command mapping, +``vllm_omni.engine.duplex.realtime_events`` for the projection state). What is +left here is the per-connection handshake policy (query-param defaults, +autostart / resume-only rules, ``session.resume`` parsing), the wire defaults +used to translate appends, and error rendering through the typed +:class:`~vllm_omni.engine.duplex.events.ErrorEvent`. +""" + from __future__ import annotations -import asyncio -import base64 -import binascii -from collections.abc import Callable -from uuid import uuid4 +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import TYPE_CHECKING -import numpy as np - -from vllm_omni.entrypoints.duplex.audio import ( - convert_input_audio_with_rate, - encode_float32_mono_wav_base64, - validate_input_sample_rate_hz, +from vllm_omni.engine.duplex.commands import DuplexCommand, DuplexCommandError +from vllm_omni.engine.duplex.events import error_event +from vllm_omni.engine.duplex.realtime_commands import ( + RealtimeInputDefaults, + translate_realtime_command, ) -from vllm_omni.entrypoints.duplex.protocol import normalize_native_duplex_key -from vllm_omni.entrypoints.duplex.realtime_state import ( - REALTIME_INPUT_AUDIO_FORMATS, - REALTIME_OUTPUT_AUDIO_FORMATS, - RealtimeStateOwner, - realtime_default_sample_rate_hz, -) -from vllm_omni.entrypoints.duplex.server_vad import ServerVADConfig, parse_session_turn_detection + +if TYPE_CHECKING: + from starlette.websockets import WebSocket + +__all__ = ["RealtimeEnvelope", "RealtimeHandshake", "ResumeRequest", "parse_resume_request"] + +_TRUE_VALUES = {"1", "true", "yes", "on"} +_FALSE_VALUES = {"0", "false", "no", "off"} + +#: Client events the envelope handles itself (never translated into commands). +#: ``conversation.item.retrieve`` is deliberately absent: the conversation items +#: it reads live engine-side, so it travels as a command like any other event. +ENVELOPE_EVENT_TYPES = frozenset({"session.resume", "session.event_ack"}) -class RealtimeInputTranslator(RealtimeStateOwner): - """Translate and validate client Realtime events for the duplex core.""" +@dataclass(frozen=True, slots=True) +class RealtimeHandshake: + """What the first client message asks for.""" - _opened: bool - _initial_session_update: bool - _input_speech_started: bool - _input_audio_buffer_has_audio: bool - _input_audio_buffer_had_non_speech: bool - _input_audio_buffer_transcript_parts: list[str] - _active_input_item_id: str | None - _last_conversation_item_id: str | None - _active_response_id: str | None - _last_response_id: str | None - _input_audio_format: str - _output_audio_format: str - _input_sample_rate_hz: int - _conversation_items: dict[str, dict[str, object]] - _item_truncation_cursors: dict[str, tuple[int, int]] - _pending_outbound: asyncio.Queue[dict[str, object]] + #: ``"open"`` (session object in ``session_payload``) or ``"resume"`` (resume payload). + kind: str + session_payload: dict[str, object] = field(default_factory=dict) + resume_payload: dict[str, object] = field(default_factory=dict) + #: A client event that arrived before the session existed and must run right after opening. + pending_command_payload: dict[str, object] | None = None - async def _send_realtime_payload(self, payload: dict[str, object]) -> None: - raise NotImplementedError - def _realtime_error_payload( - self, +@dataclass(frozen=True, slots=True) +class ResumeRequest: + """A validated ``session.resume`` client event.""" + + session_id: str + resume_token: str + last_received_server_event_seq: int + + +def parse_resume_request(event: Mapping[str, object]) -> ResumeRequest | None: + """Validate the shape of ``session.resume``; ``None`` when any field is missing or malformed.""" + session_id = event.get("session_id") + resume_token = event.get("resume_token") + last_received = event.get("last_received_server_event_seq", 0) + if ( + not isinstance(session_id, str) + or not session_id + or not isinstance(resume_token, str) + or not resume_token + or not isinstance(last_received, int) + or last_received < 0 + ): + return None + return ResumeRequest( + session_id=session_id, + resume_token=resume_token, + last_received_server_event_seq=last_received, + ) + + +@dataclass +class RealtimeEnvelope: + """Per-connection Realtime wire policy (no session state).""" + + default_model: str | None = None + #: ``?resume=1`` / ``?autostart=0``: the first message must be ``session.resume``/``session.update``. + resume_only: bool = False + opened: bool = False + autostarted_default_session: bool = False + defaults: RealtimeInputDefaults = field(default_factory=RealtimeInputDefaults) + + @classmethod + def from_query_params(cls, query_params: Mapping[str, str] | WebSocket) -> RealtimeEnvelope: + if hasattr(query_params, "query_params"): + query_params = query_params.query_params + getter = query_params.get if hasattr(query_params, "get") else None + resume_only = getter("resume") if getter is not None else None + autostart = getter("autostart") if getter is not None else None + model = getter("model") if getter is not None else None + return cls( + default_model=model if isinstance(model, str) and model else None, + resume_only=( + str(resume_only).strip().lower() in _TRUE_VALUES or str(autostart).strip().lower() in _FALSE_VALUES + ), + ) + + # ---- handshake ---- + + def default_session_payload(self) -> dict[str, object]: + return {"model": self.default_model} + + def initial_open_payload(self) -> dict[str, object] | None: + """Session object to open with *before* any client message (``?model=`` autostart), else None.""" + if self.opened or self.resume_only or self.autostarted_default_session or not self.default_model: + return None + self.opened = True + self.autostarted_default_session = True + return self.default_session_payload() + + def first_message(self, payload: Mapping[str, object]) -> RealtimeHandshake: + """Classify the first client message (call only while ``not opened``). + + ``session.update`` opens with its session object; ``session.resume`` + resumes; any other event autostarts the default session and is then + treated as a command (``pending_command_payload``). + """ + event_type = payload.get("type") + self.opened = True + if event_type == "session.resume": + return RealtimeHandshake(kind="resume", resume_payload=dict(payload)) + if event_type == "session.update": + session = payload.get("session") + session_payload = dict(session) if isinstance(session, dict) else dict(payload) + session_payload.pop("type", None) + self.note_session_payload(session_payload) + return RealtimeHandshake(kind="open", session_payload=session_payload) + self.autostarted_default_session = True + return RealtimeHandshake( + kind="open", + session_payload=self.default_session_payload(), + pending_command_payload=dict(payload), + ) + + def note_session_payload(self, session_payload: Mapping[str, object]) -> None: + """Track wire defaults declared by a session object (open or session.update).""" + self.defaults = self.defaults.with_session_payload(session_payload) + + # ---- commands ---- + + def is_envelope_event(self, payload: Mapping[str, object]) -> bool: + return payload.get("type") in ENVELOPE_EVENT_TYPES + + def translate(self, payload: Mapping[str, object]) -> DuplexCommand: + """Wire event -> command (raises :class:`DuplexCommandError`).""" + if payload.get("type") == "session.update": + session = payload.get("session") + if isinstance(session, dict): + self.note_session_payload(session) + return translate_realtime_command(payload, defaults=self.defaults) + + # ---- outbound helpers ---- + + @staticmethod + def error_payload( code: str, message: str, *, event_id: object | None = None, param: object | None = None, ) -> dict[str, object]: - raise NotImplementedError - - _response_is_done: Callable[[object], bool] - - @property - def _turn_based_server_vad(self) -> bool: - return self._turn_detection is not None and not self._native_input_append - - def _reset_realtime_input_buffer_state(self) -> None: - self._input_speech_started = False - self._active_input_item_id = None - self._input_audio_buffer_has_audio = False - self._input_audio_buffer_had_non_speech = False - self._input_audio_buffer_transcript_parts.clear() - - async def discard_pending_input_audio( - self, - *, - audio_end_ms: int | None = None, - ) -> None: - """Drop Realtime input-buffer state that was consumed as overlap. - - The serving overlap policy may classify a short acknowledgement such - as "continue" as a side-channel signal instead of user turn input. The - Realtime protocol has already observed the append and may have emitted - ``speech_started``; reset that transient buffer so a later commit does - not accidentally attach the acknowledgement transcript/audio to the - next real user turn. - """ - if self._input_speech_started and self._active_input_item_id is not None: - await self._send_realtime_payload( - { - "type": "input_audio_buffer.speech_stopped", - "audio_end_ms": max(0, int(audio_end_ms or 0)), - "item_id": self._active_input_item_id, - } - ) - self._reset_realtime_input_buffer_state() - - async def _send_realtime_input_ack(self, event: dict[str, object]) -> None: - if event.get("type") != "conversation.item.create": - return - item = event.get("item") - if not isinstance(item, dict): - return - previous_item_id = event.get("previous_item_id") - if isinstance(previous_item_id, str): - item["_previous_item_id"] = previous_item_id - if item.get("role") != "user": - return - self._conversation_items[str(item["id"])] = item - for payload in self._conversation_item_added_events(item): - await self._send_realtime_payload(payload) - - async def _to_duplex_event(self, event: dict[str, object]) -> dict[str, object] | None: - event_type = event.get("type") - if event_type == "session.update": - session = event.get("session") - session_payload: dict[str, object] = session if isinstance(session, dict) else event - try: - turn_detection_configured, server_vad = parse_session_turn_detection(session_payload) - except ValueError as exc: - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_turn_detection", - str(exc), - event_id=event.get("event_id"), - param="turn_detection", - ) - ) - return None - normalized_turn_detection = server_vad.as_dict() if server_vad is not None else None - previous_protocol_config = self._realtime_protocol_config() - try: - self._apply_realtime_session_defaults(session_payload) - except ValueError as exc: - self._set_realtime_protocol_config(previous_protocol_config) - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_audio_format", - str(exc), - event_id=event.get("event_id"), - param="audio.input.format", - ) - ) - return None - turn_detection_error = self._validate_realtime_turn_detection( - session_payload, - configured=turn_detection_configured, - normalized=normalized_turn_detection, - ) - if turn_detection_error is not None: - self._set_realtime_protocol_config(previous_protocol_config) - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_turn_detection", - turn_detection_error, - event_id=event.get("event_id"), - param="turn_detection", - ) - ) - return None - turn_detection_candidate = self._prepare_realtime_turn_detection_update( - session_payload, - configured=turn_detection_configured, - normalized=normalized_turn_detection, - ) - try: - session_payload.update(self._realtime_overlap_fields(session_payload)) - if not self._opened: - self._opened = True - self._initial_session_update = True - translated = self._session_create_from_realtime(session_payload) - event_id = event.get("event_id") - if isinstance(event_id, str) and event_id: - translated["_realtime_event_id"] = event_id - else: - translated = { - "type": "turn.signal", - "event": "session.update", - "payload": session_payload, - "realtime_event_id": event.get("event_id"), - } - candidate_protocol_config = self._realtime_protocol_config() - finally: - self._set_realtime_protocol_config(previous_protocol_config) - if turn_detection_candidate[0] or candidate_protocol_config != previous_protocol_config: - self._pending_turn_detection_update = ( - *turn_detection_candidate, - asyncio.Event(), - candidate_protocol_config, - ) - return translated - if event_type == "conversation.item.create": - return await self._conversation_item_to_duplex(event) - if event_type == "conversation.item.delete": - item_id = event.get("item_id") - if not isinstance(item_id, str) or not item_id: - await self._send_realtime_payload( - self._realtime_error_payload( - "missing_item_id", - "conversation.item.delete requires item_id", - event_id=event.get("event_id"), - ) - ) - return None - if item_id not in self._conversation_items: - await self._send_realtime_payload( - self._realtime_error_payload( - "item_not_found", - f"Conversation item not found: {item_id}", - event_id=event.get("event_id"), - ) - ) - return None - self._remove_conversation_item(item_id) - await self._pending_outbound.put( - { - "type": "turn.signal", - "event": "conversation.item.delete", - "payload": {"item_id": item_id}, - } - ) - return None - if event_type == "conversation.item.retrieve": - item_id = event.get("item_id") - if not isinstance(item_id, str) or not item_id: - await self._send_realtime_payload( - self._realtime_error_payload( - "missing_item_id", - "conversation.item.retrieve requires item_id", - event_id=event.get("event_id"), - ) - ) - return None - item = self._conversation_items.get(item_id) - if item is None: - await self._send_realtime_payload( - self._realtime_error_payload( - "item_not_found", - f"Conversation item not found: {item_id}", - event_id=event.get("event_id"), - ) - ) - return None - await self._send_realtime_payload({"type": "conversation.item.retrieved", "item": item}) - return None - if event_type == "conversation.item.truncate": - item_id = event.get("item_id") - audio_end_ms = event.get("audio_end_ms") - content_index = event.get("content_index", 0) - if not isinstance(item_id, str) or not item_id: - await self._send_realtime_payload( - self._realtime_error_payload( - "missing_item_id", - "conversation.item.truncate requires item_id", - event_id=event.get("event_id"), - ) - ) - return None - if not isinstance(audio_end_ms, int | float): - await self._send_realtime_payload( - self._realtime_error_payload( - "bad_event", - "conversation.item.truncate requires numeric audio_end_ms", - event_id=event.get("event_id"), - param="audio_end_ms", - ) - ) - return None - item = self._conversation_items.get(item_id) - if not isinstance(item, dict): - await self._send_realtime_payload( - self._realtime_error_payload( - "item_not_found", - f"Conversation item not found: {item_id}", - event_id=event.get("event_id"), - ) - ) - return None - truncate_error = self._validate_realtime_item_truncate( - item, - content_index=int(content_index) if isinstance(content_index, int | float) else 0, - audio_end_ms=int(audio_end_ms), - ) - if truncate_error is not None: - await self._send_realtime_payload( - self._realtime_error_payload( - "bad_event", - truncate_error, - event_id=event.get("event_id"), - ) - ) - return None - self._item_truncation_cursors[item_id] = ( - int(content_index) if isinstance(content_index, int | float) else 0, - int(audio_end_ms), - ) - self._truncate_realtime_item_content( - item, - content_index=int(content_index) if isinstance(content_index, int | float) else 0, - audio_end_ms=int(audio_end_ms), - ) - await self._pending_outbound.put( - { - "type": "turn.signal", - "event": "conversation.item.truncate", - "payload": { - "item_id": item_id, - "content_index": content_index, - "audio_end_ms": audio_end_ms, - }, - } - ) - await self._pending_outbound.put( - { - "type": "playback.ack", - "item_id": item_id, - "committed_ms": int(audio_end_ms), - "played_ms": int(audio_end_ms), - "truncate": True, - } - ) - return None - if event_type == "input_audio_buffer.append": - audio = event.get("audio") or event.get("delta") - event_format = event.get("format") or event.get("input_audio_format") - native_server_vad = self._turn_detection is not None and self._native_input_append - turn_based_server_vad = self._turn_based_server_vad - try: - fmt, format_rate = self._parse_realtime_audio_format(event_format or self._input_audio_format) - sample_rate_hz = event.get("sample_rate_hz") or event.get("sample_rate") or format_rate - if sample_rate_hz is None: - sample_rate_hz = ( - realtime_default_sample_rate_hz(fmt) - if event_format is not None and fmt != self._input_audio_format - else self._input_sample_rate_hz - ) - if turn_based_server_vad: - sample_rate_hz = validate_input_sample_rate_hz(sample_rate_hz) - except ValueError as exc: - await self._send_realtime_payload( - self._realtime_error_payload( - "bad_event", - str(exc), - event_id=event.get("event_id"), - param="sample_rate_hz", - ) - ) - return None - channels = event.get("channels", 1) - if not self._is_supported_realtime_input_format(fmt): - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_audio_format", - f"Unsupported input_audio_format: {fmt}", - event_id=event.get("event_id"), - ) - ) - return None - if turn_based_server_vad and (fmt != "pcm16" or channels != 1 or sample_rate_hz not in {16_000, 24_000}): - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_audio_format", - "server_vad requires mono PCM16 input at 16 kHz or 24 kHz", - event_id=event.get("event_id"), - param="audio.input.format", - ) - ) - return None - looks_like_speech: bool | None = None - if not turn_based_server_vad: - try: - audio, fmt, sample_rate_hz = convert_input_audio_with_rate( - audio, - fmt, - sample_rate_hz=sample_rate_hz if isinstance(sample_rate_hz, int | float) else None, - ) - except ValueError as exc: - await self._send_realtime_payload( - self._realtime_error_payload( - "bad_event", - str(exc), - event_id=event.get("event_id"), - param="sample_rate_hz", - ) - ) - return None - if not turn_based_server_vad and not native_server_vad: - looks_like_speech = self._input_looks_like_speech(event, audio=audio, fmt=fmt) - has_audio_payload = isinstance(audio, str) and bool(audio) - self._input_audio_buffer_has_audio = self._input_audio_buffer_has_audio or ( - looks_like_speech and has_audio_payload - ) - self._input_audio_buffer_had_non_speech = self._input_audio_buffer_had_non_speech or ( - not looks_like_speech and has_audio_payload - ) - if looks_like_speech: - await self._emit_input_speech_started(event) - self._remember_input_transcript_hint(event) - payload = { - "type": "input_audio_buffer.append", - "audio": audio, - "format": fmt, - "sample_rate_hz": sample_rate_hz, - } - event_id = event.get("event_id") - if self._turn_detection is not None and isinstance(event_id, str) and event_id: - payload["_realtime_event_id"] = event_id - video_frames = event.get("video_frames") - if video_frames is not None: - frames_error = self._validate_realtime_video_frames(video_frames, event.get("max_slice_nums")) - if frames_error is not None: - await self._send_realtime_payload( - self._realtime_error_payload( - "invalid_video_frames", - frames_error, - event_id=event.get("event_id"), - param="video_frames", - ) - ) - return None - if isinstance(video_frames, list): - payload["video_frames"] = [frame for frame in video_frames if isinstance(frame, str) and frame] - self._copy_realtime_input_hints(event, payload) - if looks_like_speech is not None: - payload["is_speech"] = looks_like_speech - return payload - if event_type == "input_audio_buffer.commit": - if self._turn_based_server_vad: - await self._send_realtime_payload( - self._realtime_error_payload( - "server_vad_manual_commit_unsupported", - ( - "input_audio_buffer.commit is not supported while server_vad is enabled; " - "wait for server endpointing, or start a session with turn_detection=null " - "to use manual commit" - ), - event_id=event.get("event_id"), - ) - ) - return None - if not self._input_audio_buffer_has_audio: - if self._input_audio_buffer_had_non_speech: - self._input_audio_buffer_had_non_speech = False - self._active_input_item_id = None - return { - "type": "input_audio_buffer.commit", - "final": event.get("final", True), - "response_create": False, - "is_speech": False, - } - await self._send_realtime_payload( - self._realtime_error_payload( - "input_audio_buffer_empty", - "input_audio_buffer.commit requires a non-empty input audio buffer", - event_id=event.get("event_id"), - ) - ) - return None - item_id = self._active_input_item_id or f"item_{uuid4().hex}" - await self._emit_input_speech_stopped(event, item_id=item_id) - transcript = self._consume_input_transcript_hint() - self._active_input_item_id = None - self._input_audio_buffer_has_audio = False - self._input_audio_buffer_had_non_speech = False - payload = { - "type": "input_audio_buffer.commit", - "final": event.get("final", True), - "item_id": item_id, - "response_create": bool(event.get("response_create", False)), - } - if transcript: - payload["transcript"] = transcript - return payload - if event_type == "input_audio_buffer.clear": - self._reset_realtime_input_buffer_state() - return {"type": "input_audio_buffer.clear", "reason": event_type} - if event_type == "output_audio_buffer.clear": - payload = {"type": "output_audio_buffer.clear", "reason": event_type} - response_id = event.get("response_id") - if not isinstance(response_id, str) or not response_id: - response_id = self._active_response_id or self._last_response_id - if self._response_is_done(response_id): - await self._send_realtime_payload( - { - "type": "output_audio_buffer.cleared", - "response_id": response_id, - } - ) - return None - if isinstance(response_id, str) and response_id: - payload["response_id"] = response_id - return payload - if event_type == "response.cancel": - payload = {"type": "response.cancel", "reason": event_type} - response_id = event.get("response_id") - if not isinstance(response_id, str) or not response_id: - response_id = self._active_response_id or self._last_response_id - if self._response_is_done(response_id): - await self._send_realtime_payload( - self._realtime_error_payload( - "response_not_active", - f"Response is already complete: {response_id}", - event_id=event.get("event_id"), - ) - ) - return None - if isinstance(response_id, str) and response_id: - payload["response_id"] = response_id - return payload - if event_type == "response.create": - response_payload = event.get("response") - if isinstance(response_payload, dict): - format_error = self._validate_realtime_response_audio_formats(response_payload) - if format_error is not None: - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_audio_format", - format_error, - event_id=event.get("event_id"), - ) - ) - return None - return { - "type": "response.create", - "response": response_payload if isinstance(response_payload, dict) else {}, - } - if event_type in {"session.close", "close"}: - return {"type": "session.close"} - return event + """Wire JSON of a transport-level error (derived from the typed ``ErrorEvent``).""" + return error_event(code, message, event_id=event_id, param=param).to_realtime() @staticmethod - def _normalize_conversation_item(item: dict[str, object]) -> dict[str, object]: - normalized = dict(item) - normalized.setdefault("id", f"item_{uuid4().hex}") - normalized.setdefault("object", "realtime.item") - normalized.setdefault("type", "message") - normalized.setdefault("status", "completed") - if "role" not in normalized and normalized.get("type") == "message": - normalized["role"] = "user" - if not isinstance(normalized.get("content"), list): - normalized["content"] = [] - return normalized - - @staticmethod - def _truncate_realtime_item_content( - item: dict[str, object], - *, - content_index: int, - audio_end_ms: int, - ) -> None: - content = item.get("content") - if not isinstance(content, list) or not content: - return - index = max(0, int(content_index)) - if index >= len(content): - return - part = content[index] - if not isinstance(part, dict): - return - transcript = part.get("transcript") - if not isinstance(transcript, str) or not transcript: - return - marks = part.get("audio_text_marks") - if isinstance(marks, list): - keep_chars = RealtimeInputTranslator._text_chars_for_audio_ms_from_marks( - audio_end_ms, - len(transcript), - marks, - final_ms=part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms"), - ) - part["transcript"] = transcript[:keep_chars].rstrip() - return - duration_ms = part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms") - if isinstance(duration_ms, int | float) and duration_ms > 0: - keep_chars = int(len(transcript) * max(0.0, min(1.0, int(audio_end_ms) / float(duration_ms)))) - part["transcript"] = transcript[:keep_chars].rstrip() - elif audio_end_ms <= 0: - part["transcript"] = "" - - @staticmethod - def _validate_realtime_item_truncate( - item: dict[str, object], - *, - content_index: int, - audio_end_ms: int, - ) -> str | None: - if item.get("type") != "message" or item.get("role") != "assistant": - return "conversation.item.truncate only supports assistant message items" - if audio_end_ms < 0: - return "conversation.item.truncate requires non-negative audio_end_ms" - content = item.get("content") - if not isinstance(content, list) or not content: - return None - index = max(0, int(content_index)) - if index >= len(content): - return f"conversation.item.truncate content_index out of range: {content_index}" - part = content[index] - if not isinstance(part, dict): - return "conversation.item.truncate target content part is invalid" - if part.get("type") not in {"audio", "output_audio"}: - return "conversation.item.truncate target content part must be audio" - duration_ms = part.get("audio_duration_ms") or part.get("duration_ms") or part.get("audio_ms") - if isinstance(duration_ms, int | float) and int(duration_ms) >= 0 and audio_end_ms > int(duration_ms): - return "conversation.item.truncate audio_end_ms exceeds item audio duration" - return None - - @staticmethod - def _text_chars_for_audio_ms_from_marks( - audio_end_ms: int, - text_len: int, - marks: list[object], - *, - final_ms: object | None = None, - ) -> int: - if text_len <= 0: - return 0 - clean_marks: list[tuple[int, int]] = [] - for mark in marks: - if not isinstance(mark, dict): - continue - raw_text_chars = mark.get("text_chars") - raw_audio_end_ms = mark.get("audio_end_ms", mark.get("audio_ms")) - if not isinstance(raw_text_chars, int | float) or not isinstance(raw_audio_end_ms, int | float): - continue - clean_marks.append((max(0, int(raw_audio_end_ms)), min(text_len, max(0, int(raw_text_chars))))) - if not clean_marks: - return 0 if audio_end_ms <= 0 else text_len - clean_marks.sort(key=lambda item: item[0]) - audio_end_ms = max(0, int(audio_end_ms)) - if audio_end_ms <= 0: - return 0 - previous_ms = 0 - previous_chars = 0 - for mark_ms, mark_chars in clean_marks: - mark_ms = max(previous_ms, mark_ms) - mark_chars = max(previous_chars, min(text_len, mark_chars)) - if audio_end_ms <= mark_ms: - if mark_ms <= previous_ms: - return mark_chars - ratio = (audio_end_ms - previous_ms) / max(1, mark_ms - previous_ms) - return int(previous_chars + (mark_chars - previous_chars) * max(0.0, min(1.0, ratio))) - previous_ms = mark_ms - previous_chars = mark_chars - if isinstance(final_ms, int | float) and int(final_ms) > previous_ms: - if audio_end_ms >= int(final_ms): - return text_len - ratio = (audio_end_ms - previous_ms) / max(1, int(final_ms) - previous_ms) - return int(previous_chars + (text_len - previous_chars) * max(0.0, min(1.0, ratio))) - return text_len if audio_end_ms >= previous_ms else previous_chars - - def _conversation_item_added_events(self, item: dict[str, object]) -> list[dict[str, object]]: - item_id = item.get("id") - explicit_previous_item_id = item.pop("_previous_item_id", None) - previous_item_id = ( - explicit_previous_item_id if isinstance(explicit_previous_item_id, str) else self._last_conversation_item_id - ) - if isinstance(item_id, str) and item_id: - self._last_conversation_item_id = item_id - return [ - { - "type": "conversation.item.added", - "previous_item_id": previous_item_id, - "item": item, - }, - ] - - def _conversation_item_done_event(self, item: dict[str, object]) -> dict[str, object]: - item_id = item.get("id") - previous_item_id = self._previous_item_id(item_id) if isinstance(item_id, str) else None - if isinstance(item_id, str) and item_id: - already_known = item_id in self._conversation_items - self._conversation_items[item_id] = item - # Completing an item that was already added must not move it to the - # end of the conversation. A later user item may have been added - # while this response was still in progress. - if not already_known: - self._last_conversation_item_id = item_id - return { - "type": "conversation.item.done", - "previous_item_id": previous_item_id, - "item": item, - } - - def _remove_conversation_item(self, item_id: str) -> bool: - removed = self._conversation_items.pop(item_id, None) is not None - if self._last_conversation_item_id == item_id: - remaining_ids = list(self._conversation_items) - self._last_conversation_item_id = remaining_ids[-1] if remaining_ids else None - self._item_truncation_cursors.pop(item_id, None) - return removed - - def _response_output_item_added_events( - self, - *, - response_id: object, - item: dict[str, object], - ) -> list[dict[str, object]]: - return [ - { - "type": "response.output_item.added", - "response_id": response_id, - "output_index": 0, - "item": item, - }, - ] - - @staticmethod - def _response_content_part(*, transcript: str = "") -> dict[str, object]: - return { - "type": "audio", - "transcript": transcript, - } - - @staticmethod - def _response_text_content_part(*, text: str = "") -> dict[str, object]: - return { - "type": "text", - "text": text, - } - - @staticmethod - def _response_item_content_part( - *, - transcript: str = "", - audio_duration_ms: int | None = None, - audio_text_marks: list[dict[str, int]] | None = None, - ) -> dict[str, object]: - part: dict[str, object] = { - "type": "output_audio", - "transcript": transcript, - } - if audio_duration_ms is not None: - part["audio_duration_ms"] = int(audio_duration_ms) - if audio_text_marks: - part["audio_text_marks"] = [dict(mark) for mark in audio_text_marks] - return part - - @staticmethod - def _response_item_text_content_part(*, text: str = "") -> dict[str, object]: - return { - "type": "output_text", - "text": text, - } - - def _previous_item_id(self, item_id: str) -> str | None: - previous: str | None = None - for known_id in self._conversation_items: - if known_id == item_id: - return previous - previous = known_id - if self._last_conversation_item_id == item_id: - return previous - return self._last_conversation_item_id - - def _input_audio_buffer_committed_event( - self, - *, - item_id: str, - ) -> dict[str, object]: - return { - "type": "input_audio_buffer.committed", - "previous_item_id": self._previous_item_id(item_id), - "item_id": item_id, - } - - def _session_create_from_realtime(self, session_payload: dict[str, object]) -> dict[str, object]: - model = session_payload.get("model") - audio_config = session_payload.get("audio") - audio_input = audio_config.get("input") if isinstance(audio_config, dict) else None - audio_output = audio_config.get("output") if isinstance(audio_config, dict) else None - extra_body_payload = session_payload.get("extra_body") - extra_body = ( - normalize_native_duplex_key(dict(extra_body_payload)) if isinstance(extra_body_payload, dict) else {} - ) - extra_body["realtime_session_payload"] = self._json_safe_realtime_payload(session_payload) - if isinstance(session_payload.get("tools"), list): - extra_body["realtime_tools"] = session_payload["tools"] - if isinstance(session_payload.get("tool_choice"), str | dict): - extra_body["realtime_tool_choice"] = session_payload["tool_choice"] - metadata = session_payload.get("metadata") - if isinstance(metadata, dict): - extra_body["realtime_metadata"] = dict(metadata) - include = session_payload.get("include") - if isinstance(include, list): - extra_body["realtime_include"] = list(include) - prompt = session_payload.get("prompt") - if isinstance(prompt, dict): - extra_body["realtime_prompt"] = dict(prompt) - input_audio_transcription = self._input_audio_transcription_config(session_payload) - if isinstance(input_audio_transcription, dict): - extra_body["realtime_input_audio_transcription"] = dict(input_audio_transcription) - noise_reduction = session_payload.get("input_audio_noise_reduction") - if isinstance(noise_reduction, dict): - extra_body["realtime_input_audio_noise_reduction"] = dict(noise_reduction) - audio_noise_reduction = audio_input.get("noise_reduction") if isinstance(audio_input, dict) else None - if isinstance(audio_noise_reduction, dict): - extra_body["realtime_input_audio_noise_reduction"] = dict(audio_noise_reduction) - audio = session_payload.get("audio") - if isinstance(audio, dict): - extra_body["realtime_audio"] = dict(audio) - if isinstance(session_payload.get("tracing"), str | dict): - extra_body["realtime_tracing"] = session_payload["tracing"] - response_format = self._duplex_response_format(self._output_audio_format) - extra_body.setdefault("realtime_output_audio_format", self._output_audio_format) - overlap_fields = self._realtime_overlap_fields(session_payload) - voice = session_payload.get("voice") - if not isinstance(voice, str) and isinstance(audio_output, dict): - voice = audio_output.get("voice") - speed = session_payload.get("speed") - if not isinstance(speed, int | float) and isinstance(audio_output, dict): - speed = audio_output.get("speed") - return { - "type": "session.create", - "session_id": session_payload.get("session_id") or session_payload.get("id"), - "session": { - "model": model, - "modalities": ( - session_payload.get("modalities") or session_payload.get("output_modalities") or ["text", "audio"] - ), - "instructions": session_payload.get("instructions"), - "voice": voice, - "ref_audio": session_payload.get("ref_audio"), - "response_format": response_format, - "temperature": session_payload.get("temperature"), - "max_tokens": self.realtime_max_output_tokens( - session_payload.get("max_response_output_tokens") - or session_payload.get("max_output_tokens") - or session_payload.get("max_tokens") - ), - "speed": speed, - "idle_timeout_s": session_payload.get("idle_timeout_s") or 300.0, - **overlap_fields, - "extra_body": extra_body, - }, - } - - def _apply_realtime_session_defaults(self, session_payload: dict[str, object]) -> None: - input_format: object = session_payload.get("input_audio_format") - audio_config = session_payload.get("audio") - if input_format is None and isinstance(audio_config, dict): - audio_input = audio_config.get("input") - if isinstance(audio_input, dict): - input_format = audio_input.get("format") - input_format_configured = input_format is not None - input_format, input_rate = self._parse_realtime_audio_format(input_format) - if input_format is not None and not ( - isinstance(input_format, str) and input_format.lower() in REALTIME_INPUT_AUDIO_FORMATS - ): - raise ValueError(f"Unsupported input_audio_format: {input_format}") - if isinstance(input_format, str): - self._input_audio_format = input_format - output_format: object = session_payload.get("output_audio_format") or session_payload.get("response_format") - output_rate_raw: object | None = None - if output_format is None and isinstance(audio_config, dict): - audio_output = audio_config.get("output") - if isinstance(audio_output, dict): - output_format = audio_output.get("format") - output_rate_raw = audio_output.get("sample_rate_hz") or audio_output.get("sample_rate") - output_format, output_rate = self._parse_realtime_audio_format(output_format) - if output_rate is None and output_rate_raw is not None: - output_rate = validate_input_sample_rate_hz(output_rate_raw) - if output_format is not None and not ( - isinstance(output_format, str) and output_format.lower() in REALTIME_OUTPUT_AUDIO_FORMATS - ): - raise ValueError(f"Unsupported output_audio_format: {output_format}") - if isinstance(output_format, str): - self._output_audio_format = self._realtime_output_format(output_format) - sample_rate = session_payload.get("sample_rate_hz") or session_payload.get("sample_rate") - if sample_rate is None and isinstance(audio_config, dict): - audio_input = audio_config.get("input") - if isinstance(audio_input, dict): - sample_rate = audio_input.get("sample_rate_hz") or audio_input.get("sample_rate") - if sample_rate is None: - sample_rate = input_rate - if sample_rate is not None: - self._input_sample_rate_hz = validate_input_sample_rate_hz(sample_rate) - elif input_format_configured: - default_input_rate = realtime_default_sample_rate_hz(input_format) - if default_input_rate is not None: - self._input_sample_rate_hz = default_input_rate - if output_rate is not None: - self._output_sample_rate_hz = output_rate - overlap_fields = self._realtime_overlap_fields(session_payload) - overlap_silence_rms = overlap_fields.get("overlap_silence_rms") - if isinstance(overlap_silence_rms, int | float): - self._overlap_silence_rms = max(0.0, float(overlap_silence_rms)) - - @classmethod - def _validate_realtime_response_audio_formats(cls, response_payload: dict[str, object]) -> str | None: - output_format: object = response_payload.get("output_audio_format") or response_payload.get("response_format") - audio_config = response_payload.get("audio") - if output_format is None and isinstance(audio_config, dict): - audio_output = audio_config.get("output") - if isinstance(audio_output, dict): - output_format = audio_output.get("format") - try: - parsed_output, _ = cls._parse_realtime_audio_format(output_format) - except ValueError as exc: - return str(exc) - if output_format is not None and not ( - isinstance(parsed_output, str) and parsed_output.lower() in REALTIME_OUTPUT_AUDIO_FORMATS - ): - return f"Unsupported output_audio_format: {output_format}" - return None - - @staticmethod - def _json_safe_realtime_payload(payload: dict[str, object]) -> dict[str, object]: - clean: dict[str, object] = {} - for key, value in payload.items(): - if key == "extra_body": - continue - if isinstance(value, str | int | float | bool) or value is None: - clean[key] = value - elif isinstance(value, dict): - clean[key] = RealtimeInputTranslator._json_safe_realtime_payload(value) - elif isinstance(value, list): - clean[key] = [ - (RealtimeInputTranslator._json_safe_realtime_payload(item) if isinstance(item, dict) else item) - for item in value - if isinstance(item, str | int | float | bool | dict) or item is None - ] - return clean - - @staticmethod - def _parse_realtime_audio_format(raw_format: object) -> tuple[object, int | None]: - def normalize_format(fmt: str) -> str: - normalized = fmt.lower() - if normalized in {"audio/pcm", "pcm"}: - return "pcm16" - if normalized in {"audio/wav", "wav"}: - return "wav" - if normalized in {"audio/pcm16", "pcm16", "pcm_s16le", "s16le"}: - return "pcm16" - if normalized in {"audio/pcm_f32le", "pcm_f32le", "f32le"}: - return "pcm_f32le" - if normalized in {"audio/g711_ulaw", "g711_ulaw", "g711-ulaw", "ulaw", "mulaw"}: - return "g711_ulaw" - if normalized in {"audio/g711_alaw", "g711_alaw", "g711-alaw", "alaw"}: - return "g711_alaw" - return fmt - - if isinstance(raw_format, str): - return normalize_format(raw_format), None - if not isinstance(raw_format, dict): - return raw_format, None - rate = raw_format.get("rate") or raw_format.get("sample_rate_hz") or raw_format.get("sample_rate") - sample_rate_hz = validate_input_sample_rate_hz(rate) if rate is not None else None - fmt = raw_format.get("type") or raw_format.get("format") - if not isinstance(fmt, str): - return raw_format, sample_rate_hz - return normalize_format(fmt), sample_rate_hz - - @staticmethod - def _duplex_response_format(realtime_format: str) -> str: - normalized = realtime_format.lower() - if normalized in {"pcm16", "pcm_s16le", "s16le"}: - return "pcm" - if normalized in {"g711_ulaw", "g711_alaw"}: - return "pcm" - if normalized in {"wav", "pcm"}: - return normalized - return "wav" - - @staticmethod - def _realtime_output_format(duplex_format: object) -> str: - if isinstance(duplex_format, str) and duplex_format.lower() in {"g711_ulaw", "g711_alaw"}: - return duplex_format.lower() - if isinstance(duplex_format, str) and duplex_format.lower() == "pcm": - return "pcm16" - return str(duplex_format or "wav") - - @staticmethod - def _is_supported_realtime_input_format(fmt: object) -> bool: - return isinstance(fmt, str) and fmt.lower() in REALTIME_INPUT_AUDIO_FORMATS - - @staticmethod - def _input_explicitly_non_speech(event: dict[str, object]) -> bool: - for key in ("is_speech", "speech"): - value = event.get(key) - if isinstance(value, bool): - return not value - vad = event.get("vad") - if isinstance(vad, dict): - value = vad.get("is_speech") - if isinstance(value, bool): - return not value - probability = vad.get("speech_probability", vad.get("probability")) - if isinstance(probability, int | float): - return float(probability) < 0.5 - probability = event.get("speech_probability") - return isinstance(probability, int | float) and float(probability) < 0.5 - - def _input_looks_like_speech(self, event: dict[str, object], *, audio: object, fmt: object) -> bool: - if RealtimeInputTranslator._input_explicitly_non_speech(event): - return False - for key in ("is_speech", "speech"): - value = event.get(key) - if isinstance(value, bool): - return value - vad = event.get("vad") - if isinstance(vad, dict): - probability = vad.get("speech_probability", vad.get("probability")) - if isinstance(probability, int | float): - return float(probability) >= 0.5 - probability = event.get("speech_probability") - if isinstance(probability, int | float): - return float(probability) >= 0.5 - if fmt != "pcm_f32le" or not isinstance(audio, str): - return True - try: - raw = base64.b64decode(audio, validate=True) - except (binascii.Error, ValueError): - return True - if len(raw) < 4 or len(raw) % 4 != 0: - return True - samples = np.frombuffer(raw, dtype=np.float32) - if samples.size == 0: - return False - rms = float(np.sqrt(np.mean(np.square(samples.astype(np.float32))))) - threshold = event.get("overlap_silence_rms") - if not isinstance(threshold, int | float): - vad = event.get("vad") - if isinstance(vad, dict): - threshold = vad.get("silence_rms") - silence_rms = float(threshold) if isinstance(threshold, int | float) else self._overlap_silence_rms - return rms >= max(0.0, silence_rms) - - @staticmethod - def _copy_realtime_input_hints(source: dict[str, object], target: dict[str, object]) -> None: - for key in ( - "duration_ms", - "audio_duration_ms", - "audio_start_ms", - "audio_end_ms", - "is_speech", - "speech", - "speech_probability", - "vad", - "overlap_action", - "overlap", - "force_barge_in", - "force_listen", - "text", - "transcript", - ): - if key in source: - target[key] = source[key] - - @classmethod - def _realtime_overlap_fields(cls, session_payload: dict[str, object]) -> dict[str, object]: - fields: dict[str, object] = {} - if isinstance(session_payload.get("overlap_policy"), str): - fields["overlap_policy"] = session_payload["overlap_policy"] - for key in ("overlap_short_ack_ms", "overlap_barge_in_ms", "overlap_silence_rms"): - value = session_payload.get(key) - if isinstance(value, int | float): - fields[key] = value - - if isinstance(session_payload.get("playback_commit_policy"), str): - fields["playback_commit_policy"] = session_payload["playback_commit_policy"] - return fields - - @staticmethod - def _validate_realtime_video_frames(video_frames: object, max_slice_nums: object) -> str | None: - """Validate omni-duplex camera frames on input_audio_buffer.append. - - Wire contract matches the official MiniCPM-o duplex loop: one base - base64 JPEG per ~1 s audio chunk, optionally followed by that unit's - stacked composite tiling the sub-frames captured inside it (at most 2 - images either way). A caller-supplied ``max_slice_nums`` is rejected - rather than silently ignored: slicing is Stage 0's decision here, and - Stage 0 already applies the official HD suggestion for a stacked unit - (``max_slice_nums=[2, 1]``). The wire simply does not let the client - choose it. - """ - if max_slice_nums not in (None, 1): - return "max_slice_nums is not selectable on the wire; Stage0 slices stacked units itself" - if not isinstance(video_frames, list): - return "video_frames must be a list of base64-encoded images" - frames = [frame for frame in video_frames if frame is not None] - if len(frames) > 2: - return "video_frames carries more than 2 frames for one append; send ~1 frame per 1 s chunk" - for frame in frames: - if not isinstance(frame, str) or not frame: - return "video_frames entries must be non-empty base64 strings" - if len(frame) > 4_000_000: - return "video_frames entry exceeds 4MB base64; reduce capture resolution or JPEG quality" - try: - header = base64.b64decode(frame[:64] + "=" * (-len(frame[:64]) % 4)) - except (binascii.Error, ValueError): - return "video_frames entries must be valid base64" - if not (header.startswith(b"\xff\xd8") or header.startswith(b"\x89PNG")): - return "video_frames entries must be JPEG or PNG images" - return None - - def _validate_realtime_turn_detection( - self, - session_payload: dict[str, object], - *, - configured: bool, - normalized: dict[str, object] | None, - ) -> str | None: - effective_turn_detection = normalized if configured else self._turn_detection - if effective_turn_detection is None: - if session_payload.get("overlap_policy") == "barge_in_on_speech": - return "overlap_policy='barge_in_on_speech' requires turn_detection.type='server_vad'" - return None - - interrupt_response = effective_turn_detection.get("interrupt_response") - if interrupt_response is None: - # Serving resolves the default and validates overlap policy once - # the session's runtime capabilities are known. - return None - desired_policy = "barge_in_on_speech" if interrupt_response else "listen_only" - overlap_policy = session_payload.get("overlap_policy") - if isinstance(overlap_policy, str) and overlap_policy != desired_policy: - return f"overlap_policy={overlap_policy!r} conflicts with turn_detection; expected {desired_policy!r}" - return None - - def validate_server_vad_input_audio(self, session_payload: dict[str, object]) -> str | None: - """Validate turn-based VAD audio after serving resolves runtime capabilities.""" - candidate = self._pending_turn_detection_update - input_format, input_rate, *_ = candidate[3] if candidate is not None else self._realtime_protocol_config() - audio_config = session_payload.get("audio") - audio_input = audio_config.get("input") if isinstance(audio_config, dict) else None - channels = audio_input.get("channels") if isinstance(audio_input, dict) else None - raw_format = session_payload.get("input_audio_format") - if raw_format is None and isinstance(audio_input, dict): - raw_format = audio_input.get("format") - if channels is None and isinstance(raw_format, dict): - channels = raw_format.get("channels") - if channels is None: - channels = 1 - if input_format != "pcm16" or channels != 1 or input_rate not in {16_000, 24_000}: - return "server_vad requires mono PCM16 input at 16 kHz or 24 kHz" - return None - - def _prepare_realtime_turn_detection_update( - self, - session_payload: dict[str, object], - *, - configured: bool, - normalized: dict[str, object] | None, - ) -> tuple[bool, dict[str, object] | None]: - if self._pending_turn_detection_update is not None: - raise RuntimeError("A Realtime session update is already pending") - - turn_detection = dict(normalized) if configured and normalized is not None else None - if configured: - if turn_detection is None: - session_payload.setdefault("overlap_policy", "listen_only") - elif "interrupt_response" in turn_detection: - session_payload.setdefault( - "overlap_policy", "barge_in_on_speech" if turn_detection["interrupt_response"] else "listen_only" - ) - session_payload["turn_detection"] = turn_detection - return configured, turn_detection - - def _realtime_protocol_config(self) -> tuple[str, int, str, int | None, float]: - return ( - self._input_audio_format, - self._input_sample_rate_hz, - self._output_audio_format, - self._output_sample_rate_hz, - self._overlap_silence_rms, - ) - - def _set_realtime_protocol_config( - self, - config: tuple[str, int, str, int | None, float], - ) -> None: - ( - self._input_audio_format, - self._input_sample_rate_hz, - self._output_audio_format, - self._output_sample_rate_hz, - self._overlap_silence_rms, - ) = config - - def pending_realtime_input_audio_config_changed(self) -> bool: - candidate = self._pending_turn_detection_update - if candidate is None: - return False - candidate_protocol_config = candidate[3] - return candidate_protocol_config[:2] != self._realtime_protocol_config()[:2] - - async def wait_for_realtime_turn_detection_update(self) -> None: - candidate = self._pending_turn_detection_update - if candidate is not None: - await candidate[2].wait() - - def commit_realtime_turn_detection_update(self, server_vad: ServerVADConfig | None = None) -> None: - self._resolve_realtime_turn_detection_update(commit=True, server_vad=server_vad) - - def reject_realtime_turn_detection_update(self) -> None: - self._resolve_realtime_turn_detection_update(commit=False) - - def _resolve_realtime_turn_detection_update( - self, *, commit: bool, server_vad: ServerVADConfig | None = None - ) -> None: - candidate = self._pending_turn_detection_update - if candidate is None: - return - configured, turn_detection, resolved, protocol_config = candidate - if commit: - self._set_realtime_protocol_config(protocol_config) - if configured: - self._turn_detection = server_vad.as_dict() if server_vad is not None else turn_detection - self._turn_detection_configured = True - self._pending_turn_detection_update = None - resolved.set() - - @staticmethod - def _input_audio_transcription_config(session_payload: dict[str, object]) -> dict[str, object] | None: - transcription = session_payload.get("input_audio_transcription") - if isinstance(transcription, dict): - return transcription - audio_config = session_payload.get("audio") - if not isinstance(audio_config, dict): - return None - audio_input = audio_config.get("input") - if not isinstance(audio_input, dict): - return None - transcription = audio_input.get("transcription") - return transcription if isinstance(transcription, dict) else None - - @staticmethod - def realtime_max_output_tokens(value: object) -> int | None: - """Normalize Realtime max output tokens. - - OpenAI Realtime clients commonly send ``"inf"`` for unbounded output. - The duplex core represents that as ``None`` because model-specific - defaults remain in force. - """ - if value is None: - return None - if isinstance(value, str) and value.strip().lower() in {"inf", "infinity", "unlimited"}: - return None - if isinstance(value, int) and value > 0: - return int(value) - return None - - async def _conversation_item_to_duplex(self, event: dict[str, object]) -> dict[str, object] | None: - raw_item = event.get("item") - if not isinstance(raw_item, dict): - return None - public_item = self._normalize_conversation_item(raw_item) - previous_item_id = event.get("previous_item_id") - if isinstance(previous_item_id, str): - public_item["_previous_item_id"] = previous_item_id - item_id = str(public_item["id"]) - item_type = public_item.get("type") - role = public_item.get("role") - if item_type == "function_call_output": - call_id = public_item.get("call_id") - output = public_item.get("output") - matching_call = any( - known.get("type") == "function_call" and known.get("call_id") == call_id - for known in self._conversation_items.values() - ) - duplicate_output = any( - known.get("type") == "function_call_output" and known.get("call_id") == call_id - for known in self._conversation_items.values() - ) - if not isinstance(call_id, str) or not call_id or not matching_call: - await self._send_realtime_payload( - self._realtime_error_payload( - "invalid_function_call_output", - "function_call_output requires the call_id of a completed function call", - event_id=event.get("event_id"), - param="item.call_id", - ) - ) - return None - if not isinstance(output, str) or duplicate_output: - await self._send_realtime_payload( - self._realtime_error_payload( - "invalid_function_call_output", - ( - "function_call_output requires a string output" - if not isinstance(output, str) - else f"function_call_output already exists for call_id {call_id}" - ), - event_id=event.get("event_id"), - param="item.output", - ) - ) - return None - # Do not retain the result until SessionRunner has delivered its - # runtime update. On success, the resulting - # ``conversation.item.created`` event records it in the shared - # input/output protocol state; on failure the client may retry the - # same call_id instead of being rejected by provisional history. - content = public_item.get("content") - if not isinstance(content, list): - return None - parsed_audio_formats: dict[int, tuple[object, int | None]] = {} - for index, part in enumerate(content): - if not isinstance(part, dict) or part.get("type") not in {"input_audio", "audio"}: - continue - raw_format = part.get("format") - try: - parsed_format = self._parse_realtime_audio_format(raw_format or self._input_audio_format) - except ValueError as exc: - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_audio_format", - str(exc), - event_id=event.get("event_id"), - param="item.content.audio", - ) - ) - return None - if not self._is_supported_realtime_input_format(parsed_format[0]): - await self._send_realtime_payload( - self._realtime_error_payload( - "unsupported_audio_format", - f"Unsupported input_audio format in conversation.item.create: {raw_format}", - event_id=event.get("event_id"), - param="item.content.audio", - ) - ) - return None - parsed_audio_formats[index] = parsed_format - if item_type != "message" or role in {"assistant", "system"}: - event["item"] = public_item - return { - "type": "turn.signal", - "event": "conversation.item.create", - "payload": {"item": public_item}, - } - server_vad_full_audio_item = self._turn_based_server_vad - text_chunks: list[str] = [] - audio_events: list[dict[str, object]] = [] - normalized_content: list[object] = [] - normalized_full_audio = False - for index, part in enumerate(content): - if not isinstance(part, dict): - normalized_content.append(part) - continue - if part.get("type") in {"input_text", "text"} and isinstance(part.get("text"), str): - text_chunks.append(str(part["text"])) - if part.get("type") in {"input_audio", "audio"}: - audio = part.get("audio") or part.get("data") - if not isinstance(audio, str) or not audio: - await self._send_realtime_payload( - self._realtime_error_payload( - "bad_audio", - "conversation.item.create input_audio must not be empty", - event_id=event.get("event_id"), - param="item.content.audio", - ) - ) - return None - raw_part_format = part.get("format") - fmt, format_rate = parsed_audio_formats[index] - explicit_sample_rate_hz = part.get("sample_rate_hz") - if explicit_sample_rate_hz is None: - explicit_sample_rate_hz = part.get("sample_rate") - if explicit_sample_rate_hz is None and isinstance(raw_part_format, dict): - explicit_sample_rate_hz = ( - raw_part_format.get("rate") - if "rate" in raw_part_format - else raw_part_format.get("sample_rate_hz", raw_part_format.get("sample_rate")) - ) - sample_rate_hz = explicit_sample_rate_hz or format_rate or self._input_sample_rate_hz - if server_vad_full_audio_item and ( - fmt != self._input_audio_format - or (explicit_sample_rate_hz is not None and explicit_sample_rate_hz != self._input_sample_rate_hz) - ): - await self._send_realtime_payload( - self._realtime_error_payload( - "bad_audio", - "conversation.item.create input_audio must use the Session input format and sample rate", - event_id=event.get("event_id"), - param="item.content.audio", - ) - ) - return None - if server_vad_full_audio_item: - try: - normalized_audio = self._normalize_complete_audio_item_part(audio, self._input_sample_rate_hz) - except ValueError as exc: - await self._send_realtime_payload( - self._realtime_error_payload( - "bad_audio", - str(exc), - event_id=event.get("event_id"), - param="item.content.audio", - ) - ) - return None - normalized_part = dict(part) - normalized_part.pop("data", None) - normalized_part.update(audio=normalized_audio, format="wav", sample_rate_hz=16_000) - normalized_content.append(normalized_part) - normalized_full_audio = True - continue - try: - audio, fmt, sample_rate_hz = convert_input_audio_with_rate( - audio, - fmt, - sample_rate_hz=sample_rate_hz if isinstance(sample_rate_hz, int | float) else None, - ) - except ValueError as exc: - await self._send_realtime_payload( - self._realtime_error_payload( - "bad_audio", - str(exc), - event_id=event.get("event_id"), - param="item.content.audio", - ) - ) - return None - speech_hints = dict(event) - speech_hints.update(part) - if not self._input_looks_like_speech(speech_hints, audio=audio, fmt=fmt): - continue - self._input_speech_started = True - payload = { - "type": "input_audio_buffer.append", - "audio": audio, - "format": fmt, - "sample_rate_hz": sample_rate_hz, - } - self._copy_realtime_input_hints(part, payload) - self._copy_realtime_input_hints(event, payload) - audio_events.append(payload) - continue - normalized_content.append(part) - if server_vad_full_audio_item and normalized_full_audio: - history_item = dict(public_item) - history_item["content"] = normalized_content - history_item["status"] = "completed" - event["item"] = public_item - return { - "type": "turn.signal", - "event": "conversation.item.create", - "payload": { - "item": public_item, - "history_item": history_item, - }, - } - if audio_events: - transcript = self._input_transcript_from_item(public_item) - for extra_event in audio_events[1:]: - self._pending_outbound.put_nowait(extra_event) - commit_payload: dict[str, object] = { - "type": "input_audio_buffer.commit", - "final": True, - "item_id": item_id, - "response_create": False, - } - if transcript: - commit_payload["transcript"] = transcript - self._pending_outbound.put_nowait(commit_payload) - # This complete item is already committed locally; a later ack - # must not be needed to start accepting the next input turn. - self._reset_realtime_input_buffer_state() - event["item"] = public_item - return audio_events[0] - if not text_chunks: - return None - text_item = dict(public_item) - text_item["status"] = "completed" - event["item"] = text_item - return { - "type": "turn.signal", - "event": "conversation.item.create", - "payload": {"item": text_item}, - } - - @staticmethod - def _normalize_complete_audio_item_part( - audio: str, - sample_rate_hz: int, - ) -> str: - """Package a complete Realtime audio item for the chat fallback path.""" - try: - raw = base64.b64decode(audio.strip(), validate=True) - except (binascii.Error, ValueError) as exc: - raise ValueError("conversation.item.create input_audio is not valid Base64") from exc - if not raw or len(raw) % np.dtype(" None: - transcript = event.get("transcript") - if not isinstance(transcript, str): - transcript = event.get("text") if isinstance(event.get("text"), str) else None - if not isinstance(transcript, str): - hints = event.get("hints") - if isinstance(hints, dict): - transcript = hints.get("transcript") - if not isinstance(transcript, str): - transcript = hints.get("text") if isinstance(hints.get("text"), str) else None - if isinstance(transcript, str) and transcript: - if ( - self._input_audio_buffer_transcript_parts - and self._input_audio_buffer_transcript_parts[-1] == transcript - ): - return - self._input_audio_buffer_transcript_parts.append(transcript) - - def _consume_input_transcript_hint(self) -> str: - transcript = "".join(self._input_audio_buffer_transcript_parts).strip() - self._input_audio_buffer_transcript_parts.clear() - return transcript - - @staticmethod - def _input_transcript_from_item(item: dict[str, object]) -> str: - content = item.get("content") - if not isinstance(content, list): - return "" - parts: list[str] = [] - for part in content: - if not isinstance(part, dict): - continue - for key in ("transcript", "text"): - value = part.get(key) - if isinstance(value, str) and value: - if part.get("type") in {"input_audio", "audio", "audio_transcript", "transcript"}: - parts.append(value) - break - return "".join(parts).strip() - - @staticmethod - def _input_audio_transcription_completed_event( - item_id: str, - item: dict[str, object], - ) -> dict[str, object] | None: - content = item.get("content") - if not isinstance(content, list): - return None - transcript_parts: list[str] = [] - for index, part in enumerate(content): - if not isinstance(part, dict): - continue - if part.get("type") != "input_audio": - continue - transcript = part.get("transcript") - if isinstance(transcript, str) and transcript: - transcript_parts.append(transcript) - return { - "type": "conversation.item.input_audio_transcription.completed", - "item_id": item_id, - "content_index": index, - "transcript": "".join(transcript_parts), - } - return None - - async def _emit_input_speech_started(self, event: dict[str, object]) -> None: - if self._input_speech_started: - return - self._input_speech_started = True - if self._active_input_item_id is None: - self._active_input_item_id = f"item_{uuid4().hex}" - audio_start_ms = event.get("audio_start_ms", 0) - await self._send_realtime_payload( - { - "type": "input_audio_buffer.speech_started", - "audio_start_ms": int(audio_start_ms) if isinstance(audio_start_ms, int | float) else 0, - "item_id": self._active_input_item_id, - } - ) - - async def _emit_input_speech_stopped(self, event: dict[str, object], *, item_id: str) -> None: - if not self._input_speech_started: - return - self._input_speech_started = False - audio_end_ms = event.get("audio_end_ms", event.get("audio_ms", 0)) - await self._send_realtime_payload( - { - "type": "input_audio_buffer.speech_stopped", - "audio_end_ms": int(audio_end_ms) if isinstance(audio_end_ms, int | float) else 0, - "item_id": item_id, - } - ) + def command_error_payload(exc: DuplexCommandError) -> dict[str, object]: + return error_event(exc.code, str(exc), event_id=exc.event_id).to_realtime() diff --git a/vllm_omni/entrypoints/duplex/realtime_output.py b/vllm_omni/entrypoints/duplex/realtime_output.py deleted file mode 100644 index fcbdf87d7..000000000 --- a/vllm_omni/entrypoints/duplex/realtime_output.py +++ /dev/null @@ -1,1007 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -from typing import Any -from uuid import uuid4 - -from vllm_omni.entrypoints.duplex.audio import convert_output_audio -from vllm_omni.entrypoints.duplex.realtime_state import ( - _RealtimeResponseState, -) - - -class RealtimeOutputProjector: - """Project internal duplex events onto the OpenAI Realtime schema.""" - - def _from_duplex_event(self, event: dict[str, Any]) -> list[dict[str, object]]: - event_type = event.get("type") - if event_type == "session.created": - session = self._realtime_session_payload(event.get("session")) - created: dict[str, object] = {"type": "session.created", "session": session} - for key in ( - "incarnation", - "attachment_generation", - "resume_token", - ): - if key in event: - created[key] = event[key] - payloads: list[dict[str, object]] = [created] - if self._initial_session_update: - payloads.append({"type": "session.updated", "session": session}) - self._initial_session_update = False - self._hold_realtime_output_until_session_created = False - if self._held_realtime_payloads: - payloads.extend(self._held_realtime_payloads) - self._held_realtime_payloads = [] - return payloads - if event_type == "session.updated": - session = self._realtime_session_payload(event.get("session")) - return [{"type": "session.updated", "session": session}] - if event_type == "input_audio_buffer.speech_started": - item_id = event.get("item_id") - if isinstance(item_id, str): - self._active_input_item_id = item_id - self._input_speech_started = True - self._input_audio_buffer_has_audio = True - return [ - { - "type": "input_audio_buffer.speech_started", - "audio_start_ms": max(0, int(event.get("audio_start_ms", 0) or 0)), - "item_id": item_id, - } - ] - if event_type == "input_audio_buffer.speech_stopped": - self._input_speech_started = False - return [ - { - "type": "input_audio_buffer.speech_stopped", - "audio_end_ms": max(0, int(event.get("audio_end_ms", 0) or 0)), - "item_id": event.get("item_id") or self._active_input_item_id, - } - ] - if event_type in { - "session.resumed", - "session.heartbeat_ack", - "session.replaced", - "session.expired", - "session.resync_required", - }: - if event_type == "session.resumed": - self._hold_realtime_output_until_session_created = False - return [dict(event)] - if event_type == "response.created": - response_id = event.get("response_id") - if isinstance(response_id, str) and response_id: - self._active_response_id = response_id - self._last_response_id = response_id - item_id = self._response_item_id(response_id) - modalities = event.get("modalities") - has_audio_modality = not isinstance(modalities, list) or "audio" in modalities - item = { - "id": item_id, - "object": "realtime.item", - "type": "message", - "role": "assistant", - "status": "in_progress", - "content": [], - } - self._conversation_items[item_id] = item - payloads = [ - self._response_created_event(event), - *self._conversation_item_added_events(item), - *self._response_output_item_added_events(response_id=response_id, item=item), - ] - if has_audio_modality: - payloads.extend(self._ensure_response_audio_part_added(response_id)) - return payloads - if event_type == "response.listen": - listen_response: dict[str, object] = { - "object": "realtime.response", - "status": "listening", - "metadata": event, - } - listen_payload: dict[str, object] = { - "type": "response.listen", - "session_id": event.get("session_id"), - "epoch": event.get("epoch"), - "response": listen_response, - } - listen_response_id = event.get("response_id") - if isinstance(listen_response_id, str) and listen_response_id: - # Keep the response identity visible on the Realtime dialect: - # a listen that terminates a precreated response is matched to - # it by this id (clients demultiplex on response_id). - listen_response["id"] = listen_response_id - listen_payload["response_id"] = listen_response_id - return [listen_payload] - if event_type == "response.speak": - response_id = event.get("response_id") - state = self._response_state(response_id) - if state is not None: - if state.speak_emitted: - return [] - state.speak_emitted = True - return [ - { - "type": "response.speak", - "response_id": response_id, - "item_id": self._response_item_id(response_id), - "output_index": 0, - "content_index": 0, - "metadata": self._response_speak_metadata(event), - } - ] - if event_type == "overlap.decision": - return [ - { - "type": "overlap.decision", - "session_id": event.get("session_id"), - "epoch": event.get("epoch"), - "policy": event.get("policy"), - "action": event.get("action"), - "reason": event.get("reason"), - "metadata": event, - } - ] - if event_type == "response.output_audio.delta": - response_id = event.get("response_id") - audio = event.get("audio", "") - payloads: list[dict[str, object]] = [] - if isinstance(audio, str) and audio: - payloads.extend(self._ensure_response_audio_part_added(response_id)) - payloads.extend(self._realtime_audio_delta_events(event, response_id, audio)) - self._refresh_in_progress_response_item(response_id) - text = event.get("text") - has_text = isinstance(text, str) and bool(text) - if has_text: - self._append_response_transcript(response_id, text) - self._refresh_in_progress_response_item(response_id) - # Keep the audio.delta + transcript.delta pair invariant even for - # text-less units (deduplicated continuations, turn-end flush): - # clients that treat the pair as unit-complete would otherwise - # wait on a transcript that never comes. - emit_transcript = has_text or (isinstance(audio, str) and bool(audio)) - if emit_transcript: - if not has_text: - text = "" - payloads.append( - { - "type": "response.output_audio_transcript.delta", - "response_id": response_id, - "item_id": self._response_item_id(response_id), - "output_index": 0, - "content_index": 0, - "delta": text, - } - ) - if event.get("end_of_turn") is True: - payloads.extend(self._realtime_audio_done_events(event, response_id)) - payloads.extend( - self._realtime_response_terminal_events( - event, - response_id, - status="completed", - status_details={ - "type": "completed", - "reason": event.get("finish_reason") or "stop", - }, - ) - ) - return payloads - if event_type == "response.text.delta": - response_id = event.get("response_id") - text = event.get("delta", "") - state = self._response_state(response_id) - if state is not None and isinstance(text, str) and text: - state.text_parts.append(text) - self._refresh_in_progress_response_item(response_id) - payloads = self._ensure_response_text_part_added(response_id) - payloads.append( - { - "type": "response.output_text.delta", - "response_id": response_id, - "item_id": self._response_item_id(response_id), - "output_index": 0, - "content_index": 1 if state is not None and state.audio_part_added else 0, - "delta": text, - } - ) - return payloads - if event_type == "response.done": - response_id = event.get("response_id") - status = event.get("status") if isinstance(event.get("status"), str) else "completed" - status_details = event.get("status_details") if isinstance(event.get("status_details"), dict) else None - return [ - *self._realtime_audio_done_events(event, response_id), - *self._realtime_response_terminal_events( - event, - response_id, - status=status, - status_details=status_details, - ), - ] - if event_type == "input.committed": - event_item_id = event.get("item_id") - # Acknowledgements carry their own identity; untracked silence - # commits must not consume another turn's item ID. - item_id = event_item_id if isinstance(event_item_id, str) and event_item_id else f"item_{uuid4().hex}" - # The reader may already have accepted audio for the next turn. - if self._active_input_item_id == item_id: - self._reset_realtime_input_buffer_state() - item = self._conversation_items.get(item_id) - payloads = [self._input_audio_buffer_committed_event(item_id=item_id)] - if item is None: - message = event.get("message") - no_response = event.get("no_response") is True - is_speech = event.get("is_speech") - item = { - "id": item_id, - "object": "realtime.item", - "type": "message", - "role": "user", - "status": "completed", - "content": ( - [{"type": "input_audio", "transcript": "", "is_speech": False}] - if no_response and is_speech is False - else self._user_item_content_from_duplex_message(message) - ), - } - self._conversation_items[item_id] = item - payloads.extend(self._conversation_item_added_events(item)) - item["status"] = "completed" - transcription_event = self._input_audio_transcription_completed_event(item_id, item) - if transcription_event is not None: - payloads.append(transcription_event) - payloads.append(self._conversation_item_done_event(item)) - return payloads - if event_type == "input.cancelled": - self._reset_realtime_input_buffer_state() - return [{"type": "input_audio_buffer.cleared"}] - if event_type == "input_audio_buffer.cleared": - self._reset_realtime_input_buffer_state() - return [{"type": "input_audio_buffer.cleared"}] - if event_type == "audio.cancelled": - response_id = event.get("response_id") - payloads: list[dict[str, object]] = [] - if event.get("reason") == "output_audio_buffer_clear": - if not isinstance(response_id, str) or not response_id: - response_id = self._active_response_id or self._last_response_id - payloads.append( - { - "type": "output_audio_buffer.cleared", - "response_id": response_id, - } - ) - if not isinstance(response_id, str) or not response_id: - return payloads - elif not isinstance(response_id, str) or not response_id: - response_id = self._active_response_id - if not isinstance(response_id, str) or not response_id: - return payloads - if self._response_is_done(response_id): - return payloads - committed_ms = event.get("committed_ms") - if isinstance(committed_ms, int | float): - item_id = self._response_item_id(response_id) - committed_audio_ms = max(0, int(committed_ms)) - self._item_truncation_cursors[item_id] = (0, committed_audio_ms) - item = self._conversation_items.get(item_id) - if item is not None: - self._truncate_realtime_item_content( - item, - content_index=0, - audio_end_ms=committed_audio_ms, - ) - # Align server history with acknowledged playback; only explicit - # client truncation emits ``conversation.item.truncated``. - payloads.extend(self._realtime_audio_done_events(event, response_id)) - payloads.extend( - self._realtime_response_terminal_events( - event, - response_id, - status="cancelled", - status_details={ - "type": "cancelled", - "reason": event.get("reason") or "client_cancelled", - }, - ) - ) - if isinstance(response_id, str) and response_id == self._active_response_id: - self._active_response_id = None - return payloads - if event_type == "playback.acknowledged": - return [{"type": "playback.acknowledged", "event": event}] - if event_type == "conversation.item.created": - item = event.get("item") - if isinstance(item, dict) and isinstance(item.get("id"), str): - item_id = str(item["id"]) - already_known = item_id in self._conversation_items - self._conversation_items[item_id] = item - if already_known: - if item.get("status") == "completed": - return [self._conversation_item_done_event(item)] - return [] - payloads = self._conversation_item_added_events(item) - if item.get("status") == "completed": - payloads.append(self._conversation_item_done_event(item)) - return payloads - return [{"type": "conversation.item.created", "item": item, "event": event}] - if event_type == "conversation.item.deleted": - item_id = event.get("item_id") - if isinstance(item_id, str): - self._remove_conversation_item(item_id) - return [ - { - "type": "conversation.item.deleted", - "item_id": item_id, - "event": event, - } - ] - if event_type == "conversation.item.truncated": - item_id = event.get("item_id") - audio_end_ms = event.get("audio_end_ms") - content_index = event.get("content_index", 0) - if isinstance(item_id, str): - item = self._conversation_items.get(item_id) - if item is not None: - self._truncate_realtime_item_content( - item, - content_index=int(content_index) if isinstance(content_index, int | float) else 0, - audio_end_ms=int(audio_end_ms) if isinstance(audio_end_ms, int | float) else 0, - ) - return [ - { - "type": "conversation.item.truncated", - "item_id": item_id, - "content_index": content_index, - "audio_end_ms": audio_end_ms, - "event": event, - } - ] - if event_type == "response.output_item.done": - response_id = event.get("response_id") - return self._realtime_response_terminal_events( - event, - response_id, - status="completed", - status_details={ - "type": "completed", - "reason": event.get("finish_reason") or "stop", - }, - ) - if event_type == "function_call.done": - call_id = event.get("call_id") - name = event.get("name") - arguments = event.get("arguments", "") - if not isinstance(call_id, str) or not call_id or not isinstance(name, str) or not name: - return [{"type": "duplex.function_call.done", "event": event}] - if not isinstance(arguments, str): - arguments = str(arguments) - response_id = f"resp_{uuid4().hex}" - item = { - "id": f"item_{uuid4().hex}", - "object": "realtime.item", - "type": "function_call", - "status": "completed", - "name": name, - "call_id": call_id, - "arguments": arguments, - } - self._conversation_items[str(item["id"])] = item - response = { - "id": response_id, - "object": "realtime.response", - "status": "completed", - "status_details": {"type": "completed", "reason": "completed"}, - "output": [dict(item)], - "metadata": {}, - } - return [ - {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}}, - *self._conversation_item_added_events(item), - { - "type": "response.output_item.added", - "response_id": response_id, - "output_index": 0, - "item": dict(item), - }, - { - "type": "response.function_call_arguments.delta", - "response_id": response_id, - "item_id": item["id"], - "output_index": 0, - "call_id": call_id, - "delta": arguments, - }, - { - "type": "response.function_call_arguments.done", - "response_id": response_id, - "item_id": item["id"], - "output_index": 0, - "call_id": call_id, - "arguments": arguments, - }, - { - "type": "response.output_item.done", - "response_id": response_id, - "output_index": 0, - "item": dict(item), - }, - self._conversation_item_done_event(item), - {"type": "response.done", "response": response}, - ] - if event_type == "error": - raw_error = event.get("error") - if isinstance(raw_error, dict): - return [event] - message = str(raw_error or event.get("message") or "Duplex runtime error") - code = str(event.get("code") or "duplex_error") - return [ - self._realtime_error_payload( - code, - message, - event_id=event.get("realtime_event_id") or event.get("event_id"), - param=event.get("param"), - ) - ] - if event_type == "session.closed": - return [{"type": "session.closed", "event": event}] - return [{"type": f"duplex.{event_type}", "event": event}] - - @staticmethod - def _user_item_content_from_duplex_message(message: object) -> list[dict[str, object]]: - if not isinstance(message, dict): - return [{"type": "input_audio", "transcript": ""}] - message_transcript = message.get("transcript") - content = message.get("content") - if isinstance(content, str): - return [{"type": "input_text", "text": content}] - if not isinstance(content, list): - return [ - { - "type": "input_audio", - "transcript": message_transcript if isinstance(message_transcript, str) else "", - } - ] - parts: list[dict[str, object]] = [] - for part in content: - if not isinstance(part, dict): - continue - part_type = part.get("type") - if part_type in {"text", "input_text"} and isinstance(part.get("text"), str): - parts.append({"type": "input_text", "text": part["text"]}) - continue - if part_type == "audio_url": - transcript = ( - part.get("transcript") - if isinstance(part.get("transcript"), str) - else message_transcript - if isinstance(message_transcript, str) - else "" - ) - parts.append({"type": "input_audio", "transcript": transcript}) - continue - if part_type in {"audio", "input_audio"}: - transcript = part.get("transcript") if isinstance(part.get("transcript"), str) else "" - parts.append({"type": "input_audio", "transcript": transcript}) - return parts or [{"type": "input_audio", "transcript": ""}] - - def _realtime_session_payload(self, session: object) -> object: - if not isinstance(session, dict): - return session - payload = dict(session) - payload.setdefault("object", "realtime.session") - payload.setdefault("type", "realtime") - payload.setdefault("id", payload.get("id") or self._default_session_id) - payload.setdefault("model", payload.get("model") or self._default_model) - payload.setdefault("input_audio_format", self._input_audio_format) - payload.setdefault("output_audio_format", self._output_audio_format) - payload.setdefault("modalities", payload.get("modalities") or ["text", "audio"]) - payload.setdefault("output_modalities", payload.get("output_modalities") or payload.get("modalities")) - payload.setdefault("object", "realtime.session") - payload.setdefault( - "audio", - { - "input": { - "format": self._realtime_audio_format_object( - self._input_audio_format, - sample_rate_hz=self._input_sample_rate_hz, - ), - "sample_rate_hz": self._input_sample_rate_hz, - }, - "output": { - "format": self._realtime_audio_format_object( - self._output_audio_format, - sample_rate_hz=self._output_sample_rate_hz, - ), - }, - }, - ) - if self._turn_detection_configured or "turn_detection" in payload: - turn_detection = payload.get("turn_detection", self._turn_detection) - audio = dict(payload["audio"]) if isinstance(payload.get("audio"), dict) else {} - audio_input = dict(audio["input"]) if isinstance(audio.get("input"), dict) else {} - audio_input["turn_detection"] = turn_detection - audio["input"] = audio_input - payload["audio"] = audio - payload["turn_detection"] = turn_detection - payload.setdefault("input_audio_transcription", payload.get("input_audio_transcription")) - payload.setdefault("tracing", payload.get("tracing")) - return payload - - @staticmethod - def _realtime_audio_format_object(fmt: object, *, sample_rate_hz: int | None = None) -> dict[str, object]: - if isinstance(fmt, str) and fmt.lower() in {"pcm16", "pcm_s16le", "s16le", "pcm"}: - payload: dict[str, object] = {"type": "audio/pcm"} - elif isinstance(fmt, str) and fmt.lower() == "pcm_f32le": - payload = {"type": "audio/pcm_f32le"} - elif isinstance(fmt, str) and fmt.lower() == "g711_ulaw": - payload = {"type": "audio/g711_ulaw"} - elif isinstance(fmt, str) and fmt.lower() == "g711_alaw": - payload = {"type": "audio/g711_alaw"} - else: - payload = {"type": "audio/wav"} - if sample_rate_hz is not None: - payload["rate"] = int(sample_rate_hz) - return payload - - def _response_state( - self, - response_id: object, - *, - event: dict[str, Any] | None = None, - create: bool = True, - ) -> _RealtimeResponseState | None: - if isinstance(response_id, str) and response_id: - key: str | int = response_id - item_id = f"item_{response_id}" - elif event is not None: - key = id(event) - item_id = f"item_{uuid4().hex}" - else: - return None - state = self._response_states.get(key) - if state is None and create: - state = _RealtimeResponseState(item_id=item_id) - self._response_states[key] = state - return state - - def _response_is_done(self, response_id: object) -> bool: - state = self._response_state(response_id, create=False) - return state is not None and state.done_emitted - - def _response_item_id(self, response_id: object) -> str: - state = self._response_state(response_id) - if state is not None: - return state.item_id - return f"item_{uuid4().hex}" - - def _response_done_output_item( - self, - response_id: object, - *, - status: str, - ) -> dict[str, object]: - item_id = self._response_item_id(response_id) - state = self._response_state(response_id) - transcript = state.transcript if state is not None else "" - text = state.text if state is not None else "" - audio_duration_ms = state.audio_duration_ms if state is not None else None - audio_text_marks = state.audio_text_marks if state is not None else None - content: list[dict[str, object]] = [] - if (state is not None and state.audio_part_added) or transcript or audio_duration_ms is not None: - content.append( - self._response_item_content_part( - transcript=transcript, - audio_duration_ms=audio_duration_ms, - audio_text_marks=audio_text_marks, - ) - ) - if text: - content.append(self._response_item_text_content_part(text=text)) - item = { - "id": item_id, - "object": "realtime.item", - "type": "message", - "role": "assistant", - "status": status, - "content": content, - } - self._apply_pending_item_truncation(item) - return item - - def _apply_pending_item_truncation(self, item: dict[str, object]) -> None: - item_id = item.get("id") - if not isinstance(item_id, str) or not item_id: - return - cursor = self._item_truncation_cursors.get(item_id) - if cursor is None: - return - content_index, audio_end_ms = cursor - self._truncate_realtime_item_content( - item, - content_index=content_index, - audio_end_ms=audio_end_ms, - ) - - def _refresh_in_progress_response_item(self, response_id: object) -> None: - if not isinstance(response_id, str) or not response_id: - return - item_id = self._response_item_id(response_id) - item = self._conversation_items.get(item_id) - if not isinstance(item, dict): - return - content = item.get("content") - if not isinstance(content, list): - content = [] - item["content"] = content - - state = self._response_state(response_id) - if state is None: - return - transcript = state.transcript - audio_duration_ms = state.audio_duration_ms - audio_text_marks = state.audio_text_marks - has_audio = state.audio_part_added or bool(transcript) or audio_duration_ms is not None - if has_audio: - audio_part = self._response_item_content_part( - transcript=transcript, - audio_duration_ms=audio_duration_ms, - audio_text_marks=audio_text_marks, - ) - if content and isinstance(content[0], dict) and content[0].get("type") in {"audio", "output_audio"}: - content[0] = audio_part - else: - content.insert(0, audio_part) - - text = state.text - if text: - text_index = ( - 1 - if content - and isinstance(content[0], dict) - and content[0].get("type") - in { - "audio", - "output_audio", - } - else 0 - ) - text_part = self._response_item_text_content_part(text=text) - if ( - len(content) > text_index - and isinstance(content[text_index], dict) - and content[text_index].get("type") in {"text", "output_text"} - ): - content[text_index] = text_part - else: - content.insert(text_index, text_part) - - self._apply_pending_item_truncation(item) - - def _append_response_transcript(self, response_id: object, text: str) -> None: - state = self._response_state(response_id) - if state is None or not text: - return - state.transcript_parts.append(text) - - def _ensure_response_text_part_added(self, response_id: object) -> list[dict[str, object]]: - state = self._response_state(response_id) - if state is None: - return [] - if state.text_part_added: - return [] - state.text_part_added = True - content_index = 1 if state.audio_part_added else 0 - return [ - { - "type": "response.content_part.added", - "response_id": response_id, - "item_id": self._response_item_id(response_id), - "output_index": 0, - "content_index": content_index, - "part": self._response_text_content_part(), - } - ] - - def _ensure_response_audio_part_added(self, response_id: object) -> list[dict[str, object]]: - state = self._response_state(response_id) - if state is None: - return [] - if state.audio_part_added: - return [] - state.audio_part_added = True - return [ - { - "type": "response.content_part.added", - "response_id": response_id, - "item_id": self._response_item_id(response_id), - "output_index": 0, - "content_index": 0, - "part": self._response_content_part(), - } - ] - - def _remember_response_audio_metadata(self, response_id: object, event: dict[str, Any]) -> None: - state = self._response_state(response_id) - if state is None: - return - duration = event.get("audio_duration_ms") - playback = event.get("playback") - if not isinstance(duration, int | float) and isinstance(playback, dict): - duration = playback.get("sent_ms") or playback.get("generated_ms") - if isinstance(duration, int | float): - state.audio_duration_ms = max(state.audio_duration_ms or 0, int(duration)) - marks = event.get("audio_text_marks") - if not isinstance(marks, list): - return - clean_marks: list[dict[str, int]] = [] - for mark in marks: - if not isinstance(mark, dict): - continue - text_chars = mark.get("text_chars") - audio_end_ms = mark.get("audio_end_ms", mark.get("audio_ms")) - if not isinstance(text_chars, int | float) or not isinstance(audio_end_ms, int | float): - continue - clean_marks.append( - { - "text_chars": max(0, int(text_chars)), - "audio_end_ms": max(0, int(audio_end_ms)), - } - ) - if clean_marks: - merged = list(state.audio_text_marks) - merged.extend(clean_marks) - deduped: dict[tuple[int, int], dict[str, int]] = {} - for mark in merged: - deduped[(int(mark["audio_end_ms"]), int(mark["text_chars"]))] = mark - state.audio_text_marks = sorted( - deduped.values(), - key=lambda mark: (mark["audio_end_ms"], mark["text_chars"]), - ) - - def _response_created_event(self, event: dict[str, Any]) -> dict[str, object]: - response_id = event.get("response_id") - metadata = event.get("metadata") - if not isinstance(metadata, dict): - metadata = {} - metadata = {**metadata, "duplex_event": event} - return { - "type": "response.created", - "response_id": response_id, - "response": { - "id": response_id, - "object": "realtime.response", - "status": "in_progress", - "status_details": None, - "output": [], - "modalities": event.get("modalities") or ["audio", "text"], - "metadata": metadata, - }, - } - - def _realtime_audio_delta_events( - self, - event: dict[str, Any], - response_id: object, - audio: str, - ) -> list[dict[str, object]]: - item_id = self._response_item_id(response_id) - fmt, format_rate = self._parse_realtime_audio_format(event.get("format", "wav")) - source_fmt = self._realtime_output_format(fmt) - source_sample_rate_hz = event.get("sample_rate_hz") or format_rate - target_sample_rate_hz = ( - self._output_sample_rate_hz if self._output_audio_format in {"g711_ulaw", "g711_alaw"} else None - ) - audio, fmt, converted_sample_rate_hz = convert_output_audio( - audio, - source_fmt=source_fmt, - target_fmt=self._output_audio_format, - source_sample_rate_hz=( - int(source_sample_rate_hz) if isinstance(source_sample_rate_hz, int | float) else None - ), - target_sample_rate_hz=target_sample_rate_hz, - ) - sample_rate_hz = ( - converted_sample_rate_hz - if fmt in {"g711_ulaw", "g711_alaw"} - else event.get("sample_rate_hz") or format_rate or self._output_sample_rate_hz - ) - metadata: dict[str, object] = {} - for key in ("session_id", "epoch", "model_speak", "end_of_turn", "playback", "vllm_omni"): - if key in event: - metadata[key] = event[key] - duration_ms = event.get("audio_duration_ms") - if isinstance(duration_ms, int | float): - metadata["audio_duration_ms"] = int(duration_ms) - marks = event.get("audio_text_marks") - if isinstance(marks, list): - metadata["audio_text_marks"] = marks - state = self._response_state(response_id) - if state is not None: - state.audio_delta_emitted = True - self._remember_response_audio_metadata(response_id, event) - payloads: list[dict[str, object]] = [] - if state is not None and not state.speak_emitted and metadata.get("model_speak") is True: - state.speak_emitted = True - payloads.insert( - 0, - { - "type": "response.speak", - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "content_index": 0, - "metadata": self._response_speak_metadata(event), - }, - ) - payloads.append( - { - "type": "response.output_audio.delta", - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "content_index": 0, - "delta": audio, - "format": fmt, - **({"sample_rate_hz": int(sample_rate_hz)} if isinstance(sample_rate_hz, int | float) else {}), - **({"metadata": metadata} if metadata else {}), - } - ) - return payloads - - @staticmethod - def _response_speak_metadata(event: dict[str, Any]) -> dict[str, object]: - return {key: event[key] for key in ("session_id", "epoch", "model_speak") if key in event} - - def _realtime_audio_done_events( - self, - event: dict[str, Any], - response_id: object, - ) -> list[dict[str, object]]: - item_id = self._response_item_id(response_id) - state = self._response_state(response_id, event=event) - transcript = state.transcript if state is not None else "" - payloads: list[dict[str, object]] = [] - if isinstance(response_id, str) and state is not None and not state.audio_delta_emitted and not transcript: - return payloads - if state is None or not state.audio_done_emitted: - if state is not None: - state.audio_done_emitted = True - payloads.append( - { - "type": "response.output_audio.done", - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "content_index": 0, - } - ) - if transcript: - payloads.append( - { - "type": "response.output_audio_transcript.done", - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "content_index": 0, - "transcript": transcript, - } - ) - return payloads - - def _realtime_response_terminal_events( - self, - event: dict[str, Any], - response_id: object, - *, - status: str = "completed", - status_details: dict[str, object] | None = None, - ) -> list[dict[str, object]]: - item_id = self._response_item_id(response_id) - state = self._response_state(response_id, event=event) - transcript = state.transcript if state is not None else "" - payloads: list[dict[str, object]] = [] - if state is not None and state.audio_part_added and not state.audio_part_done: - state.audio_part_done = True - payloads.append( - { - "type": "response.content_part.done", - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "content_index": 0, - "part": self._response_content_part(transcript=transcript), - } - ) - if state is not None and state.text_parts and not state.output_text_done: - state.output_text_done = True - content_index = 1 if state.audio_part_added else 0 - payloads.append( - { - "type": "response.output_text.done", - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "content_index": content_index, - "text": state.text, - } - ) - if state is not None and state.text_parts and not state.text_part_done: - state.text_part_done = True - content_index = 1 if state.audio_part_added else 0 - payloads.append( - { - "type": "response.content_part.done", - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "content_index": content_index, - "part": self._response_text_content_part(text=state.text), - } - ) - if state is None or not state.output_item_done: - if state is not None: - state.output_item_done = True - item = self._response_done_output_item(response_id, status=status) - payloads.append( - { - "type": "response.output_item.done", - "response_id": response_id, - "output_index": 0, - "item": item, - } - ) - if state is None or not state.conversation_item_done: - if state is not None: - state.conversation_item_done = True - payloads.append(self._conversation_item_done_event(item)) - done_event = self._realtime_response_done_event( - {**event, "response_id": response_id}, - state=state, - status=status, - status_details=status_details, - ) - if done_event is not None: - payloads.append(done_event) - payloads.append(self._rate_limits_updated_event()) - if isinstance(response_id, str) and response_id == self._active_response_id: - self._active_response_id = None - return payloads - - def _realtime_response_done_event( - self, - event: dict[str, Any], - *, - state: _RealtimeResponseState | None = None, - status: str = "completed", - status_details: dict[str, object] | None = None, - ) -> dict[str, object] | None: - response_id = event.get("response_id") - if state is None: - state = self._response_state(response_id, event=event) - if state is not None: - if state.done_emitted: - return None - state.done_emitted = True - return { - "type": "response.done", - "response_id": response_id, - "response": { - "id": response_id, - "object": "realtime.response", - "status": status, - "status_details": status_details, - "output": [self._response_done_output_item(response_id, status=status)], - "metadata": event, - }, - } - - @staticmethod - def _rate_limits_updated_event() -> dict[str, object]: - # vLLM-Omni does not yet expose a Realtime-specific quota budget. Emit - # the terminal event with an empty list so clients that sequence on - # rate_limits.updated do not need a private vLLM branch. - return {"type": "rate_limits.updated", "rate_limits": []} diff --git a/vllm_omni/entrypoints/duplex/realtime_session.py b/vllm_omni/entrypoints/duplex/realtime_session.py deleted file mode 100644 index 122a7ab52..000000000 --- a/vllm_omni/entrypoints/duplex/realtime_session.py +++ /dev/null @@ -1,147 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import json -from typing import Any -from uuid import uuid4 - -from fastapi import WebSocket - -from vllm_omni.entrypoints.duplex import realtime_state as _state -from vllm_omni.entrypoints.duplex.realtime_input import ( - RealtimeInputTranslator, -) -from vllm_omni.entrypoints.duplex.realtime_output import ( - RealtimeOutputProjector, -) - -REALTIME_ERROR_TYPES_BY_CODE = _state.REALTIME_ERROR_TYPES_BY_CODE -REALTIME_INPUT_AUDIO_FORMATS = _state.REALTIME_INPUT_AUDIO_FORMATS -REALTIME_OUTPUT_AUDIO_FORMATS = _state.REALTIME_OUTPUT_AUDIO_FORMATS -RealtimeSessionState = _state.RealtimeSessionState -_RealtimeResponseState = _state._RealtimeResponseState - - -class NativeRealtimeSessionProtocol( - RealtimeInputTranslator, - RealtimeOutputProjector, - _state.RealtimeStateOwner, -): - """Compatibility facade for one native Realtime protocol session. - - Input translation, output projection, and mutable state have separate - owners. The facade retains the established import path and private helper - surface used by serving and protocol contract tests. - """ - - def __init__(self, query_params: Any) -> None: - self._state = RealtimeSessionState.from_query_params(query_params) - - def bind_sender(self, send_realtime_json) -> None: - self._send_realtime_json = send_realtime_json - - def bind_native_input_append(self, enabled: bool) -> None: - """Use the serving-selected input path, independent of interruption policy.""" - self._native_input_append = enabled - - def _default_session_payload(self) -> dict[str, object]: - payload: dict[str, object] = { - "model": self._default_model, - "session_id": self._default_session_id, - } - if self._default_extra_body: - payload["extra_body"] = dict(self._default_extra_body) - return payload - - async def receive_internal_event_text(self, websocket: WebSocket) -> str: - if not self._pending_outbound.empty(): - return json.dumps(await self._pending_outbound.get()) - if not self._opened and not self._resume_only and not self._autostarted_default_session and self._default_model: - self._opened = True - self._autostarted_default_session = True - return json.dumps(self._session_create_from_realtime(self._default_session_payload())) - while True: - if not self._pending_outbound.empty(): - return json.dumps(await self._pending_outbound.get()) - await self.wait_for_realtime_turn_detection_update() - raw = await websocket.receive_text() - try: - event = json.loads(raw) - except json.JSONDecodeError: - return raw - if not isinstance(event, dict): - return raw - if not self._opened and event.get("type") == "session.resume": - self._opened = True - translated = await self._to_duplex_event(event) - if translated is None: - if not self._pending_outbound.empty(): - return json.dumps(await self._pending_outbound.get()) - continue - return json.dumps(translated) - if not self._opened and self._resume_only: - translated = await self._to_duplex_event(event) - if translated is None: - if not self._pending_outbound.empty(): - return json.dumps(await self._pending_outbound.get()) - continue - return json.dumps(translated) - if not self._opened and event.get("type") != "session.update": - self._opened = True - await self._pending_outbound.put(self._session_create_from_realtime(self._default_session_payload())) - translated = await self._to_duplex_event(event) - if translated is not None: - await self._send_realtime_input_ack(event) - await self._pending_outbound.put(translated) - return json.dumps(await self._pending_outbound.get()) - translated = await self._to_duplex_event(event) - if translated is None: - if not self._pending_outbound.empty(): - return json.dumps(await self._pending_outbound.get()) - continue - await self._send_realtime_input_ack(event) - return json.dumps(translated) - - def encode_outbound_event(self, data: dict[str, Any]) -> list[dict[str, object]]: - payloads = self._from_duplex_event(data) - for payload in payloads: - self._attach_event_id(payload) - return payloads - - @staticmethod - def _attach_event_id(payload: dict[str, object]) -> None: - payload.setdefault("event_id", f"event_{uuid4().hex}") - - async def _send_realtime_payload(self, payload: dict[str, object]) -> None: - self._attach_event_id(payload) - if ( - self._hold_realtime_output_until_session_created - and payload.get("type") != "error" - and payload.get("type") != "session.created" - ): - self._held_realtime_payloads.append(payload) - return - if self._send_realtime_json is None: - raise RuntimeError("Native Realtime sender is not bound") - await self._send_realtime_json(payload) - - @staticmethod - def _realtime_error_payload( - code: str, - message: str, - *, - event_id: object | None = None, - param: object | None = None, - ) -> dict[str, object]: - error: dict[str, object] = { - "type": REALTIME_ERROR_TYPES_BY_CODE.get(code, "invalid_request_error"), - "code": code, - "message": message, - } - if isinstance(event_id, str) and event_id: - error["event_id"] = event_id - if isinstance(param, str) and param: - error["param"] = param - return {"type": "error", "error": error} diff --git a/vllm_omni/entrypoints/duplex/realtime_state.py b/vllm_omni/entrypoints/duplex/realtime_state.py deleted file mode 100644 index babb470a5..000000000 --- a/vllm_omni/entrypoints/duplex/realtime_state.py +++ /dev/null @@ -1,245 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -from dataclasses import dataclass, field -from typing import Any, Generic, TypeVar - -REALTIME_INPUT_AUDIO_FORMATS = { - "pcm16", - "pcm_s16le", - "s16le", - "pcm_f32le", - "g711_ulaw", - "g711_alaw", -} -REALTIME_OUTPUT_AUDIO_FORMATS = { - "pcm16", - "pcm_s16le", - "s16le", - "wav", - "pcm", - "g711_ulaw", - "g711_alaw", -} -REALTIME_PCM_DEFAULT_SAMPLE_RATE_HZ = 24_000 -REALTIME_PCM_F32_DEFAULT_SAMPLE_RATE_HZ = 16_000 -REALTIME_G711_DEFAULT_SAMPLE_RATE_HZ = 8_000 - -_RealtimeProtocolConfig = tuple[str, int, str, int | None, float] -_PendingTurnDetectionUpdate = tuple[ - bool, - dict[str, object] | None, - asyncio.Event, - _RealtimeProtocolConfig, -] - - -def realtime_default_sample_rate_hz(fmt: object) -> int | None: - if not isinstance(fmt, str): - return None - normalized = fmt.lower() - if normalized in {"pcm16", "pcm_s16le", "s16le", "pcm", "audio/pcm"}: - return REALTIME_PCM_DEFAULT_SAMPLE_RATE_HZ - if normalized in {"pcm_f32le", "audio/pcm_f32le"}: - return REALTIME_PCM_F32_DEFAULT_SAMPLE_RATE_HZ - if normalized in {"g711_ulaw", "g711_alaw", "audio/g711_ulaw", "audio/g711_alaw"}: - return REALTIME_G711_DEFAULT_SAMPLE_RATE_HZ - return None - - -REALTIME_ERROR_TYPES_BY_CODE = { - "bad_event": "invalid_request_error", - "bad_audio": "invalid_request_error", - "config_timeout": "invalid_request_error", - "invalid_json": "invalid_request_error", - "event_too_large": "invalid_request_error", - "unknown_event": "invalid_request_error", - "internal_error": "server_error", - "runtime_open_failed": "server_error", - "runtime_open_unsupported": "server_error", - "runtime_append_failed": "server_error", - "runtime_append_task_failed": "server_error", - "runtime_signal_failed": "server_error", - "runtime_close_failed": "server_error", - "runtime_abort_failed": "server_error", - "runtime_data_plane_stream_failed": "server_error", - "runtime_data_plane_text_without_audio": "server_error", - "response_error": "server_error", - "chat_error": "server_error", - "server_vad_initialization_failed": "server_error", - "server_vad_inference_failed": "server_error", - "duplex_session_busy": "rate_limit_error", - "resource_exhausted": "rate_limit_error", - "input_backpressure": "rate_limit_error", - "response_already_active": "invalid_request_error", - "response_not_active": "invalid_request_error", - "response_create_without_input": "invalid_request_error", - "input_audio_buffer_empty": "invalid_request_error", - "server_vad_manual_commit_unsupported": "invalid_request_error", - "missing_item_id": "invalid_request_error", - "item_not_found": "invalid_request_error", - "playback_item_mismatch": "invalid_request_error", - "playback_item_not_found": "invalid_request_error", - "playback_ack_too_late": "invalid_request_error", - "unsupported_audio_format": "invalid_request_error", - "unsupported_ref_audio_path": "invalid_request_error", - "ref_audio_required": "invalid_request_error", - "model_update_unsupported": "invalid_request_error", - "voice_update_after_audio_unsupported": "invalid_request_error", - "ref_audio_update_unsupported": "invalid_request_error", - "native_text_append_unsupported": "invalid_request_error", - "runtime_native_stage_role_required": "server_error", - "runtime_native_runner_kv_required": "server_error", -} - - -@dataclass(slots=True) -class _RealtimeResponseState: - item_id: str - transcript_parts: list[str] = field(default_factory=list) - text_parts: list[str] = field(default_factory=list) - audio_duration_ms: int | None = None - audio_text_marks: list[dict[str, int]] = field(default_factory=list) - audio_delta_emitted: bool = False - audio_done_emitted: bool = False - audio_part_added: bool = False - audio_part_done: bool = False - text_part_added: bool = False - text_part_done: bool = False - output_text_done: bool = False - output_item_done: bool = False - conversation_item_done: bool = False - speak_emitted: bool = False - done_emitted: bool = False - - @property - def transcript(self) -> str: - return "".join(self.transcript_parts) - - @property - def text(self) -> str: - return "".join(self.text_parts) - - -@dataclass(slots=True) -class RealtimeSessionState: - """Single mutable authority for one Realtime protocol session.""" - - _opened: bool = False - _autostarted_default_session: bool = False - _resume_only: bool = False - _pending_outbound: asyncio.Queue[dict[str, object]] = field(default_factory=asyncio.Queue) - _held_realtime_payloads: list[dict[str, object]] = field(default_factory=list) - _hold_realtime_output_until_session_created: bool = True - _default_model: object | None = None - _default_session_id: object | None = None - _default_extra_body: dict[str, object] = field(default_factory=dict) - _input_audio_format: str = "pcm16" - _input_sample_rate_hz: int = REALTIME_PCM_DEFAULT_SAMPLE_RATE_HZ - _output_audio_format: str = "pcm16" - _overlap_silence_rms: float = 0.003 - _turn_detection: dict[str, object] | None = None - _native_input_append: bool = False - _pending_turn_detection_update: _PendingTurnDetectionUpdate | None = None - _send_realtime_json: Any = None - _initial_session_update: bool = False - _input_speech_started: bool = False - _response_states: dict[str | int, _RealtimeResponseState] = field(default_factory=dict) - _item_truncation_cursors: dict[str, tuple[int, int]] = field(default_factory=dict) - _active_response_id: str | None = None - _last_response_id: str | None = None - _conversation_items: dict[str, dict[str, object]] = field(default_factory=dict) - _last_conversation_item_id: str | None = None - _output_sample_rate_hz: int | None = None - _active_input_item_id: str | None = None - _input_audio_buffer_has_audio: bool = False - _input_audio_buffer_had_non_speech: bool = False - _input_audio_buffer_transcript_parts: list[str] = field(default_factory=list) - _turn_detection_configured: bool = False - - @classmethod - def from_query_params(cls, query_params: Any) -> RealtimeSessionState: - if hasattr(query_params, "query_params"): - query_params = query_params.query_params - getter = query_params.get if hasattr(query_params, "get") else None - # Canonical query param plus the deprecated model-prefixed alias; - # both seed the canonical extra_body key. - native_duplex = None - if getter is not None: - native_duplex = getter("native_duplex") - if native_duplex is None: - native_duplex = getter("minicpmo45_native_duplex") - resume_only = getter("resume") if getter is not None else None - autostart = getter("autostart") if getter is not None else None - default_extra_body = {} - if str(native_duplex).strip().lower() in {"1", "true", "yes", "on"}: - default_extra_body["native_duplex"] = True - return cls( - _resume_only=( - str(resume_only).strip().lower() in {"1", "true", "yes", "on"} - or str(autostart).strip().lower() in {"0", "false", "no", "off"} - ), - _default_model=getter("model") if getter is not None else None, - _default_session_id=(getter("session_id") if getter is not None else None), - _default_extra_body=default_extra_body, - ) - - -_StateValue = TypeVar("_StateValue") - - -class _RealtimeStateField(Generic[_StateValue]): - """Explicit descriptor binding one protocol attribute to session state.""" - - def __set_name__(self, owner: type, name: str) -> None: - del owner - self._name = name - - def __get__(self, instance: Any, owner: type | None = None) -> _StateValue | _RealtimeStateField[_StateValue]: - del owner - if instance is None: - return self - return getattr(instance._state, self._name) - - def __set__(self, instance: Any, value: _StateValue) -> None: - setattr(instance._state, self._name, value) - - -class RealtimeStateOwner: - """Declare the mutable state surface shared by input/output projectors.""" - - _state: RealtimeSessionState - _opened: bool = _RealtimeStateField() - _autostarted_default_session: bool = _RealtimeStateField() - _resume_only: bool = _RealtimeStateField() - _pending_outbound: asyncio.Queue[dict[str, object]] = _RealtimeStateField() - _held_realtime_payloads: list[dict[str, object]] = _RealtimeStateField() - _hold_realtime_output_until_session_created: bool = _RealtimeStateField() - _default_model: object | None = _RealtimeStateField() - _default_session_id: object | None = _RealtimeStateField() - _default_extra_body: dict[str, object] = _RealtimeStateField() - _input_audio_format: str = _RealtimeStateField() - _input_sample_rate_hz: int = _RealtimeStateField() - _output_audio_format: str = _RealtimeStateField() - _overlap_silence_rms: float = _RealtimeStateField() - _turn_detection: dict[str, object] | None = _RealtimeStateField() - _native_input_append: bool = _RealtimeStateField() - _pending_turn_detection_update: _PendingTurnDetectionUpdate | None = _RealtimeStateField() - _send_realtime_json: Any = _RealtimeStateField() - _initial_session_update: bool = _RealtimeStateField() - _input_speech_started: bool = _RealtimeStateField() - _response_states: dict[str | int, _RealtimeResponseState] = _RealtimeStateField() - _item_truncation_cursors: dict[str, tuple[int, int]] = _RealtimeStateField() - _active_response_id: str | None = _RealtimeStateField() - _last_response_id: str | None = _RealtimeStateField() - _conversation_items: dict[str, dict[str, object]] = _RealtimeStateField() - _last_conversation_item_id: str | None = _RealtimeStateField() - _output_sample_rate_hz: int | None = _RealtimeStateField() - _active_input_item_id: str | None = _RealtimeStateField() - _input_audio_buffer_has_audio: bool = _RealtimeStateField() - _input_audio_buffer_had_non_speech: bool = _RealtimeStateField() - _input_audio_buffer_transcript_parts: list[str] = _RealtimeStateField() - _turn_detection_configured: bool = _RealtimeStateField() diff --git a/vllm_omni/entrypoints/duplex/runtime_adapter.py b/vllm_omni/entrypoints/duplex/runtime_adapter.py deleted file mode 100644 index a873f33a8..000000000 --- a/vllm_omni/entrypoints/duplex/runtime_adapter.py +++ /dev/null @@ -1,331 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -from collections.abc import Awaitable, Callable, Iterable, Mapping, MutableMapping -from dataclasses import dataclass, field -from importlib import import_module -from typing import Any, Protocol - - -class ServingRuntimeConfigError(ValueError): - """A model serving plugin rejected client-visible runtime configuration.""" - - def __init__(self, message: str, *, code: str = "invalid_duplex_runtime_config") -> None: - super().__init__(message) - self.code = code - - -def reject_changed_runtime_value( - new_value: object, - current_value: object, - *, - message: str, - code: str, - error_cls: type[ServingRuntimeConfigError] = ServingRuntimeConfigError, -) -> None: - if new_value != current_value: - raise error_cls(message, code=code) - - -class PcmAppendReservation(Protocol): - operation_id: str - payload: dict[str, object] | None - - @property - def active(self) -> bool: ... - - @property - def byte_count(self) -> int: ... - - def commit(self) -> None: ... - - def rollback(self) -> None: ... - - -class PcmAppendBuffer(Protocol): - @property - def pending_byte_count(self) -> int: ... - - def clear(self) -> None: ... - - def clear_force_listen(self) -> None: ... - - def has_pending(self) -> bool: ... - - def has_reserved(self) -> bool: ... - - def prepare_append( - self, - payload: dict[str, object], - *, - operation_id: str, - chunk_period_ms: int, - allow_emit: bool, - ) -> PcmAppendReservation | None: ... - - def prepare_commit( - self, - *, - operation_id: str, - chunk_period_ms: int, - ) -> PcmAppendReservation: ... - - def flush(self, *, chunk_period_ms: int) -> dict[str, object] | None: ... - - -class ServingRuntimeSessionState(Protocol): - audio_buffer: PcmAppendBuffer - input_since_commit: bool - speech_since_commit: bool - native_context_locked: bool - committed_audio_payload: dict[str, object] | None - committed_audio_operation_id: str | None - committed_audio_reserved_bytes: int - deferred_response_create: bool - deferred_precreate_response: bool - data_plane_task: asyncio.Task[None] | None - data_plane_restart_requested: bool - continuation_owner_id: str | None - continuation_units: int - pending_silence_task: asyncio.Task[bool] | None - pending_silence_owner_id: str | None - silence_continuation_scheduler: Callable[..., Awaitable[bool]] | None - - def retain_committed_audio( - self, - payload: dict[str, object], - *, - operation_id: str | None, - reserved_bytes: int = 0, - ) -> None: ... - - def clear_committed_audio(self) -> int: ... - - def clear_continuation(self) -> None: ... - - -class _TurnBasedPcmAppendBuffer: - """Empty native-audio buffer for chat-fallback sessions. - - Turn-based sessions keep input audio in ``DuplexSession`` and never use - the model-native append path. The runner still shares lifecycle and - cancellation code with native sessions, so it needs an inert buffer for - those common operations. - """ - - @property - def pending_byte_count(self) -> int: - return 0 - - def clear(self) -> None: - return None - - def clear_force_listen(self) -> None: - return None - - def has_pending(self) -> bool: - return False - - def has_reserved(self) -> bool: - return False - - def prepare_append( - self, - payload: dict[str, object], - *, - operation_id: str, - chunk_period_ms: int, - allow_emit: bool, - ) -> PcmAppendReservation | None: - del payload, operation_id, chunk_period_ms, allow_emit - raise RuntimeError("Model-native audio append is not enabled for this session") - - def prepare_commit( - self, - *, - operation_id: str, - chunk_period_ms: int, - ) -> PcmAppendReservation: - del operation_id, chunk_period_ms - raise RuntimeError("Model-native audio append is not enabled for this session") - - def flush(self, *, chunk_period_ms: int) -> dict[str, object] | None: - del chunk_period_ms - return None - - -@dataclass(slots=True) -class TurnBasedServingSessionState: - """Runner-local state for sessions served through chat fallback.""" - - audio_buffer: PcmAppendBuffer = field(default_factory=_TurnBasedPcmAppendBuffer) - input_since_commit: bool = False - speech_since_commit: bool = False - native_context_locked: bool = False - committed_audio_payload: dict[str, object] | None = None - committed_audio_operation_id: str | None = None - committed_audio_reserved_bytes: int = 0 - deferred_response_create: bool = False - deferred_precreate_response: bool = False - data_plane_task: asyncio.Task[None] | None = None - data_plane_restart_requested: bool = False - continuation_owner_id: str | None = None - continuation_units: int = 0 - pending_silence_task: asyncio.Task[bool] | None = None - pending_silence_owner_id: str | None = None - silence_continuation_scheduler: Callable[..., Awaitable[bool]] | None = None - - def retain_committed_audio( - self, - payload: dict[str, object], - *, - operation_id: str | None, - reserved_bytes: int = 0, - ) -> None: - self.committed_audio_payload = payload - self.committed_audio_operation_id = operation_id - self.committed_audio_reserved_bytes += max(0, int(reserved_bytes)) - - def clear_committed_audio(self) -> int: - reserved_bytes = self.committed_audio_reserved_bytes - self.committed_audio_payload = None - self.committed_audio_operation_id = None - self.committed_audio_reserved_bytes = 0 - self.deferred_response_create = False - self.deferred_precreate_response = False - return reserved_bytes - - def clear_continuation(self) -> None: - self.continuation_owner_id = None - self.continuation_units = 0 - self.pending_silence_task = None - self.pending_silence_owner_id = None - - -class RuntimeDataPlane(Protocol): - def begin_request(self, request_id: str) -> None: ... - - def is_terminal(self, request_id: str | None) -> bool: ... - - def mark_terminal(self, request_id: str) -> None: ... - - def close_stream(self, request_id: str) -> None: ... - - def close_session(self, session_id: str, *, active_request_id: str | None = None) -> None: ... - - def project(self, result: object, *, context: object | None = None) -> Iterable[dict[str, object]]: ... - - -class ServingRuntimeAdapter(Protocol): - adapter_id: str - session_states: MutableMapping[str, ServingRuntimeSessionState] - data_plane: RuntimeDataPlane - clean_response_done_prefix: str - interrupted_tts_prefix: str - private_runtime_config_keys: frozenset[str] - collect_outputs_on_append: bool - - def create_session_state(self) -> ServingRuntimeSessionState: ... - - def session_state(self, session_id: str) -> ServingRuntimeSessionState: ... - - def remove_session_state(self, session_id: str) -> None: ... - - def is_enabled(self, config: object) -> bool: ... - - def capabilities(self, *, max_sessions: int) -> object: ... - - def validate_client_extra_body(self, extra_body: object) -> None: ... - - async def prepare_runtime_config(self, config: object, *, model_config: Any) -> dict[str, object]: ... - - def runtime_config_for_update( - self, - config: object, - current: Mapping[str, object], - ) -> dict[str, object]: ... - - def data_plane_context( - self, - *, - epoch: int, - turn_id: int, - active_response_turn_id: int | None, - active_response_id: str | None, - auto_responds: bool, - response_format: str, - speed: float | None, - modalities: tuple[str, ...], - ) -> object: ... - - -def load_serving_runtime_adapter( - path: str, - encode_audio, -) -> ServingRuntimeAdapter: - module_name, separator, attribute_name = path.rpartition(".") - if not separator: - raise ValueError(f"Invalid duplex serving runtime adapter path: {path!r}") - adapter_type = getattr(import_module(module_name), attribute_name) - return validate_serving_runtime_adapter(adapter_type(encode_audio)) - - -def validate_serving_runtime_adapter(adapter: object) -> ServingRuntimeAdapter: - required_methods = ( - "create_session_state", - "session_state", - "remove_session_state", - "is_enabled", - "capabilities", - "validate_client_extra_body", - "prepare_runtime_config", - "runtime_config_for_update", - "data_plane_context", - ) - missing = [name for name in required_methods if not callable(getattr(adapter, name, None))] - if missing: - raise TypeError(f"Duplex serving runtime adapter is missing callable method(s): {', '.join(missing)}") - if not isinstance(getattr(adapter, "adapter_id", None), str) or not adapter.adapter_id: - raise TypeError("Duplex serving runtime adapter must declare adapter_id") - if not isinstance(getattr(adapter, "session_states", None), MutableMapping): - raise TypeError("Duplex serving runtime adapter must declare mutable session_states") - for prefix_name in ("clean_response_done_prefix", "interrupted_tts_prefix"): - if not isinstance(getattr(adapter, prefix_name, None), str): - raise TypeError(f"Duplex serving runtime adapter must declare {prefix_name}") - private_keys = getattr(adapter, "private_runtime_config_keys", None) - if not isinstance(private_keys, frozenset) or any(not isinstance(key, str) for key in private_keys): - raise TypeError("Duplex serving runtime adapter must declare private_runtime_config_keys as frozenset[str]") - data_plane = getattr(adapter, "data_plane", None) - if data_plane is None: - raise TypeError("Duplex serving runtime adapter must declare data_plane") - data_plane_methods = ( - "begin_request", - "is_terminal", - "mark_terminal", - "close_stream", - "close_session", - "project", - ) - missing_data_plane_methods = [name for name in data_plane_methods if not callable(getattr(data_plane, name, None))] - if missing_data_plane_methods: - raise TypeError( - "Duplex serving runtime adapter data_plane is missing callable method(s): " - + ", ".join(missing_data_plane_methods) - ) - return adapter # type: ignore[return-value] - - -def payload_turn_id(payload: object) -> int | None: - if not isinstance(payload, Mapping): - return None - return coerce_int(payload.get("duplex_turn_id", payload.get("model_turn_id"))) - - -def coerce_int(value: object) -> int | None: - try: - return int(value) if value is not None else None - except (TypeError, ValueError): - return None diff --git a/vllm_omni/entrypoints/duplex/runtime_bridge.py b/vllm_omni/entrypoints/duplex/runtime_bridge.py deleted file mode 100644 index 96a82c061..000000000 --- a/vllm_omni/entrypoints/duplex/runtime_bridge.py +++ /dev/null @@ -1,1454 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -import base64 -import inspect -import time -from collections.abc import Mapping -from dataclasses import asdict, is_dataclass -from enum import Enum - -import numpy as np -from vllm.logger import init_logger - -from vllm_omni.engine.duplex.control_client import DuplexControlRequestError -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.runtime import duplex_data_plane_request_info -from vllm_omni.entrypoints.duplex.protocol import ( - DuplexSession, - DuplexSessionState, -) -from vllm_omni.entrypoints.duplex.runtime_adapter import ( - coerce_int, - payload_turn_id, -) - -logger = init_logger(__name__) - - -class NativeRuntimeBridgeMixin: - """Bridge serving sessions to native runtime control and data-plane APIs.""" - - _NATIVE_RUNTIME_FENCE_CONTRACT = { - "open_duplex_session_async": ("fence",), - "append_duplex_input_async": ("fence",), - "signal_duplex_turn_async": ("fence", "next_fence"), - "close_duplex_session_async": ("fence",), - } - - async def _open_runtime_session(self, session: DuplexSession, send_json) -> dict[str, object] | bool: - # Turn-based sessions use the chat-fallback path and do not own a - # model-native duplex runtime. Do not send control messages merely - # because the shared engine client exposes the duplex RPC methods. - if not self._uses_serving_runtime_adapter(session.config): - return True - contract_error = self._native_runtime_contract_error(session) - if contract_error is not None: - await send_json( - { - "type": "error", - "code": "runtime_contract_invalid", - "error": contract_error, - } - ) - return False - open_session = getattr(self._chat_service.engine_client, "open_duplex_session_async", None) - if not callable(open_session): - return True - try: - open_kwargs = { - "session_mode": "duplex", - "capabilities": session.capabilities.as_dict(), - "session_config": session.config.as_dict(), - "timeout": self._runtime_control_timeout_s(session), - } - if self._callable_accepts_keyword(open_session, "runtime_config"): - open_kwargs["runtime_config"] = dict(session.runtime_config) - if self._callable_accepts_keyword(open_session, "fence"): - open_kwargs["fence"] = DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ) - result = await open_session(session.session_id, **open_kwargs) - except Exception as exc: - if isinstance(exc, DuplexControlRequestError) and exc.code == "resource_exhausted": - logger.info("Duplex runtime session admission rejected: %s", exc) - else: - logger.exception("Failed to open duplex runtime session: %s", exc) - await self._send_runtime_error(send_json, "runtime_open_failed", exc, session=session) - return False - if ( - isinstance(result, dict) - and session.capabilities.implementation_level == "model_native_duplex" - and self._runtime_control_failed(result) - ): - await send_json( - { - "type": "error", - "error": "Native duplex runtime is not available for this session", - "code": ("runtime_open_unsupported" if result.get("unsupported_count") else "runtime_open_failed"), - "runtime_control": self._redact_runtime_control_result(result), - } - ) - return False - return result if isinstance(result, dict) else True - - async def _append_runtime_input( - self, - session: DuplexSession, - payload: object, - *, - operation_id: str | None = None, - final: bool, - send_json, - mode: str = "append_tokens", - expected_epoch: int | None = None, - ) -> tuple[bool, bool]: - if not session.capabilities.supports_input_append: - return True, False - append_input = getattr(self._chat_service.engine_client, "append_duplex_input_async", None) - if not callable(append_input): - return True, False - if expected_epoch is not None and session.epoch != expected_epoch: - return True, False - try: - append_kwargs = { - "mode": mode, - "payload": payload, - "final": final, - "timeout": self._runtime_control_timeout_s(session), - } - if operation_id is not None and self._callable_accepts_keyword(append_input, "operation_id"): - append_kwargs["operation_id"] = operation_id - if expected_epoch is not None and self._callable_accepts_keyword(append_input, "expected_epoch"): - append_kwargs["expected_epoch"] = expected_epoch - if self._callable_accepts_keyword(append_input, "fence"): - payload_turn = payload_turn_id(payload) - append_kwargs["fence"] = DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=( - payload_turn - if payload_turn is not None - else ( - session.active_response_turn_id - if session.active_response_turn_id is not None - else session.turn_id - ) - ), - incarnation=session.incarnation, - ) - if self._callable_accepts_keyword(append_input, "collect_outputs"): - append_kwargs["collect_outputs"] = bool( - getattr(self._serving_runtime_adapter, "collect_outputs_on_append", False) - ) - result = await append_input(session.session_id, **append_kwargs) - except Exception as exc: - logger.exception("Failed to append duplex runtime input: %s", exc) - await self._send_runtime_error(send_json, "runtime_append_failed", exc, session=session) - return False, False - if isinstance(result, dict) and self._runtime_control_failed(result): - await self._send_runtime_control_error( - send_json, - "runtime_append_failed", - "Duplex runtime append failed", - result, - session=session, - ) - return False, False - if expected_epoch is not None and session.epoch != expected_epoch: - return True, False - await self._send_runtime_control_if_needed(send_json, result, session=session) - request_id, _ = self._data_plane_request_info(result) if isinstance(result, dict) else (None, None) - if request_id is not None: - self._serving_runtime_adapter.data_plane.begin_request(request_id) - close_reason, emitted_response = await self._send_native_duplex_events( - send_json, - result, - session=session, - expected_epoch=expected_epoch, - ) - # A resumable data-plane append remains active even when its persistent - # drain task already exists. Treating only a newly-created drain as - # activity clears active_request_id after later appends and prevents a - # terminal TTS segment from scheduling the next model decision. - emitted_response = emitted_response or request_id is not None - if close_reason is None and await self._start_native_data_plane_stream_task( - send_json, - result, - session=session, - expected_epoch=expected_epoch, - ): - emitted_response = True - if close_reason is not None: - if not await self._close_runtime_session(session, reason=close_reason, send_json=send_json): - return False, emitted_response - session.close() - await send_json( - { - "type": "session.closed", - "session_id": session.session_id, - "reason": close_reason, - } - ) - return False, emitted_response - return True, emitted_response - - @staticmethod - def _callable_accepts_keyword(fn, name: str) -> bool: - try: - params = inspect.signature(fn).parameters.values() - except (TypeError, ValueError): - return True - return any(param.kind == inspect.Parameter.VAR_KEYWORD or param.name == name for param in params) - - def _native_runtime_contract_error(self, session: DuplexSession) -> str | None: - if session.capabilities.implementation_level != "model_native_duplex": - return None - engine_client = self._chat_service.engine_client - invalid: list[str] = [] - for method_name, required_keywords in self._NATIVE_RUNTIME_FENCE_CONTRACT.items(): - method = getattr(engine_client, method_name, None) - if not callable(method): - invalid.append(f"{method_name} is missing") - continue - try: - parameters = tuple(inspect.signature(method).parameters.values()) - except (TypeError, ValueError): - invalid.append(f"{method_name} signature is not inspectable") - continue - accepts_kwargs = any(parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters) - declared = {parameter.name for parameter in parameters} - missing = [name for name in required_keywords if not accepts_kwargs and name not in declared] - if missing: - invalid.append(f"{method_name} is missing keyword(s): {', '.join(missing)}") - if not invalid: - return None - return "Native duplex runtime requires the fenced control contract: " + "; ".join(invalid) - - async def _start_native_data_plane_stream_task( - self, - send_json, - result: object, - *, - session: DuplexSession, - expected_epoch: int | None = None, - ) -> bool: - request_id, _ = self._data_plane_request_info(result) - if request_id is None or self._data_plane_outputs_finished(result): - return False - session.bind_request(request_id) - - native = self._runtime_session_state(session) - old_task = native.data_plane_task - if old_task is not None and not old_task.done(): - if self._session_auto_responds(session): - # One persistent drain per session, like the official worker's - # single synchronous loop: the resumable data-plane request id - # is stable, and cancel/restart on every append orphans any - # decision that lands in the swap window. - native.data_plane_restart_requested = True - return False - old_task.cancel() - try: - await asyncio.wait_for(asyncio.gather(old_task, return_exceptions=True), timeout=0.25) - except asyncio.TimeoutError: - pass - - async def _run() -> None: - close_reason: str | None = None - try: - close_reason = await self._drain_native_data_plane_stream( - send_json, - result, - session=session, - expected_epoch=expected_epoch, - ) - if close_reason is not None: - if await self._close_runtime_session(session, reason=close_reason, send_json=send_json): - session.close() - await send_json( - { - "type": "session.closed", - "session_id": session.session_id, - "reason": close_reason, - } - ) - except asyncio.CancelledError: - raise - except Exception as exc: - logger.exception("Failed to drain duplex data-plane stream: %s", exc) - await self._send_runtime_error(send_json, "runtime_data_plane_stream_failed", exc, session=session) - if session.state != DuplexSessionState.CLOSED: - close_reason = "runtime_data_plane_stream_failed" - if await self._close_runtime_session(session, reason=close_reason, send_json=send_json): - session.close() - await send_json( - { - "type": "session.closed", - "session_id": session.session_id, - "reason": close_reason, - } - ) - finally: - if native.data_plane_task is task: - native.data_plane_task = None - restart_requested = ( - native.data_plane_restart_requested - and close_reason is None - and session.state != DuplexSessionState.CLOSED - and session.active_request_id == request_id - and (expected_epoch is None or session.epoch == expected_epoch) - ) - native.data_plane_restart_requested = False - if close_reason is None and not self._session_auto_responds(session): - # Auto-respond sessions keep one resumable stage-1 stream - # whose audio accumulates across speak units; the offset - # must survive drain-task turnover or every unit re-sends - # the reply audio from the start. - self._serving_runtime_adapter.data_plane.close_stream(request_id) - if restart_requested: - await self._start_native_data_plane_stream_task( - send_json, - result, - session=session, - expected_epoch=expected_epoch, - ) - elif close_reason is None: - await self._maybe_continue_native_response(send_json, session=session, expected_epoch=expected_epoch) - - task = asyncio.create_task(_run()) - native.data_plane_task = task - return True - - # One MiniCPM model unit (1 s at 16 kHz) is the compatibility default. - # Other native-duplex adapters may use a different frame-locked unit. - _NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO = base64.b64encode(bytes(16000 * 4)).decode("ascii") - _NATIVE_RESPONSE_MAX_CONTINUATION_UNITS = 8 - # Each unit is one second of silence. Auto-response needs substantially - # more room than the eight decision beats above, but must still terminate - # well before the 300-second benchmark playback horizon. The final unit - # asks the model to LISTEN; a protocol fallback closes the response if that - # logit hint is ignored. - _NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS = 64 - - def _native_silence_unit_payload(self) -> dict[str, object]: - samples = int(getattr(self._serving_runtime_adapter, "silence_continuation_samples", 16000)) - audio = ( - self._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO - if samples == 16000 - else base64.b64encode(bytes(samples * 4)).decode("ascii") - ) - return { - "type": "audio", - "audio": audio, - "format": "pcm_f32le", - "sample_rate_hz": 16000, - } - - def _native_response_continuations_remaining(self, session: DuplexSession, response_id: str) -> bool: - native = self._runtime_session_state(session) - owner_id = f"response:{response_id}" - count = native.continuation_units if native.continuation_owner_id == owner_id else 0 - limit = ( - self._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS - if self._session_auto_responds(session) - else self._NATIVE_RESPONSE_MAX_CONTINUATION_UNITS - ) - return count < limit - - async def _finish_bounded_native_auto_response( - self, - send_json, - *, - session: DuplexSession, - expected_epoch: int | None, - model_turn_id: int | None = None, - ) -> None: - """Close an auto-response whose final silence unit did not end it.""" - native = self._runtime_session_state(session) - 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: - return - native.clear_continuation() - payload = { - "type": "response.listen", - "session_id": session.session_id, - "epoch": response_epoch, - "reason": "continuation_limit", - "model_listen": False, - } - if response_id is not None: - # Clients demultiplex a listen against a precreated response by - # its id; every listen that terminates a response must carry it. - payload["response_id"] = response_id - await send_json(payload) - if response_id is None: - if session.epoch == response_epoch and response_turn_id is not None: - session.complete_model_turn(response_turn_id) - return - if session.epoch != response_epoch or session.active_response_id != response_id: - return - if response_turn_id is not None: - session.complete_model_turn(response_turn_id) - session.end_response(commit_text=False, preserve_request=True) - await send_json( - { - "type": "response.done", - "session_id": session.session_id, - "response_id": response_id, - "epoch": response_epoch, - "committed": False, - "playback": session.playback.as_dict(), - } - ) - - def _native_silence_continuation_is_stale( - self, - session: DuplexSession, - *, - request_id: str, - response_id: str | None, - response_owned: bool, - expected_epoch: int | None, - expected_incarnation: int, - expected_model_turn_id: int | None, - ) -> bool: - stale_common_owner = ( - session.state == DuplexSessionState.CLOSED - or session.incarnation != expected_incarnation - or session.active_request_id != request_id - or (expected_epoch is not None and session.epoch != expected_epoch) - ) - stale_response_owner = response_owned and session.active_response_id != response_id - stale_model_turn_owner = not response_owned and ( - session.active_response_id is not None or session.turn_id != expected_model_turn_id - ) - return stale_common_owner or stale_response_owner or stale_model_turn_owner - - async def _maybe_continue_native_response( - self, - send_json, - *, - session: DuplexSession, - expected_epoch: int | None, - expected_model_turn_id: int | None = None, - ) -> None: - """Give an active model turn another silence unit. - - A model turn can need another unit before it has emitted visible - text/audio and therefore before a Realtime response exists. Bind that - continuation to the model-turn identity; once a response exists, its - response/turn ownership remains the continuation fence. - """ - response_id = session.active_response_id - native = self._runtime_session_state(session) - if session.state == DuplexSessionState.CLOSED: - native.clear_continuation() - return - request_id = session.active_request_id - if request_id is None: - native.clear_continuation() - return - if expected_epoch is not None and session.epoch != expected_epoch: - return - auto_response = self._session_auto_responds(session) - response_owned = response_id is not None - if response_owned: - owner_id = f"response:{response_id}" - payload_turn_id = ( - session.active_response_turn_id if session.active_response_turn_id is not None else session.turn_id - ) - else: - if not auto_response or expected_model_turn_id is None or session.turn_id != expected_model_turn_id: - native.clear_continuation() - return - owner_id = f"model-turn:{expected_model_turn_id}" - payload_turn_id = expected_model_turn_id - count = native.continuation_units if native.continuation_owner_id == owner_id else 0 - continuation_limit = ( - self._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS - if auto_response - else self._NATIVE_RESPONSE_MAX_CONTINUATION_UNITS - ) - if count >= continuation_limit: - if auto_response: - await self._finish_bounded_native_auto_response( - send_json, - session=session, - expected_epoch=expected_epoch, - model_turn_id=payload_turn_id, - ) - return - payload = self._native_silence_unit_payload() - if auto_response and count + 1 == continuation_limit: - payload["force_listen"] = True - payload["duplex_turn_id"] = payload_turn_id - - scheduler = native.silence_continuation_scheduler - if scheduler is None: - return - try: - scheduled = await scheduler( - payload, - request_id=request_id, - owner_id=owner_id, - response_id=response_id, - response_owned=response_owned, - expected_epoch=expected_epoch, - expected_incarnation=session.incarnation, - expected_model_turn_id=expected_model_turn_id, - send_json=send_json, - ) - except Exception as exc: - logger.exception("Failed to schedule duplex native response continuation: %s", exc) - scheduled = False - if scheduled: - native.continuation_owner_id = owner_id - native.continuation_units = count + 1 - - async def _cancel_native_data_plane_stream(self, session: DuplexSession) -> bool: - native = self._runtime_session_state(session) - task = native.data_plane_task - native.data_plane_task = None - native.data_plane_restart_requested = False - if task is None or task.done(): - return False - task.cancel() - try: - await asyncio.wait_for(asyncio.gather(task, return_exceptions=True), timeout=0.25) - except asyncio.TimeoutError: - # Keep barge-in/cancel responsive. The task still carries the old - # expected_epoch and late model output is filtered by the writer. - pass - return True - - async def _signal_runtime_session( - self, - session: DuplexSession, - event: str, - send_json=None, - *, - fence: DuplexFence | None = None, - next_fence: DuplexFence | None = None, - session_config: dict[str, object] | None = None, - runtime_config: dict[str, object] | None = None, - ) -> bool: - if not self._uses_serving_runtime_adapter(session.config): - return True - signal_turn = getattr(self._chat_service.engine_client, "signal_duplex_turn_async", None) - if not callable(signal_turn): - return True - try: - signal_kwargs = { - "event": event, - "timeout": self._runtime_control_timeout_s(session), - } - if self._callable_accepts_keyword(signal_turn, "fence"): - signal_kwargs["fence"] = fence or DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ) - if next_fence is not None and self._callable_accepts_keyword(signal_turn, "next_fence"): - signal_kwargs["next_fence"] = next_fence - if session_config is not None and self._callable_accepts_keyword(signal_turn, "session_config"): - signal_kwargs["session_config"] = session_config - if runtime_config is not None and self._callable_accepts_keyword(signal_turn, "runtime_config"): - signal_kwargs["runtime_config"] = runtime_config - result = await signal_turn(session.session_id, **signal_kwargs) - except Exception as exc: - logger.exception("Failed to signal duplex runtime session: %s", exc) - if send_json is not None: - await self._send_runtime_error(send_json, "runtime_signal_failed", exc, session=session) - return False - if isinstance(result, dict) and self._runtime_signal_failed(result): - logger.warning( - "Duplex runtime signal failed: session=%s event=%s result=%s", - session.session_id, - event, - self._redact_runtime_control_result(result), - ) - if send_json is not None: - await self._send_runtime_control_error( - send_json, - "runtime_signal_failed", - "Duplex runtime signal failed", - result, - session=session, - ) - return False - await self._send_runtime_control_if_needed(send_json, result, session=session) - return True - - async def _close_runtime_session(self, session: DuplexSession, *, reason: str, send_json=None) -> bool: - if not self._uses_serving_runtime_adapter(session.config): - return True - close_session = getattr(self._chat_service.engine_client, "close_duplex_session_async", None) - if not callable(close_session): - return True - try: - close_kwargs = { - "reason": reason, - "timeout": self._runtime_control_timeout_s(session), - } - if self._callable_accepts_keyword(close_session, "fence"): - close_kwargs["fence"] = DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ) - result = await close_session(session.session_id, **close_kwargs) - except Exception as exc: - logger.exception("Failed to close duplex runtime session: %s", exc) - if send_json is not None: - await self._send_runtime_error(send_json, "runtime_close_failed", exc, session=session) - return False - if isinstance(result, dict) and self._runtime_control_failed(result): - if send_json is not None: - await self._send_runtime_control_error( - send_json, - "runtime_close_failed", - "Duplex runtime close failed", - result, - session=session, - ) - return False - await self._send_runtime_control_if_needed(send_json, result, session=session) - return True - - @staticmethod - def _runtime_control_timeout_s(session: DuplexSession) -> float: - raw = session.config.extra_body.get("duplex_control_timeout_s") or session.config.extra_body.get( - "runtime_control_timeout_s" - ) - if isinstance(raw, int | float) and raw > 0: - return float(raw) - if session.capabilities.implementation_level == "model_native_duplex": - return 60.0 - return 10.0 - - async def _send_runtime_control_if_needed( - self, - send_json, - result: object, - *, - session: DuplexSession, - ) -> None: - if send_json is None: - return - if not isinstance(result, dict): - return - if not result.get("unsupported_count") and not result.get("error_count"): - return - await send_json( - { - "type": "runtime.control", - "session_id": session.session_id, - "epoch": session.epoch, - "result": self._redact_runtime_control_result(result), - } - ) - - @staticmethod - def _runtime_control_failed(result: dict[str, object]) -> bool: - if result.get("ok") is False: - return True - for key in ("unsupported_count", "error_count"): - value = result.get(key) - if isinstance(value, int | float) and value > 0: - return True - return False - - @classmethod - def _runtime_signal_failed(cls, result: dict[str, object]) -> bool: - error_count = result.get("error_count") - if isinstance(error_count, int | float) and error_count > 0: - return True - if result.get("ok") is not False: - return False - return not cls._runtime_signal_has_data_plane_ack(result) - - @classmethod - def _runtime_signal_has_data_plane_ack(cls, value: object) -> bool: - if isinstance(value, dict): - if value.get("data_plane_signal") is True and value.get("supported") is True: - return True - return any(cls._runtime_signal_has_data_plane_ack(child) for child in value.values()) - if isinstance(value, list | tuple): - return any(cls._runtime_signal_has_data_plane_ack(item) for item in value) - return False - - @classmethod - def _redact_runtime_control_result(cls, value: object) -> object: - if isinstance(value, dict): - redacted = { - key: cls._redact_runtime_control_result(child) - for key, child in value.items() - if key - not in { - "stage_handoff", - "tts_handoff", - "omni_payload", - "tts_hidden_states", - "tts_token_ids", - "traceback", - } - } - native_result = redacted.get("native_result") - if isinstance(native_result, dict): - if native_result.get("requires_stage_handoff") is True: - native_result.pop("requires_stage_handoff", None) - if native_result.get("requires_tts_stage") is True: - native_result.pop("requires_tts_stage", None) - return redacted - if isinstance(value, list | tuple): - return [cls._redact_runtime_control_result(item) for item in value] - if is_dataclass(value) and not isinstance(value, type): - # Control results carry typed values straight from the engine: - # `fence` and `accepted_fence` are `DuplexFence` instances added by - # `DuplexControlClient.execute`. This redacted value is emitted on - # the wire (for example `session.created.runtime_control` in - # `session_runner`), where `WebSocket.send_json` would raise - # `TypeError: Object of type DuplexFence is not JSON serializable` - # and kill the session during the handshake. - return {key: cls._redact_runtime_control_result(child) for key, child in asdict(value).items()} - if isinstance(value, Enum): - return value.value - return value - - async def _send_native_duplex_events( - self, - send_json, - result: object, - *, - session: DuplexSession, - expected_epoch: int | None = None, - ) -> tuple[str | None, bool]: - if send_json is None: - return None, False - if expected_epoch is not None and session.epoch != expected_epoch: - return None, False - close_reason: str | None = None - emitted_response = False - request_id, _ = self._data_plane_request_info(result) - if self._serving_runtime_adapter.data_plane.is_terminal(request_id): - return None, False - if request_id is not None and session.active_request_id is None: - session.bind_request(request_id) - context = self._runtime_data_plane_context(session) - for native_result in self._serving_runtime_adapter.data_plane.project(result, context=context): - close_reason_for_result, did_emit = await self._send_one_native_duplex_event( - send_json, - native_result, - session=session, - expected_epoch=expected_epoch, - ) - 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: - return None, emitted_response - return close_reason, emitted_response - - async def _drain_native_data_plane_stream( - self, - send_json, - result: object, - *, - session: DuplexSession, - expected_epoch: int | None = None, - ) -> str | None: - request_id, response_stage_id = self._data_plane_request_info(result) - if request_id is None: - return None - if self._data_plane_outputs_finished(result): - return None - collect_outputs = getattr( - self._chat_service.engine_client, - "collect_duplex_data_plane_outputs_async", - None, - ) - if not callable(collect_outputs): - return None - - close_reason: str | None = None - empty_polls = 0 - while close_reason is None: - if self._serving_runtime_adapter.data_plane.is_terminal(request_id): - return None - if self._session_auto_responds(session): - active_request_id = session.active_request_id - if active_request_id is not None and active_request_id != request_id: - return None - if expected_epoch is not None and session.epoch != expected_epoch: - return None - if session.state == DuplexSessionState.CLOSED: - return None - outputs = await collect_outputs( - request_id, - response_stage_id=response_stage_id, - timeout=self._runtime_control_timeout_s(session), - ) - if expected_epoch is not None and session.epoch != expected_epoch: - return None - if not outputs: - # An empty poll means no output arrived within one control - # window, not that the stream is over. Exiting here orphans a - # decision that lands moments later: it would sit queued until - # the NEXT append starts a fresh drain task, adding one full - # chunk of latency to every model decision. - empty_polls += 1 - if empty_polls >= 3: - return None - continue - empty_polls = 0 - drain_result = { - "data_plane_outputs": outputs, - } - close_reason, emitted_response = await self._send_native_duplex_events( - send_json, - drain_result, - session=session, - expected_epoch=expected_epoch, - ) - if ( - self._data_plane_outputs_finished(drain_result) - and emitted_response - and not self._session_auto_responds(session) - ): - # Auto-respond sessions are resumable: every segment ends - # with finished=True but the stream continues with the next - # audio chunk. Exiting here would make each decision wait for - # the next append to start a fresh drain task (one full chunk - # of added latency); keep draining until epoch change, - # session close, or cancellation by a newer drain task. - return close_reason - # A batch without a client-visible event (e.g. the generic - # stage-0 final-output message that precedes the listen/speak - # decision in the same segment) must not terminate the drain; - # the real decision is still in flight. - return close_reason - - @staticmethod - def _data_plane_outputs_finished(result: object) -> bool: - if not isinstance(result, dict): - return False - outputs = result.get("data_plane_outputs") - if not isinstance(outputs, list) or not outputs: - return False - return bool(getattr(outputs[-1], "finished", False)) - - @staticmethod - def _data_plane_request_info(result: object) -> tuple[str | None, int | None]: - if not isinstance(result, dict): - return None, None - return duplex_data_plane_request_info(result) - - def _runtime_data_plane_context(self, session: DuplexSession) -> object: - response_config = session.response_config - return self._serving_runtime_adapter.data_plane_context( - epoch=session.epoch, - turn_id=session.turn_id, - active_response_turn_id=session.active_response_turn_id, - active_response_id=session.active_response_id, - auto_responds=self._session_auto_responds(session), - response_format=response_config.response_format, - speed=response_config.speed, - modalities=tuple(response_config.modalities), - ) - - async def _send_one_native_duplex_event( - self, - send_json, - native_result: dict[str, object], - *, - session: DuplexSession, - expected_epoch: int | None = None, - ) -> tuple[str | None, bool]: - close_reason: str | None = None - emitted_response = False - if expected_epoch is not None and session.epoch != expected_epoch: - return close_reason, emitted_response - data_plane_request_id = native_result.get("data_plane_request_id") - if isinstance(data_plane_request_id, str) and self._serving_runtime_adapter.data_plane.is_terminal( - data_plane_request_id - ): - return close_reason, emitted_response - active_request_matches = session.active_request_id == data_plane_request_id or ( - self._session_auto_responds(session) and session.active_request_id is None - ) - if isinstance(data_plane_request_id, str) and not active_request_matches: - return close_reason, emitted_response - if isinstance(native_result.get("error_code"), str): - response_id = session.active_response_id - await send_json( - { - "type": "error", - "code": native_result.get("error_code"), - "session_id": session.session_id, - "epoch": session.epoch, - "error": str(native_result.get("error") or "Duplex native data-plane error"), - } - ) - if response_id is not None: - session.end_response(commit_text=False) - await send_json( - { - "type": "response.done", - "session_id": session.session_id, - "response_id": response_id, - "epoch": session.epoch, - "committed": False, - "status": "failed", - "status_details": { - "type": "failed", - "reason": native_result.get("error_code"), - }, - "playback": session.playback.as_dict(), - } - ) - return close_reason, True - if native_result.get("function_call") is True: - await send_json( - { - "type": "function_call.done", - "session_id": session.session_id, - "epoch": session.epoch, - "call_id": native_result.get("call_id"), - "name": native_result.get("name"), - "arguments": native_result.get("arguments", ""), - } - ) - return close_reason, True - if native_result.get("requires_stage_handoff") is True or native_result.get("requires_tts_stage") is True: - # Stage0 announces a talker handoff before Stage1 has produced - # client-visible audio. A single client input stream may contain - # several model turns, including terminal-only turns. Reserve the - # protocol response only when Stage1 emits text/audio so empty - # model turns cannot create empty responses or steal later audio. - return close_reason, False - is_listen = native_result.get("is_listen") - model_turn_id = coerce_int(native_result.get("model_turn_id")) - if native_result.get("is_buffering") is True or native_result.get("prefill_success") is False: - if native_result.get("data_plane_request_id") == session.active_request_id: - session.clear_request() - # No response_id here: a buffering listen is informational, not a - # response terminal — clients must not close a handle over it. - payload = { - "type": "response.listen", - "session_id": session.session_id, - "epoch": session.epoch, - "reason": native_result.get("reason") or "buffering", - "model_listen": False, - "buffering": True, - } - self._attach_native_runtime_metadata(payload, native_result) - await send_json(payload) - return close_reason, emitted_response - if is_listen is True: - await self._end_active_response_before_future_model_turn( - send_json, - session=session, - model_turn_id=model_turn_id, - ) - if ( - session.active_response_id is not None - and model_turn_id is not None - and not session.active_response_accepts_model_turn(model_turn_id) - ): - return close_reason, emitted_response - auto_response = self._session_auto_responds(session) - active_response_id = session.active_response_id - auto_continuations_remaining = active_response_id is None or self._native_response_continuations_remaining( - session, active_response_id - ) - non_terminal_auto_listen = ( - auto_response - and active_response_id is not None - and session.active_request_id is not None - and native_result.get("end_of_turn") is not True - and auto_continuations_remaining - ) - if non_terminal_auto_listen: - await self._maybe_continue_native_response( - send_json, - session=session, - expected_epoch=expected_epoch, - ) - return close_reason, emitted_response - if not auto_response and data_plane_request_id == session.active_request_id: - session.clear_request() - model_listen = native_result.get("model_listen") - if not isinstance(model_listen, bool): - model_listen = native_result.get("reason") in {None, "", "model_listen"} - response_id = session.active_response_id - if isinstance(data_plane_request_id, str) and not auto_response: - self._serving_runtime_adapter.data_plane.mark_terminal(data_plane_request_id) - emitted_response = True - payload = { - "type": "response.listen", - "session_id": session.session_id, - "epoch": session.epoch, - "reason": native_result.get("reason") or "model_listen", - "model_listen": model_listen, - } - if response_id is not None: - # The listen is this response's terminal decision (a - # response.done follows below with the same id); carry the id - # so clients can bind the decision to the precreated response. - payload["response_id"] = response_id - self._attach_native_runtime_metadata(payload, native_result) - await send_json(payload) - if native_result.get("abort_data_plane_request") is True and isinstance(data_plane_request_id, str): - await self._abort_request_background( - session, - data_plane_request_id, - send_json, - notify=False, - ) - if response_id is not None: - if not self._session_auto_responds(session) and self._native_response_continuations_remaining( - session, response_id - ): - # The official model often listens for a silence beat or - # two before answering; keep the response open and give it - # the next silence unit as a decision point. - await self._maybe_continue_native_response( - send_json, - session=session, - expected_epoch=expected_epoch, - ) - return close_reason, emitted_response - if auto_response: - self._runtime_session_state(session).clear_continuation() - if not auto_continuations_remaining: - completed_turn_id = model_turn_id - if completed_turn_id is None: - completed_turn_id = session.active_response_turn_id - if completed_turn_id is not None: - session.complete_model_turn(completed_turn_id) - session.end_response(commit_text=False, preserve_request=auto_response) - await send_json( - { - "type": "response.done", - "session_id": session.session_id, - "response_id": response_id, - "epoch": session.epoch, - "committed": False, - "playback": session.playback.as_dict(), - } - ) - return close_reason, emitted_response - - text = native_result.get("text") - audio = native_result.get("audio_data", native_result.get("audio")) - end_of_turn = bool(native_result.get("end_of_turn", False)) - has_text = isinstance(text, str) and bool(text) - has_audio = isinstance(audio, str) and bool(audio) - if not has_text and not has_audio and not end_of_turn: - tts_segment_ended = ( - native_result.get("stage_role") == "tts" and native_result.get("abort_data_plane_request") is True - ) - if ( - tts_segment_ended - and self._session_auto_responds(session) - and (model_turn_id is not None or session.active_response_id is not None) - ): - await self._maybe_continue_native_response( - send_json, - session=session, - expected_epoch=expected_epoch, - expected_model_turn_id=model_turn_id, - ) - return close_reason, emitted_response - if end_of_turn and not has_text and not has_audio and session.active_response_id is None: - if isinstance(data_plane_request_id, str): - if not self._session_auto_responds(session) and data_plane_request_id == session.active_request_id: - session.clear_request() - if not self._session_auto_responds(session): - self._serving_runtime_adapter.data_plane.mark_terminal(data_plane_request_id) - if model_turn_id is not None: - session.complete_model_turn(model_turn_id) - if self._session_auto_responds(session): - self._runtime_session_state(session).clear_continuation() - emitted_response = True - payload = { - "type": "response.listen", - "session_id": session.session_id, - "epoch": session.epoch, - "reason": "model_turn_completed_without_output", - "model_listen": True, - } - self._attach_native_runtime_metadata(payload, native_result) - await send_json(payload) - return close_reason, emitted_response - if session.active_response_id is None and model_turn_id is not None and model_turn_id < session.turn_id: - # A continuation append can already be in flight when the prior - # response reaches turn EOS. Its late audio still carries the - # completed model turn, so it must not reserve a second Realtime - # response for that turn. - return close_reason, emitted_response - await self._end_active_response_before_future_model_turn( - send_json, - session=session, - model_turn_id=model_turn_id, - ) - if ( - session.active_response_id is not None - and model_turn_id is not None - and not session.active_response_accepts_model_turn(model_turn_id) - ): - return close_reason, emitted_response - emitted_response = True - response_created = False - response_id = session.active_response_id - if response_id is None: - response_id = session.begin_response(turn_id=model_turn_id) - response_created = True - response_request_metrics = session.mark_response_first_outputs( - observed_at_s=time.monotonic(), - has_text=has_text, - has_audio=has_audio, - ) - if response_created: - response_created_payload = self._response_created_payload( - session, - response_id, - epoch=session.epoch, - ) - if response_request_metrics: - response_created_payload["response_request_metrics"] = response_request_metrics - await send_json(response_created_payload) - response_stage_metrics = session.accumulate_response_stage_metrics( - native_result.get("stage_metrics") if isinstance(native_result.get("stage_metrics"), Mapping) else None - ) - if response_created: - speak_payload = { - "type": "response.speak", - "session_id": session.session_id, - "response_id": response_id, - "epoch": session.epoch, - "text": text if isinstance(text, str) else "", - "end_of_turn": end_of_turn, - "model_speak": True, - } - self._attach_native_runtime_metadata( - speak_payload, - native_result, - stage_metrics=response_stage_metrics, - response_request_metrics=response_request_metrics, - ) - await send_json(speak_payload) - previous_sent_ms = session.playback.sent_ms - text_chars_before_append = len("".join(session.assistant_text_buffer)) - if isinstance(text, str): - session.append_assistant_text(text) - duration_ms = native_result.get("audio_duration_ms") - text_chars = len("".join(session.assistant_text_buffer)) - mark_duration_ms = None - mark_text_chars: int | None = text_chars - if native_result.get("audio_text_mark") is False: - mark_text_chars = None - if isinstance(duration_ms, int | float): - mark_duration_ms = int(duration_ms) - if native_result.get("audio_duration_is_cumulative") is not True: - mark_duration_ms += session.playback.sent_ms - audio_text_marks = native_result.get("audio_text_marks") - audio_text_marks = self._normalize_native_audio_text_marks( - audio_text_marks if isinstance(audio_text_marks, list) else None, - audio_offset_ms=( - 0 - if native_result.get("audio_text_marks_are_cumulative") is True - or native_result.get("audio_duration_is_cumulative") is True - else previous_sent_ms - ), - text_offset_chars=( - 0 if native_result.get("audio_text_marks_are_cumulative") is True else text_chars_before_append - ), - ) - session.mark_audio_sent( - mark_duration_ms, - text_chars=mark_text_chars if mark_duration_ms is not None else None, - audio_text_marks=audio_text_marks, - ) - payload = { - "type": "response.output_audio.delta", - "session_id": session.session_id, - "response_id": response_id, - "epoch": session.epoch, - "text": text if isinstance(text, str) else "", - "audio": audio if isinstance(audio, str) else "", - "format": ( - native_result.get("audio_format") - if isinstance(native_result.get("audio_format"), str) - else session.response_config.response_format - ), - "end_of_turn": end_of_turn, - "model_speak": True, - } - if mark_duration_ms is not None: - payload["audio_duration_ms"] = mark_duration_ms - if audio_text_marks: - payload["audio_text_marks"] = audio_text_marks - elif mark_duration_ms is not None and mark_text_chars is not None: - payload["audio_text_marks"] = [ - { - "text_chars": max(0, int(mark_text_chars)), - "audio_end_ms": max(0, int(mark_duration_ms)), - } - ] - payload["playback"] = session.playback.as_dict() - sample_rate_hz = native_result.get("sample_rate_hz") or native_result.get("audio_sample_rate_hz") - if isinstance(sample_rate_hz, int | float) and int(sample_rate_hz) > 0: - payload["sample_rate_hz"] = int(sample_rate_hz) - self._attach_native_runtime_metadata( - payload, - native_result, - stage_metrics=response_stage_metrics, - response_request_metrics=response_request_metrics, - ) - await send_json(payload) - if ( - not end_of_turn - and native_result.get("stage_role") == "tts" - and native_result.get("abort_data_plane_request") is True - and self._session_auto_responds(session) - ): - await self._maybe_continue_native_response( - send_json, - session=session, - expected_epoch=expected_epoch, - ) - if end_of_turn: - data_plane_request_id = native_result.get("data_plane_request_id") - if isinstance(data_plane_request_id, str) and not self._session_auto_responds(session): - self._serving_runtime_adapter.data_plane.close_stream(data_plane_request_id) - if isinstance(data_plane_request_id, str): - if not self._session_auto_responds(session) and data_plane_request_id == session.active_request_id: - session.clear_request() - if not self._session_auto_responds(session): - self._serving_runtime_adapter.data_plane.mark_terminal(data_plane_request_id) - should_commit = self._should_commit_response_to_history(session, response_id) - committed_message = session.end_response( - commit_text=should_commit, - preserve_request=self._session_auto_responds(session), - ) - model_turn_id = coerce_int(native_result.get("model_turn_id")) - if model_turn_id is not None: - session.complete_model_turn(model_turn_id) - if should_commit: - session.register_history_item(f"item_{response_id}", committed_message) - await send_json( - { - "type": "response.done", - "session_id": session.session_id, - "response_id": response_id, - "epoch": session.epoch, - "committed": committed_message is not None, - "playback": session.playback.as_dict(), - } - ) - return close_reason, emitted_response - - async def _end_active_response_before_future_model_turn( - self, - send_json, - *, - session: DuplexSession, - model_turn_id: int | None, - ) -> None: - if not self._session_auto_responds(session): - return - response_id = session.active_response_id - active_turn_id = session.active_response_turn_id - if response_id is None or model_turn_id is None or active_turn_id is None: - return - if int(model_turn_id) <= int(active_turn_id): - return - session.complete_model_turn(int(model_turn_id) - 1) - should_commit = self._should_commit_response_to_history(session, response_id) - committed_message = session.end_response( - commit_text=should_commit, - preserve_request=True, - ) - if should_commit: - session.register_history_item(f"item_{response_id}", committed_message) - await send_json( - { - "type": "response.done", - "session_id": session.session_id, - "response_id": response_id, - "epoch": session.epoch, - "committed": committed_message is not None, - "playback": session.playback.as_dict(), - } - ) - - @staticmethod - def _normalize_native_audio_text_marks( - audio_text_marks: list[object] | None, - *, - audio_offset_ms: int, - text_offset_chars: int, - ) -> list[dict[str, int]] | None: - if not audio_text_marks: - return None - normalized: list[dict[str, int]] = [] - for raw_mark in audio_text_marks: - if not isinstance(raw_mark, dict): - continue - raw_text_chars = raw_mark.get("text_chars") - raw_audio_end_ms = raw_mark.get("audio_end_ms", raw_mark.get("audio_ms")) - if not isinstance(raw_text_chars, int | float) or not isinstance(raw_audio_end_ms, int | float): - continue - normalized.append( - { - "text_chars": max(0, int(raw_text_chars) + int(text_offset_chars)), - "audio_end_ms": max(0, int(raw_audio_end_ms) + int(audio_offset_ms)), - } - ) - return normalized or None - - def _cleanup_duplex_session_state(self, session: DuplexSession) -> None: - session_id = session.session_id - pipeline = self._server_vad_pipelines.pop(session_id, None) - if pipeline is not None: - pipeline.reset() - self._realtime_vad_metrics.session_finished() - adapter = self._serving_runtime_adapter - if adapter is None: - self._serving_session_states.pop(session_id, None) - else: - adapter.remove_session_state(session_id) - adapter.data_plane.close_session( - session_id, - active_request_id=session.active_request_id, - ) - - def _encode_native_data_plane_audio( - self, - audio_data: object, - sample_rate_hz: int, - response_format: str, - speed: float | None, - ) -> str | None: - if audio_data is None: - return None - try: - import torch - - from vllm_omni.entrypoints.openai.protocol.audio import CreateAudio - - if isinstance(audio_data, torch.Tensor): - audio_tensor = audio_data.detach().cpu().float().numpy() - else: - audio_tensor = np.asarray(audio_data, dtype=np.float32) - if audio_tensor.ndim > 1: - audio_tensor = audio_tensor.reshape(-1) - audio_response = self._chat_service.create_audio( - CreateAudio( - audio_tensor=audio_tensor, - sample_rate=sample_rate_hz, - response_format=response_format, - speed=float(speed) if isinstance(speed, int | float) and speed > 0 else 1.0, - stream_format="audio", - base64_encode=True, - ) - ) - return str(audio_response.audio_data) - except Exception: - logger.exception("Failed to encode duplex data-plane audio output") - return None - - @staticmethod - def _attach_native_runtime_metadata( - payload: dict[str, object], - native_result: dict[str, object], - *, - stage_metrics: Mapping[str, object] | None = None, - response_request_metrics: Mapping[str, object] | None = None, - ) -> None: - metadata: dict[str, object] = {} - runtime_impl = native_result.get("runtime_impl") - if isinstance(runtime_impl, str) and runtime_impl: - metadata["runtime_impl"] = runtime_impl - owned_runtime = native_result.get("owned_runtime") - if isinstance(owned_runtime, bool): - metadata["owned_runtime"] = owned_runtime - model_turn_id = coerce_int(native_result.get("model_turn_id")) - if model_turn_id is not None: - metadata["model_turn_id"] = model_turn_id - for name in ( - "uses_model_runner_scheduler", - "runner_kv_backed", - ): - value = native_result.get(name) - if isinstance(value, bool): - metadata[name] = value - effective_stage_metrics = stage_metrics if stage_metrics is not None else native_result.get("stage_metrics") - if isinstance(effective_stage_metrics, Mapping): - metadata["stage_metrics"] = { - str(stage_id): dict(values) - for stage_id, values in effective_stage_metrics.items() - if isinstance(values, Mapping) - } - if response_request_metrics: - metadata["response_request_metrics"] = dict(response_request_metrics) - if metadata: - payload["vllm_omni"] = metadata - - async def _send_runtime_error( - self, - send_json, - code: str, - exc: Exception, - *, - session: DuplexSession | None = None, - ) -> None: - retryable = False - message = str(exc) - if isinstance(exc, DuplexControlRequestError): - code = exc.code - retryable = exc.retryable - error_data = exc.result.get("error") - if isinstance(error_data, dict) and isinstance(error_data.get("message"), str): - message = error_data["message"] - payload: dict[str, object] = { - "type": "error", - "code": code, - "error": message, - "retryable": retryable, - } - if session is not None: - payload["session_id"] = session.session_id - payload["epoch"] = session.epoch - await send_json(payload) - - async def _send_runtime_control_error( - self, - send_json, - code: str, - message: str, - result: dict[str, object], - *, - session: DuplexSession, - ) -> None: - await send_json( - { - "type": "error", - "code": code, - "error": message, - "session_id": session.session_id, - "epoch": session.epoch, - "runtime_control": self._redact_runtime_control_result(result), - } - ) diff --git a/vllm_omni/entrypoints/duplex/server_vad.py b/vllm_omni/entrypoints/duplex/server_vad.py deleted file mode 100644 index 610e51c8c..000000000 --- a/vllm_omni/entrypoints/duplex/server_vad.py +++ /dev/null @@ -1,528 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -import hashlib -import threading -import time -from collections.abc import Mapping -from dataclasses import dataclass, replace -from pathlib import Path -from typing import TYPE_CHECKING, Literal, Protocol - -import numpy as np -from vllm.transformers_utils.repo_utils import try_get_local_file - -if TYPE_CHECKING: - from vllm_omni.entrypoints.openai.audio_utils_mixin import StreamingAudioResampler - -SILERO_VAD_REPO_ID = "istupakov/silero-vad-onnx" -SILERO_VAD_REVISION = "8b14476858ef240c50b3884bb38cc67290c1cc70" -SILERO_VAD_FILENAME = "silero_vad.onnx" -SILERO_VAD_SHA256 = "1a153a22f4509e292a94e67d6f9b85e8deb25b4988682b7e174c65279d8788e3" -SILERO_VAD_MIN_THRESHOLD = 0.15 -SILERO_VAD_DEFAULT_MIN_SPEECH_DURATION_MS = 96 - -_SERVER_VAD_FIELDS = { - "type", - "threshold", - "prefix_padding_ms", - "silence_duration_ms", - "create_response", - "interrupt_response", - "min_speech_duration_ms", -} - - -@dataclass(frozen=True, slots=True) -class ServerVADConfig: - """Validated OpenAI-compatible ``server_vad`` session configuration.""" - - type: Literal["server_vad"] = "server_vad" - threshold: float = 0.5 - prefix_padding_ms: int = 300 - silence_duration_ms: int = 500 - create_response: bool = True - # Omission is resolved from runtime capabilities before the session is opened. - interrupt_response: bool | None = None - min_speech_duration_ms: int | None = None - - @classmethod - def from_value(cls, value: object) -> ServerVADConfig: - if not isinstance(value, dict): - raise ValueError("turn_detection must be null or an object") - unknown = sorted(set(value) - _SERVER_VAD_FIELDS) - if unknown: - raise ValueError(f"Unknown server_vad field(s): {', '.join(unknown)}") - vad_type = value.get("type") - if vad_type != "server_vad": - raise ValueError("turn_detection.type must be 'server_vad'") - - threshold = value.get("threshold", 0.5) - if isinstance(threshold, bool) or not isinstance(threshold, int | float) or not 0 <= threshold <= 1: - raise ValueError("server_vad.threshold must be a number between 0 and 1") - - prefix_padding_ms = value.get("prefix_padding_ms", 300) - if isinstance(prefix_padding_ms, bool) or not isinstance(prefix_padding_ms, int) or prefix_padding_ms < 0: - raise ValueError("server_vad.prefix_padding_ms must be a non-negative integer") - - silence_duration_ms = value.get("silence_duration_ms", 500) - if ( - isinstance(silence_duration_ms, bool) - or not isinstance(silence_duration_ms, int) - or silence_duration_ms <= 0 - ): - raise ValueError("server_vad.silence_duration_ms must be a positive integer") - - create_response = value.get("create_response", True) - if not isinstance(create_response, bool): - raise ValueError("server_vad.create_response must be a boolean") - - interrupt_response = value.get("interrupt_response") - if "interrupt_response" in value and not isinstance(interrupt_response, bool): - raise ValueError("server_vad.interrupt_response must be a boolean") - - min_speech_duration_ms = value.get("min_speech_duration_ms") - if min_speech_duration_ms is not None and ( - isinstance(min_speech_duration_ms, bool) - or not isinstance(min_speech_duration_ms, int | float) - or not np.isfinite(float(min_speech_duration_ms)) - or min_speech_duration_ms < 0 - ): - raise ValueError("server_vad.min_speech_duration_ms must be a non-negative number") - - return cls( - type=vad_type, - threshold=float(threshold), - prefix_padding_ms=prefix_padding_ms, - silence_duration_ms=silence_duration_ms, - create_response=create_response, - interrupt_response=interrupt_response, - min_speech_duration_ms=(int(min_speech_duration_ms) if min_speech_duration_ms is not None else None), - ) - - def as_dict(self) -> dict[str, object]: - result: dict[str, object] = { - "type": self.type, - "threshold": self.threshold, - "prefix_padding_ms": self.prefix_padding_ms, - "silence_duration_ms": self.silence_duration_ms, - "create_response": self.create_response, - } - if self.interrupt_response is not None: - result["interrupt_response"] = self.interrupt_response - if self.min_speech_duration_ms is not None: - result["min_speech_duration_ms"] = self.min_speech_duration_ms - return result - - -def parse_session_turn_detection( - payload: Mapping[str, object], -) -> tuple[bool, ServerVADConfig | None]: - """Extract and validate OpenAI-compatible turn detection aliases.""" - configured_values: list[tuple[str, object]] = [] - if "turn_detection" in payload: - configured_values.append(("turn_detection", payload["turn_detection"])) - audio = payload.get("audio") - if isinstance(audio, Mapping): - audio_input = audio.get("input") - if isinstance(audio_input, Mapping) and "turn_detection" in audio_input: - configured_values.append(("audio.input.turn_detection", audio_input["turn_detection"])) - if not configured_values: - return False, None - - first_path, first_value = configured_values[0] - first_config = None if first_value is None else ServerVADConfig.from_value(first_value) - for field_path, value in configured_values[1:]: - config = None if value is None else ServerVADConfig.from_value(value) - if first_config is not None and config is not None: - # An omitted option must not conflict with an explicit alias value. - if first_config.interrupt_response is None: - first_config = replace(first_config, interrupt_response=config.interrupt_response) - if config.interrupt_response is None: - config = replace(config, interrupt_response=first_config.interrupt_response) - if config != first_config: - raise ValueError(f"{first_path} and {field_path} must not specify conflicting values") - return True, first_config - - -class SpeechDetectorBackend(Protocol): - """Detector contract; per-stream state is owned by the pipeline.""" - - frame_samples: int - - def new_state(self) -> object: ... - - def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: ... - - -class SpeechDetectorBackendProvider(Protocol): - def get(self) -> SpeechDetectorBackend: ... - - -@dataclass(frozen=True, slots=True) -class SpeechEndpointDecision: - speech_started: bool = False - speech_stopped: bool = False - audio_start_ms: int | None = None - audio_end_ms: int | None = None - endpoint_delay_ms: int | None = None - - -class SpeechEndpointPolicy(Protocol): - @property - def speech_active(self) -> bool: ... - - def update( - self, - probability: float, - *, - frame_start_sample: int, - frame_samples: int, - ) -> SpeechEndpointDecision: ... - - def reset(self) -> None: ... - - -class ThresholdEndpointPolicy: - """Apply threshold and trailing-silence endpoint rules.""" - - def __init__(self, config: ServerVADConfig, *, sample_rate_hz: int) -> None: - self.config = config - self.sample_rate_hz = sample_rate_hz - self._speech_active = False - self._candidate_samples = 0 - self._candidate_start_sample: int | None = None - self._silence_samples = 0 - - @property - def speech_active(self) -> bool: - return self._speech_active - - def update( - self, - probability: float, - *, - frame_start_sample: int, - frame_samples: int, - ) -> SpeechEndpointDecision: - if probability >= self.config.threshold: - self._silence_samples = 0 - if not self._speech_active: - if self._candidate_samples == 0: - self._candidate_start_sample = frame_start_sample - self._candidate_samples += frame_samples - min_speech_duration_ms = self.config.min_speech_duration_ms or 0 - min_speech_samples = max(1, round(min_speech_duration_ms * self.sample_rate_hz / 1000)) - if self._candidate_samples < min_speech_samples: - return SpeechEndpointDecision() - self._speech_active = True - detected_start_sample = ( - self._candidate_start_sample if self._candidate_start_sample is not None else frame_start_sample - ) - detected_start_ms = round(detected_start_sample * 1000 / self.sample_rate_hz) - self._candidate_samples = 0 - self._candidate_start_sample = None - return SpeechEndpointDecision( - speech_started=True, - audio_start_ms=max(0, detected_start_ms - self.config.prefix_padding_ms), - ) - return SpeechEndpointDecision() - - if not self._speech_active: - self._candidate_samples = 0 - self._candidate_start_sample = None - return SpeechEndpointDecision() - - # Match Silero v6.2's streaming VAD hysteresis: activation uses the - # configured threshold, while potential speech end starts only below - # ``max(threshold - 0.15, 0.01)``. Once a silence candidate exists, - # intermediate frames keep elapsed time moving but cannot themselves - # close the turn. - negative_threshold = max(self.config.threshold - SILERO_VAD_MIN_THRESHOLD, 0.01) - below_negative_threshold = probability < negative_threshold - if not below_negative_threshold and self._silence_samples == 0: - return SpeechEndpointDecision() - - self._silence_samples += frame_samples - if not below_negative_threshold: - return SpeechEndpointDecision() - silence_limit = max(1, round(self.config.silence_duration_ms * self.sample_rate_hz / 1000)) - if self._silence_samples < silence_limit: - return SpeechEndpointDecision() - - # OpenAI defines ``audio_end_ms`` as the end of the audio sent to the - # model, including the trailing silence used for endpoint detection. - audio_end_sample = frame_start_sample + frame_samples - endpoint_delay_ms = round(self._silence_samples * 1000 / self.sample_rate_hz) - self.reset() - return SpeechEndpointDecision( - speech_stopped=True, - audio_end_ms=max(0, round(audio_end_sample * 1000 / self.sample_rate_hz)), - endpoint_delay_ms=endpoint_delay_ms, - ) - - def reset(self) -> None: - self._speech_active = False - self._candidate_samples = 0 - self._candidate_start_sample = None - self._silence_samples = 0 - - -@dataclass(frozen=True, slots=True) -class ServerVADFrame: - samples: np.ndarray - decision: SpeechEndpointDecision - - -@dataclass(frozen=True, slots=True) -class ServerVADBatch: - frames: tuple[ServerVADFrame, ...] - inference_ms: float - - -class ServerVADPipeline: - """Frame continuous 16 kHz audio and apply acoustic endpointing. - - The pipeline retains only input-resampling residuals, a partial frame, - detector state, and timing counters. Commit-eligible prefix and utterance - audio remain session-owned. - """ - - sample_rate_hz = 16_000 - - def __init__( - self, - backend: SpeechDetectorBackend, - config: ServerVADConfig, - *, - endpoint_policy: SpeechEndpointPolicy | None = None, - ) -> None: - self.backend = backend - self.config = config - self.endpoint_policy = endpoint_policy or ThresholdEndpointPolicy( - config, - sample_rate_hz=self.sample_rate_hz, - ) - self._input_resampler: StreamingAudioResampler | None = None - self._source_sample_rate_hz: int | None = None - self._scratch = np.empty(0, dtype=np.float32) - self._detector_state = backend.new_state() - self._processed_samples = 0 - self._stream_start_ms = 0 - - @property - def speech_active(self) -> bool: - return self.endpoint_policy.speech_active - - @property - def source_sample_rate_hz(self) -> int | None: - return self._source_sample_rate_hz - - @property - def scratch_bytes(self) -> int: - resampler_bytes = self._input_resampler.scratch_bytes if self._input_resampler is not None else 0 - return int(self._scratch.nbytes) + resampler_bytes - - async def push_pcm16( - self, - samples: bytes, - *, - source_sample_rate_hz: int, - ) -> ServerVADBatch: - """Normalize chunked mono PCM16 input and run endpoint detection.""" - if isinstance(source_sample_rate_hz, bool) or source_sample_rate_hz not in {16_000, 24_000}: - raise ValueError("server_vad PCM16 input sample rate must be 16000 or 24000 Hz") - - if len(samples) % np.dtype(" ServerVADBatch: - normalized = np.ascontiguousarray(samples, dtype=np.float32).reshape(-1) - return await asyncio.to_thread(self._push_sync, normalized) - - def _push_sync(self, samples: np.ndarray) -> ServerVADBatch: - started_at = time.perf_counter() - if self._scratch.size: - samples = np.concatenate((self._scratch, samples)) - frame_samples = int(self.backend.frame_samples) - complete_samples = samples.size - samples.size % frame_samples - # Retained state must not keep the whole input chunk alive through a view. - self._scratch = samples[complete_samples:].copy() - - results: list[ServerVADFrame] = [] - for offset in range(0, complete_samples, frame_samples): - frame = np.ascontiguousarray(samples[offset : offset + frame_samples], dtype=np.float32) - frame_start = self._processed_samples - probability, self._detector_state = self.backend.infer(frame, self._detector_state) - probability = min(1.0, max(0.0, float(probability))) - decision = self.endpoint_policy.update( - probability, - frame_start_sample=frame_start, - frame_samples=frame_samples, - ) - if decision.speech_started and decision.audio_start_ms is not None: - # Prefix audio before the most recent reset is no longer available. - decision = replace(decision, audio_start_ms=max(self._stream_start_ms, decision.audio_start_ms)) - - self._processed_samples += frame_samples - results.append( - ServerVADFrame( - samples=frame, - decision=decision, - ) - ) - return ServerVADBatch( - frames=tuple(results), - inference_ms=(time.perf_counter() - started_at) * 1000, - ) - - def reset(self) -> None: - """Discard stream state while preserving the session-relative audio clock.""" - self._stream_start_ms = round(self._processed_samples * 1000 / self.sample_rate_hz) - self._input_resampler = None - self._source_sample_rate_hz = None - self._scratch = np.empty(0, dtype=np.float32) - self._detector_state = self.backend.new_state() - self.endpoint_policy.reset() - - -@dataclass(frozen=True, slots=True) -class _SileroVADState: - model_state: np.ndarray - context: np.ndarray - - -class SileroVADBackend: - """Shared ONNX Runtime Silero v6.2 detector running on CPU.""" - - sample_rate_hz = 16_000 - frame_samples = 512 - context_samples = 64 - model_state_shape = (2, 1, 128) - - def __init__(self, model_path: str | Path) -> None: - try: - import onnxruntime as ort - except ImportError as exc: # pragma: no cover - platform packaging supplies ORT. - raise RuntimeError("server_vad requires ONNX Runtime") from exc - - session_options = ort.SessionOptions() - session_options.inter_op_num_threads = 1 - session_options.intra_op_num_threads = 1 - - self.model_path = Path(model_path) - self._session = ort.InferenceSession( - str(self.model_path), - providers=["CPUExecutionProvider"], - sess_options=session_options, - ) - input_names = {item.name for item in self._session.get_inputs()} - if "input" not in input_names or "sr" not in input_names or "state" not in input_names: - raise RuntimeError(f"Unsupported Silero ONNX input contract: {sorted(input_names)}") - self._warm_up() - - def _warm_up(self) -> None: - self.infer(np.zeros(self.frame_samples, dtype=np.float32), self.new_state()) - - def new_state(self) -> _SileroVADState: - return _SileroVADState( - model_state=np.zeros(self.model_state_shape, dtype=np.float32), - context=np.zeros((1, self.context_samples), dtype=np.float32), - ) - - def infer(self, frame: np.ndarray, state: object) -> tuple[float, object]: - if not isinstance(state, _SileroVADState): - raise TypeError("Silero detector state must be created by SileroVADBackend.new_state()") - model_state = np.ascontiguousarray(state.model_state, dtype=np.float32) - context = np.ascontiguousarray(state.context, dtype=np.float32) - if model_state.shape != self.model_state_shape: - raise ValueError(f"Silero model state must have shape {self.model_state_shape}, got {model_state.shape}") - expected_context_shape = (1, self.context_samples) - if context.shape != expected_context_shape: - raise ValueError(f"Silero context must have shape {expected_context_shape}, got {context.shape}") - - audio = np.ascontiguousarray(frame, dtype=np.float32).reshape(1, -1) - if audio.shape[1] != self.frame_samples: - raise ValueError( - f"Silero detector frame must contain exactly {self.frame_samples} samples, got {audio.shape[1]}" - ) - model_input = np.concatenate((context, audio), axis=1) - inputs = { - "input": model_input, - "state": model_state, - "sr": np.asarray(self.sample_rate_hz, dtype=np.int64), - } - # ONNX Runtime permits concurrent Run calls on one CPU session. Stream - # state is explicit, so independent pipelines can safely share the session. - output = self._session.run(None, inputs) - if len(output) < 2: - raise RuntimeError("Silero ONNX model did not return probability and model state") - probability = float(np.asarray(output[0]).reshape(-1)[0]) - next_state = _SileroVADState( - model_state=np.ascontiguousarray(output[1], dtype=np.float32), - context=np.ascontiguousarray(model_input[:, -self.context_samples :], dtype=np.float32), - ) - return probability, next_state - - -class SileroVADBackendProvider: - """Resolve, verify, and load one Silero model instance per process.""" - - def __init__(self, *, model_path: str | None = None) -> None: - self.model_path = model_path - self._backend: SileroVADBackend | None = None - self._lock = threading.Lock() - - def get(self) -> SileroVADBackend: - if self._backend is not None: - return self._backend - with self._lock: - if self._backend is None: - path = self._resolve_local_artifact() - self._verify_checksum(path) - self._backend = SileroVADBackend(path) - return self._backend - - def _resolve_local_artifact(self) -> Path: - if self.model_path: - path = Path(self.model_path).expanduser() - if path.is_file(): - return path - raise RuntimeError(f"Configured Silero VAD model does not exist: {path}") - cached = try_get_local_file( - model=SILERO_VAD_REPO_ID, - file_name=SILERO_VAD_FILENAME, - revision=SILERO_VAD_REVISION, - ) - if isinstance(cached, Path) and cached.is_file(): - return cached - raise RuntimeError( - "Silero VAD artifact is not available locally. Pre-download the pinned artifact " - "or configure duplex_session.server_vad_model_path." - ) - - def _verify_checksum(self, path: Path) -> None: - digest = hashlib.sha256(path.read_bytes()).hexdigest() - if digest != SILERO_VAD_SHA256: - raise RuntimeError(f"Silero VAD checksum mismatch for {path}: expected {SILERO_VAD_SHA256}, got {digest}") diff --git a/vllm_omni/entrypoints/duplex/serving.py b/vllm_omni/entrypoints/duplex/serving.py index 7bd690841..364f715c0 100644 --- a/vllm_omni/entrypoints/duplex/serving.py +++ b/vllm_omni/entrypoints/duplex/serving.py @@ -1,2130 +1,543 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Thin websocket serving for ``/v1/realtime?duplex=1``. + +Everything about a duplex *session* lives inside the engine (see +``vllm_omni.entrypoints.duplex_omni``). This handler only does transport: + +* accept the socket and run the Realtime handshake (``session.update`` opens a + session through ``DuplexOmni.open_session``; ``session.resume`` re-attaches + to an existing one); +* translate each client JSON event into a ``DuplexCommand`` and submit it in + arrival order; +* pump the session's typed events to the wire, journaling them for resume; +* manage attachments: resume tokens, replay, takeover, detach on disconnect. +""" + from __future__ import annotations import asyncio -import base64 -import binascii -import inspect import json -from collections.abc import Mapping, MutableMapping +from collections.abc import Mapping from contextlib import suppress from dataclasses import dataclass, replace -from typing import TYPE_CHECKING -import numpy as np -from fastapi import WebSocket +from fastapi import WebSocket, WebSocketDisconnect from vllm.logger import init_logger -from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig -from vllm_omni.engine.duplex.messages import DuplexFence, DuplexSessionLifecycleMessage -from vllm_omni.engine.duplex.runtime import duplex_resource_request_id -from vllm_omni.entrypoints.duplex.capability import ( - should_enable_duplex_endpoint, -) -from vllm_omni.entrypoints.duplex.chat_fallback import ( - ChatFallbackProjectorMixin, -) -from vllm_omni.entrypoints.duplex.protocol import ( - NATIVE_DUPLEX_KEY, - DuplexCapabilities, - DuplexCommittedInput, - DuplexOverlapPolicy, - DuplexPlaybackCommitPolicy, - DuplexSession, - DuplexSessionConfig, - DuplexSessionRegistry, - DuplexSessionState, - DuplexTurnController, - ResponseCreateOptions, - native_duplex_opt_in, - normalize_native_duplex_key, -) -from vllm_omni.entrypoints.duplex.realtime_session import ( - REALTIME_OUTPUT_AUDIO_FORMATS, - NativeRealtimeSessionProtocol, -) -from vllm_omni.entrypoints.duplex.runtime_adapter import ( - ServingRuntimeAdapter, - ServingRuntimeConfigError, - ServingRuntimeSessionState, - TurnBasedServingSessionState, - load_serving_runtime_adapter, - validate_serving_runtime_adapter, -) -from vllm_omni.entrypoints.duplex.runtime_bridge import ( - NativeRuntimeBridgeMixin, -) -from vllm_omni.entrypoints.duplex.server_vad import ( - SILERO_VAD_DEFAULT_MIN_SPEECH_DURATION_MS, - SILERO_VAD_MIN_THRESHOLD, - ServerVADPipeline, - SileroVADBackendProvider, - SpeechDetectorBackendProvider, - parse_session_turn_detection, +from vllm_omni.engine.duplex.commands import DuplexCommand, DuplexCommandError +from vllm_omni.engine.duplex.events import ( + DuplexEvent, + SessionClosed, + SessionCreated, + SessionReplaced, + SessionResumed, + SessionResyncRequired, ) +from vllm_omni.engine.duplex.messages import DuplexSessionError +from vllm_omni.engine.duplex.realtime_commands import RealtimeInputDefaults +from vllm_omni.entrypoints.duplex.realtime_input import RealtimeEnvelope, parse_resume_request from vllm_omni.entrypoints.duplex.session_attachment import ( DuplexJournalGapError, + DuplexJournalOverflowError, DuplexSessionAttachmentRegistry, InvalidResumeTokenError, -) -from vllm_omni.entrypoints.duplex.session_runner import ( - DuplexSessionRunnerMixin, + ResumeToken, ) from vllm_omni.entrypoints.duplex.websocket import ( - DOMAIN_TERMINAL_EVENTS, - DuplexSessionTasks, - DuplexWebSocketActor, + MAX_EVENT_BYTES, + SendJson, + attachment_callbacks, + receive_text_with_timeout, ) -from vllm_omni.metrics.realtime import RealtimeVADMetrics - -if TYPE_CHECKING: - from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat +from vllm_omni.entrypoints.duplex_omni import DuplexOmni, DuplexSessionHandle logger = init_logger(__name__) -__all__ = ["OmniDuplexSessionHandler", "should_enable_duplex_endpoint"] +__all__ = ["OmniDuplexSessionHandler"] _DEFAULT_CONFIG_TIMEOUT_S = 10.0 _DEFAULT_IDLE_TIMEOUT_S = 300.0 +#: How long the endpoint waits for the session pump to deliver its terminal +#: event once the session is closed. Bounded so a wedged pump cannot pin the +#: connection open; the pump itself closes the socket on its way out. +_PUMP_DRAIN_TIMEOUT_S = 5.0 +#: session.created is journaled by the registry when it hands out the resume credential. +_UNJOURNALED_EVENTS = (SessionCreated,) + + +@dataclass +class _Attachment: + """What one websocket connection knows about the session it is attached to.""" + + handle: DuplexSessionHandle + generation: int @dataclass(frozen=True) -class _DuplexSessionHandshake: - session: DuplexSession - resumed: bool = False - attachment_generation: int | None = None - event_id: str | None = None +class _ResumeCredentials: + """Transport credentials stamped onto ``session.created`` for a resumable session.""" + + attachment_generation: int + resume_token: str -class OmniDuplexSessionHandler( - DuplexSessionRunnerMixin, - NativeRuntimeBridgeMixin, - ChatFallbackProjectorMixin, -): - """WebSocket handler for RFC-style full-duplex session control. - - This owns the serving-side session actor, ordered inbound mailbox, - barge-in epoch, turn controller, and playback commit state. Generic sessions - can still fall back to chat requests, while MiniCPM-o 4.5 native sessions - route audio appends through scheduler data-plane stage requests. It - deliberately does not claim core persistent KV lease support. - """ +class OmniDuplexSessionHandler: + """WebSocket transport for engine-resident duplex sessions.""" def __init__( self, *, - chat_service: OmniOpenAIServingChat, - served_model_name: str | None = None, + duplex_omni: DuplexOmni, config_timeout_s: float = _DEFAULT_CONFIG_TIMEOUT_S, idle_timeout_s: float = _DEFAULT_IDLE_TIMEOUT_S, - log_stats: bool = True, - duplex_session_config: DuplexSessionRuntimeConfig | None = None, - serving_runtime_adapter: ServingRuntimeAdapter | None = None, - serving_runtime_adapter_path: str | None = None, - server_vad_backend_provider: SpeechDetectorBackendProvider | None = None, ) -> None: - self._chat_service = chat_service + self._omni = duplex_omni self._config_timeout_s = config_timeout_s self._idle_timeout_s = idle_timeout_s - self._duplex_session_config = duplex_session_config or DuplexSessionRuntimeConfig() - self._server_vad_backend_provider = server_vad_backend_provider or SileroVADBackendProvider( - model_path=self._duplex_session_config.server_vad_model_path, - ) - self._server_vad_pipelines: dict[str, ServerVADPipeline] = {} - self._realtime_vad_metrics = RealtimeVADMetrics( - served_model_name or chat_service.model_config.model, - log_stats=log_stats, - ) - self._registry = DuplexSessionRegistry( - DuplexCapabilities( - supports_model_native_turn_policy=False, - supports_input_append=False, - supports_replace_latest_chunk=True, - supports_reencode_context=True, - supports_turn_commit_only=True, - supports_kv_lease=False, - ) - ) - self._turn_controller = DuplexTurnController() - adapter_path = serving_runtime_adapter_path or getattr( - chat_service, - "duplex_serving_adapter_path", - None, - ) - self._serving_runtime_adapter: ServingRuntimeAdapter | None - if serving_runtime_adapter is not None: - self._serving_runtime_adapter = validate_serving_runtime_adapter(serving_runtime_adapter) - elif isinstance(adapter_path, str) and adapter_path: - self._serving_runtime_adapter = load_serving_runtime_adapter( - adapter_path, - self._encode_native_data_plane_audio, - ) - else: - self._serving_runtime_adapter = None - self._serving_session_states: MutableMapping[str, ServingRuntimeSessionState] = ( - self._serving_runtime_adapter.session_states if self._serving_runtime_adapter is not None else {} - ) - self._session_tasks: dict[str, DuplexSessionTasks] = {} - self._realtime_protocols: dict[str, NativeRealtimeSessionProtocol] = {} - self._lease_generations: dict[str, int] = {} - self._resync_required_sessions: set[str] = set() - self._lifecycle_queue = getattr(self._chat_service.engine_client, "duplex_lifecycle_events", None) - self._lifecycle_task: asyncio.Task[None] | None = None + runtime_config = duplex_omni.duplex_session_config self._attachment_registry = DuplexSessionAttachmentRegistry( - replay_ttl_s=self._duplex_session_config.resume_replay_ttl_s, - replay_max_bytes_per_session=self._duplex_session_config.resume_replay_max_bytes_per_session, - disconnect_grace_s=self._duplex_session_config.disconnect_grace_s, + replay_ttl_s=runtime_config.resume_replay_ttl_s, + replay_max_bytes_per_session=runtime_config.resume_replay_max_bytes_per_session, ) + self._resync_required_sessions: set[str] = set() + #: Wire defaults (input/output audio format and rate) per session. They + #: are negotiated on the session but live on the per-connection + #: envelope, so a reconnect has to be handed them back. + self._input_defaults: dict[str, RealtimeInputDefaults] = {} + self._pumps: dict[str, asyncio.Task[None]] = {} + + # ------------------------------------------------------------------ # + # Entry point # + # ------------------------------------------------------------------ # async def handle_realtime_session(self, websocket: WebSocket) -> None: - await self.handle_session( - websocket, - realtime_protocol=NativeRealtimeSessionProtocol(websocket.query_params), - ) + await websocket.accept() + envelope = RealtimeEnvelope.from_query_params(websocket.query_params) - def _ensure_lifecycle_listener(self) -> None: - if not isinstance(self._lifecycle_queue, asyncio.Queue): - return - if self._lifecycle_task is None or self._lifecycle_task.done(): - self._lifecycle_task = asyncio.create_task( - self._run_lifecycle_listener(), - name="duplex-serving-lifecycle", - ) + async def send_json(payload: Mapping[str, object]) -> None: + await websocket.send_json(dict(payload)) - async def _run_lifecycle_listener(self) -> None: - queue = self._lifecycle_queue - if not isinstance(queue, asyncio.Queue): - return + attachment: _Attachment | None = None + pending_command: dict[str, object] | None = None try: - while True: - message = await queue.get() - try: - if isinstance(message, DuplexSessionLifecycleMessage): - await self._apply_runtime_lifecycle(message) - finally: - queue.task_done() - if self._registry.active_count() == 0: + open_payload = envelope.initial_open_payload() + if open_payload is None: + first = await self._receive_first_message(websocket, envelope, send_json) + if first is None: return - finally: - if self._lifecycle_task is asyncio.current_task(): - self._lifecycle_task = None - - async def _apply_runtime_lifecycle(self, message: DuplexSessionLifecycleMessage) -> None: - session = self._registry.get(message.session_id) - if session is None: - return - fence = message.fence - if ( - fence.session_id != session.session_id - or fence.incarnation != session.incarnation - or fence.epoch < session.epoch - ): - return - current_generation = self._lease_generations.get(session.session_id, 0) - if message.lease_generation < current_generation: - return - protocol = self._realtime_protocols.get(session.session_id) - expired_payload: dict[str, object] = { - "type": "session.expired", - "session_id": session.session_id, - "incarnation": session.incarnation, - "reason": message.reason, - } - if protocol is not None: - expired_payload = protocol.encode_outbound_event(expired_payload)[0] - with suppress(Exception): - await self._attachment_registry.send_event( - session.session_id, - expired_payload, - journal=False, - ) - - tasks = self._session_tasks.pop(session.session_id, None) - if tasks is not None: - await tasks.cancel_append_tasks() - active_response_task = tasks.active_response_task - if active_response_task is not None and not active_response_task.done(): - active_response_task.cancel() - with suppress(asyncio.CancelledError): - await active_response_task - native = self._serving_session_states.get(session.session_id) - if native is not None and native.data_plane_task is not None: - data_plane_task = native.data_plane_task - native.data_plane_task = None - data_plane_task.cancel() - with suppress(asyncio.CancelledError): - await data_plane_task - if session.active_response_id is not None: - session.end_response(commit_text=False) - self._cleanup_duplex_session_state(session) - self._registry.close(session.session_id) - self._realtime_protocols.pop(session.session_id, None) - self._lease_generations.pop(session.session_id, None) - self._resync_required_sessions.discard(session.session_id) - attachment = await self._attachment_registry.close(session.session_id) - if attachment is not None: + handshake = envelope.first_message(first) + if handshake.kind == "resume": + attachment = await self._resume(websocket, envelope, handshake.resume_payload, send_json) + if attachment is None: + return + else: + open_payload = handshake.session_payload + pending_command = handshake.pending_command_payload + if attachment is None: + assert open_payload is not None + attachment = await self._open(websocket, envelope, open_payload, send_json) + if attachment is None: + return + if pending_command is not None: + await self._submit_wire_event(attachment, envelope, pending_command, send_json) + self._input_defaults[attachment.handle.session_id] = envelope.defaults + await self._read_loop(websocket, envelope, attachment, send_json) + await self._drain_terminal_pump(attachment) + except WebSocketDisconnect: + if attachment is not None: + await self._on_disconnect(attachment) + except Exception as exc: + logger.exception("Duplex websocket session failed: %s", exc) + if attachment is not None and await self._attachment_registry.is_current_attachment( + attachment.handle.session_id, attachment.generation + ): + with suppress(Exception): + await self._attachment_registry.close(attachment.handle.session_id) + with suppress(Exception): + await attachment.handle.close(reason="transport_error") with suppress(Exception): - await attachment.close("session_expired") + await send_json(envelope.error_payload("realtime_input_failed", str(exc))) - def _stop_lifecycle_listener_if_idle(self) -> None: - task = self._lifecycle_task - if self._registry.active_count() != 0 or task is None or task.done(): - return - if task is not asyncio.current_task(): - task.cancel() - self._lifecycle_task = None - - @staticmethod - def _native_audio_payload_size_bytes(payload: Mapping[str, object]) -> int: - audio = payload.get("audio") or payload.get("data") - if not isinstance(audio, str): - return 0 + async def _receive_first_message( + self, websocket: WebSocket, envelope: RealtimeEnvelope, send_json: SendJson + ) -> dict[str, object] | None: + raw = await receive_text_with_timeout(websocket, self._config_timeout_s) + if raw is None: + await send_json(envelope.error_payload("config_timeout", "Timeout waiting for session.update")) + return None try: - return len(base64.b64decode(audio, validate=True)) - except (ValueError, binascii.Error): - return 0 + first = json.loads(raw) + except json.JSONDecodeError: + await send_json(envelope.error_payload("invalid_json", "Invalid JSON in session.update")) + return None + if not isinstance(first, dict): + await send_json(envelope.error_payload("bad_event", "Realtime event must be a JSON object")) + return None + return first - async def _apply_outbound_session_event( - self, - payload: dict[str, object], - *, - session: DuplexSession | None, - actor: DuplexWebSocketActor, - native: ServingRuntimeSessionState, - realtime_protocol: NativeRealtimeSessionProtocol | None, - ) -> tuple[bool, dict[str, object] | None]: - """Apply domain transitions before an event is queued for transport.""" - payload_type = payload.get("type") - is_terminal = payload_type in DOMAIN_TERMINAL_EVENTS - if is_terminal and session is not None: - payload_epoch = payload.get("epoch") - if isinstance(payload_epoch, int) and payload_epoch != session.epoch: - return False, None - if payload_type in {"response.done", "response.listen"} and ( - actor.closing or session.state == DuplexSessionState.CLOSED - ): - return False, None - elif actor._is_stale_model_output(payload): - return False, None + # ------------------------------------------------------------------ # + # Handshake # + # ------------------------------------------------------------------ # - if payload_type == "session.closed": - actor.close_reason = actor.close_reason or str(payload.get("reason") or "closed") - if session is not None: - session.mark_closing() - - if not is_terminal or session is None: - return True, None - - terminal_status = payload.get("status") - terminal_status_details = payload.get("status_details") - if terminal_status is None and isinstance(terminal_status_details, dict): - terminal_status = terminal_status_details.get("type") - response_terminal = payload_type == "response.done" or ( - payload_type == "response.listen" and session.active_response_id is not None - ) - can_promote_overlap = response_terminal and terminal_status not in {"cancelled", "failed"} - deferred_overlap_payload: dict[str, object] | None = None - continuous_input_crosses_terminal = ( - can_promote_overlap - and realtime_protocol is not None - and self._session_auto_responds(session) - and native.input_since_commit - and not native.deferred_response_create - ) - if continuous_input_crosses_terminal: - # Auto-response is a model-owned continuous stream. A response - # terminal advances the model turn without requiring the browser - # to close its Realtime input item. Keep any partial model unit so - # later PCM can complete it. Serving does not create or promote a - # new model generation here; Stage0 advances generation identity - # only after the model's terminal token. - session.reset_overlap_speech() - return True, None - realtime_input_still_open = ( - can_promote_overlap - and realtime_protocol is not None - and native.input_since_commit - and not native.deferred_response_create - ) - if realtime_input_still_open: - # response.done only closes the assistant response. It does not - # close the current Realtime input item, and it does not reserve - # or advance a model generation. - session.reset_overlap_speech() - return True, None - if can_promote_overlap and session.overlap_speech_ms > 0: - has_deferred_overlap = native.audio_buffer.has_pending() or native.committed_audio_payload is not None - should_promote_overlap = ( - session.state == DuplexSessionState.OPEN - and self._uses_native_input_append(session) - and has_deferred_overlap - and session.overlap_speech_ms > session.config.overlap_short_ack_ms - ) - if should_promote_overlap: - flushed_reserved_bytes = native.audio_buffer.pending_byte_count - deferred_overlap_payload = native.audio_buffer.flush( - chunk_period_ms=session.capabilities.chunk_period_ms or 1000 - ) - if native.committed_audio_payload is not None: - if deferred_overlap_payload is not None: - deferred_overlap_payload = self._merge_native_audio_payloads( - native.committed_audio_payload, - deferred_overlap_payload, - ) - else: - deferred_overlap_payload = native.committed_audio_payload - if self._session_auto_responds(session) and deferred_overlap_payload is not None: - deferred_overlap_payload = dict(deferred_overlap_payload) - deferred_overlap_payload["force_listen"] = False - if deferred_overlap_payload is not None: - native.retain_committed_audio( - deferred_overlap_payload, - operation_id=native.committed_audio_operation_id, - reserved_bytes=flushed_reserved_bytes, - ) - native.input_since_commit = deferred_overlap_payload is not None - if realtime_protocol is not None: - native.deferred_response_create = True - native.deferred_precreate_response = False - deferred_overlap_payload = None - else: - had_pending_overlap_audio = native.audio_buffer.has_pending() - native.audio_buffer.clear() - native.input_since_commit = False - native.speech_since_commit = False - if had_pending_overlap_audio and realtime_protocol is not None: - await realtime_protocol.discard_pending_input_audio(audio_end_ms=session.overlap_speech_ms) - if payload_type in {"audio.cancelled", "input.cancelled", "session.closed"}: - session.release_input_bytes(native.clear_committed_audio()) - - session.reset_overlap_speech() - if can_promote_overlap and native.deferred_response_create and native.committed_audio_payload is not None: - deferred_overlap_payload = native.committed_audio_payload - native.deferred_response_create = False - native.input_since_commit = False - native.speech_since_commit = False - return True, deferred_overlap_payload - - @staticmethod - def _advance_barge_in_epoch(session: DuplexSession) -> tuple[int, dict[str, int]]: - old_playback = session.playback.as_dict() - new_epoch = session.barge_in() - session.clear_playback_cursor() - return new_epoch, old_playback - - @staticmethod - def _commit_played_response_history( - session: DuplexSession, - response_id: str | None, - committed_ms: int, - ) -> None: - if not response_id or committed_ms < 0: - return - session.truncate_history_item( - f"item_{response_id}", - audio_end_ms=committed_ms, - playback=session.playback_for_response(response_id), - ) - - @staticmethod - def _should_commit_response_to_history(session: DuplexSession, response_id: str | None) -> bool: - if response_id is not None and response_id != session.active_response_id: - return True - mode = session.response_config.extra_body.get("realtime_response_conversation") - return not isinstance(mode, str) or mode.strip().lower() != "none" - - def _response_created_payload( - self, - session: DuplexSession, - response_id: str, - *, - epoch: int, - request_id: str | None = None, - ) -> dict[str, object]: - response_config = session.response_config - payload: dict[str, object] = { - "type": "response.created", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "modalities": list(response_config.modalities), - } - if request_id is not None: - payload["request_id"] = request_id - metadata = response_config.extra_body.get("realtime_response_metadata") - if not isinstance(metadata, dict): - metadata = response_config.extra_body.get("realtime_metadata") - if isinstance(metadata, dict): - payload["metadata"] = dict(metadata) - conversation = response_config.extra_body.get("realtime_response_conversation") - if isinstance(conversation, str): - payload["conversation"] = conversation - prompt = response_config.extra_body.get("realtime_response_prompt") - if isinstance(prompt, dict): - payload["prompt"] = dict(prompt) - return payload - - def _overlap_decision( - self, - session: DuplexSession, - event: dict[str, object], - payload: dict[str, object], - ) -> dict[str, object]: - """Classify input that arrives while assistant audio is active. - - This is a serving-side policy. The model still owns listen/speak - decisions for normal chunks; overlap policy only decides whether the - current assistant response should be interrupted before the new audio is - appended. - """ - duration_ms = self._input_audio_duration_ms(event, payload) - is_speech = self._input_looks_like_speech(event, payload, session=session) - if not session.capabilities.supports_barge_in and self._event_requests_barge_in(event): - return self._defer_unsupported_barge_in(session, duration_ms=duration_ms, is_speech=is_speech) - explicit = event.get("overlap_action") or event.get("overlap") - if isinstance(explicit, str): - normalized = explicit.strip().lower() - if normalized in {"barge_in", "interrupt", "cancel"}: - return { - "action": "barge_in", - "reason": "client_overlap_action", - "duration_ms": duration_ms, - "buffer_audio": True, - } - if normalized in {"listen", "continue", "continue_output", "ack"}: - session.reset_overlap_speech() - return { - "action": "listen", - "reason": "client_overlap_action", - "duration_ms": duration_ms, - "buffer_audio": ( - normalized == "listen" and is_speech and duration_ms > session.config.overlap_short_ack_ms - ), - "defer_runtime_append": True, - } - if normalized in {"drop", "ignore", "silence"}: - session.reset_overlap_speech() - return { - "action": "drop", - "reason": "client_overlap_action", - "duration_ms": duration_ms, - "buffer_audio": False, - } - - if bool(event.get("force_barge_in", False)): - return { - "action": "barge_in", - "reason": "client_force_barge_in", - "duration_ms": duration_ms, - "buffer_audio": True, - } - if self._session_auto_responds(session): - # Full-duplex input remains model-owned while output is active. Feed - # complete model units into the existing Stage0 stream immediately; - # playback only controls history ACKs, not model admission. - if is_speech: - session.accumulate_overlap_speech(duration_ms) - vad_speech_started = self._vad_speech_started(event, payload) - if ( - self._uses_native_input_append(session) - and session.capabilities.supports_barge_in - and is_speech - and session.config.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value - and vad_speech_started is not False - ): - return { - "action": "barge_in", - "reason": ("server_vad_speech_started" if vad_speech_started is True else "barge_in_on_speech"), - "cancel_reason": "turn_detected" if vad_speech_started is True else "barge_in", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - } - return { - "action": "listen", - "reason": "auto_response_continuous", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - "defer_runtime_append": False, - "force_listen": event.get("force_listen") is True or payload.get("force_listen") is True, - "preserve_realtime_input": True, - } - if bool(event.get("force_listen", False)): - session.reset_overlap_speech() - return { - "action": "listen", - "reason": "client_force_listen", - "duration_ms": duration_ms, - "buffer_audio": is_speech, - "defer_runtime_append": True, - } - - policy = session.config.overlap_policy - if not is_speech: - if session.overlap_speech_ms <= 0: - session.reset_overlap_speech() - return { - "action": "drop", - "reason": "silence_or_noise", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": False, - } - - if self._is_short_ack_transcript_hint(event, payload): - session.reset_overlap_speech() - return { - "action": "listen", - "reason": "short_ack_transcript", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": False, - "defer_runtime_append": True, - } - - if policy == DuplexOverlapPolicy.LISTEN_ONLY.value: - session.accumulate_overlap_speech(duration_ms) - return { - "action": "listen", - "reason": "policy_listen_only", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - "defer_runtime_append": True, - } - - if policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value and not session.capabilities.supports_barge_in: - return self._defer_unsupported_barge_in(session, duration_ms=duration_ms, is_speech=True) - - session.accumulate_overlap_speech(duration_ms) - if policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value: - vad_speech_started = self._vad_speech_started(event, payload) - if vad_speech_started is False: - return { - "action": "listen", - "reason": "server_vad_utterance_active", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - "defer_runtime_append": False, - "force_listen": True, - "preserve_realtime_input": True, - } - return { - "action": "barge_in", - "reason": ("server_vad_speech_started" if vad_speech_started is True else "policy_barge_in_on_speech"), - "cancel_reason": "turn_detected" if vad_speech_started is True else "barge_in", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - } - - if ( - duration_ms <= session.config.overlap_short_ack_ms - and session.overlap_speech_ms <= session.config.overlap_short_ack_ms - ): - return { - "action": "listen", - "reason": "short_ack", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - "defer_runtime_append": True, - } - if session.overlap_speech_ms >= session.config.overlap_barge_in_ms: - if not session.capabilities.supports_barge_in: - return { - "action": "listen", - "reason": "barge_in_unsupported", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - "defer_runtime_append": True, - } - return { - "action": "barge_in", - "reason": "long_overlap_speech", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - } - return { - "action": "listen", - "reason": "accumulating_overlap_speech", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": True, - "defer_runtime_append": True, - } - - @staticmethod - def _vad_speech_started( - event: Mapping[str, object], - payload: Mapping[str, object], - ) -> bool | None: - for source in (event, payload): - vad = source.get("vad") - if isinstance(vad, Mapping) and isinstance(vad.get("speech_started"), bool): - return bool(vad["speech_started"]) - return None - - @staticmethod - def _event_requests_barge_in(event: Mapping[str, object]) -> bool: - if event.get("force_barge_in") is True: - return True - explicit = event.get("overlap_action") or event.get("overlap") - return isinstance(explicit, str) and explicit.strip().lower() in { - "barge_in", - "interrupt", - "cancel", - } - - @staticmethod - def _defer_unsupported_barge_in( - session: DuplexSession, - *, - duration_ms: int, - is_speech: bool, - ) -> dict[str, object]: - if is_speech: - session.accumulate_overlap_speech(duration_ms) - return { - "action": "listen", - "reason": "barge_in_unsupported", - "duration_ms": duration_ms, - "overlap_speech_ms": session.overlap_speech_ms, - "buffer_audio": is_speech, - "defer_runtime_append": True, - } - - @staticmethod - def _barge_in_unsupported_error(session: DuplexSession) -> dict[str, object]: - return { - "type": "error", - "session_id": session.session_id, - "code": "barge_in_unsupported", - "error": "Barge-in is not supported by this duplex model", - } - - @staticmethod - def _is_short_ack_transcript_hint(event: dict[str, object], payload: dict[str, object]) -> bool: - raw_text = event.get("transcript") or event.get("text") or payload.get("transcript") or payload.get("text") - if not isinstance(raw_text, str): - return False - normalized = raw_text.strip().lower() - if not normalized: - return False - compact = "".join(ch for ch in normalized if ch.isalnum() or "\u4e00" <= ch <= "\u9fff") - if compact in { - "嗯", - "嗯嗯", - "对", - "对的", - "好", - "好的", - "继续", - "继续说", - "可以", - "是的", - "yes", - "yeah", - "yep", - "ok", - "okay", - "continue", - "goon", - "right", - }: - return True - return normalized in {"go on", "keep going", "please continue"} - - @staticmethod - def _input_audio_duration_ms(event: dict[str, object], payload: dict[str, object]) -> int: - for key in ("duration_ms", "audio_duration_ms"): - value = event.get(key) - if isinstance(value, int | float): - return max(0, int(value)) - fmt = payload.get("format") - sample_rate_hz = payload.get("sample_rate_hz") - audio = payload.get("audio") - if fmt == "pcm_f32le" and isinstance(sample_rate_hz, int) and sample_rate_hz > 0 and isinstance(audio, str): - try: - raw = base64.b64decode(audio, validate=True) - except (binascii.Error, ValueError): - return 0 - return int((len(raw) // 4) * 1000 / sample_rate_hz) - return 0 - - @staticmethod - def _merge_native_audio_payloads( - first: dict[str, object], - second: dict[str, object], - ) -> dict[str, object]: - if first.get("format") != "pcm_f32le" or second.get("format") != "pcm_f32le": - return second - first_rate = first.get("sample_rate_hz") - second_rate = second.get("sample_rate_hz") - if not isinstance(first_rate, int) or not isinstance(second_rate, int) or first_rate != second_rate: - return second - first_audio = first.get("audio") - second_audio = second.get("audio") - if not isinstance(first_audio, str) or not isinstance(second_audio, str): - return second - try: - first_raw = base64.b64decode(first_audio, validate=True) - second_raw = base64.b64decode(second_audio, validate=True) - except (binascii.Error, ValueError): - return second - merged = dict(second) - merged["audio"] = base64.b64encode(first_raw + second_raw).decode("ascii") - merged["sample_rate_hz"] = first_rate - merged_frames = [ - frame - for source in (first.get("video_frames"), second.get("video_frames")) - if isinstance(source, list) - for frame in source - if isinstance(frame, str) and frame - ] - if merged_frames: - merged["video_frames"] = merged_frames - else: - merged.pop("video_frames", None) - merged["force_listen"] = bool(first.get("force_listen", False)) or bool(second.get("force_listen", False)) - merged.pop("force_speak", None) - merged["is_speech"] = bool(first.get("is_speech", False)) or bool(second.get("is_speech", False)) - return merged - - @classmethod - def _should_force_listen_for_short_commit( - cls, - session: DuplexSession, - event: dict[str, object], - payload: dict[str, object], - ) -> bool: - """Keep very short committed Realtime chunks in listen mode. - - Realtime VAD can emit a commit for a short pause even when the user has - not actually yielded the turn. For MiniCPM-o native duplex, make that - policy explicit by steering the scheduler path to the model listen - token instead of letting a sub-second chunk start a response. - """ - if event.get("force_listen") is True or payload.get("force_listen") is True: - return True - if event.get("force_barge_in") is True: - return False - if event.get("response_create") is not True: - return False - duration_ms = cls._input_audio_duration_ms(event, payload) - return 0 < duration_ms <= session.config.overlap_short_ack_ms - - def _create_runtime_session_state(self) -> ServingRuntimeSessionState: - adapter = self._serving_runtime_adapter - if adapter is not None: - return adapter.create_session_state() - return TurnBasedServingSessionState() - - def _runtime_session_state(self, session: DuplexSession) -> ServingRuntimeSessionState: - adapter = self._serving_runtime_adapter - if adapter is not None: - return adapter.session_state(session.session_id) - state = self._serving_session_states.get(session.session_id) - if state is None: - state = TurnBasedServingSessionState() - self._serving_session_states[session.session_id] = state - return state - - def _require_serving_runtime_adapter(self) -> ServingRuntimeAdapter: - adapter = self._serving_runtime_adapter - if adapter is None: - raise RuntimeError("Model-native duplex runtime is not configured for this session") - return adapter - - def _should_force_listen_for_auto_response_overlap( - self, - session: DuplexSession, - event: dict[str, object], - payload: dict[str, object], - ) -> bool: - if not self._session_auto_responds(session): - return False - if event.get("force_barge_in") is True: - return False - return event.get("force_listen") is True or payload.get("force_listen") is True - - @staticmethod - def _assistant_playback_active(session: DuplexSession) -> bool: - return ( - session.config.playback_commit_policy == DuplexPlaybackCommitPolicy.ACK_ONLY.value - and session.playback.sent_ms > session.playback.committed_ms - ) - - @staticmethod - def _input_looks_like_speech( - event: dict[str, object], - payload: dict[str, object], - *, - session: DuplexSession, - ) -> bool: - # Prefer serving-derived fields in ``payload`` over optional client - # hints retained in ``event``. In server-VAD mode the pipeline, rather - # than the client or the RMS fallback below, owns speech classification. - for source in (payload, event): - for key in ("is_speech", "speech"): - value = source.get(key) - if isinstance(value, bool): - return value - vad = source.get("vad") - if isinstance(vad, dict): - value = vad.get("is_speech") - if isinstance(value, bool): - return value - probability = vad.get("speech_probability", vad.get("probability")) - if isinstance(probability, int | float): - return float(probability) >= 0.5 - probability = source.get("speech_probability") - if isinstance(probability, int | float): - return float(probability) >= 0.5 - - fmt = payload.get("format") - audio = payload.get("audio") - if fmt in {"pcm_f32le", "pcm16"} and isinstance(audio, str): - try: - raw = base64.b64decode(audio, validate=True) - except (binascii.Error, ValueError): - return True - if fmt == "pcm_f32le": - if len(raw) < 4 or len(raw) % 4 != 0: - return True - samples = np.frombuffer(raw, dtype=np.float32) - else: - if len(raw) < 2 or len(raw) % 2 != 0: - return True - samples = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 - if samples.size == 0: - return False - rms = float(np.sqrt(np.mean(np.square(samples.astype(np.float32))))) - return rms >= session.config.overlap_silence_rms - return True - - async def _emit_overlap_decision( - self, - send_json, - session: DuplexSession, - decision: dict[str, object], - ) -> None: - await send_json( - { - "type": "overlap.decision", - "session_id": session.session_id, - "epoch": session.epoch, - "policy": session.config.overlap_policy, - **decision, - } - ) - - async def _open_session( + async def _open( self, websocket: WebSocket, - send_json, - *, - realtime_protocol: NativeRealtimeSessionProtocol | None = None, - attachment_send=None, - attachment_close=None, - ) -> _DuplexSessionHandshake | None: - raw = await self._receive_text( - websocket, - self._config_timeout_s, - realtime_protocol=realtime_protocol, - ) - if raw is None: - await send_json({"type": "error", "error": "Timeout waiting for session.create", "code": "config_timeout"}) - return None + envelope: RealtimeEnvelope, + session_payload: Mapping[str, object], + send_json: SendJson, + ) -> _Attachment | None: + # Any ``session_id`` / ``id`` in the payload is ignored: the server + # allocates the id and announces it in ``session.created``. try: - event = json.loads(raw) - except json.JSONDecodeError: - await send_json({"type": "error", "error": "Invalid JSON in session.create", "code": "invalid_json"}) + handle = await self._omni.open_session(session_payload) + except DuplexSessionError as exc: + await send_json(envelope.error_payload(exc.code, str(exc))) return None - if isinstance(event, dict) and event.get("type") == "session.resume": - if realtime_protocol is None: - await send_json( - { - "type": "error", - "error": "session.resume requires the Realtime protocol", - "code": "unsupported_session_resume", - } - ) - return None - return await self._resume_session_handshake( - event, - send_json=send_json, - realtime_protocol=realtime_protocol, - attachment_send=attachment_send, - attachment_close=attachment_close, - ) - if not isinstance(event, dict) or event.get("type") not in {"session.create", "open_session", "session.config"}: - await send_json( - { - "type": "error", - "error": f"Expected session.create, got: {event.get('type') if isinstance(event, dict) else None}", - "code": "bad_event", - } - ) + except (TypeError, ValueError) as exc: + await send_json(envelope.error_payload("invalid_session_config", str(exc))) return None - - try: - config = DuplexSessionConfig.from_event(event) - except ValueError as exc: - error: dict[str, object] = { - "type": "error", - "error": str(exc), - "code": "invalid_request_error", - } - event_id = event.get("_realtime_event_id") - if isinstance(event_id, str) and event_id: - error["event_id"] = event_id - await send_json(error) - return None - if config.idle_timeout_s == _DEFAULT_IDLE_TIMEOUT_S: - config.idle_timeout_s = self._idle_timeout_s - runtime_adapter = self._serving_runtime_adapter if self._uses_serving_runtime_adapter(config) else None - runtime_config: dict[str, object] = {} - if runtime_adapter is not None: - try: - runtime_config = await runtime_adapter.prepare_runtime_config( - config, - model_config=getattr(self._chat_service, "model_config", None), - ) - except ServingRuntimeConfigError as exc: - await send_json({"type": "error", "error": str(exc), "code": exc.code}) - return None - except ValueError as exc: - await send_json({"type": "error", "error": str(exc), "code": "unsupported_ref_audio_path"}) - return None - session_id = event.get("session_id") if isinstance(event.get("session_id"), str) else None - session = self._registry.create(config=config, session_id=session_id) - if runtime_adapter is not None: - session.replace_capabilities( - runtime_adapter.capabilities(max_sessions=self._duplex_session_config.max_sessions) - ) - session.replace_runtime_config(runtime_config) - session_payload = event.get("session") - self._resolve_server_vad_defaults( - session, config, session_payload if isinstance(session_payload, dict) else event - ) - event_id = event.get("_realtime_event_id") - return _DuplexSessionHandshake( - session=session, - event_id=event_id if isinstance(event_id, str) and event_id else None, + attachment_send, attachment_close = attachment_callbacks(websocket) + created = await self._attachment_registry.create( + handle.session_id, + send=attachment_send, + close=attachment_close, ) + resume_supported = bool(handle.capabilities.supports_session_resume) + attachment = _Attachment(handle=handle, generation=created.attachment_generation) + credentials: _ResumeCredentials | None = None + if resume_supported: + credentials = _ResumeCredentials( + attachment_generation=created.attachment_generation, + resume_token=created.resume_token.plaintext, + ) + self._start_pump(handle, credentials) + return attachment - async def _resume_session_handshake( + async def _resume( self, - event: dict[str, object], - *, - send_json, - realtime_protocol: NativeRealtimeSessionProtocol, - attachment_send, - attachment_close, - ) -> _DuplexSessionHandshake | None: - session_id = event.get("session_id") - incarnation = event.get("incarnation") - resume_token = event.get("resume_token") - last_received = event.get("last_received_server_event_seq", 0) - if ( - not isinstance(session_id, str) - or not session_id - or not isinstance(incarnation, int) - or incarnation < 0 - or not isinstance(resume_token, str) - or not resume_token - or not isinstance(last_received, int) - or last_received < 0 - ): + websocket: WebSocket, + envelope: RealtimeEnvelope, + event: Mapping[str, object], + send_json: SendJson, + ) -> _Attachment | None: + request = parse_resume_request(event) + if request is None: await send_json( - { - "type": "error", - "error": ( - "session.resume requires session_id, incarnation, resume_token, " - "and a non-negative event sequence" - ), - "code": "invalid_session_resume", - } + envelope.error_payload( + "invalid_session_resume", + "session.resume requires session_id, resume_token, and a non-negative event sequence", + ) ) return None - session = self._registry.get(session_id) - if session is None or session.state != DuplexSessionState.OPEN: + session_id = request.session_id + handle = self._omni.get_session(session_id) + if handle is None or handle.closed: await send_json( - { - "type": "error", - "error": f"Unknown or expired duplex session: {session_id}", - "code": "session_resume_expired", - } + envelope.error_payload("session_resume_expired", f"Unknown or expired duplex session: {session_id}") ) return None - if not session.capabilities.supports_session_resume: + if not handle.capabilities.supports_session_resume: await send_json( - { - "type": "error", - "error": f"Session does not support resume: {session_id}", - "code": "unsupported_session_resume", - } + envelope.error_payload("unsupported_session_resume", f"Session does not support resume: {session_id}") ) return None try: await self._attachment_registry.authenticate_resume( session_id, - incarnation=incarnation, - resume_token=resume_token, - last_received_server_event_seq=last_received, + resume_token=request.resume_token, + last_received_server_event_seq=request.last_received_server_event_seq, ) except InvalidResumeTokenError: - await send_json( - { - "type": "error", - "error": "Invalid duplex session resume token", - "code": "invalid_resume_token", - } - ) + await send_json(envelope.error_payload("invalid_resume_token", "Invalid duplex session resume token")) return None except DuplexJournalGapError: - await send_json( - { - "type": "session.resync_required", - "session_id": session_id, - "reason": "journal_gap", - } - ) + await send_json(SessionResyncRequired(session_id=session_id, reason="journal_gap").to_realtime()) return None except (KeyError, ValueError) as exc: - await send_json( - { - "type": "error", - "error": str(exc), - "code": "session_resume_conflict", - } - ) + await send_json(envelope.error_payload("session_resume_conflict", str(exc))) return None - - resume_runtime = getattr(self._chat_service.engine_client, "resume_duplex_session_async", None) - if not callable(resume_runtime): - await send_json( - { - "type": "error", - "error": "Duplex runtime does not expose session resume control", - "code": "runtime_resume_unsupported", - } - ) - return None - expected_generation = self._lease_generations.get(session_id, 0) try: - runtime_result = await resume_runtime( + await self._omni.resume_session( session_id, - fence=DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ), - expected_lease_generation=expected_generation, - ) - lease_generation = self._runtime_lease_generation(runtime_result) - if lease_generation is None: - raise RuntimeError("runtime resume result omitted lease_generation") - except Exception as exc: - await send_json( - { - "type": "error", - "error": str(exc), - "code": "runtime_resume_failed", - } + expected_lease_generation=handle.lease_generation, ) + except DuplexSessionError as exc: + await send_json(envelope.error_payload("runtime_resume_failed", str(exc))) return None - if not callable(attachment_send) or not callable(attachment_close): - raise RuntimeError("session.resume requires transport attachment callbacks") + attachment_send, attachment_close = attachment_callbacks(websocket) - def activation_payload_factory(token, generation: int) -> dict[str, object]: - internal = { - "type": "session.resumed", - "session_id": session_id, - "incarnation": incarnation, - "attachment_generation": generation, - "resume_token": token.plaintext, - } - return realtime_protocol.encode_outbound_event(internal)[0] + def activation_payload_factory(token: ResumeToken, generation: int) -> dict[str, object]: + return SessionResumed( + session_id=session_id, + session=dict(handle.public_session), + attachment_generation=generation, + resume_token=token.plaintext, + ).to_realtime() try: resumed = await self._attachment_registry.resume( session_id, - incarnation=incarnation, - resume_token=resume_token, - last_received_server_event_seq=last_received, + resume_token=request.resume_token, + last_received_server_event_seq=request.last_received_server_event_seq, send=attachment_send, close=attachment_close, activation_payload_factory=activation_payload_factory, ) except Exception as exc: - # The engine-side CAS already advanced even if the transport - # vanished before it received the rotated token. Keep that - # generation so the registry's one-shot recovery token can retry. - self._lease_generations[session_id] = lease_generation - await send_json( - { - "type": "error", - "error": str(exc), - "code": "session_resume_conflict", - } - ) + # The engine lease was already resumed above, so a failed activation + # would leave the session attached to nothing -- unless another + # connection won the race and is attached right now. Detaching then + # would start the disconnect grace for the *winner*, which ordinary + # heartbeats do not clear. Roll back only what this attempt owns. + if not await self._attachment_registry.has_attachment(session_id): + with suppress(DuplexSessionError): + await self._omni.detach_session(session_id) + await send_json(envelope.error_payload("session_resume_conflict", str(exc))) return None - self._lease_generations[session_id] = lease_generation replaced = resumed.replaced_attachment if replaced is not None: - replaced_payload = realtime_protocol.encode_outbound_event( - { - "type": "session.replaced", - "session_id": session_id, - "attachment_generation": replaced.generation, - } - )[0] with suppress(Exception): - await replaced.send(replaced_payload) + await replaced.send( + SessionReplaced(session_id=session_id, attachment_generation=replaced.generation).to_realtime() + ) with suppress(Exception): await replaced.close("session_replaced") - return _DuplexSessionHandshake( - session=session, - resumed=True, - attachment_generation=resumed.attachment_generation, - ) + # A reconnect brings a fresh envelope carrying pcm16/16 kHz wire + # defaults. The negotiated input format is a wire default, not part of + # the public session object, so it has to be carried over explicitly: + # otherwise the first append that omits format/rate is decoded as pcm16. + remembered = self._input_defaults.get(session_id) + if remembered is not None: + envelope.defaults = remembered + self._start_pump(handle, None) + return _Attachment(handle=handle, generation=resumed.attachment_generation) - @classmethod - def _runtime_lease_generation(cls, result: object) -> int | None: - if isinstance(result, dict): - generation = result.get("lease_generation") - if isinstance(generation, int): - return generation - for key in ("stage_results", "result"): - generation = cls._runtime_lease_generation(result.get(key)) - if generation is not None: - return generation - return None - if isinstance(result, list | tuple): - for item in result: - generation = cls._runtime_lease_generation(item) - if generation is not None: - return generation - return None + # ------------------------------------------------------------------ # + # Outbound pump (session-scoped, survives reconnects) # + # ------------------------------------------------------------------ # - def _uses_serving_runtime_adapter(self, config: DuplexSessionConfig) -> bool: - adapter = self._serving_runtime_adapter - return adapter is not None and adapter.is_enabled(config) - - def _runtime_session_update_error( - self, - session: DuplexSession, - payload: dict[str, object], - ) -> dict[str, object] | None: - if not self._uses_native_input_append(session): - return None - try: - self._require_serving_runtime_adapter().validate_client_extra_body(payload.get("extra_body")) - except ServingRuntimeConfigError as exc: - return { - "type": "error", - "session_id": session.session_id, - "code": exc.code, - "error": str(exc), - } - return None - - def _runtime_config_for_session_update( - self, - session: DuplexSession, - candidate_config: DuplexSessionConfig, - ) -> dict[str, object]: - if not self._uses_native_input_append(session): - return dict(session.runtime_config) - return self._require_serving_runtime_adapter().runtime_config_for_update( - candidate_config, - dict(session.runtime_config), - ) - - def _resolve_server_vad_defaults( - self, - session: DuplexSession, - config: DuplexSessionConfig, - payload: dict[str, object], - ) -> None: - """Resolve omitted options without overriding explicit client policies.""" - if config.server_vad is None: + def _start_pump(self, handle: DuplexSessionHandle, credentials: _ResumeCredentials | None) -> None: + existing = self._pumps.get(handle.session_id) + if existing is not None and not existing.done(): return - if config.server_vad.interrupt_response is None: - config.server_vad = replace(config.server_vad, interrupt_response=self._uses_native_input_append(session)) - extra = payload.get("extra_body") - overlap_policy = extra.get("overlap_policy") if isinstance(extra, dict) else None - if not isinstance(overlap_policy, str): - overlap_policy = payload.get("overlap_policy") - config.overlap_policy = ( - overlap_policy - if isinstance(overlap_policy, str) - else "barge_in_on_speech" - if config.server_vad.interrupt_response - else "listen_only" + self._pumps[handle.session_id] = asyncio.create_task( + self._pump_events(handle, credentials), + name=f"duplex-session-pump-{handle.session_id}", ) - def _runtime_session_candidate_update_error( - self, - session: DuplexSession, - candidate_config: DuplexSessionConfig, - *, - realtime_protocol: NativeRealtimeSessionProtocol | None = None, - ) -> dict[str, object] | None: - if candidate_config.server_vad is not None: - desired_policy = "barge_in_on_speech" if candidate_config.server_vad.interrupt_response else "listen_only" - if candidate_config.overlap_policy != desired_policy: - return { - "type": "error", - "code": "unsupported_turn_detection", - "error": f"overlap_policy conflicts with turn_detection; expected {desired_policy!r}", - "param": "overlap_policy", - } - if not self._uses_native_input_append(session): - if ( - candidate_config.server_vad is not None - and candidate_config.server_vad.min_speech_duration_ms is not None - ): - return { - "type": "error", - "code": "unsupported_turn_detection", - "error": "server_vad.min_speech_duration_ms is only supported by model-native duplex runtimes", - } - requests_interruption = ( - candidate_config.server_vad is not None and candidate_config.server_vad.interrupt_response - ) or candidate_config.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value - if requests_interruption: - return { - "type": "error", - "code": "server_vad_requires_native_duplex", - "error": ( - "Realtime server_vad interruption requires a model-native duplex runtime; " - "set interrupt_response=false for turn-based models" - ), - } - if candidate_config.server_vad is not None and realtime_protocol is not None: - session_payload = candidate_config.extra_body.get("realtime_session_payload") - audio_error = realtime_protocol.validate_server_vad_input_audio( - session_payload if isinstance(session_payload, dict) else {} - ) - if audio_error is not None: - return { - "type": "error", - "code": "unsupported_turn_detection", - "error": audio_error, - "param": "audio.input.format", - } - return None - if candidate_config.server_vad is not None: - if not candidate_config.server_vad.interrupt_response: - return { - "type": "error", - "session_id": session.session_id, - "code": "unsupported_turn_detection", - "error": "Model-native duplex server_vad requires interrupt_response=true", - } - if candidate_config.server_vad.threshold <= SILERO_VAD_MIN_THRESHOLD: - return { - "type": "error", - "session_id": session.session_id, - "code": "unsupported_turn_detection", - "error": ( - f"Model-native duplex server_vad.threshold must be greater than {SILERO_VAD_MIN_THRESHOLD}" - ), - } - if not candidate_config.server_vad.create_response: - return { - "type": "error", - "session_id": session.session_id, - "code": "unsupported_turn_detection", - "error": "Model-native duplex server_vad requires create_response=true", - } - if ( - candidate_config.instructions != session.config.instructions - and self._runtime_session_state(session).native_context_locked - ): - return { - "type": "error", - "session_id": session.session_id, - "code": "instructions_update_unsupported", - "error": "session.update cannot change instructions after the native duplex context is initialized", - } - if native_duplex_opt_in(session.config.extra_body) is not True: - return None - if not self._config_requests_audio_output(candidate_config): - return None - if "ref_audio_data" in session.runtime_config: - return None - return { - "type": "error", - "session_id": session.session_id, - "code": "ref_audio_required", - "error": "MiniCPM-o native duplex audio output requires ref_audio", - } + async def _pump_events(self, handle: DuplexSessionHandle, credentials: _ResumeCredentials | None) -> None: + session_id = handle.session_id + close_reason = "session_closed" + try: + async for event in handle.events(): + if isinstance(event, SessionCreated) and credentials is not None: + event = replace( + event, + attachment_generation=credentials.attachment_generation, + resume_token=credentials.resume_token, + ) + await self._send_event(session_id, event) + if isinstance(event, SessionClosed): + close_reason = event.reason or event.type + break + except asyncio.CancelledError: + raise + except Exception as exc: + logger.exception("Duplex event pump failed for %s: %s", session_id, exc) + finally: + self._pumps.pop(session_id, None) + self._resync_required_sessions.discard(session_id) + self._input_defaults.pop(session_id, None) + attachment = None + with suppress(Exception): + attachment = await self._attachment_registry.close(session_id) + if attachment is not None: + with suppress(Exception): + await attachment.close(close_reason) - @staticmethod - def _uses_native_input_append(session: DuplexSession) -> bool: - return ( - session.capabilities.implementation_level == "model_native_duplex" - and session.capabilities.supports_input_append - and native_duplex_opt_in(session.config.extra_body) is not False - ) + async def _drain_terminal_pump(self, attachment: _Attachment) -> None: + """Let the pump deliver ``session.closed`` before the endpoint returns. - @staticmethod - def _config_requests_audio_output(config: DuplexSessionConfig) -> bool: - return any(str(modality).lower() == "audio" for modality in config.modalities) + ``DuplexSessionHandle._deliver`` queues ``SessionClosed`` on the outbox + and marks the handle closed in the same synchronous step, so + ``handle.closed`` is already true while that event is still sitting in + the queue. ``_read_loop`` stops on exactly that flag: if the engine + closes the session between two reads, the loop returns, this endpoint + returns, and the ASGI server tears the socket down with the terminal + event unsent -- the client sees an abrupt close instead of + ``session.closed``. Waiting for the pump keeps the two in order; it is + the pump that sends the event and then closes with code 1000. - @staticmethod - def _native_stage0_request_id(session: DuplexSession, epoch: int) -> str: - return duplex_resource_request_id( - DuplexFence( - session.session_id, - epoch=epoch, - incarnation=session.incarnation, - ), - "stage0", - ) - - @staticmethod - def _session_auto_responds(session: DuplexSession) -> bool: - """Full-duplex / model-driven mode. - - When set, the server runs per-chunk speak-generation continuously (like - the official MiniCPM-o ``duplex_generate`` loop) instead of waiting for an - explicit ``response.create``: each ~chunk_period of appended audio is - emitted and fed to the stage0 stream so the model itself decides to speak - or listen. Signaled by the client via ``extra_body.auto_response`` (or - ``extra_body.full_duplex``). + Only for a closed session. A takeover or a resumable disconnect leaves + the pump running for the next attachment, and must not be waited on. """ - extra = getattr(session.config, "extra_body", None) - if not isinstance(extra, dict): - return False - return extra.get("auto_response") is True or extra.get("full_duplex") is True + if not attachment.handle.closed: + return + pump = self._pumps.get(attachment.handle.session_id) + if pump is None or pump.done(): + return + # shield: a drain timeout must not cancel a pump that is still the + # session's only writer. + with suppress(Exception): + await asyncio.wait_for(asyncio.shield(pump), _PUMP_DRAIN_TIMEOUT_S) - async def _receive_text( + async def _send_event(self, session_id: str, event: DuplexEvent) -> None: + payload = event.to_realtime() + journal = not isinstance(event, _UNJOURNALED_EVENTS) and session_id not in self._resync_required_sessions + try: + try: + await self._attachment_registry.send_event(session_id, payload, journal=journal) + 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) + except KeyError: + # Attachment already closed (takeover or teardown); the journal is gone. + pass + except asyncio.CancelledError: + raise + except Exception as exc: + # The socket died mid-send. The session is alive in the engine: + # treat it as a disconnect of the current attachment (the pump + # keeps journaling for a resume) instead of tearing the session down. + logger.info("Duplex transport send failed for %s: %s", session_id, exc) + await self._detach_current_attachment(session_id) + + async def _detach_current_attachment(self, session_id: str) -> None: + """Disconnect semantics for the socket currently attached to ``session_id``. + + The pump is session-scoped and outlives any one connection, so it can + only name the session: the registry detaches whichever socket is + attached, which is the one whose send just failed. + """ + handle = self._omni.get_session(session_id) + if handle is None or handle.closed: + return + if handle.capabilities.supports_session_resume: + if await self._attachment_registry.detach(session_id, attachment_generation=None): + with suppress(DuplexSessionError): + await self._omni.detach_session(session_id) + return + with suppress(DuplexSessionError): + await handle.close(reason="disconnect") + + # ------------------------------------------------------------------ # + # Inbound # + # ------------------------------------------------------------------ # + + async def _read_loop( self, websocket: WebSocket, - timeout_s: float, - *, - realtime_protocol: NativeRealtimeSessionProtocol | None = None, - ) -> str | None: - try: - if realtime_protocol is not None: - return await asyncio.wait_for( - realtime_protocol.receive_internal_event_text(websocket), - timeout=max(0.1, timeout_s), - ) - return await asyncio.wait_for(websocket.receive_text(), timeout=max(0.1, timeout_s)) - except asyncio.TimeoutError: - return None - - @staticmethod - def _input_committed_payload( - session: DuplexSession, - committed: DuplexCommittedInput, - *, - item_id: object | None = None, - ) -> dict[str, object]: - payload: dict[str, object] = { - "type": "input.committed", - "session_id": session.session_id, - "turn_id": committed.turn_id, - "epoch": committed.epoch, - "history_len": len(session.history), - "message": committed.message, - } - if isinstance(item_id, str) and item_id: - payload["item_id"] = item_id - return payload - - @staticmethod - def _commit_native_audio_input( - session: DuplexSession, - *, - item_id: object | None = None, - transcript: object | None = None, - turn_id: int | None = None, - ) -> DuplexCommittedInput: - clean_transcript = transcript.strip() if isinstance(transcript, str) else None - committed = session.commit_native_audio_input( - transcript=clean_transcript or None, - turn_id=turn_id, - ) - if isinstance(item_id, str) and item_id: - session.register_history_item(item_id, committed.message) - return committed - - @staticmethod - def _native_audio_committed_payload( - session: DuplexSession, - *, - committed: DuplexCommittedInput | None = None, - item_id: object | None = None, - transcript: object | None = None, - ) -> dict[str, object]: - message = committed.message if committed is not None else None - if not isinstance(message, dict): - input_audio_part: dict[str, object] = { - "type": "audio_url", - "audio_url": {"url": "native-duplex:input-audio"}, - } - if isinstance(transcript, str) and transcript: - input_audio_part["transcript"] = transcript - message = { - "role": "user", - "content": [input_audio_part], - } - payload: dict[str, object] = { - "type": "input.committed", - "session_id": session.session_id, - "turn_id": committed.turn_id if committed is not None else session.turn_id, - "epoch": committed.epoch if committed is not None else session.epoch, - "history_len": len(session.history), - "native_audio": True, - "message": message, - } - if isinstance(transcript, str) and transcript: - payload["transcript"] = transcript - if isinstance(item_id, str) and item_id: - payload["item_id"] = item_id - return payload - - def _apply_session_update( - self, - session: DuplexSession, - payload: dict[str, object], - ) -> dict[str, object] | None: - model = payload.get("model") - audio_config = payload.get("audio") - audio_input = audio_config.get("input") if isinstance(audio_config, dict) else None - audio_output = audio_config.get("output") if isinstance(audio_config, dict) else None - voice = payload.get("voice") - if not isinstance(voice, str) and isinstance(audio_output, dict): - voice = audio_output.get("voice") - if isinstance(model, str) and session.config.model is not None and model != session.config.model: - return { - "type": "error", - "session_id": session.session_id, - "code": "model_update_unsupported", - "error": "session.update cannot change model for an open realtime duplex session", - } - if isinstance(model, str) and session.config.model is None: - session.config.model = model - if isinstance(voice, str) and (session.playback.generated_ms > 0 or session.playback.sent_ms > 0): - return { - "type": "error", - "session_id": session.session_id, - "code": "voice_update_after_audio_unsupported", - "error": "session.update cannot change voice after audio output has started", - } - if isinstance(payload.get("ref_audio"), str): - return { - "type": "error", - "session_id": session.session_id, - "code": "ref_audio_update_unsupported", - "error": "session.update cannot change ref_audio after the native duplex runtime is open", - } - extra_body_payload = payload.get("extra_body") - if isinstance(extra_body_payload, dict): - # Fold the deprecated model-prefixed alias into the canonical - # native_duplex key before comparing or merging. - extra_body_payload = normalize_native_duplex_key(dict(extra_body_payload)) - payload = {**payload, "extra_body": extra_body_payload} - if isinstance(extra_body_payload, dict) and NATIVE_DUPLEX_KEY in extra_body_payload: - requested_native_duplex = extra_body_payload[NATIVE_DUPLEX_KEY] - if not isinstance(requested_native_duplex, bool): - return { - "type": "error", - "session_id": session.session_id, - "code": "invalid_request_error", - "error": "extra_body.native_duplex must be a boolean", - } - current_native_duplex = native_duplex_opt_in(session.config.extra_body) - if current_native_duplex is not None and extra_body_payload[NATIVE_DUPLEX_KEY] != current_native_duplex: - return { - "type": "error", - "session_id": session.session_id, - "code": "native_duplex_mode_update_unsupported", - "error": "session.update cannot change native_duplex after the session is created", - } - if isinstance(payload.get("instructions"), str): - session.config.instructions = str(payload["instructions"]) - elif "instructions" in payload and payload.get("instructions") is None: - session.config.instructions = None - if isinstance(voice, str): - session.config.voice = str(voice) - elif "voice" in payload and payload.get("voice") is None: - session.config.voice = None - response_format = payload.get("output_audio_format") or payload.get("response_format") - if response_format is None and isinstance(audio_config, dict): - if isinstance(audio_output, dict): - response_format = audio_output.get("format") - response_format, _ = NativeRealtimeSessionProtocol._parse_realtime_audio_format(response_format) - if isinstance(response_format, str) and response_format.lower() in REALTIME_OUTPUT_AUDIO_FORMATS: - session.config.response_format = NativeRealtimeSessionProtocol._duplex_response_format(response_format) - if isinstance(payload.get("temperature"), int | float): - session.config.temperature = float(payload["temperature"]) - speed = payload.get("speed") - if not isinstance(speed, int | float) and isinstance(audio_output, dict): - speed = audio_output.get("speed") - if isinstance(speed, int | float): - session.config.speed = float(speed) - max_tokens = ( - payload.get("max_response_output_tokens") - if "max_response_output_tokens" in payload - else payload.get("max_output_tokens") - if "max_output_tokens" in payload - else payload.get("max_tokens") - ) - if "max_response_output_tokens" in payload or "max_output_tokens" in payload or "max_tokens" in payload: - session.config.max_tokens = NativeRealtimeSessionProtocol.realtime_max_output_tokens(max_tokens) - if isinstance(payload.get("overlap_policy"), str): - session.config.overlap_policy = DuplexSessionConfig._normalize_overlap_policy( - str(payload["overlap_policy"]) - ) - if isinstance(payload.get("overlap_short_ack_ms"), int | float): - session.config.overlap_short_ack_ms = max(0, int(payload["overlap_short_ack_ms"])) - if isinstance(payload.get("overlap_barge_in_ms"), int | float): - session.config.overlap_barge_in_ms = max(0, int(payload["overlap_barge_in_ms"])) - if isinstance(payload.get("overlap_silence_rms"), int | float): - session.config.overlap_silence_rms = max(0.0, float(payload["overlap_silence_rms"])) - if isinstance(payload.get("playback_commit_policy"), str): - session.config.playback_commit_policy = DuplexSessionConfig._normalize_playback_commit_policy( - str(payload["playback_commit_policy"]) - ) - modalities = payload.get("modalities") or payload.get("output_modalities") - if isinstance(modalities, list) and all(isinstance(item, str) for item in modalities): - session.config.modalities = list(modalities) - if isinstance(payload.get("extra_body"), dict): - session.config.extra_body.update(payload["extra_body"]) - extra = payload["extra_body"] - if isinstance(extra.get("overlap_policy"), str): - session.config.overlap_policy = DuplexSessionConfig._normalize_overlap_policy( - str(extra["overlap_policy"]) - ) - if isinstance(extra.get("playback_commit_policy"), str): - session.config.playback_commit_policy = DuplexSessionConfig._normalize_playback_commit_policy( - str(extra["playback_commit_policy"]) - ) - if isinstance(payload.get("tools"), list): - session.config.extra_body["realtime_tools"] = payload["tools"] - elif "tools" in payload and payload.get("tools") is None: - session.config.extra_body.pop("realtime_tools", None) - if isinstance(payload.get("tool_choice"), str | dict): - session.config.extra_body["realtime_tool_choice"] = payload["tool_choice"] - elif "tool_choice" in payload and payload.get("tool_choice") is None: - session.config.extra_body.pop("realtime_tool_choice", None) - if isinstance(payload.get("metadata"), dict): - session.config.extra_body["realtime_metadata"] = dict(payload["metadata"]) - elif "metadata" in payload and payload.get("metadata") is None: - session.config.extra_body.pop("realtime_metadata", None) - if isinstance(payload.get("include"), list): - session.config.extra_body["realtime_include"] = list(payload["include"]) - elif "include" in payload and payload.get("include") is None: - session.config.extra_body.pop("realtime_include", None) - if isinstance(payload.get("prompt"), dict): - session.config.extra_body["realtime_prompt"] = dict(payload["prompt"]) - elif "prompt" in payload and payload.get("prompt") is None: - session.config.extra_body.pop("realtime_prompt", None) - input_audio_transcription = NativeRealtimeSessionProtocol._input_audio_transcription_config(payload) - if isinstance(input_audio_transcription, dict): - session.config.extra_body["realtime_input_audio_transcription"] = dict(input_audio_transcription) - elif "input_audio_transcription" in payload and payload.get("input_audio_transcription") is None: - session.config.extra_body.pop("realtime_input_audio_transcription", None) - if isinstance(payload.get("input_audio_noise_reduction"), dict): - session.config.extra_body["realtime_input_audio_noise_reduction"] = dict( - payload["input_audio_noise_reduction"] - ) - elif "input_audio_noise_reduction" in payload and payload.get("input_audio_noise_reduction") is None: - session.config.extra_body.pop("realtime_input_audio_noise_reduction", None) - if isinstance(audio_input, dict) and isinstance(audio_input.get("noise_reduction"), dict): - session.config.extra_body["realtime_input_audio_noise_reduction"] = dict(audio_input["noise_reduction"]) - elif isinstance(audio_input, dict) and audio_input.get("noise_reduction") is None: - session.config.extra_body.pop("realtime_input_audio_noise_reduction", None) - if isinstance(payload.get("audio"), dict): - session.config.extra_body["realtime_audio"] = dict(payload["audio"]) - elif "audio" in payload and payload.get("audio") is None: - session.config.extra_body.pop("realtime_audio", None) - if isinstance(payload.get("tracing"), str | dict): - session.config.extra_body["realtime_tracing"] = payload["tracing"] - elif "tracing" in payload and payload.get("tracing") is None: - session.config.extra_body.pop("realtime_tracing", None) - session.config.extra_body["realtime_session_payload"] = ( - NativeRealtimeSessionProtocol._json_safe_realtime_payload(payload) - ) - try: - turn_detection_configured, server_vad = parse_session_turn_detection(payload) - if turn_detection_configured: - vad_config = replace(session.config, server_vad=server_vad) - self._resolve_server_vad_defaults(session, vad_config, payload) - server_vad = vad_config.server_vad - if ( - not self._uses_native_input_append(session) - and session.turn_detection_config_locked - and (not session.config.turn_detection_configured or server_vad != session.config.server_vad) - ): - return { - "type": "error", - "session_id": session.session_id, - "code": "unsupported_turn_detection", - "error": "turn_detection cannot be changed after the first audio append", - } - session.config.turn_detection_configured = True - session.config.server_vad = server_vad - session.config.overlap_policy = vad_config.overlap_policy - except ValueError as exc: - return { - "type": "error", - "session_id": session.session_id, - "code": "invalid_request_error", - "error": str(exc), - } - return None - - async def _configure_server_vad(self, session: DuplexSession) -> None: - pipeline = await self._prepare_server_vad_pipeline(session, session.config) - self._install_server_vad_pipeline(session, pipeline) - - def _install_server_vad_pipeline( - self, - session: DuplexSession, - pipeline: ServerVADPipeline | None, + envelope: RealtimeEnvelope, + attachment: _Attachment, + send_json: SendJson, ) -> None: - existing = self._server_vad_pipelines.pop(session.session_id, None) - if existing is not None: - existing.reset() - self._realtime_vad_metrics.session_finished() - if pipeline is None: - return - self._server_vad_pipelines[session.session_id] = pipeline - self._realtime_vad_metrics.session_started() - - async def _prepare_server_vad_pipeline( - self, - session: DuplexSession, - config: DuplexSessionConfig, - ) -> ServerVADPipeline | None: - if config.server_vad is None: - return None - backend = await asyncio.to_thread(self._server_vad_backend_provider.get) - pipeline_config = config.server_vad - if self._uses_native_input_append(session) and pipeline_config.min_speech_duration_ms is None: - pipeline_config = replace( - pipeline_config, - min_speech_duration_ms=SILERO_VAD_DEFAULT_MIN_SPEECH_DURATION_MS, - ) - return ServerVADPipeline(backend, pipeline_config) - - def _apply_response_create_options( - self, - session: DuplexSession, - payload: dict[str, object], - ) -> str | None: - """Reserve options that apply only to the next response lifecycle.""" - audio_config = payload.get("audio") - audio_output = audio_config.get("output") if isinstance(audio_config, dict) else None - if session.capabilities.implementation_level == "model_native_duplex": - nested_voice = audio_output.get("voice") if isinstance(audio_output, dict) else None - unsupported = ( - payload.get("instructions") is not None - or payload.get("voice") is not None - or nested_voice is not None - or payload.get("temperature") is not None - or any( - payload.get(field_name) is not None - for field_name in ("max_response_output_tokens", "max_output_tokens", "max_tokens") - ) - or payload.get("tools") is not None - or payload.get("tool_choice") is not None - ) - if unsupported: - return "unsupported_native_response_options" - - instructions = str(payload["instructions"]) if isinstance(payload.get("instructions"), str) else None - voice = payload.get("voice") - if not isinstance(voice, str) and isinstance(audio_output, dict): - voice = audio_output.get("voice") - voice = str(voice) if isinstance(voice, str) else None - response_format = payload.get("output_audio_format") or payload.get("response_format") - if response_format is None and isinstance(audio_config, dict): - if isinstance(audio_output, dict): - response_format = audio_output.get("format") - response_format, _ = NativeRealtimeSessionProtocol._parse_realtime_audio_format(response_format) - if isinstance(response_format, str) and response_format.lower() in REALTIME_OUTPUT_AUDIO_FORMATS: - response_format = NativeRealtimeSessionProtocol._duplex_response_format(response_format) - else: - response_format = None - temperature = float(payload["temperature"]) if isinstance(payload.get("temperature"), int | float) else None - speed = payload.get("speed") - if not isinstance(speed, int | float) and isinstance(audio_output, dict): - speed = audio_output.get("speed") - speed = float(speed) if isinstance(speed, int | float) else None - max_tokens = ( - payload.get("max_response_output_tokens") - if "max_response_output_tokens" in payload - else payload.get("max_output_tokens") - if "max_output_tokens" in payload - else payload.get("max_tokens") - ) - if "max_response_output_tokens" in payload or "max_output_tokens" in payload or "max_tokens" in payload: - max_tokens = NativeRealtimeSessionProtocol.realtime_max_output_tokens(max_tokens) - else: - max_tokens = None - modalities = payload.get("modalities") or payload.get("output_modalities") - if isinstance(modalities, list) and all(isinstance(item, str) for item in modalities): - modalities = tuple(modalities) - else: - modalities = None - response_extra: dict[str, object] = {} - conversation = payload.get("conversation") - if isinstance(conversation, str): - response_extra["realtime_response_conversation"] = conversation - metadata = payload.get("metadata") - if isinstance(metadata, dict): - response_extra["realtime_response_metadata"] = dict(metadata) - prompt = payload.get("prompt") - if isinstance(prompt, dict): - response_extra["realtime_response_prompt"] = dict(prompt) - if isinstance(payload.get("tools"), list): - response_extra["realtime_response_tools"] = payload["tools"] - if isinstance(payload.get("tool_choice"), str | dict): - response_extra["realtime_response_tool_choice"] = payload["tool_choice"] - extra_body = payload.get("extra_body") - if isinstance(extra_body, dict): - if session.capabilities.implementation_level == "model_native_duplex": - response_extra.update( - (key, value) - for key, value in extra_body.items() - if key not in self._require_serving_runtime_adapter().private_runtime_config_keys - ) - else: - response_extra.update(extra_body) - try: - session.reserve_response_options( - ResponseCreateOptions( - instructions=instructions, - voice=voice, - response_format=response_format, - temperature=temperature, - max_tokens=max_tokens, - speed=speed, - modalities=modalities, - extra_body=response_extra, - ) - ) - except RuntimeError: - return "response_already_active" - return None - - @staticmethod - def _realtime_item_to_history_message(item: object) -> dict[str, object] | None: - if not isinstance(item, dict): - return None - role = item.get("role") - if role not in {"system", "user", "assistant"}: - return None - content = item.get("content") - if isinstance(content, str): - text = content.strip() - return {"role": role, "content": text} if text else None - if not isinstance(content, list): - return None - text_chunks: list[str] = [] - audio_chunks: list[dict[str, object]] = [] - for part in content: - if not isinstance(part, dict): + handle = attachment.handle + idle_timeout_s = self._idle_timeout_s + config_timeout = handle.public_session.get("idle_timeout_s") if handle.public_session else None + if isinstance(config_timeout, int | float) and config_timeout > 0: + idle_timeout_s = float(config_timeout) + while not handle.closed: + raw = await receive_text_with_timeout(websocket, idle_timeout_s) + if not await self._attachment_registry.is_current_attachment(handle.session_id, attachment.generation): + # A newer connection took the session over; this socket is done. + return + if raw is None: + # Nothing from this socket for the idle window: same as a + # disconnect (engine grace for resumable sessions); the engine + # lease decides whether the session itself expires. + await self._on_disconnect(attachment) + return + if len(raw.encode("utf-8")) > MAX_EVENT_BYTES: + await send_json(envelope.error_payload("event_too_large", "Duplex event too large")) continue - part_type = part.get("type") - if part_type in {"input_text", "text", "output_text"} and isinstance(part.get("text"), str): - text_chunks.append(str(part["text"])) - elif part_type in {"input_audio", "audio"}: - audio = part.get("audio") or part.get("data") - fmt = part.get("format") if isinstance(part.get("format"), str) else "wav" - if isinstance(audio, str) and audio: - audio_chunks.append( - { - "type": "audio_url", - "audio_url": { - "url": f"data:audio/{fmt};base64,{audio}", - }, - } - ) - elif part_type in {"audio_transcript", "transcript"} and isinstance(part.get("text"), str): - text_chunks.append(str(part["text"])) - text = "".join(text_chunks).strip() - if audio_chunks: - content_items: list[dict[str, object]] = [] - if text: - content_items.append({"type": "text", "text": text}) - content_items.extend(audio_chunks) - return {"role": role, "content": content_items} - if text: - return {"role": role, "content": text} - return None - - async def _handle_playback_ack(self, session: DuplexSession, event: dict[str, object], send_json) -> None: - played_ms = event.get("played_ms", event.get("audio_ms", 0)) - committed_ms = event.get("committed_ms") - if not isinstance(played_ms, int | float): - await send_json({"type": "error", "error": "playback.ack requires played_ms", "code": "bad_event"}) - return - committed_cursor = int(committed_ms) if isinstance(committed_ms, int | float) else int(played_ms) - item_id = event.get("item_id") - response_id = event.get("response_id") - response_id = response_id if isinstance(response_id, str) and response_id else None - if not isinstance(item_id, str) or not item_id: - item_id = f"item_{response_id}" if response_id is not None else None - elif response_id is None and item_id.startswith("item_"): - response_id = item_id.removeprefix("item_") - if response_id is None and item_id is None and len(session.pending_history_item_ids) == 1: - item_id = next(iter(session.pending_history_item_ids)) - if item_id.startswith("item_"): - response_id = item_id.removeprefix("item_") - if response_id is None and item_id is None and session.active_response_id is not None: - response_id = session.active_response_id - item_id = f"item_{response_id}" - if response_id is not None: - expected_item_id = f"item_{response_id}" - if item_id is None: - item_id = expected_item_id - elif item_id != expected_item_id: - await send_json( - { - "type": "error", - "session_id": session.session_id, - "epoch": session.epoch, - "code": "playback_item_mismatch", - "error": "playback.ack item_id must match item_.", - } - ) - return - if not session.has_assistant_response_item(response_id, item_id): - await send_json( - { - "type": "error", - "session_id": session.session_id, - "epoch": session.epoch, - "code": "playback_item_not_found", - "error": f"No assistant response item is registered for {response_id}.", - } - ) - return - if session.playback_ack_is_too_late(response_id, item_id): - await send_json( - { - "type": "error", - "session_id": session.session_id, - "epoch": session.epoch, - "code": "playback_ack_too_late", - "error": "playback.ack arrived after a later user input was committed.", - } - ) - return - # Reserve the response's current history position before any later - # input commit can append a user item. A 0 ms ACK is intentional: - # the response may be active but have no audio delta yet. - session.reserve_history_item(item_id) - elif item_id is not None: - await send_json( - { - "type": "error", - "session_id": session.session_id, - "epoch": session.epoch, - "code": "playback_item_not_found", - "error": "playback.ack requires a response-owned assistant item.", - } - ) - return - hard_truncate = event.get("truncate") is True - if hard_truncate: - playback = session.acknowledge_playback( - int(played_ms), - committed_cursor, - response_id=response_id, - ) - playback = session.truncate_playback_commit( - committed_cursor, - response_id=response_id, - ) - else: - playback = session.acknowledge_playback( - int(played_ms), - committed_cursor, - response_id=response_id, - ) - committed_history = False - if isinstance(item_id, str) and item_id: - expected_item_id = f"item_{response_id}" if response_id is not None else None - if ( - expected_item_id == item_id - and item_id not in session.history_item_ids - and item_id not in session.pending_history_item_ids - ): - # ACK_ONLY responses are normally registered as pending when - # the response ends. Recover the response-local item if that - # registration was lost across the response/ack boundary. - session.register_history_item(item_id, None) - committed_history = session.truncate_history_item( - item_id, - audio_end_ms=committed_cursor, - playback=playback, - hard=hard_truncate, - ) - elif session.pending_history_item_ids: - # A plain playback ack has no OpenAI item id. Commit the only - # uncommitted assistant candidate if the session has an unambiguous - # pending response; otherwise wait for conversation.item.truncate. - pending_ids = list(session.pending_history_item_ids) - if len(pending_ids) == 1: - item_id = pending_ids[0] - committed_history = session.truncate_history_item( - item_id, - audio_end_ms=committed_cursor, - playback=playback, - hard=hard_truncate, - ) - elif session.active_response_id is not None: - item_id = f"item_{session.active_response_id}" - committed_history = session.truncate_history_item( - item_id, - audio_end_ms=committed_cursor, - playback=playback, - hard=hard_truncate, - ) - elif session.last_assistant_full_message is not None: - if item_id is None and session.history_item_ids: - assistant_item_ids = [ - known_item_id - for known_item_id, message in session.history_item_ids.items() - if message.get("role") == "assistant" - ] - if len(assistant_item_ids) == 1: - item_id = assistant_item_ids[0] - if isinstance(item_id, str) and item_id: - committed_history = session.truncate_history_item( - item_id, - audio_end_ms=committed_cursor, - playback=playback, - hard=hard_truncate, - ) - await send_json( - { - "type": "playback.acknowledged", - "session_id": session.session_id, - "epoch": session.epoch, - "item_id": item_id, - "played_ms": int(played_ms), - "committed_ms": committed_cursor, - "truncate": event.get("truncate") is True, - "playback": playback.as_dict(), - "history_committed": committed_history, - } - ) - if committed_history and committed_cursor >= max(playback.sent_ms, playback.generated_ms): - session.release_response_playback(response_id) - session.release_response_history_snapshot(response_id) - - async def _cancel_active_response( - self, - session: DuplexSession, - active_task: asyncio.Task[None] | None, - send_json, - *, - reason: str, - notify: bool = True, - ) -> bool: - has_running_task = active_task is not None and not active_task.done() - if not has_running_task and session.active_request_id is None and session.active_response_id is None: - return False - - old_epoch = session.epoch - old_request_id = session.active_request_id - old_response_id = session.active_response_id - committed_ms = session.playback.committed_ms - committed_message = session.end_response( - commit_text=self._should_commit_response_to_history(session, old_response_id), - playback_commit_policy=DuplexPlaybackCommitPolicy.ACK_ONLY.value, - ) - if old_response_id is not None: - item_id = f"item_{old_response_id}" - if committed_message is not None: - session.register_history_item(item_id, committed_message) - elif committed_ms > 0 and not session.playback_ack_is_too_late(old_response_id, item_id): - session.truncate_history_item(item_id, audio_end_ms=committed_ms) - new_epoch, old_playback = self._advance_barge_in_epoch(session) - if old_request_id is not None: - # Epoch-scoped request ids prevent state reuse, but explicitly - # release projector/parser cursors so cancelled epochs do not - # accumulate until the whole session closes. - if self._serving_runtime_adapter is not None: - self._serving_runtime_adapter.data_plane.close_stream(old_request_id) - await self._abort_request_background( - session, - old_request_id, - send_json, - notify=notify, - ) - if has_running_task and active_task is not None: - active_task.cancel() try: - await asyncio.wait_for(asyncio.gather(active_task, return_exceptions=True), timeout=0.25) - except asyncio.TimeoutError: - pass - if notify: - await send_json( - { - "type": "audio.cancelled", - "session_id": session.session_id, - "response_id": old_response_id, - "reason": reason, - "cancelled_epoch": old_epoch, - "epoch": new_epoch, - "committed_ms": committed_ms, - "playback": old_playback, - } - ) - return True + payload = json.loads(raw) + except json.JSONDecodeError: + await send_json(envelope.error_payload("invalid_json", "Invalid JSON event")) + continue + if not isinstance(payload, dict) or not isinstance(payload.get("type"), str): + await send_json(envelope.error_payload("bad_event", "Duplex event must be a JSON object with a type")) + continue + if envelope.is_envelope_event(payload): + await self._handle_envelope_event(attachment, envelope, payload, send_json) + continue + await self._submit_wire_event(attachment, envelope, payload, send_json) - async def _abort_request_background( + async def _handle_envelope_event( self, - session: DuplexSession, - request_id: str, - send_json, - *, - notify: bool, + attachment: _Attachment, + envelope: RealtimeEnvelope, + payload: dict[str, object], + send_json: SendJson, + ) -> None: + handle = attachment.handle + event_type = payload.get("type") + event_id = payload.get("event_id") + if event_type == "session.event_ack": + acknowledged = payload.get("server_event_seq") + if not isinstance(acknowledged, int) or acknowledged < 0: + message = "session.event_ack requires a non-negative server_event_seq" + await send_json(envelope.error_payload("invalid_event_ack", message, event_id=event_id)) + return + try: + await self._attachment_registry.acknowledge(handle.session_id, acknowledged) + except ValueError as exc: + await send_json(envelope.error_payload("invalid_event_ack", str(exc), event_id=event_id)) + return + message = "session.resume must be the first message of a new connection" + await send_json(envelope.error_payload("unsupported_session_resume", message, event_id=event_id)) + + async def _submit_wire_event( + self, + attachment: _Attachment, + envelope: RealtimeEnvelope, + payload: dict[str, object], + send_json: SendJson, ) -> None: try: - abort_internal = getattr(self._chat_service.engine_client, "_abort_internal_requests", None) - if callable(abort_internal): - result = abort_internal([request_id]) - else: - result = self._chat_service.engine_client.abort([request_id]) - if inspect.isawaitable(result): - await result - except Exception as exc: - logger.exception("Failed to abort duplex request %s: %s", request_id, exc) - if notify and session.state != DuplexSessionState.CLOSED: - await self._send_runtime_error(send_json, "runtime_abort_failed", exc, session=session) + command = envelope.translate(payload) + except DuplexCommandError as exc: + await send_json(envelope.command_error_payload(exc)) + return + # ``translate`` folds a session.update's audio settings into the + # envelope; remember them so a later reconnect starts where this left off. + self._input_defaults[attachment.handle.session_id] = envelope.defaults + await self._submit_command(attachment, envelope, command, send_json) - async def _cancel_pending_input(self, session: DuplexSession, send_json, *, reason: str) -> None: - cancelled = session.cancel_pending_input() - self._advance_barge_in_epoch(session) - await send_json( - { - "type": "input.cancelled", - "session_id": session.session_id, - "reason": reason, - "epoch": session.epoch, - "cancelled": cancelled, - } - ) + async def _submit_command( + self, + attachment: _Attachment, + envelope: RealtimeEnvelope, + command: DuplexCommand, + send_json: SendJson, + ) -> None: + try: + await attachment.handle.submit(command) + except DuplexSessionError as exc: + await send_json(envelope.error_payload(exc.code, str(exc), event_id=command.event_id)) + + # ------------------------------------------------------------------ # + # Disconnect # + # ------------------------------------------------------------------ # + + async def _on_disconnect(self, attachment: _Attachment) -> None: + handle = attachment.handle + if handle.closed: + return + if handle.capabilities.supports_session_resume: + if await self._attachment_registry.detach(handle.session_id, attachment_generation=attachment.generation): + # Engine-owned disconnect grace: expiry arrives as session.expired. + with suppress(DuplexSessionError): + await self._omni.detach_session(handle.session_id) + # Otherwise a newer connection already took the session over; the + # replaced socket's disconnect must not touch it. + return + with suppress(DuplexSessionError): + await handle.close(reason="disconnect") diff --git a/vllm_omni/entrypoints/duplex/session_attachment.py b/vllm_omni/entrypoints/duplex/session_attachment.py index 242169a6f..f2dc09f1c 100644 --- a/vllm_omni/entrypoints/duplex/session_attachment.py +++ b/vllm_omni/entrypoints/duplex/session_attachment.py @@ -180,7 +180,6 @@ class DuplexEventJournal: @dataclass(frozen=True) class DuplexSessionAttachmentCreated: session_id: str - incarnation: int attachment_generation: int resume_token: ResumeToken = field(repr=False) @@ -188,7 +187,6 @@ class DuplexSessionAttachmentCreated: @dataclass(frozen=True) class DuplexSessionResumeResult: session_id: str - incarnation: int attachment_generation: int resume_token: ResumeToken = field(repr=False) replay_entries: tuple[JournalEntry, ...] = () @@ -198,14 +196,12 @@ class DuplexSessionResumeResult: @dataclass class _DuplexSessionAttachmentState: session_id: str - incarnation: int credential: DuplexResumeCredential journal: DuplexEventJournal attachment: DuplexTransportAttachment | None attachment_generation: int outbound_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) recovery_token_digest: bytes | None = field(default=None, repr=False) - grace_task: asyncio.Task[None] | None = field(default=None, repr=False) class DuplexSessionAttachmentRegistry: @@ -214,18 +210,14 @@ class DuplexSessionAttachmentRegistry: *, replay_ttl_s: float, replay_max_bytes_per_session: int, - disconnect_grace_s: float = 30.0, clock: Callable[[], float] | None = None, ) -> None: if replay_ttl_s <= 0: raise ValueError("replay_ttl_s must be positive") if replay_max_bytes_per_session <= 0: raise ValueError("replay_max_bytes_per_session must be positive") - if disconnect_grace_s <= 0: - raise ValueError("disconnect_grace_s must be positive") self._replay_ttl_s = replay_ttl_s self._replay_max_bytes_per_session = replay_max_bytes_per_session - self._disconnect_grace_s = disconnect_grace_s self._clock = clock or time.monotonic self._sessions: dict[str, _DuplexSessionAttachmentState] = {} self._lock = asyncio.Lock() @@ -237,7 +229,6 @@ class DuplexSessionAttachmentRegistry: self, session_id: str, *, - incarnation: int, send: Callable[[dict[str, object]], Awaitable[None]], close: Callable[[str], Awaitable[None]], ) -> DuplexSessionAttachmentCreated: @@ -248,7 +239,6 @@ class DuplexSessionAttachmentRegistry: generation = 1 self._sessions[session_id] = _DuplexSessionAttachmentState( session_id=session_id, - incarnation=incarnation, credential=DuplexResumeCredential.from_token(token), journal=DuplexEventJournal( max_bytes=self._replay_max_bytes_per_session, @@ -264,7 +254,6 @@ class DuplexSessionAttachmentRegistry: ) return DuplexSessionAttachmentCreated( session_id=session_id, - incarnation=incarnation, attachment_generation=generation, resume_token=token, ) @@ -309,33 +298,39 @@ class DuplexSessionAttachmentRegistry: async with self._lock: return self._require(session_id).journal.acknowledge(sequence) - async def detach( - self, - session_id: str, - *, - attachment_generation: int, - on_grace_expired: Callable[[], Awaitable[None]] | None = None, - ) -> bool: + async def detach(self, session_id: str, *, attachment_generation: int | None = None) -> bool: + """Drop the current transport; the engine lease owns the disconnect grace. + + ``attachment_generation`` names the connection asking to detach, so a + socket that already lost a takeover cannot detach the winner. ``None`` + means "whichever connection is attached right now" and is for callers + that only know the session (the outbound pump, whose send just failed). + + Returns whether this call is the one that detached: an already-detached + session answers ``False`` so a second disconnect signal for the same + socket cannot restart the engine's disconnect grace window. + """ async with self._lock: state = self._sessions.get(session_id) - if state is None or state.attachment_generation != attachment_generation: + if state is None or state.attachment is None: + return False + if attachment_generation is not None and state.attachment_generation != attachment_generation: return False state.attachment = None - if state.grace_task is not None: - state.grace_task.cancel() - state.grace_task = ( - asyncio.create_task( - self._run_disconnect_grace( - state, - attachment_generation=attachment_generation, - callback=on_grace_expired, - ) - ) - if on_grace_expired is not None - else None - ) return True + async def has_attachment(self, session_id: str) -> bool: + """Whether some connection is attached right now, whoever it is. + + A resume that failed to activate has no generation of its own, so it + cannot ask ``is_current_attachment``. What it needs to know before + rolling the engine lease back into its disconnect grace is only whether + it would be rolling back somebody else's live attachment. + """ + async with self._lock: + state = self._sessions.get(session_id) + return state is not None and state.attachment is not None + async def is_current_attachment(self, session_id: str, attachment_generation: int) -> bool: async with self._lock: state = self._sessions.get(session_id) @@ -349,25 +344,19 @@ class DuplexSessionAttachmentRegistry: self, session_id: str, *, - incarnation: int, resume_token: str, last_received_server_event_seq: int, ) -> None: """Validate transport credentials before any engine resume control.""" async with self._lock: state = self._require(session_id) - self._validate_resume_identity( - state, - incarnation=incarnation, - resume_token=resume_token, - ) + self._validate_resume_identity(state, resume_token=resume_token) state.journal.replay_after(last_received_server_event_seq) async def resume( self, session_id: str, *, - incarnation: int, resume_token: str, last_received_server_event_seq: int, send: Callable[[dict[str, object]], Awaitable[None]], @@ -380,18 +369,11 @@ class DuplexSessionAttachmentRegistry: async with self._lock: if self._sessions.get(session_id) is not state: raise KeyError(f"unknown duplex attachment session: {session_id}") - used_recovery = self._validate_resume_identity( - state, - incarnation=incarnation, - resume_token=resume_token, - ) + used_recovery = self._validate_resume_identity(state, resume_token=resume_token) replay_entries = state.journal.replay_after(last_received_server_event_seq) accepted_token_digest = ( state.recovery_token_digest if used_recovery else bytes(state.credential.token_digest) ) - if state.grace_task is not None: - state.grace_task.cancel() - state.grace_task = None state.recovery_token_digest = None rotated_token = state.credential.rotate() replaced = state.attachment @@ -407,53 +389,53 @@ class DuplexSessionAttachmentRegistry: await send(dict(activation_payload_factory(rotated_token, attachment_generation))) for entry in replay_entries: await send(dict(entry.payload)) - except Exception: - async with self._lock: - if ( - self._sessions.get(session_id) is state - and state.attachment_generation == attachment_generation - ): - state.attachment = None - state.recovery_token_digest = accepted_token_digest + except (Exception, asyncio.CancelledError): + # CancelledError derives from BaseException, so an + # ``except Exception`` here would let a cancellation during + # activation or replay skip the rollback: the dead + # attachment would stay current while the token it rotated + # past is gone, leaving the session attached to an unusable + # transport with no way back. Both failures take the same + # path deliberately, so they cannot drift apart. + await self._rollback_resume( + session_id, + state=state, + attachment_generation=attachment_generation, + accepted_token_digest=accepted_token_digest, + ) raise return DuplexSessionResumeResult( session_id=session_id, - incarnation=incarnation, attachment_generation=attachment_generation, resume_token=rotated_token, replay_entries=replay_entries, replaced_attachment=replaced, ) + async def _rollback_resume( + self, + session_id: str, + *, + state: _DuplexSessionAttachmentState, + attachment_generation: int, + accepted_token_digest: bytes, + ) -> None: + """Undo a resume whose activation never reached the client. + + Generation-checked: a later resume that already took over must keep its + attachment. Restoring ``recovery_token_digest`` is what lets the caller + retry with the token it presented, which ``resume`` had rotated past. + """ + async with self._lock: + if self._sessions.get(session_id) is state and state.attachment_generation == attachment_generation: + state.attachment = None + state.recovery_token_digest = accepted_token_digest + async def close(self, session_id: str) -> DuplexTransportAttachment | None: async with self._lock: state = self._sessions.pop(session_id, None) - if state is not None and state.grace_task is not None: - state.grace_task.cancel() return state.attachment if state is not None else None - async def _run_disconnect_grace( - self, - state: _DuplexSessionAttachmentState, - *, - attachment_generation: int, - callback: Callable[[], Awaitable[None]] | None, - ) -> None: - try: - await asyncio.sleep(self._disconnect_grace_s) - except asyncio.CancelledError: - return - async with self._lock: - if ( - self._sessions.get(state.session_id) is not state - or state.attachment is not None - or state.attachment_generation != attachment_generation - ): - return - state.grace_task = None - if callback is not None: - await callback() - def _require(self, session_id: str) -> _DuplexSessionAttachmentState: state = self._sessions.get(session_id) if state is None: @@ -461,14 +443,7 @@ class DuplexSessionAttachmentRegistry: return state @staticmethod - def _validate_resume_identity( - state: _DuplexSessionAttachmentState, - *, - incarnation: int, - resume_token: str, - ) -> bool: - if incarnation != state.incarnation: - raise ValueError(f"duplex attachment incarnation mismatch: expected {state.incarnation}, got {incarnation}") + def _validate_resume_identity(state: _DuplexSessionAttachmentState, *, resume_token: str) -> bool: if state.credential.verify(resume_token): return False recovery_digest = state.recovery_token_digest diff --git a/vllm_omni/entrypoints/duplex/session_runner.py b/vllm_omni/entrypoints/duplex/session_runner.py deleted file mode 100644 index 90251e998..000000000 --- a/vllm_omni/entrypoints/duplex/session_runner.py +++ /dev/null @@ -1,2520 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -import binascii -import json -import time -import uuid -from collections.abc import Callable -from contextlib import suppress -from copy import deepcopy - -import numpy as np -import pybase64 as base64 -from fastapi import WebSocket, WebSocketDisconnect -from vllm.logger import init_logger - -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.entrypoints.duplex.audio import ( - convert_input_audio_with_rate, - validate_input_sample_rate_hz, -) -from vllm_omni.entrypoints.duplex.commit_policy import ( - CommitAction, - CommitSnapshot, - decide_commit_action, -) -from vllm_omni.entrypoints.duplex.protocol import ( - DuplexOverlapPolicy, - DuplexPlaybackCommitPolicy, - DuplexSession, - DuplexSessionState, - DuplexTurnEventType, -) -from vllm_omni.entrypoints.duplex.realtime_session import ( - NativeRealtimeSessionProtocol, -) -from vllm_omni.entrypoints.duplex.runtime_adapter import ( - PcmAppendReservation, - ServingRuntimeConfigError, - ServingRuntimeSessionState, - payload_turn_id, -) -from vllm_omni.entrypoints.duplex.server_vad import ServerVADPipeline -from vllm_omni.entrypoints.duplex.session_attachment import ( - DuplexJournalOverflowError, -) -from vllm_omni.entrypoints.duplex.websocket import ( - DuplexWebSocketActor, - is_input_event, - normalize_duplex_input_event, -) - -logger = init_logger(__name__) - -_MAX_EVENT_BYTES = 15 * 1024 * 1024 - - -class DuplexSessionRunnerMixin: - """Run one ordered WebSocket mailbox for a duplex session.""" - - async def handle_session( - self, - websocket: WebSocket, - *, - realtime_protocol: NativeRealtimeSessionProtocol | None = None, - ) -> None: - await websocket.accept() - session: DuplexSession | None = None - actor = DuplexWebSocketActor( - websocket, - current_epoch=lambda: session.epoch if session is not None else None, - session_closed=lambda: session is not None and session.state == DuplexSessionState.CLOSED, - outbound_protocol=realtime_protocol, - ) - runtime_opened = False - runtime_closed = False - attachment_ready = False - attachment_generation: int | None = None - resume_credential_delivered = False - transport_detached = False - pending_turn_reservations = 0 - - async def attachment_send(payload: dict[str, object]) -> None: - nonlocal resume_credential_delivered - try: - await websocket.send_json(payload) - except RuntimeError as exc: - message = str(exc) - if "after sending 'websocket.close'" in message or "response already completed" in message: - raise WebSocketDisconnect(code=1006) from exc - raise - if payload.get("type") == "session.created" and isinstance(payload.get("resume_token"), str): - resume_credential_delivered = True - - async def attachment_close(reason: str) -> None: - close = getattr(websocket, "close", None) - if callable(close): - await close(code=1000, reason=reason) - - async def send_attachment_payload( - payload: dict[str, object], *, journal: bool, on_accepted: Callable[[], None] | None = None - ) -> None: - assert session is not None - session_id = session.session_id - should_journal = journal and session_id not in self._resync_required_sessions - try: - await self._attachment_registry.send_event( - session_id, - payload, - journal=should_journal, - on_accepted=on_accepted, - ) - except DuplexJournalOverflowError: - first_overflow = session_id not in self._resync_required_sessions - self._resync_required_sessions.add(session_id) - if first_overflow: - resync_payload: dict[str, object] = { - "type": "session.resync_required", - "session_id": session_id, - "reason": "journal_overflow", - } - if realtime_protocol is not None: - resync_payload = realtime_protocol.encode_outbound_event(resync_payload)[0] - await self._attachment_registry.send_event( - session_id, - resync_payload, - journal=False, - ) - await self._attachment_registry.send_event( - session_id, - payload, - journal=False, - on_accepted=on_accepted, - ) - - if realtime_protocol is not None: - - async def send_realtime_raw(payload: dict[str, object]) -> None: - if attachment_ready and session is not None: - await send_attachment_payload( - payload, - journal=payload.get("type") not in {"session.created", "session.resumed"}, - ) - return - raw_payload = dict(payload) - raw_payload["_realtime_raw"] = True - await actor.send_json(raw_payload) - - realtime_protocol.bind_sender(send_realtime_raw) - native: ServingRuntimeSessionState = self._create_runtime_session_state() - - def begin_close(reason: str) -> None: - actor.closing = True - actor.close_reason = reason - if session is not None: - session.mark_closing() - - event_emit_lock = asyncio.Lock() - - async def send_outbound(payload: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: - if not attachment_ready or session is None: - await actor.send_json(dict(payload)) - if on_accepted is not None: - on_accepted() - return - if realtime_protocol is None: - await send_attachment_payload( - payload, - journal=payload.get("type") not in {"session.created", "session.resumed"}, - on_accepted=on_accepted, - ) - return - for projected in realtime_protocol.encode_outbound_event(payload): - await send_attachment_payload( - projected, - journal=projected.get("type") not in {"session.created", "session.resumed"}, - # Item/part announcements may precede the audio. Only the - # actual audio event establishes acceptance of this chunk. - on_accepted=on_accepted if projected.get("type") == "response.audio.delta" else None, - ) - - async def emit_event(payload: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: - deferred_precreate_response = False - async with event_emit_lock: - accepted, deferred_overlap_payload = await self._apply_outbound_session_event( - payload, - session=session, - actor=actor, - native=native, - realtime_protocol=realtime_protocol, - ) - if not accepted: - return - await send_outbound(payload, on_accepted=on_accepted) - if deferred_overlap_payload is not None: - deferred_precreate_response = native.deferred_precreate_response - native.deferred_precreate_response = False - if deferred_overlap_payload is not None and session is not None and not actor.closing: - await start_native_append( - deferred_overlap_payload, - final=True, - precreate_response=deferred_precreate_response, - operation_id=native.committed_audio_operation_id, - retained_committed_payload=( - deferred_overlap_payload if native.committed_audio_payload is deferred_overlap_payload else None - ), - ) - - writer_task = asyncio.create_task(actor.writer_loop(), name="duplex-session-writer") - reader_task: asyncio.Task[None] | None = None - - async def read_event_loop() -> None: - nonlocal pending_turn_reservations - assert session is not None - try: - while not actor.closing: - raw = await self._receive_text( - websocket, - session.config.idle_timeout_s, - realtime_protocol=realtime_protocol, - ) - if raw is None: - await actor.enqueue_event({"type": "__timeout__"}) - return - if len(raw.encode("utf-8")) > _MAX_EVENT_BYTES: - await emit_event( - {"type": "error", "error": "Duplex event too large", "code": "event_too_large"} - ) - continue - try: - event = json.loads(raw) - except json.JSONDecodeError: - await emit_event({"type": "error", "error": "Invalid JSON event", "code": "invalid_json"}) - continue - if not isinstance(event, dict): - await emit_event( - { - "type": "error", - "error": "Duplex event must be a JSON object", - "code": "bad_event", - } - ) - continue - event = normalize_duplex_input_event(event) - event_type = event.get("type") - if not isinstance(event_type, str): - await emit_event( - {"type": "error", "error": "Duplex event missing string type", "code": "bad_event"} - ) - continue - if event_type in {"input.commit", "input_audio_buffer.commit"}: - if not session.reserve_pending_turn( - limit=self._duplex_session_config.max_pending_turns_per_session - ): - await emit_event( - { - "type": "error", - "error": "Duplex session has too many pending input turns", - "code": "input_backpressure", - } - ) - continue - event["_duplex_pending_turn_reserved"] = True - pending_turn_reservations += 1 - if is_input_event(event_type) and native_response_in_progress(): - event["_duplex_overlap_candidate"] = True - await actor.enqueue_event(event) - except WebSocketDisconnect: - await actor.enqueue_event({"type": "__disconnect__"}) - except Exception as exc: - await emit_event({"type": "error", "error": str(exc), "code": "realtime_input_failed"}) - await actor.enqueue_event({"type": "__disconnect__"}) - - async def next_actor_event() -> dict[str, object]: - nonlocal pending_turn_reservations, transport_detached - event = await actor.next_event() - if event.pop("_duplex_pending_turn_reserved", False) and session is not None: - session.release_pending_turn() - pending_turn_reservations = max(0, pending_turn_reservations - 1) - if ( - attachment_ready - and attachment_generation is not None - and session is not None - and not await self._attachment_registry.is_current_attachment( - session.session_id, - attachment_generation, - ) - ): - transport_detached = True - return {"type": "__replaced_attachment__"} - return event - - def native_response_in_progress() -> bool: - if session is None: - return False - if session.active_response_id is not None: - return True - if ( - session.config.playback_commit_policy == DuplexPlaybackCommitPolicy.ACK_ONLY.value - and session.playback.sent_ms > session.playback.committed_ms - ): - return True - if actor.active_response_task is not None and not actor.active_response_task.done(): - return True - if actor.has_response_bound_append_tasks(): - return True - return False - - def start_chat_response() -> None: - assert session is not None - - async def run_and_notify() -> None: - task = asyncio.current_task() - try: - await self._run_response(session, emit_event) - finally: - await actor.enqueue_event( - { - "type": "__chat_response_finished__", - "_response_task": task, - } - ) - - actor.active_response_task = asyncio.create_task(run_and_notify()) - - def activate_waiting_server_vad_turn() -> None: - assert session is not None - pending_turn = session.activate_server_vad_turn() - if pending_turn is not None and pending_turn.create_response: - start_chat_response() - - def correlate_realtime_error( - payload: dict[str, object], - event_id: object, - ) -> dict[str, object]: - if isinstance(event_id, str) and event_id: - payload["event_id"] = event_id - return payload - - async def reject_and_clear_server_vad_input( - pipeline: ServerVADPipeline, - *, - code: str, - error: str, - metric_reason: str, - event_id: object, - ) -> None: - assert session is not None - session.cancel_pending_input() - pipeline.reset() - self._realtime_vad_metrics.error(metric_reason) - await emit_event( - correlate_realtime_error( - { - "type": "error", - "session_id": session.session_id, - "code": code, - "error": error, - }, - event_id, - ) - ) - await emit_event( - { - "type": "input_audio_buffer.cleared", - "session_id": session.session_id, - } - ) - - async def commit_server_vad_turn( - *, - item_id: str, - create_response: bool, - retain_input_bytes: int, - event_id: object = None, - ) -> None: - assert session is not None - # Retain ownership until the actor consumes the terminal event so - # queued audio cannot overtake pending-turn activation. - active_response = actor.active_response_task is not None - if active_response and session.pending_server_vad_turn is not None: - released = session.discard_uncommitted_server_vad_utterance() - session.release_input_bytes(released) - self._realtime_vad_metrics.error("pending_turn_backpressure") - await emit_event( - correlate_realtime_error( - { - "type": "error", - "session_id": session.session_id, - "code": "input_backpressure", - "error": "Only one committed Server VAD turn may wait for the active response", - }, - event_id, - ) - ) - return - if not session.stage_server_vad_audio_for_commit(): - return - committed = session.commit_user_input( - append_history=not active_response, - retain_input_bytes=retain_input_bytes, - ) - if committed is None: - return - if active_response: - session.stage_server_vad_turn( - committed, - create_response=create_response, - item_id=item_id, - ) - else: - session.register_history_item(item_id, committed.message) - await emit_event( - self._input_committed_payload( - session, - committed, - item_id=item_id, - ) - ) - if create_response and not active_response: - start_chat_response() - - def clear_completed_pending_silence() -> None: - task = native.pending_silence_task - if task is not None and task.done(): - native.pending_silence_task = None - native.pending_silence_owner_id = None - - def mark_pending_silence_superseded() -> None: - task = native.pending_silence_task - if task is None: - return - if task.done(): - native.pending_silence_task = None - # Do not cancel the task or clear native_append_tail here. A - # silence append may already have reached the Engine, and local - # cancellation cannot retract that RPC. The sequencer preserves - # wire order; before_append will skip silence that has not started. - native.pending_silence_owner_id = None - - def real_native_input_waiting() -> bool: - clear_completed_pending_silence() - return ( - native.audio_buffer.has_pending() - or native.audio_buffer.has_reserved() - or actor.has_queued_input_events() - ) - - async def start_native_append( - payload: object, - *, - final: bool, - precreate_response: bool = False, - pcm_reservation: PcmAppendReservation | None = None, - operation_id: str | None = None, - retained_committed_payload: dict[str, object] | None = None, - silence_continuation: bool = False, - before_append=None, - ) -> asyncio.Task[bool] | None: - if session is None: - return - if not silence_continuation: - mark_pending_silence_superseded() - append_epoch = session.epoch - append_turn_id = payload_turn_id(payload) - if append_turn_id is None: - append_turn_id = session.turn_id - request_id = self._native_stage0_request_id(session, append_epoch) - if final or precreate_response: - session.bind_request(request_id) - if precreate_response: - session.bind_response_turn(append_turn_id) - if precreate_response and session.active_response_id is None: - response_id = session.begin_response(turn_id=append_turn_id) - await emit_event( - self._response_created_payload( - session, - response_id, - epoch=append_epoch, - ) - ) - precreated_response_id = session.active_response_id if precreate_response else None - - def _discard_retained_committed_audio() -> None: - if ( - retained_committed_payload is not None - and native.committed_audio_payload is retained_committed_payload - ): - session.release_input_bytes(native.clear_committed_audio()) - - async def _run() -> bool: - nonlocal runtime_closed - try: - session.mark_model_turn_request_started(append_turn_id, time.monotonic()) - append_ok, emitted_response = await self._append_runtime_input( - session, - payload, - operation_id=(pcm_reservation.operation_id if pcm_reservation is not None else operation_id), - final=final, - send_json=emit_event, - mode="append_audio_chunk", - expected_epoch=append_epoch, - ) - if append_ok: - native.native_context_locked = True - if pcm_reservation is not None: - pcm_reservation.commit() - session.release_input_bytes(pcm_reservation.byte_count) - if ( - retained_committed_payload is not None - and native.committed_audio_payload is retained_committed_payload - ): - session.release_input_bytes(native.clear_committed_audio()) - else: - if pcm_reservation is not None: - pcm_reservation.rollback() - _discard_retained_committed_audio() - if ( - not append_ok - and precreated_response_id is not None - and session.active_response_id == precreated_response_id - ): - session.end_response(commit_text=False) - await emit_event( - { - "type": "response.done", - "session_id": session.session_id, - "response_id": precreated_response_id, - "epoch": session.epoch, - "committed": False, - "status": "failed", - "status_details": { - "type": "failed", - "reason": "runtime_append_failed", - }, - "playback": session.playback.as_dict(), - } - ) - if not append_ok and session.state == DuplexSessionState.CLOSED: - runtime_closed = True - return False - if not emitted_response and session.epoch == append_epoch: - if session.active_request_id == self._native_stage0_request_id(session, append_epoch): - session.clear_request(request_id) - if final: - await emit_event( - self._turn_controller.signal(session, DuplexTurnEventType.USER_STARTED.value) - ) - return append_ok - except asyncio.CancelledError: - if pcm_reservation is not None: - pcm_reservation.rollback() - raise - except Exception as exc: - if pcm_reservation is not None: - pcm_reservation.rollback() - _discard_retained_committed_audio() - logger.exception("Native duplex append task failed: %s", exc) - await self._send_runtime_error(emit_event, "runtime_append_task_failed", exc, session=session) - if session.state != DuplexSessionState.CLOSED: - begin_close("runtime_append_task_failed") - if await self._close_runtime_session( - session, - reason="runtime_append_task_failed", - send_json=emit_event, - ): - runtime_closed = True - session.close() - await emit_event( - { - "type": "session.closed", - "session_id": session.session_id, - "reason": "runtime_append_task_failed", - } - ) - return False - - async def _run_in_wire_order(predecessor: asyncio.Task[bool] | None) -> bool: - if predecessor is not None: - try: - predecessor_ok = await predecessor - except asyncio.CancelledError: - current = asyncio.current_task() - if current is not None and current.cancelling(): - raise - predecessor_ok = False - except Exception: - predecessor_ok = False - if not predecessor_ok: - if pcm_reservation is not None: - pcm_reservation.rollback() - _discard_retained_committed_audio() - return False - if actor.closing or runtime_closed or session.state != DuplexSessionState.OPEN: - if pcm_reservation is not None: - pcm_reservation.rollback() - _discard_retained_committed_audio() - return False - if before_append is not None and not before_append(): - if pcm_reservation is not None: - pcm_reservation.rollback() - _discard_retained_committed_audio() - return True - if pcm_reservation is not None and not pcm_reservation.active: - _discard_retained_committed_audio() - return False - return await _run() - - def _release_cancelled_retained_audio(done: asyncio.Task[bool]) -> None: - if done.cancelled(): - _discard_retained_committed_audio() - return - try: - append_ok = done.result() - except Exception: - append_ok = False - if not append_ok: - _discard_retained_committed_audio() - - predecessor = actor.native_append_tail - if predecessor is not None and predecessor.done(): - try: - predecessor_ok = predecessor.result() - except (asyncio.CancelledError, Exception): - predecessor_ok = False - if not predecessor_ok: - # Appends already queued behind a failed predecessor keep - # their captured dependency and stop. A later wire command - # is an explicit retry and starts a new chain. - predecessor = None - task = asyncio.create_task(_run_in_wire_order(predecessor)) - task.add_done_callback(_release_cancelled_retained_audio) - actor.native_append_tail = task - actor.track_append_task( - task, - epoch=append_epoch, - mode="append_audio_chunk", - final=final, - response_bound=final or precreate_response, - ) - if silence_continuation: - native.pending_silence_task = task - - def _clear_done_pending_silence(done: asyncio.Task[bool]) -> None: - if native.pending_silence_task is done: - native.pending_silence_task = None - native.pending_silence_owner_id = None - - task.add_done_callback(_clear_done_pending_silence) - # Let this wire-order effect start before the next mailbox event can - # cancel it. Later native appends still serialize on predecessor. - await asyncio.sleep(0) - return task - - async def schedule_native_silence_continuation( - payload: object, - *, - request_id: str, - owner_id: str, - response_id: str | None, - response_owned: bool, - expected_epoch: int | None, - expected_incarnation: int, - expected_model_turn_id: int | None, - send_json, - ) -> bool: - del send_json - if session is None: - return False - clear_completed_pending_silence() - pending_silence = native.pending_silence_task - if pending_silence is not None and not pending_silence.done(): - if pending_silence is asyncio.current_task(): - return False - try: - if not await pending_silence: - return False - except asyncio.CancelledError: - current = asyncio.current_task() - if current is not None and current.cancelling(): - raise - return False - except Exception: - return False - clear_completed_pending_silence() - pending_silence = native.pending_silence_task - if pending_silence is not None and not pending_silence.done(): - return False - append_tail = actor.native_append_tail - if (append_tail is None or append_tail.done()) and real_native_input_waiting(): - return False - continuation_delay_s = max( - 0.0, - float(session.capabilities.chunk_period_ms or 1000) / 1000.0, - ) - if continuation_delay_s > 0: - await asyncio.sleep(continuation_delay_s) - if ( - actor.native_append_tail is not append_tail - or ((append_tail is None or append_tail.done()) and real_native_input_waiting()) - or self._native_silence_continuation_is_stale( - session, - request_id=request_id, - response_id=response_id, - response_owned=response_owned, - expected_epoch=expected_epoch, - expected_incarnation=expected_incarnation, - expected_model_turn_id=expected_model_turn_id, - ) - ): - return False - - def _still_valid() -> bool: - return not real_native_input_waiting() and not self._native_silence_continuation_is_stale( - session, - request_id=request_id, - response_id=response_id, - response_owned=response_owned, - expected_epoch=expected_epoch, - expected_incarnation=expected_incarnation, - expected_model_turn_id=expected_model_turn_id, - ) - - native.pending_silence_owner_id = owner_id - task = await start_native_append( - payload, - final=False, - silence_continuation=True, - before_append=_still_valid, - ) - if task is None: - native.pending_silence_owner_id = None - return False - return True - - async def wait_for_native_append_tail() -> bool: - predecessor = actor.native_append_tail - if predecessor is None: - return True - try: - return await predecessor - except asyncio.CancelledError: - current = asyncio.current_task() - if current is not None and current.cancelling(): - raise - return False - except Exception: - # The append task already emitted its runtime error. Preserve - # wire order before the update path decides whether to continue. - return False - - async def start_runtime_append( - payload: object, - *, - final: bool, - mode: str = "append_tokens", - ) -> None: - """Schedule non-native runtime appends without blocking WS input. - - Model-native appends use ``start_native_append`` because they can - create data-plane response streams. Generic runtime appends still - need the same control-plane isolation so cancel/close can preempt a - slow engine append. - """ - if session is None: - return - append_epoch = session.epoch - - async def _run() -> None: - nonlocal runtime_closed - try: - append_ok, emitted_response = await self._append_runtime_input( - session, - payload, - final=final, - send_json=emit_event, - mode=mode, - expected_epoch=append_epoch, - ) - if not append_ok: - if session.state == DuplexSessionState.CLOSED: - runtime_closed = True - return - if not emitted_response and session.epoch == append_epoch: - await emit_event(self._turn_controller.signal(session, DuplexTurnEventType.USER_STARTED.value)) - except asyncio.CancelledError: - raise - except Exception as exc: - logger.exception("Duplex runtime append task failed: %s", exc) - await self._send_runtime_error(emit_event, "runtime_append_task_failed", exc, session=session) - - task = asyncio.create_task(_run()) - actor.track_append_task( - task, - epoch=append_epoch, - mode=mode, - final=final, - response_bound=final, - ) - - try: - handshake = await self._open_session( - websocket, - emit_event, - realtime_protocol=realtime_protocol, - attachment_send=attachment_send, - attachment_close=attachment_close, - ) - if handshake is None: - return - session = handshake.session - if initial_update_error := self._runtime_session_candidate_update_error( - session, session.config, realtime_protocol=realtime_protocol - ): - await emit_event( - correlate_realtime_error( - initial_update_error, - handshake.event_id, - ) - ) - return - if handshake.resumed: - native = self._serving_session_states[session.session_id] - actor.tasks = self._session_tasks[session.session_id] - persisted_protocol = self._realtime_protocols.get(session.session_id) - if persisted_protocol is None: - raise RuntimeError(f"Missing Realtime protocol state for resumed session {session.session_id}") - realtime_protocol = persisted_protocol - realtime_protocol.bind_sender(send_realtime_raw) - attachment_generation = handshake.attachment_generation - attachment_ready = True - resume_credential_delivered = True - runtime_opened = True - self._ensure_lifecycle_listener() - reader_task = asyncio.create_task(read_event_loop(), name="duplex-session-reader") - else: - self._serving_session_states[session.session_id] = native - self._session_tasks[session.session_id] = actor.tasks - if realtime_protocol is not None: - self._realtime_protocols[session.session_id] = realtime_protocol - native.silence_continuation_scheduler = schedule_native_silence_continuation - if realtime_protocol is not None: - realtime_protocol.bind_native_input_append(self._uses_native_input_append(session)) - session.config.playback_commit_policy = DuplexPlaybackCommitPolicy.ACK_ONLY.value - if not handshake.resumed: - open_result = await self._open_runtime_session(session, emit_event) - if open_result is False: - return - runtime_opened = True - try: - await self._configure_server_vad(session) - except Exception as exc: - logger.exception("Server VAD initialization failed: %s", exc) - self._realtime_vad_metrics.error("initialization") - await emit_event( - correlate_realtime_error( - { - "type": "error", - "session_id": session.session_id, - "code": "server_vad_initialization_failed", - "error": str(exc), - }, - handshake.event_id, - ) - ) - return - created_attachment = await self._attachment_registry.create( - session.session_id, - incarnation=session.incarnation, - send=attachment_send, - close=attachment_close, - ) - attachment_generation = created_attachment.attachment_generation - attachment_ready = True - self._lease_generations[session.session_id] = 0 - self._ensure_lifecycle_listener() - created_payload: dict[str, object] = { - "type": "session.created", - "session": session.as_public_dict(), - } - if session.capabilities.supports_session_resume: - created_payload.update( - { - "incarnation": session.incarnation, - "attachment_generation": attachment_generation, - "resume_token": created_attachment.resume_token.plaintext, - } - ) - if isinstance(open_result, dict): - created_payload["runtime_control"] = self._redact_runtime_control_result(open_result) - if realtime_protocol is not None: - realtime_protocol.commit_realtime_turn_detection_update(session.config.server_vad) - await emit_event(created_payload) - resume_credential_delivered = session.capabilities.supports_session_resume - reader_task = asyncio.create_task(read_event_loop(), name="duplex-session-reader") - - while True: - event = await next_actor_event() - event_type = event.get("type") - - if event_type == "__replaced_attachment__": - return - - if event_type == "__chat_response_finished__": - completed_task = event.get("_response_task") - if completed_task is not actor.active_response_task: - continue - actor.active_response_task = None - activate_waiting_server_vad_turn() - continue - - if event_type == "__timeout__": - begin_close("timeout") - native.audio_buffer.clear() - session.release_all_input_bytes() - native.input_since_commit = False - native.speech_since_commit = False - native.clear_committed_audio() - await actor.cancel_append_tasks() - await self._cancel_native_data_plane_stream(session) - await self._cancel_active_response( - session, - actor.active_response_task, - emit_event, - reason="timeout", - notify=True, - ) - actor.active_response_task = None - runtime_closed = await self._close_runtime_session( - session, - reason="timeout", - send_json=emit_event, - ) - if not runtime_closed: - return - session.close() - await emit_event( - { - "type": "session.closed", - "session_id": session.session_id, - "reason": "timeout", - } - ) - return - - if event_type == "__disconnect__": - transport_detached = True - return - - if not isinstance(event_type, str): - continue - - if event_type == "session.heartbeat": - touch_session = getattr(self._chat_service.engine_client, "touch_duplex_session_async", None) - if not callable(touch_session): - await emit_event( - { - "type": "error", - "error": "Duplex runtime does not expose heartbeat control", - "code": "runtime_touch_unsupported", - } - ) - continue - try: - await touch_session( - session.session_id, - fence=DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ), - activity=DuplexLeaseActivity.HEARTBEAT, - ) - except Exception as exc: - await emit_event( - { - "type": "error", - "error": str(exc), - "code": "runtime_touch_failed", - } - ) - continue - await emit_event( - { - "type": "session.heartbeat_ack", - "session_id": session.session_id, - } - ) - continue - - if event_type == "session.event_ack": - acknowledged = event.get("server_event_seq") - if not isinstance(acknowledged, int) or acknowledged < 0: - await emit_event( - { - "type": "error", - "error": "session.event_ack requires a non-negative server_event_seq", - "code": "invalid_event_ack", - } - ) - continue - try: - await self._attachment_registry.acknowledge(session.session_id, acknowledged) - except ValueError as exc: - await emit_event( - { - "type": "error", - "error": str(exc), - "code": "invalid_event_ack", - } - ) - continue - - if event_type == "turn.signal" and event.get("event") in { - "input.cancel", - "response.cancel", - "barge_in", - }: - payload = event.get("payload") - normalized_event = dict(payload) if isinstance(payload, dict) else {} - normalized_event.update(event) - normalized_event["type"] = event["event"] - event = normalized_event - event_type = str(event["type"]) - - if event_type == "session.close": - with suppress(asyncio.TimeoutError): - await asyncio.wait_for(actor.output_queue.join(), timeout=1.0) - begin_close("session_close") - if session.state == DuplexSessionState.CLOSED: - runtime_closed = True - return - native.audio_buffer.clear() - session.release_all_input_bytes() - native.input_since_commit = False - native.speech_since_commit = False - native.clear_committed_audio() - await actor.cancel_append_tasks() - await self._cancel_native_data_plane_stream(session) - await self._cancel_active_response( - session, - actor.active_response_task, - emit_event, - reason="session_close", - ) - runtime_closed = await self._close_runtime_session( - session, - reason="session_close", - send_json=emit_event, - ) - actor.active_response_task = None - if not runtime_closed: - return - await emit_event({"type": "session.closed", "session_id": session.session_id}) - session.close() - return - - if event_type == "input_audio_buffer.clear": - native.audio_buffer.clear() - session.release_all_input_bytes() - server_vad_pipeline = self._server_vad_pipelines.get(session.session_id) - if server_vad_pipeline is not None: - server_vad_pipeline.reset() - native.input_since_commit = False - native.speech_since_commit = False - native.clear_committed_audio() - cancelled = session.cancel_pending_input() - await emit_event( - { - "type": "input_audio_buffer.cleared", - "session_id": session.session_id, - "epoch": session.epoch, - "drained_input_events": 0, - "cancelled": cancelled, - } - ) - continue - - if event_type == "barge_in" and not session.capabilities.supports_barge_in: - await emit_event(self._barge_in_unsupported_error(session)) - continue - - if event_type in {"input.cancel", "response.cancel", "barge_in", "output_audio_buffer.clear"}: - cancel_reason = ( - "output_audio_buffer_clear" - if event_type == "output_audio_buffer.clear" - else "client_cancelled" - if event_type == "response.cancel" - else "barge_in" - ) - cancelled_fence = DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ) - if event_type == "response.cancel": - requested_response_id = event.get("response_id") - has_active_response_work = native_response_in_progress() - if isinstance(requested_response_id, str) and ( - session.active_response_id is None or requested_response_id != session.active_response_id - ): - await emit_event( - { - "type": "error", - "session_id": session.session_id, - "code": "response_not_active", - "error": f"Response is not active: {requested_response_id}", - } - ) - continue - if event_type == "response.cancel": - if not has_active_response_work: - if realtime_protocol is not None and isinstance(requested_response_id, str): - continue - await emit_event( - { - "type": "error", - "session_id": session.session_id, - "code": "response_not_active", - "error": "response.cancel requires an active response", - } - ) - continue - had_native_unbuffered_append = ( - self._uses_native_input_append(session) - and native.input_since_commit - and not native.audio_buffer.has_pending() - ) - playback_was_active = self._assistant_playback_active(session) - if event_type in {"input.cancel", "barge_in"}: - native.audio_buffer.clear() - session.release_all_input_bytes() - native.input_since_commit = False - native.speech_since_commit = False - native.clear_committed_audio() - had_native_append = await actor.cancel_append_tasks( - response_bound_only=event_type in {"response.cancel", "output_audio_buffer.clear"}, - ) - if event_type == "response.cancel": - session.release_input_bytes(native.clear_committed_audio()) - had_native_stream = native.data_plane_task is not None - cancelled = await self._cancel_active_response( - session, - actor.active_response_task, - emit_event, - reason=cancel_reason, - ) - had_native_stream = await self._cancel_native_data_plane_stream(session) or had_native_stream - if not cancelled and (had_native_stream or had_native_append or had_native_unbuffered_append): - old_epoch = session.epoch - old_response_id = session.active_response_id - committed_ms = session.playback.committed_ms - self._commit_played_response_history(session, old_response_id, committed_ms) - new_epoch, old_playback = self._advance_barge_in_epoch(session) - await emit_event( - { - "type": "audio.cancelled", - "session_id": session.session_id, - "response_id": old_response_id, - "reason": cancel_reason, - "cancelled_epoch": old_epoch, - "epoch": new_epoch, - "committed_ms": committed_ms, - "playback": old_playback, - } - ) - cancelled = True - if not cancelled and playback_was_active: - old_epoch = session.epoch - committed_ms = session.playback.committed_ms - self._commit_played_response_history(session, session.last_response_id, committed_ms) - new_epoch, old_playback = self._advance_barge_in_epoch(session) - await emit_event( - { - "type": "audio.cancelled", - "session_id": session.session_id, - "response_id": session.last_response_id, - "reason": cancel_reason, - "cancelled_epoch": old_epoch, - "epoch": new_epoch, - "committed_ms": committed_ms, - "playback": old_playback, - } - ) - cancelled = True - if not cancelled and event_type == "response.cancel": - old_epoch = session.epoch - old_response_id = session.active_response_id - committed_ms = session.playback.committed_ms - self._commit_played_response_history(session, old_response_id, committed_ms) - new_epoch, old_playback = self._advance_barge_in_epoch(session) - await emit_event( - { - "type": "audio.cancelled", - "session_id": session.session_id, - "response_id": old_response_id, - "reason": cancel_reason, - "cancelled_epoch": old_epoch, - "epoch": new_epoch, - "committed_ms": committed_ms, - "playback": old_playback, - } - ) - cancelled = True - if not cancelled and event_type == "output_audio_buffer.clear": - old_playback = session.playback.as_dict() - committed_ms = session.playback.committed_ms - session.clear_playback_cursor() - await emit_event( - { - "type": "audio.cancelled", - "session_id": session.session_id, - "response_id": session.active_response_id, - "reason": cancel_reason, - "cancelled_epoch": session.epoch, - "epoch": session.epoch, - "committed_ms": committed_ms, - "playback": old_playback, - } - ) - actor.active_response_task = None - activate_waiting_server_vad_turn() - continue - if not cancelled: - await self._cancel_pending_input(session, emit_event, reason="barge_in") - if not await self._signal_runtime_session( - session, - "barge_in", - emit_event, - fence=cancelled_fence, - next_fence=( - DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ) - if session.epoch > cancelled_fence.epoch - else None - ), - ): - continue - actor.active_response_task = None - activate_waiting_server_vad_turn() - continue - - if event_type == "turn.signal": - turn_event = event.get("event") - if isinstance(turn_event, str): - if turn_event == "barge_in" and not session.capabilities.supports_barge_in: - await emit_event(self._barge_in_unsupported_error(session)) - continue - if turn_event == "session.update": - realtime_event_id = event.get("realtime_event_id") - - async def emit_update_event(update_event: dict[str, object]) -> None: - if isinstance(realtime_event_id, str) and realtime_event_id: - update_event = {**update_event, "realtime_event_id": realtime_event_id} - await emit_event(update_event) - - def reject_update() -> None: - if realtime_protocol is not None: - realtime_protocol.reject_realtime_turn_detection_update() - - payload = event.get("payload") - if not isinstance(payload, dict): - await emit_update_event( - { - "type": "error", - "session_id": session.session_id, - "code": "bad_event", - "error": "session.update requires a session payload", - } - ) - reject_update() - continue - if not await wait_for_native_append_tail(): - await emit_update_event( - { - "type": "error", - "code": "session_update_aborted", - "error": "session.update was not applied because the preceding append failed", - } - ) - reject_update() - continue - if ( - realtime_protocol is not None - and not self._uses_native_input_append(session) - and session.config.server_vad is not None - and session.turn_detection_config_locked - and realtime_protocol.pending_realtime_input_audio_config_changed() - ): - await emit_update_event( - { - "type": "error", - "session_id": session.session_id, - "code": "unsupported_audio_format", - "error": ( - "session.update cannot change Server VAD input format or sample rate " - "after the first audio append" - ), - "param": "audio.input.format", - } - ) - reject_update() - continue - runtime_update_error = self._runtime_session_update_error(session, payload) - if runtime_update_error is not None: - await emit_update_event(runtime_update_error) - reject_update() - continue - previous_config = session.config - candidate_config = deepcopy(previous_config) - session.replace_config(candidate_config) - try: - update_error = self._apply_session_update(session, payload) - finally: - session.replace_config(previous_config) - server_vad_changed = candidate_config.server_vad != previous_config.server_vad - if update_error is not None: - await emit_update_event(update_error) - reject_update() - continue - runtime_update_error = self._runtime_session_candidate_update_error( - session, - candidate_config, - realtime_protocol=realtime_protocol, - ) - if runtime_update_error is not None: - await emit_update_event(runtime_update_error) - reject_update() - continue - try: - candidate_runtime_config = self._runtime_config_for_session_update( - session, - candidate_config, - ) - except ServingRuntimeConfigError as exc: - await emit_update_event( - { - "type": "error", - "session_id": session.session_id, - "code": exc.code, - "error": str(exc), - } - ) - reject_update() - continue - candidate_server_vad_pipeline = None - if server_vad_changed: - try: - candidate_server_vad_pipeline = await self._prepare_server_vad_pipeline( - session, - candidate_config, - ) - except Exception as exc: - logger.exception("Server VAD session update failed: %s", exc) - self._realtime_vad_metrics.error("initialization") - await emit_update_event( - { - "type": "error", - "session_id": session.session_id, - "code": "server_vad_initialization_failed", - "error": str(exc), - } - ) - reject_update() - continue - runtime_updated = await self._signal_runtime_session( - session, - turn_event, - emit_update_event, - session_config=candidate_config.as_dict(), - runtime_config=candidate_runtime_config, - ) - if not runtime_updated: - reject_update() - continue - session.replace_config(candidate_config) - session.replace_runtime_config(candidate_runtime_config) - if server_vad_changed: - self._install_server_vad_pipeline( - session, - candidate_server_vad_pipeline, - ) - if realtime_protocol is not None: - realtime_protocol.bind_native_input_append(self._uses_native_input_append(session)) - realtime_protocol.commit_realtime_turn_detection_update(session.config.server_vad) - await emit_update_event( - { - "type": "session.updated", - "session": session.as_public_dict(), - } - ) - continue - if turn_event == "conversation.item.create": - payload = event.get("payload") - item = payload.get("item") if isinstance(payload, dict) else None - item_type = item.get("type") if isinstance(item, dict) else None - prepare_function_output = getattr( - self._serving_runtime_adapter, - "runtime_config_for_function_output", - None, - ) - if item_type == "function_call_output" and callable(prepare_function_output): - if not await wait_for_native_append_tail(): - continue - try: - candidate_runtime_config = prepare_function_output( - item, - dict(session.runtime_config), - ) - except ServingRuntimeConfigError as exc: - await emit_event( - { - "type": "error", - "session_id": session.session_id, - "code": exc.code, - "error": str(exc), - } - ) - continue - if not await self._signal_runtime_session( - session, - "session.update", - emit_event, - runtime_config=candidate_runtime_config, - ): - continue - session.replace_runtime_config(candidate_runtime_config) - await emit_event( - { - "type": "conversation.item.created", - "session_id": session.session_id, - "item": item, - "created": True, - } - ) - await self._maybe_continue_native_response( - emit_event, - session=session, - expected_epoch=session.epoch, - expected_model_turn_id=session.turn_id, - ) - continue - history_item = payload.get("history_item") if isinstance(payload, dict) else None - if not isinstance(history_item, dict): - history_item = item - message = self._realtime_item_to_history_message(history_item) - item_id = item.get("id") if isinstance(item, dict) else None - if message is not None: - session.append_history_message(message) - session.register_history_item(item_id if isinstance(item_id, str) else None, message) - await emit_event( - { - "type": "conversation.item.created", - "session_id": session.session_id, - "item": item, - "created": message is not None, - } - ) - continue - if turn_event == "conversation.item.delete": - payload = event.get("payload") - item_id = payload.get("item_id") if isinstance(payload, dict) else None - deleted = session.delete_history_item(item_id) if isinstance(item_id, str) else False - await emit_event( - { - "type": "conversation.item.deleted", - "session_id": session.session_id, - "item_id": item_id, - "deleted": deleted, - } - ) - continue - if turn_event == "conversation.item.truncate": - payload = event.get("payload") - item_id = payload.get("item_id") if isinstance(payload, dict) else None - audio_end_ms = payload.get("audio_end_ms") if isinstance(payload, dict) else None - truncated = ( - session.truncate_history_item( - item_id, - audio_end_ms=int(audio_end_ms) if isinstance(audio_end_ms, int | float) else 0, - hard=True, - ) - if isinstance(item_id, str) - else False - ) - await emit_event( - { - "type": "conversation.item.truncated", - "session_id": session.session_id, - "item_id": item_id, - "content_index": ( - payload.get("content_index", 0) if isinstance(payload, dict) else 0 - ), - "audio_end_ms": audio_end_ms, - "truncated": truncated, - } - ) - continue - await emit_event(self._turn_controller.signal(session, turn_event, event)) - else: - await emit_event({"type": "error", "error": "turn.signal requires event", "code": "bad_event"}) - continue - - if event_type == "playback.ack": - await self._handle_playback_ack(session, event, emit_event) - continue - - if event_type == "input.text.append": - session.mark_user_input_activity() - text = event.get("text") - if not isinstance(text, str): - await emit_event( - { - "type": "error", - "error": "input.text.append requires text", - "code": "bad_event", - } - ) - continue - if self._uses_native_input_append(session): - await emit_event( - { - "type": "error", - "error": "The selected native duplex runtime accepts audio append only", - "code": "native_text_append_unsupported", - } - ) - continue - else: - session.append_text(text) - if session.capabilities.supports_input_append: - await start_runtime_append(text, final=False, mode="append_tokens") - continue - - if event_type == "input_audio_buffer.append": - session.mark_user_input_activity() - audio = event.get("audio") or event.get("data") - if not isinstance(audio, str): - await emit_event( - { - "type": "error", - "error": "input.audio.append requires audio", - "code": "bad_event", - } - ) - continue - if not session.capabilities.supports_barge_in and self._event_requests_barge_in(event): - await emit_event(self._barge_in_unsupported_error(session)) - event = dict(event) - event.pop("force_barge_in", None) - for key in ("overlap_action", "overlap"): - value = event.get(key) - if isinstance(value, str) and value.strip().lower() in { - "barge_in", - "interrupt", - "cancel", - }: - event.pop(key, None) - fmt = event.get("format") if isinstance(event.get("format"), str) else "pcm16" - default_sample_rate_hz = 16000 - sr_raw = event.get("sample_rate_hz") or event.get("sample_rate") - sample_rate_hz = sr_raw if isinstance(sr_raw, int | float) else default_sample_rate_hz - native_input = self._uses_native_input_append(session) - server_vad_config = session.config.server_vad - turn_based_server_vad = server_vad_config is not None and not native_input - if not turn_based_server_vad: - try: - audio, fmt, sample_rate_hz = convert_input_audio_with_rate( - audio, - fmt, - sample_rate_hz=sample_rate_hz, - ) - except ValueError as exc: - await emit_event({"type": "error", "error": str(exc), "code": "bad_event"}) - continue - if isinstance(fmt, str) and fmt.lower() in {"pcm16", "pcm_s16le", "s16le"}: - await emit_event( - { - "type": "error", - "error": "input_audio_buffer.append pcm16 audio could not be decoded", - "code": "bad_audio", - } - ) - continue - if server_vad_config is None: - session.lock_turn_detection_config() - force_listen = bool(event.get("force_listen", False)) - payload = { - "type": "audio", - "audio": audio, - "format": fmt, - "sample_rate_hz": sample_rate_hz, - "force_listen": force_listen, - } - video_frames = event.get("video_frames") - if isinstance(video_frames, list): - frames = [frame for frame in video_frames if isinstance(frame, str) and frame] - if frames: - payload["video_frames"] = frames - if server_vad_config is not None: - realtime_event_id = event.get("_realtime_event_id") - pipeline = self._server_vad_pipelines.get(session.session_id) - if pipeline is None: - await emit_event( - correlate_realtime_error( - { - "type": "error", - "session_id": session.session_id, - "code": "internal_error", - "error": "server_vad is configured but its detector pipeline is unavailable", - }, - realtime_event_id, - ) - ) - continue - speech_was_active = pipeline.speech_active - reserved_bytes = 0 - previous_scratch_bytes = pipeline.scratch_bytes - try: - if native_input: - if fmt != "pcm_f32le" or sample_rate_hz != pipeline.sample_rate_hz: - raise ValueError("native server_vad input must be mono PCM float32 at 16000 Hz") - raw_audio = base64.b64decode(audio, validate=True) - if len(raw_audio) % np.dtype(" 0 - else: - if fmt != "pcm16": - raise ValueError("server_vad input must be mono PCM16") - source_sample_rate_hz = validate_input_sample_rate_hz(sample_rate_hz) - if source_sample_rate_hz not in {16_000, 24_000}: - raise ValueError("server_vad PCM16 input sample rate must be 16000 or 24000 Hz") - if ( - pipeline.source_sample_rate_hz is not None - and source_sample_rate_hz != pipeline.source_sample_rate_hz - ): - raise ValueError( - "server_vad input sample rate cannot change within a continuous audio stream" - ) - try: - raw_audio = base64.b64decode(audio, validate=True) - except (binascii.Error, ValueError) as exc: - raise ValueError( - "server_vad input audio must be valid base64-encoded PCM16" - ) from exc - if len(raw_audio) % np.dtype(" 0 - except (binascii.Error, ValueError) as exc: - await emit_event( - correlate_realtime_error( - { - "type": "error", - "session_id": session.session_id, - "code": "bad_audio", - "error": str(exc), - "param": "audio", - }, - realtime_event_id, - ) - ) - continue - if not native_input: - normalized_sample_upper_bound = ( - source_samples - if source_sample_rate_hz == pipeline.sample_rate_hz - else ((source_samples + 2) // 3) * 2 - ) - reserved_bytes = normalized_sample_upper_bound * np.dtype(np.float32).itemsize - if not session.reserve_input_bytes( - reserved_bytes, - limit=self._duplex_session_config.max_pending_input_bytes_per_session, - ): - await reject_and_clear_server_vad_input( - pipeline, - code="input_backpressure", - error="Server VAD input exceeds the per-session input limit", - metric_reason="input_backpressure", - event_id=realtime_event_id, - ) - continue - try: - vad_batch = ( - await pipeline.push(normalized_audio) - if native_input - else await pipeline.push_pcm16( - raw_audio, - source_sample_rate_hz=source_sample_rate_hz, - ) - ) - except Exception as exc: - logger.exception("Server VAD inference failed: %s", exc) - await reject_and_clear_server_vad_input( - pipeline, - code="server_vad_inference_failed", - error=str(exc), - metric_reason="inference", - event_id=realtime_event_id, - ) - continue - if not native_input: - retained_delta = ( - sum(frame.samples.nbytes for frame in vad_batch.frames) - + pipeline.scratch_bytes - - previous_scratch_bytes - ) - session.release_input_bytes(max(0, reserved_bytes - retained_delta)) - self._realtime_vad_metrics.observe_inference(vad_batch.inference_ms) - prefix_samples = round(server_vad_config.prefix_padding_ms * pipeline.sample_rate_hz / 1000) - speech_started = False - speech_stopped = False - for frame_index, vad_frame in enumerate(vad_batch.frames): - decision = vad_frame.decision - if not native_input: - released_bytes = session.append_server_vad_frame( - vad_frame.samples, - speech_started=decision.speech_started, - speech_stopped=decision.speech_stopped, - prefix_samples=prefix_samples, - ) - session.release_input_bytes(released_bytes) - if decision.speech_started: - speech_started = True - item_id = f"item_{uuid.uuid4().hex}" - session.begin_server_vad_speech(item_id) - await emit_event( - { - "type": "input_audio_buffer.speech_started", - "session_id": session.session_id, - "item_id": item_id, - "audio_start_ms": decision.audio_start_ms or 0, - } - ) - if decision.speech_stopped: - speech_stopped = True - item_id = session.finish_server_vad_speech() or f"item_{uuid.uuid4().hex}" - await emit_event( - { - "type": "input_audio_buffer.speech_stopped", - "session_id": session.session_id, - "item_id": item_id, - "audio_end_ms": decision.audio_end_ms or 0, - } - ) - if decision.endpoint_delay_ms is not None: - self._realtime_vad_metrics.observe_endpoint_delay(decision.endpoint_delay_ms) - if not native_input: - remaining_frame_bytes = sum( - frame.samples.nbytes for frame in vad_batch.frames[frame_index + 1 :] - ) - await commit_server_vad_turn( - item_id=item_id, - create_response=server_vad_config.create_response, - retain_input_bytes=remaining_frame_bytes + pipeline.scratch_bytes, - event_id=realtime_event_id, - ) - if accepted_audio: - session.lock_turn_detection_config() - if not native_input: - continue - speech_active = pipeline.speech_active - contained_speech = speech_was_active or speech_started or speech_stopped or speech_active - payload["is_speech"] = contained_speech - payload["vad"] = { - "backend": "silero", - "is_speech": contained_speech, - "speech_active": speech_active, - "speech_started": speech_started, - "speech_stopped": speech_stopped, - } - if speech_active: - payload["force_listen"] = True - # Speech/silence tag for the Stage0 turn-ended latch. - payload["is_speech"] = self._input_looks_like_speech(event, payload, session=session) - defer_native_append = False - buffer_overlap_audio = True - if native_input: - mark_pending_silence_superseded() - overlap_active = native_response_in_progress() and ( - not self._session_auto_responds(session) - or ( - session.capabilities.supports_barge_in - and ( - session.config.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value - or self._event_requests_barge_in(event) - ) - ) - ) - if overlap_active: - decision = self._overlap_decision(session, event, payload) - await self._emit_overlap_decision(emit_event, session, decision) - action = decision.get("action") - if action == "drop": - if realtime_protocol is not None: - await realtime_protocol.discard_pending_input_audio( - audio_end_ms=self._input_audio_duration_ms(event, payload) - ) - continue - if action == "listen": - buffer_overlap_audio = bool(decision.get("buffer_audio", True)) - defer_native_append = bool(decision.get("defer_runtime_append", True)) - if ( - not buffer_overlap_audio - and realtime_protocol is not None - and decision.get("preserve_realtime_input") is not True - ): - await realtime_protocol.discard_pending_input_audio( - audio_end_ms=self._input_audio_duration_ms(event, payload) - ) - if decision.get("force_listen", True) is True: - payload["force_listen"] = True - else: - event["force_barge_in"] = True - cancelled_fence = DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ) - playback_was_active = self._assistant_playback_active(session) - buffer_overlap_audio = True - defer_native_append = False - native.audio_buffer.clear_force_listen() - session.reset_overlap_speech() - native.input_since_commit = False - native.speech_since_commit = False - await actor.cancel_append_tasks() - had_native_stream = native.data_plane_task is not None - cancel_reason = str(decision.get("cancel_reason") or "barge_in") - cancelled = await self._cancel_active_response( - session, - actor.active_response_task, - emit_event, - reason=cancel_reason, - ) - had_native_stream = ( - await self._cancel_native_data_plane_stream(session) or had_native_stream - ) - if not cancelled and had_native_stream: - old_epoch = session.epoch - old_response_id = session.active_response_id - committed_ms = session.playback.committed_ms - self._commit_played_response_history(session, old_response_id, committed_ms) - new_epoch, old_playback = self._advance_barge_in_epoch(session) - await emit_event( - { - "type": "audio.cancelled", - "session_id": session.session_id, - "response_id": old_response_id, - "reason": cancel_reason, - "cancelled_epoch": old_epoch, - "epoch": new_epoch, - "committed_ms": committed_ms, - "playback": old_playback, - } - ) - cancelled = True - if not cancelled and playback_was_active: - old_epoch = session.epoch - committed_ms = session.playback.committed_ms - self._commit_played_response_history( - session, session.last_response_id, committed_ms - ) - new_epoch, old_playback = self._advance_barge_in_epoch(session) - await emit_event( - { - "type": "audio.cancelled", - "session_id": session.session_id, - "response_id": session.last_response_id, - "reason": cancel_reason, - "cancelled_epoch": old_epoch, - "epoch": new_epoch, - "committed_ms": committed_ms, - "playback": old_playback, - } - ) - cancelled = True - if session.epoch > cancelled_fence.epoch: - if not await self._signal_runtime_session( - session, - "barge_in", - emit_event, - fence=cancelled_fence, - next_fence=DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ), - ): - continue - actor.active_response_task = None - elif not self._session_auto_responds(session) and not self._input_looks_like_speech( - event, payload, session=session - ): - # Turn-mode only: skip silent chunks so they don't open a - # response. In auto-respond (full-duplex) mode the model owns - # the speak/listen decision and MUST receive silence units -- - # the official model typically starts speaking during the - # silence after a question. - # No response_id: a silence-skip listen can arrive - # while a response is still streaming and must not - # read as that response's terminal decision. - await emit_event( - { - "type": "response.listen", - "session_id": session.session_id, - "epoch": session.epoch, - "reason": "silence_or_noise", - } - ) - continue - if self._should_force_listen_for_auto_response_overlap(session, event, payload): - # Auto-response keeps a long-lived native Stage0 stream. - # While assistant audio is still active, silence from the - # browser should advance the model in listen mode rather - # than letting the same stream start another assistant - # segment. Speech/barge-in chunks remain model-owned. - payload["force_listen"] = True - if not buffer_overlap_audio: - continue - session.mark_user_input_activity() - native.input_since_commit = True - native.speech_since_commit = native.speech_since_commit or self._input_looks_like_speech( - event, payload, session=session - ) - try: - raw_audio_bytes = self._native_audio_payload_size_bytes(payload) - if not session.reserve_input_bytes( - raw_audio_bytes, - limit=self._duplex_session_config.max_pending_input_bytes_per_session, - ): - await emit_event( - { - "type": "error", - "error": "Duplex session pending input exceeds server limit", - "code": "input_backpressure", - } - ) - continue - allow_emit = not defer_native_append and ( - realtime_protocol is None - or event_type != "input_audio_buffer.append" - # Full-duplex: emit each ~chunk_period of audio so the model - # runs per-chunk generation (speak/listen) without an explicit - # response.create, matching the official duplex_generate loop. - or self._session_auto_responds(session) - ) - pcm_reservation = native.audio_buffer.prepare_append( - payload, - operation_id=uuid.uuid4().hex, - chunk_period_ms=session.capabilities.chunk_period_ms or 1000, - allow_emit=allow_emit, - ) - except ValueError as exc: - session.release_input_bytes(raw_audio_bytes) - await emit_event({"type": "error", "error": str(exc), "code": "bad_event"}) - continue - if pcm_reservation is None: - continue - if pcm_reservation.byte_count == 0: - session.release_input_bytes(raw_audio_bytes) - payload = pcm_reservation.payload - else: - session.append_audio(audio, fmt=fmt, sample_rate_hz=sample_rate_hz) - if native_input: - await start_native_append( - payload, - final=False, - pcm_reservation=pcm_reservation, - ) - continue - if session.capabilities.supports_input_append: - await start_runtime_append(payload, final=False, mode="append_audio_chunk") - continue - - if event_type in {"input.commit", "input_audio_buffer.commit", "response.create"}: - if ( - event_type == "response.create" - and session.config.server_vad is not None - and actor.active_response_task is not None - ): - await emit_event( - { - "type": "error", - "session_id": session.session_id, - "code": "response_already_active", - "error": "response.create must wait for the active response to complete", - } - ) - continue - item_id = event.get("item_id") - realtime_validated_audio_commit = ( - event_type == "input_audio_buffer.commit" and isinstance(item_id, str) and bool(item_id) - ) - if ( - self._uses_native_input_append(session) - and event_type in {"input.commit", "input_audio_buffer.commit"} - and not await wait_for_native_append_tail() - ): - continue - if event_type == "input_audio_buffer.commit" and event.get("is_speech") is False: - native.input_since_commit = False - native.speech_since_commit = False - native.audio_buffer.clear() - session.release_all_input_bytes() - native.clear_committed_audio() - await emit_event( - { - "type": "input.committed", - "session_id": session.session_id, - "item_id": item_id, - "turn_id": session.turn_id, - "epoch": session.epoch, - "empty": True, - "is_speech": False, - "no_response": True, - } - ) - # No response_id: this silence listen answers an empty - # commit, it is not a response terminal. - await emit_event( - { - "type": "response.listen", - "session_id": session.session_id, - "epoch": session.epoch, - "reason": "silence_or_noise", - } - ) - continue - should_create_response = ( - event_type == "response.create" - or bool(event.get("response_create", event_type == "input.commit")) - or (event_type == "input_audio_buffer.commit" and self._session_auto_responds(session)) - ) - precreate_response_requested = event_type == "response.create" or bool( - event.get("response_create", event_type == "input.commit") - ) - if event_type == "response.create": - response_payload = event.get("response") - if isinstance(response_payload, dict): - response_options_error = self._apply_response_create_options(session, response_payload) - if response_options_error is not None: - error_message = ( - "The selected native duplex runtime does not support generation " - "overrides for instructions, voice, temperature, max tokens, tools, or " - "tool_choice." - if response_options_error == "unsupported_native_response_options" - else "response.create cannot reserve options while another response is active." - ) - await emit_event( - { - "type": "error", - "session_id": session.session_id, - "epoch": session.epoch, - "code": response_options_error, - "error": error_message, - } - ) - continue - if self._uses_native_input_append(session) and event_type == "input_audio_buffer.commit": - has_pending_native_audio = ( - native.input_since_commit - or native.audio_buffer.has_pending() - or native.committed_audio_payload is not None - or realtime_validated_audio_commit - ) - if ( - not has_pending_native_audio - and not actor.native_append_tasks - and native.data_plane_task is None - ): - await emit_event( - { - "type": "error", - "session_id": session.session_id, - "epoch": session.epoch, - "code": "input_audio_buffer_empty", - "error": "input_audio_buffer.commit requires a non-empty input audio buffer.", - } - ) - continue - commit_action = decide_commit_action( - CommitSnapshot( - auto_responds=self._session_auto_responds(session), - speech_since_commit=native.speech_since_commit, - active_response_id=session.active_response_id, - overlap_speech_ms=session.overlap_speech_ms, - native_response_in_progress=native_response_in_progress(), - playback_active=self._assistant_playback_active(session), - ) - ) - if commit_action is CommitAction.DEFER_ACTIVE_RESPONSE: - if session.overlap_speech_ms <= session.config.overlap_short_ack_ms: - native.audio_buffer.clear() - session.release_all_input_bytes() - native.input_since_commit = False - native.speech_since_commit = False - native.clear_committed_audio() - if realtime_protocol is not None: - await realtime_protocol.discard_pending_input_audio( - audio_end_ms=session.overlap_speech_ms - ) - await emit_event( - { - "type": "input.committed", - "session_id": session.session_id, - "item_id": item_id, - "turn_id": session.turn_id, - "epoch": session.epoch, - "empty": True, - "is_speech": False, - "overlap_ack": True, - "no_response": True, - } - ) - session.reset_overlap_speech() - session.discard_response_options() - continue - - commit_reservation = native.audio_buffer.prepare_commit( - operation_id=uuid.uuid4().hex, - chunk_period_ms=session.capabilities.chunk_period_ms or 1000, - ) - deferred_payload = commit_reservation.payload - if deferred_payload is None: - commit_reservation.commit() - else: - if native.committed_audio_payload is not None: - deferred_payload = self._merge_native_audio_payloads( - native.committed_audio_payload, - deferred_payload, - ) - native.retain_committed_audio( - deferred_payload, - operation_id=commit_reservation.operation_id, - reserved_bytes=commit_reservation.byte_count, - ) - commit_reservation.commit() - native.deferred_response_create = should_create_response - native.deferred_precreate_response = precreate_response_requested - native.input_since_commit = False - native.speech_since_commit = False - committed = self._commit_native_audio_input( - session, - item_id=event.get("item_id"), - transcript=event.get("transcript"), - ) - committed_payload = self._native_audio_committed_payload( - session, - committed=committed, - item_id=event.get("item_id"), - transcript=event.get("transcript"), - ) - committed_payload["overlap_deferred"] = True - committed_payload["response_create_deferred"] = should_create_response - await emit_event(committed_payload) - continue - if commit_action is CommitAction.START_AUTO_RESPONSE: - commit_reservation = native.audio_buffer.prepare_commit( - operation_id=uuid.uuid4().hex, - chunk_period_ms=session.capabilities.chunk_period_ms or 1000, - ) - committed_input = commit_reservation.payload - final_payload = committed_input - if native.committed_audio_payload is not None: - if final_payload is not None: - final_payload = self._merge_native_audio_payloads( - native.committed_audio_payload, - final_payload, - ) - else: - final_payload = native.committed_audio_payload - commit_reservation.commit() - if final_payload is not None: - native.retain_committed_audio( - final_payload, - operation_id=commit_reservation.operation_id, - reserved_bytes=commit_reservation.byte_count, - ) - native.deferred_response_create = False - native.input_since_commit = False - native.speech_since_commit = False - data_plane_turn_id = session.turn_id - committed = self._commit_native_audio_input( - session, - item_id=event.get("item_id"), - transcript=event.get("transcript"), - turn_id=data_plane_turn_id, - ) - await emit_event( - self._native_audio_committed_payload( - session, - committed=committed, - item_id=event.get("item_id"), - transcript=event.get("transcript"), - ) - ) - if final_payload is not None: - await start_native_append( - { - **final_payload, - "duplex_turn_id": data_plane_turn_id, - }, - final=True, - precreate_response=False, - operation_id=commit_reservation.operation_id, - retained_committed_payload=final_payload, - ) - continue - if self._uses_native_input_append(session) and event_type == "response.create": - if ( - native_response_in_progress() - or actor.native_append_tasks - or native.data_plane_task is not None - ): - if session.active_response_id is None and ( - session.active_request_id is not None - or actor.native_append_tasks - or native.data_plane_task is not None - ): - continue - await emit_event( - { - "type": "error", - "session_id": session.session_id, - "epoch": session.epoch, - "code": "response_already_active", - "error": "response.create cannot start while another response is active.", - } - ) - session.discard_response_options() - continue - if native.committed_audio_payload is not None: - committed_payload = native.committed_audio_payload - operation_id = native.committed_audio_operation_id - if operation_id is None: - operation_id = uuid.uuid4().hex - native.committed_audio_operation_id = operation_id - await start_native_append( - committed_payload, - final=True, - precreate_response=True, - operation_id=operation_id, - retained_committed_payload=committed_payload, - ) - continue - await emit_event( - { - "type": "error", - "session_id": session.session_id, - "epoch": session.epoch, - "code": "response_create_without_input", - "error": "Native duplex response.create requires committed audio input.", - } - ) - session.discard_response_options() - continue - if self._uses_native_input_append(session) and not native_response_in_progress(): - commit_reservation = ( - native.audio_buffer.prepare_commit( - operation_id=uuid.uuid4().hex, - chunk_period_ms=session.capabilities.chunk_period_ms or 1000, - ) - if event_type in {"input.commit", "input_audio_buffer.commit"} - else None - ) - flushed_buffer_reserved_bytes = ( - native.audio_buffer.pending_byte_count if commit_reservation is None else 0 - ) - flushed = ( - commit_reservation.payload - if commit_reservation is not None - else native.audio_buffer.flush(chunk_period_ms=session.capabilities.chunk_period_ms or 1000) - ) - if native.committed_audio_payload is not None: - if flushed is not None: - flushed = self._merge_native_audio_payloads( - native.committed_audio_payload, - flushed, - ) - else: - flushed = native.committed_audio_payload - if commit_reservation is not None: - commit_reservation.commit() - if flushed is not None: - if self._should_force_listen_for_short_commit(session, event, flushed): - flushed = dict(flushed) - flushed["force_listen"] = True - native.input_since_commit = False - committed = self._commit_native_audio_input( - session, - item_id=event.get("item_id"), - transcript=event.get("transcript"), - ) - await emit_event( - self._native_audio_committed_payload( - session, - committed=committed, - item_id=event.get("item_id"), - transcript=event.get("transcript"), - ) - ) - if should_create_response: - native.retain_committed_audio( - flushed, - operation_id=( - commit_reservation.operation_id - if commit_reservation is not None - else uuid.uuid4().hex - ), - reserved_bytes=( - commit_reservation.byte_count - if commit_reservation is not None - else flushed_buffer_reserved_bytes - ), - ) - await start_native_append( - flushed, - final=True, - precreate_response=True, - operation_id=native.committed_audio_operation_id, - retained_committed_payload=flushed, - ) - else: - native.retain_committed_audio( - flushed, - operation_id=( - commit_reservation.operation_id - if commit_reservation is not None - else uuid.uuid4().hex - ), - reserved_bytes=( - commit_reservation.byte_count - if commit_reservation is not None - else flushed_buffer_reserved_bytes - ), - ) - native.deferred_response_create = False - continue - native_had_uncommitted_audio = self._uses_native_input_append(session) and ( - native.input_since_commit - or native.audio_buffer.has_pending() - or native.committed_audio_payload is not None - or realtime_validated_audio_commit - ) - committed = session.commit_user_input() - if self._uses_native_input_append(session) and event_type in { - "input_audio_buffer.commit", - "input.commit", - }: - native.input_since_commit = False - native.speech_since_commit = False - if committed is None and event_type != "response.create": - if self._uses_native_input_append(session): - committed = ( - self._commit_native_audio_input( - session, - item_id=event.get("item_id"), - transcript=event.get("transcript"), - ) - if native_had_uncommitted_audio - else None - ) - await emit_event( - self._native_audio_committed_payload( - session, - committed=committed, - item_id=event.get("item_id"), - transcript=event.get("transcript"), - ) - ) - else: - await emit_event( - { - "type": "input.committed", - "session_id": session.session_id, - "item_id": item_id, - "empty": True, - } - ) - continue - if committed is not None: - item_id = event.get("item_id") - if isinstance(item_id, str): - session.register_history_item(item_id, committed.message) - await emit_event( - self._input_committed_payload( - session, - committed, - item_id=item_id, - ) - ) - if not should_create_response: - continue - if actor.active_response_task is not None and not actor.active_response_task.done(): - await self._cancel_active_response( - session, - actor.active_response_task, - emit_event, - reason="new_response", - ) - if session.config.server_vad is not None and not self._uses_native_input_append(session): - start_chat_response() - else: - actor.active_response_task = asyncio.create_task(self._run_response(session, emit_event)) - continue - - await emit_event( - { - "type": "error", - "error": f"Unknown duplex event: {event_type}", - "code": "unknown_event", - } - ) - - except WebSocketDisconnect: - logger.info("Duplex session disconnected") - transport_detached = True - except Exception as exc: - logger.exception("Duplex session failed: %s", exc) - with suppress(Exception): - await emit_event({"type": "error", "error": str(exc), "code": "internal_error"}) - finally: - if realtime_protocol is not None: - realtime_protocol.reject_realtime_turn_detection_update() - if session is not None: - if reader_task is not None and not reader_task.done(): - reader_task.cancel() - with suppress(asyncio.CancelledError): - await reader_task - while pending_turn_reservations > 0: - session.release_pending_turn() - pending_turn_reservations -= 1 - resumable_detach = ( - transport_detached - and runtime_opened - and not runtime_closed - and session.state == DuplexSessionState.OPEN - and session.capabilities.supports_session_resume - and attachment_generation is not None - and resume_credential_delivered - ) - if resumable_detach: - - async def cancel_orphan_response_after_grace() -> None: - tasks = self._session_tasks.get(session.session_id) - current_session = self._registry.get(session.session_id) - if tasks is None or current_session is not session or session.state != DuplexSessionState.OPEN: - return - await tasks.cancel_append_tasks(response_bound_only=True) - await self._cancel_native_data_plane_stream(session) - await self._cancel_active_response( - session, - tasks.active_response_task, - emit_event, - reason="disconnect_grace_expired", - notify=False, - ) - tasks.active_response_task = None - - detached = await self._attachment_registry.detach( - session.session_id, - attachment_generation=attachment_generation, - on_grace_expired=cancel_orphan_response_after_grace, - ) - touch_session = getattr(self._chat_service.engine_client, "touch_duplex_session_async", None) - if detached and callable(touch_session): - with suppress(Exception): - await touch_session( - session.session_id, - fence=DuplexFence( - session.session_id, - epoch=session.epoch, - turn_id=session.turn_id, - incarnation=session.incarnation, - ), - activity=DuplexLeaseActivity.DETACH, - ) - else: - begin_close(actor.close_reason or "disconnect") - await actor.cancel_append_tasks() - await self._cancel_native_data_plane_stream(session) - await self._cancel_active_response( - session, - actor.active_response_task, - emit_event, - reason="disconnect", - notify=False, - ) - if runtime_opened and not runtime_closed and session.state != DuplexSessionState.CLOSED: - await self._close_runtime_session(session, reason="disconnect") - self._cleanup_duplex_session_state(session) - self._registry.close(session.session_id) - self._session_tasks.pop(session.session_id, None) - self._realtime_protocols.pop(session.session_id, None) - self._lease_generations.pop(session.session_id, None) - self._resync_required_sessions.discard(session.session_id) - with suppress(Exception): - await self._attachment_registry.close(session.session_id) - self._stop_lifecycle_listener_if_idle() - await actor.close_writer() - with suppress(Exception): - await asyncio.wait_for(actor.output_queue.join(), timeout=2.0) - if not writer_task.done(): - writer_task.cancel() - with suppress(asyncio.CancelledError): - await writer_task diff --git a/vllm_omni/entrypoints/duplex/warmup.py b/vllm_omni/entrypoints/duplex/warmup.py index 5051a1d31..2c138096d 100644 --- a/vllm_omni/entrypoints/duplex/warmup.py +++ b/vllm_omni/entrypoints/duplex/warmup.py @@ -43,10 +43,11 @@ async def _warmup_duplex_realtime(app, args, warmup_frames: int) -> None: model_name = served else: model_name = args.model - adapter = getattr(getattr(app.state, "openai_serving_duplex", None), "_serving_runtime_adapter", None) - frame_samples = int(getattr(adapter, "silence_continuation_samples", 16000)) + # One silence unit as the engine-side plugin defines it (DuplexOmniEngine.plugin). + plugin = getattr(getattr(app.state.engine_client, "engine", None), "plugin", None) + frame_samples = int(getattr(plugin, "silence_continuation_samples", 16000)) silence = base64.b64encode(bytes(frame_samples * 4)).decode("ascii") - from vllm_omni.experimental.fullduplex.client import build_realtime_url + from vllm_omni.clients.duplex import build_realtime_url url = ( build_realtime_url(f"ws://127.0.0.1:{args.port}/v1/realtime", model_name, autostart=False) diff --git a/vllm_omni/entrypoints/duplex/websocket.py b/vllm_omni/entrypoints/duplex/websocket.py index d51a7ed57..8ec9cc780 100644 --- a/vllm_omni/entrypoints/duplex/websocket.py +++ b/vllm_omni/entrypoints/duplex/websocket.py @@ -1,266 +1,60 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""WebSocket transport helpers for the duplex Realtime handler. + +The ordered session mailbox and the per-session task handles that used to be +defined here now live engine-side (``vllm_omni.engine.duplex.session.runner``). +""" + from __future__ import annotations import asyncio -from collections.abc import Callable -from dataclasses import dataclass, field -from typing import Any +from collections.abc import Awaitable, Callable +from contextlib import suppress -from fastapi import WebSocketDisconnect +from fastapi import WebSocket, WebSocketDisconnect -INPUT_EVENTS = frozenset( - { - "input.text.append", - "input_audio_buffer.append", - "input.commit", - "input_audio_buffer.commit", - "response.create", - } -) +#: Largest client frame accepted before it is parsed (base64 audio + video frames). +MAX_EVENT_BYTES = 15 * 1024 * 1024 -_INPUT_EVENT_ALIASES = { - "signal_turn": "turn.signal", - "close_session": "session.close", - "audio.playback_ack": "playback.ack", - "input_text.append": "input.text.append", - "push_text": "input.text.append", - "input.audio.append": "input_audio_buffer.append", - "push_chunk": "input_audio_buffer.append", -} -_WAV_AUDIO_ALIASES = frozenset({"input.audio.append", "push_chunk"}) +SendJson = Callable[[dict[str, object]], Awaitable[None]] +CloseFn = Callable[[str], Awaitable[None]] -def normalize_duplex_input_event(event: dict[str, object]) -> dict[str, object]: - """Normalize compatibility aliases before an event enters the mailbox.""" - event_type = event.get("type") - if not isinstance(event_type, str): - return event - canonical_type = _INPUT_EVENT_ALIASES.get(event_type) - if canonical_type is None: - return event - normalized = dict(event) - normalized["type"] = canonical_type - if event_type in _WAV_AUDIO_ALIASES: - normalized.setdefault("format", "wav") - return normalized +async def receive_text_with_timeout(websocket: WebSocket, timeout_s: float | None) -> str | None: + """Receive one text frame; ``None`` on idle timeout. Disconnects propagate.""" + if timeout_s is None or timeout_s <= 0: + return await websocket.receive_text() + try: + return await asyncio.wait_for(websocket.receive_text(), timeout=timeout_s) + except TimeoutError: + return None -MODEL_OUTPUT_EVENTS = frozenset( - { - "response.created", - "response.listen", - "response.speak", - "response.output_item.added", - "response.content_part.added", - "response.output_audio.delta", - "response.output_audio.done", - "response.output_text.delta", - "response.output_text.done", - "response.text.delta", - "response.text.done", - "response.message", - "response.output_item.done", - "response.content_part.done", - "response.done", - "runtime.control", - } -) -DOMAIN_TERMINAL_EVENTS = frozenset( - { - "response.done", - "response.listen", - "audio.cancelled", - "input.cancelled", - "session.closed", - } -) +def attachment_callbacks(websocket: WebSocket) -> tuple[SendJson, CloseFn]: + """``send`` / ``close`` callables for a ``DuplexSessionAttachmentRegistry`` attachment. + A send on a socket that already closed surfaces as ``WebSocketDisconnect`` + so the caller's disconnect path runs instead of a generic failure. + """ -def is_input_event(event_type: object) -> bool: - return isinstance(event_type, str) and event_type in INPUT_EVENTS - - -@dataclass(frozen=True, slots=True) -class DuplexAppendTaskMeta: - epoch: int - mode: str - final: bool - response_bound: bool - - -@dataclass -class DuplexSessionTasks: - """Connection-independent task handles for one resumable session.""" - - native_append_tasks: dict[asyncio.Task[bool], DuplexAppendTaskMeta] = field(default_factory=dict) - native_append_tail: asyncio.Task[bool] | None = None - active_response_task: asyncio.Task[None] | None = None - - def track_append_task( - self, - task: asyncio.Task[bool], - *, - epoch: int, - mode: str, - final: bool, - response_bound: bool, - ) -> None: - self.native_append_tasks[task] = DuplexAppendTaskMeta(epoch, mode, final, response_bound) - task.add_done_callback(self.native_append_tasks.pop) - - def has_response_bound_append_tasks(self) -> bool: - return any(meta.response_bound for meta in self.native_append_tasks.values()) - - async def cancel_append_tasks(self, timeout_s: float = 0.25, *, response_bound_only: bool = False) -> bool: - tasks = [ - task for task, meta in self.native_append_tasks.items() if not response_bound_only or meta.response_bound - ] - if not tasks: - return False - cancelled_tail = self.native_append_tail if self.native_append_tail in tasks else None - for task in tasks: - task.cancel() + async def attachment_send(payload: dict[str, object]) -> None: try: - await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=timeout_s) - except TimeoutError: - pass - if cancelled_tail is not None and self.native_append_tail is cancelled_tail: - self.native_append_tail = None - return True + await websocket.send_json(payload) + except RuntimeError as exc: + message = str(exc) + if "after sending 'websocket.close'" in message or "response already completed" in message: + raise WebSocketDisconnect(code=1006) from exc + raise + + async def attachment_close(reason: str) -> None: + close = getattr(websocket, "close", None) + if callable(close): + with suppress(Exception): + await close(code=1000, reason=reason) + + return attachment_send, attachment_close -@dataclass -class DuplexWebSocketActor: - """Own ordered WebSocket I/O queues, but no domain identity.""" - - websocket: Any - current_epoch: Callable[[], int | None] | None = None - session_closed: Callable[[], bool] | None = None - output_queue: asyncio.Queue[dict[str, object] | None] = field(default_factory=asyncio.Queue) - mailbox: asyncio.Queue[dict[str, object]] = field(default_factory=asyncio.Queue) - outbound_protocol: Any | None = None - tasks: DuplexSessionTasks = field(default_factory=DuplexSessionTasks) - closing: bool = False - close_reason: str | None = None - stale_output_dropped: int = 0 - _queued_input_events: int = 0 - - @property - def native_append_tasks(self) -> dict[asyncio.Task[bool], DuplexAppendTaskMeta]: - return self.tasks.native_append_tasks - - @property - def active_response_task(self) -> asyncio.Task[None] | None: - return self.tasks.active_response_task - - @active_response_task.setter - def active_response_task(self, task: asyncio.Task[None] | None) -> None: - self.tasks.active_response_task = task - - @property - def native_append_tail(self) -> asyncio.Task[bool] | None: - return self.tasks.native_append_tail - - @native_append_tail.setter - def native_append_tail(self, task: asyncio.Task[bool] | None) -> None: - self.tasks.native_append_tail = task - - async def enqueue_event(self, event: dict[str, object]) -> None: - if is_input_event(event.get("type")): - self._queued_input_events += 1 - await self.mailbox.put(event) - - async def next_event(self) -> dict[str, object]: - event = await self.mailbox.get() - if is_input_event(event.get("type")): - self._queued_input_events = max(0, self._queued_input_events - 1) - self.mailbox.task_done() - return event - - def has_queued_input_events(self) -> bool: - return self._queued_input_events > 0 - - async def send_json(self, payload: dict[str, object]) -> None: - await self.output_queue.put(payload) - - async def close_writer(self) -> None: - await self.output_queue.put(None) - - async def writer_loop(self) -> None: - while True: - payload = await self.output_queue.get() - try: - if payload is None: - return - raw_realtime = payload.pop("_realtime_raw", False) is True - if not raw_realtime and self._is_stale_model_output(payload): - self.stale_output_dropped += 1 - continue - try: - if raw_realtime: - await self.websocket.send_json(payload) - elif self.outbound_protocol is not None: - for projected in self.outbound_protocol.encode_outbound_event(payload): - await self.websocket.send_json(projected) - else: - await self.websocket.send_json(payload) - except (WebSocketDisconnect, RuntimeError): - return - finally: - self.output_queue.task_done() - - def _is_stale_model_output(self, payload: dict[str, object]) -> bool: - event_type = payload.get("type") - if event_type in DOMAIN_TERMINAL_EVENTS: - return False - if event_type not in MODEL_OUTPUT_EVENTS: - return False - if self.closing and event_type not in {"response.listen", "runtime.control"}: - return True - expected_epoch = self.current_epoch() if self.current_epoch is not None else None - if ( - self.session_closed is not None - and self.session_closed() - and event_type - not in { - "response.listen", - "runtime.control", - } - ): - return True - epoch = payload.get("epoch") - return isinstance(epoch, int) and isinstance(expected_epoch, int) and epoch != expected_epoch - - def track_append_task( - self, - task: asyncio.Task[bool], - *, - epoch: int, - mode: str, - final: bool, - response_bound: bool, - ) -> None: - self.tasks.track_append_task( - task, - epoch=epoch, - mode=mode, - final=final, - response_bound=response_bound, - ) - - def has_response_bound_append_tasks(self) -> bool: - return self.tasks.has_response_bound_append_tasks() - - async def cancel_append_tasks(self, timeout_s: float = 0.25, *, response_bound_only: bool = False) -> bool: - return await self.tasks.cancel_append_tasks(timeout_s, response_bound_only=response_bound_only) - - -__all__ = [ - "DOMAIN_TERMINAL_EVENTS", - "DuplexAppendTaskMeta", - "DuplexSessionTasks", - "DuplexWebSocketActor", - "is_input_event", -] +__all__ = ["MAX_EVENT_BYTES", "CloseFn", "SendJson", "attachment_callbacks", "receive_text_with_timeout"] diff --git a/vllm_omni/entrypoints/duplex_omni.py b/vllm_omni/entrypoints/duplex_omni.py new file mode 100644 index 000000000..968b69f74 --- /dev/null +++ b/vllm_omni/entrypoints/duplex_omni.py @@ -0,0 +1,476 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""DuplexOmni: the Python API for full-duplex models. + +Sessions run inside the engine (``DuplexSessionRunner`` on the orchestrator +loop of ``DuplexOrchestrator``). This class opens / resumes / closes them and +pipes typed ``DuplexCommand`` objects in and typed ``DuplexEvent`` objects out +through a ``DuplexSessionHandle``. It holds no session state beyond the +handle registry, so the websocket handler and ``InlineDuplexClient`` are both +thin consumers of the same surface. + +Example:: + + omni = DuplexOmni(model="openbmb/MiniCPM-o-4_5", trust_remote_code=True) + async with await omni.open_session({"modalities": ["audio", "text"]}) as session: + async def consume(): + async for event in session.events(): + if isinstance(event, AudioDelta): + play(event.audio) + task = asyncio.create_task(consume()) + await session.submit(AppendAudio(audio=pcm_bytes, format="pcm16", sample_rate_hz=16000)) + await session.submit(Commit()) +""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator, Mapping, Sequence +from contextlib import suppress +from typing import Any +from uuid import uuid4 + +from vllm.logger import init_logger + +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.events import DuplexEvent, SessionClosed +from vllm_omni.engine.duplex.messages import ( + DuplexControlResultMessage, + DuplexSessionError, + DuplexSessionEventMessage, +) +from vllm_omni.engine.duplex_omni_engine import DuplexOmniEngine +from vllm_omni.entrypoints.async_omni import AsyncOmni + +logger = init_logger(__name__) + +_DEFAULT_CONTROL_TIMEOUT_S = 10.0 + + +class DuplexSessionHandle: + """Client-side view of one engine-resident duplex session. + + Connection-independent: a websocket may detach and a later one resume the + same handle. ``events()`` is single-consumer at any given time but may be + re-entered after the previous iterator was closed (resume). + """ + + def __init__(self, omni: DuplexOmni, session_id: str) -> None: + self._omni = omni + #: Allocated by ``DuplexOmni.open_session``; unique for the engine's lifetime. + self.session_id = session_id + self.capabilities: DuplexCapabilities = DuplexCapabilities() + self.public_session: dict[str, object] = {} + self.lease_generation: int = 0 + self._outbox: asyncio.Queue[DuplexEvent | None] = asyncio.Queue() + self._closed = False + self._close_reason: str | None = None + self._closed_event = asyncio.Event() + self._consumer_active = False + + # ---- state ---- + + @property + def closed(self) -> bool: + return self._closed + + @property + def close_reason(self) -> str | None: + return self._close_reason + + def __repr__(self) -> str: # pragma: no cover - debugging aid + return f"DuplexSessionHandle(session_id={self.session_id!r}, closed={self._closed})" + + # ---- input ---- + + async def submit(self, command: DuplexCommand) -> None: + """Enqueue one command in caller order; rejections arrive as ``ErrorEvent`` on ``events()``.""" + if self._closed: + raise DuplexSessionError( + f"duplex session {self.session_id} is closed", code="session_closed", session_id=self.session_id + ) + await self._omni.engine.submit_command_async(self.session_id, command) + + async def append_audio( + self, + audio: bytes, + *, + format: str = "pcm16", + sample_rate_hz: int | None = None, + is_speech: bool | None = None, + video_frames: Sequence[str] | None = None, + duration_ms: int | None = None, + audio_end_ms: int | None = None, + hints: Mapping[str, object] | None = None, + ) -> None: + await self.submit( + duplex_commands.AppendAudio( + audio=audio, + format=format, + sample_rate_hz=sample_rate_hz, + is_speech=is_speech, + video_frames=tuple(video_frames or ()), + duration_ms=duration_ms, + audio_end_ms=audio_end_ms, + hints=dict(hints or {}), + ) + ) + + async def append_text(self, text: str) -> None: + await self.submit(duplex_commands.AppendText(text=text)) + + async def commit( + self, + *, + final: bool = True, + create_response: bool | None = None, + is_speech: bool | None = None, + ) -> None: + await self.submit(duplex_commands.Commit(final=final, create_response=create_response, is_speech=is_speech)) + + async def create_response(self, options: ResponseCreateOptions | Mapping[str, object] | None = None) -> None: + if isinstance(options, ResponseCreateOptions): + payload: dict[str, object] = { + key: value + for key, value in ( + ("instructions", options.instructions), + ("voice", options.voice), + ("output_audio_format", options.response_format), + ("temperature", options.temperature), + ("max_output_tokens", options.max_tokens), + ("speed", options.speed), + ("modalities", list(options.modalities) if options.modalities is not None else None), + ) + if value is not None + } + if options.extra_body: + payload["extra_body"] = dict(options.extra_body) + else: + payload = dict(options or {}) + await self.submit(duplex_commands.CreateResponse(options=payload)) + + async def clear_input(self) -> None: + await self.submit(duplex_commands.ClearInput()) + + async def cancel_input(self) -> None: + await self.submit(duplex_commands.CancelInput()) + + async def cancel_response(self, response_id: str | None = None) -> None: + await self.submit(duplex_commands.CancelResponse(response_id=response_id)) + + async def barge_in(self) -> None: + await self.submit(duplex_commands.BargeIn()) + + async def clear_output_audio(self, response_id: str | None = None) -> None: + await self.submit(duplex_commands.ClearOutputAudio(response_id=response_id)) + + async def signal_turn(self, event: str, payload: Mapping[str, object] | None = None) -> None: + await self.submit(duplex_commands.SignalTurn(event=event, signal_payload=dict(payload or {}))) + + async def update(self, session_patch: Mapping[str, object]) -> None: + await self.submit(duplex_commands.UpdateSession(patch=dict(session_patch))) + + async def ack_playback( + self, + played_ms: int, + *, + response_id: str | None = None, + item_id: str | None = None, + committed_ms: int | None = None, + ) -> None: + await self.submit( + duplex_commands.AckPlayback( + played_ms=int(played_ms), + committed_ms=int(committed_ms) if committed_ms is not None else None, + response_id=response_id, + item_id=item_id, + ) + ) + + async def heartbeat(self) -> None: + await self.submit(duplex_commands.Heartbeat()) + + async def create_item(self, item: Mapping[str, object], *, previous_item_id: str | None = None) -> None: + await self.submit(duplex_commands.CreateItem(item=dict(item), previous_item_id=previous_item_id)) + + async def delete_item(self, item_id: str) -> None: + await self.submit(duplex_commands.DeleteItem(item_id=item_id)) + + async def truncate_item(self, item_id: str, *, audio_end_ms: int, content_index: int = 0) -> None: + await self.submit( + duplex_commands.TruncateItem(item_id=item_id, audio_end_ms=int(audio_end_ms), content_index=content_index) + ) + + async def close(self, *, reason: str = "client_close", timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S) -> None: + if self._closed: + return + await self._omni.close_session(self.session_id, reason=reason, timeout=timeout) + + # ---- output ---- + + async def events(self) -> AsyncIterator[DuplexEvent]: + """Ordered public events; ends after ``session.closed`` / ``session.expired``.""" + if self._consumer_active: + raise RuntimeError(f"duplex session {self.session_id} already has an active events() consumer") + 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 + continue + yield event + if event.is_terminal: + return + finally: + self._consumer_active = False + + async def wait_closed(self) -> str: + await self._closed_event.wait() + return self._close_reason or "closed" + + async def __aenter__(self) -> DuplexSessionHandle: + return self + + async def __aexit__(self, exc_type, exc, traceback) -> None: + try: + await self.close() + except DuplexSessionError: + pass + + # ---- internals used by DuplexOmni ---- + + def _adopt(self, result: DuplexControlResultMessage) -> None: + if result.capabilities is not None: + self.capabilities = result.capabilities + if result.public_session: + self.public_session = dict(result.public_session) + if result.lease_generation is not None: + self.lease_generation = int(result.lease_generation) + + def _deliver(self, event: DuplexEvent) -> None: + self._outbox.put_nowait(event) + if isinstance(event, SessionClosed): # SessionExpired is a SessionClosed + self._mark_closed(event.reason or event.type) + + def _mark_closed(self, reason: str) -> None: + if self._closed: + return + self._closed = True + self._close_reason = reason + self._closed_event.set() + self._outbox.put_nowait(None) + + +class DuplexOmni(AsyncOmni): + """Async Python API for full-duplex models (see module docstring). + + Construct it like ``AsyncOmni``. The pipeline must declare + ``duplex_plugin`` and the deploy config ``session_mode: duplex``; the + engine raises at startup otherwise. There is no ``generate()``: duplex + serving is session-only. + """ + + engine: DuplexOmniEngine + + def _create_engine(self, **engine_kwargs: Any) -> DuplexOmniEngine: + from vllm_omni.entrypoints.duplex.audio_encoding import encode_audio + + return DuplexOmniEngine(duplex_audio_encoder=encode_audio, **engine_kwargs) + + def __init__(self, model: str = "", *args: Any, **kwargs: Any) -> None: + super().__init__(*args, model=model, **kwargs) + self._handles: dict[str, DuplexSessionHandle] = {} + + # ---- deployment facts ---- + + @property + def duplex_session_config(self) -> DuplexSessionRuntimeConfig: + return self.engine.duplex_session_config + + @property + def duplex_capabilities(self) -> DuplexCapabilities: + return self.engine.duplex_capabilities + + @property + def sessions(self) -> Mapping[str, DuplexSessionHandle]: + return dict(self._handles) + + def active_session_count(self) -> int: + return sum(1 for handle in self._handles.values() if not handle.closed) + + # ---- session lifecycle ---- + + @staticmethod + def _resolve_session_config( + config: DuplexSessionConfig | Mapping[str, object] | None, + *, + model: str, + ) -> DuplexSessionConfig: + """Build the normalized session config the engine will own (the one parse/normalize site). + + A ``session_id`` / ``id`` key inside a Realtime session object is + ignored: Realtime clients echo the session object back, and the server + allocates every session id. + """ + if config is None: + resolved = DuplexSessionConfig(model=model) + elif isinstance(config, DuplexSessionConfig): + resolved = config + elif isinstance(config, Mapping): + resolved = DuplexSessionConfig.from_realtime(config, model=model) + else: + raise TypeError(f"unsupported duplex session config: {type(config).__name__}") + if resolved.model is None: + resolved.model = model + return resolved.normalized() + + async def open_session( + self, + config: DuplexSessionConfig | Mapping[str, object] | None = None, + *, + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexSessionHandle: + """Open an engine-resident session and return its handle. + + The session id is always allocated here (``duplex-``; never + reused, so the id alone identifies a session) and the handle is + registered before the open RPC, so the first event (``session.created``, + which announces the id) can never arrive before it exists. + """ + session_id = f"duplex-{uuid4().hex}" + session_config = self._resolve_session_config(config, model=self.model) + handle = DuplexSessionHandle(self, session_id) + self._handles[session_id] = handle + self._final_output_handler() + try: + result = await self.engine.open_session_async(session_id, session_config, timeout=timeout) + except BaseException: + 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 + # be admitting this session -- a control timeout unregisters the RPC + # waiter but cannot cancel the manager's open -- so a late admission + # would hold an admission slot and a Stage0 reservation until idle + # expiry. Send the compensating close; the manager applies it + # immediately if the session already landed, and records it for the + # open to honour on arrival if not. + with suppress(Exception): + await self.engine.close_session_async(session_id, reason="open_abandoned", timeout=timeout) + raise + handle._adopt(result) + return handle + + def get_session(self, session_id: str) -> DuplexSessionHandle | None: + return self._handles.get(session_id) + + async def resume_session( + self, + session_id: str, + *, + expected_lease_generation: int, + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> DuplexSessionHandle: + """Engine lease resume (CAS on the lease generation); returns the existing handle.""" + handle = self._require_handle(session_id) + result = await self.engine.resume_session_async( + session_id, + expected_lease_generation=expected_lease_generation, + timeout=timeout, + ) + # The result carries the engine's current public session (state, + # epoch, turn ...), which a resumed client must see, not the open-time snapshot. + handle._adopt(result) + return handle + + async def touch_session( + self, + session_id: str, + *, + activity: str = "heartbeat", + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> None: + self._require_handle(session_id) + await self.engine.touch_session_async(session_id, activity=activity, timeout=timeout) + + async def detach_session(self, session_id: str, *, timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S) -> None: + """Start the engine-owned disconnect grace; expiry arrives as ``session.expired``.""" + await self.touch_session(session_id, activity="detach", timeout=timeout) + + async def close_session( + self, + session_id: str, + *, + reason: str = "client_close", + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> None: + handle = self._handles.get(session_id) + if handle is None or handle.closed: + return + try: + await self.engine.close_session_async(session_id, reason=reason, timeout=timeout) + except Exception: + # The engine refused or lost the close: the handle is unusable + # either way, and it must not shadow the id in ``sessions``. + handle._mark_closed(reason) + self._handles.pop(session_id, None) + raise + # The manager emits session.closed after the stage cleanup, but the + # event travels on the output queue: wait for it so ``events()`` ends + # with the typed SessionClosed rather than an abrupt stop. + try: + await asyncio.wait_for(handle.wait_closed(), timeout=timeout) + except asyncio.TimeoutError: + logger.warning("[DuplexOmni] session.closed for %s did not arrive in time", session_id) + handle._mark_closed(reason) + self._handles.pop(session_id, None) + + async def close_all_sessions(self, *, reason: str = "shutdown") -> None: + for session_id in list(self._handles): + try: + await self.close_session(session_id, reason=reason) + except Exception: + logger.exception("[DuplexOmni] failed to close session %s", session_id) + + def _require_handle(self, session_id: str) -> DuplexSessionHandle: + handle = self._handles.get(session_id) + if handle is None or handle.closed: + raise DuplexSessionError( + f"unknown or closed duplex session: {session_id}", code="unknown_session", session_id=session_id + ) + return handle + + # ---- engine output routing ---- + + def _route_engine_message(self, msg: object) -> bool: + if not isinstance(msg, DuplexSessionEventMessage): + return False + handle = self._handles.get(msg.session_id) + if handle is None: + logger.debug("[DuplexOmni] dropping event for unknown session %s", msg.session_id) + return True + handle._deliver(msg.event) + if handle.closed: + self._handles.pop(msg.session_id, None) + return True + + def _on_engine_dead(self, error: str) -> None: + for handle in list(self._handles.values()): + handle._mark_closed(f"engine_dead: {error}") + self._handles.clear() + + def shutdown(self, timeout: float | None = None) -> None: + for handle in list(self._handles.values()): + handle._mark_closed("shutdown") + self._handles.clear() + super().shutdown(timeout) + + +__all__ = ["DuplexOmni", "DuplexSessionHandle"] diff --git a/vllm_omni/entrypoints/duplex_request_client.py b/vllm_omni/entrypoints/duplex_request_client.py deleted file mode 100644 index d9e70f672..000000000 --- a/vllm_omni/entrypoints/duplex_request_client.py +++ /dev/null @@ -1,383 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import asyncio -import time -from collections.abc import Callable, MutableMapping -from typing import Protocol - -from vllm_omni.engine.duplex.contracts import ( - duplex_data_plane_request_info, - duplex_resource_request_id, -) -from vllm_omni.engine.duplex.lease import DuplexLeaseActivity -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.messages import ErrorMessage, OutputMessage -from vllm_omni.entrypoints.client_request_state import ClientRequestState -from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics -from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.outputs.duplex import get_duplex_output_decision - - -class DuplexEnginePort(Protocol): - """Engine control methods required by the experimental request client. - - The protocol keeps request preregistration and output collection testable - without coupling this component to ``AsyncOmniEngine``. - """ - - async def open_duplex_session_async(self, session_id: str, **kwargs) -> dict[str, object]: ... - - async def append_duplex_input_async(self, session_id: str, **kwargs) -> dict[str, object]: ... - - async def signal_duplex_turn_async(self, session_id: str, **kwargs) -> dict[str, object]: ... - - async def close_duplex_session_async(self, session_id: str, **kwargs) -> dict[str, object]: ... - - async def touch_duplex_session_async(self, session_id: str, **kwargs) -> dict[str, object]: ... - - async def resume_duplex_session_async(self, session_id: str, **kwargs) -> dict[str, object]: ... - - -class DuplexRequestOutputPort: - """Narrow output-processing surface owned by the AsyncOmni entrypoint.""" - - def __init__( - self, - *, - request_states: MutableMapping[str, ClientRequestState], - num_stages: int, - log_stats: bool, - start_output_handler: Callable[[], object], - process_single_result: Callable[..., OmniRequestOutput | None], - ) -> None: - self.request_states = request_states - self.num_stages = num_stages - self.log_stats = log_stats - self._start_output_handler = start_output_handler - self._process_single_result = process_single_result - - def start_output_handler(self) -> None: - self._start_output_handler() - - def process_single_result(self, *args, **kwargs) -> OmniRequestOutput | None: - return self._process_single_result(*args, **kwargs) - - -class DuplexRequestClient: - """Owns duplex request identity, preregistration and output collection.""" - - def __init__(self, engine: DuplexEnginePort, output_port: DuplexRequestOutputPort) -> None: - self.engine = engine - self.output_port = output_port - - @staticmethod - def request_info(result: dict[str, object]) -> tuple[str | None, int | None]: - return duplex_data_plane_request_info(result) - - async def open( - self, - session_id: str, - *, - session_mode: str, - capabilities: dict[str, object] | None, - session_config: dict[str, object] | None, - runtime_config: dict[str, object] | None, - fence: DuplexFence, - timeout: float | None, - ) -> dict[str, object]: - # Lifecycle expiry messages share the engine output channel. Start its - # sole dispatcher before open so an idle session can be reaped even if - # it never submits an append request. - self.output_port.start_output_handler() - kwargs: dict[str, object] = { - "session_mode": session_mode, - "capabilities": capabilities, - "session_config": session_config, - "fence": fence, - "timeout": timeout, - } - if runtime_config is not None: - kwargs["runtime_config"] = runtime_config - return await self.engine.open_duplex_session_async(session_id, **kwargs) - - async def append( - self, - session_id: str, - *, - mode: str, - payload: object, - operation_id: str | None, - final: bool, - expected_epoch: int | None, - fence: DuplexFence, - timeout: float | None, - collect_outputs: bool, - ) -> dict[str, object]: - kwargs: dict[str, object] = { - "mode": mode, - "payload": payload, - "final": final, - "expected_epoch": expected_epoch, - "timeout": timeout, - "fence": fence, - } - if operation_id is not None: - kwargs["operation_id"] = operation_id - - expected_request_id = duplex_resource_request_id(fence, "stage0") - request_state = self.output_port.request_states.get(expected_request_id) - created_request_state = request_state is None - if request_state is None: - request_state = ClientRequestState(expected_request_id) - self.output_port.request_states[expected_request_id] = request_state - if request_state.metrics is None: - wall_start_ts = time.time() - request_state.metrics = OrchestratorMetrics( - self.output_port.num_stages, - self.output_port.log_stats, - wall_start_ts, - max(0, self.output_port.num_stages - 1), - ) - request_state.request_arrival_ts = wall_start_ts - self.output_port.start_output_handler() - try: - result = await self.engine.append_duplex_input_async(session_id, **kwargs) - except BaseException: - if created_request_state and request_state.queue.empty(): - self.output_port.request_states.pop(expected_request_id, None) - raise - - request_id, response_stage_id = duplex_data_plane_request_info(result) - if request_id is None: - if created_request_state and request_state.queue.empty(): - self.output_port.request_states.pop(expected_request_id, None) - return result - if request_id != expected_request_id: - if created_request_state and request_state.queue.empty(): - self.output_port.request_states.pop(expected_request_id, None) - raise RuntimeError( - f"duplex data-plane request id mismatch: expected {expected_request_id!r}, got {request_id!r}" - ) - if not collect_outputs: - return result - outputs = await self.collect_outputs( - request_id, - request_state, - response_stage_id=response_stage_id, - timeout=timeout, - ) - if outputs: - result = dict(result) - result["data_plane_outputs"] = outputs - return result - - async def collect_registered_outputs( - self, - request_id: str, - *, - response_stage_id: int | None, - timeout: float | None, - ) -> list[OmniRequestOutput]: - self.output_port.start_output_handler() - request_state = self.output_port.request_states.get(request_id) - if request_state is None: - return [] - return await self.collect_outputs( - request_id, - request_state, - response_stage_id=response_stage_id, - timeout=timeout, - ) - - async def signal( - self, - session_id: str, - *, - event: str, - fence: DuplexFence, - next_fence: DuplexFence | None, - session_config: dict[str, object] | None, - runtime_config: dict[str, object] | None, - timeout: float | None, - ) -> dict[str, object]: - kwargs: dict[str, object] = {"event": event, "fence": fence, "timeout": timeout} - if next_fence is not None: - kwargs["next_fence"] = next_fence - if session_config is not None: - kwargs["session_config"] = session_config - if runtime_config is not None: - kwargs["runtime_config"] = runtime_config - result = await self.engine.signal_duplex_turn_async(session_id, **kwargs) - if event in {"barge_in", "input.cancel", "response.cancel"}: - self.output_port.request_states.pop(duplex_resource_request_id(fence, "stage0"), None) - return result - - async def close( - self, - session_id: str, - *, - reason: str, - fence: DuplexFence, - timeout: float | None, - ) -> dict[str, object]: - result = await self.engine.close_duplex_session_async( - session_id, - reason=reason, - fence=fence, - timeout=timeout, - ) - self.output_port.request_states.pop(duplex_resource_request_id(fence, "stage0"), None) - return result - - async def touch( - self, - session_id: str, - *, - fence: DuplexFence, - activity: DuplexLeaseActivity, - timeout: float | None, - ) -> dict[str, object]: - return await self.engine.touch_duplex_session_async( - session_id, - fence=fence, - activity=activity, - timeout=timeout, - ) - - async def resume( - self, - session_id: str, - *, - fence: DuplexFence, - expected_lease_generation: int, - timeout: float | None, - ) -> dict[str, object]: - return await self.engine.resume_duplex_session_async( - session_id, - fence=fence, - expected_lease_generation=expected_lease_generation, - timeout=timeout, - ) - - async def collect_outputs( - self, - request_id: str, - request_state: ClientRequestState, - *, - response_stage_id: int | None, - timeout: float | None, - ) -> list[OmniRequestOutput]: - deadline = None if timeout is None else time.monotonic() + timeout - wall_start_ts = request_state.request_arrival_ts or time.time() - final_stage_id = response_stage_id if response_stage_id is not None else max(0, self.output_port.num_stages - 1) - num_stages = max(self.output_port.num_stages, final_stage_id + 1) - metrics = getattr(request_state, "metrics", None) or OrchestratorMetrics( - num_stages, - self.output_port.log_stats, - wall_start_ts, - final_stage_id, - ) - stage_event_cursor = metrics.stage_event_cursor(request_id) - request_start_ts = {request_id: wall_start_ts} - outputs: list[OmniRequestOutput] = [] - while True: - remaining = None if deadline is None else max(0.0, deadline - time.monotonic()) - if remaining == 0.0: - break - try: - message = await asyncio.wait_for(request_state.queue.get(), timeout=remaining) - except asyncio.TimeoutError: - break - if isinstance(message, ErrorMessage): - raise RuntimeError(message.error) - if not isinstance(message, OutputMessage): - continue - engine_outputs = message.engine_outputs - is_response_stage = response_stage_id is None or message.stage_id >= response_stage_id - is_direct_response = self.is_direct_response(engine_outputs) - output_to_collect = None - if isinstance(engine_outputs, OmniRequestOutput): - output_to_collect = engine_outputs - if message.metrics is not None: - existing_metrics = dict(output_to_collect.metrics) - metrics.process_stage_metrics( - result={"metrics": message.metrics}, - stage_type="", - stage_id=message.stage_id, - req_id=request_id, - engine_outputs=engine_outputs, - finished=message.finished, - final_output_type=engine_outputs.final_output_type, - output_to_yield=output_to_collect, - event_cursor=stage_event_cursor, - ) - generated_metrics = dict(output_to_collect.metrics) - existing_stage_metrics = existing_metrics.get("stage_metrics") - generated_stage_metrics = generated_metrics.get("stage_metrics") - if isinstance(existing_stage_metrics, dict) and isinstance(generated_stage_metrics, dict): - generated_metrics["stage_metrics"] = { - **existing_stage_metrics, - **generated_stage_metrics, - } - output_to_collect.metrics = { - **existing_metrics, - **generated_metrics, - } - elif is_response_stage: - output_to_collect = self.output_port.process_single_result( - message, - message.stage_id, - metrics, - request_start_ts, - wall_start_ts, - final_stage_id, - stage_event_cursor, - ) - if output_to_collect is not None and (is_response_stage or is_direct_response): - if message.finished: - output_to_collect.finished = True - outputs.append(output_to_collect) - if message.finished or outputs: - break - return outputs - - @classmethod - def is_direct_response(cls, output: object) -> bool: - multimodal_output = cls.multimodal_output(output) - if not multimodal_output: - return False - return multimodal_output.get("duplex_direct_response") is True or multimodal_output.get( - "duplex_native_decision" - ) in {"listen", "speak"} - - @classmethod - def multimodal_output(cls, output: object) -> dict[str, object]: - if isinstance(output, OmniRequestOutput): - decision = get_duplex_output_decision(output) - if decision is not None: - return dict(decision.metadata) - private_output = getattr(output, "_multimodal_output", None) - multimodal_output = output.multimodal_output - if isinstance(multimodal_output, dict) and multimodal_output: - return multimodal_output - inner_output = getattr(output, "request_output", None) - if inner_output is not None and inner_output is not output: - inner_multimodal_output = cls.multimodal_output(inner_output) - if inner_multimodal_output: - return inner_multimodal_output - if isinstance(private_output, dict): - return private_output - return multimodal_output if isinstance(multimodal_output, dict) else {} - multimodal_output = getattr(output, "multimodal_output", None) - if isinstance(multimodal_output, dict): - return multimodal_output - outputs = getattr(output, "outputs", None) - completion = outputs[0] if isinstance(outputs, list) and outputs else None - multimodal_output = getattr(completion, "multimodal_output", None) if completion is not None else None - return multimodal_output if isinstance(multimodal_output, dict) else {} - - -__all__ = ["DuplexEnginePort", "DuplexRequestClient", "DuplexRequestOutputPort"] diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index cf4da3b7a..ca9fc0196 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -7,12 +7,13 @@ import copy import time import uuid from collections.abc import Callable, Generator, Iterable, Sequence -from typing import TYPE_CHECKING, Literal, overload +from typing import TYPE_CHECKING, Any, Literal, overload from tqdm.auto import tqdm from vllm.logger import init_logger from vllm.sampling_params import RequestOutputKind +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.messages import OutputMessage from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.omni_base import OmniBase @@ -28,6 +29,9 @@ logger = init_logger(__name__) class Omni(OmniBase): """Synchronous entrypoint for offline generation.""" + def _create_engine(self, **engine_kwargs: Any) -> AsyncOmniEngine: + return AsyncOmniEngine(**engine_kwargs) + def _maybe_force_final_only_for_llm_stages( self, sampling_params_list: Sequence[OmniSamplingParams], diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index ead0e6f1d..79217ca8b 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -19,13 +19,13 @@ from vllm.transformers_utils.runai_utils import is_runai_obj_uri from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.config.stage_config import merge_sampling_constraints -from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.messages import ( EngineQueueMessage, ErrorMessage, OutputMessage, StageMetricsMessage, ) +from vllm_omni.engine.omni_engine_base import OmniEngineBase from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin from vllm_omni.entrypoints.utils import coerce_param_message_types, get_final_stage_id_for_e2e @@ -66,7 +66,7 @@ class OmniEngineDeadError(EngineDeadError): self.error_stage_id = error_stage_id -def _weak_shutdown_engine(engine: AsyncOmniEngine) -> None: +def _weak_shutdown_engine(engine: OmniEngineBase) -> None: """Best-effort engine cleanup for GC finalization.""" try: engine.shutdown() @@ -206,13 +206,13 @@ class OmniBase(PDDisaggregationMixin): self.tts_batch_max_items: int = kwargs.pop("tts_batch_max_items", 32) logger.info("[%s] Initializing with model %s", self.__class__.__name__, model) - # Construct transfer_metrics first so we can hand it to AsyncOmniEngine + # Construct transfer_metrics first so we can hand it to the engine # (which forwards it to the Orchestrator background thread for - # TX-side emit; see Orchestrator._forward_to_next_stage). + # TX-side emit; see OrchestratorBase._forward_to_next_stage). self.transfer_metrics = OmniTransferMetrics(model_name=model, log_stats=log_stats) self.prom_metrics = OmniPrometheusMetrics(model_name=model, log_stats=log_stats) st = time.time() - self.engine = AsyncOmniEngine( + self.engine = self._create_engine( model=model, init_timeout=init_timeout, stage_init_timeout=stage_init_timeout, @@ -224,8 +224,8 @@ class OmniBase(PDDisaggregationMixin): self._shutdown_called = False self._weak_finalizer = weakref.finalize(self, _weak_shutdown_engine, self.engine) et = time.time() - logger.info("[%s] AsyncOmniEngine initialized in %.2f seconds", self.__class__.__name__, et - st) - # Authoritative: ``AsyncOmniEngine`` resolves (pipeline + deploy YAML + + logger.info("[%s] %s initialized in %.2f seconds", self.__class__.__name__, type(self.engine).__name__, et - st) + # Authoritative: the engine resolves (pipeline + deploy YAML + # CLI overrides) through ``StageConfigFactory`` and stores the final # value on ``engine.async_chunk``; mirror it here so ``--no-async-chunk`` # (explicit ``False``) is not fallen-back-through by ``or``. @@ -253,6 +253,10 @@ class OmniBase(PDDisaggregationMixin): # PD disaggregation state (detects if a prefill/decode stage pair is configured) self._init_pd_state() + def _create_engine(self, **engine_kwargs: Any) -> OmniEngineBase: + """Construct this entrypoint's engine (``AsyncOmniEngine`` for turn-based use).""" + raise NotImplementedError + @property def num_stages(self) -> int: return self.engine.num_stages diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 36b7544ee..376a9ccd2 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -87,9 +87,9 @@ from vllm_omni.config.endpoint_policy import ( shutdown_unsupported_routes, ) from vllm_omni.entrypoints.async_omni import ABORT_TIMEOUT_S, AsyncOmni -from vllm_omni.entrypoints.duplex.capability import should_enable_duplex_endpoint from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler from vllm_omni.entrypoints.duplex.warmup import _warmup_duplex_realtime +from vllm_omni.entrypoints.duplex_omni import DuplexOmni from vllm_omni.entrypoints.openai import app_state as openai_app_state from vllm_omni.entrypoints.openai.app_state import ( ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL, @@ -440,6 +440,23 @@ async def build_async_omni( yield async_omni +def _is_duplex_model(model: str, kwargs: dict[str, Any]) -> bool: + """Whether ``model``'s pipeline declares a ``duplex_plugin`` (served by ``DuplexOmni``). + + Resolution errors propagate: a duplex model whose pipeline or deploy + config cannot be resolved must fail startup rather than silently start a + turn-based server. + """ + from vllm_omni.config.config_factory import StageConfigFactory + + pipeline_config = StageConfigFactory.get_pipeline_config( + model=model, + trust_remote_code=bool(kwargs.get("trust_remote_code")), + deploy_config_path=kwargs.get("deploy_config"), + ) + return bool(pipeline_config is not None and getattr(pipeline_config, "duplex_plugin", None)) + + @asynccontextmanager async def build_async_omni_from_stage_config( args: TrackingNamespace, @@ -501,7 +518,12 @@ async def build_async_omni_from_stage_config( kwargs.pop("robot_openpi_idle_timeout", None) model = kwargs.pop("model", None) or args.model kwargs.setdefault("log_stats", not args.disable_log_stats) - async_omni = AsyncOmni(model=model, **kwargs) + if _is_duplex_model(model, kwargs): + # A duplex model is always served in duplex mode: sessions over + # /v1/realtime?duplex=1, no turn-based HTTP routes. + async_omni = DuplexOmni(model=model, **kwargs) + else: + async_omni = AsyncOmni(model=model, **kwargs) # # Don't keep the dummy data in memory # await async_llm.reset_mm_cache() @@ -512,6 +534,134 @@ async def build_async_omni_from_stage_config( async_omni.shutdown() +async def _init_duplex_app_state( + engine_client: DuplexOmni, + state: State, + args: Namespace, + base_model_paths: list[BaseModelPath], + vllm_config: Any, + request_logger: RequestLogger | None, +) -> None: + """Minimal app state for a duplex server: only the surfaces a duplex session backs.""" + state.vllm_config = vllm_config + state.diffusion_engine = None + state.openai_serving_models = OpenAIServingModels( + engine_client=engine_client, # type: ignore[arg-type] + base_model_paths=base_model_paths, + lora_modules=None, + ) + state.serving_tokenization = None + state.serving_tokens = None + # Replaced by the chat init below when the model serves chat. + state.online_renderer = None + for attribute in ( + "openai_serving_chat_batch", + "openai_serving_completion", + "openai_serving_responses", + "openai_serving_embedding", + "openai_serving_pooling", + "openai_serving_classification", + "openai_serving_scores", + "openai_serving_transcription", + "openai_serving_translation", + "openai_serving_speech", + "openai_serving_audio_generate", + "openai_serving_video", + "openai_streaming_speech", + "openai_streaming_video", + "openai_streaming_video_output", + "openai_serving_realtime", + "openai_serving_realtime_robot", + "anthropic_serving_messages", + ): + setattr(state, attribute, None) + state.openai_serving_duplex = OmniDuplexSessionHandler(duplex_omni=engine_client) + # One engine, both surfaces. ``DuplexOmni`` extends ``AsyncOmni``, so the + # ordinary chat service runs on it unchanged: a chat request is a turn-based + # request on the same pipeline, not a session, and costs no admission slot. + state.openai_serving_chat = await _init_duplex_chat(engine_client, state, args, request_logger) + state.enable_server_load_tracking = getattr(args, "enable_server_load_tracking", False) + state.server_load_metrics = 0 + if state.openai_serving_chat is not None: + logger.info( + "Duplex mode: serving %s over /v1/realtime?duplex=1 and /v1/chat/completions", + engine_client.model, + ) + else: + logger.info( + "Duplex mode: serving %s over /v1/realtime?duplex=1 only " + "(/v1/chat/completions unavailable: the model does not declare supports_chat_completions)", + engine_client.model, + ) + + +async def _init_duplex_chat( + engine_client: DuplexOmni, + state: State, + args: Namespace, + request_logger: RequestLogger | None, +) -> OmniOpenAIServingChat | None: + """The ordinary chat service, on a duplex engine, when the model allows it. + + Gated on ``DuplexCapabilities.supports_chat_completions`` so the decision + stays the model's: a duplex model that should not answer chat requests says + so in its plugin, and the route reports "not available" rather than + answering badly. ``endpoint_restrictions`` remains the per-deployment + opt-out on top of this. + """ + if not engine_client.duplex_capabilities.supports_chat_completions: + return None + supported_tasks: set[str] = {"generate"} + if hasattr(engine_client, "get_supported_tasks"): + supported_tasks = set(await engine_client.get_supported_tasks()) + if "generate" not in supported_tasks: + return None + + resolved_chat_template = load_chat_template(args.chat_template) + if resolved_chat_template is None: + try: + tokenizer = await engine_client.get_tokenizer() + except Exception as exc: + logger.debug("Could not inspect tokenizer chat_template before duplex chat init: %s", exc) + tokenizer = None + if tokenizer is None or getattr(tokenizer, "chat_template", None) is None: + resolved_chat_template = _load_model_chat_template_json(args.model) + + state.online_renderer = OnlineRenderer( + model_config=engine_client.model_config, + renderer=engine_client.renderer, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + enable_auto_tools=args.enable_auto_tool_choice, + exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, + tool_parser=args.tool_call_parser, + reasoning_parser=args.structured_outputs_config.reasoning_parser, + default_chat_template_kwargs=args.default_chat_template_kwargs, + ) + return OmniOpenAIServingChat( + engine_client=engine_client, + models=state.openai_serving_models, + response_role=args.response_role, + online_renderer=state.online_renderer, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + default_chat_template_kwargs=args.default_chat_template_kwargs, + trust_request_chat_template=args.trust_request_chat_template, + return_tokens_as_token_ids=args.return_tokens_as_token_ids, + enable_auto_tools=args.enable_auto_tool_choice, + exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, + tool_parser=args.tool_call_parser, + reasoning_parser=args.structured_outputs_config.reasoning_parser, + enable_prompt_tokens_details=args.enable_prompt_tokens_details, + enable_force_include_usage=args.enable_force_include_usage, + enable_log_outputs=args.enable_log_outputs, + enable_log_deltas=args.enable_log_deltas, + ) + + async def omni_init_app_state( engine_client: EngineClient, state: State, @@ -562,6 +712,13 @@ async def omni_init_app_state( state.stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None model_name = served_model_names[0] if served_model_names else args.model + # Duplex mode: a duplex model is served through DuplexOmni only. Sessions + # run over /v1/realtime?duplex=1; every turn-based HTTP route reports + # "not available". + if isinstance(engine_client, DuplexOmni): + await _init_duplex_app_state(engine_client, state, args, base_model_paths, vllm_config, request_logger) + return + # Pure Diffusion mode: use simplified initialization logic if is_pure_diffusion: state.vllm_config = None @@ -942,17 +1099,6 @@ async def omni_init_app_state( else None ) state.openai_serving_duplex = None - if state.openai_serving_chat is not None and should_enable_duplex_endpoint( - state.stage_configs, - config_path=getattr(engine_client, "config_path", None) or getattr(args, "deploy_config", None), - ): - state.openai_serving_duplex = OmniDuplexSessionHandler( - chat_service=state.openai_serving_chat, - served_model_name=model_name, - log_stats=state.log_stats, - duplex_session_config=getattr(engine_client, "duplex_session_config", None), - serving_runtime_adapter_path=getattr(engine_client, "duplex_serving_adapter_path", None), - ) state.openai_serving_realtime = OpenAIServingRealtime( engine_client=engine_client, models=state.openai_serving_models, @@ -1477,14 +1623,7 @@ async def streaming_video_output(websocket: WebSocket): @router.websocket("/v1/realtime") async def realtime_websocket(websocket: WebSocket): """WebSocket endpoint for OpenAI-style realtime interactions.""" - # Hold real clients until the startup duplex warmup finishes (the warmup - # connection marks itself with vllm_omni_warmup=1 and passes through). - warmup_done = getattr(websocket.app.state, "duplex_warmup_done", None) - if warmup_done is not None and not warmup_done.is_set() and websocket.query_params.get("vllm_omni_warmup") != "1": - try: - await asyncio.wait_for(warmup_done.wait(), timeout=120) - except (TimeoutError, asyncio.TimeoutError): - logger.warning("Duplex warmup still running after 120 s; admitting the client anyway.") + await _wait_for_duplex_warmup(websocket) duplex_handler = getattr(websocket.app.state, "openai_serving_duplex", None) duplex_query = websocket.query_params.get("duplex") use_duplex_realtime = duplex_handler is not None and ( @@ -1504,6 +1643,32 @@ async def realtime_websocket(websocket: WebSocket): await connection.handle_connection() +async def _wait_for_duplex_warmup(websocket: WebSocket) -> None: + """Hold real clients until the startup duplex warmup finishes. + + The warmup connection marks itself with ``vllm_omni_warmup=1`` and passes through. + """ + warmup_done = getattr(websocket.app.state, "duplex_warmup_done", None) + if warmup_done is not None and not warmup_done.is_set() and websocket.query_params.get("vllm_omni_warmup") != "1": + try: + await asyncio.wait_for(warmup_done.wait(), timeout=120) + except (TimeoutError, asyncio.TimeoutError): + logger.warning("Duplex warmup still running after 120 s; admitting the client anyway.") + + +@router.websocket("/v1/duplex") +async def duplex_websocket(websocket: WebSocket): + """Alias of ``/v1/realtime?duplex=1``: the same Realtime duplex session protocol.""" + await _wait_for_duplex_warmup(websocket) + handler = getattr(websocket.app.state, "openai_serving_duplex", None) + if handler is None: + await websocket.accept() + await websocket.send_json({"type": "error", "error": "Duplex API is not available", "code": "unsupported"}) + await websocket.close() + return + await handler.handle_realtime_session(websocket) + + @router.websocket("/v1/realtime/robot/openpi") async def realtime_robot_openpi(websocket: WebSocket): """WebSocket endpoint for robot policy inference via OpenPI messages.""" @@ -1534,18 +1699,6 @@ async def realtime_robot_openpi(websocket: WebSocket): await connection.handle_connection() -@router.websocket("/v1/duplex") -async def duplex_websocket(websocket: WebSocket): - """WebSocket endpoint for vLLM-Omni duplex session control.""" - handler = getattr(websocket.app.state, "openai_serving_duplex", None) - if handler is None: - await websocket.accept() - await websocket.send_json({"type": "error", "error": "Duplex API is not available", "code": "unsupported"}) - await websocket.close() - return - await handler.handle_session(websocket) - - # RL Rollout serving (RFC #3747, P0) diff --git a/vllm_omni/entrypoints/openai/audio_utils_mixin.py b/vllm_omni/entrypoints/openai/audio_utils_mixin.py index a13318ad4..b82842c42 100644 --- a/vllm_omni/entrypoints/openai/audio_utils_mixin.py +++ b/vllm_omni/entrypoints/openai/audio_utils_mixin.py @@ -21,7 +21,6 @@ LONGEVITY: See ``openai/README.md`` (utils vs helpers, no overlap). """ -import math from io import BytesIO import numpy as np @@ -36,6 +35,9 @@ from vllm_omni.entrypoints.openai.protocol.audio import ( CreateAudio, ) +# Re-exported: serving_speech and the audio tests import it from here. +from vllm_omni.utils.audio_resample import StreamingAudioResampler # noqa: F401 + try: import soundfile except ImportError: @@ -44,155 +46,6 @@ except ImportError: logger = init_logger(__name__) -class StreamingAudioResampler: - """Stateful polyphase resampler for streaming mono float audio. - - Retains filter state so output is invariant to input chunk boundaries. - """ - - _half_filter_width = 10 - _kaiser_beta = 5.0 - _max_polyphase_factor = 2_048 - - def __init__(self, source_rate: int, target_rate: int): - if source_rate <= 0 or target_rate <= 0: - raise ValueError("Audio sample rates must be positive") - self.source_rate = source_rate - self.target_rate = target_rate - rate_gcd = math.gcd(source_rate, target_rate) - self._up = target_rate // rate_gcd - self._down = source_rate // rate_gcd - if max(self._up, self._down) > self._max_polyphase_factor: - raise ValueError("source and target sample rates have an unsupported resampling ratio") - self._half_len, self._phase_kernels = self._design_polyphase_filter() - self._history_samples = self._phase_kernels.shape[1] - 1 - self.reset() - - def _design_polyphase_filter(self) -> tuple[int, np.ndarray]: - if self._up == self._down: - return 0, np.ones((1, 1), dtype=np.float32) - - max_rate = max(self._up, self._down) - half_len = self._half_filter_width * max_rate - offsets = np.arange(-half_len, half_len + 1, dtype=np.float64) - cutoff = 1.0 / max_rate - taps = cutoff * np.sinc(cutoff * offsets) - taps *= np.kaiser(taps.size, self._kaiser_beta) - taps /= np.sum(taps) - taps *= self._up - taps = np.ascontiguousarray(taps, dtype=np.float32) - - phases = [taps[phase :: self._up] for phase in range(self._up)] - phase_width = max(phase.size for phase in phases) - kernels = np.zeros((self._up, phase_width), dtype=np.float32) - for phase_index, phase in enumerate(phases): - kernels[phase_index, -phase.size :] = phase[::-1] - return half_len, kernels - - @property - def scratch_bytes(self) -> int: - pending_samples = self._ceil_div(self._input_samples * self._up, self._down) - self._output_samples - return max(0, pending_samples) * np.dtype(np.float32).itemsize - - @staticmethod - def _ceil_div(numerator: int, denominator: int) -> int: - return -(-numerator // denominator) - - def _render_outputs( - self, - combined: np.ndarray, - *, - first_output: int, - output_count: int, - input_start: int, - ) -> np.ndarray: - if output_count <= 0: - return np.empty(0, dtype=np.float32) - - output_indexes = np.arange(first_output, first_output + output_count, dtype=np.int64) - filter_indexes = output_indexes * self._down + self._half_len - source_indexes = filter_indexes // self._up - phases = filter_indexes % self._up - windows = np.lib.stride_tricks.sliding_window_view( - combined, - self._phase_kernels.shape[1], - ) - selected_windows = windows[source_indexes - input_start] - return np.sum( - selected_windows * self._phase_kernels[phases], - axis=1, - dtype=np.float32, - ) - - def _push(self, chunk: np.ndarray) -> np.ndarray: - if chunk.size == 0: - return np.empty(0, dtype=np.float32) - if self._flushed: - raise RuntimeError("cannot process audio after the resampler has been finalized; call reset first") - - combined = np.concatenate((self._history, chunk)) - next_input_samples = self._input_samples + chunk.size - stable_output_samples = max( - 0, - self._ceil_div(next_input_samples * self._up - self._half_len, self._down), - ) - output = self._render_outputs( - combined, - first_output=self._output_samples, - output_count=stable_output_samples - self._output_samples, - input_start=self._input_samples, - ) - if self._history_samples: - self._history = combined[-self._history_samples :].copy() - self._input_samples = next_input_samples - self._output_samples = stable_output_samples - return output - - def _flush(self) -> np.ndarray: - if self._flushed: - return np.empty(0, dtype=np.float32) - - total_output_samples = self._ceil_div(self._input_samples * self._up, self._down) - output_count = total_output_samples - self._output_samples - if output_count: - last_output = total_output_samples - 1 - last_source_index = (last_output * self._down + self._half_len) // self._up - right_padding = max(0, last_source_index - self._input_samples + 1) - combined = np.concatenate((self._history, np.zeros(right_padding, dtype=np.float32))) - output = self._render_outputs( - combined, - first_output=self._output_samples, - output_count=output_count, - input_start=self._input_samples, - ) - else: - output = np.empty(0, dtype=np.float32) - self._output_samples = total_output_samples - self._flushed = True - return output - - def process(self, audio: np.ndarray, *, final: bool = False) -> np.ndarray: - chunk = np.asarray(audio, dtype=np.float32) - if chunk.ndim != 1: - raise ValueError(f"Streaming audio resampling only supports mono audio, got shape {chunk.shape}") - chunk = np.ascontiguousarray(chunk, dtype=np.float32) - output = self._push(chunk) - if not final: - return output - tail = self._flush() - if not output.size: - return tail - if not tail.size: - return output - return np.concatenate((output, tail)) - - def reset(self) -> None: - self._history = np.zeros(self._history_samples, dtype=np.float32) - self._input_samples = 0 - self._output_samples = 0 - self._flushed = False - - class AudioMixin: """Mixin class to add audio-related utilities.""" diff --git a/vllm_omni/experimental/fullduplex/README.md b/vllm_omni/experimental/fullduplex/README.md index b6c918fd6..b710924f1 100644 --- a/vllm_omni/experimental/fullduplex/README.md +++ b/vllm_omni/experimental/fullduplex/README.md @@ -16,16 +16,20 @@ video_stacking.py camera-frame tiling for omni duplex video input To run JoyVL, see [`recipes/JD/JoyAI-VL-Interaction.md`](../../../recipes/JD/JoyAI-VL-Interaction.md). -The MiniCPM-o 4.5 and PersonaPlex native full-duplex runtimes graduated out -of this package. They now live in the stable tree: +The native full-duplex runtimes graduated out of this package. They now live +in the stable tree, where sessions are engine-resident (RFC +[vllm-omni#7181](https://github.com/vllm-project/vllm-omni/issues/7181)): ```text -vllm_omni/engine/duplex/ engine control plane, sessions, leases -vllm_omni/entrypoints/duplex/ WebSocket serving and Realtime projection -vllm_omni/entrypoints/duplex_request_client.py request/output lifecycle -vllm_omni/model_executor/models/minicpmo_4_5/duplex/ MiniCPM adapter -vllm_omni/model_executor/models/personaplex/duplex/ PersonaPlex adapter -vllm_omni/model_executor/models/nemotron_voicechat/duplex/ Nemotron VoiceChat adapter +vllm_omni/engine/duplex/ sessions, manager, runner, leases, typed contract +vllm_omni/engine/duplex_omni_engine.py DuplexOmniEngine (session message surface) +vllm_omni/engine/duplex_orchestrator.py DuplexOrchestrator (hosts the session manager) +vllm_omni/entrypoints/duplex_omni.py DuplexOmni + DuplexSessionHandle (Python API) +vllm_omni/entrypoints/duplex/ WebSocket transport (attachments, resume, replay) +vllm_omni/clients/duplex.py DuplexClient / DuplexClientBase +vllm_omni/model_executor/models/minicpmo_4_5/duplex/ MiniCPM-o 4.5 DuplexModelPlugin +vllm_omni/model_executor/models/personaplex/duplex/ PersonaPlex (pre-framework, not ported yet) +vllm_omni/model_executor/models/nemotron_voicechat/duplex/ Nemotron VoiceChat (pre-framework, not ported yet) vllm_omni/model_executor/duplex_sampling.py AR-runner sampling hook helper vllm_omni/outputs/duplex.py typed output decision envelope ``` @@ -51,6 +55,9 @@ only model policy. 3. Promote a helper from a model package up into `core/` only once a second model actually needs it. -For production serving, prefer the stable plugin seams instead -(`duplex_serving_adapter` / `duplex_runtime_extension` dotted strings in the -model's `pipeline.py`), as MiniCPM-o 4.5 and PersonaPlex do. +For production serving, prefer the stable plugin seam instead: implement one +`vllm_omni.engine.duplex.plugin.DuplexModelPlugin` and name it in the model's +`pipeline.py` as `duplex_plugin`, as MiniCPM-o 4.5 does. PersonaPlex and +Nemotron VoiceChat still carry their pre-framework duplex code and are ported +to the plugin contract in follow-up PRs. The contract is documented in +[`docs/design/fullduplex.md`](../../../docs/design/fullduplex.md). diff --git a/vllm_omni/model_executor/duplex_sampling.py b/vllm_omni/model_executor/duplex_sampling.py index 7dbc185d1..dd48431f6 100644 --- a/vllm_omni/model_executor/duplex_sampling.py +++ b/vllm_omni/model_executor/duplex_sampling.py @@ -19,7 +19,6 @@ class DuplexSamplingRow: row_idx: int request_id: str session_id: str | None - incarnation: int seq: int | None payload: dict[str, object] | None max_tokens: int | None @@ -73,10 +72,6 @@ class DuplexSamplingHelper: session_id = duplex.get("session_id") if not isinstance(session_id, str) or not session_id: session_id = None - try: - incarnation = int(duplex.get("incarnation", 0)) - except (TypeError, ValueError): - incarnation = 0 try: seq = int(duplex.get("seq")) except (TypeError, ValueError): @@ -95,7 +90,6 @@ class DuplexSamplingHelper: row_idx=row_idx, request_id=req_id, session_id=session_id, - incarnation=incarnation, seq=seq, payload=payload, max_tokens=max_tokens if max_tokens > 0 else None, diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/__init__.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/__init__.py index b681ad51e..e38cc12ef 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/__init__.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/__init__.py @@ -1,15 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -from .adapter import MiniCPMO45ClientRuntimeConfigError, MiniCPMO45NativeDuplexServingAdapter from .input import MiniCPMO45PcmAppendBuffer +from .plugin import MiniCPMO45ClientRuntimeConfigError, MiniCPMO45DuplexPlugin from .policy import MiniCPMO45DuplexPolicy from .stage0 import MiniCPMO45Stage0DuplexRuntime __all__ = [ - "MiniCPMO45DuplexPolicy", - "MiniCPMO45NativeDuplexServingAdapter", "MiniCPMO45ClientRuntimeConfigError", + "MiniCPMO45DuplexPlugin", + "MiniCPMO45DuplexPolicy", "MiniCPMO45PcmAppendBuffer", "MiniCPMO45Stage0DuplexRuntime", ] diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/adapter.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/adapter.py deleted file mode 100644 index 758beafc9..000000000 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/adapter.py +++ /dev/null @@ -1,335 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import base64 -from copy import deepcopy -from typing import Any - -import numpy as np -from vllm.multimodal.media import MediaConnector - -from vllm_omni.entrypoints.duplex.protocol import DuplexSessionConfig, native_duplex_opt_in -from vllm_omni.entrypoints.duplex.runtime_adapter import ( - ServingRuntimeConfigError, - reject_changed_runtime_value, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import MiniCPMO45DuplexPolicy - - -class MiniCPMO45ClientRuntimeConfigError(ServingRuntimeConfigError): - pass - - -class MiniCPMO45NativeDuplexServingAdapter: - """Serving-side MiniCPM-o 4.5 native duplex session preparation. - - The generic duplex WebSocket handler should not let client-supplied local - paths reach workers. This adapter follows the existing media connector - boundary: serving resolves client media URIs, then workers receive only - normalized PCM payloads. Server-owned model assets remain local paths here. - """ - - PRIVATE_RUNTIME_CONFIG_KEYS = frozenset( - { - "duplex_stage_sampling_params", - "duplex_stage_max_tokens", - "duplex_stage0_max_tokens", - "duplex_scheduler_token_id", - "duplex_first_append_context_tokens", - "ref_audio_data", - "ref_audio_format", - "ref_audio_sample_rate_hz", - "initial_user_text", - } - ) - - @classmethod - def is_enabled(cls, config: DuplexSessionConfig) -> bool: - return native_duplex_opt_in(config.extra_body) is True - - @classmethod - def validate_client_config(cls, config: DuplexSessionConfig) -> None: - cls.validate_client_extra_body(config.extra_body) - - @classmethod - def validate_client_extra_body(cls, extra_body: object) -> None: - if not isinstance(extra_body, dict): - return - private_keys = sorted(cls.PRIVATE_RUNTIME_CONFIG_KEYS.intersection(extra_body)) - if private_keys: - raise MiniCPMO45ClientRuntimeConfigError( - "native duplex runtime configuration is server-owned: " + ", ".join(private_keys) - ) - - @classmethod - def runtime_config_for_update( - cls, - config: DuplexSessionConfig, - current: object, - ) -> dict[str, object]: - runtime_config = deepcopy(dict(current)) if isinstance(current, dict) else {} - reject_changed_runtime_value( - config.instructions, - runtime_config.get("instructions"), - message="instructions cannot be changed after the session is created", - code="instructions_update_unsupported", - error_cls=MiniCPMO45ClientRuntimeConfigError, - ) - stage_max_tokens = runtime_config.get("duplex_stage_max_tokens") - stage_max_tokens = deepcopy(stage_max_tokens) if isinstance(stage_max_tokens, dict) else {} - stage_max_tokens["0"] = ( - config.max_tokens if isinstance(config.max_tokens, int) and config.max_tokens > 0 else 20 - ) - stage_max_tokens.setdefault("1", 8192) - runtime_config["duplex_stage_max_tokens"] = stage_max_tokens - - stage_sampling = runtime_config.get("duplex_stage_sampling_params") - stage_sampling = deepcopy(stage_sampling) if isinstance(stage_sampling, dict) else {} - stage0 = stage_sampling.get("0") - stage0 = deepcopy(stage0) if isinstance(stage0, dict) else {} - stage0["temperature"] = config.temperature if config.temperature is not None else 0.7 - stage_sampling["0"] = stage0 - runtime_config["duplex_stage_sampling_params"] = stage_sampling - return runtime_config - - @classmethod - async def prepare_runtime_config(cls, config: DuplexSessionConfig, *, model_config: Any) -> dict[str, object]: - extra_body = dict(config.extra_body) - if any(key in extra_body for key in ("ref_audio_path", "tts_ref_audio_path")): - raise ValueError("ref_audio_path is not accepted by native duplex; use ref_audio URI instead") - cls.validate_client_config(config) - runtime_config: dict[str, object] = {"instructions": config.instructions} - initial_user_text = extra_body.pop("duplex_initial_user_text", None) - if isinstance(initial_user_text, str) and initial_user_text: - runtime_config["initial_user_text"] = initial_user_text - cls._apply_default_scheduler_policy(runtime_config, config=config, model_config=model_config) - - ref_audio = config.ref_audio - extra_ref_audio = extra_body.get("ref_audio") - if ref_audio is None and isinstance(extra_ref_audio, str): - extra_body.pop("ref_audio") - ref_audio = extra_ref_audio - extra_tts_ref_audio = extra_body.get("tts_ref_audio") - if ref_audio is None and isinstance(extra_tts_ref_audio, str): - extra_body.pop("tts_ref_audio") - ref_audio = extra_tts_ref_audio - - if ref_audio is None: - if any(str(modality).lower() == "audio" for modality in config.modalities): - raise MiniCPMO45ClientRuntimeConfigError( - "MiniCPM-o native duplex audio output requires ref_audio", - code="ref_audio_required", - ) - cls._apply_first_append_context_tokens( - runtime_config, - model_config=model_config, - instructions=config.instructions, - initial_user_text=initial_user_text, - ref_sample_count=None, - ) - config.extra_body = extra_body - return runtime_config - else: - wav_np, sr = await cls.resolve_ref_audio(ref_audio, model_config=model_config) - - wav_np = cls.normalize_ref_audio(wav_np, int(sr), target_sr=16000) - # Trim to a whole number of pooled audio embeddings (100 ms frames) so - # the first-append scheduler reserve can count them exactly. - usable = (len(wav_np) // MiniCPMO45DuplexPolicy.SAMPLES_PER_AUDIO_TOKEN) * ( - MiniCPMO45DuplexPolicy.SAMPLES_PER_AUDIO_TOKEN - ) - wav_np = wav_np[:usable] - ref_audio_bytes = np.ascontiguousarray(wav_np, dtype=np.float32).tobytes() - runtime_config["ref_audio_data"] = base64.b64encode(ref_audio_bytes).decode("ascii") - runtime_config["ref_audio_format"] = "pcm_f32le" - runtime_config["ref_audio_sample_rate_hz"] = 16000 - cls._apply_first_append_context_tokens( - runtime_config, - model_config=model_config, - instructions=config.instructions, - initial_user_text=initial_user_text, - ref_sample_count=len(wav_np), - ) - config.extra_body = extra_body - config.ref_audio = None - return runtime_config - - @classmethod - def _apply_default_scheduler_policy( - cls, - runtime_config: dict[str, object], - *, - config: DuplexSessionConfig, - model_config: Any, - ) -> None: - stage0_max_tokens = config.max_tokens if isinstance(config.max_tokens, int) and config.max_tokens > 0 else 20 - runtime_config["duplex_stage_max_tokens"] = {"0": stage0_max_tokens, "1": 8192} - stage0_params: dict[str, object] = { - "temperature": config.temperature if config.temperature is not None else 0.7, - "top_p": 0.8, - "top_k": 20, - "repetition_penalty": 1.05, - } - stop_token_ids = cls._native_stage0_stop_token_ids(model_config) - if stop_token_ids: - stage0_params["stop_token_ids"] = stop_token_ids - # Stage 1 keeps the deploy YAML's codec knobs, minus upstream's - # min_new_token=50: that floor is for whole-utterance chat TTS, while a - # duplex chunk is 26 codec samples, so check_stop would never release - # the Talker and Thinker would stall on the next model turn. - runtime_config["duplex_stage_sampling_params"] = {"0": stage0_params, "1": {"min_tokens": 0}} - scheduler_token_id = cls._native_scheduler_token_id(model_config) - if scheduler_token_id is not None: - runtime_config["duplex_scheduler_token_id"] = scheduler_token_id - - @classmethod - def _apply_first_append_context_tokens( - cls, - runtime_config: dict[str, object], - *, - model_config: Any, - instructions: object, - initial_user_text: object, - ref_sample_count: int | None, - ) -> None: - """Precompute the exact session-context token count for the engine. - - The first data-plane append carries the system template and optional - reference-audio embeddings ahead of the first unit. The engine - reserves scheduler slots from this count; an inexact count turns into - pad embeddings inside the model KV (surplus) or truncated context - (deficit), so it is computed with the same template and pooling math - the worker uses. - """ - if "duplex_first_append_context_tokens" in runtime_config: - return - tokenizer = cls._load_native_tokenizer(model_config) - if tokenizer is None: - return - prefix, suffix = MiniCPMO45DuplexPolicy.session_context_texts( - instructions, - ref_sample_count is not None, - initial_user_text, - ) - try: - prefix_ids = tokenizer.encode(prefix, add_special_tokens=False) - suffix_ids = tokenizer.encode(suffix, add_special_tokens=False) - except Exception: - return - ref_tokens = MiniCPMO45DuplexPolicy.audio_token_count(ref_sample_count or 0) - runtime_config["duplex_first_append_context_tokens"] = len(prefix_ids) + ref_tokens + len(suffix_ids) - - @staticmethod - def _native_stage0_stop_token_ids(model_config: Any) -> list[int]: - tokenizer = MiniCPMO45NativeDuplexServingAdapter._load_native_tokenizer(model_config) - if tokenizer is None: - return [] - out: list[int] = [] - stop_token_fields = ( - "chunk_eos_token_id", - "chunk_tts_eos_token_id", - "listen_token_id", - "turn_eos_token_id", - ) - for field in stop_token_fields: - token = MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS[field] - token_id = MiniCPMO45NativeDuplexServingAdapter._convert_token_to_id(tokenizer, token) - if token_id is not None and token_id not in out: - out.append(token_id) - return out - - @staticmethod - def _native_scheduler_token_id(model_config: Any) -> int | None: - tokenizer = MiniCPMO45NativeDuplexServingAdapter._load_native_tokenizer(model_config) - if tokenizer is None: - return None - scheduler_tokens = ( - MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS["unit_token_id"], - MiniCPMO45DuplexPolicy.OPTIONAL_TOKEN_FIELDS["audio_placeholder_token_id"], - ) - for token in scheduler_tokens: - token_id = MiniCPMO45NativeDuplexServingAdapter._convert_token_to_id(tokenizer, token) - if token_id is not None: - return token_id - eos_id = getattr(tokenizer, "eos_token_id", None) - if eos_id is None: - return None - try: - return int(eos_id) - except (TypeError, ValueError): - return None - - @staticmethod - def _load_native_tokenizer(model_config: Any) -> Any | None: - model_path = getattr(model_config, "model", None) - if not isinstance(model_path, str) or not model_path: - return None - try: - from transformers import AutoTokenizer - - return AutoTokenizer.from_pretrained( - model_path, - trust_remote_code=True, - local_files_only=True, - ) - except Exception: - return None - - @staticmethod - def _convert_token_to_id(tokenizer: Any, token: str) -> int | None: - convert = getattr(tokenizer, "convert_tokens_to_ids", None) - value = None - if callable(convert): - value = convert(token) - if isinstance(value, list): - value = value[0] if len(value) == 1 else None - if value is None: - token_id = -1 - else: - try: - token_id = int(value) - except (TypeError, ValueError): - token_id = -1 - unk_token_id = getattr(tokenizer, "unk_token_id", None) - if token_id >= 0 and token_id != unk_token_id: - return token_id - encode = getattr(tokenizer, "encode", None) - if callable(encode): - try: - ids = list(encode(token, add_special_tokens=False)) - except TypeError: - ids = list(encode(token)) - if len(ids) == 1: - try: - token_id = int(ids[0]) - except (TypeError, ValueError): - token_id = -1 - if token_id >= 0 and token_id != unk_token_id: - return token_id - return None - - @staticmethod - async def resolve_ref_audio(ref_audio: str, *, model_config: Any) -> tuple[np.ndarray, int]: - connector = MediaConnector( - allowed_local_media_path=getattr(model_config, "allowed_local_media_path", None), - allowed_media_domains=getattr(model_config, "allowed_media_domains", None), - ) - wav_np, sr = await connector.fetch_audio_async(ref_audio) - return np.asarray(wav_np, dtype=np.float32), int(sr) - - @staticmethod - def normalize_ref_audio(wav_np: np.ndarray, sample_rate: int, *, target_sr: int) -> np.ndarray: - wav_np = np.asarray(wav_np, dtype=np.float32) - if wav_np.ndim > 1: - wav_np = wav_np.mean(axis=-1) - wav_np = wav_np.reshape(-1) - if sample_rate <= 0 or sample_rate == target_sr or wav_np.size == 0: - return wav_np.astype(np.float32, copy=False) - import torch - import torchaudio - - audio = torch.from_numpy(wav_np).to(dtype=torch.float32).unsqueeze(0) - resampled = torchaudio.functional.resample(audio, int(sample_rate), int(target_sr)) - return resampled.squeeze(0).cpu().numpy().astype(np.float32, copy=False) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/capabilities.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/capabilities.py index 668405892..d850f4c3c 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/capabilities.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/capabilities.py @@ -3,7 +3,12 @@ from __future__ import annotations -from vllm_omni.entrypoints.duplex.protocol import DuplexCapabilities +from vllm_omni.engine.duplex.config import DuplexCapabilities + +#: Silence units a text-seeded turn is given to generate on. The model speaks +#: per audio unit, so a seeded turn still needs a clock; three units is what +#: the benchmark harness uses for the same purpose. +MINICPMO45_TEXT_TURN_PRIMING_UNITS = 3 def minicpmo45_native_capabilities(*, max_sessions: int = 1) -> DuplexCapabilities: @@ -31,11 +36,14 @@ def minicpmo45_native_capabilities(*, max_sessions: int = 1) -> DuplexCapabiliti supports_session_resume=True, session_admission_mode="engine_managed", supports_audio_truncate=True, + # The session template ends at the assistant turn, so seeded text puts + # the model in position to answer; it still generates per audio unit, + # hence the priming units. + supports_chat_completions=True, # via initial_user_text + text_turn_priming_units=MINICPMO45_TEXT_TURN_PRIMING_UNITS, requires_model_runner_kv=True, requires_native_stage_role=True, - implementation_level="model_native_duplex", adapter_patterns=["scheduler_data_plane"], - input_modes=["append_audio_chunk"], signal_sources=["model_native", "client_event", "server_policy"], stage_handoff_transport="scheduler_data_plane", chunk_period_ms=1000, diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/compat.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/compat.py index 9a4bcf9b0..5b6b5c909 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/compat.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/compat.py @@ -3,14 +3,12 @@ from __future__ import annotations -from typing import Any - -def patch_minicpmo_remote_config(config: Any) -> None: +def patch_minicpmo_remote_config(config: object) -> None: """Fill config fields expected by the vLLM MiniCPM implementation.""" seen: set[int] = set() - def patch_cfg(cfg: Any) -> None: + def patch_cfg(cfg: object) -> None: if cfg is None or id(cfg) in seen: return seen.add(id(cfg)) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/data_plane.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/data_plane.py index d4063ff3f..7bf51abfd 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/data_plane.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/data_plane.py @@ -12,6 +12,7 @@ from vllm.logger import init_logger from vllm_omni.engine.duplex.contracts import ( duplex_resource_request_belongs_to_session, ) +from vllm_omni.engine.duplex.plugin import DuplexDataPlane from vllm_omni.outputs.duplex import get_duplex_output_decision logger = init_logger(__name__) @@ -44,6 +45,9 @@ class _TurnState: @dataclass(slots=True) class _RequestState: audio_offset: int = 0 + #: A snapshot below ``audio_offset`` that has not been accepted as a restart + #: yet. See :meth:`MiniCPMO45DataPlaneSession.slice_cumulative_audio`. + audio_restart_candidate: int | None = None uses_segment_text_metadata: bool = False pending_audio_without_text: list[dict[str, object]] = field(default_factory=list) terminal: bool = False @@ -53,7 +57,7 @@ class _RequestState: return self.turns.setdefault(turn_id, _TurnState()) -class MiniCPMO45DataPlaneSession: +class MiniCPMO45DataPlaneSession(DuplexDataPlane): """MiniCPM output projector and request/turn cursor owner. The scheduler request owns cumulative Stage1 audio, while model turns own @@ -487,13 +491,33 @@ class MiniCPMO45DataPlaneSession: prev_samples = state.audio_offset if prev_samples <= 0: state.audio_offset = num_samples + state.audio_restart_candidate = None return audio_data if num_samples == prev_samples: return None if num_samples < prev_samples: + # Below the cursor means one of two things, and they are not + # distinguishable from this snapshot alone: the producer restarted + # its buffer, or this is a unit that carries no cumulative audio at + # all (a continuation placeholder arrives as a single sample). + # Rewinding the cursor for the second case is expensive and silent: + # the next real snapshot is then sliced from ~0 and the whole + # session's audio goes out again as one delta. So defer the + # decision — hold the candidate and leave the cursor where it is. + candidate = state.audio_restart_candidate + if candidate is None or num_samples <= candidate: + state.audio_restart_candidate = num_samples + return audio_data + # A second below-cursor snapshot that grew from the first: the + # producer really did restart, and this is its delta. + state.audio_offset = num_samples + state.audio_restart_candidate = num_samples + prev_samples = candidate + else: + # Above the cursor: whatever the dip was, this stream never + # restarted, so it is sliced from the cursor as usual. + state.audio_restart_candidate = None state.audio_offset = num_samples - return audio_data - state.audio_offset = num_samples try: import torch diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/input.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/input.py index 1b30ec4c2..4feb8cb97 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/input.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/input.py @@ -6,9 +6,11 @@ from __future__ import annotations import base64 import binascii from dataclasses import dataclass -from typing import Any import numpy as np +from numpy.typing import NDArray + +from vllm_omni.engine.duplex.plugin import PcmAppendBuffer, PcmAppendReservation @dataclass(frozen=True, slots=True) @@ -18,7 +20,7 @@ class _PcmSpan: is_speech: bool -class MiniCPMO45PcmAppendReservation: +class MiniCPMO45PcmAppendReservation(PcmAppendReservation): __slots__ = ( "_active", "_force_listen", @@ -74,7 +76,7 @@ class MiniCPMO45PcmAppendReservation: self._owner._rollback_reservation(self) -def validate_native_ref_audio_config(session_config: dict[str, Any]) -> None: +def validate_native_ref_audio_config(session_config: dict[str, object]) -> None: extra_body = session_config.get("extra_body") if not isinstance(extra_body, dict): extra_body = {} @@ -84,7 +86,7 @@ def validate_native_ref_audio_config(session_config: dict[str, Any]) -> None: raise ValueError("native duplex ref_audio_path is not accepted; resolve ref_audio in serving first") -def decode_native_ref_audio_from_config(session_config: dict[str, Any]) -> np.ndarray | None: +def decode_native_ref_audio_from_config(session_config: dict[str, object]) -> NDArray[np.float32] | None: validate_native_ref_audio_config(session_config) extra_body = session_config.get("extra_body") if not isinstance(extra_body, dict): @@ -106,7 +108,7 @@ def decode_native_ref_audio_from_config(session_config: dict[str, Any]) -> np.nd return np.frombuffer(raw, dtype=" None: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py new file mode 100644 index 000000000..673640461 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py @@ -0,0 +1,773 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""MiniCPM-o 4.5 full-duplex model plugin: engine policy and session policy in one class. + +Engine policy (scheduler token budget, resumable Stage0 prompt, listen +decision) and session policy (client config validation, reference-audio +resolution, per-stage sampling knobs) both run engine-side, so one class owns +them. Client-supplied media URIs are resolved here through the media +connector; workers only ever receive normalized PCM payloads. +""" + +from __future__ import annotations + +import asyncio +import base64 +from base64 import b64decode +from binascii import Error as BinasciiError +from collections.abc import Mapping +from copy import deepcopy +from typing import TYPE_CHECKING, Any, cast + +import numpy as np +from numpy.typing import NDArray +from vllm.multimodal.media import MediaConnector +from vllm.sampling_params import SamplingParams + +from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig +from vllm_omni.engine.duplex.contracts import ( + DuplexAppendPlan, + DuplexFence, + DuplexOutputAction, + DuplexOutputDecision, +) +from vllm_omni.engine.duplex.plugin import ( + DuplexModelPlugin, + DuplexRuntimeConfigError, + EncodeAudio, + reject_changed_runtime_value, +) +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.capabilities import ( + minicpmo45_native_capabilities, +) +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.data_plane import ( + MiniCPMO45DataPlaneContext, + MiniCPMO45DataPlaneSession, +) +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import MiniCPMO45DuplexPolicy +from vllm_omni.model_executor.models.minicpmo_4_5.duplex.session import ( + MiniCPMO45ServingSessionState, +) + +if TYPE_CHECKING: + import torch + from transformers import PreTrainedTokenizerBase + from vllm.config import ModelConfig + +_DUPLEX_CHUNK_SAMPLES = 16000 +_DUPLEX_SAMPLES_PER_AUDIO_TOKEN = 1600 +# + 64 resampler embeddings + per frame (max_slice_nums=1), +# matching MiniCPMO45DuplexPolicy.VISION_TOKENS_PER_FRAME. +_DUPLEX_VISION_TOKENS_PER_FRAME = 66 +# Official stacked pair uses max_slice_nums=[2, 1]: the current frame is HD +# sliced (1 source + 2 patches on 960x540) and the composite is not. +_DUPLEX_HD_SLICES_PER_BASE_FRAME = 3 + +PRIVATE_RUNTIME_CONFIG_KEYS = frozenset( + { + "duplex_stage_sampling_params", + "duplex_stage_max_tokens", + "duplex_stage0_max_tokens", + "duplex_scheduler_token_id", + "duplex_first_append_context_tokens", + "ref_audio_data", + "ref_audio_format", + "ref_audio_sample_rate_hz", + "initial_user_text", + } +) + + +class MiniCPMO45ClientRuntimeConfigError(DuplexRuntimeConfigError): + pass + + +# ---- engine policy helpers: scheduler token budget ---- + + +def _duplex_frame_count(payload: object) -> int: + if not isinstance(payload, dict): + return 0 + frames = payload.get("video_frames") + if not isinstance(frames, list): + return 0 + return sum(1 for frame in frames if isinstance(frame, str) and frame) + + +def _duplex_vision_tokens(payload: object) -> int: + """Scheduler slots for this append's camera track. + + Audio is never stacked: a unit still carries one second of soundtrack. + ``stack_frames`` only adds a second *image*. Official HD on that pair is + ``[2, 1]``, so the base frame reserves three 66-token blocks and every + extra frame reserves one. + """ + count = _duplex_frame_count(payload) + if count <= 0: + return 0 + if count >= 2: + return (_DUPLEX_HD_SLICES_PER_BASE_FRAME + (count - 1)) * _DUPLEX_VISION_TOKENS_PER_FRAME + return count * _DUPLEX_VISION_TOKENS_PER_FRAME + + +def _duplex_pcm_sample_count(payload: object) -> int | None: + if not isinstance(payload, dict): + return None + audio = payload.get("audio") or payload.get("data") + if payload.get("format") != "pcm_f32le" or not isinstance(audio, str): + return None + try: + raw = b64decode(audio, validate=True) + except (BinasciiError, ValueError): + return None + return len(raw) // 4 + + +def duplex_payload_is_exact_chunks(payload: object) -> bool: + sample_count = _duplex_pcm_sample_count(payload) + return sample_count is not None and sample_count != 0 and sample_count % _DUPLEX_CHUNK_SAMPLES == 0 + + +def duplex_first_append_unit_count(payload: object) -> int | None: + sample_count = _duplex_pcm_sample_count(payload) + if not sample_count or sample_count % _DUPLEX_CHUNK_SAMPLES != 0: + return None + return max(1, sample_count // _DUPLEX_CHUNK_SAMPLES - 1) + + +def duplex_scheduler_token_budget(payload: object, *, default: int = 64) -> int: + vision_tokens = _duplex_vision_tokens(payload) + sample_count = _duplex_pcm_sample_count(payload) + if sample_count is None: + return max(1, int(default)) + vision_tokens + sample_count = max(1, sample_count) + if sample_count % _DUPLEX_CHUNK_SAMPLES == 0: + units = sample_count // _DUPLEX_CHUNK_SAMPLES + return units * (2 + _DUPLEX_CHUNK_SAMPLES // _DUPLEX_SAMPLES_PER_AUDIO_TOKEN) + vision_tokens + return max(16, min(768, sample_count // _DUPLEX_SAMPLES_PER_AUDIO_TOKEN + 8)) + vision_tokens + + +def duplex_first_append_context_reserve(runtime_config: object) -> int: + if not isinstance(runtime_config, dict): + return 48 + exact = runtime_config.get("duplex_first_append_context_tokens") + if isinstance(exact, int) and exact >= 0: + return exact + reserve = 48 + ref = runtime_config.get("ref_audio_data") + if isinstance(ref, str) and ref: + try: + raw = b64decode(ref, validate=True) + except (BinasciiError, ValueError): + raw = b"" + if raw: + reserve += max(0, (len(raw) // 4) // _DUPLEX_SAMPLES_PER_AUDIO_TOKEN + 8) + return reserve + + +def _duplex_force_listen_count(extra_body: object) -> int: + raw = extra_body.get("force_listen_count") if isinstance(extra_body, dict) else None + try: + return 0 if raw is None else max(0, int(raw)) + except (TypeError, ValueError): + return 0 + + +def build_duplex_data_plane_prompt( + *, + request_id: str, + fence: DuplexFence, + session_config: dict[str, object], + runtime_config: dict[str, object], + seq: int, + turn_seq: int, + payload: object, + final: bool, +) -> dict[str, object]: + token_budget = duplex_scheduler_token_budget(payload) + if seq <= 1: + context_reserve = duplex_first_append_context_reserve(runtime_config) + token_budget += context_reserve + first_units = duplex_first_append_unit_count(payload) + if first_units is not None: + token_budget = context_reserve + first_units * 12 - 1 + _duplex_vision_tokens(payload) + if seq > 1 and duplex_payload_is_exact_chunks(payload): + token_budget += 1 + if final and duplex_payload_is_exact_chunks(payload): + token_budget += 12 + extra_body = session_config.get("extra_body") + raw_token_id = runtime_config.get("duplex_scheduler_token_id") + try: + token_id = 0 if raw_token_id is None else max(0, int(raw_token_id)) + except (TypeError, ValueError): + token_id = 0 + force_listen_count = _duplex_force_listen_count(extra_body) + if ( + force_listen_count > 0 + and turn_seq <= force_listen_count + and isinstance(payload, dict) + and payload.get("force_listen") is not True + ): + payload = {**payload, "force_listen": True} + return { + "prompt_token_ids": [token_id] * token_budget, + "model_intermediate_buffer": { + "request_id": request_id, + "global_request_id": [fence.session_id], + "duplex": { + "fence": fence, + "session_id": fence.session_id, + "epoch": fence.epoch, + "seq": seq, + "turn_id": fence.turn_id, + "turn_seq": turn_seq, + "mode": "append_audio_chunk", + "payload": payload, + "final": final, + "data_plane": True, + "session_config": dict(session_config), + "runtime_config": dict(runtime_config), + "scheduler_token_budget": token_budget, + "scheduler_token_id": token_id, + }, + }, + } + + +# ---- engine policy helpers: listen decision ---- + + +def _coerce_int(value: object) -> int | None: + detach = getattr(value, "detach", None) + if callable(detach): + try: + flat: torch.Tensor = detach().cpu().reshape(-1) + if flat.numel() == 0: + return None + value = flat[0].item() + except Exception: + return None + try: + return int(cast(Any, value)) # Any: duck-typed scalar (int/float/str/tensor item) + except (TypeError, ValueError): + return None + + +def _coerce_int_list(value: object) -> list[int]: + if value is None: + return [] + if hasattr(value, "detach"): + try: + value = value.detach().cpu().reshape(-1).tolist() + except Exception: + return [] + if not isinstance(value, (list, tuple)): + return [] + return [token_id for item in value if (token_id := _coerce_int(item)) is not None] + + +def _first_completion(output: object) -> object | None: + outputs = getattr(output, "outputs", None) + return outputs[0] if isinstance(outputs, list) and outputs else None + + +def _multimodal_output(output: object, completion: object | None) -> dict[str, object]: + metadata = getattr(output, "multimodal_output", None) + if isinstance(metadata, dict): + return metadata + metadata = getattr(completion, "multimodal_output", None) if completion is not None else None + return metadata if isinstance(metadata, dict) else {} + + +def _special_token_ids(metadata: dict[str, object]) -> dict[str, int]: + sources: list[object] = [metadata.get("special_token_ids"), metadata.get("meta")] + sources.append( + { + key.removeprefix("meta."): value + for key, value in metadata.items() + if isinstance(key, str) and key.startswith("meta.") + } + ) + token_ids: dict[str, int] = {} + for source in sources: + if not isinstance(source, dict): + continue + for key, value in source.items(): + token_id = _coerce_int(value) + if isinstance(key, str) and token_id is not None and token_id >= 0: + token_ids[key] = token_id + return token_ids + + +def _completion_token_ids(completion: object | None) -> list[int]: + if completion is None: + return [] + for attribute in ("token_ids", "cumulative_token_ids"): + token_ids = _coerce_int_list(getattr(completion, attribute, None)) + if token_ids: + return token_ids + return [] + + +def _stage_config_value(runtime_config: dict[str, object], key: str, stage_id: int) -> object | None: + raw = runtime_config.get(key) + if isinstance(raw, dict): + value = raw.get(stage_id) + return raw.get(str(stage_id)) if value is None else value + if isinstance(raw, (list, tuple)) and stage_id < len(raw): + return raw[stage_id] + return None + + +# ---- session policy helpers: tokenizer / reference audio ---- + + +def _load_tokenizer(model_config: ModelConfig | None) -> PreTrainedTokenizerBase | None: + model_path = getattr(model_config, "model", None) + if not isinstance(model_path, str) or not model_path: + return None + try: + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained( + model_path, + trust_remote_code=True, + local_files_only=True, + ) + except Exception: + return None + + +def _convert_token_to_id(tokenizer: PreTrainedTokenizerBase, token: str) -> int | None: + convert = getattr(tokenizer, "convert_tokens_to_ids", None) + value = None + if callable(convert): + value = convert(token) + if isinstance(value, list): + value = value[0] if len(value) == 1 else None + if value is None: + token_id = -1 + else: + try: + token_id = int(value) + except (TypeError, ValueError): + token_id = -1 + unk_token_id = getattr(tokenizer, "unk_token_id", None) + if token_id >= 0 and token_id != unk_token_id: + return token_id + encode = getattr(tokenizer, "encode", None) + if callable(encode): + try: + ids = list(encode(token, add_special_tokens=False)) + except TypeError: + ids = list(encode(token)) + if len(ids) == 1: + try: + token_id = int(ids[0]) + except (TypeError, ValueError): + token_id = -1 + if token_id >= 0 and token_id != unk_token_id: + return token_id + return None + + +def _stage0_stop_token_ids(tokenizer: PreTrainedTokenizerBase | None) -> list[int]: + if tokenizer is None: + return [] + out: list[int] = [] + stop_token_fields = ( + "chunk_eos_token_id", + "chunk_tts_eos_token_id", + "listen_token_id", + "turn_eos_token_id", + ) + for field in stop_token_fields: + token = MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS[field] + token_id = _convert_token_to_id(tokenizer, token) + if token_id is not None and token_id not in out: + out.append(token_id) + return out + + +def _scheduler_token_id(tokenizer: PreTrainedTokenizerBase | None) -> int | None: + if tokenizer is None: + return None + scheduler_tokens = ( + MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS["unit_token_id"], + MiniCPMO45DuplexPolicy.OPTIONAL_TOKEN_FIELDS["audio_placeholder_token_id"], + ) + for token in scheduler_tokens: + token_id = _convert_token_to_id(tokenizer, token) + if token_id is not None: + return token_id + eos_id = getattr(tokenizer, "eos_token_id", None) + if eos_id is None: + return None + try: + return int(eos_id) + except (TypeError, ValueError): + return None + + +async def resolve_ref_audio(ref_audio: str, *, model_config: ModelConfig | None) -> tuple[NDArray[np.float32], int]: + connector = MediaConnector( + allowed_local_media_path=getattr(model_config, "allowed_local_media_path", None), + allowed_media_domains=getattr(model_config, "allowed_media_domains", None), + ) + wav_np, sr = await connector.fetch_audio_async(ref_audio) + return np.asarray(wav_np, dtype=np.float32), int(sr) + + +def normalize_ref_audio(wav_np: NDArray[np.float32], sample_rate: int, *, target_sr: int) -> NDArray[np.float32]: + wav_np = np.asarray(wav_np, dtype=np.float32) + if wav_np.ndim > 1: + wav_np = wav_np.mean(axis=-1) + wav_np = wav_np.reshape(-1) + if sample_rate <= 0 or sample_rate == target_sr or wav_np.size == 0: + return wav_np.astype(np.float32, copy=False) + import torch + import torchaudio + + audio = torch.from_numpy(wav_np).to(dtype=torch.float32).unsqueeze(0) + resampled = torchaudio.functional.resample(audio, int(sample_rate), int(target_sr)) + return resampled.squeeze(0).cpu().numpy().astype(np.float32, copy=False) + + +def _apply_first_append_context_tokens( + runtime_config: dict[str, object], + *, + tokenizer: PreTrainedTokenizerBase | None, + instructions: object, + initial_user_text: object, + ref_sample_count: int | None, +) -> None: + """Precompute the exact session-context token count for the engine. + + The first data-plane append carries the system template and optional + reference-audio embeddings ahead of the first unit. The engine reserves + scheduler slots from this count; an inexact count turns into pad + embeddings inside the model KV (surplus) or truncated context (deficit), + so it is computed with the same template and pooling math the worker uses. + """ + if "duplex_first_append_context_tokens" in runtime_config or tokenizer is None: + return + prefix, suffix = MiniCPMO45DuplexPolicy.session_context_texts( + instructions, + ref_sample_count is not None, + initial_user_text, + ) + try: + prefix_ids = tokenizer.encode(prefix, add_special_tokens=False) + suffix_ids = tokenizer.encode(suffix, add_special_tokens=False) + except Exception: + return + ref_tokens = MiniCPMO45DuplexPolicy.audio_token_count(ref_sample_count or 0) + runtime_config["duplex_first_append_context_tokens"] = len(prefix_ids) + ref_tokens + len(suffix_ids) + + +def _apply_default_scheduler_policy( + runtime_config: dict[str, object], + *, + config: DuplexSessionConfig, + tokenizer: PreTrainedTokenizerBase | None, +) -> None: + stage0_max_tokens = config.max_tokens if isinstance(config.max_tokens, int) and config.max_tokens > 0 else 20 + runtime_config["duplex_stage_max_tokens"] = {"0": stage0_max_tokens, "1": 8192} + stage0_params: dict[str, object] = { + "temperature": config.temperature if config.temperature is not None else 0.7, + "top_p": 0.8, + "top_k": 20, + "repetition_penalty": 1.05, + } + stop_token_ids = _stage0_stop_token_ids(tokenizer) + if stop_token_ids: + stage0_params["stop_token_ids"] = stop_token_ids + # Stage 1 keeps the deploy YAML's codec knobs, minus upstream's + # min_new_token=50: that floor is for whole-utterance chat TTS, while a + # duplex chunk is 26 codec samples, so check_stop would never release + # the Talker and Thinker would stall on the next model turn. + runtime_config["duplex_stage_sampling_params"] = {"0": stage0_params, "1": {"min_tokens": 0}} + scheduler_token_id = _scheduler_token_id(tokenizer) + if scheduler_token_id is not None: + runtime_config["duplex_scheduler_token_id"] = scheduler_token_id + + +class MiniCPMO45DuplexPlugin(DuplexModelPlugin): + """MiniCPM-owned sampling policy, append planning, session state and output projection.""" + + plugin_id = "minicpmo45" + private_runtime_config_keys = PRIVATE_RUNTIME_CONFIG_KEYS + silence_continuation_samples = 16000 + + def __init__(self, encode_audio: EncodeAudio) -> None: + super().__init__(encode_audio) + self.data_plane = MiniCPMO45DataPlaneSession(encode_audio) + # The tokenizer is only needed for special-token ids and the session + # context token count; it is the same for every session, so load it + # once per model path, off the orchestrator loop. + self._tokenizers: dict[str, PreTrainedTokenizerBase | None] = {} + self._tokenizer_lock = asyncio.Lock() + + async def _tokenizer_for(self, model_config: ModelConfig | None) -> PreTrainedTokenizerBase | None: + model_path = getattr(model_config, "model", None) + key = model_path if isinstance(model_path, str) else "" + if key in self._tokenizers: + return self._tokenizers[key] + async with self._tokenizer_lock: + if key not in self._tokenizers: + self._tokenizers[key] = await asyncio.to_thread(_load_tokenizer, model_config) + return self._tokenizers[key] + + # ---- engine policy (the resumable Stage0 request) ---- + + def configure_sampling_params( + self, + *, + runtime_config: dict[str, object], + defaults: tuple[object, ...], + ) -> tuple[object, ...]: + configured: list[object] = [] + for stage_id, default in enumerate(defaults): + max_tokens = _coerce_int(_stage_config_value(runtime_config, "duplex_stage_max_tokens", stage_id)) + raw_overrides = _stage_config_value(runtime_config, "duplex_stage_sampling_params", stage_id) + overrides = dict(raw_overrides) if isinstance(raw_overrides, dict) else {} + if not isinstance(default, SamplingParams) or (not overrides and (max_tokens is None or max_tokens <= 0)): + configured.append(default) + continue + params = default.clone() + if max_tokens is not None and max_tokens > 0: + params.max_tokens = max_tokens + for name, value in overrides.items(): + if not hasattr(params, name): + continue + setattr(params, name, value) + if name == "stop_token_ids": + all_stop_token_ids = getattr(params, "_all_stop_token_ids", None) + if isinstance(all_stop_token_ids, set): + all_stop_token_ids.update(int(token_id) for token_id in value) + configured.append(params) + return tuple(configured) + + def plan_append( + self, + *, + request_id: str, + fence: DuplexFence, + session_config: dict[str, object], + runtime_config: dict[str, object], + seq: int, + turn_seq: int, + payload: object, + final: bool, + sampling_params: object, + ) -> DuplexAppendPlan: + del sampling_params + return DuplexAppendPlan( + prompt=build_duplex_data_plane_prompt( + request_id=request_id, + fence=fence, + session_config=session_config, + runtime_config=runtime_config, + seq=seq, + turn_seq=turn_seq, + payload=payload, + final=final, + ) + ) + + def decide_output( + self, + *, + stage_id: int, + final_stage_id: int, + segment_finished: bool, + segment_token_ids: tuple[int, ...], + segment_output_metadata: dict[str, object], + output: object, + ) -> DuplexOutputDecision | None: + if stage_id >= final_stage_id or not segment_finished: + return None + + completion = _first_completion(output) + output_metadata = _multimodal_output(output, completion) + special_token_ids = _special_token_ids(segment_output_metadata) + special_token_ids.update(_special_token_ids(output_metadata)) + listen_id = special_token_ids.get("listen_token_id") + if listen_id is None: + return None + + stop_reason = getattr(completion, "stop_reason", None) if completion is not None else None + token_ids = _completion_token_ids(completion) or list(segment_token_ids) + if _coerce_int(stop_reason) != listen_id and (not token_ids or token_ids[-1] != listen_id): + return None + + metadata = dict(output_metadata) + for key, value in special_token_ids.items(): + metadata.setdefault(f"meta.{key}", value) + metadata.update( + { + "duplex_direct_response": True, + "duplex_native_decision": "listen", + "model_listen": True, + "listen_source": "model_listen", + } + ) + return DuplexOutputDecision( + action=DuplexOutputAction.DIRECT_RESPONSE, + metadata=metadata, + ) + + # ---- session policy ---- + + def create_session_state(self) -> MiniCPMO45ServingSessionState: + return MiniCPMO45ServingSessionState() + + def capabilities(self, *, max_sessions: int) -> DuplexCapabilities: + return minicpmo45_native_capabilities(max_sessions=max_sessions) + + def validate_client_extra_body(self, extra_body: object) -> None: + if not isinstance(extra_body, dict): + return + private_keys = sorted(PRIVATE_RUNTIME_CONFIG_KEYS.intersection(extra_body)) + if private_keys: + raise MiniCPMO45ClientRuntimeConfigError( + "duplex runtime configuration is server-owned: " + ", ".join(private_keys) + ) + + async def prepare_runtime_config( + self, config: DuplexSessionConfig, *, model_config: ModelConfig | None + ) -> dict[str, object]: + extra_body = dict(config.extra_body) + if any(key in extra_body for key in ("ref_audio_path", "tts_ref_audio_path")): + raise MiniCPMO45ClientRuntimeConfigError( + "ref_audio_path is not accepted by duplex sessions; use ref_audio URI instead", + code="unsupported_ref_audio_path", + ) + runtime_config: dict[str, object] = {"instructions": config.instructions} + # ``duplex_initial_user_text`` is the older extra_body spelling and + # still works; the session field is the framework-level one. + initial_user_text = extra_body.pop("duplex_initial_user_text", None) + if not (isinstance(initial_user_text, str) and initial_user_text): + initial_user_text = config.initial_user_text + if isinstance(initial_user_text, str) and initial_user_text: + runtime_config["initial_user_text"] = initial_user_text + tokenizer = await self._tokenizer_for(model_config) + _apply_default_scheduler_policy(runtime_config, config=config, tokenizer=tokenizer) + + ref_audio = config.ref_audio + extra_ref_audio = extra_body.get("ref_audio") + if ref_audio is None and isinstance(extra_ref_audio, str): + extra_body.pop("ref_audio") + ref_audio = extra_ref_audio + extra_tts_ref_audio = extra_body.get("tts_ref_audio") + if ref_audio is None and isinstance(extra_tts_ref_audio, str): + extra_body.pop("tts_ref_audio") + ref_audio = extra_tts_ref_audio + + if ref_audio is None: + if any(str(modality).lower() == "audio" for modality in config.modalities): + raise MiniCPMO45ClientRuntimeConfigError( + "MiniCPM-o duplex audio output requires ref_audio", + code="ref_audio_required", + ) + _apply_first_append_context_tokens( + runtime_config, + tokenizer=tokenizer, + instructions=config.instructions, + initial_user_text=initial_user_text, + ref_sample_count=None, + ) + config.extra_body = extra_body + return runtime_config + + wav_np, sr = await resolve_ref_audio(ref_audio, model_config=model_config) + # torchaudio resampling is CPU work: keep it off the orchestrator loop. + wav_np = await asyncio.to_thread(normalize_ref_audio, wav_np, int(sr), target_sr=16000) + # Trim to a whole number of pooled audio embeddings (100 ms frames) so + # the first-append scheduler reserve can count them exactly. + usable = (len(wav_np) // MiniCPMO45DuplexPolicy.SAMPLES_PER_AUDIO_TOKEN) * ( + MiniCPMO45DuplexPolicy.SAMPLES_PER_AUDIO_TOKEN + ) + wav_np = wav_np[:usable] + ref_audio_bytes = np.ascontiguousarray(wav_np, dtype=np.float32).tobytes() + runtime_config["ref_audio_data"] = base64.b64encode(ref_audio_bytes).decode("ascii") + runtime_config["ref_audio_format"] = "pcm_f32le" + runtime_config["ref_audio_sample_rate_hz"] = 16000 + _apply_first_append_context_tokens( + runtime_config, + tokenizer=tokenizer, + instructions=config.instructions, + initial_user_text=initial_user_text, + ref_sample_count=len(wav_np), + ) + config.extra_body = extra_body + config.ref_audio = None + return runtime_config + + def runtime_config_for_update( + self, + config: DuplexSessionConfig, + current: Mapping[str, object], + ) -> dict[str, object]: + runtime_config = deepcopy(dict(current)) + reject_changed_runtime_value( + config.instructions, + runtime_config.get("instructions"), + message="instructions cannot be changed after the session is created", + code="instructions_update_unsupported", + error_cls=MiniCPMO45ClientRuntimeConfigError, + ) + stage_max_tokens = runtime_config.get("duplex_stage_max_tokens") + stage_max_tokens = deepcopy(stage_max_tokens) if isinstance(stage_max_tokens, dict) else {} + stage_max_tokens["0"] = ( + config.max_tokens if isinstance(config.max_tokens, int) and config.max_tokens > 0 else 20 + ) + stage_max_tokens.setdefault("1", 8192) + runtime_config["duplex_stage_max_tokens"] = stage_max_tokens + + stage_sampling = runtime_config.get("duplex_stage_sampling_params") + stage_sampling = deepcopy(stage_sampling) if isinstance(stage_sampling, dict) else {} + stage0 = stage_sampling.get("0") + stage0 = deepcopy(stage0) if isinstance(stage0, dict) else {} + stage0["temperature"] = config.temperature if config.temperature is not None else 0.7 + stage_sampling["0"] = stage0 + runtime_config["duplex_stage_sampling_params"] = stage_sampling + return runtime_config + + def data_plane_context( + self, + *, + epoch: int, + turn_id: int, + active_response_turn_id: int | None, + active_response_id: str | None, + auto_responds: bool, + response_format: str, + speed: float | None, + modalities: tuple[str, ...], + ) -> MiniCPMO45DataPlaneContext: + return MiniCPMO45DataPlaneContext( + epoch=epoch, + turn_id=turn_id, + active_response_turn_id=active_response_turn_id, + active_response_id=active_response_id, + auto_responds=auto_responds, + response_format=response_format, + speed=speed, + modalities=modalities, + ) + + +__all__ = [ + "PRIVATE_RUNTIME_CONFIG_KEYS", + "MiniCPMO45ClientRuntimeConfigError", + "MiniCPMO45DuplexPlugin", + "build_duplex_data_plane_prompt", + "duplex_first_append_context_reserve", + "duplex_first_append_unit_count", + "duplex_payload_is_exact_chunks", + "duplex_scheduler_token_budget", + "normalize_ref_audio", + "resolve_ref_audio", +] diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py index 9a2643521..0332c81c0 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/policy.py @@ -3,7 +3,10 @@ from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from transformers import PreTrainedTokenizerBase class MiniCPMO45DuplexPolicy: @@ -89,7 +92,7 @@ class MiniCPMO45DuplexPolicy: } @classmethod - def token_ids_from_tokenizer(cls, tokenizer: Any) -> dict[str, int]: + def token_ids_from_tokenizer(cls, tokenizer: PreTrainedTokenizerBase) -> dict[str, int]: convert = getattr(tokenizer, "convert_tokens_to_ids", None) def token_id(token: str) -> int: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/runtime.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/runtime.py deleted file mode 100644 index d5e98d723..000000000 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/runtime.py +++ /dev/null @@ -1,439 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -import math -from base64 import b64decode -from binascii import Error as BinasciiError -from typing import Any, cast - -from vllm.sampling_params import SamplingParams - -from vllm_omni.engine.duplex.messages import DuplexFence -from vllm_omni.engine.duplex.runtime import ( - DuplexAppendPlan, - DuplexInputMode, - DuplexOutputAction, - DuplexOutputDecision, -) - -_DUPLEX_CHUNK_SAMPLES = 16000 -_DUPLEX_SAMPLES_PER_AUDIO_TOKEN = 1600 -# + 64 resampler embeddings + per frame (max_slice_nums=1), -# matching MiniCPMO45DuplexPolicy.VISION_TOKENS_PER_FRAME. -_DUPLEX_VISION_TOKENS_PER_FRAME = 66 -# Official stacked pair uses max_slice_nums=[2, 1]: the current frame is HD -# sliced (1 source + up to 2 patches, e.g. on 960x540) and the composite is -# not. The patch count depends on the frame size (see -# ``_duplex_hd_slice_count``); the constant is the fallback when the frame -# header cannot be read. -_DUPLEX_HD_SLICES_PER_BASE_FRAME = 3 -_DUPLEX_HD_MAX_SLICE_NUMS = 2 -_DUPLEX_SCALE_RESOLUTION = 448 - - -def _duplex_frames(payload: object) -> list[str]: - if not isinstance(payload, dict): - return [] - frames = payload.get("video_frames") - if not isinstance(frames, list): - return [] - return [frame for frame in frames if isinstance(frame, str) and frame] - - -def _duplex_frame_count(payload: object) -> int: - return len(_duplex_frames(payload)) - - -def _duplex_frame_size(frame_b64: str) -> tuple[int, int] | None: - """Pixel size of a base64 JPEG/PNG frame from its header, or ``None``.""" - from io import BytesIO - - try: - from PIL import Image - - with Image.open(BytesIO(b64decode(frame_b64, validate=True))) as image: - width, height = image.size - except Exception: # noqa: BLE001 - Stage0 rejects bad frames with a reason - return None - if width <= 0 or height <= 0: - return None - return int(width), int(height) - - -def _duplex_hd_slice_count(image_size: tuple[int, int], max_slice_nums: int) -> int: - """Number of HD patches ``MiniCPMVImageProcessor.get_sliced_grid`` adds. - - Port of the official grid selection (``scale_resolution=448``): an image - whose area fits in one 448x448 tile is not sliced at all, otherwise the - grid closest to the image aspect ratio among ``multiple-1 .. multiple+1`` - splits is used. Stage0 runs the same processor, so the scheduler budget - and the worker-built embeddings agree. - """ - width, height = image_size - ratio = width * height / (_DUPLEX_SCALE_RESOLUTION * _DUPLEX_SCALE_RESOLUTION) - multiple = min(math.ceil(ratio), max_slice_nums) - if multiple <= 1: - return 0 - log_ratio = math.log(width / height) - best_grid = (1, 1) - min_error = float("inf") - for split_grids_nums in (multiple - 1, multiple, multiple + 1): - if split_grids_nums == 1 or split_grids_nums > max_slice_nums: - continue - for m in range(1, split_grids_nums + 1): - if split_grids_nums % m: - continue - grid = (m, split_grids_nums // m) - error = abs(log_ratio - math.log(grid[0] / grid[1])) - if error < min_error: - best_grid = grid - min_error = error - return best_grid[0] * best_grid[1] - - -def _duplex_vision_tokens(payload: object) -> int: - """Scheduler slots for this append's camera track. - - Audio is never stacked: a unit still carries one second of soundtrack. - ``stack_frames`` only adds a second *image*. Official HD on that pair is - ``max_slice_nums=[2, 1]``: the base frame keeps its source block plus the - HD patches the processor cuts for its size (two for a 960x540 camera - frame, none for anything that fits in one 448x448 tile) and every extra - frame is one block. A single frame is never sliced. The count must match - the embeddings Stage0 builds exactly: surplus slots become pad - embeddings inside the KV. - """ - frames = _duplex_frames(payload) - count = len(frames) - if count <= 0: - return 0 - if count == 1: - return _DUPLEX_VISION_TOKENS_PER_FRAME - base_size = _duplex_frame_size(frames[0]) - if base_size is None: - base_blocks = _DUPLEX_HD_SLICES_PER_BASE_FRAME - else: - base_blocks = 1 + _duplex_hd_slice_count(base_size, _DUPLEX_HD_MAX_SLICE_NUMS) - return (base_blocks + (count - 1)) * _DUPLEX_VISION_TOKENS_PER_FRAME - - -def _duplex_pcm_sample_count(payload: object) -> int | None: - if not isinstance(payload, dict): - return None - audio = payload.get("audio") or payload.get("data") - if payload.get("format") != "pcm_f32le" or not isinstance(audio, str): - return None - try: - raw = b64decode(audio, validate=True) - except (BinasciiError, ValueError): - return None - return len(raw) // 4 - - -def duplex_payload_is_exact_chunks(payload: object) -> bool: - sample_count = _duplex_pcm_sample_count(payload) - return sample_count is not None and sample_count != 0 and sample_count % _DUPLEX_CHUNK_SAMPLES == 0 - - -def duplex_first_append_unit_count(payload: object) -> int | None: - sample_count = _duplex_pcm_sample_count(payload) - if not sample_count or sample_count % _DUPLEX_CHUNK_SAMPLES != 0: - return None - return max(1, sample_count // _DUPLEX_CHUNK_SAMPLES - 1) - - -def duplex_scheduler_token_budget(payload: object, *, default: int = 64) -> int: - vision_tokens = _duplex_vision_tokens(payload) - sample_count = _duplex_pcm_sample_count(payload) - if sample_count is None: - return max(1, int(default)) + vision_tokens - sample_count = max(1, sample_count) - if sample_count % _DUPLEX_CHUNK_SAMPLES == 0: - units = sample_count // _DUPLEX_CHUNK_SAMPLES - return units * (2 + _DUPLEX_CHUNK_SAMPLES // _DUPLEX_SAMPLES_PER_AUDIO_TOKEN) + vision_tokens - return max(16, min(768, sample_count // _DUPLEX_SAMPLES_PER_AUDIO_TOKEN + 8)) + vision_tokens - - -def duplex_first_append_context_reserve(runtime_config: object) -> int: - if not isinstance(runtime_config, dict): - return 48 - exact = runtime_config.get("duplex_first_append_context_tokens") - if isinstance(exact, int) and exact >= 0: - return exact - reserve = 48 - ref = runtime_config.get("ref_audio_data") - if isinstance(ref, str) and ref: - try: - raw = b64decode(ref, validate=True) - except (BinasciiError, ValueError): - raw = b"" - if raw: - reserve += max(0, (len(raw) // 4) // _DUPLEX_SAMPLES_PER_AUDIO_TOKEN + 8) - return reserve - - -def _duplex_force_listen_count(extra_body: object) -> int: - raw = extra_body.get("force_listen_count") if isinstance(extra_body, dict) else None - try: - return 0 if raw is None else max(0, int(raw)) - except (TypeError, ValueError): - return 0 - - -def build_duplex_data_plane_prompt( - *, - request_id: str, - fence: DuplexFence, - session_config: dict[str, Any], - runtime_config: dict[str, Any], - seq: int, - turn_seq: int, - mode: DuplexInputMode, - payload: object, - final: bool, -) -> dict[str, Any]: - token_budget = duplex_scheduler_token_budget(payload) - if seq <= 1: - context_reserve = duplex_first_append_context_reserve(runtime_config) - token_budget += context_reserve - first_units = duplex_first_append_unit_count(payload) - if first_units is not None: - token_budget = context_reserve + first_units * 12 - 1 + _duplex_vision_tokens(payload) - if seq > 1 and duplex_payload_is_exact_chunks(payload): - token_budget += 1 - if final and duplex_payload_is_exact_chunks(payload): - token_budget += 12 - extra_body = session_config.get("extra_body") - raw_token_id = runtime_config.get("duplex_scheduler_token_id") - try: - token_id = 0 if raw_token_id is None else max(0, int(raw_token_id)) - except (TypeError, ValueError): - token_id = 0 - force_listen_count = _duplex_force_listen_count(extra_body) - if ( - force_listen_count > 0 - and turn_seq <= force_listen_count - and isinstance(payload, dict) - and payload.get("force_listen") is not True - ): - payload = {**payload, "force_listen": True} - return { - "prompt_token_ids": [token_id] * token_budget, - "model_intermediate_buffer": { - "request_id": request_id, - "global_request_id": [fence.session_id], - "duplex": { - "fence": fence, - "session_id": fence.session_id, - "incarnation": fence.incarnation, - "epoch": fence.epoch, - "seq": seq, - "turn_id": fence.turn_id, - "response_seq": fence.response_seq, - "turn_seq": turn_seq, - "mode": mode.value, - "payload": payload, - "final": final, - "data_plane": True, - "session_config": dict(session_config), - "runtime_config": dict(runtime_config), - "scheduler_token_budget": token_budget, - "scheduler_token_id": token_id, - }, - }, - } - - -def _coerce_int(value: object) -> int | None: - detach = getattr(value, "detach", None) - if callable(detach): - try: - flat: Any = detach().cpu().reshape(-1) - if flat.numel() == 0: - return None - value = flat[0].item() - except Exception: - return None - try: - return int(cast(Any, value)) - except (TypeError, ValueError): - return None - - -def _coerce_int_list(value: object) -> list[int]: - if value is None: - return [] - if hasattr(value, "detach"): - try: - value = value.detach().cpu().reshape(-1).tolist() - except Exception: - return [] - if not isinstance(value, (list, tuple)): - return [] - return [token_id for item in value if (token_id := _coerce_int(item)) is not None] - - -def _first_completion(output: object) -> object | None: - outputs = getattr(output, "outputs", None) - return outputs[0] if isinstance(outputs, list) and outputs else None - - -def _multimodal_output(output: object, completion: object | None) -> dict[str, Any]: - metadata = getattr(output, "multimodal_output", None) - if isinstance(metadata, dict): - return metadata - metadata = getattr(completion, "multimodal_output", None) if completion is not None else None - return metadata if isinstance(metadata, dict) else {} - - -def _special_token_ids(metadata: dict[str, Any]) -> dict[str, int]: - sources: list[object] = [metadata.get("special_token_ids"), metadata.get("meta")] - sources.append( - { - key.removeprefix("meta."): value - for key, value in metadata.items() - if isinstance(key, str) and key.startswith("meta.") - } - ) - token_ids: dict[str, int] = {} - for source in sources: - if not isinstance(source, dict): - continue - for key, value in source.items(): - token_id = _coerce_int(value) - if isinstance(key, str) and token_id is not None and token_id >= 0: - token_ids[key] = token_id - return token_ids - - -def _completion_token_ids(completion: object | None) -> list[int]: - if completion is None: - return [] - for attribute in ("token_ids", "cumulative_token_ids"): - token_ids = _coerce_int_list(getattr(completion, attribute, None)) - if token_ids: - return token_ids - return [] - - -def _stage_config_value(runtime_config: dict[str, Any], key: str, stage_id: int) -> object | None: - raw = runtime_config.get(key) - if isinstance(raw, dict): - value = raw.get(stage_id) - return raw.get(str(stage_id)) if value is None else value - if isinstance(raw, (list, tuple)) and stage_id < len(raw): - return raw[stage_id] - return None - - -class MiniCPMO45DuplexRuntimeExtension: - def configure_sampling_params( - self, - *, - runtime_config: dict[str, Any], - defaults: tuple[object, ...], - ) -> tuple[object, ...]: - configured: list[object] = [] - for stage_id, default in enumerate(defaults): - max_tokens = _coerce_int(_stage_config_value(runtime_config, "duplex_stage_max_tokens", stage_id)) - raw_overrides = _stage_config_value(runtime_config, "duplex_stage_sampling_params", stage_id) - overrides = dict(raw_overrides) if isinstance(raw_overrides, dict) else {} - if not isinstance(default, SamplingParams) or (not overrides and (max_tokens is None or max_tokens <= 0)): - configured.append(default) - continue - params = default.clone() - if max_tokens is not None and max_tokens > 0: - params.max_tokens = max_tokens - for name, value in overrides.items(): - if not hasattr(params, name): - continue - setattr(params, name, value) - if name == "stop_token_ids": - all_stop_token_ids = getattr(params, "_all_stop_token_ids", None) - if isinstance(all_stop_token_ids, set): - all_stop_token_ids.update(int(token_id) for token_id in value) - configured.append(params) - return tuple(configured) - - def plan_append( - self, - *, - request_id: str, - fence: DuplexFence, - session_config: dict[str, Any], - runtime_config: dict[str, Any], - seq: int, - turn_seq: int, - mode: DuplexInputMode, - payload: object, - final: bool, - sampling_params: object, - ) -> DuplexAppendPlan: - del sampling_params - return DuplexAppendPlan( - prompt=build_duplex_data_plane_prompt( - request_id=request_id, - fence=fence, - session_config=session_config, - runtime_config=runtime_config, - seq=seq, - turn_seq=turn_seq, - mode=mode, - payload=payload, - final=final, - ) - ) - - def decide_output( - self, - *, - stage_id: int, - final_stage_id: int, - segment_finished: bool, - segment_token_ids: tuple[int, ...], - segment_output_metadata: dict[str, Any], - output: object, - ) -> DuplexOutputDecision | None: - if stage_id >= final_stage_id or not segment_finished: - return None - - completion = _first_completion(output) - output_metadata = _multimodal_output(output, completion) - special_token_ids = _special_token_ids(segment_output_metadata) - special_token_ids.update(_special_token_ids(output_metadata)) - listen_id = special_token_ids.get("listen_token_id") - if listen_id is None: - return None - - stop_reason = getattr(completion, "stop_reason", None) if completion is not None else None - token_ids = _completion_token_ids(completion) or list(segment_token_ids) - if _coerce_int(stop_reason) != listen_id and (not token_ids or token_ids[-1] != listen_id): - return None - - metadata = dict(output_metadata) - for key, value in special_token_ids.items(): - metadata.setdefault(f"meta.{key}", value) - metadata.update( - { - "duplex_direct_response": True, - "duplex_native_decision": "listen", - "model_listen": True, - "listen_source": "model_listen", - } - ) - return DuplexOutputDecision( - action=DuplexOutputAction.DIRECT_RESPONSE, - metadata=metadata, - ) - - -__all__ = [ - "MiniCPMO45DuplexRuntimeExtension", - "build_duplex_data_plane_prompt", - "duplex_first_append_context_reserve", - "duplex_first_append_unit_count", - "duplex_payload_is_exact_chunks", - "duplex_scheduler_token_budget", -] diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/serving_adapter.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/serving_adapter.py deleted file mode 100644 index ea51662a9..000000000 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/serving_adapter.py +++ /dev/null @@ -1,102 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project - -from __future__ import annotations - -from collections.abc import Callable, Mapping -from typing import Any - -from vllm_omni.entrypoints.duplex.protocol import DuplexCapabilities -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.adapter import ( - MiniCPMO45NativeDuplexServingAdapter, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.capabilities import ( - minicpmo45_native_capabilities, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.data_plane import ( - MiniCPMO45DataPlaneContext, - MiniCPMO45DataPlaneSession, -) -from vllm_omni.model_executor.models.minicpmo_4_5.duplex.session import ( - MiniCPMO45ServingSessionState, -) - -EncodeAudio = Callable[[object, int, str, float | None], str | None] - - -class MiniCPMO45ServingRuntimeAdapter: - """MiniCPM-owned serving state, input packing, and output projection.""" - - adapter_id = "minicpmo45" - clean_response_done_prefix = "" - interrupted_tts_prefix = "" - private_runtime_config_keys = MiniCPMO45NativeDuplexServingAdapter.PRIVATE_RUNTIME_CONFIG_KEYS - - def __init__(self, encode_audio: EncodeAudio) -> None: - self.session_states: dict[str, MiniCPMO45ServingSessionState] = {} - self.data_plane = MiniCPMO45DataPlaneSession(encode_audio) - - def create_session_state(self) -> MiniCPMO45ServingSessionState: - return MiniCPMO45ServingSessionState() - - def session_state(self, session_id: str) -> MiniCPMO45ServingSessionState: - state = self.session_states.get(session_id) - if state is None: - state = self.create_session_state() - self.session_states[session_id] = state - return state - - def remove_session_state(self, session_id: str) -> None: - self.session_states.pop(session_id, None) - - @staticmethod - def is_enabled(config: object) -> bool: - return MiniCPMO45NativeDuplexServingAdapter.is_enabled(config) # type: ignore[arg-type] - - @staticmethod - def capabilities(*, max_sessions: int) -> DuplexCapabilities: - return minicpmo45_native_capabilities(max_sessions=max_sessions) - - @staticmethod - def validate_client_extra_body(extra_body: object) -> None: - MiniCPMO45NativeDuplexServingAdapter.validate_client_extra_body(extra_body) - - @staticmethod - async def prepare_runtime_config(config: object, *, model_config: Any) -> dict[str, object]: - return await MiniCPMO45NativeDuplexServingAdapter.prepare_runtime_config( - config, # type: ignore[arg-type] - model_config=model_config, - ) - - @staticmethod - def runtime_config_for_update( - config: object, - current: Mapping[str, object], - ) -> dict[str, object]: - return MiniCPMO45NativeDuplexServingAdapter.runtime_config_for_update( - config, # type: ignore[arg-type] - dict(current), - ) - - @staticmethod - def data_plane_context( - *, - epoch: int, - turn_id: int, - active_response_turn_id: int | None, - active_response_id: str | None, - auto_responds: bool, - response_format: str, - speed: float | None, - modalities: tuple[str, ...], - ) -> MiniCPMO45DataPlaneContext: - return MiniCPMO45DataPlaneContext( - epoch=epoch, - turn_id=turn_id, - active_response_turn_id=active_response_turn_id, - active_response_id=active_response_id, - auto_responds=auto_responds, - response_format=response_format, - speed=speed, - modalities=modalities, - ) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/session.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/session.py index 40c42cb81..49d0caab6 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/session.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/session.py @@ -4,34 +4,31 @@ from __future__ import annotations import asyncio -from collections.abc import Awaitable, Callable from dataclasses import dataclass, field +from vllm_omni.engine.duplex.plugin import DuplexModelSessionState from vllm_omni.model_executor.models.minicpmo_4_5.duplex.input import ( MiniCPMO45PcmAppendBuffer, ) @dataclass(slots=True) -class MiniCPMO45ServingSessionState: - """Mutable serving state owned by one MiniCPM duplex session.""" +class MiniCPMO45ServingSessionState(DuplexModelSessionState): + """Mutable model-owned state of one MiniCPM duplex session (owned by the session runner).""" audio_buffer: MiniCPMO45PcmAppendBuffer = field(default_factory=MiniCPMO45PcmAppendBuffer) input_since_commit: bool = False speech_since_commit: bool = False - native_context_locked: bool = False + context_locked: bool = False committed_audio_payload: dict[str, object] | None = None committed_audio_operation_id: str | None = None committed_audio_reserved_bytes: int = 0 deferred_response_create: bool = False deferred_precreate_response: bool = False - data_plane_task: asyncio.Task[None] | None = None - data_plane_restart_requested: bool = False continuation_owner_id: str | None = None continuation_units: int = 0 pending_silence_task: asyncio.Task[bool] | None = None pending_silence_owner_id: str | None = None - silence_continuation_scheduler: Callable[..., Awaitable[bool]] | None = None def retain_committed_audio( self, diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py index 6d9800025..bc13b58f2 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/stage0.py @@ -6,15 +6,22 @@ from __future__ import annotations import base64 import copy import time +from collections.abc import Callable from contextlib import suppress from dataclasses import dataclass, field from threading import Lock -from typing import Any +from typing import TYPE_CHECKING, Any, cast import numpy as np +from numpy.typing import NDArray from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import MiniCPMO45DuplexPolicy +if TYPE_CHECKING: + import torch + from PIL import Image + from torch import nn + _MINICPMO45_SPECIAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS _MINICPMO45_OPTIONAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.OPTIONAL_TOKEN_FIELDS _MINICPMO45_PROCESSOR_LOAD_LOCK = Lock() @@ -23,17 +30,19 @@ _MINICPMO45_PROCESSOR_LOAD_LOCK = Lock() @dataclass class _MiniCPMO45Stage0SessionState: session_id: str + #: Any: MiniCPM-o remote-code processor (transformers, no importable type). streaming_processor: Any | None = None - audio_buffer: np.ndarray = field(default_factory=lambda: np.array([], dtype=np.float32)) + audio_buffer: NDArray[np.float32] = field(default_factory=lambda: np.array([], dtype=np.float32)) audio_chunk_idx: int = 0 - context_embeds: list[Any] = field(default_factory=list) + context_embeds: list[torch.Tensor] = field(default_factory=list) context_token_ids: list[int] = field(default_factory=list) current_turn_ended: bool = True prepared_append_identity: tuple[int | None, int] | None = None - prepared_inputs_embeds: Any | None = None + prepared_inputs_embeds: torch.Tensor | None = None prepared_input_token_ids: list[int] = field(default_factory=list) - prepared_result: dict[str, Any] = field(default_factory=dict) - audio_past_key_values: Any | None = None + prepared_result: dict[str, object] = field(default_factory=dict) + #: Opaque transformers cache, handed back to the model unchanged. + audio_past_key_values: object | None = None pending_terminator_token: int | None = None last_terminator_token: int | None = None pending_speech_context: bool = False @@ -59,7 +68,7 @@ class MiniCPMO45Stage0DuplexRuntime: image_start_token_id: int = -1 image_end_token_id: int = -1 - def __init__(self, stage_model: Any, *, model_path: str | None = None, device: str = "cuda") -> None: + def __init__(self, stage_model: nn.Module, *, model_path: str | None = None, device: str = "cuda") -> None: self.stage_model = stage_model self.model_path = model_path self.device = device @@ -85,7 +94,7 @@ class MiniCPMO45Stage0DuplexRuntime: def _configure_streaming_processor( self, state: _MiniCPMO45Stage0SessionState | None = None, - ) -> Any | None: + ) -> Any | None: # Any: MiniCPM-o remote-code processor processor = self.processor if processor is None: return None @@ -129,9 +138,9 @@ class MiniCPMO45Stage0DuplexRuntime: def _prepare_session_context( self, state: _MiniCPMO45Stage0SessionState, - session_config: dict[str, Any], + session_config: dict[str, object], *, - runtime_config: dict[str, Any] | None = None, + runtime_config: dict[str, object] | None = None, ) -> None: if not self._stage_runtime_ready(): return @@ -163,14 +172,14 @@ class MiniCPMO45Stage0DuplexRuntime: def _stage_prefill_embeddings_only( self, state: _MiniCPMO45Stage0SessionState, - audio_waveform: Any, + audio_waveform: NDArray[np.float32] | None, *, - video_frames: list[Any] | None = None, + video_frames: list[Image.Image] | None = None, epoch: int | None = None, seq: int | None = None, is_speech: bool = False, final: bool = False, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Build scheduler-owned Stage0 input embeddings for one audio append. Unlike the legacy worker-control path, this method never calls an eager @@ -217,7 +226,7 @@ class MiniCPMO45Stage0DuplexRuntime: # Omni duplex: encode this append's camera frames so the first unit can # carry them, mirroring official streaming_prefill (feed , then # every frame_list entry as its own block, then the audio). - frame_blocks: list[Any] = [] + frame_blocks: list[list[torch.Tensor]] = [] if video_frames: self._require_vision_token_ids() encoded_frames = self._stage_vision_embeddings(video_frames) @@ -227,7 +236,7 @@ class MiniCPMO45Stage0DuplexRuntime: return self._stage_prefill_result(False, start_time, "streaming vision embedding failed") frame_blocks = encoded_frames - embed_parts: list[Any] = [] + embed_parts: list[torch.Tensor] = [] token_ids: list[int] = [] if state.audio_chunk_idx == 0 and state.context_embeds: embed_parts.extend(state.context_embeds) @@ -336,7 +345,7 @@ class MiniCPMO45Stage0DuplexRuntime: return result @staticmethod - def _stage_prefill_result(success: bool, start_time: float, reason: str = "") -> dict[str, Any]: + def _stage_prefill_result(success: bool, start_time: float, reason: str = "") -> dict[str, object]: return { "success": success, "prefill_success": success, @@ -347,14 +356,14 @@ class MiniCPMO45Stage0DuplexRuntime: } @staticmethod - def _as_2d_tensor(value: Any) -> Any: + def _as_2d_tensor(value: torch.Tensor) -> torch.Tensor: if value.ndim == 1: return value.unsqueeze(0) if value.ndim == 3 and value.shape[0] == 1: return value.squeeze(0) return value - def _embed_token(self, token_id: int) -> Any: + def _embed_token(self, token_id: int) -> torch.Tensor: import torch token = torch.tensor([int(token_id)], dtype=torch.long, device=self._model_device()) @@ -362,7 +371,7 @@ class MiniCPMO45Stage0DuplexRuntime: embeds = embedder(token) return self._as_2d_tensor(embeds) - def _token_embedding_dtype(self) -> Any: + def _token_embedding_dtype(self) -> torch.dtype | None: """dtype of the decoder token embeddings, the unit's reference dtype. Official ``get_vllm_embedding`` casts vision hidden states to the token @@ -376,7 +385,7 @@ class MiniCPMO45Stage0DuplexRuntime: return dtype return getattr(next(self.thinker.parameters()), "dtype", None) - def _token_embedder(self) -> Any: + def _token_embedder(self) -> Callable[[torch.Tensor], torch.Tensor]: nested_embed = getattr(getattr(getattr(self.thinker, "llm", None), "model", None), "embed_tokens", None) if callable(nested_embed): return nested_embed @@ -391,7 +400,7 @@ class MiniCPMO45Stage0DuplexRuntime: return embedder raise AttributeError("MiniCPM-o stage0 model does not expose token embeddings") - def _model_device(self) -> Any: + def _model_device(self) -> torch.device | str: try: return next(self.thinker.parameters()).device except Exception: @@ -402,14 +411,14 @@ class MiniCPMO45Stage0DuplexRuntime: pass return self.device - def _streaming_chunk_size(self, processor: Any | None = None) -> int: + def _streaming_chunk_size(self, processor: Any | None = None) -> int: # Any: remote-code processor processor = processor or self.processor get_chunk = getattr(processor, "get_streaming_chunk_size", None) if callable(get_chunk): return int(get_chunk()) return 16000 - def _sample_rate(self, processor: Any | None = None) -> int: + def _sample_rate(self, processor: Any | None = None) -> int: # Any: remote-code processor processor = processor or self.processor return int( self._stage_param( @@ -418,7 +427,7 @@ class MiniCPMO45Stage0DuplexRuntime: ) ) - def _first_chunk_samples(self, default_chunk_size: int, processor: Any | None = None) -> int: + def _first_chunk_samples(self, default_chunk_size: int, processor: Any | None = None) -> int: # Any: processor processor = processor or self.processor if getattr(processor, "_streaming_mel_processor", None) is None: return default_chunk_size @@ -427,7 +436,7 @@ class MiniCPMO45Stage0DuplexRuntime: def _pad_first_audio_chunk_if_needed( self, state: _MiniCPMO45Stage0SessionState, - processor: Any | None = None, + processor: Any | None = None, # Any: remote-code processor ) -> None: if state.audio_chunk_idx != 0 or len(state.audio_buffer) == 0: return @@ -440,7 +449,7 @@ class MiniCPMO45Stage0DuplexRuntime: padding: np.ndarray = np.zeros(first_chunk_samples - len(state.audio_buffer), dtype=np.float32) state.audio_buffer = np.concatenate([padding, state.audio_buffer]) - def _stage_param(self, name: str, default: Any) -> Any: + def _stage_param(self, name: str, default: object) -> Any: # Any: attribute read off the remote-code model for target in (self.stage_model, self.thinker, getattr(self.thinker, "llm", None)): value = getattr(target, name, None) if value is not None: @@ -455,7 +464,7 @@ class MiniCPMO45Stage0DuplexRuntime: chunk_idx: int, default_chunk_size: int, *, - processor: Any | None = None, + processor: Any | None = None, # Any: remote-code processor ) -> int: processor = processor or self.processor if chunk_idx != 0: @@ -472,11 +481,11 @@ class MiniCPMO45Stage0DuplexRuntime: def _process_streaming_audio( self, - audio_chunk: Any, + audio_chunk: NDArray[np.float32], chunk_idx: int, *, - processor: Any | None = None, - ) -> Any: + processor: Any | None = None, # Any: remote-code processor + ) -> Any: # Any: transformers BatchFeature (or the dict fallback below) processor = processor or self.processor process = getattr(processor, "process_audio_streaming", None) if callable(process): @@ -488,10 +497,10 @@ class MiniCPMO45Stage0DuplexRuntime: def _stage_audio_embeddings( self, - batch_feature: Any, + batch_feature: Any, # Any: transformers BatchFeature *, state: _MiniCPMO45Stage0SessionState | None = None, - ) -> Any | None: + ) -> torch.Tensor | None: if hasattr(batch_feature, "to"): batch_feature = batch_feature.to(self.device) self._ensure_dynamic_cache_compat() @@ -534,17 +543,17 @@ class MiniCPMO45Stage0DuplexRuntime: self.thinker.audio_past_key_values = previous_audio_past_key_values @staticmethod - def _decode_ref_audio_from_session_config(session_config: dict[str, Any]) -> Any | None: + def _decode_ref_audio_from_session_config(session_config: dict[str, object]) -> NDArray[np.float32] | None: from vllm_omni.model_executor.models.minicpmo_4_5.duplex.input import decode_native_ref_audio_from_config return decode_native_ref_audio_from_config({"extra_body": session_config}) def _stage_ref_audio_embeddings( self, - ref_audio: Any, + ref_audio: NDArray[np.float32], *, state: _MiniCPMO45Stage0SessionState | None = None, - ) -> Any | None: + ) -> torch.Tensor | None: process_audio = getattr(self.processor, "process_audio", None) if callable(process_audio): batch_feature = process_audio([ref_audio]) @@ -606,16 +615,16 @@ class MiniCPMO45Stage0DuplexRuntime: DynamicCache.get_usable_length = get_usable_length # type: ignore[attr-defined] @staticmethod - def _cat_nested_tensors(value: Any) -> Any | None: + def _cat_nested_tensors(value: object) -> torch.Tensor | None: import torch - tensors = [] + tensors: list[torch.Tensor] = [] - def collect(item: Any) -> None: + def collect(item: object) -> None: if item is None: return if hasattr(item, "detach"): - tensors.append(item) + tensors.append(cast("torch.Tensor", item)) return if isinstance(item, dict): for child in item.values(): @@ -655,7 +664,7 @@ class MiniCPMO45Stage0DuplexRuntime: } @staticmethod - def _load_processor_from_path(model_path: str | None) -> Any | None: + def _load_processor_from_path(model_path: str | None) -> Any | None: # Any: remote-code AutoProcessor if not model_path: return None try: @@ -752,7 +761,7 @@ class MiniCPMO45Stage0DuplexRuntime: return self.stage_padding_token_id() @staticmethod - def _decode_audio_payload(payload: dict[str, Any]) -> Any: + def _decode_audio_payload(payload: dict[str, object]) -> NDArray[np.float32]: audio = payload.get("audio") or payload.get("data") if not isinstance(audio, str): raise ValueError("audio append payload requires base64 audio") @@ -762,7 +771,7 @@ class MiniCPMO45Stage0DuplexRuntime: return np.frombuffer(base64.b64decode(audio), dtype=np.float32) @staticmethod - def _decode_video_frames_payload(payload: dict[str, Any]) -> list[Any]: + def _decode_video_frames_payload(payload: dict[str, object]) -> list[Image.Image]: """Decode omni-duplex camera frames (base64 JPEG/PNG) to PIL images.""" frames = payload.get("video_frames") if not isinstance(frames, list) or not frames: @@ -771,7 +780,7 @@ class MiniCPMO45Stage0DuplexRuntime: from PIL import Image - decoded: list[Any] = [] + decoded: list[Image.Image] = [] for frame_b64 in frames: if not isinstance(frame_b64, str) or not frame_b64: continue @@ -826,7 +835,7 @@ class MiniCPMO45Stage0DuplexRuntime: return [1] return [2] + [1] * (frame_count - 1) - def _encode_processed_vision(self, processed: Any) -> list[Any] | None: + def _encode_processed_vision(self, processed: Any) -> list[torch.Tensor] | None: # Any: BatchFeature targets = (self.stage_model, self.thinker, getattr(self.stage_model, "model", None)) for target in targets: if target is None: @@ -842,8 +851,8 @@ class MiniCPMO45Stage0DuplexRuntime: device, dtype = vpm_param.device, vpm_param.dtype pixel_nested = processed["pixel_values"] tgt_nested = processed["tgt_sizes"] - flat_pixels: list[Any] = [] - flat_tgt: list[Any] = [] + flat_pixels: list[torch.Tensor] = [] + flat_tgt: list[torch.Tensor] = [] for image_slices, image_tgt in zip(pixel_nested, tgt_nested): for slice_pixels in image_slices: flat_pixels.append(slice_pixels.to(device=device, dtype=dtype)) @@ -858,7 +867,7 @@ class MiniCPMO45Stage0DuplexRuntime: return None expected = MiniCPMO45DuplexPolicy.VISION_EMBEDS_PER_FRAME embed_dtype = self._token_embedding_dtype() - out: list[Any] = [] + out: list[torch.Tensor] = [] for block in hidden: block_2d = self._as_2d_tensor(block) if int(block_2d.shape[0]) != expected: @@ -869,7 +878,7 @@ class MiniCPMO45Stage0DuplexRuntime: return out return None - def _stage_vision_embeddings(self, frames: list[Any]) -> list[list[Any]] | None: + def _stage_vision_embeddings(self, frames: list[Image.Image]) -> list[list[torch.Tensor]] | None: """Encode camera frames for omni duplex via the loaded vision tower. Official ``streaming_prefill`` encodes each ``frame_list`` entry as its @@ -883,7 +892,7 @@ class MiniCPMO45Stage0DuplexRuntime: process_image = getattr(self.processor, "process_image", None) if not callable(process_image): return None - out: list[list[Any]] = [] + out: list[list[torch.Tensor]] = [] for frame, max_slices in zip(frames, self._official_max_slice_nums(len(frames))): try: processed = process_image([frame], max_slice_nums=max_slices) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index 60252e78c..b1c8db89c 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -157,15 +157,13 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp del sampling_metadata self._minicpmo45_active_duplex_rows = [row.row_idx for row in rows] self._minicpmo45_duplex_row_sessions = { - row.row_idx: (row.session_id, row.incarnation) for row in rows if row.session_id is not None + row.row_idx: row.session_id for row in rows if row.session_id is not None } request_sessions = getattr(self, "_minicpmo45_duplex_request_sessions", None) if not isinstance(request_sessions, dict): request_sessions = {} self._minicpmo45_duplex_request_sessions = request_sessions - request_sessions.update( - {row.request_id: (row.session_id, row.incarnation) for row in rows if row.session_id is not None} - ) + request_sessions.update({row.request_id: row.session_id for row in rows if row.session_id is not None}) self._minicpmo45_duplex_row_payloads = {row.row_idx: row.payload for row in rows if row.payload is not None} self._minicpmo45_duplex_row_max_tokens = { row.row_idx: row.max_tokens for row in rows if row.max_tokens is not None @@ -199,7 +197,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp force_listen = payload.get("force_listen") is True is_speech = payload.get("is_speech") segment_key = (row.request_id, row.seq if row.seq is not None else -1) - session_key = (row.session_id, row.incarnation) if row.session_id is not None else None + session_key = row.session_id if turn_eos_id >= 0 and session_key is not None: state = helper_sessions.get(session_key) if isinstance(helper_sessions, dict) else None pending_speech_context = ( @@ -326,16 +324,12 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp helper = self._duplex_data_plane_helper() session_id = str(duplex.get("session_id") or "") - try: - incarnation = int(duplex.get("incarnation", 0)) - except (TypeError, ValueError): - incarnation = 0 payload = duplex.get("payload") if not session_id or not isinstance(payload, dict): embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids) return input_ids, embeds, {"duplex": {"prefill_success": False, "reason": "bad_duplex_payload"}} - session_key = (session_id, incarnation) + session_key = session_id state = helper.sessions.get(session_key) if state is None: from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import ( diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index e48328d0b..e75393845 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -24,13 +24,7 @@ MINICPMO_4_5_PIPELINE = PipelineConfig( model_type="minicpmo_4_5", default_deploy_config_name="minicpmo_4_5.yaml", model_arch="MiniCPMO45OmniForConditionalGeneration", - duplex_runtime_extension=( - "vllm_omni.model_executor.models.minicpmo_4_5.duplex.runtime.MiniCPMO45DuplexRuntimeExtension" - ), - duplex_serving_adapter=( - "vllm_omni.model_executor.models.minicpmo_4_5.duplex.serving_adapter.MiniCPMO45ServingRuntimeAdapter" - ), - duplex_control_enabled=True, + duplex_plugin="vllm_omni.model_executor.models.minicpmo_4_5.duplex.plugin.MiniCPMO45DuplexPlugin", # MiniCPM-o 4.5's HF config.json reports `model_type="minicpmo"` and # `architectures=["MiniCPMO"]` — both shared verbatim with older MiniCPM-o # 1.0 / 2.6 checkpoints. The only field distinguishing the generations is diff --git a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py index 80995f92e..1c374be85 100644 --- a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py @@ -690,9 +690,6 @@ def _native_duplex_data_plane_metadata(streaming_context) -> dict[str, object] | session_id = duplex_state.get("session_id") if isinstance(session_id, str) and session_id: metadata["session_id"] = session_id - incarnation = duplex_state.get("incarnation") - if isinstance(incarnation, int): - metadata["incarnation"] = incarnation epoch = duplex_state.get("epoch") if isinstance(epoch, int): metadata["epoch"] = epoch diff --git a/vllm_omni/utils/audio_resample.py b/vllm_omni/utils/audio_resample.py new file mode 100644 index 000000000..fc6230f6e --- /dev/null +++ b/vllm_omni/utils/audio_resample.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Streaming resampling, shared by the engine and the entrypoints. + +Lives here rather than beside the OpenAI audio mixin because the duplex +engine needs it too, and engine code must not import the entrypoints layer. +Pure numpy on purpose: importing this must not pull in torch or the API +server. +""" + +from __future__ import annotations + +import math + +import numpy as np + +__all__ = ["StreamingAudioResampler"] + + +class StreamingAudioResampler: + """Stateful polyphase resampler for streaming mono float audio. + + Retains filter state so output is invariant to input chunk boundaries. + """ + + _half_filter_width = 10 + _kaiser_beta = 5.0 + max_polyphase_factor = 2_048 + + def __init__(self, source_rate: int, target_rate: int): + if source_rate <= 0 or target_rate <= 0: + raise ValueError("Audio sample rates must be positive") + self.source_rate = source_rate + self.target_rate = target_rate + rate_gcd = math.gcd(source_rate, target_rate) + self._up = target_rate // rate_gcd + self._down = source_rate // rate_gcd + if max(self._up, self._down) > self.max_polyphase_factor: + raise ValueError("source and target sample rates have an unsupported resampling ratio") + self._half_len, self._phase_kernels = self._design_polyphase_filter() + self._history_samples = self._phase_kernels.shape[1] - 1 + self.reset() + + def _design_polyphase_filter(self) -> tuple[int, np.ndarray]: + if self._up == self._down: + return 0, np.ones((1, 1), dtype=np.float32) + + max_rate = max(self._up, self._down) + half_len = self._half_filter_width * max_rate + offsets = np.arange(-half_len, half_len + 1, dtype=np.float64) + cutoff = 1.0 / max_rate + taps = cutoff * np.sinc(cutoff * offsets) + taps *= np.kaiser(taps.size, self._kaiser_beta) + taps /= np.sum(taps) + taps *= self._up + taps = np.ascontiguousarray(taps, dtype=np.float32) + + phases = [taps[phase :: self._up] for phase in range(self._up)] + phase_width = max(phase.size for phase in phases) + kernels = np.zeros((self._up, phase_width), dtype=np.float32) + for phase_index, phase in enumerate(phases): + kernels[phase_index, -phase.size :] = phase[::-1] + return half_len, kernels + + @property + def scratch_bytes(self) -> int: + pending_samples = self._ceil_div(self._input_samples * self._up, self._down) - self._output_samples + return max(0, pending_samples) * np.dtype(np.float32).itemsize + + @staticmethod + def _ceil_div(numerator: int, denominator: int) -> int: + return -(-numerator // denominator) + + def _render_outputs( + self, + combined: np.ndarray, + *, + first_output: int, + output_count: int, + input_start: int, + ) -> np.ndarray: + if output_count <= 0: + return np.empty(0, dtype=np.float32) + + output_indexes = np.arange(first_output, first_output + output_count, dtype=np.int64) + filter_indexes = output_indexes * self._down + self._half_len + source_indexes = filter_indexes // self._up + phases = filter_indexes % self._up + windows = np.lib.stride_tricks.sliding_window_view( + combined, + self._phase_kernels.shape[1], + ) + selected_windows = windows[source_indexes - input_start] + return np.sum( + selected_windows * self._phase_kernels[phases], + axis=1, + dtype=np.float32, + ) + + def _push(self, chunk: np.ndarray) -> np.ndarray: + if chunk.size == 0: + return np.empty(0, dtype=np.float32) + if self._flushed: + raise RuntimeError("cannot process audio after the resampler has been finalized; call reset first") + + combined = np.concatenate((self._history, chunk)) + next_input_samples = self._input_samples + chunk.size + stable_output_samples = max( + 0, + self._ceil_div(next_input_samples * self._up - self._half_len, self._down), + ) + output = self._render_outputs( + combined, + first_output=self._output_samples, + output_count=stable_output_samples - self._output_samples, + input_start=self._input_samples, + ) + if self._history_samples: + self._history = combined[-self._history_samples :].copy() + self._input_samples = next_input_samples + self._output_samples = stable_output_samples + return output + + def _flush(self) -> np.ndarray: + if self._flushed: + return np.empty(0, dtype=np.float32) + + total_output_samples = self._ceil_div(self._input_samples * self._up, self._down) + output_count = total_output_samples - self._output_samples + if output_count: + last_output = total_output_samples - 1 + last_source_index = (last_output * self._down + self._half_len) // self._up + right_padding = max(0, last_source_index - self._input_samples + 1) + combined = np.concatenate((self._history, np.zeros(right_padding, dtype=np.float32))) + output = self._render_outputs( + combined, + first_output=self._output_samples, + output_count=output_count, + input_start=self._input_samples, + ) + else: + output = np.empty(0, dtype=np.float32) + self._output_samples = total_output_samples + self._flushed = True + return output + + def process(self, audio: np.ndarray, *, final: bool = False) -> np.ndarray: + chunk = np.asarray(audio, dtype=np.float32) + if chunk.ndim != 1: + raise ValueError(f"Streaming audio resampling only supports mono audio, got shape {chunk.shape}") + chunk = np.ascontiguousarray(chunk, dtype=np.float32) + output = self._push(chunk) + if not final: + return output + tail = self._flush() + if not output.size: + return tail + if not tail.size: + return output + return np.concatenate((output, tail)) + + def reset(self) -> None: + self._history = np.zeros(self._history_samples, dtype=np.float32) + self._input_samples = 0 + self._output_samples = 0 + self._flushed = False