feat(server): /v1/realtime WebSocket endpoint for Nemotron VoiceChat sessions (PCM16 base64, one active session, dedicated worker)
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Summary
Expose the streaming session from #1378 over a WebSocket at `/v1/realtime` on `mlxcel-server`: JSON control messages in, base64 PCM16 audio both ways, one active session at a time, all MLX work on one dedicated thread. Add a small microphone/speaker client so the loop can be exercised end to end from a terminal.
## Current behavior
- `src/server/app.rs:104-113` registers the audio HTTP routes (`/v1/audio/speech`, `/v1/audio/transcriptions`, `/v1/audio/translations`) as request/response handlers; there is no WebSocket route anywhere and `Cargo.toml:219` enables axum with `["json", "macros", "multipart"]` only (no `ws` feature).
- `src/server/startup.rs:2141-2160` wires a Whisper checkpoint into `AudioModelProvider` through `WhisperSttProvider::load` on its own worker thread; `src/server/audio_worker.rs:139-215` (`AudioWorker::spawn`, `transcribe`, `synthesize`) serves one command at a time over a bounded channel and documents the thread-affinity rule. Neither has a notion of a long-lived per-connection session.
- `src/server/model_provider.rs:130-138` (`uses_single_stream_queue_admission`) lists the model types that bypass the batch scheduler; a VoiceChat checkpoint must not start the decoder-only scheduler either.
## Expected behavior
Server start: `./target/release/mlxcel-server -m models/NemotronLabs-VoiceChat-11B-4bit --port 8080`. The startup path detects `ModelType::NemotronVoiceChat`, loads the model on a dedicated realtime worker thread (not the batch scheduler, not the audio worker), skips the text warmup, and registers `/v1/realtime`. Chat/completions requests against this checkpoint return `501` with a message naming `/v1/realtime`.
### Protocol (client -> server, JSON text frames)
| type | fields | behavior |
|---|---|---|
| `session.update` | `session.model` (optional when the server has one model), `session.system_prompt`, `session.seed` (int, default 0), `session.max_streaming_seconds` (float, optional) | creates the streaming session; answered with `session.updated`; a second `session.update` on a configured session is an error |
| `input_audio_buffer.append` | `audio` (base64 little-endian PCM16 mono), `sample_rate` (must be 16000) | decoded to f32 in `[-1, 1)` (`/ 32768`), pushed in 1280-sample slices so events flow while a large chunk is still being processed |
| `input_audio_buffer.commit` | `pad_partial` (bool, default true) | answered with `input_audio_buffer.committed`, then the session is flushed and the connection loop ends after `response.done` |
| `session.cancel` / `response.cancel` | | cancels; emits `response.cancelled`; loop ends |
| `session.ping` | | `session.pong` |
Any other `type` yields an `error` event with `code: "invalid_request"`; a non-JSON message yields `error` and the loop continues. Audio before `session.update` yields `error` ("send session.update before audio").
### Protocol (server -> client)
Every event carries `event_id` (`event_<16 hex>`). On accept: `session.created` with `session.id` (`sess_<16 hex>`), `state: "configuring"`, `input_audio_format {type: "pcm16", sample_rate: 16000}`, `output_audio_format {type: "pcm16", sample_rate: 22050}`. After configuration: `session.updated` with `state: "ready"`, `model`, `frame_samples: 1280`, both formats. Model events map 1:1 from the session events:
| session event | wire `type` | fields |
|---|---|---|
| AssistantTextDelta | `response.text.delta` | `frame_index`, `token_id`, `delta`, `text` |
| FunctionDelta | `response.function.delta` | `frame_index`, `token_id`, `delta`, `text` |
| UserTranscriptDelta | `conversation.item.input_audio_transcription.delta` | `frame_index`, `delta`, `transcript` |
| Audio | `response.audio.delta` | `frame_index`, `delta` (base64 PCM16 of `round(clip(x, -1, 1) * 32767)`), `format: "pcm16"`, `sample_rate: 22050`, `channels: 1`, `audio_codes` (31 ints) |
| Done | `response.done` | `frame_index` |
| Cancelled | `response.cancelled` | `frame_index` |
Errors: `{"type": "error", "error": {"code": <"invalid_request" | "server_busy" | "session_initialization_failed" | "inference_error">, "message": ...}}`. When a session is already active, a new connection gets `server_busy` and is closed with WebSocket code 1013.
### Concurrency
One realtime engine per process owns one thread; a reservation (`try_reserve(session_id)`) admits exactly one connection. Commands (`open`, `push`, `flush`, `cancel`, `close`) are sent over a channel and answered on a one-shot reply channel; the WebSocket task awaits them with `spawn_blocking` / `tokio::sync::oneshot` so the async runtime never runs MLX code. On disconnect the task sends `close` (which cancels an unflushed session) and releases the reservation. The MLX buffer cache is not cleared per frame (clearing on every push raises per-frame latency); it is cleared on `close`.
## Implementation plan
1. **Dependencies.** Enable `axum` feature `ws` (`Cargo.toml:219`); `base64` (`Cargo.toml:181`) and `uuid` with `v4` (`Cargo.toml:223`) are already dependencies and cover the event and session ids; `tokio-tungstenite` and `cpal` are new and needed only by the example client (step 6).
2. **Engine (`src/server/realtime_engine.rs`).** `RealtimeVoiceChatEngine::spawn(model_path) -> Self` starts the worker thread (install the thread-local default stream as `AudioWorker::spawn` does at `src/server/audio_worker.rs:139`), loads `NemotronVoiceChatModel` on it, and loops on `RealtimeCommand { kind: Open{system_prompt, seed, max_streaming_seconds} | Push{samples, sample_rate} | Flush{pad_partial} | Cancel | Close, session_id, reply }`. `open` creates `VoiceChatStreamingSession` (errors if one is active), `push` / `flush` / `cancel` forward to it and return `Vec`, `close` drops it. Public blocking methods `open`, `push`, `flush`, `cancel`, `close`, plus `try_reserve` / `release` / `is_reserved`. Unit tests with a fake session trait object: `second_reserve_fails_until_release`, `push_on_unopened_session_errors`, `close_cancels_unflushed_session`.
3. **Wire types (`src/server/realtime_protocol.rs`).** `serde` enums for the client messages and the server events above; `pcm16_from_base64(&str) -> Result>` (reject odd byte counts and invalid base64), `audio_to_base64(&[f32]) -> String`, `fn serialize_event(VoiceChatEvent) -> WireEvent`. Tests: `pcm16_round_trip`, `odd_byte_count_rejected`, `event_mapping_names_and_fields`.
4. **Route (`src/server/routes/realtime.rs`).** `async fn realtime_ws(ws: WebSocketUpgrade, State(app)) -> Response` implementing the loop above; registered in `src/server/app.rs` as `.route("/v1/realtime", get(routes::realtime_ws))` only when `AppState.realtime_engine` is `Some`. Keep the `--api-key` auth layer in front of it (the upgrade request carries the header). `input_audio_buffer.append` slices the decoded audio into 1280-sample pieces and awaits one engine call per piece, forwarding that piece's events before the next call.
5. **Startup (`src/server/startup.rs`).** Next to the Whisper branch (line 2141): for `ModelType::NemotronVoiceChat`, spawn the realtime engine, store it in `AppState`, skip the chat `ModelProvider` load and the text warmup (extend the `is_florence2`-style guard at line 1695), and add the type to `uses_single_stream_queue_admission` so the scheduler never starts. `/v1/chat/completions` and `/v1/completions` answer `501 Not Implemented` with a pointer to `/v1/realtime` (reuse the `AudioModelError::KindNotLoaded` mapping pattern).
6. **Client example (`examples/voicechat_microphone.rs` or a `mlxcel voicechat` CLI subcommand).** Connect with `tokio-tungstenite`, send `session.update`, capture the default input device with `cpal` at 16 kHz mono PCM16 (resample when the device does not offer 16 kHz), send `input_audio_buffer.append` every 80 ms, play `response.audio.delta` through a 22.05 kHz output stream via a shared ring buffer, print text deltas, and send `input_audio_buffer.commit` on Ctrl-C. `--list-devices`, `--input-device`, `--output-device` flags. Document that there is no echo cancellation (headphones recommended). If adding `cpal` is not acceptable for the main crate, ship the example behind an `examples`-only dependency.
7. **Docs.** `docs/nemotron-voicechat.md`: the protocol tables above, a `websocat` transcript, and the single-session limit; README server section lists `/v1/realtime`.
## Validation
(a) Unit tests from steps 2 and 3, plus an integration test `tests/realtime_ws.rs` using a fake engine: `session_created_then_updated`, `append_before_update_errors`, `append_streams_events_per_1280_sample_slice`, `commit_emits_committed_then_done`, `second_connection_gets_server_busy_1013`, `cancel_emits_cancelled_and_releases`, `wrong_sample_rate_is_inference_error`.
(b) Real checkpoint: `mlx-community/NemotronLabs-VoiceChat-11B-4bit`.
```
./target/release/mlxcel-server -m models/NemotronLabs-VoiceChat-11B-4bit --port 8080 &
# drive a WAV through the socket with any small WebSocket client script or the Rust example:
cargo run --example voicechat_file_client -- ws://127.0.0.1:8080/v1/realtime /tmp/question.wav /tmp/response_ws.wav \
--system-prompt "Be concise and answer in one sentence."
```
Acceptance: the client receives `session.created`, `session.updated`, a transcript delta containing the spoken question, text deltas whose cumulative text answers it, `response.audio.delta` events with 1764-sample payloads and increasing `frame_index`, and `response.done`; the concatenated audio equals the `--stream` CLI output for the same WAV and prompt byte for byte (same seed); a second `websocat` connection during the session receives `server_busy` and close code 1013; after the first closes, the second connects normally; the microphone example holds a live exchange with headphones.
## Acceptance criteria
- [ ] `/v1/realtime` WebSocket registered only for VoiceChat checkpoints; chat endpoints answer 501 for this model type
- [ ] Protocol messages and events implemented exactly as tabulated, including error codes and the 1013 busy close
- [ ] One active session; reservation released on commit, cancel, or disconnect
- [ ] All MLX work on the engine thread; the async task only awaits channel replies
- [ ] WAV-through-socket output matches the CLI `--stream` output; microphone example works
- [ ] docs/supported-models.md updated
- [ ] cargo test --workspace --profile test-fast --features metal,accelerate passes
- [ ] cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass
## Out of scope
- Multiple concurrent sessions, session resumption, or serving chat and realtime from one process.
- Opus / WebRTC transports; PCM16 over WebSocket only.
## Dependencies
Blocked by #1378
- Part of epic #1372
Contributor guide
Research direction
Start with the blocked dependency #1378, then read src/server/audio_worker.rs, src/server/app.rs, src/server/startup.rs, and the proposed realtime_engine.rs, realtime_protocol.rs, and routes/realtime.rs boundaries. Use the listed unit and integration tests to drive protocol, reservation, streaming, and cancellation behavior. Done means the VoiceChat checkpoint exposes the specified WebSocket, passes the validation commands, and matches the CLI stream output.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- api, audio-video-rtc, backend
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100