rig_core/providers/openai/responses_api/
websocket.rs

1//! Stateful Responses WebSocket sessions over caller-supplied connections.
2//! Sessions permit one in-flight turn and chain completed or incomplete response IDs.
3//!
4//! ```
5//! use rig_core::providers::openai::responses_api::websocket::ResponsesWebSocketCreateOptions;
6//! let options = ResponsesWebSocketCreateOptions::warmup();
7//! assert_eq!(options.generate, Some(false));
8//! ```
9
10use crate::completion;
11use crate::completion::options::FinalBody;
12use crate::driver::Model;
13use crate::error::{EncodeError, ProviderError};
14use crate::http_client::{self, NoBody};
15use crate::json_utils::remove_preserving_order;
16use crate::operation::{Completion, ReplyFold};
17use crate::providers::openai::responses_api::streaming::{
18    ResponsesDecoder, ResponsesEvent, classify_responses_payload, document, is_lifecycle_event,
19};
20use crate::providers::openai::responses_api::wire::Responses;
21use crate::streaming::Item;
22use crate::wire::{Call, Flow, Mode, Operation, Reply, Shared, Wire, WireFrame};
23use crate::ws_client::{
24    BoxedWebSocketConnection, ConnectOptions, Frame, WebSocketClientExt, WebSocketConnection,
25};
26use serde::{Deserialize, Serialize};
27use serde_json::{Map, Value};
28use std::time::Duration;
29
30/// The websocket endpoint's path, appended to the client's configured base URL.
31const WEBSOCKET_PATH: &str = "responses";
32
33const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(30);
34
35/// Request-ID header read from rejected WebSocket upgrades.
36const REQUEST_ID_HEADER: Option<&'static str> =
37    crate::providers::openai::wire::OPENAI.request_id_header;
38
39/// Options for a `response.create` message sent over OpenAI WebSocket mode.
40#[derive(Debug, Clone, Default, Serialize, Deserialize)]
41pub struct ResponsesWebSocketCreateOptions {
42    /// When set to `false`, OpenAI prepares request state without generating a model output.
43    ///
44    /// This is the "warmup" mode described in the OpenAI WebSocket mode guide.
45    #[serde(skip_serializing_if = "Option::is_none")]
46    pub generate: Option<bool>,
47}
48
49impl ResponsesWebSocketCreateOptions {
50    /// Creates warmup options equivalent to `generate: false`.
51    #[must_use]
52    pub fn warmup() -> Self {
53        Self {
54            generate: Some(false),
55        }
56    }
57}
58
59#[derive(Debug, Clone, Serialize)]
60struct ResponsesWebSocketClientEvent {
61    #[serde(rename = "type")]
62    kind: ResponsesWebSocketClientEventKind,
63    #[serde(flatten)]
64    request: FinalBody,
65}
66
67#[derive(Debug, Clone, Serialize)]
68enum ResponsesWebSocketClientEventKind {
69    #[serde(rename = "response.create")]
70    ResponseCreate,
71}
72
73/// A protocol error event emitted by OpenAI WebSocket mode.
74#[derive(Debug, Clone, Serialize, Deserialize)]
75pub struct ResponsesWebSocketErrorEvent {
76    /// The event type.
77    #[serde(rename = "type")]
78    pub kind: ResponsesWebSocketErrorEventKind,
79    /// The provider error payload.
80    pub error: ResponsesWebSocketErrorPayload,
81}
82
83impl std::fmt::Display for ResponsesWebSocketErrorEvent {
84    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
85        self.error.fmt(f)
86    }
87}
88
89/// The event kind for an OpenAI WebSocket protocol error.
90#[derive(Debug, Clone, Serialize, Deserialize)]
91pub enum ResponsesWebSocketErrorEventKind {
92    #[serde(rename = "error")]
93    Error,
94}
95
96/// The payload carried by an OpenAI WebSocket protocol error event.
97#[derive(Debug, Clone, Default, Serialize, Deserialize)]
98pub struct ResponsesWebSocketErrorPayload {
99    /// Provider-specific error code when supplied.
100    #[serde(skip_serializing_if = "Option::is_none")]
101    pub code: Option<String>,
102    /// Human-readable error message.
103    #[serde(skip_serializing_if = "Option::is_none")]
104    pub message: Option<String>,
105    /// Any extra fields supplied by the provider.
106    #[serde(flatten, default)]
107    pub extra: Map<String, Value>,
108}
109
110impl std::fmt::Display for ResponsesWebSocketErrorPayload {
111    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        match (&self.code, &self.message) {
113            (Some(code), Some(message)) => write!(f, "{code}: {message}"),
114            (None, Some(message)) => f.write_str(message),
115            (Some(code), None) => f.write_str(code),
116            (None, None) => f.write_str("OpenAI websocket error"),
117        }
118    }
119}
120
121/// The optional `response.done` event emitted by OpenAI WebSocket mode.
122#[derive(Debug, Clone, Serialize, Deserialize)]
123pub struct ResponsesWebSocketDoneEvent {
124    /// The event type.
125    #[serde(rename = "type")]
126    pub kind: ResponsesWebSocketDoneEventKind,
127    /// The provider payload for the finished response.
128    pub response: Value,
129}
130
131impl ResponsesWebSocketDoneEvent {
132    /// Returns the response ID if the payload includes one.
133    #[must_use]
134    pub fn response_id(&self) -> Option<&str> {
135        self.response.get("id").and_then(Value::as_str)
136    }
137
138    fn status(&self) -> Option<&str> {
139        self.response.get("status").and_then(Value::as_str)
140    }
141}
142
143/// The event kind for the terminal websocket event.
144#[derive(Debug, Clone, Serialize, Deserialize)]
145pub enum ResponsesWebSocketDoneEventKind {
146    #[serde(rename = "response.done")]
147    ResponseDone,
148}
149
150/// A server event emitted by OpenAI WebSocket mode.
151#[derive(Debug, Clone)]
152pub enum ResponsesWebSocketEvent {
153    /// A response lifecycle event such as `response.created` or
154    /// `response.completed`, with the response object it carries.
155    Response {
156        /// The event's `type`.
157        kind: String,
158        /// The response object, as the provider sent it.
159        response: Value,
160    },
161    /// A streaming item or delta event such as `response.output_text.delta`,
162    /// as the provider sent it.
163    Item(Value),
164    /// A protocol-level websocket error event.
165    Error(ResponsesWebSocketErrorEvent),
166    /// An optional `response.done` event emitted by OpenAI over WebSockets.
167    Done(ResponsesWebSocketDoneEvent),
168    /// Unrecognized event retained for [`Item::Unknown`] passthrough.
169    Unknown(crate::streaming::UnknownPayload),
170}
171
172impl ResponsesWebSocketEvent {
173    /// Returns the response ID when the event includes one.
174    #[must_use]
175    pub fn response_id(&self) -> Option<&str> {
176        match self {
177            Self::Response { response, .. } => response.get("id").and_then(Value::as_str),
178            Self::Done(done) => done.response_id(),
179            Self::Item(_) | Self::Error(_) | Self::Unknown(_) => None,
180        }
181    }
182
183    /// Returns `true` when this event ends the current in-flight websocket turn.
184    #[must_use]
185    pub fn is_terminal(&self) -> bool {
186        match self {
187            Self::Response { kind, .. } => is_terminal_kind(kind),
188            Self::Error(_) | Self::Done(_) => true,
189            Self::Item(_) | Self::Unknown(_) => false,
190        }
191    }
192}
193
194/// Whether a lifecycle event of `kind` ends the turn.
195fn is_terminal_kind(kind: &str) -> bool {
196    matches!(
197        kind,
198        "response.completed" | "response.failed" | "response.incomplete"
199    )
200}
201
202/// A builder for an OpenAI Responses WebSocket session.
203///
204/// The default builder applies a 30 second connection timeout and leaves the
205/// per-event timeout disabled.
206pub struct ResponsesWebSocketSessionBuilder {
207    wire: Responses,
208    connect_timeout: Option<Duration>,
209    event_timeout: Option<Duration>,
210}
211
212impl ResponsesWebSocketSessionBuilder {
213    pub(crate) fn new(wire: Responses) -> Self {
214        Self {
215            wire,
216            connect_timeout: Some(DEFAULT_CONNECT_TIMEOUT),
217            event_timeout: None,
218        }
219    }
220
221    /// Sets the timeout for establishing the websocket connection.
222    #[must_use]
223    pub fn connect_timeout(mut self, timeout: Duration) -> Self {
224        self.connect_timeout = Some(timeout);
225        self
226    }
227
228    /// Waits for the websocket connection with no deadline. The event
229    /// timeout is off unless [`Self::event_timeout`] sets it.
230    #[must_use]
231    pub fn without_connect_timeout(mut self) -> Self {
232        self.connect_timeout = None;
233        self
234    }
235
236    /// Sets the timeout for waiting on the next websocket event.
237    #[must_use]
238    pub fn event_timeout(mut self, timeout: Duration) -> Self {
239        self.event_timeout = Some(timeout);
240        self
241    }
242
243    /// Waits for each websocket event with no deadline, the default.
244    #[must_use]
245    pub fn without_event_timeout(mut self) -> Self {
246        self.event_timeout = None;
247        self
248    }
249}
250
251impl ResponsesWebSocketSessionBuilder {
252    /// Open a session over the bundled tungstenite backend with the configured
253    /// timeouts. Return handshake construction, transport, or provider errors.
254    #[cfg(all(feature = "tungstenite", not(target_family = "wasm")))]
255    #[cfg_attr(docsrs, doc(cfg(feature = "tungstenite")))]
256    pub async fn connect(self) -> Result<ResponsesWebSocketSession, ProviderError> {
257        self.connect_with(&rig_tungstenite::TungsteniteClient::new())
258            .await
259    }
260
261    /// Open a session over `backend` with the configured timeouts.
262    /// Return handshake construction, transport, or provider errors.
263    pub async fn connect_with<W>(
264        self,
265        backend: &W,
266    ) -> Result<ResponsesWebSocketSession, ProviderError>
267    where
268        W: WebSocketClientExt,
269    {
270        ResponsesWebSocketSession::connect_with_timeouts(
271            backend,
272            self.wire,
273            self.connect_timeout,
274            self.event_timeout,
275        )
276        .await
277    }
278}
279
280/// Sequential Responses session with automatic response-ID chaining.
281/// Completed and incomplete responses update the chain unless a request supplies
282/// its own `previous_response_id`. Call [`Self::close`] to perform a close handshake.
283pub struct ResponsesWebSocketSession {
284    wire: Responses,
285    previous_response_id: Option<String>,
286    pending_done_response_id: Option<String>,
287    socket: BoxedWebSocketConnection,
288    in_flight: bool,
289    /// The fold of the turn in flight, opened when it was sent.
290    turn: Option<ReplyFold>,
291    event_timeout: Option<Duration>,
292    closed: bool,
293    failed: bool,
294}
295
296impl ResponsesWebSocketSession {
297    async fn connect_with_timeouts<W>(
298        backend: &W,
299        wire: Responses,
300        connect_timeout: Option<Duration>,
301        event_timeout: Option<Duration>,
302    ) -> Result<Self, ProviderError>
303    where
304        W: WebSocketClientExt,
305    {
306        let request = websocket_request(&wire)?;
307        let socket = backend
308            .connect(request, ConnectOptions::new().with_timeout(connect_timeout))
309            .await
310            .map_err(websocket_provider_error)?;
311
312        Ok(Self::from_connection(wire, socket, event_timeout))
313    }
314
315    /// Build a session over an already-open, authenticated connection.
316    /// `event_timeout: None` waits indefinitely for each event.
317    pub fn from_connection(
318        wire: Responses,
319        connection: BoxedWebSocketConnection,
320        event_timeout: Option<Duration>,
321    ) -> Self {
322        Self {
323            wire,
324            previous_response_id: None,
325            pending_done_response_id: None,
326            socket: connection,
327            in_flight: false,
328            turn: None,
329            event_timeout,
330            closed: false,
331            failed: false,
332        }
333    }
334
335    /// Return the response ID retained for automatic chaining, if any.
336    #[must_use]
337    pub fn previous_response_id(&self) -> Option<&str> {
338        self.previous_response_id.as_deref()
339    }
340
341    /// Clears the cached `previous_response_id` so the next turn starts a fresh chain.
342    pub fn clear_previous_response_id(&mut self) {
343        self.previous_response_id = None;
344    }
345
346    /// Sends a `response.create` event for a Rig completion request.
347    pub async fn send(
348        &mut self,
349        completion_request: crate::completion::CompletionRequest,
350    ) -> Result<(), ProviderError> {
351        self.send_with_options(
352            completion_request,
353            ResponsesWebSocketCreateOptions::default(),
354        )
355        .await
356    }
357
358    /// Sends a `response.create` event with explicit websocket-mode options.
359    pub async fn send_with_options(
360        &mut self,
361        completion_request: crate::completion::CompletionRequest,
362        options: ResponsesWebSocketCreateOptions,
363    ) -> Result<(), ProviderError> {
364        self.ensure_open()?;
365
366        if self.in_flight {
367            return Err(ProviderError::Provider(
368                "An OpenAI websocket response is already in flight on this session".to_string(),
369            ));
370        }
371
372        // A request that names no `previous_response_id` continues the
373        // session's chain, so the history is shaped as a stored
374        // continuation's.
375        let mut completion_request = completion_request;
376        if let Some(generate) = options.generate {
377            let params = completion_request
378                .additional_params
379                .get_or_insert_with(|| Value::Object(Map::new()));
380            if params.is_null() {
381                *params = Value::Object(Map::new());
382            }
383            if let Some(params) = params.as_object_mut() {
384                params.insert("generate".to_owned(), Value::Bool(generate));
385            }
386        }
387        if let Some(previous) = &self.previous_response_id {
388            let params = completion_request
389                .additional_params
390                .get_or_insert_with(|| Value::Object(Map::new()));
391            if let Some(params) = params.as_object_mut()
392                && params
393                    .get("previous_response_id")
394                    .is_none_or(Value::is_null)
395            {
396                params.insert(
397                    "previous_response_id".to_owned(),
398                    Value::String(previous.clone()),
399                );
400            }
401        }
402        // The session sends without the driver, so it shapes the history
403        // for the model the way the driver does.
404        let describe = self.wire.describe();
405        let completion_request = Completion::prepare(completion_request, &describe)?;
406        let turn = Completion::fold(
407            &completion_request,
408            &mut Call::new(&describe, Mode::Streaming),
409        );
410        let payload = ResponsesWebSocketClientEvent {
411            kind: ResponsesWebSocketClientEventKind::ResponseCreate,
412            request: self.prepare_request(completion_request)?,
413        };
414
415        crate::providers::internal::trace_json(
416            crate::providers::internal::LogTarget::Completions,
417            "OpenAI websocket request",
418            &payload,
419        );
420
421        let payload = serde_json::to_string(&payload)?;
422
423        if let Err(error) = self.socket.send(Frame::Text(payload)).await {
424            return Err(self.fail_session(websocket_provider_error(error)));
425        }
426        self.in_flight = true;
427        self.turn = Some(turn);
428
429        Ok(())
430    }
431
432    /// Reads the next server event for the current in-flight turn.
433    pub async fn next_event(&mut self) -> Result<ResponsesWebSocketEvent, ProviderError> {
434        self.next_event_with_payload().await.map(|(event, _)| event)
435    }
436
437    /// Reads the next lifecycle event and retains its payload for content decoding
438    /// by [`ResponsesDecoder`]. Returns the same session errors as [`Self::next_event`].
439    async fn next_event_with_payload(
440        &mut self,
441    ) -> Result<(ResponsesWebSocketEvent, String), ProviderError> {
442        self.ensure_open()?;
443
444        if !self.in_flight {
445            return Err(ProviderError::Provider(
446                "No OpenAI websocket response is currently in flight on this session".to_string(),
447            ));
448        }
449
450        loop {
451            let message = match self.read_next_frame().await? {
452                Ok(message) => message,
453                Err(error) => return Err(self.fail_session(websocket_provider_error(error))),
454            };
455
456            let Some(message) = message else {
457                self.mark_closed();
458                return Err(ProviderError::Provider(
459                    "The OpenAI websocket connection closed before the turn finished".to_string(),
460                ));
461            };
462
463            let payload = match websocket_frame_to_text(message) {
464                Ok(Some(payload)) => payload,
465                Ok(None) => continue,
466                Err(error) => return Err(self.fail_session(error)),
467            };
468            let event = match parse_server_event(&payload) {
469                Ok(Some(event)) => event,
470                Ok(None) => continue,
471                Err(error) => return Err(self.fail_session(error)),
472            };
473            if let ResponsesWebSocketEvent::Done(done) = &event {
474                // OpenAI may emit `response.done` after the turn has already ended at
475                // `response.completed`. Ignore that trailing event on the next turn.
476                if self.pending_done_response_id.as_deref() == done.response_id() {
477                    self.pending_done_response_id = None;
478                    continue;
479                }
480            }
481            self.update_state_for_event(&event);
482            return Ok((event, payload));
483        }
484    }
485
486    /// Sends a warmup turn (`generate: false`) and returns the resulting response ID.
487    pub async fn warmup(
488        &mut self,
489        completion_request: crate::completion::CompletionRequest,
490    ) -> Result<String, ProviderError> {
491        self.send_with_options(
492            completion_request,
493            ResponsesWebSocketCreateOptions::warmup(),
494        )
495        .await?;
496        let response = self.wait_for_terminal_response().await?.0;
497        Ok(response
498            .get("id")
499            .and_then(Value::as_str)
500            .unwrap_or_default()
501            .to_owned())
502    }
503
504    /// Sends a completion turn and collects the final OpenAI response,
505    /// normalized; its `raw` is the `Response` the turn's events rebuild.
506    pub async fn completion(
507        &mut self,
508        completion_request: crate::completion::CompletionRequest,
509    ) -> Result<completion::CompletionResponse, ProviderError> {
510        self.send(completion_request).await?;
511        Ok(self.wait_for_terminal_response().await?.1)
512    }
513
514    /// Closes the websocket connection.
515    ///
516    /// Call this when you are finished with the session so the websocket can
517    /// terminate with a clean close handshake.
518    pub async fn close(&mut self) -> Result<(), ProviderError> {
519        if self.closed {
520            return Ok(());
521        }
522
523        let result = self
524            .socket
525            .close(None)
526            .await
527            .map_err(websocket_provider_error);
528        self.mark_closed();
529        result
530    }
531
532    /// The `response.create` body: the HTTP body, without the `stream` and
533    /// `background` flags an event-driven session ignores.
534    fn prepare_request(
535        &self,
536        completion_request: crate::completion::CompletionRequest,
537    ) -> Result<FinalBody, ProviderError> {
538        Ok(self
539            .wire
540            .responses_request(&completion_request, super::Delivery::WebSocket)?)
541    }
542
543    /// Decode the turn's events as they arrive, and return the provider's
544    /// completed or incomplete response with the reply they fold into.
545    /// Transport, protocol, and decoder failures return an error. A terminal
546    /// event without a response body is an error.
547    async fn wait_for_terminal_response(
548        &mut self,
549    ) -> Result<(Value, completion::CompletionResponse), ProviderError> {
550        let provider = self.wire.describe().name.to_owned();
551        let Some(turn) = self.turn.take() else {
552            return Err(ProviderError::Provider(
553                "No OpenAI websocket response is currently in flight on this session".to_string(),
554            ));
555        };
556        // The reply's state and its decoder live for this turn only.
557        let reply = std::sync::Mutex::new(Shared::new(turn));
558        let mut decoder = self.wire.decoder();
559        let mut document = self.wire.reassembler();
560        loop {
561            let (event, payload) = self.next_event_with_payload().await?;
562            match event {
563                ResponsesWebSocketEvent::Response { kind, response } => {
564                    if !is_terminal_kind(&kind) {
565                        feed(&mut decoder, &mut document, &reply, payload)?;
566                        continue;
567                    }
568                    // A failed turn is reported from its own envelope; only a
569                    // completed or incomplete one reaches the decoder, whose
570                    // end closes the turn.
571                    let response = terminal_response_result(response)?;
572                    let ended = feed(&mut decoder, &mut document, &reply, payload)?;
573                    let folded = fold_reply(reply, ended, &provider, document)?;
574                    return Ok((response, folded));
575                }
576                ResponsesWebSocketEvent::Done(done) => {
577                    if done.response.is_object() {
578                        // A failed turn is reported from its own envelope, as
579                        // on the `response.failed` path.
580                        let response = terminal_response_result(done.response)?;
581                        // `response.done` carries the response object itself,
582                        // which is the decoder's whole-body shape: hand it
583                        // over as the frame it is.
584                        let body = serde_json::to_string(&response)?;
585                        let ended = feed(&mut decoder, &mut document, &reply, body)?;
586                        let folded = fold_reply(reply, ended, &provider, document)?;
587                        return Ok((response, folded));
588                    }
589
590                    let message = if let Some(response_id) = done.response_id() {
591                        format!(
592                            "OpenAI websocket turn ended with response.done before a terminal response body was available (response_id={response_id})"
593                        )
594                    } else {
595                        "OpenAI websocket turn ended with response.done before a terminal response body was available"
596                            .to_string()
597                    };
598
599                    return Err(ProviderError::Provider(message));
600                }
601                ResponsesWebSocketEvent::Error(error) => {
602                    // Genuine provider error event: preserve the serialized payload
603                    // (code + message + any extra fields) so provider_response_json()
604                    // parses it, matching the response.failed path. No HTTP status on
605                    // the websocket stream, so status: None.
606                    return Err(provider_error_from_event(&error));
607                }
608                // Unknown frames keep their raw payload.
609                ResponsesWebSocketEvent::Item(_) | ResponsesWebSocketEvent::Unknown(_) => {
610                    feed(&mut decoder, &mut document, &reply, payload)?;
611                }
612            }
613        }
614    }
615
616    fn update_state_for_event(&mut self, event: &ResponsesWebSocketEvent) {
617        match event {
618            ResponsesWebSocketEvent::Response { kind, response } => {
619                let response_id = response
620                    .get("id")
621                    .and_then(Value::as_str)
622                    .map(str::to_owned);
623                match kind.as_str() {
624                    // An incomplete turn still produced a response the next
625                    // turn can chain from, so it keeps
626                    // `previous_response_id` like a completed one.
627                    "response.completed" | "response.incomplete" => {
628                        self.previous_response_id.clone_from(&response_id);
629                        self.pending_done_response_id = response_id;
630                        self.in_flight = false;
631                    }
632                    "response.failed" => {
633                        self.pending_done_response_id = response_id;
634                        self.previous_response_id = None;
635                        self.in_flight = false;
636                    }
637                    _ => {}
638                }
639            }
640            ResponsesWebSocketEvent::Done(done) => {
641                match done.status() {
642                    Some("completed" | "incomplete") => {
643                        if let Some(response_id) = done.response_id() {
644                            self.previous_response_id = Some(response_id.to_string());
645                        }
646                    }
647                    Some("in_progress" | "queued") | None => {}
648                    Some(_) => {
649                        self.previous_response_id = None;
650                    }
651                }
652                self.pending_done_response_id = None;
653                self.in_flight = false;
654            }
655            ResponsesWebSocketEvent::Error(_) => {
656                self.previous_response_id = None;
657                self.pending_done_response_id = None;
658                self.in_flight = false;
659            }
660            // An unknown frame carries no turn-lifecycle signal.
661            ResponsesWebSocketEvent::Item(_) | ResponsesWebSocketEvent::Unknown(_) => {}
662        }
663    }
664
665    fn abort_turn(&mut self) {
666        self.turn = None;
667        self.previous_response_id = None;
668        self.pending_done_response_id = None;
669        self.in_flight = false;
670    }
671
672    fn mark_closed(&mut self) {
673        self.abort_turn();
674        self.closed = true;
675        self.failed = false;
676    }
677
678    fn mark_failed(&mut self) {
679        self.abort_turn();
680        self.failed = true;
681    }
682
683    fn ensure_open(&self) -> Result<(), ProviderError> {
684        if self.closed || self.failed {
685            return Err(ProviderError::Provider(
686                "The OpenAI websocket session is closed".to_string(),
687            ));
688        }
689
690        Ok(())
691    }
692
693    fn fail_session(&mut self, error: ProviderError) -> ProviderError {
694        self.mark_failed();
695        error
696    }
697
698    /// Read a frame with a WASM-compatible timeout.
699    /// Timeout failure marks the session failed; transport results remain nested.
700    async fn read_next_frame(
701        &mut self,
702    ) -> Result<http_client::Result<Option<Frame>>, ProviderError> {
703        let Some(timeout_duration) = self.event_timeout else {
704            return Ok(self.socket.recv().await);
705        };
706
707        match crate::wasm_compat::timeout(timeout_duration, self.socket.recv()).await {
708            Ok(message) => Ok(message),
709            Err(_) => Err(self.fail_session(event_timeout_error(timeout_duration))),
710        }
711    }
712}
713
714impl Drop for ResponsesWebSocketSession {
715    fn drop(&mut self) {
716        if !self.closed {
717            tracing::warn!(
718                target: "rig::completions",
719                in_flight = self.in_flight,
720                "Dropping an OpenAI websocket session without calling close(); the connection will end without a close handshake"
721            );
722        }
723    }
724}
725
726/// Feed one message to the turn's reassembler and decoder. Returns whether
727/// it ended the turn.
728fn feed(
729    decoder: &mut ResponsesDecoder,
730    document: &mut document::Response,
731    reply: &std::sync::Mutex<Shared<Completion>>,
732    payload: String,
733) -> Result<bool, ProviderError> {
734    crate::driver::step(
735        decoder,
736        Some(document),
737        reply,
738        WireFrame::Text(payload),
739        None,
740    )
741    .map(|step| matches!(step, Flow::Ended(_)))
742}
743
744/// Fold the turn's events and end into the normalized response, whose raw
745/// is the `Response` the turn's events rebuild, as a stream's is.
746fn fold_reply(
747    reply: std::sync::Mutex<Shared<Completion>>,
748    ended: bool,
749    provider: &str,
750    document: document::Response,
751) -> Result<completion::CompletionResponse, ProviderError> {
752    let fed = if ended {
753        Ok(())
754    } else {
755        Err(ProviderError::Truncated)
756    };
757    let reply_of = Reply {
758        provider: provider.to_owned(),
759        raw: crate::wire::document::Reassemble::finish(document),
760        // The websocket carries no reply headers past the handshake.
761        provider_request_id: None,
762    };
763    crate::driver::settle(reply, fed, reply_of).outcome
764}
765
766fn terminal_response_result(response: Value) -> Result<Value, ProviderError> {
767    match response.get("status").and_then(Value::as_str) {
768        // An incomplete response (e.g. hitting `max_output_tokens`) is a
769        // genuine terminal: the partial output and usage are kept, and the
770        // decoder maps its status and reason to the finish.
771        Some("completed" | "incomplete") | None => Ok(response),
772        // Preserve provider error envelopes as reserialized JSON without an HTTP status.
773        // Without an error object, return a local diagnostic instead.
774        Some("failed") => match response.get("error").filter(|error| !error.is_null()) {
775            Some(_) => Err(ProviderError::from_provider_body(response.to_string())),
776            None => Err(ProviderError::Provider(response_error_message(
777                "failed response",
778            ))),
779        },
780        Some(other) => Err(ProviderError::Provider(format!(
781            "OpenAI websocket response ended in state `{other}`"
782        ))),
783    }
784}
785
786fn response_error_message(fallback: &str) -> String {
787    format!("OpenAI websocket returned a {fallback}")
788}
789
790/// Preserve an error event as reserialized provider JSON without an HTTP status.
791/// Fall back to its display text if serialization fails.
792fn provider_error_from_event(error: &ResponsesWebSocketErrorEvent) -> ProviderError {
793    ProviderError::from_provider_body(
794        serde_json::to_string(&error).unwrap_or_else(|_| error.to_string()),
795    )
796}
797
798/// Decode WebSocket error and done events or delegate to Responses classification.
799/// Return parsing and triage errors; preserve unknown payloads.
800fn parse_server_event(payload: &str) -> Result<Option<ResponsesWebSocketEvent>, ProviderError> {
801    #[derive(Deserialize)]
802    struct EventType {
803        #[serde(rename = "type")]
804        kind: String,
805    }
806
807    let event_type = serde_json::from_str::<EventType>(payload)?;
808    match event_type.kind.as_str() {
809        "error" => serde_json::from_str(payload)
810            .map(|e| Some(ResponsesWebSocketEvent::Error(e)))
811            .map_err(ProviderError::from),
812        "response.done" => serde_json::from_str(payload)
813            .map(|d| Some(ResponsesWebSocketEvent::Done(d)))
814            .map_err(ProviderError::from),
815        _ => Ok(
816            match crate::driver::triage(classify_responses_payload(payload))? {
817                Item::Event(ResponsesEvent::Frame {
818                    kind, mut frame, ..
819                }) => {
820                    if is_lifecycle_event(&kind) {
821                        let response = frame
822                            .as_object_mut()
823                            .and_then(|frame| remove_preserving_order(frame, "response"))
824                            .unwrap_or_else(|| Value::Object(Map::new()));
825                        Some(ResponsesWebSocketEvent::Response { kind, response })
826                    } else {
827                        Some(ResponsesWebSocketEvent::Item(frame))
828                    }
829                }
830                // Only the event types above reach here: a frame without a
831                // `type` is not a websocket event.
832                Item::Event(ResponsesEvent::Whole(value)) => {
833                    Some(ResponsesWebSocketEvent::Unknown(value.into()))
834                }
835                Item::Event(ResponsesEvent::Failure(body)) => {
836                    return Err(ProviderError::from_provider_body(body));
837                }
838                Item::Event(ResponsesEvent::Sentinel) => None,
839                Item::Unknown(value) => Some(ResponsesWebSocketEvent::Unknown(value)),
840            },
841        ),
842    }
843}
844
845/// Lower one websocket frame onto the JSON payload the protocol carries.
846///
847/// `Ok(None)` is a frame with no protocol payload (a keepalive), which the
848/// session skips; a close frame mid-turn is an error naming the peer's reason.
849fn websocket_frame_to_text(frame: Frame) -> Result<Option<String>, ProviderError> {
850    match frame {
851        Frame::Text(text) => Ok(Some(text)),
852        Frame::Binary(bytes) => String::from_utf8(bytes.to_vec())
853            .map(Some)
854            .map_err(|error| ProviderError::Response(error.to_string())),
855        Frame::Ping(_) | Frame::Pong(_) => Ok(None),
856        Frame::Close(frame) => {
857            let reason = frame
858                .map(|frame| frame.reason)
859                .filter(|reason| !reason.is_empty())
860                .unwrap_or_else(|| "without a close reason".to_string());
861            Err(ProviderError::Provider(format!(
862                "The OpenAI websocket connection closed {reason}"
863            )))
864        }
865    }
866}
867
868/// Build the handshake request: the websocket URL derived from the client's
869/// base URL, carrying the client's own auth headers.
870///
871/// The backend supplies the websocket-specific handshake headers; this only
872/// states where to connect and who is connecting.
873fn websocket_request(wire: &Responses) -> Result<http_client::Request<NoBody>, EncodeError> {
874    let url = crate::ws_client::websocket_url(&wire.provider.base_url, WEBSOCKET_PATH)
875        .map_err(EncodeError::request)?;
876
877    let request = wire.provider.headers_in_session(
878        http_client::Request::builder()
879            .method(http::Method::GET)
880            .uri(url),
881        None,
882    );
883
884    request.body(NoBody).map_err(|error| {
885        EncodeError::request(format!("Failed to build OpenAI websocket request: {error}"))
886    })
887}
888
889fn event_timeout_error(timeout: Duration) -> ProviderError {
890    ProviderError::Provider(format!(
891        "Timed out waiting for the next OpenAI websocket event after {timeout:?}"
892    ))
893}
894
895/// Convert transport errors, retaining rejected-upgrade status, body, and request ID.
896/// Failures without a provider response retain transport error classification.
897fn websocket_provider_error(error: http_client::Error) -> ProviderError {
898    let provider_request_id = error.non_success_headers().and_then(|headers| {
899        crate::driver::http_transport::request_id_from(headers, REQUEST_ID_HEADER)
900    });
901    ProviderError::from_transport_error(error).with_provider_request_id(provider_request_id)
902}
903
904impl<T> Model<Responses, T> {
905    /// Start configuring a Responses WebSocket session for this model. Open
906    /// it with [`connect_with`](ResponsesWebSocketSessionBuilder::connect_with)
907    /// and a backend, or, under the `tungstenite` feature, with
908    /// [`connect`](ResponsesWebSocketSessionBuilder::connect).
909    pub fn responses_websocket(&self) -> ResponsesWebSocketSessionBuilder {
910        ResponsesWebSocketSessionBuilder::new(self.wire.clone())
911    }
912}
913
914/// Native sessions and builders satisfy Send and Sync.
915#[cfg(not(target_family = "wasm"))]
916const _: fn() = || {
917    fn assert_send_sync<T: Send + Sync>() {}
918    assert_send_sync::<ResponsesWebSocketSession>();
919    assert_send_sync::<ResponsesWebSocketSessionBuilder>();
920};
921
922#[cfg(test)]
923#[allow(clippy::expect_used, clippy::panic)]
924mod tests;
rig_core/providers/openai/responses_api/websocket.rs

rig_core/providers/openai/responses_api/
websocket.rs
