diff --git a/examples/online_serving/minicpmo/README.md b/examples/online_serving/minicpmo/README.md index e1c084224..11f6f8195 100644 --- a/examples/online_serving/minicpmo/README.md +++ b/examples/online_serving/minicpmo/README.md @@ -1,159 +1,139 @@ # MiniCPM-o 4.5: Online serving -OpenAI-compatible `/v1/chat/completions` serving for **MiniCPM-o 4.5**, plus a -Gradio UI and curl / Python clients. +This directory contains MiniCPM-o 4.5 online serving demos for vLLM-Omni. +Inputs can include text, image, audio, or video; outputs are text and optional +24 kHz speech. -Inputs: text / image / audio / video. Outputs: text and optional **24 kHz** speech. +For the experimental native duplex runtime architecture, lifecycle invariants, +capability boundary, and validation scope, see +[`vllm_omni/experimental/fullduplex/DESIGN.md`](../../../vllm_omni/experimental/fullduplex/DESIGN.md). ## Installation -Please refer to [README.md](../../../README.md). Install the talker extra: +Install vLLM-Omni with the MiniCPM-o talker dependencies: ```bash pip install 'vllm-omni[minicpmo]' + +# From a source checkout: +pip install -e '.[minicpmo]' ``` -## Launch the server +The `minicpmo` extra installs `stepaudio2-minicpmo` and its audio dependencies, +including `librosa`. -The deploy config auto-loads via `--omni`; the default -`vllm_omni/deploy/minicpmo_4_5.yaml` targets a single-GPU layout (thinker -and talker + t2w co-located on GPU 0). For other hardware layouts pick -one of the deploy variants below. +## Start the backend server -| deploy config | GPUs | Notes | -|---|---|---| -| `minicpmo_4_5.yaml` (default) | 1 | Thinker and talker+t2w co-located on GPU0. | -| `minicpmo_4_5_2gpu.yaml` | 2 | Thinker on GPU0, talker+t2w on GPU1. | -| `minicpmo_4_5_3gpu.yaml` | 3 | Thinker 2-way TP on GPU0/1, talker+t2w share GPU2. | -| `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. | +Pick a deploy config that matches your GPU layout: -Default (single-GPU): +| config | GPUs | TP | Notes | +|---|---|---|---| +| `minicpmo_4_5.yaml` | 1 | 1 | Thinker and talker+t2w co-located on GPU0. | +| `minicpmo_4_5_2gpu.yaml` | 2 | 1 | Thinker on GPU0, talker+t2w on GPU1. | +| `minicpmo_4_5_3gpu.yaml` | 3 | 2 | Thinker 2-way TP on GPU0/1, talker+t2w share GPU2. | +| `minicpmo_4_5_8x4090.yaml` | 8 | 4 | Thinker 4-way TP on GPU0-3, talker+t2w on GPU4. | +| `minicpmo_4_5_3gpu_stage1_replicas.yaml` | 3 | 1 | Experimental validation profile: Thinker on GPU0, two talker+Token2wav replicas on GPU1/2. | +| `minicpmo_4_5_4gpu_stage1_replicas.yaml` | 4 | 1 | Experimental validation profile: Thinker on GPU0, three talker+Token2wav replicas on GPU1/2/3. | +| `minicpmo_4_5_8x4090_stage1_replicas.yaml` | 8 | 4 | Experimental validation profile: Thinker 4-way TP on GPU0-3, four talker+Token2wav replicas on GPU4-7. | +| `minicpmo_4_5_duplex.yaml` | 2 | 1 | Experimental native duplex profile. | ```bash -vllm serve openbmb/MiniCPM-o-4_5 --omni \ +vllm-omni serve openbmb/MiniCPM-o-4_5 \ + --omni \ + --deploy-config vllm_omni/deploy/minicpmo_4_5.yaml \ --trust-remote-code \ --host 0.0.0.0 --port 8099 ``` -Other layouts: +For local ModelScope checkpoints, replace `openbmb/MiniCPM-o-4_5` with the +checkpoint path. To start the experimental native duplex backend, use +`vllm_omni/deploy/minicpmo_4_5_duplex.yaml`. -```bash -vllm serve openbmb/MiniCPM-o-4_5 --omni \ - --deploy-config vllm_omni/deploy/minicpmo_4_5_3gpu.yaml \ - --trust-remote-code \ - --host 0.0.0.0 --port 8099 -``` +The `*_stage1_replicas.yaml` files exercise composite Stage1 replica routing +and failure recovery. They are validation profiles, not recommended production +entrypoints. -### Stage-based CLI (optional) +## Send chat requests -Stage 0 (thinker + API) and stage 1 (talker) can run in separate processes: - -```bash -# Stage 0 -CUDA_VISIBLE_DEVICES=0 vllm serve openbmb/MiniCPM-o-4_5 --omni \ - --trust-remote-code --port 8099 --stage-id 0 \ - --omni-master-address 127.0.0.1 --omni-master-port 26000 - -# Stage 1 (headless) -CUDA_VISIBLE_DEVICES=1 vllm serve openbmb/MiniCPM-o-4_5 --omni \ - --trust-remote-code --stage-id 1 --headless \ - --omni-master-address 127.0.0.1 --omni-master-port 26000 -``` - -### Per-stage overrides - -```bash -vllm serve openbmb/MiniCPM-o-4_5 --omni --trust-remote-code --port 8099 \ - --stage-overrides '{"0": {"gpu_memory_utilization": 0.65}}' -``` - -## Send multimodal requests - -```bash -cd examples/online_serving/minicpmo -``` - -### curl +From this directory, run the MiniCPM-specific curl or Python clients: ```bash bash run_curl_multimodal_generation.sh text bash run_curl_multimodal_generation.sh use_image -bash run_curl_multimodal_generation.sh use_audio '["text"]' # text-only -``` +bash run_curl_multimodal_generation.sh use_audio '["text"]' -Text + speech smoke test (TTS needs top-level `chat_template_kwargs`): - -```bash -curl http://localhost:8099/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{ - "model": "openbmb/MiniCPM-o-4_5", - "messages": [{"role": "user", "content": "Say hello, then introduce vLLM in one sentence."}], - "modalities": ["text", "audio"], - "chat_template_kwargs": {"use_tts_template": true} - }' -``` - -### Python client - -```bash python openai_chat_completion_client_for_multimodal_generation.py \ --query-type use_image \ - --port 8099 \ - --host localhost - -# Text-only (faster; no <|tts_bos|>) -python openai_chat_completion_client_for_multimodal_generation.py \ - --query-type text \ - --modalities text \ - --prompt "Briefly introduce yourself." -``` - -Shared helpers also work if you pass MiniCPM defaults yourself: - -```bash -python ../openai_chat_completion_client_for_multimodal_generation.py \ - --model openbmb/MiniCPM-o-4_5 \ - --query-type text \ + --host localhost \ --port 8099 ``` -(Note: the shared client does **not** set `use_tts_template`; prefer the -MiniCPM-specific client above for speech.) +Speech output requires `chat_template_kwargs.use_tts_template=true`. Put that +field at the request root for curl; the OpenAI Python SDK can merge it from +`extra_body`. -### Gradio demo +## Launch the Gradio demo ```bash -bash run_gradio_demo.sh +bash examples/online_serving/minicpmo/run_gradio_demo.sh -# Or: -python gradio_demo.py \ +# Or run the Python entry point directly: +python examples/online_serving/minicpmo/gradio_demo.py \ --minicpmo45-api-base http://localhost:8099/v1 \ --minicpmo45-model openbmb/MiniCPM-o-4_5 \ --port 7862 ``` -Open `http://:7862`. Uncheck **"Generate speech output (TTS)"** for -text-only responses. +Open `http://:7862` in a browser. -## Modality control +## Run the Realtime duplex CLI demo -| Modalities | Output | -|---|---| -| `["text"]` | Text only (no TTS bos) | -| `["text", "audio"]` / unset | Text + 24 kHz speech | +After the duplex backend is running, stream one WAV through the Realtime +WebSocket endpoint: -Speech requires `chat_template_kwargs.use_tts_template=true` so the chat -template appends `<|tts_bos|>`. For **curl**, put that field at the request -root; nested `extra_body` is ignored. The OpenAI Python SDK may use -`extra_body` because it merges those fields into the root. +```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 +``` -## Notes +## Open the experimental browser client -- Stage 1 is capped at `max_num_seqs: 1` in the deploy YAML (talker shares - request-0 audio metadata). -- Output audio is base64 WAV in `message.audio.data` (24 kHz mono). -- Offline counterpart: - [`examples/offline_inference/minicpmo/`](../../offline_inference/minicpmo/) -- Recipe: - [`recipes/OpenBMB/MiniCPM-o-4_5.md`](../../../recipes/OpenBMB/MiniCPM-o-4_5.md) +The browser UI serves the page and proxies the same-origin Realtime WebSocket to +the backend: + +```bash +python -m examples.online_serving.minicpmo.realtime_web \ + --port 7862 \ + --ws-backend ws://127.0.0.1:8099 \ + --ref-audio /path/to/MiniCPM-o-Demo/assets/ref_audio/ref_minicpm_signature.wav +``` + +Open `http://:7862/`. When using a reverse proxy, open the URL mapped to +port `7862`; the browser derives its WebSocket endpoint relative to that URL. + +If the page proxy serves HTTP but does not forward WebSocket upgrades, point the +browser at a separately exposed Realtime endpoint: + +```bash +python -m examples.online_serving.minicpmo.realtime_web \ + --port 7862 \ + --ws-backend ws://127.0.0.1:8099 \ + --public-realtime-url wss://public.example/v1/realtime +``` + +## Validate soft-interrupt behavior + +The soft-interrupt E2E driver defaults to `--validation-mode model-policy`, +which checks lifecycle and streaming invariants for arbitrary input audio. The +stronger `response-required` mode is diagnostic: it requires a purpose-built +two-response WAV, its `--input-sha256`, and an +`--expect-second-response-substring` value. + +## Related examples + +- [Offline MiniCPM-o inference](../../offline_inference/minicpmo/) +- [MiniCPM-o 4.5 recipe](../../../recipes/OpenBMB/MiniCPM-o-4_5.md) diff --git a/examples/online_serving/minicpmo/realtime_duplex_demo.py b/examples/online_serving/minicpmo/realtime_duplex_demo.py new file mode 100644 index 000000000..fcb7042a0 --- /dev/null +++ b/examples/online_serving/minicpmo/realtime_duplex_demo.py @@ -0,0 +1,399 @@ +"""Minimal single-input MiniCPM-o 4.5 Realtime duplex demo. + +Run this after starting the duplex server. Strict lifecycle, overlap, and +multi-session validation lives under ``tests/e2e/online_serving``. +""" + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import json +import sys +import time +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[3] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from vllm_omni.experimental.fullduplex.client import ( # noqa: E402 + PCM16_BYTES_PER_SAMPLE, + PCM16_SAMPLE_RATE, + RealtimeDuplexClient, + RealtimeEventCollector, + build_realtime_url, + read_pcm16_wav, + wait_for, + write_pcm16_wav, +) + + +class _StreamingOutputWriter: + """Persist and report output deltas as the client receives them.""" + + def __init__(self, output_dir: Path) -> None: + self.output_dir = output_dir + self.audio_chunk_dir = output_dir / "audio_chunks" + self.audio_chunk_paths: list[Path] = [] + self.output_dir.mkdir(parents=True, exist_ok=True) + self.audio_chunk_dir.mkdir(parents=True, exist_ok=True) + (self.output_dir / "output.pcm").write_bytes(b"") + + def handle(self, event: dict[str, object]) -> None: + event_type = event.get("type") + if event_type in { + "response.audio_transcript.delta", + "response.output_text.delta", + }: + delta = event.get("delta") + if isinstance(delta, str) and delta: + print(delta, end="", file=sys.stderr, flush=True) + return + if event_type != "response.audio.delta": + return + + delta = event.get("delta") or event.get("audio") + if not isinstance(delta, str) or not delta: + return + try: + pcm16 = base64.b64decode(delta) + except ValueError: + return + if not pcm16: + return + + chunk_index = len(self.audio_chunk_paths) + 1 + chunk_path = self.audio_chunk_dir / f"chunk_{chunk_index:04d}.wav" + sample_rate_hz = event.get("sample_rate_hz") + if not isinstance(sample_rate_hz, int) or sample_rate_hz <= 0: + sample_rate_hz = 24_000 + with (self.output_dir / "output.pcm").open("ab") as output_pcm: + output_pcm.write(pcm16) + write_pcm16_wav(chunk_path, pcm16, sample_rate_hz=sample_rate_hz) + self.audio_chunk_paths.append(chunk_path) + print( + f"\n[audio chunk {chunk_index}: {len(pcm16)} bytes -> {chunk_path}]", + file=sys.stderr, + flush=True, + ) + + +class _StreamingEventCollector(RealtimeEventCollector): + def __init__(self, writer: _StreamingOutputWriter) -> None: + super().__init__() + self._writer = writer + + def add(self, event: dict[str, object], *, received_at_s: float | None = None) -> None: + super().add(event, received_at_s=received_at_s) + self._writer.handle(self.events[-1]) + + +def _input_committed_index( + events: list[dict[str, object]], + after_index: int, +) -> int | None: + for index, event in enumerate(events[max(after_index, 0) :], start=max(after_index, 0)): + if event.get("type") == "input_audio_buffer.committed": + return index + return None + + +def _post_commit_model_decision( + events: list[dict[str, object]], + committed_index: int | None, +) -> str | None: + if committed_index is None: + return None + for event in events[committed_index + 1 :]: + event_type = event.get("type") + if event_type == "response.listen": + return "listen" + if event_type == "response.done": + response = event.get("response") + if not isinstance(response, dict) or response.get("status") != "cancelled": + return "speak" + return None + + +def _latest_model_decision( + events: list[dict[str, object]], + after_index: int, +) -> str | None: + decision: str | None = None + for event in events[max(after_index, 0) :]: + event_type = event.get("type") + if event_type == "response.listen": + decision = "listen" + elif event_type == "response.done": + response = event.get("response") + if not isinstance(response, dict) or response.get("status") != "cancelled": + decision = "speak" + return decision + + +def _chunk_period_ms(events: list[dict[str, object]]) -> int: + for event in reversed(events): + session = event.get("session") + if not isinstance(session, dict): + continue + capabilities = session.get("capabilities") + if not isinstance(capabilities, dict): + continue + chunk_period_ms = capabilities.get("chunk_period_ms") + if isinstance(chunk_period_ms, int) and chunk_period_ms > 0: + return chunk_period_ms + return 1000 + + +def _has_residual_model_unit(pcm16: bytes, *, chunk_period_ms: int) -> bool: + unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * chunk_period_ms // 1000 + return bool(unit_bytes > 0 and len(pcm16) % unit_bytes) + + +def _response_in_progress(events: list[dict[str, object]]) -> bool: + return sum(event.get("type") == "response.created" for event in events) > sum( + event.get("type") == "response.done" for event in events + ) + + +def _event_count_after( + events: list[dict[str, object]], + event_type: str, + index: int | None, +) -> int: + if index is None: + return 0 + return sum(event.get("type") == event_type for event in events[index + 1 :]) + + +def _ref_audio_data_url(path: str | None) -> str | None: + if path is None: + return None + ref_path = Path(path).expanduser() + return "data:audio/wav;base64," + base64.b64encode(ref_path.read_bytes()).decode("ascii") + + +async def run_demo(args: argparse.Namespace) -> dict[str, object]: + input_pcm16 = read_pcm16_wav(Path(args.input_wav)) + if not input_pcm16: + raise ValueError("input WAV has no audio") + + output_dir = Path(args.output_dir) + stream_writer = _StreamingOutputWriter(output_dir) + + url = build_realtime_url( + args.url, + args.model, + autostart=False if args.ref_audio else None, + session_id=args.session_id, + ) + client = RealtimeDuplexClient(url) + client.events = _StreamingEventCollector(stream_writer) + async with client: + await client.configure( + args.model, + ref_audio=_ref_audio_data_url(args.ref_audio), + session_id=args.session_id, + timeout_s=args.timeout_s, + ) + stream_event_cursor = len(client.events.events) + await client.stream_pcm16( + input_pcm16, + chunk_ms=args.chunk_ms, + realtime=not args.no_realtime_pacing, + ) + commit_event_cursor = len(client.events.events) + stream_decision = _latest_model_decision(client.events.events, stream_event_cursor) + input_has_residual_model_unit = _has_residual_model_unit( + input_pcm16, + chunk_period_ms=_chunk_period_ms(client.events.events), + ) + wait_for_post_commit_decision = False + commit_sent_at_s = time.monotonic() + await client.commit() + wait_error: str | None = None + committed_index: int | None = None + post_commit_decision: str | None = None + try: + await wait_for( + lambda: _input_committed_index(client.events.events, commit_event_cursor) is not None, + timeout_s=args.timeout_s, + label="input_audio_buffer.committed", + ) + committed_index = _input_committed_index(client.events.events, commit_event_cursor) + stream_decision = _latest_model_decision(client.events.events[: committed_index + 1], stream_event_cursor) + wait_for_post_commit_decision = input_has_residual_model_unit or _response_in_progress( + client.events.events[: committed_index + 1] + ) + if wait_for_post_commit_decision: + await wait_for( + lambda: _post_commit_model_decision(client.events.events, committed_index) is not None, + timeout_s=args.timeout_s, + label="post-commit model decision or response drain", + ) + post_commit_decision = _post_commit_model_decision(client.events.events, committed_index) + except TimeoutError as exc: + wait_error = str(exc) + await client.acknowledge_playback() + close_error: str | None = None + try: + await client.close_session(timeout_s=args.timeout_s) + except TimeoutError as exc: + close_error = str(exc) + + audio = client.events.audio_bytes() + first_text_at_s = client.events.first_received_at( + "response.audio_transcript.delta", + "response.output_text.delta", + after_s=commit_sent_at_s, + ) + first_audio_at_s = client.events.first_received_at( + "response.audio.delta", + after_s=commit_sent_at_s, + ) + response_created_at_s = client.events.first_received_at( + "response.created", + after_s=commit_sent_at_s, + ) + response_done_at_s = client.events.first_received_at( + "response.done", + after_s=commit_sent_at_s, + ) + audio_duration_s = len(audio) / (client.events.output_sample_rate_hz * 2) + response_generation_s = ( + response_done_at_s - response_created_at_s + if response_done_at_s is not None and response_created_at_s is not None + else None + ) + transcript_deltas = [ + str(event.get("delta", "")) + for event in client.events.events + if event.get("type") + in { + "response.audio_transcript.delta", + "response.output_text.delta", + } + ] + response_id = client.events.response_ids[0] if client.events.response_ids else None + timing = client.events.timing_summary( + after_s=commit_sent_at_s, + input_committed_at_s=commit_sent_at_s, + response_id=response_id, + ) + errors = client.events.errors() + if wait_error: + errors.append({"type": "client.timeout", "message": wait_error}) + if close_error: + errors.append({"type": "client.timeout", "message": close_error}) + model_decision = post_commit_decision or stream_decision + (output_dir / "events.jsonl").write_text( + "".join(json.dumps(event, ensure_ascii=False) + "\n" for event in client.events.events), + encoding="utf-8", + ) + (output_dir / "output.pcm").write_bytes(audio) + if audio: + write_pcm16_wav( + output_dir / "output.wav", + audio, + sample_rate_hz=client.events.output_sample_rate_hz, + ) + + result = { + "ok": ( + client.events.count("session.created") > 0 + and client.events.count("session.closed") > 0 + and not errors + and model_decision is not None + and (bool(audio) or not args.require_audio) + ), + "model_decision": model_decision, + "post_commit": { + "input_committed_event_index": committed_index, + "decision": post_commit_decision, + "decision_required": wait_for_post_commit_decision, + "input_had_residual_model_unit": input_has_residual_model_unit, + "response_listen_count": _event_count_after( + client.events.events, + "response.listen", + committed_index, + ), + "response_done_count": _event_count_after( + client.events.events, + "response.done", + committed_index, + ), + }, + "audio_bytes": len(audio), + "audio_chunk_count": len(stream_writer.audio_chunk_paths), + "audio_chunk_files": [str(path) for path in stream_writer.audio_chunk_paths], + "output_sample_rate_hz": client.events.output_sample_rate_hz, + "latency": { + "ttft_ms": ( + round((first_text_at_s - commit_sent_at_s) * 1000, 2) if first_text_at_s is not None else None + ), + "ttfp_ms": ( + round((first_audio_at_s - commit_sent_at_s) * 1000, 2) if first_audio_at_s is not None else None + ), + "rtf": ( + round(response_generation_s / audio_duration_s, 4) + if response_generation_s is not None and audio_duration_s > 0 + else None + ), + "response_generation_ms": ( + round(response_generation_s * 1000, 2) if response_generation_s is not None else None + ), + "text_stream_ms": ( + round((response_done_at_s - first_text_at_s) * 1000, 2) + if response_done_at_s is not None and first_text_at_s is not None + else None + ), + "transcript_delta_count": len(transcript_deltas), + "audio_duration_s": round(audio_duration_s, 3), + "measurement_origin": "input_audio_buffer.commit send", + }, + "timing": timing, + "response_ids": client.events.response_ids, + "transcript": "".join(transcript_deltas), + "errors": errors, + "output_dir": str(output_dir), + } + (output_dir / "result.json").write_text( + json.dumps(result, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + return result + + +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", required=True) + parser.add_argument( + "--ref-audio", + required=True, + help=( + "Reference WAV for the MiniCPM-o duplex assistant voice. " + "This demo matches the official flow by always providing a reference audio clip." + ), + ) + parser.add_argument("--output-dir", default="/tmp/minicpmo_realtime_duplex_demo") + parser.add_argument("--chunk-ms", type=int, default=200) + parser.add_argument("--timeout-s", type=float, default=60.0) + parser.add_argument("--no-realtime-pacing", action="store_true") + parser.add_argument("--require-audio", action="store_true") + return parser.parse_args() + + +def main() -> None: + result = asyncio.run(run_demo(parse_args())) + print(json.dumps(result, ensure_ascii=False, indent=2)) + if not result["ok"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/minicpmo/realtime_web/__init__.py b/examples/online_serving/minicpmo/realtime_web/__init__.py new file mode 100644 index 000000000..9a5de5e18 --- /dev/null +++ b/examples/online_serving/minicpmo/realtime_web/__init__.py @@ -0,0 +1,5 @@ +"""Browser demo for the experimental full-duplex runtime.""" + +from .server import build_app + +__all__ = ["build_app"] diff --git a/examples/online_serving/minicpmo/realtime_web/__main__.py b/examples/online_serving/minicpmo/realtime_web/__main__.py new file mode 100644 index 000000000..f5f6e402d --- /dev/null +++ b/examples/online_serving/minicpmo/realtime_web/__main__.py @@ -0,0 +1,4 @@ +from .server import main + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/minicpmo/realtime_web/app/index.html b/examples/online_serving/minicpmo/realtime_web/app/index.html new file mode 100644 index 000000000..46ab15907 --- /dev/null +++ b/examples/online_serving/minicpmo/realtime_web/app/index.html @@ -0,0 +1,79 @@ + + + + + + + MiniCPM-o 4.5 Voice + + + +
+
+

Experimental full-duplex runtime

+

MiniCPM-o 4.5 Voice

+
+ Offline +
+ +
+
+
+ + + + +
+
+ + +
+
+ Mic +
+
+
+
Model
Idle
+
Playback
Idle
+
Session
00:00
+
+
+ +
+
+
+

Live transcript

+

Conversation

+
+ Model-controlled listen / speak +
+
+

Start a session and speak naturally.

+
+
+ +
+ + Event log + 0 events + +
+ No active connection + +
+

+      
+
+ + + + + diff --git a/examples/online_serving/minicpmo/realtime_web/app/static/app.js b/examples/online_serving/minicpmo/realtime_web/app/static/app.js new file mode 100644 index 000000000..e8a3047c6 --- /dev/null +++ b/examples/online_serving/minicpmo/realtime_web/app/static/app.js @@ -0,0 +1,671 @@ +(() => { + 'use strict'; + + const config = window.FULL_DUPLEX_CONFIG || {}; + const callButton = document.getElementById('callButton'); + const muteButton = document.getElementById('muteButton'); + const cameraButton = document.getElementById('cameraButton'); + const cameraPreview = document.getElementById('cameraPreview'); + const promptPreset = document.getElementById('promptPreset'); + const systemPromptInput = document.getElementById('systemPrompt'); + const connectionState = document.getElementById('connectionState'); + const modelState = document.getElementById('modelState'); + const playbackState = document.getElementById('playbackState'); + const sessionTimer = document.getElementById('sessionTimer'); + const meterFill = document.getElementById('meterFill'); + const conversation = document.getElementById('conversation'); + const emptyConversation = document.getElementById('emptyConversation'); + const eventLog = document.getElementById('eventLog'); + const eventCount = document.getElementById('eventCount'); + const runtimeDetail = document.getElementById('runtimeDetail'); + const clearLogButton = document.getElementById('clearLogButton'); + + const INPUT_RATE = 16000; + const OUTPUT_RATE = 24000; + const SEND_INTERVAL_MS = 200; + const ECHO_GUARD_MS = 300; + const INITIAL_PLAYBACK_BUFFER_MS = 400; + const SESSION_CLOSE_TIMEOUT_MS = 1000; + + // Default prompts mirroring the official MiniCPM-o-Demo presets + // (assets/presets/{omni,audio_duplex}/*.yaml). + const PROMPT_PRESETS = { + omni: 'Streaming Omni Conversation.', + chinese_call: '扮演一个具有以上声音特征的助手。请认真、高质量地回复用户的问题。' + + '请用高自然度的方式和用户聊天。你处于双工模式,可以一边听、一边说。' + + '你是由面壁智能开发的人工智能助手:面壁小钢炮。', + english_call: 'Replicate the tone and style from the input audio. Your task is to be ' + + 'a helpful assistant using this voice pattern. Please answer the user\'s questions ' + + 'seriously and in a high quality. Please chat with the user in a high naturalness ' + + 'style. You are in duplex mode, where you can listen and speak at the same time.', + }; + + let socket = null; + let mediaStream = null; + let captureContext = null; + let captureNode = null; + let playbackContext = null; + let playbackNode = null; + let sendTimer = null; + let clockTimer = null; + let startedAt = 0; + let running = false; + let muted = false; + let assistantActive = false; + let captureRate = INPUT_RATE; + let cameraStream = null; + let cameraTimer = null; + let cameraPendingFrame = null; + const cameraCanvas = document.createElement('canvas'); + let playbackRate = OUTPUT_RATE; + let pendingCapture = []; + let currentResponseId = null; + let responseHasAudio = false; + let logCount = 0; + let liveUserTurn = null; + let liveAssistantTurn = null; + let sessionCloseResolver = null; + + function staticAssetUrl(path) { + const version = String(config.appVersion || '').trim(); + return version ? `${path}?v=${encodeURIComponent(version)}` : path; + } + + if (promptPreset && systemPromptInput) { + promptPreset.addEventListener('change', () => { + const preset = PROMPT_PRESETS[promptPreset.value]; + if (preset !== undefined) systemPromptInput.value = preset; + }); + systemPromptInput.addEventListener('input', () => { + promptPreset.value = 'custom'; + }); + } + + function realtimeUrl() { + const url = new URL(config.realtimePath, window.location.href); + 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('minicpmo45_native_duplex', '1'); + url.searchParams.set('autostart', '0'); + return url.toString(); + } + + function setConnection(label, kind) { + connectionState.textContent = label; + connectionState.className = `status status-${kind}`; + } + + function setModel(label) { + modelState.textContent = label; + } + + function setPlayback(label) { + playbackState.textContent = label; + } + + function compactEvent(event) { + const fields = []; + const payload = event.event || event; + const responseId = responseIdOf(event); + if (responseId) fields.push(`response=${responseId}`); + if (event.item_id) fields.push(`item=${event.item_id}`); + if (event.code) fields.push(`code=${event.code}`); + if (event.response && event.response.status) fields.push(`status=${event.response.status}`); + if (payload.played_ms !== undefined) fields.push(`played=${payload.played_ms}ms`); + if (payload.committed_ms !== undefined) fields.push(`committed=${payload.committed_ms}ms`); + if (payload.history_committed !== undefined) fields.push(`history=${payload.history_committed}`); + return fields.join(' '); + } + + function appendLog(message, error = false) { + const time = new Date().toLocaleTimeString([], { hour12: false }); + const line = document.createElement('span'); + if (error) line.className = 'log-error'; + line.textContent = `${time} ${message}\n`; + eventLog.appendChild(line); + eventLog.scrollTop = eventLog.scrollHeight; + logCount += 1; + eventCount.textContent = `${logCount} ${logCount === 1 ? 'event' : 'events'}`; + } + + function appendEventLog(event) { + const detail = compactEvent(event); + appendLog(`${event.type || 'unknown'}${detail ? ` ${detail}` : ''}`, event.type === 'error'); + } + + function responseIdOf(event) { + return event.response_id || (event.response && event.response.id) || null; + } + + function ensureTurn(role) { + const existing = role === 'user' ? liveUserTurn : liveAssistantTurn; + if (existing) return existing; + if (emptyConversation) emptyConversation.remove(); + const row = document.createElement('div'); + row.className = `turn turn-${role} turn-live`; + const label = document.createElement('div'); + label.className = 'turn-role'; + label.textContent = role === 'user' ? 'You' : 'Assistant'; + const text = document.createElement('div'); + text.className = 'turn-text'; + row.append(label, text); + conversation.appendChild(row); + conversation.scrollTop = conversation.scrollHeight; + const turn = { row, text, value: '' }; + if (role === 'user') liveUserTurn = turn; + else liveAssistantTurn = turn; + return turn; + } + + function addTranscript(role, delta) { + if (!delta) return; + const turn = ensureTurn(role); + turn.value += delta; + turn.text.textContent = turn.value; + conversation.scrollTop = conversation.scrollHeight; + } + + function finishTranscript(role, finalText = '') { + const turn = role === 'user' ? liveUserTurn : liveAssistantTurn; + if (!turn && !finalText) return; + const current = turn || ensureTurn(role); + if (finalText) { + current.value = finalText; + current.text.textContent = finalText; + } + current.row.classList.remove('turn-live'); + if (role === 'user') liveUserTurn = null; + else liveAssistantTurn = null; + } + + function bytesToBase64(bytes) { + let binary = ''; + const chunkSize = 0x8000; + for (let offset = 0; offset < bytes.length; offset += chunkSize) { + binary += String.fromCharCode(...bytes.subarray(offset, offset + chunkSize)); + } + return btoa(binary); + } + + function int16ToBase64(pcm) { + return bytesToBase64(new Uint8Array(pcm.buffer, pcm.byteOffset, pcm.byteLength)); + } + + function base64ToBytes(encoded) { + const binary = atob(encoded); + const bytes = new Uint8Array(binary.length); + for (let index = 0; index < binary.length; index += 1) bytes[index] = binary.charCodeAt(index); + return bytes; + } + + function resampleInt16(input, sourceRate, targetRate) { + if (sourceRate === targetRate) return input; + const ratio = sourceRate / targetRate; + const output = new Int16Array(Math.floor(input.length / ratio)); + for (let index = 0; index < output.length; index += 1) { + const start = Math.floor(index * ratio); + const end = Math.max(start + 1, Math.min(input.length, Math.floor((index + 1) * ratio))); + let sum = 0; + for (let source = start; source < end; source += 1) sum += input[source]; + output[index] = sum / (end - start); + } + return output; + } + + async function decodeAudioDelta(event) { + const encoded = event.delta || (event.response && event.response.audio); + if (!encoded) return null; + const bytes = base64ToBytes(encoded); + const format = String(event.format || event.audio_format || 'pcm16').toLowerCase(); + const sourceRate = Number(event.sample_rate_hz || event.sample_rate || OUTPUT_RATE); + if (format.includes('f32')) { + const floats = new Float32Array(bytes.buffer, bytes.byteOffset, Math.floor(bytes.byteLength / 4)); + const pcm = new Int16Array(floats.length); + for (let index = 0; index < floats.length; index += 1) { + const sample = Math.max(-1, Math.min(1, floats[index])); + pcm[index] = sample < 0 ? sample * 32768 : sample * 32767; + } + return { pcm, sourceRate }; + } + if (format.includes('wav')) { + const decoded = await playbackContext.decodeAudioData(bytes.buffer.slice(0)); + const channel = decoded.getChannelData(0); + const pcm = new Int16Array(channel.length); + for (let index = 0; index < channel.length; index += 1) { + const sample = Math.max(-1, Math.min(1, channel[index])); + pcm[index] = sample < 0 ? sample * 32768 : sample * 32767; + } + return { pcm, sourceRate: decoded.sampleRate }; + } + return { + pcm: new Int16Array(bytes.buffer, bytes.byteOffset, Math.floor(bytes.byteLength / 2)), + sourceRate, + }; + } + + function updateMeter(pcm) { + let peak = 0; + for (let index = 0; index < pcm.length; index += 8) peak = Math.max(peak, Math.abs(pcm[index])); + meterFill.style.width = `${Math.min(100, (peak / 32768) * 150).toFixed(0)}%`; + } + + function microphoneUploadEnabled() { + return running && !muted; + } + + function flushCapture() { + if (!socket || socket.readyState !== WebSocket.OPEN || pendingCapture.length === 0) return; + if (!microphoneUploadEnabled()) { + pendingCapture = []; + return; + } + const length = pendingCapture.reduce((total, chunk) => total + chunk.length, 0); + const merged = new Int16Array(length); + let offset = 0; + for (const chunk of pendingCapture) { + merged.set(chunk, offset); + offset += chunk.length; + } + pendingCapture = []; + const pcm = resampleInt16(merged, captureRate, INPUT_RATE); + const appendEvent = { + type: 'input_audio_buffer.append', + audio: int16ToBase64(pcm), + format: 'pcm16', + sample_rate_hz: INPUT_RATE, + }; + // Omni duplex: ~1 fps camera frame rides the audio append (official + // MiniCPM-o-Demo contract: one base64 JPEG per ~1 s chunk). + if (cameraPendingFrame) { + appendEvent.video_frames = [cameraPendingFrame]; + cameraPendingFrame = null; + } + socket.send(JSON.stringify(appendEvent)); + } + + function beginAssistant(responseId) { + currentResponseId = responseId || currentResponseId; + responseHasAudio = false; + assistantActive = true; + setModel('Speaking'); + } + + function feedPlayback(decoded, responseId) { + if (!decoded || !decoded.pcm || decoded.pcm.length === 0 || !playbackNode) return; + const pcm = resampleInt16(decoded.pcm, decoded.sourceRate, playbackRate); + responseHasAudio = true; + assistantActive = true; + setPlayback('Buffering'); + playbackNode.port.postMessage({ + type: 'audio', + pcm, + responseId: responseId || currentResponseId, + initialBufferMs: INITIAL_PLAYBACK_BUFFER_MS, + }, [pcm.buffer]); + } + + function requestPlaybackDrain(responseId) { + if (!playbackNode) return; + playbackNode.port.postMessage({ type: 'drain', responseId: responseId || currentResponseId }); + } + + function sendPlaybackAck(responseId, playedMs) { + if (!responseId || !socket || socket.readyState !== WebSocket.OPEN || playedMs <= 0) { + if (!responseId && playedMs > 0) appendLog('playback ack skipped: missing response id', true); + return; + } + socket.send(JSON.stringify({ + type: 'playback.ack', + response_id: responseId, + item_id: `item_${responseId}`, + played_ms: playedMs, + committed_ms: playedMs, + })); + } + + function playbackDrained(message) { + const responseId = message.responseId || currentResponseId; + sendPlaybackAck(responseId, Number(message.playedMs) || 0); + setPlayback('Idle'); + if (message.underrunMs > 0) { + appendLog(`playback underrun ${message.underrunMs} ms`); + } + window.setTimeout(() => { + assistantActive = false; + currentResponseId = null; + responseHasAudio = false; + if (running) setModel('Listening'); + }, ECHO_GUARD_MS); + } + + function handleEvent(event) { + appendEventLog(event); + const responseId = responseIdOf(event); + switch (event.type) { + case 'session.created': + case 'session.updated': + setConnection('Connected', 'online'); + setModel('Listening'); + break; + case 'response.listen': + assistantActive = false; + setModel('Listening'); + break; + case 'response.created': + case 'response.speak': + beginAssistant(responseId); + break; + case 'response.audio.delta': + currentResponseId = responseId || currentResponseId; + assistantActive = true; + setModel('Speaking'); + decodeAudioDelta(event) + .then((decoded) => feedPlayback(decoded, responseId)) + .catch((error) => appendLog(`audio decode failed: ${error.message || error}`, true)); + break; + case 'response.audio.done': + requestPlaybackDrain(responseId); + break; + case 'response.audio_transcript.delta': + addTranscript('assistant', event.delta || ''); + break; + case 'response.audio_transcript.done': + finishTranscript('assistant', event.transcript || ''); + break; + case 'conversation.item.input_audio_transcription.delta': + addTranscript('user', event.delta || ''); + break; + case 'conversation.item.input_audio_transcription.completed': + finishTranscript('user', event.transcript || ''); + break; + case 'response.done': + finishTranscript('assistant'); + if (!responseHasAudio) requestPlaybackDrain(responseId); + break; + case 'playback.acknowledged': + { + const acknowledgement = event.event || event; + runtimeDetail.textContent = `Playback committed ${acknowledgement.committed_ms || 0} ms`; + } + break; + case 'session.closed': + if (sessionCloseResolver) sessionCloseResolver(); + sessionCloseResolver = null; + break; + case 'error': + setConnection('Error', 'error'); + runtimeDetail.textContent = String(event.error || event.code || 'Server error'); + break; + default: + break; + } + } + + async function openPlayback() { + playbackContext = new (window.AudioContext || window.webkitAudioContext)({ sampleRate: OUTPUT_RATE }); + playbackRate = playbackContext.sampleRate; + await playbackContext.audioWorklet.addModule(staticAssetUrl('static/playback_worklet.js')); + playbackNode = new AudioWorkletNode(playbackContext, 'fullduplex-pcm-playback'); + playbackNode.port.onmessage = (message) => { + if (message.data.type === 'playback-started') setPlayback('Playing'); + else if (message.data.type === 'playback-drained') playbackDrained(message.data); + else if (message.data.type === 'playback-underrun') { + runtimeDetail.textContent = `Playback underrun ${message.data.underrunMs || 0} ms`; + } + }; + playbackNode.connect(playbackContext.destination); + await playbackContext.resume(); + } + + async function openCapture() { + mediaStream = await navigator.mediaDevices.getUserMedia({ + audio: { + channelCount: 1, + echoCancellation: true, + noiseSuppression: true, + autoGainControl: true, + sampleRate: { ideal: INPUT_RATE }, + }, + }); + try { + captureContext = new (window.AudioContext || window.webkitAudioContext)({ sampleRate: INPUT_RATE }); + } catch (_error) { + captureContext = new (window.AudioContext || window.webkitAudioContext)(); + } + captureRate = captureContext.sampleRate; + await captureContext.audioWorklet.addModule(staticAssetUrl('static/pcm_worklet.js')); + const source = captureContext.createMediaStreamSource(mediaStream); + captureNode = new AudioWorkletNode(captureContext, 'fullduplex-pcm-capture'); + captureNode.port.onmessage = (message) => { + const pcm = new Int16Array(message.data); + updateMeter(pcm); + if (microphoneUploadEnabled()) pendingCapture.push(pcm); + }; + const silentSink = captureContext.createGain(); + silentSink.gain.value = 0; + source.connect(captureNode); + captureNode.connect(silentSink).connect(captureContext.destination); + await captureContext.resume(); + } + + function openSocket() { + return new Promise((resolve, reject) => { + const url = realtimeUrl(); + socket = new WebSocket(url); + let settled = false; + socket.onopen = () => { + settled = true; + const extraBody = { + auto_response: true, + minicpmo45_native_duplex: true, + }; + const session = { + modalities: ['audio', 'text'], + voice: 'default', + extra_body: extraBody, + }; + // Reference voice for TTS cloning, provided by the server via + // --ref-audio (mirrors the official demo's default ref audio). + if (config.refAudio) session.ref_audio = config.refAudio; + const instructions = systemPromptInput ? systemPromptInput.value.trim() : ''; + if (instructions) session.instructions = instructions; + socket.send(JSON.stringify({ type: 'session.update', session })); + runtimeDetail.textContent = `${captureRate} Hz capture / ${playbackRate} Hz playback`; + appendLog(`websocket open ${url}`); + resolve(); + }; + socket.onmessage = (message) => { + if (typeof message.data !== 'string') return; + try { + handleEvent(JSON.parse(message.data)); + } catch (error) { + appendLog(`invalid server event: ${error.message || error}`, true); + } + }; + socket.onerror = () => { + if (!settled) { + settled = true; + reject(new Error(`WebSocket connection failed: ${url}`)); + } + }; + socket.onclose = (event) => { + appendLog(`websocket closed code=${event.code}`); + if (running) { + setConnection('Disconnected', 'error'); + stopSession(); + } + }; + }); + } + + function formatElapsed(seconds) { + const minutes = Math.floor(seconds / 60); + return `${String(minutes).padStart(2, '0')}:${String(seconds % 60).padStart(2, '0')}`; + } + + function startClock() { + startedAt = Date.now(); + clockTimer = window.setInterval(() => { + sessionTimer.textContent = formatElapsed(Math.floor((Date.now() - startedAt) / 1000)); + }, 1000); + } + + async function startSession() { + if (running) return; + callButton.disabled = true; + setConnection('Connecting', 'connecting'); + runtimeDetail.textContent = 'Requesting microphone access'; + try { + await openPlayback(); + await openCapture(); + await openSocket(); + running = true; + muted = false; + assistantActive = false; + sendTimer = window.setInterval(flushCapture, SEND_INTERVAL_MS); + startClock(); + callButton.textContent = 'End session'; + callButton.classList.add('is-active'); + muteButton.disabled = false; + cameraButton.disabled = false; + setConnection('Connected', 'online'); + setModel('Listening'); + appendLog('session started'); + } catch (error) { + appendLog(`start failed: ${error.message || error}`, true); + setConnection('Error', 'error'); + runtimeDetail.textContent = String(error.message || error); + await stopSession(); + } finally { + callButton.disabled = false; + } + } + + async function startCamera() { + if (cameraStream) return; + cameraStream = await navigator.mediaDevices.getUserMedia({ video: true, audio: false }); + cameraPreview.srcObject = cameraStream; + cameraPreview.style.display = ''; + await cameraPreview.play().catch(() => {}); + // Official omni-duplex cadence: one JPEG (quality 0.7) per ~1 s chunk, + // no client-side resize (the server normalizes at scale_resolution=448). + cameraTimer = window.setInterval(() => { + if (!cameraStream || cameraPreview.videoWidth === 0) return; + cameraCanvas.width = cameraPreview.videoWidth; + cameraCanvas.height = cameraPreview.videoHeight; + cameraCanvas.getContext('2d').drawImage(cameraPreview, 0, 0); + cameraPendingFrame = cameraCanvas.toDataURL('image/jpeg', 0.7).split(',')[1]; + }, 1000); + cameraButton.textContent = 'Camera off'; + cameraButton.classList.add('is-active'); + appendLog('camera on (1 fps omni frames)'); + } + + function stopCamera() { + if (cameraTimer !== null) clearInterval(cameraTimer); + cameraTimer = null; + if (cameraStream) { + for (const track of cameraStream.getTracks()) track.stop(); + } + cameraStream = null; + cameraPendingFrame = null; + cameraPreview.srcObject = null; + cameraPreview.style.display = 'none'; + cameraButton.textContent = 'Camera'; + cameraButton.classList.remove('is-active'); + } + + cameraButton.addEventListener('click', () => { + if (cameraStream) { + stopCamera(); + appendLog('camera off'); + return; + } + startCamera().catch((error) => appendLog(`camera failed: ${error.message || error}`, true)); + }); + + function waitForSessionClosed(targetSocket, timeoutMs) { + if (!targetSocket || targetSocket.readyState !== WebSocket.OPEN) return Promise.resolve(); + return new Promise((resolve) => { + let done = false; + const finish = () => { + if (done) return; + done = true; + if (sessionCloseResolver === finish) sessionCloseResolver = null; + resolve(); + }; + sessionCloseResolver = finish; + window.setTimeout(finish, timeoutMs); + }); + } + + async function stopSession({ terminal = true } = {}) { + running = false; + assistantActive = false; + pendingCapture = []; + if (sendTimer !== null) clearInterval(sendTimer); + if (clockTimer !== null) clearInterval(clockTimer); + sendTimer = null; + clockTimer = null; + if (socket) { + const closingSocket = socket; + socket = null; + closingSocket.onclose = null; + if (terminal && closingSocket.readyState === WebSocket.OPEN) { + const closed = waitForSessionClosed(closingSocket, SESSION_CLOSE_TIMEOUT_MS); + closingSocket.send(JSON.stringify({ type: 'session.close' })); + await closed; + } + closingSocket.close(1000, 'client stop'); + } + if (playbackNode) playbackNode.port.postMessage({ type: 'clear' }); + if (mediaStream) { + for (const track of mediaStream.getTracks()) track.stop(); + } + mediaStream = null; + stopCamera(); + cameraButton.disabled = true; + if (captureContext) await captureContext.close().catch(() => {}); + if (playbackContext) await playbackContext.close().catch(() => {}); + captureContext = null; + captureNode = null; + playbackContext = null; + playbackNode = null; + currentResponseId = null; + responseHasAudio = false; + meterFill.style.width = '0%'; + sessionTimer.textContent = '00:00'; + callButton.textContent = 'Start session'; + callButton.classList.remove('is-active'); + muteButton.textContent = 'Mute'; + muteButton.classList.remove('is-active'); + muteButton.disabled = true; + setConnection('Offline', 'offline'); + setModel('Idle'); + setPlayback('Idle'); + if (runtimeDetail.textContent.startsWith('Playback committed')) return; + if (!runtimeDetail.textContent.startsWith('start failed')) runtimeDetail.textContent = 'No active connection'; + } + + function toggleMute() { + if (!running) return; + muted = !muted; + pendingCapture = []; + muteButton.textContent = muted ? 'Unmute' : 'Mute'; + muteButton.classList.toggle('is-active', muted); + appendLog(muted ? 'microphone muted' : 'microphone unmuted'); + } + + callButton.addEventListener('click', () => { + if (running) stopSession(); + else startSession(); + }); + muteButton.addEventListener('click', toggleMute); + clearLogButton.addEventListener('click', () => { + eventLog.textContent = ''; + logCount = 0; + eventCount.textContent = '0 events'; + }); + window.addEventListener('beforeunload', () => { stopSession({ terminal: false }); }); +})(); diff --git a/examples/online_serving/minicpmo/realtime_web/app/static/pcm_worklet.js b/examples/online_serving/minicpmo/realtime_web/app/static/pcm_worklet.js new file mode 100644 index 000000000..608ac6b01 --- /dev/null +++ b/examples/online_serving/minicpmo/realtime_web/app/static/pcm_worklet.js @@ -0,0 +1,15 @@ +class FullDuplexPcmCapture extends AudioWorkletProcessor { + process(inputs) { + const channel = inputs[0] && inputs[0][0]; + if (!channel || channel.length === 0) return true; + const pcm = new Int16Array(channel.length); + for (let index = 0; index < channel.length; index += 1) { + const sample = Math.max(-1, Math.min(1, channel[index])); + pcm[index] = sample < 0 ? sample * 32768 : sample * 32767; + } + this.port.postMessage(pcm.buffer, [pcm.buffer]); + return true; + } +} + +registerProcessor('fullduplex-pcm-capture', FullDuplexPcmCapture); diff --git a/examples/online_serving/minicpmo/realtime_web/app/static/playback_worklet.js b/examples/online_serving/minicpmo/realtime_web/app/static/playback_worklet.js new file mode 100644 index 000000000..c3216fd7f --- /dev/null +++ b/examples/online_serving/minicpmo/realtime_web/app/static/playback_worklet.js @@ -0,0 +1,153 @@ +class FullDuplexPcmPlayback extends AudioWorkletProcessor { + constructor() { + super(); + this.queue = []; + this.offset = 0; + this.playedFrames = 0; + this.underrunFrames = 0; + this.drain = null; + this.started = false; + this.activeResponseId = null; + this.initialBufferFrames = Math.round(sampleRate * 0.2); + this.bufferWaitFrames = this.initialBufferFrames; + this.rebuffering = false; + this.fadeFrames = Math.max(1, Math.round(sampleRate * 0.005)); + this.fadeInFrames = 0; + this.port.onmessage = (event) => this.handleMessage(event.data || {}); + } + + handleMessage(message) { + if (message.type === 'audio' && message.pcm) { + const wasEmpty = this.queue.length === 0; + if (!this.started && !this.activeResponseId) { + this.activeResponseId = message.responseId || null; + } + if (!this.started && Number.isFinite(message.initialBufferMs)) { + this.initialBufferFrames = Math.max(0, Math.round((sampleRate * message.initialBufferMs) / 1000)); + } + this.queue.push(message.pcm); + if (!this.started && wasEmpty && !this.rebuffering) { + this.bufferWaitFrames = this.initialBufferFrames; + } + } else if (message.type === 'drain') { + this.drain = { responseId: message.responseId || null }; + if (!this.started && this.bufferedFrames() > 0) { + this.bufferWaitFrames = 0; + this.startPlayback(); + } + this.notifyIfDrained(); + } else if (message.type === 'clear') { + this.queue = []; + this.offset = 0; + this.playedFrames = 0; + this.underrunFrames = 0; + this.drain = null; + this.started = false; + this.activeResponseId = null; + this.bufferWaitFrames = this.initialBufferFrames; + this.rebuffering = false; + this.fadeInFrames = 0; + } + } + + bufferedFrames() { + return this.queue.reduce((total, pcm, index) => ( + total + pcm.length - (index === 0 ? this.offset : 0) + ), 0); + } + + startPlayback() { + if (this.started) return; + this.started = true; + this.fadeInFrames = this.fadeFrames; + if (this.playedFrames === 0) { + this.port.postMessage({ type: 'playback-started', responseId: this.activeResponseId }); + } + } + + notifyIfDrained() { + if (!this.drain || this.queue.length > 0) return; + this.port.postMessage({ + type: 'playback-drained', + responseId: this.drain.responseId, + playedMs: Math.round((this.playedFrames * 1000) / sampleRate), + underrunMs: Math.round((this.underrunFrames * 1000) / sampleRate), + }); + this.playedFrames = 0; + this.underrunFrames = 0; + this.drain = null; + this.started = false; + this.activeResponseId = null; + this.bufferWaitFrames = this.initialBufferFrames; + this.rebuffering = false; + this.fadeInFrames = 0; + } + + reportUnderrun() { + this.port.postMessage({ + type: 'playback-underrun', + responseId: this.activeResponseId, + underrunMs: Math.round((this.underrunFrames * 1000) / sampleRate), + }); + } + + process(_inputs, outputs) { + const output = outputs[0][0]; + output.fill(0); + if (!this.started) { + if (this.rebuffering && !this.drain) { + this.underrunFrames += output.length; + this.reportUnderrun(); + } + if (this.queue.length > 0 && this.bufferWaitFrames > 0) { + this.bufferWaitFrames = Math.max(0, this.bufferWaitFrames - output.length); + return true; + } + if (this.queue.length > 0) { + this.startPlayback(); + this.rebuffering = false; + } + } + if (!this.started) { + this.notifyIfDrained(); + return true; + } + let target = 0; + while (target < output.length && this.queue.length > 0) { + const pcm = this.queue[0]; + const count = Math.min(output.length - target, pcm.length - this.offset); + for (let index = 0; index < count; index += 1) { + let sample = pcm[this.offset + index] / 32768; + if (this.fadeInFrames > 0) { + const elapsed = this.fadeFrames - this.fadeInFrames; + sample *= elapsed / this.fadeFrames; + this.fadeInFrames -= 1; + } + output[target + index] = sample; + } + target += count; + this.offset += count; + this.playedFrames += count; + if (this.offset >= pcm.length) { + this.queue.shift(); + this.offset = 0; + } + } + if (target < output.length && !this.drain) { + const fadeCount = Math.min(target, this.fadeFrames); + for (let index = 0; index < fadeCount; index += 1) { + output[target - fadeCount + index] *= (fadeCount - index - 1) / fadeCount; + } + this.underrunFrames += output.length - target; + this.reportUnderrun(); + this.started = false; + this.rebuffering = true; + this.bufferWaitFrames = this.initialBufferFrames; + this.fadeInFrames = 0; + } + this.notifyIfDrained(); + return true; + } +} + +registerProcessor('fullduplex-pcm-playback', FullDuplexPcmPlayback); diff --git a/examples/online_serving/minicpmo/realtime_web/app/static/styles.css b/examples/online_serving/minicpmo/realtime_web/app/static/styles.css new file mode 100644 index 000000000..9b080176b --- /dev/null +++ b/examples/online_serving/minicpmo/realtime_web/app/static/styles.css @@ -0,0 +1,127 @@ +:root { + color-scheme: light; + font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; + color: #20252b; + background: #f4f6f7; + font-synthesis: none; +} + +* { box-sizing: border-box; } + +body { margin: 0; min-width: 320px; } + +button, summary { font: inherit; } + +button:focus-visible, summary:focus-visible { + outline: 3px solid rgba(21, 101, 192, .25); + outline-offset: 2px; +} + +.app-header { + min-height: 104px; + padding: 24px clamp(20px, 5vw, 64px); + background: #ffffff; + border-bottom: 1px solid #d9dee3; + display: flex; + align-items: center; + justify-content: space-between; + gap: 20px; +} + +h1, h2, p { margin: 0; } +h1 { font-size: 24px; line-height: 1.2; letter-spacing: 0; } +h2 { font-size: 18px; line-height: 1.3; letter-spacing: 0; } + +.eyebrow { + margin-bottom: 5px; + color: #68717b; + font-size: 11px; + font-weight: 700; + text-transform: uppercase; + letter-spacing: .08em; +} + +main { width: min(1040px, calc(100% - 40px)); margin: 28px auto 48px; } + +.status { + min-width: 88px; + padding: 7px 10px; + border: 1px solid; + border-radius: 6px; + text-align: center; + font-size: 13px; + font-weight: 700; +} + +.status-offline { color: #616971; background: #f0f2f3; border-color: #cbd1d6; } +.status-connecting { color: #8a4b00; background: #fff5df; border-color: #e6b76a; } +.status-online { color: #176a3a; background: #e9f7ef; border-color: #7dc49b; } +.status-error { color: #a12222; background: #fff0f0; border-color: #df8f8f; } + +.call-bar { + min-height: 86px; + padding: 16px 18px; + background: #252b31; + color: #ffffff; + border-radius: 6px; + display: grid; + grid-template-columns: auto minmax(150px, 1fr) auto; + align-items: center; + gap: 24px; +} + +.call-actions { display: flex; gap: 8px; } +button { border-radius: 5px; border: 1px solid transparent; cursor: pointer; } +button:disabled { opacity: .46; cursor: default; } +.primary { padding: 10px 15px; color: #ffffff; background: #16834b; font-weight: 700; } +.primary.is-active { background: #b43636; } +.secondary { padding: 10px 13px; color: #ffffff; background: #39414a; border-color: #59636d; } +.secondary.is-active { color: #2b2525; background: #ffc9c9; border-color: #e58f8f; } + +.meter { min-width: 0; display: flex; align-items: center; gap: 10px; color: #c8cfd6; font-size: 12px; } +.meter-track { height: 8px; flex: 1; overflow: hidden; background: #444d56; border-radius: 4px; } +.meter-fill { width: 0; height: 100%; background: #54bd7d; transition: width 70ms linear; } + +.runtime-status { margin: 0; display: flex; gap: 24px; } +.runtime-status div { min-width: 68px; } +.runtime-status dt { color: #9fa8b1; font-size: 10px; text-transform: uppercase; } +.runtime-status dd { margin: 3px 0 0; font-size: 13px; font-weight: 700; font-variant-numeric: tabular-nums; } + +.conversation-section { margin-top: 22px; padding: 22px 0 10px; } +.section-heading { display: flex; justify-content: space-between; align-items: end; gap: 16px; padding: 0 2px 14px; border-bottom: 1px solid #cfd5da; } +.policy-label { color: #5e6871; font-size: 12px; } +.conversation { min-height: 260px; max-height: 440px; overflow-y: auto; padding: 18px 2px; } +.empty-state { color: #7a838b; font-size: 14px; } + +.turn { display: grid; grid-template-columns: 92px minmax(0, 1fr); gap: 12px; padding: 12px 0; border-bottom: 1px solid #e0e4e7; } +.turn-role { padding-top: 2px; color: #68717b; font-size: 11px; font-weight: 800; text-transform: uppercase; } +.turn-user .turn-role { color: #1565c0; } +.turn-assistant .turn-role { color: #a45117; } +.turn-text { min-width: 0; white-space: pre-wrap; overflow-wrap: anywhere; line-height: 1.55; } +.turn-live .turn-text::after { content: ""; display: inline-block; width: 6px; height: 14px; margin-left: 4px; background: #9098a0; vertical-align: -2px; animation: blink 1s step-end infinite; } +@keyframes blink { 50% { opacity: 0; } } + +.diagnostics { margin-top: 18px; border-top: 1px solid #bfc6cc; border-bottom: 1px solid #bfc6cc; background: #ffffff; } +.diagnostics summary { padding: 13px 14px; display: flex; justify-content: space-between; cursor: pointer; font-weight: 700; list-style: none; } +.diagnostics summary::-webkit-details-marker { display: none; } +.event-count { color: #6f7880; font-size: 12px; font-weight: 500; } +.log-toolbar { min-height: 40px; padding: 7px 12px; border-top: 1px solid #e0e4e7; display: flex; align-items: center; justify-content: space-between; gap: 12px; color: #69727a; font-size: 12px; } +.icon-button { padding: 5px 9px; color: #394149; background: #f4f6f7; border-color: #cdd3d8; } +.event-log { height: 200px; margin: 0; padding: 12px 14px; overflow: auto; background: #171b1f; color: #d9e0e6; font: 11px/1.55 ui-monospace, SFMono-Regular, Menlo, Consolas, monospace; white-space: pre-wrap; } +.log-error { color: #ff9f9f; } + +@media (max-width: 760px) { + .app-header { min-height: 88px; padding: 18px 20px; } + h1 { font-size: 20px; } + main { width: min(100% - 24px, 620px); margin-top: 16px; } + .call-bar { grid-template-columns: 1fr; gap: 15px; } + .call-actions button { flex: 1; } + .runtime-status { justify-content: space-between; gap: 10px; } + .section-heading { align-items: start; flex-direction: column; gap: 7px; } + .conversation { min-height: 230px; } + .turn { grid-template-columns: 72px minmax(0, 1fr); } +} + +@media (prefers-reduced-motion: reduce) { + *, *::before, *::after { animation-duration: .01ms !important; animation-iteration-count: 1 !important; } +} diff --git a/examples/online_serving/minicpmo/realtime_web/server.py b/examples/online_serving/minicpmo/realtime_web/server.py new file mode 100644 index 000000000..ee78fff74 --- /dev/null +++ b/examples/online_serving/minicpmo/realtime_web/server.py @@ -0,0 +1,186 @@ +"""Static host and same-origin Realtime proxy for the full-duplex demo.""" + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import contextlib +import hashlib +import json +import logging +from pathlib import Path +from urllib.parse import urlencode + +import uvicorn +import websockets +from fastapi import FastAPI, WebSocket, WebSocketDisconnect +from fastapi.responses import HTMLResponse, Response +from fastapi.staticfiles import StaticFiles + +logger = logging.getLogger(__name__) +APP_DIR = Path(__file__).parent / "app" +STATIC_DIR = APP_DIR / "static" + + +def _join_ws_url(base: str, path: str, query: str) -> str: + return base.rstrip("/") + path + (("?" + query) if query else "") + + +async def _pump_client_to_backend(client: WebSocket, backend) -> None: + try: + while True: + message = await client.receive() + if message["type"] == "websocket.disconnect": + await backend.close() + return + if message.get("text") is not None: + await backend.send(message["text"]) + elif message.get("bytes") is not None: + await backend.send(message["bytes"]) + except (WebSocketDisconnect, websockets.ConnectionClosed): + with contextlib.suppress(Exception): + await backend.close() + + +async def _pump_backend_to_client(client: WebSocket, backend) -> None: + try: + async for message in backend: + if isinstance(message, bytes): + await client.send_bytes(message) + else: + await client.send_text(message) + except (WebSocketDisconnect, websockets.ConnectionClosed): + return + except RuntimeError as exc: + if "websocket.send" in str(exc) and "websocket.close" in str(exc): + return + raise + + +def _expected_proxy_close(exc: BaseException) -> bool: + if isinstance(exc, (WebSocketDisconnect, websockets.ConnectionClosed, asyncio.CancelledError)): + return True + return isinstance(exc, RuntimeError) and "websocket.send" in str(exc) and "websocket.close" in str(exc) + + +def build_app( + *, + ws_backend: str = "ws://127.0.0.1:8099", + model: str = "openbmb/MiniCPM-o-4_5", + public_realtime_url: str | None = None, + ref_audio: str | None = None, +) -> FastAPI: + app = FastAPI(title="Experimental Full-Duplex Web Demo") + index_path = APP_DIR / "index.html" + app_version_hash = hashlib.sha256() + for asset_path in ( + STATIC_DIR / "app.js", + STATIC_DIR / "pcm_worklet.js", + STATIC_DIR / "playback_worklet.js", + ): + app_version_hash.update(asset_path.read_bytes()) + app_version = app_version_hash.hexdigest()[:12] + + ref_audio_uri: str | None = None + if ref_audio: + ref_path = Path(ref_audio) + if not ref_path.is_file(): + raise SystemExit(f"--ref-audio not found: {ref_audio}") + encoded = base64.b64encode(ref_path.read_bytes()).decode("ascii") + ref_audio_uri = f"data:audio/wav;base64,{encoded}" + + @app.get("/", response_class=HTMLResponse) + def index() -> HTMLResponse: + config = json.dumps( + { + "model": model, + "realtimePath": public_realtime_url or "v1/realtime", + "refAudio": ref_audio_uri, + "appVersion": app_version, + }, + ensure_ascii=True, + ) + html = ( + index_path.read_text(encoding="utf-8") + .replace("__FULL_DUPLEX_CONFIG__", config) + .replace("__FULL_DUPLEX_APP_VERSION__", app_version) + ) + return HTMLResponse(html, headers={"Cache-Control": "no-store"}) + + @app.get("/healthz") + def healthz() -> Response: + return Response(content="ok", media_type="text/plain") + + @app.websocket("/v1/realtime") + async def realtime_proxy(websocket: WebSocket) -> None: + await websocket.accept() + query = urlencode(websocket.query_params.multi_items()) + backend_url = _join_ws_url(ws_backend, "/v1/realtime", query) + logger.info("Proxying Realtime WebSocket to %s", backend_url) + try: + async with websockets.connect( + backend_url, + max_size=64 * 1024 * 1024, + ) as backend: + tasks = { + asyncio.create_task(_pump_client_to_backend(websocket, backend)), + asyncio.create_task(_pump_backend_to_client(websocket, backend)), + } + done, pending = await asyncio.wait( + tasks, + return_when=asyncio.FIRST_COMPLETED, + ) + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + for result in await asyncio.gather(*done, return_exceptions=True): + if isinstance(result, BaseException) and not _expected_proxy_close(result): + raise result + except (WebSocketDisconnect, websockets.ConnectionClosed): + return + except Exception: + logger.exception("Realtime WebSocket proxy failed") + with contextlib.suppress(Exception): + await websocket.close(code=1011) + + app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static") + return app + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=7862) + parser.add_argument("--ws-backend", default="ws://127.0.0.1:8099") + parser.add_argument( + "--public-realtime-url", + help="Browser-visible ws:// or wss:// Realtime URL; defaults to the same-origin proxy.", + ) + parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5") + parser.add_argument( + "--ref-audio", + required=True, + help=( + "Reference voice wav for TTS voice cloning, e.g. the " + "official MiniCPM-o-Demo assets/ref_audio/ref_minicpm_signature.wav" + ), + ) + args = parser.parse_args() + + logging.basicConfig(level=logging.INFO) + uvicorn.run( + build_app( + ws_backend=args.ws_backend, + model=args.model, + public_realtime_url=args.public_realtime_url, + ref_audio=args.ref_audio, + ), + host=args.host, + port=args.port, + log_level="info", + ) + + +if __name__ == "__main__": + main() diff --git a/tests/core/sched/test_sched_package_imports.py b/tests/core/sched/test_sched_package_imports.py new file mode 100644 index 000000000..bc4ec9d44 --- /dev/null +++ b/tests/core/sched/test_sched_package_imports.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_importing_scheduler_output_does_not_eagerly_import_scheduler_implementations() -> None: + package_init = Path(__file__).parents[3] / "vllm_omni/core/sched/__init__.py" + module = ast.parse(package_init.read_text()) + eager_modules = { + node.module + for node in ast.walk(module) + if isinstance(node, ast.ImportFrom) + and node.level == 1 + and node.module in {"omni_ar_scheduler", "omni_generation_scheduler"} + } + + assert eager_modules == set() + assert any(isinstance(node, ast.FunctionDef) and node.name == "__getattr__" for node in module.body) diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 9c6a86ed3..76d0cb550 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -677,6 +677,28 @@ def test_finish_requests_removes_zombies_from_chunk_waiting_deques(build_adapter assert "req-zombie" not in adapter.finished_requests +def test_finish_requests_releases_active_stream_slot(build_adapter): + adapter, _ = build_adapter(stage_id=1, max_num_seqs=1, active_stream_window=1) + aborted = _req("req-aborted", RequestStatus.RUNNING) + waiting = _req("req-waiting", RequestStatus.WAITING) + waiting_queue = DummyWaitingQueue([waiting]) + running_queue = [] + adapter._active_streams[aborted.request_id] = aborted + adapter._held_non_active.append(aborted) + + adapter.finish_requests( + [aborted.request_id], + RequestStatus.FINISHED_ABORTED, + {aborted.request_id: aborted}, + ) + adapter.process_pending_chunks(waiting_queue, running_queue) + + assert aborted.request_id not in adapter._active_streams + assert [request.request_id for request in adapter._held_non_active] == [] + assert list(adapter._active_streams) == [waiting.request_id] + assert waiting.status == RequestStatus.WAITING_FOR_CHUNK + + def test_restore_queues_skips_requests_missing_from_scheduler_requests(build_adapter): adapter, _ = build_adapter(stage_id=1) zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK) diff --git a/tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py b/tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py new file mode 100644 index 000000000..8b521c378 --- /dev/null +++ b/tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py @@ -0,0 +1,1673 @@ +"""Strict MiniCPM-o 4.5 Realtime duplex E2E scenario harness. + +This script is intentionally scenario-based instead of a generic chat client. +It has two explicit validation contracts: + +1. ``model-policy`` accepts either a model-owned listen decision or a complete + spoken response for each streamed user turn; +2. ``response-required`` uses a known fixture and requires a complete audio + response with an independent transcript for every requested turn. + +Explicit serving-side barge-in is intentionally not part of this smoke test. +MiniCPM-o native duplex currently exposes model-owned listen/speak switching, +not a separate model-level barge-in contract. + +Run only after a MiniCPM-o 4.5 vLLM-Omni server is up: + + python tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py \ + --url ws://localhost:8099/v1/realtime?duplex=1 \ + --model openbmb/MiniCPM-o-4_5 \ + --input-wav input_16k_mono_pcm16.wav \ + --validation-mode model-policy \ + --output-dir /tmp/minicpmo_duplex_demo +""" + +from __future__ import annotations + +import argparse +import array +import asyncio +import base64 +import hashlib +import inspect +import json +import sys +import time +import wave +from dataclasses import dataclass, field +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[3] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +try: + import websockets + from websockets.exceptions import ConnectionClosed +except ImportError as exc: # pragma: no cover - demo dependency. + raise SystemExit("Install websockets first: pip install websockets") from exc + +from vllm_omni.experimental.fullduplex.client import ( # noqa: E402 + PCM16_BYTES_PER_SAMPLE, + PCM16_SAMPLE_RATE, + RealtimeEventCollector, + build_realtime_url, + read_pcm16_wav, +) + +_url_with_model = build_realtime_url +_read_wav_pcm16 = read_pcm16_wav + + +def _ref_audio_data_url(path: str | None) -> str | None: + if path is None: + return None + ref_path = Path(path).expanduser() + return "data:audio/wav;base64," + base64.b64encode(ref_path.read_bytes()).decode("ascii") + + +@dataclass +class DemoState: + events: list[dict[str, object]] = field(default_factory=list) + timing_events: RealtimeEventCollector = field(default_factory=RealtimeEventCollector) + audio_deltas: list[bytes] = field(default_factory=list) + response_audio_deltas: dict[str, list[bytes]] = field(default_factory=dict) + response_ids: list[str] = field(default_factory=list) + assistant_item_ids: list[str] = field(default_factory=list) + done_count: int = 0 + cancelled_count: int = 0 + listen_count: int = 0 + model_listen_count: int = 0 + buffering_listen_count: int = 0 + model_speak_event_count: int = 0 + model_speak_delta_count: int = 0 + playback_ack_count: int = 0 + playback_history_committed_count: int = 0 + truncate_count: int = 0 + input_transcription_count: int = 0 + audio_marks_seen: bool = False + overlap_decisions: list[dict[str, object]] = field(default_factory=list) + output_sample_rate_hz: int = 24000 + input_commit_sent_at_s: list[float] = field(default_factory=list) + + def add(self, event: dict[str, object], *, received_at_s: float | None = None) -> None: + self.events.append(event) + self.timing_events.add(event, received_at_s=received_at_s) + event_type = event.get("type") + if event_type == "response.created": + response = event.get("response") + response_id = response.get("id") if isinstance(response, dict) else event.get("response_id") + if isinstance(response_id, str) and response_id not in self.response_ids: + self.response_ids.append(response_id) + elif event_type == "conversation.item.added": + item = event.get("item") + if isinstance(item, dict) and item.get("role") == "assistant": + item_id = item.get("id") + if isinstance(item_id, str) and item_id not in self.assistant_item_ids: + self.assistant_item_ids.append(item_id) + elif event_type == "response.audio.delta": + delta = event.get("delta") or event.get("audio") + if isinstance(delta, str) and delta: + try: + decoded = base64.b64decode(delta) + self.audio_deltas.append(decoded) + response_id = self._event_response_id(event) + if isinstance(response_id, str): + self.response_audio_deltas.setdefault(response_id, []).append(decoded) + except Exception: + pass + metadata = event.get("metadata") + if isinstance(metadata, dict): + if metadata.get("model_speak") is True: + self.model_speak_delta_count += 1 + if isinstance(metadata.get("audio_text_marks"), list): + self.audio_marks_seen = True + sample_rate_hz = event.get("sample_rate_hz") + if isinstance(sample_rate_hz, int) and sample_rate_hz > 0: + self.output_sample_rate_hz = sample_rate_hz + elif event_type == "response.done": + self.done_count += 1 + response = event.get("response") + if isinstance(response, dict) and response.get("status") == "cancelled": + self.cancelled_count += 1 + elif event_type in {"audio.cancelled", "input.cancelled"}: + self.cancelled_count += 1 + elif event_type == "response.listen": + self.listen_count += 1 + response = event.get("response") + metadata = response.get("metadata") if isinstance(response, dict) else None + if isinstance(metadata, dict) and metadata.get("model_listen") is True: + self.model_listen_count += 1 + if isinstance(metadata, dict) and metadata.get("buffering") is True: + self.buffering_listen_count += 1 + elif event_type == "response.speak": + self.model_speak_event_count += 1 + elif event_type == "overlap.decision": + self.overlap_decisions.append(event) + elif event_type == "playback.acknowledged": + self.playback_ack_count += 1 + payload = event.get("event") + if isinstance(payload, dict) and payload.get("history_committed") is True: + self.playback_history_committed_count += 1 + elif event_type == "conversation.item.truncated": + self.truncate_count += 1 + elif event_type == "conversation.item.input_audio_transcription.completed": + self.input_transcription_count += 1 + + def count(self, event_type: str) -> int: + return sum(1 for event in self.events if event.get("type") == event_type) + + def response_timing_summaries(self) -> dict[str, dict[str, object]]: + summaries: dict[str, dict[str, object]] = {} + for response_id in self.response_ids: + created_at_s = next( + ( + received_at_s + for event, received_at_s in zip( + self.timing_events.events, + self.timing_events.event_received_at_s, + strict=True, + ) + if event.get("type") == "response.created" and self._event_response_id(event) == response_id + ), + None, + ) + if created_at_s is None: + continue + input_committed_at_s = next( + ( + committed_at_s + for committed_at_s in reversed(self.input_commit_sent_at_s) + if committed_at_s <= created_at_s + ), + None, + ) + timing = self.timing_events.timing_summary( + after_s=created_at_s, + input_committed_at_s=input_committed_at_s, + response_id=response_id, + ) + timing["measurement_origin"] = { + "response": "response.created client receive", + "input": "input_audio_buffer.commit client send" if input_committed_at_s is not None else None, + } + summaries[response_id] = timing + return summaries + + def first_index(self, event_type: str, predicate=None) -> int | None: + for index, event in enumerate(self.events): + if event.get("type") != event_type: + continue + if predicate is not None and not predicate(event): + continue + return index + return None + + @staticmethod + def _event_response_id(event: dict[str, object]) -> str | None: + response_id = event.get("response_id") + if isinstance(response_id, str) and response_id: + return response_id + response = event.get("response") + if isinstance(response, dict): + response_id = response.get("id") + if isinstance(response_id, str) and response_id: + return response_id + return None + + @staticmethod + def _event_item_id(event: dict[str, object]) -> str | None: + item_id = event.get("item_id") + if isinstance(item_id, str) and item_id: + return item_id + item = event.get("item") + if isinstance(item, dict): + item_id = item.get("id") + if isinstance(item_id, str) and item_id: + return item_id + return None + + def first_response_lifecycle_indices(self) -> dict[str, int]: + response_created_index = self.first_index("response.created") + if response_created_index is None: + return {} + response_id = self._event_response_id(self.events[response_created_index]) + if not response_id: + return {} + item_id = f"item_{response_id}" + indices: dict[str, int] = {"response.created": response_created_index} + for event_type in ( + "conversation.item.added", + "response.output_item.added", + "response.content_part.added", + "response.speak", + "response.audio.delta", + "response.audio.done", + "response.content_part.done", + "response.output_item.done", + "response.done", + ): + index = self.first_index( + event_type, + lambda event, event_type=event_type: ( + self._event_item_id(event) == item_id + if event_type == "conversation.item.added" + else self._event_response_id(event) == response_id + ), + ) + if index is None and event_type not in { + "response.speak", + "response.audio.delta", + "response.audio.done", + }: + return {} + if index is not None: + indices[event_type] = index + return indices + + def event_order_ok(self, *, require_input_commit: bool = True) -> bool: + if not self.events or self.events[0].get("type") != "session.created": + return False + first_commit_index = self.first_index("input_audio_buffer.committed") + first_response_index = self.first_index("response.created") + if first_response_index is None: + return False + if require_input_commit and (first_commit_index is None or first_commit_index > first_response_index): + return False + indices_by_type = self.first_response_lifecycle_indices() + if not indices_by_type: + return False + required_types = [ + "response.created", + "conversation.item.added", + "response.output_item.added", + "response.content_part.added", + "response.content_part.done", + "response.output_item.done", + "response.done", + ] + if any(event_type not in indices_by_type for event_type in required_types): + return False + ordered_types = [ + "response.created", + "conversation.item.added", + "response.output_item.added", + "response.content_part.added", + "response.speak", + "response.audio.delta", + "response.audio.done", + "response.content_part.done", + "response.output_item.done", + "response.done", + ] + if "response.audio.delta" not in indices_by_type: + listen_index = self.first_index( + "response.listen", + lambda event: ( + isinstance(event.get("response"), dict) + and isinstance(event["response"].get("metadata"), dict) + and event["response"]["metadata"].get("model_listen") is True + ), + ) + return ( + listen_index is not None + and indices_by_type["response.created"] < listen_index < indices_by_type["response.done"] + ) + if any(event_type not in indices_by_type for event_type in ordered_types): + return False + indices = [indices_by_type[event_type] for event_type in ordered_types] + return indices == sorted(indices) + + def model_policy_event_order_ok( + self, + *, + expected_turns: int, + require_input_commit: bool = True, + ) -> bool: + if not self.events or self.events[0].get("type") != "session.created": + return False + commit_indices = [ + index for index, event in enumerate(self.events) if event.get("type") == "input_audio_buffer.committed" + ] + decision_indices = [ + index + for index, event in enumerate(self.events) + if event.get("type") == "response.created" + or ( + event.get("type") == "response.listen" + and isinstance(event.get("response"), dict) + and isinstance(event["response"].get("metadata"), dict) + and event["response"]["metadata"].get("model_listen") is True + ) + ] + first_input_index = self.first_index("input_audio_buffer.speech_started") + if first_input_index is None: + first_input_index = commit_indices[0] if commit_indices else None + if first_input_index is None and not require_input_commit: + first_input_index = 0 + if ( + first_input_index is None + or (require_input_commit and len(commit_indices) < expected_turns) + or len(decision_indices) < expected_turns + ): + return False + return decision_indices[0] > first_input_index + + def model_speak_before_audio_ok(self) -> bool: + speak_index = self.first_index("response.speak") + audio_index = self.first_index("response.audio.delta") + return speak_index is not None and audio_index is not None and speak_index < audio_index + + def response_done(self, response_id: str | None) -> bool: + if not response_id: + return False + return any( + event.get("type") == "response.done" and self._event_response_id(event) == response_id + for event in self.events + ) + + def response_audio_delta_count(self, response_id: str | None) -> int: + if not response_id: + return 0 + return sum( + 1 + for event in self.events + if event.get("type") == "response.audio.delta" and self._event_response_id(event) == response_id + ) + + def response_playback_sent_ms(self, response_id: str | None) -> int: + if not response_id: + return 0 + for event in reversed(self.events): + if event.get("type") != "response.done" or self._event_response_id(event) != response_id: + continue + response = event.get("response") + metadata = response.get("metadata") if isinstance(response, dict) else None + playback = metadata.get("playback") if isinstance(metadata, dict) else event.get("playback") + sent_ms = playback.get("sent_ms") if isinstance(playback, dict) else None + if isinstance(sent_ms, int | float): + return max(0, int(sent_ms)) + return 0 + + def response_playback_history_committed(self, response_id: str) -> bool: + item_id = f"item_{response_id}" + for event in self.events: + if event.get("type") != "playback.acknowledged": + continue + payload = event.get("event") + if not isinstance(payload, dict): + continue + if payload.get("item_id") == item_id and payload.get("history_committed") is True: + return True + return False + + def response_transcript_delta(self, response_id: str) -> str: + return "".join( + str(event.get("delta", "")) + for event in self.events + if event.get("type") == "response.audio_transcript.delta" and self._event_response_id(event) == response_id + ) + + def response_transcript_done(self, response_id: str) -> list[str]: + return [ + str(event.get("transcript", "")) + for event in self.events + if event.get("type") == "response.audio_transcript.done" and self._event_response_id(event) == response_id + ] + + def completed_response_ids(self) -> list[str]: + response_ids: list[str] = [] + for event in self.events: + if event.get("type") != "response.done": + continue + response_id = self._event_response_id(event) + if isinstance(response_id, str): + response_ids.append(response_id) + return response_ids + + def stale_audio_delta_count(self) -> int: + cancelled_epochs_by_index: list[tuple[int, int]] = [] + for index, event in enumerate(self.events): + if event.get("type") != "response.done": + continue + response = event.get("response") + if not isinstance(response, dict) or response.get("status") != "cancelled": + continue + metadata = response.get("metadata") + if not isinstance(metadata, dict): + continue + cancelled_epoch = metadata.get("cancelled_epoch") + if isinstance(cancelled_epoch, int): + cancelled_epochs_by_index.append((index, cancelled_epoch)) + if not cancelled_epochs_by_index: + return 0 + stale = 0 + for index, event in enumerate(self.events): + if event.get("type") != "response.audio.delta": + continue + metadata = event.get("metadata") + if not isinstance(metadata, dict): + continue + event_epoch = metadata.get("epoch") + for cancel_index, cancelled_epoch in cancelled_epochs_by_index: + if index > cancel_index and event_epoch == cancelled_epoch: + stale += 1 + break + return stale + + +def _session_update_event(args: argparse.Namespace) -> dict[str, object]: + session_payload: dict[str, object] = { + "model": args.model, + "modalities": ["audio", "text"], + "input_audio_format": "pcm16", + "output_audio_format": args.output_audio_format, + "turn_detection": None, + "overlap_policy": "listen_only", + "overlap_short_ack_ms": args.short_ack_ms, + "playback_commit_policy": "ack_only", + "extra_body": { + "auto_response": True, + "minicpmo45_native_duplex": True, + "force_listen_count": 0, + }, + } + temperature = getattr(args, "temperature", None) + if temperature is None and getattr(args, "validation_mode", None) == "response-required": + temperature = 0.0 + if temperature is not None: + session_payload["temperature"] = temperature + ref_audio = _ref_audio_data_url(getattr(args, "ref_audio", None)) + if ref_audio is not None: + session_payload["ref_audio"] = ref_audio + event: 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 + + +def _turn_input_paths(primary: Path, additional: list[str], *, turns: int) -> list[Path]: + turn_count = max(1, turns) + if not additional: + return [primary] * turn_count + if len(additional) != turn_count - 1: + raise ValueError( + f"provide exactly one --turn-input-wav for each turn after the first " + f"(expected {turn_count - 1}, got {len(additional)})" + ) + return [primary, *(Path(path) for path in additional)] + + +def _turn_inputs_are_distinct(paths: list[Path], pcm16_inputs: list[bytes]) -> bool: + if len(paths) != len(pcm16_inputs): + return False + distinct_paths = len({str(path.resolve()) for path in paths}) == len(paths) + distinct_audio = len({hashlib.sha256(pcm16).digest() for pcm16 in pcm16_inputs}) == len(pcm16_inputs) + return distinct_paths and distinct_audio + + +def _turn_durations( + explicit: list[int], + *, + turns: int, + first_turn_ms: int, +) -> list[int | None]: + turn_count = max(1, turns) + if not explicit: + return [first_turn_ms, *([min(first_turn_ms, 1200)] * (turn_count - 1))] + if len(explicit) != turn_count: + raise ValueError( + f"provide exactly one --turn-duration-ms for each turn (expected {turn_count}, got {len(explicit)})" + ) + if any(duration_ms < 0 for duration_ms in explicit): + raise ValueError("--turn-duration-ms values must be non-negative") + return [None if duration_ms == 0 else duration_ms for duration_ms in explicit] + + +def _turn_transcripts(first: str, *, turns: int) -> list[str]: + turn_count = max(1, turns) + transcripts = [first, "继续", "再说一次"][:turn_count] + transcripts.extend(f"turn-{turn_index + 1}" for turn_index in range(len(transcripts), turn_count)) + return transcripts + + +def _write_wav(path: Path, pcm_bytes: bytes, *, sample_rate_hz: int) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with wave.open(str(path), "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(PCM16_BYTES_PER_SAMPLE) + wf.setframerate(sample_rate_hz) + wf.writeframes(pcm_bytes) + + +def _write_demo_artifacts(state: DemoState, output_dir: Path, *, output_audio_format: str) -> None: + output_dir.mkdir(parents=True, exist_ok=True) + if state.audio_deltas and output_audio_format == "pcm16": + _write_wav( + output_dir / "joined_audio_deltas.wav", + b"".join(state.audio_deltas), + sample_rate_hz=state.output_sample_rate_hz, + ) + elif state.audio_deltas: + (output_dir / "joined_audio_deltas.bin").write_bytes(b"".join(state.audio_deltas)) + for index, response_id in enumerate(state.response_ids, start=1): + response_audio = state.response_audio_deltas.get(response_id, []) + if not response_audio: + continue + payload = b"".join(response_audio) + if output_audio_format == "pcm16": + _write_wav( + output_dir / f"response_{index:02d}.wav", + payload, + sample_rate_hz=state.output_sample_rate_hz, + ) + else: + (output_dir / f"response_{index:02d}.bin").write_bytes(payload) + (output_dir / "events.jsonl").write_text( + "\n".join(json.dumps(event, ensure_ascii=False) for event in state.events) + "\n", + encoding="utf-8", + ) + + +def _pcm16_silence(duration_ms: int) -> bytes: + samples = PCM16_SAMPLE_RATE * max(0, duration_ms) // 1000 + return b"\x00\x00" * samples + + +def _pcm16_slice(pcm16: bytes, duration_ms: int) -> bytes: + byte_count = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * max(1, duration_ms) // 1000 + return pcm16[: min(len(pcm16), byte_count)] + + +def _pcm16_active_slice(pcm16: bytes, duration_ms: int) -> bytes: + byte_count = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * max(1, duration_ms) // 1000 + byte_count = min(len(pcm16), max(PCM16_BYTES_PER_SAMPLE, byte_count)) + byte_count -= byte_count % PCM16_BYTES_PER_SAMPLE + if byte_count <= 0: + return _pcm16_slice(pcm16, duration_ms) + step = max(PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * 20 // 1000, PCM16_BYTES_PER_SAMPLE) + step -= step % PCM16_BYTES_PER_SAMPLE + best_offset = 0 + best_energy = -1.0 + for offset in range(0, max(1, len(pcm16) - byte_count + 1), max(PCM16_BYTES_PER_SAMPLE, step)): + chunk = pcm16[offset : offset + byte_count] + samples = array.array("h") + samples.frombytes(chunk) + if not samples: + continue + energy = sum(abs(sample) for sample in samples) / len(samples) + if energy > best_energy: + best_energy = energy + best_offset = offset + return pcm16[best_offset : best_offset + byte_count] + + +def _select_turn_audio(pcm16: bytes, duration_ms: int | None) -> bytes: + if duration_ms is None: + return pcm16 + return _pcm16_active_slice(pcm16, duration_ms) + + +def _canonical_transcript(text: str) -> str: + return "".join(text.split()) + + +def _reuses_previous_turn_tail(previous: str, current: str) -> bool: + if not previous or not current: + return False + if len(previous) >= 4 and previous in current: + return True + max_overlap = min(len(previous), len(current)) + return any(previous.endswith(current[:overlap]) for overlap in range(max_overlap, 2, -1)) + + +def _has_terminal_punctuation(text: str) -> bool: + stripped = text.rstrip("\"'”’))]} ") + return bool(stripped) and (len(stripped) <= 4 or stripped[-1] in "。!?!?…") + + +def _all_audio_responses_have_transcript(state: DemoState, response_ids: list[str]) -> bool: + return all( + state.response_audio_delta_count(response_id) == 0 + or bool(_canonical_transcript(state.response_transcript_delta(response_id))) + for response_id in response_ids + ) + + +def _all_response_playback_history_committed(state: DemoState, response_ids: list[str]) -> bool: + return all( + state.response_playback_history_committed(response_id) + for response_id in response_ids + if state.response_playback_sent_ms(response_id) > 0 + ) + + +def _response_cardinality_ok( + completed_response_ids: list[str], + *, + expected_turns: int, + validation_mode: str, +) -> bool: + if validation_mode == "response-required": + return len(completed_response_ids) == expected_turns + return True + + +def _requires_cross_turn_independence( + *, + validation_mode: str, + distinct_inputs_required: bool, +) -> bool: + return distinct_inputs_required and validation_mode == "response-required" + + +def _input_transcription_ok(count: int, *, transcript_hints_enabled: bool) -> bool: + return count > 0 if transcript_hints_enabled else True + + +def _unexpected_error_events(state: DemoState) -> list[dict[str, object]]: + return [event for event in state.events if event.get("type") == "error"] + + +def _evaluate_response_speak_contract(state: DemoState) -> dict[str, object]: + speak_counts: dict[str, int] = {} + invalid_response_speak_count = 0 + text_bearing_response_speak_count = 0 + + for event in state.events: + if event.get("type") != "response.speak": + continue + response_id = state._event_response_id(event) + if not response_id or response_id not in state.response_ids: + invalid_response_speak_count += 1 + else: + speak_counts[response_id] = speak_counts.get(response_id, 0) + 1 + if "text" in event: + text_bearing_response_speak_count += 1 + + duplicate_response_speak_ids = sorted(response_id for response_id, count in speak_counts.items() if count > 1) + return { + "response_speak_contract_ok": not duplicate_response_speak_ids + and invalid_response_speak_count == 0 + and text_bearing_response_speak_count == 0, + "duplicate_response_speak_ids": duplicate_response_speak_ids, + "invalid_response_speak_count": invalid_response_speak_count, + "text_bearing_response_speak_count": text_bearing_response_speak_count, + } + + +def _evaluate_transcript_integrity( + state: DemoState, + response_ids: list[str], + *, + expected_empty_response_ids: set[str], + require_cross_turn_independence: bool, + require_terminal_punctuation: bool = False, +) -> dict[str, object]: + details: list[dict[str, object]] = [] + transcripts: list[str] = [] + transcript_delta_done_ok = True + empty_turns_ok = True + nonempty_audio_has_transcript_ok = True + terminal_punctuation_ok = True + + for response_id in response_ids: + transcript = state.response_transcript_delta(response_id) + done_transcripts = state.response_transcript_done(response_id) + canonical_transcript = _canonical_transcript(transcript) + response_delta_done = ( + len(done_transcripts) == 1 and _canonical_transcript(done_transcripts[0]) == canonical_transcript + ) or (not done_transcripts and not canonical_transcript) + expected_empty = response_id in expected_empty_response_ids + audio_delta_count = state.response_audio_delta_count(response_id) + response_empty_ok = not expected_empty or (not canonical_transcript and audio_delta_count == 0) + response_audio_has_transcript = expected_empty or audio_delta_count == 0 or bool(canonical_transcript) + response_terminal_punctuation_ok = ( + expected_empty + or not canonical_transcript + or not require_terminal_punctuation + or _has_terminal_punctuation(canonical_transcript) + ) + transcript_delta_done_ok = transcript_delta_done_ok and response_delta_done + empty_turns_ok = empty_turns_ok and response_empty_ok + nonempty_audio_has_transcript_ok = nonempty_audio_has_transcript_ok and response_audio_has_transcript + terminal_punctuation_ok = terminal_punctuation_ok and response_terminal_punctuation_ok + transcripts.append(canonical_transcript) + details.append( + { + "response_id": response_id, + "transcript": transcript, + "delta_done_ok": response_delta_done, + "expected_empty": expected_empty, + "empty_ok": response_empty_ok, + "audio_delta_count": audio_delta_count, + "audio_has_transcript": response_audio_has_transcript, + "terminal_punctuation_ok": response_terminal_punctuation_ok, + } + ) + + cross_turn_independent_ok = True + if require_cross_turn_independence: + for index, current in enumerate(transcripts): + for previous in transcripts[:index]: + if _reuses_previous_turn_tail(previous, current): + cross_turn_independent_ok = False + break + if not cross_turn_independent_ok: + break + + return { + "transcript_delta_done_ok": transcript_delta_done_ok, + "cross_turn_independent_ok": cross_turn_independent_ok, + "empty_turns_ok": empty_turns_ok, + "nonempty_audio_has_transcript_ok": nonempty_audio_has_transcript_ok, + "terminal_punctuation_ok": terminal_punctuation_ok, + "transcript_integrity": details, + } + + +async def _reader(ws, state: DemoState, stop: asyncio.Event) -> None: + try: + while not stop.is_set(): + raw = await ws.recv() + if not isinstance(raw, str): + continue + event = json.loads(raw) + if isinstance(event, dict): + state.add(event) + except ConnectionClosed: + return + + +async def _send_pcm16( + ws, + pcm16: bytes, + *, + chunk_ms: int, + realtime_delay: bool, + hints: dict[str, object] | None = None, + first_chunk_hints: dict[str, object] | None = None, + on_model_unit_ready=None, + frame_b64: str | None = None, +) -> None: + hints = hints or {} + first_chunk_hints = first_chunk_hints or {} + chunk_bytes = max(PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * chunk_ms // 1000, PCM16_BYTES_PER_SAMPLE) + model_unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE + audio_ms = 0 + next_model_unit_bytes = model_unit_bytes + frames_sent = 0 + for offset in range(0, len(pcm16), chunk_bytes): + chunk = pcm16[offset : offset + chunk_bytes] + duration_ms = int(len(chunk) / (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE) * 1000) + audio_ms += duration_ms + chunk_hints = dict(hints) + if offset == 0: + chunk_hints.update(first_chunk_hints) + if frame_b64 is not None and audio_ms > frames_sent * 1000: + # Omni duplex cadence: one camera frame per 1 s of audio. + chunk_hints["video_frames"] = [frame_b64] + frames_sent += 1 + await ws.send( + json.dumps( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(chunk).decode("ascii"), + "input_audio_format": "pcm16", + "sample_rate_hz": PCM16_SAMPLE_RATE, + "duration_ms": duration_ms, + "audio_end_ms": audio_ms, + **chunk_hints, + } + ) + ) + while offset + len(chunk) >= next_model_unit_bytes: + next_model_unit_bytes += model_unit_bytes + if on_model_unit_ready is None: + continue + should_continue = on_model_unit_ready() + if inspect.isawaitable(should_continue): + should_continue = await should_continue + if should_continue is False: + return + if realtime_delay: + await asyncio.sleep(duration_ms / 1000) + + +async def _wait_for( + state: DemoState, + predicate, + *, + timeout_s: float, + label: str, +) -> None: + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if predicate(): + return + await asyncio.sleep(0.02) + raise TimeoutError(f"Timed out waiting for {label}") + + +async def _wait_for_model_policy_outcome( + state: DemoState, + *, + before_created: int, + before_model_listen: int, + timeout_s: float, + listen_settle_s: float, + label: str, +) -> str: + deadline = time.monotonic() + timeout_s + listen_deadline: float | None = None + while time.monotonic() < deadline: + if state.count("response.created") > before_created: + return "speak" + if state.model_listen_count > before_model_listen: + if listen_deadline is None: + listen_deadline = time.monotonic() + max(0.0, listen_settle_s) + if time.monotonic() >= listen_deadline: + return "listen" + await asyncio.sleep(0.02) + raise TimeoutError(f"Timed out waiting for {label}") + + +async def _send_clean_turn( + ws, + state: DemoState, + pcm16: bytes, + *, + transcript: str, + duration_ms: int | None, + chunk_ms: int, + timeout_s: float, + require_audio: bool, + validation_mode: str, + send_transcript_hint: bool = True, + realtime_input: bool = False, + model_policy_settle_s: float = 2.0, + commit_input: bool = True, + frame_b64: str | None = None, +) -> tuple[str | None, str]: + before_created = state.count("response.created") + before_model_listen = state.model_listen_count + + def eligible_response_id() -> str | None: + for response_id in state.response_ids[before_created:]: + if not state.response_done(response_id): + continue + if require_audio and state.response_audio_delta_count(response_id) == 0: + continue + if require_audio and not _canonical_transcript(state.response_transcript_delta(response_id)): + continue + return response_id + return None + + await _send_pcm16( + ws, + _select_turn_audio(pcm16, duration_ms), + chunk_ms=chunk_ms, + realtime_delay=realtime_input, + hints={"transcript": transcript} if send_transcript_hint else {}, + frame_b64=frame_b64, + ) + if commit_input: + state.input_commit_sent_at_s.append(time.monotonic()) + await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True})) + if validation_mode == "model-policy": + outcome = await _wait_for_model_policy_outcome( + state, + before_created=before_created, + before_model_listen=before_model_listen, + timeout_s=timeout_s, + listen_settle_s=model_policy_settle_s, + label=f"{transcript} model speak/listen decision", + ) + if outcome == "listen": + return None, "listen" + else: + await _wait_for( + state, + lambda: state.count("response.created") > before_created, + timeout_s=timeout_s, + label=f"{transcript} response.created", + ) + await _wait_for( + state, + lambda: eligible_response_id() is not None, + timeout_s=timeout_s, + label=f"{transcript} completed response", + ) + response_id = ( + eligible_response_id() + if validation_mode == "response-required" + else (state.response_ids[before_created] if len(state.response_ids) > before_created else None) + ) + if require_audio: + await _wait_for( + state, + lambda: state.response_audio_delta_count(response_id) > 0, + timeout_s=timeout_s, + label=f"{transcript} response.audio.delta", + ) + await _wait_for( + state, + lambda: state.response_done(response_id), + timeout_s=timeout_s, + label=f"{transcript} response.done", + ) + if not isinstance(response_id, str): + raise TimeoutError(f"Missing response id for {transcript}") + await _ack_response_playback( + ws, + state, + response_id, + timeout_s=timeout_s, + label=transcript, + ) + return response_id, "speak" + + +async def _ack_response_playback( + ws, + state: DemoState, + response_id: str, + *, + timeout_s: float, + label: str, +) -> None: + played_ms = state.response_playback_sent_ms(response_id) + if played_ms <= 0: + return + before_ack = state.playback_ack_count + await ws.send( + json.dumps( + { + "type": "playback.ack", + "response_id": response_id, + "item_id": f"item_{response_id}", + "played_ms": played_ms, + "committed_ms": played_ms, + } + ) + ) + await _wait_for( + state, + lambda: state.playback_ack_count > before_ack and state.response_playback_history_committed(response_id), + timeout_s=timeout_s, + label=f"{label} playback history committed", + ) + + +async def _ack_all_completed_response_playback( + ws, + state: DemoState, + *, + timeout_s: float, +) -> None: + for response_id in state.completed_response_ids(): + if state.response_playback_sent_ms(response_id) <= 0: + continue + if state.response_playback_history_committed(response_id): + continue + await _ack_response_playback( + ws, + state, + response_id, + timeout_s=timeout_s, + label=response_id, + ) + + +async def _drain_model_policy_responses( + ws, + state: DemoState, + *, + timeout_s: float, + settle_s: float, +) -> None: + deadline = time.monotonic() + timeout_s + settle_s = max(0.0, settle_s) + observed_created_count = state.count("response.created") + quiet_since: float | None = None + + while True: + remaining_s = deadline - time.monotonic() + active_response_ids = [ + response_id for response_id in state.response_ids if not state.response_done(response_id) + ] + if remaining_s <= 0: + raise TimeoutError(f"Timed out draining model-policy responses; active response ids: {active_response_ids}") + + await _ack_all_completed_response_playback( + ws, + state, + timeout_s=remaining_s, + ) + + now = time.monotonic() + created_count = state.count("response.created") + if created_count != observed_created_count: + observed_created_count = created_count + quiet_since = None + + active_response_ids = [ + response_id for response_id in state.response_ids if not state.response_done(response_id) + ] + if active_response_ids: + quiet_since = None + else: + if quiet_since is None: + quiet_since = now + if now - quiet_since >= settle_s: + return + + await asyncio.sleep(min(0.02, max(0.0, deadline - time.monotonic()))) + + +def _event_index_for_response(state: DemoState, event_type: str, response_id: str) -> int | None: + return state.first_index( + event_type, + lambda event: state._event_response_id(event) == response_id, + ) + + +def _nth_event_index(state: DemoState, event_type: str, occurrence: int) -> int | None: + seen = 0 + for index, event in enumerate(state.events): + if event.get("type") != event_type: + continue + seen += 1 + if seen == occurrence: + return index + return None + + +def _continuous_overlap_terminal_is_outcome( + state: DemoState, + *, + validation_mode: str, + model_unit_ready_while_active: bool, + before_created: int, + before_model_listen: int, +) -> bool: + """Accept current-response termination as a model-policy overlap outcome.""" + return ( + validation_mode == "model-policy" + and model_unit_ready_while_active + and state.count("response.created") == before_created + and state.model_listen_count == before_model_listen + ) + + +async def _send_listen_only_overlap_pair( + ws, + state: DemoState, + first_pcm16: bytes, + second_pcm16: bytes, + *, + transcripts: tuple[str, str], + durations_ms: tuple[int | None, int | None], + chunk_ms: int, + timeout_s: float, + send_transcript_hint: bool, + realtime_input: bool, + model_policy_settle_s: float, + validation_mode: str, + silence_ms: int, + frame_b64: str | None = None, +) -> tuple[list[str | None], list[str], bool]: + if not realtime_input: + raise ValueError("listen-only-overlap requires --realtime-input") + + before_first_created = state.count("response.created") + observed_first_created = before_first_created + observed_first_listen = state.model_listen_count + + async def continue_until_first_speak() -> bool: + nonlocal observed_first_created, observed_first_listen + outcome = await _wait_for_model_policy_outcome( + state, + before_created=observed_first_created, + before_model_listen=observed_first_listen, + timeout_s=timeout_s, + listen_settle_s=model_policy_settle_s, + label=f"{transcripts[0]} model-unit speak/listen decision", + ) + observed_first_created = state.count("response.created") + observed_first_listen = state.model_listen_count + return outcome != "speak" + + first_input = _select_turn_audio(first_pcm16, durations_ms[0]) + _pcm16_silence(silence_ms) + await _send_pcm16( + ws, + first_input, + chunk_ms=chunk_ms, + realtime_delay=True, + hints={"transcript": transcripts[0]} if send_transcript_hint else {}, + on_model_unit_ready=continue_until_first_speak, + frame_b64=frame_b64, + ) + state.input_commit_sent_at_s.append(time.monotonic()) + await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True})) + await _wait_for( + state, + lambda: state.count("response.created") > before_first_created, + timeout_s=timeout_s, + label=f"{transcripts[0]} response.created", + ) + first_response_id = state.response_ids[before_first_created] + await _wait_for( + state, + lambda: state.response_audio_delta_count(first_response_id) > 0, + timeout_s=timeout_s, + label=f"{transcripts[0]} response.audio.delta", + ) + if state.response_done(first_response_id): + raise RuntimeError("first response completed before overlap input could start") + + before_second_created = state.count("response.created") + before_second_listen = state.model_listen_count + overlap_started_while_active = not state.response_done(first_response_id) + model_unit_ready_while_active = False + + def record_model_unit_ready() -> None: + nonlocal model_unit_ready_while_active + model_unit_ready_while_active = not state.response_done(first_response_id) + + await _send_pcm16( + ws, + _select_turn_audio(second_pcm16, durations_ms[1]), + chunk_ms=chunk_ms, + realtime_delay=True, + hints={"transcript": transcripts[1]} if send_transcript_hint else {}, + on_model_unit_ready=record_model_unit_ready, + frame_b64=frame_b64, + ) + state.input_commit_sent_at_s.append(time.monotonic()) + await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True})) + + await _wait_for( + state, + lambda: state.response_done(first_response_id), + timeout_s=timeout_s, + label=f"{transcripts[0]} response.done", + ) + await _ack_response_playback( + ws, + state, + first_response_id, + timeout_s=timeout_s, + label=transcripts[0], + ) + + if _continuous_overlap_terminal_is_outcome( + state, + validation_mode=validation_mode, + model_unit_ready_while_active=model_unit_ready_while_active, + before_created=before_second_created, + before_model_listen=before_second_listen, + ): + second_outcome = "continuous_terminal" + else: + second_outcome = await _wait_for_model_policy_outcome( + state, + before_created=before_second_created, + before_model_listen=before_second_listen, + timeout_s=timeout_s, + listen_settle_s=model_policy_settle_s, + label=f"{transcripts[1]} model speak/listen decision", + ) + second_response_id: str | None = None + if second_outcome == "speak": + second_response_id = state.response_ids[before_second_created] + await _wait_for( + state, + lambda: state.response_done(second_response_id), + timeout_s=timeout_s, + label=f"{transcripts[1]} response.done", + ) + await _ack_response_playback( + ws, + state, + second_response_id, + timeout_s=timeout_s, + label=transcripts[1], + ) + + second_speech_index = _nth_event_index(state, "input_audio_buffer.speech_started", 2) + first_done_index = _event_index_for_response(state, "response.done", first_response_id) + overlap_ok = ( + overlap_started_while_active + and model_unit_ready_while_active + and second_speech_index is not None + and first_done_index is not None + and second_speech_index < first_done_index + ) + return [first_response_id, second_response_id], ["speak", second_outcome], overlap_ok + + +async def run_demo(args: argparse.Namespace) -> dict[str, object]: + demo_frame_b64: str | None = None + if getattr(args, "frame_image", None): + demo_frame_b64 = base64.b64encode(Path(args.frame_image).read_bytes()).decode("ascii") + turn_input_paths = _turn_input_paths( + Path(args.input_wav), + list(getattr(args, "turn_input_wav", []) or []), + turns=args.turns, + ) + turn_pcm16 = [_read_wav_pcm16(path) for path in turn_input_paths] + if any(not pcm16 for pcm16 in turn_pcm16): + raise ValueError("input WAV has no audio") + turn_durations = _turn_durations( + list(getattr(args, "turn_duration_ms", []) or []), + turns=args.turns, + first_turn_ms=args.first_turn_ms, + ) + expected_empty_turns = set(getattr(args, "expect_empty_turn", []) or []) + invalid_empty_turns = sorted( + turn_number for turn_number in expected_empty_turns if turn_number < 1 or turn_number > max(1, args.turns) + ) + if invalid_empty_turns: + raise ValueError( + f"--expect-empty-turn values are 1-based and must refer to an existing turn: {invalid_empty_turns}" + ) + distinct_turn_inputs = _turn_inputs_are_distinct(turn_input_paths, turn_pcm16) + if getattr(args, "require_distinct_inputs", False) and not distinct_turn_inputs: + raise ValueError("--require-distinct-inputs requires a different WAV path and audio payload for every turn") + url = _url_with_model( + 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() + output_dir = Path(args.output_dir) + turn_response_ids: list[str | None] = [] + turn_outcomes: list[str] = [] + validation_mode = getattr(args, "validation_mode", "response-required") + scenario = getattr(args, "scenario", "sequential") + transcript_hints_enabled = not getattr(args, "omit_transcript_hints", False) + realtime_input = getattr(args, "realtime_input", False) + continuous_input = getattr(args, "continuous_input", False) + listen_only_overlap_ok = scenario != "listen-only-overlap" + if scenario == "listen-only-overlap" and args.turns < 2: + raise ValueError("listen-only-overlap requires at least two turns") + + async with websockets.connect(url, max_size=64 * 1024 * 1024) as ws: + reader = asyncio.create_task(_reader(ws, state, stop)) + try: + await ws.send(json.dumps(_session_update_event(args))) + await _wait_for(state, lambda: state.count("session.created") > 0, timeout_s=20, label="session.created") + + turn_transcripts = _turn_transcripts(args.first_turn_transcript, turns=args.turns) + turn_specs = list(zip(turn_transcripts, turn_durations, strict=True)) + if scenario == "listen-only-overlap": + turn_response_ids, turn_outcomes, listen_only_overlap_ok = await _send_listen_only_overlap_pair( + ws, + state, + turn_pcm16[0], + turn_pcm16[1], + transcripts=(turn_specs[0][0], turn_specs[1][0]), + durations_ms=(turn_specs[0][1], turn_specs[1][1]), + chunk_ms=args.chunk_ms, + timeout_s=args.timeout_s, + send_transcript_hint=transcript_hints_enabled, + realtime_input=realtime_input, + model_policy_settle_s=max(0.0, args.model_policy_settle_ms / 1000), + validation_mode=validation_mode, + silence_ms=max(0, args.silence_ms), + frame_b64=demo_frame_b64, + ) + for turn_index in range(2, len(turn_specs)): + transcript, duration_ms = turn_specs[turn_index] + response_id, outcome = await _send_clean_turn( + ws, + state, + turn_pcm16[turn_index], + transcript=transcript, + duration_ms=duration_ms, + chunk_ms=args.chunk_ms, + timeout_s=args.timeout_s, + require_audio=args.require_audio and (turn_index + 1) not in expected_empty_turns, + validation_mode=validation_mode, + send_transcript_hint=transcript_hints_enabled, + realtime_input=realtime_input, + model_policy_settle_s=max(0.0, args.model_policy_settle_ms / 1000), + commit_input=not continuous_input, + frame_b64=demo_frame_b64, + ) + turn_response_ids.append(response_id) + turn_outcomes.append(outcome) + else: + for turn_index, (transcript, duration_ms) in enumerate(turn_specs): + response_id, outcome = await _send_clean_turn( + ws, + state, + turn_pcm16[turn_index], + transcript=transcript, + duration_ms=duration_ms, + chunk_ms=args.chunk_ms, + timeout_s=args.timeout_s, + require_audio=args.require_audio and (turn_index + 1) not in expected_empty_turns, + validation_mode=validation_mode, + send_transcript_hint=transcript_hints_enabled, + realtime_input=realtime_input, + model_policy_settle_s=max(0.0, args.model_policy_settle_ms / 1000), + commit_input=not continuous_input, + frame_b64=demo_frame_b64, + ) + turn_response_ids.append(response_id) + turn_outcomes.append(outcome) + + if validation_mode == "model-policy": + await _drain_model_policy_responses( + ws, + state, + timeout_s=args.timeout_s, + settle_s=max(0.0, args.model_policy_settle_ms / 1000), + ) + await _ack_all_completed_response_playback( + ws, + state, + timeout_s=args.timeout_s, + ) + await ws.send(json.dumps({"type": "session.close"})) + await _wait_for(state, lambda: state.count("session.closed") > 0, timeout_s=20, label="session.closed") + finally: + stop.set() + reader.cancel() + try: + await reader + except asyncio.CancelledError: + pass + _write_demo_artifacts(state, output_dir, output_audio_format=args.output_audio_format) + + overlap_barge_in = any(decision.get("action") == "barge_in" for decision in state.overlap_decisions) + expected_turns = max(1, args.turns) + event_order_ok = ( + state.model_policy_event_order_ok( + expected_turns=expected_turns, + require_input_commit=not continuous_input, + ) + if validation_mode == "model-policy" + else state.event_order_ok(require_input_commit=not continuous_input) + ) + input_transcription_ok = _input_transcription_ok( + state.input_transcription_count, + transcript_hints_enabled=transcript_hints_enabled and not continuous_input, + ) + model_speak_event_ok = state.model_speak_before_audio_ok() + realtime_audio_lifecycle_ok = state.count("response.audio.delta") > 0 and state.count("response.audio.done") > 0 + completed_response_ids = state.completed_response_ids() + observed_turn_response_ids = [response_id for response_id in turn_response_ids if isinstance(response_id, str)] + expected_empty_response_ids = { + response_id + for turn_number in expected_empty_turns + if turn_number <= len(turn_response_ids) + and isinstance((response_id := turn_response_ids[turn_number - 1]), str) + } + transcript_integrity = _evaluate_transcript_integrity( + state, + observed_turn_response_ids, + expected_empty_response_ids=expected_empty_response_ids, + require_cross_turn_independence=_requires_cross_turn_independence( + validation_mode=validation_mode, + distinct_inputs_required=getattr(args, "require_distinct_inputs", False), + ), + require_terminal_punctuation=getattr(args, "require_distinct_inputs", False), + ) + response_speak_contract = _evaluate_response_speak_contract(state) + expected_audio_turns = expected_turns - len(expected_empty_turns) + lifecycle_counts_ok = ( + state.count("response.created") == state.count("response.done") + and len(state.response_ids) == len(completed_response_ids) + and len(completed_response_ids) == len(set(completed_response_ids)) + and state.count("response.audio.done") <= state.count("response.done") + and _response_cardinality_ok( + completed_response_ids, + expected_turns=expected_turns, + validation_mode=validation_mode, + ) + ) + turn_outcomes_ok = len(turn_outcomes) == expected_turns and ( + validation_mode == "model-policy" or all(outcome == "speak" for outcome in turn_outcomes) + ) + clean_turn_audio_ok = all( + state.response_audio_delta_count(response_id) > 0 + for response_id in observed_turn_response_ids + if response_id not in expected_empty_response_ids + ) + playback_history_committed_ok = _all_response_playback_history_committed( + state, + completed_response_ids, + ) + if validation_mode == "response-required": + full_audio_response_ok = ( + len(observed_turn_response_ids) == expected_audio_turns + len(expected_empty_turns) + and clean_turn_audio_ok + and model_speak_event_ok + and state.model_speak_delta_count > 0 + and realtime_audio_lifecycle_ok + and state.audio_marks_seen + ) + else: + full_audio_response_ok = not observed_turn_response_ids or ( + realtime_audio_lifecycle_ok and state.model_speak_delta_count > 0 + ) + stale_audio_delta_count = state.stale_audio_delta_count() + terminal_activity_ok = state.count("response.done") > 0 or ( + validation_mode == "model-policy" and state.model_listen_count > 0 + ) + all_audio_responses_have_transcript = ( + validation_mode != "response-required" + or _all_audio_responses_have_transcript( + state, + completed_response_ids, + ) + ) + unexpected_error_events = _unexpected_error_events(state) + continuous_input_ok = not continuous_input or state.count("input_audio_buffer.committed") == 0 + result = { + "ok": terminal_activity_ok + and state.count("session.closed") > 0 + and state.cancelled_count == 0 + and not overlap_barge_in + and state.truncate_count == 0 + and event_order_ok + and input_transcription_ok + and stale_audio_delta_count == 0 + and lifecycle_counts_ok + and turn_outcomes_ok + and full_audio_response_ok, + "event_counts": { + event_type: state.count(event_type) + for event_type in sorted({str(event.get("type")) for event in state.events}) + }, + "audio_delta_count": len(state.audio_deltas), + "done_count": state.done_count, + "cancelled_count": state.cancelled_count, + "listen_count": state.listen_count, + "model_listen_count": state.model_listen_count, + "buffering_listen_count": state.buffering_listen_count, + "model_speak_event_count": state.model_speak_event_count, + "model_speak_delta_count": state.model_speak_delta_count, + "playback_ack_count": state.playback_ack_count, + "playback_history_committed_count": state.playback_history_committed_count, + "playback_history_committed_ok": playback_history_committed_ok, + "truncate_count": state.truncate_count, + "input_transcription_count": state.input_transcription_count, + "audio_marks_seen": state.audio_marks_seen, + "overlap_decisions": state.overlap_decisions, + "overlap_barge_in": overlap_barge_in, + "event_order_ok": event_order_ok, + "terminal_activity_ok": terminal_activity_ok, + "input_transcription_ok": input_transcription_ok, + "completed_response_ids": completed_response_ids, + "response_timings": state.response_timing_summaries(), + "lifecycle_counts_ok": lifecycle_counts_ok, + "validation_mode": validation_mode, + "scenario": scenario, + "listen_only_overlap_ok": listen_only_overlap_ok, + "transcript_hints_enabled": transcript_hints_enabled, + "realtime_input": realtime_input, + "continuous_input": continuous_input, + "continuous_input_ok": continuous_input_ok, + "input_chunk_ms": args.chunk_ms, + "model_policy_settle_ms": args.model_policy_settle_ms, + "turn_outcomes": turn_outcomes, + "turn_outcomes_ok": turn_outcomes_ok, + "clean_turn_audio_ok": clean_turn_audio_ok, + "full_audio_response_ok": full_audio_response_ok, + "model_speak_event_ok": model_speak_event_ok, + "realtime_audio_lifecycle_ok": realtime_audio_lifecycle_ok, + "stale_audio_delta_count": stale_audio_delta_count, + "all_audio_responses_have_transcript": all_audio_responses_have_transcript, + "distinct_turn_inputs": distinct_turn_inputs, + "error_count": len(unexpected_error_events), + "errors": unexpected_error_events, + **response_speak_contract, + **transcript_integrity, + "turn_inputs": [ + { + "path": str(path), + "source_duration_ms": len(pcm16) * 1000 // (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE), + "requested_duration_ms": duration_ms, + "sent_duration_ms": len(_select_turn_audio(pcm16, duration_ms)) + * 1000 + // (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE), + } + for path, pcm16, duration_ms in zip(turn_input_paths, turn_pcm16, turn_durations, strict=True) + ], + "output_dir": str(output_dir), + } + result["ok"] = bool( + result["ok"] + and transcript_integrity["transcript_delta_done_ok"] + and transcript_integrity["cross_turn_independent_ok"] + and transcript_integrity["empty_turns_ok"] + and transcript_integrity["nonempty_audio_has_transcript_ok"] + and transcript_integrity["terminal_punctuation_ok"] + and all_audio_responses_have_transcript + and playback_history_committed_ok + and listen_only_overlap_ok + and not unexpected_error_events + and response_speak_contract["response_speak_contract_ok"] + and continuous_input_ok + and (distinct_turn_inputs or not getattr(args, "require_distinct_inputs", False)) + ) + return result + + +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( + "--frame-image", + default=None, + help="Optional image file sent as an omni-duplex camera frame (one per 1 s of audio)", + ) + parser.add_argument( + "--turn-input-wav", + action="append", + default=[], + help="WAV for each turn after the first; repeat exactly turns-1 times.", + ) + parser.add_argument("--output-dir", default="/tmp/minicpmo_realtime_duplex_demo") + parser.add_argument( + "--output-audio-format", + default="pcm16", + choices=["pcm16", "wav", "g711_ulaw", "g711_alaw"], + ) + parser.add_argument("--chunk-ms", type=int, default=200) + parser.add_argument( + "--model-policy-settle-ms", + type=int, + default=2000, + help="Wait for a delayed response.created after model-listen and before closing the session.", + ) + parser.add_argument("--first-turn-ms", type=int, default=1400) + parser.add_argument( + "--turn-duration-ms", + action="append", + type=int, + default=[], + help="Audio duration for each turn; repeat turns times. Use 0 to send the complete WAV.", + ) + parser.add_argument("--first-turn-transcript", default="demo input speech") + parser.add_argument( + "--omit-transcript-hints", + action="store_true", + help="Keep local turn labels but do not send transcript hints to the server.", + ) + parser.add_argument( + "--realtime-input", + action="store_true", + help="Pace each input chunk according to its audio duration.", + ) + parser.add_argument( + "--continuous-input", + action="store_true", + help="Model-owned browser mode: stream PCM continuously without input_audio_buffer.commit.", + ) + parser.add_argument( + "--scenario", + choices=["sequential", "listen-only-overlap"], + default="sequential", + help="Run sequential turns or stream turn 2 while turn 1 audio is active without interruption.", + ) + parser.add_argument( + "--validation-mode", + choices=["model-policy", "response-required"], + default="response-required", + help="Allow model listen decisions, or require at least one response for every input turn.", + ) + parser.add_argument( + "--temperature", + type=float, + default=None, + help="Stage0 sampling temperature; response-required defaults to deterministic sampling.", + ) + parser.add_argument( + "--require-distinct-inputs", + action="store_true", + help="Require distinct WAV paths and audio payloads, and reject cross-turn transcript reuse.", + ) + parser.add_argument( + "--expect-empty-turn", + action="append", + type=int, + default=[], + help="1-based turn expected to end without text or audio; repeat for multiple turns.", + ) + parser.add_argument("--short-ack-ms", type=int, default=350) + parser.add_argument("--silence-ms", type=int, default=500) + parser.add_argument("--playback-ack-ms", type=int, default=500) + parser.add_argument("--turns", type=int, default=3) + parser.add_argument("--timeout-s", type=float, default=60.0) + parser.add_argument( + "--require-audio", + action="store_true", + help="Fail if the first native response listens instead of producing audio.", + ) + return parser.parse_args() + + +def main() -> None: + result = asyncio.run(run_demo(parse_args())) + print(json.dumps(result, ensure_ascii=False, indent=2)) + if not result["ok"]: + raise SystemExit(1) + + +if __name__ == "__main__": + main() 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 new file mode 100644 index 000000000..f5539968b --- /dev/null +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py @@ -0,0 +1,388 @@ +"""Run two MiniCPM-o Realtime duplex demo processes concurrently. + +This E2E driver covers the same process boundary users exercise manually: +two independent ``realtime_duplex_demo.py`` invocations, two distinct input +WAVs, and two distinct output directories. +""" + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import hashlib +import io +import json +import sys +import uuid +import wave +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[3] +DEMO_PATH = REPO_ROOT / "examples/online_serving/minicpmo/realtime_duplex_demo.py" +AUDIO_DELTA_EVENTS = {"response.audio.delta", "response.output_audio.delta"} + + +def _canonical_path(path: str) -> Path: + return Path(path).expanduser().resolve(strict=False) + + +def _sha256_file(path: Path) -> str: + hasher = hashlib.sha256() + with path.open("rb") as f: + for chunk in iter(lambda: f.read(1024 * 1024), b""): + hasher.update(chunk) + return hasher.hexdigest() + + +def _validate_pair_args(args: argparse.Namespace) -> dict[str, object]: + input_a = _canonical_path(args.input_wav_a) + input_b = _canonical_path(args.input_wav_b) + output_a = _canonical_path(args.output_dir_a) + output_b = _canonical_path(args.output_dir_b) + if output_a == output_b: + raise ValueError("output directories must be different") + if not input_a.is_file(): + raise ValueError(f"input WAV does not exist: {input_a}") + if not input_b.is_file(): + raise ValueError(f"input WAV does not exist: {input_b}") + sha_a = _sha256_file(input_a) + sha_b = _sha256_file(input_b) + if sha_a == sha_b: + raise ValueError("input WAV files must have different content") + return { + "input_wav_a": str(input_a), + "input_wav_b": str(input_b), + "input_sha256_a": sha_a, + "input_sha256_b": sha_b, + "output_dir_a": str(output_a), + "output_dir_b": str(output_b), + } + + +def _event_response_id(event: dict[str, object]) -> str | None: + response_id = event.get("response_id") + if isinstance(response_id, str): + return response_id + response = event.get("response") + if isinstance(response, dict) and isinstance(response.get("id"), str): + return response["id"] + return None + + +def _read_json(path: Path) -> dict[str, object]: + if not path.is_file(): + return {} + data = json.loads(path.read_text(encoding="utf-8")) + return data if isinstance(data, dict) else {} + + +def _read_events(path: Path) -> list[dict[str, object]]: + if not path.is_file(): + return [] + events: list[dict[str, object]] = [] + for line in path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + event = json.loads(line) + if isinstance(event, dict): + events.append(event) + return events + + +def _response_ids_from_artifacts(result: dict[str, object], events: list[dict[str, object]]) -> list[str]: + response_ids = result.get("response_ids") + if isinstance(response_ids, list): + ids = [item for item in response_ids if isinstance(item, str)] + if ids: + return ids + seen: list[str] = [] + for event in events: + response_id = _event_response_id(event) + if response_id is not None and response_id not in seen: + seen.append(response_id) + return seen + + +def _audio_delta_duration_ms(event: dict[str, object]) -> float | None: + encoded = event.get("delta") or event.get("audio") + if not isinstance(encoded, str) or not encoded: + return None + try: + raw = base64.b64decode(encoded) + except ValueError: + return None + fmt = str(event.get("format") or event.get("audio_format") or "pcm16").lower() + sample_rate = event.get("sample_rate_hz") or event.get("sample_rate") or 24_000 + if not isinstance(sample_rate, int | float) or sample_rate <= 0: + return None + if "wav" in fmt: + try: + with wave.open(io.BytesIO(raw), "rb") as wav: + return round((wav.getnframes() * 1000.0) / wav.getframerate(), 3) + except (EOFError, wave.Error): + return None + bytes_per_sample = 4 if "f32" in fmt else 2 + return round((len(raw) * 1000.0) / (bytes_per_sample * float(sample_rate)), 3) + + +def _audio_delta_received_at_s(event: dict[str, object]) -> float | None: + raw = event.get("_client_received_at_s") + if isinstance(raw, int | float): + return float(raw) + return None + + +def _session_audio_contract_ok( + *, + audio_delta_count: int, + delta_response_ids: list[str], + response_ids: list[str], + first_audio_delta_before_done: bool, +) -> bool: + if audio_delta_count <= 0 or not response_ids: + return False + if len(delta_response_ids) != audio_delta_count: + return False + delta_id_set = set(delta_response_ids) + if len(delta_id_set) != 1: + return False + return first_audio_delta_before_done and next(iter(delta_id_set)) in set(response_ids) + + +def _summarize_session( + *, + label: str, + input_wav: str, + output_dir: str, + returncode: int, + stdout: str, + stderr: str, +) -> dict[str, object]: + output_path = Path(output_dir) + result = _read_json(output_path / "result.json") + events = _read_events(output_path / "events.jsonl") + audio_delta_indices = [index for index, event in enumerate(events) if event.get("type") in AUDIO_DELTA_EVENTS] + done_indices = [index for index, event in enumerate(events) if event.get("type") == "response.done"] + response_ids = _response_ids_from_artifacts(result, events) + delta_response_ids = [ + response_id + for response_id in (_event_response_id(events[index]) for index in audio_delta_indices) + if response_id is not None + ] + first_audio_delta_before_done = bool(audio_delta_indices and done_indices) and min(audio_delta_indices) < min( + done_indices + ) + audio_delta_events = [events[index] for index in audio_delta_indices] + first_received_at_s = next( + ( + received_at + for received_at in (_audio_delta_received_at_s(event) for event in audio_delta_events) + if received_at is not None + ), + None, + ) + audio_delta_timings = [] + for event in audio_delta_events: + received_at_s = _audio_delta_received_at_s(event) + audio_delta_timings.append( + { + "response_id": _event_response_id(event), + "received_at_s": received_at_s, + "arrival_offset_ms": ( + round((received_at_s - first_received_at_s) * 1000.0, 3) + if received_at_s is not None and first_received_at_s is not None + else None + ), + "audio_duration_ms": _audio_delta_duration_ms(event), + } + ) + all_audio_deltas_same_response = ( + len(delta_response_ids) == len(audio_delta_indices) and len(set(delta_response_ids)) == 1 + ) + response_id_nonempty = bool(response_ids) + response_audio_contract_ok = _session_audio_contract_ok( + audio_delta_count=len(audio_delta_indices), + delta_response_ids=delta_response_ids, + response_ids=response_ids, + first_audio_delta_before_done=first_audio_delta_before_done, + ) + return { + "label": label, + "input_wav": input_wav, + "output_dir": output_dir, + "returncode": returncode, + "ok": returncode == 0 and result.get("ok") is True and response_audio_contract_ok, + "result_ok": result.get("ok"), + "response_ids": response_ids, + "response_id_nonempty": response_id_nonempty, + "audio_bytes": result.get("audio_bytes", 0), + "audio_delta_count": len(audio_delta_indices), + "audio_delta_response_ids": delta_response_ids, + "audio_delta_timings": audio_delta_timings, + "all_audio_deltas_same_response": all_audio_deltas_same_response, + "first_audio_delta_before_done": first_audio_delta_before_done, + "response_audio_contract_ok": response_audio_contract_ok, + "event_count": len(events), + "stdout_tail": stdout[-4000:], + "stderr_tail": stderr[-4000:], + } + + +def _identity_isolation_ok(sessions: list[dict[str, object]]) -> bool: + seen: set[str] = set() + for session in sessions: + ids = {response_id for response_id in session.get("response_ids", []) if isinstance(response_id, str)} + if not ids: + return False + if ids & seen: + return False + seen.update(ids) + return True + + +async def _run_demo_process( + *, + label: str, + args: argparse.Namespace, + input_wav: str, + output_dir: str, +) -> dict[str, object]: + command = [ + sys.executable, + str(DEMO_PATH), + "--url", + args.url, + "--model", + args.model, + "--input-wav", + input_wav, + "--output-dir", + output_dir, + "--chunk-ms", + 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") + if args.require_audio: + command.append("--require-audio") + command.extend(["--ref-audio", str(args.ref_audio)]) + + process = await asyncio.create_subprocess_exec( + *command, + cwd=str(REPO_ROOT), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout_bytes, stderr_bytes = await asyncio.wait_for( + process.communicate(), + timeout=args.timeout_s + 15.0, + ) + except TimeoutError: + process.kill() + stdout_bytes, stderr_bytes = await process.communicate() + return _summarize_session( + label=label, + input_wav=input_wav, + output_dir=output_dir, + returncode=-1, + stdout=stdout_bytes.decode(errors="replace"), + stderr=(stderr_bytes.decode(errors="replace") + "\nclient timeout"), + ) + return _summarize_session( + label=label, + input_wav=input_wav, + output_dir=output_dir, + returncode=process.returncode if process.returncode is not None else -1, + stdout=stdout_bytes.decode(errors="replace"), + stderr=stderr_bytes.decode(errors="replace"), + ) + + +async def run_pair(args: argparse.Namespace) -> dict[str, object]: + validated = _validate_pair_args(args) + Path(validated["output_dir_a"]).mkdir(parents=True, exist_ok=True) + Path(validated["output_dir_b"]).mkdir(parents=True, exist_ok=True) + + sessions = await asyncio.gather( + _run_demo_process( + label="a", + args=args, + input_wav=str(validated["input_wav_a"]), + output_dir=str(validated["output_dir_a"]), + ), + _run_demo_process( + label="b", + args=args, + input_wav=str(validated["input_wav_b"]), + output_dir=str(validated["output_dir_b"]), + ), + ) + identity_isolation_ok = _identity_isolation_ok(sessions) + min_audio_delta_ok = all( + int(session.get("audio_delta_count", 0)) >= args.min_audio_deltas_per_session for session in sessions + ) + response_audio_contract_ok = all(session.get("response_audio_contract_ok") is True for session in sessions) + result = { + "ok": ( + all(session.get("ok") is True for session in sessions) + and identity_isolation_ok + and min_audio_delta_ok + and response_audio_contract_ok + ), + "input_wavs_distinct": True, + "output_dirs_distinct": True, + "identity_isolation_ok": identity_isolation_ok, + "min_audio_deltas_per_session": args.min_audio_deltas_per_session, + "min_audio_delta_ok": min_audio_delta_ok, + "response_audio_contract_ok": response_audio_contract_ok, + "inputs": { + "a": validated["input_wav_a"], + "b": validated["input_wav_b"], + "sha256_a": validated["input_sha256_a"], + "sha256_b": validated["input_sha256_b"], + }, + "sessions": sessions, + } + summary_output = ( + Path(args.summary_output) if args.summary_output else Path(validated["output_dir_a"]).parent / "summary.json" + ) + summary_output.parent.mkdir(parents=True, exist_ok=True) + summary_output.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8") + return result + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime?duplex=1") + parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5") + parser.add_argument("--input-wav-a", required=True) + parser.add_argument("--input-wav-b", required=True) + parser.add_argument("--ref-audio", required=True) + parser.add_argument("--output-dir-a", required=True) + parser.add_argument("--output-dir-b", required=True) + parser.add_argument("--summary-output") + parser.add_argument("--chunk-ms", type=int, default=200) + parser.add_argument("--timeout-s", type=float, default=120.0) + parser.add_argument("--no-realtime-pacing", action="store_true") + parser.add_argument("--require-audio", action="store_true") + parser.add_argument("--min-audio-deltas-per-session", type=int, default=2) + args = parser.parse_args() + if args.min_audio_deltas_per_session < 0: + parser.error("--min-audio-deltas-per-session must be non-negative") + return args + + +def main() -> None: + result = asyncio.run(run_pair(parse_args())) + print(json.dumps(result, ensure_ascii=False, indent=2)) + raise SystemExit(0 if result["ok"] else 1) + + +if __name__ == "__main__": + main() 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 new file mode 100644 index 000000000..61a466913 --- /dev/null +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py @@ -0,0 +1,556 @@ +"""Concurrent MiniCPM-o Realtime duplex and resumable-session E2E driver.""" + +from __future__ import annotations + +import argparse +import asyncio +import hashlib +import json +import sys +import uuid +from pathlib import Path +from types import SimpleNamespace +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit + +import websockets +from websockets.exceptions import ConnectionClosed + +SCRIPT_DIR = Path(__file__).resolve().parent +if str(SCRIPT_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPT_DIR)) + +from minicpmo_realtime_duplex_scenarios import _url_with_model, run_demo # noqa: E402 + + +def _with_resume_mode(url: str) -> str: + parts = urlsplit(url) + query = dict(parse_qsl(parts.query, keep_blank_values=True)) + query["resume"] = "1" + return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment)) + + +def _response_ids(result: dict[str, object]) -> set[str]: + values = result.get("completed_response_ids") + return {value for value in values if isinstance(value, str)} if isinstance(values, list) else set() + + +def _event_session_id(event: dict[str, object]) -> str | None: + session_id = event.get("session_id") + if isinstance(session_id, str): + return session_id + inner = event.get("event") + if isinstance(inner, dict): + nested_session_id = inner.get("session_id") + if isinstance(nested_session_id, str): + return nested_session_id + return None + + +def _error_code(event: dict[str, object]) -> str | None: + error = event.get("error") + if isinstance(error, dict) and isinstance(error.get("code"), str): + return error["code"] + code = event.get("code") + return code if isinstance(code, str) else None + + +def _validate_identity_isolation(results: list[dict[str, object]]) -> bool: + seen: set[str] = set() + for result in results: + current = _response_ids(result) + if current & seen: + return False + seen.update(current) + return True + + +def _validate_semantic_isolation( + results: list[dict[str, object]], + *, + input_wavs: list[str], + expected_tokens: list[str], +) -> bool: + if not input_wavs: + return True + if len(results) != len(input_wavs): + return False + input_hashes = [hashlib.sha256(Path(path).read_bytes()).digest() for path in input_wavs] + if len(set(input_hashes)) != len(input_hashes): + return False + if not expected_tokens: + return True + if len(expected_tokens) != len(results): + return False + for result, expected_token in zip(results, expected_tokens, strict=True): + details = result.get("transcript_integrity") + transcripts = ( + [str(item.get("transcript", "")) for item in details if isinstance(item, dict)] + if isinstance(details, list) + else [] + ) + if expected_token not in "".join(transcripts): + return False + return True + + +async def _receive_until(ws, event_type: str, *, timeout_s: float) -> tuple[dict[str, object], list[dict[str, object]]]: + async def receive() -> tuple[dict[str, object], list[dict[str, object]]]: + events: list[dict[str, object]] = [] + while True: + raw = await ws.recv() + if not isinstance(raw, str): + continue + event = json.loads(raw) + if not isinstance(event, dict): + continue + events.append(event) + if event.get("type") == event_type: + return event, events + + return await asyncio.wait_for(receive(), timeout=timeout_s) + + +async def _open_admission_session( + args: argparse.Namespace, + session_id: str, +) -> tuple[object, dict[str, object]]: + url = _url_with_model(args.url, args.model, session_id=session_id) + ws = await websockets.connect(url, max_size=64 * 1024 * 1024) + await ws.send( + json.dumps( + { + "type": "session.update", + "session": { + "session_id": session_id, + "model": args.model, + "modalities": ["audio", "text"], + "extra_body": {"minicpmo45_native_duplex": True}, + }, + } + ) + ) + created, _ = await _receive_until(ws, "session.created", timeout_s=args.timeout_s) + return ws, created + + +async def _close_admission_session(ws, *, timeout_s: float) -> None: + await ws.send(json.dumps({"type": "session.close"})) + await _receive_until(ws, "session.closed", timeout_s=timeout_s) + await ws.close() + + +async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str, object]: + if limit < 1: + raise ValueError("admission limit must be positive") + prefix = f"admission-{uuid.uuid4().hex}" + accepted: list[tuple[object, dict[str, object]]] = [] + replacement: tuple[object, dict[str, object]] | None = None + overflow_code = None + try: + 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, session_id=overflow_id) + async with websockets.connect(overflow_url, max_size=64 * 1024 * 1024) as overflow: + await overflow.send( + json.dumps( + { + "type": "session.update", + "session": { + "session_id": overflow_id, + "model": args.model, + "modalities": ["audio", "text"], + "extra_body": {"minicpmo45_native_duplex": True}, + }, + } + ) + ) + error, _ = await _receive_until(overflow, "error", timeout_s=args.timeout_s) + overflow_code = _error_code(error) + + first_ws, _ = accepted.pop(0) + await _close_admission_session(first_ws, timeout_s=args.timeout_s) + replacement = await _open_admission_session(args, f"{prefix}-replacement") + + first_capabilities = accepted[0][1].get("session") if accepted else replacement[1].get("session") + capabilities = first_capabilities.get("capabilities") if isinstance(first_capabilities, dict) else None + advertised_multi = capabilities.get("supports_multi_session") if isinstance(capabilities, dict) else None + admission_mode = capabilities.get("session_admission_mode") if isinstance(capabilities, dict) else None + return { + "ok": ( + overflow_code == "resource_exhausted" + and replacement[1].get("type") == "session.created" + and admission_mode == "engine_managed" + ), + "configured_limit": limit, + "accepted_before_rejection": limit, + "overflow_error_code": overflow_code, + "replacement_accepted": True, + "advertised_multi_session": advertised_multi, + "session_admission_mode": admission_mode, + } + finally: + cleanup = list(accepted) + if replacement is not None: + cleanup.append(replacement) + for ws, _ in cleanup: + try: + await _close_admission_session(ws, timeout_s=args.timeout_s) + except Exception: + await ws.close() + + +async def _resume_probe( + args: argparse.Namespace, + *, + session_id: str, + expect_expired: bool = False, +) -> dict[str, object]: + url = _url_with_model(args.url, args.model, session_id=session_id) + async with websockets.connect(url, max_size=64 * 1024 * 1024) as first: + await first.send( + json.dumps( + { + "type": "session.update", + "session": { + "model": args.model, + "modalities": ["audio", "text"], + "extra_body": {"minicpmo45_native_duplex": True}, + }, + } + ) + ) + created, first_events = await _receive_until(first, "session.created", timeout_s=args.timeout_s) + token = created.get("resume_token") + incarnation = created.get("incarnation") + generation = created.get("attachment_generation") + if not isinstance(token, str) or not isinstance(incarnation, int): + raise RuntimeError("session.created omitted resumable credentials") + last_seq = max( + ( + event.get("server_event_seq", 0) + for event in first_events + if isinstance(event.get("server_event_seq"), int) + ), + default=0, + ) + + delay_s = args.expire_after_s if expect_expired else args.resume_after_ms / 1000 + if delay_s > 0: + await asyncio.sleep(delay_s) + + resume_url = _with_resume_mode(url) + async with websockets.connect(resume_url, max_size=64 * 1024 * 1024) as second: + await second.send( + json.dumps( + { + "type": "session.resume", + "session_id": session_id, + "incarnation": incarnation, + "resume_token": token, + "last_received_server_event_seq": last_seq, + } + ) + ) + if expect_expired: + error, error_events = await _receive_until(second, "error", timeout_s=args.timeout_s) + error_payload = error.get("error") + code = error_payload.get("code") if isinstance(error_payload, dict) else error.get("code") + return { + "ok": code == "session_resume_expired", + "session_id": session_id, + "expired": True, + "error_code": code, + "event_count": len(error_events), + } + resumed, replay = await _receive_until(second, "session.resumed", timeout_s=args.timeout_s) + rotated = resumed.get("resume_token") + if not isinstance(rotated, str) or rotated == token: + raise RuntimeError("session.resume did not rotate the resume token") + await second.send(json.dumps({"type": "session.heartbeat"})) + heartbeat, heartbeat_events = await _receive_until( + second, + "session.heartbeat_ack", + timeout_s=args.timeout_s, + ) + await second.send(json.dumps({"type": "session.close"})) + closed, close_events = await _receive_until(second, "session.closed", timeout_s=args.timeout_s) + + replay_sequences = [event["server_event_seq"] for event in replay if isinstance(event.get("server_event_seq"), int)] + return { + "ok": ( + resumed.get("session_id") == session_id + and isinstance(generation, int) + and resumed.get("attachment_generation") == generation + 1 + and _event_session_id(heartbeat) == session_id + and _event_session_id(closed) == session_id + and replay_sequences == sorted(replay_sequences) + ), + "session_id": session_id, + "resumed_session_id": resumed.get("session_id"), + "heartbeat_session_id": _event_session_id(heartbeat), + "closed_session_id": _event_session_id(closed), + "initial_attachment_generation": generation, + "resumed_attachment_generation": resumed.get("attachment_generation"), + "replayed_event_count": len(replay_sequences), + "replayed_event_sequences": replay_sequences, + "heartbeat_event_count": len(heartbeat_events), + "close_event_count": len(close_events), + "token_rotated": True, + } + + +async def _takeover_probe( + args: argparse.Namespace, + *, + session_id: str, +) -> dict[str, object]: + url = _url_with_model(args.url, args.model, session_id=session_id) + first = await websockets.connect(url, max_size=64 * 1024 * 1024) + second = None + try: + await first.send( + json.dumps( + { + "type": "session.update", + "session": { + "model": args.model, + "modalities": ["audio", "text"], + "extra_body": {"minicpmo45_native_duplex": True}, + }, + } + ) + ) + created, first_events = await _receive_until(first, "session.created", timeout_s=args.timeout_s) + 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): + raise RuntimeError("session.created omitted takeover credentials") + last_seq = max( + ( + event.get("server_event_seq", 0) + for event in first_events + if isinstance(event.get("server_event_seq"), int) + ), + default=0, + ) + + second = await websockets.connect(_with_resume_mode(url), max_size=64 * 1024 * 1024) + await second.send( + json.dumps( + { + "type": "session.resume", + "session_id": session_id, + "incarnation": incarnation, + "resume_token": token, + "last_received_server_event_seq": last_seq, + } + ) + ) + resumed, replay_events = await _receive_until(second, "session.resumed", timeout_s=args.timeout_s) + replaced, replaced_events = await _receive_until(first, "session.replaced", timeout_s=args.timeout_s) + await asyncio.wait_for(first.wait_closed(), timeout=args.timeout_s) + + rejected_old_writes = 0 + for _ in range(4): + try: + await first.send(json.dumps({"type": "session.heartbeat"})) + except ConnectionClosed: + rejected_old_writes += 1 + + await second.send(json.dumps({"type": "session.heartbeat"})) + heartbeat, heartbeat_events = await _receive_until( + second, + "session.heartbeat_ack", + timeout_s=args.timeout_s, + ) + await second.send(json.dumps({"type": "session.close"})) + closed, close_events = await _receive_until(second, "session.closed", timeout_s=args.timeout_s) + rotated_token = resumed.get("resume_token") + return { + "ok": ( + resumed.get("session_id") == session_id + and resumed.get("attachment_generation") == generation + 1 + and isinstance(rotated_token, str) + and rotated_token != token + and _event_session_id(replaced) == session_id + and replaced.get("attachment_generation") == generation + and rejected_old_writes == 4 + and _event_session_id(heartbeat) == session_id + and _event_session_id(closed) == session_id + ), + "session_id": session_id, + "initial_attachment_generation": generation, + "resumed_attachment_generation": resumed.get("attachment_generation"), + "replaced_attachment_generation": replaced.get("attachment_generation"), + "token_rotated": isinstance(rotated_token, str) and rotated_token != token, + "old_attachment_closed": True, + "rejected_old_writes": rejected_old_writes, + "replay_event_count": len(replay_events), + "replaced_event_count": len(replaced_events), + "heartbeat_event_count": len(heartbeat_events), + "close_event_count": len(close_events), + } + finally: + if second is not None: + await second.close() + await first.close() + + +def _demo_args(args: argparse.Namespace, index: int) -> SimpleNamespace: + validation_mode = "response-required" if args.response_required else "model-policy" + input_wav = args.session_input_wav[index] if args.session_input_wav else args.input_wav + return SimpleNamespace( + url=args.url, + model=args.model, + session_id=f"multi-{index}-{uuid.uuid4().hex}", + input_wav=input_wav, + ref_audio=args.ref_audio, + turn_input_wav=list(args.turn_input_wav), + output_dir=str(Path(args.output_dir) / f"session_{index:02d}"), + output_audio_format="pcm16", + chunk_ms=args.chunk_ms, + realtime_input=args.realtime_input, + first_turn_ms=args.first_turn_ms, + turn_duration_ms=list(args.turn_duration_ms), + first_turn_transcript=f"session {index} input", + omit_transcript_hints=True, + validation_mode=validation_mode, + temperature=args.temperature, + scenario="sequential", + require_audio=args.response_required, + require_distinct_inputs=False, + expect_empty_turn=[], + short_ack_ms=350, + turns=args.turns, + timeout_s=args.timeout_s, + model_policy_settle_ms=args.model_policy_settle_ms, + ) + + +async def run_multi_session(args: argparse.Namespace) -> dict[str, object]: + if args.sessions < 1: + raise ValueError("--sessions must be positive") + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + resume_result = None + if args.disconnect_session_index is not None: + if not 0 <= args.disconnect_session_index < args.sessions: + raise ValueError("--disconnect-session-index is outside the session range") + resume_result = await _resume_probe( + args, + session_id=f"resume-{args.disconnect_session_index}-{uuid.uuid4().hex}", + ) + takeover_result = None + if args.takeover_session_index is not None: + if not 0 <= args.takeover_session_index < args.sessions: + raise ValueError("--takeover-session-index is outside the session range") + takeover_result = await _takeover_probe( + args, + session_id=f"takeover-{args.takeover_session_index}-{uuid.uuid4().hex}", + ) + expiry_result = None + if args.expire_session_index is not None: + if not 0 <= args.expire_session_index < args.sessions: + raise ValueError("--expire-session-index is outside the session range") + expiry_result = await _resume_probe( + args, + session_id=f"expire-{args.expire_session_index}-{uuid.uuid4().hex}", + expect_expired=True, + ) + admission_result = ( + await _admission_probe(args, limit=args.verify_admission_limit) + if args.verify_admission_limit is not None + else None + ) + + session_results = await asyncio.gather( + *(run_demo(_demo_args(args, index)) for index in range(args.sessions)), + return_exceptions=True, + ) + failures = [repr(result) for result in session_results if isinstance(result, BaseException)] + completed = [result for result in session_results if isinstance(result, dict)] + identity_isolation_ok = _validate_identity_isolation(completed) + semantic_isolation_ok = _validate_semantic_isolation( + completed, + input_wavs=list(args.session_input_wav), + expected_tokens=list(args.session_expected_token), + ) + result = { + "ok": ( + not failures + and len(completed) == args.sessions + and all(item.get("ok") is True for item in completed) + and identity_isolation_ok + and semantic_isolation_ok + and (resume_result is None or resume_result.get("ok") is True) + and (takeover_result is None or takeover_result.get("ok") is True) + and (expiry_result is None or expiry_result.get("ok") is True) + and (admission_result is None or admission_result.get("ok") is True) + ), + "session_count": args.sessions, + "identity_isolation_ok": identity_isolation_ok, + "semantic_isolation_ok": semantic_isolation_ok, + "resume": resume_result, + "takeover": takeover_result, + "expiry": expiry_result, + "admission": admission_result, + "failures": failures, + "sessions": completed, + } + (output_dir / "summary.json").write_text( + json.dumps(result, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + return result + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--url", default="ws://127.0.0.1:8113/v1/realtime?duplex=1") + parser.add_argument("--base-url", help="Deprecated alias; /v1/realtime is appended when supplied.") + parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5") + parser.add_argument("--sessions", type=int, default=2) + parser.add_argument("--input-wav", required=True) + parser.add_argument("--ref-audio", help="Optional WAV used as the MiniCPM-o voice prompt for every session.") + parser.add_argument("--session-input-wav", action="append", default=[]) + parser.add_argument("--session-expected-token", action="append", default=[]) + parser.add_argument("--turn-input-wav", action="append", default=[]) + parser.add_argument("--output-dir", default="/tmp/minicpmo_pr3907_multi_session_e2e") + parser.add_argument("--realtime-input", action="store_true") + parser.add_argument("--chunk-ms", type=int, default=200) + parser.add_argument("--turns", type=int, default=1) + parser.add_argument("--first-turn-ms", type=int, default=1400) + parser.add_argument("--turn-duration-ms", type=int, action="append", default=[]) + parser.add_argument("--response-required", action="store_true") + parser.add_argument("--temperature", type=float, default=None) + parser.add_argument("--disconnect-session-index", type=int) + parser.add_argument("--takeover-session-index", type=int) + parser.add_argument("--resume-after-ms", type=int, default=1000) + parser.add_argument("--expire-session-index", type=int) + parser.add_argument("--expire-after-s", type=float, default=6.0) + parser.add_argument("--verify-admission-limit", type=int) + parser.add_argument("--model-policy-settle-ms", type=int, default=600) + parser.add_argument("--timeout-s", type=float, default=120.0) + args = parser.parse_args() + if args.base_url: + args.url = args.base_url.rstrip("/") + "/v1/realtime?duplex=1" + if args.session_input_wav and len(args.session_input_wav) != args.sessions: + parser.error("provide exactly one --session-input-wav per session") + if args.session_expected_token and len(args.session_expected_token) != args.sessions: + parser.error("provide exactly one --session-expected-token per session") + if args.session_expected_token and not args.session_input_wav: + parser.error("--session-expected-token requires --session-input-wav") + return args + + +def main() -> None: + result = asyncio.run(run_multi_session(parse_args())) + print(json.dumps(result, ensure_ascii=False, indent=2)) + raise SystemExit(0 if result["ok"] else 1) + + +if __name__ == "__main__": + main() 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 new file mode 100644 index 000000000..3eb9405a6 --- /dev/null +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py @@ -0,0 +1,467 @@ +"""Validate MiniCPM-o Realtime duplex soft-interrupt delta streaming. + +This E2E driver runs the public ``realtime_duplex_demo.py`` against a live +duplex backend. Arbitrary audio defaults to the model-policy lifecycle contract. +The stronger response-required mode binds the two-response contract to a known +input checksum and expected second response. +""" + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import hashlib +import io +import json +import sys +import uuid +import wave +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[3] +DEMO_PATH = REPO_ROOT / "examples/online_serving/minicpmo/realtime_duplex_demo.py" +AUDIO_DELTA_EVENTS = {"response.audio.delta", "response.output_audio.delta"} +TRANSCRIPT_DELTA_EVENTS = { + "response.audio_transcript.delta", + "response.output_text.delta", +} + + +def _canonical_path(path: str) -> Path: + return Path(path).expanduser().resolve(strict=False) + + +def _validate_input_sha256(path: Path, expected: str) -> str: + actual = hashlib.sha256(path.read_bytes()).hexdigest() + if actual != expected.lower(): + raise ValueError(f"input WAV SHA256 mismatch: expected {expected.lower()}, got {actual}") + return actual + + +def _read_json(path: Path) -> dict[str, object]: + if not path.is_file(): + return {} + data = json.loads(path.read_text(encoding="utf-8")) + return data if isinstance(data, dict) else {} + + +def _read_events(path: Path) -> list[dict[str, object]]: + if not path.is_file(): + return [] + events: list[dict[str, object]] = [] + for line in path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + event = json.loads(line) + if isinstance(event, dict): + events.append(event) + return events + + +def _event_response_id(event: dict[str, object]) -> str | None: + response_id = event.get("response_id") + if isinstance(response_id, str): + return response_id + response = event.get("response") + if isinstance(response, dict) and isinstance(response.get("id"), str): + return response["id"] + return None + + +def _event_received_at_s(event: dict[str, object]) -> float | None: + received_at = event.get("_client_received_at_s") + if isinstance(received_at, int | float): + return float(received_at) + return None + + +def _first_event_index(events: list[dict[str, object]], event_type: str) -> int | None: + for index, event in enumerate(events): + if event.get("type") == event_type: + return index + return None + + +def _event_counts(events: list[dict[str, object]]) -> dict[str, int]: + counts: dict[str, int] = {} + for event in events: + event_type = str(event.get("type")) + counts[event_type] = counts.get(event_type, 0) + 1 + return counts + + +def _response_created_ids(events: list[dict[str, object]]) -> list[str]: + response_ids: list[str] = [] + for event in events: + if event.get("type") != "response.created": + continue + response_id = _event_response_id(event) + if response_id and response_id not in response_ids: + response_ids.append(response_id) + return response_ids + + +def _audio_delta_duration_ms(event: dict[str, object]) -> float | None: + encoded = event.get("delta") or event.get("audio") + if not isinstance(encoded, str) or not encoded: + return None + try: + raw = base64.b64decode(encoded) + except ValueError: + return None + fmt = str(event.get("format") or event.get("audio_format") or "pcm16").lower() + sample_rate = event.get("sample_rate_hz") or event.get("sample_rate") or 24_000 + if not isinstance(sample_rate, int | float) or sample_rate <= 0: + return None + if "wav" in fmt: + try: + with wave.open(io.BytesIO(raw), "rb") as wav_file: + return round((wav_file.getnframes() * 1000.0) / wav_file.getframerate(), 3) + except (EOFError, wave.Error): + return None + bytes_per_sample = 4 if "f32" in fmt else 2 + return round((len(raw) * 1000.0) / (bytes_per_sample * float(sample_rate)), 3) + + +def _normalize_text(text: str) -> str: + return "".join(text.split()) + + +def _compact_sequence(events: list[dict[str, object]]) -> list[str]: + sequence: list[str] = [] + for event in events: + event_type = event.get("type") + response_id = _event_response_id(event) + suffix = f" {response_id[-8:]}" if response_id else "" + if event_type == "response.listen": + sequence.append("listen") + elif event_type == "response.created": + sequence.append(f"response.created{suffix}") + elif event_type == "response.speak": + sequence.append(f"speak{suffix}") + elif event_type in AUDIO_DELTA_EVENTS: + sequence.append(f"audio.delta{suffix}") + elif event_type in TRANSCRIPT_DELTA_EVENTS: + delta = str(event.get("delta") or "") + sequence.append(f"text.delta{suffix}> {delta}") + elif event_type == "response.done": + sequence.append(f"response.done{suffix}") + elif event_type == "input_audio_buffer.committed": + sequence.append("input.committed") + return sequence + + +def _response_summary( + events: list[dict[str, object]], + response_id: str, + *, + t0: float | None, +) -> dict[str, object]: + created_indices = [ + index + for index, event in enumerate(events) + if event.get("type") == "response.created" and _event_response_id(event) == response_id + ] + done_indices = [ + index + for index, event in enumerate(events) + if event.get("type") == "response.done" and _event_response_id(event) == response_id + ] + audio_indices = [ + index + for index, event in enumerate(events) + if event.get("type") in AUDIO_DELTA_EVENTS and _event_response_id(event) == response_id + ] + transcript_indices = [ + index + for index, event in enumerate(events) + if event.get("type") in TRANSCRIPT_DELTA_EVENTS and _event_response_id(event) == response_id + ] + done_index = done_indices[0] if done_indices else None + audio_before_done_ok = done_index is not None and all(index < done_index for index in audio_indices) + stale_audio_count = sum(index > done_index for index in audio_indices) if done_index is not None else 0 + + def offset_ms(index: int) -> float | None: + if t0 is None: + return None + received_at = _event_received_at_s(events[index]) + return round((received_at - t0) * 1000.0, 3) if received_at is not None else None + + return { + "response_id": response_id, + "created_index": created_indices[0] if created_indices else None, + "done_indices": done_indices, + "created_offset_ms": offset_ms(created_indices[0]) if created_indices else None, + "done_offset_ms": offset_ms(done_indices[0]) if done_indices else None, + "audio_delta_count": len(audio_indices), + "transcript_delta_count": len(transcript_indices), + "audio_delta_offsets_ms": [offset_ms(index) for index in audio_indices], + "audio_duration_ms": [_audio_delta_duration_ms(events[index]) for index in audio_indices], + "transcript": "".join(str(events[index].get("delta") or "") for index in transcript_indices), + "one_done": len(done_indices) == 1, + "audio_before_done_ok": audio_before_done_ok, + "stale_audio_count": stale_audio_count, + } + + +def summarize_artifacts( + *, + output_dir: Path, + validation_mode: str, + min_responses: int, + min_audio_deltas_per_response: int, + expect_second_response_substring: str | None, +) -> dict[str, object]: + if validation_mode not in {"model-policy", "response-required"}: + raise ValueError(f"unsupported validation mode: {validation_mode}") + result = _read_json(output_dir / "result.json") + events = _read_events(output_dir / "events.jsonl") + response_ids = _response_created_ids(events) + times = [_event_received_at_s(event) for event in events] + clean_times = [received_at for received_at in times if received_at is not None] + t0 = min(clean_times) if clean_times else None + response_summaries = [_response_summary(events, response_id, t0=t0) for response_id in response_ids] + commit_index = _first_event_index(events, "input_audio_buffer.committed") + first_created_index = next( + (summary["created_index"] for summary in response_summaries if isinstance(summary.get("created_index"), int)), + None, + ) + first_done_index = ( + response_summaries[0]["done_indices"][0] + if response_summaries and response_summaries[0]["done_indices"] + else None + ) + second_created_index = ( + response_summaries[1]["created_index"] + if len(response_summaries) >= 2 and isinstance(response_summaries[1].get("created_index"), int) + else None + ) + last_done_index = next( + (summary["done_indices"][0] for summary in reversed(response_summaries) if summary.get("done_indices")), + None, + ) + listen_indices = [index for index, event in enumerate(events) if event.get("type") == "response.listen"] + effective_min_responses = min_responses if validation_mode == "response-required" else 1 + enough_responses = len(response_summaries) >= effective_min_responses + response_lifecycle_ok = enough_responses and all(bool(summary.get("one_done")) for summary in response_summaries) + multi_delta_ok = enough_responses and all( + int(summary.get("audio_delta_count", 0)) >= min_audio_deltas_per_response for summary in response_summaries + ) + response_audio_contract_ok = enough_responses and all( + summary.get("audio_before_done_ok") is True and int(summary.get("stale_audio_count", 0)) == 0 + for summary in response_summaries + ) + response_before_final_commit = ( + first_created_index is not None and commit_index is not None and first_created_index < commit_index + ) + second_response_before_final_commit = ( + second_created_index is not None and commit_index is not None and second_created_index < commit_index + ) + listen_before_first_response = first_created_index is not None and any( + index < first_created_index for index in listen_indices + ) + listen_between_responses = ( + first_done_index is not None + and second_created_index is not None + and any(first_done_index < index < second_created_index for index in listen_indices) + ) + listen_after_last_done = last_done_index is not None and any(index > last_done_index for index in listen_indices) + listen_after_response_before_commit = ( + last_done_index is not None + and commit_index is not None + and any(last_done_index < index < commit_index for index in listen_indices) + ) + final_listen_after_commit = commit_index is not None and any(index > commit_index for index in listen_indices) + transcript = "".join(str(summary.get("transcript") or "") for summary in response_summaries) + second_response_transcript = ( + str(response_summaries[1].get("transcript") or "") if len(response_summaries) >= 2 else "" + ) + second_response_transcript_expectation_ok = not expect_second_response_substring or _normalize_text( + expect_second_response_substring + ) in _normalize_text(second_response_transcript) + error_events = [ + event + for event in events + if event.get("type") == "error" + or str(event.get("type") or "").endswith(".cancelled") + or str(event.get("type") or "").endswith(".truncated") + ] + cancelled_count = sum( + 1 + for event in events + if event.get("type") == "response.done" + and isinstance(event.get("response"), dict) + and event["response"].get("status") == "cancelled" + ) + result_ok = result.get("ok") is True + common_contract_ok = bool( + result_ok + and not error_events + and cancelled_count == 0 + and enough_responses + and response_lifecycle_ok + and multi_delta_ok + and response_audio_contract_ok + and response_before_final_commit + and listen_after_response_before_commit + and final_listen_after_commit + ) + mode_contract_ok = validation_mode == "model-policy" or ( + second_response_before_final_commit + and listen_before_first_response + and listen_between_responses + and listen_after_last_done + and second_response_transcript_expectation_ok + ) + ok = common_contract_ok and mode_contract_ok + return { + "ok": ok, + "result_ok": result_ok, + "validation_mode": validation_mode, + "event_counts": _event_counts(events), + "response_ids": response_ids, + "response_summaries": response_summaries, + "min_responses": min_responses, + "effective_min_responses": effective_min_responses, + "min_audio_deltas_per_response": min_audio_deltas_per_response, + "enough_responses": enough_responses, + "response_lifecycle_ok": response_lifecycle_ok, + "multi_delta_ok": multi_delta_ok, + "response_audio_contract_ok": response_audio_contract_ok, + "response_before_final_commit": response_before_final_commit, + "second_response_before_final_commit": second_response_before_final_commit, + "listen_before_first_response": listen_before_first_response, + "listen_between_responses": listen_between_responses, + "listen_after_last_done": listen_after_last_done, + "listen_after_response_before_commit": listen_after_response_before_commit, + "final_listen_after_commit": final_listen_after_commit, + "transcript": transcript, + "second_response_transcript": second_response_transcript, + "expect_second_response_substring": expect_second_response_substring, + "second_response_transcript_expectation_ok": second_response_transcript_expectation_ok, + "error_count": len(error_events), + "cancelled_count": cancelled_count, + "compact_sequence": _compact_sequence(events), + "output_dir": str(output_dir), + } + + +async def run_soft_interrupt(args: argparse.Namespace) -> dict[str, object]: + input_wav = _canonical_path(args.input_wav) + output_dir = _canonical_path(args.output_dir) + if not input_wav.is_file(): + raise ValueError(f"input WAV does not exist: {input_wav}") + input_sha256 = hashlib.sha256(input_wav.read_bytes()).hexdigest() + if args.input_sha256: + input_sha256 = _validate_input_sha256(input_wav, args.input_sha256) + command = [ + sys.executable, + str(DEMO_PATH), + "--url", + args.url, + "--model", + args.model, + "--input-wav", + str(input_wav), + "--output-dir", + str(output_dir), + "--chunk-ms", + 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))]) + if args.require_audio: + command.append("--require-audio") + if args.no_realtime_pacing: + command.append("--no-realtime-pacing") + + process = await asyncio.create_subprocess_exec( + *command, + cwd=str(REPO_ROOT), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout_bytes, stderr_bytes = await asyncio.wait_for( + process.communicate(), + timeout=args.timeout_s + 30.0, + ) + except TimeoutError: + process.kill() + stdout_bytes, stderr_bytes = await process.communicate() + return { + "ok": False, + "returncode": -1, + "stdout_tail": stdout_bytes.decode(errors="replace")[-4000:], + "stderr_tail": (stderr_bytes.decode(errors="replace") + "\nclient timeout")[-4000:], + "output_dir": str(output_dir), + } + + summary = summarize_artifacts( + output_dir=output_dir, + validation_mode=args.validation_mode, + min_responses=args.min_responses, + min_audio_deltas_per_response=args.min_audio_deltas_per_response, + expect_second_response_substring=args.expect_second_response_substring, + ) + summary.update( + { + "returncode": process.returncode if process.returncode is not None else -1, + "input_wav": str(input_wav), + "input_sha256": input_sha256, + "stdout_tail": stdout_bytes.decode(errors="replace")[-4000:], + "stderr_tail": stderr_bytes.decode(errors="replace")[-4000:], + } + ) + summary["ok"] = bool(summary["ok"] and summary["returncode"] == 0) + summary_output = Path(args.summary_output) if args.summary_output else output_dir / "summary.json" + summary_output.parent.mkdir(parents=True, exist_ok=True) + summary_output.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8") + return summary + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime?duplex=1") + parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5") + parser.add_argument("--input-wav", required=True) + parser.add_argument("--ref-audio", required=True) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--summary-output") + parser.add_argument("--chunk-ms", type=int, default=200) + parser.add_argument("--timeout-s", type=float, default=120.0) + parser.add_argument("--require-audio", action="store_true") + parser.add_argument("--no-realtime-pacing", action="store_true") + parser.add_argument( + "--validation-mode", + choices=("model-policy", "response-required"), + default="model-policy", + ) + parser.add_argument("--min-responses", type=int, default=2) + parser.add_argument("--min-audio-deltas-per-response", type=int, default=2) + parser.add_argument("--input-sha256") + parser.add_argument("--expect-second-response-substring") + args = parser.parse_args() + if args.validation_mode == "response-required" and args.min_responses < 2: + parser.error("--min-responses must be at least 2") + if args.validation_mode == "response-required" and not args.input_sha256: + parser.error("--input-sha256 is required in response-required mode") + if args.validation_mode == "response-required" and not args.expect_second_response_substring: + parser.error("--expect-second-response-substring is required in response-required mode") + if args.min_audio_deltas_per_response < 1: + parser.error("--min-audio-deltas-per-response must be positive") + return args + + +def main() -> None: + result = asyncio.run(run_soft_interrupt(parse_args())) + print(json.dumps(result, ensure_ascii=False, indent=2)) + raise SystemExit(0 if result["ok"] else 1) + + +if __name__ == "__main__": + main() diff --git a/tests/engine/test_async_omni_engine_input.py b/tests/engine/test_async_omni_engine_input.py index a1ec472f6..91fdc752d 100644 --- a/tests/engine/test_async_omni_engine_input.py +++ b/tests/engine/test_async_omni_engine_input.py @@ -1,4 +1,5 @@ import pytest +import torch from pytest_mock import MockerFixture from vllm.sampling_params import SamplingParams from vllm.v1.engine import EngineCoreRequest @@ -64,6 +65,48 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock output_processor.add_request.assert_not_called() +def test_build_add_request_message_preserves_model_intermediate_buffer(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + params = SamplingParams(max_tokens=8) + engine.default_sampling_params_list = [params] + engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")] + engine.supported_tasks = ("speech",) + + input_processor = mocker.Mock() + input_processor.process_inputs.return_value = _make_engine_core_request() + engine.input_processor = input_processor + + output_processor = mocker.Mock() + engine.output_processors = [output_processor] + + hidden = torch.arange(6, dtype=torch.float32).reshape(2, 3) + prompt = { + "prompt_token_ids": [1, 1, 1], + "model_intermediate_buffer": { + "ids": {"tts": [11, 12]}, + "hidden_states": {"tts": hidden}, + }, + } + + msg = engine._build_add_request_message( + request_id="req-1", + prompt=prompt, + sampling_params_list=[params], + final_stage_id=0, + arrival_time=0.0, + ) + + request = msg.prompt + assert isinstance(request, OmniEngineCoreRequest) + assert request.additional_information is not None + assert request.additional_information.entries["global_request_id"].list_data == ["req-1"] + assert request.additional_information.entries["omni_final_stage_id"].scalar_data == 0 + assert isinstance(request.model_intermediate_buffer, dict) + info = request.model_intermediate_buffer + assert info["ids"]["tts"] == [11, 12] + assert torch.equal(info["hidden_states"]["tts"], hidden) + + def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixture): engine = object.__new__(AsyncOmniEngine) params = SamplingParams(max_tokens=8) diff --git a/tests/engine/test_async_omni_engine_outputs.py b/tests/engine/test_async_omni_engine_outputs.py index ecf3ee9ee..a1b7a536e 100644 --- a/tests/engine/test_async_omni_engine_outputs.py +++ b/tests/engine/test_async_omni_engine_outputs.py @@ -4,13 +4,33 @@ 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 +from types import SimpleNamespace import pytest 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.messages import ErrorMessage, OutputMessage +from vllm_omni.engine.messages import ( + CollectiveRPCResultMessage, + ErrorMessage, + OutputMessage, +) +from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient +from vllm_omni.experimental.fullduplex.engine.duplex_control_client import ( + DuplexControlClient, + DuplexControlRequestError, +) +from vllm_omni.experimental.fullduplex.engine.messages import ( + AppendDuplexInputMessage, + DuplexControlResultMessage, + DuplexFence, + SignalDuplexTurnMessage, +) from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -26,6 +46,33 @@ def _make_engine(output_queue, mocker: MockerFixture, *, thread_alive: bool = Tr return engine +def test_weak_shutdown_closes_rpc_router_before_joining_orchestrator(mocker: MockerFixture): + request_queue = mocker.MagicMock() + output_queue = mocker.MagicMock() + rpc_output_queue = mocker.MagicMock() + router = mocker.MagicMock() + orchestrator_thread = mocker.MagicMock() + orchestrator_thread.is_alive.return_value = True + + def assert_router_closed_before_join(*args, **kwargs): + router.close.assert_called_once_with() + + orchestrator_thread.join.side_effect = assert_router_closed_before_join + + weak_shutdown_async_omni_engine( + orchestrator_thread, + request_queue, + output_queue, + rpc_output_queue, + router, + ) + + request_queue.sync_q.put.assert_called_once() + assert request_queue.sync_q.put.call_args.kwargs["timeout"] > 0 + orchestrator_thread.join.assert_called_once() + assert orchestrator_thread.join.call_args.kwargs["timeout"] > 0 + + def test_try_get_output_raises_after_orchestrator_dies(mocker: MockerFixture): """Draining remaining results then hitting an empty queue with a dead orchestrator must raise RuntimeError so callers know the pipeline is gone.""" @@ -56,18 +103,17 @@ def test_try_get_output_raises_after_orchestrator_dies(mocker: MockerFixture): @pytest.mark.asyncio async def test_try_get_output_async_raises_after_orchestrator_dies(mocker: MockerFixture): """Same scenario as above but for the async variant.""" - mock_queue = mocker.MagicMock() - mock_queue.sync_q.get_nowait.side_effect = [ + raw_queue = queue.Queue() + raw_queue.put_nowait( OutputMessage( request_id="r1", stage_id=0, engine_outputs=OmniRequestOutput(request_id="r1"), finished=False, - ), - queue.Empty, - ] + ) + ) - engine = _make_engine(mock_queue, mocker, thread_alive=True) + engine = _make_engine(SimpleNamespace(sync_q=raw_queue), mocker, thread_alive=True) assert (await engine.try_get_output_async()).request_id == "r1" @@ -102,12 +148,495 @@ async def test_fatal_error_message_surfaces_through_try_get_output_async(mocker: """Async variant of the fatal error message test.""" fatal_msg = ErrorMessage(error="Orchestrator thread crashed", fatal=True) - mock_queue = mocker.MagicMock() - mock_queue.sync_q.get_nowait.return_value = fatal_msg + raw_queue = queue.Queue() + raw_queue.put_nowait(fatal_msg) - engine = _make_engine(mock_queue, mocker, thread_alive=False) + engine = _make_engine(SimpleNamespace(sync_q=raw_queue), mocker, thread_alive=False) msg = await engine.try_get_output_async() assert msg is not None assert msg.type == "error" assert msg.fatal is True + + +def test_output_remains_on_shared_output_path(mocker: MockerFixture): + raw_queue = queue.Queue() + raw_queue.put_nowait( + OutputMessage( + request_id="legacy-duplex-request", + stage_id=0, + engine_outputs=OmniRequestOutput(request_id="legacy-duplex-request"), + finished=False, + ) + ) + engine = _make_engine(SimpleNamespace(sync_q=raw_queue), mocker) + + 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() + + +def test_collective_rpc_preserves_request_queue_backpressure(mocker: MockerFixture): + class SignallingQueue(queue.Queue): + def __init__(self) -> None: + super().__init__(maxsize=1) + self.put_attempted = threading.Event() + + def put(self, item, block=True, timeout=None): + self.put_attempted.set() + return super().put(item, block=block, timeout=timeout) + + request_q = SignallingQueue() + request_q.put("queue-is-full") + request_q.put_attempted.clear() + 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 + mocker.patch("vllm_omni.engine.async_omni_engine.uuid.uuid4", return_value=SimpleNamespace(hex="blocked-rpc")) + + with ThreadPoolExecutor(max_workers=1) as executor: + pending = executor.submit(engine.collective_rpc, "health", timeout=1) + assert request_q.put_attempted.wait(timeout=1) + assert not pending.done() + + assert request_q.get(timeout=1) == "queue-is-full" + assert request_q.get(timeout=1).rpc_id == "blocked-rpc" + rpc_q.put( + CollectiveRPCResultMessage( + rpc_id="blocked-rpc", + method="health", + stage_ids=[0], + results=["healthy"], + ) + ) + assert pending.result(timeout=1) == ["healthy"] + + 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 + engine._weak_finalizer = None + engine.request_queue = mocker.MagicMock() + engine.request_queue.sync_q.put.side_effect = queue.Full + engine.output_queue = mocker.MagicMock() + engine.rpc_output_queue = mocker.MagicMock() + engine._correlated_rpc_client = mocker.MagicMock() + engine.orchestrator_thread = mocker.MagicMock() + engine.orchestrator_thread.is_alive.return_value = True + engine._runtime = mocker.MagicMock() + + engine.shutdown() + + engine.request_queue.sync_q.put.assert_called_once() + assert engine.request_queue.sync_q.put.call_args.kwargs["timeout"] > 0 + assert engine.orchestrator_thread.join.call_args_list[0].kwargs["timeout"] > 0 + engine._correlated_rpc_client.close.assert_called_once_with() + engine.request_queue.close.assert_called_once_with() + engine.output_queue.close.assert_called_once_with() + engine.rpc_output_queue.close.assert_called_once_with() + engine._runtime.shutdown.assert_not_called() + + +def test_shutdown_releases_runtime_after_orchestrator_stops(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine._weak_finalizer = None + engine.request_queue = mocker.MagicMock() + engine.output_queue = mocker.MagicMock() + engine.rpc_output_queue = mocker.MagicMock() + engine._correlated_rpc_client = mocker.MagicMock() + engine.orchestrator_thread = mocker.MagicMock() + engine.orchestrator_thread.is_alive.side_effect = [True, False] + engine._runtime = mocker.MagicMock() + + engine.shutdown() + + engine.orchestrator_thread.join.assert_called_once() + engine._runtime.shutdown.assert_called_once_with() + + +def test_shutdown_defers_runtime_release_until_live_orchestrator_stops(mocker: MockerFixture): + orchestrator_stopped = threading.Event() + runtime_released = threading.Event() + + class ControllableOrchestratorThread: + def __init__(self) -> None: + self.join_timeouts: list[float | None] = [] + + def is_alive(self) -> bool: + return not orchestrator_stopped.is_set() + + def join(self, timeout: float | None = None) -> None: + self.join_timeouts.append(timeout) + if timeout is None: + orchestrator_stopped.wait() + + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine._weak_finalizer = None + engine.request_queue = mocker.MagicMock() + engine.output_queue = mocker.MagicMock() + engine.rpc_output_queue = mocker.MagicMock() + engine._correlated_rpc_client = mocker.MagicMock() + engine.orchestrator_thread = ControllableOrchestratorThread() + engine._runtime = mocker.MagicMock() + engine._runtime.shutdown.side_effect = runtime_released.set + + try: + engine.shutdown() + + engine._runtime.shutdown.assert_not_called() + orchestrator_stopped.set() + assert runtime_released.wait(timeout=1) + engine._runtime.shutdown.assert_called_once_with() + assert engine.orchestrator_thread.join_timeouts[0] is not None + assert engine.orchestrator_thread.join_timeouts[-1] is None + finally: + orchestrator_stopped.set() + + +def test_weak_shutdown_is_bounded_when_request_queue_is_full(mocker: MockerFixture): + request_queue = mocker.MagicMock() + request_queue.sync_q.put.side_effect = queue.Full + output_queue = mocker.MagicMock() + rpc_output_queue = mocker.MagicMock() + router = mocker.MagicMock() + orchestrator_thread = mocker.MagicMock() + orchestrator_thread.is_alive.return_value = True + + weak_shutdown_async_omni_engine( + orchestrator_thread, + request_queue, + output_queue, + rpc_output_queue, + router, + ) + + request_queue.sync_q.put.assert_called_once() + assert request_queue.sync_q.put.call_args.kwargs["timeout"] > 0 + orchestrator_thread.join.assert_called_once() + assert orchestrator_thread.join.call_args.kwargs["timeout"] > 0 + request_queue.close.assert_called_once_with() + output_queue.close.assert_called_once_with() + rpc_output_queue.close.assert_called_once_with() diff --git a/tests/engine/test_correlated_rpc_client.py b/tests/engine/test_correlated_rpc_client.py new file mode 100644 index 000000000..2f470a996 --- /dev/null +++ b/tests/engine/test_correlated_rpc_client.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import queue +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from vllm_omni.engine.messages import ( + CollectiveRPCRequestMessage, + CollectiveRPCResultMessage, + ErrorMessage, +) +from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _request(rpc_id: str) -> CollectiveRPCRequestMessage: + return CollectiveRPCRequestMessage( + rpc_id=rpc_id, + method="health", + timeout=1, + args=(), + kwargs={}, + stage_ids=None, + ) + + +def test_correlated_rpc_client_unregisters_timeout_before_late_result() -> None: + request_queue: queue.Queue = queue.Queue() + result_queue: queue.Queue = queue.Queue() + client = CorrelatedRpcClient(request_queue, result_queue) + + try: + with pytest.raises(TimeoutError, match="first timed out"): + client.execute( + ("collective", "first"), + _request("first"), + timeout=0.01, + timeout_message="first timed out", + ) + assert request_queue.get_nowait().rpc_id == "first" + + result_queue.put( + CollectiveRPCResultMessage( + rpc_id="first", + method="health", + stage_ids=[0], + results=["late"], + ) + ) + result_queue.put( + CollectiveRPCResultMessage( + rpc_id="second", + method="health", + stage_ids=[0], + results=["current"], + ) + ) + current = client.execute( + ("collective", "second"), + _request("second"), + timeout=1, + timeout_message="second timed out", + ) + + assert isinstance(current, CollectiveRPCResultMessage) + assert current.results == ["current"] + finally: + client.close() + + +def test_correlated_rpc_client_rejects_after_fatal_without_enqueuing() -> None: + request_queue: queue.Queue = queue.Queue() + result_queue: queue.Queue = queue.Queue() + client = CorrelatedRpcClient(request_queue, result_queue) + + try: + with ThreadPoolExecutor(max_workers=1) as executor: + pending = executor.submit( + client.execute, + ("collective", "pending"), + _request("pending"), + timeout=1, + timeout_message="unexpected timeout", + ) + assert request_queue.get(timeout=1).rpc_id == "pending" + result_queue.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"): + client.execute( + ("collective", "after-fatal"), + _request("after-fatal"), + timeout=1, + timeout_message="unexpected timeout", + ) + assert request_queue.empty() + finally: + client.close() diff --git a/tests/engine/test_duplex_import_boundary.py b/tests/engine/test_duplex_import_boundary.py new file mode 100644 index 000000000..6eae2e698 --- /dev/null +++ b/tests/engine/test_duplex_import_boundary.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def _assert_isolated_import_succeeds(script: str) -> None: + 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_stable_engine_imports_do_not_load_experimental_duplex() -> None: + _assert_isolated_import_succeeds(""" +import sys + +import vllm_omni.engine.async_omni_engine +import vllm_omni.engine.orchestrator +import vllm_omni.entrypoints.async_omni + +loaded = sorted( + name + for name in sys.modules + if name == "vllm_omni.experimental.fullduplex" + or name.startswith("vllm_omni.experimental.fullduplex.") +) +if loaded: + raise SystemExit("stable imports loaded experimental duplex: " + ", ".join(loaded)) +""") + + +def test_stable_engine_does_not_expose_duplex_contract_modules() -> None: + _assert_isolated_import_succeeds(""" +import importlib.util + +from vllm_omni.engine import messages + +legacy_modules = ( + "vllm_omni.engine.duplex_contracts", + "vllm_omni.engine.duplex_lease", + "vllm_omni.engine.resumable", +) +present = [name for name in legacy_modules if importlib.util.find_spec(name) is not None] +if present: + raise SystemExit("stable duplex modules still exist: " + ", ".join(present)) + +duplex_exports = sorted(name for name in vars(messages) if name.startswith("Duplex")) +if duplex_exports: + raise SystemExit("stable messages still expose duplex contracts: " + ", ".join(duplex_exports)) +""") + + +def test_stable_outputs_do_not_declare_duplex_decision_field() -> None: + _assert_isolated_import_succeeds(""" +import dataclasses + +from vllm_omni.outputs import OmniRequestOutput + +fields = {field.name for field in dataclasses.fields(OmniRequestOutput)} +if "duplex_output_decision" in fields: + raise SystemExit("stable output declares duplex_output_decision") +""") + + +def test_stable_model_executor_does_not_expose_duplex_helper_module() -> None: + _assert_isolated_import_succeeds(""" +import importlib.util + +if importlib.util.find_spec("vllm_omni.model_executor.duplex") is not None: + raise SystemExit("stable model_executor still exposes duplex helper module") +""") + + +def test_runtime_package_does_not_bundle_the_browser_demo() -> None: + assert not (REPO_ROOT / "vllm_omni" / "experimental" / "fullduplex" / "web").exists() + + +def test_experimental_engine_uses_canonical_contract_module_names() -> None: + engine_dir = REPO_ROOT / "vllm_omni" / "experimental" / "fullduplex" / "engine" + core_dir = REPO_ROOT / "vllm_omni" / "experimental" / "fullduplex" / "core" + + assert (engine_dir / "contracts.py").is_file() + assert (engine_dir / "lease.py").is_file() + assert (engine_dir / "messages.py").is_file() + assert not (engine_dir / "duplex_lease.py").exists() + assert not (engine_dir / "duplex_types.py").exists() + assert not (core_dir / "identity.py").exists() diff --git a/tests/engine/test_membership_controller.py b/tests/engine/test_membership_controller.py index c00a7f646..1b8bf2adb 100644 --- a/tests/engine/test_membership_controller.py +++ b/tests/engine/test_membership_controller.py @@ -11,6 +11,7 @@ import pytest from vllm_omni.distributed.omni_coordinator.messages import ReplicaStatus from vllm_omni.engine.membership_controller import MembershipController +from vllm_omni.engine.stage_pool import StagePool pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -45,10 +46,10 @@ class FakePool: def attach_load_balancer(self, lb): self.lb = lb - def add_client(self, input_addr, client): - self.added.append((input_addr, client)) + def add_client(self, input_addr, client, *, replica_id=None): + self.added.append((input_addr, client, replica_id)) self.clients.append(client) - return len(self.clients) - 1 + return len(self.clients) - 1 if replica_id is None else replica_id def invalidate_addr(self, input_addr): self.invalidated.append(input_addr) @@ -155,4 +156,35 @@ async def test_do_register_offloads_remote_factory(monkeypatch): assert factory_thread_ids assert factory_thread_ids[0] != threading.get_ident() - assert pool.added[0][0] == "tcp://stage-0-replica-1" + assert pool.added[0] == ("tcp://stage-0-replica-1", pool.clients[0], 1) + + +@pytest.mark.asyncio +async def test_unregister_then_register_restores_coordinator_replica_slot(monkeypatch): + class Client: + def __init__(self, input_addr: str) -> None: + self.client_addresses = {"input_address": input_addr} + self.shutdown_calls = 0 + + def shutdown(self) -> None: + self.shutdown_calls += 1 + + old_clients = [Client("tcp://old-0"), Client("tcp://old-1")] + replacement = Client("tcp://new-1") + pool = StagePool(0, old_clients) + + def factory(stage_id: int, replica_id: int): + assert (stage_id, replica_id) == (0, 1) + return replacement + + controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=factory) + await controller.handle_unregister(0, "tcp://old-0") + await controller.handle_unregister(0, "tcp://old-1") + + await controller.handle_register(0, 1) + await controller.drain_tasks(timeout=1) + + assert pool.clients == [None, replacement] + assert pool.get_replica_id_by_addr("tcp://new-1") == 1 + assert pool.get_replica_id_by_addr("tcp://old-1") is None + assert pool.available_replica_ids() == [1] diff --git a/tests/engine/test_multimodal_accumulation.py b/tests/engine/test_multimodal_accumulation.py new file mode 100644 index 000000000..ca449154c --- /dev/null +++ b/tests/engine/test_multimodal_accumulation.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm_omni.outputs.mm_outputs import MultimodalPayload +from vllm_omni.outputs.multimodal_accumulation import ( + drain_delta_payload, + is_non_final_delta_audio_chunk, + replace_snapshot_keys, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_chunk_accumulation_policy_replaces_snapshots_and_drains_delta_state(): + accumulated = MultimodalPayload.from_dict( + { + "audio": torch.tensor([1.0]), + "meta.tts_is_last_chunk": torch.tensor([0]), + "meta.turn_end": torch.tensor([0]), + "meta.stable_request_value": "keep", + } + ) + incoming = MultimodalPayload.from_dict( + { + "audio": torch.tensor([2.0]), + "meta.tts_is_last_chunk": torch.tensor([1]), + "meta.turn_end": torch.tensor([1]), + } + ) + assert accumulated is not None + assert incoming is not None + + replace_snapshot_keys(accumulated, incoming) + merged = accumulated.merged_with(incoming) + + assert not is_non_final_delta_audio_chunk(merged, "audio") + + drain_delta_payload(merged) + + assert "audio" not in merged + assert "meta.tts_is_last_chunk" not in merged + assert "meta.turn_end" not in merged + assert merged.metadata["meta.stable_request_value"] == "keep" diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 23f8cfaa2..68d9c2c58 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -14,12 +14,14 @@ import janus import pytest from vllm.outputs import CompletionOutput, RequestOutput from vllm.sampling_params import SamplingParams +from vllm.v1.engine.exceptions import EngineDeadError from vllm_omni.engine.messages import ( AbortRequestMessage, AddCompanionRequestMessage, CollectiveRPCRequestMessage, CollectiveRPCResultMessage, + ErrorMessage, OutputMessage, ShutdownRequestMessage, StageSubmissionMessage, @@ -31,12 +33,68 @@ from vllm_omni.engine.orchestrator import ( _infer_stage_audio_sample_rate, ) from vllm_omni.engine.stage_pool import StagePool +from vllm_omni.experimental.fullduplex.engine.duplex_control_plane import DuplexControlPlane +from vllm_omni.experimental.fullduplex.engine.duplex_runtime import ( + DuplexInputMode, + DuplexRuntimeCapabilities, + DuplexSessionRuntimeState, + duplex_resource_request_id, +) +from vllm_omni.experimental.fullduplex.engine.messages import ( + AppendDuplexInputMessage, + CloseDuplexSessionMessage, + DuplexFence, + OpenDuplexSessionMessage, + SignalDuplexTurnMessage, +) +from vllm_omni.experimental.fullduplex.minicpmo45.runtime import MiniCPMO45DuplexRuntimeExtension from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +class FakeRunningCounter: + def __init__(self) -> None: + self.value = 0 + + def increment(self) -> None: + self.value += 1 + + def decrement(self) -> None: + self.value -= 1 + + +@pytest.mark.asyncio +async def test_engine_dead_broadcasts_fatal_to_rpc_waiters(monkeypatch: pytest.MonkeyPatch) -> None: + rpc_queue: asyncio.Queue = asyncio.Queue() + orchestrator = Orchestrator( + request_async_queue=asyncio.Queue(), + output_async_queue=asyncio.Queue(), + rpc_async_queue=rpc_queue, + stage_pools=[], + ) + orchestrator._fatal_error = "stage engine died" + orchestrator._fatal_error_stage_id = 2 + + async def wait_for_requests() -> None: + await asyncio.Event().wait() + + async def fail_outputs() -> None: + raise EngineDeadError("stage engine died") + + monkeypatch.setattr(orchestrator, "_request_handler", wait_for_requests) + monkeypatch.setattr(orchestrator, "_orchestration_output_handler", fail_outputs) + + await orchestrator.run() + + fatal = rpc_queue.get_nowait() + assert isinstance(fatal, ErrorMessage) + assert fatal.fatal is True + assert fatal.error == "stage engine died" + assert fatal.stage_id == 2 + + @dataclass class OrchestratorFixture: orchestrator: Orchestrator @@ -724,6 +782,778 @@ async def test_run_abort(orchestrator_factory) -> None: await _shutdown_orchestrator(orchestrator_fixture) +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 # --------------------------------------------------------------------------- @@ -898,6 +1728,50 @@ async def test_stage_pool_submit_update_reuses_existing_binding() -> None: assert stage0_r1.add_request_calls[1][0].request_id == "req-1" +@pytest.mark.asyncio +async def test_stage_pool_failed_replica_releases_distributed_affinity_and_stops_polling() -> None: + stage0 = FakeStageClient(stage_type="llm", final_output=False) + pool = StagePool( + 0, + [stage0], + output_processor=FakeOutputProcessor(), + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ) + pool._addr_to_replica_id["tcp://replica-0"] = 0 + pool.bind("distributed-request", "tcp://replica-0") + pool._request_bindings["legacy-request"] = 0 + + affected = pool.mark_replica_unavailable(0) + + assert set(affected) == {"distributed-request", "legacy-request"} + assert pool.get_bound_replica_id("distributed-request") is None + assert pool.get_bound_replica_id("legacy-request") is None + assert await pool.poll_llm_raw_output(0) is None + + +def test_stage_pool_reattached_replica_becomes_available_again() -> None: + failed_client = FakeStageClient(stage_type="llm", final_output=False) + replacement_client = FakeStageClient(stage_type="llm", final_output=False) + pool = StagePool( + 0, + [failed_client], + output_processor=FakeOutputProcessor(), + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ) + input_addr = "tcp://replica-0" + pool._addr_to_replica_id[input_addr] = 0 + pool.mark_replica_unavailable(0) + removed = pool.remove_client(input_addr) + + replica_id = pool.add_client(input_addr, replacement_client) + + assert removed is failed_client + assert replica_id == 0 + assert len(pool.clients) == 1 + assert pool.clients[0] is replacement_client + assert pool.available_replica_ids() == [0] + + @pytest.mark.asyncio async def test_stage_pool_submit_update_refreshes_output_processor_state() -> None: output_processor = FakeOutputProcessor() @@ -980,6 +1854,73 @@ async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> Non assert orchestrator.request_states["req-stream"].streaming.enabled is True +@pytest.mark.asyncio +async def test_resumable_segment_boundary_builds_stage_metrics() -> None: + built_metrics = SimpleNamespace(pipeline_timings={}) + + class RecordingPool: + def __init__(self) -> None: + self.calls: list[list[Any]] = [] + + def build_stage_metrics(self, outputs, **_kwargs): + self.calls.append(outputs) + return built_metrics + + pool = RecordingPool() + orchestrator = object.__new__(Orchestrator) + req_state = OrchestratorRequestState( + request_id="req-stream", + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + req_state.streaming.enabled = True + req_state.streaming.segment_finished = True + req_state.stage_submit_ts[0] = time.time() + orchestrator.request_states = {"req-stream": req_state} + orchestrator.stage_pools = [pool] + routed: list[Any] = [] + + async def record_route(_stage_id, _replica_id, _output, _req_state, stage_metrics): + routed.append(stage_metrics) + + orchestrator._route_output = record_route + output = SimpleNamespace(request_id="req-stream", error=None, finished=False) + + await orchestrator._handle_processed_outputs(0, 0, [output]) + + assert pool.calls == [[output]] + assert routed == [built_metrics] + + +def test_stage_pool_metrics_use_resumable_segment_token_count() -> None: + class SegmentMetricsOutputProcessor(FakeOutputProcessor): + def pop_native_text_metrics(self, request_id: str) -> dict[str, Any]: + assert request_id == "req-stream" + return {"num_generation_tokens": 3} + + stage0 = FakeStageClient(stage_type="llm", final_output=False) + pool = StagePool( + 0, + [stage0], + output_processor=SegmentMetricsOutputProcessor(), + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ) + output = SimpleNamespace( + request_id="req-stream", + outputs=[SimpleNamespace(cumulative_token_ids=list(range(11)))], + ) + + metrics = pool.build_stage_metrics( + [output], + submit_ts=time.time(), + request_timestamp=time.time(), + replica_id=0, + ) + + assert metrics.num_tokens_out == 3 + assert metrics.output_unit_count == 3 + + @pytest.mark.asyncio async def test_stage_pool_submit_initial_rolls_back_output_processor_when_client_submit_fails() -> None: class FailingStageClient(FakeStageClient): @@ -1179,3 +2120,132 @@ def test_orchestrator_does_not_re_introduce_global_stats_throttle() -> None: "raw_outputs.scheduler_stats being non-None — the per-scheduler 1Hz " "throttle in OmniSchedulerMixin.make_stats() is the only gate needed." ) + + +@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 +async def test_duplex_reaper_loop_survives_one_cleanup_failure(): + class _Plane: + def __init__(self) -> None: + self.calls = 0 + + async def reap_expired(self) -> int: + self.calls += 1 + if self.calls == 1: + raise RuntimeError("transient cleanup failure") + 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 orchestrator.duplex_control_plane.calls >= 2 + + +@pytest.mark.asyncio +async def test_abort_retry_does_not_repeat_successful_stage_abort(): + class _Pool: + def __init__(self, *, fail_once: bool = False) -> None: + self.bound = {"req-a"} + self.fail_once = fail_once + self.physical_abort_calls = 0 + + async def abort_requests(self, request_ids: list[str]) -> None: + active = self.bound.intersection(request_ids) + if not active: + return + self.physical_abort_calls += 1 + if self.fail_once: + self.fail_once = False + raise RuntimeError("stage abort failed") + + def release_bindings(self, request_ids: list[str]) -> None: + self.bound.difference_update(request_ids) + + first = _Pool() + second = _Pool(fail_once=True) + orchestrator = object.__new__(Orchestrator) + orchestrator.stage_pools = [first, second] + + with pytest.raises(RuntimeError, match="stage abort failed"): + await orchestrator._abort_request_ids(["req-a"]) + + assert first.bound == set() + assert first.physical_abort_calls == 1 + assert second.bound == {"req-a"} + + await orchestrator._abort_request_ids(["req-a"]) + + 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) + + orchestrator = object.__new__(Orchestrator) + orchestrator.stage_pools = [_FailingPool()] + orchestrator.duplex_control_plane = _Plane() + orchestrator._pd_kv_params = {} + orchestrator.request_states = {} + orchestrator._running_counter = None + + 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_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 27ee11df4..3abb244eb 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -38,6 +38,7 @@ class FakeStageClient: self.custom_process_input_func = None self.next_inputs = list(next_inputs or []) self.add_request_calls: list[tuple[Any, ...]] = [] + self.decoded_source_tokens: str | None = None self._engine_core_outputs = queue.Queue() async def add_request_async(self, *args, **_kwargs) -> None: @@ -50,6 +51,9 @@ class FakeStageClient: return SimpleNamespace(outputs=[]) def process_engine_inputs(self, _source_outputs, prompt=None, streaming_context=None): + decoder = getattr(streaming_context, "source_token_decoder", None) + if callable(decoder): + self.decoded_source_tokens = decoder([11, 12], skip_special_tokens=True) return list(self.next_inputs) async def abort_requests_async(self, _request_ids: list[str]) -> None: @@ -66,6 +70,9 @@ class FakeStageClient: class FakeOutputProcessor: + def __init__(self, tokenizer=None) -> None: + self.tokenizer = tokenizer + def add_request(self, *args, **kwargs) -> None: return None @@ -108,6 +115,11 @@ def _request_output(request_id: str) -> RequestOutput: @pytest.mark.asyncio async def test_forward_text_prompt_uses_target_stage_input_processor() -> None: + class SourceTokenizer: + def decode(self, token_ids, *, skip_special_tokens): + assert skip_special_tokens is True + return ":".join(str(token_id) for token_id in token_ids) + stage0 = FakeStageClient(final_output=True) stage1 = FakeStageClient( final_output=True, @@ -117,7 +129,7 @@ async def test_forward_text_prompt_uses_target_stage_input_processor() -> None: StagePool( 0, [stage0], - output_processor=FakeOutputProcessor(), + output_processor=FakeOutputProcessor(tokenizer=SourceTokenizer()), stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), ), StagePool( @@ -150,6 +162,8 @@ async def test_forward_text_prompt_uses_target_stage_input_processor() -> None: assert input_processor.calls assert input_processor.calls[0]["prompt"] == {"prompt": "hello", "multi_modal_data": {"video": ["frame"]}} + assert stage1.decoded_source_tokens == "11:12" + assert req_state.streaming.source_token_decoder is None assert stage1.add_request_calls submitted_request = stage1.add_request_calls[0][0] assert submitted_request.prompt_token_ids == [101, 102] diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index 86e71fcbc..aef78e47e 100644 --- a/tests/engine/test_output_processor.py +++ b/tests/engine/test_output_processor.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Regression tests for OmniRequestState multimodal DELTA drain and consolidation guard.""" +from types import SimpleNamespace from unittest.mock import MagicMock import numpy as np @@ -10,9 +11,11 @@ import torch from vllm.outputs import PoolingRequestOutput from vllm.sampling_params import RequestOutputKind from vllm.v1.engine import FinishReason +from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor +from vllm_omni.outputs import output_processor from vllm_omni.outputs.output_modality import OutputModalityNames -from vllm_omni.outputs.output_processor import OmniRequestState +from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -62,6 +65,65 @@ def test_init_empty_dict(): assert _make_state(RequestOutputKind.DELTA).mm_accumulated == {} +def test_streaming_update_resets_native_text_metrics_for_next_segment(): + state = _make_state(RequestOutputKind.DELTA) + state.native_text_stats.num_generation_tokens = 12 + state.native_text_stats.first_token_ts = 10.0 + state.native_text_stats.last_token_ts = 10.2 + + state.apply_streaming_update( + SimpleNamespace( + final=False, + prompt=None, + prompt_token_ids=[], + arrival_time=20.0, + ) + ) + + assert state.native_text_stats.arrival_time == 20.0 + assert state.native_text_stats.num_generation_tokens == 0 + assert state.native_text_stats.first_token_ts == 0.0 + assert state.native_text_stats.last_token_ts == 0.0 + + +def test_native_text_tpot_snapshot_matches_vllm_finished_metric_definition(): + calculate = getattr(output_processor, "_mean_time_per_output_token_ms", None) + assert calculate is not None + stats = SimpleNamespace( + num_generation_tokens=4, + first_token_ts=10.0, + last_token_ts=10.045, + ) + + assert calculate(stats) == pytest.approx(15.0) + + +def test_native_text_metrics_include_segment_generation_token_count(monkeypatch): + monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_output", lambda *args, **kwargs: None) + processor = object.__new__(MultimodalOutputProcessor) + processor._native_text_metrics_by_request = {} + processor.lora_states = {} + state = _make_state(RequestOutputKind.DELTA) + state.is_prefilling = False + iteration_stats = MagicMock() + + def update_native_stats(_output, _timestamp, _was_prefilling, native_stats, *_args): + native_stats.num_generation_tokens = 27 + native_stats.first_token_ts = 10.0 + native_stats.last_token_ts = 10.26 + + iteration_stats.update_from_output.side_effect = update_native_stats + + processor._update_stats_from_output( + state, + SimpleNamespace(), + 10.26, + iteration_stats, + ) + + assert processor.pop_native_text_metrics("r")["num_generation_tokens"] == 27 + + def test_delta_drains_output_modality_per_step(): """DELTA drains the mm_type key (output modality) but preserves hidden-state keys.""" s = _make_state(RequestOutputKind.DELTA) @@ -172,6 +234,145 @@ def test_finish_consolidation_drains_mm_delta(): assert AUDIO not in s.mm_accumulated # drained +def test_delta_drains_audio_chunk_metadata_per_step(): + s = _make_state(RequestOutputKind.DELTA) + segment_text = "你好" + segment_utf8 = torch.tensor(list(segment_text.encode("utf-8")), dtype=torch.uint8) + + s.add_multimodal_tensor( + { + "model_outputs": torch.ones(5), + "meta.duplex_epoch": torch.tensor([7], dtype=torch.int32), + "meta.duplex_turn_id": torch.tensor([1], dtype=torch.int32), + "meta.llm_output_text_utf8": segment_utf8, + "meta.audio_text_total_chars": torch.tensor([len(segment_text)], dtype=torch.int32), + "meta.tts_is_last_chunk": torch.tensor([0], dtype=torch.int32), + }, + mm_type=AUDIO, + ) + result1 = s.make_request_output([1], None, FinishReason.STOP, None) + assert result1 is not None and not isinstance(result1, PoolingRequestOutput) + meta1 = result1.outputs[0].multimodal_output["meta"] + assert bytes(meta1["llm_output_text_utf8"].tolist()).decode("utf-8") == segment_text + + s.add_multimodal_tensor( + { + "model_outputs": torch.ones(3), + "meta.duplex_epoch": torch.tensor([8], dtype=torch.int32), + "meta.duplex_turn_id": torch.tensor([2], dtype=torch.int32), + "meta.llm_output_text_utf8": segment_utf8, + "meta.audio_text_total_chars": torch.tensor([len(segment_text)], dtype=torch.int32), + "meta.tts_is_last_chunk": torch.tensor([1], dtype=torch.int32), + }, + mm_type=AUDIO, + ) + result2 = s.make_request_output([2], None, FinishReason.STOP, None) + + assert result2 is not None and not isinstance(result2, PoolingRequestOutput) + meta2 = result2.outputs[0].multimodal_output["meta"] + assert bytes(meta2["llm_output_text_utf8"].tolist()).decode("utf-8") == segment_text + assert meta2["audio_text_total_chars"].tolist() == [len(segment_text)] + assert meta2["duplex_epoch"].tolist() == [8] + assert meta2["duplex_turn_id"].tolist() == [2] + + +def test_cumulative_audio_replaces_chunk_metadata_per_step(): + s = _make_state(RequestOutputKind.CUMULATIVE) + segment_text = "你好" + segment_utf8 = torch.tensor(list(segment_text.encode("utf-8")), dtype=torch.uint8) + + s.add_multimodal_tensor( + { + "model_outputs": torch.ones(5), + "meta.duplex_epoch": torch.tensor([7], dtype=torch.int32), + "meta.duplex_turn_id": torch.tensor([1], dtype=torch.int32), + "meta.llm_output_text_utf8": segment_utf8, + "meta.audio_text_total_chars": torch.tensor([len(segment_text)], dtype=torch.int32), + "meta.tts_is_last_chunk": torch.tensor([0], dtype=torch.int32), + }, + mm_type=AUDIO, + ) + result1 = s.make_request_output([1], None, None, None) + assert result1 is not None and not isinstance(result1, PoolingRequestOutput) + assert result1.outputs[0].multimodal_output["audio"].shape[0] == 5 + + s.add_multimodal_tensor( + { + "model_outputs": torch.ones(3), + "meta.duplex_epoch": torch.tensor([8], dtype=torch.int32), + "meta.duplex_turn_id": torch.tensor([2], dtype=torch.int32), + "meta.llm_output_text_utf8": segment_utf8, + "meta.audio_text_total_chars": torch.tensor([len(segment_text)], dtype=torch.int32), + "meta.tts_is_last_chunk": torch.tensor([1], dtype=torch.int32), + }, + mm_type=AUDIO, + ) + result2 = s.make_request_output([2], None, None, None) + + assert result2 is not None and not isinstance(result2, PoolingRequestOutput) + output = result2.outputs[0].multimodal_output + assert output["audio"].shape[0] == 8 + meta = output["meta"] + assert bytes(meta["llm_output_text_utf8"].tolist()).decode("utf-8") == segment_text + assert meta["audio_text_total_chars"].tolist() == [len(segment_text)] + assert meta["duplex_epoch"].tolist() == [8] + assert meta["duplex_turn_id"].tolist() == [2] + assert meta["tts_is_last_chunk"].tolist() == [1] + + +def test_delta_audio_non_final_tts_chunk_overrides_spurious_finish(): + s = _make_state(RequestOutputKind.DELTA) + s.add_multimodal_tensor( + { + "model_outputs": torch.ones(5), + "meta.tts_is_last_chunk": torch.tensor([0]), + }, + mm_type=AUDIO, + ) + + result = s.make_request_output([1], None, FinishReason.STOP, None) + + assert result is not None and not isinstance(result, PoolingRequestOutput) + assert result.finished is False + assert result.outputs[0].finish_reason is None + assert result.outputs[0].multimodal_output["meta"]["tts_is_last_chunk"].tolist() == [0] + assert torch.equal(result.outputs[0].multimodal_output[AUDIO], torch.ones(5)) + + +def test_delta_audio_final_tts_chunk_keeps_finish(): + s = _make_state(RequestOutputKind.DELTA) + s.add_multimodal_tensor( + { + "model_outputs": torch.ones(5), + "meta.tts_is_last_chunk": torch.tensor([1]), + }, + mm_type=AUDIO, + ) + + result = s.make_request_output([151645], None, FinishReason.STOP, None) + + assert result is not None and not isinstance(result, PoolingRequestOutput) + assert result.finished is True + assert result.outputs[0].finish_reason == "stop" + + +def test_delta_audio_non_final_tts_chunk_accepts_nested_meta(): + s = _make_state(RequestOutputKind.DELTA) + s.add_multimodal_tensor( + { + AUDIO: torch.ones(5), + "meta": {"tts_is_last_chunk": torch.tensor([0])}, + }, + mm_type=AUDIO, + ) + + result = s.make_request_output([1], None, FinishReason.STOP, None) + + assert result is not None and not isinstance(result, PoolingRequestOutput) + assert result.finished is False + assert result.outputs[0].multimodal_output["meta"]["tts_is_last_chunk"].tolist() == [0] + + @pytest.mark.parametrize("mm_type", [AUDIO, "hidden"]) def test_final_only_consolidates_drainable_keys(mm_type): """FINAL_ONLY never drains per-step, so modality keys and hidden state diff --git a/tests/engine/test_rpc_result_router.py b/tests/engine/test_rpc_result_router.py new file mode 100644 index 000000000..ec983392b --- /dev/null +++ b/tests/engine/test_rpc_result_router.py @@ -0,0 +1,109 @@ +import queue +from typing import Literal + +import pytest + +from vllm_omni.engine.messages import ( + CollectiveRPCResultMessage, + EngineQueueMessage, + ErrorMessage, +) +from vllm_omni.engine.rpc_result_router import RpcResultRouter + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class GenericCorrelatedResultMessage(EngineQueueMessage, kw_only=True): + type: Literal["generic_correlated_result"] = "generic_correlated_result" + namespace: str + correlation_id: str + + @property + def rpc_correlation_key(self) -> tuple[str, str]: + return (self.namespace, self.correlation_id) + + +def _generic_result(correlation_id: str) -> GenericCorrelatedResultMessage: + return GenericCorrelatedResultMessage( + namespace="plugin", + correlation_id=correlation_id, + ) + + +def test_rpc_result_router_routes_out_of_order_results_by_correlation_id(): + source: queue.Queue = queue.Queue() + router = RpcResultRouter(source) + first = router.register(("plugin", "first")) + second = router.register(("plugin", "second")) + + source.put(_generic_result("second")) + source.put(_generic_result("first")) + + assert first.get(timeout=1).correlation_id == "first" + assert second.get(timeout=1).correlation_id == "second" + router.close() + + +def test_rpc_result_router_drops_only_the_late_result_after_unregister(): + source: queue.Queue = queue.Queue() + router = RpcResultRouter(source) + expired = router.register(("plugin", "expired")) + active = router.register(("collective", "active")) + router.unregister(("plugin", "expired"), expired) + + source.put(_generic_result("expired")) + source.put( + CollectiveRPCResultMessage( + rpc_id="active", + method="health", + stage_ids=[0], + results=["ok"], + ) + ) + + assert active.get(timeout=1).rpc_id == "active" + assert expired.empty() + router.close() + + +def test_rpc_result_router_broadcasts_fatal_errors_to_pending_waiters(): + source: queue.Queue = queue.Queue() + router = RpcResultRouter(source) + plugin = router.register(("plugin", "one")) + collective = router.register(("collective", "two")) + + source.put(ErrorMessage(error="orchestrator failed", fatal=True)) + + assert plugin.get(timeout=1).error == "orchestrator failed" + assert collective.get(timeout=1).error == "orchestrator failed" + with pytest.raises(RuntimeError, match="orchestrator failed"): + router.register(("plugin", "after-failure")) + router.close() + + +def test_rpc_result_router_does_not_broadcast_uncorrelated_nonfatal_errors(): + source: queue.Queue = queue.Queue() + router = RpcResultRouter(source) + waiter = router.register(("plugin", "active")) + + source.put(ErrorMessage(error="request failed", fatal=False, request_id="other")) + source.put(_generic_result("active")) + + assert waiter.get(timeout=1).correlation_id == "active" + router.close() + + +def test_rpc_result_router_close_unblocks_waiters_and_stops_consumer(): + source: queue.Queue = queue.Queue() + router = RpcResultRouter(source) + waiter = router.register(("plugin", "pending")) + + router.close() + + result = waiter.get(timeout=1) + assert isinstance(result, ErrorMessage) + assert result.fatal is True + assert result.error == "RPC result router closed" + assert not router._thread.is_alive() + with pytest.raises(RuntimeError, match="router is closed"): + router.register(("plugin", "after-close")) diff --git a/tests/entrypoints/openai/test_duplex_protocol.py b/tests/entrypoints/openai/test_duplex_protocol.py new file mode 100644 index 000000000..3c3666026 --- /dev/null +++ b/tests/entrypoints/openai/test_duplex_protocol.py @@ -0,0 +1,474 @@ +import asyncio +import json + +import pytest + +from vllm_omni.experimental.fullduplex.openai.protocol import ( + DuplexCapabilities, + DuplexOverlapPolicy, + DuplexSession, + DuplexSessionConfig, + DuplexSessionRegistry, + DuplexTurnController, + DuplexTurnEventType, + DuplexTurnState, + ResponseCreateOptions, +) +from vllm_omni.experimental.fullduplex.openai.realtime_session import ( + NativeRealtimeSessionProtocol, +) +from vllm_omni.experimental.fullduplex.openai.realtime_state import RealtimeStateOwner + +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 "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": {"minicpmo45_native_duplex": True}, + } + ) + + assert event["session"]["extra_body"]["minicpmo45_native_duplex"] is True + + +def test_realtime_explicit_query_native_duplex_flag_is_available_before_autostart(): + 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"]["minicpmo45_native_duplex"] is True + + +@pytest.mark.asyncio +async def test_realtime_resume_heartbeat_and_event_ack_translate_without_session_autostart(): + protocol = NativeRealtimeSessionProtocol( + { + "model": "openbmb/MiniCPM-o-4_5", + "minicpmo45_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", + "minicpmo45_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_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 + + +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 = DuplexCapabilities.minicpmo45_native(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 False + 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 = DuplexCapabilities.minicpmo45_native(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 new file mode 100644 index 000000000..24a6f5365 --- /dev/null +++ b/tests/entrypoints/openai/test_duplex_session_attachment.py @@ -0,0 +1,387 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import asyncio +import base64 +import json + +import pytest + +from vllm_omni.experimental.fullduplex.openai.session_attachment import ( + DuplexEventJournal, + DuplexJournalGapError, + DuplexJournalOverflowError, + DuplexResumeCredential, + DuplexSessionAttachmentRegistry, + InvalidResumeTokenError, + ResumeToken, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +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 + + +def _decoded_token_bytes(token: ResumeToken) -> bytes: + padding = "=" * (-len(token.plaintext) % 4) + return base64.urlsafe_b64decode(token.plaintext + padding) + + +def test_resume_credential_uses_256_bit_token_digest_and_redacted_repr(monkeypatch) -> None: + token = ResumeToken.generate() + credential = DuplexResumeCredential.from_token(token) + calls = [] + + import vllm_omni.experimental.fullduplex.openai.session_attachment as attachment_module + + original_compare = attachment_module.hmac.compare_digest + + def recording_compare(left, right): + calls.append((left, right)) + return original_compare(left, right) + + monkeypatch.setattr(attachment_module.hmac, "compare_digest", recording_compare) + + assert len(_decoded_token_bytes(token)) == 32 + assert credential.verify(token.plaintext) is True + assert credential.verify("invalid-token") is False + assert len(calls) == 2 + assert token.plaintext not in repr(token) + assert token.plaintext not in repr(credential) + assert not hasattr(credential, "plaintext") + assert len(credential.token_digest) == 32 + + +def test_resume_credential_rotation_revokes_old_token() -> None: + first = ResumeToken.generate() + credential = DuplexResumeCredential.from_token(first) + + second = credential.rotate() + + assert credential.verify(first.plaintext) is False + assert credential.verify(second.plaintext) is True + assert second.plaintext != first.plaintext + + +def test_event_journal_sequences_acknowledges_and_replays_exact_payloads() -> None: + clock = _Clock() + journal = DuplexEventJournal(max_bytes=4096, ttl_s=60.0, clock=clock) + + first = journal.record({"type": "response.audio.delta", "delta": "AAAA"}) + second = journal.record({"type": "response.done", "response_id": "resp-1"}) + + assert first.sequence == 1 + assert second.sequence == 2 + assert first.payload["server_event_seq"] == 1 + assert second.payload["server_event_seq"] == 2 + assert first.encoded_bytes == len( + json.dumps(dict(first.payload), separators=(",", ":"), ensure_ascii=False).encode("utf-8") + ) + assert [entry.sequence for entry in journal.replay_after(0)] == [1, 2] + assert journal.acknowledge(1) == 1 + assert [entry.sequence for entry in journal.replay_after(1)] == [2] + assert journal.acknowledge(1) == 0 + + +def test_event_journal_ttl_pruning_reports_replay_gap() -> None: + clock = _Clock() + journal = DuplexEventJournal(max_bytes=4096, ttl_s=5.0, clock=clock) + journal.record({"type": "one"}) + clock.advance(3.0) + journal.record({"type": "two"}) + clock.advance(3.0) + + assert journal.prune() == 1 + with pytest.raises(DuplexJournalGapError, match="older than retained"): + journal.replay_after(0) + assert [entry.payload["type"] for entry in journal.replay_after(1)] == ["two"] + + +def test_event_journal_overflow_is_explicit_and_does_not_evict_silently() -> None: + clock = _Clock() + journal = DuplexEventJournal(max_bytes=100, ttl_s=60.0, clock=clock) + first = journal.record({"type": "small", "value": "x"}) + + with pytest.raises(DuplexJournalOverflowError, match="byte limit"): + journal.record({"type": "large", "value": "x" * 200}) + + assert journal.overflowed is True + assert journal.retained_bytes == first.encoded_bytes + with pytest.raises(DuplexJournalGapError, match="overflowed"): + journal.replay_after(0) + + +@pytest.mark.asyncio +async def test_registry_resume_rotates_token_replays_and_atomically_replaces_attachment() -> None: + clock = _Clock() + registry = DuplexSessionAttachmentRegistry( + replay_ttl_s=60.0, + replay_max_bytes_per_session=4096, + clock=clock, + ) + sends_a = [] + closes_a = [] + sends_b = [] + closes_b = [] + + async def send_a(payload): + sends_a.append(payload) + + async def close_a(reason): + closes_a.append(reason) + + async def send_b(payload): + sends_b.append(payload) + + async def close_b(reason): + closes_b.append(reason) + + created = await registry.create("sid", incarnation=2, 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, + close=close_b, + ) + + assert resumed.attachment_generation == 2 + assert resumed.resume_token.plaintext != created.resume_token.plaintext + assert [entry.sequence for entry in resumed.replay_entries] == [2] + assert resumed.replaced_attachment is not None + assert resumed.replaced_attachment.generation == 1 + assert await registry.is_current_attachment("sid", created.attachment_generation) is False + assert await registry.is_current_attachment("sid", resumed.attachment_generation) is True + assert await registry.detach("sid", attachment_generation=1) is False + assert await registry.detach("sid", attachment_generation=2) is True + assert [(payload["type"], payload["server_event_seq"]) for payload in sends_a] == [ + ("event-1", 1), + ("event-2", 2), + ] + assert closes_a == [] and sends_b == [] and closes_b == [] + + +@pytest.mark.asyncio +async def test_registry_resume_sends_activation_then_replay_before_new_live_events() -> None: + registry = DuplexSessionAttachmentRegistry( + replay_ttl_s=60.0, + replay_max_bytes_per_session=4096, + ) + wire = [] + activation_started = asyncio.Event() + release_activation = asyncio.Event() + + async def old_send(payload): + del payload + + async def close(reason): + del reason + + async def new_send(payload): + wire.append(payload["type"]) + if payload["type"] == "session.resumed": + activation_started.set() + await release_activation.wait() + + created = await registry.create("sid-order", incarnation=0, 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, + close=close, + activation_payload_factory=lambda _token, _generation: {"type": "session.resumed"}, + ) + ) + await activation_started.wait() + live_task = asyncio.create_task(registry.send_event("sid-order", {"type": "live"})) + + release_activation.set() + await asyncio.gather(resume_task, live_task) + + assert wire == ["session.resumed", "replayed", "live"] + + +@pytest.mark.asyncio +async def test_registry_resume_delivery_failure_keeps_one_shot_old_token_recovery() -> None: + registry = DuplexSessionAttachmentRegistry( + replay_ttl_s=60.0, + replay_max_bytes_per_session=4096, + ) + + async def send(payload): + del payload + + async def failing_send(payload): + del payload + raise RuntimeError("transport lost before rotated token arrived") + + async def close(reason): + del reason + + created = await registry.create("sid-recovery", incarnation=0, 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, + close=close, + activation_payload_factory=lambda token, generation: { + "type": "session.resumed", + "resume_token": token.plaintext, + "attachment_generation": generation, + }, + ) + + recovered = await registry.resume( + "sid-recovery", + incarnation=0, + resume_token=created.resume_token.plaintext, + last_received_server_event_seq=0, + send=send, + close=close, + ) + + assert recovered.attachment_generation == 3 + 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, + close=close, + ) + + +@pytest.mark.asyncio +async def test_registry_concurrent_resume_allows_exactly_one_rotated_token_winner() -> None: + 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-race", incarnation=0, 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, + close=close, + ) + except InvalidResumeTokenError as exc: + return exc + + results = await asyncio.gather(attempt(), attempt()) + + assert sum(not isinstance(result, Exception) for result in results) == 1 + assert sum(isinstance(result, InvalidResumeTokenError) for result in results) == 1 + + +@pytest.mark.asyncio +async def test_registry_keeps_sessions_journals_and_incarnations_isolated() -> None: + 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_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) + 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"}) + await registry.send_event("sid-b", {"type": "b-only"}) + + resumed_a = await registry.resume( + "sid-a", + incarnation=1, + resume_token=created_a.resume_token.plaintext, + last_received_server_event_seq=0, + send=send, + close=close, + ) + resumed_b = await registry.resume( + "sid-b", + incarnation=4, + resume_token=created_b.resume_token.plaintext, + last_received_server_event_seq=0, + send=send, + close=close, + ) + + 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"): + await registry.resume( + "sid-a", + incarnation=2, + resume_token=resumed_a.resume_token.plaintext, + last_received_server_event_seq=1, + send=send, + close=close, + ) + + +@pytest.mark.asyncio +async def test_registry_repr_never_contains_plaintext_tokens() -> None: + 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-repr", incarnation=0, send=send, close=close) + await registry.detach("sid-repr", attachment_generation=1) + entry = await registry.send_event( + "sid-repr", + {"type": "session.resumed", "resume_token": created.resume_token.plaintext}, + ) + + assert entry is not 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) diff --git a/tests/entrypoints/openai_api/test_duplex_capability.py b/tests/entrypoints/openai_api/test_duplex_capability.py new file mode 100644 index 000000000..153563cbb --- /dev/null +++ b/tests/entrypoints/openai_api/test_duplex_capability.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_api_server_does_not_import_experimental_duplex_at_module_load() -> None: + api_server = Path(__file__).parents[3] / "vllm_omni/entrypoints/openai/api_server.py" + module = ast.parse(api_server.read_text()) + + top_level_imports = { + node.module for node in module.body if isinstance(node, ast.ImportFrom) and isinstance(node.module, str) + } + + assert not any(module.startswith("vllm_omni.experimental.fullduplex") for module in top_level_imports) diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py new file mode 100644 index 000000000..3ed30e7e8 --- /dev/null +++ b/tests/entrypoints/openai_api/test_duplex_handler.py @@ -0,0 +1,6890 @@ +from __future__ import annotations + +import asyncio +import base64 +import json +import struct +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +from starlette.websockets import WebSocketDisconnect + +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.experimental.fullduplex.engine.duplex_control_client import DuplexControlRequestError +from vllm_omni.experimental.fullduplex.engine.duplex_runtime import duplex_resource_request_id +from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseActivity +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence, DuplexSessionLifecycleMessage +from vllm_omni.experimental.fullduplex.minicpmo45 import ( + MiniCPMO45NativeDuplexServingAdapter, + MiniCPMO45PcmAppendBuffer, +) +from vllm_omni.experimental.fullduplex.minicpmo45.data_plane import ( + MiniCPMO45DataPlaneContext, + MiniCPMO45DataPlaneSession, +) +from vllm_omni.experimental.fullduplex.minicpmo45.runtime import ( + MiniCPMO45DuplexRuntimeExtension, +) +from vllm_omni.experimental.fullduplex.minicpmo45.session import ( + MiniCPMO45ServingSessionState, +) +from vllm_omni.experimental.fullduplex.openai.protocol import ( + DuplexCapabilities, + DuplexOverlapPolicy, + DuplexPlaybackCommitPolicy, + DuplexSession, + DuplexSessionConfig, + ResponseCreateOptions, +) +from vllm_omni.experimental.fullduplex.openai.realtime_session import NativeRealtimeSessionProtocol +from vllm_omni.experimental.fullduplex.openai.runtime_adapter import ServingRuntimeConfigError +from vllm_omni.experimental.fullduplex.openai.serving import ( + OmniDuplexSessionHandler, + should_enable_duplex_endpoint, +) +from vllm_omni.experimental.fullduplex.openai.websocket import DuplexWebSocketActor +from vllm_omni.experimental.fullduplex.output import attach_duplex_output_decision +from vllm_omni.outputs import OmniRequestOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +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, object, bool]] = [] + self.append_operation_ids: list[str | None] = [] + self.appended_fences: list[DuplexFence | None] = [] + self.opened_configs: list[dict[str, object]] = [] + self.opened_runtime_configs: list[dict[str, object]] = [] + 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, object] | None] = [] + self.signal_runtime_configs: list[dict[str, object] | 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, + ) -> None: + 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, + ) -> None: + 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, + ) -> None: + 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, + ) -> None: + 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.experimental.fullduplex.minicpmo45.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])}") + + +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, object]]: + 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] + + +def test_native_realtime_protocol_emits_speak_once_per_response(): + protocol = NativeRealtimeSessionProtocol({}) + response_id = "resp-speak-once" + + projected = protocol.encode_outbound_event( + { + "type": "response.created", + "response_id": response_id, + "modalities": ["audio", "text"], + } + ) + projected.extend( + protocol.encode_outbound_event( + { + "type": "response.speak", + "response_id": response_id, + "text": "hello", + "model_speak": True, + } + ) + ) + projected.extend( + protocol.encode_outbound_event( + { + "type": "response.output_audio.delta", + "response_id": response_id, + "audio": base64.b64encode(b"\x00\x00").decode("ascii"), + "format": "pcm16", + "sample_rate_hz": 24000, + "model_speak": True, + } + ) + ) + + assert [event["type"] for event in projected].count("response.speak") == 1 + speak = next(event for event in projected if event["type"] == "response.speak") + assert "text" not in speak + assert speak["metadata"] == {"model_speak": True} + + +def test_native_realtime_protocol_speak_does_not_expose_internal_event(): + protocol = NativeRealtimeSessionProtocol({}) + + projected = protocol.encode_outbound_event( + { + "type": "response.speak", + "response_id": "resp-curated-speak", + "session_id": "sid-curated-speak", + "epoch": 3, + "text": "transcript belongs on the delta channel", + "model_speak": True, + "data_plane_request_id": "internal-request-id", + "uses_model_runner_scheduler": True, + } + ) + + assert len(projected) == 1 + projected[0].pop("event_id") + assert projected == [ + { + "type": "response.speak", + "response_id": "resp-curated-speak", + "item_id": "item_resp-curated-speak", + "output_index": 0, + "content_index": 0, + "metadata": { + "session_id": "sid-curated-speak", + "epoch": 3, + "model_speak": True, + }, + } + ] + + +@pytest.mark.asyncio +async def test_native_realtime_protocol_drains_internal_conversation_item_control_event(): + ws = TimedWebSocket() + protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] + protocol.bind_sender(ws.send_json) + + ws.put({"type": "session.update", "model": "test-model", "session_id": "rt-delete"}) + session_create = json.loads(await protocol.receive_internal_event_text(ws)) + assert session_create["type"] == "session.create" + + ws.put( + { + "type": "conversation.item.create", + "item": { + "id": "item-a", + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}], + }, + } + ) + assert json.loads(await protocol.receive_internal_event_text(ws)) == { + "type": "turn.signal", + "event": "conversation.item.create", + "payload": { + "item": { + "id": "item-a", + "object": "realtime.item", + "type": "message", + "status": "completed", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}], + } + }, + } + + ws.put({"type": "conversation.item.delete", "item_id": "item-a"}) + translated = json.loads(await protocol.receive_internal_event_text(ws)) + + assert translated == { + "type": "turn.signal", + "event": "conversation.item.delete", + "payload": {"item_id": "item-a"}, + } + + +@pytest.mark.asyncio +async def test_native_realtime_protocol_conversation_item_create_commits_user_text(): + ws = TimedWebSocket() + protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] + protocol.bind_sender(ws.send_json) + + ws.put({"type": "session.update", "model": "test-model", "session_id": "rt-item-create"}) + session_create = json.loads(await protocol.receive_internal_event_text(ws)) + assert session_create["type"] == "session.create" + + ws.put( + { + "type": "conversation.item.create", + "item": { + "id": "item-user-text", + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}], + }, + } + ) + + translated = json.loads(await protocol.receive_internal_event_text(ws)) + + assert translated == { + "type": "turn.signal", + "event": "conversation.item.create", + "payload": { + "item": { + "id": "item-user-text", + "object": "realtime.item", + "type": "message", + "status": "completed", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}], + } + }, + } + + +@pytest.mark.asyncio +async def test_native_realtime_protocol_audio_commit_requires_non_empty_buffer(): + ws = TimedWebSocket() + protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] + protocol.bind_sender(ws.send_json) + + ws.put({"type": "session.update", "model": "test-model", "session_id": "rt-audio-commit"}) + session_create = json.loads(await protocol.receive_internal_event_text(ws)) + assert session_create["type"] == "session.create" + + ws.put({"type": "input_audio_buffer.commit", "final": True}) + with pytest.raises(WebSocketDisconnect): + await protocol.receive_internal_event_text(ws) + + assert ws.sent[-1]["type"] == "error" + assert ws.sent[-1]["error"]["code"] == "input_audio_buffer_empty" + + +@pytest.mark.asyncio +async def test_native_realtime_protocol_audio_commit_does_not_auto_create_response(): + ws = TimedWebSocket() + protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] + protocol.bind_sender(ws.send_json) + + ws.put({"type": "session.update", "model": "test-model", "session_id": "rt-audio-commit-full"}) + session_create = json.loads(await protocol.receive_internal_event_text(ws)) + assert session_create["type"] == "session.create" + + pcm = struct.pack("<2h", 1024, -1024) + ws.put({"type": "input_audio_buffer.append", "audio": base64.b64encode(pcm).decode("ascii"), "format": "pcm16"}) + append_event = json.loads(await protocol.receive_internal_event_text(ws)) + assert append_event["type"] == "input_audio_buffer.append" + + ws.put({"type": "input_audio_buffer.commit", "final": True}) + commit_event = json.loads(await protocol.receive_internal_event_text(ws)) + + assert commit_event["type"] == "input_audio_buffer.commit" + assert commit_event["response_create"] is False + + +@pytest.mark.asyncio +async def test_native_realtime_protocol_rejects_invalid_sample_rate_without_closing(): + ws = TimedWebSocket() + protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] + protocol.bind_sender(ws.send_json) + + ws.put({"type": "session.update", "model": "test-model", "session_id": "rt-invalid-rate"}) + session_create = json.loads(await protocol.receive_internal_event_text(ws)) + assert session_create["type"] == "session.create" + + ws.put( + { + "type": "input_audio_buffer.append", + "audio": _pcm16_b64(1), + "format": "pcm16", + "sample_rate_hz": 192_001, + } + ) + ws.put({"type": "input_audio_buffer.clear"}) + translated = json.loads(await protocol.receive_internal_event_text(ws)) + + assert translated == {"type": "input_audio_buffer.clear", "reason": "input_audio_buffer.clear"} + error = next(event for event in ws.sent if event.get("type") == "error") + assert error["error"]["code"] == "bad_event" + assert error["error"]["param"] == "sample_rate_hz" + + +@pytest.mark.asyncio +async def test_native_realtime_protocol_audio_clear_is_not_barge_in_cancel(): + ws = TimedWebSocket() + protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] + protocol.bind_sender(ws.send_json) + + ws.put({"type": "session.update", "model": "test-model", "session_id": "rt-audio-clear"}) + session_create = json.loads(await protocol.receive_internal_event_text(ws)) + assert session_create["type"] == "session.create" + + ws.put({"type": "input_audio_buffer.clear"}) + clear_event = json.loads(await protocol.receive_internal_event_text(ws)) + + assert clear_event == {"type": "input_audio_buffer.clear", "reason": "input_audio_buffer.clear"} + + +@pytest.mark.asyncio +async def test_native_realtime_protocol_preserves_input_turn_policy_hints(): + ws = TimedWebSocket() + protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] + protocol.bind_sender(ws.send_json) + + translated = await protocol._to_duplex_event( + { + "type": "input_audio_buffer.append", + "audio": "AAAA", + "format": "pcm_f32le", + "duration_ms": 240, + "vad": {"is_speech": True, "speech_probability": 0.9}, + "overlap_action": "listen", + } + ) + + assert translated is not None + assert translated["duration_ms"] == 240 + assert translated["vad"] == {"is_speech": True, "speech_probability": 0.9} + assert translated["overlap_action"] == "listen" + + +@pytest.mark.asyncio +async def test_native_realtime_protocol_rejects_unimplemented_server_vad(): + ws = TimedWebSocket() + protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] + protocol.bind_sender(ws.send_json) + + translated = await protocol._to_duplex_event( + { + "type": "session.update", + "model": "test-model", + "session_id": "rt-turn-detection", + "turn_detection": { + "type": "server_vad", + "interrupt_response": False, + "silence_duration_ms": 900, + "threshold": 0.4, + }, + } + ) + + assert translated is None + error = ws.sent[-1] + assert error["type"] == "error" + assert error["error"]["code"] == "unsupported_turn_detection" + assert "server_vad" in error["error"]["message"] + + +@pytest.mark.asyncio +async def test_native_realtime_protocol_accepts_disabled_turn_detection(): + ws = TimedWebSocket() + protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] + protocol.bind_sender(ws.send_json) + + ws.put( + { + "type": "session.update", + "model": "test-model", + "session_id": "rt-no-turn-detection", + "turn_detection": None, + } + ) + translated = json.loads(await protocol.receive_internal_event_text(ws)) + + assert translated["type"] == "session.create" + assert translated["session"]["extra_body"]["realtime_session_payload"]["turn_detection"] is None + + +@pytest.mark.asyncio +async def test_native_realtime_protocol_rejects_nested_vad_even_when_top_level_is_disabled(): + ws = TimedWebSocket() + protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] + protocol.bind_sender(ws.send_json) + + translated = await protocol._to_duplex_event( + { + "type": "session.update", + "model": "test-model", + "turn_detection": None, + "audio": { + "input": { + "turn_detection": { + "type": "server_vad", + } + } + }, + } + ) + + assert translated is None + assert ws.sent[-1]["error"]["code"] == "unsupported_turn_detection" + + +def test_native_duplex_handler_has_no_fixed_session_admission_cap(): + handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) + client_config = DuplexSessionConfig( + extra_body={"max_duplex_sessions": 99, "duplex_max_sessions": 99}, + ) + + assert not hasattr(handler, "_max_native_duplex_sessions") + assert not hasattr(handler, "_max_duplex_sessions") + assert client_config.extra_body["max_duplex_sessions"] == 99 + + +def test_native_response_options_ignore_private_runtime_config(): + session = DuplexSession( + session_id="sid-response-private-config", + config=DuplexSessionConfig(), + capabilities=DuplexCapabilities.minicpmo45_native(), + ) + + handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) + handler._apply_response_create_options( + session, + { + "extra_body": { + "duplex_stage_max_tokens": {"0": 999}, + "client_trace_id": "trace-1", + } + }, + ) + session.begin_response() + + assert "duplex_stage_max_tokens" not in session.config.extra_body + assert "client_trace_id" not in session.config.extra_body + assert session.response_config.extra_body["client_trace_id"] == "trace-1" + + +@pytest.mark.parametrize( + "payload", + [ + {"instructions": "not wired"}, + {"voice": "not-wired"}, + {"temperature": 0.2}, + {"max_output_tokens": 8}, + {"tools": [{"type": "function", "name": "lookup"}]}, + ], +) +def test_native_response_create_rejects_unwired_generation_options(payload): + session = DuplexSession( + session_id="sid-response-unwired-options", + config=DuplexSessionConfig(), + capabilities=DuplexCapabilities.minicpmo45_native(), + ) + + handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) + error = handler._apply_response_create_options(session, payload) + + assert error == "unsupported_native_response_options" + session.begin_response() + assert session.response_config == session.config + + +def test_chat_fallback_excludes_realtime_response_metadata_from_model_request(): + handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) + session = DuplexSession( + session_id="sid-response-protocol-metadata", + config=DuplexSessionConfig(), + ) + session.reserve_response_options( + ResponseCreateOptions( + extra_body={ + "user": "trace-1", + "realtime_response_conversation": "none", + "realtime_response_metadata": {"source": "test"}, + "realtime_response_prompt": {"id": "prompt-1"}, + "realtime_response_tools": [ + { + "type": "function", + "function": {"name": "lookup", "parameters": {}}, + } + ], + "realtime_response_tool_choice": "auto", + } + ) + ) + session.begin_response() + + request = handler._build_chat_request(session, "request-1") + created = handler._response_created_payload(session, session.active_response_id or "response-1", epoch=0) + dumped = request.model_dump() + + assert dumped["user"] == "trace-1" + assert dumped["tools"][0]["function"]["name"] == "lookup" + assert dumped["tool_choice"] == "auto" + assert not any(key.startswith("realtime_response_") for key in dumped) + assert created["conversation"] == "none" + assert created["metadata"] == {"source": "test"} + assert created["prompt"] == {"id": "prompt-1"} + assert session.response_config.extra_body["realtime_response_conversation"] == "none" + + +@pytest.mark.asyncio +async def test_native_response_create_surfaces_unwired_generation_options_error(): + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + ws = TimedWebSocket() + ws.put(_native_session_create("sid-response-unwired-options-event")) + ws.put( + { + "type": "response.create", + "response": {"instructions": "not wired"}, + } + ) + 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"] == "unsupported_native_response_options" + assert "response_create_without_input" not in {message.get("code") for message in ws.sent} + + +@pytest.mark.asyncio +async def test_realtime_session_update_preserves_tools_and_metadata(): + 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-rt-update-fields")) + ws.put( + { + "type": "turn.signal", + "event": "session.update", + "payload": { + "tools": [{"type": "function", "name": "lookup"}], + "tool_choice": "auto", + "metadata": {"demo": "yes"}, + }, + } + ) + ws.put({"type": "session.close"}) + + await handler.handle_session(ws) + + updated = next(m for m in ws.sent if m.get("type") == "session.updated") + session = updated["session"] + assert session["tools"] == [{"type": "function", "name": "lookup"}] + assert session["tool_choice"] == "auto" + assert session["metadata"] == {"demo": "yes"} + assert "turn_detection" not in session + runtime_config = engine.signal_session_configs[-1] + assert runtime_config is not None + assert runtime_config["extra_body"]["realtime_tools"] == [{"type": "function", "name": "lookup"}] + assert runtime_config["extra_body"]["realtime_metadata"] == {"demo": "yes"} + + +@pytest.mark.asyncio +async def test_native_realtime_protocol_non_speech_append_does_not_emit_speech_started(): + ws = TimedWebSocket() + protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] + protocol.bind_sender(ws.send_json) + + translated = await protocol._to_duplex_event( + { + "type": "input_audio_buffer.append", + "audio": "AAAA", + "format": "pcm_f32le", + "vad": {"is_speech": False}, + } + ) + + assert translated is not None + assert translated["vad"] == {"is_speech": False} + assert not any(event.get("type") == "input_audio_buffer.speech_started" for event in ws.sent) + + +@pytest.mark.asyncio +async def test_native_session_update_rebuilds_server_runtime_policy(): + 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")) + ws.put( + { + "type": "turn.signal", + "event": "session.update", + "payload": {"temperature": 0.2, "max_output_tokens": 7}, + } + ) + ws.put({"type": "session.close"}) + + await handler.handle_session(ws) + + assert "session.updated" in ws.sent_types() + public_config = engine.signal_session_configs[-1] + runtime_config = engine.signal_runtime_configs[-1] + assert public_config is not None + assert runtime_config is not None + assert "duplex_stage_sampling_params" not in public_config["extra_body"] + assert runtime_config["duplex_stage_max_tokens"]["0"] == 7 + assert runtime_config["duplex_stage_sampling_params"]["0"]["temperature"] == 0.2 + assert runtime_config["duplex_stage_sampling_params"]["0"]["top_k"] == 20 + assert runtime_config["duplex_stage_sampling_params"]["0"]["top_p"] == 0.8 + + +@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_native_realtime_protocol_audio_delta_preserves_sample_rate_hz(): + 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": 24000, + } + ) + + 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.audio.delta"} + assert {payload["format"] for payload in audio_events} == {"pcm16"} + assert {payload["sample_rate_hz"] for payload in audio_events} == {24000} + + +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.audio.delta" in event_types + assert "response.audio.done" in event_types + assert "response.audio_transcript.delta" in event_types + assert "response.audio_transcript.done" in event_types + assert "response.output_audio.delta" not in event_types + assert "response.output_audio.done" not in event_types + assert "response.output_audio_transcript.delta" not in event_types + assert "response.output_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.audio.done" in by_type + assert "response.audio_transcript.delta" in by_type + assert "response.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.audio_transcript.delta") + done = [str(event["transcript"]) for event in events if event["type"] == "response.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.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,) + + +@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"] = {"minicpmo45_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": {"minicpmo45_native_duplex": True}, + }, + } + + +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._minicpmo_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_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.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_session_mode_duplex(): + 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_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( + request_id="duplex-direct-listen", + finished=True, + stage_id=0, + final_output_type=decision.final_output_type, + request_output=inner_output, + 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_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 = DuplexCapabilities.minicpmo45_native() + 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 = DuplexCapabilities.minicpmo45_native() + 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 = DuplexCapabilities.minicpmo45_native() + 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 = DuplexCapabilities.minicpmo45_native() + session.bind_request(request_id) + model_turn_id = session.turn_id + + native = handler._minicpmo_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_has_no_semantic_unit_cap(): + 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 = DuplexCapabilities.minicpmo45_native() + 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(9): + await handler._maybe_continue_native_response( + ws.send_json, + session=session, + expected_epoch=session.epoch, + ) + await asyncio.sleep(0.02) + + assert len(engine.appended) == 9 + assert all(payload["duplex_turn_id"] == 0 for _, _, payload, _ in engine.appended) + assert all("force_speak" not in payload for _, _, payload, _ in engine.appended) + + +@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 = DuplexCapabilities.minicpmo45_native() + 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 = DuplexCapabilities.minicpmo45_native() + 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._minicpmo_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( + request_id=request_id, + finished=True, + stage_id=0, + final_output_type=decision.final_output_type, + request_output=inner_output, + ), + 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 = DuplexCapabilities.minicpmo45_native() + 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 = DuplexCapabilities.minicpmo45_native() + 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 = DuplexCapabilities.minicpmo45_native() + 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 = DuplexCapabilities.minicpmo45_native() + 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._minicpmo_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._minicpmo_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._minicpmo_data_plane.slice_cumulative_audio(request_id, np.zeros(24000, dtype=np.float32)) + handler._minicpmo_data_plane.segment_text_delta(request_id, "hello") + handler._minicpmo_data_plane.mark_terminal(request_id) + native = MiniCPMO45ServingSessionState() + handler._minicpmo_sessions[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._minicpmo_data_plane.has_request(request_id) + assert session_id not in handler._minicpmo_sessions + + +@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]) -> None: + sent.append(data) + + await handler._emit_chat_payload( + session, + { + "modality": "audio", + "choices": [ + { + "delta": { + "content": "AAAA", + } + } + ], + }, + 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", + } + ] + + +@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(_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(_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_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_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(_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._minicpmo_sessions + + +@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._minicpmo_sessions + + +@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._minicpmo_sessions + + +@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", + "minicpmo45_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") + session = handler._registry.get("sid-resume-append-order") + assert session is not None + native_state = handler._minicpmo_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", + "minicpmo45_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") + + 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") + + 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._minicpmo_sessions + + +@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._minicpmo_sessions + 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(_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(_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(_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(_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_duplex_barge_in_aborts_active_response_when_runtime_signal_fails(): + engine = FakeEngineClient(fail_signal_events={"barge_in"}) + 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-barge-signal-fail")) + ws.put({"type": "input.text.append", "text": "hello"}) + ws.put({"type": "input.commit"}) + + await handler.handle_session(ws) + + assert "audio.cancelled" in ws.sent_types() + assert engine.aborted == ["chatcmpl-duplex-sid-barge-signal-fail-0-1"] + assert engine.internal_abort_batches == [["chatcmpl-duplex-sid-barge-signal-fail-0-1"]] + error = next(m for m in ws.sent if m.get("type") == "error") + assert error["code"] == "runtime_signal_failed" + + +@pytest.mark.asyncio +async def test_duplex_handler_surfaces_stage_unsupported_result_to_client(): + control_result = { + "operation": "open", + "session_id": "sid-unsupported", + "ok": True, + "unsupported_count": 1, + "error_count": 0, + "stage_results": [ + { + "stage_id": 0, + "replica_id": 0, + "result": {"supported": False, "reason": "not implemented"}, + } + ], + } + engine = FakeEngineClient(open_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(_session_create("sid-unsupported")) + ws.put({"type": "session.close"}) + + await handler.handle_session(ws) + + created = next(m for m in ws.sent if m.get("type") == "session.created") + assert created["runtime_control"]["unsupported_count"] == 1 + + +@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, + final: bool = False, + timeout: float | None = None, + collect_outputs: bool = True, + fence: DuplexFence | None = None, + ): + 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, + final=final, + timeout=timeout, + collect_outputs=collect_outputs, + fence=fence, + ) + + 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() + 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(_session_create("sid-control-signal-fail")) + ws.put( + { + "type": "turn.signal", + "event": "session.update", + "payload": {"instructions": "must roll back"}, + } + ) + 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 + assert engine.signal_session_configs[-1]["instructions"] == "must roll back" + + +@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, + ) + ws = TimedWebSocket() + ws.put(_session_create("sid-update-two-phase")) + ws.put( + { + "type": "turn.signal", + "event": "session.update", + "payload": {"instructions": "candidate instructions"}, + } + ) + + handler_task = asyncio.create_task(handler.handle_session(ws)) + await asyncio.wait_for(engine.update_started.wait(), timeout=1) + live_session = handler._registry.get("sid-update-two-phase") + assert live_session is not None + assert live_session.config.instructions != "candidate instructions" + + engine.release_update.set() + ws.put({"type": "session.close"}) + await asyncio.wait_for(handler_task, timeout=2) + + updated = next(message for message in ws.sent if message.get("type") == "session.updated") + assert updated["session"]["instructions"] == "candidate instructions" + + +@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_keeps_committed_audio_for_response_retry(): + failed_result = { + "operation": "append", + "session_id": "sid-final-append-retry", + "ok": False, + "unsupported_count": 0, + "error_count": 1, + "stage_results": [], + } + successful_result = { + "operation": "append", + "session_id": "sid-final-append-retry", + "ok": True, + "unsupported_count": 0, + "error_count": 0, + "stage_results": [], + } + engine = FakeEngineClient(append_results=[failed_result, successful_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): + if len(engine.appended) >= 2: + break + await asyncio.sleep(0.01) + ws.put({"type": "session.close"}) + await asyncio.wait_for(handler_task, timeout=2) + + assert len(engine.appended) == 2 + assert engine.appended[0][2]["audio"] == engine.appended[1][2]["audio"] + assert engine.append_operation_ids[0] == engine.append_operation_ids[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(_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( + 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, + ) -> None: + 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.experimental.fullduplex.minicpmo45.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 + 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.experimental.fullduplex.minicpmo45.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"] = { + "minicpmo45_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_is_deferred_as_listen(): + 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=DuplexCapabilities.minicpmo45_native(), + ) + 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": "listen", + "reason": "barge_in_unsupported", + "duration_ms": 1000, + "overlap_speech_ms": 1000, + "buffer_audio": True, + "defer_runtime_append": True, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "barge_in_event", + [ + {"type": "barge_in"}, + {"type": "turn.signal", "event": "barge_in"}, + ], +) +async def test_minicpmo_native_duplex_rejects_unadvertised_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) + + error = next(message for message in ws.sent if message.get("code") == "barge_in_unsupported") + assert error["session_id"] == "sid-native-barge-control-disabled" + assert ("sid-native-barge-control-disabled", "barge_in") not 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 = DuplexCapabilities.minicpmo45_native() + session.bind_request(request_id) + native = handler._minicpmo_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, + "realtime_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["realtime_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, + ) + ws = TimedWebSocket() + event = _native_session_create("sid-native-auto-continuation") + 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}) + + 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_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, + ) + ws = TimedWebSocket(receive_timeout_s=2.0) + event = _native_session_create("sid-native-silence-tail") + 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(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_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._minicpmo_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.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._minicpmo_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 + + +@pytest.mark.asyncio +async def test_continuous_response_metrics_accumulate_only_owned_model_units(): + handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) + session = DuplexSession( + session_id="sid-response-metrics", + 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(*, text: str, audio: str, tokens: int, ttft_ms: float, itls: list[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": 9.0, + "vllm_itls_ms": itls, + } + }, + }, + session=session, + ) + + await emit(text="first", audio="audio-a", tokens=4, ttft_ms=50.0, itls=[10.0, 11.0]) + await emit(text="second", audio="audio-b", tokens=6, ttft_ms=70.0, itls=[12.0, 13.0, 14.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"] == 4 + assert second_metrics["num_tokens_out"] == 10 + assert second_metrics["vllm_ttft_ms"] == 50.0 + assert second_metrics["vllm_itls_ms"] == [10.0, 11.0, 12.0, 13.0, 14.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]) + + 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] + + +@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 = DuplexCapabilities.minicpmo45_native() + 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._minicpmo_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_serving_video_stream.py b/tests/entrypoints/openai_api/test_serving_video_stream.py index de2fcc24d..fc8d83da6 100644 --- a/tests/entrypoints/openai_api/test_serving_video_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_stream.py @@ -138,6 +138,112 @@ async def test_receive_config_accepts_client_legacy_aliases(): assert config.frame_filter_threshold == 0.87 +@pytest.mark.asyncio +async def test_video_frame_ack_reports_receiver_buffer_state(): + ws = MockWebSocket( + [ + json.dumps( + { + "type": "session.config", + "model": "test", + "enable_frame_filter": False, + } + ), + json.dumps( + { + "type": "video.frame", + "data": _b64(_make_jpeg()), + "frame_id": "frame-7", + "pts_ms": 700, + "capture_ts_ms": 1234.5, + } + ), + json.dumps({"type": "video.done"}), + ] + ) + handler = QwenOmniStreamingVideoHandler(chat_service=object()) + + await handler.handle_session(ws) + + ack = next(message for message in ws.sent if message.get("type") == "video.frame.ack") + assert ack["frame_id"] == "frame-7" + assert ack["pts_ms"] == 700 + assert ack["capture_ts_ms"] == 1234.5 + assert ack["accepted"] is True + assert ack["buffered_frames"] == 1 + assert ack["server_receive_ts_ms"] > 0 + + +@pytest.mark.asyncio +async def test_video_frames_consumed_is_emitted_after_engine_uses_frame_prompt(): + class OneOutputEngine: + def generate(self, **_kwargs): + async def _gen(): + await asyncio.sleep(0.05) + yield _text_result("visible") + + return _gen() + + class CapturingHandler(QwenOmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + return {"prompt": "with-video"} + + ws = MockWebSocket( + [ + json.dumps( + { + "type": "session.config", + "model": "test", + "modalities": ["text"], + "num_frames": 1, + "enable_frame_filter": False, + } + ), + json.dumps( + { + "type": "video.frame", + "data": _b64(_make_jpeg()), + "frame_id": "frame-9", + "pts_ms": 900, + "source_pts_ms": 880, + "quality_profile": "balanced", + } + ), + json.dumps({"type": "video.query", "text": "describe"}), + json.dumps({"type": "video.done"}), + ] + ) + handler = CapturingHandler( + chat_service=object(), + engine_client=OneOutputEngine(), + idle_timeout=2.0, + ) + + await handler.handle_session(ws) + + consumed = next(message for message in ws.sent if message.get("type") == "video.frames.consumed") + assert consumed["frame_ids"] == ["frame-9"] + assert consumed["latest_pts_ms"] == 900 + assert consumed["request_id"].startswith("video-") + assert consumed["model_selected_ts_ms"] > 0 + assert consumed["frames"] == [ + { + "frame_id": "frame-9", + "pts_ms": 900, + "source_pts_ms": 880, + "quality_profile": "balanced", + "receiver_received_ts_ms": consumed["frames"][0]["receiver_received_ts_ms"], + "decoded_ready_ts_ms": consumed["frames"][0]["decoded_ready_ts_ms"], + } + ] + assert consumed["frames"][0]["receiver_received_ts_ms"] > 0 + assert consumed["frames"][0]["decoded_ready_ts_ms"] >= consumed["frames"][0]["receiver_received_ts_ms"] + assert consumed["model_selected_ts_ms"] >= consumed["frames"][0]["decoded_ready_ts_ms"] + assert ws.sent.index(consumed) < next( + index for index, message in enumerate(ws.sent) if message.get("type") == "response.text.delta" + ) + + @pytest.mark.asyncio async def test_audio_in_video_sets_mm_processor_kwargs(): captured_requests = [] diff --git a/tests/entrypoints/test_async_omni_duplex.py b/tests/entrypoints/test_async_omni_duplex.py new file mode 100644 index 000000000..62f3f49ad --- /dev/null +++ b/tests/entrypoints/test_async_omni_duplex.py @@ -0,0 +1,600 @@ +from types import SimpleNamespace + +import pytest + +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.experimental.fullduplex.engine.duplex_runtime import ( + DuplexOutputAction, + DuplexOutputDecision, + duplex_resource_request_id, +) +from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseActivity +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence +from vllm_omni.experimental.fullduplex.output import attach_duplex_output_decision +from vllm_omni.experimental.fullduplex.request_client import DuplexRequestClient +from vllm_omni.metrics.stats import OrchestratorAggregator, StageRequestStats, StageStats +from vllm_omni.outputs import OmniRequestOutput + + +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, + request_output=SimpleNamespace( + outputs=[], + multimodal_output=SimpleNamespace(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].request_output is raw_stage1_output + + +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_duplex_fence_propagation.py b/tests/entrypoints/test_duplex_fence_propagation.py new file mode 100644 index 000000000..f7b8cdcfc --- /dev/null +++ b/tests/entrypoints/test_duplex_fence_propagation.py @@ -0,0 +1,123 @@ +import pytest + +from tests.entrypoints.openai_api.test_duplex_handler import FakeChatService, TimedWebSocket +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence +from vllm_omni.experimental.fullduplex.openai.protocol import DuplexSession, DuplexSessionConfig +from vllm_omni.experimental.fullduplex.openai.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(), + 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()) + 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()) + 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_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 4da82390b..5d9b85938 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -108,6 +108,8 @@ 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]] = [] diff --git a/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py b/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py new file mode 100644 index 000000000..93fc4db38 --- /dev/null +++ b/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py @@ -0,0 +1,211 @@ +import base64 +import importlib.util +import sys +import wave +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +DEMO_PATH = Path(__file__).resolve().parents[2] / "examples/online_serving/minicpmo/realtime_duplex_demo.py" + + +def _load_demo_module(): + spec = importlib.util.spec_from_file_location( + "minicpmo_realtime_duplex_simple_demo_test", + DEMO_PATH, + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def test_post_commit_decision_ignores_streaming_listens_before_commit_ack(): + demo = _load_demo_module() + events = [ + {"type": "session.created"}, + {"type": "response.listen"}, + {"type": "response.listen"}, + {"type": "input_audio_buffer.committed"}, + ] + + commit_index = demo._input_committed_index(events, after_index=0) + + assert commit_index == 3 + assert demo._post_commit_model_decision(events, commit_index) is None + + +def test_post_commit_decision_accepts_listen_after_commit_ack(): + demo = _load_demo_module() + events = [ + {"type": "response.listen"}, + {"type": "input_audio_buffer.committed"}, + {"type": "response.listen"}, + ] + + commit_index = demo._input_committed_index(events, after_index=0) + + assert demo._post_commit_model_decision(events, commit_index) == "listen" + + +def test_post_commit_decision_accepts_drained_speak_after_commit_ack(): + demo = _load_demo_module() + events = [ + {"type": "input_audio_buffer.committed"}, + {"type": "response.created", "response": {"id": "resp-a"}}, + {"type": "response.audio.delta", "response_id": "resp-a", "delta": "AAAA"}, + {"type": "response.done", "response_id": "resp-a"}, + ] + + commit_index = demo._input_committed_index(events, after_index=0) + + assert demo._post_commit_model_decision(events, commit_index) == "speak" + + +def test_post_commit_decision_ignores_cancelled_response(): + demo = _load_demo_module() + events = [ + {"type": "input_audio_buffer.committed"}, + {"type": "response.done", "response": {"status": "cancelled"}}, + ] + + commit_index = demo._input_committed_index(events, after_index=0) + + assert demo._post_commit_model_decision(events, commit_index) is None + + +def test_commit_ack_search_starts_after_stream_cursor(): + demo = _load_demo_module() + events = [ + {"type": "input_audio_buffer.committed"}, + {"type": "response.listen"}, + {"type": "input_audio_buffer.committed"}, + ] + + assert demo._input_committed_index(events, after_index=1) == 2 + + +def test_exact_model_unit_boundary_does_not_require_an_extra_decision(): + demo = _load_demo_module() + unit_bytes = demo.PCM16_SAMPLE_RATE * demo.PCM16_BYTES_PER_SAMPLE + + assert not demo._has_residual_model_unit(b"\0" * (unit_bytes * 6), chunk_period_ms=1000) + + +def test_partial_model_unit_requires_a_post_commit_decision(): + demo = _load_demo_module() + unit_bytes = demo.PCM16_SAMPLE_RATE * demo.PCM16_BYTES_PER_SAMPLE + + assert demo._has_residual_model_unit(b"\0" * (unit_bytes * 6 + 512), chunk_period_ms=1000) + + +def test_latest_streaming_decision_is_used_at_exact_boundary(): + demo = _load_demo_module() + events = [ + {"type": "response.listen"}, + {"type": "response.listen"}, + {"type": "input_audio_buffer.committed"}, + ] + + assert demo._latest_model_decision(events, after_index=0) == "listen" + + +def test_ref_audio_data_url_encodes_explicit_wav(tmp_path): + demo = _load_demo_module() + ref = tmp_path / "ref.wav" + ref.write_bytes(b"RIFFfake") + + assert demo._ref_audio_data_url(str(ref)) == "data:audio/wav;base64,UklGRmZha2U=" + assert demo._ref_audio_data_url(None) is None + + +def test_realtime_duplex_demo_requires_ref_audio(monkeypatch): + demo = _load_demo_module() + monkeypatch.setattr( + demo.sys, + "argv", + [ + "realtime_duplex_demo.py", + "--input-wav", + "input.wav", + ], + ) + + with pytest.raises(SystemExit): + demo.parse_args() + + +def test_realtime_duplex_demo_accepts_explicit_ref_audio(monkeypatch): + demo = _load_demo_module() + monkeypatch.setattr( + demo.sys, + "argv", + [ + "realtime_duplex_demo.py", + "--input-wav", + "input.wav", + "--ref-audio", + "ref.wav", + ], + ) + + args = demo.parse_args() + + assert args.ref_audio == "ref.wav" + + +def test_open_streaming_response_requires_post_commit_drain(): + demo = _load_demo_module() + + assert demo._response_in_progress([{"type": "response.created"}]) + assert not demo._response_in_progress( + [ + {"type": "response.created"}, + {"type": "response.done"}, + ] + ) + + +def test_streaming_output_writer_persists_audio_deltas_as_they_arrive(tmp_path, capsys): + demo = _load_demo_module() + writer = demo._StreamingOutputWriter(tmp_path) + first_pcm = b"\x01\x00\x02\x00" + second_pcm = b"\x03\x00" + + writer.handle( + { + "type": "response.audio.delta", + "response_id": "resp-a", + "delta": base64.b64encode(first_pcm).decode("ascii"), + "sample_rate_hz": 24_000, + } + ) + writer.handle( + { + "type": "response.audio_transcript.delta", + "response_id": "resp-a", + "delta": "你好", + } + ) + writer.handle( + { + "type": "response.audio.delta", + "response_id": "resp-a", + "delta": base64.b64encode(second_pcm).decode("ascii"), + "sample_rate_hz": 24_000, + } + ) + + assert (tmp_path / "output.pcm").read_bytes() == first_pcm + second_pcm + chunk_paths = sorted((tmp_path / "audio_chunks").glob("*.wav")) + assert [path.name for path in chunk_paths] == ["chunk_0001.wav", "chunk_0002.wav"] + with wave.open(str(chunk_paths[0]), "rb") as chunk_wav: + assert chunk_wav.getframerate() == 24_000 + assert chunk_wav.readframes(chunk_wav.getnframes()) == first_pcm + assert writer.audio_chunk_paths == chunk_paths + stderr = capsys.readouterr().err + assert "audio chunk 1" in stderr + assert "你好" in stderr diff --git a/tests/examples/test_minicpmo_realtime_web.py b/tests/examples/test_minicpmo_realtime_web.py new file mode 100644 index 000000000..03fbc6d6d --- /dev/null +++ b/tests/examples/test_minicpmo_realtime_web.py @@ -0,0 +1,1663 @@ +import asyncio +import base64 +import importlib.util +import subprocess +import sys +import wave +from pathlib import Path +from types import SimpleNamespace +from urllib.parse import parse_qs, urlsplit + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +DEMO_PATH = Path(__file__).resolve().parents[1] / "e2e/online_serving/minicpmo_realtime_duplex_scenarios.py" +MULTI_DEMO_PATH = ( + Path(__file__).resolve().parents[1] / "e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py" +) +PAIR_DEMO_PATH = Path(__file__).resolve().parents[1] / "e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py" +SOFT_INTERRUPT_DEMO_PATH = ( + Path(__file__).resolve().parents[1] / "e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py" +) + + +def _load_demo_module(): + spec = importlib.util.spec_from_file_location("minicpmo_realtime_duplex_demo_test", DEMO_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _load_multi_demo_module(): + spec = importlib.util.spec_from_file_location("minicpmo_realtime_duplex_multi_session_test", MULTI_DEMO_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _load_pair_demo_module(): + spec = importlib.util.spec_from_file_location("minicpmo_realtime_duplex_demo_pair_test", PAIR_DEMO_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _load_soft_interrupt_demo_module(): + spec = importlib.util.spec_from_file_location( + "minicpmo_realtime_duplex_soft_interrupt_test", + SOFT_INTERRUPT_DEMO_PATH, + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def test_realtime_duplex_multi_session_script_is_directly_executable(): + result = subprocess.run( + [sys.executable, str(MULTI_DEMO_PATH), "--help"], + cwd=MULTI_DEMO_PATH.parents[3], + capture_output=True, + text=True, + timeout=20, + check=False, + ) + + assert result.returncode == 0, result.stderr + + +def test_realtime_duplex_demo_pair_script_is_directly_executable(): + result = subprocess.run( + [sys.executable, str(PAIR_DEMO_PATH), "--help"], + cwd=PAIR_DEMO_PATH.parents[3], + capture_output=True, + text=True, + timeout=20, + check=False, + ) + + assert result.returncode == 0, result.stderr + + +def test_realtime_duplex_soft_interrupt_script_is_directly_executable(): + result = subprocess.run( + [sys.executable, str(SOFT_INTERRUPT_DEMO_PATH), "--help"], + cwd=SOFT_INTERRUPT_DEMO_PATH.parents[3], + capture_output=True, + text=True, + timeout=20, + check=False, + ) + + assert result.returncode == 0, result.stderr + + +def test_realtime_duplex_demo_pair_requires_distinct_inputs_and_outputs(tmp_path): + demo = _load_pair_demo_module() + wav_a = tmp_path / "request_a.wav" + wav_b = tmp_path / "request_b.wav" + wav_a.write_bytes(b"same") + wav_b.write_bytes(b"same") + + args = SimpleNamespace( + input_wav_a=str(wav_a), + input_wav_b=str(wav_b), + output_dir_a=str(tmp_path / "session_a"), + output_dir_b=str(tmp_path / "session_a"), + ) + + with pytest.raises(ValueError, match="output directories must be different"): + demo._validate_pair_args(args) + + args.output_dir_b = str(tmp_path / "session_b") + with pytest.raises(ValueError, match="input WAV files must have different content"): + demo._validate_pair_args(args) + + +def test_realtime_duplex_demo_pair_launches_demo_processes_concurrently(tmp_path, monkeypatch): + demo = _load_pair_demo_module() + fake_demo = tmp_path / "fake_realtime_duplex_demo.py" + start_log = tmp_path / "starts.log" + fake_demo.write_text( + "\n".join( + [ + "import argparse, json, os, time", + "from pathlib import Path", + "parser = argparse.ArgumentParser()", + "parser.add_argument('--url')", + "parser.add_argument('--model')", + "parser.add_argument('--input-wav')", + "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')", + "args = parser.parse_args()", + "if args.ref_audio != os.environ.get('EXPECTED_REF_AUDIO'):", + " raise SystemExit(f'unexpected ref audio: {args.ref_audio}')", + "output = Path(args.output_dir)", + "output.mkdir(parents=True, exist_ok=True)", + "response_id = 'resp-' + output.name", + "with open(os.environ['PAIR_START_LOG'], 'a', encoding='utf-8') as f:", + " f.write(f'{output.name} {time.time()}\\n')", + "time.sleep(0.25)", + "events = [", + " {'type': 'session.created'},", + " {'type': 'response.created', 'response': {'id': response_id}},", + " {'type': 'response.audio.delta', 'response_id': response_id, " + "'delta': 'AAAA', 'sample_rate_hz': 24000, '_client_received_at_s': 10.1},", + " {'type': 'response.audio.delta', 'response_id': response_id, " + "'delta': 'AAAA', 'sample_rate_hz': 24000, '_client_received_at_s': 10.6},", + " {'type': 'response.done', 'response_id': response_id},", + " {'type': 'session.closed'},", + "]", + "Path(output / 'events.jsonl').write_text(", + " ''.join(json.dumps(event) + '\\n' for event in events),", + " encoding='utf-8',", + ")", + "result = {'ok': True, 'response_ids': [response_id], 'audio_bytes': 1, 'output_dir': str(output)}", + "Path(output / 'result.json').write_text(json.dumps(result), encoding='utf-8')", + "print(json.dumps(result))", + ] + ), + encoding="utf-8", + ) + monkeypatch.setattr(demo, "DEMO_PATH", fake_demo) + monkeypatch.setenv("PAIR_START_LOG", str(start_log)) + monkeypatch.setenv("EXPECTED_REF_AUDIO", str(tmp_path / "ref.wav")) + wav_a = tmp_path / "request_a.wav" + wav_b = tmp_path / "request_b.wav" + wav_a.write_bytes(b"request a") + wav_b.write_bytes(b"request b") + + result = asyncio.run( + demo.run_pair( + SimpleNamespace( + url="ws://localhost:28889/v1/realtime?duplex=1", + model="/data/why/MiniCPM-o-4_5", + input_wav_a=str(wav_a), + input_wav_b=str(wav_b), + output_dir_a=str(tmp_path / "duplex_session_a"), + output_dir_b=str(tmp_path / "duplex_session_b"), + ref_audio=str(tmp_path / "ref.wav"), + summary_output=str(tmp_path / "summary.json"), + chunk_ms=200, + timeout_s=5.0, + no_realtime_pacing=False, + require_audio=True, + min_audio_deltas_per_session=2, + ) + ) + ) + + starts = [float(line.split()[1]) for line in start_log.read_text(encoding="utf-8").splitlines()] + assert len(starts) == 2 + assert max(starts) - min(starts) < 0.2 + assert result["ok"] is True + assert result["input_wavs_distinct"] is True + assert result["output_dirs_distinct"] is True + assert result["identity_isolation_ok"] is True + assert result["response_audio_contract_ok"] is True + assert {session["audio_delta_count"] for session in result["sessions"]} == {2} + for session in result["sessions"]: + assert session["all_audio_deltas_same_response"] is True + assert session["response_id_nonempty"] is True + assert session["first_audio_delta_before_done"] is True + assert len(session["audio_delta_timings"]) == 2 + + +def test_realtime_duplex_demo_pair_rejects_false_green_audio_contract(tmp_path): + demo = _load_pair_demo_module() + output = tmp_path / "session" + output.mkdir() + response_id = "resp-a" + events = [ + {"type": "response.created", "response": {"id": response_id}}, + {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "sample_rate_hz": 24000}, + {"type": "response.audio.delta", "response_id": "resp-other", "delta": "AAAA", "sample_rate_hz": 24000}, + {"type": "response.done", "response_id": response_id}, + ] + (output / "events.jsonl").write_text( + "".join(demo.json.dumps(event) + "\n" for event in events), + encoding="utf-8", + ) + (output / "result.json").write_text( + demo.json.dumps({"ok": True, "response_ids": [response_id], "audio_bytes": 4}), + encoding="utf-8", + ) + + session = demo._summarize_session( + label="a", + input_wav="request.wav", + output_dir=str(output), + returncode=0, + stdout="", + stderr="", + ) + + assert session["ok"] is False + assert session["all_audio_deltas_same_response"] is False + assert session["response_audio_contract_ok"] is False + + +def test_realtime_duplex_demo_pair_default_requires_multiple_audio_deltas(monkeypatch): + demo = _load_pair_demo_module() + monkeypatch.setattr( + demo.sys, + "argv", + [ + "pair.py", + "--input-wav-a", + "a.wav", + "--input-wav-b", + "b.wav", + "--output-dir-a", + "a", + "--output-dir-b", + "b", + "--ref-audio", + "ref.wav", + ], + ) + + args = demo.parse_args() + + assert args.ref_audio == "ref.wav" + assert args.min_audio_deltas_per_session == 2 + + +def test_realtime_duplex_demo_pair_requires_ref_audio(monkeypatch): + demo = _load_pair_demo_module() + monkeypatch.setattr( + demo.sys, + "argv", + [ + "pair.py", + "--input-wav-a", + "a.wav", + "--input-wav-b", + "b.wav", + "--output-dir-a", + "a", + "--output-dir-b", + "b", + ], + ) + + with pytest.raises(SystemExit): + demo.parse_args() + + +def test_realtime_duplex_soft_interrupt_requires_ref_audio(monkeypatch): + demo = _load_soft_interrupt_demo_module() + monkeypatch.setattr( + demo.sys, + "argv", + [ + "soft.py", + "--input-wav", + "input.wav", + "--output-dir", + "out", + ], + ) + + with pytest.raises(SystemExit): + demo.parse_args() + + +def test_realtime_duplex_soft_interrupt_accepts_explicit_ref_audio(monkeypatch): + demo = _load_soft_interrupt_demo_module() + monkeypatch.setattr( + demo.sys, + "argv", + [ + "soft.py", + "--input-wav", + "input.wav", + "--output-dir", + "out", + "--ref-audio", + "ref.wav", + ], + ) + + args = demo.parse_args() + + assert args.ref_audio == "ref.wav" + assert args.validation_mode == "model-policy" + + +@pytest.mark.parametrize( + "extra_args", + [ + ["--expect-second-response-substring", "一加一等于二"], + ["--input-sha256", "a" * 64], + ], +) +def test_realtime_duplex_soft_interrupt_response_required_needs_bound_fixture(monkeypatch, extra_args): + demo = _load_soft_interrupt_demo_module() + monkeypatch.setattr( + demo.sys, + "argv", + [ + "soft.py", + "--input-wav", + "input.wav", + "--output-dir", + "out", + "--ref-audio", + "ref.wav", + "--validation-mode", + "response-required", + *extra_args, + ], + ) + + with pytest.raises(SystemExit): + demo.parse_args() + + +def test_realtime_duplex_soft_interrupt_rejects_input_checksum_mismatch(tmp_path): + demo = _load_soft_interrupt_demo_module() + input_wav = tmp_path / "input.wav" + input_wav.write_bytes(b"two-turn fixture") + + with pytest.raises(ValueError, match="input WAV SHA256 mismatch"): + demo._validate_input_sha256(input_wav, "0" * 64) + + +def test_realtime_duplex_soft_interrupt_accepts_multi_delta_handoff_sequence(tmp_path): + demo = _load_soft_interrupt_demo_module() + output = tmp_path / "soft_interrupt" + output.mkdir() + first_response_id = "resp-first" + second_response_id = "resp-second" + events = [ + {"type": "response.listen", "_client_received_at_s": 1.0}, + {"type": "response.created", "response": {"id": first_response_id}, "_client_received_at_s": 2.0}, + {"type": "response.speak", "response_id": first_response_id, "_client_received_at_s": 2.0}, + { + "type": "response.audio.delta", + "response_id": first_response_id, + "delta": "AAAA", + "_client_received_at_s": 2.1, + }, + { + "type": "response.audio_transcript.delta", + "response_id": first_response_id, + "delta": "中国古代四大发明", + "_client_received_at_s": 2.1, + }, + { + "type": "response.audio.delta", + "response_id": first_response_id, + "delta": "AAAA", + "_client_received_at_s": 2.6, + }, + { + "type": "response.audio_transcript.delta", + "response_id": first_response_id, + "delta": "是造纸术。", + "_client_received_at_s": 2.6, + }, + {"type": "response.done", "response_id": first_response_id, "_client_received_at_s": 3.0}, + {"type": "response.listen", "_client_received_at_s": 3.2}, + {"type": "response.created", "response": {"id": second_response_id}, "_client_received_at_s": 4.0}, + {"type": "response.speak", "response_id": second_response_id, "_client_received_at_s": 4.0}, + { + "type": "response.audio.delta", + "response_id": second_response_id, + "delta": "AAAA", + "_client_received_at_s": 4.1, + }, + { + "type": "response.audio_transcript.delta", + "response_id": second_response_id, + "delta": "一加一等于", + "_client_received_at_s": 4.1, + }, + { + "type": "response.audio.delta", + "response_id": second_response_id, + "delta": "AAAA", + "_client_received_at_s": 4.6, + }, + { + "type": "response.audio_transcript.delta", + "response_id": second_response_id, + "delta": "二。", + "_client_received_at_s": 4.6, + }, + {"type": "response.done", "response_id": second_response_id, "_client_received_at_s": 5.0}, + {"type": "response.listen", "_client_received_at_s": 5.2}, + {"type": "input_audio_buffer.committed", "_client_received_at_s": 6.0}, + {"type": "response.listen", "_client_received_at_s": 6.1}, + ] + (output / "events.jsonl").write_text( + "".join(demo.json.dumps(event) + "\n" for event in events), + encoding="utf-8", + ) + (output / "result.json").write_text( + demo.json.dumps({"ok": True, "response_ids": [first_response_id, second_response_id]}), + encoding="utf-8", + ) + + summary = demo.summarize_artifacts( + output_dir=output, + validation_mode="response-required", + min_responses=2, + min_audio_deltas_per_response=2, + expect_second_response_substring="一加一等于二", + ) + + assert summary["ok"] is True + assert summary["second_response_before_final_commit"] is True + assert summary["final_listen_after_commit"] is True + assert summary["response_audio_contract_ok"] is True + assert summary["response_summaries"][0]["audio_delta_count"] == 2 + assert summary["response_summaries"][1]["transcript"] == "一加一等于二。" + + +def test_realtime_duplex_soft_interrupt_model_policy_accepts_single_response(tmp_path): + demo = _load_soft_interrupt_demo_module() + output = tmp_path / "model_policy" + output.mkdir() + response_id = "resp-only" + events = [ + {"type": "response.listen", "_client_received_at_s": 1.0}, + {"type": "response.created", "response": {"id": response_id}, "_client_received_at_s": 2.0}, + {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 2.1}, + {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 2.15}, + {"type": "response.done", "response_id": response_id, "_client_received_at_s": 2.2}, + {"type": "response.listen", "_client_received_at_s": 2.3}, + {"type": "input_audio_buffer.committed", "_client_received_at_s": 3.0}, + {"type": "response.listen", "_client_received_at_s": 3.1}, + ] + (output / "events.jsonl").write_text( + "".join(demo.json.dumps(event) + "\n" for event in events), + encoding="utf-8", + ) + (output / "result.json").write_text( + demo.json.dumps({"ok": True, "response_ids": [response_id]}), + encoding="utf-8", + ) + + summary = demo.summarize_artifacts( + output_dir=output, + validation_mode="model-policy", + min_responses=2, + min_audio_deltas_per_response=2, + expect_second_response_substring=None, + ) + + assert summary["ok"] is True + assert summary["validation_mode"] == "model-policy" + assert summary["enough_responses"] is True + assert summary["response_before_final_commit"] is True + assert summary["listen_after_response_before_commit"] is True + + +def test_realtime_duplex_soft_interrupt_response_required_rejects_single_response(tmp_path): + demo = _load_soft_interrupt_demo_module() + output = tmp_path / "response_required" + output.mkdir() + response_id = "resp-only" + events = [ + {"type": "response.listen", "_client_received_at_s": 1.0}, + {"type": "response.created", "response": {"id": response_id}, "_client_received_at_s": 2.0}, + {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 2.1}, + {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 2.15}, + {"type": "response.done", "response_id": response_id, "_client_received_at_s": 2.2}, + {"type": "response.listen", "_client_received_at_s": 2.3}, + {"type": "input_audio_buffer.committed", "_client_received_at_s": 3.0}, + {"type": "response.listen", "_client_received_at_s": 3.1}, + ] + (output / "events.jsonl").write_text( + "".join(demo.json.dumps(event) + "\n" for event in events), + encoding="utf-8", + ) + (output / "result.json").write_text( + demo.json.dumps({"ok": True, "response_ids": [response_id]}), + encoding="utf-8", + ) + + summary = demo.summarize_artifacts( + output_dir=output, + validation_mode="response-required", + min_responses=2, + min_audio_deltas_per_response=2, + expect_second_response_substring="一加一等于二", + ) + + assert summary["ok"] is False + assert summary["enough_responses"] is False + + +def test_realtime_duplex_soft_interrupt_checks_expected_text_in_second_response(tmp_path): + demo = _load_soft_interrupt_demo_module() + output = tmp_path / "second_response_text" + output.mkdir() + first_response_id = "resp-first" + second_response_id = "resp-second" + events = [ + {"type": "response.listen", "_client_received_at_s": 1.0}, + {"type": "response.created", "response": {"id": first_response_id}, "_client_received_at_s": 2.0}, + { + "type": "response.audio.delta", + "response_id": first_response_id, + "delta": "AAAA", + "_client_received_at_s": 2.1, + }, + { + "type": "response.audio.delta", + "response_id": first_response_id, + "delta": "AAAA", + "_client_received_at_s": 2.2, + }, + { + "type": "response.audio_transcript.delta", + "response_id": first_response_id, + "delta": "一加一等于二", + "_client_received_at_s": 2.2, + }, + {"type": "response.done", "response_id": first_response_id, "_client_received_at_s": 2.3}, + {"type": "response.listen", "_client_received_at_s": 2.4}, + {"type": "response.created", "response": {"id": second_response_id}, "_client_received_at_s": 3.0}, + { + "type": "response.audio.delta", + "response_id": second_response_id, + "delta": "AAAA", + "_client_received_at_s": 3.1, + }, + { + "type": "response.audio.delta", + "response_id": second_response_id, + "delta": "AAAA", + "_client_received_at_s": 3.2, + }, + { + "type": "response.audio_transcript.delta", + "response_id": second_response_id, + "delta": "不知道", + "_client_received_at_s": 3.2, + }, + {"type": "response.done", "response_id": second_response_id, "_client_received_at_s": 3.3}, + {"type": "response.listen", "_client_received_at_s": 3.4}, + {"type": "input_audio_buffer.committed", "_client_received_at_s": 4.0}, + {"type": "response.listen", "_client_received_at_s": 4.1}, + ] + (output / "events.jsonl").write_text( + "".join(demo.json.dumps(event) + "\n" for event in events), + encoding="utf-8", + ) + (output / "result.json").write_text( + demo.json.dumps({"ok": True, "response_ids": [first_response_id, second_response_id]}), + encoding="utf-8", + ) + + summary = demo.summarize_artifacts( + output_dir=output, + validation_mode="response-required", + min_responses=2, + min_audio_deltas_per_response=2, + expect_second_response_substring="一加一等于二", + ) + + assert summary["ok"] is False + assert summary["second_response_transcript_expectation_ok"] is False + + +def test_realtime_duplex_multi_session_resume_url_disables_autostart(): + demo = _load_multi_demo_module() + + url = demo._with_resume_mode("ws://localhost:8113/v1/realtime?duplex=1&model=openbmb%2FMiniCPM-o-4_5") + + query = parse_qs(urlsplit(url).query) + assert query["model"] == ["openbmb/MiniCPM-o-4_5"] + assert query["resume"] == ["1"] + + +def test_realtime_duplex_multi_session_rejects_cross_session_response_identity(): + demo = _load_multi_demo_module() + + assert demo._validate_identity_isolation( + [ + {"completed_response_ids": ["resp-a"]}, + {"completed_response_ids": ["resp-b"]}, + ] + ) + assert not demo._validate_identity_isolation( + [ + {"completed_response_ids": ["resp-shared"]}, + {"completed_response_ids": ["resp-shared"]}, + ] + ) + + +def test_realtime_duplex_multi_session_reads_nested_terminal_identity(): + demo = _load_multi_demo_module() + + assert demo._event_session_id({"type": "session.heartbeat_ack", "session_id": "sid"}) == "sid" + assert ( + demo._event_session_id({"type": "session.closed", "event": {"type": "session.closed", "session_id": "sid"}}) + == "sid" + ) + + +def test_realtime_duplex_demo_resolves_distinct_turn_inputs(): + demo = _load_demo_module() + primary = Path("first.wav") + + assert demo._turn_input_paths(primary, [], turns=3) == [primary, primary, primary] + assert demo._turn_input_paths(primary, ["second.wav", "third.wav"], turns=3) == [ + primary, + Path("second.wav"), + Path("third.wav"), + ] + with pytest.raises(ValueError, match="one --turn-input-wav"): + demo._turn_input_paths(primary, ["second.wav"], turns=3) + + +def test_realtime_duplex_demo_explicitly_enables_native_runtime_before_connect(): + demo = _load_demo_module() + + url = demo._url_with_model( + "ws://localhost:8099/v1/realtime?duplex=1", + "openbmb/MiniCPM-o-4_5", + ) + + query = parse_qs(urlsplit(url).query) + assert query["minicpmo45_native_duplex"] == ["1"] + + +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(): + demo = _load_demo_module() + + event = demo._session_update_event( + SimpleNamespace( + model="openbmb/MiniCPM-o-4_5", + output_audio_format="pcm16", + short_ack_ms=1200, + session_id="reopen-e2e", + ) + ) + + assert event["type"] == "session.update" + assert "session_id" not in event + assert event["session"]["session_id"] == "reopen-e2e" + assert event["session"]["extra_body"]["minicpmo45_native_duplex"] is True + + +def test_realtime_duplex_demo_response_required_uses_deterministic_sampling(): + demo = _load_demo_module() + + event = demo._session_update_event( + SimpleNamespace( + model="openbmb/MiniCPM-o-4_5", + output_audio_format="pcm16", + short_ack_ms=1200, + session_id=None, + validation_mode="response-required", + temperature=None, + ) + ) + + assert event["session"]["temperature"] == 0.0 + + +def test_realtime_duplex_demo_model_policy_preserves_default_sampling(): + demo = _load_demo_module() + + event = demo._session_update_event( + SimpleNamespace( + model="openbmb/MiniCPM-o-4_5", + output_audio_format="pcm16", + short_ack_ms=1200, + session_id=None, + validation_mode="model-policy", + temperature=None, + ) + ) + + assert "temperature" not in event["session"] + + +def test_realtime_duplex_demo_resolves_explicit_turn_durations(): + demo = _load_demo_module() + + assert demo._turn_durations([], turns=3, first_turn_ms=3000) == [3000, 1200, 1200] + assert demo._turn_durations([0, 900, 1500], turns=3, first_turn_ms=3000) == [None, 900, 1500] + with pytest.raises(ValueError, match="one --turn-duration-ms"): + demo._turn_durations([0, 900], turns=3, first_turn_ms=3000) + with pytest.raises(ValueError, match="non-negative"): + demo._turn_durations([0, -1, 900], turns=3, first_turn_ms=3000) + + +def test_realtime_duplex_demo_resolves_transcript_labels_for_requested_turns(): + demo = _load_demo_module() + + assert demo._turn_transcripts("first", turns=1) == ["first"] + assert demo._turn_transcripts("first", turns=4) == ["first", "继续", "再说一次", "turn-4"] + + +def test_realtime_duplex_demo_reads_response_playback_cursor(): + demo = _load_demo_module() + state = demo.DemoState() + state.add( + { + "type": "response.done", + "response_id": "resp-1", + "response": { + "id": "resp-1", + "metadata": { + "playback": { + "sent_ms": 27920, + "played_ms": 0, + } + }, + }, + } + ) + + assert state.response_playback_sent_ms("resp-1") == 27920 + + +def test_realtime_duplex_demo_partitions_timing_by_response_identity(): + demo = _load_demo_module() + state = demo.DemoState() + for response_id, created_at_s, audio_at_s, token_count in ( + ("resp-1", 10.0, 10.1, 3), + ("resp-2", 20.0, 20.2, 5), + ): + state.add( + {"type": "response.created", "response": {"id": response_id}}, + received_at_s=created_at_s, + ) + state.add( + { + "type": "response.audio.delta", + "response_id": response_id, + "delta": base64.b64encode(b"audio").decode("ascii"), + "metadata": { + "audio_duration_ms": 80, + "vllm_omni": { + "stage_metrics": { + "0": { + "num_tokens_out": token_count, + "vllm_itls_ms": [10.0], + } + } + }, + }, + }, + received_at_s=audio_at_s, + ) + + timings = state.response_timing_summaries() + + assert timings["resp-1"]["stage0_tokens"]["output_token_count"] == 3 + assert timings["resp-1"]["audio_output"]["response_created_to_first_audio_ms"] == 100.0 + assert timings["resp-1"]["audio_output"]["commit_to_first_audio_ms"] is None + assert timings["resp-2"]["stage0_tokens"]["output_token_count"] == 5 + assert timings["resp-2"]["audio_output"]["response_created_to_first_audio_ms"] == 200.0 + assert timings["resp-2"]["audio_output"]["commit_to_first_audio_ms"] is None + + +def test_realtime_duplex_demo_model_policy_accepts_one_listen_per_streamed_turn(): + demo = _load_demo_module() + state = demo.DemoState() + state.add({"type": "session.created"}) + for turn in range(3): + state.add({"type": "input_audio_buffer.speech_started", "turn": turn}) + state.add( + { + "type": "response.listen", + "response": { + "status": "listening", + "metadata": {"model_listen": True}, + }, + } + ) + state.add({"type": "input_audio_buffer.committed", "turn": turn}) + + assert state.model_policy_event_order_ok(expected_turns=3) + + +def test_realtime_duplex_demo_model_policy_accepts_continuous_input_without_commits(): + demo = _load_demo_module() + state = demo.DemoState() + state.add({"type": "session.created"}) + state.add({"type": "input_audio_buffer.speech_started"}) + for _ in range(3): + state.add( + { + "type": "response.listen", + "response": {"metadata": {"model_listen": True}}, + } + ) + + assert not state.model_policy_event_order_ok(expected_turns=3) + assert state.model_policy_event_order_ok(expected_turns=3, require_input_commit=False) + + +def test_realtime_duplex_demo_accepts_overlap_unit_terminating_continuous_response(): + demo = _load_demo_module() + state = demo.DemoState() + state.add({"type": "response.created", "response_id": "resp-1"}) + state.add({"type": "response.done", "response_id": "resp-1"}) + + assert demo._continuous_overlap_terminal_is_outcome( + state, + validation_mode="model-policy", + model_unit_ready_while_active=True, + before_created=1, + before_model_listen=0, + ) + assert not demo._continuous_overlap_terminal_is_outcome( + state, + validation_mode="response-required", + model_unit_ready_while_active=True, + before_created=1, + before_model_listen=0, + ) + + +def test_realtime_duplex_demo_full_turn_duration_does_not_slice_audio(): + demo = _load_demo_module() + pcm16 = b"\x01\x00" * (demo.PCM16_SAMPLE_RATE * 2) + + assert demo._select_turn_audio(pcm16, None) == pcm16 + assert len(demo._select_turn_audio(pcm16, 1000)) == demo.PCM16_SAMPLE_RATE * demo.PCM16_BYTES_PER_SAMPLE + + +def test_realtime_duplex_demo_distinct_inputs_compare_audio_content(): + demo = _load_demo_module() + paths = [Path("first.wav"), Path("second.wav"), Path("third.wav")] + + assert demo._turn_inputs_are_distinct(paths, [b"first", b"second", b"third"]) is True + assert demo._turn_inputs_are_distinct(paths, [b"same", b"same", b"third"]) is False + assert demo._turn_inputs_are_distinct([paths[0], paths[0]], [b"first", b"second"]) is False + + +def _add_response_transcript(state, response_id, *, transcript, audio=True): + state.add( + { + "type": "response.audio.delta", + "response_id": response_id, + "delta": "YQ==" if audio else "", + } + ) + if transcript: + state.add( + { + "type": "response.audio_transcript.delta", + "response_id": response_id, + "delta": transcript, + } + ) + state.add( + { + "type": "response.audio_transcript.done", + "response_id": response_id, + "transcript": transcript, + } + ) + state.add( + { + "type": "response.done", + "response_id": response_id, + "response": {"id": response_id, "status": "completed"}, + } + ) + + +def test_realtime_duplex_demo_gate_checks_delta_done_and_turn_independence(): + demo = _load_demo_module() + state = demo.DemoState() + _add_response_transcript( + state, + "resp-1", + transcript="第一轮回答", + ) + _add_response_transcript( + state, + "resp-2", + transcript="第二轮回答", + ) + + result = demo._evaluate_transcript_integrity( + state, + ["resp-1", "resp-2"], + expected_empty_response_ids=set(), + require_cross_turn_independence=True, + ) + + assert result["transcript_delta_done_ok"] is True + assert result["cross_turn_independent_ok"] is True + + +def test_realtime_duplex_demo_speak_gate_accepts_one_decision_event_per_response(): + demo = _load_demo_module() + state = demo.DemoState() + state.add({"type": "response.created", "response": {"id": "resp-1"}}) + state.add( + { + "type": "response.speak", + "response_id": "resp-1", + "metadata": {"model_speak": True}, + } + ) + + assert demo._evaluate_response_speak_contract(state)["response_speak_contract_ok"] is True + + +def test_realtime_duplex_demo_speak_gate_rejects_duplicate_event(): + demo = _load_demo_module() + state = demo.DemoState() + state.add({"type": "response.created", "response": {"id": "resp-1"}}) + state.add({"type": "response.speak", "response_id": "resp-1"}) + state.add({"type": "response.speak", "response_id": "resp-1"}) + + result = demo._evaluate_response_speak_contract(state) + + assert result["response_speak_contract_ok"] is False + assert result["duplicate_response_speak_ids"] == ["resp-1"] + + +def test_realtime_duplex_demo_speak_gate_rejects_text_channel(): + demo = _load_demo_module() + state = demo.DemoState() + state.add({"type": "response.created", "response": {"id": "resp-1"}}) + state.add( + { + "type": "response.speak", + "response_id": "resp-1", + "text": "must be emitted through response.audio_transcript.delta", + } + ) + + result = demo._evaluate_response_speak_contract(state) + + assert result["response_speak_contract_ok"] is False + assert result["text_bearing_response_speak_count"] == 1 + + +def test_realtime_duplex_demo_gate_rejects_cross_turn_tail_reuse(): + demo = _load_demo_module() + state = demo.DemoState() + _add_response_transcript(state, "resp-1", transcript="第一轮回答") + _add_response_transcript(state, "resp-2", transcript="第二轮仍然带着第一轮回答") + + result = demo._evaluate_transcript_integrity( + state, + ["resp-1", "resp-2"], + expected_empty_response_ids=set(), + require_cross_turn_independence=True, + ) + + assert result["cross_turn_independent_ok"] is False + + +def test_realtime_duplex_demo_gate_rejects_terminal_only_previous_tail(): + demo = _load_demo_module() + state = demo.DemoState() + _add_response_transcript(state, "resp-1", transcript="上一轮回答结尾是的吗?") + _add_response_transcript(state, "resp-2", transcript="的吗?") + + result = demo._evaluate_transcript_integrity( + state, + ["resp-1", "resp-2"], + expected_empty_response_ids=set(), + require_cross_turn_independence=True, + ) + + assert result["cross_turn_independent_ok"] is False + + +def test_realtime_duplex_demo_gate_rejects_delta_done_mismatch(): + demo = _load_demo_module() + state = demo.DemoState() + _add_response_transcript(state, "resp-1", transcript="delta文本") + done = next(event for event in state.events if event.get("type") == "response.audio_transcript.done") + done["transcript"] = "另一个done文本" + + result = demo._evaluate_transcript_integrity( + state, + ["resp-1"], + expected_empty_response_ids=set(), + require_cross_turn_independence=False, + ) + + assert result["transcript_delta_done_ok"] is False + + +def test_realtime_duplex_demo_gate_rejects_terminal_only_stale_tail(): + demo = _load_demo_module() + state = demo.DemoState() + _add_response_transcript( + state, + "resp-empty", + transcript="的吗?", + audio=False, + ) + + result = demo._evaluate_transcript_integrity( + state, + ["resp-empty"], + expected_empty_response_ids={"resp-empty"}, + require_cross_turn_independence=False, + ) + + assert result["empty_turns_ok"] is False + + +def test_realtime_duplex_demo_gate_rejects_audio_without_transcript(): + demo = _load_demo_module() + state = demo.DemoState() + _add_response_transcript( + state, + "resp-audio-no-text", + transcript="", + audio=True, + ) + + result = demo._evaluate_transcript_integrity( + state, + ["resp-audio-no-text"], + expected_empty_response_ids=set(), + require_cross_turn_independence=False, + ) + + assert result["nonempty_audio_has_transcript_ok"] is False + + +def test_realtime_duplex_demo_response_required_rejects_unselected_audio_only_response(): + demo = _load_demo_module() + state = demo.DemoState() + _add_response_transcript(state, "resp-audio-only", transcript="", audio=True) + _add_response_transcript(state, "resp-with-text", transcript="你好", audio=True) + + assert ( + demo._all_audio_responses_have_transcript( + state, + ["resp-audio-only", "resp-with-text"], + ) + is False + ) + + +def test_realtime_duplex_demo_gate_rejects_incomplete_model_turn_sentence(): + demo = _load_demo_module() + state = demo.DemoState() + _add_response_transcript(state, "resp-1", transcript="哎,不是说好不") + + result = demo._evaluate_transcript_integrity( + state, + ["resp-1"], + expected_empty_response_ids=set(), + require_cross_turn_independence=False, + require_terminal_punctuation=True, + ) + + assert result["terminal_punctuation_ok"] is False + + +def test_realtime_duplex_demo_gate_accepts_short_interjection_without_punctuation(): + demo = _load_demo_module() + state = demo.DemoState() + _add_response_transcript(state, "resp-1", transcript="哈哈") + + result = demo._evaluate_transcript_integrity( + state, + ["resp-1"], + expected_empty_response_ids=set(), + require_cross_turn_independence=False, + require_terminal_punctuation=True, + ) + + assert result["terminal_punctuation_ok"] is True + + +def test_realtime_duplex_demo_response_required_skips_audio_only_model_turn(monkeypatch): + demo = _load_demo_module() + state = demo.DemoState() + + async def fake_send_pcm16(*args, **kwargs): + state.add({"type": "response.created", "response": {"id": "resp-audio-only"}}) + _add_response_transcript(state, "resp-audio-only", transcript="", audio=True) + state.add({"type": "response.created", "response": {"id": "resp-with-text"}}) + _add_response_transcript(state, "resp-with-text", transcript="你好", audio=True) + + class FakeWebSocket: + async def send(self, payload): + return None + + monkeypatch.setattr(demo, "_send_pcm16", fake_send_pcm16) + + response_id, outcome = asyncio.run( + demo._send_clean_turn( + FakeWebSocket(), + state, + b"\x00\x00", + transcript="fixture", + duration_ms=None, + chunk_ms=200, + timeout_s=0.1, + require_audio=True, + validation_mode="response-required", + ) + ) + + assert response_id == "resp-with-text" + assert outcome == "speak" + + +def test_realtime_duplex_demo_no_hint_realtime_turn_omits_transcript_hint(monkeypatch): + demo = _load_demo_module() + state = demo.DemoState() + captured = {} + + async def fake_send_pcm16(*args, **kwargs): + del args + captured.update(kwargs) + state.add( + { + "type": "response.listen", + "response": {"metadata": {"model_listen": True}}, + } + ) + + class FakeWebSocket: + async def send(self, payload): + del payload + + monkeypatch.setattr(demo, "_send_pcm16", fake_send_pcm16) + + response_id, outcome = asyncio.run( + demo._send_clean_turn( + FakeWebSocket(), + state, + b"\x00\x00", + transcript="turn label only", + duration_ms=None, + chunk_ms=200, + timeout_s=0.1, + require_audio=False, + validation_mode="model-policy", + send_transcript_hint=False, + realtime_input=True, + model_policy_settle_s=0.01, + ) + ) + + assert response_id is None + assert outcome == "listen" + assert captured["hints"] == {} + assert captured["realtime_delay"] is True + + +def test_realtime_duplex_demo_continuous_input_omits_browser_commit(monkeypatch): + demo = _load_demo_module() + state = demo.DemoState() + + async def fake_send_pcm16(*args, **kwargs): + del args, kwargs + state.add({"type": "response.created", "response": {"id": "resp-continuous"}}) + _add_response_transcript(state, "resp-continuous", transcript="连续输入", audio=True) + + class FakeWebSocket: + def __init__(self): + self.messages = [] + + async def send(self, payload): + self.messages.append(demo.json.loads(payload)) + + monkeypatch.setattr(demo, "_send_pcm16", fake_send_pcm16) + ws = FakeWebSocket() + + response_id, outcome = asyncio.run( + demo._send_clean_turn( + ws, + state, + b"\x00\x00", + transcript="fixture", + duration_ms=None, + chunk_ms=200, + timeout_s=0.1, + require_audio=True, + validation_mode="response-required", + commit_input=False, + ) + ) + + assert response_id == "resp-continuous" + assert outcome == "speak" + assert not any(message.get("type") == "input_audio_buffer.commit" for message in ws.messages) + + +def test_realtime_duplex_demo_no_hint_gate_does_not_require_transcription_event(): + demo = _load_demo_module() + + assert demo._input_transcription_ok(0, transcript_hints_enabled=False) + assert not demo._input_transcription_ok(0, transcript_hints_enabled=True) + assert demo._input_transcription_ok(3, transcript_hints_enabled=True) + + +def test_realtime_duplex_demo_gate_rejects_server_errors(): + demo = _load_demo_module() + state = demo.DemoState() + state.add({"type": "error", "code": "runtime_signal_failed", "error": "signal failed"}) + + errors = demo._unexpected_error_events(state) + + assert errors == [{"type": "error", "code": "runtime_signal_failed", "error": "signal failed"}] + + +def test_realtime_duplex_demo_model_policy_waits_for_speak_after_intermediate_listen(monkeypatch): + demo = _load_demo_module() + state = demo.DemoState() + + async def fake_send_pcm16(*args, **kwargs): + del args, kwargs + state.add( + { + "type": "response.listen", + "response": {"metadata": {"model_listen": True}}, + } + ) + + async def delayed_speak(): + await asyncio.sleep(0.01) + state.add({"type": "response.created", "response": {"id": "resp-delayed"}}) + _add_response_transcript(state, "resp-delayed", transcript="延迟回答", audio=False) + + asyncio.create_task(delayed_speak()) + + class FakeWebSocket: + async def send(self, payload): + del payload + + monkeypatch.setattr(demo, "_send_pcm16", fake_send_pcm16) + + response_id, outcome = asyncio.run( + demo._send_clean_turn( + FakeWebSocket(), + state, + b"\x00\x00", + transcript="fixture", + duration_ms=None, + chunk_ms=200, + timeout_s=0.2, + require_audio=False, + validation_mode="model-policy", + model_policy_settle_s=0.05, + ) + ) + + assert response_id == "resp-delayed" + assert outcome == "speak" + + +def test_realtime_duplex_demo_playback_gate_covers_unassigned_completed_response(): + demo = _load_demo_module() + state = demo.DemoState() + for response_id in ("resp-assigned", "resp-unassigned"): + state.add({"type": "response.created", "response": {"id": response_id}}) + state.add( + { + "type": "response.done", + "response": { + "id": response_id, + "metadata": {"playback": {"sent_ms": 1200}}, + }, + } + ) + state.add( + { + "type": "playback.acknowledged", + "event": { + "item_id": "item_resp-assigned", + "history_committed": True, + }, + } + ) + + assert not demo._all_response_playback_history_committed( + state, + state.completed_response_ids(), + ) + + +def test_realtime_duplex_demo_acks_unassigned_completed_response(): + demo = _load_demo_module() + state = demo.DemoState() + for response_id in ("resp-assigned", "resp-unassigned"): + state.add({"type": "response.created", "response": {"id": response_id}}) + state.add( + { + "type": "response.done", + "response": { + "id": response_id, + "metadata": {"playback": {"sent_ms": 1200}}, + }, + } + ) + state.add( + { + "type": "playback.acknowledged", + "event": { + "item_id": "item_resp-assigned", + "history_committed": True, + }, + } + ) + + class FakeWebSocket: + def __init__(self): + self.messages = [] + + async def send(self, payload): + message = demo.json.loads(payload) + self.messages.append(message) + state.add( + { + "type": "playback.acknowledged", + "event": { + "item_id": message["item_id"], + "history_committed": True, + }, + } + ) + + ws = FakeWebSocket() + asyncio.run(demo._ack_all_completed_response_playback(ws, state, timeout_s=0.1)) + + assert [message["response_id"] for message in ws.messages] == ["resp-unassigned"] + assert demo._all_response_playback_history_committed(state, state.completed_response_ids()) + + +def test_realtime_duplex_demo_model_policy_does_not_assume_one_response_per_physical_input(): + demo = _load_demo_module() + + assert demo._response_cardinality_ok( + ["resp-1", "resp-2", "resp-3", "resp-4"], + expected_turns=3, + validation_mode="model-policy", + ) + assert not demo._response_cardinality_ok( + ["resp-1", "resp-2", "resp-3", "resp-4"], + expected_turns=3, + validation_mode="response-required", + ) + + +def test_realtime_duplex_demo_only_requires_cross_turn_independence_for_response_required(): + demo = _load_demo_module() + + assert not demo._requires_cross_turn_independence( + validation_mode="model-policy", + distinct_inputs_required=True, + ) + assert demo._requires_cross_turn_independence( + validation_mode="response-required", + distinct_inputs_required=True, + ) + assert not demo._requires_cross_turn_independence( + validation_mode="response-required", + distinct_inputs_required=False, + ) + + +def test_realtime_duplex_demo_drains_late_model_response_before_close(): + demo = _load_demo_module() + state = demo.DemoState() + state.add({"type": "response.created", "response": {"id": "resp-first"}}) + state.add({"type": "response.done", "response": {"id": "resp-first"}}) + + class FakeWebSocket: + def __init__(self): + self.messages = [] + + async def send(self, payload): + self.messages.append(demo.json.loads(payload)) + + async def run_fixture(): + async def finish_late_response(): + await asyncio.sleep(0.01) + state.add({"type": "response.created", "response": {"id": "resp-late"}}) + await asyncio.sleep(0.01) + state.add({"type": "response.done", "response": {"id": "resp-late"}}) + + late_response = asyncio.create_task(finish_late_response()) + ws = FakeWebSocket() + await demo._drain_model_policy_responses( + ws, + state, + timeout_s=0.2, + settle_s=0.03, + ) + await late_response + return ws + + ws = asyncio.run(run_fixture()) + + assert state.completed_response_ids() == ["resp-first", "resp-late"] + assert not [message for message in ws.messages if message.get("type") == "session.close"] + + +def test_realtime_duplex_demo_model_response_drain_times_out_with_active_ids(): + demo = _load_demo_module() + state = demo.DemoState() + state.add({"type": "response.created", "response": {"id": "resp-active"}}) + + class FakeWebSocket: + async def send(self, payload): + del payload + + with pytest.raises(TimeoutError, match="resp-active"): + asyncio.run( + demo._drain_model_policy_responses( + FakeWebSocket(), + state, + timeout_s=0.03, + settle_s=0.01, + ) + ) + + +def test_realtime_duplex_demo_waits_at_each_model_unit_and_stops_after_speak(): + demo = _load_demo_module() + + class FakeWebSocket: + def __init__(self): + self.messages = [] + + async def send(self, payload): + self.messages.append(demo.json.loads(payload)) + + async def run_fixture(): + ws = FakeWebSocket() + model_unit_message_counts = [] + + async def on_model_unit_ready(): + model_unit_message_counts.append(len(ws.messages)) + await asyncio.sleep(0) + return len(model_unit_message_counts) < 2 + + await demo._send_pcm16( + ws, + b"\x01\x00" * (demo.PCM16_SAMPLE_RATE * 3), + chunk_ms=200, + realtime_delay=False, + on_model_unit_ready=on_model_unit_ready, + ) + return ws, model_unit_message_counts + + ws, model_unit_message_counts = asyncio.run(run_fixture()) + + assert model_unit_message_counts == [5, 10] + assert len(ws.messages) == 10 + + +def test_realtime_duplex_demo_listen_only_overlap_sends_next_turn_before_first_done(monkeypatch): + demo = _load_demo_module() + state = demo.DemoState() + send_calls = 0 + + async def fake_send_pcm16(*args, **kwargs): + nonlocal send_calls + send_calls += 1 + if send_calls == 1: + assert len(args[1]) == 2 + demo.PCM16_SAMPLE_RATE * demo.PCM16_BYTES_PER_SAMPLE + state.add({"type": "input_audio_buffer.speech_started", "turn": 1}) + state.add({"type": "response.created", "response": {"id": "resp-first"}}) + state.add( + { + "type": "response.audio.delta", + "response_id": "resp-first", + "delta": "YQ==", + } + ) + assert await kwargs["on_model_unit_ready"]() is False + return + + del args + assert not state.response_done("resp-first") + kwargs["on_model_unit_ready"]() + state.add({"type": "input_audio_buffer.speech_started", "turn": 2}) + _add_response_transcript(state, "resp-first", transcript="第一轮完成", audio=False) + state.add({"type": "response.created", "response": {"id": "resp-second"}}) + _add_response_transcript(state, "resp-second", transcript="第二轮完成", audio=True) + + class FakeWebSocket: + def __init__(self): + self.messages = [] + + async def send(self, payload): + self.messages.append(payload) + + monkeypatch.setattr(demo, "_send_pcm16", fake_send_pcm16) + ws = FakeWebSocket() + + response_ids, outcomes, overlap_ok = asyncio.run( + demo._send_listen_only_overlap_pair( + ws, + state, + b"\x00\x00", + b"\x01\x00", + transcripts=("first", "second"), + durations_ms=(None, None), + chunk_ms=200, + timeout_s=0.2, + send_transcript_hint=False, + realtime_input=True, + model_policy_settle_s=0.01, + validation_mode="response-required", + silence_ms=1000, + ) + ) + + assert response_ids == ["resp-first", "resp-second"] + assert outcomes == ["speak", "speak"] + assert overlap_ok is True + assert sum(demo.json.loads(message)["type"] == "input_audio_buffer.commit" for message in ws.messages) == 2 + + +def test_realtime_duplex_demo_playback_ack_identifies_and_commits_response(): + demo = _load_demo_module() + state = demo.DemoState() + response_id = "resp-overlap-second" + state.add( + { + "type": "response.done", + "response": { + "id": response_id, + "metadata": {"playback": {"sent_ms": 1200}}, + }, + } + ) + + class FakeWebSocket: + def __init__(self): + self.messages = [] + + async def send(self, payload): + message = demo.json.loads(payload) + self.messages.append(message) + state.add( + { + "type": "playback.acknowledged", + "event": { + "item_id": message.get("item_id"), + "history_committed": message.get("item_id") == f"item_{response_id}", + }, + } + ) + + ws = FakeWebSocket() + asyncio.run( + demo._ack_response_playback( + ws, + state, + response_id, + timeout_s=0.1, + label="overlap second", + ) + ) + + assert ws.messages == [ + { + "type": "playback.ack", + "response_id": response_id, + "item_id": f"item_{response_id}", + "played_ms": 1200, + "committed_ms": 1200, + } + ] + assert state.playback_history_committed_count == 1 + + +def test_realtime_duplex_demo_writes_audio_per_response(tmp_path): + demo = _load_demo_module() + state = demo.DemoState() + for response_id, payload in (("resp-1", b"\x01\x00"), ("resp-2", b"\x02\x00")): + state.add({"type": "response.created", "response": {"id": response_id}}) + state.add( + { + "type": "response.audio.delta", + "response_id": response_id, + "delta": demo.base64.b64encode(payload).decode(), + "sample_rate_hz": 24000, + } + ) + + demo._write_demo_artifacts(state, tmp_path, output_audio_format="pcm16") + + for index, expected in enumerate((b"\x01\x00", b"\x02\x00"), start=1): + with wave.open(str(tmp_path / f"response_{index:02d}.wav"), "rb") as wf: + assert wf.getframerate() == 24000 + assert wf.readframes(wf.getnframes()) == expected diff --git a/tests/examples/test_minicpmo_realtime_web_server.py b/tests/examples/test_minicpmo_realtime_web_server.py new file mode 100644 index 000000000..6c35760a7 --- /dev/null +++ b/tests/examples/test_minicpmo_realtime_web_server.py @@ -0,0 +1,99 @@ +import hashlib +import importlib.util +import sys +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient + +SERVER_PATH = Path(__file__).resolve().parents[2] / "examples/online_serving/minicpmo/realtime_web/server.py" +spec = importlib.util.spec_from_file_location( + "minicpmo_realtime_web_server_test", + SERVER_PATH, +) +assert spec is not None and spec.loader is not None +server = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = server +spec.loader.exec_module(server) + +APP_DIR = server.APP_DIR +STATIC_DIR = server.STATIC_DIR +_join_ws_url = server._join_ws_url +build_app = server.build_app + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_join_ws_url_preserves_realtime_query(): + assert ( + _join_ws_url( + "ws://127.0.0.1:8099/", + "/v1/realtime", + "duplex=1&model=openbmb%2FMiniCPM-o-4_5", + ) + == "ws://127.0.0.1:8099/v1/realtime?duplex=1&model=openbmb%2FMiniCPM-o-4_5" + ) + + +def test_build_app_serves_health_and_injected_client_config(): + app = build_app( + ws_backend="ws://127.0.0.1:9001", + model="local/MiniCPM-o-4_5", + ) + client = TestClient(app) + + health = client.get("/healthz") + assert health.status_code == 200 + assert health.text == "ok" + + index = client.get("/") + assert index.status_code == 200 + assert '"model": "local/MiniCPM-o-4_5"' in index.text + assert '"realtimePath": "v1/realtime"' in index.text + assert "__FULL_DUPLEX_CONFIG__" not in index.text + + +def test_build_app_injects_public_realtime_url(): + public_realtime_url = "wss://proxy.example.test/backend/v1/realtime" + app = build_app( + ws_backend="ws://127.0.0.1:9001", + public_realtime_url=public_realtime_url, + ) + client = TestClient(app) + + index = client.get("/") + + assert index.status_code == 200 + assert f'"realtimePath": "{public_realtime_url}"' in index.text + + +def test_build_app_versions_client_bundle_and_disables_index_cache(): + client = TestClient(build_app()) + expected_version_hash = hashlib.sha256() + for asset_name in ("app.js", "pcm_worklet.js", "playback_worklet.js"): + expected_version_hash.update((STATIC_DIR / asset_name).read_bytes()) + expected_version = expected_version_hash.hexdigest()[:12] + + assert (APP_DIR / "index.html").is_file() + assert (STATIC_DIR / "app.js").is_file() + assert not (APP_DIR.parent / "static").exists() + assert not (STATIC_DIR / "index.html").exists() + + index = client.get("/") + + assert index.status_code == 200 + assert f'src="static/app.js?v={expected_version}"' in index.text + assert index.headers["cache-control"] == "no-store" + + +def test_build_app_exposes_realtime_websocket_and_static_assets(): + app = build_app() + paths = {route.path for route in app.routes} + + assert "/v1/realtime" in paths + assert "/static" in paths + + client = TestClient(app) + assert client.get("/static/app.js").status_code == 200 + assert client.get("/static/pcm_worklet.js").status_code == 200 + assert client.get("/static/playback_worklet.js").status_code == 200 diff --git a/tests/examples/test_minicpmo_realtime_web_static.py b/tests/examples/test_minicpmo_realtime_web_static.py new file mode 100644 index 000000000..aa9f60fe2 --- /dev/null +++ b/tests/examples/test_minicpmo_realtime_web_static.py @@ -0,0 +1,209 @@ +import re +import shutil +import subprocess +import textwrap +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[2] / "examples" / "online_serving" / "minicpmo" / "realtime_web" +APP_ROOT = ROOT / "app" +STATIC_ROOT = APP_ROOT / "static" + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_page_exposes_focused_call_conversation_and_log_surfaces(): + html = (APP_ROOT / "index.html").read_text(encoding="utf-8") + + assert 'id="callButton"' in html + assert 'id="muteButton"' in html + assert 'id="connectionState"' in html + assert 'id="modelState"' in html + assert 'id="conversation"' in html + assert 'id="eventLog"' in html + assert ".*?)\n\s*\)", + source, + re.DOTALL, + ) + assert ref_audio_arg is not None + assert "required=True" in ref_audio_arg.group("body") + assert "Optional reference voice" not in ref_audio_arg.group("body") + + +def test_client_sends_ref_audio_in_realtime_session_contract(): + source = (STATIC_ROOT / "app.js").read_text(encoding="utf-8") + + assert "if (config.refAudio) session.ref_audio = config.refAudio;" in source + assert "extraBody.ref_audio" not in source + + +def test_client_has_transactional_cleanup_and_visible_event_logging(): + source = (STATIC_ROOT / "app.js").read_text(encoding="utf-8") + + assert "async function stopSession" in source + assert "type: 'session.close'" in source + assert "waitForSessionClosed" in source + assert "SESSION_CLOSE_TIMEOUT_MS" in source + assert "case 'session.closed':" in source + assert "track.stop()" in source + assert "clearInterval(sendTimer)" in source + assert "appendEventLog(event)" in source + assert "stopSession({ terminal: false })" in source + + +def test_client_keeps_microphone_upload_active_during_assistant_playback(): + source = (STATIC_ROOT / "app.js").read_text(encoding="utf-8") + upload_gate = re.search(r"function microphoneUploadEnabled\(\) \{(?P.*?)\n \}", source, re.DOTALL) + begin_assistant = re.search(r"function beginAssistant\(responseId\) \{(?P.*?)\n \}", source, re.DOTALL) + + assert upload_gate is not None + assert "return running && !muted;" in upload_gate.group("body") + assert "assistantActive" not in upload_gate.group("body") + assert begin_assistant is not None + assert "pendingCapture = []" not in begin_assistant.group("body") + + +def test_audio_worklets_define_capture_and_playback_processors(): + capture = (STATIC_ROOT / "pcm_worklet.js").read_text(encoding="utf-8") + playback = (STATIC_ROOT / "playback_worklet.js").read_text(encoding="utf-8") + + assert "registerProcessor('fullduplex-pcm-capture'" in capture + assert "Int16Array" in capture + assert "registerProcessor('fullduplex-pcm-playback'" in playback + assert "playback-drained" in playback + assert "clear" in playback + + +def test_audio_worklet_urls_use_the_static_asset_version(): + server = (ROOT / "server.py").read_text(encoding="utf-8") + app = (STATIC_ROOT / "app.js").read_text(encoding="utf-8") + + assert '"appVersion": app_version' in server + assert 'STATIC_DIR / "playback_worklet.js"' in server + assert 'STATIC_DIR / "pcm_worklet.js"' in server + assert "staticAssetUrl('static/playback_worklet.js')" in app + assert "staticAssetUrl('static/pcm_worklet.js')" in app + + +def test_playback_worklet_buffers_first_400ms_and_reports_underruns(): + app = (STATIC_ROOT / "app.js").read_text(encoding="utf-8") + playback = (STATIC_ROOT / "playback_worklet.js").read_text(encoding="utf-8") + + assert "INITIAL_PLAYBACK_BUFFER_MS = 400" in app + assert "initialBufferMs" in app + assert "responseId" in app + assert "playback-underrun" in app + assert "underrunMs" in app + assert "initialBufferFrames" in playback + assert "playback-underrun" in playback + assert "underrunFrames" in playback + assert "underrunMs" in playback + + +def test_playback_worklet_waits_before_playing_and_rebuffers_after_underrun(): + node = shutil.which("node") + if node is None: + pytest.skip("node is required for the AudioWorklet regression test") + + script = textwrap.dedent( + """ + const fs = require('fs'); + const vm = require('vm'); + + global.sampleRate = 1000; + global.AudioWorkletProcessor = class { + constructor() { + this.port = { onmessage: null, postMessage: () => {} }; + } + }; + let Processor = null; + global.registerProcessor = (_name, processor) => { Processor = processor; }; + vm.runInThisContext(fs.readFileSync(process.argv[1], 'utf8')); + + const processor = new Processor(); + const render = () => { + const output = new Float32Array(100); + processor.process([], [[output]]); + return output; + }; + const assert = (condition, message) => { + if (!condition) throw new Error(message); + }; + + const first = new Int16Array(150); + first.fill(16384); + processor.handleMessage({ + type: 'audio', + pcm: first, + responseId: 'response-1', + initialBufferMs: 400, + }); + assert(!processor.started, 'large first delta must not bypass wall-clock prebuffer'); + assert(render().every((sample) => sample === 0), 'first prebuffer render must be silent'); + assert(render().every((sample) => sample === 0), 'second prebuffer render must be silent'); + assert(render().every((sample) => sample === 0), 'third prebuffer render must be silent'); + assert(render().every((sample) => sample === 0), 'fourth prebuffer render must be silent'); + const firstPlayback = render(); + assert(firstPlayback.some((sample) => sample !== 0), 'playback must start after prebuffer'); + + const underrun = render(); + assert(!processor.started, 'an empty queue must return to buffering'); + assert(underrun[underrun.length - 1] === 0, 'underrun boundary must fade to zero'); + + const resumed = new Int16Array(300); + resumed.fill(8192); + processor.handleMessage({ + type: 'audio', + pcm: resumed, + responseId: 'response-1', + initialBufferMs: 400, + }); + assert(render().every((sample) => sample === 0), 'resume must rebuild jitter buffer'); + assert(render().every((sample) => sample === 0), 'resume must keep rebuilding jitter buffer'); + assert(render().every((sample) => sample === 0), 'resume must keep waiting'); + assert(render().every((sample) => sample === 0), 'resume must wait the full buffer interval'); + const resumedPlayback = render(); + assert(resumedPlayback.some((sample) => sample !== 0), 'playback must resume after rebuffer'); + assert( + Math.abs(resumedPlayback[0]) < Math.abs(resumedPlayback[50]), + 'resumed playback must fade in instead of jumping from silence', + ); + """ + ) + subprocess.run( + [node, "-e", script, str(STATIC_ROOT / "playback_worklet.js")], + check=True, + capture_output=True, + text=True, + ) diff --git a/tests/fullduplex/engine/test_duplex_control_client.py b/tests/fullduplex/engine/test_duplex_control_client.py new file mode 100644 index 000000000..c1fa9a592 --- /dev/null +++ b/tests/fullduplex/engine/test_duplex_control_client.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import pytest + +from vllm_omni.experimental.fullduplex.engine.duplex_control_client import ( + DuplexControlClient, + DuplexControlRequestError, +) +from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseActivity +from vllm_omni.experimental.fullduplex.engine.messages import ( + DuplexControlError, + DuplexControlResultMessage, + DuplexFence, + ResumeDuplexSessionMessage, + TouchDuplexSessionMessage, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Transport: + def __init__(self) -> None: + self.calls = [] + + def execute(self, key, message, **kwargs): + self.calls.append((key, message, kwargs)) + return DuplexControlResultMessage( + control_id=message.control_id, + fence=message.fence, + operation="touch" if isinstance(message, TouchDuplexSessionMessage) else "resume", + session_id=message.session_id, + ok=True, + stage_results=[], + ) + + +class _UnsupportedTransport: + def execute(self, key, message, **kwargs): + del key, kwargs + return DuplexControlResultMessage( + control_id=message.control_id, + fence=message.fence, + operation="touch", + session_id=message.session_id, + ok=False, + stage_results=[{"result": {"supported": False}}], + unsupported_count=1, + error=DuplexControlError(code="invalid_capability", message="unsupported", retryable=False), + accepted_fence=DuplexFence(message.session_id, epoch=2), + lease_generation=7, + ) + + +def test_control_client_routes_touch_and_resume_by_control_id() -> None: + transport = _Transport() + control_ids = iter(("touch-id", "resume-id")) + client = DuplexControlClient(transport, control_id_factory=lambda: next(control_ids)) + fence = DuplexFence("sid-client") + + assert ( + client.touch( + fence.session_id, + fence=fence, + activity=DuplexLeaseActivity.PLAYBACK_ACK, + timeout=2.0, + )["ok"] + is True + ) + assert ( + client.resume( + fence.session_id, + fence=fence, + expected_lease_generation=7, + timeout=3.0, + )["ok"] + is True + ) + + touch_key, touch_message, touch_kwargs = transport.calls[0] + assert touch_key == ("duplex", "touch-id") + assert isinstance(touch_message, TouchDuplexSessionMessage) + assert touch_message.activity == DuplexLeaseActivity.PLAYBACK_ACK.value + assert touch_kwargs["timeout"] == 2.0 + + resume_key, resume_message, resume_kwargs = transport.calls[1] + assert resume_key == ("duplex", "resume-id") + assert isinstance(resume_message, ResumeDuplexSessionMessage) + assert resume_message.expected_lease_generation == 7 + assert resume_kwargs["timeout"] == 3.0 + + +def test_control_client_treats_ok_false_as_authoritative_failure() -> None: + client = DuplexControlClient(_UnsupportedTransport(), control_id_factory=lambda: "unsupported") + fence = DuplexFence("sid-unsupported") + + with pytest.raises(DuplexControlRequestError, match="duplex touch failed") as exc_info: + client.touch( + fence.session_id, + fence=fence, + activity=DuplexLeaseActivity.HEARTBEAT, + timeout=1.0, + ) + + assert exc_info.value.code == "invalid_capability" + assert exc_info.value.retryable is False + assert exc_info.value.accepted_fence == DuplexFence(fence.session_id, epoch=2) + assert exc_info.value.lease_generation == 7 diff --git a/tests/fullduplex/engine/test_duplex_control_plane.py b/tests/fullduplex/engine/test_duplex_control_plane.py new file mode 100644 index 000000000..46df33ffa --- /dev/null +++ b/tests/fullduplex/engine/test_duplex_control_plane.py @@ -0,0 +1,1270 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import FrozenInstanceError + +import msgspec +import pytest + +from vllm_omni.experimental.fullduplex.engine.duplex_control_plane import ( + DuplexControlPlane, + DuplexOutputContext, + DuplexRequestIdentity, + DuplexStageRequestContext, + DuplexStageSubmission, + DuplexStageSubmissionResult, +) +from vllm_omni.experimental.fullduplex.engine.duplex_runtime import ( + DuplexAppendPlan, + DuplexInputMode, + DuplexRuntimeCapabilities, +) +from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseActivity, DuplexLeaseConfig +from vllm_omni.experimental.fullduplex.engine.messages import ( + AppendDuplexInputMessage, + CloseDuplexSessionMessage, + DuplexFence, + DuplexSessionLifecycleMessage, + OpenDuplexSessionMessage, + ResumeDuplexSessionMessage, + SignalDuplexTurnMessage, + TouchDuplexSessionMessage, +) + +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/fullduplex/engine/test_duplex_deploy_config.py b/tests/fullduplex/engine/test_duplex_deploy_config.py new file mode 100644 index 000000000..abb387573 --- /dev/null +++ b/tests/fullduplex/engine/test_duplex_deploy_config.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Focused tests for typed deployment-only runtime configuration.""" + +from __future__ import annotations + +from dataclasses import FrozenInstanceError + +import pytest + +from vllm_omni.config.stage_config import ( + DuplexSessionRuntimeConfig, + load_deploy_config, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_duplex_session_runtime_defaults_are_typed_and_immutable(tmp_path) -> None: + deploy_path = tmp_path / "duplex.yaml" + deploy_path.write_text("duplex_session: {}\nstages: []\n", encoding="utf-8") + + deploy = load_deploy_config(deploy_path) + + assert isinstance(deploy.duplex_session, DuplexSessionRuntimeConfig) + assert deploy.duplex_session.idle_ttl_s == 300.0 + assert deploy.duplex_session.disconnect_grace_s == 30.0 + assert deploy.duplex_session.reaper_interval_s == 5.0 + assert deploy.duplex_session.resume_replay_ttl_s == 60.0 + assert deploy.duplex_session.resume_replay_max_bytes_per_session == 8 * 1024 * 1024 + assert deploy.duplex_session.max_pending_input_bytes_per_session == 16 * 1024 * 1024 + assert deploy.duplex_session.max_pending_turns_per_session == 4 + assert deploy.duplex_session.max_sessions == 1 + assert deploy.duplex_session.completed_append_cache_size == 256 + with pytest.raises(FrozenInstanceError): + deploy.duplex_session.idle_ttl_s = 1.0 # type: ignore[misc] + + +def test_duplex_session_runtime_accepts_disabled_idle_expiry(tmp_path) -> None: + deploy_path = tmp_path / "duplex.yaml" + deploy_path.write_text("duplex_session:\n idle_ttl_s: null\nstages: []\n", encoding="utf-8") + + deploy = load_deploy_config(deploy_path) + + assert deploy.duplex_session.idle_ttl_s is None + + +@pytest.mark.parametrize( + ("name", "value"), + [ + ("idle_ttl_s", 0), + ("idle_ttl_s", -1), + ("disconnect_grace_s", 0), + ("reaper_interval_s", -1), + ("resume_replay_ttl_s", 0), + ("resume_replay_max_bytes_per_session", -1), + ("max_pending_input_bytes_per_session", 0), + ("max_pending_turns_per_session", -1), + ("max_sessions", 0), + ("completed_append_cache_size", 0), + ], +) +def test_duplex_session_runtime_rejects_non_positive_values(tmp_path, name: str, value: int) -> None: + deploy_path = tmp_path / "duplex.yaml" + deploy_path.write_text(f"duplex_session:\n {name}: {value}\nstages: []\n", encoding="utf-8") + + with pytest.raises(ValueError, match=rf"duplex_session\.{name} must be positive"): + load_deploy_config(deploy_path) diff --git a/tests/fullduplex/engine/test_duplex_intermediate.py b/tests/fullduplex/engine/test_duplex_intermediate.py new file mode 100644 index 000000000..37602e62e --- /dev/null +++ b/tests/fullduplex/engine/test_duplex_intermediate.py @@ -0,0 +1,14 @@ +import pytest + +from vllm_omni.experimental.fullduplex.engine.intermediate import get_stream_request_key + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_get_stream_request_key_requires_stable_identifier(): + with pytest.raises(ValueError, match="stable request id"): + get_stream_request_key({"ids": {"tts": [1, 2, 3]}}) + + +def test_get_stream_request_key_accepts_global_request_id(): + assert get_stream_request_key({"global_request_id": ["duplex-sid-stage0"]}) == "duplex-sid-stage0" diff --git a/tests/fullduplex/engine/test_duplex_lease.py b/tests/fullduplex/engine/test_duplex_lease.py new file mode 100644 index 000000000..80e7a73eb --- /dev/null +++ b/tests/fullduplex/engine/test_duplex_lease.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from dataclasses import FrozenInstanceError + +import pytest + +from vllm_omni.experimental.fullduplex.engine.duplex_runtime import DuplexInputMode +from vllm_omni.experimental.fullduplex.engine.duplex_session import DuplexSessionRuntimeManager +from vllm_omni.experimental.fullduplex.engine.lease import ( + DuplexLeaseActivity, + DuplexLeaseConfig, +) +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class FakeMonotonicClock: + def __init__(self, value: float) -> None: + self.value = value + + def __call__(self) -> float: + return self.value + + def advance(self, seconds: float) -> None: + self.value += seconds + + +def _lease_config(*, idle_ttl_s: float | None = 300.0) -> DuplexLeaseConfig: + return DuplexLeaseConfig( + idle_ttl_s=idle_ttl_s, + disconnect_grace_s=30.0, + ) + + +def test_open_touch_detach_and_idle_expiry_use_monotonic_time() -> None: + clock = FakeMonotonicClock(100.0) + manager = DuplexSessionRuntimeManager(clock=clock) + fence = DuplexFence("sid-expiry") + session = manager.open_session(fence, lease_config=_lease_config()) + + assert session.lease.last_activity == 100.0 + clock.advance(10.0) + session.touch(fence, DuplexLeaseActivity.HEARTBEAT) + session.detach(fence) + clock.advance(29.0) + assert session.lease.disconnect_grace_expired(clock()) is False + assert manager.collect_expired() == [] + + clock.advance(272.0) + expired = manager.collect_expired() + + assert [item.session_id for item in expired] == ["sid-expiry"] + assert expired[0].reason == "disconnect_grace_expired" + assert manager.get("sid-expiry") is session + manager.finalize_close_session(session) + assert manager.get("sid-expiry") is None + + +def test_detach_grace_and_resume_advance_token_independent_lease_generation() -> None: + clock = FakeMonotonicClock(10.0) + manager = DuplexSessionRuntimeManager(clock=clock) + fence = DuplexFence("sid-resume") + session = manager.open_session(fence, lease_config=_lease_config()) + + session.detach(fence) + clock.advance(29.0) + assert session.lease.disconnect_grace_expired(clock()) is False + clock.advance(2.0) + assert session.lease.disconnect_grace_expired(clock()) is True + + generation = session.resume(fence, expected_lease_generation=0) + + assert generation == 1 + assert session.lease.detached_at is None + assert session.lease.disconnect_grace_expired(clock()) is False + with pytest.raises(ValueError, match="lease generation mismatch"): + session.resume(fence, expected_lease_generation=0) + + +def test_active_operation_prevents_mid_transaction_expiry() -> None: + clock = FakeMonotonicClock(0.0) + manager = DuplexSessionRuntimeManager(clock=clock) + fence = DuplexFence("sid-operation") + session = manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=5.0)) + + session.begin_operation(fence, "append-1") + clock.advance(10.0) + assert manager.collect_expired() == [] + + session.end_operation(fence, "append-1") + assert manager.collect_expired() == [] + clock.advance(6.0) + assert [item.session_id for item in manager.collect_expired()] == ["sid-operation"] + + +def test_sessions_have_independent_activity_deadlines_and_resources() -> None: + clock = FakeMonotonicClock(0.0) + manager = DuplexSessionRuntimeManager(clock=clock) + fence_a = DuplexFence("sid-a") + fence_b = DuplexFence("sid-b") + session_a = manager.open_session(fence_a, lease_config=_lease_config(idle_ttl_s=10.0)) + session_b = manager.open_session(fence_b, lease_config=_lease_config(idle_ttl_s=10.0)) + session_a.reserve_stage_request(0, "req-a-reserved", fence=fence_a) + session_a.bind_stage_request(1, "req-a-submitted", fence=fence_a) + session_b.bind_stage_request(0, "req-b", fence=fence_b) + + clock.advance(6.0) + session_b.touch(fence_b, DuplexLeaseActivity.MODEL_OUTPUT) + clock.advance(5.0) + expired = manager.collect_expired() + + assert len(expired) == 1 + assert expired[0].session_id == "sid-a" + assert expired[0].reserved_request_ids == ("req-a-reserved",) + assert expired[0].submitted_request_ids == ("req-a-submitted",) + assert manager.get("sid-b") is session_b + assert session_b.resource_request_ids() == ["req-b"] + + +def test_close_and_reaper_select_exactly_one_terminal_transition() -> None: + clock = FakeMonotonicClock(0.0) + manager = DuplexSessionRuntimeManager(clock=clock) + fence = DuplexFence("sid-race") + manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=1.0)) + clock.advance(2.0) + + assert len(manager.collect_expired()) == 1 + assert manager.close_session(fence, reason="explicit_close") is None + assert manager.collect_expired() == [] + + fence_2 = DuplexFence("sid-race-2") + manager.open_session(fence_2, lease_config=_lease_config(idle_ttl_s=1.0)) + assert manager.close_session(fence_2, reason="explicit_close") is not None + clock.advance(2.0) + assert manager.collect_expired() == [] + + +def test_stale_fence_cannot_touch_detach_or_resume_lease() -> None: + clock = FakeMonotonicClock(0.0) + manager = DuplexSessionRuntimeManager(clock=clock) + current = DuplexFence("sid-fence", epoch=1) + stale = DuplexFence("sid-fence", epoch=0) + session = manager.open_session(current, lease_config=_lease_config()) + + with pytest.raises(RuntimeError, match="fence mismatch"): + session.touch(stale, DuplexLeaseActivity.APPEND) + with pytest.raises(RuntimeError, match="fence mismatch"): + session.detach(stale) + with pytest.raises(RuntimeError, match="fence mismatch"): + session.resume(stale, expected_lease_generation=0) + + +def test_disabled_idle_expiry_never_collects_session() -> None: + clock = FakeMonotonicClock(0.0) + manager = DuplexSessionRuntimeManager(clock=clock) + manager.open_session(DuplexFence("sid-no-expiry"), lease_config=_lease_config(idle_ttl_s=None)) + + clock.advance(1_000_000.0) + + assert manager.collect_expired() == [] + + +def test_detached_session_expires_at_disconnect_grace_when_idle_ttl_is_disabled() -> None: + clock = FakeMonotonicClock(0.0) + manager = DuplexSessionRuntimeManager(clock=clock) + fence = DuplexFence("sid-disconnect-grace") + session = manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=None)) + + session.detach(fence) + clock.advance(29.0) + assert manager.collect_expired() == [] + + clock.advance(1.0) + expired = manager.collect_expired() + + assert [item.session_id for item in expired] == [fence.session_id] + assert expired[0].reason == "disconnect_grace_expired" + + +def test_runtime_manager_enforces_server_owned_session_admission() -> None: + manager = DuplexSessionRuntimeManager(max_sessions=2) + manager.open_session(DuplexFence("sid-admission-a")) + manager.open_session(DuplexFence("sid-admission-b")) + + with pytest.raises(RuntimeError, match="duplex_session_capacity_exhausted"): + manager.open_session(DuplexFence("sid-admission-c")) + + +def test_completed_append_cache_is_bounded() -> None: + manager = DuplexSessionRuntimeManager(completed_append_limit=2) + fence = DuplexFence("sid-completed-cache") + session = manager.open_session(fence) + + for index in range(3): + session.record_completed_append( + f"operation-{index}", + fence=fence, + mode=DuplexInputMode.TURN_COMMIT_ONLY, + final=True, + stage_results=[{"index": index}], + ) + + assert list(session.completed_appends) == ["operation-1", "operation-2"] + + session.accept_fence(DuplexFence(fence.session_id, epoch=1)) + assert session.completed_appends == {} + + +def test_lease_config_and_expiry_record_are_immutable() -> None: + config = _lease_config() + with pytest.raises(FrozenInstanceError): + config.idle_ttl_s = 1.0 # type: ignore[misc] + + clock = FakeMonotonicClock(0.0) + manager = DuplexSessionRuntimeManager(clock=clock) + manager.open_session(DuplexFence("sid-record"), lease_config=_lease_config(idle_ttl_s=1.0)) + clock.advance(2.0) + record = manager.collect_expired()[0] + with pytest.raises(FrozenInstanceError): + record.reason = "changed" # type: ignore[misc] diff --git a/tests/fullduplex/engine/test_duplex_runtime.py b/tests/fullduplex/engine/test_duplex_runtime.py new file mode 100644 index 000000000..7cce93161 --- /dev/null +++ b/tests/fullduplex/engine/test_duplex_runtime.py @@ -0,0 +1,480 @@ +from dataclasses import FrozenInstanceError +from types import SimpleNamespace + +import pytest +from vllm.sampling_params import SamplingParams + +from vllm_omni.experimental.fullduplex.engine import duplex_runtime +from vllm_omni.experimental.fullduplex.engine.contracts import ( + duplex_data_plane_request_info, + duplex_resource_request_belongs_to_session, + duplex_resource_request_id, +) +from vllm_omni.experimental.fullduplex.engine.duplex_runtime import ( + DuplexInputMode, + DuplexOutputAction, + DuplexOutputDecision, + DuplexRuntimeCapabilities, + DuplexSessionRuntimeManager, +) +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence +from vllm_omni.experimental.fullduplex.minicpmo45.runtime import ( + MiniCPMO45DuplexRuntimeExtension, + build_duplex_data_plane_prompt, + duplex_scheduler_token_budget, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_duplex_fence_is_immutable(): + fence = DuplexFence("session") + + with pytest.raises(FrozenInstanceError): + fence.epoch = 1 # type: ignore[misc] + + assert not hasattr(fence, "__dict__") + + +def test_duplex_runtime_tracks_stage_bindings_and_barge_in_epoch(): + manager = DuplexSessionRuntimeManager() + session = manager.open_session( + DuplexFence("sid-1"), + capabilities=DuplexRuntimeCapabilities( + input_modes={DuplexInputMode.APPEND_TOKENS}, + ), + ) + session.bind_stage_request(stage_id=0, request_id="req-stage0", fence=session.fence) + session.bind_stage_request(stage_id=1, request_id="req-stage1", fence=session.fence) + + update = session.append_input(mode=DuplexInputMode.APPEND_TOKENS, fence=session.fence) + next_fence = DuplexFence("sid-1", epoch=1) + stale_request_ids = session.release_fence(session.fence) + session.accept_fence(next_fence) + + assert update.seq == 1 + assert session.fence == next_fence + assert stale_request_ids == ["req-stage0", "req-stage1"] + assert session.stage_bindings == {} + + +def test_duplex_runtime_tracks_same_request_id_for_each_pipeline_stage(): + manager = DuplexSessionRuntimeManager() + fence = DuplexFence("sid-shared-pipeline-request") + session = manager.open_session(fence) + + session.reserve_stage_request(0, "req-shared", fence=fence) + session.bind_stage_request(0, "req-shared", fence=fence) + session.bind_stage_request(1, "req-shared", fence=fence) + + assert session.stage_bindings[0].request_id == "req-shared" + assert session.stage_bindings[1].request_id == "req-shared" + assert session.resource_request_ids() == ["req-shared"] + assert session.input_seq == 0 + + +def test_duplex_runtime_extension_validation_rejects_missing_methods(): + class IncompleteExtension: + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return defaults + + with pytest.raises(TypeError, match="plan_append"): + duplex_runtime.validate_duplex_runtime_extension(IncompleteExtension()) + + +def test_duplex_runtime_extension_validation_rejects_stage_count_mismatch(): + class WrongStageCountExtension(MiniCPMO45DuplexRuntimeExtension): + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return defaults[:1] + + with pytest.raises(ValueError, match="one sampling parameter per stage"): + duplex_runtime.validate_duplex_runtime_extension( + WrongStageCountExtension(), + sampling_defaults=(SamplingParams(), SamplingParams()), + ) + + +def test_duplex_runtime_extension_validation_rejects_sampling_type_mismatch(): + class WrongSamplingTypeExtension(MiniCPMO45DuplexRuntimeExtension): + def configure_sampling_params(self, *, runtime_config, defaults): + del runtime_config + return tuple(object() for _ in defaults) + + with pytest.raises(TypeError, match="sampling parameter type"): + duplex_runtime.validate_duplex_runtime_extension( + WrongSamplingTypeExtension(), + sampling_defaults=(SamplingParams(), SamplingParams()), + ) + + +def test_duplex_output_decision_metadata_is_immutable(): + decision = DuplexOutputDecision( + action=DuplexOutputAction.DIRECT_RESPONSE, + metadata={"model_listen": True}, + ) + + with pytest.raises(TypeError): + decision.metadata["model_listen"] = False + + +def _decide_minicpmo_output( + output: object, + *, + segment_token_ids: tuple[int, ...] = (), + segment_output_metadata: dict | None = None, +): + return MiniCPMO45DuplexRuntimeExtension().decide_output( + stage_id=0, + final_stage_id=1, + segment_finished=True, + segment_token_ids=segment_token_ids, + segment_output_metadata=segment_output_metadata or {}, + output=output, + ) + + +def test_minicpmo_extension_owns_stage_sampling_overrides(): + defaults = ( + SamplingParams(max_tokens=4), + SamplingParams(max_tokens=8), + ) + + configured = MiniCPMO45DuplexRuntimeExtension().configure_sampling_params( + runtime_config={ + "duplex_stage_max_tokens": {"0": 20}, + "duplex_stage_sampling_params": {"1": {"stop_token_ids": [151645]}}, + }, + defaults=defaults, + ) + + assert configured[0].max_tokens == 20 + assert configured[1].stop_token_ids == [151645] + assert defaults[0].max_tokens == 4 + assert 151645 not in (defaults[1].stop_token_ids or []) + + +def test_minicpmo_output_decision_uses_raw_streaming_token_snapshot(): + decision = _decide_minicpmo_output( + SimpleNamespace(outputs=[SimpleNamespace()]), + segment_token_ids=(151705,), + segment_output_metadata={"special_token_ids": {"listen_token_id": 151705}}, + ) + + assert decision is not None + assert decision.action is DuplexOutputAction.DIRECT_RESPONSE + assert decision.metadata["duplex_native_decision"] == "listen" + assert decision.metadata["model_listen"] is True + + +@pytest.mark.parametrize("attr", ["token_ids", "cumulative_token_ids"]) +def test_minicpmo_output_decision_ignores_output_level_token_history(attr): + output = SimpleNamespace( + multimodal_output={"special_token_ids": {"listen_token_id": 151705}}, + outputs=[SimpleNamespace()], + **{attr: [42, 151705]}, + ) + + assert _decide_minicpmo_output(output) is None + + +@pytest.mark.parametrize("attr", ["token_ids", "cumulative_token_ids"]) +def test_minicpmo_output_decision_uses_completion_token_ids(attr): + output = SimpleNamespace( + multimodal_output={"special_token_ids": {"listen_token_id": 151705}}, + outputs=[SimpleNamespace(**{attr: [42, 151705]})], + ) + + assert _decide_minicpmo_output(output) is not None + + +def test_minicpmo_output_decision_uses_completion_stop_reason(): + output = SimpleNamespace( + multimodal_output={"special_token_ids": {"listen_token_id": 151705}}, + outputs=[SimpleNamespace(stop_reason=151705)], + ) + + assert _decide_minicpmo_output(output) is not None + + +def test_duplex_runtime_cancel_fence_rejects_late_append_and_accepts_next_epoch(): + manager = DuplexSessionRuntimeManager() + cancelled_fence = DuplexFence("sid-cancel-race") + next_fence = DuplexFence("sid-cancel-race", epoch=1) + session = manager.open_session( + cancelled_fence, + capabilities=DuplexRuntimeCapabilities( + input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, + ), + ) + session.bind_stage_request( + stage_id=0, + request_id="req-cancelled", + fence=cancelled_fence, + ) + + stale_request_ids = session.cancel_fence(cancelled_fence, next_fence) + + assert stale_request_ids == ["req-cancelled"] + assert session.fence == next_fence + assert session.stage_bindings == {} + with pytest.raises(RuntimeError, match="fence mismatch"): + session.append_input( + mode=DuplexInputMode.APPEND_AUDIO_CHUNK, + fence=cancelled_fence, + ) + update = session.append_input( + mode=DuplexInputMode.APPEND_AUDIO_CHUNK, + fence=next_fence, + ) + assert update.seq == 1 + + +def test_duplex_runtime_stale_close_preserves_live_session_and_bindings(): + manager = DuplexSessionRuntimeManager() + current_fence = DuplexFence("sid-stale-close", epoch=1) + session = manager.open_session( + current_fence, + capabilities=DuplexRuntimeCapabilities( + input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, + ), + ) + session.bind_stage_request(0, "req-live", fence=current_fence) + + with pytest.raises(RuntimeError, match="fence mismatch"): + manager.close_session(DuplexFence("sid-stale-close")) + + assert manager.get("sid-stale-close") is session + assert session.stage_request_ids() == ["req-live"] + + +def test_duplex_runtime_reopen_rejects_late_append_from_old_incarnation(): + manager = DuplexSessionRuntimeManager() + old_fence = DuplexFence("sid-reopen", incarnation=0) + old_session = manager.open_session( + old_fence, + capabilities=DuplexRuntimeCapabilities( + input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, + ), + ) + manager.close_session(old_fence) + + new_fence = DuplexFence("sid-reopen", incarnation=1) + new_session = manager.open_session( + new_fence, + capabilities=old_session.capabilities, + ) + + with pytest.raises(RuntimeError, match="fence mismatch"): + new_session.append_input( + mode=DuplexInputMode.APPEND_AUDIO_CHUNK, + fence=old_fence, + ) + assert ( + new_session.append_input( + mode=DuplexInputMode.APPEND_AUDIO_CHUNK, + fence=new_fence, + ).seq + == 1 + ) + + +def test_duplex_prompt_expands_incarnation_metadata(): + fence = DuplexFence("sid-incarnation", incarnation=3) + + prompt = build_duplex_data_plane_prompt( + request_id=duplex_resource_request_id(fence, "stage0"), + fence=fence, + session_config={}, + runtime_config={}, + seq=1, + turn_seq=1, + mode=DuplexInputMode.APPEND_AUDIO_CHUNK, + payload={"is_speech": True}, + final=False, + ) + + assert prompt["model_intermediate_buffer"]["duplex"]["incarnation"] == 3 + + +def test_duplex_runtime_tracks_turn_local_append_sequence(): + manager = DuplexSessionRuntimeManager() + session = manager.open_session( + DuplexFence("sid-turn-seq"), + capabilities=DuplexRuntimeCapabilities( + input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}, + ), + ) + + first = session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK, fence=session.fence) + second = session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK, fence=session.fence) + next_turn = DuplexFence("sid-turn-seq", turn_id=1, response_seq=1) + third = session.append_input( + mode=DuplexInputMode.APPEND_AUDIO_CHUNK, + fence=next_turn, + ) + fourth = session.append_input( + mode=DuplexInputMode.APPEND_AUDIO_CHUNK, + fence=next_turn, + ) + + assert [first.seq, second.seq, third.seq, fourth.seq] == [1, 2, 3, 4] + assert [first.turn_seq, second.turn_seq, third.turn_seq, fourth.turn_seq] == [1, 2, 1, 2] + assert [first.turn_id, second.turn_id, third.turn_id, fourth.turn_id] == [0, 0, 1, 1] + + +def test_duplex_runtime_rejects_unsupported_append_mode(): + manager = DuplexSessionRuntimeManager() + session = manager.open_session( + DuplexFence("sid-2"), + capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.TURN_COMMIT_ONLY}), + ) + + with pytest.raises(ValueError, match="not supported"): + session.append_input(mode=DuplexInputMode.APPEND_TOKENS, fence=session.fence) + + +def test_duplex_data_plane_request_info_extracts_structured_stage_result(): + request_id, response_stage_id = duplex_data_plane_request_info( + { + "stage_results": [ + {"result": {"supported": True}}, + { + "result": { + "data_plane_append": True, + "request_id": "duplex-sid-e0-stage0-s1", + "response_stage_id": 1, + } + }, + ] + } + ) + + assert request_id == "duplex-sid-e0-stage0-s1" + assert response_stage_id == 1 + + +def test_duplex_data_plane_request_info_rejects_missing_request_id(): + assert duplex_data_plane_request_info( + { + "stage_results": [ + { + "result": { + "data_plane_append": True, + "request_id": "", + "response_stage_id": 1, + } + } + ] + } + ) == (None, None) + + +def test_duplex_scheduler_token_budget_estimates_pcm_slots(): + assert ( + duplex_scheduler_token_budget( + { + "audio": "AAAAAA==", + "format": "pcm_f32le", + "sample_rate_hz": 16000, + } + ) + == 16 + ) + + +def test_duplex_scheduler_token_budget_ignores_client_budget_fields(): + assert ( + duplex_scheduler_token_budget( + { + "audio": "AAAAAA==", + "format": "pcm_f32le", + "duplex_num_input_tokens": 999, + "num_input_tokens": 999, + } + ) + == 16 + ) + + +def test_resource_state_rejects_fence_regression_and_requires_explicit_fence(): + current = DuplexFence("sid", epoch=2, turn_id=3, response_seq=4) + manager = DuplexSessionRuntimeManager() + session = manager.open_session( + current, + capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}), + ) + + for stale in ( + DuplexFence("sid", epoch=1, turn_id=99, response_seq=99), + DuplexFence("sid", epoch=2, turn_id=2, response_seq=4), + DuplexFence("sid", epoch=2, turn_id=3, response_seq=3), + ): + with pytest.raises(RuntimeError, match="fence mismatch"): + session.accept_fence(stale) + assert session.fence == current + + with pytest.raises(TypeError, match="fence"): + session.bind_stage_request(0, "request") + with pytest.raises(TypeError, match="fence"): + session.append_input(mode=DuplexInputMode.APPEND_AUDIO_CHUNK) + with pytest.raises(TypeError, match="DuplexFence"): + manager.open_session("legacy-session") + with pytest.raises(TypeError, match="DuplexFence"): + manager.close_session("legacy-session") + + +def test_resource_request_id_is_derived_from_fence_and_role(): + fence = DuplexFence("sid-with-dashes", epoch=7, turn_id=11, response_seq=13) + + assert duplex_resource_request_id(fence, "stage0") == "duplex-s.c2lkLXdpdGgtZGFzaGVz.i.0.e.7.r.stage0" + assert duplex_resource_request_id(fence, "stage1") == "duplex-s.c2lkLXdpdGgtZGFzaGVz.i.0.e.7.r.stage1" + + +def test_resource_request_id_codec_separates_session_id_from_incarnation(): + embedded_incarnation = duplex_resource_request_id( + DuplexFence("foo-i1", incarnation=0), + "stage0", + ) + actual_incarnation = duplex_resource_request_id( + DuplexFence("foo", incarnation=1), + "stage0", + ) + + assert embedded_incarnation != actual_incarnation + + +def test_resource_request_id_session_membership_uses_encoded_identity(): + request_id = duplex_resource_request_id( + DuplexFence("sid-with-dashes", incarnation=2, epoch=7), + "stage0", + ) + + assert duplex_resource_request_belongs_to_session(request_id, "sid-with-dashes") is True + assert duplex_resource_request_belongs_to_session(request_id, "sid") is False + assert duplex_resource_request_belongs_to_session("duplex-s.invalid.i.x.e.7.r.stage0", "sid") is False + + +def test_placeholder_budget_is_planned_inside_omni_engine_boundary(): + fence = DuplexFence("sid", turn_id=1, response_seq=1) + prompt = build_duplex_data_plane_prompt( + request_id=duplex_resource_request_id(fence, "stage0"), + fence=fence, + session_config={}, + runtime_config={}, + seq=2, + turn_seq=1, + mode=DuplexInputMode.APPEND_AUDIO_CHUNK, + payload={ + "audio": "AAAAAA==", + "format": "pcm_f32le", + "duplex_num_input_tokens": 999, + "num_input_tokens": 999, + }, + final=False, + ) + + assert len(prompt["prompt_token_ids"]) == 16 + assert prompt["model_intermediate_buffer"]["duplex"]["fence"] == fence + assert prompt["model_intermediate_buffer"]["duplex"]["scheduler_token_budget"] == 16 diff --git a/tests/fullduplex/minicpmo45/__init__.py b/tests/fullduplex/minicpmo45/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/fullduplex/minicpmo45/test_commit_policy.py b/tests/fullduplex/minicpmo45/test_commit_policy.py new file mode 100644 index 000000000..5ee66b41a --- /dev/null +++ b/tests/fullduplex/minicpmo45/test_commit_policy.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import pytest + +from vllm_omni.experimental.fullduplex.openai.commit_policy import ( + CommitAction, + CommitSnapshot, + decide_commit_action, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize( + ("case", "snapshot", "expected"), + [ + ( + "first idle commit", + CommitSnapshot( + auto_responds=True, + speech_since_commit=True, + active_response_id=None, + overlap_speech_ms=0, + native_response_in_progress=False, + playback_active=False, + ), + CommitAction.START_AUTO_RESPONSE, + ), + ( + "native auto-response remains model-owned during playback", + CommitSnapshot( + auto_responds=True, + speech_since_commit=True, + active_response_id=None, + overlap_speech_ms=800, + native_response_in_progress=True, + playback_active=True, + ), + CommitAction.START_AUTO_RESPONSE, + ), + ( + "native auto-response commit without speech", + CommitSnapshot( + auto_responds=True, + speech_since_commit=False, + active_response_id=None, + overlap_speech_ms=0, + native_response_in_progress=True, + playback_active=False, + ), + CommitAction.COMMIT_ONLY, + ), + ( + "non-auto response overlap", + CommitSnapshot( + auto_responds=False, + speech_since_commit=True, + active_response_id=None, + overlap_speech_ms=800, + native_response_in_progress=True, + playback_active=False, + ), + CommitAction.DEFER_ACTIVE_RESPONSE, + ), + ( + "non-auto playback overlap", + CommitSnapshot( + auto_responds=False, + speech_since_commit=True, + active_response_id=None, + overlap_speech_ms=800, + native_response_in_progress=False, + playback_active=True, + ), + CommitAction.DEFER_ACTIVE_RESPONSE, + ), + ( + "non-auto commit without overlap", + CommitSnapshot( + auto_responds=False, + speech_since_commit=True, + active_response_id=None, + overlap_speech_ms=0, + native_response_in_progress=False, + playback_active=False, + ), + CommitAction.COMMIT_ONLY, + ), + ], + ids=lambda value: value if isinstance(value, str) else None, +) +def test_decide_commit_action_covers_realtime_lifecycle_sequences( + case: str, + snapshot: CommitSnapshot, + expected: CommitAction, +) -> None: + del case + + assert decide_commit_action(snapshot) is expected diff --git a/tests/fullduplex/minicpmo45/test_input.py b/tests/fullduplex/minicpmo45/test_input.py new file mode 100644 index 000000000..035ab47ef --- /dev/null +++ b/tests/fullduplex/minicpmo45/test_input.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import base64 + +import numpy as np +import pytest + +from vllm_omni.experimental.fullduplex.minicpmo45.input import MiniCPMO45PcmAppendBuffer + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def pcm_payload(samples: int, *, speech: bool = True) -> dict[str, object]: + audio = np.ones(samples, dtype=np.float32).tobytes() + return { + "type": "audio", + "audio": base64.b64encode(audio).decode("ascii"), + "format": "pcm_f32le", + "sample_rate_hz": 16_000, + "is_speech": speech, + } + + +def test_commit_does_not_add_silence_after_incremental_audio_was_drained(): + buffer = MiniCPMO45PcmAppendBuffer() + + emitted = buffer.append(pcm_payload(16_000), chunk_period_ms=1_000) + committed = buffer.commit(chunk_period_ms=1_000) + + assert emitted is not None + assert not buffer.has_pending() + assert committed is None + + +def test_append_emits_one_model_unit_when_multiple_units_are_buffered(): + buffer = MiniCPMO45PcmAppendBuffer() + + emitted = buffer.append(pcm_payload(32_000), chunk_period_ms=1_000) + + assert emitted is not None + assert len(base64.b64decode(emitted["audio"])) == 16_000 * 4 + assert buffer.pending_byte_count == 16_000 * 4 + + +def test_commit_without_speech_does_not_synthesize_terminal_audio(): + buffer = MiniCPMO45PcmAppendBuffer() + buffer.append(pcm_payload(8_000, speech=False), chunk_period_ms=1_000) + + committed = buffer.commit(chunk_period_ms=1_000) + + assert committed is None + + +def test_commit_resets_cumulative_turn_accounting(): + buffer = MiniCPMO45PcmAppendBuffer() + buffer.append(pcm_payload(16_000), chunk_period_ms=1_000) + buffer.commit(chunk_period_ms=1_000) + + empty = buffer.commit(chunk_period_ms=1_000) + + assert empty is None + + +def test_pcm_append_reservation_rollback_restores_emitted_audio(): + buffer = MiniCPMO45PcmAppendBuffer() + original = pcm_payload(16_000) + + reservation = buffer.prepare_append( + original, + operation_id="append-1", + chunk_period_ms=1_000, + ) + + assert reservation is not None + assert not buffer.has_pending() + reservation.rollback() + assert buffer.has_pending() + + retried = buffer.flush(chunk_period_ms=1_000) + assert retried is not None + assert base64.b64decode(retried["audio"]) == base64.b64decode(original["audio"]) + + +def test_pcm_commit_keeps_prior_append_reservation_active(): + buffer = MiniCPMO45PcmAppendBuffer() + append_reservation = buffer.prepare_append( + pcm_payload(16_000), + operation_id="append-before-commit", + chunk_period_ms=1_000, + ) + + commit_reservation = buffer.prepare_commit( + operation_id="commit-after-append", + chunk_period_ms=1_000, + ) + + assert append_reservation is not None + assert append_reservation.active + assert commit_reservation.payload is None + append_reservation.commit() + commit_reservation.commit() + + +def test_pcm_commit_reservation_rollback_restores_residual_audio(): + buffer = MiniCPMO45PcmAppendBuffer() + original = pcm_payload(8_000) + assert ( + buffer.prepare_append( + original, + operation_id="buffer-half-chunk", + chunk_period_ms=1_000, + ) + is None + ) + + reservation = buffer.prepare_commit( + operation_id="final-half-chunk", + chunk_period_ms=1_000, + ) + + assert reservation.payload is not None + assert reservation.payload["final"] is True + reservation.rollback() + + retried = buffer.prepare_commit( + operation_id="retry-final-half-chunk", + chunk_period_ms=1_000, + ) + assert retried.payload is not None + assert base64.b64decode(retried.payload["audio"]) == (base64.b64decode(original["audio"]) + b"\x00" * (8_000 * 4)) diff --git a/tests/fullduplex/openai/test_duplex_audio.py b/tests/fullduplex/openai/test_duplex_audio.py new file mode 100644 index 000000000..535856976 --- /dev/null +++ b/tests/fullduplex/openai/test_duplex_audio.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import base64 + +import numpy as np +import pytest + +from vllm_omni.experimental.fullduplex.openai.audio import convert_input_audio_with_rate + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_pcm16_input_conversion_is_pure_and_reports_target_rate(): + pcm16 = np.array([-32768, 0, 32767], dtype=" None: + repo_root = Path(__file__).resolve().parents[3] + script = f""" +import sys + +import {module_name} + +loaded = sorted( + name + for name in sys.modules + if name == "vllm_omni.experimental.fullduplex.minicpmo45" + or name.startswith("vllm_omni.experimental.fullduplex.minicpmo45.") +) +if loaded: + raise SystemExit("generic OpenAI runtime loaded MiniCPM-o: " + ", ".join(loaded)) +""" + + result = subprocess.run( + [sys.executable, "-c", script], + cwd=repo_root, + capture_output=True, + text=True, + timeout=60, + check=False, + ) + + assert result.returncode == 0, result.stderr + + +def test_generic_handler_requires_explicit_serving_runtime_adapter() -> None: + repo_root = Path(__file__).resolve().parents[3] + script = """ +from types import SimpleNamespace + +from vllm_omni.experimental.fullduplex.openai.serving import OmniDuplexSessionHandler + +chat_service = SimpleNamespace(engine_client=SimpleNamespace()) +try: + OmniDuplexSessionHandler(chat_service=chat_service) +except ValueError as exc: + if "serving runtime adapter" not in str(exc).lower(): + raise +else: + raise SystemExit("generic handler silently selected a model serving adapter") +""" + + result = subprocess.run( + [sys.executable, "-c", script], + cwd=repo_root, + capture_output=True, + text=True, + timeout=60, + check=False, + ) + + assert result.returncode == 0, result.stderr + + +def _valid_runtime_adapter() -> SimpleNamespace: + data_plane = SimpleNamespace( + begin_request=lambda request_id: None, + is_terminal=lambda request_id: False, + mark_terminal=lambda request_id: None, + close_stream=lambda request_id: None, + close_session=lambda session_id, **kwargs: None, + project=lambda result, **kwargs: (), + ) + return SimpleNamespace( + adapter_id="test-adapter", + session_states={}, + data_plane=data_plane, + clean_response_done_prefix="clean", + interrupted_tts_prefix="interrupted", + private_runtime_config_keys=frozenset({"private_key"}), + create_session_state=lambda: object(), + session_state=lambda session_id: object(), + remove_session_state=lambda session_id: None, + is_enabled=lambda config: True, + capabilities=lambda **kwargs: object(), + validate_client_extra_body=lambda extra_body: None, + prepare_runtime_config=lambda config, **kwargs: {}, + runtime_config_for_update=lambda config, current: {}, + data_plane_context=lambda **kwargs: object(), + ) + + +@pytest.mark.parametrize( + ("attribute", "invalid_value"), + [ + ("session_states", None), + ("clean_response_done_prefix", None), + ("interrupted_tts_prefix", None), + ("private_runtime_config_keys", {"private_key"}), + ], +) +def test_runtime_adapter_validator_rejects_invalid_protocol_attributes( + attribute: str, + invalid_value: object, +) -> None: + adapter = _valid_runtime_adapter() + setattr(adapter, attribute, invalid_value) + + with pytest.raises(TypeError, match=attribute): + validate_serving_runtime_adapter(adapter) + + +def test_runtime_adapter_validator_rejects_incomplete_data_plane() -> None: + adapter = _valid_runtime_adapter() + adapter.data_plane.close_session = None + + with pytest.raises(TypeError, match="data_plane.*close_session"): + validate_serving_runtime_adapter(adapter) diff --git a/tests/fullduplex/openai/test_websocket_actor.py b/tests/fullduplex/openai/test_websocket_actor.py new file mode 100644 index 000000000..37924af52 --- /dev/null +++ b/tests/fullduplex/openai/test_websocket_actor.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio + +import pytest + +from vllm_omni.experimental.fullduplex.openai.websocket import ( + DuplexWebSocketActor, + normalize_duplex_input_event, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class FakeWebSocket: + def __init__(self) -> None: + self.sent: list[dict[str, object]] = [] + + async def send_json(self, payload: dict[str, object]) -> None: + self.sent.append(dict(payload)) + + +@pytest.mark.parametrize( + ("event", "expected"), + [ + ({"type": "signal_turn", "event": "barge_in"}, {"type": "turn.signal", "event": "barge_in"}), + ({"type": "close_session"}, {"type": "session.close"}), + ({"type": "audio.playback_ack", "played_ms": 1}, {"type": "playback.ack", "played_ms": 1}), + ({"type": "input_text.append", "text": "a"}, {"type": "input.text.append", "text": "a"}), + ({"type": "push_text", "text": "b"}, {"type": "input.text.append", "text": "b"}), + ( + {"type": "input.audio.append", "audio": "wav"}, + {"type": "input_audio_buffer.append", "audio": "wav", "format": "wav"}, + ), + ( + {"type": "push_chunk", "audio": "wav", "format": "pcm_f32le"}, + {"type": "input_audio_buffer.append", "audio": "wav", "format": "pcm_f32le"}, + ), + ], +) +def test_input_aliases_normalize_once_at_mailbox_boundary(event, expected): + assert normalize_duplex_input_event(event) == expected + + +def test_actor_uses_one_fifo_mailbox_for_inbound_events(): + actor = DuplexWebSocketActor(FakeWebSocket()) + + assert isinstance(actor.mailbox, asyncio.Queue) + assert not isinstance(actor.mailbox, asyncio.PriorityQueue) + assert not hasattr(actor, "control_queue") + assert not hasattr(actor, "input_queue") + assert not hasattr(actor, "event_queue") + assert not hasattr(actor, "lifecycle_state") + assert not hasattr(actor, "last_response_id") + assert not hasattr(actor, "overlap_speech_ms") + assert not hasattr(actor, "transition") + assert not hasattr(actor, "output_generation_in_flight") + assert not hasattr(actor, "session") + assert not hasattr(actor, "runtime_opened") + assert not hasattr(actor, "runtime_closed") + assert not hasattr(actor, "drain_input_queue") + assert not hasattr(actor, "control_events_seen") + assert not hasattr(actor, "input_events_seen") + assert not hasattr(actor, "cancel_count") + + +@pytest.mark.asyncio +async def test_control_event_does_not_overtake_earlier_audio_input(): + actor = DuplexWebSocketActor(FakeWebSocket()) + await actor.enqueue_event({"type": "input_audio_buffer.append", "audio": "pcm"}) + await actor.enqueue_event({"type": "response.cancel"}) + + assert await actor.next_event() == {"type": "input_audio_buffer.append", "audio": "pcm"} + assert await actor.next_event() == {"type": "response.cancel"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("event_type", ["input.cancel", "session.close", "close_session"]) +async def test_terminal_control_preserves_wire_order(event_type: str): + actor = DuplexWebSocketActor(FakeWebSocket()) + await actor.enqueue_event({"type": "input_audio_buffer.append", "audio": "pcm"}) + await actor.enqueue_event({"type": event_type}) + + assert await actor.next_event() == {"type": "input_audio_buffer.append", "audio": "pcm"} + assert await actor.next_event() == {"type": event_type} + + +@pytest.mark.asyncio +async def test_mailbox_delivers_every_enqueued_event_exactly_once(): + actor = DuplexWebSocketActor(FakeWebSocket()) + events = [ + {"type": "input_audio_buffer.append", "audio": str(index)} + if index % 2 + else {"type": "response.created", "response_id": f"resp-{index}"} + for index in range(20) + ] + + await asyncio.gather(*(actor.enqueue_event(dict(event)) for event in events)) + received = [await actor.next_event() for _ in events] + + assert sorted(received, key=repr) == sorted(events, key=repr) + + +@pytest.mark.asyncio +async def test_writer_is_single_owner_of_websocket_send(): + websocket = FakeWebSocket() + actor = DuplexWebSocketActor(websocket) + writer = asyncio.create_task(actor.writer_loop()) + + await actor.send_json({"type": "one"}) + await actor.send_json({"type": "two"}) + await actor.close_writer() + await writer + + assert websocket.sent == [{"type": "one"}, {"type": "two"}] + + +@pytest.mark.asyncio +async def test_stale_fence_payload_is_dropped_before_websocket_send(): + websocket = FakeWebSocket() + actor = DuplexWebSocketActor(websocket, current_epoch=lambda: 2) + writer = asyncio.create_task(actor.writer_loop()) + + await actor.send_json({"type": "response.audio.delta", "epoch": 1}) + await actor.send_json({"type": "response.audio.delta", "epoch": 2}) + await actor.close_writer() + await writer + + assert websocket.sent == [{"type": "response.audio.delta", "epoch": 2}] + assert actor.stale_output_dropped == 1 + + +@pytest.mark.asyncio +async def test_writer_does_not_revoke_accepted_terminal_after_close_starts(): + websocket = FakeWebSocket() + actor = DuplexWebSocketActor(websocket, current_epoch=lambda: 1) + await actor.send_json({"type": "response.done", "epoch": 1, "response_id": "resp-1"}) + actor.closing = True + await actor.close_writer() + + await actor.writer_loop() + + assert websocket.sent == [{"type": "response.done", "epoch": 1, "response_id": "resp-1"}] + assert actor.stale_output_dropped == 0 diff --git a/tests/fullduplex/test_client.py b/tests/fullduplex/test_client.py new file mode 100644 index 000000000..218050eb1 --- /dev/null +++ b/tests/fullduplex/test_client.py @@ -0,0 +1,240 @@ +import base64 +import wave +from urllib.parse import parse_qs, urlsplit + +import pytest + +from vllm_omni.experimental.fullduplex.client import ( + RealtimeDuplexClient, + RealtimeEventCollector, + build_realtime_url, + read_pcm16_wav, + write_pcm16_wav, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_realtime_client_builds_explicit_native_duplex_url(): + url = build_realtime_url( + "ws://localhost:8099/v1/realtime?custom=1", + "openbmb/MiniCPM-o-4_5", + session_id="session-a", + ) + + query = parse_qs(urlsplit(url).query) + assert query == { + "custom": ["1"], + "duplex": ["1"], + "model": ["openbmb/MiniCPM-o-4_5"], + "minicpmo45_native_duplex": ["1"], + "session_id": ["session-a"], + } + + +def test_realtime_client_builds_resume_only_url_when_autostart_disabled(): + url = build_realtime_url( + "ws://localhost:8099/v1/realtime?duplex=1", + "openbmb/MiniCPM-o-4_5", + autostart=False, + ) + + query = parse_qs(urlsplit(url).query) + assert query["autostart"] == ["0"] + assert query["minicpmo45_native_duplex"] == ["1"] + + +@pytest.mark.asyncio +async def test_realtime_client_configure_omits_ref_audio_by_default(): + class Client(RealtimeDuplexClient): + def __init__(self): + super().__init__("ws://unused") + self.sent = [] + + async def send(self, event): + self.sent.append(event) + self.events.add({"type": "session.created"}) + + client = Client() + + await client.configure("openbmb/MiniCPM-o-4_5", timeout_s=1) + + session = client.sent[0]["session"] + assert "ref_audio" not in session + + +@pytest.mark.asyncio +async def test_realtime_client_configure_sends_explicit_ref_audio(): + class Client(RealtimeDuplexClient): + def __init__(self): + super().__init__("ws://unused") + self.sent = [] + + async def send(self, event): + self.sent.append(event) + self.events.add({"type": "session.created"}) + + client = Client() + + await client.configure( + "openbmb/MiniCPM-o-4_5", + ref_audio="data:audio/wav;base64,AAAA", + timeout_s=1, + ) + + session = client.sent[0]["session"] + assert session["ref_audio"] == "data:audio/wav;base64,AAAA" + + +def test_realtime_event_collector_partitions_audio_by_response(): + collector = RealtimeEventCollector() + collector.add({"type": "response.created", "response": {"id": "resp-a"}}) + collector.add( + { + "type": "response.audio.delta", + "response_id": "resp-a", + "delta": base64.b64encode(b"audio-a").decode("ascii"), + "sample_rate_hz": 16_000, + } + ) + + assert collector.response_ids == ["resp-a"] + assert collector.audio_bytes("resp-a") == b"audio-a" + assert collector.output_sample_rate_hz == 16_000 + assert collector.first_received_at("response.created") is not None + assert collector.last_received_at("response.audio.delta") is not None + + +def test_realtime_event_collector_reports_engine_token_and_audio_intervals(): + collector = RealtimeEventCollector() + collector.add( + {"type": "response.created", "response": {"id": "resp-a"}}, + received_at_s=10.0, + ) + stage_metrics = { + "0": { + "num_tokens_out": 4, + "vllm_ttft_ms": 120.0, + "vllm_tpot_ms": 15.0, + "vllm_itl_ms": 14.0, + "vllm_itls_ms": [10.0, 14.0, 18.0], + } + } + for received_at_s, cumulative_audio_ms in ((10.2, 80), (10.25, 160), (10.36, 240)): + collector.add( + { + "type": "response.audio.delta", + "response_id": "resp-a", + "delta": base64.b64encode(b"audio").decode("ascii"), + "sample_rate_hz": 16_000, + "metadata": { + "audio_duration_ms": cumulative_audio_ms, + "vllm_omni": {"stage_metrics": stage_metrics}, + }, + }, + received_at_s=received_at_s, + ) + + timing = collector.timing_summary( + after_s=10.0, + input_committed_at_s=9.9, + response_id="resp-a", + ) + + assert timing["stage0_tokens"] == { + "source": "engine_stage_metrics", + "output_token_count": 4, + "ttft_ms": 120.0, + "tpot_ms": 15.0, + "inter_token_interval_ms": { + "count": 3, + "mean": 14.0, + "p50": 14.0, + "p95": 18.0, + "max": 18.0, + }, + } + assert timing["audio_output"] == { + "source": "client_monotonic_receive", + "chunk_count": 3, + "response_created_to_first_audio_ms": 200.0, + "commit_to_first_audio_ms": 300.0, + "inter_chunk_interval_ms": { + "count": 2, + "mean": 80.0, + "p50": 50.0, + "p95": 110.0, + "max": 110.0, + }, + "chunk_duration_ms": { + "count": 3, + "mean": 80.0, + "p50": 80.0, + "p95": 80.0, + "max": 80.0, + }, + "max_chunk_gap_ms": 110.0, + } + + +def test_response_timing_ignores_unowned_session_level_metrics(): + collector = RealtimeEventCollector() + collector.add( + {"type": "response.created", "response": {"id": "resp-a"}}, + received_at_s=10.0, + ) + collector.add( + { + "type": "response.audio.delta", + "response_id": "resp-a", + "delta": base64.b64encode(b"audio").decode("ascii"), + "metadata": { + "vllm_omni": { + "stage_metrics": { + "0": { + "num_tokens_out": 20, + "vllm_ttft_ms": 157.0, + "vllm_tpot_ms": 16.0, + "vllm_itls_ms": [15.0, 17.0], + } + } + } + }, + }, + received_at_s=10.2, + ) + collector.add( + { + "type": "response.listen", + "metadata": { + "vllm_omni": { + "stage_metrics": { + "0": { + "num_tokens_out": 2, + "vllm_ttft_ms": 106.0, + "vllm_tpot_ms": 0.0, + "vllm_itls_ms": [], + } + } + } + }, + }, + received_at_s=10.3, + ) + + timing = collector.timing_summary(after_s=10.0, response_id="resp-a") + + assert timing["stage0_tokens"]["output_token_count"] == 20 + assert timing["stage0_tokens"]["ttft_ms"] == 157.0 + + +def test_realtime_client_pcm16_wav_round_trip(tmp_path): + path = tmp_path / "audio.wav" + pcm16 = b"\x01\x00\x02\x00" + + write_pcm16_wav(path, pcm16, sample_rate_hz=16_000) + + with wave.open(str(path), "rb") as wav_file: + assert wav_file.getnchannels() == 1 + assert wav_file.getframerate() == 16_000 + assert read_pcm16_wav(path) == pcm16 diff --git a/tests/model_executor/models/minicpmo_4_5/test_hf_model_path.py b/tests/model_executor/models/minicpmo_4_5/test_hf_model_path.py index dc1599e6b..834b7b0e0 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_hf_model_path.py +++ b/tests/model_executor/models/minicpmo_4_5/test_hf_model_path.py @@ -5,7 +5,6 @@ from types import SimpleNamespace import pytest -import torch from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( MiniCPMO45OmniTTSForConditionalGeneration, @@ -14,10 +13,19 @@ from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +class _FakeMiniCPMTTS: + def __init__(self, config, audio_tokenizer): + self.config = config + self.audio_tokenizer = audio_tokenizer + self.emb_text = object() + self.projector_semantic = object() + + def generate(self): + raise NotImplementedError + + def test_hf_model_id_path_caching(mocker, tmp_path): - """Verify HF model-ID path is resolved via download_weights_from_hf_specific - once during _lazy_init_tts, cached in self._model_path, and reused in generate_speech. - """ + """Verify the resolved HF path is cached and reused for default ref audio.""" resolved_path = str(tmp_path / "resolved_hf_model") assets_dir = tmp_path / "resolved_hf_model" / "assets" assets_dir.mkdir(parents=True, exist_ok=True) @@ -29,11 +37,9 @@ def test_hf_model_id_path_caching(mocker, tmp_path): return_value=resolved_path, ) - # Mock dynamic class loading - fake_tts_class = mocker.MagicMock() mocker.patch( "transformers.dynamic_module_utils.get_class_from_dynamic_module", - return_value=fake_tts_class, + return_value=_FakeMiniCPMTTS, ) # Build fake vllm_config with HF model ID @@ -59,40 +65,10 @@ def test_hf_model_id_path_caching(mocker, tmp_path): tts_model = MiniCPMO45OmniTTSForConditionalGeneration(vllm_config=vllm_config) - # 1. Trigger lazy init + tts_model._lazy_init_tts() tts_model._lazy_init_tts() - # Verify download_weights_from_hf_specific was called with HF model ID and cached mock_download.assert_called_once_with(model_id, None, ["*"]) assert tts_model._model_path == resolved_path assert tts_model._assets_loaded is True - - # 2. Mock TTS object generate method and audio tokenizer for generate_speech - fake_tts_obj = mocker.MagicMock() - fake_tts_obj.emb_text.weight.device = torch.device("cpu") - fake_tts_obj.emb_text.side_effect = lambda t: torch.zeros(t.shape[0], 768) - fake_tts_obj.projector_semantic.side_effect = lambda h, **kw: torch.zeros(h.shape[0], 768) - fake_tts_obj.config = tts_cfg - fake_tts_obj.audio_bos_token_id = 151687 - - fake_output = SimpleNamespace(new_ids=torch.tensor([[[100]]])) - fake_tts_obj.generate.return_value = fake_output - - tts_model.tts_obj = fake_tts_obj - fake_tokenizer = mocker.MagicMock(return_value=b"dummy_wav_bytes") - tts_model.audio_tokenizer = fake_tokenizer - - # Mock sf.read to return fake numpy audio - mocker.patch( - "soundfile.read", - return_value=(torch.zeros(16000).numpy(), 24000), - ) - - # Call generate_speech - tts_token_ids = torch.tensor([1, 2, 3]) - tts_hidden_states = torch.zeros(3, 768) - waveform = tts_model.generate_speech(tts_token_ids, tts_hidden_states) - - # Verify download_weights_from_hf_specific was NOT called a second time - assert mock_download.call_count == 1 - assert waveform is not None + assert tts_model._resolve_prompt_wav_path(None, None) == (str(ref_audio_file), None) diff --git a/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py b/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py index d52bcafa1..7ec8e8114 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py +++ b/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py @@ -7,11 +7,9 @@ Covers ``vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.llm2t - empty ``source_outputs`` raises - latent fallback to ``hidden_states`` when ``multimodal_output`` is empty - both inputs missing -> raises - - additional_information payload carries the keys the talker expects - (prompt_embeds, prompt_token_ids, llm_output_token_ids, llm_output_text) - - dummy talker prompt ``[BOS, PAD, EOS] = [1, 0, 2]`` (single prefill step) + - structured model_intermediate_buffer carries the thinker ids and text + - scheduler prompt tokens follow the selected TTS region or output tokens - MiniCPM-o 4.5 TTS region detection on 151703 / 151704 tokens - - MiniCPM-o 2.6 fallback detection on 151691 / 151692 when no 4.5 markers - No TTS markers present -> no ``tts_token_ids`` / ``tts_hidden_states`` keys - prompt arg is normalized to a list and ``multi_modal_data`` is gated by ``requires_multimodal_data`` @@ -24,6 +22,9 @@ from types import SimpleNamespace import pytest import torch +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, +) from vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni import llm2tts pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -32,6 +33,55 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] _HIDDEN_DIM = 4 +def test_native_duplex_talker_emits_empty_text_metadata_to_clear_previous_segment() -> None: + class _Talker: + _ar_last_chunk_flags = [True] + _ar_turn_end_flags = [True] + + def __call__(self, **kwargs): + del kwargs + return None, torch.zeros(0, dtype=torch.float32) + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + torch.nn.Module.__init__(model) + model.model_stage = "tts" + model.config = SimpleNamespace(hidden_size=_HIDDEN_DIM) + model.talker = _Talker() + + output = model.forward( + input_ids=torch.tensor([1]), + positions=torch.tensor([0]), + runtime_additional_information=[ + { + "native_duplex": True, + "duplex": {"turn_id": 1, "epoch": 0}, + "meta": {"native_duplex_segment_text": ""}, + } + ], + ) + + assert output.multimodal_outputs is not None + assert output.multimodal_outputs["meta.llm_output_text_utf8"].numel() == 0 + assert output.multimodal_outputs["meta.audio_text_total_chars"].tolist() == [0] + + +def test_tts_embed_input_ids_uses_active_talker_embedding(monkeypatch) -> None: + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + torch.nn.Module.__init__(model) + model.model_stage = "tts" + expected = torch.ones((2, _HIDDEN_DIM)) + + monkeypatch.setattr( + model, + "get_input_embeddings", + lambda input_ids: expected, + ) + + actual = model.embed_input_ids(torch.tensor([11, 12])) + + assert actual is expected + + def _make_thinker_output( *, prompt_token_ids: list[int], @@ -40,6 +90,7 @@ def _make_thinker_output( request_id: str = "req-0", latent: torch.Tensor | None = None, hidden_states: torch.Tensor | None = None, + multimodal_output: dict[str, object] | None = None, ): """Construct a minimal mock of a thinker engine output entry. @@ -48,8 +99,11 @@ def _make_thinker_output( - per-output: ``multimodal_output`` (dict), ``hidden_states`` (opt), ``text``, ``token_ids`` """ + mm_output = dict(multimodal_output or {}) + if latent is not None: + mm_output["latent"] = latent output = SimpleNamespace( - multimodal_output={"latent": latent} if latent is not None else {}, + multimodal_output=mm_output, token_ids=output_token_ids, text=text, ) @@ -90,17 +144,16 @@ class TestBasicShape: ) assert len(out) == 2 - def test_talker_prompt_token_ids_dummy_bos_pad_eos(self) -> None: + def test_talker_scheduler_prompt_uses_output_tokens(self) -> None: hidden = torch.zeros((2, _HIDDEN_DIM)) out = llm2tts( [_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden)], prompt=None, ) - # The talker AR framework needs *some* tokens to do a single prefill; - # this is the agreed [BOS, PAD, EOS] minimal payload. - assert out[0]["prompt_token_ids"] == [1, 0, 2] + assert out[0]["prompt_token_ids"] == [20] + assert "stream_output" not in out[0]["model_intermediate_buffer"] - def test_additional_information_carries_thinker_outputs(self) -> None: + def test_model_intermediate_buffer_carries_thinker_outputs(self) -> None: prompt_ids = [10, 11, 12] out_ids = [20, 21] hidden = torch.randn(len(prompt_ids) + len(out_ids), _HIDDEN_DIM) @@ -116,21 +169,17 @@ class TestBasicShape: ], prompt=None, ) - ai = result[0]["additional_information"] - assert ai["prompt_token_ids"] == prompt_ids - assert ai["llm_output_token_ids"] == out_ids - assert ai["llm_output_text"] == ["hello"] - # prompt_embeds = float32 view of the prompt-portion of hidden states. - assert ai["prompt_embeds"].dtype == torch.float32 - assert ai["prompt_embeds"].shape == (len(prompt_ids), _HIDDEN_DIM) - assert torch.equal(ai["prompt_embeds"], hidden[: len(prompt_ids)].to(torch.float32)) + buffer = result[0]["model_intermediate_buffer"] + assert buffer["ids"]["prompt"] == prompt_ids + assert buffer["ids"]["output"] == out_ids + assert buffer["llm_output_text"] == ["hello"] def test_latent_in_multimodal_output_takes_precedence(self) -> None: # When both ``multimodal_output["latent"]`` and ``hidden_states`` are # present, the latent payload must win (this is the steady-state path # produced by the thinker stage). prompt_ids = [10, 11] - out_ids = [20] + out_ids = [151703, 30, 151704] latent = torch.ones((len(prompt_ids) + len(out_ids), _HIDDEN_DIM)) hidden = torch.zeros_like(latent) @@ -145,19 +194,17 @@ class TestBasicShape: ], prompt=None, ) - ai = result[0]["additional_information"] + buffer = result[0]["model_intermediate_buffer"] # latent (ones) won over hidden_states (zeros) - assert torch.equal(ai["prompt_embeds"], latent[: len(prompt_ids)].to(torch.float32)) + assert torch.equal(torch.tensor(buffer["hidden_states"]["tts"]), latent[3:4].to(torch.float32)) class TestTtsRegionDetection: - """The bridge auto-detects which MiniCPM-o generation produced the output - by sniffing for 4.5-specific BOS/EOS token ids (151703 / 151704); when - those are absent it falls back to the 2.6 ids (151691 / 151692). + """The bridge detects MiniCPM-o 4.5 BOS/EOS ids (151703 / 151704). This is a regression guard: a single off-by-one or wrong branch - precedence here causes the talker to receive an empty / wrong-region - slice and emit silent audio. + causes the talker to receive an empty / wrong-region slice and emit + silent audio. """ def _run(self, prompt_token_ids, output_token_ids): @@ -173,7 +220,7 @@ class TestTtsRegionDetection: ], prompt=None, ) - return result[0]["additional_information"], hidden + return result[0]["model_intermediate_buffer"], hidden def test_4_5_markers_detected(self) -> None: # prompt: [10, 11] @@ -181,45 +228,26 @@ class TestTtsRegionDetection: # full sequence: [10, 11, 151703, 30, 31, 151704, 40] # 0 1 2 3 4 5 6 # 4.5 BOS at idx 2 -> slice starts at 3; EOS at idx 5 -> slice ends at 5. - ai, hidden = self._run([10, 11], [151703, 30, 31, 151704, 40]) - assert "tts_token_ids" in ai and "tts_hidden_states" in ai - assert ai["tts_token_ids"].tolist() == [30, 31] - assert torch.equal(ai["tts_hidden_states"], hidden[3:5]) - - def test_4_5_takes_precedence_when_both_markers_present(self) -> None: - # If both 2.6 and 4.5 markers appear (shouldn't really happen on a - # real checkpoint, but exercises the loop's break), the 4.5 markers - # win because the source loop pins on the first 4.5 id seen. - # sequence index: 0 1 2 3 4 5 6 7 8 - ai, hidden = self._run([10], [151691, 70, 151692, 151703, 30, 151704, 40]) - assert ai["tts_token_ids"].tolist() == [30] - assert torch.equal(ai["tts_hidden_states"], hidden[5:6]) - - def test_2_6_fallback(self) -> None: - # No 4.5 markers anywhere -> use 2.6 BOS/EOS pair. - # full sequence: [10, 151691, 30, 31, 151692, 40] - # 0 1 2 3 4 5 - # BOS at idx 1 -> slice starts at 2; EOS at idx 4 -> slice ends at 4. - ai, hidden = self._run([10], [151691, 30, 31, 151692, 40]) - assert ai["tts_token_ids"].tolist() == [30, 31] - assert torch.equal(ai["tts_hidden_states"], hidden[2:4]) + buffer, hidden = self._run([10, 11], [151703, 30, 31, 151704, 40]) + assert buffer["ids"]["tts"] == [30, 31] + assert torch.equal(torch.tensor(buffer["hidden_states"]["tts"]), hidden[3:5]) def test_bos_without_eos_runs_to_end(self) -> None: # When BOS is found but EOS is missing (typical for an in-flight or # truncated decode), the slice should extend to the end of the # hidden-state matrix instead of being dropped. # sequence: [10, 11, 151703, 30, 31] - ai, hidden = self._run([10, 11], [151703, 30, 31]) - assert ai["tts_token_ids"].tolist() == [30, 31] - assert torch.equal(ai["tts_hidden_states"], hidden[3:5]) + buffer, hidden = self._run([10, 11], [151703, 30, 31]) + assert buffer["ids"]["tts"] == [30, 31] + assert torch.equal(torch.tensor(buffer["hidden_states"]["tts"]), hidden[3:5]) def test_no_tts_markers_omits_slice_keys(self) -> None: # If neither marker pair is present, the bridge should NOT populate # ``tts_token_ids`` / ``tts_hidden_states`` — the talker should fall # through to the dummy path. - ai, _ = self._run([10, 11], [20, 21, 22]) - assert "tts_token_ids" not in ai - assert "tts_hidden_states" not in ai + buffer, _ = self._run([10, 11], [20, 21, 22]) + assert "tts" not in buffer["ids"] + assert "tts" not in buffer.get("hidden_states", {}) class TestPromptAndMultiModal: @@ -251,14 +279,11 @@ class TestPromptAndMultiModal: ) assert out[0]["multi_modal_data"] == mm - def test_streaming_context_is_accepted_and_ignored(self) -> None: - # The framework passes ``streaming_context`` for streaming-capable - # pipelines; MiniCPM-o 4.5 is turn-taking so the bridge must accept - # but discard it without crashing. + def test_internal_streaming_context_is_accepted(self) -> None: hidden = torch.zeros((2, _HIDDEN_DIM)) out = llm2tts( [_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden)], prompt=None, - streaming_context=object(), + _streaming_context=object(), ) assert len(out) == 1 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 48db0d497..b9c7938f0 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -44,6 +44,20 @@ 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: + pipeline = OMNI_PIPELINES[_PIPELINE_KEY] + assert pipeline.duplex_control_enabled is True + assert pipeline.duplex_serving_adapter == ( + "vllm_omni.experimental.fullduplex.minicpmo45.serving_adapter.MiniCPMO45ServingRuntimeAdapter" + ) + assert not hasattr(pipeline, "max_native_duplex_sessions") + + def test_ordinary_pipeline_defaults_to_no_duplex_control(self) -> None: + pipeline = PipelineConfig(model_type="ordinary") + assert pipeline.duplex_control_enabled is False + assert pipeline.duplex_serving_adapter is None + assert not hasattr(pipeline, "max_native_duplex_sessions") + class TestPipelineTopology: @pytest.fixture(scope="class") diff --git a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py new file mode 100644 index 000000000..453b3a595 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py @@ -0,0 +1,1347 @@ +import base64 +import os +import struct +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + MiniCPMO45OmniTTSForConditionalGeneration, + _drain_native_duplex_emitted_text, + _queue_native_duplex_segment_text, + _TalkerTurnState, +) +from vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni import ( + _extract_first_audio_ref, + _native_duplex_segment_output_ids, + llm2tts, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_native_duplex_transcript_waits_for_buffered_vocoder_audio(): + state = _TalkerTurnState(None, None, turn_id=1) + + _queue_native_duplex_segment_text(state, "好") + assert _drain_native_duplex_emitted_text(state, has_audio=False) == "" + + _queue_native_duplex_segment_text(state, "的,刚") + assert _drain_native_duplex_emitted_text(state, has_audio=True) == "好的,刚" + assert _drain_native_duplex_emitted_text(state, has_audio=True) == "" + + +def test_native_duplex_token2wav_cache_isolated_across_interleaved_sessions(): + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + model.audio_tokenizer = SimpleNamespace(stream_cache=None, hift_cache_dict={}) + observed = [] + + def stream_window(token_list, _prompt, *, last_chunk): + del token_list, last_chunk + marker = model.audio_tokenizer.stream_cache["marker"] + observed.append(marker) + prefix = marker[0] + index = int(marker[1:]) + 1 + model.audio_tokenizer.stream_cache = {"marker": f"{prefix}{index}"} + model.audio_tokenizer.hift_cache_dict = {"marker": f"{prefix}{index}"} + return torch.ones(1) + + model._t2w_stream_window = stream_window + a = _TalkerTurnState(None, None) + b = _TalkerTurnState(None, None) + a.stream_cache = {"marker": "A0"} + a.hift_cache_dict = {"marker": "A0"} + a.vocoder_initialized = True + b.stream_cache = {"marker": "B0"} + b.hift_cache_dict = {"marker": "B0"} + b.vocoder_initialized = True + + model._run_vocoder_window(a, [1], last_chunk=False) + model._run_vocoder_window(b, [1], last_chunk=False) + model._run_vocoder_window(a, [1], last_chunk=False) + model._run_vocoder_window(b, [1], last_chunk=True) + + assert observed == ["A0", "B0", "A1", "B1"] + assert a.stream_cache == {"marker": "A2"} + assert b.stream_cache == {"marker": "B2"} + assert model.audio_tokenizer.stream_cache is None + assert model.audio_tokenizer.hift_cache_dict == {} + + +def test_extract_first_audio_ref_accepts_dict_stereo_audio(): + ref = _extract_first_audio_ref( + { + "audio": { + "array": [[1.0, 3.0, 5.0], [2.0, 4.0, 6.0]], + "sampling_rate": 16000, + } + } + ) + + assert ref is not None + waveform, sample_rate = ref + assert sample_rate == 16000 + assert torch.allclose(waveform, torch.tensor([1.5, 3.5, 5.5])) + + +def test_llm2tts_carries_request_ref_audio_to_talker_payload(): + latent = torch.arange(20, dtype=torch.float16).reshape(5, 4) + output = SimpleNamespace( + token_ids=[11, 12, 9002], + text="hello", + multimodal_output={ + "latent": latent, + "meta": { + "tts_bos_token_id": 9001, + "tts_eos_token_id": 9002, + }, + }, + ) + llm_output = SimpleNamespace( + request_id="req-1", + prompt_token_ids=[101, 9001], + outputs=[output], + ) + ref_waveform = torch.tensor([0.1, 0.2, 0.3]) + + converted = llm2tts( + [llm_output], + prompt=[ + { + "multi_modal_data": { + "audio": (ref_waveform, 22050), + } + } + ], + ) + + assert "additional_information" not in converted[0] + info = converted[0]["model_intermediate_buffer"] + assert info["codes"]["ref"] == ref_waveform.tolist() + assert info["meta"]["ref_audio_sr"] == 22050 + assert info["ids"]["tts"] == [11, 12] + assert torch.equal(torch.as_tensor(info["hidden_states"]["tts"]), latent[2:4].float()) + assert "tts_token_ids" not in info + assert "tts_hidden_states" not in info + + +def test_llm2tts_uses_duplex_prompt_token_ids_for_tts_boundary(): + latent = torch.arange(24, dtype=torch.float16).reshape(6, 4) + output = SimpleNamespace( + token_ids=[11, 12, 9002], + text="hello", + multimodal_output={ + "latent": latent, + "duplex_prompt_token_ids": [101, 102, 9001], + "meta": { + "tts_bos_token_id": 9001, + "tts_eos_token_id": 9002, + }, + }, + ) + llm_output = SimpleNamespace( + request_id="duplex-req", + prompt_token_ids=[0, 0, 0], + outputs=[output], + ) + + converted = llm2tts([llm_output], prompt=[{}]) + + assert "additional_information" not in converted[0] + info = converted[0]["model_intermediate_buffer"] + assert info["prompt_token_ids"] == [101, 102, 9001] + assert converted[0]["prompt_token_ids"] == [11, 12] + assert info["ids"]["tts"] == [11, 12] + assert torch.equal(torch.as_tensor(info["hidden_states"]["tts"]), latent[3:5].float()) + assert "tts_token_ids" not in info + assert "tts_hidden_states" not in info + + +def test_llm2tts_accepts_flat_tokenizer_derived_special_token_metadata(): + latent = torch.arange(24, dtype=torch.float16).reshape(6, 4) + output = SimpleNamespace( + token_ids=[11, 12, 9302], + text="hello", + multimodal_output={ + "latent": latent, + "duplex_prompt_token_ids": [101, 102, 9301], + "meta.tts_bos_token_id": torch.tensor([9301]), + "meta.tts_eos_token_id": torch.tensor([9302]), + "meta.listen_token_id": torch.tensor([9303]), + "meta.speak_token_id": torch.tensor([9304]), + "meta.chunk_eos_token_id": torch.tensor([9305]), + "meta.chunk_tts_eos_token_id": torch.tensor([9306]), + "meta.turn_eos_token_id": torch.tensor([9307]), + }, + ) + llm_output = SimpleNamespace( + request_id="duplex-flat-meta", + prompt_token_ids=[0, 0, 0], + outputs=[output], + ) + + converted = llm2tts([llm_output], prompt=[{}]) + + assert "additional_information" not in converted[0] + info = converted[0]["model_intermediate_buffer"] + assert converted[0]["prompt_token_ids"] == [11, 12] + assert info["ids"]["tts"] == [11, 12] + assert info["stream_output"] is True + assert info["native_duplex"] is True + assert torch.equal(torch.as_tensor(info["hidden_states"]["tts"]), latent[3:5].float()) + assert "tts_token_ids" not in info + assert "tts_hidden_states" not in info + + +def test_llm2tts_stops_at_tokenizer_derived_duplex_state_tokens(): + latent = torch.arange(32, dtype=torch.float16).reshape(8, 4) + output = SimpleNamespace( + token_ids=[11, 12, 9303, 13, 9302], + text="hello", + multimodal_output={ + "latent": latent, + "duplex_prompt_token_ids": [101, 102, 9301], + "meta": { + "tts_bos_token_id": 9301, + "tts_eos_token_id": 9302, + "listen_token_id": 9303, + "speak_token_id": 9304, + "tts_pad_token_id": 9305, + "unit_token_id": 9306, + "unit_end_token_id": 9307, + "chunk_eos_token_id": 9308, + "chunk_tts_eos_token_id": 9309, + "turn_eos_token_id": 9310, + }, + }, + ) + llm_output = SimpleNamespace( + request_id="duplex-state-boundary", + prompt_token_ids=[0, 0, 0], + outputs=[output], + ) + + converted = llm2tts([llm_output], prompt=[{}]) + + info = converted[0]["model_intermediate_buffer"] + assert converted[0]["prompt_token_ids"] == [11, 12] + assert info["ids"]["tts"] == [11, 12] + assert torch.equal(torch.as_tensor(info["hidden_states"]["tts"]), latent[3:5].float()) + assert "tts_token_ids" not in info + assert "tts_hidden_states" not in info + + +def test_llm2tts_native_duplex_uses_speak_region_without_tts_bos_prompt(): + latent = torch.arange(36, dtype=torch.float16).reshape(9, 4) + output = SimpleNamespace( + token_ids=[9304, 11, 12, 9308, 13], + text="hello", + multimodal_output={ + "latent": latent, + "duplex_prompt_token_ids": [101, 102, 9306, 9306], + "meta": { + "tts_bos_token_id": 9301, + "tts_eos_token_id": 9302, + "listen_token_id": 9303, + "speak_token_id": 9304, + "tts_pad_token_id": 9305, + "unit_token_id": 9306, + "unit_end_token_id": 9307, + "chunk_eos_token_id": 9308, + "chunk_tts_eos_token_id": 9309, + "turn_eos_token_id": 9310, + }, + }, + ) + llm_output = SimpleNamespace( + request_id="duplex-speak-region", + prompt_token_ids=[0, 0, 0, 0], + outputs=[output], + ) + + converted = llm2tts([llm_output], prompt=[{}]) + + info = converted[0]["model_intermediate_buffer"] + assert converted[0]["prompt_token_ids"] == [11, 12] + assert info["ids"]["tts"] == [11, 12] + assert torch.equal(torch.as_tensor(info["hidden_states"]["tts"]), latent[5:7].float()) + assert "tts_token_ids" not in info + assert "tts_hidden_states" not in info + + +def test_llm2tts_native_duplex_uses_plain_text_region_without_speak_marker(): + latent = torch.arange(32, dtype=torch.float16).reshape(8, 4) + output = SimpleNamespace( + token_ids=[11, 12, 13, 9308], + text="hello", + multimodal_output={ + "latent": latent, + "duplex_prompt_token_ids": [101, 102, 9306, 9306], + "meta": { + "tts_bos_token_id": 9301, + "tts_eos_token_id": 9302, + "listen_token_id": 9303, + "speak_token_id": 9304, + "tts_pad_token_id": 9305, + "unit_token_id": 9306, + "unit_end_token_id": 9307, + "chunk_eos_token_id": 9308, + "chunk_tts_eos_token_id": 9309, + "turn_eos_token_id": 9310, + }, + }, + ) + llm_output = SimpleNamespace( + request_id="duplex-plain-text-region", + prompt_token_ids=[0, 0, 0, 0], + outputs=[output], + ) + + converted = llm2tts([llm_output], prompt=[{}]) + + info = converted[0]["model_intermediate_buffer"] + assert converted[0]["prompt_token_ids"] == [12, 13] + assert info["ids"]["tts"] == [12, 13] + assert torch.equal(torch.as_tensor(info["hidden_states"]["tts"]), latent[5:7].float()) + + +def test_llm2tts_native_duplex_conditions_on_turn_eos_and_midunit_speak(): + latent = torch.arange(36, dtype=torch.float16).reshape(9, 4) + # Official duplex includes mid-unit <|speak|> tokens AND the <|turn_eos|> + # token+hidden in the talker condition (its embedding is the trained + # stop signal); only chunk terminators bound the slice. + output = SimpleNamespace( + token_ids=[9304, 11, 9304, 12, 9310, 9308], + text="hello", + multimodal_output={ + "latent": latent, + "duplex_prompt_token_ids": [101, 102, 9306], + "meta": { + "tts_bos_token_id": 9301, + "tts_eos_token_id": 9302, + "listen_token_id": 9303, + "speak_token_id": 9304, + "tts_pad_token_id": 9305, + "unit_token_id": 9306, + "unit_end_token_id": 9307, + "chunk_eos_token_id": 9308, + "chunk_tts_eos_token_id": 9309, + "turn_eos_token_id": 9310, + }, + }, + ) + llm_output = SimpleNamespace( + request_id="duplex-turn-eos", + prompt_token_ids=[0, 0, 0], + outputs=[output], + ) + + converted = llm2tts([llm_output], prompt=[{}]) + + info = converted[0]["model_intermediate_buffer"] + assert info["ids"]["tts"] == [11, 9304, 12, 9310] + assert torch.equal(torch.as_tensor(info["hidden_states"]["tts"]), latent[4:8].float()) + assert info["meta"]["turn_eos_token_id"] == 9310 + + +def test_llm2tts_native_duplex_ignores_stale_tts_bos_inside_folded_prompt(): + latent = torch.arange(36, dtype=torch.float16).reshape(9, 4) + # A continuation unit: the resumable prompt has an EARLIER reply's + # <|tts_bos|> folded mid-prompt, and the model decided to LISTEN (no new + # speak text). The stale boundary must not re-slice already-spoken text. + output = SimpleNamespace( + token_ids=[9303, 9309], + text="", + multimodal_output={ + "latent": latent, + "duplex_prompt_token_ids": [101, 9301, 41, 42, 9306, 9306], + "meta": { + "tts_bos_token_id": 9301, + "tts_eos_token_id": 9302, + "listen_token_id": 9303, + "speak_token_id": 9304, + "tts_pad_token_id": 9305, + "unit_token_id": 9306, + "unit_end_token_id": 9307, + "chunk_eos_token_id": 9308, + "chunk_tts_eos_token_id": 9309, + "turn_eos_token_id": 9310, + }, + }, + ) + llm_output = SimpleNamespace( + request_id="duplex-stale-bos", + prompt_token_ids=[0, 0, 0, 0, 0, 0], + outputs=[output], + ) + + converted = llm2tts([llm_output], prompt=[{}]) + + assert converted == [] + + +def test_llm2tts_rejects_native_duplex_output_without_special_token_metadata(): + latent = torch.arange(16, dtype=torch.float16).reshape(4, 4) + output = SimpleNamespace( + token_ids=[11, 12], + text="hello", + multimodal_output={ + "latent": latent, + "duplex_prompt_token_ids": [101, 102, 9301], + }, + ) + llm_output = SimpleNamespace( + request_id="duplex-missing-meta", + prompt_token_ids=[0, 0, 0], + outputs=[output], + ) + + with pytest.raises(ValueError, match=r"<\|tts_bos\|>"): + llm2tts([llm_output], prompt=[{}]) + + +def test_minicpmo_talker_normalizes_list_handoff_payload(monkeypatch): + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + captured = {} + + monkeypatch.setattr(model, "_should_stream_output", lambda *args, **kwargs: False) + + def fake_generate_speech(tts_token_ids, tts_hidden_states, **kwargs): + captured["tts_token_ids"] = tts_token_ids + captured["tts_hidden_states"] = tts_hidden_states + return torch.tensor([0.0, 0.1], dtype=torch.float32) + + monkeypatch.setattr(model, "generate_speech", fake_generate_speech) + + waveform, mel = model.forward( + additional_information={ + "tts_token_ids": [11, 12], + "tts_hidden_states": [[0.1, 0.2], [0.3, 0.4]], + } + ) + + assert mel is None + assert torch.equal(waveform, torch.tensor([0.0, 0.1], dtype=torch.float32)) + assert torch.equal(captured["tts_token_ids"], torch.tensor([11, 12])) + assert captured["tts_hidden_states"].dtype == torch.float32 + assert captured["tts_hidden_states"].shape == (2, 2) + + +def test_minicpmo_talker_accepts_structured_omni_payload_handoff(monkeypatch): + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + captured = {} + + monkeypatch.setattr(model, "_should_stream_output", lambda *args, **kwargs: False) + + def fake_generate_speech(tts_token_ids, tts_hidden_states, **kwargs): + captured["tts_token_ids"] = tts_token_ids + captured["tts_hidden_states"] = tts_hidden_states + return torch.tensor([0.0, 0.1], dtype=torch.float32) + + monkeypatch.setattr(model, "generate_speech", fake_generate_speech) + + waveform, mel = model.forward( + additional_information={ + "ids": {"tts": [11, 12]}, + "hidden_states": {"tts": torch.tensor([[0.1, 0.2], [0.3, 0.4]])}, + } + ) + + assert mel is None + assert torch.equal(waveform, torch.tensor([0.0, 0.1], dtype=torch.float32)) + assert torch.equal(captured["tts_token_ids"], torch.tensor([11, 12])) + assert captured["tts_hidden_states"].dtype == torch.float32 + assert captured["tts_hidden_states"].shape == (2, 2) + + +def test_token2wav_prompt_cache_resets_when_reference_changes(): + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + model.audio_tokenizer = SimpleNamespace(cache="old") + + model._reset_token2wav_cache_if_needed("/tmp/ref_a.wav") + assert model.audio_tokenizer.cache is None + + model.audio_tokenizer.cache = "prepared" + model._reset_token2wav_cache_if_needed("/tmp/ref_a.wav") + assert model.audio_tokenizer.cache == "prepared" + + model._reset_token2wav_cache_if_needed("/tmp/ref_b.wav") + assert model.audio_tokenizer.cache is None + + +def test_request_ref_audio_prompt_wav_uses_content_cache(): + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + ref_audio = torch.tensor([0.0, 0.1, 0.2], dtype=torch.float32) + + path1 = model._write_ref_audio_prompt_wav(ref_audio, 16000) + path2 = model._write_ref_audio_prompt_wav(ref_audio.clone(), 16000) + + try: + assert path1 == path2 + assert os.path.exists(path1) + assert model._is_cached_ref_audio_prompt_wav(path1) + finally: + for path in set(getattr(model, "_ref_audio_prompt_files", {}).values()): + try: + os.unlink(path) + except OSError: + pass + + +def test_native_duplex_talker_uses_generate_chunk_continuation(monkeypatch): + class _FakeTTS: + audio_bos_token_id = 88 + + def __init__(self): + self.config = SimpleNamespace(num_audio_tokens=128) + self.model = SimpleNamespace(config=SimpleNamespace()) + self.emb_text = torch.nn.Embedding(256, 4) + self.calls = [] + self._responses = [ + (torch.tensor([[[10], [11], [12]]]), "kv1"), + (torch.tensor([[[20]]]), "kv2"), + ] + + def generate_chunk(self, **kwargs): + self.calls.append(kwargs) + return self._responses[len(self.calls) - 1] + + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + tts = _FakeTTS() + windows = [] + model.tts_obj = tts + model.audio_tokenizer = SimpleNamespace() + model._talker_turn_states = {} + model._talker_consumed_tokens = {} + + monkeypatch.setattr(model, "_lazy_init_tts", lambda: None) + monkeypatch.setattr( + model, "_build_tts_sampling_params", lambda: SimpleNamespace(temperature=0.8, repetition_penalty=1.05) + ) + monkeypatch.setattr(model, "_tts_runtime_config", lambda: SimpleNamespace(streaming_generator_chunk=2)) + monkeypatch.setattr(model, "_resolve_prompt_wav_path", lambda ref, sr: ("/tmp/ref.wav", None)) + monkeypatch.setattr(model, "_begin_turn_vocoder_cache", lambda prompt, **kwargs: None) + monkeypatch.setattr(model, "_t2w_pre_lookahead", lambda: 1) + monkeypatch.setattr( + model, + "_build_tts_condition_embeds", + lambda ids, hidden: torch.ones((ids.numel(), 4), dtype=torch.float32), + ) + + def fake_stream_window(token_list, prompt_wav_path, *, last_chunk): + windows.append((list(token_list), last_chunk)) + return torch.tensor([float(len(token_list))]) + + monkeypatch.setattr(model, "_t2w_stream_window", fake_stream_window) + + base_info = { + "request_id": "duplex-generate-chunk", + "meta": {"turn_eos_token_id": 9310}, + "ids": {"tts": [21]}, + "hidden_states": {"tts": [[0.1, 0.2, 0.3, 0.4]]}, + "codes": {"ref": None}, + "native_duplex": True, + } + list(model._create_native_duplex_stream_gen(base_info)) + + turn_end_info = { + **base_info, + "ids": {"tts": [21, 9310]}, + "hidden_states": {"tts": [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8]]}, + } + list(model._create_native_duplex_stream_gen(turn_end_info)) + + assert len(tts.calls) == 2 + first, second = tts.calls + assert first["past_key_values"] is None + assert first["text_start_pos"] == 0 + assert first["min_new_tokens"] == 0 + assert first["max_new_token"] == 3 + assert torch.equal(first["eos_token"], torch.tensor([127])) + assert float(first["temperature"].item()) == pytest.approx(0.8) + assert first["repetition_penalty"] == pytest.approx(1.05) + + assert second["past_key_values"] == "kv1" + assert second["text_start_pos"] == 5 + assert second["min_new_tokens"] == 0 + assert second["max_new_token"] == 3 + assert "duplex-generate-chunk" not in model._talker_turn_states + assert windows[-1] == ([20], True) + + +def test_minicpmo_native_duplex_talker_turn_end_metadata_flushes_tail(monkeypatch): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + MiniCPMO45OmniTTSForConditionalGeneration, + ) + + class _FakeEmbText: + weight = torch.ones((200, 4), dtype=torch.float32) + + def __call__(self, ids): + return torch.ones((ids.numel(), 4), dtype=torch.float32) + + class _FakeTTS: + audio_bos_token_id = 1 + config = SimpleNamespace(num_audio_tokens=128) + emb_text = _FakeEmbText() + model = SimpleNamespace(config=SimpleNamespace(rope_theta=10000.0)) + + def generate_chunk(self, **kwargs): + del kwargs + return torch.tensor([[[33]]]), None + + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + model.tts_obj = _FakeTTS() + model.audio_tokenizer = SimpleNamespace() + model._talker_turn_states = {} + model._talker_consumed_tokens = {} + + monkeypatch.setattr(model, "_lazy_init_tts", lambda: None) + monkeypatch.setattr( + model, "_build_tts_sampling_params", lambda: SimpleNamespace(temperature=0.8, repetition_penalty=1.05) + ) + monkeypatch.setattr(model, "_tts_runtime_config", lambda: SimpleNamespace(streaming_generator_chunk=25)) + monkeypatch.setattr(model, "_resolve_prompt_wav_path", lambda ref, sr: ("/tmp/ref.wav", None)) + monkeypatch.setattr(model, "_begin_turn_vocoder_cache", lambda prompt, **kwargs: None) + monkeypatch.setattr(model, "_t2w_pre_lookahead", lambda: 5) + monkeypatch.setattr( + model, + "_build_tts_condition_embeds", + lambda ids, hidden: torch.ones((ids.numel(), 4), dtype=torch.float32), + ) + + windows = [] + + def fake_stream_window(token_list, prompt_wav_path, *, last_chunk): + del prompt_wav_path + windows.append((list(token_list), last_chunk)) + return torch.tensor([float(len(token_list))]) + + monkeypatch.setattr(model, "_t2w_stream_window", fake_stream_window) + + info = { + "request_id": "duplex-metadata-tail-flush", + "meta": {"turn_eos_token_id": 9310}, + "ids": {"tts": [21]}, + "hidden_states": {"tts": [[0.1, 0.2, 0.3, 0.4]]}, + "codes": {"ref": None}, + "end_of_turn": True, + "native_duplex": True, + } + + list(model._create_native_duplex_stream_gen(info)) + + assert windows[-1][1] is True + assert "duplex-metadata-tail-flush" not in model._talker_turn_states + assert "duplex-metadata-tail-flush" not in model._talker_consumed_tokens + + +def test_minicpmo_native_duplex_terminal_only_new_turn_does_not_start_tts(monkeypatch): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + MiniCPMO45OmniTTSForConditionalGeneration, + ) + + class _FakeEmbText: + weight = torch.ones((200, 4), dtype=torch.float32) + + def __call__(self, ids): + return torch.ones((ids.numel(), 4), dtype=torch.float32) + + class _FakeTTS: + audio_bos_token_id = 1 + config = SimpleNamespace(num_audio_tokens=128) + emb_text = _FakeEmbText() + model = SimpleNamespace(config=SimpleNamespace(rope_theta=10000.0)) + + def __init__(self): + self.calls = 0 + + def generate_chunk(self, **kwargs): + del kwargs + self.calls += 1 + return torch.tensor([[[33]]]), None + + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + tts = _FakeTTS() + model.tts_obj = tts + model.audio_tokenizer = SimpleNamespace() + model._talker_turn_states = {} + model._talker_consumed_tokens = {} + + monkeypatch.setattr(model, "_lazy_init_tts", lambda: None) + monkeypatch.setattr( + model, "_build_tts_sampling_params", lambda: SimpleNamespace(temperature=0.8, repetition_penalty=1.05) + ) + monkeypatch.setattr(model, "_tts_runtime_config", lambda: SimpleNamespace(streaming_generator_chunk=25)) + monkeypatch.setattr(model, "_resolve_prompt_wav_path", lambda ref, sr: ("/tmp/ref.wav", None)) + monkeypatch.setattr(model, "_begin_turn_vocoder_cache", lambda prompt, **kwargs: None) + monkeypatch.setattr(model, "_t2w_pre_lookahead", lambda: 5) + monkeypatch.setattr( + model, + "_build_tts_condition_embeds", + lambda ids, hidden: torch.ones((ids.numel(), 4), dtype=torch.float32), + ) + monkeypatch.setattr( + model, + "_t2w_stream_window", + lambda token_list, prompt_wav_path, *, last_chunk: torch.tensor([float(len(token_list))]), + ) + + outputs = list( + model._create_native_duplex_stream_gen( + { + "request_id": "duplex-terminal-only", + "duplex": {"epoch": 0, "turn_id": 1}, + "meta": {"turn_eos_token_id": 9310}, + "ids": {"tts": [9310]}, + "hidden_states": {"tts": [[0.1, 0.2, 0.3, 0.4]]}, + "codes": {"ref": None}, + "native_duplex": True, + } + ) + ) + + assert tts.calls == 0 + assert len(outputs) == 1 + assert outputs[0][0].numel() == 0 + assert outputs[0][1] is True + assert model._talker_turn_states == {} + assert model._talker_consumed_tokens == {} + + +def test_minicpmo_native_duplex_talker_finished_request_closes_session_keyed_turn_state(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + MiniCPMO45OmniTTSForConditionalGeneration, + _TalkerTurnState, + ) + + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + model.audio_tokenizer = SimpleNamespace(stream_cache={"stale": True}, hift_cache_dict={"stale": True}) + model._stream_gens = {} + model._talker_turn_states = { + "sid-stage1": _TalkerTurnState(prompt_wav_path=None, temp_prompt_wav_path=None), + } + model._talker_consumed_tokens = {"sid-stage1": 12} + model._talker_request_keys = {"req-stage1": "sid-stage1"} + + model.on_requests_finished({"req-stage1"}) + + assert "sid-stage1" not in model._talker_turn_states + assert "sid-stage1" not in model._talker_consumed_tokens + assert "req-stage1" not in model._talker_request_keys + assert model.audio_tokenizer.stream_cache is None + assert model.audio_tokenizer.hift_cache_dict == {} + + +def test_minicpmo_native_duplex_talker_new_turn_reopens_session_keyed_state(monkeypatch): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + MiniCPMO45OmniTTSForConditionalGeneration, + _TalkerTurnState, + ) + + class _FakeTTS: + audio_bos_token_id = 88 + + def __init__(self): + self.config = SimpleNamespace(num_audio_tokens=128) + self.model = SimpleNamespace(config=SimpleNamespace()) + self.emb_text = torch.nn.Embedding(256, 4) + self.calls = [] + + def generate_chunk(self, **kwargs): + self.calls.append(kwargs) + return torch.tensor([[[10]]]), "new-kv" + + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + tts = _FakeTTS() + model.tts_obj = tts + model.audio_tokenizer = SimpleNamespace(stream_cache={"old": True}, hift_cache_dict={"old": True}) + model._talker_turn_states = { + "sid-stage1": _TalkerTurnState( + prompt_wav_path=None, + temp_prompt_wav_path=None, + epoch=0, + turn_id=1, + ), + } + model._talker_turn_states["sid-stage1"].past_key_values = "old-kv" + model._talker_consumed_tokens = {"sid-stage1": 9} + model._talker_request_keys = {} + + monkeypatch.setattr(model, "_lazy_init_tts", lambda: None) + monkeypatch.setattr( + model, "_build_tts_sampling_params", lambda: SimpleNamespace(temperature=0.8, repetition_penalty=1.05) + ) + monkeypatch.setattr(model, "_tts_runtime_config", lambda: SimpleNamespace(streaming_generator_chunk=25)) + monkeypatch.setattr(model, "_resolve_prompt_wav_path", lambda ref, sr: ("/tmp/ref.wav", None)) + monkeypatch.setattr(model, "_begin_turn_vocoder_cache", lambda prompt, **kwargs: None) + monkeypatch.setattr(model, "_t2w_pre_lookahead", lambda: 5) + monkeypatch.setattr( + model, + "_build_tts_condition_embeds", + lambda ids, hidden: torch.ones((ids.numel(), 4), dtype=torch.float32), + ) + monkeypatch.setattr( + model, + "_t2w_stream_window", + lambda token_list, prompt_wav_path, *, last_chunk: torch.tensor([float(len(token_list))]), + ) + + info = { + "request_id": "sid-stage1", + "duplex": {"epoch": 0, "turn_id": 2}, + "meta": {"turn_eos_token_id": 9310}, + "ids": {"tts": [21]}, + "hidden_states": {"tts": [[0.1, 0.2, 0.3, 0.4]]}, + "codes": {"ref": None}, + "native_duplex": True, + } + + list(model._create_native_duplex_stream_gen(info)) + + assert tts.calls[0]["past_key_values"] is None + assert tts.calls[0]["text_start_pos"] == 0 + assert model._talker_consumed_tokens["sid-stage1"] == 1 + assert model._talker_turn_states["sid-stage1"].turn_id == 2 + + +def test_minicpmo_native_duplex_talker_segment_end_preserves_token2wav_stream_until_turn_end(monkeypatch): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + MiniCPMO45OmniTTSForConditionalGeneration, + ) + + class _FakeEmbText: + weight = torch.ones((200, 4), dtype=torch.float32) + + def __call__(self, ids): + return torch.ones((ids.numel(), 4), dtype=torch.float32) + + class _FakeTTS: + audio_bos_token_id = 1 + config = SimpleNamespace(num_audio_tokens=128) + emb_text = _FakeEmbText() + model = SimpleNamespace(config=SimpleNamespace(rope_theta=10000.0)) + + def generate_chunk(self, **kwargs): + del kwargs + return torch.tensor([[[33]]]), "kv" + + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + model.tts_obj = _FakeTTS() + model.audio_tokenizer = SimpleNamespace() + model._talker_turn_states = {} + model._talker_consumed_tokens = {} + + monkeypatch.setattr(model, "_lazy_init_tts", lambda: None) + monkeypatch.setattr( + model, "_build_tts_sampling_params", lambda: SimpleNamespace(temperature=0.8, repetition_penalty=1.05) + ) + monkeypatch.setattr(model, "_tts_runtime_config", lambda: SimpleNamespace(streaming_generator_chunk=25)) + monkeypatch.setattr(model, "_resolve_prompt_wav_path", lambda ref, sr: ("/tmp/ref.wav", None)) + monkeypatch.setattr(model, "_begin_turn_vocoder_cache", lambda prompt, **kwargs: None) + monkeypatch.setattr(model, "_t2w_pre_lookahead", lambda: 5) + monkeypatch.setattr( + model, + "_build_tts_condition_embeds", + lambda ids, hidden: torch.ones((ids.numel(), 4), dtype=torch.float32), + ) + + windows = [] + + def fake_stream_window(token_list, prompt_wav_path, *, last_chunk): + del prompt_wav_path + windows.append((list(token_list), last_chunk)) + return torch.tensor([float(len(token_list))]) + + monkeypatch.setattr(model, "_t2w_stream_window", fake_stream_window) + + info = { + "request_id": "duplex-segment-tail-flush", + "meta": {"turn_eos_token_id": 9310, "segment_end": True}, + "ids": {"tts": [21]}, + "hidden_states": {"tts": [[0.1, 0.2, 0.3, 0.4]]}, + "codes": {"ref": None}, + "native_duplex": True, + } + + list(model._create_native_duplex_stream_gen(info)) + + state = model._talker_turn_states["duplex-segment-tail-flush"] + assert windows == [] + assert state.token2wav_buffer == [4218, 4218, 4218, 33] + + turn_end_info = { + **info, + "ids": {"tts": [21, 9310]}, + "hidden_states": { + "tts": [ + [0.1, 0.2, 0.3, 0.4], + [0.5, 0.6, 0.7, 0.8], + ] + }, + "meta": {"turn_eos_token_id": 9310, "turn_end": True}, + } + list(model._create_native_duplex_stream_gen(turn_end_info)) + + assert windows == [([4218, 4218, 4218, 33, 33], True)] + assert "duplex-segment-tail-flush" not in model._talker_turn_states + + +def test_tts_scheduler_eos_uses_tokenizer_im_end_when_config_has_no_eos(): + class _Tokenizer: + unk_token_id = 0 + + def convert_tokens_to_ids(self, token): + return {"<|im_end|>": 77}.get(token, self.unk_token_id) + + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return [77] if text == "<|im_end|>" else [self.unk_token_id] + + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + model.config = SimpleNamespace(eos_token_id=None, vocab_size=100) + model.vllm_config = SimpleNamespace(model_config=SimpleNamespace()) + model._text_tokenizer = _Tokenizer() + model._ar_last_chunk_flags = [True] + + logits = model.compute_logits(torch.zeros(1, 4)) + + assert int(torch.argmax(logits[0]).item()) == 77 + + +def test_tts_streaming_non_final_chunk_does_not_mark_extra_rows_eos(): + class _Tokenizer: + unk_token_id = 0 + + def convert_tokens_to_ids(self, token): + return {"<|im_end|>": 77}.get(token, self.unk_token_id) + + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return [77] if text == "<|im_end|>" else [self.unk_token_id] + + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + model.config = SimpleNamespace(eos_token_id=None, vocab_size=100) + model.vllm_config = SimpleNamespace(model_config=SimpleNamespace()) + model._text_tokenizer = _Tokenizer() + model._ar_last_chunk_flags = [False] + + logits = model.compute_logits(torch.zeros(3, 4)) + sampled = torch.argmax(logits, dim=-1).tolist() + + assert sampled == [1, 1, 1] + + +def test_native_duplex_turn_end_detection_is_not_segment_end_detection(): + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + turn_eos_id = 9310 + + assert model._native_duplex_input_ends_turn( + { + "native_duplex": True, + "ids": {"tts": [9304, 42, turn_eos_id]}, + "meta": {"turn_eos_token_id": turn_eos_id}, + } + ) + assert not model._native_duplex_input_ends_turn( + { + "native_duplex": True, + "ids": {"tts": [9304, 42, 9308]}, + "meta": { + "turn_eos_token_id": turn_eos_id, + "segment_end": True, + }, + } + ) + + +def _native_duplex_meta(): + return { + "tts_bos_token_id": 9301, + "tts_eos_token_id": 9302, + "listen_token_id": 9303, + "speak_token_id": 9304, + "tts_pad_token_id": 9305, + "unit_token_id": 9306, + "unit_end_token_id": 9307, + "chunk_eos_token_id": 9308, + "chunk_tts_eos_token_id": 9309, + "turn_eos_token_id": 9310, + } + + +def _native_duplex_handoff(request_id, prompt_ids, cumulative_output_ids, *, text="hello"): + latent = torch.arange( + (len(prompt_ids) + len(cumulative_output_ids)) * 4, + dtype=torch.float16, + ).reshape(-1, 4) + output = SimpleNamespace( + token_ids=cumulative_output_ids, + text=text, + multimodal_output={ + "latent": latent, + "duplex_prompt_token_ids": list(prompt_ids), + "meta": _native_duplex_meta(), + }, + ) + return SimpleNamespace( + request_id=request_id, + prompt_token_ids=[0, 0, 0], + outputs=[output], + ) + + +def _install_native_duplex_decoder(streaming_context, token_text): + def _decode(token_ids): + return "".join(token_text.get(int(token_id), "") for token_id in token_ids) + + streaming_context.source_token_decoder = _decode + return streaming_context + + +def test_native_duplex_segment_text_comes_from_token_delta_decode(): + streaming_context = _install_native_duplex_decoder( + SimpleNamespace(bridge_states={}), + { + 21: "今早", + 22: "吃东。", + 31: "西了吗?", + }, + ) + + ids1, text1 = _native_duplex_segment_output_ids( + [9304, 21, 22, 9308], + "今早吃东。", + streaming_context, + request_id="duplex-token-text", + ) + assert ids1 == [9304, 21, 22, 9308] + assert text1 == "今早吃东。" + + ids2, text2 = _native_duplex_segment_output_ids( + [9304, 21, 22, 9308, 9304, 31, 9308], + "今早吃东。西了吗?如需陪聊,", + streaming_context, + request_id="duplex-token-text", + ) + assert ids2 == [9304, 31, 9308] + assert text2 == "西了吗?" + + +def test_native_duplex_segment_restart_does_not_reuse_character_cursor(): + streaming_context = _install_native_duplex_decoder( + SimpleNamespace(bridge_states={}), + { + 21: "你好", + 31: "你好", + 32: ",有什么想聊的吗?", + }, + ) + + _native_duplex_segment_output_ids( + [9304, 21, 9310], + "你好,第一轮尾巴", + streaming_context, + request_id="duplex-restart-text", + ) + + ids2, text2 = _native_duplex_segment_output_ids( + [9304, 31, 32, 9310], + "你好,第一轮尾巴你好,有什么想聊的吗?", + streaming_context, + request_id="duplex-restart-text", + ) + assert ids2 == [9304, 31, 32, 9310] + assert text2 == "你好,有什么想聊的吗?" + + +def test_native_duplex_segment_turn_id_change_keeps_cumulative_token_cursor(): + streaming_context = _install_native_duplex_decoder( + SimpleNamespace(bridge_states={"duplex": {"turn_id": 1}}), + { + 21: "你好", + 32: ",有什么想聊的吗?", + }, + ) + + ids1, text1 = _native_duplex_segment_output_ids( + [9304, 21], + "你好", + streaming_context, + request_id="duplex-turn-id-reset", + ) + assert ids1 == [9304, 21] + assert text1 == "你好" + + streaming_context.bridge_states["duplex"]["turn_id"] = 2 + ids2, text2 = _native_duplex_segment_output_ids( + [9304, 21, 32], + "你好,有什么想聊的吗?", + streaming_context, + request_id="duplex-turn-id-reset", + ) + assert ids2 == [32] + assert text2 == ",有什么想聊的吗?" + + +def test_llm2tts_native_duplex_hands_off_segment_deltas(): + """Each duplex handoff must carry only the current segment's tokens. + + The thinker's resumable request reports cumulative output ids while + earlier segments are already folded into the prompt, so re-forwarding + the full list misaligns the hidden states and grows the talker prompt + without bound. + """ + streaming_context = SimpleNamespace(bridge_states={}) + + # Segment 1: speak + two text tokens + chunk_eos. + seg1_output = [9304, 21, 22, 9308] + handoff1 = _native_duplex_handoff("duplex-delta", [101, 102], list(seg1_output)) + converted1 = llm2tts([handoff1], prompt=[{}], _streaming_context=streaming_context) + info1 = converted1[0]["model_intermediate_buffer"] + assert info1["ids"]["output"] == seg1_output + assert info1["llm_output_text"] == ["hello"] + assert info1["meta"]["native_duplex_segment_text"] == "hello" + assert "llm_output_text" in info1["meta"]["override_keys"] + assert ["meta", "native_duplex_segment_text"] in info1["meta"]["override_keys"] + assert converted1[0]["prompt_token_ids"] == [21, 22] + assert info1["meta"]["segment_end"] is True + + # Segment 2: prompt now contains the folded segment-1 tokens; the + # cumulative output list still carries segment 1 ahead of segment 2. + seg2_output = [9304, 31, 9308] + prompt2 = [101, 102, *seg1_output, 555] + handoff2 = _native_duplex_handoff("duplex-delta", prompt2, [*seg1_output, *seg2_output]) + converted2 = llm2tts([handoff2], prompt=[{}], _streaming_context=streaming_context) + info2 = converted2[0]["model_intermediate_buffer"] + assert info2["ids"]["output"] == seg2_output + # Cumulative text "hello" was fully delivered with segment 1; only the + # delta (empty here) rides along with segment 2. + assert info2["llm_output_text"] == [""] + assert info2["meta"]["native_duplex_segment_text"] == "" + assert converted2[0]["prompt_token_ids"] == [31] + assert info2["meta"]["segment_end"] is True + + +def test_llm2tts_native_duplex_marks_talker_handoff_as_data_plane(): + streaming_context = SimpleNamespace( + bridge_states={ + "duplex": { + "session_id": "sid-stage1-stop", + "incarnation": 2, + "epoch": 0, + "turn_id": 7, + "session_config": { + "extra_body": { + "duplex_stage_sampling_params": { + "1": {"stop_token_ids": [151645]}, + }, + }, + }, + }, + }, + ) + handoff = _native_duplex_handoff("duplex-stage1-stop", [101, 102], [9304, 21, 9308]) + + converted = llm2tts([handoff], prompt=[{}], _streaming_context=streaming_context) + + duplex = converted[0]["model_intermediate_buffer"]["duplex"] + assert duplex["data_plane"] is True + assert duplex["session_id"] == "sid-stage1-stop" + assert duplex["incarnation"] == 2 + assert duplex["epoch"] == 0 + assert duplex["turn_id"] == 7 + assert duplex["session_config"]["extra_body"]["duplex_stage_sampling_params"]["1"]["stop_token_ids"] == [151645] + + +def test_llm2tts_native_duplex_carries_runtime_ref_audio_to_talker_payload(): + ref_audio = base64.b64encode(struct.pack("<3f", 0.1, 0.2, 0.3)).decode("ascii") + streaming_context = SimpleNamespace( + bridge_states={ + "duplex": { + "session_id": "sid-runtime-ref", + "runtime_config": { + "ref_audio_data": ref_audio, + "ref_audio_format": "pcm_f32le", + "ref_audio_sample_rate_hz": 16000, + }, + }, + }, + ) + handoff = _native_duplex_handoff("duplex-runtime-ref", [101, 102], [9304, 21, 9308]) + + converted = llm2tts([handoff], prompt=[{}], _streaming_context=streaming_context) + + info = converted[0]["model_intermediate_buffer"] + assert torch.allclose(torch.tensor(info["codes"]["ref"]), torch.tensor([0.1, 0.2, 0.3])) + assert info["meta"]["ref_audio_sr"] == 16000 + assert info["duplex"]["runtime_config"]["ref_audio_data"] == ref_audio + + +def test_llm2tts_native_duplex_advances_model_turn_only_after_turn_eos(): + streaming_context = SimpleNamespace( + bridge_states={ + "duplex": { + "session_id": "sid-model-turn", + "epoch": 0, + "turn_id": 9, + "model_turn_id": 0, + }, + } + ) + + first = _native_duplex_handoff( + "duplex-model-turn", + [101], + [9304, 21, 9310], + text="first", + ) + converted_first = llm2tts([first], prompt=[{}], _streaming_context=streaming_context) + + first_duplex = converted_first[0]["model_intermediate_buffer"]["duplex"] + assert first_duplex["turn_id"] == 0 + assert streaming_context.bridge_states["duplex"]["model_turn_id"] == 1 + + second = _native_duplex_handoff( + "duplex-model-turn", + [101, 9304, 21, 9310, 555], + [9304, 21, 9310, 9304, 22, 9308], + text="second", + ) + converted_second = llm2tts([second], prompt=[{}], _streaming_context=streaming_context) + + second_duplex = converted_second[0]["model_intermediate_buffer"]["duplex"] + assert second_duplex["turn_id"] == 1 + assert streaming_context.bridge_states["duplex"]["model_turn_id"] == 1 + + +def test_llm2tts_native_duplex_accumulates_tts_condition_across_handoffs(): + """Every handoff must carry the FULL accumulated tts condition. + + The runner's resume-prefill path REPLACES the streaming buffer (only + in-place updates merge), so per-segment tts payloads were silently lost + for alternating segments; the talker then vocalized text it never saw. + Handing the complete history per handoff makes replacement lossless. + """ + streaming_context = SimpleNamespace(bridge_states={}) + + seg1_output = [9304, 21, 22, 9308] + handoff1 = _native_duplex_handoff("duplex-acc", [101, 102], list(seg1_output)) + converted1 = llm2tts([handoff1], prompt=[{}], _streaming_context=streaming_context) + info1 = converted1[0]["model_intermediate_buffer"] + assert info1["ids"]["tts"] == [21, 22] + assert len(info1["hidden_states"]["tts"]) == 2 + + seg2_output = [9304, 31, 9308] + prompt2 = [101, 102, *seg1_output, 555] + handoff2 = _native_duplex_handoff("duplex-acc", prompt2, [*seg1_output, *seg2_output]) + converted2 = llm2tts([handoff2], prompt=[{}], _streaming_context=streaming_context) + info2 = converted2[0]["model_intermediate_buffer"] + assert info2["ids"]["tts"] == [21, 22, 31] + assert len(info2["hidden_states"]["tts"]) == 3 + + +def test_llm2tts_native_duplex_preserves_text_cursor_across_turn_reset(): + streaming_context = _install_native_duplex_decoder( + SimpleNamespace(bridge_states={}), + { + 21: "hello", + 31: "hello", + }, + ) + + seg1_output = [9304, 21, 9310] + handoff1 = _native_duplex_handoff("duplex-turn-text", [101, 102], seg1_output, text="hello") + converted1 = llm2tts([handoff1], prompt=[{}], _streaming_context=streaming_context) + info1 = converted1[0]["model_intermediate_buffer"] + assert info1["meta"]["native_duplex_segment_text"] == "hello" + + seg2_output = [9304, 31, 9310] + handoff2 = _native_duplex_handoff( + "duplex-turn-text", + [101, 102, *seg1_output], + seg2_output, + text="hellohello", + ) + converted2 = llm2tts([handoff2], prompt=[{}], _streaming_context=streaming_context) + info2 = converted2[0]["model_intermediate_buffer"] + assert info2["ids"]["tts"] == [31, 9310] + assert info2["meta"]["native_duplex_segment_text"] == "hello" + + +def test_llm2tts_native_duplex_preserves_output_cursor_across_turn_reset(): + streaming_context = _install_native_duplex_decoder( + SimpleNamespace(bridge_states={}), + { + 21: "hello", + 31: "hello", + }, + ) + + seg1_output = [9304, 21, 9310] + handoff1 = _native_duplex_handoff("duplex-turn-output", [101, 102], seg1_output, text="hello") + converted1 = llm2tts([handoff1], prompt=[{}], _streaming_context=streaming_context) + info1 = converted1[0]["model_intermediate_buffer"] + assert info1["ids"]["output"] == seg1_output + assert info1["ids"]["tts"] == [21, 9310] + + seg2_output = [9304, 31, 9310] + handoff2 = _native_duplex_handoff( + "duplex-turn-output", + [101, 102, *seg1_output], + [*seg1_output, *seg2_output], + text="hellohello", + ) + converted2 = llm2tts([handoff2], prompt=[{}], _streaming_context=streaming_context) + info2 = converted2[0]["model_intermediate_buffer"] + assert info2["ids"]["output"] == seg2_output + assert info2["ids"]["tts"] == [31, 9310] + assert info2["meta"]["native_duplex_segment_text"] == "hello" + + +def test_llm2tts_native_duplex_resets_tts_condition_on_new_turn_without_resetting_output_cursor(): + streaming_context = _install_native_duplex_decoder( + SimpleNamespace(bridge_states={"duplex": {"turn_id": 0}}), + { + 21: "first", + 31: "second", + }, + ) + + seg1_output = [9304, 21, 9308] + handoff1 = _native_duplex_handoff("duplex-turn-condition", [101, 102], seg1_output, text="first") + converted1 = llm2tts([handoff1], prompt=[{}], _streaming_context=streaming_context) + info1 = converted1[0]["model_intermediate_buffer"] + assert info1["ids"]["output"] == seg1_output + assert info1["ids"]["tts"] == [21] + + streaming_context.bridge_states["duplex"]["turn_id"] = 1 + seg2_output = [9304, 31, 9308] + handoff2 = _native_duplex_handoff( + "duplex-turn-condition", + [101, 102, *seg1_output], + [*seg1_output, *seg2_output], + text="firstsecond", + ) + converted2 = llm2tts([handoff2], prompt=[{}], _streaming_context=streaming_context) + info2 = converted2[0]["model_intermediate_buffer"] + assert info2["ids"]["output"] == seg2_output + assert info2["ids"]["tts"] == [31] + assert info2["meta"]["native_duplex_segment_text"] == "second" + + +def test_llm2tts_never_aliases_thinker_token_list(): + """The talker prompt must never be the thinker's live token list object. + + CompletionOutput.token_ids can alias the upstream detokenizer's internal + list; forwarding that object lets the stage-1 streaming update extend the + list with itself, doubling the recorded output every segment. + """ + live_token_list = [9304, 21, 22, 9308] + handoff = _native_duplex_handoff("duplex-alias", [101, 102], live_token_list) + converted = llm2tts([handoff], prompt=[{}]) + scheduler_prompt = converted[0]["prompt_token_ids"] + assert scheduler_prompt is not live_token_list + scheduler_prompt.extend(scheduler_prompt) + assert live_token_list == [9304, 21, 22, 9308] diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index c70b8c487..13a491896 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -131,6 +131,18 @@ class TestStageConfig: assert "max_num_seqs" not in omega_config.engine_args # Legacy field name for backward compatibility assert omega_config.engine_input_source == [] + assert omega_config.session_mode == "turn" + + def test_to_omegaconf_duplex_session_mode(self): + """Test that session mode is preserved as stage metadata.""" + config = StageConfig( + stage_id=0, + model_stage="thinker", + session_mode="duplex", + ) + omega_config = config.to_omegaconf() + + assert omega_config.session_mode == "duplex" def test_to_omegaconf_with_runtime_overrides(self): """Test that runtime overrides are applied to OmegaConf output.""" @@ -1107,6 +1119,57 @@ class TestResolveScheduler: class TestDeployConfigLoading: + def test_load_minicpmo_duplex_deploy_config(self): + deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "minicpmo_4_5_duplex.yaml" + + deploy = load_deploy_config(deploy_path) + pipeline = resolve_pipeline_config("minicpmo_4_5") + assert isinstance(pipeline, PipelineConfig) + + stages = merge_pipeline_deploy(pipeline, deploy) + + assert deploy.session_mode == "duplex" + assert deploy.async_chunk is False + assert deploy.active_stream_window == 1 + assert [stage.session_mode for stage in stages] == ["duplex", "duplex"] + assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex"] + assert [stage.yaml_engine_args["async_scheduling"] for stage in stages] == [False, False] + assert all("Async" not in (stage.scheduler_cls or "") for stage in stages) + assert deploy.stages[1].enforce_eager is False + assert deploy.stages[1].compilation_config == {"cudagraph_mode": "PIECEWISE"} + assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 4096 + assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [151645] + + @pytest.mark.parametrize( + ("filename", "stage0_devices", "stage1_devices", "stage1_replicas"), + [ + ("minicpmo_4_5_3gpu_stage1_replicas.yaml", "0", "1,2", 2), + ("minicpmo_4_5_4gpu_stage1_replicas.yaml", "0", "1,2,3", 3), + ("minicpmo_4_5_8x4090_stage1_replicas.yaml", "0,1,2,3", "4,5,6,7", 4), + ], + ) + def test_load_minicpmo_replica_deploy_configs( + self, + filename: str, + stage0_devices: str, + stage1_devices: str, + stage1_replicas: int, + ): + deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / filename + + deploy = load_deploy_config(deploy_path) + pipeline = resolve_pipeline_config("minicpmo_4_5") + assert isinstance(pipeline, PipelineConfig) + + stages = merge_pipeline_deploy(pipeline, deploy) + + assert deploy.session_mode == "duplex" + assert [stage.yaml_runtime["devices"] for stage in stages] == [stage0_devices, stage1_devices] + assert stages[0].yaml_runtime["num_replicas"] == 1 + assert stages[1].yaml_runtime["num_replicas"] == stage1_replicas + assert all(stage.yaml_engine_args["async_scheduling"] is False for stage in stages) + assert all("Async" not in (stage.scheduler_cls or "") for stage in stages) + def test_custom_voice_dir_is_pipeline_wide_engine_arg(self, tmp_path): deploy_path = tmp_path / "qwen3_tts_custom_voice.yaml" custom_voice_dir = tmp_path / "voices" diff --git a/tests/test_data_entry_keys.py b/tests/test_data_entry_keys.py index 73dcc82c2..4720cd8ea 100644 --- a/tests/test_data_entry_keys.py +++ b/tests/test_data_entry_keys.py @@ -21,6 +21,8 @@ from vllm_omni.data_entry_keys import ( ) from vllm_omni.engine import AdditionalInformationPayload +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + class TestOmniPayloadStruct: """Runtime-validated mirror of OmniPayload (msgspec.Struct).""" @@ -43,6 +45,16 @@ class TestOmniPayloadStruct: with pytest.raises(msgspec.ValidationError, match="unknown field"): to_struct({"meta": {"finisheed": True}}) + def test_to_struct_rejects_minicpmo_duplex_runtime_handoff_fields(self): + with pytest.raises(msgspec.ValidationError, match="unknown field"): + to_struct({"hidden_states": {"tts": torch.zeros(1)}}) + with pytest.raises(msgspec.ValidationError, match="unknown field"): + to_struct({"ids": {"tts": [1]}}) + with pytest.raises(msgspec.ValidationError, match="unknown field"): + to_struct({"meta": {"ref_audio_sr": 16000}}) + with pytest.raises(msgspec.ValidationError, match="unknown field"): + to_struct({"meta": {"listen_token_id": 151685}}) + def test_to_struct_rejects_wrong_type(self): with pytest.raises(msgspec.ValidationError, match="Expected"): to_struct({"meta": {"left_context_size": "not_an_int"}}) @@ -66,10 +78,10 @@ class TestOmniPayloadStruct: def test_struct_with_all_categories(self): d = { - "hidden_states": {"output": torch.zeros(1)}, + "hidden_states": {"output": torch.zeros(1), "last": torch.ones(2, 4)}, "embed": {"prefill": torch.zeros(1), "tts_bos": torch.zeros(1)}, - "ids": {"all": [1, 2], "prompt": [1]}, - "codes": {"audio": torch.zeros(1)}, + "ids": {"all": [1, 2], "prompt": [1], "speech_token": [3, 4]}, + "codes": {"audio": torch.zeros(1), "ref": torch.ones(4)}, "meta": {"left_context_size": 3, "num_processed_tokens": 7}, } s = to_struct(d) @@ -79,7 +91,10 @@ class TestOmniPayloadStruct: assert isinstance(s.codes, CodesStruct) assert isinstance(s.meta, MetaStruct) assert s.ids.all == [1, 2] + assert s.ids.speech_token == [3, 4] + assert torch.equal(s.hidden_states.last, torch.ones(2, 4)) assert s.meta.num_processed_tokens == 7 + assert torch.equal(s.codes.ref, torch.ones(4)) class TestValidatePayload: @@ -332,7 +347,7 @@ class TestSerializeDeserializePayload: "hidden_states": {"output": torch.tensor([1.0, 2.0])}, "ids": {"all": [1, 2, 3]}, "meta": {"finished": torch.tensor(False, dtype=torch.bool), "ar_width": 4}, - "codes": {"audio": torch.tensor([3.0])}, + "codes": {"audio": torch.tensor([3.0]), "ref": torch.tensor([0.1, 0.2])}, } wire = serialize_payload(original) restored = deserialize_payload(wire) @@ -341,6 +356,7 @@ class TestSerializeDeserializePayload: assert restored["meta"]["finished"].item() is False assert restored["meta"]["ar_width"] == 4 assert torch.equal(restored["codes"]["audio"], original["codes"]["audio"]) + assert torch.equal(restored["codes"]["ref"], original["codes"]["ref"]) def test_hidden_states_layers_round_trip(self): original = { diff --git a/tests/test_omni_request.py b/tests/test_omni_request.py index 09f893a03..4f4718f09 100644 --- a/tests/test_omni_request.py +++ b/tests/test_omni_request.py @@ -20,7 +20,7 @@ from vllm.sampling_params import SamplingParams from vllm.v1.request import Request from vllm_omni.engine import PromptEmbedsPayload -from vllm_omni.request import OmniRequest +from vllm_omni.request import OmniRequest, OmniStreamingUpdate pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -32,7 +32,12 @@ def test_omni_params_are_keyword_only(): came first positionally and broke positional ``Request(...)`` construction. """ params = inspect.signature(OmniRequest.__init__).parameters - for name in ("prompt_embeds", "external_req_id", "additional_information"): + for name in ( + "prompt_embeds", + "external_req_id", + "additional_information", + "model_intermediate_buffer", + ): assert params[name].kind is inspect.Parameter.KEYWORD_ONLY, name @@ -87,3 +92,21 @@ def test_keyword_omni_params_round_trip(): # The serialized payload is retained and decoded into a tensor on the base. assert req.prompt_embeds_payload is payload assert torch.equal(req.prompt_embeds, torch.from_numpy(arr)) + + +def test_model_intermediate_buffer_round_trips_to_streaming_update(): + buffer = {"duplex": {"turn_id": 3}} + request = OmniRequest( + request_id="req-model-buffer", + prompt_token_ids=[1, 2], + sampling_params=SamplingParams(max_tokens=16), + pooling_params=None, + resumable=True, + model_intermediate_buffer=buffer, + ) + + update = OmniStreamingUpdate.from_request(request) + + assert update is not None + assert request.model_intermediate_buffer is buffer + assert update.model_intermediate_buffer is buffer diff --git a/tests/utils/test_mm_outputs_partition.py b/tests/utils/test_mm_outputs_partition.py index f9affa471..d70b95a79 100644 --- a/tests/utils/test_mm_outputs_partition.py +++ b/tests/utils/test_mm_outputs_partition.py @@ -1,7 +1,10 @@ +import pytest import torch from vllm_omni.utils.mm_outputs import partition_flat_payload, partition_payload_list +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def test_partition_thinker_latent_payload(): payload = { @@ -42,3 +45,27 @@ def test_partition_payload_list_preserves_request_alignment(): inter_list, client_list = partition_payload_list(payloads) assert inter_list == [payloads[0], None] assert client_list == [None, payloads[1]] + + +def test_partition_duplex_audio_transcript_metadata_to_client_mm(): + payload = { + "model_outputs": torch.zeros(1, 2400), + "meta.duplex_epoch": torch.tensor([3], dtype=torch.int32), + "meta.duplex_turn_id": torch.tensor([2], dtype=torch.int32), + "meta.llm_output_text_utf8": torch.tensor([104, 105], dtype=torch.uint8), + "meta.audio_text_total_chars": torch.tensor([2], dtype=torch.int32), + "meta.tts_is_last_chunk": torch.tensor([1], dtype=torch.int32), + "meta.turn_end": torch.tensor([1], dtype=torch.int32), + "meta.native_duplex_segment_text": "hi", + } + + inter, client = partition_flat_payload(payload) + + assert "meta.llm_output_text_utf8" in client + assert "meta.audio_text_total_chars" in client + assert "meta.duplex_epoch" in client + assert "meta.duplex_turn_id" in client + assert "meta.tts_is_last_chunk" in client + assert "meta.turn_end" in client + assert "meta.native_duplex_segment_text" not in client + assert "meta.native_duplex_segment_text" in inter diff --git a/tests/worker/test_native_duplex_hooks.py b/tests/worker/test_native_duplex_hooks.py new file mode 100644 index 000000000..a5a81a1de --- /dev/null +++ b/tests/worker/test_native_duplex_hooks.py @@ -0,0 +1,2576 @@ +from __future__ import annotations + +import inspect +import os +import threading +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _minicpmo_duplex_policy_case( + state: SimpleNamespace, + payload: dict[str, object], +): + from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingRow + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + session_key = ("sid-policy", 1) + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model._minicpmo45_native_duplex_token_ids_cache = { + "listen_token_id": 7, + "tts_bos_token_id": 8, + "turn_eos_token_id": 9, + } + model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state}) + row = DuplexSamplingRow( + row_idx=0, + request_id="req-policy", + session_id=session_key[0], + incarnation=session_key[1], + seq=3, + payload=payload, + max_tokens=20, + ) + return model, row + + +def test_minicpmo_model_hook_owns_duplex_sampling_rows_and_force_listen(): + from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingRow + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + listen_id = 7 + state = SimpleNamespace( + current_turn_ended=True, + last_terminator_token=None, + pending_terminator_token=None, + ) + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model._minicpmo45_native_duplex_token_ids_cache = { + "listen_token_id": listen_id, + "tts_bos_token_id": 8, + "turn_eos_token_id": 9, + } + model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={("sid-hook", 2): 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, + ) + + 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_payloads == {0: row.payload} + assert model._minicpmo45_duplex_row_max_tokens == {0: 20} + assert logits[0, listen_id].item() == 0.0 + assert torch.isneginf(logits[0, :listen_id]).all() + assert torch.isneginf(logits[0, listen_id + 1 :]).all() + + +def test_minicpmo_model_hook_turn_end_latch_forces_silence_to_listen(): + state = SimpleNamespace( + current_turn_ended=True, + last_terminator_token=9, + pending_terminator_token=None, + ) + model, row = _minicpmo_duplex_policy_case(state, {"is_speech": False}) + logits = torch.zeros((1, 16), dtype=torch.float32) + logits[0, 10] = 20.0 + + model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,)) + + assert logits[0, 7].item() == 0.0 + assert torch.isneginf(logits[0, :7]).all() + assert torch.isneginf(logits[0, 8:]).all() + + +def test_minicpmo_model_hook_pending_speech_after_turn_eos_allows_silence_sampling(): + state = SimpleNamespace( + current_turn_ended=True, + last_terminator_token=9, + pending_terminator_token=None, + pending_speech_context=True, + ) + model, row = _minicpmo_duplex_policy_case(state, {"is_speech": False}) + logits = torch.zeros((1, 16), dtype=torch.float32) + logits[0, 7] = 30.0 + logits[0, 10] = 20.0 + original_logits = logits.clone() + + model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,)) + + assert torch.equal(logits, original_logits) + assert state.pending_speech_context is True + + +def test_minicpmo_model_hook_speech_row_does_not_create_pending_context(): + state = SimpleNamespace( + current_turn_ended=True, + last_terminator_token=9, + pending_terminator_token=None, + pending_speech_context=False, + ) + model, row = _minicpmo_duplex_policy_case(state, {"is_speech": True}) + logits = torch.zeros((1, 16), dtype=torch.float32) + logits[0, 7] = 30.0 + logits[0, 8] = -2.0 + logits[0, 10] = 20.0 + original_logits = logits.clone() + + model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,)) + + assert state.current_turn_ended is True + assert state.last_terminator_token is None + assert state.pending_speech_context is False + assert torch.equal(logits, original_logits) + + +def test_minicpmo_model_hook_old_response_output_does_not_clear_pending_speech_context(): + state = SimpleNamespace( + current_turn_ended=False, + last_terminator_token=None, + pending_terminator_token=None, + 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_payloads = {0: row.payload} + + model._record_minicpmo45_duplex_terminator( + 0, + 10, + {"listen_token_id": 7, "chunk_eos_token_id": -1, "chunk_tts_eos_token_id": -1, "turn_eos_token_id": 9}, + ) + + assert state.pending_speech_context is True + assert state.current_turn_ended is False + + +def test_minicpmo_model_hook_new_response_output_clears_pending_speech_context(): + state = SimpleNamespace( + current_turn_ended=True, + last_terminator_token=7, + pending_terminator_token=7, + 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_payloads = {0: row.payload} + + model._record_minicpmo45_duplex_terminator( + 0, + 8, + {"listen_token_id": 7, "chunk_eos_token_id": -1, "chunk_tts_eos_token_id": -1, "turn_eos_token_id": 9}, + ) + + assert state.pending_speech_context is False + assert state.current_turn_ended is False + + +def test_minicpmo_model_hook_empty_speak_envelope_preserves_pending_speech_context(): + state = SimpleNamespace( + current_turn_ended=True, + last_terminator_token=9, + pending_terminator_token=None, + pending_speech_context=True, + 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_payloads = {0: row.payload} + token_ids = { + "listen_token_id": 7, + "tts_bos_token_id": 8, + "chunk_eos_token_id": -1, + "chunk_tts_eos_token_id": -1, + "turn_eos_token_id": 9, + } + + model._record_minicpmo45_duplex_terminator(0, 8, token_ids) + model._record_minicpmo45_duplex_terminator(0, 9, token_ids) + + assert state.current_turn_ended is True + assert state.pending_speech_context is True + assert state.pending_speech_response_open is False + + logits = torch.zeros((1, 16), dtype=torch.float32) + logits[0, 7] = 30.0 + logits[0, 10] = 20.0 + original_logits = logits.clone() + + model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,)) + + assert torch.equal(logits, original_logits) + + +def test_minicpmo_model_hook_second_new_response_step_is_not_forced_to_listen(): + state = SimpleNamespace( + current_turn_ended=True, + last_terminator_token=9, + pending_terminator_token=None, + pending_speech_context=True, + ) + model, row = _minicpmo_duplex_policy_case(state, {"is_speech": False}) + logits = torch.zeros((1, 16), dtype=torch.float32) + logits[0, 7] = 30.0 + logits[0, 10] = 20.0 + 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_payloads = {0: row.payload} + model._record_minicpmo45_duplex_terminator( + 0, + 8, + {"listen_token_id": 7, "chunk_eos_token_id": -1, "chunk_tts_eos_token_id": -1, "turn_eos_token_id": 9}, + ) + next_logits = torch.zeros((1, 16), dtype=torch.float32) + next_logits[0, 7] = 30.0 + next_logits[0, 10] = 20.0 + next_original_logits = next_logits.clone() + + model.prepare_duplex_sampling(next_logits, SimpleNamespace(), (row,)) + + assert torch.equal(logits, original_logits) + assert torch.equal(next_logits, next_original_logits) + assert state.pending_speech_context is False + assert state.current_turn_ended is False + + +def test_minicpmo_model_hook_same_speech_row_second_step_does_not_rearm_pending_context(): + state = SimpleNamespace( + current_turn_ended=True, + last_terminator_token=9, + pending_terminator_token=None, + pending_speech_context=True, + ) + model, row = _minicpmo_duplex_policy_case(state, {"is_speech": True}) + logits = torch.zeros((1, 16), dtype=torch.float32) + logits[0, 7] = 30.0 + 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_payloads = {0: row.payload} + model._record_minicpmo45_duplex_terminator( + 0, + 10, + {"listen_token_id": 7, "chunk_eos_token_id": -1, "chunk_tts_eos_token_id": -1, "turn_eos_token_id": 9}, + ) + next_logits = torch.zeros((1, 16), dtype=torch.float32) + next_logits[0, 7] = 30.0 + next_logits[0, 8] = -2.0 + next_logits[0, 10] = 20.0 + + model.prepare_duplex_sampling(next_logits, SimpleNamespace(), (row,)) + + assert state.pending_speech_context is False + assert state.current_turn_ended is False + assert torch.isneginf(next_logits[0, 7]) + assert next_logits[0, 8].item() == 30.0 + + +def test_minicpmo_model_hook_mid_turn_speech_redirects_listen_to_tts_bos(): + state = SimpleNamespace( + current_turn_ended=False, + last_terminator_token=None, + pending_terminator_token=None, + ) + model, row = _minicpmo_duplex_policy_case(state, {"is_speech": True}) + logits = torch.zeros((1, 16), dtype=torch.float32) + logits[0, 7] = 30.0 + logits[0, 8] = -2.0 + logits[0, 10] = 20.0 + + model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,)) + + assert torch.isneginf(logits[0, 7]) + assert logits[0, 8].item() == 30.0 + assert logits[0, 10].item() == 20.0 + + +def test_minicpmo_model_hook_ignores_serving_new_user_turn_marker(): + state = SimpleNamespace( + current_turn_ended=False, + last_terminator_token=8, + pending_terminator_token=8, + ) + model, row = _minicpmo_duplex_policy_case( + state, + {"is_speech": True, "new_user_turn": True, "force_speak": True}, + ) + logits = torch.zeros((1, 16), dtype=torch.float32) + logits[0, 7] = 5.0 + logits[0, 10] = 20.0 + model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,)) + + assert state.current_turn_ended is False + assert state.last_terminator_token is None + assert torch.isneginf(logits[0, 7]) + assert logits[0, 8].item() == 5.0 + + +def test_generic_ar_runner_builds_typed_duplex_sampling_rows(): + from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingHelper + from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner + + runner = GPUARModelRunner.__new__(GPUARModelRunner) + runner.input_batch = SimpleNamespace(req_ids=["req-duplex", "req-plain"]) + runner.model_intermediate_buffer = { + "req-duplex": { + "duplex": { + "data_plane": True, + "session_id": "sid-runner-hook", + "incarnation": 4, + "seq": 4, + "payload": {"is_speech": True}, + } + } + } + runner.requests = { + "req-duplex": SimpleNamespace( + sampling_params=SimpleNamespace(max_tokens=32), + ) + } + helper = DuplexSamplingHelper() + helper.active_request_ids = {"req-duplex"} + + rows = helper.rows(runner) + + assert len(rows) == 1 + 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 + + +def test_generic_ar_runner_skips_duplex_rows_without_model_hook(): + from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner + + runner = GPUARModelRunner.__new__(GPUARModelRunner) + runner.input_batch = SimpleNamespace( + sampling_metadata=object(), + update_async_output_token_ids=lambda: None, + ) + runner.model = SimpleNamespace(sample=lambda *_args, **_kwargs: None, prefer_model_sampler=False) + runner.sampler = lambda **_kwargs: "standard-sampler" + runner._duplex_sampling_helper = SimpleNamespace( + active_request_ids={"req-duplex"}, + rows=lambda *_args: pytest.fail("duplex rows built without a model hook"), + ) + + assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "standard-sampler" + + +def test_plain_model_hook_resolution_does_not_allocate_duplex_tracking(): + from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner + + runner = GPUARModelRunner.__new__(GPUARModelRunner) + runner.model = SimpleNamespace() + runner._duplex_sampling_hook = None + runner._duplex_sampling_hook_resolved = False + + assert runner._resolve_duplex_sampling_hook() is None + assert not hasattr(runner, "_duplex_sampling_helper") + + +def test_minicpmo_non_duplex_sample_skips_duplex_row_scan(): + from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner + + calls = [] + runner = GPUARModelRunner.__new__(GPUARModelRunner) + runner.input_batch = SimpleNamespace( + sampling_metadata=SimpleNamespace(output_token_ids=[]), + update_async_output_token_ids=lambda: None, + ) + runner.model = SimpleNamespace( + sample=lambda *_args, **_kwargs: "model-sampler", + prefer_model_sampler=True, + skips_model_sampler_output_token_history=True, + prepare_duplex_sampling=lambda *_args: calls.append("prepare"), + ) + runner.sampler = SimpleNamespace() + runner._duplex_sampling_helper = SimpleNamespace( + active_request_ids=set(), + hook_active=False, + rows=lambda *_args: pytest.fail("non-duplex MiniCPM scanned request rows"), + ) + + assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "model-sampler" + assert calls == [] + + +def test_minicpmo_duplex_sample_clears_stale_rows_once_without_scanning(): + from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner + + calls = [] + runner = GPUARModelRunner.__new__(GPUARModelRunner) + runner.input_batch = SimpleNamespace( + sampling_metadata=SimpleNamespace(output_token_ids=[]), + update_async_output_token_ids=lambda: None, + ) + runner.model = SimpleNamespace( + sample=lambda *_args, **_kwargs: "model-sampler", + prefer_model_sampler=True, + skips_model_sampler_output_token_history=True, + prepare_duplex_sampling=lambda _logits, _metadata, rows: calls.append(rows), + ) + runner.sampler = SimpleNamespace() + runner._duplex_sampling_helper = SimpleNamespace( + active_request_ids=set(), + hook_active=True, + rows=lambda *_args: pytest.fail("duplex cleanup scanned request rows"), + ) + + assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "model-sampler" + assert runner._sample(torch.zeros((1, 4)), spec_decode_metadata=None) == "model-sampler" + assert calls == [()] + + +def test_generic_ar_runner_has_no_minicpmo_sampler_state_or_typeerror_probe(): + from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner + + source = inspect.getsource(GPUARModelRunner) + + assert "_minicpmo45_duplex_row" not in source + assert "_minicpmo45_native_duplex_token_ids" not in source + assert 'if "duplex_rows" not in str(exc)' not in source + + +def test_minicpmo_model_cleans_incarnation_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) + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model = SimpleNamespace() + model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: object()}) + model._minicpmo45_duplex_request_sessions = {request_id: session_key} + model._minicpmo45_force_listen_applied_segments = { + (request_id, 1), + ("duplex-sid-other-e0-stage0", 2), + } + + model.on_requests_finished({request_id}) + + assert model._minicpmo45_duplex_data_plane_helper.sessions == {} + assert model._minicpmo45_duplex_request_sessions == {} + assert model._minicpmo45_force_listen_applied_segments == { + ("duplex-sid-other-e0-stage0", 2), + } + + +def test_minicpmo_stage0_routes_duplex_metadata_per_batched_request(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + from vllm_omni.utils.mm_outputs import to_payload_element + + class _Thinker(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("device_anchor", torch.zeros(1)) + + def forward(self, *, input_ids, **kwargs): + del kwargs + return torch.arange(input_ids.numel() * 4, dtype=torch.float32).reshape(input_ids.numel(), 4) + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + torch.nn.Module.__init__(model) + model.model_stage = "llm" + model.thinker = _Thinker() + output = model.forward( + input_ids=torch.tensor([1, 2, 3, 4]), + positions=torch.arange(4), + runtime_additional_information=[ + { + "global_request_id": ["req-a"], + "duplex": { + "duplex_prompt_token_ids": [101, 102], + "special_token_ids": {"listen_token_id": 701}, + }, + }, + { + "global_request_id": ["req-b"], + "duplex": { + "duplex_prompt_token_ids": [201, 202, 203], + "special_token_ids": {"listen_token_id": 702}, + }, + }, + ], + ) + + prompt_rows = output.multimodal_outputs["duplex_prompt_token_ids"] + listen_rows = output.multimodal_outputs["meta"]["listen_token_id"] + assert to_payload_element(prompt_rows, 0, 0, 2) == [101, 102] + assert to_payload_element(prompt_rows, 1, 2, 4) == [201, 202, 203] + assert int(to_payload_element(listen_rows, 0, 0, 2).reshape(-1)[0]) == 701 + assert int(to_payload_element(listen_rows, 1, 2, 4).reshape(-1)[0]) == 702 + + +def test_minicpmo_stage0_rejects_invalid_resolved_ref_audio(): + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + ) + + with pytest.raises(ValueError, match="invalid native duplex ref_audio_data"): + MiniCPMO45Stage0DuplexRuntime._decode_ref_audio_from_session_config( + { + "ref_audio_data": "a", + "ref_audio_format": "pcm_f32le", + } + ) + + +def test_minicpmo_tts_native_duplex_exports_segment_text_not_accumulated_condition_text(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Talker: + _ar_last_chunk_flags = [False] + _ar_last_emitted_text = "你好,你有什莫想聊的吗?" + + def __call__(self, **kwargs): + del kwargs + return None, torch.zeros(8, dtype=torch.float32) + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "tts" + model.config = SimpleNamespace(hidden_size=4) + model.talker = _Talker() + + output = model.forward( + input_ids=torch.zeros(1, dtype=torch.long), + positions=torch.zeros(1, dtype=torch.long), + runtime_additional_information=[ + { + "native_duplex": True, + "llm_output_text": ["你好,你有什莫想聊的吗?你好,你有什莫想聊的吗?"], + "meta": { + "native_duplex_segment_text": "你好,你有什莫想聊的吗?", + }, + } + ], + ) + + text_bytes = output.multimodal_outputs["meta.llm_output_text_utf8"].detach().cpu().tolist() + assert bytes(text_bytes).decode("utf-8") == "你好,你有什莫想聊的吗?" + assert int(output.multimodal_outputs["meta.audio_text_total_chars"].item()) == len("你好,你有什莫想聊的吗?") + + +def test_minicpmo_tts_native_duplex_exports_model_turn_end_metadata(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Talker: + _ar_last_chunk_flags = [True] + _ar_turn_end_flags = [True] + + def __call__(self, **kwargs): + del kwargs + return None, torch.zeros(0, dtype=torch.float32) + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "tts" + model.config = SimpleNamespace(hidden_size=4) + model.talker = _Talker() + + output = model.forward( + input_ids=torch.zeros(1, dtype=torch.long), + positions=torch.zeros(1, dtype=torch.long), + runtime_additional_information=[ + { + "native_duplex": True, + "meta": {"native_duplex_segment_text": ""}, + } + ], + ) + + assert int(output.multimodal_outputs["meta.turn_end"].item()) == 1 + + +def test_minicpmo_tts_routes_batched_requests_and_terminal_flags_independently(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Talker: + def __init__(self): + self.calls = [] + self._ar_last_chunk_flags = [True] + self._ar_turn_end_flags = [False] + self._ar_last_emitted_text = "" + + def __call__(self, **kwargs): + info = kwargs["additional_information"] + request_id = info["global_request_id"][0] + self.calls.append(request_id) + is_last = request_id == "req-b" + self._ar_last_chunk_flags = [is_last] + self._ar_turn_end_flags = [is_last] + self._ar_last_emitted_text = request_id + value = 1.0 if request_id == "req-a" else 2.0 + return None, torch.full((2,), value, dtype=torch.float32) + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "tts" + model.config = SimpleNamespace(hidden_size=4) + model.talker = _Talker() + + output = model.forward( + input_ids=torch.zeros(3, dtype=torch.long), + positions=torch.zeros(3, dtype=torch.long), + runtime_additional_information=[ + { + "global_request_id": ["req-a"], + "native_duplex": True, + "duplex": {"epoch": 3, "turn_id": 7}, + "meta": {"native_duplex_segment_text": "first"}, + }, + { + "global_request_id": ["req-idle"], + "native_duplex": True, + "duplex": {"epoch": 3, "turn_id": 9}, + "meta": {"native_duplex_segment_text": "idle"}, + }, + { + "global_request_id": ["req-b"], + "native_duplex": True, + "duplex": {"epoch": 4, "turn_id": 8}, + "meta": {"native_duplex_segment_text": "second"}, + }, + ], + request_token_spans=[(0, 2), (2, 2), (2, 3)], + ) + + assert model.talker.calls == ["req-a", "req-b"] + assert model.talker._ar_last_chunk_flags == [False, False, True] + assert model.talker._ar_turn_end_flags == [False, False, True] + assert output.multimodal_outputs["meta.req_id"] == ["req-a", "req-b"] + assert output.multimodal_outputs["meta.sparse_audio"] == ["1"] + assert [waveform.tolist() for waveform in output.multimodal_outputs["model_outputs"]] == [ + [1.0, 1.0], + [2.0, 2.0], + ] + assert [int(value.item()) for value in output.multimodal_outputs["meta.tts_is_last_chunk"]] == [0, 1] + assert [int(value.item()) for value in output.multimodal_outputs["meta.turn_end"]] == [0, 1] + assert [int(value.item()) for value in output.multimodal_outputs["meta.duplex_epoch"]] == [3, 4] + assert [int(value.item()) for value in output.multimodal_outputs["meta.duplex_turn_id"]] == [7, 8] + + +@pytest.mark.parametrize( + ("runtime_info", "request_token_spans", "expected_rows", "expected_ids"), + [ + ( + [ + { + "global_request_id": ["req-a"], + "native_duplex": True, + "duplex": {"epoch": 3, "turn_id": 7}, + "meta": {"native_duplex_segment_text": "first"}, + } + ], + [(0, 3)], + [3], + ["req-a"], + ), + ( + [ + { + "global_request_id": ["req-a"], + "native_duplex": True, + "duplex": {"epoch": 3, "turn_id": 7}, + "meta": {"native_duplex_segment_text": "first"}, + }, + { + "global_request_id": ["req-b"], + "native_duplex": True, + "duplex": {"epoch": 4, "turn_id": 8}, + "meta": {"native_duplex_segment_text": "second"}, + }, + ], + [(0, 2), (2, 3)], + [2, 1], + ["req-a", "req-b"], + ), + ], + ids=["single-request", "batched-requests"], +) +def test_minicpmo_tts_ignores_cuda_graph_padding_after_request_spans( + runtime_info, + request_token_spans, + expected_rows, + expected_ids, +): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Talker: + def __init__(self): + self.input_rows = [] + self._ar_last_chunk_flags = [True] + self._ar_turn_end_flags = [False] + self._ar_last_emitted_text = "" + + def __call__(self, **kwargs): + self.input_rows.append(int(kwargs["input_ids"].shape[0])) + return None, torch.ones(2, dtype=torch.float32) + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "tts" + model.config = SimpleNamespace(hidden_size=4) + model.talker = _Talker() + + output = model.forward( + input_ids=torch.zeros(4, dtype=torch.long), + positions=torch.zeros(4, dtype=torch.long), + runtime_additional_information=runtime_info, + request_token_spans=request_token_spans, + ) + + assert model.talker.input_rows == expected_rows + assert output.text_hidden_states.shape == (4, 4) + assert output.multimodal_outputs["meta.req_id"] == expected_ids + + +def test_minicpmo_stage0_special_token_ids_are_tokenizer_derived(): + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + ) + + class _Tokenizer: + unk_token_id = 0 + ids = { + "": 101, + "": 102, + "<|listen|>": 103, + "<|speak|>": 104, + "<|tts_bos|>": 105, + "<|tts_eos|>": 106, + "<|tts_pad|>": 107, + "<|chunk_eos|>": 108, + "<|chunk_tts_eos|>": 109, + "<|turn_eos|>": 110, + } + + def convert_tokens_to_ids(self, token): + return self.ids.get(token, self.unk_token_id) + + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return [self.ids[text]] if text in self.ids else [201, self.ids["<|tts_bos|>"]] + + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.tokenizer = _Tokenizer() + runtime._init_token_ids() + + runtime._require_special_token_ids() + assert runtime.tts_bos_token_id == 105 + assert runtime.stage_padding_token_id() == 102 + assert runtime._special_token_ids()["chunk_tts_eos_token_id"] == 109 + + +def test_minicpmo_stage0_rejects_unknown_special_token_fallbacks(): + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + ) + + class _Tokenizer: + unk_token_id = 0 + ids = { + "": 101, + "": 102, + "<|listen|>": 103, + "<|speak|>": 104, + "<|tts_eos|>": 106, + "<|tts_pad|>": 107, + "<|chunk_eos|>": 108, + "<|chunk_tts_eos|>": 109, + "<|turn_eos|>": 110, + } + + def convert_tokens_to_ids(self, token): + return self.ids.get(token, self.unk_token_id) + + def encode(self, text, add_special_tokens=False): + del text, add_special_tokens + return [self.unk_token_id] + + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.tokenizer = _Tokenizer() + runtime._init_token_ids() + + with pytest.raises(ValueError, match=r"<\|tts_bos\|>"): + runtime._require_special_token_ids() + + +def test_minicpmo_stage0_data_plane_prefill_matches_official_unit_format(): + import torch + + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + _MiniCPMO45Stage0SessionState, + ) + + class _StageModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.embed = torch.nn.Embedding(256, 2) + + def get_input_embeddings(self): + return self.embed + + def get_audio_hidden_states(self, _data): + return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)] + + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.stage_model = _StageModel() + runtime.thinker = runtime.stage_model + runtime.tokenizer = SimpleNamespace( + unk_token_id=0, + convert_tokens_to_ids=lambda token: { + "": 1, + "": 2, + "<|listen|>": 3, + "<|speak|>": 4, + "<|tts_bos|>": 5, + "<|tts_eos|>": 6, + "<|tts_pad|>": 7, + "<|chunk_eos|>": 8, + "<|chunk_tts_eos|>": 9, + "<|turn_eos|>": 10, + "<|audio|>": 11, + }.get(token, 0), + encode=lambda text, add_special_tokens=False: [201, 5], + ) + runtime.processor = SimpleNamespace(get_streaming_chunk_size=lambda: 4) + runtime.device = "cpu" + runtime._init_token_ids() + state = _MiniCPMO45Stage0SessionState(session_id="sid-data-plane-prefill") + + # Official duplex format: each unit is + audio embeddings with no + # per-chunk assistant header or <|tts_bos|> boundary. Decoding starts right + # after the audio so the first sampled token is the listen/speak decision. + result = runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=1) + + assert result["success"] is True + assert result["input_token_ids"] == [1, 11] + assert result["prompt_suffix_len"] == 0 + + # Subsequent units must close the previous unit with first, + # mirroring the official finalize_unit() feed. + result = runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=2) + + assert result["success"] is True + assert result["input_token_ids"] == [2, 1, 11] + assert result["prompt_suffix_len"] == 0 + + +def test_minicpmo_stage0_speech_append_sets_pending_context_once(): + import torch + + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + _MiniCPMO45Stage0SessionState, + ) + + class _StageModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.embed = torch.nn.Embedding(256, 2) + + def get_input_embeddings(self): + return self.embed + + def get_audio_hidden_states(self, _data): + return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)] + + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.stage_model = _StageModel() + runtime.thinker = runtime.stage_model + runtime.tokenizer = SimpleNamespace( + unk_token_id=0, + convert_tokens_to_ids=lambda token: { + "": 1, + "": 2, + "<|listen|>": 3, + "<|speak|>": 4, + "<|tts_bos|>": 5, + "<|tts_eos|>": 6, + "<|tts_pad|>": 7, + "<|chunk_eos|>": 8, + "<|chunk_tts_eos|>": 9, + "<|turn_eos|>": 10, + "<|audio|>": 11, + }.get(token, 0), + encode=lambda text, add_special_tokens=False: [], + ) + runtime.processor = SimpleNamespace(get_streaming_chunk_size=lambda: 4) + runtime.device = "cpu" + runtime._init_token_ids() + state = _MiniCPMO45Stage0SessionState(session_id="sid-speech-append") + + result = runtime._stage_prefill_embeddings_only( + state, + np.zeros(4, dtype=np.float32), + epoch=0, + seq=1, + is_speech=True, + ) + + assert result["success"] is True + assert state.pending_speech_context is True + assert state.pending_speech_append_identity == (0, 1) + + state.pending_speech_context = False + cached = runtime._stage_prefill_embeddings_only( + state, + np.zeros(4, dtype=np.float32), + epoch=0, + seq=1, + is_speech=True, + ) + + assert cached["success"] is True + assert state.pending_speech_context is False + assert state.pending_speech_append_identity == (0, 1) + + next_seq = runtime._stage_prefill_embeddings_only( + state, + np.zeros(4, dtype=np.float32), + epoch=0, + seq=2, + is_speech=True, + ) + + assert next_seq["success"] is True + assert state.pending_speech_context is True + assert state.pending_speech_append_identity == (0, 2) + + state.pending_speech_context = False + next_epoch = runtime._stage_prefill_embeddings_only( + state, + np.zeros(4, dtype=np.float32), + epoch=1, + seq=1, + is_speech=True, + ) + + assert next_epoch["success"] is True + assert state.pending_speech_context is True + assert state.pending_speech_append_identity == (1, 1) + + state.pending_speech_context = False + silence = runtime._stage_prefill_embeddings_only( + state, + np.zeros(4, dtype=np.float32), + epoch=1, + seq=2, + is_speech=False, + ) + + assert silence["success"] is True + assert state.pending_speech_context is False + assert state.pending_speech_append_identity == (1, 1) + + +def test_minicpmo_stage0_failed_append_does_not_set_pending_speech_context(): + import torch + + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + _MiniCPMO45Stage0SessionState, + ) + + class _StageModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.embed = torch.nn.Embedding(256, 2) + + def get_input_embeddings(self): + return self.embed + + def get_audio_hidden_states(self, _data): + return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)] + + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.stage_model = _StageModel() + runtime.thinker = runtime.stage_model + runtime.tokenizer = SimpleNamespace( + unk_token_id=0, + convert_tokens_to_ids=lambda token: { + "": 1, + "": 2, + "<|listen|>": 3, + "<|speak|>": 4, + "<|tts_bos|>": 5, + "<|tts_eos|>": 6, + "<|tts_pad|>": 7, + "<|chunk_eos|>": 8, + "<|chunk_tts_eos|>": 9, + "<|turn_eos|>": 10, + "<|audio|>": 11, + }.get(token, 0), + encode=lambda text, add_special_tokens=False: [], + ) + runtime.processor = SimpleNamespace(get_streaming_chunk_size=lambda: 4) + runtime.device = "cpu" + runtime._init_token_ids() + state = _MiniCPMO45Stage0SessionState(session_id="sid-failed-speech") + + result = runtime._stage_prefill_embeddings_only( + state, + np.zeros(0, dtype=np.float32), + epoch=0, + seq=1, + is_speech=True, + ) + + assert result["success"] is False + assert state.pending_speech_context is False + assert state.pending_speech_append_identity is None + + +def test_minicpmo_stage0_streaming_processor_is_isolated_per_session(): + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + _MiniCPMO45Stage0SessionState, + ) + + class _Mel: + def __init__(self): + self.counter = 0 + + class _Processor: + def __init__(self): + self._streaming_mel_processor = _Mel() + + def set_streaming_mode(self, **_kwargs): + return None + + def reset_streaming(self): + self._streaming_mel_processor.counter = 0 + + def process_audio_streaming(self, _audio, **_kwargs): + self._streaming_mel_processor.counter += 1 + return self._streaming_mel_processor.counter + + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.processor = _Processor() + runtime.stage_model = SimpleNamespace() + runtime.thinker = runtime.stage_model + a = _MiniCPMO45Stage0SessionState(session_id="a") + b = _MiniCPMO45Stage0SessionState(session_id="b") + processor_a = runtime._configure_streaming_processor(a) + processor_b = runtime._configure_streaming_processor(b) + + observed = [ + runtime._process_streaming_audio([], 0, processor=processor_a), + runtime._process_streaming_audio([], 0, processor=processor_b), + runtime._process_streaming_audio([], 1, processor=processor_a), + runtime._process_streaming_audio([], 1, processor=processor_b), + ] + + assert observed == [1, 1, 2, 2] + assert processor_a is not processor_b + assert runtime.processor._streaming_mel_processor.counter == 0 + + +def test_minicpmo_stage0_data_plane_next_append_reinjects_previous_listen(): + import torch + + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + _MiniCPMO45Stage0SessionState, + ) + + class _StageModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.embed = torch.nn.Embedding(256, 2) + + def get_input_embeddings(self): + return self.embed + + def get_audio_hidden_states(self, data): + return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)] + + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.stage_model = _StageModel() + runtime.thinker = runtime.stage_model + runtime.tokenizer = SimpleNamespace( + unk_token_id=0, + convert_tokens_to_ids=lambda token: { + "": 1, + "": 2, + "<|listen|>": 3, + "<|speak|>": 4, + "<|tts_bos|>": 5, + "<|tts_eos|>": 6, + "<|tts_pad|>": 7, + "<|chunk_eos|>": 8, + "<|chunk_tts_eos|>": 9, + "<|turn_eos|>": 10, + "<|audio|>": 11, + }.get(token, 0), + encode=lambda text, add_special_tokens=False: [], + ) + runtime.processor = SimpleNamespace(get_streaming_chunk_size=lambda: 4) + runtime.device = "cpu" + runtime._init_token_ids() + state = _MiniCPMO45Stage0SessionState( + session_id="sid-new-speech-prefill", + audio_chunk_idx=1, + pending_terminator_token=3, + last_terminator_token=3, + current_turn_ended=True, + ) + + result = runtime._stage_prefill_embeddings_only( + state, + np.zeros(4, dtype=np.float32), + seq=2, + ) + + assert result["success"] is True + assert result["input_token_ids"] == [3, 2, 1, 11] + assert state.pending_terminator_token is None + assert state.last_terminator_token == 3 + assert state.current_turn_ended is True + + +def test_minicpmo_stage0_data_plane_turn_eos_closes_previous_unit(): + import torch + + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + _MiniCPMO45Stage0SessionState, + ) + + class _StageModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.embed = torch.nn.Embedding(256, 2) + + def get_input_embeddings(self): + return self.embed + + def get_audio_hidden_states(self, data): + return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)] + + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.stage_model = _StageModel() + runtime.thinker = runtime.stage_model + runtime.tokenizer = SimpleNamespace( + unk_token_id=0, + convert_tokens_to_ids=lambda token: { + "": 1, + "": 2, + "<|listen|>": 3, + "<|speak|>": 4, + "<|tts_bos|>": 5, + "<|tts_eos|>": 6, + "<|tts_pad|>": 7, + "<|chunk_eos|>": 8, + "<|chunk_tts_eos|>": 9, + "<|turn_eos|>": 10, + "<|audio|>": 11, + }.get(token, 0), + encode=lambda text, add_special_tokens=False: [], + ) + runtime.processor = SimpleNamespace(get_streaming_chunk_size=lambda: 4) + runtime.device = "cpu" + runtime._init_token_ids() + state = _MiniCPMO45Stage0SessionState( + session_id="sid-new-user-turn-prefill", + audio_chunk_idx=1, + pending_terminator_token=10, + last_terminator_token=10, + current_turn_ended=True, + ) + + result = runtime._stage_prefill_embeddings_only( + state, + np.zeros(4, dtype=np.float32), + seq=2, + ) + + assert result["success"] is True + assert result["input_token_ids"] == [10, 2, 1, 11] + assert state.pending_terminator_token is None + assert state.last_terminator_token == 10 + assert state.current_turn_ended is True + + +def test_minicpmo_stage0_data_plane_model_owned_turn_boundary_preserves_audio_cache(): + import torch + + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + _MiniCPMO45Stage0SessionState, + ) + + stale_cache = object() + + class _StageModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.embed = torch.nn.Embedding(256, 2) + self.audio_past_key_values = stale_cache + + def get_input_embeddings(self): + return self.embed + + def get_audio_hidden_states(self, data): + return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)] + + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.stage_model = _StageModel() + runtime.thinker = runtime.stage_model + runtime.tokenizer = SimpleNamespace( + unk_token_id=0, + convert_tokens_to_ids=lambda token: { + "": 1, + "": 2, + "<|listen|>": 3, + "<|speak|>": 4, + "<|tts_bos|>": 5, + "<|tts_eos|>": 6, + "<|tts_pad|>": 7, + "<|chunk_eos|>": 8, + "<|chunk_tts_eos|>": 9, + "<|turn_eos|>": 10, + "<|audio|>": 11, + }.get(token, 0), + encode=lambda text, add_special_tokens=False: [], + ) + runtime.processor = SimpleNamespace(get_streaming_chunk_size=lambda: 4) + runtime.device = "cpu" + runtime._init_token_ids() + state = _MiniCPMO45Stage0SessionState( + session_id="sid-new-user-turn-audio-cache", + audio_chunk_idx=1, + audio_past_key_values=stale_cache, + pending_terminator_token=8, + last_terminator_token=8, + current_turn_ended=True, + ) + + result = runtime._stage_prefill_embeddings_only( + state, + np.zeros(4, dtype=np.float32), + seq=2, + ) + + assert result["success"] is True + assert state.audio_past_key_values is stale_cache + assert runtime.thinker.audio_past_key_values is stale_cache + + +def test_minicpmo_stage0_data_plane_final_first_chunk_does_not_add_silence_unit(): + import torch + + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + _MiniCPMO45Stage0SessionState, + ) + + class _StageModel(torch.nn.Module): + first_chunk_ms = 10 + sample_rate = 1000 + + def __init__(self): + super().__init__() + self.seen_audio = None + self.embed = torch.nn.Embedding(256, 2) + + def get_input_embeddings(self): + return self.embed + + def get_audio_hidden_states(self, data): + self.seen_audio = np.asarray(data["audio_features"], dtype=np.float32) + return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)] + + class _MelProcessor: + sample_rate = 1000 + + def get_config(self): + return {"effective_first_chunk_ms": 10} + + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.stage_model = _StageModel() + runtime.thinker = runtime.stage_model + runtime.tokenizer = SimpleNamespace( + unk_token_id=0, + convert_tokens_to_ids=lambda token: { + "": 1, + "": 2, + "<|listen|>": 3, + "<|speak|>": 4, + "<|tts_bos|>": 5, + "<|tts_eos|>": 6, + "<|tts_pad|>": 7, + "<|chunk_eos|>": 8, + "<|chunk_tts_eos|>": 9, + "<|turn_eos|>": 10, + "<|audio|>": 11, + }.get(token, 0), + encode=lambda text, add_special_tokens=False: [], + ) + runtime.processor = SimpleNamespace( + _streaming_mel_processor=_MelProcessor(), + get_streaming_chunk_size=lambda: 10, + ) + runtime.device = "cpu" + runtime._init_token_ids() + state = _MiniCPMO45Stage0SessionState(session_id="sid-first-chunk-padding") + + result = runtime._stage_prefill_embeddings_only( + state, + np.arange(8, dtype=np.float32), + seq=1, + final=True, + ) + + assert result["success"] is True + assert result["input_token_ids"] == [1, 11] + assert runtime.stage_model.seen_audio is not None + np.testing.assert_allclose( + runtime.stage_model.seen_audio.reshape(-1), + np.array([0, 0, 0, 1, 2, 3, 4, 5, 6, 7], dtype=np.float32), + ) + + +def test_minicpmo_stage0_runtime_uses_loaded_vllm_embed_tokens_when_get_input_embeddings_is_broken(): + import torch + + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + ) + + class _Embed(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(128, 2)) + self.calls = [] + + def forward(self, input_ids): + ids = input_ids.reshape(-1).tolist() + self.calls.append(ids) + return torch.tensor([[float(i), 0.0] for i in ids], dtype=torch.float32) + + class _Thinker: + def __init__(self): + self.llm = SimpleNamespace(model=SimpleNamespace(embed_tokens=_Embed())) + + def get_input_embeddings(self, input_ids, multimodal_embeddings=None): + raise AttributeError("'Qwen3ForCausalLM' object has no attribute 'get_input_embeddings'") + + thinker = _Thinker() + stage_model = SimpleNamespace(model_stage="llm", thinker=thinker, processor=None) + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.stage_model = stage_model + runtime.thinker = thinker + runtime.device = "cpu" + + embeds = runtime._embed_token(11) + + assert embeds.shape == (1, 2) + assert thinker.llm.model.embed_tokens.calls == [[11]] + + +def test_minicpmo_remote_config_patch_handles_nested_and_dict_configs(): + from vllm_omni.experimental.fullduplex.minicpmo45.compat import ( + patch_minicpmo_remote_config, + ) + + nested = SimpleNamespace(base_model_tp_plan=None) + config = SimpleNamespace( + base_model_tp_plan=None, + text_config=nested, + tts_config={}, + ) + + patch_minicpmo_remote_config(config) + + assert config.base_model_tp_plan == {} + assert nested.base_model_tp_plan == {} + assert config.tts_config["top_p"] == 0.8 + assert config.tts_config["top_k"] == 100 + assert config.tts_config["temperature"] == 0.8 + assert config.tts_config["repetition_penalty"] == 1.05 + + +def test_minicpmo_stage0_short_audio_buffers_without_context_mutation(): + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + _MiniCPMO45Stage0SessionState, + ) + + class _Processor: + def get_streaming_chunk_size(self): + return 16000 + + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.stage_model = SimpleNamespace() + runtime.thinker = SimpleNamespace() + runtime.processor = _Processor() + runtime._require_special_token_ids = lambda: None + state = _MiniCPMO45Stage0SessionState(session_id="sid") + + result = runtime._stage_prefill_embeddings_only(state, np.zeros(1600, dtype=np.float32)) + + assert result["success"] is False + assert result["reason"] + assert len(state.audio_buffer) >= 1600 + assert state.context_embeds == [] + + +def test_minicpmo_stage0_native_sampler_penalizes_repeated_text_token(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Tokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids = [] + all_special_ids = [] + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) + vocab_size = 151723 + repeated = 198 + alternative = 1234 + logits = torch.full((1, vocab_size), -100.0) + logits[0, repeated] = 20.0 + logits[0, alternative] = 19.5 + sampling_metadata = SimpleNamespace( + all_greedy=False, + all_random=True, + temperature=torch.tensor([1.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[[repeated] * 8], + ) + + sampled = model.sample(logits, sampling_metadata) + + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[alternative]] + + +def test_minicpmo_stage0_native_sampler_penalizes_text_from_prior_chunk(): + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + _MiniCPMO45Stage0SessionState, + ) + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Tokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids = [] + all_special_ids = [] + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + 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]) + state.generated_text_tokens = [198] * 8 + model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state}) + model._minicpmo45_duplex_row_sessions = {0: session_key} + + vocab_size = 151723 + repeated = 198 + alternative = 1234 + logits = torch.full((1, vocab_size), -100.0) + logits[0, repeated] = 20.0 + logits[0, alternative] = 19.5 + sampling_metadata = SimpleNamespace( + all_greedy=False, + all_random=True, + temperature=torch.tensor([1.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[[]], + ) + + sampled = model.sample(logits, sampling_metadata) + + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[alternative]] + + +def test_minicpmo_stage0_native_sampler_matches_official_negative_logit_penalty(): + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + _MiniCPMO45Stage0SessionState, + ) + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Tokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids = [] + all_special_ids = [] + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + 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]) + repeated = 198 + alternative = 1234 + state.generated_text_tokens = [repeated] + model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state}) + model._minicpmo45_duplex_row_sessions = {0: session_key} + + logits = torch.full((1, 151723), -100.0) + logits[0, repeated] = -1.0 + logits[0, alternative] = -0.97 + sampling_metadata = SimpleNamespace( + all_greedy=False, + all_random=True, + temperature=torch.tensor([1.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[[]], + ) + + sampled = model.sample(logits, sampling_metadata) + + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[repeated]] + + +def test_minicpmo_stage0_records_only_bounded_text_repetition_history(): + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + _MiniCPMO45Stage0SessionState, + ) + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + session_key = ("sid-text-history", 0) + state = _MiniCPMO45Stage0SessionState(session_id=session_key[0]) + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state}) + model._minicpmo45_duplex_row_sessions = {0: session_key} + token_ids = { + "unit_token_id": 1, + "unit_end_token_id": 2, + "listen_token_id": 3, + "speak_token_id": 4, + "tts_bos_token_id": 5, + "tts_eos_token_id": 6, + "tts_pad_token_id": 7, + "chunk_eos_token_id": 8, + "chunk_tts_eos_token_id": 9, + "turn_eos_token_id": 10, + } + + for sampled in range(1000, 1520): + model._record_minicpmo45_duplex_terminator(0, sampled, token_ids) + + expected_history = list(range(1008, 1520)) + assert state.generated_text_tokens == expected_history + + for sampled in token_ids.values(): + model._record_minicpmo45_duplex_terminator(0, sampled, token_ids) + + assert state.generated_text_tokens == expected_history + + +def test_minicpmo_stage0_native_sampler_does_not_override_model_at_punctuation(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Tokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids = [] + all_special_ids = [] + + text = { + 200: "我", + 201: "喜", + 202: "欢", + 203: "。", + } + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + def decode(self, ids, skip_special_tokens=True): + del skip_special_tokens + return "".join(self.text.get(int(token_id), "") for token_id in ids) + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) + model.min_new_speak_tokens_before_chunk_boundary = 4 + model.max_new_speak_tokens_per_chunk = 64 + vocab_size = 151723 + alternative = 1234 + logits = torch.full((1, vocab_size), -100.0) + logits[0, alternative] = 20.0 + sampling_metadata = SimpleNamespace( + all_greedy=False, + all_random=True, + temperature=torch.tensor([1.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[[151706, 200, 201, 202, 203]], + ) + + sampled = model.sample(logits, sampling_metadata) + + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[alternative]] + + +def test_minicpmo_stage0_native_sampler_does_not_cut_before_natural_boundary_minimum(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Tokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids = [] + all_special_ids = [] + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + def decode(self, ids, skip_special_tokens=True): + del ids, skip_special_tokens + return "。" + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) + model.min_new_speak_tokens_before_chunk_boundary = 4 + model.max_new_speak_tokens_per_chunk = 64 + vocab_size = 151723 + alternative = 1234 + logits = torch.full((1, vocab_size), -100.0) + logits[0, alternative] = 20.0 + sampling_metadata = SimpleNamespace( + all_greedy=False, + all_random=True, + temperature=torch.tensor([1.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[[151706, 200, 201, 202]], + ) + + sampled = model.sample(logits, sampling_metadata) + + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[alternative]] + + +def test_minicpmo_stage0_native_sampler_does_not_rewrite_model_punctuation(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Tokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids = [] + all_special_ids = [] + + text = { + 108386: "你", + 104256: "好", + 3837: ",", + 100644: "今", + 99172: "天", + 100281: "想", + 27442: "聊", + 99217: "什", + 1773: "。", + 99218: "么", + } + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return {"。": [1773]}.get(text, []) + + def decode(self, ids, skip_special_tokens=True): + del skip_special_tokens + return "".join(self.text.get(int(token_id), "") for token_id in ids) + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) + model.min_new_speak_tokens_before_chunk_boundary = 4 + model.max_new_speak_tokens_per_chunk = 64 + vocab_size = 151723 + period = 1773 + continuation = 99218 + logits = torch.full((1, vocab_size), -100.0) + logits[0, period] = 30.0 + logits[0, continuation] = 20.0 + sampling_metadata = SimpleNamespace( + all_greedy=False, + all_random=True, + temperature=torch.tensor([1.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[[151706, 108386, 104256, 3837, 100644, 99172, 100281, 27442, 99217]], + ) + + sampled = model.sample(logits, sampling_metadata) + + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[period]] + + +def test_minicpmo_stage0_native_sampler_preserves_model_chunk_eos_decision(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Tokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids = [] + all_special_ids = [] + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) + logits = torch.full((1, 151723), -100.0) + logits[0, 151718] = 30.0 + logits[0, 1234] = 20.0 + sampling_metadata = SimpleNamespace( + all_greedy=True, + all_random=False, + temperature=torch.tensor([0.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[[151706, 200, 201]], + ) + + sampled = model.sample(logits, sampling_metadata) + + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[151718]] + + +def test_minicpmo_stage0_native_sampler_preserves_early_model_turn_eos_decision(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Tokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids = [] + all_special_ids = [] + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) + model.min_new_speak_tokens_before_chunk_boundary = 8 + logits = torch.full((1, 151723), -100.0) + logits[0, 151717] = 30.0 + logits[0, 1234] = 20.0 + sampling_metadata = SimpleNamespace( + all_greedy=True, + all_random=False, + temperature=torch.tensor([0.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[[151706, 200, 201]], + ) + + sampled = model.sample(logits, sampling_metadata) + + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[151717]] + + +def test_minicpmo_stage0_native_sampler_char_cap_does_not_override_special_token(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Tokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids = [] + all_special_ids = [151706, 151717] + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + def decode(self, ids, skip_special_tokens=True): + del skip_special_tokens + return "已经超过二十八个字符的当前语音分段文本" + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) + logits = torch.full((1, 151723), -100.0) + logits[0, 151717] = 30.0 + sampling_metadata = SimpleNamespace( + all_greedy=True, + all_random=False, + temperature=torch.tensor([0.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[[151706, 200, 201]], + ) + + sampled = model.sample(logits, sampling_metadata) + + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[151717]] + + +def test_minicpmo_stage0_native_sampler_preserves_model_chunk_eos(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Tokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids = [] + all_special_ids = [] + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) + model.min_new_speak_tokens_before_chunk_boundary = 8 + model.max_new_speak_tokens_per_chunk = 64 + vocab_size = 151723 + logits = torch.full((1, vocab_size), -100.0) + logits[0, 151718] = 30.0 + logits[0, 1234] = 20.0 + sampling_metadata = SimpleNamespace( + all_greedy=False, + all_random=True, + temperature=torch.tensor([1.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[[151706, 200, 201, 202]], + ) + + sampled = model.sample(logits, sampling_metadata) + + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[151718]] + + +def test_minicpmo_stage0_native_sampler_keeps_hard_chunk_cap(): + from vllm_omni.experimental.fullduplex.minicpmo45.policy import ( + MiniCPMO45DuplexPolicy, + ) + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Tokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids = [] + all_special_ids = [] + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + def decode(self, ids, skip_special_tokens=True): + del ids, skip_special_tokens + return "没有自然边界" + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) + vocab_size = 151723 + alternative = 1234 + logits = torch.full((1, vocab_size), -100.0) + logits[0, alternative] = 20.0 + sampling_metadata = SimpleNamespace( + all_greedy=False, + all_random=True, + temperature=torch.tensor([1.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[[200] * (MiniCPMO45DuplexPolicy.DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK - 1)], + ) + + sampled = model.sample(logits, sampling_metadata) + + assert MiniCPMO45DuplexPolicy.DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK == 20 + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[151718]] + + +def test_minicpmo_stage0_native_sampler_cuts_before_request_length_cap(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Tokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids = [] + all_special_ids = [] + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + def decode(self, ids, skip_special_tokens=True): + del ids, skip_special_tokens + return "没有自然边界" + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) + model.max_new_speak_tokens_per_chunk = 64 + model._minicpmo45_duplex_row_max_tokens = {0: 20} + vocab_size = 151723 + alternative = 1234 + logits = torch.full((1, vocab_size), -100.0) + logits[0, alternative] = 20.0 + sampling_metadata = SimpleNamespace( + all_greedy=False, + all_random=True, + temperature=torch.tensor([1.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[[151706] + [200] * 18], + ) + + sampled = model.sample(logits, sampling_metadata) + + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[151718]] + + +def test_minicpmo_stage0_native_sampler_cuts_on_decoded_text_length(): + from vllm_omni.experimental.fullduplex.minicpmo45.policy import ( + MiniCPMO45DuplexPolicy, + ) + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Tokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids = [] + all_special_ids = [151706, 151718, 151717, 151705] + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + def decode(self, ids, skip_special_tokens=True): + token_text = { + 200: "一二三四五六七八九十", + 201: "十一十二十三十四十五", + 202: "十六十七十八十九", + } + special = set(self.all_special_ids) if skip_special_tokens else set() + return "".join(token_text.get(int(token_id), "") for token_id in ids if int(token_id) not in special) + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) + vocab_size = 151723 + candidate = 202 + logits = torch.full((1, vocab_size), -100.0) + logits[0, candidate] = 20.0 + sampling_metadata = SimpleNamespace( + all_greedy=True, + all_random=False, + temperature=torch.tensor([0.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[[151706, 200, 201]], + ) + + sampled = model.sample(logits, sampling_metadata) + + assert MiniCPMO45DuplexPolicy.DEFAULT_MAX_SPEAK_CHARS_PER_CHUNK == 28 + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[151718]] + + +def test_minicpmo_stage0_native_sampler_ignores_pending_placeholders(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Tokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids = [] + all_special_ids = [] + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) + vocab_size = 151723 + newline = 198 + alternative = 1234 + logits = torch.full((1, vocab_size), -100.0) + logits[0, newline] = 20.0 + logits[0, alternative] = 19.5 + sampling_metadata = SimpleNamespace( + all_greedy=False, + all_random=True, + temperature=torch.tensor([1.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[[-1, -1, -1]], + ) + + sampled = model.sample(logits, sampling_metadata) + + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[newline]] + + +def test_minicpmo_stage0_native_sampler_converts_mid_turn_listen_to_tts_bos(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Tokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids = [] + all_special_ids = [] + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + state = SimpleNamespace(current_turn_ended=False) + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) + model._minicpmo45_duplex_row_sessions = {0: ("sid-native", 0)} + model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={("sid-native", 0): state}) + vocab_size = 151723 + logits = torch.full((1, vocab_size), -100.0) + logits[0, 151705] = 30.0 + sampling_metadata = SimpleNamespace( + all_greedy=True, + all_random=False, + temperature=torch.tensor([0.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[[]], + ) + + sampled = model.sample(logits, sampling_metadata) + + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[151703]] + assert state.current_turn_ended is False + + +def test_minicpmo_stage0_native_sampler_forced_listen_yields_floor(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + class _Tokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids = [] + all_special_ids = [] + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + state = SimpleNamespace(current_turn_ended=False) + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer()) + model._minicpmo45_duplex_row_sessions = {0: ("sid-native", 0)} + model._minicpmo45_duplex_row_payloads = {0: {"force_listen": True}} + model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={("sid-native", 0): state}) + vocab_size = 151723 + logits = torch.full((1, vocab_size), -100.0) + logits[0, 151705] = 30.0 + sampling_metadata = SimpleNamespace( + all_greedy=True, + all_random=False, + temperature=torch.tensor([0.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[[]], + ) + + sampled = model.sample(logits, sampling_metadata) + + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[151705]] + assert state.current_turn_ended is True + + +def test_minicpmo_stage0_native_sampler_uses_runner_duplex_rows(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + metadata = SimpleNamespace( + prompt_token_ids=torch.tensor([[1, 2, 3]]), + ) + + rows = model._minicpmo45_native_duplex_prompt_rows( + metadata, + unit_id=151683, + batch_size=1, + duplex_rows=[0], + ) + + assert rows == [0] + + +def test_minicpmo_stage0_session_context_includes_resolved_ref_audio(): + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + _MiniCPMO45Stage0SessionState, + ) + + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.unit_token_id = 151683 + runtime.processor = SimpleNamespace() + runtime.stage_model = SimpleNamespace() + runtime.thinker = SimpleNamespace() + runtime.device = "cpu" + token_map = { + "<|im_start|>system\nUse speech.\n<|audio_start|>": [1, 2, 3], + "<|audio_end|><|im_end|>": [4, 5], + } + runtime._stage_runtime_ready = lambda: True + runtime._require_special_token_ids = lambda: None + runtime._decode_ref_audio_from_session_config = lambda _config: np.array([0.1, -0.1], dtype=np.float32) + runtime._encode_text = lambda text: token_map[text] + runtime._embed_token = lambda token_id: torch.full((1, 2), float(token_id)) + runtime._stage_ref_audio_embeddings = lambda ref_audio, state=None: torch.tensor([[10.0, 11.0], [12.0, 13.0]]) + + state = _MiniCPMO45Stage0SessionState(session_id="sid-ref") + runtime._prepare_session_context(state, {"instructions": "Use speech.", "extra_body": {"ref_audio_data": "x"}}) + + assert state.context_token_ids == [1, 2, 3, 151683, 151683, 4, 5] + assert len(state.context_embeds) == 6 + + +def test_minicpmo_tts_ref_audio_lru_evicts_matching_vocoder_base_cache(tmp_path): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + MiniCPMO45OmniTTSForConditionalGeneration, + MiniCPMO45TTSRuntimeConfig, + ) + + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + model._runtime_config = MiniCPMO45TTSRuntimeConfig(ref_audio_file_cache_size=1) + model._t2w_base_caches = {} + model._token2wav_state_lock = threading.RLock() + first_path = model._write_ref_audio_prompt_wav(np.array([0.1, -0.1], dtype=np.float32), 16_000) + assert first_path is not None + model._t2w_base_caches[first_path] = (torch.tensor([1]), {"cache": torch.tensor([2])}) + + second_path = model._write_ref_audio_prompt_wav(np.array([0.2, -0.2], dtype=np.float32), 16_000) + + assert second_path is not None + assert second_path != first_path + assert first_path not in model._t2w_base_caches + assert not os.path.exists(first_path) + + +def test_minicpmo_tts_prompt_wav_path_does_not_default_to_model_asset(tmp_path): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + MiniCPMO45OmniTTSForConditionalGeneration, + ) + + model_dir = tmp_path / "MiniCPM-o-4_5" + assets_dir = model_dir / "assets" + assets_dir.mkdir(parents=True) + (assets_dir / "HT_ref_audio.wav").write_bytes(b"placeholder") + + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + model.vllm_config = SimpleNamespace(model_config=SimpleNamespace(model=str(model_dir))) + model._write_ref_audio_prompt_wav = lambda ref_audio, ref_audio_sr: None + + assert model._resolve_prompt_wav_path(None, None) == (None, None) + + +def test_minicpmo_tts_stream_turn_id_prefers_duplex_identity(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + MiniCPMO45OmniTTSForConditionalGeneration, + ) + + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + + assert ( + model._stream_turn_id( + { + "duplex": {"turn_id": 0, "model_turn_id": 4}, + "meta": {"turn_id": 3}, + } + ) + == 4 + ) + assert ( + model._stream_turn_id( + { + "duplex": {"turn_id": 0}, + "meta": {"turn_id": 3}, + } + ) + == 0 + ) + assert model._stream_turn_id({"meta": {"turn_id": 3}, "turn_id": 2}) == 3 + assert model._stream_turn_id({"turn_id": 2}) == 2 + + +def test_minicpmo_tts_no_ref_audio_initializes_empty_vocoder_cache(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + MiniCPMO45OmniTTSForConditionalGeneration, + _TalkerTurnState, + ) + + class AudioTokenizer: + def set_stream_cache(self, prompt_wav_path): + raise AssertionError(f"unexpected reference-audio cache load: {prompt_wav_path!r}") + + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + model._t2w_base_caches = {} + model._token2wav_state_lock = threading.RLock() + model.audio_tokenizer = AudioTokenizer() + state = _TalkerTurnState(None, None) + + model._begin_turn_vocoder_cache(None, state=state) + + assert state.stream_cache is None + assert state.hift_cache_dict == {} + assert state.vocoder_initialized is True + assert model._t2w_base_caches[""] == (None, {}) + + +def test_minicpmo_tts_native_duplex_merges_and_left_pads_nonfinal_unit_audio(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + _native_duplex_unit_waveform, + ) + + waveform = _native_duplex_unit_waveform( + [torch.tensor([1.0, 2.0]), torch.tensor([3.0])], + turn_end=False, + target_samples=6, + ) + + assert waveform is not None + assert waveform.tolist() == [0.0, 0.0, 0.0, 1.0, 2.0, 3.0] + + +def test_minicpmo_tts_native_duplex_does_not_pad_final_unit_audio(): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + _native_duplex_unit_waveform, + ) + + waveform = _native_duplex_unit_waveform( + [torch.tensor([1.0, 2.0]), torch.tensor([3.0])], + turn_end=True, + target_samples=6, + ) + + assert waveform is not None + assert waveform.tolist() == [1.0, 2.0, 3.0] + + +def test_minicpmo_tts_turn_cleanup_removes_deleted_ref_audio_vocoder_cache(tmp_path): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + MiniCPMO45OmniTTSForConditionalGeneration, + _TalkerTurnState, + ) + + prompt_path = tmp_path / "session-ref.wav" + prompt_path.write_bytes(b"wav") + model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) + model._talker_turn_states = { + "request": _TalkerTurnState(str(prompt_path), str(prompt_path)), + } + model._talker_consumed_tokens = {"request": 1} + model._talker_request_keys = {} + model._t2w_base_caches = { + str(prompt_path): (torch.tensor([1]), {"cache": torch.tensor([2])}), + } + model._token2wav_state_lock = threading.RLock() + model.audio_tokenizer = None + + assert model._close_turn_state("request") is True + assert str(prompt_path) not in model._t2w_base_caches + assert not prompt_path.exists() diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index 48c0bc6de..cff36b9f6 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -196,8 +196,7 @@ def _make_runner_for_mimo(req_id="r_mimo"): def test_talker_mtp_forward_cpu_updates_inputs_and_info(monkeypatch): - # `_talker_mtp_forward` calls `current_omni_platform.set_forward_context`, - # which would otherwise dispatch to the real device implementation. + # `_talker_mtp_forward` dispatches through the active platform. import vllm_omni.worker.gpu_model_runner as mod # Must be the same module that defines OmniGPUModelRunner monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context) @@ -459,28 +458,33 @@ def test_update_intermediate_buffer_skips_unknown_req_id(): assert "unknown_req" not in runner.model_intermediate_buffer -def test_maybe_run_batch_preprocess_calls_model_hook(): - runner = object.__new__(OmniGPUModelRunner) - runner.model_intermediate_buffer = {"r1": {"text": ["hello"]}} - calls = [] +def test_streaming_input_update_merges_model_intermediate_buffer(): + runner = _make_runner(req_ids=("r1",), hidden_size=4) + runner.model_intermediate_buffer["r1"] = { + "duplex": { + "session_id": "sid", + "seq": 1, + } + } + runner.requests["r1"].additional_information_cpu = runner.model_intermediate_buffer["r1"] + new_req_data = SimpleNamespace( + model_intermediate_buffer={ + "duplex": { + "session_id": "sid", + "seq": 2, + "payload": {"type": "audio"}, + } + }, + additional_information=None, + ) - class DummyModel: - def preprocess_batch(self, *, req_ids, model_intermediate_buffer, device): - calls.append((req_ids, model_intermediate_buffer, device)) + OmniGPUModelRunner._update_streaming_input_additional_info(runner, new_req_data, "r1") - runner.model = DummyModel() - - OmniGPUModelRunner._maybe_run_batch_preprocess(runner, ["r1"], torch.device("cpu")) - - assert calls == [(["r1"], runner.model_intermediate_buffer, torch.device("cpu"))] - - -def test_maybe_run_batch_preprocess_skips_missing_hook(): - runner = object.__new__(OmniGPUModelRunner) - runner.model_intermediate_buffer = {} - runner.model = object() - - OmniGPUModelRunner._maybe_run_batch_preprocess(runner, ["r1"], torch.device("cpu")) + info = runner.model_intermediate_buffer["r1"] + assert info["duplex"]["session_id"] == "sid" + assert info["duplex"]["seq"] == 2 + assert info["duplex"]["payload"] == {"type": "audio"} + assert runner.requests["r1"].additional_information_cpu is info def _make_full_payload_accumulation_runner( @@ -501,8 +505,6 @@ def _make_full_payload_accumulation_runner( runner._custom_process_func = object() runner._pending_full_payload_send = {} runner._stage_id = 1 - # Non-None sentinel: the gate short-circuits to False when no connector - # is configured at all (terminal stages in pipelines with no connector). runner._omni_connector = object() return runner @@ -519,7 +521,7 @@ def test_accumulate_full_payload_output_preserves_aligned_all_zero_qwen3_omni_co def test_accumulate_full_payload_output_keeps_misaligned_all_zero_qwen3_omni_codec_rows(): - # After removing the sender-side zero filter, the full-payload accumulator keeps every + # After removing the sender-side zero filter, the accumulator keeps every # codec row including misaligned all-zero rows. The downstream consumer # (_extract_qwen3_full_payload_codec_rows) is the authoritative crop and # filters by output_token_ids. @@ -566,45 +568,25 @@ def test_accumulate_full_payload_output_keeps_all_zero_qwen3_omni_prefill_placeh def test_full_payload_output_accumulation_hook_matrix(): - """Producer-side gate: fires iff an explicit next-stage payload hook is loaded. - - A derived `*_full_payload` helper from `custom_process_input_func` is not - enough: terminal/input-only consumer stages must not enqueue orphan - downstream payloads. - """ - # Thinker / talker producer stages: explicit next-stage payload hook -> gate fires. assert _make_full_payload_accumulation_runner(model_stage="thinker")._should_accumulate_full_payload_output() assert _make_full_payload_accumulation_runner(model_stage="talker")._should_accumulate_full_payload_output() - - # Terminal stage: even if _load_custom_func derived a builder from - # custom_process_input_func, final output stages are not producers. - runner = _make_full_payload_accumulation_runner(model_stage="code2wav", final_output=True) - assert not runner._should_accumulate_full_payload_output() - - # Input-only consumer stage without an explicit producer hook must not - # accumulate/send just because a same-module *_full_payload helper exists. - runner = _make_full_payload_accumulation_runner( + assert not _make_full_payload_accumulation_runner( + model_stage="code2wav", final_output=True + )._should_accumulate_full_payload_output() + assert not _make_full_payload_accumulation_runner( model_stage="token2audio", custom_process_next_stage_input_func=None, - ) - assert not runner._should_accumulate_full_payload_output() - - # async_chunk mode -> gate off. + )._should_accumulate_full_payload_output() assert not _make_full_payload_accumulation_runner( model_stage="talker", async_chunk=True )._should_accumulate_full_payload_output() - - # Non-qwen3 arches: gate is arch-agnostic, but if the fixture's arch - # does not configure a connector payload builder, its runtime - # `_custom_process_func` is None. Emulate that. - runner = _make_full_payload_accumulation_runner(model_arch="Qwen3TTSForConditionalGeneration") - runner._custom_process_func = None - runner._should_accumulate_full_payload_output_cached = None - assert not runner._should_accumulate_full_payload_output() - runner = _make_full_payload_accumulation_runner(model_arch="Qwen2_5OmniForConditionalGeneration") - runner._custom_process_func = None - runner._should_accumulate_full_payload_output_cached = None - assert not runner._should_accumulate_full_payload_output() + for model_arch in ( + "Qwen3TTSForConditionalGeneration", + "Qwen2_5OmniForConditionalGeneration", + ): + runner = _make_full_payload_accumulation_runner(model_arch=model_arch) + runner._custom_process_func = None + assert not runner._should_accumulate_full_payload_output() def test_sync_local_stage_payloads_retains_payload_until_request_is_active(): diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 367e1b317..5e16140b8 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -270,6 +270,15 @@ class PipelineConfig: # multi-component repos (e.g. GLM-Image). ``None`` = not a diffusers model. diffusers_class_name: str | None = None endpoint_restrictions: tuple[EndpointRestriction, ...] = () + # Optional model-owned duplex planner loaded by the stable engine runtime. + 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. default_deploy_config_name: str | None = None @@ -406,6 +415,38 @@ class StageDeployConfig: engine_extras: dict[str, Any] = field(default_factory=dict) +@dataclass(frozen=True) +class DuplexSessionRuntimeConfig: + """Server-owned lifecycle and per-session buffering limits.""" + + idle_ttl_s: float | None = 300.0 + disconnect_grace_s: float = 30.0 + reaper_interval_s: float = 5.0 + resume_replay_ttl_s: float = 60.0 + resume_replay_max_bytes_per_session: int = 8 * 1024 * 1024 + max_pending_input_bytes_per_session: int = 16 * 1024 * 1024 + max_pending_turns_per_session: int = 4 + max_sessions: int = 1 + completed_append_cache_size: int = 256 + + def __post_init__(self) -> None: + positive = { + "disconnect_grace_s": self.disconnect_grace_s, + "reaper_interval_s": self.reaper_interval_s, + "resume_replay_ttl_s": self.resume_replay_ttl_s, + "resume_replay_max_bytes_per_session": self.resume_replay_max_bytes_per_session, + "max_pending_input_bytes_per_session": self.max_pending_input_bytes_per_session, + "max_pending_turns_per_session": self.max_pending_turns_per_session, + "max_sessions": self.max_sessions, + "completed_append_cache_size": self.completed_append_cache_size, + } + if self.idle_ttl_s is not None and self.idle_ttl_s <= 0: + raise ValueError("duplex_session.idle_ttl_s must be positive or null") + for name, value in positive.items(): + if value <= 0: + raise ValueError(f"duplex_session.{name} must be positive") + + @dataclass class DeployConfig: """Loaded from deploy/.yaml — the only config file users edit. @@ -419,8 +460,10 @@ class DeployConfig: """ async_chunk: bool = True + session_mode: str = "turn" # Stage-1 active stream slots; 0 preserves legacy all-stream cycling. active_stream_window: int = 0 + duplex_session: DuplexSessionRuntimeConfig = field(default_factory=DuplexSessionRuntimeConfig) connectors: dict[str, Any] | None = None edges: list[dict[str, Any]] | None = None stages: list[StageDeployConfig] = field(default_factory=list) @@ -610,7 +653,9 @@ def load_deploy_config(path: str | Path) -> DeployConfig: kwargs: dict[str, Any] = { "async_chunk": raw_dict.get("async_chunk", True), + "session_mode": raw_dict.get("session_mode", "turn"), "active_stream_window": int(raw_dict.get("active_stream_window", 0) or 0), + "duplex_session": DuplexSessionRuntimeConfig(**(raw_dict.get("duplex_session") or {})), "connectors": raw_dict.get("connectors", None), "edges": raw_dict.get("edges", None), "stages": stages, @@ -896,6 +941,7 @@ def merge_pipeline_deploy( StageConfig( stage_id=ps.stage_id, model_stage=ps.model_stage, + session_mode=deploy.session_mode, stage_type=stage_type, input_sources=list(ps.input_sources), custom_process_input_func=input_proc, @@ -922,6 +968,7 @@ class StageConfig: stage_id: int model_stage: str + session_mode: str = "turn" stage_type: StageType = StageType.LLM input_sources: list[int] = field(default_factory=list) custom_process_input_func: str | None = None @@ -984,6 +1031,7 @@ class StageConfig: config_dict: dict[str, Any] = { "stage_id": self.stage_id, "stage_type": StageType(self.stage_type).value, + "session_mode": self.session_mode, "engine_args": create_config(engine_args), "runtime": create_config(runtime), "engine_input_source": self.input_sources, # Legacy field name diff --git a/vllm_omni/core/sched/__init__.py b/vllm_omni/core/sched/__init__.py index 0f80769d3..ce968893d 100644 --- a/vllm_omni/core/sched/__init__.py +++ b/vllm_omni/core/sched/__init__.py @@ -2,9 +2,28 @@ Scheduling components for vLLM-Omni. """ -from .omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler -from .omni_generation_scheduler import OmniGenerationScheduler -from .output import OmniNewRequestData +from __future__ import annotations + +from importlib import import_module +from typing import Any + +_LAZY_EXPORTS = { + "OmniARAsyncScheduler": (".omni_ar_scheduler", "OmniARAsyncScheduler"), + "OmniARScheduler": (".omni_ar_scheduler", "OmniARScheduler"), + "OmniGenerationScheduler": (".omni_generation_scheduler", "OmniGenerationScheduler"), + "OmniNewRequestData": (".output", "OmniNewRequestData"), +} + + +def __getattr__(name: str) -> Any: + target = _LAZY_EXPORTS.get(name) + if target is None: + raise AttributeError(name) + module_name, attribute_name = target + value = getattr(import_module(module_name, __name__), attribute_name) + globals()[name] = value + return value + __all__ = [ "OmniARAsyncScheduler", diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 912b67bd7..e88730353 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -259,6 +259,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None), prompt_is_token_ids=nr.prompt_is_token_ids, additional_information=(getattr(request, "additional_information", None) if request else None), + model_intermediate_buffer=( + getattr(request, "model_intermediate_buffer", None) if request else None + ), ) new_list.append(omni_nr) @@ -718,6 +721,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): session.record_event(EngineCoreEventType.QUEUED) return super()._update_request_as_session(session, update) + if hasattr(update, "model_intermediate_buffer"): + session.model_intermediate_buffer = update.model_intermediate_buffer def _free_request(self, request: Request, delay_free_blocks: bool = False) -> dict[str, Any] | None: # TODO(wzliu)! for offline mode, we should not end process until all data is transferred diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 77a07be7d..c6f82ec23 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -319,6 +319,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None), prompt_is_token_ids=nr.prompt_is_token_ids, additional_information=(getattr(request, "additional_information", None) if request else None), + model_intermediate_buffer=( + getattr(request, "model_intermediate_buffer", None) if request else None + ), ) new_list.append(omni_nr) diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index fc7e1b626..4eaa07cbb 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -40,16 +40,16 @@ except ValueError: class OmniSchedulerMixin: """Shared scheduler helpers for omni-specific request handling.""" + # ------------------------------------------------------------------ # + # Shared scheduler/output helpers (lift the AR / generation duplicates) + # ------------------------------------------------------------------ # + def _free_input_coordinator_request(self, request_id: str) -> None: """Prune full-payload coordinator state for a completed request.""" input_coordinator = getattr(self, "input_coordinator", None) if input_coordinator is not None: input_coordinator.free_finished_request(request_id) - # ------------------------------------------------------------------ # - # Shared scheduler/output helpers (lift the AR / generation duplicates) - # ------------------------------------------------------------------ # - def _consume_pending_connector_output(self, model_mode: str) -> None: """Drain ``self._latest_omni_connector_output`` into the coordinator. diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index ec017b11c..be9c867b1 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -20,10 +20,13 @@ class OmniNewRequestData(NewRequestData): external_req_id: Optional external request ID for tracking additional_information: Optional serialized additional information dictionary containing tensors or lists + model_intermediate_buffer: Optional runner-owned payload for + GPUModelRunner.model_intermediate_buffer """ external_req_id: str | None = None additional_information: AdditionalInformationPayload | None = None + model_intermediate_buffer: dict[str, object] | None = None @classmethod def from_request( @@ -56,6 +59,7 @@ class OmniNewRequestData(NewRequestData): prompt_is_token_ids=getattr(request, "prompt_is_token_ids", None), prefill_token_ids=prefill_token_ids, additional_information=getattr(request, "additional_information", None), + model_intermediate_buffer=getattr(request, "model_intermediate_buffer", None), ) diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml new file mode 100644 index 000000000..17cd014e9 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml @@ -0,0 +1,17 @@ +# Experimental validation profile: one thinker and two Stage1 replicas. +# This is not a recommended production entrypoint. +# +# Stage 0: GPU 0 +# Stage 1: replicas on GPU 1 and 2 +base_config: minicpmo_4_5.yaml +pipeline: minicpmo_4_5 +session_mode: duplex + +stages: + - stage_id: 0 + async_scheduling: false + + - stage_id: 1 + devices: "1,2" + num_replicas: 2 + async_scheduling: false diff --git a/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml b/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml new file mode 100644 index 000000000..e4fe45922 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml @@ -0,0 +1,17 @@ +# Experimental validation profile: one thinker and three Stage1 replicas. +# This is not a recommended production entrypoint. +# +# Stage 0: GPU 0 +# Stage 1: replicas on GPU 1, 2, and 3 +base_config: minicpmo_4_5.yaml +pipeline: minicpmo_4_5 +session_mode: duplex + +stages: + - stage_id: 0 + async_scheduling: false + + - stage_id: 1 + devices: "1,2,3" + num_replicas: 3 + async_scheduling: false diff --git a/vllm_omni/deploy/minicpmo_4_5_8x4090_stage1_replicas.yaml b/vllm_omni/deploy/minicpmo_4_5_8x4090_stage1_replicas.yaml new file mode 100644 index 000000000..f4ad2c6b4 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_8x4090_stage1_replicas.yaml @@ -0,0 +1,20 @@ +# Experimental validation profile: 4-way thinker TP and four Stage1 replicas. +# This is not a recommended production entrypoint. +# +# Stage 0: GPU 0, 1, 2, and 3 +# Stage 1: replicas on GPU 4, 5, 6, and 7 +base_config: minicpmo_4_5_8x4090.yaml +pipeline: minicpmo_4_5 +session_mode: duplex + +stages: + - stage_id: 0 + async_scheduling: false + gpu_memory_utilization: 0.40 + max_num_batched_tokens: 16384 + + - stage_id: 1 + devices: "4,5,6,7" + num_replicas: 4 + async_scheduling: false + gpu_memory_utilization: 0.8 diff --git a/vllm_omni/deploy/minicpmo_4_5_duplex.yaml b/vllm_omni/deploy/minicpmo_4_5_duplex.yaml new file mode 100644 index 000000000..c227f2817 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_duplex.yaml @@ -0,0 +1,36 @@ +# MiniCPM-o 4.5 native duplex deployment. +# +# The registered minicpmo_4_5 pipeline owns the fixed Stage0 -> Stage1 +# topology and llm2tts bridge. This overlay only changes session lifetime and +# the Stage1 token budget needed to drain streaming audio segments. +base_config: minicpmo_4_5.yaml +pipeline: minicpmo_4_5 +session_mode: duplex +active_stream_window: 1 +duplex_session: + idle_ttl_s: 300 + disconnect_grace_s: 30 + reaper_interval_s: 5 + resume_replay_ttl_s: 60 + resume_replay_max_bytes_per_session: 8388608 + max_pending_input_bytes_per_session: 16777216 + max_pending_turns_per_session: 4 + # Matches the configured Stage0/Stage1 sequence capacity below. + max_sessions: 2 + completed_append_cache_size: 256 + +stages: + - stage_id: 0 + max_num_seqs: 2 + async_scheduling: false + + - stage_id: 1 + max_num_seqs: 2 + async_scheduling: false + # The streaming talker advances Python-owned per-request generators on + # every step, so only fixed-topology model regions are graph-capturable. + compilation_config: + cudagraph_mode: PIECEWISE + default_sampling_params: + max_tokens: 4096 + stop_token_ids: [151645] diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index 04368842c..26be84012 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -771,8 +771,12 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.waiting_for_chunk_running_requests = deque( request for request in self.waiting_for_chunk_running_requests if request.request_id not in request_ids ) + self._held_non_active = deque( + request for request in self._held_non_active if request.request_id not in request_ids + ) for req_id in request_ids: + self._active_streams.pop(req_id, None) self.requests_with_ready_chunks.discard(req_id) self.finished_requests.discard(req_id) self._finished_load_reqs.discard(req_id) diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index 89bb41908..e72ce9fed 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -75,6 +75,10 @@ class OmniEngineCoreRequest(EngineCoreRequest): # Optional additional information dictionary (serialized) additional_information: AdditionalInformationPayload | None = None + # Runner-owned runtime payload. This is materialized directly into + # GPUModelRunner.model_intermediate_buffer instead of using the deprecated + # additional_information request transport. + model_intermediate_buffer: dict[str, Any] | None = None @classmethod def from_request( @@ -83,6 +87,7 @@ class OmniEngineCoreRequest(EngineCoreRequest): *, prompt_embeds: torch.Tensor | None = None, additional_information: AdditionalInformationPayload | None = None, + model_intermediate_buffer: dict[str, Any] | None = None, ) -> "OmniEngineCoreRequest": """Clone an EngineCoreRequest into an OmniEngineCoreRequest with optional payload overrides.""" @@ -90,6 +95,8 @@ class OmniEngineCoreRequest(EngineCoreRequest): prompt_embeds = request.prompt_embeds if additional_information is None: additional_information = getattr(request, "additional_information", None) + if model_intermediate_buffer is None: + model_intermediate_buffer = getattr(request, "model_intermediate_buffer", None) return cls( request_id=request.request_id, @@ -113,6 +120,7 @@ class OmniEngineCoreRequest(EngineCoreRequest): reasoning_parser_kwargs=request.reasoning_parser_kwargs, abort_immediately=request.abort_immediately, additional_information=additional_information, + model_intermediate_buffer=model_intermediate_buffer, ) diff --git a/vllm_omni/engine/async_engine_utils.py b/vllm_omni/engine/async_engine_utils.py new file mode 100644 index 000000000..58f8351b0 --- /dev/null +++ b/vllm_omni/engine/async_engine_utils.py @@ -0,0 +1,160 @@ +"""Stateless request and shutdown helpers for :mod:`async_omni_engine`.""" + +from __future__ import annotations + +import threading +from typing import Any + +import janus +import torch +from vllm.logger import init_logger +from vllm.v1.engine import EngineCoreRequest + +from vllm_omni.engine import OmniEngineCoreRequest +from vllm_omni.engine.messages import EngineQueueMessage, ShutdownRequestMessage +from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient +from vllm_omni.engine.serialization import ( + deserialize_additional_information, + serialize_additional_information, +) +from vllm_omni.engine.stage_runtime import StageRuntime + +logger = init_logger(__name__) + +SHUTDOWN_ENQUEUE_TIMEOUT_S = 1.0 +SHUTDOWN_JOIN_TIMEOUT_S = 30.0 +_WEAK_SHUTDOWN_JOIN_TIMEOUT_S = 1.0 + + +def inject_global_id(target: Any, request_id: str) -> None: + """Inject ``global_request_id`` into a prompt's additional information.""" + if isinstance(target, dict): + if "additional_information" not in target: + target["additional_information"] = {} + if target["additional_information"] is None: + target["additional_information"] = {} + if isinstance(target["additional_information"], dict): + target["additional_information"]["global_request_id"] = [str(request_id)] + + +def upgrade_to_omni_request( + request: EngineCoreRequest, + raw_prompt: Any, +) -> EngineCoreRequest: + """Restore omni-only fields omitted by the upstream input processor.""" + prompt_embeds = request.prompt_embeds + additional_information = None + wire_payload = None + model_intermediate_buffer = None + + if isinstance(raw_prompt, dict): + if prompt_embeds is None: + raw_prompt_embeds = raw_prompt.get("prompt_embeds") + if isinstance(raw_prompt_embeds, torch.Tensor): + prompt_embeds = raw_prompt_embeds + raw_info = raw_prompt.get("additional_information") + raw_buffer = raw_prompt.get("model_intermediate_buffer") + if isinstance(raw_info, dict): + wire_payload = dict(raw_info) + if isinstance(raw_buffer, dict): + model_intermediate_buffer = raw_buffer + additional_information = serialize_additional_information( + wire_payload, + log_prefix="AsyncOmniEngine", + ) + + if prompt_embeds is None and additional_information is None and model_intermediate_buffer is None: + return request + + return OmniEngineCoreRequest.from_request( + request, + prompt_embeds=prompt_embeds, + additional_information=additional_information, + model_intermediate_buffer=model_intermediate_buffer, + ) + + +def apply_omni_final_stage_metadata( + request: EngineCoreRequest, + final_stage_id: int, +) -> EngineCoreRequest: + """Tag a request so the AR scheduler can skip DiT KV for stage zero.""" + merged: dict[str, Any] = {} + if isinstance(request, OmniEngineCoreRequest) and request.additional_information is not None: + merged = deserialize_additional_information(request.additional_information) + merged["omni_final_stage_id"] = final_stage_id + payload = serialize_additional_information(merged) + return OmniEngineCoreRequest.from_request( + request, + additional_information=payload, + ) + + +def weak_shutdown_async_omni_engine( + orchestrator_thread: threading.Thread | None, + request_queue: janus.Queue[EngineQueueMessage] | None, + output_queue: janus.Queue[EngineQueueMessage] | None, + rpc_output_queue: janus.Queue[EngineQueueMessage] | None, + rpc_client: CorrelatedRpcClient | None, +) -> None: + """Best-effort orchestrator cleanup for garbage-collection finalization.""" + request_queue_closed = False + shutdown_enqueued = enqueue_orchestrator_shutdown( + request_queue, + timeout=SHUTDOWN_ENQUEUE_TIMEOUT_S, + ) + if request_queue is not None and not shutdown_enqueued: + try: + request_queue.close() + request_queue_closed = True + except Exception: + pass + + if rpc_client is not None: + rpc_client.close() + + try: + if orchestrator_thread is not None and orchestrator_thread.is_alive(): + orchestrator_thread.join(timeout=_WEAK_SHUTDOWN_JOIN_TIMEOUT_S) + except Exception: + pass + + for q in (request_queue, output_queue, rpc_output_queue): + try: + if q is not None and not (q is request_queue and request_queue_closed): + q.close() + except Exception: + pass + + +def enqueue_orchestrator_shutdown( + request_queue: janus.Queue[EngineQueueMessage] | None, + *, + timeout: float, +) -> bool: + """Deliver shutdown without waiting forever behind a full request queue.""" + if request_queue is None: + return False + try: + request_queue.sync_q.put(ShutdownRequestMessage(), timeout=timeout) + except Exception: + return False + return True + + +def shutdown_runtime_after_orchestrator( + orchestrator_thread: threading.Thread, + runtime: StageRuntime, +) -> None: + """Release a stage runtime once its orchestrator can no longer use it.""" + try: + orchestrator_thread.join() + except Exception: + logger.exception("[AsyncOmniEngine] Failed while awaiting deferred Orchestrator shutdown") + if orchestrator_thread.is_alive(): + logger.error("[AsyncOmniEngine] Orchestrator is still alive; StageRuntime cleanup remains deferred") + return + try: + runtime.shutdown() + except Exception: + logger.exception("[AsyncOmniEngine] Failed to shutdown deferred StageRuntime") diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 459057486..dbebd7e6d 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -18,7 +18,7 @@ import uuid import weakref from collections.abc import Mapping, Sequence from dataclasses import asdict -from typing import Any, Literal, cast +from typing import TYPE_CHECKING, Any, Literal, cast import janus import torch @@ -31,10 +31,23 @@ 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.stage_config import strip_parent_engine_args +from vllm_omni.config.stage_config import ( + DuplexSessionRuntimeConfig, + load_deploy_config, + strip_parent_engine_args, +) from vllm_omni.diffusion.data import DiffusionParallelConfig, parse_attention_config from vllm_omni.diffusion.diffusion_engine import supports_audio_output -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, + shutdown_runtime_after_orchestrator, + upgrade_to_omni_request, + weak_shutdown_async_omni_engine, +) from vllm_omni.engine.messages import ( AbortRequestMessage, AddCompanionRequestMessage, @@ -42,14 +55,10 @@ from vllm_omni.engine.messages import ( CollectiveRPCResultMessage, EngineQueueMessage, ErrorMessage, - ShutdownRequestMessage, StageSubmissionMessage, ) from vllm_omni.engine.orchestrator import Orchestrator -from vllm_omni.engine.serialization import ( - deserialize_additional_information, - serialize_additional_information, -) +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 @@ -67,10 +76,13 @@ from vllm_omni.metrics.prometheus import OmniRequestCounter logger = init_logger(__name__) +if TYPE_CHECKING: + from vllm_omni.experimental.fullduplex.engine.duplex_control_client import DuplexControlClient + from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseActivity + from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence + _STARTUP_POLL_INTERVAL_S = 1.0 _REQUEST_QUEUE_MAXSIZE = 256 - - # ============================================================================ # Parent-EngineArgs field-routing contracts (consumed by # AsyncOmniEngine._strip_parent_engine_args when ``stage_configs_path`` is set). @@ -106,88 +118,6 @@ _PARENT_ARGS_STRIP: frozenset[str] = frozenset({"stage_configs_path"}) _PARENT_ARGS_NO_WARN: frozenset[str] = frozenset({"model"}) -def _inject_global_id(target: Any, request_id: str) -> None: - """Inject global_request_id into a prompt dict's additional_information.""" - if isinstance(target, dict): - if "additional_information" not in target: - target["additional_information"] = {} - if target["additional_information"] is None: - target["additional_information"] = {} - if isinstance(target["additional_information"], dict): - target["additional_information"]["global_request_id"] = [str(request_id)] - - -def _upgrade_to_omni_request( - request: EngineCoreRequest, - raw_prompt: Any, -) -> EngineCoreRequest: - """Restore omni-only fields omitted by upstream InputProcessor.""" - prompt_embeds = request.prompt_embeds - additional_information = None - - if isinstance(raw_prompt, dict): - if prompt_embeds is None: - raw_prompt_embeds = raw_prompt.get("prompt_embeds") - if isinstance(raw_prompt_embeds, torch.Tensor): - prompt_embeds = raw_prompt_embeds - additional_information = serialize_additional_information( - raw_prompt.get("additional_information"), - log_prefix="AsyncOmniEngine", - ) - - if prompt_embeds is None and additional_information is None: - return request - - return OmniEngineCoreRequest.from_request( - request, - prompt_embeds=prompt_embeds, - additional_information=additional_information, - ) - - -def _apply_omni_final_stage_metadata( - request: EngineCoreRequest, - final_stage_id: int, -) -> EngineCoreRequest: - """Tag EngineCoreRequest so OmniARScheduler can skip DiT KV when final_stage_id is 0.""" - merged: dict[str, Any] = {} - if isinstance(request, OmniEngineCoreRequest) and request.additional_information is not None: - merged = deserialize_additional_information(request.additional_information) - merged["omni_final_stage_id"] = final_stage_id - payload = serialize_additional_information(merged) - return OmniEngineCoreRequest.from_request( - request, - additional_information=payload, - ) - - -def _weak_shutdown_async_omni_engine( - orchestrator_thread: threading.Thread | None, - request_queue: janus.Queue[EngineQueueMessage] | None, - output_queue: janus.Queue[EngineQueueMessage] | None, - rpc_output_queue: janus.Queue[EngineQueueMessage] | None, -) -> None: - """Best-effort orchestrator cleanup for GC finalization.""" - try: - if request_queue is not None: - request_queue.sync_q.put_nowait(ShutdownRequestMessage()) - except Exception: - pass - - try: - if orchestrator_thread is not None and orchestrator_thread.is_alive(): - orchestrator_thread.join() - except Exception: - pass - - for q in (request_queue, output_queue, rpc_output_queue): - try: - if q is not None: - q.close() - except Exception: - pass - - class AsyncOmniEngine: """Thin proxy that launches an Orchestrator in a background thread. @@ -282,11 +212,22 @@ class AsyncOmniEngine: # restriction beforehand. TODO (Alex) make this cleaner and refactor # stage config resolution to remove kwargs hacks. deploy_config_path = kwargs.get("deploy_config") - self.endpoint_restrictions = StageConfigFactory.get_pipeline_endpoint_restrictions( + pipeline_config = StageConfigFactory.get_pipeline_config( model=model, trust_remote_code=trust_remote_code, deploy_config_path=deploy_config_path, ) + 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 deploy_config_path is not None: + self.duplex_session_config = load_deploy_config(deploy_config_path).duplex_session kwargs["trust_remote_code"] = trust_remote_code self.config_path, self.stage_configs = self._resolve_stage_configs( @@ -315,7 +256,8 @@ class AsyncOmniEngine: self.rpc_output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() self._shutdown_called = False self._weak_finalizer: weakref.finalize | None = None - self._rpc_lock = threading.Lock() + self._correlated_rpc_client: CorrelatedRpcClient | None = None + self._duplex_control_client: DuplexControlClient | None = None self._running_counter = OmniRequestCounter() logger.info(f"[AsyncOmniEngine] Launching Orchestrator thread with {self.num_stages} stages") @@ -334,15 +276,20 @@ class AsyncOmniEngine: ) 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, + 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") @@ -435,6 +382,21 @@ class AsyncOmniEngine: pd_config = self._detect_pd_config() membership_controller = self._runtime.create_membership_controller() + duplex_runtime_extension = None + if self._duplex_control_enabled: + from vllm_omni.experimental.fullduplex.engine.duplex_runtime import ( + load_duplex_runtime_extension, + validate_duplex_runtime_extension, + ) + + duplex_runtime_extension = load_duplex_runtime_extension( + getattr(self, "_duplex_runtime_extension_path", None) + ) + if duplex_runtime_extension is not None: + validate_duplex_runtime_extension( + duplex_runtime_extension, + sampling_defaults=tuple(pool.stage_client.default_sampling_params for pool in self.stage_pools), + ) orchestrator = Orchestrator( request_async_queue=self.request_queue.async_q, @@ -448,6 +410,9 @@ class AsyncOmniEngine: transfer_emitter=self._transfer_emitter, 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()) @@ -718,10 +683,10 @@ class AsyncOmniEngine: if stage_type != "diffusion" and not isinstance(prompt, EngineCoreRequest): # Inject global_request_id into the raw prompt. if isinstance(prompt, dict): - _inject_global_id(prompt, request_id) + inject_global_id(prompt, request_id) elif isinstance(prompt, list): for item in prompt: - _inject_global_id(item, request_id) + inject_global_id(item, request_id) preselected_stage0_replica = self._scope_stage0_multimodal_cache_to_replica( request_id, @@ -751,7 +716,7 @@ class AsyncOmniEngine: _preprocess_ms = (time.perf_counter() - _t_preprocess) * 1000.0 # TODO (Peiqi): add this for Qwen3-TTS only. Other models don't have # additional_information field in the prompt. - request = _upgrade_to_omni_request(request, prompt) + request = upgrade_to_omni_request(request, prompt) if reasoning_ended is not None: request.reasoning_ended = reasoning_ended @@ -763,7 +728,7 @@ class AsyncOmniEngine: # to match the key used in Orchestrator.request_states so that # output routing (output.request_id lookup) can find the req_state. request.external_req_id = request_id - request = _apply_omni_final_stage_metadata(request, final_stage_id) + request = apply_omni_final_stage_metadata(request, final_stage_id) # Registration with stage 0's output processor is deferred to the # orchestrator thread (see Orchestrator._handle_add_request), which @@ -811,7 +776,7 @@ class AsyncOmniEngine: companion_params, companion_spl = ep.apply_overrides(stage0_params, sampling_params_list) if isinstance(companion_prompt, dict): - _inject_global_id(companion_prompt, cid) + inject_global_id(companion_prompt, cid) request = self.input_processor.process_inputs( request_id=cid, @@ -1434,6 +1399,268 @@ 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: + from vllm_omni.experimental.fullduplex.engine.duplex_control_client import 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: @@ -1489,36 +1716,19 @@ class AsyncOmniEngine: stage_ids=stage_ids, ) - with self._rpc_lock: - self.request_queue.sync_q.put(msg) - deadline = None if timeout is None else time.monotonic() + timeout - - while True: - remaining = None if deadline is None else max(0.0, deadline - time.monotonic()) - try: - result_msg = self.rpc_output_queue.sync_q.get(timeout=remaining) - except queue.Empty as exc: - raise TimeoutError(f"collective_rpc timed out after {timeout} seconds") from exc - - if isinstance(result_msg, ErrorMessage): - raise RuntimeError(result_msg.error) - - if not isinstance(result_msg, CollectiveRPCResultMessage): - logger.warning( - "[AsyncOmniEngine] Dropping unexpected rpc queue message type=%s", - getattr(result_msg, "type", type(result_msg).__name__), - ) - continue - - if result_msg.rpc_id != rpc_id: - logger.warning( - "[AsyncOmniEngine] Dropping mismatched rpc result rpc_id=%s expected=%s", - result_msg.rpc_id, - rpc_id, - ) - continue - - return list(result_msg.results) + 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, @@ -1555,22 +1765,61 @@ class AsyncOmniEngine: finalizer.detach() logger.info("[AsyncOmniEngine] Shutting down Orchestrator") - if self.request_queue is not None: - self.request_queue.sync_q.put_nowait(ShutdownRequestMessage()) - if self.is_alive(): - self.orchestrator_thread.join() + 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: - q.close() + if not (q is self.request_queue and request_queue_closed): + q.close() except Exception: pass - if hasattr(self, "_runtime") and self._runtime is not 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 diff --git a/vllm_omni/engine/membership_controller.py b/vllm_omni/engine/membership_controller.py index 48b6fe559..a7676de54 100644 --- a/vllm_omni/engine/membership_controller.py +++ b/vllm_omni/engine/membership_controller.py @@ -157,7 +157,7 @@ class MembershipController: input_addr = StagePool._client_input_addr(client) if input_addr is None: raise RuntimeError(f"remote replica factory for stage {stage_id} produced a client without input address") - pool.add_client(input_addr, client) + pool.add_client(input_addr, client, replica_id=replica_id) logger.info( "[MembershipController] attached remote replica stage=%d replica=%d addr=%s", stage_id, diff --git a/vllm_omni/engine/messages.py b/vllm_omni/engine/messages.py index bb73f74a3..384d8a055 100644 --- a/vllm_omni/engine/messages.py +++ b/vllm_omni/engine/messages.py @@ -106,3 +106,7 @@ class CollectiveRPCResultMessage(EngineQueueMessage, kw_only=True): method: str stage_ids: list[int] results: list[object] + + @property + def rpc_correlation_key(self) -> tuple[str, str]: + return ("collective", self.rpc_id) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index bdd678f27..9372fee28 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -13,8 +13,9 @@ from __future__ import annotations import asyncio import time as _time +from collections.abc import Callable from dataclasses import dataclass, field -from typing import Any +from typing import TYPE_CHECKING, Any import janus import torch @@ -22,11 +23,12 @@ from vllm.config import ModelConfig from vllm.logger import init_logger from vllm.outputs import CompletionOutput, RequestOutput from vllm.pooling_params import PoolingParams -from vllm.sampling_params import SamplingParams +from vllm.sampling_params import RequestOutputKind, SamplingParams 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.engine import OmniEngineCoreRequest from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.engine.membership_controller import MembershipController @@ -53,6 +55,23 @@ from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) +if TYPE_CHECKING: + from vllm_omni.experimental.fullduplex.engine.contracts import ( + DuplexControlPlanePort, + DuplexOutputContext, + DuplexOutputDecision, + DuplexRequestIdentity, + DuplexRuntimeExtension, + DuplexStageRequestContext, + DuplexStageSubmission, + DuplexStageSubmissionResult, + ) + from vllm_omni.experimental.fullduplex.engine.duplex_session import ( + DuplexSessionRuntimeManager, + DuplexSessionRuntimeState, + ) + from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence + def _build_terminal_empty_output( request_id: str, @@ -141,8 +160,13 @@ def build_engine_core_request_from_tokens( pooling_params = params.clone() prompt_embeds: torch.Tensor | None = prompt.get("prompt_embeds") + raw_additional_information = prompt.get("additional_information") + model_intermediate_buffer = prompt.get("model_intermediate_buffer") + wire_payload: dict[str, Any] | None = None + if isinstance(raw_additional_information, dict): + wire_payload = dict(raw_additional_information) additional_info_payload = serialize_additional_information( - prompt.get("additional_information"), + wire_payload, log_prefix=f"build_engine_core_request_from_tokens req={request_id}", ) @@ -159,6 +183,7 @@ def build_engine_core_request_from_tokens( prompt_embeds=prompt_embeds, resumable=resumable, additional_information=additional_info_payload, + model_intermediate_buffer=model_intermediate_buffer if isinstance(model_intermediate_buffer, dict) else None, ) @@ -186,6 +211,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 + running_counter_registered: bool = False @dataclass @@ -194,10 +223,124 @@ class StreamingInputState: enabled: bool = False # Flag of segment of streaming input finished segment_finished: bool = False + # Tokens from the current raw segment boundary. The vLLM output processor + # does not guarantee that EngineCoreOutput.new_token_ids survives on the + # processed RequestOutput used by the routing layer. + segment_token_ids: list[int] = field(default_factory=list) + segment_output_metadata: dict[str, Any] = field(default_factory=dict) # Streaming update prompt length new_prompt_len_snapshot: int | None = None # Model/bridge-specific runtime states (e.g., thinker->talker) bridge_states: dict[str, Any] = field(default_factory=dict) + # Synchronous stage-transition capability installed by the orchestrator + # while the downstream input processor consumes upstream token output. + source_token_decoder: Callable[..., str] | None = None + + +class _OrchestratorDuplexStagePort: + """Adapts generic stage pools to the model-neutral duplex control plane.""" + + def __init__( + self, + *, + stage_pools: list[StagePool], + request_states: dict[str, OrchestratorRequestState], + running_counter: OmniRequestCounter | None, + cleanup_request_ids: Callable[..., Any], + ) -> None: + self._stage_pools = stage_pools + self._request_states = request_states + self._running_counter = running_counter + self._cleanup_request_ids = cleanup_request_ids + + @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") + if not isinstance(duplex_state.get("model_turn_id"), int) or previous_epoch != context.fence.epoch: + 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: + from vllm_omni.experimental.fullduplex.engine.contracts import DuplexRequestIdentity + + 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: + from vllm_omni.experimental.fullduplex.engine.contracts import 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) + 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: @@ -208,6 +351,7 @@ class Orchestrator: _running_counter: OmniRequestCounter | None = None _transfer_emitter: Any = None _stat_logger: OmniPrometheusStatLogger | None = None + duplex_control_plane: DuplexControlPlanePort | None = None def __init__( self, @@ -223,6 +367,9 @@ class Orchestrator: transfer_emitter: 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 @@ -236,7 +383,7 @@ class Orchestrator: replica_sampler=self._sample_replica_metrics, ) for stage_id, pool in enumerate(self.stage_pools): - for replica_id in pool.live_replica_ids(): + for replica_id in pool.available_replica_ids(): self._orch_monitor.register_replica(stage_id, replica_id) # PD disaggregation state @@ -254,6 +401,33 @@ 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: + from vllm_omni.experimental.fullduplex.engine.duplex_control_plane import DuplexControlPlane + from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseConfig + + 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, + ), + 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._fatal_error: str | None = None @@ -327,6 +501,17 @@ 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") @@ -347,6 +532,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")) if membership_watcher is not None: tasks.append(membership_watcher) @@ -361,6 +548,15 @@ class Orchestrator: # During teardown this is expected; the finally block handles # proper cleanup. Do not re-raise. logger.info("[Orchestrator] Engine dead during shutdown: %s", e) + if self._fatal_error is None: + self._fatal_error = str(e) or "Stage engine died" + await self.rpc_async_queue.put( + ErrorMessage( + error=self._fatal_error or str(e), + fatal=True, + stage_id=self._fatal_error_stage_id, + ) + ) except Exception: logger.exception("[Orchestrator] Fatal error in orchestrator tasks") raise @@ -373,6 +569,8 @@ 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() if self._fatal_error is not None: await self._drain_pending_requests_on_fatal() @@ -405,6 +603,8 @@ class Orchestrator: 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": await self._handle_abort(msg) elif msg_type == "collective_rpc": @@ -432,6 +632,21 @@ 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") + async def _handle_add_request(self, msg: StageSubmissionMessage) -> None: """Handle an add_request message from the main thread.""" stage_id = 0 @@ -466,8 +681,7 @@ class Orchestrator: mm_features=getattr(prompt, "mm_features", None), ) self.request_states[request_id] = req_state - if self._running_counter is not None: - self._running_counter.increment() + 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 @@ -590,6 +804,7 @@ class Orchestrator: return for pool in self.stage_pools: await pool.abort_requests(request_ids) + pool.release_bindings(request_ids) def _release_request_bindings(self, request_ids: list[str]) -> None: """Release all stage-local route bindings for the given request ids.""" @@ -662,7 +877,7 @@ class Orchestrator: idle = True for stage_id in range(self.num_stages): pool = self.stage_pools[stage_id] - for replica_id in pool.live_replica_ids(): + for replica_id in pool.available_replica_ids(): if self._shutdown_event.is_set(): return @@ -686,6 +901,17 @@ class Orchestrator: if req_state is None or not req_state.streaming.enabled: continue req_state.streaming.segment_finished = bool(getattr(eco, "is_segment_finished", False)) + req_state.streaming.segment_token_ids = ( + self._coerce_int_list(getattr(eco, "new_token_ids", None)) + if req_state.streaming.segment_finished + else [] + ) + raw_mm = self._completion_multimodal_output(eco, None) + req_state.streaming.segment_output_metadata = ( + dict(raw_mm) + if req_state.streaming.segment_finished and isinstance(raw_mm, dict) + else {} + ) req_state.streaming.new_prompt_len_snapshot = getattr( eco, "new_prompt_len_snapshot", @@ -715,32 +941,62 @@ class Orchestrator: raise except EngineDeadError as e: logger.error( - "[Orchestrator] Stage-%s is dead: %s", + "[Orchestrator] Stage-%s replica-%s is dead: %s", stage_id, + replica_id, e, ) - # TODO: Fault handling is intentionally fail-stop at - # the orchestrator level today. If one replica in a - # logical stage dies, we promote it to `_fatal_error`, - # notify requests already admitted to that stage, and - # re-raise so `run()` shuts down all stages. This is - # conservative but means a single unhealthy replica in - # a multi-replica deployment can take down otherwise - # healthy replicas in other stages. Revisit this when - # adding per-replica fault isolation / eviction. - self._fatal_error = str(e) - self._fatal_error_stage_id = stage_id - for req_id, req_state in list(self.request_states.items()): - if stage_id in req_state.stage_submit_ts: + affected_request_ids = pool.mark_replica_unavailable(replica_id) + closed_sessions = ( + self.duplex_control_plane.close_sessions_for_request_ids( + affected_request_ids, + abort=False, + ) + if self.duplex_control_plane is not None + else {} + ) + for session_id, stale_request_ids in closed_sessions.items(): + affected_request_ids.extend(stale_request_ids) + logger.warning( + "[Orchestrator] closed duplex session %s after stage-%s replica-%s died; " + "stale_request_ids=%s", + session_id, + stage_id, + replica_id, + stale_request_ids, + ) + affected_request_ids = list(dict.fromkeys(affected_request_ids)) + if pool.available_replica_ids(): + for req_id in affected_request_ids: await self.output_async_queue.put( ErrorMessage( error=str(e), - fatal=True, + fatal=False, request_id=req_id, stage_id=stage_id, ) ) - self.request_states.pop(req_id, None) + await self._cleanup_request_ids( + affected_request_ids, + close_duplex_sessions=True, + ) + continue + + self._fatal_error = str(e) + self._fatal_error_stage_id = stage_id + for req_id in affected_request_ids: + await self.output_async_queue.put( + ErrorMessage( + error=str(e), + fatal=True, + request_id=req_id, + stage_id=stage_id, + ) + ) + await self._cleanup_request_ids( + affected_request_ids, + close_duplex_sessions=True, + ) self._shutdown_event.set() raise except Exception: @@ -781,7 +1037,8 @@ class Orchestrator: continue stage_metrics = None - if output.finished: + segment_finished = req_state.streaming.enabled and req_state.streaming.segment_finished + if output.finished or segment_finished: stage_metrics = pool.build_stage_metrics( [output], submit_ts=req_state.stage_submit_ts.get(stage_id, _time.time()), @@ -811,22 +1068,56 @@ class Orchestrator: await self._cleanup_request_ids( [parent_id, *self._cfg_tracker.cleanup_parent(parent_id)], abort=True, + close_duplex_sessions=True, ) # ---- Shared helpers ---- - async def _cleanup_request_ids(self, request_ids: list[str], *, abort: bool = False) -> None: + async def _cleanup_request_ids( + self, + request_ids: list[str], + *, + abort: bool = False, + close_duplex_sessions: bool = False, + ) -> None: """Release pool bindings and logical request state for the given ids.""" if not request_ids: return - if abort: - await self._abort_request_ids(request_ids) - self._release_request_bindings(request_ids) - for request_id in request_ids: - self._pd_kv_params.pop(request_id, None) - if self.request_states.pop(request_id, None) is not None and self._running_counter is not None: - self._running_counter.decrement() + cleanup_ids = list(dict.fromkeys(request_ids)) + 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)) + + try: + if abort: + 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 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) async def _apply_raw_terminal_stage_finish( self, @@ -873,6 +1164,34 @@ 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 + self._running_counter.increment() + req_state.running_counter_registered = True + async def _route_output( self, stage_id: int, @@ -885,7 +1204,6 @@ class Orchestrator: req_id = output.request_id finished = output.finished submit_ts = req_state.stage_submit_ts.get(stage_id) - # CFG companion: stash output so parent can bundle [parent, *companions] # into source_outputs for the bridge (e.g. thinker2imagegen). if finished and self._cfg_tracker.is_companion(req_id): @@ -899,7 +1217,17 @@ 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) - if self.stage_pools[stage_id].final_output: + # 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_finished + ) + if self.stage_pools[stage_id].final_output and not is_duplex_stage0_segment: await self.output_async_queue.put( OutputMessage( request_id=req_id, @@ -907,7 +1235,10 @@ class Orchestrator: replica_id=replica_id, engine_outputs=output, metrics=stage_metrics, - finished=request_finished, + finished=( + request_finished + or (self._is_duplex_session_request(req_state) and req_state.streaming.segment_finished) + ), stage_submit_ts=submit_ts, ) ) @@ -928,6 +1259,18 @@ 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, + ) + return + if ( (finished or (req_state.streaming.enabled and req_state.streaming.segment_finished)) and stage_id < req_state.final_stage_id @@ -941,6 +1284,12 @@ class Orchestrator: ): self._cfg_tracker.defer_parent(req_id, output, stage_id) else: + stage_params = req_state.sampling_params_list[stage_id] + final_only_finished = ( + req_state.streaming.enabled + and finished + and getattr(stage_params, "output_kind", None) == RequestOutputKind.FINAL_ONLY + ) await self._forward_to_next_stage( req_id, stage_id, @@ -948,9 +1297,14 @@ class Orchestrator: req_state, src_replica_id=replica_id, is_streaming_session=req_state.streaming.enabled, - is_final_update=False, + is_final_update=final_only_finished, ) - if req_state.streaming.enabled and finished: + if ( + req_state.streaming.enabled + and finished + and not final_only_finished + and not self._is_duplex_session_request(req_state) + ): # For streaming sessions, send the terminal (resumable=False) update only on a finish await self._forward_to_next_stage( req_id, @@ -962,7 +1316,7 @@ class Orchestrator: is_final_update=True, ) - if request_finished: + if request_finished and not self._is_duplex_session_request(req_state): 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: @@ -1039,6 +1393,137 @@ class Orchestrator: request.external_req_id = req_id return request + @staticmethod + def _duplex_output_context( + req_state: OrchestratorRequestState, + *, + stage_id: int | None = None, + ) -> DuplexOutputContext | None: + identity = req_state.duplex_identity + if identity is None: + return None + from vllm_omni.experimental.fullduplex.engine.contracts import ( + DuplexOutputContext, + DuplexRequestIdentity, + ) + + fence = req_state.duplex_stage_fences.get(stage_id, identity.fence) if stage_id is not None else identity.fence + 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 req_state.streaming.segment_finished, + segment_token_ids=tuple(req_state.streaming.segment_token_ids), + segment_output_metadata=req_state.streaming.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}") + from vllm_omni.experimental.fullduplex.output import attach_duplex_output_decision + + engine_output = attach_duplex_output_decision( + OmniRequestOutput( + request_id=req_id, + finished=True, + stage_id=stage_id, + final_output_type=decision.final_output_type, + request_output=output, + ), + 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) + if isinstance(mm_output, dict): + return mm_output + mm_output = getattr(completion, "multimodal_output", None) if completion is not None else None + return mm_output if isinstance(mm_output, dict) else {} + + @classmethod + def _coerce_int_list(cls, value: Any) -> 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 [] + out: list[int] = [] + for item in value: + token_id = cls._coerce_int(item) + if token_id is not None: + out.append(token_id) + return out + + @staticmethod + def _coerce_int(value: Any) -> int | None: + if hasattr(value, "detach"): + try: + value = value.detach().cpu().reshape(-1) + if value.numel() == 0: + return None + value = value[0].item() + except Exception: + return None + try: + return int(value) + except (TypeError, ValueError): + return None + async def _handle_cfg_companion_ready(self, req_id: str) -> None: """Mark a CFG companion as done; if all companions are done, flush deferred parent.""" parent_id = self._cfg_tracker.on_companion_completed(req_id) @@ -1293,9 +1778,9 @@ class Orchestrator: diffusion_prompt = req_state.prompt if already_submitted: - await next_pool.submit_update(req_id, req_state, diffusion_prompt) + replica_id = await next_pool.submit_update(req_id, req_state, diffusion_prompt) else: - await next_pool.submit_initial( + replica_id = await next_pool.submit_initial( req_id, req_state, diffusion_prompt, @@ -1307,6 +1792,12 @@ class Orchestrator: }, params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params), ) + self._record_duplex_stage_submission( + next_logical, + req_id, + replica_id, + req_state, + ) req_state.stage_submit_ts[next_logical] = _time.time() _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 self._emit_tx_edge( @@ -1351,9 +1842,15 @@ class Orchestrator: ) request.external_req_id = request.request_id if already_submitted: - await next_pool.submit_update(req_id, req_state, request) + replica_id = await next_pool.submit_update(req_id, req_state, request) else: - await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) + replica_id = await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) + self._record_duplex_stage_submission( + next_logical, + req_id, + replica_id, + req_state, + ) req_state.stage_submit_ts[next_logical] = _time.time() _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 @@ -1373,6 +1870,13 @@ class Orchestrator: {}, )[req_id] = req_state.pd_prefill_multimodal_output + previous_decoder = req_state.streaming.source_token_decoder + source_processor = self.stage_pools[src_stage_id].output_processor + tokenizer = getattr(source_processor, "tokenizer", None) + decode = getattr(tokenizer, "decode", None) + if callable(decode): + req_state.streaming.source_token_decoder = decode + try: next_inputs = next_client.process_engine_inputs( source_outputs, @@ -1386,6 +1890,8 @@ class Orchestrator: next_logical, ) raise + finally: + req_state.streaming.source_token_decoder = previous_decoder if not next_inputs: if not getattr(output, "finished", False): @@ -1446,9 +1952,15 @@ class Orchestrator: ) if already_submitted: - await next_pool.submit_update(req_id, req_state, request) + replica_id = await next_pool.submit_update(req_id, req_state, request) else: - await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) + replica_id = await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) + self._record_duplex_stage_submission( + next_logical, + req_id, + replica_id, + req_state, + ) req_state.stage_submit_ts[next_logical] = _time.time() _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 diff --git a/vllm_omni/engine/rpc_result_router.py b/vllm_omni/engine/rpc_result_router.py new file mode 100644 index 000000000..516b36f16 --- /dev/null +++ b/vllm_omni/engine/rpc_result_router.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import queue +import threading +from typing import TypeAlias + +from vllm.logger import init_logger + +from vllm_omni.engine.messages import ( + EngineQueueMessage, + ErrorMessage, +) + +logger = init_logger(__name__) + +RpcCorrelationKey: TypeAlias = tuple[str, str] +RpcWaiter: TypeAlias = queue.Queue[EngineQueueMessage] + + +class RpcResultRouter: + """Single consumer that dispatches shared RPC results by correlation ID.""" + + def __init__(self, source_queue) -> None: + self._source_queue = source_queue + self._pending: dict[RpcCorrelationKey, RpcWaiter] = {} + self._lock = threading.Lock() + self._stopped = threading.Event() + self._terminal_error: ErrorMessage | None = None + self._thread = threading.Thread( + target=self._run, + daemon=True, + name="omni-rpc-result-router", + ) + self._thread.start() + + def register(self, key: RpcCorrelationKey) -> RpcWaiter: + waiter: RpcWaiter = queue.Queue(maxsize=1) + with self._lock: + if self._stopped.is_set(): + raise RuntimeError("RPC result router is closed") + if self._terminal_error is not None: + raise RuntimeError(self._terminal_error.error) + if key in self._pending: + raise RuntimeError(f"duplicate pending RPC correlation key: {key!r}") + self._pending[key] = waiter + return waiter + + def unregister(self, key: RpcCorrelationKey, waiter: RpcWaiter) -> None: + with self._lock: + if self._pending.get(key) is waiter: + self._pending.pop(key, None) + + def close(self) -> None: + if self._stopped.is_set(): + return + self._stopped.set() + self._broadcast_error(ErrorMessage(error="RPC result router closed", fatal=True)) + self._thread.join(timeout=1.0) + + def _run(self) -> None: + while not self._stopped.is_set(): + try: + message = self._source_queue.get(timeout=0.1) + except queue.Empty: + continue + except Exception as exc: + if not self._stopped.is_set(): + logger.exception("RPC result router source queue failed") + self._broadcast_error(ErrorMessage(error=str(exc), fatal=True)) + return + + if isinstance(message, ErrorMessage): + if message.fatal: + self._broadcast_error(message) + else: + logger.warning( + "Dropping uncorrelated non-fatal RPC error request_id=%s stage_id=%s: %s", + message.request_id, + message.stage_id, + message.error, + ) + continue + key = self._correlation_key(message) + if key is None: + logger.warning( + "Dropping unexpected RPC result message type=%s", + getattr(message, "type", type(message).__name__), + ) + continue + with self._lock: + waiter = self._pending.pop(key, None) + if waiter is None: + logger.warning("Dropping late or unknown RPC result correlation_key=%s", key) + continue + waiter.put_nowait(message) + + def _broadcast_error(self, message: ErrorMessage) -> None: + with self._lock: + if message.fatal: + self._terminal_error = message + waiters = list(self._pending.values()) + self._pending.clear() + for waiter in waiters: + try: + waiter.put_nowait(message) + except queue.Full: + pass + + @staticmethod + def _correlation_key(message: EngineQueueMessage) -> RpcCorrelationKey | None: + key = getattr(message, "rpc_correlation_key", None) + if not isinstance(key, tuple) or len(key) != 2 or not all(isinstance(part, str) and part for part in key): + return None + return key + + +class CorrelatedRpcClient: + """Own request submission and correlated result waiting as one lifecycle.""" + + def __init__(self, request_queue, result_queue) -> None: + self._request_queue = request_queue + self._router = RpcResultRouter(result_queue) + + def execute( + self, + key: RpcCorrelationKey, + message: EngineQueueMessage, + *, + timeout: float | None, + timeout_message: str, + block_on_submit: bool = False, + ) -> EngineQueueMessage: + waiter = self._router.register(key) + try: + if block_on_submit: + self._request_queue.put(message) + else: + self._request_queue.put_nowait(message) + try: + result = waiter.get(timeout=timeout) + except queue.Empty as exc: + raise TimeoutError(timeout_message) from exc + if isinstance(result, ErrorMessage): + raise RuntimeError(result.error) + return result + finally: + self._router.unregister(key, waiter) + + def close(self) -> None: + self._router.close() + + +__all__ = [ + "CorrelatedRpcClient", + "RpcCorrelationKey", + "RpcResultRouter", + "RpcWaiter", +] diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 34af459ab..4c4d3b7ee 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -97,6 +97,7 @@ class StagePool: self._stage_vllm_config = stage_vllm_config self._next_replica_id = 0 self._request_bindings: dict[str, int] = {} + self._unavailable_replicas: set[int] = set() self._replica_metrics: list[_ReplicaMetrics] = [_ReplicaMetrics() for _ in self.clients] self._output_timestamps_by_request: dict[str, list[float]] = {} self._non_empty_first_output_timestamps_by_request: dict[str, float] = {} @@ -206,27 +207,41 @@ class StagePool: return addr return None - def add_client(self, input_addr: str, client: Any) -> int: + def add_client(self, input_addr: str, client: Any, *, replica_id: int | None = None) -> int: """Register a head-side client for ``input_addr``. Returns the assigned ``replica_id`` (index into :attr:`clients`). - If the address is already known, replaces the existing client and - returns its existing id (this should not happen in practice — the - master server assigns unique slots — but the contract is idempotent - to keep the dispatch layer robust). + A coordinator-provided ``replica_id`` restores that exact slot after + unregister/re-register. Without one, a known address keeps its slot + and a new address takes the first vacant slot. """ if not input_addr: raise ValueError("input_addr must be a non-empty string") + if replica_id is not None and replica_id < 0: + raise ValueError("replica_id must be non-negative") existing = self._addr_to_replica_id.get(input_addr) - if existing is not None: - self.clients[existing] = client - return existing + if replica_id is not None and existing is not None and existing != replica_id: + raise ValueError( + f"input address {input_addr!r} is already assigned to replica {existing}, not {replica_id}" + ) + if replica_id is None: + replica_id = existing + if replica_id is None: + replica_id = next( + (index for index, existing_client in enumerate(self.clients) if existing_client is None), + len(self.clients), + ) - replica_id = len(self.clients) - self.clients.append(client) + while len(self.clients) <= replica_id: + self.clients.append(None) + self._replica_metrics.append(_ReplicaMetrics()) + for known_addr, known_replica_id in list(self._addr_to_replica_id.items()): + if known_replica_id == replica_id and known_addr != input_addr: + self._addr_to_replica_id.pop(known_addr, None) + self.clients[replica_id] = client self._addr_to_replica_id[input_addr] = replica_id - self._replica_metrics.append(_ReplicaMetrics()) + self._unavailable_replicas.discard(replica_id) return replica_id def remove_client(self, input_addr: str) -> Any | None: @@ -234,7 +249,7 @@ class StagePool: Slot is marked ``None`` to preserve indices for outstanding bindings. """ - replica_id = self._addr_to_replica_id.pop(input_addr, None) + replica_id = self._addr_to_replica_id.get(input_addr) if replica_id is None: return None client = self.clients[replica_id] @@ -475,6 +490,39 @@ class StagePool: for request_id in request_ids: self.release_binding(request_id) + def release_replica_bindings(self, replica_id: int) -> list[str]: + """Drop all route/session bindings owned by one physical replica.""" + released_request_ids = [ + request_id + for request_id, bound_replica_id in list(self._request_bindings.items()) + if bound_replica_id == replica_id + ] + released_request_ids.extend( + request_id + for request_id, input_addr in list(self._affinity.items()) + if self._addr_to_replica_id.get(input_addr) == replica_id + ) + released_request_ids = list(dict.fromkeys(released_request_ids)) + for request_id in released_request_ids: + self.release_binding(request_id) + return released_request_ids + + def mark_replica_unavailable(self, replica_id: int) -> list[str]: + """Evict a failed replica from admission and release its bindings.""" + if 0 <= replica_id < self.num_replicas: + self._unavailable_replicas.add(replica_id) + return self.release_replica_bindings(replica_id) + + def is_replica_available(self, replica_id: int) -> bool: + return ( + 0 <= replica_id < self.num_replicas + and self.clients[replica_id] is not None + and replica_id not in self._unavailable_replicas + ) + + def available_replica_ids(self) -> list[int]: + return [replica_id for replica_id in range(self.num_replicas) if self.is_replica_available(replica_id)] + def select_replica_id( self, request_id: str, @@ -483,23 +531,27 @@ class StagePool: ) -> int: """Pick a replica id for *request_id* and cache the choice (legacy path).""" cached = self.get_bound_replica_id(request_id) - if cached is not None and self.clients[cached] is not None: + if cached is not None and self.clients[cached] is not None and self.is_replica_available(cached): return cached + if cached is not None: + # Cached replica is dead/unavailable: drop the stale binding. + self.release_binding(request_id) chosen: int | None = None if affinity_request_id is not None: parent = self.get_bound_replica_id(affinity_request_id) - if parent is not None and self.clients[parent] is not None: + if parent is not None and self.clients[parent] is not None and self.is_replica_available(parent): chosen = parent if chosen is None: - live = self.live_replica_ids() + # Prefer replicas that are both live (client up) and available. + live = [r for r in self.live_replica_ids() if self.is_replica_available(r)] if not live: raise RuntimeError(f"stage {self.stage_id} has no live replicas") if len(live) == 1: chosen = live[0] else: - # Round-robin over live replicas only. + # Round-robin over live, available replicas only. start = self._next_replica_id % len(live) chosen = live[start] self._next_replica_id = (self._next_replica_id + 1) % len(live) @@ -539,18 +591,23 @@ class StagePool: non_empty_first_output_ts = ( self._non_empty_first_output_timestamps_by_request.pop(request_id, None) if request_id else None ) - num_tokens_out = count_tokens_from_outputs(request_outputs) + native_text_metrics = {} + if request_id: + pop_native_text_metrics = getattr(self.output_processor, "pop_native_text_metrics", None) + if callable(pop_native_text_metrics): + native_text_metrics = pop_native_text_metrics(request_id) + native_generation_tokens = native_text_metrics.get("num_generation_tokens") + num_tokens_out = ( + max(int(native_generation_tokens), 0) + if isinstance(native_generation_tokens, int) and not isinstance(native_generation_tokens, bool) + else count_tokens_from_outputs(request_outputs) + ) output_unit_type = self._infer_output_unit_type(request_outputs, token_count=num_tokens_out) output_unit_count = self._count_output_units( request_outputs, unit_type=output_unit_type, fallback_token_count=num_tokens_out, ) - native_text_metrics = {} - if request_id: - pop_native_text_metrics = getattr(self.output_processor, "pop_native_text_metrics", None) - if callable(pop_native_text_metrics): - native_text_metrics = pop_native_text_metrics(request_id) current_audio_frames, current_audio_sample_rate, _ = self._collect_audio_metrics(request_outputs) accumulated_audio_frames = self._audio_frames_by_request.pop(request_id, 0) if request_id else 0 accumulated_audio_sample_rate = self._audio_sample_rate_by_request.pop(request_id, 0) if request_id else 0 @@ -1029,14 +1086,28 @@ class StagePool: # Refresh the shared output-processor state before yielding to the # stage client so streaming segments are merged against the latest # prompt/token metadata. - self.output_processor.add_request( - request=request, - prompt=prompt_text, - parent_req=None, - request_index=0, - queue=None, - ) - await self._llm_client(replica_id).add_request_async(request) + try: + self.output_processor.add_request( + request=request, + prompt=prompt_text, + parent_req=None, + request_index=0, + queue=None, + ) + await self._llm_client(replica_id).add_request_async(request) + except Exception: + rollback = getattr(self.output_processor, "remove_request", None) + if callable(rollback): + try: + rollback(request_id) + except Exception as rollback_error: + logger.warning( + "[StagePool] Failed to rollback output processor update for req=%s stage-%s: %s", + request_id, + self.stage_id, + rollback_error, + ) + raise return replica_id async def _pick_or_select( @@ -1097,6 +1168,8 @@ class StagePool: timeout_s: float = 0.001, ) -> EngineCoreOutputs | None: """Poll raw EngineCore outputs from one LLM replica once.""" + if not self.is_replica_available(replica_id): + return None raw_client = self.clients[replica_id] if raw_client is None: return None @@ -1120,6 +1193,8 @@ class StagePool: def poll_diffusion_output(self, replica_id: int) -> Any | None: """Drain one ready diffusion output from the given replica if present.""" + if not self.is_replica_available(replica_id): + return None raw_client = self.clients[replica_id] if raw_client is None: return None diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 8897043ca..4ee4ccc61 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -396,7 +396,6 @@ class StageRuntime: replica_metadata.replica_id = replica_id if launch_mode == "remote" and replica_metadata.stage_type != "diffusion": replica_metadata.runtime_cfg = None - replicas.append( ReplicaInitPlan( replica_id=replica_id, diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 74911af9e..34f5f0a91 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -44,6 +44,12 @@ if TYPE_CHECKING: from vllm.tokenizers import TokenizerLike from vllm.v1.engine import PauseMode + from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseActivity + from vllm_omni.experimental.fullduplex.engine.messages import ( + DuplexFence, + DuplexSessionLifecycleMessage, + ) + from vllm_omni.experimental.fullduplex.request_client import DuplexRequestClient from vllm_omni.inputs.data import OmniPromptType logger = init_logger(__name__) @@ -143,12 +149,16 @@ class AsyncOmni(EngineClient, OmniBase): self._paused: bool = False self._sleeping_tags: set[str] = set() 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: @@ -254,6 +264,191 @@ class AsyncOmni(EngineClient, OmniBase): 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: + from vllm_omni.experimental.fullduplex.request_client import ( + DuplexRequestClient, + DuplexRequestOutputPort, + ) + + 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]: + from vllm_omni.experimental.fullduplex.request_client import DuplexRequestClient + + 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: + from vllm_omni.experimental.fullduplex.request_client import DuplexRequestClient + + return DuplexRequestClient.is_direct_response(output) + + @classmethod + def _duplex_multimodal_output(cls, output: object) -> dict[str, object]: + from vllm_omni.experimental.fullduplex.request_client import DuplexRequestClient + + return DuplexRequestClient.multimodal_output(output) + # ==================== Generate Method ==================== async def generate( @@ -693,6 +888,10 @@ class AsyncOmni(EngineClient, OmniBase): 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) and not msg.fatal: req_state = self.request_states.get(msg.request_id) if req_state is not None: diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py index 316d4c2aa..9526df36a 100644 --- a/vllm_omni/entrypoints/client_request_state.py +++ b/vllm_omni/entrypoints/client_request_state.py @@ -4,7 +4,7 @@ from vllm_omni.metrics import OrchestratorAggregator class ClientRequestState: - """Tracks the state of an individual request in the orchestrator.""" + """Tracks one entrypoint request and its output queue.""" def __init__( self, diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index d8b2c0291..44db512b2 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -473,6 +473,7 @@ class OmniBase(PDDisaggregationMixin): req_start_ts: dict[str, float], wall_start_ts: float, final_stage_id_for_e2e: int, + stage_event_cursor: int = 0, ) -> OmniRequestOutput | None: req_id = result.request_id engine_outputs = result.engine_outputs @@ -481,7 +482,7 @@ class OmniBase(PDDisaggregationMixin): # Merge AR stage timing from OrchestratorAggregator.stage_events if self._enable_ar_profiler: - ar_events = metrics.stage_events.get(str(req_id), []) + ar_events = metrics.stage_events.get(str(req_id), [])[stage_event_cursor:] for evt in ar_events: if evt.stage_id != stage_id: stage_durations[f"ar_stage_{evt.stage_id}"] = evt.stage_gen_time_ms / 1000.0 @@ -494,7 +495,7 @@ class OmniBase(PDDisaggregationMixin): stage_durations[key] = value # Merge per-stage gen times into stage_durations - for evt in metrics.stage_events.get(str(req_id), []): + for evt in metrics.stage_events.get(str(req_id), [])[stage_event_cursor:]: key = f"stage_{evt.stage_id}_gen_ms" if key not in stage_durations: stage_durations[key] = evt.stage_gen_time_ms @@ -546,7 +547,7 @@ class OmniBase(PDDisaggregationMixin): # Token counters — aggregate across all stages for this request. _prompt_tok = 0 _gen_tok = 0 - for evt in metrics.stage_events.get(rid_key, []): + for evt in metrics.stage_events.get(rid_key, [])[stage_event_cursor:]: if evt.stage_id == 0: _prompt_tok += int(evt.num_tokens_in) _gen_tok += int(evt.num_tokens_out) @@ -581,7 +582,7 @@ class OmniBase(PDDisaggregationMixin): response_metrics: dict[str, Any] = {} stage_metrics: dict[str, dict[str, Any]] = {} rid_key = str(req_id) - for evt in metrics.stage_events.get(rid_key, []): + for evt in metrics.stage_events.get(rid_key, [])[stage_event_cursor:]: if evt.stage_id is None: continue sid = int(evt.stage_id) @@ -604,6 +605,7 @@ class OmniBase(PDDisaggregationMixin): response_metrics["num_tokens_out"] = current_stage_metrics["num_tokens_out"] return OmniRequestOutput( request_id=req_id or "", + finished=finished, stage_id=stage_id, final_output_type=output_type, request_output=engine_outputs, @@ -616,7 +618,6 @@ class OmniBase(PDDisaggregationMixin): metrics=response_metrics, stage_durations=stage_durations, peak_memory_mb=peak_memory_mb, - finished=finished, ) def shutdown(self, timeout: float | None = None) -> None: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index c9ddab2e9..56c7ccf46 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -94,6 +94,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.config.endpoint_policy import shutdown_unsupported_routes from vllm_omni.diffusion.models.interface import ReferenceVideoDecodeSpec from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.openai.duplex_capability import should_enable_duplex_endpoint from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError from vllm_omni.entrypoints.openai.image_api_utils import ( SUPPORTED_LAYERED_RESOLUTIONS, @@ -793,6 +794,7 @@ async def omni_init_app_state( model_name=model_name, stage_configs=diffusion_stage_configs, ) + state.openai_serving_duplex = None state.openai_streaming_speech = None state.openai_streaming_video = None state.openai_serving_realtime_robot = ServingRealtimeRobotOpenPI.create_policy_server( @@ -1112,6 +1114,18 @@ async def omni_init_app_state( if state.openai_serving_chat is not None 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(args, "stage_configs_path", None) or getattr(args, "deploy_config", None), + ): + from vllm_omni.experimental.fullduplex.openai.serving import OmniDuplexSessionHandler + + state.openai_serving_duplex = OmniDuplexSessionHandler( + chat_service=state.openai_serving_chat, + 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, @@ -1612,6 +1626,31 @@ async def streaming_video_output(websocket: WebSocket): @router.websocket("/v1/realtime") async def realtime_websocket(websocket: WebSocket): """WebSocket endpoint for OpenAI-style realtime interactions.""" + 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 isinstance(duplex_query, str) and duplex_query.lower() in {"1", "true", "on"} + ) + if use_duplex_realtime and duplex_handler is not None: + await duplex_handler.handle_realtime_session(websocket) + return + + engine_client = getattr(websocket.app.state, "engine_client", None) + if engine_client is not None and getattr(engine_client, "async_chunk", False): + await websocket.accept() + await websocket.send_json( + { + "type": "error", + "error": ( + "The /v1/realtime API is not supported when async_chunk is enabled on the server. " + "Use a stage configuration with async_chunk disabled and restart the server before using " + "this endpoint." + ), + "code": "unsupported", + } + ) + await websocket.close() + return serving = getattr(websocket.app.state, "openai_serving_realtime", None) if serving is None: await websocket.accept() @@ -1639,6 +1678,18 @@ 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) + + # Health and Model endpoints for diffusion mode diff --git a/vllm_omni/entrypoints/openai/duplex_capability.py b/vllm_omni/entrypoints/openai/duplex_capability.py new file mode 100644 index 000000000..c3e4e3e38 --- /dev/null +++ b/vllm_omni/entrypoints/openai/duplex_capability.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM 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 duplex routes only for deployments that explicitly opt in.""" + 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 omegaconf import OmegaConf + + raw_config = OmegaConf.load(config_path) + session_mode = raw_config.get("session_mode") if hasattr(raw_config, "get") else None + if session_mode == "duplex": + return True + except Exception as exc: + logger.warning("Failed to inspect duplex session_mode from %s: %s", config_path, exc) + return False + + +__all__ = ["should_enable_duplex_endpoint"] diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index f819c7b57..93b73ad66 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -79,7 +79,7 @@ from vllm.outputs import RequestOutput from vllm.reasoning import ReasoningParser from vllm.renderers import BaseRenderer, merge_kwargs from vllm.renderers.inputs import TokPrompt -from vllm.sampling_params import SamplingParams +from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.tokenizers import TokenizerLike from vllm.tokenizers import TokenizerLike as AnyTokenizer from vllm.tokenizers.mistral import ( @@ -187,6 +187,63 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tp_cls = self.parser_cls.tool_parser_cls if self.parser_cls is not None else None return request.tools and tp_cls is not None and self.enable_auto_tools and request.tool_choice in ["auto", None] + @staticmethod + def _stage_get(obj: Any, key: str, default: Any = None) -> Any: + if obj is None: + return default + if hasattr(obj, "get"): + return obj.get(key, default) + return getattr(obj, key, default) + + def _has_minicpmo45_stage(self) -> bool: + for stage in getattr(self.engine_client, "stage_configs", []) or []: + engine_args = self._stage_get(stage, "engine_args") + model_arch = self._stage_get(engine_args, "model_arch") + if model_arch == "MiniCPMO45OmniForConditionalGeneration": + return True + return False + + def _fix_minicpmo45_audio_stream_output_kinds( + self, + sampling_params_list: list[Any], + output_modalities: list[str] | tuple[str, ...] | None, + ) -> list[Any]: + """Keep MiniCPM-o 4.5's thinker final-only during audio streaming. + + Chat ``stream=true`` normally coerces every AR stage to DELTA output. + That is correct for token-level text streaming, but MiniCPM-o 4.5's + Talker consumes the completed TTS span plus aligned thinker hidden + states. If Stage0 is DELTA, llm2tts receives one generated token at a + time and repeatedly starts tiny, independent TTS streams. The API then + never represents one coherent audio response. + + For this model, stream the audio stage, not the thinker-to-talker + boundary. + """ + if not output_modalities or "audio" not in output_modalities: + return sampling_params_list + if not self._has_minicpmo45_stage(): + return sampling_params_list + + stage_configs = getattr(self.engine_client, "stage_configs", []) or [] + for idx, stage in enumerate(stage_configs): + if idx >= len(sampling_params_list): + break + sp = sampling_params_list[idx] + if not hasattr(sp, "output_kind"): + continue + + engine_args = self._stage_get(stage, "engine_args") + if self._stage_get(engine_args, "model_arch") != "MiniCPMO45OmniForConditionalGeneration": + continue + + model_stage = self._stage_get(engine_args, "model_stage") + if model_stage == "llm": + sp.output_kind = RequestOutputKind.FINAL_ONLY + elif model_stage == "tts": + sp.output_kind = RequestOutputKind.DELTA + return sampling_params_list + @classmethod def for_diffusion( cls, @@ -605,6 +662,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # to delta to ensure emitted outputs are correctly drained. Otherwise # convert cumulative to Final Only to ensure the output is correct. sampling_params_list = coerce_param_message_types(sampling_params_list, request.stream) + if request.stream: + sampling_params_list = self._fix_minicpmo45_audio_stream_output_kinds( + sampling_params_list, + output_modalities, + ) # Apply user-specified overrides to diffusion stage(s) for image generation for idx, sp in enumerate(sampling_params_list): diff --git a/vllm_omni/entrypoints/openai/video_stream_base.py b/vllm_omni/entrypoints/openai/video_stream_base.py index 460a66d7c..1694704f3 100644 --- a/vllm_omni/entrypoints/openai/video_stream_base.py +++ b/vllm_omni/entrypoints/openai/video_stream_base.py @@ -10,12 +10,14 @@ via :class:`VideoStreamPipelineHooks`. Protocol: Client -> Server: {"type": "session.config", ...} # Session config (sent once) - {"type": "video.frame", "data": "..."} # base64 JPEG/PNG frame + {"type": "video.frame", "data": "...", "frame_id": "...", "pts_ms": 0} {"type": "audio.chunk", "data": "..."} # base64 PCM16 16kHz mono {"type": "video.query", "text": "..."} # Submit query about buffered frames {"type": "video.done"} # End of session Server -> Client: + {"type": "video.frame.ack", ...} # when frame_id is provided + {"type": "video.frames.consumed", ...} # after first engine output {"type": "response.start"} {"type": "response.text.delta", "delta": "..."} {"type": "response.text.done", "text": "..."} @@ -215,6 +217,7 @@ class OmniStreamingVideoHandler: return frame_buffer: list[str] = [] # base64-encoded JPEG frames + frame_metadata: list[dict[str, Any]] = [] # Per-frame PIL cache + uuid for mm_hash reuse. Aligned with frame_buffer by index. frame_pil_cache: dict[str, tuple[Any, str] | object] = {} # b64 -> (PIL.Image, uuid) or _BAD_FRAME frame_filter = ( @@ -259,6 +262,8 @@ class OmniStreamingVideoHandler: if msg_type.startswith("_internal."): await self._send_error(websocket, f"Unknown type: {msg_type}") continue + if msg_type == "video.frame": + msg["_receiver_received_ts_ms"] = _time.monotonic() * 1000 await msg_queue.put(msg) if msg.get("type") == "video.done": @@ -308,6 +313,7 @@ class OmniStreamingVideoHandler: active_request_id = request_id interrupt_event.clear() query_frames = list(frame_buffer) + query_frame_metadata = list(frame_metadata) query_audio_buffer = bytearray(audio_buffer) audio_buffer.clear() query_prewarmed_frames = dict(frame_pil_cache) @@ -315,6 +321,9 @@ class OmniStreamingVideoHandler: async def _run_query() -> None: nonlocal active_request_id, prev_request_id try: + process_kwargs: dict[str, Any] = {} + if any(metadata.get("frame_id") for metadata in query_frame_metadata): + process_kwargs["frame_metadata"] = query_frame_metadata await self._process_query( websocket, config, @@ -325,6 +334,7 @@ class OmniStreamingVideoHandler: request_id, interrupt_event, query_prewarmed_frames, + **process_kwargs, ) finally: if active_request_id == request_id: @@ -349,7 +359,11 @@ class OmniStreamingVideoHandler: frame_data = msg.get("b64", "") removed = frame_data in frame_buffer if removed: - frame_buffer[:] = [f for f in frame_buffer if f != frame_data] + retained_indices = [ + index for index, frame in enumerate(frame_buffer) if frame != frame_data + ] + frame_buffer[:] = [frame_buffer[index] for index in retained_indices] + frame_metadata[:] = [frame_metadata[index] for index in retained_indices] if frame_pil_cache.get(frame_data) is _BAD_FRAME: frame_pil_cache.pop(frame_data, None) if removed: @@ -370,16 +384,43 @@ class OmniStreamingVideoHandler: if frame_filter is not None: try: if not frame_filter.should_retain(raw_bytes): + await self._send_frame_ack( + websocket, + msg, + accepted=False, + buffered_frames=len(frame_buffer), + reason="filtered", + ) continue except Exception: await self._send_error(websocket, "Invalid image data") continue max_buf = config.max_frames + dropped_frame_id: str | None = None if len(frame_buffer) >= max_buf: dropped = frame_buffer.pop(0) + dropped_metadata = frame_metadata.pop(0) + dropped_frame_id = dropped_metadata.get("frame_id") frame_pil_cache.pop(dropped, None) frame_buffer.append(frame_data) + frame_metadata.append( + { + "frame_id": msg.get("frame_id"), + "pts_ms": msg.get("pts_ms"), + "source_pts_ms": msg.get("source_pts_ms"), + "quality_profile": msg.get("quality_profile"), + "capture_ts_ms": msg.get("capture_ts_ms"), + "receiver_received_ts_ms": msg.get("_receiver_received_ts_ms"), + } + ) self.on_frame_buffered(raw_bytes, frame_data, message_history, config) + await self._send_frame_ack( + websocket, + msg, + accepted=True, + buffered_frames=len(frame_buffer), + dropped_frame_id=dropped_frame_id, + ) # Prewarm: decode PIL off the event loop so query-time chat_template # can skip base64+Image.open. uuid=md5 lets mm_cache dedupe identical frames. if frame_data not in frame_pil_cache: @@ -537,6 +578,7 @@ class OmniStreamingVideoHandler: request_id: str, interrupt_event: asyncio.Event, prewarmed_frames: dict[str, tuple[Any, str]], + frame_metadata: list[dict[str, Any]] | None = None, ) -> None: """Build prompt, run inference, stream text + audio response.""" @@ -544,6 +586,9 @@ class OmniStreamingVideoHandler: await self._send_error(websocket, "Streaming video requires an engine client") return + engine_kwargs: dict[str, Any] = {} + if frame_metadata: + engine_kwargs["frame_metadata"] = frame_metadata await self._process_query_engine( websocket, config, @@ -554,6 +599,7 @@ class OmniStreamingVideoHandler: request_id, interrupt_event, prewarmed_frames, + **engine_kwargs, ) # ------------------------------------------------------------------ @@ -571,6 +617,7 @@ class OmniStreamingVideoHandler: request_id: str, interrupt_event: asyncio.Event, prewarmed_frames: dict[str, tuple[Any, str]], + frame_metadata: list[dict[str, Any]] | None = None, ) -> None: """Direct engine_client.generate() path for async_chunk audio.""" from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -613,6 +660,9 @@ class OmniStreamingVideoHandler: except Exception as e: await self._send_error(websocket, f"Preprocess failed: {e}") return + decoded_ready_ts_ms = _time.monotonic() * 1000 + selected_metadata = self._sample_frame_metadata(frame_metadata or [], config.num_frames) + model_selected_ts_ms = _time.monotonic() * 1000 await websocket.send_json({"type": "response.start"}) text_parts: list[str] = [] @@ -623,6 +673,7 @@ class OmniStreamingVideoHandler: audio_chunks_drained = 0 previous_text = "" interrupted = False + frames_consumed_sent = False t_start = _time.monotonic() t_first_text = None t_first_audio = None @@ -652,6 +703,33 @@ class OmniStreamingVideoHandler: if not isinstance(output, OmniRequestOutput): continue + if not frames_consumed_sent and frame_metadata: + await websocket.send_json( + { + "type": "video.frames.consumed", + "request_id": request_id, + "model_selected_ts_ms": model_selected_ts_ms, + "frame_ids": [ + metadata["frame_id"] + for metadata in selected_metadata + if isinstance(metadata.get("frame_id"), str) + ], + "frames": [ + { + "frame_id": metadata.get("frame_id"), + "pts_ms": metadata.get("pts_ms"), + "source_pts_ms": metadata.get("source_pts_ms"), + "quality_profile": metadata.get("quality_profile"), + "receiver_received_ts_ms": metadata.get("receiver_received_ts_ms"), + "decoded_ready_ts_ms": decoded_ready_ts_ms, + } + for metadata in selected_metadata + ], + "latest_pts_ms": selected_metadata[-1].get("pts_ms") if selected_metadata else None, + } + ) + frames_consumed_sent = True + out_type = getattr(output, "final_output_type", "text") if out_type == "audio": @@ -971,3 +1049,42 @@ class OmniStreamingVideoHandler: await websocket.send_json({"type": "error", "message": message}) except Exception: pass + + @staticmethod + def _sample_frame_metadata( + frame_metadata: list[dict[str, Any]], + num_frames: int, + ) -> list[dict[str, Any]]: + if len(frame_metadata) <= num_frames: + return list(frame_metadata) + stride = max(1, len(frame_metadata) // num_frames) + indices = [index * stride for index in range(num_frames - 1)] + [len(frame_metadata) - 1] + return [frame_metadata[index] for index in indices] + + @staticmethod + async def _send_frame_ack( + websocket: WebSocket, + message: Mapping[str, Any], + *, + accepted: bool, + buffered_frames: int, + reason: str | None = None, + dropped_frame_id: str | None = None, + ) -> None: + frame_id = message.get("frame_id") + if not isinstance(frame_id, str) or not frame_id: + return + ack: dict[str, Any] = { + "type": "video.frame.ack", + "frame_id": frame_id, + "pts_ms": message.get("pts_ms"), + "capture_ts_ms": message.get("capture_ts_ms"), + "accepted": accepted, + "buffered_frames": buffered_frames, + "server_receive_ts_ms": message.get("_receiver_received_ts_ms", _time.monotonic() * 1000), + } + if reason is not None: + ack["reason"] = reason + if dropped_frame_id is not None: + ack["dropped_frame_id"] = dropped_frame_id + await websocket.send_json(ack) diff --git a/vllm_omni/experimental/fullduplex/DESIGN.md b/vllm_omni/experimental/fullduplex/DESIGN.md new file mode 100644 index 000000000..256214299 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/DESIGN.md @@ -0,0 +1,1029 @@ +# MiniCPM-o 4.5 Full-Duplex Runtime Architecture + +## 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. + +The checkpoint preserves the active runtime path already exercised on H20. It does +not introduce another reducer, controller, worker provider, or shadow runtime. + +## 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 experimental 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; +- existing JoyVL behavior. + +The checkpoint does not claim: + +- scheduler-native KV append; +- deterministic VAD-triggered interruption (the browser intentionally does not + run VAD; MiniCPM owns listen/speak decisions at model-unit boundaries); +- production multi-session admission, fairness, capacity, or failure recovery; +- bounded long-session KV; +- video input or audio/video synchronization. + +## Active Runtime Path + +The actual path is: + +```text +WebSocket + -> OmniDuplexSessionHandler + -> DuplexSession + MiniCPMO45ServingSessionState + -> DuplexRequestClient + -> AsyncOmni thin open/append/signal/close proxies + -> DuplexControlClient over the engine-owned correlated RPC transport + -> DuplexControlPlane + DuplexSessionRuntimeManager + -> experimental DuplexRuntimeExtension + engine session/stage bindings + -> StagePool + -> resumable scheduler request + -> MiniCPM Stage0 + -> Stage1 TTS / Token2Wav + -> output processor + -> Realtime protocol projection + -> WebSocket writer +``` + +The removed typed reducer/facade and worker-provider runtimes were not in this +path. This refactor continues to modify only the path above. + +## State Ownership + +The architecture allows several state objects only when their facts are +orthogonal. + +| Owner | Authoritative facts | +| --- | --- | +| `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 experimental 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` | +| `experimental.fullduplex.engine.contracts` | immutable duplex DTOs, extension/stage protocols, typed decisions, request-ID codec | +| `experimental.fullduplex.engine.lease` | duplex lease activity, TTL/grace configuration, and expiry records | +| `experimental.fullduplex.engine.messages` | identity fence and duplex control message/result envelopes | +| `experimental.fullduplex.engine.duplex_runtime` | extension loading, validation, and compatibility exports | +| `experimental.fullduplex.engine.duplex_session` | session resources, fences, append reservations, idempotency cache, binding cleanup | +| `experimental.fullduplex.engine.duplex_control_plane` | control algorithms and the narrow `DuplexStagePort` used to submit and clean up stage requests | +| `experimental.fullduplex.engine.duplex_control_client` | typed control-message construction and calls through the engine-owned correlated RPC transport | +| `experimental.fullduplex.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 remains experimental: + +```text +vllm_omni.experimental.fullduplex.engine.duplex_runtime + # extension loading and validation +vllm_omni.experimental.fullduplex.engine.contracts + # immutable DTOs and narrow protocols +vllm_omni.experimental.fullduplex.engine.lease + # duplex lease primitives +vllm_omni.experimental.fullduplex.engine.messages + # duplex identity and queue envelopes +vllm_omni.experimental.fullduplex.engine.duplex_session + # session transaction implementation +vllm_omni.experimental.fullduplex.engine.duplex_control_plane + # engine-side control algorithms +vllm_omni.experimental.fullduplex.engine.duplex_control_client + # correlated control client +vllm_omni.experimental.fullduplex.request_client + # entrypoint request/output lifecycle +``` + +Stable engine and entrypoint modules import only the kernel at module load. +They load the experimental ControlPlane, ControlClient, request client, and +runtime extension inside duplex-only construction or call paths. Importing +`Orchestrator`, `AsyncOmniEngine`, or `AsyncOmni` for an ordinary deployment +therefore does not import `vllm_omni.experimental.fullduplex`. Compatibility +exports remain lazy and are not part of ordinary startup. + +`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. + +## 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 inspects the deployment `session_mode` before importing or +constructing the experimental duplex handler. Ordinary deployments do not load +the full-duplex serving package or 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.minicpmo45_native_duplex=true`; repository demos do so in their +session payloads. + +`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.audio.delta`. It carries decision metadata, not transcript text. | At most once. | +| `response.audio.delta` | Carries one ordered audio chunk. Every emitted audio delta is followed by a `response.audio_transcript.delta` for the same chunk. | Zero or more. | +| `response.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.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.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.audio_transcript.delta` values to equal +`response.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 `realtime_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. This is an input +ownership correction. It does not add a VAD policy: overlap is admitted to +Stage0 and the model decides whether to listen or speak. + +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 + +The checked-in MiniCPM duplex deploy profile sets Stage0 and Stage1 +`max_num_seqs` to two and configures the engine runtime manager with +`max_sessions=2`. The Realtime capability response therefore 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 deterministic VAD-triggered interruption, +multi-session production concurrency, bounded long-session KV, +scheduler-native append, or video input. diff --git a/vllm_omni/experimental/fullduplex/README.md b/vllm_omni/experimental/fullduplex/README.md index 7f9e22fa6..5359edf87 100644 --- a/vllm_omni/experimental/fullduplex/README.md +++ b/vllm_omni/experimental/fullduplex/README.md @@ -1,80 +1,28 @@ -# Full-duplex interaction framework +# Experimental Full-Duplex Runtime -A model-agnostic framework for real-time, full-duplex (streaming-in / streaming-out) -model interaction, plus the JoyVL implementation built on it. +This package contains two experimental integrations: -To **run the JoyVL model**, see the recipe: +- the existing JoyVL framework and example integration; +- the MiniCPM-o 4.5 native audio path used by `/v1/duplex` and + `/v1/realtime?duplex=1`. + +For the MiniCPM active runtime path, lifecycle invariants, capability boundary, +Realtime response contract, and validation scope, see [`DESIGN.md`](DESIGN.md). + +To run JoyVL, see [`recipes/JD/JoyAI-VL-Interaction.md`](../../../recipes/JD/JoyAI-VL-Interaction.md). -This README covers the framework itself and how to add a new model. -## Layout +## Package boundaries -``` -vllm_omni/experimental/fullduplex/ - core/ generic full-duplex framework (model-agnostic): DuplexRuntime (event - loop + epoch barge-in), DuplexSession, DuplexAdapter (ABC), protocol - joyvl/ JoyVL implementation (model-specific): - adapter.py JoyVLDuplexAdapter (implements core.DuplexAdapter) - decision/ policy + output_parser + prompts (speak/silence/delegate) - memory/ InteractionBrain — 3-tier summary memory (async) - serving/ OpenAI-compatible HTTP orchestrator - bridges/ model backend + delegation +```text +core/ existing JoyVL framework and experimental compatibility exports +engine/ AsyncOmni/orchestrator scheduler data-plane adapter +openai/ WebSocket transport, Realtime projection, and audio codecs +minicpmo45/ MiniCPM input framing, policy, compatibility, and Stage0 state +joyvl/ JoyVL model-specific integration ``` -`core/` is the only part shared across models; data planes differ by model and are -intentionally not shared. - -## Scope - -The runnable serving path is `joyvl/serving/` driving `decision/` + `memory/` directly. -`joyvl/adapter.py` + `core/` are a **demonstration** of how a model plugs into the -generic full-duplex framework (exercised by `tests/fullduplex/`), not the serving path -the HTTP orchestrator currently uses — a fused-audio model (e.g. MiniCPM-o) is the case -`core/` is built for. - -## Adding a full-duplex model - -The seam is `core.DuplexAdapter`. `core/` owns the session lifecycle, epoch-based -barge-in, playback cursor, and the event protocol — you implement only model policy. - -1. Create a sibling package `vllm_omni/experimental/fullduplex//` next to `joyvl/`; keep - model-specific code there and do not touch `core/`. -2. Implement one `DuplexAdapter` (three required methods; the rest have defaults): - - ```python - from collections.abc import AsyncIterator - from vllm_omni.experimental.fullduplex.core.adapter import DuplexAdapter, DuplexCapability, OutputChunk - from vllm_omni.experimental.fullduplex.core.session import DuplexSession - - class MyModelAdapter(DuplexAdapter): - def capabilities(self) -> DuplexCapability: - return DuplexCapability( - input_modalities=frozenset({"audio", "text"}), - output_modalities=frozenset({"audio", "text"}), - proactive=True, # speak without being asked? - ) - - async def on_input(self, session: DuplexSession, modality: str, data) -> None: - ... # buffer/route an incoming chunk - - async def respond(self, session: DuplexSession) -> AsyncIterator[OutputChunk]: - async for piece in self._model_stream(session): - yield OutputChunk(modality="audio", data=piece) # runtime drops stale - # chunks after a barge-in - - # optional: should_respond / on_barge_in / on_playback_ack - ``` - -3. Run it through the shared runtime — no new control-plane code: - - ```python - from vllm_omni.experimental.fullduplex.core.runtime import DuplexRuntime - from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexSessionConfig - - rt = DuplexRuntime(DuplexSession("sid", DuplexSessionConfig()), MyModelAdapter()) - await rt.run(input_events, emit) - ``` - -`joyvl/adapter.py` is the worked demonstration (currently exercised by tests, not the -HTTP serving path). Promote a helper from a model package up into `core/` only once a -second model actually needs it. +MiniCPM does not run through the removed experimental `core.DuplexRuntime` +facade. Its active path uses the `openai` session controller, the experimental +engine contracts, the standard scheduler/model runners, and an injected +MiniCPM-specific runtime extension from `minicpmo45/runtime.py`. diff --git a/vllm_omni/experimental/fullduplex/client.py b/vllm_omni/experimental/fullduplex/client.py new file mode 100644 index 000000000..57cc2f700 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/client.py @@ -0,0 +1,417 @@ +"""Reusable Realtime WebSocket, PCM, and event helpers for MiniCPM-o demos.""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import math +import time +import wave +from collections.abc import Callable +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit + +try: + import websockets + from websockets.exceptions import ConnectionClosed +except ImportError as exc: # pragma: no cover - example dependency + raise SystemExit("Install websockets first: pip install websockets") from exc + +PCM16_SAMPLE_RATE = 16_000 +PCM16_BYTES_PER_SAMPLE = 2 + + +def _rounded_ms(value: float) -> float: + return round(float(value), 3) + + +def _interval_summary(values: list[float]) -> dict[str, float | int]: + clean = sorted(_rounded_ms(value) for value in values if math.isfinite(value) and value >= 0) + if not clean: + return {"count": 0, "mean": 0.0, "p50": 0.0, "p95": 0.0, "max": 0.0} + + def nearest_rank(percentile: float) -> float: + index = max(0, math.ceil(percentile * len(clean)) - 1) + return clean[min(index, len(clean) - 1)] + + return { + "count": len(clean), + "mean": _rounded_ms(sum(clean) / len(clean)), + "p50": nearest_rank(0.50), + "p95": nearest_rank(0.95), + "max": clean[-1], + } + + +def _event_stage_metrics(event: dict[str, object]) -> dict[str, object] | None: + candidates: list[object] = [event.get("vllm_omni")] + metadata = event.get("metadata") + if isinstance(metadata, dict): + candidates.extend((metadata, metadata.get("vllm_omni"))) + response = event.get("response") + if isinstance(response, dict): + response_metadata = response.get("metadata") + if isinstance(response_metadata, dict): + candidates.extend((response_metadata, response_metadata.get("vllm_omni"))) + for candidate in candidates: + if not isinstance(candidate, dict): + continue + stage_metrics = candidate.get("stage_metrics") + if isinstance(stage_metrics, dict): + return stage_metrics + return None + + +def build_realtime_url( + url: str, + model: str, + *, + autostart: bool | None = None, + session_id: str | None = None, +) -> str: + """Add the explicit native-duplex query parameters to a Realtime URL.""" + parts = urlsplit(url) + query = dict(parse_qsl(parts.query, keep_blank_values=True)) + query.setdefault("duplex", "1") + query.setdefault("model", model) + query.setdefault("minicpmo45_native_duplex", "1") + 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)) + + +def read_pcm16_wav(path: Path) -> bytes: + """Read a mono, uncompressed, 16 kHz PCM16 WAV file.""" + with wave.open(str(path), "rb") as wav_file: + if wav_file.getnchannels() != 1: + raise ValueError("input WAV must be mono") + if wav_file.getsampwidth() != PCM16_BYTES_PER_SAMPLE: + raise ValueError("input WAV must be 16-bit PCM") + if wav_file.getframerate() != PCM16_SAMPLE_RATE: + raise ValueError("input WAV must be 16 kHz") + if wav_file.getcomptype() != "NONE": + raise ValueError("input WAV must be uncompressed PCM") + return wav_file.readframes(wav_file.getnframes()) + + +def write_pcm16_wav(path: Path, pcm16: bytes, *, sample_rate_hz: int) -> None: + """Write mono PCM16 bytes as a WAV artifact.""" + with wave.open(str(path), "wb") as wav_file: + wav_file.setnchannels(1) + wav_file.setsampwidth(PCM16_BYTES_PER_SAMPLE) + wav_file.setframerate(sample_rate_hz) + wav_file.writeframes(pcm16) + + +async def wait_for( + predicate: Callable[[], bool], + *, + timeout_s: float, + label: str, +) -> None: + """Wait for a collector predicate without coupling to a scenario runner.""" + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if predicate(): + return + await asyncio.sleep(0.02) + raise TimeoutError(f"Timed out waiting for {label}") + + +@dataclass +class RealtimeEventCollector: + """Collect server events and decode response audio by response identity.""" + + events: list[dict[str, object]] = field(default_factory=list) + event_received_at_s: list[float] = field(default_factory=list) + response_audio: dict[str, list[bytes]] = field(default_factory=dict) + response_ids: list[str] = field(default_factory=list) + output_sample_rate_hz: int = 24_000 + + @staticmethod + def response_id(event: dict[str, object]) -> str | None: + response_id = event.get("response_id") + if isinstance(response_id, str): + return response_id + response = event.get("response") + if isinstance(response, dict): + response_id = response.get("id") + if isinstance(response_id, str): + return response_id + return None + + def add(self, event: dict[str, object], *, received_at_s: float | None = None) -> None: + received_at = time.monotonic() if received_at_s is None else float(received_at_s) + stored_event = dict(event) + stored_event.setdefault("_client_received_at_s", received_at) + self.events.append(stored_event) + self.event_received_at_s.append(received_at) + response_id = self.response_id(stored_event) + event_type = stored_event.get("type") + if event_type == "response.created" and response_id and response_id not in self.response_ids: + self.response_ids.append(response_id) + if event_type == "response.audio.delta": + delta = stored_event.get("delta") or stored_event.get("audio") + if isinstance(delta, str) and response_id: + try: + self.response_audio.setdefault(response_id, []).append(base64.b64decode(delta)) + except ValueError: + pass + sample_rate_hz = stored_event.get("sample_rate_hz") + if isinstance(sample_rate_hz, int) and sample_rate_hz > 0: + self.output_sample_rate_hz = sample_rate_hz + + def count(self, event_type: str) -> int: + return sum(event.get("type") == event_type for event in self.events) + + def audio_bytes(self, response_id: str | None = None) -> bytes: + if response_id is not None: + return b"".join(self.response_audio.get(response_id, ())) + return b"".join( + chunk for response_id in self.response_ids for chunk in self.response_audio.get(response_id, ()) + ) + + def errors(self) -> list[dict[str, object]]: + return [event for event in self.events if event.get("type") == "error"] + + def first_received_at( + self, + *event_types: str, + after_s: float = 0.0, + ) -> float | None: + for event, received_at_s in zip(self.events, self.event_received_at_s, strict=True): + if received_at_s >= after_s and event.get("type") in event_types: + return received_at_s + return None + + def last_received_at(self, event_type: str) -> float | None: + for event, received_at_s in zip( + reversed(self.events), + reversed(self.event_received_at_s), + strict=True, + ): + if event.get("type") == event_type: + return received_at_s + return None + + def timing_summary( + self, + *, + after_s: float, + input_committed_at_s: float | None = None, + response_id: str | None = None, + ) -> dict[str, object]: + """Summarize engine token metrics and client-observed audio cadence.""" + stage0_metrics: dict[str, object] | None = None + response_created_at_s: float | None = None + audio_received_at_s: list[float] = [] + cumulative_audio_ms: list[float] = [] + for event, received_at_s in zip(self.events, self.event_received_at_s, strict=True): + if received_at_s < after_s: + continue + event_response_id = self.response_id(event) + if response_id is not None and event_response_id != response_id: + continue + if event.get("type") == "response.created" and response_created_at_s is None: + response_created_at_s = received_at_s + + stage_metrics = _event_stage_metrics(event) + stage0 = stage_metrics.get("0") if isinstance(stage_metrics, dict) else None + if isinstance(stage0, dict): + stage0_metrics = stage0 + + if event.get("type") != "response.audio.delta" or ( + response_id is not None and event_response_id != response_id + ): + continue + delta = event.get("delta") or event.get("audio") + if not isinstance(delta, str) or not delta: + continue + audio_received_at_s.append(received_at_s) + metadata = event.get("metadata") + duration_ms = metadata.get("audio_duration_ms") if isinstance(metadata, dict) else None + if isinstance(duration_ms, int | float) and math.isfinite(float(duration_ms)): + cumulative_audio_ms.append(max(0.0, float(duration_ms))) + + result: dict[str, object] = {} + if stage0_metrics is not None: + raw_itls = stage0_metrics.get("vllm_itls_ms") + itls = ( + [float(value) for value in raw_itls if isinstance(value, int | float)] + if isinstance(raw_itls, list) + else [] + ) + result["stage0_tokens"] = { + "source": "engine_stage_metrics", + "output_token_count": int(stage0_metrics.get("num_tokens_out") or 0), + "ttft_ms": float(stage0_metrics.get("vllm_ttft_ms") or 0.0), + "tpot_ms": float(stage0_metrics.get("vllm_tpot_ms") or 0.0), + "inter_token_interval_ms": _interval_summary(itls), + } + + if audio_received_at_s: + intervals_ms = [ + (current - previous) * 1000.0 for previous, current in zip(audio_received_at_s, audio_received_at_s[1:]) + ] + chunk_durations_ms: list[float] = [] + previous_duration_ms = 0.0 + for duration_ms in cumulative_audio_ms: + chunk_durations_ms.append( + duration_ms - previous_duration_ms if duration_ms >= previous_duration_ms else duration_ms + ) + previous_duration_ms = duration_ms + interval_summary = _interval_summary(intervals_ms) + result["audio_output"] = { + "source": "client_monotonic_receive", + "chunk_count": len(audio_received_at_s), + "response_created_to_first_audio_ms": ( + _rounded_ms((audio_received_at_s[0] - response_created_at_s) * 1000.0) + if response_created_at_s is not None + else None + ), + "commit_to_first_audio_ms": ( + _rounded_ms((audio_received_at_s[0] - input_committed_at_s) * 1000.0) + if input_committed_at_s is not None + else None + ), + "inter_chunk_interval_ms": interval_summary, + "chunk_duration_ms": _interval_summary(chunk_durations_ms), + "max_chunk_gap_ms": interval_summary["max"], + } + return result + + +class RealtimeDuplexClient: + """Small async client used by the user demo and reusable smoke probes.""" + + def __init__(self, url: str, *, max_size: int = 64 * 1024 * 1024) -> None: + self.url = url + self.max_size = max_size + self.events = RealtimeEventCollector() + self._ws: Any = None + self._reader_task: asyncio.Task[None] | None = None + + async def __aenter__(self) -> RealtimeDuplexClient: + self._ws = await websockets.connect(self.url, max_size=self.max_size) + self._reader_task = asyncio.create_task(self._read_events()) + return self + + async def __aexit__(self, exc_type, exc, traceback) -> None: + if self._ws is not None: + await self._ws.close() + if self._reader_task is not None: + self._reader_task.cancel() + try: + await self._reader_task + except asyncio.CancelledError: + pass + + async def _read_events(self) -> None: + try: + while True: + raw = await self._ws.recv() + if not isinstance(raw, str): + continue + event = json.loads(raw) + if isinstance(event, dict): + self.events.add(event) + except ConnectionClosed: + return + + async def send(self, event: dict[str, object]) -> None: + await self._ws.send(json.dumps(event)) + + async def configure( + self, + model: str, + *, + output_audio_format: str = "pcm16", + ref_audio: str | None = None, + session_id: str | None = None, + timeout_s: float = 20.0, + ) -> None: + session: dict[str, object] = { + "model": model, + "modalities": ["audio", "text"], + "input_audio_format": "pcm16", + "output_audio_format": output_audio_format, + "turn_detection": None, + "overlap_policy": "listen_only", + "playback_commit_policy": "ack_only", + "extra_body": { + "auto_response": True, + "minicpmo45_native_duplex": True, + "force_listen_count": 0, + }, + } + if ref_audio is not None: + session["ref_audio"] = ref_audio + if session_id: + session["session_id"] = session_id + await self.send({"type": "session.update", "session": session}) + await wait_for( + lambda: self.events.count("session.created") > 0, + timeout_s=timeout_s, + label="session.created", + ) + + async def stream_pcm16( + self, + pcm16: bytes, + *, + chunk_ms: int = 200, + realtime: bool = True, + ) -> None: + chunk_bytes = max( + PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * chunk_ms // 1000, + PCM16_BYTES_PER_SAMPLE, + ) + audio_end_ms = 0 + for offset in range(0, len(pcm16), chunk_bytes): + chunk = pcm16[offset : offset + chunk_bytes] + duration_ms = len(chunk) * 1000 // (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE) + audio_end_ms += duration_ms + await self.send( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(chunk).decode("ascii"), + "input_audio_format": "pcm16", + "sample_rate_hz": PCM16_SAMPLE_RATE, + "duration_ms": duration_ms, + "audio_end_ms": audio_end_ms, + } + ) + if realtime: + await asyncio.sleep(duration_ms / 1000) + + async def commit(self) -> None: + await self.send({"type": "input_audio_buffer.commit", "final": True}) + + async def acknowledge_playback(self) -> None: + for response_id in self.events.response_ids: + pcm16 = self.events.audio_bytes(response_id) + if not pcm16: + continue + played_ms = len(pcm16) * 1000 // (self.events.output_sample_rate_hz * PCM16_BYTES_PER_SAMPLE) + await self.send( + { + "type": "playback.ack", + "response_id": response_id, + "item_id": f"item_{response_id}", + "played_ms": played_ms, + "committed_ms": played_ms, + } + ) + + async def close_session(self, *, timeout_s: float = 20.0) -> None: + await self.send({"type": "session.close"}) + await wait_for( + lambda: self.events.count("session.closed") > 0, + timeout_s=timeout_s, + label="session.closed", + ) diff --git a/vllm_omni/experimental/fullduplex/engine/__init__.py b/vllm_omni/experimental/fullduplex/engine/__init__.py new file mode 100644 index 000000000..c8e20437f --- /dev/null +++ b/vllm_omni/experimental/fullduplex/engine/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Engine-side helpers for the experimental full-duplex runtime.""" diff --git a/vllm_omni/experimental/fullduplex/engine/contracts.py b/vllm_omni/experimental/fullduplex/engine/contracts.py new file mode 100644 index 000000000..ae449e82f --- /dev/null +++ b/vllm_omni/experimental/fullduplex/engine/contracts.py @@ -0,0 +1,276 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Model-neutral contracts for the experimental duplex engine plugin. + +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. +""" + +from __future__ import annotations + +import base64 +from collections.abc import Iterable, Mapping +from dataclasses import dataclass, field +from enum import Enum +from types import MappingProxyType +from typing import Any, Protocol + +from vllm_omni.engine.messages import EngineQueueMessage +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence + + +class SessionMode(str, Enum): + TURN = "turn" + DUPLEX = "duplex" + + +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" + + +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] + + +@dataclass(frozen=True) +class DuplexOutputDecision: + action: DuplexOutputAction + metadata: Mapping[str, Any] = 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 + fence: DuplexFence + + +@dataclass(frozen=True) +class DuplexStageRequestContext: + request_id: str + session_id: str + fence: DuplexFence + stage_id: int + 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) + + def __post_init__(self) -> None: + object.__setattr__(self, "sampling_params", tuple(self.sampling_params)) + object.__setattr__(self, "session_config", MappingProxyType(dict(self.session_config))) + object.__setattr__(self, "runtime_config", MappingProxyType(dict(self.runtime_config))) + + @property + def stage_sampling_params(self) -> object: + return self.sampling_params[self.stage_id] + + +@dataclass(frozen=True) +class DuplexStageSubmission: + context: DuplexStageRequestContext + prompt: Mapping[str, Any] + already_submitted: bool + + def __post_init__(self) -> None: + object.__setattr__(self, "prompt", MappingProxyType(dict(self.prompt))) + + +@dataclass(frozen=True) +class DuplexStageSubmissionResult: + request_id: str + stage_id: int + replica_id: int + + +@dataclass(frozen=True) +class DuplexOutputContext: + identity: DuplexRequestIdentity + final_stage_id: int + segment_finished: bool + segment_token_ids: tuple[int, ...] = () + segment_output_metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__(self, "segment_token_ids", tuple(self.segment_token_ids)) + object.__setattr__( + self, + "segment_output_metadata", + MappingProxyType(dict(self.segment_output_metadata)), + ) + + +class DuplexStagePort(Protocol): + @property + def stage_count(self) -> int: ... + + def sampling_defaults(self) -> tuple[object, ...]: ... + + def ensure_request(self, context: DuplexStageRequestContext) -> None: ... + + async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: ... + + 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: ... + + +def duplex_data_plane_request_info(result: dict[str, object]) -> tuple[str | None, int | None]: + stage_results = result.get("stage_results") + if not isinstance(stage_results, list): + return None, None + for item in stage_results: + if not isinstance(item, dict): + continue + inner = item.get("result") + if not isinstance(inner, dict) or inner.get("data_plane_append") is not True: + continue + request_id = inner.get("request_id") + if not isinstance(request_id, str) or not request_id: + continue + response_stage_id = inner.get("response_stage_id") + return request_id, response_stage_id if isinstance(response_stage_id, int) else None + return None, None + + +def duplex_resource_request_id(fence: DuplexFence, role: str) -> str: + if not role or any( + character not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-" for character in role + ): + 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}" + + +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": + return False + try: + int(parts[3]) + int(parts[5]) + except ValueError: + return False + role = parts[7] + if not role or any( + character not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-" for character in role + ): + return False + encoded_session_id = base64.urlsafe_b64encode(session_id.encode("utf-8")).decode("ascii").rstrip("=") + return parts[1] == encoded_session_id + + +__all__ = [ + "CorrelatedRpcTransport", + "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/experimental/fullduplex/engine/duplex_control_client.py b/vllm_omni/experimental/fullduplex/engine/duplex_control_client.py new file mode 100644 index 000000000..b37effe32 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/engine/duplex_control_client.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import uuid +from collections.abc import Callable + +from vllm_omni.experimental.fullduplex.engine.contracts import CorrelatedRpcTransport +from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseActivity +from vllm_omni.experimental.fullduplex.engine.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/experimental/fullduplex/engine/duplex_control_plane.py b/vllm_omni/experimental/fullduplex/engine/duplex_control_plane.py new file mode 100644 index 000000000..a1d3dac00 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/engine/duplex_control_plane.py @@ -0,0 +1,1063 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM 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.experimental.fullduplex.engine.contracts import ( + DuplexAppendPlan, + DuplexInputMode, + DuplexOutputContext, + DuplexOutputDecision, + DuplexRequestIdentity, + DuplexRuntimeCapabilities, + DuplexRuntimeExtension, + DuplexStagePort, + DuplexStageRequestContext, + DuplexStageSubmission, + DuplexStageSubmissionResult, + SessionMode, + duplex_resource_request_id, +) +from vllm_omni.experimental.fullduplex.engine.duplex_session import ( + DuplexAppendReservation, + DuplexFenceMismatchError, + DuplexSessionExpiry, + DuplexSessionRuntimeManager, + DuplexSessionRuntimeState, +) +from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseActivity, DuplexLeaseConfig +from vllm_omni.experimental.fullduplex.engine.messages import ( + AppendDuplexInputMessage, + CloseDuplexSessionMessage, + DuplexControlError, + DuplexControlResultMessage, + DuplexFence, + DuplexSessionLifecycleMessage, + OpenDuplexSessionMessage, + ResumeDuplexSessionMessage, + SignalDuplexTurnMessage, + TouchDuplexSessionMessage, +) + +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/experimental/fullduplex/engine/duplex_runtime.py b/vllm_omni/experimental/fullduplex/engine/duplex_runtime.py new file mode 100644 index 000000000..acd977448 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/engine/duplex_runtime.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Experimental duplex runtime loading and compatibility exports.""" + +from __future__ import annotations + +from importlib import import_module + +from vllm_omni.experimental.fullduplex.engine.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.experimental.fullduplex.engine.duplex_session import ( # noqa: E402, F401 + DuplexAppendReservation, + DuplexCompletedAppend, + DuplexFenceMismatchError, + DuplexInputAppend, + DuplexRequestResource, + DuplexSessionRuntimeManager, + DuplexSessionRuntimeState, + DuplexStageBinding, +) +from vllm_omni.experimental.fullduplex.engine.lease import ( # noqa: E402, F401 + DuplexLeaseActivity, + DuplexLeaseConfig, + DuplexLeaseState, + DuplexSessionExpiry, +) + +__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/experimental/fullduplex/engine/duplex_session.py b/vllm_omni/experimental/fullduplex/engine/duplex_session.py new file mode 100644 index 000000000..244005f06 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/engine/duplex_session.py @@ -0,0 +1,534 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM 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.experimental.fullduplex.engine.contracts import ( + DuplexInputMode, + DuplexRuntimeCapabilities, +) +from vllm_omni.experimental.fullduplex.engine.lease import ( + DuplexLeaseActivity, + DuplexLeaseConfig, + DuplexLeaseState, + DuplexSessionExpiry, +) +from vllm_omni.experimental.fullduplex.engine.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/experimental/fullduplex/engine/intermediate.py b/vllm_omni/experimental/fullduplex/engine/intermediate.py new file mode 100644 index 000000000..07c6ab668 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/engine/intermediate.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +from typing import TypedDict + + +class DuplexIntermediateBuffer(TypedDict, total=False): + """Structured keys carried through ``model_intermediate_buffer``. + + The buffer remains a dict for scheduler and msgspec compatibility, but + duplex-specific producers and consumers should use the helpers in this + module instead of scattering nested string keys across serving, runner, and + model code. + """ + + request_id: str + global_request_id: list[str] + prompt_token_ids: list[int] + llm_output_token_ids: list[int] + llm_output_text: list[str] + stream_output: bool + native_duplex: bool + ids: dict[str, object] + hidden_states: dict[str, object] + codes: dict[str, object] + meta: dict[str, object] + duplex: dict[str, object] + omni_payload: object + waveform: object + mel_spec: object + + +def build_duplex_intermediate_buffer( + *, + request_id: str, + prompt_token_ids: list[int] | None = None, + output_token_ids: list[int] | None = None, + output_text: str | None = None, + stream_output: bool = False, + native_duplex: bool = False, +) -> DuplexIntermediateBuffer: + buffer: DuplexIntermediateBuffer = { + "global_request_id": [str(request_id)], + "ids": {}, + } + if prompt_token_ids is not None: + prompt_ids = [int(token_id) for token_id in prompt_token_ids] + buffer["prompt_token_ids"] = prompt_ids + buffer["ids"]["prompt"] = prompt_ids + if output_token_ids is not None: + output_ids = [int(token_id) for token_id in output_token_ids] + buffer["llm_output_token_ids"] = output_ids + buffer["ids"]["output"] = output_ids + if output_text is not None: + buffer["llm_output_text"] = [output_text] + if stream_output: + buffer["stream_output"] = True + if native_duplex: + buffer["native_duplex"] = True + return buffer + + +def set_ref_audio(buffer: dict[str, object], waveform: object, sample_rate_hz: int) -> None: + buffer.setdefault("codes", {})["ref"] = waveform + buffer.setdefault("meta", {})["ref_audio_sr"] = int(sample_rate_hz) + + +def set_tts_handoff(buffer: dict[str, object], token_ids: object | None, hidden_states: object | None) -> None: + """Store the generic AR-stage token/hidden-state handoff consumed by TTS.""" + if token_ids is not None: + buffer.setdefault("ids", {})["tts"] = token_ids + if hidden_states is not None: + buffer.setdefault("hidden_states", {})["tts"] = hidden_states + + +def get_tts_handoff(info: dict[str, object]) -> tuple[object | None, object | None]: + """Read the generic AR-stage to TTS handoff, including legacy flat keys.""" + ids_info = info.get("ids") + hidden_info = info.get("hidden_states") + tts_token_ids = ids_info.get("tts") if isinstance(ids_info, dict) else None + tts_hidden_states = hidden_info.get("tts") if isinstance(hidden_info, dict) else None + if tts_token_ids is None: + tts_token_ids = info.get("tts_token_ids") + if tts_hidden_states is None: + tts_hidden_states = info.get("tts_hidden_states") + return tts_token_ids, tts_hidden_states + + +def get_stream_request_key(info: dict[str, object]) -> str: + key = info.get("global_request_id") or info.get("request_id") or info.get("_omni_req_id") + if isinstance(key, (list, tuple)): + key = key[0] if key else None + if isinstance(key, bytes): + key = key.decode("utf-8", errors="replace") + if key is None: + raise ValueError( + "Duplex streaming handoff requires a stable request id; " + "expected global_request_id, request_id, or _omni_req_id." + ) + return str(key) diff --git a/vllm_omni/experimental/fullduplex/engine/lease.py b/vllm_omni/experimental/fullduplex/engine/lease.py new file mode 100644 index 000000000..09430a7bb --- /dev/null +++ b/vllm_omni/experimental/fullduplex/engine/lease.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Generic engine lease primitives shared by optional session runtimes.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum + +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence + + +class DuplexLeaseActivity(str, Enum): + APPEND = "append" + SIGNAL = "signal" + PLAYBACK_ACK = "playback_ack" + HEARTBEAT = "heartbeat" + ATTACH = "attach" + DETACH = "detach" + RESUME = "resume" + MODEL_OUTPUT = "model_output" + + +@dataclass(frozen=True) +class DuplexLeaseConfig: + idle_ttl_s: float | None = 300.0 + disconnect_grace_s: float = 30.0 + + def __post_init__(self) -> None: + if self.idle_ttl_s is not None and self.idle_ttl_s <= 0: + raise ValueError("idle_ttl_s must be positive or null") + if self.disconnect_grace_s <= 0: + raise ValueError("disconnect_grace_s must be positive") + + +@dataclass +class DuplexLeaseState: + config: DuplexLeaseConfig + generation: int + last_activity: float + detached_at: float | None = None + active_operations: set[str] = field(default_factory=set) + terminal_reason: str | None = None + + @property + def expires_at(self) -> float | None: + if self.config.idle_ttl_s is None: + return None + return self.last_activity + self.config.idle_ttl_s + + def _require_open(self) -> None: + if self.terminal_reason is not None: + raise RuntimeError(f"duplex lease is terminal: {self.terminal_reason}") + + def touch(self, now: float, activity: DuplexLeaseActivity) -> None: + self._require_open() + if not isinstance(activity, DuplexLeaseActivity): + raise TypeError("lease activity must be DuplexLeaseActivity") + self.last_activity = now + + def detach(self, now: float) -> None: + self.touch(now, DuplexLeaseActivity.DETACH) + self.detached_at = now + + def resume(self, now: float, *, expected_generation: int) -> int: + self._require_open() + if expected_generation != self.generation: + raise ValueError(f"duplex lease generation mismatch: expected {self.generation}, got {expected_generation}") + self.generation += 1 + self.detached_at = None + self.last_activity = now + return self.generation + + def begin_operation(self, now: float, operation_id: str) -> None: + self._require_open() + if not operation_id: + raise ValueError("duplex lease operation_id must not be empty") + if operation_id in self.active_operations: + raise ValueError(f"duplex lease operation already active: {operation_id}") + self.active_operations.add(operation_id) + self.last_activity = now + + def end_operation(self, now: float, operation_id: str) -> None: + self._require_open() + if operation_id not in self.active_operations: + raise KeyError(f"duplex lease operation is not active: {operation_id}") + self.active_operations.remove(operation_id) + self.last_activity = now + + def disconnect_grace_expired(self, now: float) -> bool: + return ( + self.terminal_reason is None + and self.detached_at is not None + and now >= self.detached_at + self.config.disconnect_grace_s + ) + + def idle_expired(self, now: float) -> bool: + expires_at = self.expires_at + return ( + self.terminal_reason is None and not self.active_operations and expires_at is not None and now >= expires_at + ) + + def mark_terminal(self, reason: str) -> bool: + if self.terminal_reason is not None: + return False + self.terminal_reason = reason + self.generation += 1 + self.detached_at = None + self.active_operations.clear() + 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/experimental/fullduplex/engine/messages.py b/vllm_omni/experimental/fullduplex/engine/messages.py new file mode 100644 index 000000000..94f154a8a --- /dev/null +++ b/vllm_omni/experimental/fullduplex/engine/messages.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Queue contracts owned by the experimental duplex control plane.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +from vllm_omni.engine.messages import EngineQueueMessage + + +@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 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 + + +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 + session_id: str + event: str + reason: str + lease_generation: int + submitted_request_ids: list[str] + reserved_request_ids: list[str] + + +class DuplexControlError(EngineQueueMessage, kw_only=True): + code: str + message: str + retryable: bool = False + + +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 + + @property + def rpc_correlation_key(self) -> tuple[str, str]: + return ("duplex", self.control_id) + + +__all__ = [ + "AppendDuplexInputMessage", + "CloseDuplexSessionMessage", + "DuplexControlError", + "DuplexControlResultMessage", + "DuplexFence", + "DuplexSessionLifecycleMessage", + "OpenDuplexSessionMessage", + "ResumeDuplexSessionMessage", + "SignalDuplexTurnMessage", + "TouchDuplexSessionMessage", +] diff --git a/vllm_omni/experimental/fullduplex/minicpmo45/__init__.py b/vllm_omni/experimental/fullduplex/minicpmo45/__init__.py new file mode 100644 index 000000000..6df010470 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/minicpmo45/__init__.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .adapter import MiniCPMO45ClientRuntimeConfigError, MiniCPMO45NativeDuplexServingAdapter +from .input import MiniCPMO45PcmAppendBuffer +from .policy import MiniCPMO45DuplexPolicy +from .stage0 import MiniCPMO45Stage0DuplexRuntime + +__all__ = [ + "MiniCPMO45DuplexPolicy", + "MiniCPMO45NativeDuplexServingAdapter", + "MiniCPMO45ClientRuntimeConfigError", + "MiniCPMO45PcmAppendBuffer", + "MiniCPMO45Stage0DuplexRuntime", +] diff --git a/vllm_omni/experimental/fullduplex/minicpmo45/adapter.py b/vllm_omni/experimental/fullduplex/minicpmo45/adapter.py new file mode 100644 index 000000000..10d75320d --- /dev/null +++ b/vllm_omni/experimental/fullduplex/minicpmo45/adapter.py @@ -0,0 +1,304 @@ +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.experimental.fullduplex.minicpmo45.policy import MiniCPMO45DuplexPolicy +from vllm_omni.experimental.fullduplex.openai.protocol import DuplexSessionConfig +from vllm_omni.experimental.fullduplex.openai.runtime_adapter import ( + ServingRuntimeConfigError, +) + + +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", + } + ) + + @classmethod + def is_enabled(cls, config: DuplexSessionConfig) -> bool: + return config.extra_body.get("minicpmo45_native_duplex") 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 {} + runtime_config["instructions"] = config.instructions + 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} + cls._apply_default_scheduler_policy(runtime_config, config=config, model_config=model_config) + + ref_audio = config.ref_audio + if ref_audio is None and isinstance(extra_body.get("ref_audio"), str): + ref_audio = extra_body.pop("ref_audio") + if ref_audio is None and isinstance(extra_body.get("tts_ref_audio"), str): + ref_audio = extra_body.pop("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, + 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, + 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 + runtime_config["duplex_stage_sampling_params"] = {"0": stage0_params} + 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, + 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, + ) + 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) + 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 + 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/experimental/fullduplex/minicpmo45/compat.py b/vllm_omni/experimental/fullduplex/minicpmo45/compat.py new file mode 100644 index 000000000..9b44bed4c --- /dev/null +++ b/vllm_omni/experimental/fullduplex/minicpmo45/compat.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from typing import Any + + +def patch_minicpmo_remote_config(config: Any) -> None: + """Fill config fields expected by the vLLM MiniCPM implementation.""" + seen: set[int] = set() + + def patch_cfg(cfg: Any) -> None: + if cfg is None or id(cfg) in seen: + return + seen.add(id(cfg)) + if getattr(cfg, "base_model_tp_plan", None) is None: + setattr(cfg, "base_model_tp_plan", {}) + for value in getattr(cfg, "__dict__", {}).values(): + if value is cfg: + continue + if hasattr(value, "__dict__") and ( + value.__class__.__name__.endswith("Config") + or hasattr(value, "model_type") + or hasattr(value, "base_model_tp_plan") + ): + patch_cfg(value) + + patch_cfg(config) + tts_config = getattr(config, "tts_config", None) + if tts_config is None: + return + defaults = { + "top_p": 0.8, + "top_k": 100, + "temperature": 0.8, + "repetition_penalty": 1.05, + } + for name, value in defaults.items(): + if isinstance(tts_config, dict): + tts_config.setdefault(name, value) + elif not hasattr(tts_config, name): + setattr(tts_config, name, value) + + +__all__ = ["patch_minicpmo_remote_config"] diff --git a/vllm_omni/experimental/fullduplex/minicpmo45/data_plane.py b/vllm_omni/experimental/fullduplex/minicpmo45/data_plane.py new file mode 100644 index 000000000..d3b082ffc --- /dev/null +++ b/vllm_omni/experimental/fullduplex/minicpmo45/data_plane.py @@ -0,0 +1,838 @@ +from __future__ import annotations + +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass, field + +import numpy as np +from vllm.logger import init_logger + +from vllm_omni.experimental.fullduplex.engine.contracts import ( + duplex_resource_request_belongs_to_session, +) +from vllm_omni.experimental.fullduplex.output import get_duplex_output_decision + +logger = init_logger(__name__) + +EncodeAudio = Callable[[object, int, str, float | None], str | None] + + +@dataclass(frozen=True, slots=True) +class MiniCPMO45DataPlaneContext: + """Serving state needed to project one MiniCPM data-plane output.""" + + epoch: int = 0 + turn_id: int = 0 + active_response_turn_id: int | None = None + active_response_id: str | None = None + auto_responds: bool = False + response_format: str = "wav" + speed: float | None = None + modalities: tuple[str, ...] = () + + +@dataclass(slots=True) +class _TurnState: + sent_segment_text: str = "" + has_text: bool = False + tts_eos_done: bool = False + turn_eos_done: bool = False + + +@dataclass(slots=True) +class _RequestState: + audio_offset: int = 0 + uses_segment_text_metadata: bool = False + pending_audio_without_text: list[dict[str, object]] = field(default_factory=list) + terminal: bool = False + turns: dict[int | None, _TurnState] = field(default_factory=dict) + + def turn(self, turn_id: int | None) -> _TurnState: + return self.turns.setdefault(turn_id, _TurnState()) + + +class MiniCPMO45DataPlaneSession: + """MiniCPM output projector and request/turn cursor owner. + + The scheduler request owns cumulative Stage1 audio, while model turns own + transcript and terminal cursors. Keeping both levels in one object avoids + leaking MiniCPM-specific state into the generic Realtime handler. + """ + + def __init__(self, encode_audio: EncodeAudio) -> None: + self._encode_audio = encode_audio + self._requests: dict[str, _RequestState] = {} + + def begin_request(self, request_id: str) -> None: + state = self._requests.setdefault(request_id, _RequestState()) + state.terminal = False + # Preserve turn-scoped cursors for resumable requests. The legacy path + # only reset request-scoped terminal markers on append. + request_turn = state.turns.get(None) + if request_turn is not None: + request_turn.tts_eos_done = False + request_turn.turn_eos_done = False + + def is_terminal(self, request_id: str | None) -> bool: + if request_id is None: + return False + state = self._requests.get(request_id) + return state is not None and state.terminal + + def mark_terminal(self, request_id: str) -> None: + self._requests.setdefault(request_id, _RequestState()).terminal = True + + def close_stream(self, request_id: str) -> None: + """Release non-resumable request cursors after its drain finishes.""" + state = self._requests.get(request_id) + if state is None: + return + state.audio_offset = 0 + state.turns.pop(None, None) + + def close_request(self, request_id: str) -> None: + self._requests.pop(request_id, None) + + def close_session(self, session_id: str, *, active_request_id: str | None = None) -> None: + if active_request_id is not None: + self.close_request(active_request_id) + for request_id in list(self._requests): + if self.request_belongs_to_session(request_id, session_id): + self.close_request(request_id) + + def has_request(self, request_id: str) -> bool: + return request_id in self._requests + + def has_pending_audio(self, request_id: str) -> bool: + state = self._requests.get(request_id) + return bool(state and state.pending_audio_without_text) + + def audio_offset(self, request_id: str | None) -> int | None: + if request_id is None: + return None + state = self._requests.get(request_id) + return state.audio_offset if state is not None and state.audio_offset > 0 else None + + def project( + self, + result: object, + *, + context: MiniCPMO45DataPlaneContext | None = None, + ) -> Iterator[dict[str, object]]: + if not isinstance(result, dict): + return + outputs = result.get("data_plane_outputs") + if not isinstance(outputs, list): + return + for output in outputs: + yield from self.project_output(output, context=context) + + def project_output( + self, + output: object, + *, + context: MiniCPMO45DataPlaneContext | None = None, + ) -> Iterator[dict[str, object]]: + context = context or MiniCPMO45DataPlaneContext() + stage_metrics = _output_stage_metrics(output) + + def runtime_result(**values: object) -> dict[str, object]: + result = _runtime_result(**values) + if stage_metrics is not None: + result["stage_metrics"] = stage_metrics + return result + + request_id = getattr(output, "request_id", None) + if not isinstance(request_id, str) or not request_id: + request_id = None + request_state = self._requests.setdefault(request_id, _RequestState()) if request_id is not None else None + + outputs = getattr(output, "outputs", None) + completion = outputs[0] if isinstance(outputs, list) and outputs else None + text = getattr(completion, "text", "") if completion is not None else "" + direct_decision = get_duplex_output_decision(output) + direct_metadata = getattr(direct_decision, "metadata", None) + if isinstance(direct_metadata, Mapping): + mm_output = direct_metadata + else: + mm_output = getattr(output, "multimodal_output", None) + if not isinstance(mm_output, Mapping): + mm_output = getattr(completion, "multimodal_output", {}) if completion is not None else {} + if not mm_output: + inner_output = getattr(output, "request_output", None) + if inner_output is not None and inner_output is not output: + inner_mm_output = getattr(inner_output, "multimodal_output", None) + if isinstance(inner_mm_output, Mapping) and inner_mm_output: + mm_output = inner_mm_output + else: + inner_outputs = getattr(inner_output, "outputs", None) + inner_completion = inner_outputs[0] if isinstance(inner_outputs, list) and inner_outputs else None + inner_mm_output = ( + getattr(inner_completion, "multimodal_output", None) if inner_completion is not None else None + ) + if completion is None and inner_completion is not None: + completion = inner_completion + text = getattr(inner_completion, "text", "") or text + if isinstance(inner_mm_output, Mapping): + mm_output = inner_mm_output + mm_output = dict(mm_output) if isinstance(mm_output, Mapping) else {} + + output_turn_id = output_turn_id_from_metadata(mm_output) + output_epoch = output_epoch_from_metadata(mm_output) + expected_turn_id = context.active_response_turn_id + stale_turn = expected_turn_id is not None and output_turn_id is not None and output_turn_id < expected_turn_id + if expected_turn_id is None and output_turn_id is not None: + stale_turn = output_turn_id < context.turn_id + stale_epoch = output_epoch is not None and output_epoch != context.epoch + if context.auto_responds and (stale_turn or stale_epoch): + return + + mm_text = _llm_output_text(mm_output) + if mm_text: + text = mm_text + if request_state is not None: + request_state.uses_segment_text_metadata = True + elif request_state is not None and request_state.uses_segment_text_metadata: + text = "" + + finished = bool(getattr(output, "finished", False)) + tts_is_last_chunk = _bool_metadata(mm_output, ("tts_is_last_chunk",), default=False) + token_ids = _completion_token_ids(completion) + native_decision = _native_decision(completion, mm_output, token_ids=token_ids, finished=finished) + if native_decision == "listen": + listen_result = runtime_result( + stage_role="llm", + is_listen=True, + model_listen=True, + listen_source="model_listen", + data_plane_request_id=request_id, + end_of_turn=False, + ) + if output_turn_id is not None: + listen_result["model_turn_id"] = output_turn_id + yield listen_result + return + + raw_audio = next((mm_output[key] for key in ("audio", "model_outputs", "latent") if key in mm_output), None) + raw_audio_samples = _audio_num_samples(raw_audio) + offset_before = self.audio_offset(request_id) + audio_chunks = list( + self._encode_audio_chunks_with_duration( + mm_output, + request_id=request_id, + response_format=context.response_format, + speed=context.speed, + ) + ) + stage_turn_end = _bool_metadata(mm_output, ("turn_end", "end_of_turn"), default=False) + terminal_turn_state = request_state.turn(output_turn_id) if request_state is not None else None + stage_tts_eos = ( + context.auto_responds + and 151645 in token_ids + and not audio_chunks + and raw_audio_samples is not None + and (raw_audio_samples == 0 or raw_audio_samples == offset_before) + and (terminal_turn_state is None or not terminal_turn_state.tts_eos_done) + ) + tts_segment_end = bool(tts_is_last_chunk or stage_tts_eos) and ( + terminal_turn_state is None or not terminal_turn_state.tts_eos_done + ) + if tts_segment_end and terminal_turn_state is not None: + terminal_turn_state.tts_eos_done = True + stage_turn_end_new = bool(stage_turn_end) and ( + terminal_turn_state is None or not terminal_turn_state.turn_eos_done + ) + if stage_turn_end_new and terminal_turn_state is not None: + terminal_turn_state.turn_eos_done = True + unit_end_of_turn = stage_turn_end_new or (finished and not context.auto_responds) + + text_turn_id = output_turn_id if output_turn_id is not None else context.turn_id + text_turn_state = request_state.turn(text_turn_id) if request_state is not None else None + if audio_chunks: + delta_text = self.segment_text_delta(request_id, text, turn_id=text_turn_id) + last_idx = len(audio_chunks) - 1 + sample_rate_hz = _sample_rate_hz(mm_output) + audio_text_marks = _audio_text_marks(mm_output) + fallback_marks = _fallback_audio_text_marks(audio_chunks, delta_text) + audio_results: list[dict[str, object]] = [] + for idx, (audio, duration_ms) in enumerate(audio_chunks): + native_result = runtime_result( + stage_role="tts", + is_listen=False, + data_plane_request_id=request_id, + text=delta_text if idx == 0 else "", + audio_data=audio, + audio_format=context.response_format, + audio_duration_ms=duration_ms, + audio_text_mark=idx == last_idx, + sample_rate_hz=sample_rate_hz, + end_of_turn=unit_end_of_turn and idx == last_idx, + abort_data_plane_request=tts_segment_end and idx == last_idx, + ) + if output_turn_id is not None: + native_result["model_turn_id"] = output_turn_id + if audio_text_marks and idx == last_idx: + native_result["audio_text_marks"] = audio_text_marks + native_result["audio_text_marks_are_cumulative"] = True + elif idx < len(fallback_marks) and fallback_marks[idx]: + native_result["audio_text_marks"] = fallback_marks[idx] + native_result["audio_text_marks_are_cumulative"] = True + audio_results.append(native_result) + + if context.auto_responds: + if delta_text and text_turn_state is not None: + text_turn_state.has_text = True + future_model_turn = ( + context.active_response_turn_id is not None + and output_turn_id is not None + and output_turn_id > context.active_response_turn_id + ) + response_turn_bound = context.active_response_id is not None and ( + output_turn_id is None + or context.active_response_turn_id is None + or context.active_response_turn_id == output_turn_id + ) + turn_has_text = text_turn_state is not None and text_turn_state.has_text + if not future_model_turn and not response_turn_bound and not turn_has_text: + if request_state is not None: + if tts_segment_end: + request_state.pending_audio_without_text.clear() + else: + request_state.pending_audio_without_text.extend(audio_results) + if tts_segment_end: + terminal_result = dict(audio_results[-1]) + terminal_result.update( + audio_data="", + audio_duration_ms=0, + audio_text_mark=False, + end_of_turn=True, + ) + yield terminal_result + return + if request_state is not None and request_state.pending_audio_without_text: + pending = request_state.pending_audio_without_text + request_state.pending_audio_without_text = [] + yield from pending + yield from audio_results + return + + if context.auto_responds and request_state is not None and isinstance(text, str) and text: + pending_audio = request_state.pending_audio_without_text + request_state.pending_audio_without_text = [] + if pending_audio: + delta_text = self.segment_text_delta(request_id, text, turn_id=text_turn_id) + if delta_text: + pending_audio[0]["text"] = delta_text + if text_turn_state is not None: + text_turn_state.has_text = True + total_duration_ms = sum( + max(0, int(result.get("audio_duration_ms", 0) or 0)) for result in pending_audio + ) + if total_duration_ms > 0 and not pending_audio[-1].get("audio_text_marks"): + pending_audio[-1]["audio_text_marks"] = [ + {"text_chars": len(delta_text), "audio_end_ms": total_duration_ms} + ] + pending_audio[-1]["audio_text_marks_are_cumulative"] = True + if unit_end_of_turn: + pending_audio[-1]["end_of_turn"] = True + if tts_segment_end: + pending_audio[-1]["abort_data_plane_request"] = True + yield from pending_audio + return + request_state.pending_audio_without_text = pending_audio + + if tts_segment_end: + if unit_end_of_turn and request_state is not None and request_state.pending_audio_without_text: + request_state.pending_audio_without_text[-1]["end_of_turn"] = True + request_state.pending_audio_without_text[-1]["abort_data_plane_request"] = True + terminal_result = runtime_result( + stage_role="tts", + is_listen=False, + data_plane_request_id=request_id, + text="", + audio_data="", + audio_format=context.response_format, + audio_text_mark=False, + end_of_turn=unit_end_of_turn, + abort_data_plane_request=True, + ) + if output_turn_id is not None: + terminal_result["model_turn_id"] = output_turn_id + yield terminal_result + return + + if context.active_response_id is not None and unit_end_of_turn: + terminal_result = runtime_result( + stage_role="tts", + is_listen=False, + data_plane_request_id=request_id, + text="", + audio_data="", + audio_format=context.response_format, + audio_text_mark=False, + end_of_turn=True, + ) + if output_turn_id is not None: + terminal_result["model_turn_id"] = output_turn_id + yield terminal_result + return + + if request_id is not None and context.auto_responds and unit_end_of_turn and context.active_response_id is None: + # The legacy projector only removed the request-scoped pending + # buffer here when no explicit model-turn key was present. + if request_state is not None and output_turn_id is None: + request_state.pending_audio_without_text.clear() + terminal_result = runtime_result( + stage_role="tts", + is_listen=False, + data_plane_request_id=request_id, + text="", + audio_data="", + audio_format=context.response_format, + audio_text_mark=False, + end_of_turn=True, + abort_data_plane_request=True, + ) + if output_turn_id is not None: + terminal_result["model_turn_id"] = output_turn_id + yield terminal_result + return + + if ( + finished + and context.auto_responds + and context.active_response_id is not None + and request_id is not None + and not unit_end_of_turn + ): + listen_result = runtime_result( + 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, + ) + if output_turn_id is not None: + listen_result["model_turn_id"] = output_turn_id + yield listen_result + return + + if not text: + if context.auto_responds: + return + if finished: + listen_result = runtime_result( + stage_role="llm", + is_listen=True, + model_listen=False, + listen_source="data_plane_finished_without_output", + reason="data_plane_finished_without_output", + data_plane_request_id=request_id, + end_of_turn=False, + ) + if output_turn_id is not None: + listen_result["model_turn_id"] = output_turn_id + yield listen_result + return + if context.auto_responds: + return + if "audio" in context.modalities: + yield runtime_result( + stage_role="tts", + error_code="runtime_data_plane_text_without_audio", + error="MiniCPM-o native duplex data-plane produced text without audio.", + data_plane_request_id=request_id, + ) + return + yield runtime_result( + stage_role="llm", + is_listen=False, + data_plane_request_id=request_id, + text=text if isinstance(text, str) else "", + audio_data="", + end_of_turn=unit_end_of_turn, + ) + + def segment_text_delta(self, request_id: str | None, text: object, *, turn_id: int | None = None) -> str: + if not isinstance(text, str) or not text: + return "" + if request_id is None: + return text + turn_state = self._requests.setdefault(request_id, _RequestState()).turn(turn_id) + sent_text = turn_state.sent_segment_text + if not sent_text: + delta_text = text + elif text == sent_text: + delta_text = "" + elif text.startswith(sent_text): + delta_text = text[len(sent_text) :] + else: + # The producer sends text for the current thinker segment, not a + # cumulative turn snapshot. Distinct adjacent segments need not + # have a prefix relationship and must both remain visible. + delta_text = text + turn_state.sent_segment_text = text + return delta_text + + def slice_cumulative_audio(self, request_id: str | None, audio_data: object) -> object: + if request_id is None: + return audio_data + num_samples = _audio_num_samples(audio_data) + if num_samples is None or num_samples <= 0: + return audio_data + state = self._requests.setdefault(request_id, _RequestState()) + prev_samples = state.audio_offset + if prev_samples <= 0: + state.audio_offset = num_samples + return audio_data + if num_samples == prev_samples: + return None + if num_samples < prev_samples: + state.audio_offset = num_samples + return audio_data + state.audio_offset = num_samples + try: + import torch + + if isinstance(audio_data, torch.Tensor): + return audio_data.reshape(-1)[prev_samples:].contiguous() + return np.asarray(audio_data, dtype=np.float32).reshape(-1)[prev_samples:] + except Exception: + logger.exception("Failed to slice cumulative duplex audio output") + return audio_data + + def _encode_audio_chunks_with_duration( + self, + mm_output: dict[str, object], + *, + request_id: str | None, + response_format: str, + speed: float | None, + ) -> Iterator[tuple[str, int]]: + sample_rate_hz = _sample_rate_hz(mm_output) + audio_data = next((mm_output[key] for key in ("audio", "model_outputs", "latent") if key in mm_output), None) + if isinstance(audio_data, list): + for value in audio_data: + encoded = self._encode_audio(value, sample_rate_hz, response_format, speed) + if encoded: + duration_ms = int((_audio_num_samples(value) or 0) * 1000 / max(1, sample_rate_hz)) + yield encoded, duration_ms + return + sliced = self.slice_cumulative_audio(request_id, audio_data) + encoded = self._encode_audio(sliced, sample_rate_hz, response_format, speed) + if encoded: + duration_ms = int((_audio_num_samples(sliced) or 0) * 1000 / max(1, sample_rate_hz)) + yield encoded, duration_ms + + @staticmethod + def request_belongs_to_session(request_id: str, session_id: str) -> bool: + return ( + duplex_resource_request_belongs_to_session(request_id, session_id) + or request_id.startswith(f"duplex-{session_id}-") + or request_id.startswith(f"chatcmpl-duplex-{session_id}-") + ) + + +def coerce_int(value: object) -> int | None: + if hasattr(value, "detach"): + try: + value = value.detach().cpu().reshape(-1) + if value.numel() == 0: + return None + value = value[0].item() + except Exception: + return None + elif hasattr(value, "reshape") and hasattr(value, "size"): + try: + value = value.reshape(-1) + if int(value.size) == 0: + return None + item = value[0] + value = item.item() if hasattr(item, "item") else item + except Exception: + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + +def payload_turn_id(payload: object) -> int | None: + if not isinstance(payload, Mapping): + return None + return coerce_int(payload.get("duplex_turn_id")) + + +def output_turn_id_from_metadata(mm_output: dict[str, object]) -> int | None: + return _first_metadata_int(mm_output, "duplex_turn_id", "turn_id") + + +def output_epoch_from_metadata(mm_output: dict[str, object]) -> int | None: + return _first_metadata_int(mm_output, "duplex_epoch", "epoch") + + +def _first_metadata_int(mm_output: dict[str, object], *names: str) -> int | None: + candidates: list[object] = [] + meta = mm_output.get("meta") + for name in names: + candidates.extend((mm_output.get(name), mm_output.get(f"meta.{name}"))) + if isinstance(meta, dict): + candidates.append(meta.get(name)) + for value in candidates: + result = coerce_int(value) + if result is not None: + return result + return None + + +def _runtime_result(**values: object) -> dict[str, object]: + return { + "supported": True, + **values, + "uses_model_runner_scheduler": True, + "runner_kv_backed": True, + "runtime_impl": "scheduler_data_plane", + "owned_runtime": False, + } + + +def _output_stage_metrics(output: object) -> dict[str, dict[str, object]] | None: + metrics = getattr(output, "metrics", None) + if not isinstance(metrics, Mapping): + return None + stage_metrics = metrics.get("stage_metrics") + if not isinstance(stage_metrics, Mapping): + return None + snapshot = { + str(stage_id): dict(values) for stage_id, values in stage_metrics.items() if isinstance(values, Mapping) + } + return snapshot or None + + +def _native_decision( + completion: object, + mm_output: dict[str, object], + *, + token_ids: list[int], + finished: bool, +) -> str | None: + if not finished: + return None + if mm_output.get("duplex_native_decision") == "listen" or mm_output.get("model_listen") is True: + return "listen" + listen_id = _special_token_ids(mm_output).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 + if coerce_int(stop_reason) == listen_id: + return "listen" + return "listen" if token_ids and token_ids[-1] == listen_id else None + + +def _special_token_ids(mm_output: dict[str, object]) -> dict[str, int]: + sources: list[object] = [] + raw_special = mm_output.get("special_token_ids") + if isinstance(raw_special, dict): + sources.append(raw_special) + raw_meta = mm_output.get("meta") + if isinstance(raw_meta, dict): + sources.append(raw_meta) + sources.append( + { + key.removeprefix("meta."): value + for key, value in mm_output.items() + if isinstance(key, str) and key.startswith("meta.") + } + ) + out: dict[str, int] = {} + for source in sources: + if not isinstance(source, dict): + continue + for key, value in source.items(): + if not isinstance(key, str): + continue + token_id = coerce_int(value) + if token_id is not None and token_id >= 0: + out[key] = token_id + return out + + +def _completion_token_ids(completion: object) -> list[int]: + if completion is None: + return [] + for candidate in ( + getattr(completion, "token_ids", None), + getattr(completion, "cumulative_token_ids", None), + ): + token_ids = _coerce_int_list(candidate) + if token_ids: + return token_ids + return [] + + +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 [item for raw in value if (item := coerce_int(raw)) is not None] + + +def _audio_num_samples(audio_data: object) -> int | None: + try: + import torch + + if isinstance(audio_data, torch.Tensor): + return int(audio_data.numel()) + return int(np.asarray(audio_data, dtype=np.float32).size) + except Exception: + return None + + +def _bool_metadata( + mm_output: dict[str, object], + names: tuple[str, ...], + *, + default: bool, +) -> bool: + def coerce(value: object) -> bool | None: + if value is None: + return None + if isinstance(value, bool): + return value + try: + import torch + + if isinstance(value, torch.Tensor): + return bool(value.reshape(-1)[-1].item()) if value.numel() else None + except Exception: + pass + if isinstance(value, np.ndarray): + return bool(value.reshape(-1)[-1].item()) if value.size else None + if isinstance(value, np.generic): + return bool(value.item()) + if isinstance(value, (list, tuple)): + return coerce(value[-1]) if value else None + if isinstance(value, (int, float)): + return bool(value) + return None + + meta = mm_output.get("meta") + for name in names: + for key in (name, f"meta.{name}"): + result = coerce(mm_output.get(key)) + if result is not None: + return result + if isinstance(meta, dict): + result = coerce(meta.get(name)) + if result is not None: + return result + return default + + +def _sample_rate_hz(mm_output: dict[str, object]) -> int: + sr_raw = mm_output.get("sr") + if sr_raw is None: + sr_raw = mm_output.get("sample_rate_hz", mm_output.get("sample_rate")) + meta = mm_output.get("meta") + if sr_raw is None and isinstance(meta, dict): + sr_raw = meta.get("sr") or meta.get("sample_rate_hz") or meta.get("sample_rate") + if sr_raw is None: + sr_raw = mm_output.get("meta.sr") or mm_output.get("meta.sample_rate_hz") or mm_output.get("meta.sample_rate") + if hasattr(sr_raw, "item"): + try: + return int(sr_raw.item()) + except Exception: + return 24000 + return int(sr_raw) if isinstance(sr_raw, (int, float)) else 24000 + + +def _audio_text_marks(mm_output: dict[str, object]) -> list[dict[str, object]]: + names = ("audio_text_marks", "text_audio_marks", "audio_text_alignment", "alignment_marks") + candidates = [mm_output.get(name) for name in names] + meta = mm_output.get("meta") + if isinstance(meta, dict): + candidates.extend(meta.get(name) for name in names) + candidates.extend( + value + for key, value in mm_output.items() + if isinstance(key, str) and key.startswith("meta.") and key.rsplit(".", 1)[-1] in names + ) + for raw in candidates: + if not isinstance(raw, list): + continue + marks: list[dict[str, object]] = [] + for item in raw: + if not isinstance(item, dict): + continue + text_chars = item.get("text_chars") + audio_end_ms = item.get("audio_end_ms", item.get("audio_ms")) + if isinstance(text_chars, (int, float)) and isinstance(audio_end_ms, (int, float)): + marks.append({"text_chars": max(0, int(text_chars)), "audio_end_ms": max(0, int(audio_end_ms))}) + if marks: + return marks + return [] + + +def _llm_output_text(mm_output: dict[str, object]) -> str: + candidates: list[object] = [ + mm_output.get("llm_output_text"), + mm_output.get("text"), + mm_output.get("llm_output_text_utf8"), + mm_output.get("meta.llm_output_text_utf8"), + ] + meta = mm_output.get("meta") + if isinstance(meta, dict): + candidates.extend((meta.get("llm_output_text"), meta.get("text"), meta.get("llm_output_text_utf8"))) + candidates.extend(mm_output.get(key) for key in ("meta.llm_output_text", "meta.text")) + for value in candidates: + if isinstance(value, str) and value: + return value + decoded = _decode_text_tensor(value) + if decoded: + return decoded + if isinstance(value, list): + text_chunks = [item for item in value if isinstance(item, str)] + if text_chunks: + return "".join(text_chunks) + return "" + + +def _decode_text_tensor(value: object) -> str: + if value is None: + return "" + try: + if isinstance(value, np.ndarray): + raw = value.astype(np.uint8, copy=False).reshape(-1).tobytes() + elif hasattr(value, "detach"): + raw = value.detach().cpu().numpy().astype(np.uint8, copy=False).reshape(-1).tobytes() + else: + return "" + return raw.decode("utf-8", errors="ignore") + except Exception: + return "" + + +def _fallback_audio_text_marks( + audio_chunks: list[tuple[str, int]], + delta_text: str, +) -> list[list[dict[str, int]] | None]: + if not delta_text: + return [] + total_duration_ms = sum(max(0, int(duration_ms)) for _, duration_ms in audio_chunks) + cumulative_duration_ms = 0 + marks: list[list[dict[str, int]] | None] = [] + for _, duration_ms in audio_chunks: + cumulative_duration_ms += max(0, int(duration_ms)) + if total_duration_ms <= 0: + marks.append(None) + continue + text_chars = int(len(delta_text) * min(1.0, cumulative_duration_ms / float(total_duration_ms))) + marks.append([{"text_chars": max(0, text_chars), "audio_end_ms": max(0, cumulative_duration_ms)}]) + return marks diff --git a/vllm_omni/experimental/fullduplex/minicpmo45/input.py b/vllm_omni/experimental/fullduplex/minicpmo45/input.py new file mode 100644 index 000000000..2a4eac3ea --- /dev/null +++ b/vllm_omni/experimental/fullduplex/minicpmo45/input.py @@ -0,0 +1,384 @@ +from __future__ import annotations + +import base64 +import binascii +from typing import Any + +import numpy as np + + +class MiniCPMO45PcmAppendReservation: + __slots__ = ( + "_active", + "_force_listen", + "_is_speech", + "_owner", + "_raw", + "_sample_rate_hz", + "_turn_had_speech", + "_video_frames", + "operation_id", + "payload", + ) + + def __init__( + self, + *, + owner: MiniCPMO45PcmAppendBuffer, + operation_id: str, + payload: dict[str, object] | None, + raw: bytes, + sample_rate_hz: int, + force_listen: bool, + is_speech: bool, + turn_had_speech: bool = False, + video_frames: list[str] | None = None, + ) -> None: + self._owner = owner + self.operation_id = operation_id + self.payload = payload + self._raw = raw + self._sample_rate_hz = sample_rate_hz + self._force_listen = force_listen + self._is_speech = is_speech + self._turn_had_speech = turn_had_speech + self._video_frames = list(video_frames or []) + self._active = True + + @property + def active(self) -> bool: + return self._active + + @property + def byte_count(self) -> int: + return len(self._raw) + + def commit(self) -> None: + self._owner._commit_reservation(self) + + def rollback(self) -> None: + self._owner._rollback_reservation(self) + + +def validate_native_ref_audio_config(session_config: dict[str, Any]) -> None: + extra_body = session_config.get("extra_body") + if not isinstance(extra_body, dict): + extra_body = {} + if any(key in session_config for key in ("ref_audio_path", "tts_ref_audio_path")) or any( + key in extra_body for key in ("ref_audio_path", "tts_ref_audio_path") + ): + 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: + validate_native_ref_audio_config(session_config) + extra_body = session_config.get("extra_body") + if not isinstance(extra_body, dict): + extra_body = {} + audio_data = extra_body.get("ref_audio_data") + if audio_data is None: + return None + if not isinstance(audio_data, str): + raise TypeError("native duplex ref_audio_data must be base64 pcm_f32le") + fmt = extra_body.get("ref_audio_format") or "pcm_f32le" + if fmt != "pcm_f32le": + raise ValueError(f"unsupported native duplex ref_audio_format: {fmt!r}") + try: + raw = base64.b64decode(audio_data, validate=True) + except (ValueError, binascii.Error) as exc: + raise ValueError("invalid native duplex ref_audio_data") from exc + if len(raw) % 4: + raise ValueError("invalid native duplex ref_audio_data length") + return np.frombuffer(raw, dtype=" None: + self._buffer = bytearray() + self._sample_rate_hz: int | None = None + self._force_listen = False + self._is_speech = False + self._turn_had_speech = False + self._reservation_seq = 0 + self._reservations: list[MiniCPMO45PcmAppendReservation] = [] + # Omni duplex: queued camera frames (base64 JPEG), consumed FIFO at + # one frame per emitted model unit alongside the unit's audio. + self._frame_queue: list[str] = [] + + def clear(self) -> None: + for reservation in self._reservations: + reservation._active = False + self._reservations.clear() + self._buffer.clear() + self._frame_queue.clear() + self._sample_rate_hz = None + self._force_listen = False + self._is_speech = False + self._turn_had_speech = False + + def clear_force_listen(self) -> None: + self._force_listen = False + + def has_pending(self) -> bool: + return bool(self._buffer) + + def has_reserved(self) -> bool: + return bool(self._reservations) + + @property + def pending_byte_count(self) -> int: + return len(self._buffer) + + def _reserve_passthrough( + self, + payload: dict[str, object], + *, + operation_id: str, + ) -> MiniCPMO45PcmAppendReservation: + sample_rate_hz = payload.get("sample_rate_hz") + if isinstance(payload, dict) and "video_frames" in payload: + # Frames align to whole audio units; a passthrough payload has no + # unit framing for Stage0 to interleave against. + payload = {key: value for key, value in payload.items() if key != "video_frames"} + reservation = MiniCPMO45PcmAppendReservation( + owner=self, + operation_id=operation_id, + payload=payload, + raw=b"", + sample_rate_hz=sample_rate_hz if isinstance(sample_rate_hz, int) else 0, + force_listen=bool(payload.get("force_listen", False)), + is_speech=bool(payload.get("is_speech", False)), + turn_had_speech=bool(payload.get("is_speech", False)), + ) + self._reservations.append(reservation) + return reservation + + def prepare_append( + self, + payload: dict[str, object], + *, + operation_id: str, + chunk_period_ms: int, + flush: bool = False, + allow_emit: bool = True, + ) -> MiniCPMO45PcmAppendReservation | None: + fmt = payload.get("format") + sample_rate_hz = payload.get("sample_rate_hz") + audio = payload.get("audio") + if fmt != "pcm_f32le" or not isinstance(sample_rate_hz, int) or not isinstance(audio, str): + return self._reserve_passthrough(payload, operation_id=operation_id) + try: + raw = base64.b64decode(audio, validate=True) + except (binascii.Error, ValueError): + return self._reserve_passthrough(payload, operation_id=operation_id) + if len(raw) % 4 != 0: + return self._reserve_passthrough(payload, operation_id=operation_id) + + if self._sample_rate_hz is not None and self._sample_rate_hz != sample_rate_hz: + raise ValueError("MiniCPM-o native duplex audio append sample_rate_hz changed within a session") + self._sample_rate_hz = sample_rate_hz + self._buffer.extend(raw) + frames_in = payload.get("video_frames") + if isinstance(frames_in, list): + self._frame_queue.extend(frame for frame in frames_in if isinstance(frame, str) and frame) + self._turn_had_speech = self._turn_had_speech or bool(payload.get("is_speech", False)) + self._force_listen = self._force_listen or bool(payload.get("force_listen", False)) + self._is_speech = self._is_speech or bool(payload.get("is_speech", False)) + if not allow_emit: + return None + + min_samples = max(1, int(sample_rate_hz * max(1, int(chunk_period_ms)) / 1000)) + buffered_samples = len(self._buffer) // 4 + if not flush and buffered_samples < min_samples: + return None + + # Emit whole model chunks only: the engine reserves scheduler slots + # from the payload size and the worker consumes whole chunks per + # append, so partial chunks would turn into pad embeddings inside the + # model KV. On flush, the tail is zero-padded (silence) up to the + # chunk boundary instead. + if flush: + emit_samples = buffered_samples + remainder = emit_samples % min_samples + pad_samples = (min_samples - remainder) if remainder else 0 + else: + emit_samples = min_samples + pad_samples = 0 + emit_bytes = emit_samples * 4 + reserved_raw = bytes(self._buffer[:emit_bytes]) + emit_raw = reserved_raw + b"\x00" * (pad_samples * 4) + del self._buffer[:emit_bytes] + + out = dict(payload) + out.pop("force_speak", None) + out.pop("video_frames", None) + out["audio"] = base64.b64encode(emit_raw).decode("ascii") + out["sample_rate_hz"] = sample_rate_hz + # Omni duplex: attach at most one queued camera frame per emitted + # model unit (official cadence: one frame per 1 s chunk). The engine + # budgets 66 scheduler slots per attached frame from this payload. + # Official omni cadence is one frame per ~1 s chunk, and the first + # append consumes extra samples (1035 ms first window), so per-unit + # attachment could outrun the units Stage0 actually builds. Attach at + # most ONE frame per emitted payload; the rest stay queued. + attached_frames: list[str] = [] + if emit_samples + pad_samples >= min_samples and self._frame_queue: + attached_frames = [self._frame_queue.pop(0)] + out["video_frames"] = attached_frames + out["force_listen"] = self._force_listen + out["is_speech"] = self._is_speech + if not self._buffer: + self._force_listen = False + self._is_speech = False + reservation = MiniCPMO45PcmAppendReservation( + owner=self, + operation_id=operation_id, + payload=out, + raw=reserved_raw, + sample_rate_hz=sample_rate_hz, + force_listen=bool(out.get("force_listen", False)), + is_speech=bool(out.get("is_speech", False)), + turn_had_speech=self._turn_had_speech, + video_frames=attached_frames, + ) + self._reservations.append(reservation) + return reservation + + def prepare_commit( + self, + *, + operation_id: str, + chunk_period_ms: int, + ) -> MiniCPMO45PcmAppendReservation: + """Reserve the terminal payload without invalidating prior appends.""" + had_speech = self._turn_had_speech + reservation: MiniCPMO45PcmAppendReservation | None = None + if had_speech and self._buffer: + payload: dict[str, object] = { + "type": "audio", + "audio": "", + "format": "pcm_f32le", + "sample_rate_hz": self._sample_rate_hz or 16000, + "force_listen": self._force_listen, + "is_speech": self._is_speech, + } + reservation = self.prepare_append( + payload, + operation_id=operation_id, + chunk_period_ms=chunk_period_ms, + flush=True, + ) + assert reservation is not None + assert reservation.payload is not None + reservation.payload["final"] = True + if reservation is None: + reservation = MiniCPMO45PcmAppendReservation( + owner=self, + operation_id=operation_id, + payload=None, + raw=b"", + sample_rate_hz=self._sample_rate_hz or 0, + force_listen=self._force_listen, + is_speech=self._is_speech, + turn_had_speech=had_speech, + ) + self._reservations.append(reservation) + + # Open the next physical input generation without touching already + # reserved appends. A rollback restores this generation's metadata. + self._sample_rate_hz = None + self._force_listen = False + self._is_speech = False + self._turn_had_speech = False + return reservation + + def append( + self, + payload: dict[str, object], + *, + chunk_period_ms: int, + flush: bool = False, + allow_emit: bool = True, + ) -> dict[str, object] | None: + self._reservation_seq += 1 + reservation = self.prepare_append( + payload, + operation_id=f"immediate-{self._reservation_seq}", + chunk_period_ms=chunk_period_ms, + flush=flush, + allow_emit=allow_emit, + ) + if reservation is None: + return None + reservation.commit() + assert reservation.payload is not None + return reservation.payload + + def _commit_reservation(self, reservation: MiniCPMO45PcmAppendReservation) -> None: + if not reservation._active: + return + if not self._reservations or self._reservations[0] is not reservation: + raise RuntimeError("PCM append reservations must commit in wire order") + self._reservations.pop(0) + reservation._active = False + + def _rollback_reservation(self, reservation: MiniCPMO45PcmAppendReservation) -> None: + if not reservation._active: + return + try: + index = self._reservations.index(reservation) + except ValueError: + reservation._active = False + return + rolled_back = self._reservations[index:] + restored = b"".join(item._raw for item in rolled_back) + self._buffer[:0] = restored + restored_frames = [frame for item in rolled_back for frame in item._video_frames] + if restored_frames: + self._frame_queue[:0] = restored_frames + self._sample_rate_hz = self._sample_rate_hz or reservation._sample_rate_hz + self._force_listen = self._force_listen or any(item._force_listen for item in rolled_back) + self._is_speech = self._is_speech or any(item._is_speech for item in rolled_back) + self._turn_had_speech = self._turn_had_speech or any(item._turn_had_speech for item in rolled_back) + for item in rolled_back: + item._active = False + del self._reservations[index:] + + def flush(self, *, chunk_period_ms: int) -> dict[str, object] | None: + if not self._buffer: + return None + payload: dict[str, object] = { + "type": "audio", + "audio": "", + "format": "pcm_f32le", + "sample_rate_hz": self._sample_rate_hz or 16000, + "force_listen": self._force_listen, + "is_speech": self._is_speech, + } + return self.append(payload, chunk_period_ms=chunk_period_ms, flush=True) + + def commit(self, *, chunk_period_ms: int) -> dict[str, object] | None: + """Commit the real residual PCM and reset per-client-input state. + + Complete model units are emitted by :meth:`append` as soon as they are + available. A commit at that exact boundary therefore has no payload; + synthesizing another unit would add a model decision that does not + exist in the official continuous streaming loop. + """ + self._reservation_seq += 1 + reservation = self.prepare_commit( + operation_id=f"immediate-commit-{self._reservation_seq}", + chunk_period_ms=chunk_period_ms, + ) + reservation.commit() + return reservation.payload + + +__all__ = [ + "MiniCPMO45PcmAppendReservation", + "MiniCPMO45PcmAppendBuffer", + "decode_native_ref_audio_from_config", + "validate_native_ref_audio_config", +] diff --git a/vllm_omni/experimental/fullduplex/minicpmo45/policy.py b/vllm_omni/experimental/fullduplex/minicpmo45/policy.py new file mode 100644 index 000000000..c55b25360 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/minicpmo45/policy.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +from typing import Any + + +class MiniCPMO45DuplexPolicy: + """MiniCPM-o 4.5 native duplex model policy. + + Keep model-specific token names, input modes, and handoff type strings out + of the generic scheduler/orchestrator path. These names are part of the + MiniCPM-o 4.5 remote-code contract, not a general vLLM-Omni duplex schema. + """ + + # Audio framing contract shared by serving, orchestrator, and worker. + # MiniCPM-o consumes 1 s units at 16 kHz and pools audio to one embedding + # per 100 ms, so a unit contributes exactly 10 audio embeddings plus the + # open (and a closure for every unit after the first). + # Scheduler token budgets must match the worker-built embeddings exactly: + # surplus slots become pad embeddings inside the KV and measurably corrupt + # the model's listen/speak behavior. + SAMPLE_RATE_HZ = 16000 + CHUNK_SAMPLES = 16000 + SAMPLES_PER_AUDIO_TOKEN = 1600 + # Vision framing contract (omni duplex). Official streaming_prefill feeds + # each frame as + 64 resampler embeddings + inside the + # unit, ahead of the unit's audio embeddings (max_slice_nums=1 in + # streaming, so exactly one 64-token block per frame). + VISION_EMBEDS_PER_FRAME = 64 + VISION_TOKENS_PER_FRAME = VISION_EMBEDS_PER_FRAME + 2 # + embeds + + DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK = 20 + DEFAULT_MAX_SPEAK_CHARS_PER_CHUNK = 28 + DEFAULT_MIN_NEW_SPEAK_TOKENS_BEFORE_CHUNK_BOUNDARY = 8 + REPETITION_HISTORY_SIZE = 512 + + @classmethod + def audio_token_count(cls, sample_count: int) -> int: + """Audio embedding count for a clip of ``sample_count`` samples. + + Matches the whole-clip encoder math (hop 160 mel frames -> CNN stride 2 + -> avg-pool 5) for any multiple of ``SAMPLES_PER_AUDIO_TOKEN``; serving + normalizes clips to that boundary before this is used for budgets. + """ + return max(0, int(sample_count) // cls.SAMPLES_PER_AUDIO_TOKEN) + + @staticmethod + def session_context_texts(instructions: object, has_ref_audio: bool) -> tuple[str, str]: + """System-context prefix/suffix, matching MiniCPMODuplex.prepare().""" + system_prompt = ( + instructions if isinstance(instructions, str) and instructions else "Streaming Omni Conversation." + ) + prefix = f"<|im_start|>system\n{system_prompt}" + suffix = "<|im_end|>" + if has_ref_audio: + prefix += "\n<|audio_start|>" + suffix = "<|audio_end|>" + suffix + return prefix, suffix + + SPECIAL_TOKEN_FIELDS: dict[str, str] = { + "unit_token_id": "", + "unit_end_token_id": "", + "listen_token_id": "<|listen|>", + "speak_token_id": "<|speak|>", + "tts_bos_token_id": "<|tts_bos|>", + "tts_eos_token_id": "<|tts_eos|>", + "tts_pad_token_id": "<|tts_pad|>", + "chunk_eos_token_id": "<|chunk_eos|>", + "chunk_tts_eos_token_id": "<|chunk_tts_eos|>", + "turn_eos_token_id": "<|turn_eos|>", + } + OPTIONAL_TOKEN_FIELDS: dict[str, str] = { + "audio_placeholder_token_id": "<|audio|>", + "image_start_token_id": "", + "image_end_token_id": "", + } + + @classmethod + def token_ids_from_tokenizer(cls, tokenizer: Any) -> dict[str, int]: + convert = getattr(tokenizer, "convert_tokens_to_ids", None) + + def token_id(token: str) -> int: + if not callable(convert): + value = None + else: + value = convert(token) + if isinstance(value, list): + value = value[0] if len(value) == 1 else None + unk_token_id = getattr(tokenizer, "unk_token_id", None) + try: + candidate = int(value) + except (TypeError, ValueError): + candidate = -1 + if candidate >= 0 and candidate != unk_token_id: + return candidate + + 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: + candidate = int(ids[0]) + except (TypeError, ValueError): + candidate = -1 + if candidate >= 0 and candidate != unk_token_id: + return candidate + return -1 + + resolved = {field: token_id(token) for field, token in cls.SPECIAL_TOKEN_FIELDS.items()} + resolved.update({field: token_id(token) for field, token in cls.OPTIONAL_TOKEN_FIELDS.items()}) + if resolved.get("listen_token_id", -1) < 0: + eos_id = getattr(tokenizer, "eos_token_id", None) + try: + resolved["listen_token_id"] = int(eos_id) + except (TypeError, ValueError): + pass + return resolved + + @staticmethod + def native_forbidden_token_ids( + token_ids: dict[str, int], + *, + bad_token_ids: list[int] | tuple[int, ...] = (), + ) -> list[int]: + return [ + token_ids.get("tts_pad_token_id", -1), + *list(bad_token_ids), + token_ids.get("chunk_eos_token_id", -1), + ] + + @staticmethod + def native_special_token_ids( + token_ids: dict[str, int], + *, + tokenizer_special_ids: list[int] | tuple[int, ...] = (), + ) -> set[int]: + special = set(tokenizer_special_ids or []) + special.update(value for value in token_ids.values() if isinstance(value, int) and value >= 0) + return special diff --git a/vllm_omni/experimental/fullduplex/minicpmo45/runtime.py b/vllm_omni/experimental/fullduplex/minicpmo45/runtime.py new file mode 100644 index 000000000..b5fbd37a5 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/minicpmo45/runtime.py @@ -0,0 +1,354 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from base64 import b64decode +from binascii import Error as BinasciiError +from typing import Any + +from vllm.sampling_params import SamplingParams + +from vllm_omni.experimental.fullduplex.engine.duplex_runtime import ( + DuplexAppendPlan, + DuplexInputMode, + DuplexOutputAction, + DuplexOutputDecision, +) +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence + +_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 + + +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_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 bool(sample_count) 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_frame_count(payload) * _DUPLEX_VISION_TOKENS_PER_FRAME + 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_frame_count(payload) * _DUPLEX_VISION_TOKENS_PER_FRAME + ) + 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 = 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: + if hasattr(value, "detach"): + try: + value = value.detach().cpu().reshape(-1) + if value.numel() == 0: + return None + value = value[0].item() + except Exception: + return None + try: + return int(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/experimental/fullduplex/minicpmo45/serving_adapter.py b/vllm_omni/experimental/fullduplex/minicpmo45/serving_adapter.py new file mode 100644 index 000000000..e811d73e5 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/minicpmo45/serving_adapter.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Any + +from vllm_omni.experimental.fullduplex.minicpmo45.adapter import ( + MiniCPMO45NativeDuplexServingAdapter, +) +from vllm_omni.experimental.fullduplex.minicpmo45.data_plane import ( + MiniCPMO45DataPlaneContext, + MiniCPMO45DataPlaneSession, +) +from vllm_omni.experimental.fullduplex.minicpmo45.session import ( + MiniCPMO45ServingSessionState, +) +from vllm_omni.experimental.fullduplex.openai.protocol import DuplexCapabilities + +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 DuplexCapabilities.minicpmo45_native(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/experimental/fullduplex/minicpmo45/session.py b/vllm_omni/experimental/fullduplex/minicpmo45/session.py new file mode 100644 index 000000000..763bcf8b3 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/minicpmo45/session.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field + +from vllm_omni.experimental.fullduplex.minicpmo45.input import ( + MiniCPMO45PcmAppendBuffer, +) + + +@dataclass(slots=True) +class MiniCPMO45ServingSessionState: + """Mutable serving state owned by one MiniCPM duplex session.""" + + audio_buffer: MiniCPMO45PcmAppendBuffer = field(default_factory=MiniCPMO45PcmAppendBuffer) + input_since_commit: bool = False + speech_since_commit: 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 diff --git a/vllm_omni/experimental/fullduplex/minicpmo45/stage0.py b/vllm_omni/experimental/fullduplex/minicpmo45/stage0.py new file mode 100644 index 000000000..84aaf1cb8 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/minicpmo45/stage0.py @@ -0,0 +1,795 @@ +from __future__ import annotations + +import base64 +import copy +import sys +import time +from contextlib import suppress +from dataclasses import dataclass, field +from typing import Any + +import numpy as np + +from vllm_omni.experimental.fullduplex.minicpmo45.policy import MiniCPMO45DuplexPolicy + +_MINICPMO45_SPECIAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS +_MINICPMO45_OPTIONAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.OPTIONAL_TOKEN_FIELDS + + +@dataclass +class _MiniCPMO45Stage0SessionState: + session_id: str + streaming_processor: Any | None = None + audio_buffer: np.ndarray = field(default_factory=lambda: np.array([], dtype=np.float32)) + audio_chunk_idx: int = 0 + context_embeds: list[Any] = 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_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 + pending_terminator_token: int | None = None + last_terminator_token: int | None = None + pending_speech_context: bool = False + pending_speech_append_identity: tuple[int | None, int] | None = None + pending_speech_response_open: bool = False + generated_text_tokens: list[int] = field(default_factory=list) + + +class MiniCPMO45Stage0DuplexRuntime: + """Build scheduler-owned MiniCPM-o 4.5 Stage0 duplex inputs.""" + + def __init__(self, stage_model: Any, *, model_path: str | None = None, device: str = "cuda") -> None: + self.stage_model = stage_model + self.model_path = model_path + self.device = device + self.sessions: dict[tuple[str, int], _MiniCPMO45Stage0SessionState] = {} + self.thinker = getattr(stage_model, "thinker", None) or getattr(stage_model, "model", None) or stage_model + self.processor = ( + getattr(stage_model, "processor", None) + or getattr(self.thinker, "processor", None) + or self._load_processor_from_path(model_path) + ) + self.tokenizer = ( + getattr(self.processor, "tokenizer", None) + if self.processor is not None + else getattr(stage_model, "tokenizer", None) + ) + self._init_token_ids() + + def _stage_runtime_ready(self) -> bool: + return self.processor is not None and self.tokenizer is not None and self.thinker is not None + + def _configure_streaming_processor( + self, + state: _MiniCPMO45Stage0SessionState | None = None, + ) -> Any | None: + processor = self.processor + if processor is None: + return None + if state is not None: + if state.streaming_processor is not None: + return state.streaming_processor + processor = copy.copy(processor) + shared_mel = getattr(self.processor, "_streaming_mel_processor", None) + if shared_mel is not None: + processor._streaming_mel_processor = copy.deepcopy(shared_mel) + state.streaming_processor = processor + if processor is None: + return + set_streaming_mode = getattr(processor, "set_streaming_mode", None) + if callable(set_streaming_mode): + set_streaming_mode( + mode="exact", + chunk_ms=int(self._stage_param("chunk_ms", 1000)), + first_chunk_ms=int(self._stage_param("first_chunk_ms", 1035)), + cnn_redundancy_ms=int(self._stage_param("cnn_redundancy_ms", 20)), + enable_sliding_window=True, + slide_trigger_seconds=30.0, + slide_stride_seconds=10.0, + ) + # Match official init_streaming_processor: reset the streaming mel-processor + # buffers at session init (modeling_minicpmo_unified.py:207). + reset_streaming = getattr(processor, "reset_streaming", None) + if callable(reset_streaming): + reset_streaming() + return processor + configure_streaming = getattr(processor, "configure_streaming", None) + if callable(configure_streaming): + configure_streaming( + chunk_ms=int(self._stage_param("chunk_ms", 1000)), + enable_sliding_window=True, + slide_trigger_seconds=30.0, + slide_stride_seconds=10.0, + ) + return processor + + def _prepare_session_context( + self, + state: _MiniCPMO45Stage0SessionState, + session_config: dict[str, Any], + *, + runtime_config: dict[str, Any] | None = None, + ) -> None: + if not self._stage_runtime_ready(): + return + self._require_special_token_ids() + ref_audio = self._decode_ref_audio_from_session_config(runtime_config or {}) + # Matches MiniCPMODuplex.prepare() in the released checkpoint's + # modeling_minicpmo.py: the <|audio_start|>/<|audio_end|> markers are + # only present when reference audio is embedded between them. The + # template is shared with the serving adapter so the first-append + # scheduler reserve can count these tokens exactly. + prefix, suffix = MiniCPMO45DuplexPolicy.session_context_texts( + session_config.get("instructions"), + ref_audio is not None, + ) + for token_id in self._encode_text(prefix): + state.context_embeds.append(self._embed_token(token_id)) + state.context_token_ids.append(token_id) + if ref_audio is not None: + ref_audio_embeds = self._stage_ref_audio_embeddings(ref_audio, state=state) + if ref_audio_embeds is not None: + ref_audio_embeds = self._as_2d_tensor(ref_audio_embeds) + state.context_embeds.append(ref_audio_embeds) + state.context_token_ids.extend([self.unit_token_id] * int(ref_audio_embeds.shape[0])) + for token_id in self._encode_text(suffix): + state.context_embeds.append(self._embed_token(token_id)) + state.context_token_ids.append(token_id) + + def _stage_prefill_embeddings_only( + self, + state: _MiniCPMO45Stage0SessionState, + audio_waveform: Any, + *, + video_frames: list[Any] | None = None, + epoch: int | None = None, + seq: int | None = None, + is_speech: bool = False, + final: bool = False, + ) -> dict[str, Any]: + """Build scheduler-owned Stage0 input embeddings for one audio append. + + Unlike the legacy worker-control path, this method never calls an eager + model forward. The normal vLLM runner consumes the returned embeddings + and owns attention metadata, block tables, KV cache, and sampling. + """ + start_time = time.time() + processor = self._configure_streaming_processor(state) + append_identity = (epoch, seq) if seq is not None else None + if ( + append_identity is not None + and state.prepared_append_identity == append_identity + and state.prepared_inputs_embeds is not None + ): + result = dict(state.prepared_result) + result["inputs_embeds"] = state.prepared_inputs_embeds + result["input_token_ids"] = list(state.prepared_input_token_ids) + return result + self._require_special_token_ids() + if audio_waveform is None or len(audio_waveform) == 0: + return self._stage_prefill_result(False, start_time, "empty audio") + # Omni duplex: encode this append's camera frames up front so the unit + # loop can interleave one block per unit, mirroring official + # streaming_prefill (feed , then image embeds, then audio). + frame_blocks: list[Any] = [] + if video_frames: + self._require_vision_token_ids() + frame_blocks = self._stage_vision_embeddings(video_frames) + if frame_blocks is None or len(frame_blocks) != len(video_frames): + return self._stage_prefill_result(False, start_time, "streaming vision embedding failed") + state.audio_buffer = np.concatenate([state.audio_buffer, np.asarray(audio_waveform, dtype=np.float32)]) + chunk_size = self._streaming_chunk_size(processor) + self._pad_first_audio_chunk_if_needed(state, processor) + if len(state.audio_buffer) < chunk_size: + return self._stage_prefill_result( + False, + start_time, + f"audio not enough: need {chunk_size} samples, only {len(state.audio_buffer)}", + ) + + embed_parts: list[Any] = [] + token_ids: list[int] = [] + if state.audio_chunk_idx == 0 and state.context_embeds: + embed_parts.extend(state.context_embeds) + token_ids.extend(state.context_token_ids) + + # Consume every complete chunk in the buffer so the appended span and + # the scheduler's slot reservation for this append agree exactly. A + # final append may zero-pad a real residual chunk, but it must not add + # a whole silence unit after all input was already consumed: official + # duplex generation runs once per microphone unit and client commit is + # not an additional model decision. + units_built = 0 + while True: + if len(state.audio_buffer) < chunk_size: + if not final or len(state.audio_buffer) == 0: + break + pad = np.zeros(chunk_size - len(state.audio_buffer), dtype=np.float32) + state.audio_buffer = np.concatenate([state.audio_buffer, pad]) + audio_chunk = state.audio_buffer[:chunk_size] + batch_feature = self._process_streaming_audio(audio_chunk, state.audio_chunk_idx, processor=processor) + for name, value in ( + ("chunk_idx", state.audio_chunk_idx), + ("use_extra_context", True), + ("prefix_extra_frames", 0 if state.audio_chunk_idx == 0 else 2), + ("suffix_extra_frames", 2), + ): + with suppress(Exception): + setattr(batch_feature, name, value) + audio_embeds = self._stage_audio_embeddings(batch_feature, state=state) + if audio_embeds is None: + if units_built == 0: + return self._stage_prefill_result(False, start_time, "streaming audio embedding returned empty") + break + if state.audio_chunk_idx > 0: + # Official duplex closes every unit (finalize_unit feeds the + # sampled terminator + ) before the next opens. + # The scheduler session update discards the previous segment's + # sampled terminator token, so it is re-injected here ahead of + # the closure; the model's listen/speak policy depends on + # seeing its own past decisions in context. + pending_terminator = state.pending_terminator_token + if pending_terminator is not None and units_built == 0: + state.pending_terminator_token = None + embed_parts.append(self._embed_token(pending_terminator)) + token_ids.append(int(pending_terminator)) + embed_parts.append(self._embed_token(self.unit_end_token_id)) + token_ids.append(self.unit_end_token_id) + embed_parts.append(self._embed_token(self.unit_token_id)) + token_ids.append(self.unit_token_id) + if frame_blocks: + # Official order inside a unit: + 64 resampler + # embeddings + ahead of the unit's audio embeddings + # (max_slice_nums=1 in streaming, one frame per unit). + vision_block = self._as_2d_tensor(frame_blocks.pop(0)) + embed_parts.append(self._embed_token(self.image_start_token_id)) + token_ids.append(int(self.image_start_token_id)) + embed_parts.append(vision_block) + token_ids.extend([self._vision_embedding_placeholder_token_id()] * int(vision_block.shape[0])) + embed_parts.append(self._embed_token(self.image_end_token_id)) + token_ids.append(int(self.image_end_token_id)) + embed_parts.append(audio_embeds) + token_ids.extend( + [self._audio_embedding_placeholder_token_id()] * int(self._as_2d_tensor(audio_embeds).shape[0]) + ) + state.audio_buffer = state.audio_buffer[ + self._consumed_audio_samples(state.audio_chunk_idx, chunk_size, processor=processor) : + ] + state.audio_chunk_idx += 1 + units_built += 1 + chunk_size = self._streaming_chunk_size(processor) + if frame_blocks: + # Each frame reserves one image block in the append plan. A frame + # without a matching audio unit would desynchronize that plan. + return self._stage_prefill_result( + False, + start_time, + f"{len(frame_blocks)} video frame(s) left without a matching audio unit", + ) + # Match official streaming_prefill: per chunk feed ONLY +audio. The assistant + # turn is opened once at session init; re-emitting the turn-open prefix per chunk + # re-opened the turn each chunk -> degenerate repetition. tts_bos/listen/turn_eos are + # model-generated and tracked via current_turn_ended (mirrors streaming_generate). + prompt_suffix_len = 0 + + import torch + + inputs_embeds = torch.cat([self._as_2d_tensor(embed) for embed in embed_parts], dim=0) + result = self._stage_prefill_result(True, start_time) + result.update( + { + "inputs_embeds": inputs_embeds, + "input_token_ids": token_ids, + "special_token_ids": self._special_token_ids(), + "num_input_tokens": int(inputs_embeds.shape[0]), + "prompt_suffix_len": prompt_suffix_len, + "uses_model_runner_scheduler": True, + "runner_kv_backed": True, + "runtime_impl": "scheduler_data_plane", + } + ) + if is_speech and (append_identity is None or state.pending_speech_append_identity != append_identity): + state.pending_speech_context = True + state.pending_speech_append_identity = append_identity + if append_identity is not None: + state.prepared_append_identity = append_identity + state.prepared_inputs_embeds = inputs_embeds + state.prepared_input_token_ids = list(token_ids) + state.prepared_result = {k: v for k, v in result.items() if k not in {"inputs_embeds", "input_token_ids"}} + return result + + @staticmethod + def _stage_prefill_result(success: bool, start_time: float, reason: str = "") -> dict[str, Any]: + return { + "success": success, + "prefill_success": success, + "is_buffering": not success, + "reason": reason, + "cost_all": time.time() - start_time, + "stage_runtime_ready": True, + } + + @staticmethod + def _as_2d_tensor(value: Any) -> Any: + 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: + import torch + + token = torch.tensor([int(token_id)], dtype=torch.long, device=self._model_device()) + embedder = self._token_embedder() + embeds = embedder(token) + return self._as_2d_tensor(embeds) + + def _token_embedder(self) -> Any: + nested_embed = getattr(getattr(getattr(self.thinker, "llm", None), "model", None), "embed_tokens", None) + if callable(nested_embed): + return nested_embed + for target in (self.thinker, self.stage_model): + embedder = getattr(target, "get_input_embeddings", None) + if callable(embedder): + try: + embeddings = embedder() + if callable(embeddings): + return embeddings + except TypeError: + return embedder + raise AttributeError("MiniCPM-o stage0 model does not expose token embeddings") + + def _model_device(self) -> Any: + try: + return next(self.thinker.parameters()).device + except Exception: + pass + try: + return next(self.stage_model.parameters()).device + except Exception: + pass + return self.device + + def _streaming_chunk_size(self, processor: Any | None = None) -> int: + 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: + processor = processor or self.processor + return int( + self._stage_param( + "sample_rate", + getattr(getattr(processor, "_streaming_mel_processor", None), "sample_rate", 16000), + ) + ) + + def _first_chunk_samples(self, default_chunk_size: int, processor: Any | None = None) -> int: + processor = processor or self.processor + if getattr(processor, "_streaming_mel_processor", None) is None: + return default_chunk_size + return int(self._stage_param("first_chunk_ms", 1035) * self._sample_rate(processor) / 1000) + + def _pad_first_audio_chunk_if_needed( + self, + state: _MiniCPMO45Stage0SessionState, + processor: Any | None = None, + ) -> None: + if state.audio_chunk_idx != 0 or len(state.audio_buffer) == 0: + return + first_chunk_samples = self._first_chunk_samples( + self._streaming_chunk_size(processor), + processor, + ) + if len(state.audio_buffer) >= first_chunk_samples: + return + padding = 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: + for target in (self.stage_model, self.thinker, getattr(self.thinker, "llm", None)): + value = getattr(target, name, None) + if value is not None: + return value + value = getattr(target, name.upper(), None) + if value is not None: + return value + return default + + def _consumed_audio_samples( + self, + chunk_idx: int, + default_chunk_size: int, + *, + processor: Any | None = None, + ) -> int: + processor = processor or self.processor + if chunk_idx != 0: + chunk_ms = int(self._stage_param("chunk_ms", 1000)) + return int(chunk_ms * self._sample_rate(processor) / 1000) + mel_processor = getattr(processor, "_streaming_mel_processor", None) + get_config = getattr(mel_processor, "get_config", None) + if callable(get_config): + cfg = get_config() + if isinstance(cfg, dict): + consumed_ms = int(cfg.get("effective_first_chunk_ms", self._stage_param("first_chunk_ms", 1035))) + return int(consumed_ms * self._sample_rate(processor) / 1000) + return default_chunk_size + + def _process_streaming_audio( + self, + audio_chunk: Any, + chunk_idx: int, + *, + processor: Any | None = None, + ) -> Any: + processor = processor or self.processor + process = getattr(processor, "process_audio_streaming", None) + if callable(process): + try: + return process(audio_chunk, reset=False, return_batch_feature=True) + except TypeError: + return process(audio_chunk, chunk_idx=chunk_idx) + return {"audio_features": audio_chunk, "audio_feature_lens": [[len(audio_chunk)]]} + + def _stage_audio_embeddings( + self, + batch_feature: Any, + *, + state: _MiniCPMO45Stage0SessionState | None = None, + ) -> Any | None: + if hasattr(batch_feature, "to"): + batch_feature = batch_feature.to(self.device) + self._ensure_dynamic_cache_compat() + has_audio_cache = state is not None and hasattr(self.thinker, "audio_past_key_values") + previous_audio_past_key_values = ( + getattr(self.thinker, "audio_past_key_values", None) if has_audio_cache else None + ) + if has_audio_cache: + self.thinker.audio_past_key_values = state.audio_past_key_values + try: + for target in (self.stage_model, self.thinker): + get_streaming = getattr(target, "get_audio_embedding_streaming", None) + if callable(get_streaming): + try: + result = self._cat_nested_tensors( + get_streaming( + batch_feature, + use_extra_context=True, + prefix_extra_frames=0 if int(getattr(batch_feature, "chunk_idx", 0)) == 0 else 2, + suffix_extra_frames=2, + ) + ) + if has_audio_cache: + state.audio_past_key_values = getattr(self.thinker, "audio_past_key_values", None) + return result + except TypeError: + result = self._cat_nested_tensors(get_streaming(batch_feature)) + if has_audio_cache: + state.audio_past_key_values = getattr(self.thinker, "audio_past_key_values", None) + return result + get_hidden = getattr(target, "get_audio_hidden_states", None) + if callable(get_hidden): + result = self._cat_nested_tensors(get_hidden(batch_feature)) + if has_audio_cache: + state.audio_past_key_values = getattr(self.thinker, "audio_past_key_values", None) + return result + return None + finally: + if has_audio_cache: + 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: + from vllm_omni.experimental.fullduplex.minicpmo45.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, + *, + state: _MiniCPMO45Stage0SessionState | None = None, + ) -> Any | None: + process_audio = getattr(self.processor, "process_audio", None) + if callable(process_audio): + batch_feature = process_audio([ref_audio]) + if hasattr(batch_feature, "to"): + batch_feature = batch_feature.to(self.device) + self._ensure_dynamic_cache_compat() + for target in (self.stage_model, self.thinker): + get_audio_embedding = getattr(target, "get_audio_embedding", None) + if callable(get_audio_embedding): + try: + chunk_length = getattr(getattr(target, "config", None), "audio_chunk_length", None) + if chunk_length is not None: + return self._cat_nested_tensors( + get_audio_embedding(batch_feature, chunk_length=chunk_length) + ) + except TypeError: + pass + return self._cat_nested_tensors(get_audio_embedding(batch_feature)) + # The split vLLM stage0 wrapper ports official + # get_audio_embedding(chunk_length=...) as + # get_audio_hidden_states (chunk_length comes from config). + get_hidden = getattr(target, "get_audio_hidden_states", None) + if callable(get_hidden): + return self._cat_nested_tensors(get_hidden(batch_feature)) + # The split vLLM stage0 wrapper may only expose the streaming encoder + # path. Use it as a fallback so the server-resolved reference audio is + # still represented in the same system context location as official + # MiniCPM-o prepare(). + processor = state.streaming_processor if state is not None else self.processor + batch_feature = self._process_streaming_audio(ref_audio, 0, processor=processor) + for name, value in ( + ("chunk_idx", 0), + ("use_extra_context", True), + ("prefix_extra_frames", 0), + ("suffix_extra_frames", 2), + ): + with suppress(Exception): + setattr(batch_feature, name, value) + return self._stage_audio_embeddings(batch_feature, state=state) + + @staticmethod + def _ensure_dynamic_cache_compat() -> None: + try: + from transformers.cache_utils import DynamicCache + except Exception: + return + if hasattr(DynamicCache, "get_usable_length"): + return + + def get_usable_length(self, new_seq_length: int | None = None, layer_idx: int = 0) -> int: + get_seq_length = getattr(self, "get_seq_length", None) + if not callable(get_seq_length): + return 0 + try: + return int(get_seq_length(layer_idx)) + except TypeError: + return int(get_seq_length()) + + DynamicCache.get_usable_length = get_usable_length # type: ignore[attr-defined] + + @staticmethod + def _cat_nested_tensors(value: Any) -> Any | None: + import torch + + tensors = [] + + def collect(item: Any) -> None: + if item is None: + return + if hasattr(item, "detach"): + tensors.append(item) + return + if isinstance(item, dict): + for child in item.values(): + collect(child) + return + if isinstance(item, (list, tuple)): + for child in item: + collect(child) + + collect(value) + if not tensors: + return None + return torch.cat([tensor.reshape(-1, tensor.shape[-1]) for tensor in tensors], dim=0) + + def _encode_text(self, text: str) -> list[int]: + encode = getattr(self.tokenizer, "encode", None) + if callable(encode): + return list(encode(text, add_special_tokens=False)) + return [] + + def _special_token_ids(self) -> dict[str, int]: + return { + name: value + for name, value in { + "unit_token_id": self.unit_token_id, + "unit_end_token_id": self.unit_end_token_id, + "listen_token_id": self.listen_token_id, + "speak_token_id": self.speak_token_id, + "tts_bos_token_id": self.tts_bos_token_id, + "tts_eos_token_id": self.tts_eos_token_id, + "tts_pad_token_id": self.tts_pad_token_id, + "chunk_eos_token_id": self.chunk_eos_token_id, + "chunk_tts_eos_token_id": self.chunk_tts_eos_token_id, + "turn_eos_token_id": self.turn_eos_token_id, + }.items() + if isinstance(value, int) and value >= 0 + } + + @staticmethod + def _load_processor_from_path(model_path: str | None) -> Any | None: + if not model_path: + return None + if model_path not in sys.path: + sys.path.insert(0, model_path) + try: + from processing_minicpmo import MiniCPMOProcessor + + return MiniCPMOProcessor.from_pretrained(model_path, trust_remote_code=True) + except Exception: + return None + + def _init_token_ids(self) -> None: + if self.tokenizer is None: + for field_name in _MINICPMO45_SPECIAL_TOKEN_FIELDS: + setattr(self, field_name, -1) + for field_name in _MINICPMO45_OPTIONAL_TOKEN_FIELDS: + setattr(self, field_name, -1) + else: + for field_name, token in _MINICPMO45_SPECIAL_TOKEN_FIELDS.items(): + setattr(self, field_name, self._resolve_special_token_id(token)) + for field_name, token in _MINICPMO45_OPTIONAL_TOKEN_FIELDS.items(): + setattr(self, field_name, self._resolve_special_token_id(token)) + + def _resolve_special_token_id(self, token: str) -> int: + if self.tokenizer is None: + return -1 + + unk_token_id = getattr(self.tokenizer, "unk_token_id", None) + candidate = None + convert = getattr(self.tokenizer, "convert_tokens_to_ids", None) + if callable(convert): + value = convert(token) + if isinstance(value, list): + value = value[0] if len(value) == 1 else None + with suppress(TypeError, ValueError): + candidate = int(value) + if candidate is not None and candidate >= 0 and candidate != unk_token_id: + return candidate + + encode = getattr(self.tokenizer, "encode", None) + if callable(encode): + ids = list(encode(token, add_special_tokens=False)) + if len(ids) == 1: + value = int(ids[0]) + if value >= 0 and value != unk_token_id: + return value + return -1 + + def _require_special_token_ids(self) -> None: + missing = [ + token + for field_name, token in _MINICPMO45_SPECIAL_TOKEN_FIELDS.items() + if not isinstance(getattr(self, field_name, None), int) or getattr(self, field_name) < 0 + ] + if missing: + raise ValueError( + "MiniCPM-o 4.5 native duplex requires tokenizer-defined special " + f"tokens, missing or unknown: {', '.join(missing)}" + ) + + def _required_token_id(self, field_name: str) -> int: + token_id = getattr(self, field_name, None) + if not isinstance(token_id, int) or token_id < 0: + token = _MINICPMO45_SPECIAL_TOKEN_FIELDS.get(field_name, field_name) + raise ValueError(f"MiniCPM-o 4.5 missing required special token id for {token}") + return token_id + + def stage_padding_token_id(self) -> int: + return self._required_token_id("unit_end_token_id") + + def _audio_embedding_placeholder_token_id(self) -> int: + token_id = getattr(self, "audio_placeholder_token_id", -1) + if isinstance(token_id, int) and token_id >= 0: + return token_id + return self.stage_padding_token_id() + + @staticmethod + def _decode_audio_payload(payload: dict[str, Any]) -> Any: + audio = payload.get("audio") or payload.get("data") + if not isinstance(audio, str): + raise ValueError("audio append payload requires base64 audio") + fmt = payload.get("format") or "pcm_f32le" + if fmt != "pcm_f32le": + raise ValueError(f"MiniCPM-o stage0 expects pcm_f32le audio, got {fmt!r}") + return np.frombuffer(base64.b64decode(audio), dtype=np.float32) + + @staticmethod + def _decode_video_frames_payload(payload: dict[str, Any]) -> list[Any]: + """Decode omni-duplex camera frames (base64 JPEG/PNG) to PIL images.""" + frames = payload.get("video_frames") + if not isinstance(frames, list) or not frames: + return [] + from io import BytesIO + + from PIL import Image + + decoded: list[Any] = [] + for frame_b64 in frames: + if not isinstance(frame_b64, str) or not frame_b64: + continue + try: + raw = base64.b64decode(frame_b64, validate=True) + image = Image.open(BytesIO(raw)) + image.load() + except Exception as exc: # noqa: BLE001 - normalized below + raise ValueError("invalid omni duplex video frame payload") from exc + decoded.append(image.convert("RGB")) + return decoded + + def _vision_embedding_placeholder_token_id(self) -> int: + unk_token_id = getattr(self.tokenizer, "unk_token_id", None) + if isinstance(unk_token_id, int) and unk_token_id >= 0: + return unk_token_id + return self._required_optional_token_id("image_start_token_id") + + def _required_optional_token_id(self, field_name: str) -> int: + token_id = getattr(self, field_name, None) + if not isinstance(token_id, int) or token_id < 0: + token = _MINICPMO45_OPTIONAL_TOKEN_FIELDS.get(field_name, field_name) + raise ValueError(f"MiniCPM-o 4.5 missing required special token id for {token}") + return token_id + + def _require_vision_token_ids(self) -> None: + missing = [ + _MINICPMO45_OPTIONAL_TOKEN_FIELDS[field_name] + for field_name in ("image_start_token_id", "image_end_token_id") + if not isinstance(getattr(self, field_name, None), int) or getattr(self, field_name) < 0 + ] + if missing: + raise ValueError( + "MiniCPM-o 4.5 omni duplex requires tokenizer-defined image tokens, " + f"missing or unknown: {', '.join(missing)}" + ) + + def _stage_vision_embeddings(self, frames: list[Any]) -> list[Any] | None: + """Encode camera frames for omni duplex via the loaded vision tower. + + Semantics mirror ``MiniCPMODuplex.streaming_prefill`` (one 64-embedding + block per frame at ``max_slice_nums=1``), executed through the vLLM + wrapper's ``get_vision_hidden_states`` (vpm + resampler) since the + stage model does not expose the remote-code ``get_vision_embedding``. + """ + process_image = getattr(self.processor, "process_image", None) + if not callable(process_image): + return None + try: + processed = process_image(frames, max_slice_nums=1) + except Exception: # noqa: BLE001 - prefill fails with a reason + return None + targets = (self.stage_model, self.thinker, getattr(self.stage_model, "model", None)) + for target in targets: + if target is None: + continue + get_hidden = getattr(target, "get_vision_hidden_states", None) + vpm = getattr(target, "vpm", None) + if not callable(get_hidden) or vpm is None: + continue + try: + import torch + + vpm_param = next(vpm.parameters()) + 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] = [] + 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)) + tgt_tensor = image_tgt if hasattr(image_tgt, "reshape") else torch.tensor(image_tgt) + flat_tgt.append(tgt_tensor.reshape(-1, 2)) + tgt_sizes = torch.cat(flat_tgt, dim=0).to(device=device, dtype=torch.int64) + with torch.no_grad(): + hidden = get_hidden({"pixel_values": flat_pixels, "tgt_sizes": tgt_sizes}) + except Exception: # noqa: BLE001 - prefill fails with a reason + return None + expected = MiniCPMO45DuplexPolicy.VISION_EMBEDS_PER_FRAME + out: list[Any] = [] + for block in hidden: + block_2d = self._as_2d_tensor(block) + if int(block_2d.shape[0]) != expected: + return None + out.append(block_2d) + return out + return None diff --git a/vllm_omni/experimental/fullduplex/model_executor.py b/vllm_omni/experimental/fullduplex/model_executor.py new file mode 100644 index 000000000..eebe01876 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/model_executor.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True, slots=True) +class DuplexSamplingRow: + """Request-local context for the experimental duplex sampling hook.""" + + row_idx: int + request_id: str + session_id: str | None + incarnation: int + seq: int | None + payload: dict[str, object] | None + max_tokens: int | None + + +class DuplexSamplingHelper: + """Tracks duplex sampling rows outside the generic AR model runner.""" + + def __init__(self) -> None: + self.active_request_ids: set[str] = set() + self.hook_active = False + + @staticmethod + def _is_duplex_data_plane_info(info: object) -> bool: + duplex = info.get("duplex") if isinstance(info, dict) else None + return isinstance(duplex, dict) and duplex.get("data_plane") is True + + @staticmethod + def _request_intermediate_info(runner: object, req_id: str) -> dict[str, Any] | None: + model_intermediate_buffer = getattr(runner, "model_intermediate_buffer", {}) + info = model_intermediate_buffer.get(req_id) if isinstance(model_intermediate_buffer, dict) else None + if isinstance(info, dict): + return info + requests = getattr(runner, "requests", {}) + req_state = requests.get(req_id) if isinstance(requests, dict) else None + info = getattr(req_state, "additional_information_cpu", None) + return info if isinstance(info, dict) else None + + def refresh_active_request(self, runner: object, req_id: str) -> None: + if self._is_duplex_data_plane_info(self._request_intermediate_info(runner, req_id)): + self.active_request_ids.add(req_id) + else: + self.active_request_ids.discard(req_id) + + def update_states(self, runner: object, scheduler_output: object) -> None: + self.active_request_ids.difference_update(str(req_id) for req_id in scheduler_output.finished_req_ids) + for request in scheduler_output.scheduled_new_reqs: + self.refresh_active_request(runner, str(request.req_id)) + + def rows(self, runner: object) -> tuple[DuplexSamplingRow, ...]: + rows: list[DuplexSamplingRow] = [] + req_ids = [str(req_id) for req_id in getattr(runner.input_batch, "req_ids", [])] + requests = getattr(runner, "requests", {}) + for row_idx, req_id in enumerate(req_ids): + if req_id not in self.active_request_ids: + continue + info = self._request_intermediate_info(runner, req_id) + duplex = info.get("duplex") if isinstance(info, dict) else None + if not isinstance(duplex, dict) or duplex.get("data_plane") is not True: + continue + 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): + seq = None + payload = duplex.get("payload") + if not isinstance(payload, dict): + payload = None + request = requests.get(req_id) if isinstance(requests, dict) else None + sampling_params = getattr(request, "sampling_params", None) + try: + max_tokens = int(getattr(sampling_params, "max_tokens", 0) or 0) + except (TypeError, ValueError): + max_tokens = 0 + rows.append( + DuplexSamplingRow( + 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, + ) + ) + return tuple(rows) + + def clear(self) -> None: + self.active_request_ids.clear() + self.hook_active = False + + +__all__ = ["DuplexSamplingHelper", "DuplexSamplingRow"] diff --git a/vllm_omni/experimental/fullduplex/openai/__init__.py b/vllm_omni/experimental/fullduplex/openai/__init__.py new file mode 100644 index 000000000..1bb45f8d4 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/openai/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .websocket import DuplexWebSocketActor + +__all__ = ["DuplexWebSocketActor"] diff --git a/vllm_omni/experimental/fullduplex/openai/audio.py b/vllm_omni/experimental/fullduplex/openai/audio.py new file mode 100644 index 000000000..023bfdd9e --- /dev/null +++ b/vllm_omni/experimental/fullduplex/openai/audio.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import base64 +import binascii +import io +import math +import wave + +import numpy as np + +try: + from audioop import alaw2lin, lin2alaw, lin2ulaw, ulaw2lin +except ImportError: # pragma: no cover - audioop is removed in newer Python. + alaw2lin = lin2alaw = lin2ulaw = ulaw2lin = None + + +MIN_INPUT_SAMPLE_RATE_HZ = 8_000 +MAX_INPUT_SAMPLE_RATE_HZ = 192_000 + + +def validate_input_sample_rate_hz(sample_rate_hz: int | float) -> int: + if isinstance(sample_rate_hz, bool) or (isinstance(sample_rate_hz, float) and not math.isfinite(sample_rate_hz)): + raise ValueError("sample_rate_hz must be a finite integer") + rate = int(sample_rate_hz) + if rate != sample_rate_hz: + raise ValueError("sample_rate_hz must be an integer") + if not MIN_INPUT_SAMPLE_RATE_HZ <= rate <= MAX_INPUT_SAMPLE_RATE_HZ: + raise ValueError(f"sample_rate_hz must be between {MIN_INPUT_SAMPLE_RATE_HZ} and {MAX_INPUT_SAMPLE_RATE_HZ}") + return rate + + +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 + samples = np.frombuffer(raw, dtype=" bytes: + if ulaw2lin is not None: + return ulaw2lin(raw, 2) + data = np.frombuffer(raw, dtype=np.uint8) + value = np.bitwise_not(data).astype(np.int16) + sign = value & 0x80 + exponent = (value >> 4) & 0x07 + mantissa = value & 0x0F + sample = ((mantissa << 3) + 0x84) << exponent + return np.where(sign != 0, -(sample - 0x84), sample - 0x84).astype(" bytes: + if alaw2lin is not None: + return alaw2lin(raw, 2) + data = np.bitwise_xor(np.frombuffer(raw, dtype=np.uint8), 0x55).astype(np.int16) + sign = data & 0x80 + exponent = (data >> 4) & 0x07 + mantissa = data & 0x0F + sample = np.where(exponent == 0, (mantissa << 4) + 8, ((mantissa << 4) + 0x108) << (exponent - 1)) + return np.where(sign != 0, sample, -sample).astype(" bytes: + if lin2ulaw is not None: + return lin2ulaw(raw, 2) + pcm = np.clip(np.frombuffer(raw, dtype=" (0xFF << exp), exp + 1, exponent) + mantissa = (magnitude >> (exponent + 3)) & 0x0F + return (np.bitwise_not(sign | (exponent << 4) | mantissa) & 0xFF).astype(np.uint8).tobytes() + + +def encode_g711_alaw(raw: bytes) -> bytes: + if lin2alaw is not None: + return lin2alaw(raw, 2) + pcm = np.frombuffer(raw, dtype="= 0, 0x80, 0x00) + magnitude = np.abs(pcm) + exponent = np.zeros_like(magnitude) + for exp in range(1, 8): + exponent = np.where(magnitude >= (1 << (exp + 7)), exp, exponent) + mantissa = np.where( + exponent == 0, + (magnitude >> 4) & 0x0F, + (magnitude >> (exponent + 3)) & 0x0F, + ) + return ((sign | (exponent << 4) | mantissa) ^ 0x55).astype(np.uint8).tobytes() + + +def wav_payload_to_pcm16(raw: bytes) -> tuple[bytes | None, int | None]: + try: + with wave.open(io.BytesIO(raw), "rb") as wav_file: + channels = wav_file.getnchannels() + sample_width = wav_file.getsampwidth() + sample_rate_hz = wav_file.getframerate() + frames = wav_file.readframes(wav_file.getnframes()) + if sample_width != 2: + return None, sample_rate_hz + if channels <= 1: + return frames, sample_rate_hz + pcm = np.frombuffer(frames, dtype=" tuple[object, object, int | float | None]: + if not isinstance(audio, str) or not isinstance(fmt, str): + return audio, fmt, sample_rate_hz + normalized = fmt.lower() + if normalized not in {"pcm16", "pcm_s16le", "s16le", "g711_ulaw", "g711_alaw"}: + return audio, fmt, sample_rate_hz + if isinstance(sample_rate_hz, int | float): + sample_rate_hz = validate_input_sample_rate_hz(sample_rate_hz) + try: + raw = base64.b64decode(audio.strip(), validate=False) + except (binascii.Error, ValueError): + return audio, fmt, sample_rate_hz + if normalized == "g711_ulaw": + raw = decode_g711_ulaw(raw) + sample_rate_hz = sample_rate_hz if isinstance(sample_rate_hz, int | float) else 8_000 + elif normalized == "g711_alaw": + raw = decode_g711_alaw(raw) + sample_rate_hz = sample_rate_hz if isinstance(sample_rate_hz, int | float) else 8_000 + elif len(raw) % 2: + return audio, fmt, sample_rate_hz + if isinstance(sample_rate_hz, int | float) and int(sample_rate_hz) != target_sample_rate_hz: + raw = resample_pcm16_mono( + raw, + source_rate_hz=int(sample_rate_hz), + target_rate_hz=target_sample_rate_hz, + ) + sample_rate_hz = target_sample_rate_hz + pcm = np.frombuffer(raw, dtype=" tuple[str, str, int | None]: + target = target_fmt.lower() + if target not in {"g711_ulaw", "g711_alaw"}: + return audio, source_fmt, source_sample_rate_hz + try: + raw = base64.b64decode(audio, validate=False) + except (binascii.Error, ValueError): + return audio, source_fmt, source_sample_rate_hz + source = source_fmt.lower() + if source == "wav": + pcm_raw, wav_rate = wav_payload_to_pcm16(raw) + if pcm_raw is None: + return audio, source_fmt, source_sample_rate_hz + raw = pcm_raw + source_sample_rate_hz = source_sample_rate_hz or wav_rate + elif source not in {"pcm", "pcm16", "pcm_s16le", "s16le"} or len(raw) % 2: + return audio, source_fmt, source_sample_rate_hz + target_rate = target_sample_rate_hz or 8_000 + if source_sample_rate_hz is not None: + raw = resample_pcm16_mono(raw, source_rate_hz=source_sample_rate_hz, target_rate_hz=target_rate) + encoded = encode_g711_ulaw(raw) if target == "g711_ulaw" else encode_g711_alaw(raw) + return base64.b64encode(encoded).decode("ascii"), target, target_rate diff --git a/vllm_omni/experimental/fullduplex/openai/chat_fallback.py b/vllm_omni/experimental/fullduplex/openai/chat_fallback.py new file mode 100644 index 000000000..65d3b32fd --- /dev/null +++ b/vllm_omni/experimental/fullduplex/openai/chat_fallback.py @@ -0,0 +1,252 @@ +from __future__ import annotations + +import asyncio +import json +from collections.abc import AsyncGenerator +from typing import Any + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.logger import init_logger + +from vllm_omni.experimental.fullduplex.openai.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): + await send_json({"type": "error", "error": result.message, "code": result.type or "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: + 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) + await send_json( + { + "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(), + } + ) + 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) + 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 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: + 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.mark_audio_sent() + await send_json( + { + "type": "response.output_audio.delta", + "session_id": session.session_id, + "response_id": response_id, + "epoch": epoch, + "audio": content, + "format": session.response_config.response_format, + } + ) + else: + 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, + } + ) + + 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/experimental/fullduplex/openai/commit_policy.py b/vllm_omni/experimental/fullduplex/openai/commit_policy.py new file mode 100644 index 000000000..940509ec3 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/openai/commit_policy.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum + + +class CommitAction(str, Enum): + """How one native audio commit participates in response scheduling.""" + + START_AUTO_RESPONSE = "start_auto_response" + DEFER_ACTIVE_RESPONSE = "defer_active_response" + COMMIT_ONLY = "commit_only" + + +@dataclass(frozen=True, slots=True) +class CommitSnapshot: + auto_responds: bool + speech_since_commit: bool + active_response_id: str | None + overlap_speech_ms: int + native_response_in_progress: bool + playback_active: bool + + +def decide_commit_action(snapshot: CommitSnapshot) -> CommitAction: + """Choose response scheduling from an immutable session snapshot.""" + if snapshot.auto_responds: + 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: + return CommitAction.DEFER_ACTIVE_RESPONSE + if snapshot.playback_active and snapshot.overlap_speech_ms > 0: + return CommitAction.DEFER_ACTIVE_RESPONSE + return CommitAction.COMMIT_ONLY + + +__all__ = ["CommitAction", "CommitSnapshot", "decide_commit_action"] diff --git a/vllm_omni/experimental/fullduplex/openai/protocol.py b/vllm_omni/experimental/fullduplex/openai/protocol.py new file mode 100644 index 000000000..0a87d55b4 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/openai/protocol.py @@ -0,0 +1,1336 @@ +from __future__ import annotations + +import copy +from collections.abc import Mapping +from dataclasses import dataclass, field +from enum import Enum +from types import MappingProxyType +from uuid import uuid4 + + +class DuplexOverlapPolicy(str, Enum): + AUTO = "auto" + 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 + 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 + + @classmethod + def minicpmo45_native(cls, *, max_sessions: int = 1) -> DuplexCapabilities: + supports_multi_session = max_sessions > 1 + return cls( + supports_model_native_turn_policy=True, + supports_barge_in=False, + supports_input_append=True, + supports_replace_latest_chunk=False, + supports_reencode_context=False, + supports_turn_commit_only=False, + supports_kv_lease=False, + supports_core_kv_lease=False, + supports_model_internal_state=True, + supports_stage_resumption=True, + supports_scheduler_native_append=False, + supports_core_resumable_request=True, + supports_stage_connector_handoff=True, + supports_independent_io_streams=True, + supports_realtime_endpoint=True, + supports_multi_session=supports_multi_session, + supports_multi_session_same_replica=supports_multi_session, + supports_session_lease=True, + supports_session_resume=True, + session_admission_mode="engine_managed", + supports_audio_truncate=True, + 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, + target_barge_in_latency_ms=None, + ) + + 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 + + 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 + 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, + "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), + } + + @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"]) + 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 + + @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 DuplexAssistantAudioTextMark: + text_chars: int + audio_end_ms: int + + +@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 + + +@dataclass +class ResponseState: + active_request_id: str | None = None + active_response_id: str | None = None + active_response_turn_id: int | None = None + last_response_id: str | None = None + assistant_text_buffer: list[str] = field(default_factory=list) + assistant_audio_text_marks: list[DuplexAssistantAudioTextMark] = field(default_factory=list) + pending_options: ResponseCreateOptions | None = None + active_options: ResponseCreateOptions | None = None + active_config: DuplexSessionConfig | None = None + stage_metrics: dict[str, dict[str, object]] = field(default_factory=dict) + stage_metric_tpot_weighted_ms: dict[str, float] = field(default_factory=dict) + stage_metric_tpot_weight: dict[str, int] = field(default_factory=dict) + + +@dataclass +class PlaybackLedger: + current: DuplexPlaybackCursor = field(default_factory=DuplexPlaybackCursor) + by_response: dict[str, DuplexPlaybackCursor] = field(default_factory=dict) + + +@dataclass +class ConversationHistory: + messages: list[dict[str, object]] = field(default_factory=list) + item_ids: dict[str, dict[str, object]] = field(default_factory=dict) + item_audio_text_marks: dict[str, list[DuplexAssistantAudioTextMark]] = field(default_factory=dict) + pending_item_ids: dict[str, dict[str, object]] = field(default_factory=dict) + pending_item_audio_text_marks: dict[str, list[DuplexAssistantAudioTextMark]] = field(default_factory=dict) + pending_truncations_ms: dict[str, int] = field(default_factory=dict) + last_assistant_full_message: dict[str, object] | None = None + last_assistant_audio_text_marks: list[DuplexAssistantAudioTextMark] = field(default_factory=list) + + +@dataclass +class DuplexSession: + 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 + _input: InputBufferState = field(default_factory=InputBufferState, repr=False) + _response: ResponseState = field(default_factory=ResponseState, 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) + + @property + def response_config(self) -> DuplexSessionConfig: + """Return the immutable-for-this-lifecycle response configuration. + + ``config`` remains the session defaults. A response takes a deep + snapshot when it begins so response.create overrides and concurrent + session updates cannot mutate each other's ownership domains. + """ + return self._response.active_config or self.config + + @property + def runtime_config(self) -> Mapping[str, object]: + return MappingProxyType(dict(self._runtime_config)) + + def replace_runtime_config(self, runtime_config: Mapping[str, object]) -> None: + self._runtime_config = dict(runtime_config) + + @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], ...]: + return tuple(dict(message) for message in self._conversation.messages) + + @property + def active_request_id(self) -> str | None: + return self._response.active_request_id + + @property + def active_response_id(self) -> str | None: + return self._response.active_response_id + + @property + def active_response_turn_id(self) -> int | None: + return self._response.active_response_turn_id + + @property + def last_response_id(self) -> str | None: + return self._response.last_response_id + + @property + def overlap_speech_ms(self) -> int: + return self._input.overlap_speech_ms + + @property + def assistant_text_buffer(self) -> tuple[str, ...]: + return tuple(self._response.assistant_text_buffer) + + @property + def assistant_audio_text_marks(self) -> tuple[DuplexAssistantAudioTextMark, ...]: + return tuple(self._response.assistant_audio_text_marks) + + @property + def last_assistant_full_message(self) -> dict[str, object] | None: + message = self._conversation.last_assistant_full_message + return dict(message) if message is not None else None + + @property + def last_assistant_audio_text_marks(self) -> tuple[DuplexAssistantAudioTextMark, ...]: + return tuple(self._conversation.last_assistant_audio_text_marks) + + @property + 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()}) + + @property + def pending_history_item_ids(self) -> Mapping[str, dict[str, object]]: + return MappingProxyType({key: dict(value) for key, value in self._conversation.pending_item_ids.items()}) + + @property + def pending_history_truncations_ms(self) -> Mapping[str, int]: + return MappingProxyType(dict(self._conversation.pending_truncations_ms)) + + def replace_config(self, config: DuplexSessionConfig) -> None: + self.config = config + + def replace_capabilities(self, capabilities: DuplexCapabilities) -> None: + self.capabilities = capabilities + + def transition_turn(self, state: DuplexTurnState) -> None: + self.turn_state = state + + def transition_session(self, state: DuplexSessionState) -> None: + self.state = state + + def bind_request(self, request_id: str | None) -> None: + self._response.active_request_id = request_id + + def clear_request(self, expected_request_id: str | None = None) -> bool: + if expected_request_id is not None and self._response.active_request_id != expected_request_id: + return False + self._response.active_request_id = None + return True + + def bind_response_turn(self, turn_id: int | None) -> None: + self._response.active_response_turn_id = turn_id + + def active_response_accepts_model_turn(self, turn_id: int | None) -> bool: + if self._response.active_response_id is None: + return False + if turn_id is None: + return True + active_turn_id = self._response.active_response_turn_id + return active_turn_id is None or int(turn_id) == active_turn_id + + def append_history_message(self, message: dict[str, object]) -> None: + 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 append_text(self, text: str) -> None: + if not text: + return + self._input.pending_text.append(text) + self.turn_state = DuplexTurnState.USER_SPEAKING + + @property + def pending_input_bytes(self) -> int: + return self._input.reserved_input_bytes + + @property + def pending_input_turns(self) -> int: + return self._input.pending_turns + + def reserve_input_bytes(self, size: int, *, limit: int) -> bool: + size = max(0, int(size)) + if self._input.reserved_input_bytes + size > int(limit): + return False + self._input.reserved_input_bytes += size + return 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))) + + def release_all_input_bytes(self) -> None: + self._input.reserved_input_bytes = 0 + + def reserve_pending_turn(self, *, limit: int) -> bool: + if self._input.pending_turns >= int(limit): + return False + self._input.pending_turns += 1 + return True + + 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 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() + self._input.reserved_input_bytes = 0 + self._input.pending_turns = 0 + self.turn_state = DuplexTurnState.IDLE + return cancelled + + def commit_user_input(self) -> 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._input.commit_seq += 1 + message = {"role": "user", "content": content} + self._conversation.messages.append(message) + self._input.pending_text.clear() + self._input.pending_audio.clear() + self._input.reserved_input_bytes = 0 + 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, + ) + + def commit_native_audio_input( + self, + *, + transcript: str | None = None, + turn_id: int | None = None, + ) -> DuplexCommittedInput: + input_audio_part: dict[str, object] = { + "type": "audio_url", + "audio_url": {"url": "native-duplex:input-audio"}, + } + if transcript: + input_audio_part["transcript"] = transcript + self._input.commit_seq += 1 + message = {"role": "user", "content": [input_audio_part]} + 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, + turn_id=self.turn_id if turn_id is None else int(turn_id), + epoch=self.epoch, + input_commit_seq=self._input.commit_seq, + ) + + 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) + if completed_turn_id >= self.turn_id: + self.turn_id = completed_turn_id + 1 + + def reserve_response_options(self, options: ResponseCreateOptions) -> None: + if self._response.active_response_id is not None: + raise RuntimeError("response options cannot be reserved while a response is active") + if self._response.pending_options is not None: + raise RuntimeError("response options are already reserved") + self._response.pending_options = options + + def discard_response_options(self) -> None: + self._response.pending_options = None + + def _activate_response_options(self) -> None: + options = self._response.pending_options + self._response.active_config = copy.deepcopy(self.config) + self._response.active_options = options + self._response.pending_options = None + if options is not None: + options.apply_to(self._response.active_config) + + def _restore_response_config(self) -> None: + self._response.active_config = None + self._response.active_options = None + self._response.pending_options = None + + def begin_response(self, *, turn_id: int | None = None) -> str: + self._activate_response_options() + 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.last_response_id = response_id + self._response.assistant_text_buffer.clear() + self._response.assistant_audio_text_marks.clear() + self._clear_response_metrics() + self._conversation.last_assistant_full_message = None + self._conversation.last_assistant_audio_text_marks.clear() + self._playback.current = DuplexPlaybackCursor() + self._playback.by_response[response_id] = self._playback.current + self.turn_state = DuplexTurnState.ASSISTANT_GENERATING + return response_id + + def _clear_response_metrics(self) -> None: + self._response.stage_metrics.clear() + self._response.stage_metric_tpot_weighted_ms.clear() + self._response.stage_metric_tpot_weight.clear() + + 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): + return copy.deepcopy(self._response.stage_metrics) + + additive_fields = ( + "num_tokens_in", + "num_tokens_out", + "stage_gen_time_ms", + "postprocess_time_ms", + "audio_generated_frames", + "audio_duration_s", + "image_pixels", + "output_unit_count", + ) + first_positive_fields = ( + "serving_time_to_first_output_ms", + "vllm_ttft_ms", + ) + interval_fields = ( + ("inter_output_latencies_ms", "inter_output_latency_ms"), + ("vllm_itls_ms", "vllm_itl_ms"), + ) + handled_fields = { + *additive_fields, + *first_positive_fields, + "vllm_tpot_ms", + *(name for pair in interval_fields for name in pair), + } + + for raw_stage_id, raw_values in stage_metrics.items(): + if not isinstance(raw_values, Mapping): + continue + stage_id = str(raw_stage_id) + current = self._response.stage_metrics.setdefault(stage_id, {}) + for name in additive_fields: + value = raw_values.get(name) + if isinstance(value, int | float) and not isinstance(value, bool): + current[name] = current.get(name, 0) + value + for name in first_positive_fields: + value = raw_values.get(name) + current_value = current.get(name) + if ( + isinstance(value, int | float) + and not isinstance(value, bool) + and value > 0 + and not (isinstance(current_value, int | float) and current_value > 0) + ): + current[name] = value + for list_name, mean_name in interval_fields: + values = raw_values.get(list_name) + if isinstance(values, list): + combined = list(current.get(list_name, [])) + combined.extend( + value for value in values if isinstance(value, int | float) and not isinstance(value, bool) + ) + current[list_name] = combined + current[mean_name] = sum(combined) / len(combined) if combined else 0.0 + + tpot_ms = raw_values.get("vllm_tpot_ms") + token_count = raw_values.get("num_tokens_out") + 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] + ) + + for name, value in raw_values.items(): + if name not in handled_fields: + current[str(name)] = copy.deepcopy(value) + + return copy.deepcopy(self._response.stage_metrics) + + def accumulate_overlap_speech(self, duration_ms: int) -> int: + self._input.overlap_speech_ms += max(0, int(duration_ms)) + return self._input.overlap_speech_ms + + def reset_overlap_speech(self) -> int: + previous = self._input.overlap_speech_ms + self._input.overlap_speech_ms = 0 + return previous + + def append_assistant_text(self, text: str) -> None: + if text: + self._response.assistant_text_buffer.append(text) + + def mark_audio_sent( + self, + duration_ms: int | None = None, + *, + text_chars: int | None = None, + audio_text_marks: list[dict[str, object]] | None = None, + ) -> None: + 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 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 audio_text_marks: + 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 + self._response.assistant_audio_text_marks.append( + DuplexAssistantAudioTextMark( + text_chars=max(0, int(raw_text_chars)), + audio_end_ms=max(0, int(raw_audio_end_ms)), + ) + ) + self.turn_state = DuplexTurnState.ASSISTANT_PLAYING + + def _playback_cursor_for_response(self, response_id: str | None = None) -> DuplexPlaybackCursor: + if response_id is None: + return self._playback.current + playback = self._playback.by_response.get(response_id) + if playback is None: + # A restored or legacy session may not have response-scoped state. + # Keep its acknowledgement isolated from the active response. + playback = DuplexPlaybackCursor() + self._playback.by_response[response_id] = playback + return playback + + def playback_for_response(self, response_id: str | None = None) -> DuplexPlaybackView: + return self._playback_cursor_for_response(response_id).snapshot() + + def acknowledge_playback( + self, + played_ms: int, + committed_ms: int | None = None, + *, + response_id: str | None = None, + ) -> DuplexPlaybackView: + playback = self._playback_cursor_for_response(response_id) + playback.acknowledge(played_ms, committed_ms) + return playback.snapshot() + + def truncate_playback_commit( + self, + committed_ms: int, + *, + response_id: str | None = None, + ) -> DuplexPlaybackView: + playback = self._playback_cursor_for_response(response_id) + playback.truncate_committed(committed_ms) + return playback.snapshot() + + def release_response_playback(self, response_id: str | None) -> None: + if response_id is None or response_id == self.active_response_id: + return + self._playback.by_response.pop(response_id, None) + + def clear_playback_cursor(self) -> None: + self._playback.current = DuplexPlaybackCursor() + if self.active_response_id is not None: + self._playback.by_response[self.active_response_id] = self._playback.current + + def end_response( + self, + *, + commit_text: bool = True, + playback_commit_policy: str | None = None, + preserve_request: bool = False, + ) -> dict[str, object] | None: + assistant_text = "".join(self._response.assistant_text_buffer).strip() + message = None + if assistant_text: + self._conversation.last_assistant_full_message = {"role": "assistant", "content": assistant_text} + self._conversation.last_assistant_audio_text_marks = list(self._response.assistant_audio_text_marks) + if commit_text and assistant_text: + committed_text = self._playback_committed_text( + assistant_text, + playback_commit_policy=playback_commit_policy, + ) + else: + committed_text = "" + if commit_text and committed_text: + message = {"role": "assistant", "content": committed_text} + self._conversation.messages.append(message) + 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._clear_response_metrics() + self.turn_state = DuplexTurnState.IDLE + self._restore_response_config() + return message + + def register_history_item(self, item_id: str | None, message: dict[str, object] | None) -> None: + if not item_id: + return + if message is None: + last_message = self._conversation.last_assistant_full_message + if last_message is None: + return + self._conversation.pending_item_ids[item_id] = dict(last_message) + if self._conversation.last_assistant_audio_text_marks: + self._conversation.pending_item_audio_text_marks[item_id] = list( + self._conversation.last_assistant_audio_text_marks + ) + pending_audio_ms = self._conversation.pending_truncations_ms.pop(item_id, None) + if pending_audio_ms is not None: + self.truncate_history_item( + item_id, + audio_end_ms=pending_audio_ms, + playback=self._playback_cursor_for_item_id(item_id), + ) + return + pending_audio_ms = self._conversation.pending_truncations_ms.pop(item_id, None) + if pending_audio_ms is not None: + self._truncate_message_to_audio_ms( + message, + audio_end_ms=pending_audio_ms, + marks=self._response.assistant_audio_text_marks or self._conversation.last_assistant_audio_text_marks, + playback=self._playback_cursor_for_item_id(item_id), + ) + if self._message_text_len(message) <= 0: + try: + self._conversation.messages.remove(message) + except ValueError: + pass + return + self._conversation.item_ids[item_id] = message + self._conversation.pending_item_ids.pop(item_id, None) + self._conversation.pending_item_audio_text_marks.pop(item_id, None) + if message.get("role") == "assistant": + marks = self._response.assistant_audio_text_marks or self._conversation.last_assistant_audio_text_marks + if marks: + self._conversation.item_audio_text_marks[item_id] = list(marks) + + def delete_history_item(self, item_id: str) -> bool: + message = self._conversation.item_ids.pop(item_id, None) + self._conversation.item_audio_text_marks.pop(item_id, None) + pending = self._conversation.pending_item_ids.pop(item_id, None) + self._conversation.pending_item_audio_text_marks.pop(item_id, None) + self._conversation.pending_truncations_ms.pop(item_id, None) + if message is None: + return pending is not None + try: + self._conversation.messages.remove(message) + except ValueError: + pass + return True + + def truncate_history_item( + self, + item_id: str, + *, + audio_end_ms: int, + playback: DuplexPlaybackCursor | DuplexPlaybackView | None = None, + ) -> bool: + playback = playback or self._playback_cursor_for_item_id(item_id) + message = self._conversation.item_ids.get(item_id) + if message is None: + pending = self._conversation.pending_item_ids.get(item_id) + if pending is None: + self._conversation.pending_truncations_ms[item_id] = max(0, int(audio_end_ms)) + return False + message = dict(pending) + changed = self._truncate_message_to_audio_ms( + message, + audio_end_ms=audio_end_ms, + marks=self._conversation.pending_item_audio_text_marks.get(item_id), + playback=playback, + ) + if not changed or self._message_text_len(message) <= 0: + if changed: + self._conversation.pending_item_ids.pop(item_id, None) + self._conversation.pending_item_audio_text_marks.pop(item_id, None) + self._conversation.pending_truncations_ms.pop(item_id, None) + return changed + self._conversation.messages.append(message) + self._conversation.item_ids[item_id] = message + if item_id in self._conversation.pending_item_audio_text_marks: + self._conversation.item_audio_text_marks[item_id] = list( + self._conversation.pending_item_audio_text_marks[item_id] + ) + self._conversation.pending_item_ids.pop(item_id, None) + self._conversation.pending_item_audio_text_marks.pop(item_id, None) + self._conversation.pending_truncations_ms.pop(item_id, None) + return True + changed = self._truncate_message_to_audio_ms( + message, + audio_end_ms=audio_end_ms, + marks=self._conversation.item_audio_text_marks.get(item_id), + playback=playback, + ) + if changed and self._message_text_len(message) <= 0: + self._conversation.item_ids.pop(item_id, None) + self._conversation.item_audio_text_marks.pop(item_id, None) + try: + self._conversation.messages.remove(message) + except ValueError: + pass + return changed + + def _truncate_message_to_audio_ms( + self, + message: dict[str, object], + *, + audio_end_ms: int, + marks: list[DuplexAssistantAudioTextMark] | None = None, + playback: DuplexPlaybackCursor | DuplexPlaybackView | None = None, + ) -> bool: + content = message.get("content") + if isinstance(content, str): + keep_chars = self._text_chars_for_audio_ms( + audio_end_ms, + len(content), + marks=marks, + playback=playback, + ) + message["content"] = content[:keep_chars].rstrip() + return True + if not isinstance(content, list): + return False + changed = False + for part in content: + if not isinstance(part, dict): + continue + part_type = part.get("type") + if part_type in {"output_audio", "audio", "audio_transcript"}: + transcript = part.get("transcript") + if isinstance(transcript, str): + keep_chars = self._text_chars_for_audio_ms( + audio_end_ms, + len(transcript), + marks=marks, + playback=playback, + ) + part["transcript"] = transcript[:keep_chars].rstrip() + changed = True + if part_type in {"output_text", "text"}: + text = part.get("text") + if isinstance(text, str): + keep_chars = self._text_chars_for_audio_ms( + audio_end_ms, + len(text), + marks=marks, + playback=playback, + ) + part["text"] = text[:keep_chars].rstrip() + changed = True + if not changed: + return False + return True + + @staticmethod + def _message_text_len(message: dict[str, object]) -> int: + content = message.get("content") + if isinstance(content, str): + return len(content) + if not isinstance(content, list): + return 0 + total = 0 + for part in content: + if not isinstance(part, dict): + continue + for key in ("text", "transcript"): + value = part.get(key) + if isinstance(value, str): + total += len(value) + return total + + def _playback_committed_text( + self, + assistant_text: str, + *, + playback_commit_policy: str | None = None, + ) -> str: + sent_ms = max(self.playback.sent_ms, self.playback.generated_ms) + committed_ms = self.playback.committed_ms + policy = playback_commit_policy or self.config.playback_commit_policy + if sent_ms <= 0 or committed_ms >= sent_ms: + return assistant_text + if committed_ms <= 0: + if policy == DuplexPlaybackCommitPolicy.COMMIT_ALL_ON_DONE.value: + return assistant_text + return "" + keep_chars = self._text_chars_for_audio_ms(committed_ms, len(assistant_text)) + if keep_chars <= 0: + return "" + return assistant_text[:keep_chars].rstrip() + + def _text_chars_for_audio_ms( + self, + audio_end_ms: int, + text_len: int, + *, + marks: list[DuplexAssistantAudioTextMark] | None = None, + playback: DuplexPlaybackCursor | DuplexPlaybackView | None = None, + ) -> int: + if text_len <= 0: + return 0 + audio_end_ms = max(0, int(audio_end_ms)) + marks = marks if marks is not None else self._response.assistant_audio_text_marks + playback = playback or self._playback.current + if not marks: + sent_ms = max(1, playback.sent_ms, playback.generated_ms) + return int(text_len * max(0.0, min(1.0, audio_end_ms / sent_ms))) + marks = sorted( + (mark for mark in marks if mark.audio_end_ms >= 0 and mark.text_chars >= 0), + key=lambda mark: mark.audio_end_ms, + ) + if not marks: + return 0 + if audio_end_ms <= 0: + return 0 + previous_ms = 0 + previous_chars = 0 + for mark in marks: + mark_ms = max(previous_ms, mark.audio_end_ms) + mark_chars = min(text_len, max(previous_chars, mark.text_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 + final_ms = max(playback.sent_ms, playback.generated_ms, previous_ms) + if audio_end_ms >= final_ms: + return text_len + ratio = (audio_end_ms - previous_ms) / max(1, final_ms - previous_ms) + return int(previous_chars + (text_len - previous_chars) * max(0.0, min(1.0, ratio))) + + def _playback_cursor_for_item_id(self, item_id: str) -> DuplexPlaybackCursor | None: + if not item_id.startswith("item_"): + return None + return self._playback.by_response.get(item_id.removeprefix("item_")) + + def barge_in(self) -> int: + self.epoch += 1 + self._input.pending_text.clear() + self._input.pending_audio.clear() + 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._clear_response_metrics() + self._restore_response_config() + self.turn_state = DuplexTurnState.BARGE_IN + return self.epoch + + def mark_closing(self) -> None: + if self.state != DuplexSessionState.CLOSED: + self.state = DuplexSessionState.CLOSING + + def close(self) -> None: + self.state = DuplexSessionState.CLOSED + self.turn_state = DuplexTurnState.IDLE + self._response.active_response_turn_id = None + self._clear_response_metrics() + self._restore_response_config() + + def as_public_dict(self) -> dict[str, object]: + payload: dict[str, object] = { + "id": self.session_id, + "state": self.state.value, + "turn_state": self.turn_state.value, + "epoch": self.epoch, + "turn_id": self.turn_id, + "active_request_id": self.active_request_id, + "active_response_id": self.active_response_id, + "active_response_turn_id": self.active_response_turn_id, + "model": self.config.model, + "modalities": list(self.config.modalities), + "instructions": self.config.instructions, + "voice": self.config.voice, + "response_format": self.config.response_format, + "temperature": self.config.temperature, + "max_tokens": self.config.max_tokens, + "speed": self.config.speed, + "idle_timeout_s": self.config.idle_timeout_s, + "overlap_policy": self.config.overlap_policy, + "overlap_short_ack_ms": self.config.overlap_short_ack_ms, + "overlap_barge_in_ms": self.config.overlap_barge_in_ms, + "overlap_silence_rms": self.config.overlap_silence_rms, + "playback_commit_policy": self.config.playback_commit_policy, + "playback": self.playback.as_dict(), + "capabilities": self.capabilities.as_dict(), + } + 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): + payload["tool_choice"] = self.config.extra_body["realtime_tool_choice"] + if isinstance(self.config.extra_body.get("realtime_metadata"), dict): + payload["metadata"] = dict(self.config.extra_body["realtime_metadata"]) + if isinstance(self.config.extra_body.get("realtime_include"), list): + payload["include"] = list(self.config.extra_body["realtime_include"]) + if isinstance(self.config.extra_body.get("realtime_prompt"), dict): + payload["prompt"] = dict(self.config.extra_body["realtime_prompt"]) + if isinstance(self.config.extra_body.get("realtime_input_audio_transcription"), dict): + payload["input_audio_transcription"] = dict(self.config.extra_body["realtime_input_audio_transcription"]) + if isinstance(self.config.extra_body.get("realtime_input_audio_noise_reduction"), dict): + payload["input_audio_noise_reduction"] = dict( + self.config.extra_body["realtime_input_audio_noise_reduction"] + ) + if isinstance(self.config.extra_body.get("realtime_audio"), dict): + payload["audio"] = dict(self.config.extra_body["realtime_audio"]) + if isinstance(self.config.extra_body.get("realtime_tracing"), str | dict): + payload["tracing"] = self.config.extra_body["realtime_tracing"] + raw_realtime_session = self.config.extra_body.get("realtime_session_payload") + if isinstance(raw_realtime_session, dict): + for key, value in raw_realtime_session.items(): + if key not in payload and key != "extra_body": + payload[key] = value + 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 diff --git a/vllm_omni/experimental/fullduplex/openai/realtime_input.py b/vllm_omni/experimental/fullduplex/openai/realtime_input.py new file mode 100644 index 000000000..e3fe7f3ec --- /dev/null +++ b/vllm_omni/experimental/fullduplex/openai/realtime_input.py @@ -0,0 +1,1233 @@ +from __future__ import annotations + +import asyncio +import base64 +import binascii +from typing import Any +from uuid import uuid4 + +import numpy as np + +from vllm_omni.experimental.fullduplex.openai.audio import ( + convert_input_audio_with_rate, +) +from vllm_omni.experimental.fullduplex.openai.realtime_state import ( + REALTIME_INPUT_AUDIO_FORMATS, + REALTIME_OUTPUT_AUDIO_FORMATS, +) + + +class RealtimeInputTranslator: + """Translate and validate client Realtime events for the duplex core.""" + + 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._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 _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 + item = self._normalize_conversation_item(item) + 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_payload = event.get("session") if isinstance(event.get("session"), dict) else event + format_error = self._validate_realtime_session_audio_formats(session_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 + turn_detection_error = self._validate_realtime_turn_detection(session_payload) + if turn_detection_error is not None: + 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 + self._apply_realtime_session_defaults(session_payload) + session_payload.update(self._realtime_overlap_fields(session_payload)) + if not self._opened: + self._opened = True + self._initial_session_update = True + return self._session_create_from_realtime(session_payload) + return { + "type": "turn.signal", + "event": "session.update", + "payload": session_payload, + } + if event_type == "conversation.item.create": + item = event.get("item") + format_error = self._validate_conversation_item_audio_formats(item) + 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 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 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 + 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") + fmt, format_rate = self._parse_realtime_audio_format( + event.get("format") or event.get("input_audio_format") or self._input_audio_format + ) + sample_rate_hz = ( + event.get("sample_rate_hz") or event.get("sample_rate") or format_rate or self._input_sample_rate_hz + ) + 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 + 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 + looks_like_speech = self._input_looks_like_speech(event, audio=audio, fmt=fmt) + self._input_audio_buffer_has_audio = self._input_audio_buffer_has_audio or ( + looks_like_speech and isinstance(audio, str) and bool(audio) + ) + self._input_audio_buffer_had_non_speech = self._input_audio_buffer_had_non_speech or ( + not looks_like_speech and isinstance(audio, str) and bool(audio) + ) + 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, + } + 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 + payload["video_frames"] = [frame for frame in video_frames if isinstance(frame, str) and frame] + self._copy_realtime_input_hints(event, payload) + if not looks_like_speech: + payload["is_speech"] = False + return payload + if event_type == "input_audio_buffer.commit": + 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._pending_commit_item_ids.put(item_id) + 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), + "realtime_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._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() + 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 + + @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, + }, + { + "type": "conversation.item.created", + "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: + self._conversation_items[item_id] = item + 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: + self._last_conversation_item_id = next(reversed(self._conversation_items), 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, + event: dict[str, Any], + ) -> dict[str, object]: + return { + "type": "input_audio_buffer.committed", + "previous_item_id": self._previous_item_id(item_id), + "item_id": item_id, + "event": event, + } + + def _session_create_from_realtime(self, session_payload: dict[str, object]) -> dict[str, object]: + self._apply_realtime_session_defaults(session_payload) + 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 = ( + dict(session_payload.get("extra_body")) if isinstance(session_payload.get("extra_body"), 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"] + if isinstance(session_payload.get("metadata"), dict): + extra_body["realtime_metadata"] = dict(session_payload["metadata"]) + if isinstance(session_payload.get("include"), list): + extra_body["realtime_include"] = list(session_payload["include"]) + if isinstance(session_payload.get("prompt"), dict): + extra_body["realtime_prompt"] = dict(session_payload["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) + if isinstance(session_payload.get("input_audio_noise_reduction"), dict): + extra_body["realtime_input_audio_noise_reduction"] = dict(session_payload["input_audio_noise_reduction"]) + if isinstance(audio_input, dict) and isinstance(audio_input.get("noise_reduction"), dict): + extra_body["realtime_input_audio_noise_reduction"] = dict(audio_input["noise_reduction"]) + if isinstance(session_payload.get("audio"), dict): + extra_body["realtime_audio"] = dict(session_payload["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, input_rate = self._parse_realtime_audio_format(input_format) + if isinstance(input_format, str) and input_format.lower() in REALTIME_INPUT_AUDIO_FORMATS: + 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 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: + 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 isinstance(sample_rate, int | float) and sample_rate > 0: + self._input_sample_rate_hz = int(sample_rate) + if isinstance(output_rate, int | float) and output_rate > 0: + self._output_sample_rate_hz = int(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_session_audio_formats(cls, session_payload: dict[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, _ = cls._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, _ = cls._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 + + @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") + parsed_output, _ = cls._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 + + @classmethod + def _validate_conversation_item_audio_formats(cls, 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, _ = cls._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 + + @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 = 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 + + @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-Demo omni client: one + base64 JPEG per ~1 s audio chunk. HD slicing (max_slice_nums > 1) is + not implemented by the duplex adapter yet and is rejected explicitly + rather than silently ignored. + """ + if max_slice_nums not in (None, 1): + return "max_slice_nums > 1 (HD slicing) is not implemented by the duplex Realtime adapter" + 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 + + @staticmethod + def _validate_realtime_turn_detection(session_payload: dict[str, object]) -> str | None: + configured_values: list[tuple[str, object]] = [] + if "turn_detection" in session_payload: + configured_values.append(("turn_detection", session_payload["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: + configured_values.append(("audio.input.turn_detection", audio_input["turn_detection"])) + for field_path, turn_detection in configured_values: + if turn_detection is None: + continue + turn_detection_type = turn_detection.get("type") if isinstance(turn_detection, dict) else turn_detection + return ( + f"{field_path}={turn_detection_type!r} is not implemented by the duplex Realtime adapter; " + "set turn_detection to null and commit input explicitly, or use the model-owned duplex policy" + ) + return None + + @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: + item = event.get("item") + if not isinstance(item, dict): + return None + item = self._normalize_conversation_item(item) + previous_item_id = event.get("previous_item_id") + if isinstance(previous_item_id, str): + item["_previous_item_id"] = previous_item_id + item_id = str(item["id"]) + item_type = item.get("type") + role = item.get("role") + if item_type != "message" or role in {"assistant", "system"}: + return { + "type": "turn.signal", + "event": "conversation.item.create", + "payload": {"item": item}, + } + self._conversation_items[item_id] = item + content = item.get("content") + if not isinstance(content, list): + return None + text_chunks: list[str] = [] + audio_events: 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") + fmt, format_rate = self._parse_realtime_audio_format(part.get("format") or self._input_audio_format) + sample_rate_hz = ( + part.get("sample_rate_hz") or part.get("sample_rate") or format_rate or self._input_sample_rate_hz + ) + if not self._is_supported_realtime_input_format(fmt): + 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_event", + str(exc), + event_id=event.get("event_id"), + param="sample_rate_hz", + ) + ) + return None + if not isinstance(audio, str) or not audio: + continue + 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) + if audio_events: + self._input_audio_buffer_has_audio = True + transcript = self._input_transcript_from_item(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, + "realtime_item_id": item_id, + "response_create": False, + } + if transcript: + commit_payload["transcript"] = transcript + self._pending_outbound.put_nowait(commit_payload) + return audio_events[0] + if not text_chunks: + return None + text_item = dict(item) + text_item["status"] = "completed" + return { + "type": "turn.signal", + "event": "conversation.item.create", + "payload": {"item": text_item}, + } + + def _remember_input_transcript_hint(self, event: dict[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: + 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 + + def _pop_pending_commit_item_id(self) -> str: + try: + return self._pending_commit_item_ids.get_nowait() + except asyncio.QueueEmpty: + return f"item_{uuid4().hex}" + + 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}" + await self._send_realtime_payload( + { + "type": "input_audio_buffer.speech_started", + "audio_start_ms": int(event.get("audio_start_ms", 0) or 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, + } + ) diff --git a/vllm_omni/experimental/fullduplex/openai/realtime_output.py b/vllm_omni/experimental/fullduplex/openai/realtime_output.py new file mode 100644 index 000000000..57c07c852 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/openai/realtime_output.py @@ -0,0 +1,907 @@ +from __future__ import annotations + +from typing import Any +from uuid import uuid4 + +from vllm_omni.experimental.fullduplex.openai.audio import convert_output_audio +from vllm_omni.experimental.fullduplex.openai.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 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": + return [ + { + "type": "response.listen", + "session_id": event.get("session_id"), + "epoch": event.get("epoch"), + "response": { + "object": "realtime.response", + "status": "listening", + "metadata": event, + }, + } + ] + 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.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("realtime_item_id") + item_id = ( + event_item_id + if isinstance(event_item_id, str) and event_item_id + else self._pop_pending_commit_item_id() + ) + item = self._conversation_items.get(item_id) + created_payload: list[dict[str, object]] = [] + 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 + created_payload.extend(self._conversation_item_added_events(item)) + item["status"] = "completed" + payloads = created_payload + [ + self._input_audio_buffer_committed_event(item_id=item_id, event=event), + ] + 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": + return [{"type": "input_audio_buffer.cleared"}] + if event_type == "input_audio_buffer.cleared": + 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, + ) + payloads.append( + { + "type": "conversation.item.truncated", + "item_id": item_id, + "content_index": 0, + "audio_end_ms": committed_audio_ms, + "event": event, + } + ) + 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 == "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)] + 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, + ), + }, + }, + ) + 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 + + @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.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.audio.done", + "response_id": response_id, + "item_id": item_id, + "output_index": 0, + "content_index": 0, + } + ) + if transcript: + payloads.append( + { + "type": "response.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/experimental/fullduplex/openai/realtime_session.py b/vllm_omni/experimental/fullduplex/openai/realtime_session.py new file mode 100644 index 000000000..519c071cf --- /dev/null +++ b/vllm_omni/experimental/fullduplex/openai/realtime_session.py @@ -0,0 +1,139 @@ +from __future__ import annotations + +import json +from typing import Any +from uuid import uuid4 + +from fastapi import WebSocket + +from vllm_omni.experimental.fullduplex.openai import realtime_state as _state +from vllm_omni.experimental.fullduplex.openai.realtime_input import ( + RealtimeInputTranslator, +) +from vllm_omni.experimental.fullduplex.openai.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 _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()) + 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/experimental/fullduplex/openai/realtime_state.py b/vllm_omni/experimental/fullduplex/openai/realtime_state.py new file mode 100644 index 000000000..ac3b2a8de --- /dev/null +++ b/vllm_omni/experimental/fullduplex/openai/realtime_state.py @@ -0,0 +1,197 @@ +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_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", + "duplex_session_busy": "rate_limit_error", + "resource_exhausted": "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", + "missing_item_id": "invalid_request_error", + "item_not_found": "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 = 16000 + _output_audio_format: str = "pcm16" + _overlap_silence_rms: float = 0.003 + _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) + _pending_commit_item_ids: asyncio.Queue[str] = field(default_factory=asyncio.Queue) + _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) + + @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 + native_duplex = getter("minicpmo45_native_duplex") if getter is not None else None + 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["minicpmo45_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() + _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() + _pending_commit_item_ids: asyncio.Queue[str] = _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() diff --git a/vllm_omni/experimental/fullduplex/openai/runtime_adapter.py b/vllm_omni/experimental/fullduplex/openai/runtime_adapter.py new file mode 100644 index 000000000..f88cffd3b --- /dev/null +++ b/vllm_omni/experimental/fullduplex/openai/runtime_adapter.py @@ -0,0 +1,215 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable, Iterable, Mapping, MutableMapping +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 + + +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 + 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 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] + + 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/experimental/fullduplex/openai/runtime_bridge.py b/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py new file mode 100644 index 000000000..3d6674621 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py @@ -0,0 +1,1298 @@ +from __future__ import annotations + +import asyncio +import base64 +import inspect +from collections.abc import Mapping + +import numpy as np +from vllm.logger import init_logger + +from vllm_omni.experimental.fullduplex.engine.duplex_control_client import DuplexControlRequestError +from vllm_omni.experimental.fullduplex.engine.duplex_runtime import duplex_data_plane_request_info +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence +from vllm_omni.experimental.fullduplex.openai.protocol import ( + DuplexSession, + DuplexSessionState, +) +from vllm_omni.experimental.fullduplex.openai.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: + 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"] = 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 model unit (1 s at 16 kHz) of pcm_f32le silence, matching the + # official full-duplex behavior where the microphone keeps streaming + # silence while the assistant speaks; replies span multiple units. + _NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO = base64.b64encode(bytes(16000 * 4)).decode("ascii") + _NATIVE_RESPONSE_MAX_CONTINUATION_UNITS = 8 + + def _native_silence_unit_payload(self) -> dict[str, object]: + return { + "type": "audio", + "audio": self._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO, + "format": "pcm_f32le", + "sample_rate_hz": 16000, + } + + def _native_response_continuations_remaining(self, session: DuplexSession, response_id: str) -> bool: + if self._session_auto_responds(session): + return True + 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 + return count < self._NATIVE_RESPONSE_MAX_CONTINUATION_UNITS + + 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 + if not auto_response and count >= self._NATIVE_RESPONSE_MAX_CONTINUATION_UNITS: + return + payload = self._native_silence_unit_payload() + 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: + 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: + 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] + 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("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() + 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 + non_terminal_auto_listen = ( + self._session_auto_responds(session) + and session.active_response_id is not None + and session.active_request_id is not None + and native_result.get("end_of_turn") is not True + ) + if non_terminal_auto_listen: + await self._maybe_continue_native_response( + send_json, + session=session, + expected_epoch=expected_epoch, + ) + return close_reason, emitted_response + auto_response = self._session_auto_responds(session) + 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, + } + 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 + 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 + await send_json( + self._response_created_payload( + session, + response_id, + epoch=session.epoch, + ) + ) + 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, + ) + 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, + ) + 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 + self._serving_runtime_adapter.remove_session_state(session_id) + self._serving_runtime_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, + ) -> 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 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/experimental/fullduplex/openai/serving.py b/vllm_omni/experimental/fullduplex/openai/serving.py new file mode 100644 index 000000000..3ce23ea8b --- /dev/null +++ b/vllm_omni/experimental/fullduplex/openai/serving.py @@ -0,0 +1,1764 @@ +from __future__ import annotations + +import asyncio +import base64 +import binascii +import inspect +import json +from collections.abc import Mapping +from contextlib import suppress +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np +from fastapi import WebSocket +from vllm.logger import init_logger + +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.entrypoints.openai.duplex_capability import ( + should_enable_duplex_endpoint, +) +from vllm_omni.experimental.fullduplex.engine.duplex_runtime import duplex_resource_request_id +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence, DuplexSessionLifecycleMessage +from vllm_omni.experimental.fullduplex.openai.chat_fallback import ( + ChatFallbackProjectorMixin, +) +from vllm_omni.experimental.fullduplex.openai.protocol import ( + DuplexCapabilities, + DuplexCommittedInput, + DuplexOverlapPolicy, + DuplexPlaybackCommitPolicy, + DuplexSession, + DuplexSessionConfig, + DuplexSessionRegistry, + DuplexSessionState, + DuplexTurnController, + ResponseCreateOptions, +) +from vllm_omni.experimental.fullduplex.openai.realtime_session import ( + REALTIME_OUTPUT_AUDIO_FORMATS, + NativeRealtimeSessionProtocol, +) +from vllm_omni.experimental.fullduplex.openai.runtime_adapter import ( + ServingRuntimeAdapter, + ServingRuntimeConfigError, + ServingRuntimeSessionState, + load_serving_runtime_adapter, + validate_serving_runtime_adapter, +) +from vllm_omni.experimental.fullduplex.openai.runtime_bridge import ( + NativeRuntimeBridgeMixin, +) +from vllm_omni.experimental.fullduplex.openai.session_attachment import ( + DuplexJournalGapError, + DuplexSessionAttachmentRegistry, + InvalidResumeTokenError, +) +from vllm_omni.experimental.fullduplex.openai.session_runner import ( + DuplexSessionRunnerMixin, +) +from vllm_omni.experimental.fullduplex.openai.websocket import ( + DOMAIN_TERMINAL_EVENTS, + DuplexSessionTasks, + DuplexWebSocketActor, +) + +if TYPE_CHECKING: + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + +logger = init_logger(__name__) + +__all__ = ["OmniDuplexSessionHandler", "should_enable_duplex_endpoint"] + +_DEFAULT_CONFIG_TIMEOUT_S = 10.0 +_DEFAULT_IDLE_TIMEOUT_S = 300.0 + + +@dataclass(frozen=True) +class _DuplexSessionHandshake: + session: DuplexSession + resumed: bool = False + attachment_generation: int | None = None + + +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. + """ + + def __init__( + self, + *, + chat_service: OmniOpenAIServingChat, + config_timeout_s: float = _DEFAULT_CONFIG_TIMEOUT_S, + idle_timeout_s: float = _DEFAULT_IDLE_TIMEOUT_S, + duplex_session_config: DuplexSessionRuntimeConfig | None = None, + serving_runtime_adapter: ServingRuntimeAdapter | None = None, + serving_runtime_adapter_path: str | None = None, + ) -> None: + self._chat_service = chat_service + self._config_timeout_s = config_timeout_s + self._idle_timeout_s = idle_timeout_s + self._duplex_session_config = duplex_session_config or DuplexSessionRuntimeConfig() + 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, + ) + 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: + raise ValueError("A duplex serving runtime adapter must be explicitly configured") + 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 + 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, + ) + + async def handle_realtime_session(self, websocket: WebSocket) -> None: + await self.handle_session( + websocket, + realtime_protocol=NativeRealtimeSessionProtocol(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 _run_lifecycle_listener(self) -> None: + queue = self._lifecycle_queue + if not isinstance(queue, asyncio.Queue): + return + 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: + 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_runtime_adapter.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: + with suppress(Exception): + await attachment.close("session_expired") + + 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 + try: + return len(base64.b64decode(audio, validate=True)) + except (ValueError, binascii.Error): + return 0 + + 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 + + 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) + 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": False, + "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: + return { + "action": "barge_in", + "reason": "policy_barge_in_on_speech", + "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 _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 _runtime_session_state(self, session: DuplexSession) -> ServingRuntimeSessionState: + return self._serving_runtime_adapter.session_state(session.session_id) + + # Temporary compatibility accessors for downstream tests and extensions. + def _minicpmo_session_state(self, session: DuplexSession) -> ServingRuntimeSessionState: + return self._runtime_session_state(session) + + @property + def _minicpmo_sessions(self): + return self._serving_runtime_adapter.session_states + + @property + def _minicpmo_data_plane(self): + return self._serving_runtime_adapter.data_plane + + 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: + 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 + + 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( + 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 + try: + event = json.loads(raw) + except json.JSONDecodeError: + await send_json({"type": "error", "error": "Invalid JSON in session.create", "code": "invalid_json"}) + 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", + } + ) + return None + + config = DuplexSessionConfig.from_event(event) + if config.idle_timeout_s == _DEFAULT_IDLE_TIMEOUT_S: + config.idle_timeout_s = self._idle_timeout_s + use_native_runtime = self._uses_serving_runtime_adapter(config) + runtime_config: dict[str, object] = {} + if use_native_runtime: + try: + runtime_config = await self._serving_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 use_native_runtime: + session.replace_capabilities( + self._serving_runtime_adapter.capabilities( + max_sessions=self._duplex_session_config.max_sessions, + ) + ) + session.replace_runtime_config(runtime_config) + return _DuplexSessionHandshake(session=session) + + async def _resume_session_handshake( + 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 + ): + await send_json( + { + "type": "error", + "error": ( + "session.resume requires session_id, incarnation, resume_token, " + "and a non-negative event sequence" + ), + "code": "invalid_session_resume", + } + ) + return None + session = self._registry.get(session_id) + if session is None or session.state != DuplexSessionState.OPEN: + await send_json( + { + "type": "error", + "error": f"Unknown or expired duplex session: {session_id}", + "code": "session_resume_expired", + } + ) + return None + if not session.capabilities.supports_session_resume: + await send_json( + { + "type": "error", + "error": f"Session does not support resume: {session_id}", + "code": "unsupported_session_resume", + } + ) + return None + try: + await self._attachment_registry.authenticate_resume( + session_id, + incarnation=incarnation, + resume_token=resume_token, + last_received_server_event_seq=last_received, + ) + except InvalidResumeTokenError: + await send_json( + { + "type": "error", + "error": "Invalid duplex session resume token", + "code": "invalid_resume_token", + } + ) + return None + except DuplexJournalGapError: + await send_json( + { + "type": "session.resync_required", + "session_id": session_id, + "reason": "journal_gap", + } + ) + return None + except (KeyError, ValueError) as exc: + await send_json( + { + "type": "error", + "error": str(exc), + "code": "session_resume_conflict", + } + ) + 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( + 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", + } + ) + return None + + if not callable(attachment_send) or not callable(attachment_close): + raise RuntimeError("session.resume requires transport attachment callbacks") + + 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] + + try: + resumed = await self._attachment_registry.resume( + session_id, + incarnation=incarnation, + resume_token=resume_token, + last_received_server_event_seq=last_received, + 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", + } + ) + 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) + with suppress(Exception): + await replaced.close("session_replaced") + return _DuplexSessionHandshake( + session=session, + resumed=True, + attachment_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 + + def _uses_serving_runtime_adapter(self, config: DuplexSessionConfig) -> bool: + return self._serving_runtime_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._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._serving_runtime_adapter.runtime_config_for_update( + candidate_config, + dict(session.runtime_config), + ) + + def _runtime_session_candidate_update_error( + self, + session: DuplexSession, + candidate_config: DuplexSessionConfig, + ) -> dict[str, object] | None: + if not self._uses_native_input_append(session): + 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", + } + + @staticmethod + def _uses_native_input_append(session: DuplexSession) -> bool: + return ( + session.capabilities.implementation_level == "model_native_duplex" + and session.capabilities.supports_input_append + ) + + @staticmethod + def _config_requests_audio_output(config: DuplexSessionConfig) -> bool: + return any(str(modality).lower() == "audio" for modality in config.modalities) + + @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``). + """ + 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 + + async def _receive_text( + 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, + *, + 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 + + @staticmethod + def _commit_native_audio_input( + session: DuplexSession, + *, + 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_native_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 + + @staticmethod + def _native_audio_committed_payload( + session: DuplexSession, + *, + 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 + + @staticmethod + def _apply_session_update(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", + } + 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) + ) + return None + + 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._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): + 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 event.get("truncate") is True: + 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: + committed_history = session.truncate_history_item( + item_id, + audio_end_ms=committed_cursor, + playback=playback, + ) + 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, + ) + 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, + ) + 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, + ) + 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) + + 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: + 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: + 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 + + async def _abort_request_background( + self, + session: DuplexSession, + request_id: str, + send_json, + *, + notify: bool, + ) -> 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) + + 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, + } + ) diff --git a/vllm_omni/experimental/fullduplex/openai/session_attachment.py b/vllm_omni/experimental/fullduplex/openai/session_attachment.py new file mode 100644 index 000000000..b0faa65d2 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/openai/session_attachment.py @@ -0,0 +1,486 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import asyncio +import hashlib +import hmac +import json +import secrets +import time +from collections import deque +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass, field +from types import MappingProxyType + + +class InvalidResumeTokenError(RuntimeError): + pass + + +class DuplexJournalGapError(RuntimeError): + pass + + +class DuplexJournalOverflowError(RuntimeError): + pass + + +@dataclass(frozen=True) +class ResumeToken: + plaintext: str = field(repr=False) + + @classmethod + def generate(cls) -> ResumeToken: + return cls(secrets.token_urlsafe(32)) + + +@dataclass +class DuplexResumeCredential: + token_digest: bytes + + @classmethod + def from_token(cls, token: ResumeToken) -> DuplexResumeCredential: + return cls(token_digest=cls._digest(token.plaintext)) + + @staticmethod + def _digest(plaintext: str) -> bytes: + return hashlib.sha256(plaintext.encode("utf-8")).digest() + + def verify(self, plaintext: str) -> bool: + return hmac.compare_digest(self.token_digest, self._digest(plaintext)) + + def rotate(self) -> ResumeToken: + token = ResumeToken.generate() + self.token_digest = self._digest(token.plaintext) + return token + + +@dataclass(frozen=True) +class DuplexTransportAttachment: + generation: int + send: Callable[[dict[str, object]], Awaitable[None]] = field(repr=False) + close: Callable[[str], Awaitable[None]] = field(repr=False) + + +@dataclass(frozen=True) +class JournalEntry: + sequence: int + created_monotonic: float + encoded_bytes: int + payload: Mapping[str, object] = field(repr=False) + + def __post_init__(self) -> None: + object.__setattr__(self, "payload", MappingProxyType(dict(self.payload))) + + +class DuplexEventJournal: + def __init__( + self, + *, + max_bytes: int, + ttl_s: float, + clock: Callable[[], float] | None = None, + ) -> None: + if max_bytes <= 0: + raise ValueError("journal max_bytes must be positive") + if ttl_s <= 0: + raise ValueError("journal ttl_s must be positive") + self._max_bytes = max_bytes + self._ttl_s = ttl_s + self._clock = clock or time.monotonic + self._entries: deque[JournalEntry] = deque() + self._next_sequence = 1 + self._dropped_through = 0 + self._retained_bytes = 0 + self._overflowed = False + + @property + def retained_bytes(self) -> int: + return self._retained_bytes + + @property + def overflowed(self) -> bool: + return self._overflowed + + @property + def last_sequence(self) -> int: + return self._next_sequence - 1 + + def record(self, payload: Mapping[str, object]) -> JournalEntry: + if self._overflowed: + raise DuplexJournalOverflowError("duplex event journal already exceeded its byte limit") + self.prune() + sequence = self._next_sequence + sequenced_payload = dict(payload) + sequenced_payload["server_event_seq"] = sequence + encoded_bytes = len( + json.dumps( + sequenced_payload, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + ) + if self._retained_bytes + encoded_bytes > self._max_bytes: + self._overflowed = True + raise DuplexJournalOverflowError( + f"duplex event journal byte limit exceeded: {self._retained_bytes + encoded_bytes} > {self._max_bytes}" + ) + entry = JournalEntry( + sequence=sequence, + created_monotonic=self._clock(), + encoded_bytes=encoded_bytes, + payload=sequenced_payload, + ) + self._entries.append(entry) + self._retained_bytes += encoded_bytes + self._next_sequence += 1 + return entry + + def acknowledge(self, sequence: int) -> int: + if sequence < 0: + raise ValueError("acknowledged sequence must not be negative") + if sequence > self.last_sequence: + raise ValueError(f"acknowledged sequence {sequence} is newer than journal head {self.last_sequence}") + removed = 0 + while self._entries and self._entries[0].sequence <= sequence: + entry = self._entries.popleft() + self._retained_bytes -= entry.encoded_bytes + removed += 1 + self._dropped_through = max(self._dropped_through, sequence) + return removed + + def prune(self, now: float | None = None) -> int: + effective_now = self._clock() if now is None else now + cutoff = effective_now - self._ttl_s + removed = 0 + while self._entries and self._entries[0].created_monotonic <= cutoff: + entry = self._entries.popleft() + self._retained_bytes -= entry.encoded_bytes + self._dropped_through = max(self._dropped_through, entry.sequence) + removed += 1 + return removed + + def replay_after(self, sequence: int) -> tuple[JournalEntry, ...]: + if sequence < 0: + raise ValueError("replay sequence must not be negative") + self.prune() + if self._overflowed: + raise DuplexJournalGapError("duplex event journal overflowed; replay is incomplete") + if sequence < self._dropped_through: + raise DuplexJournalGapError( + f"requested sequence {sequence} is older than retained journal boundary {self._dropped_through}" + ) + if sequence > self.last_sequence: + raise ValueError(f"requested sequence {sequence} is newer than journal head {self.last_sequence}") + return tuple(entry for entry in self._entries if entry.sequence > sequence) + + +@dataclass(frozen=True) +class DuplexSessionAttachmentCreated: + session_id: str + incarnation: int + attachment_generation: int + resume_token: ResumeToken = field(repr=False) + + +@dataclass(frozen=True) +class DuplexSessionResumeResult: + session_id: str + incarnation: int + attachment_generation: int + resume_token: ResumeToken = field(repr=False) + replay_entries: tuple[JournalEntry, ...] = () + replaced_attachment: DuplexTransportAttachment | None = None + + +@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: + def __init__( + self, + *, + 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() + + def __repr__(self) -> str: + return f"{type(self).__name__}(session_ids={sorted(self._sessions)})" + + async def create( + self, + session_id: str, + *, + incarnation: int, + send: Callable[[dict[str, object]], Awaitable[None]], + close: Callable[[str], Awaitable[None]], + ) -> DuplexSessionAttachmentCreated: + async with self._lock: + if session_id in self._sessions: + raise ValueError(f"duplex attachment session already exists: {session_id}") + token = ResumeToken.generate() + 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, + ttl_s=self._replay_ttl_s, + clock=self._clock, + ), + attachment=DuplexTransportAttachment( + generation=generation, + send=send, + close=close, + ), + attachment_generation=generation, + ) + return DuplexSessionAttachmentCreated( + session_id=session_id, + incarnation=incarnation, + attachment_generation=generation, + resume_token=token, + ) + + async def send_event( + self, + session_id: str, + payload: Mapping[str, object], + *, + journal: bool = True, + ) -> JournalEntry | None: + """Sequence and dispatch one event to the current attachment. + + The per-session lock keeps wire order equal to journal order without + serializing unrelated sessions. A detached session still records + replayable events, but has no transport side effect. + """ + async with self._lock: + state = self._require(session_id) + async with state.outbound_lock: + async with self._lock: + if self._sessions.get(session_id) is not state: + raise KeyError(f"unknown duplex attachment session: {session_id}") + entry = state.journal.record(payload) if journal else None + attachment = state.attachment + wire_payload = dict(entry.payload) if entry is not None else dict(payload) + if attachment is not None: + await attachment.send(wire_payload) + return entry + + async def acknowledge(self, session_id: str, sequence: int) -> int: + 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 with self._lock: + state = self._sessions.get(session_id) + if state is None or 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 is_current_attachment(self, session_id: str, attachment_generation: int) -> bool: + async with self._lock: + state = self._sessions.get(session_id) + return ( + state is not None + and state.attachment is not None + and state.attachment_generation == attachment_generation + ) + + async def authenticate_resume( + 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, + ) + 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]], + close: Callable[[str], Awaitable[None]], + activation_payload_factory: Callable[[ResumeToken, int], Mapping[str, object]] | None = None, + ) -> DuplexSessionResumeResult: + async with self._lock: + state = self._require(session_id) + async with state.outbound_lock: + 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, + ) + 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 + state.attachment_generation += 1 + attachment_generation = state.attachment_generation + state.attachment = DuplexTransportAttachment( + generation=attachment_generation, + send=send, + close=close, + ) + if activation_payload_factory is not None: + try: + 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 + 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 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: + raise KeyError(f"unknown duplex attachment session: {session_id}") + 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}") + if state.credential.verify(resume_token): + return False + recovery_digest = state.recovery_token_digest + if ( + state.attachment is None + and recovery_digest is not None + and hmac.compare_digest(recovery_digest, DuplexResumeCredential._digest(resume_token)) + ): + return True + raise InvalidResumeTokenError(f"invalid resume token for duplex session {state.session_id}") + + +__all__ = [ + "DuplexEventJournal", + "DuplexJournalGapError", + "DuplexJournalOverflowError", + "DuplexResumeCredential", + "DuplexSessionAttachmentCreated", + "DuplexSessionAttachmentRegistry", + "DuplexSessionResumeResult", + "DuplexTransportAttachment", + "InvalidResumeTokenError", + "JournalEntry", + "ResumeToken", +] diff --git a/vllm_omni/experimental/fullduplex/openai/session_runner.py b/vllm_omni/experimental/fullduplex/openai/session_runner.py new file mode 100644 index 000000000..7e2d1ecaf --- /dev/null +++ b/vllm_omni/experimental/fullduplex/openai/session_runner.py @@ -0,0 +1,1912 @@ +from __future__ import annotations + +import asyncio +import json +import uuid +from contextlib import suppress +from copy import deepcopy + +from fastapi import WebSocket, WebSocketDisconnect +from vllm.logger import init_logger + +from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseActivity +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence +from vllm_omni.experimental.fullduplex.openai.audio import convert_input_audio_with_rate +from vllm_omni.experimental.fullduplex.openai.commit_policy import ( + CommitAction, + CommitSnapshot, + decide_commit_action, +) +from vllm_omni.experimental.fullduplex.openai.protocol import ( + DuplexPlaybackCommitPolicy, + DuplexSession, + DuplexSessionState, + DuplexTurnEventType, +) +from vllm_omni.experimental.fullduplex.openai.realtime_session import ( + NativeRealtimeSessionProtocol, +) +from vllm_omni.experimental.fullduplex.openai.runtime_adapter import ( + PcmAppendReservation, + ServingRuntimeSessionState, + payload_turn_id, +) +from vllm_omni.experimental.fullduplex.openai.session_attachment import ( + DuplexJournalOverflowError, +) +from vllm_omni.experimental.fullduplex.openai.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) -> 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, + ) + 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, + ) + + 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._serving_runtime_adapter.create_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]) -> None: + if not attachment_ready or session is None: + await actor.send_json(dict(payload)) + return + if realtime_protocol is None: + await send_attachment_payload( + payload, + journal=payload.get("type") not in {"session.created", "session.resumed"}, + ) + 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"}, + ) + + async def emit_event(payload: dict[str, object]) -> 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) + 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__"}) + + 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 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 + + async def _run() -> bool: + nonlocal runtime_closed + try: + 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: + 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()) + elif pcm_reservation is not None: + pcm_reservation.rollback() + 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: + raise + except Exception as exc: + if pcm_reservation is not None: + pcm_reservation.rollback() + 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() + return False + if actor.closing or runtime_closed or session.state != DuplexSessionState.OPEN: + if pcm_reservation is not None: + pcm_reservation.rollback() + return False + if before_append is not None and not before_append(): + if pcm_reservation is not None: + pcm_reservation.rollback() + return True + if pcm_reservation is not None and not pcm_reservation.active: + return False + return await _run() + + 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)) + 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(): + return False + if actor.native_append_tail is not None and not actor.native_append_tail.done(): + return False + if real_native_input_waiting(): + 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 handshake.resumed: + native = self._serving_runtime_adapter.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_runtime_adapter.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: + 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 + 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) + 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 == "__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() + 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 "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 not None + and 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 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"}, + ) + 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 + 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 + 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": + payload = event.get("payload") + if not isinstance(payload, dict): + await emit_event( + { + "type": "error", + "session_id": session.session_id, + "code": "bad_event", + "error": "session.update requires a session payload", + } + ) + continue + if not await wait_for_native_append_tail(): + continue + runtime_update_error = self._runtime_session_update_error(session, payload) + if runtime_update_error is not None: + await emit_event(runtime_update_error) + 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) + if update_error is not None: + await emit_event(update_error) + continue + runtime_update_error = self._runtime_session_candidate_update_error( + session, + candidate_config, + ) + if runtime_update_error is not None: + await emit_event(runtime_update_error) + continue + candidate_runtime_config = self._runtime_config_for_session_update( + session, + candidate_config, + ) + if not await self._signal_runtime_session( + session, + turn_event, + emit_event, + session_config=candidate_config.as_dict(), + runtime_config=candidate_runtime_config, + ): + continue + session.replace_config(candidate_config) + session.replace_runtime_config(candidate_runtime_config) + await emit_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 + message = self._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) + 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, + ) + 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 + 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 + 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 + # 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 self._uses_native_input_append(session): + mark_pending_silence_superseded() + overlap_active = not self._session_auto_responds(session) and native_response_in_progress() + 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 + cancelled = await self._cancel_active_response( + session, + actor.active_response_task, + emit_event, + reason="barge_in", + ) + 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": "barge_in", + "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": "barge_in", + "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. + 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 self._uses_native_input_append(session): + 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"}: + 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 ( + 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, + "turn_id": session.turn_id, + "epoch": session.epoch, + "empty": True, + "is_speech": False, + "no_response": True, + } + ) + 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, + "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, + realtime_item_id=event.get("realtime_item_id"), + transcript=event.get("transcript"), + ) + committed_payload = self._native_audio_committed_payload( + session, + committed=committed, + realtime_item_id=event.get("realtime_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, + realtime_item_id=event.get("realtime_item_id"), + transcript=event.get("transcript"), + turn_id=data_plane_turn_id, + ) + await emit_event( + self._native_audio_committed_payload( + session, + committed=committed, + realtime_item_id=event.get("realtime_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, + realtime_item_id=event.get("realtime_item_id"), + transcript=event.get("transcript"), + ) + await emit_event( + self._native_audio_committed_payload( + session, + committed=committed, + realtime_item_id=event.get("realtime_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, + realtime_item_id=event.get("realtime_item_id"), + transcript=event.get("transcript"), + ) + if native_had_uncommitted_audio + else None + ) + await emit_event( + self._native_audio_committed_payload( + session, + committed=committed, + realtime_item_id=event.get("realtime_item_id"), + transcript=event.get("transcript"), + ) + ) + else: + await emit_event( + { + "type": "input.committed", + "session_id": session.session_id, + "empty": True, + } + ) + continue + if committed is not None: + realtime_item_id = event.get("realtime_item_id") + if isinstance(realtime_item_id, str): + session.register_history_item(realtime_item_id, committed.message) + await emit_event( + self._input_committed_payload( + session, + committed, + realtime_item_id=realtime_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", + ) + 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 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/experimental/fullduplex/openai/websocket.py b/vllm_omni/experimental/fullduplex/openai/websocket.py new file mode 100644 index 000000000..18cd11c8d --- /dev/null +++ b/vllm_omni/experimental/fullduplex/openai/websocket.py @@ -0,0 +1,265 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +from fastapi import WebSocketDisconnect + +INPUT_EVENTS = frozenset( + { + "input.text.append", + "input_audio_buffer.append", + "input.commit", + "input_audio_buffer.commit", + "response.create", + } +) + +_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"}) + + +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 + + +MODEL_OUTPUT_EVENTS = frozenset( + { + "response.created", + "response.listen", + "response.speak", + "response.output_item.added", + "response.content_part.added", + "response.output_audio.delta", + "response.audio.delta", + "response.output_audio.done", + "response.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 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 + for task in tasks: + task.cancel() + try: + await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=timeout_s) + except TimeoutError: + pass + return True + + +@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", +] diff --git a/vllm_omni/experimental/fullduplex/output.py b/vllm_omni/experimental/fullduplex/output.py new file mode 100644 index 000000000..bfe2243f9 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/output.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from typing import Any + +from vllm_omni.outputs import OmniRequestOutput + +DUPLEX_OUTPUT_DECISION_KEY = "_vllm_omni.experimental.fullduplex.duplex_output_decision" + + +def attach_duplex_output_decision( + output: OmniRequestOutput, + decision: object, +) -> OmniRequestOutput: + custom_output = dict(output._custom_output) + custom_output[DUPLEX_OUTPUT_DECISION_KEY] = decision + output._custom_output = custom_output + return output + + +def get_duplex_output_decision(output: object) -> Any | None: + if not isinstance(output, OmniRequestOutput): + return None + custom_output = output._custom_output + if not isinstance(custom_output, dict): + return None + return custom_output.get(DUPLEX_OUTPUT_DECISION_KEY) + + +__all__ = [ + "DUPLEX_OUTPUT_DECISION_KEY", + "attach_duplex_output_decision", + "get_duplex_output_decision", +] diff --git a/vllm_omni/experimental/fullduplex/request_client.py b/vllm_omni/experimental/fullduplex/request_client.py new file mode 100644 index 000000000..d39a037a2 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/request_client.py @@ -0,0 +1,383 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import asyncio +import time +from collections.abc import Callable, MutableMapping +from typing import Protocol + +from vllm_omni.engine.messages import ErrorMessage, OutputMessage +from vllm_omni.entrypoints.client_request_state import ClientRequestState +from vllm_omni.experimental.fullduplex.engine.contracts import ( + duplex_data_plane_request_info, + duplex_resource_request_id, +) +from vllm_omni.experimental.fullduplex.engine.lease import DuplexLeaseActivity +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence +from vllm_omni.experimental.fullduplex.output import get_duplex_output_decision +from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics +from vllm_omni.outputs import OmniRequestOutput + + +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/inputs/data.py b/vllm_omni/inputs/data.py index f576b9724..1eb4dc202 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -21,15 +21,16 @@ from vllm.inputs.engine import TokensInput class OmniTextPrompt(TextPrompt): - """Text prompt with optional embeddings and additional information. + """Text prompt with optional embeddings and stage payloads. Extends TextPrompt to support prompt embeddings and additional information payloads for direct transfer between pipeline stages. Attributes: prompt_embeds: Optional tensor containing prompt embeddings - additional_information: Optional dictionary containing additional - information (tensors or lists) to pass along with the prompt + additional_information: Legacy optional dictionary containing + additional information to pass along with the prompt. + model_intermediate_buffer: Optional runner-owned stage payload. """ negative_prompt: NotRequired[str] @@ -39,18 +40,20 @@ class OmniTextPrompt(TextPrompt): prompt_embeds: NotRequired[torch.Tensor] negative_prompt_embeds: NotRequired[torch.Tensor] additional_information: NotRequired[dict[str, Any]] + model_intermediate_buffer: NotRequired[dict[str, Any]] class OmniTokensPrompt(TokensPrompt): - """Tokens prompt with optional embeddings and additional information. + """Tokens prompt with optional embeddings and stage payloads. Extends TokensPrompt to support prompt embeddings and additional information payloads for direct transfer between pipeline stages. Attributes: prompt_embeds: Optional tensor containing prompt embeddings - additional_information: Optional dictionary containing additional - information (tensors or lists) to pass along with the prompt + additional_information: Legacy optional dictionary containing + additional information to pass along with the prompt. + model_intermediate_buffer: Optional runner-owned stage payload. """ negative_prompt: NotRequired[str] @@ -61,10 +64,11 @@ class OmniTokensPrompt(TokensPrompt): # New: optional additional information dictionary # Values may be torch.Tensor or list additional_information: NotRequired[dict[str, Any]] + model_intermediate_buffer: NotRequired[dict[str, Any]] class OmniTokenInputs(TokensInput): - """Token inputs with optional embeddings and additional information. + """Token inputs with optional embeddings and stage payloads. Extends TokensInput to support prompt embeddings and additional information payloads for direct transfer between pipeline stages. @@ -72,8 +76,9 @@ class OmniTokenInputs(TokensInput): Attributes: prompt_embeds: Optional tensor containing prompt embeddings aligned with token IDs - additional_information: Optional dictionary containing additional - information (tensors or lists) to pass along with the inputs + additional_information: Legacy optional dictionary containing + additional information to pass along with the inputs. + model_intermediate_buffer: Optional runner-owned stage payload. """ # New: optional prompt embeddings aligned with token ids @@ -84,18 +89,20 @@ class OmniTokenInputs(TokensInput): # New: optional additional information dictionary # Values may be torch.Tensor or list additional_information: NotRequired[dict[str, Any]] + model_intermediate_buffer: NotRequired[dict[str, Any]] class OmniEmbedsPrompt(EmbedsPrompt): - """Embeddings prompt with optional additional information. + """Embeddings prompt with optional stage payloads. Extends EmbedsPrompt to support additional information payloads for direct transfer between pipeline stages. Attributes: prompt_embeds: Optional tensor containing prompt embeddings - additional_information: Optional dictionary containing additional - information (tensors or lists) to pass along with the prompt + additional_information: Legacy optional dictionary containing + additional information to pass along with the prompt. + model_intermediate_buffer: Optional runner-owned stage payload. """ # New: optional prompt embeddings aligned with token ids @@ -105,6 +112,7 @@ class OmniEmbedsPrompt(EmbedsPrompt): # New: optional additional information dictionary # Values may be torch.Tensor or list additional_information: NotRequired[dict[str, Any]] + model_intermediate_buffer: NotRequired[dict[str, Any]] class OmniCustomPrompt(TypedDict, total=False): @@ -143,6 +151,7 @@ def token_inputs_omni( cache_salt: str | None = None, prompt_embeds: torch.Tensor | None = None, additional_information: dict[str, Any] | None = None, + model_intermediate_buffer: dict[str, Any] | None = None, ) -> OmniTokenInputs: """Construct token inputs with optional embeddings and metadata. @@ -154,8 +163,9 @@ def token_inputs_omni( prompt: Optional prompt string cache_salt: Optional cache salt for prefix caching prompt_embeds: Optional tensor containing prompt embeddings - additional_information: Optional dictionary containing additional - information (tensors or lists) + additional_information: Legacy optional dictionary containing + additional information (tensors or lists). + model_intermediate_buffer: Optional runner-owned stage payload. Returns: OmniTokenInputs instance with the provided data @@ -170,6 +180,8 @@ def token_inputs_omni( inputs["prompt_embeds"] = prompt_embeds if additional_information is not None: inputs["additional_information"] = additional_information + if model_intermediate_buffer is not None: + inputs["model_intermediate_buffer"] = model_intermediate_buffer return inputs diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index 7282d7a52..3da1ef897 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -60,6 +60,9 @@ class OmniInputPreprocessor(InputPreprocessor): additional_information = parsed_content.get("additional_information") if additional_information is not None: inputs["additional_information"] = additional_information + model_intermediate_buffer = parsed_content.get("model_intermediate_buffer") + if model_intermediate_buffer is not None: + inputs["model_intermediate_buffer"] = model_intermediate_buffer elif "mm_processor_kwargs" in parsed_content: # Presence — not truthiness. An explicitly-set empty dict still # signals "route through the multimodal processor" (needed for @@ -81,6 +84,7 @@ class OmniInputPreprocessor(InputPreprocessor): prompt_token_ids, prompt_embeds=parsed_content.get("prompt_embeds"), additional_information=parsed_content.get("additional_information"), + model_intermediate_buffer=parsed_content.get("model_intermediate_buffer"), ) prompt_embeds = parsed_content.get("prompt_embeds") if prompt_embeds is not None: @@ -88,6 +92,9 @@ class OmniInputPreprocessor(InputPreprocessor): additional_information = parsed_content.get("additional_information") if additional_information is not None: inputs["additional_information"] = additional_information + model_intermediate_buffer = parsed_content.get("model_intermediate_buffer") + if model_intermediate_buffer is not None: + inputs["model_intermediate_buffer"] = model_intermediate_buffer if cache_salt := parsed_content.get("cache_salt"): inputs["cache_salt"] = cache_salt @@ -101,6 +108,7 @@ class OmniInputPreprocessor(InputPreprocessor): prompt_token_ids = self._truncate_inputs(parsed_content["prompt_token_ids"], tokenization_kwargs) prompt_embeds = parsed_content.get("prompt_embeds") additional_information = parsed_content.get("additional_information") + model_intermediate_buffer = parsed_content.get("model_intermediate_buffer") multi_modal_data = parsed_content.get("multi_modal_data") @@ -119,11 +127,14 @@ class OmniInputPreprocessor(InputPreprocessor): prompt_token_ids=prompt_token_ids, prompt_embeds=prompt_embeds, additional_information=additional_information, + model_intermediate_buffer=model_intermediate_buffer, ) if prompt_embeds is not None: inputs["prompt_embeds"] = prompt_embeds if additional_information is not None: inputs["additional_information"] = additional_information + if model_intermediate_buffer is not None: + inputs["model_intermediate_buffer"] = model_intermediate_buffer if prompt_text := parsed_content.get("prompt"): inputs["prompt"] = prompt_text if cache_salt := parsed_content.get("cache_salt"): @@ -147,6 +158,9 @@ class OmniInputPreprocessor(InputPreprocessor): additional_information = parsed_content.get("additional_information") if additional_information is not None: inputs["additional_information"] = additional_information # type: ignore[typeddict-unknown-key] + model_intermediate_buffer = parsed_content.get("model_intermediate_buffer") + if model_intermediate_buffer is not None: + inputs["model_intermediate_buffer"] = model_intermediate_buffer # type: ignore[typeddict-unknown-key] return inputs diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index a282d7d50..3d08adae0 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -381,6 +381,7 @@ class OrchestratorAggregator: finished: bool, final_output_type: str | None, output_to_yield: Any | None, + event_cursor: int = 0, ) -> None: """Process and record stage metrics. @@ -409,7 +410,7 @@ class OrchestratorAggregator: return rid_key = str(req_id) - stage_snapshot = self._build_stage_metrics_snapshot(rid_key) + stage_snapshot = self._build_stage_metrics_snapshot(rid_key, event_cursor=event_cursor) # 3. Not finished yet — expose incremental per-stage snapshot for streaming clients. if not finished: @@ -423,7 +424,11 @@ class OrchestratorAggregator: # token fields only for text stages (OpenAI-style completion token accounting). output_to_yield.metrics = {"stage_metrics": stage_snapshot} stage_event = next( - (evt for evt in reversed(self.stage_events.get(rid_key, [])) if evt.stage_id == stage_id), + ( + evt + for evt in reversed(self.stage_events.get(rid_key, [])[event_cursor:]) + if evt.stage_id == stage_id + ), None, ) if stage_event is not None and stage_event.final_output_type == "text": @@ -541,10 +546,18 @@ class OrchestratorAggregator: current[defs.TIME_PER_OUTPUT_UNIT_MS] = remaining_ms / float(output_count - 1) if output_count > 1 else 0.0 return current - def _build_stage_metrics_snapshot(self, req_id: str) -> dict[str, dict[str, Any]]: + def stage_event_cursor(self, req_id: str) -> int: + return len(self.stage_events.get(str(req_id), ())) + + def _build_stage_metrics_snapshot( + self, + req_id: str, + *, + event_cursor: int = 0, + ) -> dict[str, dict[str, Any]]: """Aggregate per-stage metrics for ``req_id`` (string key), for streaming/benchmark clients.""" snapshot: dict[str, dict[str, Any]] = {} - for evt in self.stage_events.get(req_id, []): + for evt in self.stage_events.get(req_id, [])[event_cursor:]: sid = int(evt.stage_id) if evt.stage_id is not None else -1 if sid < 0: continue 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 d9b48807a..ba51c9cfc 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 @@ -18,10 +18,13 @@ # limitations under the License. from collections.abc import Iterable +from contextlib import suppress from functools import cached_property +from typing import Any import torch import torch.nn as nn +import torch.nn.functional as F from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP @@ -31,6 +34,9 @@ from vllm.sequence import IntermediateTensors from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata +from vllm_omni.experimental.fullduplex.engine.intermediate import get_stream_request_key +from vllm_omni.experimental.fullduplex.minicpmo45.policy import MiniCPMO45DuplexPolicy +from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingRow from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( MiniCPMO45OmniLLMDummyInputsBuilder, MiniCPMO45OmniLLMMultiModalProcessor, @@ -52,11 +58,9 @@ logger = init_logger(__name__) class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP, SupportsMRoPE): """MiniCPM-o 4.5 Omni model for conditional generation. - Two-stage pipeline: - - thinker (model_stage="llm"): image / video / audio encoders + 3D - resampler + the omni LLM that emits text + hidden states. - - talker (model_stage="tts"): MiniCPMTTS + the in-process Token2Wav - vocoder that emits the final audio waveform directly. + This model has two pipeline stages: + - llm: multimodal thinker and text generation + - tts: talker generation followed by its built-in Token2Wav vocoder """ @classmethod @@ -80,6 +84,11 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp # Store configs self.config = config self.multimodal_config = multimodal_config + from vllm_omni.experimental.fullduplex.minicpmo45.compat import ( + patch_minicpmo_remote_config, + ) + + patch_minicpmo_remote_config(config) self.model_stage = vllm_config.model_config.model_stage @@ -89,8 +98,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp vllm_config=vllm_config, prefix=maybe_prefix(prefix, "thinker"), hf_config=config, - # Registry key — must match the entries declared in - # vllm_omni/model_executor/models/registry.py::_OMNI_MODELS. + # Use registry architecture key architectures=["MiniCPMO45OmniLLMForConditionalGeneration"], ) self.model = self.thinker @@ -98,21 +106,18 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp elif self.model_stage == "tts": self.thinker = None - # Initialize talker model — runs MiniCPMTTS + the in-process - # Token2Wav vocoder and emits the final audio waveform directly. + # Initialize talker model (LLM generation) self.talker = init_vllm_registered_model( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "talker"), hf_config=config, - # Registry key — must match the entries declared in - # vllm_omni/model_executor/models/registry.py::_OMNI_MODELS. + # Use registry architecture key architectures=["MiniCPMO45OmniTTSForConditionalGeneration"], ) # Initialize multimodal components if needed if hasattr(self.talker, "init_multi_modal"): self.talker.init_multi_modal(config) self.model = self.talker - else: raise ValueError(f"Invalid model stage: {self.model_stage}. Must be one of: 'llm', 'tts'") @@ -124,6 +129,8 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp ) self._language_model_names = ["model"] + self.prefer_model_sampler = self.model_stage in {"llm", "tts"} + self.has_preprocess = self.model_stage == "llm" @cached_property def sampler(self): @@ -133,6 +140,90 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp return Sampler() + def prepare_duplex_sampling( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + rows: tuple[DuplexSamplingRow, ...], + ) -> None: + """Apply MiniCPM duplex policy before the standard model sampler.""" + 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 + } + 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} + ) + 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 + } + if self.model_stage != "llm" or not rows or logits.ndim != 2: + return + + token_ids = self._minicpmo45_native_duplex_token_ids() + listen_id = int(token_ids.get("listen_token_id", -1)) + tts_bos_id = int(token_ids.get("tts_bos_token_id", -1)) + turn_eos_id = int(token_ids.get("turn_eos_token_id", -1)) + if listen_id < 0 or listen_id >= logits.shape[-1]: + return + + force_listen_segments = getattr( + self, + "_minicpmo45_force_listen_applied_segments", + None, + ) + if not isinstance(force_listen_segments, set): + force_listen_segments = set() + self._minicpmo45_force_listen_applied_segments = force_listen_segments + helper = getattr(self, "_minicpmo45_duplex_data_plane_helper", None) + helper_sessions = getattr(helper, "sessions", None) if helper is not None else None + for row in rows: + row_idx = row.row_idx + if row_idx < 0 or row_idx >= logits.shape[0]: + continue + payload = row.payload + if not isinstance(payload, dict): + continue + force_listen = payload.get("force_listen") is True + is_speech = payload.get("is_speech") + redirect_listen = False + 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 + 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 = ( + bool(getattr(state, "pending_speech_context", False)) if state is not None else False + ) + if is_speech is True: + if state is not None: + if not getattr(state, "current_turn_ended", True): + redirect_listen = True + with suppress(Exception): + state.last_terminator_token = None + else: + turn_ended = bool(getattr(state, "current_turn_ended", True)) if state is not None else False + if turn_ended and not pending_speech_context: + force_listen = True + + if not force_listen and not redirect_listen: + continue + if force_listen and segment_key in force_listen_segments: + continue + if force_listen: + logits[row_idx, :] = float("-inf") + logits[row_idx, listen_id] = 0.0 + force_listen_segments.add(segment_key) + elif 0 <= tts_bos_id < logits.shape[-1]: + if logits[row_idx, listen_id] > logits[row_idx, tts_bos_id]: + logits[row_idx, tts_bos_id] = logits[row_idx, listen_id] + logits[row_idx, listen_id] = float("-inf") + # -------------------- Device utilities -------------------- @staticmethod def _module_device(module: nn.Module) -> torch.device: @@ -150,7 +241,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp thinker_device: str | torch.device | None = None, talker_device: str | torch.device | None = None, ) -> None: - """Optionally move thinker / talker to different devices. + """Optionally move the thinker and talker to different devices. Example: model.move_submodules_to_devices( @@ -168,7 +259,22 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp input_ids: torch.Tensor, multimodal_embeddings=None, ) -> torch.Tensor: - return self.model.get_input_embeddings(input_ids, multimodal_embeddings) + embed_fn = getattr(self.model, "get_input_embeddings", None) + if callable(embed_fn): + try: + return embed_fn(input_ids, multimodal_embeddings) + except TypeError: + embeddings = embed_fn() + if callable(embeddings): + return embeddings(input_ids) + except AttributeError: + pass + + embed_tokens = getattr(getattr(getattr(self.model, "llm", None), "model", None), "embed_tokens", None) + if callable(embed_tokens): + return embed_tokens(input_ids) + + raise AttributeError(f"{type(self.model).__name__} does not expose token embeddings") def embed_input_ids( self, @@ -181,6 +287,182 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp return self.get_input_embeddings(input_ids) return super().embed_input_ids(input_ids, multimodal_embeddings, is_multimodal=is_multimodal) + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, object]]: + """Model-runner data-plane hook for MiniCPM-o 4.5 duplex audio. + + The scheduler owns the request, block table, attention metadata, KV, + and sampler. This hook only turns the current duplex audio append into + the prompt embeddings consumed by the normal runner forward. + """ + if self.model_stage != "llm": + embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids) + return input_ids, embeds, {} + + duplex = kwargs.get("duplex") + if not isinstance(duplex, dict) or duplex.get("data_plane") is not True: + embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids) + return input_ids, embeds, {} + + prompt_len_meta = kwargs.get("duplex_prompt_len") + token_offset_meta = kwargs.get("duplex_token_offset", 0) + if ( + isinstance(prompt_len_meta, int) + and isinstance(token_offset_meta, int) + and token_offset_meta >= prompt_len_meta + ): + # Decode step of the resumable duplex request: input_ids are the + # runner-sampled tokens and the normal embedding lookup is the + # correct input. Slicing the (prompt-only) duplex embeddings here + # would come up empty and pad-fill, feeding a embedding in + # place of every sampled token and corrupting generation. + embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids) + return input_ids, embeds, {} + + 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) + state = helper.sessions.get(session_key) + if state is None: + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + _MiniCPMO45Stage0SessionState, + ) + + state = _MiniCPMO45Stage0SessionState(session_id=session_id) + helper.sessions[session_key] = state + session_config = duplex.get("session_config") + session_config = dict(session_config) if isinstance(session_config, dict) else {} + runtime_config = duplex.get("runtime_config") + runtime_config = dict(runtime_config) if isinstance(runtime_config, dict) else {} + if hasattr(helper.thinker, "audio_past_key_values"): + helper.thinker.audio_past_key_values = None + helper._configure_streaming_processor(state) + helper._prepare_session_context(state, session_config, runtime_config=runtime_config) + + audio_waveform = helper._decode_audio_payload(payload) + try: + video_frames = helper._decode_video_frames_payload(payload) + except ValueError as exc: + 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": str(exc)}} + seq = duplex.get("seq") + try: + seq = int(seq) if seq is not None else None + except (TypeError, ValueError): + seq = None + epoch = duplex.get("epoch") + try: + epoch = int(epoch) if epoch is not None else None + except (TypeError, ValueError): + epoch = None + result = helper._stage_prefill_embeddings_only( + state, + audio_waveform, + video_frames=video_frames, + epoch=epoch, + seq=seq, + is_speech=bool(payload.get("is_speech", False)), + final=bool(duplex.get("final")), + ) + update_result = dict(result) + update_result.pop("inputs_embeds", None) + if result.get("success") is not True: + embeds = input_embeds if input_embeds is not None else self.get_input_embeddings(input_ids) + return input_ids, embeds, {"duplex": update_result} + + target_dtype = ( + input_embeds.dtype if input_embeds is not None else self.get_input_embeddings(input_ids[:1]).dtype + ) + full_req_embeds = result["inputs_embeds"].to(device=input_ids.device, dtype=target_dtype) + full_input_token_ids = list(result.get("input_token_ids") or []) + prompt_len = kwargs.get("duplex_prompt_len") + try: + prompt_len = int(prompt_len) if prompt_len is not None else int(full_req_embeds.shape[0]) + except (TypeError, ValueError): + prompt_len = int(full_req_embeds.shape[0]) + pad_token_id = helper.stage_padding_token_id() + if prompt_len > int(full_req_embeds.shape[0]): + pad_len = prompt_len - int(full_req_embeds.shape[0]) + pad_ids = torch.full( + (pad_len,), + pad_token_id, + dtype=input_ids.dtype, + device=input_ids.device, + ) + pad_embeds = self.get_input_embeddings(pad_ids).to(dtype=full_req_embeds.dtype) + # The appended duplex tokens occupy the tail of the request prompt + # and the runner schedules the span [num_computed_tokens, prompt_len). + # Padding must therefore sit in front of the real chunk embeddings; + # otherwise the audio lands outside the scheduled span, is never + # forwarded, and generation runs on pad tokens only. Keeping the + # embeddings last also places the decode position directly after the + # final audio embedding, matching the official listen/speak decision + # point. + full_req_embeds = torch.cat([pad_embeds, full_req_embeds], dim=0) + full_input_token_ids = [pad_token_id] * pad_len + full_input_token_ids + elif prompt_len < int(full_req_embeds.shape[0]): + logger.warning( + "MiniCPM-o duplex append produced %d embeddings but the scheduler " + "reserved only %d prompt slots; the tail will be truncated. " + "Increase the duplex scheduler token budget.", + int(full_req_embeds.shape[0]), + prompt_len, + ) + + span_len = int(input_ids.shape[0]) + token_offset = kwargs.get("duplex_token_offset", 0) + try: + token_offset = max(0, int(token_offset)) + except (TypeError, ValueError): + token_offset = 0 + req_embeds = full_req_embeds[token_offset : token_offset + span_len] + if req_embeds.shape[0] < span_len: + pad_ids = torch.full( + (span_len - req_embeds.shape[0],), + pad_token_id, + dtype=input_ids.dtype, + device=input_ids.device, + ) + pad_embeds = self.get_input_embeddings(pad_ids).to(dtype=req_embeds.dtype) + req_embeds = torch.cat([req_embeds, pad_embeds], dim=0) + elif req_embeds.shape[0] > span_len: + req_embeds = req_embeds[:span_len] + + input_token_ids = full_input_token_ids[token_offset : token_offset + span_len] + if len(input_token_ids) < span_len: + input_token_ids.extend([pad_token_id] * (span_len - len(input_token_ids))) + if input_token_ids: + req_input_ids = torch.tensor(input_token_ids, dtype=input_ids.dtype, device=input_ids.device) + update_result["duplex_prompt_token_ids"] = full_input_token_ids + else: + req_input_ids = torch.full_like(input_ids, helper._required_token_id("unit_token_id")) + return req_input_ids, req_embeds, {"duplex": update_result} + + def _duplex_data_plane_helper(self): + helper = getattr(self, "_minicpmo45_duplex_data_plane_helper", None) + if helper is not None: + return helper + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import MiniCPMO45Stage0DuplexRuntime + + model_path = getattr(getattr(self.vllm_config, "model_config", None), "model", None) + device = str(self._module_device(self.thinker if self.thinker is not None else self)) + helper = MiniCPMO45Stage0DuplexRuntime(self, model_path=model_path, device=device) + self._minicpmo45_duplex_data_plane_helper = helper + return helper + def get_multimodal_embeddings(self, **kwargs): # Delegate to the active stage submodule when it implements MM encoding. mm_fn = getattr(self.model, "get_multimodal_embeddings", None) @@ -192,6 +474,174 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp """vLLM V1 encoder profiling calls this; the inherited Protocol stub returns None.""" return self.get_multimodal_embeddings(**kwargs) + def _run_tts_request( + self, + *, + input_ids: torch.Tensor | None, + positions: torch.Tensor | None, + inputs_embeds: torch.Tensor | None, + talker_info: dict[str, Any], + device: torch.device, + ) -> tuple[dict[str, Any] | None, bool, bool]: + meta_info = talker_info.get("meta") if isinstance(talker_info.get("meta"), dict) else {} + if talker_info.get("native_duplex") is True: + # ids/hidden_states may be accumulated for the talker KV stream, + # but displayed transcript belongs to the current thinker segment. + tts_text = meta_info.get("native_duplex_segment_text", "") + else: + tts_text = talker_info.get("llm_output_text", "") + if isinstance(tts_text, list): + tts_text = ( + tts_text[-1] + if talker_info.get("native_duplex") is True and tts_text + else (tts_text[0] if tts_text else "") + ) + if not isinstance(tts_text, str): + tts_text = "" + + with torch.inference_mode(): + talker_result = self.talker( + input_ids=input_ids, + positions=positions, + inputs_embeds=inputs_embeds, + additional_information=talker_info, + ) + chunk_flags = getattr(self.talker, "_ar_last_chunk_flags", [True]) + chunk_is_last = bool(chunk_flags[-1]) if chunk_flags else True + turn_end_flags = getattr(self.talker, "_ar_turn_end_flags", [False]) + turn_ended = bool(turn_end_flags[-1]) if turn_end_flags else False + if talker_info.get("native_duplex") is True: + emitted_text = getattr(self.talker, "_ar_last_emitted_text", "") + tts_text = emitted_text if isinstance(emitted_text, str) else "" + + if not (isinstance(talker_result, tuple) and len(talker_result) == 2): + return None, chunk_is_last, turn_ended + + mel_spec, waveform = talker_result + mm_out: dict[str, Any] = {} + duplex_info = talker_info.get("duplex") if isinstance(talker_info.get("duplex"), dict) else {} + if talker_info.get("native_duplex") is True: + turn_id = duplex_info.get("turn_id") + if isinstance(turn_id, int): + mm_out["meta.duplex_turn_id"] = torch.tensor([turn_id], dtype=torch.int32, device=device) + epoch = duplex_info.get("epoch") + if isinstance(epoch, int): + mm_out["meta.duplex_epoch"] = torch.tensor([epoch], dtype=torch.int32, device=device) + mm_out["meta.tts_is_last_chunk"] = torch.tensor( + [int(chunk_is_last)], + dtype=torch.int32, + device=device, + ) + if talker_info.get("native_duplex") is True: + mm_out["meta.turn_end"] = torch.tensor( + [int(turn_ended)], + dtype=torch.int32, + device=device, + ) + if tts_text or talker_info.get("native_duplex") is True: + mm_out["meta.llm_output_text_utf8"] = torch.tensor( + list(tts_text.encode("utf-8")), + dtype=torch.uint8, + device=device, + ) + mm_out["meta.audio_text_total_chars"] = torch.tensor( + [len(tts_text)], + dtype=torch.int32, + device=device, + ) + if mel_spec is not None: + mm_out["mel_spec"] = [mel_spec] + if waveform is not None: + mm_out["model_outputs"] = [waveform] + elif mel_spec is not None: + mm_out["model_outputs"] = [mel_spec] + return mm_out, chunk_is_last, turn_ended + + @staticmethod + def _slice_tts_request_rows(value: torch.Tensor | None, start: int, end: int) -> torch.Tensor | None: + return value[start:end] if isinstance(value, torch.Tensor) else value + + def _run_batched_tts_requests( + self, + *, + input_ids: torch.Tensor | None, + positions: torch.Tensor | None, + inputs_embeds: torch.Tensor | None, + runtime_info: list[object], + request_token_spans: object, + num_tokens: int, + hidden_dim: int, + device: torch.device, + ) -> OmniOutput: + if not isinstance(request_token_spans, list) or len(request_token_spans) != len(runtime_info): + raise RuntimeError( + "MiniCPM-o 4.5 batched TTS requires one request_token_spans entry " + f"per request; got {request_token_spans!r} for {len(runtime_info)} requests" + ) + + request_ids: list[str] = [] + request_outputs: list[dict[str, Any]] = [] + row_chunk_flags = [True] * num_tokens + row_turn_end_flags = [False] * num_tokens + previous_end = 0 + for index, raw_info in enumerate(runtime_info): + if not isinstance(raw_info, dict): + raise RuntimeError(f"MiniCPM-o 4.5 TTS request metadata at index {index} is not a dict") + span = request_token_spans[index] + if not ( + isinstance(span, (list, tuple)) and len(span) == 2 and all(isinstance(value, int) for value in span) + ): + raise RuntimeError(f"Invalid MiniCPM-o 4.5 TTS token span at index {index}: {span!r}") + start, end = span + if start != previous_end or start < 0 or end < start or end > num_tokens: + raise RuntimeError( + f"Invalid MiniCPM-o 4.5 TTS token span at index {index}: " + f"expected start {previous_end} and 0 <= start <= end <= {num_tokens}, got {span!r}" + ) + previous_end = end + if start == end: + continue + mm_out, chunk_is_last, turn_ended = self._run_tts_request( + input_ids=self._slice_tts_request_rows(input_ids, start, end), + positions=self._slice_tts_request_rows(positions, start, end), + inputs_embeds=self._slice_tts_request_rows(inputs_embeds, start, end), + talker_info=raw_info, + device=device, + ) + row_chunk_flags[start:end] = [chunk_is_last] * (end - start) + row_turn_end_flags[start:end] = [turn_ended] * (end - start) + if mm_out is not None: + request_ids.append(get_stream_request_key(raw_info)) + request_outputs.append(mm_out) + # vLLM pads the flattened token tensor to a CUDA Graph capture size. + # The spans describe only scheduled request rows, so an uncovered tail + # is graph padding rather than a missing request. Contiguity checks + # above still reject gaps between real requests. + + self.talker._ar_last_chunk_flags = row_chunk_flags + self.talker._ar_turn_end_flags = row_turn_end_flags + dummy_hidden = torch.zeros(num_tokens, hidden_dim, device=device) + if not request_outputs: + return OmniOutput(text_hidden_states=dummy_hidden, multimodal_outputs=None) + + sparse_output: dict[str, Any] = { + "meta.req_id": request_ids, + "meta.sparse_audio": ["1"], + } + output_keys = set().union(*(output.keys() for output in request_outputs)) + for key in output_keys: + values: list[Any] = [] + for output in request_outputs: + value = output.get(key) + if isinstance(value, list) and len(value) == 1: + value = value[0] + if value is None: + dtype = torch.uint8 if key == "meta.llm_output_text_utf8" else torch.float32 + value = torch.empty(0, dtype=dtype, device=device) + values.append(value) + sparse_output[key] = values + return OmniOutput(text_hidden_states=dummy_hidden, multimodal_outputs=sparse_output) + def forward( self, input_ids: torch.Tensor, @@ -208,10 +658,8 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp Forward pass for MiniCPM-o Omni model. Workflow: - 1) Thinker (model_stage="llm"): Image / video / audio encoders + - 3D resampler + omni LLM → text + hidden states. - 2) Talker (model_stage="tts"): MiniCPMTTS + the in-process - Token2Wav vocoder → audio waveform (final pipeline output). + 1) LLM: multimodal thinker → hidden states and text tokens + 2) TTS: talker + Token2Wav → speech waveform """ if self.model_stage == "llm": # Normalize to batched inputs if caller provides 1D/2D unbatched tensors @@ -275,13 +723,62 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp text_hidden_states = text_hidden_states.squeeze(0) # Return hidden states with latent in multimodal_outputs for stage_input_processors + multimodal_outputs = {"latent": text_hidden_states} + runtime_info = kwargs.get("runtime_additional_information") + if runtime_info and isinstance(runtime_info, list) and len(runtime_info) > 0: + duplex_rows = [] + for req_info in runtime_info: + duplex_info = req_info.get("duplex") if isinstance(req_info, dict) else None + duplex_rows.append(duplex_info if isinstance(duplex_info, dict) else {}) + + prompt_rows = [] + for duplex_info in duplex_rows: + prompt_token_ids = duplex_info.get("duplex_prompt_token_ids") + # This is a complete per-handoff snapshot, not a generated + # tensor delta. Keep it as row-local metadata so output + # accumulation replaces the previous value instead of + # attempting to concatenate variable-length prompts. + prompt_rows.append(list(prompt_token_ids) if isinstance(prompt_token_ids, list) else None) + if any(row is not None for row in prompt_rows): + multimodal_outputs["duplex_prompt_token_ids"] = prompt_rows + + special_keys = { + key + for duplex_info in duplex_rows + for key, value in ( + duplex_info.get("special_token_ids", {}).items() + if isinstance(duplex_info.get("special_token_ids"), dict) + else () + ) + if isinstance(key, str) and isinstance(value, int) and value >= 0 + } + if special_keys: + multimodal_outputs["meta"] = { + key: [ + torch.tensor( + [int(value)], + dtype=torch.long, + device=text_hidden_states.device, + ) + if isinstance(value, int) and value >= 0 + else None + for duplex_info in duplex_rows + for value in [ + ( + duplex_info.get("special_token_ids", {}).get(key) + if isinstance(duplex_info.get("special_token_ids"), dict) + else None + ) + ] + ] + for key in sorted(special_keys) + } return OmniOutput( text_hidden_states=text_hidden_states, - multimodal_outputs={"latent": text_hidden_states}, + multimodal_outputs=multimodal_outputs, ) - # Talker stage: runs MiniCPMTTS + the in-process Token2Wav vocoder and - # emits the final audio waveform directly. + # Talker stage: runs TTS generation and its built-in Token2Wav vocoder. if self.model_stage == "tts": if input_ids is not None: num_tokens = input_ids.shape[0] @@ -302,34 +799,33 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp return OmniOutput(text_hidden_states=dummy_hidden, multimodal_outputs=None) runtime_info = kwargs.get("runtime_additional_information") - talker_info = {} - if runtime_info and isinstance(runtime_info, list) and len(runtime_info) > 0: - talker_info = runtime_info[0] if isinstance(runtime_info[0], dict) else {} - - with torch.inference_mode(): - talker_result = self.talker( + request_token_spans = kwargs.get("request_token_spans") + if isinstance(runtime_info, list) and (len(runtime_info) > 1 or request_token_spans is not None): + return self._run_batched_tts_requests( input_ids=input_ids, positions=positions, inputs_embeds=inputs_embeds, - additional_information=talker_info, + runtime_info=runtime_info, + request_token_spans=request_token_spans, + num_tokens=num_tokens, + hidden_dim=hidden_dim, + device=device, ) - + talker_info = {} + if runtime_info and isinstance(runtime_info, list) and len(runtime_info) > 0: + talker_info = runtime_info[0] if isinstance(runtime_info[0], dict) else {} dummy_hidden = torch.zeros(num_tokens, hidden_dim, device=device) - - # Talker returns a (mel_spec, waveform_or_None) tuple. MiniCPM-o - # 4.5 emits only the waveform (mel_spec is always None); keep the - # 2-slot unpack so a future mel-emitting variant can plug in - # without changing the wrapper. - mm_out: dict = {} - if isinstance(talker_result, tuple) and len(talker_result) == 2: - _, waveform = talker_result - if waveform is not None: - mm_out["model_outputs"] = [waveform] - - return OmniOutput( - text_hidden_states=dummy_hidden, - multimodal_outputs=mm_out if mm_out else None, + mm_out, _, _ = self._run_tts_request( + input_ids=input_ids, + positions=positions, + inputs_embeds=inputs_embeds, + talker_info=talker_info, + device=device, ) + if mm_out is not None: + return OmniOutput(text_hidden_states=dummy_hidden, multimodal_outputs=mm_out) + + return OmniOutput(text_hidden_states=dummy_hidden, multimodal_outputs=None) raise ValueError(f"Unsupported model stage: {self.model_stage}") @@ -341,26 +837,431 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp # Use model for logits computation return self.model.compute_logits(hidden_states) + def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: + request_sessions = getattr(self, "_minicpmo45_duplex_request_sessions", None) + helper = getattr(self, "_minicpmo45_duplex_data_plane_helper", None) + sessions = getattr(helper, "sessions", None) if helper is not None else None + if isinstance(request_sessions, dict): + for request_id in finished_req_ids: + session_key = request_sessions.pop(request_id, None) + if session_key is not None and isinstance(sessions, dict): + sessions.pop(session_key, None) + forced_segments = getattr(self, "_minicpmo45_force_listen_applied_segments", None) + if isinstance(forced_segments, set): + finished = set(finished_req_ids) + completed_segments = {segment for segment in forced_segments if segment[0] in finished} + forced_segments.difference_update(completed_segments) + if hasattr(self.model, "on_requests_finished"): + self.model.on_requests_finished(finished_req_ids) + def sample( self, logits: torch.Tensor, sampling_metadata: SamplingMetadata, ) -> SamplerOutput | None: - # Use model for sampling - return self.model.sample(logits, sampling_metadata) + native_duplex = self._sample_minicpmo45_native_duplex_stage0( + logits, + sampling_metadata, + duplex_rows=getattr(self, "_minicpmo45_active_duplex_rows", None), + ) + if native_duplex is not None: + return native_duplex + if self.model_stage == "tts": + return self.model.sample(logits, sampling_metadata) + return None + + def _sample_minicpmo45_native_duplex_stage0( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + *, + duplex_rows: list[int] | None = None, + ) -> SamplerOutput | None: + if self.model_stage != "llm" or logits.ndim != 2 or logits.shape[0] == 0: + return None + token_ids = self._minicpmo45_native_duplex_token_ids() + unit_id = token_ids.get("unit_token_id", -1) + if unit_id < 0: + return None + native_rows = self._minicpmo45_native_duplex_prompt_rows( + sampling_metadata, + unit_id, + logits.shape[0], + duplex_rows=duplex_rows, + ) + if not native_rows or len(native_rows) != logits.shape[0]: + return None + + sampled_ids: list[int] = [] + for row_idx in range(logits.shape[0]): + row_logits = logits[row_idx : row_idx + 1].clone() + sampled = self._sample_minicpmo45_native_duplex_row( + row_logits, + sampling_metadata, + row_idx=row_idx, + token_ids=token_ids, + ) + self._record_minicpmo45_duplex_terminator(row_idx, sampled, token_ids) + sampled_ids.append(sampled) + return SamplerOutput( + sampled_token_ids=torch.tensor(sampled_ids, device=logits.device, dtype=torch.int32).unsqueeze(-1), + logprobs_tensors=None, + ) + + def _sample_minicpmo45_native_duplex_row( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + *, + row_idx: int, + token_ids: dict[str, int], + ) -> int: + chunk_eos_id = token_ids.get("chunk_eos_token_id", -1) + generator = getattr(sampling_metadata, "generators", {}).get(row_idx) + output_token_ids = getattr(sampling_metadata, "output_token_ids", None) or [] + raw_recent_tokens = output_token_ids[row_idx] if row_idx < len(output_token_ids) else [] + recent_tokens = [int(token_id) for token_id in raw_recent_tokens if isinstance(token_id, int) and token_id >= 0] + if chunk_eos_id >= 0 and chunk_eos_id < logits.shape[-1]: + max_speak_tokens = int( + getattr( + self, + "max_new_speak_tokens_per_chunk", + MiniCPMO45DuplexPolicy.DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK, + ) + or MiniCPMO45DuplexPolicy.DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK + ) + request_max_tokens = self._minicpmo45_duplex_row_request_max_tokens(row_idx) + effective_max_speak_tokens = max_speak_tokens + if request_max_tokens is not None: + effective_max_speak_tokens = min(effective_max_speak_tokens, request_max_tokens) + if len(recent_tokens) >= max(1, effective_max_speak_tokens - 1): + return int(chunk_eos_id) + + # Match the released StreamDecoder: first sample the original + # distribution only to preserve the model's own chunk boundary. + # If it does not choose chunk_eos, mask that token before the + # normal text/listen/turn sampling pass below. + if getattr(sampling_metadata, "all_greedy", False): + boundary_sample = int(torch.argmax(logits, dim=-1).item()) + else: + boundary_probs = F.softmax(logits, dim=-1) + boundary_sample = int(torch.multinomial(boundary_probs, num_samples=1, generator=generator).item()) + if boundary_sample == chunk_eos_id: + return int(chunk_eos_id) + + logits = logits.clone() + forbidden = self._minicpmo45_native_forbidden_token_ids(token_ids) + if forbidden: + valid_forbidden = [token_id for token_id in forbidden if 0 <= token_id < logits.shape[-1]] + if valid_forbidden: + logits[:, valid_forbidden] = float("-inf") + + special_ids = self._minicpmo45_native_special_token_ids(token_ids) + state = self._minicpmo45_duplex_state_for_row(row_idx) + generated_text_tokens = getattr(state, "generated_text_tokens", None) + repetition_tokens = generated_text_tokens if generated_text_tokens else recent_tokens + repetition_penalty = 1.05 + if repetition_penalty != 1.0 and repetition_tokens: + history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE + for token_id in set(repetition_tokens[-history_size:]): + if token_id in special_ids or token_id < 0 or token_id >= logits.shape[-1]: + continue + logits[0, token_id] /= repetition_penalty + + temperature = float(self._sampling_metadata_value(sampling_metadata, "temperature", row_idx, 0.7)) + top_k = int(self._sampling_metadata_value(sampling_metadata, "top_k", row_idx, 100)) + top_p = float(self._sampling_metadata_value(sampling_metadata, "top_p", row_idx, 0.8)) + if getattr(sampling_metadata, "all_greedy", False) or temperature <= 0: + sampled = int(torch.argmax(logits, dim=-1).item()) + sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk( + sampled, + recent_tokens, + token_ids, + ) + return self._finalize_minicpmo45_native_duplex_sample(row_idx, sampled, token_ids) + + logits = logits / temperature + logits = self._top_k_top_p_filter(logits, top_k=top_k, top_p=top_p) + probs = F.softmax(logits, dim=-1) + sampled = int(torch.multinomial(probs, num_samples=1, generator=generator).item()) + sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk( + sampled, + recent_tokens, + token_ids, + ) + return self._finalize_minicpmo45_native_duplex_sample( + row_idx, + sampled, + token_ids, + ) + + def _maybe_cut_minicpmo45_native_duplex_text_chunk( + self, + sampled: int, + recent_tokens: list[int], + token_ids: dict[str, int], + ) -> int: + chunk_eos_id = token_ids.get("chunk_eos_token_id", -1) + if chunk_eos_id < 0: + return int(sampled) + special_ids = self._minicpmo45_native_special_token_ids(token_ids) + if sampled in special_ids: + return int(sampled) + max_chars = int( + getattr( + self, + "max_speak_chars_per_chunk", + MiniCPMO45DuplexPolicy.DEFAULT_MAX_SPEAK_CHARS_PER_CHUNK, + ) + or MiniCPMO45DuplexPolicy.DEFAULT_MAX_SPEAK_CHARS_PER_CHUNK + ) + if max_chars <= 0: + return int(sampled) + tokenizer = self._minicpmo45_tokenizer() + decode = getattr(tokenizer, "decode", None) + if not callable(decode): + return int(sampled) + candidate_tokens = self._minicpmo45_current_chunk_tokens(recent_tokens, token_ids) + candidate_tokens.append(int(sampled)) + try: + text = decode(candidate_tokens, skip_special_tokens=True) + except TypeError: + text = decode(candidate_tokens) + except Exception: + return int(sampled) + return int(chunk_eos_id) if isinstance(text, str) and len(text) >= max_chars else int(sampled) + + @staticmethod + def _minicpmo45_current_chunk_tokens( + tokens: list[int], + token_ids: dict[str, int], + ) -> list[int]: + boundaries = { + token_ids.get("listen_token_id", -1), + token_ids.get("chunk_eos_token_id", -1), + token_ids.get("chunk_tts_eos_token_id", -1), + token_ids.get("turn_eos_token_id", -1), + } + start = 0 + for idx, token_id in enumerate(tokens): + if token_id in boundaries: + start = idx + 1 + return list(tokens[start:]) + + def _minicpmo45_duplex_state_for_row(self, row_idx: int): + row_sessions = getattr(self, "_minicpmo45_duplex_row_sessions", None) + session_key = row_sessions.get(row_idx) if isinstance(row_sessions, dict) else None + if not session_key: + return None + helper = getattr(self, "_minicpmo45_duplex_data_plane_helper", None) + sessions = getattr(helper, "sessions", None) if helper is not None else None + return sessions.get(session_key) if isinstance(sessions, dict) else None + + def _minicpmo45_duplex_payload_for_row(self, row_idx: int) -> dict[str, Any] | None: + row_payloads = getattr(self, "_minicpmo45_duplex_row_payloads", None) + payload = row_payloads.get(row_idx) if isinstance(row_payloads, dict) else None + return payload if isinstance(payload, dict) else None + + def _minicpmo45_duplex_row_request_max_tokens(self, row_idx: int) -> int | None: + row_max_tokens = getattr(self, "_minicpmo45_duplex_row_max_tokens", None) + value = row_max_tokens.get(row_idx) if isinstance(row_max_tokens, dict) else None + try: + max_tokens = int(value) + except (TypeError, ValueError): + return None + return max_tokens if max_tokens > 0 else None + + def _finalize_minicpmo45_native_duplex_sample( + self, + row_idx: int, + sampled: int, + token_ids: dict[str, int], + ) -> int: + listen_id = token_ids.get("listen_token_id", -1) + tts_bos_id = token_ids.get("tts_bos_token_id", -1) + state = self._minicpmo45_duplex_state_for_row(row_idx) + payload = self._minicpmo45_duplex_payload_for_row(row_idx) + force_listen = isinstance(payload, dict) and payload.get("force_listen") is True + if ( + sampled == listen_id + and 0 <= tts_bos_id + and state is not None + and not getattr(state, "current_turn_ended", True) + and not force_listen + ): + return int(tts_bos_id) + return int(sampled) + + def _record_minicpmo45_duplex_terminator(self, row_idx: int, sampled: int, token_ids: dict[str, int]) -> None: + """Remember sampled unit state for the next append. + + The scheduler session update discards the final sampled token of a + segment before the next streaming update, but the official duplex + format feeds it (terminator + ) into the KV at every unit + boundary, and the model's listen/speak policy depends on seeing its own + past decisions. Non-terminators clear the turn-ended latch.""" + state = self._minicpmo45_duplex_state_for_row(row_idx) + if state is None: + return + payload = self._minicpmo45_duplex_payload_for_row(row_idx) + force_listen = isinstance(payload, dict) and payload.get("force_listen") is True + listen_id = token_ids.get("listen_token_id", -1) + tts_bos_id = token_ids.get("tts_bos_token_id", -1) + chunk_eos_id = token_ids.get("chunk_eos_token_id", -1) + chunk_tts_eos_id = token_ids.get("chunk_tts_eos_token_id", -1) + turn_eos_id = token_ids.get("turn_eos_token_id", -1) + terminators = {listen_id, chunk_eos_id, chunk_tts_eos_id, turn_eos_id} + if sampled in terminators: + state.pending_terminator_token = int(sampled) + state.last_terminator_token = int(sampled) + if sampled == turn_eos_id or (sampled == listen_id and force_listen): + state.current_turn_ended = True + with suppress(Exception): + state.pending_speech_response_open = False + return + special_ids = self._minicpmo45_native_special_token_ids(token_ids) + if sampled not in special_ids: + generated_text_tokens = getattr(state, "generated_text_tokens", None) + if not isinstance(generated_text_tokens, list): + generated_text_tokens = [] + state.generated_text_tokens = generated_text_tokens + generated_text_tokens.append(int(sampled)) + history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE + del generated_text_tokens[:-history_size] + if ( + sampled == tts_bos_id + and getattr(state, "current_turn_ended", True) + and getattr(state, "pending_speech_context", False) + ): + with suppress(Exception): + state.pending_speech_response_open = True + elif getattr(state, "pending_speech_response_open", False): + with suppress(Exception): + state.pending_speech_context = False + state.pending_speech_response_open = False + elif getattr(state, "current_turn_ended", True): + with suppress(Exception): + state.pending_speech_context = False + state.pending_terminator_token = None + state.last_terminator_token = None + state.current_turn_ended = False + + def _minicpmo45_tokenizer(self): + if hasattr(self, "_minicpmo45_tokenizer_cache"): + return self._minicpmo45_tokenizer_cache + tokenizer = None + get_tokenizer = getattr(getattr(self, "thinker", None), "get_tokenizer", None) + if callable(get_tokenizer): + tokenizer = get_tokenizer() + if tokenizer is None: + try: + from vllm.tokenizers import cached_tokenizer_from_config + + tokenizer = cached_tokenizer_from_config(self.vllm_config.model_config) + except Exception: + pass + self._minicpmo45_tokenizer_cache = tokenizer + return tokenizer + + def _minicpmo45_native_duplex_token_ids(self) -> dict[str, int]: + cached = getattr(self, "_minicpmo45_native_duplex_token_ids_cache", None) + if isinstance(cached, dict): + return cached + tokenizer = self._minicpmo45_tokenizer() + cached = MiniCPMO45DuplexPolicy.token_ids_from_tokenizer(tokenizer) + self._minicpmo45_native_duplex_token_ids_cache = cached + return cached + + def _minicpmo45_native_duplex_prompt_rows( + self, + sampling_metadata: SamplingMetadata, + unit_id: int, + batch_size: int, + *, + duplex_rows: list[int] | None = None, + ) -> list[int]: + if duplex_rows is not None: + rows: list[int] = [] + for row in duplex_rows: + try: + row_idx = int(row) + except (TypeError, ValueError): + continue + if 0 <= row_idx < batch_size: + rows.append(row_idx) + return rows + + prompt_token_ids = getattr(sampling_metadata, "prompt_token_ids", None) + if prompt_token_ids is None: + return [] + if prompt_token_ids.ndim == 1: + prompt_token_ids = prompt_token_ids.unsqueeze(0) + rows: list[int] = [] + for row_idx in range(min(batch_size, int(prompt_token_ids.shape[0]))): + row = prompt_token_ids[row_idx] + if torch.count_nonzero(row == unit_id).item() >= 2: + rows.append(row_idx) + return rows + + def _minicpmo45_native_forbidden_token_ids(self, token_ids: dict[str, int]) -> list[int]: + tokenizer = self._minicpmo45_tokenizer() + bad_token_ids = getattr(tokenizer, "bad_token_ids", []) if tokenizer is not None else [] + return MiniCPMO45DuplexPolicy.native_forbidden_token_ids(token_ids, bad_token_ids=bad_token_ids) + + def _minicpmo45_native_special_token_ids(self, token_ids: dict[str, int]) -> set[int]: + tokenizer = self._minicpmo45_tokenizer() + return MiniCPMO45DuplexPolicy.native_special_token_ids( + token_ids, + tokenizer_special_ids=getattr(tokenizer, "all_special_ids", []) if tokenizer is not None else [], + ) + + @staticmethod + def _sampling_metadata_value( + sampling_metadata: SamplingMetadata, + name: str, + row_idx: int, + default: float, + ) -> float: + value = getattr(sampling_metadata, name, None) + if value is None: + return default + if isinstance(value, torch.Tensor): + if value.numel() == 0: + return default + if value.ndim == 0: + return float(value.item()) + idx = min(row_idx, int(value.numel()) - 1) + return float(value.reshape(-1)[idx].item()) + try: + return float(value) + except (TypeError, ValueError): + return default + + @staticmethod + def _top_k_top_p_filter(logits: torch.Tensor, *, top_k: int, top_p: float) -> torch.Tensor: + if top_k > 0 and top_k < logits.shape[-1]: + kth = torch.topk(logits, top_k, dim=-1).values[..., -1, None] + logits = logits.masked_fill(logits < kth, float("-inf")) + if 0.0 < top_p < 1.0: + sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1) + cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) + sorted_remove = cumulative_probs > top_p + sorted_remove[..., 1:] = sorted_remove[..., :-1].clone() + sorted_remove[..., 0] = False + remove = torch.zeros_like(logits, dtype=torch.bool) + remove.scatter_(dim=-1, index=sorted_indices, src=sorted_remove) + logits = logits.masked_fill(remove, float("-inf")) + return logits def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights for the active stage of the omni model.""" + """Load weights for all components of the omni model.""" loaded_weights = set() thinker_weights = [] talker_weights = [] # MiniCPM-o checkpoint prefixes → stage mapping: # thinker: vpm, resampler, llm, apm, audio_projection_layer - # talker: tts (MiniCPMTTS); the Token2Wav vocoder weights load - # separately inside the talker module from the - # ``assets/token2wav`` subdirectory and do not appear - # in this iterator. + # talker: tts (ConditionalChatTTS) for k, v in weights: if k.startswith(("vpm.", "resampler.", "llm.", "apm.", "audio_projection_layer.")): thinker_weights.append((k, v)) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 34a4da1ac..e59742ad7 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -105,6 +105,11 @@ from vllm.multimodal.inputs import ( MultiModalKwargsItems, NestedTensors, ) + +try: + from vllm.multimodal.inputs import ModalityData, MultiModalDataDict +except ImportError: + from vllm.multimodal.parse import ModalityData, MultiModalDataDict from vllm.multimodal.parse import ( AudioItem, AudioProcessorItems, @@ -127,9 +132,15 @@ from vllm.multimodal.processing import ( ) from vllm.sequence import IntermediateTensors +try: + from vllm.transformers_utils.tokenizer import encode_tokens as _vllm_encode_tokens +except ImportError: + _vllm_encode_tokens = None -def encode_tokens(tokenizer, prompt: str) -> list[int]: - """Tokenize ``prompt`` without adding special tokens.""" + +def _encode_tokens(tokenizer: Any, prompt: str) -> list[int]: + if _vllm_encode_tokens is not None: + return _vllm_encode_tokens(tokenizer, prompt) return tokenizer.encode(prompt, add_special_tokens=False) @@ -2561,7 +2572,11 @@ class MiniCPMWhisperEncoder(WhisperEncoder): return_dict=None, past_key_values: EncoderDecoderCache | None = None, use_cache: bool | None = None, - ) -> BaseModelOutputWithPast | tuple: + use_extra_context: bool | None = False, + prefix_extra_frames: int | None = 1, + suffix_extra_frames: int | None = 1, + cnn_min_length: int | None = None, + ) -> BaseModelOutputWithPast | tuple[Any, ...]: r""" Forward pass of the Whisper encoder. @@ -2672,9 +2687,37 @@ class MiniCPMWhisperEncoder(WhisperEncoder): # Ignore copy input_features = input_features.to(dtype=self.conv1.weight.dtype, device=self.conv1.weight.device) + original_length = input_features.shape[2] + padded_for_cnn = False + if cnn_min_length is not None and original_length < cnn_min_length: + padded_features = torch.zeros( + input_features.shape[0], + input_features.shape[1], + cnn_min_length, + dtype=input_features.dtype, + device=input_features.device, + ) + padded_features[:, :, :original_length] = input_features + input_features = padded_features + padded_for_cnn = True + inputs_embeds = nn.functional.gelu(self.conv1(input_features)) inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds)) + if padded_for_cnn: + actual_cnn_output_length = (original_length + 1) // 2 + inputs_embeds = inputs_embeds[:, :, :actual_cnn_output_length] + + if use_extra_context: + prefix_extra_frames = int(prefix_extra_frames or 0) + suffix_extra_frames = int(suffix_extra_frames or 0) + prefix_to_remove = (prefix_extra_frames + 1) // 2 if prefix_extra_frames > 0 else 0 + suffix_to_remove = (suffix_extra_frames + 1) // 2 if suffix_extra_frames > 0 else 0 + if prefix_to_remove > 0: + inputs_embeds = inputs_embeds[:, :, prefix_to_remove:] + if 0 < suffix_to_remove < inputs_embeds.shape[2]: + inputs_embeds = inputs_embeds[:, :, :-suffix_to_remove] + inputs_embeds = inputs_embeds.permute(0, 2, 1) embed_pos = self.embed_positions.weight @@ -3328,7 +3371,7 @@ class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45Omn tokenization_kwargs=tokenization_kwargs, ) tokenizer = self.info.get_tokenizer() - prompt_ids = encode_tokens(tokenizer, prompt) + prompt_ids = _encode_tokens(tokenizer, prompt) else: prompt_ids = self._apply_hf_processor_tokens_only(prompt) if use_tts: @@ -3817,8 +3860,18 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S else: text_config = Qwen2Config.from_dict(config_dict) llm_arch = "Qwen2ForCausalLM" + from vllm_omni.experimental.fullduplex.minicpmo45.compat import ( + patch_minicpmo_remote_config, + ) + + patch_minicpmo_remote_config(text_config) + text_vllm_config = ( + vllm_config.with_hf_config(text_config) + if callable(getattr(vllm_config, "with_hf_config", None)) + else vllm_config + ) self.llm = init_vllm_registered_model( - vllm_config=vllm_config, + vllm_config=text_vllm_config, prefix=maybe_prefix(prefix, "llm"), hf_config=text_config, architectures=[llm_arch], @@ -3846,11 +3899,13 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S self.audio_avg_pooler = nn.AvgPool1d(config.audio_pool_step, stride=config.audio_pool_step) self.audio_projection_layer = MultiModalProjector(in_dim=audio_output_dim, out_dim=embed_dim) self.audio_encoder_layer = -1 + self.audio_past_key_values = None else: self.apm = None self.audio_avg_pooler = None self.audio_projection_layer = None self.audio_encoder_layer = None + self.audio_past_key_values = None self.mm_token_ids = set[int]() self.make_empty_intermediate_tensors = self.llm.make_empty_intermediate_tensors @@ -4279,6 +4334,100 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S return final_audio_embeds + def get_audio_embedding_streaming( + self, + data: MiniCPMOAudioFeatureInputs, + use_extra_context: bool = False, + prefix_extra_frames: int = 1, + suffix_extra_frames: int = 1, + cnn_min_length: int | None = None, + ) -> list[list[torch.Tensor]]: + if self.apm is None or self.audio_projection_layer is None or self.audio_avg_pooler is None: + return [] + + wavforms = data["audio_features"] + audio_feature_lens_raw = data["audio_feature_lens"] + if isinstance(wavforms, list): + if not wavforms: + return [] + batch_size = len(wavforms) + channels = wavforms[0].shape[-2] + max_len = max(item.shape[-1] for item in wavforms) + wavforms_tensor = torch.zeros( + (batch_size, channels, max_len), + dtype=wavforms[0].dtype, + device=wavforms[0].device, + ) + for idx, item in enumerate(wavforms): + wavforms_tensor[idx, ..., : item.shape[-1]] = item + wavforms = wavforms_tensor + if isinstance(audio_feature_lens_raw, torch.Tensor): + audio_feature_lens_raw = audio_feature_lens_raw.unbind(0) + if wavforms.shape[0] == 0: + return [] + audio_feature_lens = torch.hstack(audio_feature_lens_raw) + batch_size, _, max_mel_seq_len = wavforms.shape + assert batch_size == 1 + + current_seq_len = (max_mel_seq_len - 1) // 2 + 1 + if use_extra_context: + prefix_to_remove = (int(prefix_extra_frames) + 1) // 2 if prefix_extra_frames > 0 else 0 + suffix_to_remove = (int(suffix_extra_frames) + 1) // 2 if suffix_extra_frames > 0 else 0 + current_seq_len = current_seq_len - prefix_to_remove - suffix_to_remove + if current_seq_len <= 0: + return [] + + if self.audio_past_key_values is not None: + cache_length = self.audio_past_key_values[0][0].shape[2] + apm_max_len = self.apm.embed_positions.weight.shape[0] + if cache_length + current_seq_len >= apm_max_len: + logger.warning( + "audio_past_key_values length %s exceeds %s, reset.", + cache_length + current_seq_len, + apm_max_len, + ) + self.audio_past_key_values = None + + past_len = 0 + if self.audio_past_key_values is not None: + past_len = self.audio_past_key_values[0][0].shape[2] + total_seq_len = past_len + current_seq_len + audio_attention_mask = torch.zeros( + (batch_size, 1, current_seq_len, total_seq_len), + dtype=self.apm.conv1.weight.dtype, + device=wavforms.device, + ) + + audio_outputs = self.apm( + wavforms, + past_key_values=self.audio_past_key_values, + use_cache=True, + output_hidden_states=True, + attention_mask=audio_attention_mask, + use_extra_context=use_extra_context, + prefix_extra_frames=prefix_extra_frames, + suffix_extra_frames=suffix_extra_frames, + cnn_min_length=cnn_min_length, + ) + audio_states = audio_outputs.hidden_states[self.audio_encoder_layer] + self.audio_past_key_values = audio_outputs.past_key_values + + audio_embeds = self.audio_projection_layer(audio_states) + audio_embeds = audio_embeds.transpose(1, 2) + audio_embeds = self.audio_avg_pooler(audio_embeds) + audio_embeds = audio_embeds.transpose(1, 2) + + _, feature_lens_after_pooling = self._get_feat_extract_output_lengths(audio_feature_lens) + final_audio_embeds: list[list[torch.Tensor]] = [] + idx = 0 + for i in range(len(audio_feature_lens_raw)): + target_audio_embeds = [] + for _ in range(len(audio_feature_lens_raw[i])): + target_audio_embeds.append(audio_embeds[idx, : feature_lens_after_pooling[idx], :]) + idx += 1 + final_audio_embeds.append(target_audio_embeds) + return final_audio_embeds + def _process_audio_input( self, audio_input: MiniCPMOAudioInputs, diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index f109eb737..73570ef74 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -12,11 +12,19 @@ Pipeline: 5. Run Token2wav(tokens) -> waveform bytes -> numpy array """ +import hashlib import io import logging import os import sys +import tempfile +import threading +import time +from collections import OrderedDict from collections.abc import Iterable +from dataclasses import dataclass +from importlib import import_module +from typing import Any import numpy as np import soundfile as sf @@ -25,14 +33,17 @@ import torch.nn as nn import torch.nn.functional as F from vllm.config import VllmConfig from vllm.model_executor.models.interfaces import SupportsPP +from vllm.v1.outputs import SamplerOutput +from vllm_omni.experimental.fullduplex.engine.intermediate import get_stream_request_key, get_tts_handoff from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) from vllm_omni.platforms import current_omni_platform -# Preserve the established external vocoder on CUDA. Ascend uses the in-tree -# adapter because ``stepaudio2-minicpmo`` hard-codes CUDA device placement. +# The external vocoder hard-codes CUDA placement. Ascend uses the in-tree +# adapter, with the external package retained as a fallback for compatibility. +_stepaudio2_import_error: ImportError | None = None if current_omni_platform.is_npu(): try: from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_token2wav import ( @@ -40,28 +51,54 @@ if current_omni_platform.is_npu(): ) _token2wav_backend = "step_audio2_core" - except ImportError: + except ImportError as e: + _stepaudio2_import_error = e try: from stepaudio2 import Token2wav as _Token2wav _token2wav_backend = "stepaudio2_pkg" - except ImportError: + except ImportError as fallback_error: _Token2wav = None _token2wav_backend = None + _stepaudio2_import_error = fallback_error else: try: from stepaudio2 import Token2wav as _Token2wav _token2wav_backend = "stepaudio2_pkg" - except ImportError: + except ImportError as e: _Token2wav = None _token2wav_backend = None + _stepaudio2_import_error = e _stepaudio2_available = _Token2wav is not None logger = logging.getLogger(__name__) +@dataclass(slots=True) +class MiniCPMO45TTSRuntimeConfig: + """Internal MiniCPM-o 4.5 Talker runtime defaults. + + These values are deliberately not process environment knobs. If a value + needs to become user configurable, route it through stage/model config with + an explicit API contract instead of adding another ad hoc env var. + """ + + token2wav_n_timesteps: int = 10 + tts_dtype: torch.dtype = torch.float32 + token2wav_autocast_dtype: torch.dtype | None = None + use_direct_token2wav: bool = True + ref_audio_file_cache_size: int = 16 + max_token_ratio: int = 32 + min_max_new_tokens: int = 256 + hard_max_new_tokens: int = 16384 + min_new_tokens: int = 50 + streaming_generator_chunk: int = 25 + streaming_vocoder_threshold: int = 2500 + streaming_vocoder_chunk: int = 50 + + def _install_torchaudio_soundfile_shim() -> None: """Monkey-patch torchaudio.load to use soundfile instead of the default torchcodec backend, which requires libtorchcodec/ffmpeg shared libs that @@ -94,9 +131,118 @@ def _install_torchaudio_soundfile_shim() -> None: _install_torchaudio_soundfile_shim() +class _TalkerTurnState: + """Per-turn talker continuity for native duplex. + + The official duplex implementation runs ONE TTS stream per spoken turn: + each 1s unit appends its condition to the same KV (carried + past_key_values + text_start_pos) and the token2wav stream caches and + token buffer persist across units, reset only at turn end. Synthesizing + units as independent utterances resets prosody every second and garbles + the reply. + """ + + __slots__ = ( + "past_key_values", + "text_start_pos", + "token2wav_buffer", + "prompt_wav_path", + "temp_prompt_wav_path", + "epoch", + "turn_id", + "pending_text", + "stream_cache", + "hift_cache_dict", + "vocoder_initialized", + ) + + def __init__( + self, + prompt_wav_path, + temp_prompt_wav_path, + *, + epoch: int | None = None, + turn_id: int | None = None, + ): + self.past_key_values = None + self.text_start_pos = 0 + # Official seeds each turn's vocoder buffer with three silence + # tokens so the first synthesized window does not directly abut the + # reference-audio prompt cache (audible ref-voice bleed otherwise). + self.token2wav_buffer: list[int] = [_T2W_SILENCE_TOKEN] * 3 + self.prompt_wav_path = prompt_wav_path + self.temp_prompt_wav_path = temp_prompt_wav_path + self.epoch = epoch + self.turn_id = turn_id + self.pending_text = "" + self.stream_cache = None + self.hift_cache_dict: dict[str, Any] = {} + self.vocoder_initialized = False + + +def _queue_native_duplex_segment_text(state: _TalkerTurnState, text: object) -> None: + if isinstance(text, str) and text: + state.pending_text += text + + +def _drain_native_duplex_emitted_text(state: _TalkerTurnState, *, has_audio: bool) -> str: + if not has_audio: + return "" + text = state.pending_text + state.pending_text = "" + return text + + +_T2W_SILENCE_TOKEN = 4218 +_NATIVE_DUPLEX_UNIT_AUDIO_SAMPLES = 24_000 + + +def _native_duplex_unit_waveform( + waveforms: Iterable[torch.Tensor], + *, + turn_end: bool, + target_samples: int = _NATIVE_DUPLEX_UNIT_AUDIO_SAMPLES, +) -> torch.Tensor | None: + pieces = [torch.as_tensor(waveform).reshape(-1).cpu().contiguous() for waveform in waveforms] + pieces = [piece for piece in pieces if piece.numel() > 0] + if not pieces: + return None + waveform = pieces[0] if len(pieces) == 1 else torch.cat(pieces, dim=0) + if not turn_end and waveform.numel() < target_samples: + waveform = F.pad(waveform, (target_samples - waveform.numel(), 0)) + return waveform + + +def _soundfile_patched_save(orig_save): + def _patched_save(uri, src, sample_rate, **kw): + kw.pop("backend", None) + if hasattr(uri, "write"): + sf.write(uri, src.cpu().numpy().T, sample_rate, format="WAV") + return + return orig_save(uri, src, sample_rate, backend="soundfile", **kw) + + return _patched_save + + +def _torch_clone_recursive(obj): + if isinstance(obj, torch.Tensor): + return obj.clone() + if isinstance(obj, dict): + return {k: _torch_clone_recursive(v) for k, v in obj.items()} + if isinstance(obj, list): + return [_torch_clone_recursive(v) for v in obj] + if isinstance(obj, tuple): + return tuple(_torch_clone_recursive(v) for v in obj) + return obj + + class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): """MiniCPM-o 4.5 Talker: MiniCPMTTS + Token2wav in a single forward pass.""" + # llm2tts hands the FULL accumulated condition per handoff (the runner's + # resume-prefill path REPLACES the streaming buffer, so runner-side + # accumulation is lossy); the per-turn state consumes it by cursor. + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import MiniCPMOConfig @@ -104,40 +250,114 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): config: MiniCPMOConfig = vllm_config.model_config.hf_config self.config = config self.vllm_config = vllm_config + self._runtime_config = MiniCPMO45TTSRuntimeConfig() self.tts = None self.audio_tokenizer = None self._assets_loaded = False - self._model_path = None + self._model_path: str | None = None + self._stream_gens: dict[str, Any] = {} + self._talker_turn_states: dict[str, _TalkerTurnState] = {} + # Consumed-cursor into the accumulated handoff condition for the + # currently open spoken turn. + self._talker_consumed_tokens: dict[str, int] = {} + self._talker_request_keys: dict[str, str] = {} + self._t2w_base_caches: dict[str, tuple[Any, Any]] = {} + self._token2wav_state_lock = threading.RLock() + self._ar_last_chunk_flags: list[bool] = [True] + self._ar_turn_end_flags: list[bool] = [False] + self._ar_last_emitted_text = "" + self._text_tokenizer = None tts_config = getattr(config, "tts_config", None) if tts_config is not None: self._tts_config = tts_config - self._tts_bos_id = getattr(tts_config, "audio_bos_token_id", 151687) - self._text_eos_id = getattr(tts_config, "text_eos_token_id", 151692) + self._tts_bos_id = self._config_token_id(tts_config, "audio_bos_token_id") + self._text_eos_id = self._config_token_id(tts_config, "text_eos_token_id") self._num_audio_tokens = getattr(tts_config, "num_audio_tokens", 6562) self._hidden_size = getattr(tts_config, "hidden_size", 768) self._normalize = getattr(tts_config, "normalize_projected_hidden", True) else: self._tts_config = None + def _tts_runtime_config(self) -> MiniCPMO45TTSRuntimeConfig: + cfg = getattr(self, "_runtime_config", None) + if cfg is None: + cfg = MiniCPMO45TTSRuntimeConfig() + self._runtime_config = cfg + return cfg + + @staticmethod + def _config_token_id(config: Any, attr: str) -> int: + value = getattr(config, attr, None) + if isinstance(value, (list, tuple)): + value = value[0] if value else None + if value is None: + raise ValueError(f"MiniCPM-o 4.5 TTS config missing required {attr}") + return int(value) + + def _get_text_tokenizer(self) -> Any: + tokenizer = getattr(self, "_text_tokenizer", None) + if tokenizer is not None: + return tokenizer + from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config + + tokenizer = cached_tokenizer_from_config(self.vllm_config.model_config) + self._text_tokenizer = tokenizer + return tokenizer + + def _tokenizer_token_id(self, token: str) -> int | None: + tokenizer = self._get_text_tokenizer() + unk_token_id = getattr(tokenizer, "unk_token_id", None) + convert = getattr(tokenizer, "convert_tokens_to_ids", None) + if callable(convert): + value = convert(token) + if isinstance(value, list): + value = value[0] if len(value) == 1 else None + try: + candidate = int(value) + except (TypeError, ValueError): + candidate = None + if candidate is not None and candidate >= 0 and candidate != unk_token_id: + return candidate + encode = getattr(tokenizer, "encode", None) + if callable(encode): + ids = list(encode(token, add_special_tokens=False)) + if len(ids) == 1: + candidate = int(ids[0]) + if candidate >= 0 and candidate != unk_token_id: + return candidate + return None + + def _scheduler_eos_token_id(self) -> int: + eos_raw = getattr(self.config, "eos_token_id", None) + if isinstance(eos_raw, (list, tuple)): + eos_raw = eos_raw[0] if eos_raw else None + if eos_raw is not None: + return int(eos_raw) + eos_id = self._tokenizer_token_id("<|im_end|>") + if eos_id is None: + raise ValueError( + "MiniCPM-o 4.5 TTS scheduler EOS requires config.eos_token_id or tokenizer-defined <|im_end|>" + ) + return eos_id + def _lazy_init_tts(self): if self._assets_loaded or self._tts_config is None: return + self._assets_loaded = True try: model_path = download_weights_from_hf_specific(self.vllm_config.model_config.model, None, ["*"]) self._model_path = model_path - if model_path not in sys.path: sys.path.insert(0, model_path) from transformers import AutoImageProcessor from transformers.dynamic_module_utils import get_class_from_dynamic_module - # openbmb/MiniCPM-o-4_5/processing_minicpmo.py registers via a - # string: AutoImageProcessor.register("MiniCPMVImageProcessor", ...), - # which crashes on transformers>=5 (register reads key.__module__). - # Loading MiniCPMTTS imports that module, so no-op the string form - # (unused by the standalone talker) while it runs, then restore. + # The remote processing module registers an image processor by + # string, which transformers>=5 rejects. The standalone talker does + # not use that registration, so ignore only the string form while + # importing MiniCPMTTS and restore the global method immediately. original_register = AutoImageProcessor.register AutoImageProcessor.register = ( # type: ignore[method-assign] lambda key, *a, **k: None if isinstance(key, str) else original_register(key, *a, **k) @@ -147,28 +367,30 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): finally: AutoImageProcessor.register = original_register # type: ignore[method-assign] - # MiniCPMTTS.__init__ reads `config.top_p / top_k / repetition_penalty` - # directly (modeling_minicpmo.py L4112-4114), but the model repo's - # config.json `tts_config` block does not declare these fields and - # PretrainedConfig in recent transformers no longer surfaces - # generation-style params on `self.config`. Inject the defaults the - # upstream code itself ships with (modeling_minicpmo.py L2212-2214, - # L3132-3133) so attribute access does not raise. - for _attr, _default in (("top_p", 0.8), ("top_k", 100), ("repetition_penalty", 1.02)): - if not hasattr(self._tts_config, _attr): - setattr(self._tts_config, _attr, _default) - - # The copied Hugging Face flash_attention_2 setting is not valid - # for this standalone MiniCPMTTS path. Use PyTorch SDPA on every - # backend until a dedicated flash-attention implementation exists. - self._tts_config.attn_implementation = "sdpa" - prev_dtype = torch.get_default_dtype() torch.set_default_dtype(torch.float32) try: + for name, default in ( + ("top_p", 0.85), + ("top_k", 25), + ("repetition_penalty", 1.05), + ("temperature", 0.8), + ): + if not hasattr(self._tts_config, name): + setattr(self._tts_config, name, default) + self._tts_config.attn_implementation = "sdpa" self.tts_obj = MiniCPMTTS(config=self._tts_config, audio_tokenizer=None) finally: torch.set_default_dtype(prev_dtype) + tts_module = import_module(self.tts_obj.__class__.__module__) + + def get_tts_module_attr(name: str): + return self.tts_obj.generate.__globals__.get(name) or getattr(tts_module, name, None) + + self._tts_sampling_params_cls = get_tts_module_attr("TTSSamplingParams") + self._tts_gen_logits = get_tts_module_attr("gen_logits") + self._tts_parametrize = get_tts_module_attr("P") + self._tts_streaming_generator_cls = get_tts_module_attr("TTSStreamingGenerator") self.emb_text = self.tts_obj.emb_text self.projector_semantic = self.tts_obj.projector_semantic @@ -176,119 +398,914 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): if os.path.isdir(token2wav_dir): if not _stepaudio2_available: raise ImportError( - "MiniCPM-o 4.5 token2wav stage requires the `stepaudio2` Python " - "module (a MiniCPM-o-flavored Token2wav vocoder, NOT the upstream " - "stepfun-ai/Step-Audio2 — the upstream signature does not accept " - "n_timesteps and will fail at __init__). Install via:\n" - " pip install 'vllm-omni[minicpmo]' # recommended, declared as PR extra\n" - "Equivalent direct installs of the same `from stepaudio2 import Token2wav`\n" - "entry point used by openbmb/MiniCPM-o-4_5/modeling_minicpmo.py:\n" - " pip install stepaudio2-minicpmo # bare token2wav package\n" - " pip install 'minicpmo-utils[all]' # MiniCPM-o umbrella (also brings image/video deps)" - ) + "MiniCPM-o 4.5 token2wav stage requires the stepaudio2 package, " + "and all of its runtime dependencies." + ) from _stepaudio2_import_error + self._token2wav_n_timesteps = self._tts_runtime_config().token2wav_n_timesteps prev_dtype2 = torch.get_default_dtype() torch.set_default_dtype(torch.float32) try: - # NB: this must be the MiniCPM-o-flavored Token2wav from - # the `stepaudio2-minicpmo` PyPI package (or the - # `minicpmo-utils[all]` umbrella), not the upstream - # `stepfun-ai/Step-Audio2` repo. The MiniCPM-o variant's - # __init__ accepts n_timesteps; the upstream signature is - # (model_path, float16=False) and will raise - # TypeError on n_timesteps. See ImportError message below - # for installation guidance. - self.audio_tokenizer = _Token2wav(token2wav_dir, float16=False, n_timesteps=10) + self.audio_tokenizer = _Token2wav( + token2wav_dir, + float16=False, + n_timesteps=self._token2wav_n_timesteps, + ) finally: torch.set_default_dtype(prev_dtype2) self.tts_obj.audio_tokenizer = self.audio_tokenizer logger.info( - "Loaded Token2wav from %s (backend=%s)", + "Loaded Token2wav from %s (backend=%s, n_timesteps=%d)", token2wav_dir, _token2wav_backend, + self._token2wav_n_timesteps, ) - # Only mark init as complete after every step succeeds, so a - # partial failure leaves the next call free to retry the full - # init instead of short-circuiting back to a silent empty path. - self._assets_loaded = True except ImportError: # Surface missing dependencies directly so users can act on them # instead of getting a silent None waveform downstream. raise - except Exception: - # Re-raise non-import init failures (bad token2wav assets, missing - # weights, OOM during Token2wav construction, etc.) so the server - # fails loudly at startup / first request instead of returning - # silent empty audio for every subsequent request. - logger.error("Failed to init 4.5 TTS", exc_info=True) - raise + except Exception as e: + logger.error("Failed to init 4.5 TTS: %s", e, exc_info=True) - def generate_speech( - self, - tts_token_ids: torch.Tensor, - tts_hidden_states: torch.Tensor, - ) -> np.ndarray | None: - """Run full 4.5 TTS pipeline using original MiniCPMTTS.generate.""" - self._lazy_init_tts() - if not hasattr(self, "tts_obj") or self.tts_obj is None: - logger.warning("generate_speech: tts_obj not initialized") + def _build_tts_sampling_params(self): + params_cls = getattr(self, "_tts_sampling_params_cls", None) + if params_cls is None or not hasattr(self, "tts_obj"): return None tts = self.tts_obj - device = tts.emb_text.weight.device - # MiniCPMTTS AR backbone uses FlashAttention (fp16/bf16 only). The - # submodule is constructed under float32 default dtype during lazy init, - # so pin the condition embeddings to bfloat16 explicitly rather than - # inheriting the (float32) parameter dtype — a float32 condition breaks - # the CUDA FA2 path and wastes memory on the NPU sdpa path. - ar_dtype = torch.bfloat16 + top_p = getattr(tts, "top_p", getattr(tts.config, "top_p", 0.85)) + top_k = getattr(tts, "top_k", getattr(tts.config, "top_k", 25)) + repetition_penalty = getattr( + tts, + "repetition_penalty", + getattr(tts.config, "repetition_penalty", 1.05), + ) + temperature = getattr(tts.config, "temperature", 0.8) + + return params_cls( + top_p=None if top_p is not None and top_p >= 1.0 else top_p, + top_k=None if top_k is not None and top_k <= 0 else top_k, + repetition_penalty=repetition_penalty, + temperature=temperature, + ) + + def _target_tts_dtype(self) -> torch.dtype: + return self._tts_runtime_config().tts_dtype + + def _token2wav_autocast_context(self): + dtype = self._tts_runtime_config().token2wav_autocast_dtype + device_type = str(current_omni_platform.device_type) + if dtype is None: + return ( + current_omni_platform.create_autocast_context( + device_type=device_type, + dtype=torch.float32, + enabled=False, + ), + "off", + ) + if dtype is torch.bfloat16: + return ( + current_omni_platform.create_autocast_context( + device_type=device_type, + dtype=torch.bfloat16, + enabled=True, + ), + "bf16", + ) + if dtype is torch.float16: + return ( + current_omni_platform.create_autocast_context( + device_type=device_type, + dtype=torch.float16, + enabled=True, + ), + "fp16", + ) + raise ValueError("MiniCPM-o 4.5 token2wav autocast only supports None, bfloat16, or float16") + + def _should_use_direct_token2wav(self) -> bool: + return self._tts_runtime_config().use_direct_token2wav + + def _should_stream_output(self, info: dict[str, Any] | None = None) -> bool: + if isinstance(info, dict): + for key in ("stream_output", "native_duplex"): + value = info.get(key) + if isinstance(value, bool): + return value + return False + + def _token2wav_prompt_cache_key(self, prompt_wav: str | None) -> str | None: + return os.path.abspath(prompt_wav) if prompt_wav else None + + def _reset_token2wav_cache_if_needed(self, prompt_wav: str | None) -> None: + token2wav = self.audio_tokenizer + if token2wav is None: + return + cache_key = self._token2wav_prompt_cache_key(prompt_wav) + if getattr(self, "_token2wav_prompt_cache_id", None) != cache_key: + token2wav.cache = None + self._token2wav_prompt_cache_id = cache_key + + def _normalize_ref_audio_tensor(self, ref_audio) -> np.ndarray | None: + if ref_audio is None: + return None + if isinstance(ref_audio, torch.Tensor): + waveform = ref_audio.detach().float().cpu().numpy() + else: + waveform = np.asarray(ref_audio, dtype=np.float32) + if waveform.ndim > 1: + if waveform.shape[0] <= 2 and waveform.shape[-1] > waveform.shape[0]: + waveform = waveform.mean(axis=0) + else: + waveform = waveform.mean(axis=-1) + waveform = np.asarray(waveform, dtype=np.float32).reshape(-1) + return waveform if waveform.size else None + + def _write_ref_audio_prompt_wav(self, ref_audio, ref_audio_sr: int | None) -> str | None: + waveform = self._normalize_ref_audio_tensor(ref_audio) + if waveform is None: + return None + sample_rate = int(ref_audio_sr or 24000) + cache_size = self._tts_runtime_config().ref_audio_file_cache_size + if cache_size > 0: + digest = hashlib.sha256() + digest.update(str(sample_rate).encode("ascii")) + digest.update(waveform.tobytes()) + cache_key = digest.hexdigest() + cache = getattr(self, "_ref_audio_prompt_files", None) + if cache is None: + cache = OrderedDict() + self._ref_audio_prompt_files = cache + cached_path = cache.get(cache_key) + if cached_path and os.path.exists(cached_path): + cache.move_to_end(cache_key) + return cached_path + + tmp_path = os.path.join(tempfile.gettempdir(), f"minicpmo45_ref_{cache_key[:24]}_{sample_rate}.wav") + if not os.path.exists(tmp_path): + sf.write(tmp_path, waveform, sample_rate, format="WAV") + cache[cache_key] = tmp_path + cache.move_to_end(cache_key) + while len(cache) > cache_size: + _, evicted_path = cache.popitem(last=False) + with self._token2wav_lock(): + self._t2w_base_caches.pop(evicted_path, None) + try: + os.unlink(evicted_path) + except OSError: + pass + return tmp_path + + tmp = tempfile.NamedTemporaryFile(prefix="minicpmo45_ref_", suffix=".wav", delete=False) + tmp_path = tmp.name + tmp.close() + sf.write(tmp_path, waveform, sample_rate, format="WAV") + return tmp_path + + def _is_cached_ref_audio_prompt_wav(self, prompt_wav: str | None) -> bool: + cache = getattr(self, "_ref_audio_prompt_files", None) + return bool(prompt_wav and cache and prompt_wav in cache.values()) + + def _run_token2wav_direct( + self, + generated_speech_tokens: list[int] | torch.Tensor, + prompt_wav: str | None, + ) -> tuple[torch.Tensor, int]: + """Run Token2wav without WAV encode/decode round-tripping. + + ``stepaudio2.Token2wav.__call__`` renders the GPU waveform to a WAV + BytesIO object and the vLLM-Omni adapter immediately decodes it back to + a numpy waveform. The engine already expects a float waveform tensor, + so keep the same flow/HIFT computation and return the waveform directly. + This also makes the configured ``n_timesteps`` apply to one-shot + inference; upstream currently hard-codes ``10`` in ``__call__``. + """ + token2wav = self.audio_tokenizer + if token2wav is None: + raise RuntimeError("Token2wav is not initialized") + required_attrs = ("_prepare_prompt", "flow", "hift") + if any(not hasattr(token2wav, attr) for attr in required_attrs): + raise RuntimeError("Token2wav direct path is incompatible with the installed stepaudio2 package") + + self._reset_token2wav_cache_if_needed(prompt_wav) + if token2wav.cache is None: + token2wav.cache = token2wav._prepare_prompt(prompt_wav) + prompt_speech_tokens, prompt_speech_tokens_lens, spk_emb, prompt_mels, prompt_mels_lens = token2wav.cache + + device = prompt_speech_tokens.device + if isinstance(generated_speech_tokens, torch.Tensor): + generated = generated_speech_tokens + if generated.ndim == 1: + generated = generated.unsqueeze(0) + elif generated.ndim == 3 and generated.shape[-1] == 1: + generated = generated.squeeze(-1) + generated = generated.to(device=device, dtype=torch.int32) + else: + generated = torch.tensor([generated_speech_tokens], dtype=torch.int32, device=device) + generated_lens = torch.tensor([generated.shape[1]], dtype=torch.int32, device=device) + mel = token2wav.flow.inference( + generated, + generated_lens, + prompt_speech_tokens, + prompt_speech_tokens_lens, + prompt_mels, + prompt_mels_lens, + spk_emb, + self._token2wav_n_timesteps, + ) + wav, _ = token2wav.hift(speech_feat=mel) + waveform = wav.squeeze(0).detach().float().reshape(-1).cpu().contiguous() + return waveform, 24000 + + def _build_tts_condition_embeds( + self, + tts_token_ids: torch.Tensor, + tts_hidden_states: torch.Tensor, + ) -> torch.Tensor: + tts = self.tts_obj + device = tts.emb_text.weight.device + dtype = tts.emb_text.weight.dtype llm_embeds = tts.emb_text(tts_token_ids.to(device)) - hidden_embeds = tts.projector_semantic(tts_hidden_states.to(device=device, dtype=ar_dtype)) + hidden_embeds = tts.projector_semantic(tts_hidden_states.to(device=device, dtype=dtype)) if getattr(tts.config, "normalize_projected_hidden", False): hidden_embeds = F.normalize(hidden_embeds, p=2, dim=-1) - tts_embeds = (llm_embeds + hidden_embeds).to(dtype=ar_dtype) + return llm_embeds + hidden_embeds - text_eos = tts.emb_text(torch.tensor([tts.config.text_eos_token_id], device=device, dtype=torch.long)) - audio_bos = tts.emb_text(torch.tensor([tts.audio_bos_token_id], device=device, dtype=torch.long)) - spk_embeds = torch.zeros(0, tts.config.hidden_size, device=device, dtype=ar_dtype) + def _normalize_tts_handoff_tensors( + self, + tts_token_ids: Any, + tts_hidden_states: Any, + ) -> tuple[torch.Tensor, torch.Tensor]: + token_ids = torch.as_tensor(tts_token_ids, dtype=torch.long) + if token_ids.ndim == 0: + token_ids = token_ids.reshape(1) + elif token_ids.ndim > 1: + token_ids = token_ids.reshape(-1) - inputs_embeds = torch.cat([spk_embeds, tts_embeds, text_eos, audio_bos], dim=0).unsqueeze(0) - inputs_embeds = inputs_embeds.to(dtype=ar_dtype) - logger.info("generate_speech: inputs_embeds shape=%s", list(inputs_embeds.shape)) + hidden_states = torch.as_tensor(tts_hidden_states, dtype=torch.float32) + if hidden_states.ndim == 1: + hidden_states = hidden_states.unsqueeze(0) + elif hidden_states.ndim == 3 and hidden_states.shape[0] == 1: + hidden_states = hidden_states.squeeze(0) + elif hidden_states.ndim != 2: + hidden_states = hidden_states.reshape(token_ids.numel(), -1) - # Scale max_new_token with input text length to avoid mid-stream truncation on long - # responses (default 2048 can only cover ~300 text tokens at ~6x audio/text ratio). - # Empirically 511 text tokens → 1951 audio tokens (~3.8x) finishes cleanly, so use 10x - # as a safe upper bound with a floor of 2048 and a hard cap of 16384 to bound latency/mem. - num_text = int(tts_token_ids.shape[-1]) if tts_token_ids.ndim > 0 else 0 - max_new_token = max(2048, min(16384, num_text * 10)) + if hidden_states.shape[0] != token_ids.numel(): + raise ValueError( + "MiniCPM-o 4.5 TTS handoff has mismatched token/hidden lengths: " + f"tokens={token_ids.numel()} hidden_rows={hidden_states.shape[0]}" + ) + return token_ids.contiguous(), hidden_states.contiguous() - eos_token = torch.tensor([tts.config.num_audio_tokens - 1], dtype=torch.long, device=device) - outputs = tts.generate( - inputs_embeds=inputs_embeds, - eos_token=eos_token, - max_new_token=max_new_token, - show_tqdm=False, - ) - generated_tokens = outputs.new_ids.squeeze(-1) - logger.info( - "generate_speech: generated %d audio tokens (cap=%d, text_tokens=%d)", - generated_tokens.shape[-1], - max_new_token, - num_text, + def _resolve_prompt_wav_path(self, ref_audio, ref_audio_sr: int | None) -> tuple[str | None, str | None]: + temp_prompt_wav_path = self._write_ref_audio_prompt_wav(ref_audio, ref_audio_sr) + if temp_prompt_wav_path is not None: + return temp_prompt_wav_path, temp_prompt_wav_path + if (model_path := getattr(self, "_model_path", None)) is not None: + default_ref = os.path.join(model_path, "assets", "HT_ref_audio.wav") + if os.path.exists(default_ref): + return default_ref, None + return None, None + + def _max_tts_tokens_for_text(self, num_text: int) -> tuple[int, int]: + cfg = self._tts_runtime_config() + max_new_token = min( + cfg.hard_max_new_tokens, + max(cfg.min_max_new_tokens, num_text * cfg.max_token_ratio), ) + return cfg.min_new_tokens, max_new_token - if self.audio_tokenizer is None: - logger.warning("No audio_tokenizer") + def _stream_request_key(self, info: dict[str, Any]) -> str: + return get_stream_request_key(info) + + @staticmethod + def _coerce_request_key(value: Any) -> str | None: + if isinstance(value, (list, tuple)): + value = value[0] if value else None + if isinstance(value, bytes): + value = value.decode("utf-8", errors="replace") + if value is None: + return None + text = str(value) + return text if text else None + + @staticmethod + def _coerce_epoch(value: Any) -> int | None: + try: + return int(value) if value is not None else None + except (TypeError, ValueError): return None + _coerce_turn_id = _coerce_epoch + + def _stream_epoch(self, info: dict[str, Any]) -> int | None: + duplex = info.get("duplex") + if isinstance(duplex, dict): + epoch = self._coerce_epoch(duplex.get("epoch")) + if epoch is not None: + return epoch + meta = info.get("meta") + if isinstance(meta, dict): + epoch = self._coerce_epoch(meta.get("epoch")) + if epoch is not None: + return epoch + return self._coerce_epoch(info.get("epoch")) + + def _stream_turn_id(self, info: dict[str, Any]) -> int | None: + duplex = info.get("duplex") + if isinstance(duplex, dict): + turn_id = self._coerce_turn_id(duplex.get("model_turn_id")) + if turn_id is not None: + return turn_id + turn_id = self._coerce_turn_id(duplex.get("turn_id")) + if turn_id is not None: + return turn_id + meta = info.get("meta") + if isinstance(meta, dict): + turn_id = self._coerce_turn_id(meta.get("turn_id")) + if turn_id is not None: + return turn_id + return self._coerce_turn_id(info.get("turn_id")) + + def _remember_talker_request_key(self, info: dict[str, Any], key: str) -> None: + aliases = { + self._coerce_request_key(info.get("request_id")), + self._coerce_request_key(info.get("_omni_req_id")), + } + duplex = info.get("duplex") + if isinstance(duplex, dict): + aliases.add(self._coerce_request_key(duplex.get("request_id"))) + request_keys = getattr(self, "_talker_request_keys", None) + if request_keys is None: + request_keys = {} + self._talker_request_keys = request_keys + for alias in aliases: + if alias and alias != key: + request_keys[alias] = key + + def _empty_audio_chunk(self) -> torch.Tensor: + return torch.zeros((0,), dtype=torch.float32) + + @staticmethod + def _extract_tts_handoff(info: dict[str, Any]) -> tuple[Any, Any]: + return get_tts_handoff(info) + + def _native_duplex_input_ends_turn(self, info: dict[str, Any]) -> bool: + if info.get("native_duplex") is not True: + return False + meta = info.get("meta") + if not isinstance(meta, dict): + meta = {} + if bool(info.get("turn_end") or info.get("end_of_turn") or meta.get("turn_end") or meta.get("end_of_turn")): + return True + turn_eos_id = self._coerce_epoch(meta.get("turn_eos_token_id")) + if turn_eos_id is None: + return False + tts_token_ids, _ = self._extract_tts_handoff(info) + if isinstance(tts_token_ids, torch.Tensor): + return bool((tts_token_ids == turn_eos_id).any().item()) + if isinstance(tts_token_ids, np.ndarray): + return bool(np.any(tts_token_ids == turn_eos_id)) + if isinstance(tts_token_ids, (list, tuple)): + return turn_eos_id in tts_token_ids + return False + + def _t2w_pre_lookahead(self) -> int: + flow = getattr(self.audio_tokenizer, "flow", None) + try: + return int(getattr(flow, "pre_lookahead_len", 3) or 3) + except (TypeError, ValueError): + return 3 + + def _token2wav_lock(self) -> threading.RLock: + lock = getattr(self, "_token2wav_state_lock", None) + if lock is None: + lock = threading.RLock() + self._token2wav_state_lock = lock + return lock + + def _begin_turn_vocoder_cache( + self, + prompt_wav_path: str | None, + *, + state: _TalkerTurnState | None = None, + ) -> None: + """Restore a fresh per-turn clone of the ref-audio vocoder caches.""" import torchaudio - model_path = self._model_path or download_weights_from_hf_specific( - self.vllm_config.model_config.model, None, ["*"] + with self._token2wav_lock(): + cache_key = prompt_wav_path or "" + base = self._t2w_base_caches.get(cache_key) + if base is None: + if prompt_wav_path is None: + base = (None, {}) + else: + _orig_save = torchaudio.save + prev_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.float32) + try: + torchaudio.save = _soundfile_patched_save(_orig_save) + stream_cache, hift_cache_dict = self.audio_tokenizer.set_stream_cache(prompt_wav_path) + finally: + torch.set_default_dtype(prev_dtype) + torchaudio.save = _orig_save + base = ( + _torch_clone_recursive(stream_cache), + _torch_clone_recursive(hift_cache_dict), + ) + self._t2w_base_caches[cache_key] = base + stream_cache = _torch_clone_recursive(base[0]) + hift_cache_dict = _torch_clone_recursive(base[1]) + if state is None: + self.audio_tokenizer.stream_cache = stream_cache + self.audio_tokenizer.hift_cache_dict = hift_cache_dict + else: + state.stream_cache = stream_cache + state.hift_cache_dict = hift_cache_dict + state.vocoder_initialized = True + self.audio_tokenizer.stream_cache = None + self.audio_tokenizer.hift_cache_dict = {} + + def _t2w_stream_window(self, token_list: list[int], prompt_wav_path: str | None, *, last_chunk: bool): + import torchaudio + + _orig_save = torchaudio.save + prev_dtype = torch.get_default_dtype() + autocast_context, _ = self._token2wav_autocast_context() + torch.set_default_dtype(torch.float32) + try: + torchaudio.save = _soundfile_patched_save(_orig_save) + with autocast_context: + wav_np = self.audio_tokenizer.stream( + token_list, + prompt_wav_path, + last_chunk=bool(last_chunk), + return_waveform=True, + ) + finally: + torch.set_default_dtype(prev_dtype) + torchaudio.save = _orig_save + return torch.as_tensor(np.asarray(wav_np).reshape(-1), dtype=torch.float32).cpu().contiguous() + + def _run_vocoder_window( + self, + state: _TalkerTurnState, + token_list: list[int], + *, + last_chunk: bool, + ) -> torch.Tensor: + with self._token2wav_lock(): + if not state.vocoder_initialized: + self._begin_turn_vocoder_cache(state.prompt_wav_path, state=state) + self.audio_tokenizer.stream_cache = _torch_clone_recursive(state.stream_cache) + self.audio_tokenizer.hift_cache_dict = _torch_clone_recursive(state.hift_cache_dict) + try: + waveform = self._t2w_stream_window( + token_list, + state.prompt_wav_path, + last_chunk=last_chunk, + ) + state.stream_cache = _torch_clone_recursive(self.audio_tokenizer.stream_cache) + state.hift_cache_dict = _torch_clone_recursive(self.audio_tokenizer.hift_cache_dict) + return waveform + finally: + self.audio_tokenizer.stream_cache = None + self.audio_tokenizer.hift_cache_dict = {} + + def _native_duplex_vocode_tokens( + self, + state: _TalkerTurnState, + new_tokens: torch.Tensor, + *, + turn_end: bool, + force_flush: bool, + chunk_size: int, + ) -> list[torch.Tensor]: + """Run Token2wav with the same buffering policy as official duplex.""" + pre_lookahead = self._t2w_pre_lookahead() + token_list = new_tokens.reshape(-1).detach().cpu().tolist() + state.token2wav_buffer.extend(int(t) for t in token_list) + pieces: list[torch.Tensor] = [] + + if force_flush: + while len(state.token2wav_buffer) >= pre_lookahead + 5: + chunk_to_process = min(chunk_size + pre_lookahead, len(state.token2wav_buffer)) + window = state.token2wav_buffer[:chunk_to_process] + pieces.append(self._run_vocoder_window(state, window, last_chunk=False)) + state.token2wav_buffer = state.token2wav_buffer[min(chunk_size, chunk_to_process - pre_lookahead) :] + else: + while len(state.token2wav_buffer) >= chunk_size + pre_lookahead: + window = state.token2wav_buffer[: chunk_size + pre_lookahead] + pieces.append(self._run_vocoder_window(state, window, last_chunk=False)) + state.token2wav_buffer = state.token2wav_buffer[chunk_size:] + + if turn_end and state.token2wav_buffer: + pieces.append( + self._run_vocoder_window( + state, + list(state.token2wav_buffer), + last_chunk=True, + ) + ) + state.token2wav_buffer = [] + + return pieces + + def _close_turn_state( + self, + key: str, + *, + expected_epoch: int | None = None, + expected_turn_id: int | None = None, + ) -> bool: + state = self._talker_turn_states.get(key) + if ( + state is not None + and expected_epoch is not None + and state.epoch is not None + and state.epoch != expected_epoch + ): + return False + if ( + state is not None + and expected_turn_id is not None + and state.turn_id is not None + and state.turn_id != expected_turn_id + ): + return False + state = self._talker_turn_states.pop(key, None) + self._talker_consumed_tokens.pop(key, None) + request_keys = getattr(self, "_talker_request_keys", None) + if isinstance(request_keys, dict): + request_keys.pop(key, None) + for alias, mapped_key in list(request_keys.items()): + if mapped_key == key: + request_keys.pop(alias, None) + if state is None: + return True + if self.audio_tokenizer is not None: + with self._token2wav_lock(): + self.audio_tokenizer.stream_cache = None + self.audio_tokenizer.hift_cache_dict = {} + temp_path = state.temp_prompt_wav_path + if temp_path and not self._is_cached_ref_audio_prompt_wav(temp_path): + with self._token2wav_lock(): + self._t2w_base_caches.pop(temp_path, None) + try: + os.unlink(temp_path) + except OSError: + pass + return True + + @staticmethod + def _stream_identity_relation( + state: _TalkerTurnState | None, + *, + epoch: int | None, + turn_id: int | None, + ) -> str: + if state is None: + return "current" + if epoch is not None and state.epoch is not None: + if epoch < state.epoch: + return "stale" + if epoch > state.epoch: + return "newer" + if turn_id is not None and state.turn_id is not None: + if turn_id < state.turn_id: + return "stale" + if turn_id > state.turn_id: + return "newer" + return "current" + + def _warmup_duplex_vocoder(self) -> None: + """Pre-compile the two token2wav stream() modes used per turn. + + The first stream() call of each mode (mid-turn 25+pre_lookahead + window, and the variable-size last_chunk tail flush) costs ~20s of + one-time compilation; without this warmup the first spoken turn of + the first session stalls for both. Mirrors the official demo's + precompile step. + """ + if getattr(self, "_t2w_warmed", False): + return + self._t2w_warmed = True + try: + self._lazy_init_tts() + if self.audio_tokenizer is None: + return + prompt_wav_path, _ = self._resolve_prompt_wav_path(None, None) + if prompt_wav_path is None: + return + t0 = time.perf_counter() + chunk_size = self._tts_runtime_config().streaming_generator_chunk + pre_lookahead = self._t2w_pre_lookahead() + self._begin_turn_vocoder_cache(prompt_wav_path) + self._t2w_stream_window( + [_T2W_SILENCE_TOKEN] * (chunk_size + pre_lookahead), + prompt_wav_path, + last_chunk=False, + ) + self._begin_turn_vocoder_cache(prompt_wav_path) + self._t2w_stream_window( + [_T2W_SILENCE_TOKEN] * (chunk_size + pre_lookahead), + prompt_wav_path, + last_chunk=True, + ) + self.audio_tokenizer.stream_cache = None + self.audio_tokenizer.hift_cache_dict = {} + logger.info("4.5 Talker duplex vocoder warmup done in %.1fs", time.perf_counter() - t0) + except Exception: + logger.exception("4.5 Talker duplex vocoder warmup failed") + + def _create_native_duplex_stream_gen(self, info: dict[str, Any]): + """Per-segment generator over a persistent per-turn talker stream. + + Mirrors the official duplex talker: each unit calls + ``MiniCPMTTS.generate_chunk`` while carrying ``past_key_values`` and + ``text_start_pos`` across the spoken turn, then feeds the generated + audio tokens through the same per-turn Token2wav buffer/cache. + """ + key = self._stream_request_key(info) + self._remember_talker_request_key(info, key) + stream_epoch = self._stream_epoch(info) + stream_turn_id = self._stream_turn_id(info) + meta_info = info.get("meta") if isinstance(info.get("meta"), dict) else {} + codes_info = info.get("codes") if isinstance(info.get("codes"), dict) else {} + tts_token_ids, tts_hidden_states = self._extract_tts_handoff(info) + + self._lazy_init_tts() + if getattr(self, "tts_obj", None) is None or self.audio_tokenizer is None: + logger.warning("4.5 Talker duplex streaming: TTS runtime unavailable") + yield self._empty_audio_chunk(), True + return + + if isinstance(tts_token_ids, torch.Tensor): + ids_list = tts_token_ids.reshape(-1).tolist() + elif isinstance(tts_token_ids, list): + ids_list = [int(t) for t in tts_token_ids] + else: + ids_list = [] + + state = self._talker_turn_states.get(key) + identity_relation = self._stream_identity_relation( + state, + epoch=stream_epoch, + turn_id=stream_turn_id, ) - default_ref = os.path.join(model_path, "assets", "HT_ref_audio.wav") - prompt_wav_path = default_ref if os.path.exists(default_ref) else None + if identity_relation == "stale": + logger.info( + "4.5 Talker duplex drop stale handoff: key=%s state_epoch=%s state_turn_id=%s " + "handoff_epoch=%s handoff_turn_id=%s", + key, + getattr(state, "epoch", None), + getattr(state, "turn_id", None), + stream_epoch, + stream_turn_id, + ) + yield self._empty_audio_chunk(), True + return + if identity_relation == "newer": + self._close_turn_state(key) + state = None + consumed = self._talker_consumed_tokens.get(key, 0) + if consumed > len(ids_list): + consumed = 0 + pending_ids = ids_list[consumed:] + + turn_eos_raw = meta_info.get("turn_eos_token_id") + try: + turn_eos_id = int(turn_eos_raw) if turn_eos_raw is not None else None + except (TypeError, ValueError): + turn_eos_id = None + explicit_turn_end = bool( + info.get("end_of_turn") or info.get("turn_end") or meta_info.get("end_of_turn") or meta_info.get("turn_end") + ) + turn_end = explicit_turn_end or (turn_eos_id is not None and turn_eos_id in pending_ids) + terminal_only_new_turn = ( + state is None + and turn_end + and turn_eos_id is not None + and bool(pending_ids) + and all(token_id == turn_eos_id for token_id in pending_ids) + ) + if terminal_only_new_turn: + # A model-owned empty turn can be [speak, turn_eos]. There is no + # spoken state to flush, so conditioning TTS on turn_eos alone + # would synthesize an unrelated audio-only response. + yield self._empty_audio_chunk(), True + return + if state is None and not pending_ids: + # No turn open and nothing new to speak: nothing to synthesize. + yield self._empty_audio_chunk(), True + return + + tts = self.tts_obj + if not hasattr(tts.model.config, "rope_theta"): + tts.model.config.rope_theta = 10000.0 + if not callable(getattr(tts, "generate_chunk", None)): + logger.warning("4.5 Talker duplex streaming: MiniCPMTTS.generate_chunk unavailable") + yield self._empty_audio_chunk(), True + return + sampling_params = self._build_tts_sampling_params() + if sampling_params is None: + logger.warning("4.5 Talker duplex streaming: sampling params unavailable") + yield self._empty_audio_chunk(), True + return + chunk_size = self._tts_runtime_config().streaming_generator_chunk + if chunk_size <= 0: + raise ValueError("MiniCPM-o 4.5 TTS streaming generator chunk must be positive") + + if state is None: + ref_audio = codes_info.get("ref", info.get("ref_audio")) + ref_audio_sr = meta_info.get("ref_audio_sr", info.get("ref_audio_sr")) + prompt_wav_path, temp_prompt_wav_path = self._resolve_prompt_wav_path(ref_audio, ref_audio_sr) + if prompt_wav_path is None: + logger.warning("4.5 Talker duplex streaming: no ref_audio prompt; skipping audio synthesis") + yield self._empty_audio_chunk(), True + return + state = _TalkerTurnState( + prompt_wav_path, + temp_prompt_wav_path, + epoch=stream_epoch, + turn_id=stream_turn_id, + ) + self._begin_turn_vocoder_cache(prompt_wav_path, state=state) + self._talker_turn_states[key] = state + + _queue_native_duplex_segment_text( + state, + meta_info.get("native_duplex_segment_text", ""), + ) + + if pending_ids: + pending_hidden = ( + tts_hidden_states[consumed:] + if isinstance(tts_hidden_states, list) + else torch.as_tensor(tts_hidden_states)[consumed:] + ) + cond_ids, cond_hidden = self._normalize_tts_handoff_tensors(pending_ids, pending_hidden) + condition = self._build_tts_condition_embeds(cond_ids, cond_hidden).unsqueeze(0) + else: + emb_dim = int(tts.emb_text.weight.shape[1]) + condition = tts.emb_text.weight.new_zeros((1, 0, emb_dim)) + audio_bos = tts.emb_text( + torch.tensor( + [tts.audio_bos_token_id], + dtype=torch.long, + device=tts.emb_text.weight.device, + ) + ).unsqueeze(0) + condition = torch.cat([condition, audio_bos], dim=1) + max_token_per_chunk = chunk_size + 1 + min_token_per_chunk = 0 if turn_end else max_token_per_chunk + force_flush = False + if state.text_start_pos == 0: + min_token_per_chunk = 0 + force_flush = True + eos_token = torch.tensor( + [tts.config.num_audio_tokens - 1], + dtype=torch.long, + device=tts.emb_text.weight.device, + ) + temperature = torch.tensor( + [float(sampling_params.temperature)], + dtype=torch.float, + device=tts.emb_text.weight.device, + ) + new_tokens, past_key_values = tts.generate_chunk( + inputs_embeds=condition, + temperature=temperature, + repetition_penalty=sampling_params.repetition_penalty, + eos_token=eos_token, + force_no_stop=False, + max_new_token=max_token_per_chunk, + min_new_tokens=min_token_per_chunk, + past_key_values=state.past_key_values, + logits_processors=None, + text_start_pos=state.text_start_pos, + ) + if turn_end: + state.past_key_values = None + state.text_start_pos = 0 + else: + state.past_key_values = past_key_values + state.text_start_pos += int(condition.shape[1]) + int(new_tokens.shape[1]) + waveforms = self._native_duplex_vocode_tokens( + state, + new_tokens, + turn_end=turn_end, + force_flush=force_flush, + chunk_size=chunk_size, + ) + self._talker_consumed_tokens[key] = len(ids_list) + unit_waveform = _native_duplex_unit_waveform(waveforms, turn_end=turn_end) + if unit_waveform is not None: + self._ar_last_emitted_text = _drain_native_duplex_emitted_text( + state, + has_audio=True, + ) + yield unit_waveform, False + if turn_end: + self._close_turn_state(key) + self._ar_last_emitted_text = "" + yield self._empty_audio_chunk(), True + return + self._ar_last_emitted_text = "" + yield self._empty_audio_chunk(), True + + def _create_stream_gen(self, info: dict[str, Any]): + """Yield waveform chunks from MiniCPM-o remote-code TTS streaming. + + This is the real vLLM streaming path: each yielded tensor is returned + through one scheduler step. The older streaming probe still concatenates + chunks inside generate_speech(), so it cannot improve API TTFA. + """ + if info.get("native_duplex") is True: + yield from self._create_native_duplex_stream_gen(info) + return + tts_token_ids, tts_hidden_states = self._extract_tts_handoff(info) + codes_info = info.get("codes") + meta_info = info.get("meta") + if not isinstance(codes_info, dict): + codes_info = {} + if not isinstance(meta_info, dict): + meta_info = {} + + ref_audio = codes_info.get("ref", info.get("ref_audio")) + ref_audio_sr = meta_info.get("ref_audio_sr", info.get("ref_audio_sr")) + + if tts_token_ids is None or tts_hidden_states is None: + logger.warning("4.5 Talker streaming: missing tts_token_ids or tts_hidden_states") + yield self._empty_audio_chunk(), True + return + tts_token_ids, tts_hidden_states = self._normalize_tts_handoff_tensors( + tts_token_ids, + tts_hidden_states, + ) + + self._lazy_init_tts() + if not hasattr(self, "tts_obj") or self.tts_obj is None: + logger.warning("4.5 Talker streaming: tts_obj not initialized") + yield self._empty_audio_chunk(), True + return + if self.audio_tokenizer is None: + logger.warning("4.5 Talker streaming: audio_tokenizer not initialized") + yield self._empty_audio_chunk(), True + return + + generator_cls = getattr(self, "_tts_streaming_generator_cls", None) + if generator_cls is None or self._tts_gen_logits is None: + logger.warning("4.5 Talker streaming: remote-code TTSStreamingGenerator unavailable") + waveform = self.generate_speech( + tts_token_ids, + tts_hidden_states, + ref_audio=ref_audio, + ref_audio_sr=ref_audio_sr, + ) + if waveform is None: + yield self._empty_audio_chunk(), True + else: + yield torch.as_tensor(waveform, dtype=torch.float32).reshape(-1).cpu().contiguous(), True + return + + tts = self.tts_obj + if not hasattr(tts.model.config, "rope_theta"): + tts.model.config.rope_theta = 10000.0 + + tts_embeds = self._build_tts_condition_embeds(tts_token_ids, tts_hidden_states) + num_text = int(tts_token_ids.shape[-1]) if tts_token_ids.ndim > 0 else 0 + min_new_token, max_new_token = self._max_tts_tokens_for_text(num_text) + sampling_params = self._build_tts_sampling_params() + if sampling_params is None: + logger.warning("4.5 Talker streaming: sampling params unavailable") + yield self._empty_audio_chunk(), True + return + + logits_warpers, logits_processors = self._tts_gen_logits( + num_code=tts.config.num_audio_tokens, + repetition_penalty=sampling_params.repetition_penalty, + top_p=sampling_params.top_p, + top_k=sampling_params.top_k, + ) + eos_token = torch.tensor([tts.config.num_audio_tokens - 1], dtype=torch.long, device=tts.emb_text.weight.device) + chunk_size = self._tts_runtime_config().streaming_generator_chunk + if chunk_size <= 0: + raise ValueError("MiniCPM-o 4.5 TTS streaming generator chunk must be positive") + + tts_streaming_generator = generator_cls( + model=tts, + temperature=sampling_params.temperature, + eos_token=eos_token, + chunk_size=chunk_size, + logits_processors=logits_processors, + logits_warpers=logits_warpers, + ) + + prompt_wav_path, temp_prompt_wav_path = self._resolve_prompt_wav_path(ref_audio, ref_audio_sr) + stream_cache = hift_cache_dict = None + import torchaudio _orig_save = torchaudio.save @@ -299,84 +1316,255 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): return return _orig_save(uri, src, sample_rate, backend="soundfile", **kw) - torchaudio.save = _patched_save - prev_dtype = torch.get_default_dtype() - torch.set_default_dtype(torch.float32) + yielded_any = False try: - # Vocoder path is float32; use the platform abstraction because - # torch.amp.autocast validates unsupported device types even when - # autocast is disabled. - autocast_device = device.type if isinstance(device, torch.device) else str(device) - with current_omni_platform.create_autocast_context( - device_type=autocast_device, - dtype=torch.float32, - enabled=False, - ): - token_list = generated_tokens.squeeze(0).tolist() - num_tokens = len(token_list) + torchaudio.save = _patched_save + prev_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.float32) + try: + stream_cache, hift_cache_dict = self.audio_tokenizer.set_stream_cache(prompt_wav_path) + finally: + torch.set_default_dtype(prev_dtype) + torchaudio.save = _orig_save + self.audio_tokenizer.stream_cache = stream_cache + self.audio_tokenizer.hift_cache_dict = hift_cache_dict + token_iter = tts_streaming_generator.generate_with_buffer( + condition=tts_embeds.unsqueeze(0), + text_finished=True, + max_new_token=max_new_token, + ) + while True: + try: + audio_token_chunk, is_last = next(token_iter) + except StopIteration: + break + if audio_token_chunk is None: + break - # For long outputs, the one-shot vocoder path - # (Token2wav.__call__ -> flow.inference) runs full O(N^2) self- - # attention over all audio tokens and OOMs on a 24GB card once - # N exceeds a few thousand (e.g. 4964 tokens needs ~3GiB for a - # single attention matmul). Switch to the chunked / streaming - # vocoder (set_stream_cache + stream) which truncates the flow - # attention caches to prompt_len + 100 steps on every chunk, - # keeping peak memory bounded regardless of total length. - STREAM_THRESHOLD = int(os.environ.get("MINICPMO45_TTS_STREAM_THRESHOLD", "2500")) # ~100s @ 25Hz - CHUNK_SIZE = int(os.environ.get("MINICPMO45_TTS_STREAM_CHUNK", "50")) # ~2s per chunk - MIN_TAIL = 6 # must exceed flow.pre_lookahead_len (typically 3) + token_list = audio_token_chunk.reshape(-1).detach().cpu().tolist() + if not token_list: + if is_last: + yield self._empty_audio_chunk(), True + yielded_any = True + break + continue - if num_tokens <= STREAM_THRESHOLD: - wav_bytes = self.audio_tokenizer(token_list, prompt_wav_path) - waveform, sr = sf.read(io.BytesIO(wav_bytes)) - waveform = waveform.astype(np.float32) - else: - # Build chunk boundaries, merging a too-small tail into the - # previous chunk so every chunk satisfies MIN_TAIL. - boundaries = [] - i = 0 - while i < num_tokens: - end = min(i + CHUNK_SIZE, num_tokens) - if 0 < num_tokens - end < MIN_TAIL: - end = num_tokens - boundaries.append((i, end)) - i = end - - logger.info( - "generate_speech: streaming vocoder, %d tokens -> %d chunks (chunk=%d)", - num_tokens, - len(boundaries), - CHUNK_SIZE, - ) - - stream_cache, hift_cache_dict = self.audio_tokenizer.set_stream_cache(prompt_wav_path) - self.audio_tokenizer.stream_cache = stream_cache - self.audio_tokenizer.hift_cache_dict = hift_cache_dict - - try: - pieces = [] - for idx, (s, e) in enumerate(boundaries): - is_last = idx == len(boundaries) - 1 - wav_np = self.audio_tokenizer.stream( - token_list[s:e], - prompt_wav_path, - last_chunk=is_last, - return_waveform=True, - ) - pieces.append(np.asarray(wav_np).reshape(-1)) - waveform = np.concatenate(pieces, axis=0).astype(np.float32) - sr = 24000 - finally: - # Free per-request streaming state so the next request starts clean - self.audio_tokenizer.stream_cache = None - self.audio_tokenizer.hift_cache_dict = {} + autocast_context, _ = self._token2wav_autocast_context() + torchaudio.save = _patched_save + prev_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.float32) + try: + with autocast_context: + wav_np = self.audio_tokenizer.stream( + token_list, + prompt_wav_path, + last_chunk=bool(is_last), + return_waveform=True, + ) + finally: + torch.set_default_dtype(prev_dtype) + torchaudio.save = _orig_save + chunk = torch.as_tensor(np.asarray(wav_np).reshape(-1), dtype=torch.float32).cpu().contiguous() + yielded_any = True + yield chunk, bool(is_last) + if is_last: + break finally: - torch.set_default_dtype(prev_dtype) torchaudio.save = _orig_save + self.audio_tokenizer.stream_cache = None + self.audio_tokenizer.hift_cache_dict = {} + if temp_prompt_wav_path and not self._is_cached_ref_audio_prompt_wav(temp_prompt_wav_path): + try: + os.unlink(temp_prompt_wav_path) + except OSError: + pass - logger.info("generate_speech: waveform %d samples, sr=%d", waveform.shape[0], sr) - return waveform + if not yielded_any: + yield self._empty_audio_chunk(), True + + def _move_tts_modules_to_device(self) -> torch.dtype: + device = current_omni_platform.get_torch_device() + target_dtype = torch.bfloat16 if current_omni_platform.is_npu() else self._target_tts_dtype() + if target_dtype is torch.float32: + self.tts_obj = self.tts_obj.to(device) + logger.info("Moved MiniCPM-o 4.5 TTS object to %s dtype=%s", device, target_dtype) + return target_dtype + + for module_name in ( + "emb_text", + "model", + "projector_spk", + "projector_semantic", + "emb_code", + "head_code", + ): + module = getattr(self.tts_obj, module_name, None) + if module is not None: + module.to(device=device, dtype=target_dtype) + logger.info("Moved MiniCPM-o 4.5 TTS AR modules to %s dtype=%s", device, target_dtype) + return target_dtype + + def generate_speech( + self, + tts_token_ids: torch.Tensor, + tts_hidden_states: torch.Tensor, + ref_audio=None, + ref_audio_sr: int | None = None, + ) -> torch.Tensor | np.ndarray | None: + """Run full 4.5 TTS pipeline using original MiniCPMTTS.generate.""" + self._lazy_init_tts() + if not hasattr(self, "tts_obj") or self.tts_obj is None: + logger.warning("generate_speech: tts_obj not initialized") + return None + + tts = self.tts_obj + device = tts.emb_text.weight.device + dtype = tts.emb_text.weight.dtype + + llm_embeds = tts.emb_text(tts_token_ids.to(device)) + hidden_embeds = tts.projector_semantic(tts_hidden_states.to(device=device, dtype=dtype)) + if getattr(tts.config, "normalize_projected_hidden", False): + hidden_embeds = F.normalize(hidden_embeds, p=2, dim=-1) + tts_embeds = llm_embeds + hidden_embeds + + text_eos = tts.emb_text(torch.tensor([tts.config.text_eos_token_id], device=device, dtype=torch.long)) + audio_bos = tts.emb_text(torch.tensor([tts.audio_bos_token_id], device=device, dtype=torch.long)) + spk_embeds = torch.zeros(0, tts.config.hidden_size, device=device, dtype=tts_embeds.dtype) + + inputs_embeds = torch.cat([spk_embeds, tts_embeds, text_eos, audio_bos], dim=0).unsqueeze(0) + + # Scale max_new_token with input text length. A fixed 2048-token floor + # can turn an EOS miss on a very short response into ~82s of audio and + # ~18s E2E latency. Keep a conservative short-text floor while bounding + # the tail. + num_text = int(tts_token_ids.shape[-1]) if tts_token_ids.ndim > 0 else 0 + min_new_token, max_new_token = self._max_tts_tokens_for_text(num_text) + + eos_token = torch.tensor([tts.config.num_audio_tokens - 1], dtype=torch.long, device=device) + sampling_params = self._build_tts_sampling_params() + generate_kwargs = { + "inputs_embeds": inputs_embeds, + "eos_token": eos_token, + "max_new_token": max_new_token, + "min_new_token": min_new_token, + "show_tqdm": False, + } + if sampling_params is not None: + generate_kwargs["sampling_params"] = sampling_params + + if self.audio_tokenizer is None: + logger.warning("No audio_tokenizer") + return None + + prompt_wav_path, temp_prompt_wav_path = self._resolve_prompt_wav_path(ref_audio, ref_audio_sr) + + try: + outputs = tts.generate(**generate_kwargs) + generated_tokens = outputs.new_ids.squeeze(-1) + + import torchaudio + + _orig_save = torchaudio.save + + def _patched_save(uri, src, sample_rate, **kw): + kw.pop("backend", None) + if hasattr(uri, "write"): + sf.write(uri, src.cpu().numpy().T, sample_rate, format="WAV") + return + return _orig_save(uri, src, sample_rate, backend="soundfile", **kw) + + torchaudio.save = _patched_save + prev_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.float32) + try: + autocast_context, token2wav_autocast = self._token2wav_autocast_context() + with autocast_context: + num_tokens = int(generated_tokens.shape[-1]) + + # For long outputs, the one-shot vocoder path + # (Token2wav.__call__ -> flow.inference) runs full O(N^2) self- + # attention over all audio tokens and OOMs on a 24GB card once + # N exceeds a few thousand (e.g. 4964 tokens needs ~3GiB for a + # single attention matmul). Switch to the chunked / streaming + # vocoder (set_stream_cache + stream) which truncates the flow + # attention caches to prompt_len + 100 steps on every chunk, + # keeping peak memory bounded regardless of total length. + STREAM_THRESHOLD = self._tts_runtime_config().streaming_vocoder_threshold # ~100s @ 25Hz + CHUNK_SIZE = self._tts_runtime_config().streaming_vocoder_chunk # ~2s per chunk + MIN_TAIL = 6 # must exceed flow.pre_lookahead_len (typically 3) + + if num_tokens <= STREAM_THRESHOLD: + if self._should_use_direct_token2wav() and token2wav_autocast == "off": + try: + waveform, _ = self._run_token2wav_direct(generated_tokens, prompt_wav_path) + except Exception as exc: + logger.warning( + "MiniCPM-o 4.5 direct Token2wav path failed; falling back to WAV path: %s", + exc, + exc_info=True, + ) + token_list = generated_tokens.squeeze(0).tolist() + self._reset_token2wav_cache_if_needed(prompt_wav_path) + wav_bytes = self.audio_tokenizer(token_list, prompt_wav_path) + waveform, _ = sf.read(io.BytesIO(wav_bytes)) + waveform = waveform.astype(np.float32) + else: + token_list = generated_tokens.squeeze(0).tolist() + self._reset_token2wav_cache_if_needed(prompt_wav_path) + wav_bytes = self.audio_tokenizer(token_list, prompt_wav_path) + waveform, _ = sf.read(io.BytesIO(wav_bytes)) + waveform = waveform.astype(np.float32) + else: + token_list = generated_tokens.squeeze(0).tolist() + # Build chunk boundaries, merging a too-small tail into the + # previous chunk so every chunk satisfies MIN_TAIL. + boundaries = [] + i = 0 + while i < num_tokens: + end = min(i + CHUNK_SIZE, num_tokens) + if 0 < num_tokens - end < MIN_TAIL: + end = num_tokens + boundaries.append((i, end)) + i = end + + logger.info( + "generate_speech: streaming vocoder, %d tokens -> %d chunks (chunk=%d)", + num_tokens, + len(boundaries), + CHUNK_SIZE, + ) + + stream_cache, hift_cache_dict = self.audio_tokenizer.set_stream_cache(prompt_wav_path) + self.audio_tokenizer.stream_cache = stream_cache + self.audio_tokenizer.hift_cache_dict = hift_cache_dict + + try: + pieces = [] + for idx, (s, e) in enumerate(boundaries): + is_last = idx == len(boundaries) - 1 + wav_np = self.audio_tokenizer.stream( + token_list[s:e], + prompt_wav_path, + last_chunk=is_last, + return_waveform=True, + ) + pieces.append(np.asarray(wav_np).reshape(-1)) + waveform = np.concatenate(pieces, axis=0).astype(np.float32) + finally: + # Free per-request streaming state so the next request starts clean + self.audio_tokenizer.stream_cache = None + self.audio_tokenizer.hift_cache_dict = {} + finally: + torch.set_default_dtype(prev_dtype) + torchaudio.save = _orig_save + + return waveform + finally: + if temp_prompt_wav_path and not self._is_cached_ref_audio_prompt_wav(temp_prompt_wav_path): + try: + os.unlink(temp_prompt_wav_path) + except OSError: + pass def _generate_tokens(self, inputs_embeds: torch.Tensor, max_new_token: int = 2048) -> torch.Tensor | None: """Autoregressive generation of audio tokens using the TTS LlamaModel.""" @@ -423,30 +1611,6 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): return new_tokens[:, : t + 1 if finished else t, :] - def _dummy_hidden_states( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor | None, - inputs_embeds: torch.Tensor | None, - ) -> torch.Tensor: - """Shape-correct zero tensor for vllm KV cache profiling. - - vllm's gpu_model_runner._dummy_run takes forward()'s return value as - ``hidden_states`` and does ``hidden_states[logit_indices_device]``; - returning None on the dummy path crashes with - ``TypeError: 'NoneType' object is not subscriptable``. - """ - for ref in (input_ids, positions, inputs_embeds): - if isinstance(ref, torch.Tensor): - num_tokens = int(ref.shape[0]) if ref.ndim >= 1 else 1 - device = ref.device - break - else: - num_tokens = 1 - device = current_omni_platform.get_torch_device() - hidden_size = int(getattr(self, "_hidden_size", 768) or 768) - return torch.zeros((num_tokens, hidden_size), device=device, dtype=torch.bfloat16) - def forward( self, input_ids=None, @@ -458,53 +1622,133 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): ): if additional_information is None: additional_information = {} + if not additional_information: + # Profile/dummy run: use it to pre-compile the per-turn vocoder + # stream modes so the first real spoken turn does not stall. + self._warmup_duplex_vocoder() - tts_token_ids = additional_information.get("tts_token_ids") - tts_hidden_states = additional_information.get("tts_hidden_states") + tts_token_ids, tts_hidden_states = self._extract_tts_handoff(additional_information) tts_text = additional_information.get("llm_output_text", [""]) if isinstance(tts_text, list): tts_text = tts_text[0] if tts_text else "" + codes_info = additional_information.get("codes") + meta_info = additional_information.get("meta") + if not isinstance(codes_info, dict): + codes_info = {} + if not isinstance(meta_info, dict): + meta_info = {} + ref_audio = codes_info.get("ref") + if ref_audio is None: + ref_audio = additional_information.get("ref_audio") + ref_audio_sr = meta_info.get("ref_audio_sr") + if ref_audio_sr is None: + ref_audio_sr = additional_information.get("ref_audio_sr") if tts_token_ids is None or tts_hidden_states is None: - # KV cache profiling / dummy run path — no real TTS input yet. - logger.debug("4.5 Talker: dummy forward (missing tts_token_ids/tts_hidden_states)") - return self._dummy_hidden_states(input_ids, positions, inputs_embeds) + logger.warning("4.5 Talker: missing tts_token_ids or tts_hidden_states") + self._ar_last_chunk_flags = [True] + self._ar_turn_end_flags = [False] + return None, None + tts_token_ids, tts_hidden_states = self._normalize_tts_handoff_tensors( + tts_token_ids, + tts_hidden_states, + ) - logger.info("4.5 Talker: generating speech for %d tokens", tts_token_ids.shape[0]) - waveform = self.generate_speech(tts_token_ids, tts_hidden_states) - # Tuple layout: (mel_spec, waveform). 4.5 talker emits only waveform, - # so mel_spec stays None; the wrapper unpacks in this order and - # packages the waveform into ``multimodal_outputs["model_outputs"]``. + if self._should_stream_output(additional_information): + request_key = self._stream_request_key(additional_information) + input_ends_turn = self._native_duplex_input_ends_turn(additional_information) + if request_key not in self._stream_gens: + self._stream_gens[request_key] = self._create_stream_gen(additional_information) + generator = self._stream_gens[request_key] + try: + waveform_chunk, is_last = next(generator) + except StopIteration: + self._stream_gens.pop(request_key, None) + waveform_chunk = self._empty_audio_chunk() + is_last = True + if is_last: + self._stream_gens.pop(request_key, None) + self._ar_last_chunk_flags = [bool(is_last)] + # A TTS generator also ends at ordinary chunk boundaries. Export + # turn_end only on the terminal output for a condition that + # actually contains the model's <|turn_eos|> decision. + self._ar_turn_end_flags = [bool(is_last and input_ends_turn)] + return None, waveform_chunk.reshape(-1).contiguous() + + self._ar_last_chunk_flags = [True] + self._ar_turn_end_flags = [False] + waveform = self.generate_speech( + tts_token_ids, + tts_hidden_states, + ref_audio=ref_audio, + ref_audio_sr=ref_audio_sr, + ) if waveform is not None: - return None, torch.tensor(waveform, dtype=torch.float32) + waveform_tensor = torch.as_tensor(waveform, dtype=torch.float32).detach() + if waveform_tensor.device.type != "cpu": + waveform_tensor = waveform_tensor.cpu() + return waveform_tensor.reshape(-1).contiguous(), None return None, None def compute_logits(self, hidden_states, *args, **kwargs): - # Placeholder logits: one row per sampled request (the scheduler - # indexes sampled_token_ids by req_index). Hardcoding a single row - # breaks batched/concurrent decoding with IndexError. The values are - # discarded — real output is the waveform via multimodal_outputs. + device = hidden_states.device if isinstance(hidden_states, torch.Tensor) else torch.device("cuda") if isinstance(hidden_states, torch.Tensor): - device = hidden_states.device - num_reqs = hidden_states.shape[0] if hidden_states.ndim >= 1 else 1 + if hidden_states.ndim == 1: + num_rows = 1 + else: + num_rows = max(1, int(hidden_states.shape[0])) else: - device = current_omni_platform.get_torch_device() - num_reqs = 1 - return torch.zeros(num_reqs, 2, device=device) + num_rows = 1 + eos_id = self._scheduler_eos_token_id() + vocab_size = max(int(getattr(self.config, "vocab_size", eos_id + 1) or (eos_id + 1)), eos_id + 1, 3) + safe_id = 1 if eos_id != 1 else 0 + logits = torch.full((num_rows, vocab_size), -1.0e9, dtype=torch.float32, device=device) + flags = self._ar_last_chunk_flags + default_is_last = bool(flags[-1]) if flags else True + for row in range(num_rows): + is_last = bool(flags[row]) if row < len(flags) else default_is_last + if is_last: + logits[row, eos_id] = 1.0e6 + else: + logits[row, safe_id] = 1.0e6 + return logits def sample(self, logits, sampling_metadata): - return None + if logits is None or logits.numel() == 0: + return None + sampled = torch.argmax(logits, dim=-1).to(torch.int32) + return SamplerOutput(sampled_token_ids=sampled.unsqueeze(-1), logprobs_tensors=None) + + def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: + for req_id in finished_req_ids: + request_key = str(req_id) + mapped_key = getattr(self, "_talker_request_keys", {}).pop(request_key, None) + keys = {request_key} + if mapped_key: + keys.add(mapped_key) + for key in list(self._stream_gens): + if key in keys: + gen = self._stream_gens.pop(key, None) + if gen is not None: + try: + gen.close() + except Exception: + logger.exception("MiniCPM-o 4.5 failed to close stream gen for request %s", req_id) + for key in list(self._talker_turn_states): + if key in keys: + self._close_turn_state(key) + for key in list(self._talker_consumed_tokens): + if key in keys: + self._talker_consumed_tokens.pop(key, None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): loaded = set() tts_weights = {} for k, v in weights: if k.startswith("tts."): - tts_weights[k.replace("tts.", "", 1)] = v - # vllm sanity-checks `loaded` against `named_parameters()`. - # The submodule is attached at `self.tts_obj`, not `self.tts`, - # so report the loaded name under the on-module path. - loaded.add(k.replace("tts.", "tts_obj.", 1)) + local_name = k.replace("tts.", "", 1) + tts_weights[local_name] = v + loaded.add(f"tts_obj.{local_name}") if tts_weights and self._tts_config is not None: self._lazy_init_tts() @@ -514,30 +1758,20 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): logger.warning("TTS missing keys (%d): %s", len(missing), missing[:5]) if unexpected: logger.warning("TTS unexpected keys (%d): %s", len(unexpected), unexpected[:5]) - # Move the AR backbone to the active device (cuda / npu / …) and - # cast to bfloat16: MiniCPMTTS AR uses FlashAttention (fp16/bf16 - # only) and is built under a float32 default dtype during lazy - # init, so an uncast float32 backbone breaks CUDA FA2 and wastes - # memory on the NPU sdpa path. Detach the Token2wav vocoder first - # so the cast does not drag it onto the accelerator or downcast - # its float32 flow/HiFT weights: it manages its own device - # placement and may not be an nn.Module. - device = current_omni_platform.get_torch_device() - audio_tok = getattr(self.tts_obj, "audio_tokenizer", None) - if audio_tok is not None: - self.tts_obj.audio_tokenizer = None - try: - self.tts_obj = self.tts_obj.to(device=device, dtype=torch.bfloat16) - finally: - if audio_tok is not None: - self.tts_obj.audio_tokenizer = audio_tok - self.audio_tokenizer = audio_tok + tts_dtype = self._move_tts_modules_to_device() + if ( + not current_omni_platform.is_npu() + and self.audio_tokenizer is not None + and hasattr(self.audio_tokenizer, "to") + ): + self.audio_tokenizer.to("cuda") self.emb_text = self.tts_obj.emb_text self.projector_semantic = self.tts_obj.projector_semantic logger.info( - "Loaded %d TTS weights, moved to %s (bfloat16)", + "Loaded %d TTS weights, moved AR modules to %s dtype=%s", len(tts_weights), - device, + current_omni_platform.get_torch_device(), + tts_dtype, ) return loaded 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 f966555d0..57d8c4614 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -24,6 +24,11 @@ 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.experimental.fullduplex.minicpmo45.runtime.MiniCPMO45DuplexRuntimeExtension"), + duplex_serving_adapter=( + "vllm_omni.experimental.fullduplex.minicpmo45.serving_adapter.MiniCPMO45ServingRuntimeAdapter" + ), + duplex_control_enabled=True, # 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 799914846..561c6ac8f 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 @@ -2,37 +2,350 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Stage input processor for MiniCPM-o 4.5: Thinker (LLM) -> Talker (TTS). -This is the original vLLM-Omni bridge: it converts the thinker stage's -hidden states + token ids into the talker stage's prompt payload. The -talker model itself is adapted from openbmb/MiniCPM-o-4_5 (see the headers -on vllm_omni/model_executor/models/minicpmo_4_5/*.py). +This bridge converts thinker hidden states and token ids into the talker +prompt payload. The talker runs Token2Wav internally; MiniCPM-o 4.5 does not +have an independent code2wav pipeline stage. """ -from typing import Any +import logging +from collections.abc import Callable, Mapping import torch from vllm.inputs import TextPrompt +from vllm_omni.experimental.fullduplex.engine.intermediate import ( + build_duplex_intermediate_buffer, + set_ref_audio, + set_tts_handoff, +) from vllm_omni.inputs.data import OmniTokensPrompt +logger = logging.getLogger(__name__) + + +def _extract_first_audio_ref(multi_modal_data): + if not isinstance(multi_modal_data, dict): + return None + audio_data = multi_modal_data.get("audio") + if audio_data is None: + return None + if isinstance(audio_data, list): + if not audio_data: + return None + audio_data = audio_data[0] + + samples = None + sample_rate = None + if isinstance(audio_data, tuple) and len(audio_data) >= 2: + samples, sample_rate = audio_data[0], audio_data[1] + elif isinstance(audio_data, dict): + sample_rate = audio_data.get("sample_rate") or audio_data.get("sampling_rate") or audio_data.get("sr") + for key in ("audio", "wav", "samples", "array", "waveform"): + if key in audio_data: + samples = audio_data[key] + break + if samples is None or sample_rate is None: + return None + + waveform = torch.as_tensor(samples, dtype=torch.float32) + if waveform.ndim > 1: + if waveform.shape[0] <= 2 and waveform.shape[-1] > waveform.shape[0]: + waveform = waveform.mean(dim=0) + else: + waveform = waveform.mean(dim=-1) + return waveform.reshape(-1).cpu(), int(sample_rate) + + +def _extract_native_runtime_ref_audio(data_plane_metadata): + if not isinstance(data_plane_metadata, dict): + return None + runtime_config = data_plane_metadata.get("runtime_config") + if not isinstance(runtime_config, dict): + return None + from vllm_omni.experimental.fullduplex.minicpmo45.input import decode_native_ref_audio_from_config + + waveform = decode_native_ref_audio_from_config({"extra_body": runtime_config}) + if waveform is None: + return None + sample_rate = runtime_config.get("ref_audio_sample_rate_hz") or 16000 + return torch.as_tensor(waveform, dtype=torch.float32).reshape(-1).cpu(), int(sample_rate) + + +def _coerce_token_id_list(value): + if value is None: + return None + if hasattr(value, "detach"): + value = value.detach().cpu().reshape(-1).tolist() + if isinstance(value, list) and value and isinstance(value[0], list): + value = value[0] + if isinstance(value, tuple): + value = list(value) + if not isinstance(value, list): + return None + out = [] + for item in value: + try: + out.append(int(item)) + except (TypeError, ValueError): + return None + return out + + +def _to_transport_list(value): + if hasattr(value, "detach"): + value = value.detach().cpu() + if isinstance(value, torch.Tensor): + return value.tolist() + return value + + +def _coerce_int(value): + if hasattr(value, "detach"): + flat = value.detach().cpu().reshape(-1) + if flat.numel() == 0: + return None + value = flat[0].item() + try: + return int(value) + except (TypeError, ValueError): + return None + + +def _special_token_ids_from_mm_output(mm_output): + if not isinstance(mm_output, Mapping): + return {} + meta = mm_output.get("meta") + if not isinstance(meta, dict): + meta = {} + special_token_ids = mm_output.get("special_token_ids") + if not isinstance(special_token_ids, dict): + special_token_ids = {} + flat_meta = { + key.removeprefix("meta."): value + for key, value in mm_output.items() + if isinstance(key, str) and key.startswith("meta.") + } + return { + key: value + for key, value in ( + (key, _coerce_int(value)) + for source in (special_token_ids, meta, flat_meta) + for key, value in source.items() + ) + if value is not None and value >= 0 + } + + +def _has_native_duplex_prompt_metadata(mm_output): + if not isinstance(mm_output, Mapping): + return False + return mm_output.get("duplex_prompt_token_ids") is not None or mm_output.get("ids.prompt") is not None + + +def _require_native_tts_boundary_metadata(special_token_ids): + if special_token_ids.get("tts_bos_token_id") is None: + raise ValueError( + "MiniCPM-o native duplex TTS handoff requires tokenizer-derived " + "<|tts_bos|> metadata; refusing to infer special token ids." + ) + + +def _native_tts_boundary_token_ids(special_token_ids): + # Official duplex conditions the talker on mid-unit <|speak|> tokens AND + # the <|turn_eos|> token+hidden (its embedding is the trained stop + # signal); only chunk terminators and framing tokens bound the slice. + return { + token_id + for token_id in ( + special_token_ids.get("tts_eos_token_id"), + special_token_ids.get("tts_pad_token_id"), + special_token_ids.get("listen_token_id"), + special_token_ids.get("chunk_eos_token_id"), + special_token_ids.get("chunk_tts_eos_token_id"), + special_token_ids.get("unit_token_id"), + special_token_ids.get("unit_end_token_id"), + ) + if token_id is not None + } + + +def _native_duplex_segment_output_ids( + output_ids: list[int], + output_text: str, + streaming_context, + *, + request_id: str, +) -> tuple[list[int], str]: + """Slice the cumulative thinker output down to the current segment. + + Tracks how many output tokens were already handed to the talker in the + orchestrator's streaming bridge state. Segment transcript is decoded from + the same token delta, not from a character cursor over cumulative text. + """ + bridge_states = getattr(streaming_context, "bridge_states", None) + if not isinstance(bridge_states, dict): + return output_ids, output_text + state = bridge_states.setdefault("minicpmo45_tts_handoff", {}) + duplex_state = bridge_states.get("duplex") + turn_id = duplex_state.get("model_turn_id", duplex_state.get("turn_id")) if isinstance(duplex_state, dict) else None + if not isinstance(turn_id, int): + turn_id = None + if state.get("request_id") != request_id: + state["request_id"] = request_id + state["sent_output_len"] = 0 + state["sent_output_ids"] = [] + state["acc_tts_ids"] = [] + state["acc_tts_hidden"] = [] + sent_len = state.get("sent_output_len", 0) + prev_output_ids = state.get("sent_output_ids", []) + prev_turn_id = state.get("turn_id") + if not isinstance(sent_len, int) or sent_len < 0 or sent_len > len(output_ids): + # Shrunken cumulative output = epoch reset after barge-in: the talker + # condition history is stale too. + sent_len = 0 + state["acc_tts_ids"] = [] + state["acc_tts_hidden"] = [] + elif sent_len and isinstance(prev_output_ids, list): + prev_prefix = prev_output_ids[:sent_len] + current_prefix = output_ids[:sent_len] + if current_prefix != prev_prefix: + # Some runtimes restart output token ids after a turn boundary while + # others keep returning cumulative ids. Detect restart from the + # token prefix instead of clearing the cursor at turn_eos. + sent_len = 0 + state["acc_tts_ids"] = [] + state["acc_tts_hidden"] = [] + if isinstance(prev_turn_id, int) and isinstance(turn_id, int) and prev_turn_id != turn_id: + # The thinker output list is cumulative across clean turns, so the + # output cursor must stay put. The talker condition is per assistant + # turn, though; carrying it across turn_id boundaries makes stage 1 + # replay the previous turn after its consumed cursor has been reset. + state["acc_tts_ids"] = [] + state["acc_tts_hidden"] = [] + segment_ids = output_ids[sent_len:] + decode_token_ids = getattr(streaming_context, "source_token_decoder", None) + if isinstance(decode_token_ids, Callable): + decode_ids = [int(token_id) for token_id in segment_ids] + try: + try: + segment_text = str(decode_token_ids(decode_ids, skip_special_tokens=True)) + except TypeError: + segment_text = str(decode_token_ids(decode_ids)) + except Exception: + logger.exception("Failed to decode MiniCPM-o duplex token delta for request_id=%s", request_id) + segment_text = "" + elif sent_len == 0: + # Non-orchestrator unit tests and legacy callers may not install a + # decoder. Never slice cumulative text with a stale character cursor. + segment_text = output_text + else: + logger.warning( + "MiniCPM-o native duplex token delta decoder missing for request_id=%s; " + "suppressing cumulative transcript fallback", + request_id, + ) + segment_text = "" + state["sent_output_len"] = len(output_ids) + state["sent_output_ids"] = list(output_ids) + state["turn_id"] = turn_id + return segment_ids, segment_text + + +def _reset_native_tts_handoff(streaming_context) -> None: + bridge_states = getattr(streaming_context, "bridge_states", None) + if not isinstance(bridge_states, dict): + return + state = bridge_states.get("minicpmo45_tts_handoff") + if not isinstance(state, dict): + return + # Keep the output cursor: the thinker can keep reporting cumulative + # output after turn_eos on the same resumable request. If a runtime really + # restarts token ids, _native_duplex_segment_output_ids detects the prefix + # mismatch and resets the slice cursor there. + state["acc_tts_ids"] = [] + state["acc_tts_hidden"] = [] + + +def _native_duplex_data_plane_metadata(streaming_context) -> dict[str, object] | None: + bridge_states = getattr(streaming_context, "bridge_states", None) + if not isinstance(bridge_states, dict): + return None + duplex_state = bridge_states.get("duplex") + if not isinstance(duplex_state, dict): + return None + + metadata: dict[str, object] = {"data_plane": True} + 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 + turn_id = duplex_state.get("model_turn_id", duplex_state.get("turn_id")) + if isinstance(turn_id, int): + metadata["turn_id"] = turn_id + session_config = duplex_state.get("session_config") + if isinstance(session_config, dict): + metadata["session_config"] = dict(session_config) + runtime_config = duplex_state.get("runtime_config") + if isinstance(runtime_config, dict): + metadata["runtime_config"] = dict(runtime_config) + return metadata + + +def _accumulate_native_tts_handoff(streaming_context, new_ids, new_hidden): + """Hand the talker the FULL accumulated condition on every handoff. + + The runner's streaming buffer update is not merge-safe for a resumable + stage-1 request: in-place updates merge sub-keys, but a resume prefill + REPLACES the buffer, silently dropping every earlier segment's tts + tokens/hiddens (observed losing alternating reply segments — the talker + vocalized text it never saw between islands it did). Accumulate here so + the latest handoff always carries the complete history and downstream + replace semantics are lossless; the talker consumes by cursor. + """ + bridge_states = getattr(streaming_context, "bridge_states", None) + if not isinstance(bridge_states, dict): + return new_ids, new_hidden + state = bridge_states.setdefault("minicpmo45_tts_handoff", {}) + acc_ids = state.setdefault("acc_tts_ids", []) + acc_hidden = state.setdefault("acc_tts_hidden", []) + if new_ids: + acc_ids.extend(int(t) for t in new_ids) + if new_hidden: + acc_hidden.extend(new_hidden) + if not acc_ids: + return None, None + return list(acc_ids), list(acc_hidden) + + +def _build_tts_scheduler_prompt_token_ids( + tts_token_ids: torch.Tensor | None, + llm_output_ids: list[int], + prompt_token_ids: list[int], +) -> list[int]: + if tts_token_ids is not None: + ids = _coerce_token_id_list(tts_token_ids) + if ids: + return ids + if llm_output_ids: + return llm_output_ids + if prompt_token_ids: + return prompt_token_ids[-1:] + raise ValueError("MiniCPM-o TTS stage requires at least one scheduler prompt token") + def llm2tts( - source_outputs: list[Any], - prompt: OmniTokensPrompt | TextPrompt | dict | list | None = None, + source_outputs, + prompt: OmniTokensPrompt | TextPrompt = None, requires_multimodal_data: bool = False, - streaming_context: Any | None = None, + _streaming_context=None, ): """Convert thinker stage output to talker stage input for MiniCPMO Omni. - The signature matches the framework's ``custom_process_input_func`` call - convention used by ``StageEngineCoreClientBase.process_engine_inputs``: - - (source_outputs, prompt, requires_multimodal_data, streaming_context) - - ``source_outputs`` is the already-resolved list of upstream engine - outputs (one entry per request), so we do not need to look anything up - via ``stage_list[source_stage_id].engine_outputs``. - Extracts from thinker output: - Full hidden states (prompt + generated) for speaker embedding extraction - Prompt token IDs (for finding spk_bos/spk_eos positions) @@ -44,8 +357,6 @@ def llm2tts( 3. Decode generated text and extract TTS content 4. Run ConditionalChatTTS pipeline """ - del streaming_context # not used by MiniCPM-o 4.5 turn-taking pipeline - if not source_outputs: raise ValueError("source_outputs cannot be empty") @@ -55,75 +366,242 @@ def llm2tts( if not isinstance(prompt, list): prompt = [prompt] - multi_modal_data = { - llm_output.request_id: p.get("multi_modal_data", None) if isinstance(p, dict) else None - for llm_output, p in zip(llm_outputs, prompt) - } + multi_modal_data = {} + for llm_output, p in zip(llm_outputs, prompt): + if isinstance(p, dict): + multi_modal_data[llm_output.request_id] = p.get("multi_modal_data", None) + else: + multi_modal_data[llm_output.request_id] = getattr(p, "multi_modal_data", None) - for i, llm_output in enumerate(llm_outputs): + for llm_output in llm_outputs: output = llm_output.outputs[0] - prompt_token_ids = llm_output.prompt_token_ids - llm_output_ids = output.token_ids + mm_output = output.multimodal_output if isinstance(output.multimodal_output, Mapping) else {} + special_token_ids = _special_token_ids_from_mm_output(mm_output) + prompt_token_ids = ( + _coerce_token_id_list(mm_output.get("duplex_prompt_token_ids")) + or _coerce_token_id_list(mm_output.get("ids.prompt")) + or list(llm_output.prompt_token_ids) + ) + llm_output_ids = getattr(output, "token_ids", None) + if llm_output_ids is None: + llm_output_ids = getattr(output, "cumulative_token_ids", []) + # Always copy: CompletionOutput.token_ids can alias the upstream + # detokenizer's live token list. Forwarding that exact object as the + # talker prompt makes the stage-1 streaming update extend the list + # with itself (state and update bind the same object), doubling the + # thinker's recorded output every segment until the TTS engine input + # buffer overflows. + llm_output_ids = list(llm_output_ids) + thinker_text = getattr(output, "text", "") or "" + if _has_native_duplex_prompt_metadata(mm_output): + # The thinker's resumable duplex request reports cumulative + # output ids/text, but earlier segments are already folded into + # the prompt by the scheduler session update, and the forwarded + # hidden states only cover the current prompt + current segment. + # Hand stage 1 exactly one segment per handoff so token/hidden + # alignment holds, the talker prompt grows linearly with new + # tokens instead of quadratically, and downstream transcripts + # carry per-unit deltas instead of re-sending the whole reply. + llm_output_ids, thinker_text = _native_duplex_segment_output_ids( + llm_output_ids, + thinker_text, + _streaming_context, + request_id=str(llm_output.request_id), + ) prompt_token_ids_len = len(prompt_token_ids) - latent = output.multimodal_output.get("latent", None) + latent = mm_output.get("latent", None) if latent is None: latent = output.hidden_states if hasattr(output, "hidden_states") else None if latent is None: raise ValueError("No latent or hidden_states found in thinker output") - thinker_hidden_states = latent.clone().detach() - - # Split hidden states: prompt portion has speaker embedding, - # generated portion has the text content - prompt_hidden = thinker_hidden_states[:prompt_token_ids_len].to(torch.float32) - - # Extract decoded text from thinker output for TTS text extraction - thinker_text = getattr(output, "text", "") or "" + thinker_hidden_states = latent.detach() + if thinker_hidden_states.ndim == 3 and thinker_hidden_states.shape[0] == 1: + thinker_hidden_states = thinker_hidden_states.squeeze(0) # Build full token sequence and extract TTS region - full_token_ids = list(prompt_token_ids) + ( - list(llm_output_ids) if not isinstance(llm_output_ids, list) else llm_output_ids - ) - full_hidden = thinker_hidden_states.to(torch.float32) + full_token_ids = prompt_token_ids + llm_output_ids - # Detect TTS token IDs (4.5: 151703/151704, 2.6: 151691/151692) - tts_bos_id, tts_eos_id = 151691, 151692 - for _id in [151703, 151704]: - if _id in full_token_ids: - tts_bos_id, tts_eos_id = 151703, 151704 - break + tts_bos_id = special_token_ids.get("tts_bos_token_id") + is_native_duplex_handoff = _has_native_duplex_prompt_metadata(mm_output) + if is_native_duplex_handoff: + _require_native_tts_boundary_metadata(special_token_ids) + tts_end_ids = _native_tts_boundary_token_ids(special_token_ids) - tts_bos_idx = tts_eos_idx = None - for idx_t, tid in enumerate(full_token_ids): - if tid == tts_bos_id: + # Plain-chat (use_tts_template) fallback: non-duplex requests do not + # surface special_token_ids, so use MiniCPM-o 4.5's fixed boundaries. + if tts_bos_id is None and not is_native_duplex_handoff: + tts_bos_id = 151703 + tts_end_ids = set(tts_end_ids) | {151704, 151645} + + tts_bos_idx = None + # For native duplex the resumable prompt folds every earlier unit, so + # a <|tts_bos|> from an already-spoken reply can sit mid-prompt; only + # a boundary folded as the FINAL prompt token (this unit's decision) + # or one inside the current segment may start the slice, or stale + # text would be re-handed to the talker on text-less continuations. + search_start = max(0, prompt_token_ids_len - 1) if is_native_duplex_handoff else 0 + for idx_t in range(search_start, len(full_token_ids)): + if full_token_ids[idx_t] == tts_bos_id: tts_bos_idx = idx_t + 1 - elif tid == tts_eos_id: - tts_eos_idx = idx_t + + tts_eos_idx = None + if tts_bos_idx is not None: + for idx_t in range(tts_bos_idx, len(full_token_ids)): + if full_token_ids[idx_t] in tts_end_ids: + tts_eos_idx = idx_t + break tts_token_ids_slice = tts_hidden_slice = None - if tts_bos_idx is not None and full_hidden.shape[0] > tts_bos_idx: - end_idx = tts_eos_idx if tts_eos_idx is not None else full_hidden.shape[0] + native_segment_end = False + if tts_bos_idx is not None and thinker_hidden_states.shape[0] > tts_bos_idx: + end_idx = tts_eos_idx if tts_eos_idx is not None else thinker_hidden_states.shape[0] + if is_native_duplex_handoff and tts_eos_idx is not None: + boundary_token = full_token_ids[tts_eos_idx] + native_segment_end = boundary_token in { + special_token_ids.get("chunk_eos_token_id"), + special_token_ids.get("chunk_tts_eos_token_id"), + } tts_token_ids_slice = torch.tensor(full_token_ids[tts_bos_idx:end_idx], dtype=torch.long) - tts_hidden_slice = full_hidden[tts_bos_idx:end_idx] + tts_hidden_slice = thinker_hidden_states[tts_bos_idx:end_idx].to(torch.float32).contiguous() + elif is_native_duplex_handoff: + # Official MiniCPM-o duplex does not prefill an assistant + # <|tts_bos|> boundary before generation. A segment delta can + # start with SEVERAL unit decisions (forced/model listens from + # chunks that produced no handoff accumulate ahead of the speak), + # so skip the leading listen run, then the speak decision itself; + # hidden states for TTS start after it and stop at the next + # chunk terminator. + listen_id = special_token_ids.get("listen_token_id") + speak_id = special_token_ids.get("speak_token_id") + out_ids = llm_output_ids + j = 0 + while j < len(out_ids) and out_ids[j] == listen_id: + j += 1 + if j < len(out_ids) and out_ids[j] == speak_id: + out_start = j + 1 + out_end = len(out_ids) + turn_eos_id = special_token_ids.get("turn_eos_token_id") + for idx_t in range(out_start, len(out_ids)): + token_id = out_ids[idx_t] + if turn_eos_id is not None and token_id == turn_eos_id: + # Native duplex trains the talker on the <|turn_eos|> + # embedding itself, but any text sampled after it + # belongs to a stale tail and must not enter TTS. + out_end = idx_t + 1 + break + if token_id in tts_end_ids: + out_end = idx_t + native_segment_end = token_id in { + special_token_ids.get("chunk_eos_token_id"), + special_token_ids.get("chunk_tts_eos_token_id"), + } + break + # Map output indices onto hidden rows by END alignment: the + # leading decision tokens of the delta may ALSO be folded into + # the resumable prompt (they belong to earlier non-forwarded + # segments), so prompt_len + delta over-counts them and + # front-aligned indexing truncates the slice. The hidden + # tensor's last len(out_ids) rows are the delta's rows. + hidden_base = int(thinker_hidden_states.shape[0]) - len(out_ids) + if hidden_base >= 0 and out_end > out_start: + tts_token_ids_slice = torch.tensor(out_ids[out_start:out_end], dtype=torch.long) + tts_hidden_slice = ( + thinker_hidden_states[hidden_base + out_start : hidden_base + out_end] + .to(torch.float32) + .contiguous() + ) + elif j < len(out_ids) and out_ids[j] not in tts_end_ids: + # HF streaming_generate does not require an explicit <|speak|> + # marker. If a unit starts directly with text, the first token + # is fed back into the LLM but is not included in + # total_hidden_in_unit; TTS starts from the following token. + out_start = j + 1 + out_end = len(out_ids) + turn_eos_id = special_token_ids.get("turn_eos_token_id") + for idx_t in range(out_start, len(out_ids)): + token_id = out_ids[idx_t] + if turn_eos_id is not None and token_id == turn_eos_id: + out_end = idx_t + 1 + break + if token_id in tts_end_ids: + out_end = idx_t + native_segment_end = token_id in { + special_token_ids.get("chunk_eos_token_id"), + special_token_ids.get("chunk_tts_eos_token_id"), + } + break + hidden_base = int(thinker_hidden_states.shape[0]) - len(out_ids) + if hidden_base >= 0 and out_end > out_start: + tts_token_ids_slice = torch.tensor(out_ids[out_start:out_end], dtype=torch.long) + tts_hidden_slice = ( + thinker_hidden_states[hidden_base + out_start : hidden_base + out_end] + .to(torch.float32) + .contiguous() + ) + model_intermediate_buffer = build_duplex_intermediate_buffer( + request_id=str(llm_output.request_id), + prompt_token_ids=prompt_token_ids, + output_token_ids=llm_output_ids, + output_text=thinker_text, + stream_output=is_native_duplex_handoff, + native_duplex=is_native_duplex_handoff, + ) + if is_native_duplex_handoff: + turn_eos_id = special_token_ids.get("turn_eos_token_id") + meta = model_intermediate_buffer.setdefault("meta", {}) + data_plane_metadata = _native_duplex_data_plane_metadata(_streaming_context) + if data_plane_metadata is not None: + model_intermediate_buffer["duplex"] = data_plane_metadata + meta["native_duplex_segment_text"] = thinker_text + meta.setdefault("override_keys", []).extend( + [ + "llm_output_text", + ["meta", "native_duplex_segment_text"], + ] + ) + if turn_eos_id is not None: + # The talker detects turn end from <|turn_eos|> inside the + # handed condition (official conditions on its embedding). + meta["turn_eos_token_id"] = int(turn_eos_id) + if native_segment_end: + meta["segment_end"] = True + req_mm_data = multi_modal_data.get(llm_output.request_id) + ref_audio = _extract_first_audio_ref(req_mm_data) + if ref_audio is None: + ref_audio = _extract_native_runtime_ref_audio( + model_intermediate_buffer.get("duplex"), + ) + if ref_audio is not None: + ref_waveform, ref_sr = ref_audio + set_ref_audio(model_intermediate_buffer, _to_transport_list(ref_waveform), ref_sr) + handoff_ids = _coerce_token_id_list(tts_token_ids_slice) if tts_token_ids_slice is not None else None + handoff_hidden = _to_transport_list(tts_hidden_slice) if tts_hidden_slice is not None else None + native_turn_end_handoff = False + if is_native_duplex_handoff: + turn_eos_id = special_token_ids.get("turn_eos_token_id") + native_turn_end_handoff = turn_eos_id is not None and handoff_ids is not None and turn_eos_id in handoff_ids + handoff_ids, handoff_hidden = _accumulate_native_tts_handoff( + _streaming_context, + handoff_ids, + handoff_hidden, + ) + if not handoff_ids: + continue + set_tts_handoff(model_intermediate_buffer, handoff_ids, handoff_hidden) + if native_turn_end_handoff: + _reset_native_tts_handoff(_streaming_context) - additional_information = { - "prompt_embeds": prompt_hidden, - "prompt_token_ids": list(prompt_token_ids), - "llm_output_token_ids": list(llm_output_ids) if not isinstance(llm_output_ids, list) else llm_output_ids, - "llm_output_text": [thinker_text], - } - if tts_token_ids_slice is not None: - additional_information["tts_token_ids"] = tts_token_ids_slice - if tts_hidden_slice is not None: - additional_information["tts_hidden_states"] = tts_hidden_slice - - # Minimal prompt token IDs: the talker's AR framework needs *some* tokens - # to do a single prefill step. We use [BOS, PAD, EOS] as a dummy. + scheduler_prompt_token_ids = _build_tts_scheduler_prompt_token_ids( + tts_token_ids_slice, + llm_output_ids, + prompt_token_ids, + ) tts_inputs.append( OmniTokensPrompt( - prompt_token_ids=[1, 0, 2], - additional_information=additional_information, + prompt_token_ids=scheduler_prompt_token_ids, + model_intermediate_buffer=model_intermediate_buffer, multi_modal_data=( multi_modal_data[llm_output.request_id] if requires_multimodal_data and multi_modal_data.get(llm_output.request_id) is not None @@ -132,5 +610,12 @@ def llm2tts( mm_processor_kwargs=None, ) ) + if native_turn_end_handoff: + bridge_states = getattr(_streaming_context, "bridge_states", None) + duplex_state = bridge_states.get("duplex") if isinstance(bridge_states, dict) else None + if isinstance(duplex_state, dict): + current_model_turn_id = duplex_state.get("model_turn_id", duplex_state.get("turn_id", 0)) + if isinstance(current_model_turn_id, int): + duplex_state["model_turn_id"] = current_model_turn_id + 1 return tts_inputs diff --git a/vllm_omni/outputs/multimodal_accumulation.py b/vllm_omni/outputs/multimodal_accumulation.py new file mode 100644 index 000000000..08b2551bc --- /dev/null +++ b/vllm_omni/outputs/multimodal_accumulation.py @@ -0,0 +1,91 @@ +from typing import Any + +import torch + +from vllm_omni.outputs.mm_outputs import MultimodalPayload +from vllm_omni.outputs.output_modality import DRAINABLE_MODALITIES + +CHUNK_METADATA_KEYS: frozenset[str] = frozenset( + { + "audio_text_total_chars", + "duplex_epoch", + "duplex_turn_id", + "llm_output_text_utf8", + "tts_is_last_chunk", + "turn_end", + } +) + + +def _is_chunk_metadata_key(key: str) -> bool: + if key.startswith("meta."): + return key.split(".", 1)[1] in CHUNK_METADATA_KEYS + return key in CHUNK_METADATA_KEYS + + +def replace_snapshot_keys( + accumulated: MultimodalPayload, + incoming: MultimodalPayload, +) -> None: + """Replace per-chunk metadata with the latest values.""" + for key in (*incoming.tensors, *incoming.metadata): + if _is_chunk_metadata_key(key): + accumulated.tensors.pop(key, None) + accumulated.metadata.pop(key, None) + + +def drain_delta_payload(payload: MultimodalPayload) -> None: + """Remove client-facing delta data while retaining request-level state.""" + for modality_key in DRAINABLE_MODALITIES: + key = str(modality_key) + payload.tensors.pop(key, None) + payload.metadata.pop(key, None) + + for metadata_key in CHUNK_METADATA_KEYS: + flat_key = f"meta.{metadata_key}" + payload.tensors.pop(flat_key, None) + payload.metadata.pop(flat_key, None) + + meta = payload.metadata.get("meta") + if isinstance(meta, dict): + filtered = {key: value for key, value in meta.items() if key not in CHUNK_METADATA_KEYS} + if filtered: + payload.metadata["meta"] = filtered + else: + payload.metadata.pop("meta", None) + + +def _payload_meta_value(payload: MultimodalPayload, key: str) -> Any: + flat_key = f"meta.{key}" + if flat_key in payload.tensors: + return payload.tensors[flat_key] + if flat_key in payload.metadata: + return payload.metadata[flat_key] + meta = payload.metadata.get("meta") + if isinstance(meta, dict): + return meta.get(key) + return None + + +def _last_scalar_int(value: Any) -> int | None: + if value is None: + return None + if isinstance(value, bool | int): + return int(value) + if isinstance(value, list | tuple): + return _last_scalar_int(value[-1]) if value else None + if isinstance(value, torch.Tensor): + try: + if value.numel() == 0: + return None + return int(value.detach().cpu().reshape(-1)[-1].item()) + except (RuntimeError, TypeError, ValueError): + return None + return None + + +def is_non_final_delta_audio_chunk(payload: MultimodalPayload, mm_type: str | None) -> bool: + """Return whether an audio delta explicitly declares more chunks.""" + if str(mm_type or "").lower() != "audio" and "audio" not in payload: + return False + return _last_scalar_int(_payload_meta_value(payload, "tts_is_last_chunk")) == 0 diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py index e39b995b0..ebfe0689a 100644 --- a/vllm_omni/outputs/output_processor.py +++ b/vllm_omni/outputs/output_processor.py @@ -19,15 +19,15 @@ from vllm.v1.metrics.stats import IterationStats, RequestStateStats from vllm_omni.data_entry_keys import unflatten_payload from vllm_omni.outputs import OmniRequestOutput from vllm_omni.outputs.mm_outputs import MultimodalCompletionOutput, MultimodalPayload -from vllm_omni.outputs.output_modality import ( - DRAINABLE_MODALITIES, - OutputModality, - get_accumulation_strategy, +from vllm_omni.outputs.multimodal_accumulation import ( + drain_delta_payload, + is_non_final_delta_audio_chunk, + replace_snapshot_keys, ) +from vllm_omni.outputs.output_modality import OutputModality, get_accumulation_strategy logger = init_logger(__name__) - # --------------------------------------------------------------------------- # Module-level helpers # --------------------------------------------------------------------------- @@ -55,6 +55,14 @@ def _modality_to_type_string(modality: OutputModality) -> str: return "text" +def _mean_time_per_output_token_ms(stats: RequestStateStats) -> float: + output_intervals = stats.num_generation_tokens - 1 + if output_intervals <= 0: + return 0.0 + decode_time_s = max(stats.last_token_ts - stats.first_token_ts, 0.0) + return decode_time_s * 1000.0 / float(output_intervals) + + class OmniRequestState(RequestState): """Request state for omni models, tracking multimodal outputs. @@ -85,7 +93,7 @@ class OmniRequestState(RequestState): def apply_streaming_update(self, update) -> None: super().apply_streaming_update(update) - self.native_text_stats.arrival_time = update.arrival_time + self.native_text_stats = RequestStateStats(arrival_time=float(update.arrival_time or 0.0)) def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None: """Accumulate a multimodal tensor payload into the request state. @@ -104,6 +112,7 @@ class OmniRequestState(RequestState): incoming = MultimodalPayload.from_raw(payload, modality_key) if incoming is not None: + replace_snapshot_keys(self.mm_accumulated, incoming) self.mm_accumulated = self.mm_accumulated.merged_with(incoming) except (ValueError, TypeError, RuntimeError): logger.exception("Error accumulating multimodal tensor") @@ -170,9 +179,13 @@ class OmniRequestState(RequestState): kv_transfer_params, ) + is_delta = self.output_kind == RequestOutputKind.DELTA + if is_delta and finish_reason is not None and is_non_final_delta_audio_chunk(self.mm_accumulated, self.mm_type): + finish_reason = None + stop_reason = None + finished = finish_reason is not None final_only = self.output_kind == RequestOutputKind.FINAL_ONLY - is_delta = self.output_kind == RequestOutputKind.DELTA if not finished and final_only: return None @@ -280,10 +293,10 @@ class OmniRequestState(RequestState): output.cumulative_text = base_output.cumulative_text # DELTA mode: drain modality keys (e.g. audio) so the next - # step only sees freshly accumulated data for those keys. + # step only sees freshly accumulated data for those keys and + # their client-facing per-chunk metadata. if self.output_kind == RequestOutputKind.DELTA: - for modality_key in DRAINABLE_MODALITIES: - self.mm_accumulated.tensors.pop(modality_key, None) + drain_delta_payload(self.mm_accumulated) return output except (RuntimeError, TypeError, AttributeError): @@ -391,6 +404,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): return self._native_text_metrics_by_request.setdefault( request_id, { + "num_generation_tokens": 0, "vllm_ttft_ms": 0.0, "vllm_tpot_ms": 0.0, "vllm_itl_ms": 0.0, @@ -592,6 +606,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): req_state.lora_name, ) record = self._native_text_metric_record(req_state.external_req_id) + record["num_generation_tokens"] = int(native_stats.num_generation_tokens) if was_prefilling: record["vllm_ttft_ms"] = max(float(native_stats.first_token_latency) * 1000.0, 0.0) return @@ -601,6 +616,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): itls_ms = record.setdefault("vllm_itls_ms", []) itls_ms.append(itl_ms) record["vllm_itl_ms"] = sum(itls_ms) / float(len(itls_ms)) + record["vllm_tpot_ms"] = _mean_time_per_output_token_ms(native_stats) def _update_stats_from_finished( self, diff --git a/vllm_omni/request.py b/vllm_omni/request.py index b094f2d39..2f1fd5379 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -35,6 +35,7 @@ class OmniRequest(Request): # Optional external request ID for tracking external_req_id: str | None = None, additional_information: AdditionalInformationPayload | None = None, + model_intermediate_buffer: dict | None = None, **kwargs, ): if prompt_embeds is not None: @@ -48,6 +49,8 @@ class OmniRequest(Request): self.external_req_id: str | None = external_req_id # Serialized additional information payload (optional) self.additional_information: AdditionalInformationPayload | None = additional_information + # Runner-owned runtime payload. + self.model_intermediate_buffer: dict | None = model_intermediate_buffer @staticmethod def _maybe_decode_prompt_embeds( @@ -94,6 +97,7 @@ class OmniRequest(Request): trace_headers=request.trace_headers, block_hasher=block_hasher, additional_information=request.additional_information, + model_intermediate_buffer=getattr(request, "model_intermediate_buffer", None), resumable=request.resumable, reasoning_ended=request.reasoning_ended, reasoning_parser_kwargs=request.reasoning_parser_kwargs, @@ -117,6 +121,7 @@ class OmniStreamingUpdate: arrival_time: float sampling_params: SamplingParams | None additional_information: AdditionalInformationPayload | None = None + model_intermediate_buffer: dict | None = None @classmethod def from_request(cls, request: "Request") -> "OmniStreamingUpdate | None": @@ -129,4 +134,5 @@ class OmniStreamingUpdate: arrival_time=request.arrival_time, sampling_params=request.sampling_params, additional_information=request.additional_information, + model_intermediate_buffer=getattr(request, "model_intermediate_buffer", None), ) diff --git a/vllm_omni/utils/mm_outputs.py b/vllm_omni/utils/mm_outputs.py index 189cc00aa..1dc4e37f7 100644 --- a/vllm_omni/utils/mm_outputs.py +++ b/vllm_omni/utils/mm_outputs.py @@ -26,6 +26,17 @@ _CLIENT_MM_ROOT_KEYS: frozenset[str] = frozenset( } ) +_CLIENT_MM_META_KEYS: frozenset[str] = frozenset( + { + "audio_text_total_chars", + "duplex_epoch", + "duplex_turn_id", + "llm_output_text_utf8", + "tts_is_last_chunk", + "turn_end", + } +) + def partition_flat_payload( payload: Mapping[str, object], @@ -39,6 +50,13 @@ def partition_flat_payload( root = key.split(".", 1)[0] if root in _CLIENT_MM_ROOT_KEYS: client_mm[key] = value + elif root == "meta" and "." in key and key.split(".", 1)[1] in _CLIENT_MM_META_KEYS: + # Small final-output metadata needed by serving (for example + # transcript text attached to audio) must ride with client MM + # output. Keep it in inter-stage too so downstream stages that read + # metadata from the full payload are not starved. + inter_stage[key] = value + client_mm[key] = value else: inter_stage[key] = value return inter_stage, client_mm diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 67ea34b10..3e1ea6ddb 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -327,6 +327,24 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): kv_transfer_manager=self.kv_transfer_manager, ) self._downstream_payload_cache: dict[str, bool] = {} + self._duplex_sampling_hook = None + self._duplex_sampling_hook_resolved = False + + def load_model(self, *args, **kwargs) -> None: + super().load_model(*args, **kwargs) + self._resolve_duplex_sampling_hook(force=True) + + def _resolve_duplex_sampling_hook(self, *, force: bool = False): + if not force and getattr(self, "_duplex_sampling_hook_resolved", False): + return self._duplex_sampling_hook + candidate = getattr(getattr(self, "model", None), "prepare_duplex_sampling", None) + self._duplex_sampling_hook = candidate if callable(candidate) else None + self._duplex_sampling_hook_resolved = True + if self._duplex_sampling_hook is not None and not hasattr(self, "_duplex_sampling_helper"): + from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingHelper + + self._duplex_sampling_helper = DuplexSamplingHelper() + return self._duplex_sampling_hook def _make_buffer(self, *size, dtype, numpy=True): # Prevent ray from pinning the buffer due to large size @@ -390,9 +408,18 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): if getattr(self.model, "skips_model_sampler_output_token_history", False): return sampling_metadata output_token_ids = self._build_model_sampler_output_token_ids() - if output_token_ids == sampling_metadata.output_token_ids: - return sampling_metadata - return replace(sampling_metadata, output_token_ids=output_token_ids) + if output_token_ids != sampling_metadata.output_token_ids: + return replace(sampling_metadata, output_token_ids=output_token_ids) + return sampling_metadata + + def _update_states(self, scheduler_output: SchedulerOutput) -> Callable | None: + deferred_state_corrections_fn = super()._update_states(scheduler_output) + if self._resolve_duplex_sampling_hook() is None: + return deferred_state_corrections_fn + helper = getattr(self, "_duplex_sampling_helper", None) + if helper is not None: + helper.update_states(self, scheduler_output) + return deferred_state_corrections_fn def _request_final_stage_id(self, req_id: str) -> int | None: info = self.model_intermediate_buffer.get(req_id) @@ -519,6 +546,9 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): self._downstream_payload_cache.clear() if hasattr(self, "model_intermediate_buffer"): self.model_intermediate_buffer.clear() + duplex_helper = getattr(self, "_duplex_sampling_helper", None) + if duplex_helper is not None: + duplex_helper.clear() # 5. Release all CUDA graphs unconditionally (upstream only does this # on ROCm; on CUDA the graphs are only freed by Python GC during @@ -1430,10 +1460,16 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): self.input_batch.idx_mapping_np, self.input_batch.positions[self.input_batch.logits_indices], ) - sampler_output = model_sample( - logits, - self._sampling_metadata_for_model_sampler(sampling_metadata), - ) + prepared_sampling_metadata = self._sampling_metadata_for_model_sampler(sampling_metadata) + prepare_duplex_sampling = self._resolve_duplex_sampling_hook() + if prepare_duplex_sampling is not None: + helper = getattr(self, "_duplex_sampling_helper", None) + rows = helper.rows(self) if helper is not None and helper.active_request_ids else () + if rows or (helper is not None and helper.hook_active): + prepare_duplex_sampling(logits, prepared_sampling_metadata, rows) + if helper is not None: + helper.hook_active = bool(rows) + sampler_output = model_sample(logits, prepared_sampling_metadata) if sampler_output is not None: return sampler_output return self.sampler( diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index b816a36f7..0e14784ca 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -537,7 +537,7 @@ class OmniGPUModelRunner(GPUModelRunner): for new_req_data in scheduler_output.scheduled_new_reqs: req_id = new_req_data.req_id if req_id in self.requests: - self._update_streaming_input_additional_info(req_id) + self._update_streaming_input_additional_info(new_req_data, req_id) req_state = self._update_streaming_request(req_id, new_req_data) reqs_to_add.append(req_state) continue @@ -598,6 +598,15 @@ class OmniGPUModelRunner(GPUModelRunner): pass except Exception as e: logger.error(f"Error decoding prompt embeds: {e}") + # Direct runner data-plane payloads populate + # model_intermediate_buffer without going through the deprecated + # additional_information request transport. + try: + model_buffer = getattr(new_req_data, "model_intermediate_buffer", None) + if isinstance(model_buffer, dict) and model_buffer: + self._update_intermediate_buffer(req_id, model_buffer) + except Exception as e: + logger.error(f"Error updating model intermediate buffer: {e}") # Decode additional_information payloads (dictionary) try: if getattr(new_req_data, "additional_information", None) is not None: @@ -1260,6 +1269,9 @@ class OmniGPUModelRunner(GPUModelRunner): pe_cpu = self._resolve_prompt_embeds_cpu(getattr(nr, "prompt_embeds", None)) if pe_cpu is not None: setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) + model_buffer = getattr(nr, "model_intermediate_buffer", None) + if isinstance(model_buffer, dict) and model_buffer: + self._update_intermediate_buffer(req_id, model_buffer) info_payload = getattr(nr, "additional_information", None) if info_payload is not None: logger.warning_once( @@ -1466,6 +1478,9 @@ class OmniGPUModelRunner(GPUModelRunner): def _update_additional_information(self, scheduler_output: "SchedulerOutput") -> None: for new_req in scheduler_output.scheduled_new_reqs: + model_buffer = getattr(new_req, "model_intermediate_buffer", None) + if isinstance(model_buffer, dict) and model_buffer: + self._update_intermediate_buffer(new_req.req_id, model_buffer) payload_info = getattr(new_req, "additional_information", None) if isinstance(payload_info, dict): self._update_intermediate_buffer(new_req.req_id, payload_info) @@ -1704,6 +1719,7 @@ class OmniGPUModelRunner(GPUModelRunner): decode_start_offsets.extend(start_offsets_b) decode_batch_items.clear() + preprocess_input_ids = input_ids if input_ids is not None else self.input_ids.gpu[:num_input_tokens] for req_index, req_id in enumerate(self.input_batch.req_ids): req_infos = self.model_intermediate_buffer.get(req_id, {}) @@ -1718,6 +1734,8 @@ class OmniGPUModelRunner(GPUModelRunner): # call the custom process function req_infos["request_id"] = req_id + req_infos["duplex_token_offset"] = int(self.input_batch.num_computed_tokens_cpu[req_index]) + req_infos["duplex_prompt_len"] = len(req_state.prompt_token_ids) if req_state is not None else None prompt_token_ids = getattr(req_state, "prompt_token_ids", ()) if req_state is not None else () prompt_len = len(prompt_token_ids or ()) num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index]) @@ -1733,14 +1751,15 @@ class OmniGPUModelRunner(GPUModelRunner): embed_slice = inputs_embeds[s:e] if inputs_embeds is not None else None req_input_ids, req_embeds, update_dict = self.model.preprocess( - input_ids=input_ids[s:e], input_embeds=embed_slice, **req_infos + input_ids=preprocess_input_ids[s:e], input_embeds=embed_slice, **req_infos ) if inputs_embeds is None: inputs_embeds = torch.empty( - (input_ids.shape[0], req_embeds.shape[-1]), + (preprocess_input_ids.shape[0], req_embeds.shape[-1]), device=req_embeds.device, dtype=req_embeds.dtype, ) + input_ids = preprocess_input_ids if self.has_talker_mtp and span_len == 1 and not is_prefill: last_talker_hidden, text_step = update_dict.pop("mtp_inputs") @@ -1759,7 +1778,9 @@ class OmniGPUModelRunner(GPUModelRunner): seg_len = min(span_len, req_embeds.shape[0]) inputs_embeds[s : s + seg_len] = req_embeds[:seg_len] if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len: - input_ids[s : s + seg_len] = req_input_ids + preprocess_input_ids[s : s + seg_len] = req_input_ids + if input_ids is None: + input_ids = preprocess_input_ids flush_decode_batch() @@ -1977,12 +1998,26 @@ class OmniGPUModelRunner(GPUModelRunner): # Backward compatible: mirror to old setattr location setattr(req_state, "additional_information_cpu", existing) - def _update_streaming_input_additional_info(self, req_id): - # For streaming input prefill case only. Set num processed tokens = 0 for new segment input + def _update_streaming_input_additional_info(self, new_req_data, req_id): + # For streaming input prefill case only. Update buffer from last segment input. cached_additional_info = self.model_intermediate_buffer.get(req_id, {}) - if cached_additional_info: - merged_info = dict(cached_additional_info) - merged_info.setdefault("meta", {})["num_processed_tokens"] = 0 - merged_info.setdefault("meta", {})["resumable"] = True - self.model_intermediate_buffer[req_id] = merged_info - setattr(self.requests[req_id], "additional_information_cpu", merged_info) + inc_info = getattr(new_req_data, "model_intermediate_buffer", None) + if not isinstance(inc_info, dict) or not inc_info: + payload_info = getattr(new_req_data, "additional_information", None) + inc_info = deserialize_additional_information(payload_info) + if not isinstance(inc_info, dict) or not inc_info: + return + merged_info = dict(cached_additional_info) if isinstance(cached_additional_info, dict) else {} + for key, value in inc_info.items(): + if isinstance(value, dict): + existing_sub = merged_info.get(key) + merged_sub = dict(existing_sub) if isinstance(existing_sub, dict) else {} + for sk, sv in value.items(): + merged_sub[sk] = sv + merged_info[key] = merged_sub + else: + merged_info[key] = value + merged_info.setdefault("meta", {})["num_processed_tokens"] = 0 + merged_info.setdefault("meta", {})["resumable"] = True + self.model_intermediate_buffer[req_id] = merged_info + setattr(self.requests[req_id], "additional_information_cpu", merged_info)