Skip to main content

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