zai_rs/realtime/
session.rs

1//! [`RealtimeSession`] — an active realtime conversation over WebSocket.
2//!
3//! A session owns a background event-loop task that pumps client events onto
4//! the socket and fans server events (and decoded audio) out to subscribers.
5//! Callers drive it via command methods (`send_audio`, `send_text`, …) and
6//! consume the two streams: [`RealtimeSession::events`] and
7//! [`RealtimeSession::audio_stream`].
8
9use std::{pin::Pin, sync::Arc};
10
11use bytes::Bytes;
12use futures::{Stream, StreamExt};
13use tokio::{
14    sync::{broadcast, mpsc},
15    task::JoinHandle,
16};
17use tokio_stream::wrappers::BroadcastStream;
18use tracing::{debug, warn};
19
20use super::{
21    audio::{OutputAudioFormat, decode_base64, encode_wav_pcm_base64},
22    client::AuthMode,
23    events::{ClientEvent, ServerEvent},
24    jwt,
25    protocol::{ChatMode, RealtimeTool, SessionConfig, TurnDetectionType},
26    transport::{RealtimeTransport, WsMessage},
27};
28use crate::{ZaiResult, client::error::RealtimeErrorKind};
29
30/// Commands queued onto the event-loop task. `ClientEvent` is boxed to keep
31/// the enum small (the event payload is ~224 bytes vs. a tag for `Close`).
32enum Command {
33    /// Serialize + send this client event.
34    ClientEvent(Box<ClientEvent>),
35    /// Tear the session down.
36    Close,
37}
38
39/// Builder for an [`RealtimeSession`].
40///
41/// Produced by [`super::client::RealtimeClient::session`]. Configure the
42/// session defaults, then [`SessionBuilder::build`] opens the WebSocket and
43/// sends the initial `session.update`.
44pub struct SessionBuilder {
45    api_key: Arc<String>,
46    auth: AuthMode,
47    realtime_url: String,
48    model_name: String,
49    session_config: SessionConfig,
50}
51
52impl SessionBuilder {
53    pub(super) fn new(
54        api_key: Arc<String>,
55        auth: AuthMode,
56        realtime_url: String,
57        model_name: String,
58    ) -> Self {
59        Self {
60            api_key,
61            auth,
62            realtime_url,
63            model_name,
64            session_config: SessionConfig::default(),
65        }
66    }
67
68    /// System instructions guiding the model.
69    pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
70        self.session_config.instructions = Some(instructions.into());
71        self
72    }
73
74    /// VAD strategy (defaults to client-VAD).
75    pub fn turn_detection(mut self, vad: TurnDetectionType) -> Self {
76        self.session_config.turn_detection.type_ = vad;
77        self
78    }
79
80    /// Output audio format (defaults to PCM).
81    pub fn output_audio_format(mut self, format: OutputAudioFormat) -> Self {
82        self.session_config.output_audio_format = format;
83        self
84    }
85
86    /// Conversation mode under `beta_fields.chat_mode`.
87    pub fn chat_mode(mut self, mode: ChatMode) -> Self {
88        self.session_config.beta_fields.chat_mode = Some(mode);
89        self
90    }
91
92    /// Enable/disable the server-side built-in web search.
93    pub fn auto_search(mut self, enabled: bool) -> Self {
94        self.session_config.beta_fields.auto_search = Some(enabled);
95        self
96    }
97
98    /// Register function tools.
99    pub fn tools(mut self, tools: Vec<RealtimeTool>) -> Self {
100        self.session_config.tools = tools;
101        self
102    }
103
104    /// Override the entire session config.
105    pub fn session_config(mut self, config: SessionConfig) -> Self {
106        self.session_config = config;
107        self
108    }
109
110    /// Open the WebSocket, send `session.update`, and spawn the event loop.
111    #[tracing::instrument(name = "realtime.session.build", skip_all, fields(model = %self.model_name))]
112    pub async fn build(self) -> ZaiResult<RealtimeSession> {
113        let Self {
114            api_key,
115            auth,
116            realtime_url,
117            model_name,
118            session_config,
119        } = self;
120
121        let jwt_ttl = match auth {
122            AuthMode::Bearer => None,
123            AuthMode::Jwt { ttl_seconds } => Some(ttl_seconds),
124        };
125        let authorization = jwt::authorization_header(&api_key, jwt_ttl)?;
126
127        let mut transport =
128            super::transport::TungsteniteTransport::connect(&realtime_url, &authorization).await?;
129
130        // Initial session.update with the negotiated defaults.
131        let init = ClientEvent::SessionUpdate {
132            event_id: Some(new_event_id()),
133            session: session_config,
134        };
135        transport.send(serde_json::to_string(&init)?).await?;
136        debug!(model = %model_name, "Realtime session opened");
137
138        let (cmd_tx, cmd_rx) = mpsc::channel::<Command>(64);
139        let (events_tx, _) = broadcast::channel::<ServerEvent>(256);
140        let (audio_tx, _) = broadcast::channel::<Bytes>(256);
141
142        let join = tokio::spawn(run_loop(
143            transport,
144            cmd_rx,
145            events_tx.clone(),
146            audio_tx.clone(),
147        ));
148
149        Ok(RealtimeSession {
150            cmd_tx,
151            events_tx,
152            audio_tx,
153            model_name,
154            join,
155        })
156    }
157}
158
159/// Background event-loop body: drains commands onto the socket and fans server
160/// messages out to the broadcast channels. Generic over the transport so a mock
161/// can be substituted in tests.
162async fn run_loop<T: RealtimeTransport>(
163    mut transport: T,
164    mut cmd_rx: mpsc::Receiver<Command>,
165    events_tx: broadcast::Sender<ServerEvent>,
166    audio_tx: broadcast::Sender<Bytes>,
167) {
168    loop {
169        tokio::select! {
170            biased;
171            cmd = cmd_rx.recv() => match cmd {
172                Some(Command::ClientEvent(ev)) => {
173                    let json = match serde_json::to_string(&ev) {
174                        Ok(s) => s,
175                        // Drop a malformed event but keep the session alive.
176                        Err(_) => continue,
177                    };
178                    if transport.send(json).await.is_err() {
179                        break;
180                    }
181                },
182                Some(Command::Close) | None => {
183                    let _ = transport.close().await;
184                    debug!("Realtime session closed (client requested)");
185                    break;
186                },
187            },
188            msg = transport.recv() => match msg {
189                Ok(Some(WsMessage::Text(text))) => match serde_json::from_str::<ServerEvent>(&text) {
190                    Ok(ServerEvent::ResponseAudioDelta { delta, .. }) => {
191                        if let Ok(bytes) = decode_base64(&delta) {
192                            let _ = audio_tx.send(Bytes::from(bytes));
193                        }
194                    },
195                    Ok(ServerEvent::Error { error }) => {
196                        warn!(
197                            code = ?error.code,
198                            message = ?error.message,
199                            "Realtime server error event"
200                        );
201                        let _ = events_tx.send(ServerEvent::Error { error });
202                    },
203                    Ok(event) => {
204                        let _ = events_tx.send(event);
205                    },
206                    // Ignore unparseable/unknown frames; the session stays open.
207                    Err(_) => {
208                        warn!(
209                            bytes = text.len(),
210                            "Dropping unparseable realtime frame"
211                        );
212                    },
213                },
214                Ok(Some(WsMessage::Binary(_))) => {},
215                Ok(None) => {
216                    // peer closed cleanly
217                    debug!("Realtime session closed (peer disconnected)");
218                    break;
219                },
220                Err(e) => {
221                    warn!(error = %e, "Realtime event loop terminated due to transport error");
222                    break;
223                },
224            },
225        }
226    }
227}
228
229/// An active realtime session.
230///
231/// Cheap to share indirectly via the channels it owns; call
232/// [`RealtimeSession::close`] to terminate the background task.
233pub struct RealtimeSession {
234    cmd_tx: mpsc::Sender<Command>,
235    events_tx: broadcast::Sender<ServerEvent>,
236    audio_tx: broadcast::Sender<Bytes>,
237    model_name: String,
238    join: JoinHandle<()>,
239}
240
241impl RealtimeSession {
242    /// Send raw 16-bit LE mono PCM (16 kHz) audio; it is wrapped in a WAV
243    /// header, base64-encoded, and uploaded via `input_audio_buffer.append`.
244    pub async fn send_audio(&self, pcm: Bytes) -> ZaiResult<()> {
245        let audio = encode_wav_pcm_base64(&pcm, 16_000);
246        self.dispatch(ClientEvent::InputAudioBufferAppend {
247            audio,
248            client_timestamp: Some(now_ms()),
249        })
250        .await
251    }
252
253    /// Commit buffered audio for inference (client-VAD). Server-VAD commits
254    /// automatically; calling this is harmless.
255    pub async fn commit_audio(&self) -> ZaiResult<()> {
256        self.dispatch(ClientEvent::InputAudioBufferCommit {
257            client_timestamp: Some(now_ms()),
258        })
259        .await
260    }
261
262    /// Inject a user text message into the conversation history.
263    pub async fn send_text(&self, text: impl Into<String>) -> ZaiResult<()> {
264        self.dispatch(ClientEvent::ConversationItemCreate {
265            event_id: Some(new_event_id()),
266            item: super::protocol::RealtimeConversationItem::user_text(text),
267        })
268        .await
269    }
270
271    /// Feed back a function-call result.
272    pub async fn send_function_output(
273        &self,
274        call_name: impl Into<String>,
275        output: impl Into<String>,
276    ) -> ZaiResult<()> {
277        self.dispatch(ClientEvent::ConversationItemCreate {
278            event_id: Some(new_event_id()),
279            item: super::protocol::RealtimeConversationItem::function_output(call_name, output),
280        })
281        .await
282    }
283
284    /// Trigger model inference (`response.create`).
285    pub async fn create_response(&self) -> ZaiResult<()> {
286        self.dispatch(ClientEvent::ResponseCreate {
287            client_timestamp: Some(now_ms()),
288        })
289        .await
290    }
291
292    /// Cancel the in-flight response (`response.cancel`), e.g. on interruption.
293    pub async fn cancel(&self) -> ZaiResult<()> {
294        self.dispatch(ClientEvent::ResponseCancel {
295            client_timestamp: Some(now_ms()),
296        })
297        .await
298    }
299
300    /// Stream of all server events (transcripts, response lifecycle, errors,
301    /// heartbeats). Late subscribers miss events broadcast before they joined;
302    /// subscribe before driving commands when ordering matters. Transient
303    /// `Lagged` gaps (slow consumer) are dropped rather than erroring.
304    pub fn events(&self) -> Pin<Box<dyn Stream<Item = ServerEvent> + Send + '_>> {
305        let stream = BroadcastStream::new(self.events_tx.subscribe())
306            .filter_map(|res| async move { res.ok() });
307        Box::pin(stream)
308    }
309
310    /// Stream of decoded audio output chunks (PCM/MP3 bytes, per
311    /// `output_audio_format`).
312    pub fn audio_stream(&self) -> Pin<Box<dyn Stream<Item = Bytes> + Send + '_>> {
313        let stream = BroadcastStream::new(self.audio_tx.subscribe())
314            .filter_map(|res| async move { res.ok() });
315        Box::pin(stream)
316    }
317
318    /// The model id this session was opened for (metadata; the protocol does
319    /// not transmit it on the wire).
320    pub fn model_name(&self) -> &str {
321        &self.model_name
322    }
323
324    #[tracing::instrument(name = "realtime.dispatch", skip(self, event))]
325    async fn dispatch(&self, event: ClientEvent) -> ZaiResult<()> {
326        self.cmd_tx
327            .send(Command::ClientEvent(Box::new(event)))
328            .await
329            .map_err(|_| RealtimeErrorKind::Closed.into())
330    }
331
332    /// Close the session and wait for the event loop to finish.
333    pub async fn close(self) -> ZaiResult<()> {
334        let _ = self.cmd_tx.send(Command::Close).await;
335        // Dropping the last Sender closes the command channel; the event loop
336        // observes `Command::Close` (sent above) and exits.
337        let _ = self.join.await;
338        Ok(())
339    }
340}
341
342fn now_ms() -> i64 {
343    chrono::Utc::now().timestamp_millis()
344}
345
346fn new_event_id() -> String {
347    format!("evt_{}", uuid::Uuid::new_v4().simple())
348}
zai_rs/realtime/session.rs

zai_rs/realtime/
session.rs
