ferrin_core/realtime/
session.rs

1//! Connection task of a realtime session.
2//!
3//! Derived from the Vercel AI SDK (Apache-2.0, Copyright 2023 Vercel, Inc.),
4//! translated from TypeScript to Rust and modified; see `NOTICE`.
5
6use std::sync::Arc;
7use std::sync::Mutex;
8use std::sync::atomic::AtomicBool;
9use std::sync::atomic::Ordering;
10
11use ferrin_spec::DynRealtimeModel;
12use ferrin_spec::JsonValue;
13use ferrin_spec::ToolCallId;
14use ferrin_spec::ToolName;
15use ferrin_spec::realtime_model::ClientSecret;
16use ferrin_spec::realtime_model::ConversationItem;
17use ferrin_spec::realtime_model::ConversationRole;
18use ferrin_spec::realtime_model::RealtimeClientEvent;
19use ferrin_spec::realtime_model::RealtimeServerEvent;
20use ferrin_spec::realtime_model::RealtimeSessionConfig;
21use ferrin_tool::Tool;
22use ferrin_tool::ToolContext;
23use ferrin_tool::ToolSet;
24use ferrin_tool::execute_to_completion;
25use futures_util::SinkExt;
26use futures_util::StreamExt;
27use http::HeaderValue;
28use http::header::SEC_WEBSOCKET_PROTOCOL;
29use http::header::USER_AGENT;
30use tokio::net::TcpStream;
31use tokio::sync::mpsc;
32use tokio::task::JoinSet;
33use tokio_tungstenite::MaybeTlsStream;
34use tokio_tungstenite::WebSocketStream;
35use tokio_tungstenite::connect_async_tls_with_config;
36use tokio_tungstenite::tungstenite::Message;
37use tokio_tungstenite::tungstenite::client::IntoClientRequest;
38use tokio_util::sync::CancellationToken;
39
40use super::RealtimeEvent;
41use super::tools::ToolTurn;
42use crate::error::Error;
43
44type Socket = WebSocketStream<MaybeTlsStream<TcpStream>>;
45
46/// Capacity of the outbound message queue.
47const OUTBOUND_BUFFER: usize = 64;
48
49/// State shared between the connection task and the handles.
50struct Shared {
51    model: Arc<dyn DynRealtimeModel>,
52    outbound: mpsc::Sender<JsonValue>,
53    cancellation: CancellationToken,
54    closed: AtomicBool,
55    tool_turn: Mutex<ToolTurn>,
56    serialization: tokio::sync::Semaphore,
57}
58
59impl Shared {
60    fn tool_turn(&self) -> std::sync::MutexGuard<'_, ToolTurn> {
61        self.tool_turn
62            .lock()
63            .unwrap_or_else(std::sync::PoisonError::into_inner)
64    }
65}
66
67/// Sends events into a realtime session and closes it.
68///
69/// Handles are cheap to clone and remain valid until the session ends.
70#[derive(Clone)]
71pub struct RealtimeHandle {
72    shared: Arc<Shared>,
73}
74
75impl std::fmt::Debug for RealtimeHandle {
76    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77        f.debug_struct("RealtimeHandle")
78            .field("provider", self.shared.model.provider())
79            .field("model_id", self.shared.model.model_id())
80            .field("closed", &self.is_closed())
81            .finish()
82    }
83}
84
85impl RealtimeHandle {
86    /// Sends a client event.
87    ///
88    /// # Errors
89    ///
90    /// Fails when the model cannot serialize the event or the session is
91    /// closed.
92    pub async fn send(&self, event: RealtimeClientEvent) -> Result<(), Error> {
93        let _serialization = tokio::select! {
94            biased;
95            () = self.shared.cancellation.cancelled() => return Err(closed_error()),
96            permit = self.shared.serialization.acquire() => permit.map_err(|_| closed_error())?,
97        };
98        let raw = tokio::select! {
99            biased;
100            () = self.shared.cancellation.cancelled() => return Err(closed_error()),
101            result = self.shared.model.serialize_client_event(event) => result?,
102        };
103        if raw.is_null() {
104            Ok(())
105        } else {
106            self.send_raw(raw).await
107        }
108    }
109
110    /// Sends a raw provider message.
111    ///
112    /// # Errors
113    ///
114    /// Fails when the session is closed.
115    pub async fn send_raw(&self, raw: JsonValue) -> Result<(), Error> {
116        self.shared
117            .outbound
118            .send(raw)
119            .await
120            .map_err(|_| closed_error())
121    }
122
123    /// Sends a user text message and requests a response.
124    ///
125    /// # Errors
126    ///
127    /// See [`RealtimeHandle::send`].
128    pub async fn send_text(&self, text: impl Into<String>) -> Result<(), Error> {
129        self.send(RealtimeClientEvent::ConversationItemCreate {
130            item: ConversationItem::TextMessage {
131                role: ConversationRole::User,
132                text: text.into(),
133            },
134        })
135        .await?;
136        self.send(RealtimeClientEvent::ResponseCreate { options: None })
137            .await
138    }
139
140    /// Submits the output of a tool call. A follow-up response is requested
141    /// once every tool call of the response has an output.
142    ///
143    /// # Errors
144    ///
145    /// See [`RealtimeHandle::send`].
146    pub async fn add_tool_output(&self, call_id: &str, output: &JsonValue) -> Result<(), Error> {
147        let name = self
148            .shared
149            .tool_turn()
150            .name(call_id)
151            .map(|name| name.as_str().to_owned());
152        self.send(RealtimeClientEvent::ConversationItemCreate {
153            item: ConversationItem::FunctionCallOutput {
154                call_id: call_id.to_owned(),
155                name,
156                output: output.to_string(),
157            },
158        })
159        .await?;
160        let request_response = self.shared.tool_turn().output_submitted(call_id);
161        if request_response {
162            self.send(RealtimeClientEvent::ResponseCreate { options: None })
163                .await?;
164        }
165        Ok(())
166    }
167
168    /// Requests a response with the default options.
169    ///
170    /// # Errors
171    ///
172    /// See [`RealtimeHandle::send`].
173    pub async fn request_response(&self) -> Result<(), Error> {
174        self.send(RealtimeClientEvent::ResponseCreate { options: None })
175            .await
176    }
177
178    /// Asks the connection task to close the connection.
179    pub fn close(&self) {
180        self.shared.cancellation.cancel();
181    }
182
183    /// Returns `true` once the connection task has finished.
184    #[must_use]
185    pub fn is_closed(&self) -> bool {
186        self.shared.closed.load(Ordering::Acquire)
187    }
188}
189
190fn closed_error() -> Error {
191    Error::message("realtime session is closed")
192}
193
194/// Inputs of [`start`].
195pub(super) struct StartOptions {
196    pub(super) model: Arc<dyn DynRealtimeModel>,
197    pub(super) secret: ClientSecret,
198    pub(super) config: RealtimeSessionConfig,
199    pub(super) tools: Arc<ToolSet>,
200    pub(super) tools_context: Option<JsonValue>,
201    pub(super) cancellation: CancellationToken,
202    pub(super) events: mpsc::Sender<RealtimeEvent>,
203}
204
205/// Opens the WebSocket, sends the initial `session-update` and spawns the
206/// connection task into `tasks`.
207pub(super) async fn start(
208    options: StartOptions,
209    tasks: &mut JoinSet<()>,
210) -> Result<RealtimeHandle, Error> {
211    let StartOptions {
212        model,
213        secret,
214        config,
215        tools,
216        tools_context,
217        cancellation,
218        events,
219    } = options;
220    let ws_config = model.websocket_config(&secret.token, &secret.url);
221    let mut request = ws_config
222        .url
223        .as_str()
224        .into_client_request()
225        .map_err(Error::other)?;
226    if !ws_config.protocols.is_empty() {
227        let value = HeaderValue::from_str(&ws_config.protocols.join(", "))
228            .map_err(|error| Error::invalid_argument("protocols", error.to_string()))?;
229        request.headers_mut().insert(SEC_WEBSOCKET_PROTOCOL, value);
230    }
231    request
232        .headers_mut()
233        .insert(USER_AGENT, HeaderValue::from_static(crate::USER_AGENT));
234
235    let connect = connect_async_tls_with_config(request, None, false, None);
236    let (mut socket, _response) = tokio::select! {
237        result = connect => result.map_err(Error::other)?,
238        () = cancellation.cancelled() => return Err(Error::Cancelled),
239    };
240
241    let initial = model
242        .serialize_client_event(RealtimeClientEvent::SessionUpdate {
243            config: Box::new(config),
244        })
245        .await?;
246    if !initial.is_null() {
247        socket
248            .send(wire_message(initial))
249            .await
250            .map_err(Error::other)?;
251    }
252
253    let (outbound_tx, outbound_rx) = mpsc::channel(OUTBOUND_BUFFER);
254    let shared = Arc::new(Shared {
255        model,
256        outbound: outbound_tx,
257        cancellation,
258        closed: AtomicBool::new(false),
259        tool_turn: Mutex::new(ToolTurn::default()),
260        serialization: tokio::sync::Semaphore::new(1),
261    });
262    let handle = RealtimeHandle {
263        shared: Arc::clone(&shared),
264    };
265    tasks.spawn(run(Connection {
266        socket,
267        outbound: outbound_rx,
268        shared,
269        events,
270        tools,
271        tools_context,
272    }));
273    Ok(handle)
274}
275
276struct Connection {
277    socket: Socket,
278    outbound: mpsc::Receiver<JsonValue>,
279    shared: Arc<Shared>,
280    events: mpsc::Sender<RealtimeEvent>,
281    tools: Arc<ToolSet>,
282    tools_context: Option<JsonValue>,
283}
284
285async fn run(mut connection: Connection) {
286    let mut tool_tasks: JoinSet<()> = JoinSet::new();
287    loop {
288        let stop = tokio::select! {
289            biased;
290            () = connection.shared.cancellation.cancelled() => {
291                let _ = connection.socket.send(Message::Close(None)).await;
292                true
293            }
294            outbound = connection.outbound.recv() => match outbound {
295                Some(raw) => connection.send_raw(raw).await.is_err(),
296                None => {
297                    let _ = connection.socket.send(Message::Close(None)).await;
298                    true
299                }
300            },
301            frame = connection.socket.next() => match frame {
302                Some(Ok(Message::Text(text))) => {
303                    connection.handle_text(text.as_str(), &mut tool_tasks).await
304                }
305                Some(Ok(Message::Binary(bytes))) => match std::str::from_utf8(&bytes) {
306                    Ok(text) => connection.handle_text(text, &mut tool_tasks).await,
307                    Err(_) => false,
308                },
309                Some(Ok(Message::Close(_))) | None => true,
310                Some(Ok(_)) => false,
311                Some(Err(error)) => {
312                    connection.emit(Err(Error::other(error))).await;
313                    true
314                }
315            },
316            Some(joined) = tool_tasks.join_next(), if !tool_tasks.is_empty() => {
317                if let Err(error) = joined
318                    && !error.is_cancelled()
319                {
320                    connection.emit(Err(Error::other(error))).await;
321                }
322                false
323            }
324        };
325        if stop {
326            break;
327        }
328    }
329    tool_tasks.shutdown().await;
330    connection.shared.closed.store(true, Ordering::Release);
331    connection.shared.cancellation.cancel();
332}
333
334impl Connection {
335    /// Forwards an item to the consumer; returns `false` when the consumer is
336    /// gone.
337    async fn emit(&self, item: RealtimeEvent) -> bool {
338        self.events.send(item).await.is_ok()
339    }
340
341    async fn send_raw(&mut self, raw: JsonValue) -> Result<(), ()> {
342        match self.socket.send(wire_message(raw)).await {
343            Ok(()) => Ok(()),
344            Err(error) => {
345                self.emit(Err(Error::other(error))).await;
346                Err(())
347            }
348        }
349    }
350
351    /// Handles one text frame; returns `true` when the loop must stop.
352    async fn handle_text(&mut self, text: &str, tool_tasks: &mut JoinSet<()>) -> bool {
353        let raw: JsonValue = match serde_json::from_str(text) {
354            Ok(raw) => raw,
355            Err(error) => {
356                tracing::debug!(error = %error, "ignoring non-JSON realtime message");
357                return false;
358            }
359        };
360        if let Some(reply) = self.shared.model.health_check_response(&raw)
361            && self.send_raw(reply).await.is_err()
362        {
363            return true;
364        }
365        let events = match self.shared.model.parse_server_event(raw) {
366            Ok(events) => events,
367            Err(error) => return !self.emit(Err(Error::from(error))).await,
368        };
369        for event in events {
370            if !self.handle_event(event, tool_tasks).await {
371                return true;
372            }
373        }
374        false
375    }
376
377    /// Handles one standardized event; returns `false` when the consumer is
378    /// gone or sending failed.
379    async fn handle_event(
380        &mut self,
381        event: RealtimeServerEvent,
382        tool_tasks: &mut JoinSet<()>,
383    ) -> bool {
384        let follow_up = match &event {
385            RealtimeServerEvent::FunctionCallArgumentsDone {
386                call_id,
387                name,
388                arguments,
389                ..
390            } => {
391                let name = ToolName::new(name.clone());
392                self.shared.tool_turn().record_name(call_id, &name);
393                Some(ToolCall {
394                    call_id: call_id.clone(),
395                    name,
396                    arguments: arguments.clone(),
397                })
398            }
399            _ => None,
400        };
401        let response_done = matches!(event, RealtimeServerEvent::ResponseDone { .. });
402
403        if !self.emit(Ok(event)).await {
404            return false;
405        }
406        if let Some(call) = follow_up
407            && !self.start_tool_call(call, tool_tasks).await
408        {
409            return false;
410        }
411        if response_done && self.shared.tool_turn().response_done() {
412            let handle = RealtimeHandle {
413                shared: Arc::clone(&self.shared),
414            };
415            if let Err(error) = handle.request_response().await {
416                return self.emit(Err(error)).await;
417            }
418        }
419        true
420    }
421
422    /// Starts executing a tool call; returns `false` when the consumer is
423    /// gone.
424    async fn start_tool_call(&self, call: ToolCall, tool_tasks: &mut JoinSet<()>) -> bool {
425        let input: JsonValue = match serde_json::from_str(&call.arguments) {
426            Ok(input) => input,
427            Err(error) => {
428                return self
429                    .emit(Err(Error::invalid_tool_input(
430                        call.name,
431                        call.arguments,
432                        error,
433                    )))
434                    .await;
435            }
436        };
437        self.shared
438            .tool_turn()
439            .call_started(&call.call_id, &call.name);
440        let Some(tool) = self.tools.get(call.name.as_str()).map(Arc::clone) else {
441            let available = self.tools.names().cloned().collect();
442            return self
443                .emit(Err(Error::no_such_tool(call.name, available)))
444                .await;
445        };
446        if !tool.is_executable() {
447            return true;
448        }
449        let input = match tool.validate_input(&call.name, input) {
450            Ok(input) => input,
451            Err(error) => {
452                return self
453                    .emit(Err(Error::invalid_tool_input(
454                        call.name,
455                        call.arguments,
456                        error,
457                    )))
458                    .await;
459            }
460        };
461        let tools_context =
462            match tool.validate_named_context(&call.name, self.tools_context.as_ref()) {
463                Ok(context) => context,
464                Err(error) => {
465                    return self
466                        .emit(Err(Error::invalid_argument(
467                            "tools_context",
468                            error.to_string(),
469                        )))
470                        .await;
471                }
472            };
473        let ctx = ToolContext::new(ToolCallId::new(call.call_id.clone()))
474            .with_cancellation(self.shared.cancellation.child_token())
475            .with_tools_context(tools_context);
476        let handle = RealtimeHandle {
477            shared: Arc::clone(&self.shared),
478        };
479        let events = self.events.clone();
480        tool_tasks.spawn(execute_tool_call(
481            tool,
482            input,
483            ctx,
484            call.call_id,
485            handle,
486            events,
487        ));
488        true
489    }
490}
491
492fn wire_message(raw: JsonValue) -> Message {
493    match raw {
494        JsonValue::String(text) => Message::text(text),
495        other => Message::text(other.to_string()),
496    }
497}
498
499struct ToolCall {
500    call_id: String,
501    name: ToolName,
502    arguments: String,
503}
504
505async fn execute_tool_call(
506    tool: Arc<Tool>,
507    input: JsonValue,
508    ctx: ToolContext,
509    call_id: String,
510    handle: RealtimeHandle,
511    events: mpsc::Sender<RealtimeEvent>,
512) {
513    let needs_approval = tokio::select! {
514        biased;
515        () = ctx.cancellation.cancelled() => return,
516        required = tool.needs_approval().resolve(input.clone(), ctx.clone()) => required,
517    };
518    if needs_approval {
519        let _ = events.send(Err(Error::invalid_argument(
520            "tools",
521            "realtime tool requires application approval; submit its approved result or denial with add_tool_output",
522        ))).await;
523        return;
524    }
525    if ctx.cancellation.is_cancelled() {
526        let _ = events.send(Err(Error::Cancelled)).await;
527        return;
528    }
529    let Some(stream) = tool.execute(input, ctx) else {
530        return;
531    };
532    let result = match execute_to_completion(stream, |_| {}).await {
533        Ok(output) => handle.add_tool_output(&call_id, &output).await,
534        Err(error) => Err(Error::from(error)),
535    };
536    if let Err(error) = result {
537        let _ = events.send(Err(error)).await;
538    }
539}
ferrin_core/realtime/session.rs

ferrin_core/realtime/
session.rs
