Skip to main content

rig_core/providers/openai/responses_api/
websocket.rs

1//! WebSocket session support for the OpenAI Responses API.
2//!
3//! This module implements OpenAI's `/v1/responses` WebSocket mode as a stateful,
4//! sequential session. Each connection supports a single in-flight response at a
5//! time, which matches OpenAI's current protocol constraints.
6
7use crate::completion::{self, CompletionError};
8use crate::http_client::HttpClientExt;
9use crate::providers::openai::responses_api::streaming::{
10    ItemChunk, ResponseChunk, ResponseChunkKind, StreamingCompletionChunk,
11};
12use crate::wasm_compat::{WasmCompatSend, WasmCompatSync};
13use futures::{SinkExt, StreamExt};
14use serde::{Deserialize, Serialize};
15use serde_json::{Map, Value};
16use std::time::Duration;
17use tokio::net::TcpStream;
18use tokio_tungstenite::{
19    MaybeTlsStream, WebSocketStream, connect_async,
20    tungstenite::{self, Message, client::IntoClientRequest},
21};
22use tracing::Level;
23use url::Url;
24
25use super::{CompletionResponse, ResponseStatus, ResponsesCompletionModel};
26
27type OpenAIWebSocket = WebSocketStream<MaybeTlsStream<TcpStream>>;
28const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(30);
29
30/// Options for a `response.create` message sent over OpenAI WebSocket mode.
31#[derive(Debug, Clone, Default, Serialize, Deserialize)]
32pub struct ResponsesWebSocketCreateOptions {
33    /// When set to `false`, OpenAI prepares request state without generating a model output.
34    ///
35    /// This is the "warmup" mode described in the OpenAI WebSocket mode guide.
36    #[serde(skip_serializing_if = "Option::is_none")]
37    pub generate: Option<bool>,
38}
39
40impl ResponsesWebSocketCreateOptions {
41    /// Creates warmup options equivalent to `generate: false`.
42    #[must_use]
43    pub fn warmup() -> Self {
44        Self {
45            generate: Some(false),
46        }
47    }
48}
49
50#[derive(Debug, Clone, Serialize)]
51struct ResponsesWebSocketClientEvent {
52    #[serde(rename = "type")]
53    kind: ResponsesWebSocketClientEventKind,
54    #[serde(flatten)]
55    request: super::CompletionRequest,
56    #[serde(skip_serializing_if = "Option::is_none")]
57    generate: Option<bool>,
58}
59
60#[derive(Debug, Clone, Serialize)]
61enum ResponsesWebSocketClientEventKind {
62    #[serde(rename = "response.create")]
63    ResponseCreate,
64}
65
66/// A protocol error event emitted by OpenAI WebSocket mode.
67#[derive(Debug, Clone, Serialize, Deserialize)]
68pub struct ResponsesWebSocketErrorEvent {
69    /// The event type.
70    #[serde(rename = "type")]
71    pub kind: ResponsesWebSocketErrorEventKind,
72    /// The provider error payload.
73    pub error: ResponsesWebSocketErrorPayload,
74}
75
76impl std::fmt::Display for ResponsesWebSocketErrorEvent {
77    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78        self.error.fmt(f)
79    }
80}
81
82/// The event kind for an OpenAI WebSocket protocol error.
83#[derive(Debug, Clone, Serialize, Deserialize)]
84pub enum ResponsesWebSocketErrorEventKind {
85    #[serde(rename = "error")]
86    Error,
87}
88
89/// The payload carried by an OpenAI WebSocket protocol error event.
90#[derive(Debug, Clone, Default, Serialize, Deserialize)]
91pub struct ResponsesWebSocketErrorPayload {
92    /// Provider-specific error code when supplied.
93    #[serde(skip_serializing_if = "Option::is_none")]
94    pub code: Option<String>,
95    /// Human-readable error message.
96    #[serde(skip_serializing_if = "Option::is_none")]
97    pub message: Option<String>,
98    /// Any extra fields supplied by the provider.
99    #[serde(flatten, default)]
100    pub extra: Map<String, Value>,
101}
102
103impl std::fmt::Display for ResponsesWebSocketErrorPayload {
104    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105        match (&self.code, &self.message) {
106            (Some(code), Some(message)) => write!(f, "{code}: {message}"),
107            (None, Some(message)) => f.write_str(message),
108            (Some(code), None) => f.write_str(code),
109            (None, None) => f.write_str("OpenAI websocket error"),
110        }
111    }
112}
113
114/// The optional `response.done` event emitted by OpenAI WebSocket mode.
115#[derive(Debug, Clone, Serialize, Deserialize)]
116pub struct ResponsesWebSocketDoneEvent {
117    /// The event type.
118    #[serde(rename = "type")]
119    pub kind: ResponsesWebSocketDoneEventKind,
120    /// The provider payload for the finished response.
121    pub response: Value,
122}
123
124impl ResponsesWebSocketDoneEvent {
125    /// Returns the response ID if the payload includes one.
126    #[must_use]
127    pub fn response_id(&self) -> Option<&str> {
128        self.response.get("id").and_then(Value::as_str)
129    }
130
131    fn status(&self) -> Option<ResponseStatus> {
132        self.response
133            .get("status")
134            .cloned()
135            .and_then(|status| serde_json::from_value(status).ok())
136    }
137
138    fn as_completion_response(&self) -> Option<CompletionResponse> {
139        serde_json::from_value(self.response.clone()).ok()
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 `response.completed`.
154    Response(Box<ResponseChunk>),
155    /// A streaming item/delta event such as `response.output_text.delta`.
156    Item(ItemChunk),
157    /// A protocol-level websocket error event.
158    Error(ResponsesWebSocketErrorEvent),
159    /// An optional `response.done` event emitted by OpenAI over WebSockets.
160    Done(ResponsesWebSocketDoneEvent),
161}
162
163impl ResponsesWebSocketEvent {
164    /// Returns the response ID when the event includes one.
165    #[must_use]
166    pub fn response_id(&self) -> Option<&str> {
167        match self {
168            Self::Response(chunk) => Some(&chunk.response.id),
169            Self::Done(done) => done.response_id(),
170            Self::Item(_) | Self::Error(_) => None,
171        }
172    }
173
174    /// Returns `true` when this event ends the current in-flight websocket turn.
175    #[must_use]
176    pub fn is_terminal(&self) -> bool {
177        match self {
178            Self::Response(chunk) => matches!(
179                chunk.kind,
180                ResponseChunkKind::ResponseCompleted
181                    | ResponseChunkKind::ResponseFailed
182                    | ResponseChunkKind::ResponseIncomplete
183            ),
184            Self::Error(_) | Self::Done(_) => true,
185            Self::Item(_) => false,
186        }
187    }
188}
189
190/// A builder for an OpenAI Responses WebSocket session.
191///
192/// The default builder applies a 30 second connection timeout and leaves the
193/// per-event timeout disabled.
194pub struct ResponsesWebSocketSessionBuilder<H = reqwest::Client> {
195    model: ResponsesCompletionModel<H>,
196    connect_timeout: Option<Duration>,
197    event_timeout: Option<Duration>,
198}
199
200impl<H> ResponsesWebSocketSessionBuilder<H> {
201    pub(crate) fn new(model: ResponsesCompletionModel<H>) -> Self {
202        Self {
203            model,
204            connect_timeout: Some(DEFAULT_CONNECT_TIMEOUT),
205            event_timeout: None,
206        }
207    }
208
209    /// Sets the timeout for establishing the websocket connection.
210    #[must_use]
211    pub fn connect_timeout(mut self, timeout: Duration) -> Self {
212        self.connect_timeout = Some(timeout);
213        self
214    }
215
216    /// Disables the websocket connection timeout.
217    #[must_use]
218    pub fn without_connect_timeout(mut self) -> Self {
219        self.connect_timeout = None;
220        self
221    }
222
223    /// Sets the timeout for waiting on the next websocket event.
224    #[must_use]
225    pub fn event_timeout(mut self, timeout: Duration) -> Self {
226        self.event_timeout = Some(timeout);
227        self
228    }
229
230    /// Disables the websocket event timeout.
231    #[must_use]
232    pub fn without_event_timeout(mut self) -> Self {
233        self.event_timeout = None;
234        self
235    }
236}
237
238impl<H> ResponsesWebSocketSessionBuilder<H>
239where
240    H: HttpClientExt
241        + Clone
242        + std::fmt::Debug
243        + Default
244        + WasmCompatSend
245        + WasmCompatSync
246        + 'static,
247{
248    /// Opens the websocket session using the configured builder options.
249    pub async fn connect(self) -> Result<ResponsesWebSocketSession<H>, CompletionError> {
250        ResponsesWebSocketSession::connect_with_timeouts(
251            self.model,
252            self.connect_timeout,
253            self.event_timeout,
254        )
255        .await
256    }
257}
258
259/// A stateful OpenAI Responses WebSocket session.
260///
261/// This session keeps track of the most recent successful `response.id` so later
262/// turns can automatically chain via `previous_response_id` unless the request
263/// explicitly sets a different one.
264///
265/// Call [`ResponsesWebSocketSession::close`] when you are finished with the
266/// session so the websocket can complete a close handshake cleanly.
267pub struct ResponsesWebSocketSession<H = reqwest::Client> {
268    model: ResponsesCompletionModel<H>,
269    previous_response_id: Option<String>,
270    pending_done_response_id: Option<String>,
271    socket: OpenAIWebSocket,
272    in_flight: bool,
273    event_timeout: Option<Duration>,
274    closed: bool,
275    failed: bool,
276}
277
278impl<H> ResponsesWebSocketSession<H>
279where
280    H: HttpClientExt
281        + Clone
282        + std::fmt::Debug
283        + Default
284        + WasmCompatSend
285        + WasmCompatSync
286        + 'static,
287{
288    async fn connect_with_timeouts(
289        model: ResponsesCompletionModel<H>,
290        connect_timeout: Option<Duration>,
291        event_timeout: Option<Duration>,
292    ) -> Result<Self, CompletionError> {
293        let url = websocket_url(model.client.base_url())?;
294        let request = websocket_request(&url, model.client.headers())?;
295        let socket = connect_websocket(request, connect_timeout).await?;
296
297        Ok(Self {
298            model,
299            previous_response_id: None,
300            pending_done_response_id: None,
301            socket,
302            in_flight: false,
303            event_timeout,
304            closed: false,
305            failed: false,
306        })
307    }
308
309    /// Returns the most recent successful `response.id` tracked by this session.
310    #[must_use]
311    pub fn previous_response_id(&self) -> Option<&str> {
312        self.previous_response_id.as_deref()
313    }
314
315    /// Clears the cached `previous_response_id` so the next turn starts a fresh chain.
316    pub fn clear_previous_response_id(&mut self) {
317        self.previous_response_id = None;
318    }
319
320    /// Sends a `response.create` event for a Rig completion request.
321    pub async fn send(
322        &mut self,
323        completion_request: crate::completion::CompletionRequest,
324    ) -> Result<(), CompletionError> {
325        self.send_with_options(
326            completion_request,
327            ResponsesWebSocketCreateOptions::default(),
328        )
329        .await
330    }
331
332    /// Sends a `response.create` event with explicit websocket-mode options.
333    pub async fn send_with_options(
334        &mut self,
335        completion_request: crate::completion::CompletionRequest,
336        options: ResponsesWebSocketCreateOptions,
337    ) -> Result<(), CompletionError> {
338        self.ensure_open()?;
339
340        if self.in_flight {
341            return Err(CompletionError::ProviderError(
342                "An OpenAI websocket response is already in flight on this session".to_string(),
343            ));
344        }
345
346        let payload = ResponsesWebSocketClientEvent {
347            kind: ResponsesWebSocketClientEventKind::ResponseCreate,
348            request: self.prepare_request(completion_request)?,
349            generate: options.generate,
350        };
351
352        if tracing::enabled!(Level::TRACE) {
353            tracing::trace!(
354                target: "rig::completions",
355                "OpenAI websocket request: {}",
356                serde_json::to_string_pretty(&payload)?
357            );
358        }
359
360        let payload = serde_json::to_string(&payload)?;
361
362        if let Err(error) = self.socket.send(Message::text(payload)).await {
363            return Err(self.fail_session(websocket_provider_error(error)));
364        }
365        self.in_flight = true;
366
367        Ok(())
368    }
369
370    /// Reads the next server event for the current in-flight turn.
371    pub async fn next_event(&mut self) -> Result<ResponsesWebSocketEvent, CompletionError> {
372        self.ensure_open()?;
373
374        if !self.in_flight {
375            return Err(CompletionError::ProviderError(
376                "No OpenAI websocket response is currently in flight on this session".to_string(),
377            ));
378        }
379
380        loop {
381            let message = match self.read_next_message().await {
382                Ok(message) => message,
383                Err(error) => return Err(error),
384            };
385
386            let Some(message) = message else {
387                self.mark_closed();
388                return Err(CompletionError::ProviderError(
389                    "The OpenAI websocket connection closed before the turn finished".to_string(),
390                ));
391            };
392
393            let message = match message {
394                Ok(message) => message,
395                Err(error) => return Err(self.fail_session(websocket_provider_error(error))),
396            };
397            let payload = match websocket_message_to_text(message) {
398                Ok(Some(payload)) => payload,
399                Ok(None) => continue,
400                Err(error) => return Err(self.fail_session(error)),
401            };
402            let event = match parse_server_event(&payload) {
403                Ok(Some(event)) => event,
404                Ok(None) => continue,
405                Err(error) => return Err(self.fail_session(error)),
406            };
407            if let ResponsesWebSocketEvent::Done(done) = &event {
408                // OpenAI may emit `response.done` after the turn has already ended at
409                // `response.completed`. Ignore that trailing event on the next turn.
410                if self.pending_done_response_id.as_deref() == done.response_id() {
411                    self.pending_done_response_id = None;
412                    continue;
413                }
414            }
415            self.update_state_for_event(&event);
416            return Ok(event);
417        }
418    }
419
420    /// Sends a warmup turn (`generate: false`) and returns the resulting response ID.
421    pub async fn warmup(
422        &mut self,
423        completion_request: crate::completion::CompletionRequest,
424    ) -> Result<String, CompletionError> {
425        self.send_with_options(
426            completion_request,
427            ResponsesWebSocketCreateOptions::warmup(),
428        )
429        .await?;
430        let response = self.wait_for_completed_response().await?;
431        Ok(response.id)
432    }
433
434    /// Sends a completion turn and collects the final OpenAI response.
435    pub async fn completion(
436        &mut self,
437        completion_request: crate::completion::CompletionRequest,
438    ) -> Result<completion::CompletionResponse<CompletionResponse>, CompletionError> {
439        self.send(completion_request).await?;
440        let response = self.wait_for_completed_response().await?;
441        response.try_into()
442    }
443
444    /// Closes the websocket connection.
445    ///
446    /// Call this when you are finished with the session so the websocket can
447    /// terminate with a clean close handshake.
448    pub async fn close(&mut self) -> Result<(), CompletionError> {
449        if self.closed {
450            return Ok(());
451        }
452
453        let result = self
454            .socket
455            .close(None)
456            .await
457            .map_err(websocket_provider_error);
458        self.mark_closed();
459        result
460    }
461
462    fn prepare_request(
463        &self,
464        completion_request: crate::completion::CompletionRequest,
465    ) -> Result<super::CompletionRequest, CompletionError> {
466        let mut request = self.model.create_completion_request(completion_request)?;
467
468        // WebSocket mode is always event-driven, so these HTTP/SSE-specific flags
469        // are ignored by the provider and only add noise to the payload.
470        request.stream = None;
471        request.additional_parameters.background = None;
472
473        if request.additional_parameters.previous_response_id.is_none() {
474            request.additional_parameters.previous_response_id = self.previous_response_id.clone();
475        }
476
477        Ok(request)
478    }
479
480    async fn wait_for_completed_response(&mut self) -> Result<CompletionResponse, CompletionError> {
481        loop {
482            match self.next_event().await? {
483                ResponsesWebSocketEvent::Response(chunk) => {
484                    if matches!(
485                        chunk.kind,
486                        ResponseChunkKind::ResponseCompleted
487                            | ResponseChunkKind::ResponseFailed
488                            | ResponseChunkKind::ResponseIncomplete
489                    ) {
490                        return terminal_response_result(chunk.response);
491                    }
492                }
493                ResponsesWebSocketEvent::Done(done) => {
494                    if let Some(response) = done.as_completion_response() {
495                        return terminal_response_result(response);
496                    }
497
498                    let message = if let Some(response_id) = done.response_id() {
499                        format!(
500                            "OpenAI websocket turn ended with response.done before a terminal response body was available (response_id={response_id})"
501                        )
502                    } else {
503                        "OpenAI websocket turn ended with response.done before a terminal response body was available"
504                            .to_string()
505                    };
506
507                    return Err(CompletionError::ProviderError(message));
508                }
509                ResponsesWebSocketEvent::Error(error) => {
510                    // Genuine provider error event: preserve the serialized payload
511                    // (code + message + any extra fields) so provider_response_json()
512                    // parses it, matching the response.failed path. No HTTP status on
513                    // the websocket stream, so status: None.
514                    return Err(provider_error_from_event(error));
515                }
516                ResponsesWebSocketEvent::Item(_) => {}
517            }
518        }
519    }
520
521    fn update_state_for_event(&mut self, event: &ResponsesWebSocketEvent) {
522        match event {
523            ResponsesWebSocketEvent::Response(chunk) => match chunk.kind {
524                ResponseChunkKind::ResponseCompleted => {
525                    let response_id = chunk.response.id.clone();
526                    self.previous_response_id = Some(response_id.clone());
527                    self.pending_done_response_id = Some(response_id);
528                    self.in_flight = false;
529                }
530                ResponseChunkKind::ResponseFailed | ResponseChunkKind::ResponseIncomplete => {
531                    self.pending_done_response_id = Some(chunk.response.id.clone());
532                    self.previous_response_id = None;
533                    self.in_flight = false;
534                }
535                ResponseChunkKind::ResponseCreated | ResponseChunkKind::ResponseInProgress => {}
536            },
537            ResponsesWebSocketEvent::Done(done) => {
538                match done.status() {
539                    Some(ResponseStatus::Completed) => {
540                        if let Some(response_id) = done.response_id() {
541                            self.previous_response_id = Some(response_id.to_string());
542                        }
543                    }
544                    Some(ResponseStatus::Failed)
545                    | Some(ResponseStatus::Incomplete)
546                    | Some(ResponseStatus::Cancelled) => {
547                        self.previous_response_id = None;
548                    }
549                    Some(ResponseStatus::InProgress | ResponseStatus::Queued) | None => {}
550                }
551                self.pending_done_response_id = None;
552                self.in_flight = false;
553            }
554            ResponsesWebSocketEvent::Error(_) => {
555                self.previous_response_id = None;
556                self.pending_done_response_id = None;
557                self.in_flight = false;
558            }
559            ResponsesWebSocketEvent::Item(_) => {}
560        }
561    }
562
563    fn abort_turn(&mut self) {
564        self.previous_response_id = None;
565        self.pending_done_response_id = None;
566        self.in_flight = false;
567    }
568
569    fn mark_closed(&mut self) {
570        self.abort_turn();
571        self.closed = true;
572        self.failed = false;
573    }
574
575    fn mark_failed(&mut self) {
576        self.abort_turn();
577        self.failed = true;
578    }
579
580    fn ensure_open(&self) -> Result<(), CompletionError> {
581        if self.closed || self.failed {
582            return Err(CompletionError::ProviderError(
583                "The OpenAI websocket session is closed".to_string(),
584            ));
585        }
586
587        Ok(())
588    }
589
590    fn fail_session(&mut self, error: CompletionError) -> CompletionError {
591        self.mark_failed();
592        error
593    }
594
595    async fn read_next_message(
596        &mut self,
597    ) -> Result<Option<Result<Message, tungstenite::Error>>, CompletionError> {
598        if let Some(timeout_duration) = self.event_timeout {
599            match tokio::time::timeout(timeout_duration, self.socket.next()).await {
600                Ok(message) => Ok(message),
601                Err(_) => Err(self.fail_session(event_timeout_error(timeout_duration))),
602            }
603        } else {
604            Ok(self.socket.next().await)
605        }
606    }
607}
608
609impl<H> Drop for ResponsesWebSocketSession<H> {
610    fn drop(&mut self) {
611        if !self.closed {
612            tracing::warn!(
613                target: "rig::completions",
614                in_flight = self.in_flight,
615                "Dropping an OpenAI websocket session without calling close(); the connection will end without a close handshake"
616            );
617        }
618    }
619}
620
621fn terminal_response_result(
622    response: CompletionResponse,
623) -> Result<CompletionResponse, CompletionError> {
624    match response.status {
625        ResponseStatus::Completed => Ok(response),
626        // Deliberate two-tier behaviour: when the provider supplies its own error
627        // object we preserve the full failed-response envelope through
628        // `from_provider_body` (status: None, no HTTP status on the websocket
629        // stream) so `provider_response_json()` parses it — consistent with the
630        // `error` event and the streaming paths. The body is re-serialized from
631        // the parsed response (not byte-identical to the wire bytes, which aren't
632        // retained past parsing) — semantically the provider's payload. When the
633        // object is absent we have nothing provider-authored to surface, so we
634        // emit a Rig-authored `ProviderError` diagnostic (provider_response_body()
635        // is None).
636        ResponseStatus::Failed => match response.error.as_ref() {
637            Some(error) => Err(CompletionError::from_provider_body(
638                serde_json::to_string(&response).unwrap_or_else(|_| error.message.clone()),
639            )),
640            None => Err(CompletionError::ProviderError(response_error_message(
641                "failed response",
642            ))),
643        },
644        ResponseStatus::Incomplete => {
645            let reason = response
646                .incomplete_details
647                .as_ref()
648                .map(|details| details.reason.as_str())
649                .unwrap_or("unknown reason");
650            Err(CompletionError::ProviderError(format!(
651                "OpenAI websocket response was incomplete: {reason}"
652            )))
653        }
654        other => Err(CompletionError::ProviderError(format!(
655            "OpenAI websocket response ended in state {other:?}"
656        ))),
657    }
658}
659
660fn response_error_message(fallback: &str) -> String {
661    format!("OpenAI websocket returned a {fallback}")
662}
663
664/// Maps a provider `error` event into a [`CompletionError`] that preserves the
665/// raw error payload as JSON (code + message + any extra provider fields) so the
666/// `provider_response_*` helpers can inspect it. The websocket stream carries no
667/// HTTP status, so `status` is `None`. The body is the event re-serialized from
668/// the parsed representation (not byte-identical to the original wire bytes,
669/// which are not retained past parsing) — semantically the provider's payload.
670fn provider_error_from_event(error: ResponsesWebSocketErrorEvent) -> CompletionError {
671    CompletionError::from_provider_body(
672        serde_json::to_string(&error).unwrap_or_else(|_| error.to_string()),
673    )
674}
675
676fn is_known_streaming_event(kind: &str) -> bool {
677    matches!(
678        kind,
679        "response.created"
680            | "response.in_progress"
681            | "response.completed"
682            | "response.failed"
683            | "response.incomplete"
684            | "response.output_item.added"
685            | "response.output_item.done"
686            | "response.content_part.added"
687            | "response.content_part.done"
688            | "response.output_text.delta"
689            | "response.output_text.done"
690            | "response.refusal.delta"
691            | "response.refusal.done"
692            | "response.function_call_arguments.delta"
693            | "response.function_call_arguments.done"
694            | "response.reasoning_summary_part.added"
695            | "response.reasoning_summary_part.done"
696            | "response.reasoning_summary_text.delta"
697            | "response.reasoning_summary_text.done"
698    )
699}
700
701fn parse_server_event(payload: &str) -> Result<Option<ResponsesWebSocketEvent>, CompletionError> {
702    #[derive(Deserialize)]
703    struct EventType {
704        #[serde(rename = "type")]
705        kind: String,
706    }
707
708    let event_type = serde_json::from_str::<EventType>(payload)?;
709    match event_type.kind.as_str() {
710        "error" => serde_json::from_str(payload)
711            .map(|e| Some(ResponsesWebSocketEvent::Error(e)))
712            .map_err(CompletionError::from),
713        "response.done" => serde_json::from_str(payload)
714            .map(|d| Some(ResponsesWebSocketEvent::Done(d)))
715            .map_err(CompletionError::from),
716        kind if is_known_streaming_event(kind) => match serde_json::from_str(payload)? {
717            StreamingCompletionChunk::Response(response) => {
718                Ok(Some(ResponsesWebSocketEvent::Response(response)))
719            }
720            StreamingCompletionChunk::Delta(item) => Ok(Some(ResponsesWebSocketEvent::Item(item))),
721        },
722        _ => {
723            tracing::debug!(
724                target: "rig::completions",
725                event_type = event_type.kind.as_str(),
726                "Skipping unrecognised OpenAI websocket event"
727            );
728            Ok(None)
729        }
730    }
731}
732
733fn websocket_message_to_text(message: Message) -> Result<Option<String>, CompletionError> {
734    match message {
735        Message::Text(text) => Ok(Some(text.to_string())),
736        Message::Binary(bytes) => String::from_utf8(bytes.to_vec())
737            .map(Some)
738            .map_err(|error| CompletionError::ResponseError(error.to_string())),
739        Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => Ok(None),
740        Message::Close(frame) => {
741            let reason = frame
742                .map(|frame| frame.reason.to_string())
743                .filter(|reason| !reason.is_empty())
744                .unwrap_or_else(|| "without a close reason".to_string());
745            Err(CompletionError::ProviderError(format!(
746                "The OpenAI websocket connection closed {reason}"
747            )))
748        }
749    }
750}
751
752fn websocket_url(base_url: &str) -> Result<String, CompletionError> {
753    let mut url = Url::parse(base_url)?;
754    match url.scheme() {
755        "https" => {
756            url.set_scheme("wss").map_err(|_| {
757                CompletionError::ProviderError("Failed to convert https URL to wss".to_string())
758            })?;
759        }
760        "http" => {
761            url.set_scheme("ws").map_err(|_| {
762                CompletionError::ProviderError("Failed to convert http URL to ws".to_string())
763            })?;
764        }
765        scheme => {
766            return Err(CompletionError::ProviderError(format!(
767                "Unsupported base URL scheme for OpenAI websocket mode: {scheme}"
768            )));
769        }
770    }
771
772    let path = format!("{}/responses", url.path().trim_end_matches('/'));
773    url.set_path(&path);
774    Ok(url.to_string())
775}
776
777fn websocket_request(
778    url: &str,
779    headers: &http::HeaderMap,
780) -> Result<http::Request<()>, CompletionError> {
781    let mut request = url.into_client_request().map_err(|error| {
782        CompletionError::ProviderError(format!("Failed to build OpenAI websocket request: {error}"))
783    })?;
784
785    for (name, value) in headers {
786        request.headers_mut().insert(name, value.clone());
787    }
788
789    Ok(request)
790}
791
792async fn connect_websocket(
793    request: http::Request<()>,
794    connect_timeout: Option<Duration>,
795) -> Result<OpenAIWebSocket, CompletionError> {
796    if let Some(timeout_duration) = connect_timeout {
797        match tokio::time::timeout(timeout_duration, connect_async(request)).await {
798            Ok(result) => result
799                .map(|(socket, _)| socket)
800                .map_err(websocket_provider_error),
801            Err(_) => Err(connect_timeout_error(timeout_duration)),
802        }
803    } else {
804        connect_async(request)
805            .await
806            .map(|(socket, _)| socket)
807            .map_err(websocket_provider_error)
808    }
809}
810
811fn connect_timeout_error(timeout: Duration) -> CompletionError {
812    CompletionError::ProviderError(format!(
813        "Timed out connecting to the OpenAI websocket after {timeout:?}"
814    ))
815}
816
817fn event_timeout_error(timeout: Duration) -> CompletionError {
818    CompletionError::ProviderError(format!(
819        "Timed out waiting for the next OpenAI websocket event after {timeout:?}"
820    ))
821}
822
823fn websocket_provider_error(error: tungstenite::Error) -> CompletionError {
824    CompletionError::ProviderError(error.to_string())
825}
826
827#[cfg(test)]
828mod tests {
829    use super::{
830        ResponsesWebSocketCreateOptions, ResponsesWebSocketDoneEvent, ResponsesWebSocketEvent,
831        parse_server_event, terminal_response_result, websocket_url,
832    };
833    use crate::client::CompletionClient;
834    use crate::completion::CompletionModel;
835    use crate::providers::openai::responses_api::{
836        CompletionResponse, ResponseError, ResponseObject, ResponseStatus, ResponsesUsage,
837    };
838    use futures::{SinkExt, StreamExt};
839    use serde_json::json;
840    use std::time::Duration;
841    use tokio::net::TcpListener;
842    use tokio::time::sleep;
843    use tokio_tungstenite::{accept_async, tungstenite::Message};
844
845    #[test]
846    fn websocket_error_event_preserves_provider_payload_as_json() {
847        let mut extra = serde_json::Map::new();
848        extra.insert(
849            "type".to_string(),
850            serde_json::Value::String("invalid_request_error".to_string()),
851        );
852        let event = super::ResponsesWebSocketErrorEvent {
853            kind: super::ResponsesWebSocketErrorEventKind::Error,
854            error: super::ResponsesWebSocketErrorPayload {
855                code: Some("rate_limit_exceeded".to_string()),
856                message: Some("slow down".to_string()),
857                extra,
858            },
859        };
860
861        let err = super::provider_error_from_event(event);
862
863        // No HTTP status on the websocket stream, and the raw payload round-trips
864        // through provider_response_json() (code + message + extra all preserved).
865        assert_eq!(err.provider_response_status(), None);
866        let json = err
867            .provider_response_json()
868            .expect("preserved body should be valid JSON")
869            .expect("provider response body should be present");
870        assert_eq!(json["error"]["code"], "rate_limit_exceeded");
871        assert_eq!(json["error"]["message"], "slow down");
872        assert_eq!(json["error"]["type"], "invalid_request_error");
873    }
874
875    fn sample_response(status: ResponseStatus) -> CompletionResponse {
876        CompletionResponse {
877            id: "resp_123".to_string(),
878            object: ResponseObject::Response,
879            created_at: 0,
880            status,
881            error: None,
882            incomplete_details: None,
883            instructions: None,
884            max_output_tokens: None,
885            model: "gpt-5.4".to_string(),
886            usage: Some(ResponsesUsage {
887                input_tokens: 1,
888                input_tokens_details: None,
889                output_tokens: 2,
890                output_tokens_details: Some(
891                    crate::providers::openai::responses_api::OutputTokensDetails {
892                        reasoning_tokens: 0,
893                    },
894                ),
895                total_tokens: 3,
896            }),
897            output: Vec::new(),
898            tools: Vec::new(),
899            additional_parameters: Default::default(),
900            provider_reasoning: None,
901        }
902    }
903
904    #[test]
905    fn warmup_options_serialize_generate_false() {
906        let options = ResponsesWebSocketCreateOptions::warmup();
907        let json = serde_json::to_value(options).expect("options should serialize");
908
909        assert_eq!(json, json!({ "generate": false }));
910    }
911
912    #[test]
913    fn websocket_url_converts_https_to_wss() {
914        let url = websocket_url("https://api.openai.com/v1").expect("url should convert");
915        assert_eq!(url, "wss://api.openai.com/v1/responses");
916    }
917
918    #[test]
919    fn parse_done_event_exposes_response_id() {
920        let payload = json!({
921            "type": "response.done",
922            "response": {
923                "id": "resp_done_1",
924                "status": "completed"
925            }
926        });
927
928        let event = parse_server_event(&payload.to_string())
929            .expect("done event should deserialize")
930            .expect("done event should not be skipped");
931
932        assert!(matches!(
933            event,
934            ResponsesWebSocketEvent::Done(ResponsesWebSocketDoneEvent { .. })
935        ));
936        assert_eq!(event.response_id(), Some("resp_done_1"));
937        assert!(event.is_terminal());
938    }
939
940    #[test]
941    fn parse_response_completed_event_is_terminal() {
942        let payload = json!({
943            "type": "response.completed",
944            "sequence_number": 12,
945            "response": {
946                "id": "resp_completed_1",
947                "object": "response",
948                "created_at": 0,
949                "status": "completed",
950                "error": null,
951                "incomplete_details": null,
952                "instructions": null,
953                "max_output_tokens": null,
954                "model": "gpt-5.4",
955                "usage": null,
956                "output": [],
957                "tools": []
958            }
959        });
960
961        let event = parse_server_event(&payload.to_string())
962            .expect("response event should deserialize")
963            .expect("response event should not be skipped");
964
965        assert!(matches!(event, ResponsesWebSocketEvent::Response(_)));
966        assert!(event.is_terminal());
967        assert_eq!(event.response_id(), Some("resp_completed_1"));
968    }
969
970    #[test]
971    fn parse_live_output_item_added_event() {
972        let payload = json!({
973            "type": "response.output_item.added",
974            "item": {
975                "id": "msg_036471c3a72c147b0069ae7848d68881959773fd2d99e3d98a",
976                "type": "message",
977                "status": "in_progress",
978                "content": [],
979                "role": "assistant"
980            },
981            "output_index": 0,
982            "sequence_number": 2
983        });
984
985        let event = parse_server_event(&payload.to_string())
986            .expect("output item event should parse")
987            .expect("output item event should not be skipped");
988
989        assert!(matches!(event, ResponsesWebSocketEvent::Item(_)));
990    }
991
992    #[test]
993    fn parse_live_content_part_added_event() {
994        let payload = json!({
995            "type": "response.content_part.added",
996            "content_index": 0,
997            "item_id": "msg_036471c3a72c147b0069ae7848d68881959773fd2d99e3d98a",
998            "output_index": 0,
999            "part": {
1000                "type": "output_text",
1001                "annotations": [],
1002                "logprobs": [],
1003                "text": ""
1004            },
1005            "sequence_number": 3
1006        });
1007
1008        let event = parse_server_event(&payload.to_string())
1009            .expect("content part event should parse")
1010            .expect("content part event should not be skipped");
1011
1012        assert!(matches!(event, ResponsesWebSocketEvent::Item(_)));
1013    }
1014
1015    #[test]
1016    fn parse_live_output_text_delta_event() {
1017        let payload = json!({
1018            "type": "response.output_text.delta",
1019            "content_index": 0,
1020            "delta": "Web",
1021            "item_id": "msg_023af0f0a91bc2a90069ae788612e881958345bb156915ba29",
1022            "logprobs": [],
1023            "obfuscation": "2YYErYq7jkqqM",
1024            "output_index": 0,
1025            "sequence_number": 4
1026        });
1027
1028        let event = parse_server_event(&payload.to_string())
1029            .expect("output text delta event should parse")
1030            .expect("output text delta event should not be skipped");
1031
1032        assert!(matches!(event, ResponsesWebSocketEvent::Item(_)));
1033    }
1034
1035    #[test]
1036    fn terminal_response_requires_completed_status() {
1037        let completed = terminal_response_result(sample_response(ResponseStatus::Completed))
1038            .expect("completed response should succeed");
1039        assert_eq!(completed.id, "resp_123");
1040
1041        let failed = terminal_response_result(sample_response(ResponseStatus::Failed))
1042            .expect_err("failed response should error");
1043        assert!(failed.to_string().contains("failed response"));
1044    }
1045
1046    #[test]
1047    fn terminal_failed_response_with_error_preserves_raw_payload() {
1048        let mut response = sample_response(ResponseStatus::Failed);
1049        response.error = Some(ResponseError {
1050            code: "server_error".to_string(),
1051            message: "the model failed to generate a response".to_string(),
1052        });
1053
1054        let err = match terminal_response_result(response) {
1055            Ok(_) => panic!("failed response with an error object should fail"),
1056            Err(e) => e,
1057        };
1058
1059        // The full failed-response envelope is preserved as a ProviderResponse with
1060        // no HTTP status (the websocket stream carries none), so the raw JSON parses
1061        // back with the provider error nested under `error` — proving the whole
1062        // envelope is kept, not just the error object.
1063        assert_eq!(err.provider_response_status(), None);
1064
1065        let json = err
1066            .provider_response_json()
1067            .expect("preserved body should parse as JSON")
1068            .expect("preserved body should not be empty");
1069        assert_eq!(
1070            json["error"]["message"],
1071            "the model failed to generate a response"
1072        );
1073        assert_eq!(json["error"]["code"], "server_error");
1074    }
1075
1076    #[test]
1077    fn terminal_failed_response_without_error_is_rig_diagnostic() {
1078        let err = match terminal_response_result(sample_response(ResponseStatus::Failed)) {
1079            Ok(_) => panic!("failed response should fail"),
1080            Err(e) => e,
1081        };
1082
1083        // No provider error object, so this is a Rig-authored diagnostic and exposes
1084        // no preserved provider response body.
1085        assert_eq!(err.provider_response_body(), None);
1086        assert!(err.to_string().contains("failed response"));
1087    }
1088
1089    #[tokio::test]
1090    async fn malformed_known_event_rejects_reuse_and_allows_close() {
1091        let listener = TcpListener::bind("127.0.0.1:0")
1092            .await
1093            .expect("listener should bind");
1094        let address = listener.local_addr().expect("listener should have address");
1095
1096        let server = tokio::spawn(async move {
1097            let (stream, _) = listener.accept().await.expect("server should accept");
1098            let mut socket = accept_async(stream)
1099                .await
1100                .expect("server should upgrade websocket");
1101
1102            let request = socket
1103                .next()
1104                .await
1105                .expect("request should exist")
1106                .expect("request should be valid");
1107            let payload = request.into_text().expect("request should be text");
1108            assert!(
1109                payload.contains("\"type\":\"response.create\""),
1110                "expected response.create payload, got {payload}"
1111            );
1112
1113            socket
1114                .send(Message::text(
1115                    json!({
1116                        "type": "response.completed"
1117                    })
1118                    .to_string(),
1119                ))
1120                .await
1121                .expect("malformed known event should send");
1122
1123            let message = socket
1124                .next()
1125                .await
1126                .expect("close frame should arrive")
1127                .expect("close frame should be valid");
1128            assert!(
1129                matches!(message, Message::Close(_)),
1130                "expected close frame, got {message:?}"
1131            );
1132        });
1133
1134        let base_url = format!("http://{address}/v1");
1135        let client = crate::providers::openai::Client::builder()
1136            .api_key("test-key")
1137            .base_url(&base_url)
1138            .build()
1139            .expect("client should build");
1140        let model = client.completion_model("gpt-4o");
1141        let mut session = client
1142            .responses_websocket("gpt-4o")
1143            .await
1144            .expect("session should connect");
1145
1146        session
1147            .send(model.completion_request("hello").build())
1148            .await
1149            .expect("request should send");
1150
1151        let error = session
1152            .next_event()
1153            .await
1154            .expect_err("malformed known event should fail");
1155        assert!(
1156            error.to_string().contains("StreamingCompletionChunk"),
1157            "expected strict decode failure, got {error}"
1158        );
1159
1160        let closed = session
1161            .send(model.completion_request("retry").build())
1162            .await
1163            .expect_err("session should close after fatal parse error");
1164        assert!(
1165            closed.to_string().contains("session is closed"),
1166            "expected closed-session error, got {closed}"
1167        );
1168
1169        session
1170            .close()
1171            .await
1172            .expect("explicit close after fatal parse error should succeed");
1173
1174        server.await.expect("server task should finish");
1175    }
1176
1177    #[tokio::test]
1178    async fn event_timeout_rejects_reuse_and_allows_close() {
1179        let listener = TcpListener::bind("127.0.0.1:0")
1180            .await
1181            .expect("listener should bind");
1182        let address = listener.local_addr().expect("listener should have address");
1183
1184        let server = tokio::spawn(async move {
1185            let (stream, _) = listener.accept().await.expect("server should accept");
1186            let mut socket = accept_async(stream)
1187                .await
1188                .expect("server should upgrade websocket");
1189
1190            let request = socket
1191                .next()
1192                .await
1193                .expect("request should exist")
1194                .expect("request should be valid");
1195            let payload = request.into_text().expect("request should be text");
1196            assert!(
1197                payload.contains("\"type\":\"response.create\""),
1198                "expected response.create payload, got {payload}"
1199            );
1200
1201            sleep(Duration::from_millis(60)).await;
1202            let message = socket
1203                .next()
1204                .await
1205                .expect("close frame should arrive")
1206                .expect("close frame should be valid");
1207            assert!(
1208                matches!(message, Message::Close(_)),
1209                "expected close frame, got {message:?}"
1210            );
1211        });
1212
1213        let base_url = format!("http://{address}/v1");
1214        let client = crate::providers::openai::Client::builder()
1215            .api_key("test-key")
1216            .base_url(&base_url)
1217            .build()
1218            .expect("client should build");
1219        let model = client.completion_model("gpt-4o");
1220        let mut session = client
1221            .responses_websocket_builder("gpt-4o")
1222            .event_timeout(Duration::from_millis(20))
1223            .connect()
1224            .await
1225            .expect("session should connect");
1226
1227        session
1228            .send(model.completion_request("hello").build())
1229            .await
1230            .expect("request should send");
1231
1232        let error = session
1233            .next_event()
1234            .await
1235            .expect_err("next_event should time out");
1236        assert!(
1237            error
1238                .to_string()
1239                .contains("Timed out waiting for the next OpenAI websocket event"),
1240            "expected timeout error, got {error}"
1241        );
1242
1243        let closed = session
1244            .send(model.completion_request("retry").build())
1245            .await
1246            .expect_err("timed-out session should close");
1247        assert!(
1248            closed.to_string().contains("session is closed"),
1249            "expected closed-session error, got {closed}"
1250        );
1251
1252        session
1253            .close()
1254            .await
1255            .expect("explicit close after timeout should succeed");
1256
1257        server.await.expect("server task should finish");
1258    }
1259
1260    #[tokio::test]
1261    async fn late_response_done_is_ignored_on_next_turn() {
1262        let listener = TcpListener::bind("127.0.0.1:0")
1263            .await
1264            .expect("listener should bind");
1265        let address = listener.local_addr().expect("listener should have address");
1266
1267        let server = tokio::spawn(async move {
1268            let (stream, _) = listener.accept().await.expect("server should accept");
1269            let mut socket = accept_async(stream)
1270                .await
1271                .expect("server should upgrade websocket");
1272
1273            for (index, response_id) in ["resp_1", "resp_2"].iter().enumerate() {
1274                let request = socket
1275                    .next()
1276                    .await
1277                    .expect("request should exist")
1278                    .expect("request should be valid");
1279                let payload = request.into_text().expect("request should be text");
1280                assert!(
1281                    payload.contains("\"type\":\"response.create\""),
1282                    "expected response.create payload, got {payload}"
1283                );
1284
1285                let response = sample_response(ResponseStatus::Completed);
1286                let response = serde_json::to_value(CompletionResponse {
1287                    id: (*response_id).to_string(),
1288                    ..response
1289                })
1290                .expect("response should serialize");
1291
1292                socket
1293                    .send(Message::text(
1294                        json!({
1295                            "type": "response.completed",
1296                            "sequence_number": (index * 2) + 1,
1297                            "response": response,
1298                        })
1299                        .to_string(),
1300                    ))
1301                    .await
1302                    .expect("completed event should send");
1303                socket
1304                    .send(Message::text(
1305                        json!({
1306                            "type": "response.done",
1307                            "response": {
1308                                "id": response_id,
1309                                "status": "completed",
1310                            },
1311                        })
1312                        .to_string(),
1313                    ))
1314                    .await
1315                    .expect("done event should send");
1316            }
1317        });
1318
1319        let base_url = format!("http://{address}/v1");
1320        let client = crate::providers::openai::Client::builder()
1321            .api_key("test-key")
1322            .base_url(&base_url)
1323            .build()
1324            .expect("client should build");
1325        let model = client.completion_model("gpt-4o");
1326        let mut session = client
1327            .responses_websocket("gpt-4o")
1328            .await
1329            .expect("session should connect");
1330
1331        session
1332            .send(model.completion_request("first").build())
1333            .await
1334            .expect("first request should send");
1335        let first = session
1336            .wait_for_completed_response()
1337            .await
1338            .expect("first response should complete");
1339        assert_eq!(first.id, "resp_1");
1340        assert_eq!(session.previous_response_id(), Some("resp_1"));
1341
1342        session
1343            .send(model.completion_request("second").build())
1344            .await
1345            .expect("second request should send");
1346        let second = session
1347            .wait_for_completed_response()
1348            .await
1349            .expect("second response should complete");
1350        assert_eq!(second.id, "resp_2");
1351        assert_eq!(session.previous_response_id(), Some("resp_2"));
1352
1353        server.await.expect("server task should finish");
1354    }
1355
1356    #[tokio::test]
1357    async fn clearing_previous_response_id_does_not_disable_late_done_filter() {
1358        let listener = TcpListener::bind("127.0.0.1:0")
1359            .await
1360            .expect("listener should bind");
1361        let address = listener.local_addr().expect("listener should have address");
1362
1363        let server = tokio::spawn(async move {
1364            let (stream, _) = listener.accept().await.expect("server should accept");
1365            let mut socket = accept_async(stream)
1366                .await
1367                .expect("server should upgrade websocket");
1368
1369            for response_id in ["resp_1", "resp_2"] {
1370                let request = socket
1371                    .next()
1372                    .await
1373                    .expect("request should exist")
1374                    .expect("request should be valid");
1375                let payload = request.into_text().expect("request should be text");
1376                assert!(
1377                    payload.contains("\"type\":\"response.create\""),
1378                    "expected response.create payload, got {payload}"
1379                );
1380
1381                let response = sample_response(ResponseStatus::Completed);
1382                let response = serde_json::to_value(CompletionResponse {
1383                    id: response_id.to_string(),
1384                    ..response
1385                })
1386                .expect("response should serialize");
1387
1388                socket
1389                    .send(Message::text(
1390                        json!({
1391                            "type": "response.completed",
1392                            "sequence_number": 1,
1393                            "response": response,
1394                        })
1395                        .to_string(),
1396                    ))
1397                    .await
1398                    .expect("completed event should send");
1399                socket
1400                    .send(Message::text(
1401                        json!({
1402                            "type": "response.done",
1403                            "response": {
1404                                "id": response_id,
1405                                "status": "completed",
1406                            },
1407                        })
1408                        .to_string(),
1409                    ))
1410                    .await
1411                    .expect("done event should send");
1412            }
1413        });
1414
1415        let base_url = format!("http://{address}/v1");
1416        let client = crate::providers::openai::Client::builder()
1417            .api_key("test-key")
1418            .base_url(&base_url)
1419            .build()
1420            .expect("client should build");
1421        let model = client.completion_model("gpt-4o");
1422        let mut session = client
1423            .responses_websocket("gpt-4o")
1424            .await
1425            .expect("session should connect");
1426
1427        session
1428            .send(model.completion_request("first").build())
1429            .await
1430            .expect("first request should send");
1431        let first = session
1432            .wait_for_completed_response()
1433            .await
1434            .expect("first response should complete");
1435        assert_eq!(first.id, "resp_1");
1436
1437        session.clear_previous_response_id();
1438        assert_eq!(session.previous_response_id(), None);
1439
1440        session
1441            .send(model.completion_request("second").build())
1442            .await
1443            .expect("second request should send");
1444        let second = session
1445            .wait_for_completed_response()
1446            .await
1447            .expect("second response should complete");
1448        assert_eq!(second.id, "resp_2");
1449
1450        server.await.expect("server task should finish");
1451    }
1452
1453    #[tokio::test]
1454    async fn failed_turn_keeps_late_done_out_of_next_request() {
1455        let listener = TcpListener::bind("127.0.0.1:0")
1456            .await
1457            .expect("listener should bind");
1458        let address = listener.local_addr().expect("listener should have address");
1459
1460        let server = tokio::spawn(async move {
1461            let (stream, _) = listener.accept().await.expect("server should accept");
1462            let mut socket = accept_async(stream)
1463                .await
1464                .expect("server should upgrade websocket");
1465
1466            let first_request = socket
1467                .next()
1468                .await
1469                .expect("request should exist")
1470                .expect("request should be valid");
1471            let payload = first_request
1472                .into_text()
1473                .expect("failed request should be text");
1474            assert!(
1475                payload.contains("\"type\":\"response.create\""),
1476                "expected response.create payload, got {payload}"
1477            );
1478
1479            let failed_response = serde_json::to_value(CompletionResponse {
1480                id: "resp_failed".to_string(),
1481                status: ResponseStatus::Failed,
1482                ..sample_response(ResponseStatus::Completed)
1483            })
1484            .expect("failed response should serialize");
1485
1486            socket
1487                .send(Message::text(
1488                    json!({
1489                        "type": "response.failed",
1490                        "sequence_number": 1,
1491                        "response": failed_response,
1492                    })
1493                    .to_string(),
1494                ))
1495                .await
1496                .expect("failed event should send");
1497            socket
1498                .send(Message::text(
1499                    json!({
1500                        "type": "response.done",
1501                        "response": {
1502                            "id": "resp_failed",
1503                            "status": "failed",
1504                        },
1505                    })
1506                    .to_string(),
1507                ))
1508                .await
1509                .expect("done event should send");
1510
1511            let second_request = socket
1512                .next()
1513                .await
1514                .expect("request should exist")
1515                .expect("request should be valid");
1516            let payload = second_request
1517                .into_text()
1518                .expect("second request should be text");
1519            assert!(
1520                payload.contains("\"type\":\"response.create\""),
1521                "expected response.create payload, got {payload}"
1522            );
1523
1524            let response = sample_response(ResponseStatus::Completed);
1525            let response = serde_json::to_value(CompletionResponse {
1526                id: "resp_2".to_string(),
1527                ..response
1528            })
1529            .expect("response should serialize");
1530
1531            socket
1532                .send(Message::text(
1533                    json!({
1534                        "type": "response.completed",
1535                        "sequence_number": 2,
1536                        "response": response,
1537                    })
1538                    .to_string(),
1539                ))
1540                .await
1541                .expect("completed event should send");
1542            socket
1543                .send(Message::text(
1544                    json!({
1545                        "type": "response.done",
1546                        "response": {
1547                            "id": "resp_2",
1548                            "status": "completed",
1549                        },
1550                    })
1551                    .to_string(),
1552                ))
1553                .await
1554                .expect("done event should send");
1555        });
1556
1557        let base_url = format!("http://{address}/v1");
1558        let client = crate::providers::openai::Client::builder()
1559            .api_key("test-key")
1560            .base_url(&base_url)
1561            .build()
1562            .expect("client should build");
1563        let model = client.completion_model("gpt-4o");
1564        let mut session = client
1565            .responses_websocket("gpt-4o")
1566            .await
1567            .expect("session should connect");
1568
1569        session
1570            .send(model.completion_request("first").build())
1571            .await
1572            .expect("first request should send");
1573        let error = session
1574            .wait_for_completed_response()
1575            .await
1576            .expect_err("failed response should error");
1577        assert!(error.to_string().contains("failed response"));
1578        assert_eq!(session.previous_response_id(), None);
1579
1580        session
1581            .send(model.completion_request("second").build())
1582            .await
1583            .expect("second request should send");
1584        let second = session
1585            .wait_for_completed_response()
1586            .await
1587            .expect("second response should complete");
1588        assert_eq!(second.id, "resp_2");
1589
1590        server.await.expect("server task should finish");
1591    }
1592
1593    #[tokio::test]
1594    async fn done_first_completed_turn_updates_previous_response_id() {
1595        let listener = TcpListener::bind("127.0.0.1:0")
1596            .await
1597            .expect("listener should bind");
1598        let address = listener.local_addr().expect("listener should have address");
1599
1600        let server = tokio::spawn(async move {
1601            let (stream, _) = listener.accept().await.expect("server should accept");
1602            let mut socket = accept_async(stream)
1603                .await
1604                .expect("server should upgrade websocket");
1605
1606            for response_id in ["resp_1", "resp_2"] {
1607                let request = socket
1608                    .next()
1609                    .await
1610                    .expect("request should exist")
1611                    .expect("request should be valid");
1612                let payload = request.into_text().expect("request should be text");
1613                assert!(
1614                    payload.contains("\"type\":\"response.create\""),
1615                    "expected response.create payload, got {payload}"
1616                );
1617
1618                if response_id == "resp_2" {
1619                    assert!(
1620                        payload.contains("\"previous_response_id\":\"resp_1\""),
1621                        "expected chained previous_response_id in payload, got {payload}"
1622                    );
1623                }
1624
1625                let response = serde_json::to_value(CompletionResponse {
1626                    id: response_id.to_string(),
1627                    ..sample_response(ResponseStatus::Completed)
1628                })
1629                .expect("response should serialize");
1630
1631                socket
1632                    .send(Message::text(
1633                        json!({
1634                            "type": "response.done",
1635                            "response": response,
1636                        })
1637                        .to_string(),
1638                    ))
1639                    .await
1640                    .expect("done event should send");
1641            }
1642        });
1643
1644        let base_url = format!("http://{address}/v1");
1645        let client = crate::providers::openai::Client::builder()
1646            .api_key("test-key")
1647            .base_url(&base_url)
1648            .build()
1649            .expect("client should build");
1650        let model = client.completion_model("gpt-4o");
1651        let mut session = client
1652            .responses_websocket("gpt-4o")
1653            .await
1654            .expect("session should connect");
1655
1656        session
1657            .send(model.completion_request("first").build())
1658            .await
1659            .expect("first request should send");
1660        let first = session
1661            .wait_for_completed_response()
1662            .await
1663            .expect("first response should complete");
1664        assert_eq!(first.id, "resp_1");
1665        assert_eq!(session.previous_response_id(), Some("resp_1"));
1666
1667        session
1668            .send(model.completion_request("second").build())
1669            .await
1670            .expect("second request should send");
1671        let second = session
1672            .wait_for_completed_response()
1673            .await
1674            .expect("second response should complete");
1675        assert_eq!(second.id, "resp_2");
1676        assert_eq!(session.previous_response_id(), Some("resp_2"));
1677
1678        server.await.expect("server task should finish");
1679    }
1680
1681    #[tokio::test]
1682    async fn done_first_failed_turn_does_not_chain_next_request() {
1683        let listener = TcpListener::bind("127.0.0.1:0")
1684            .await
1685            .expect("listener should bind");
1686        let address = listener.local_addr().expect("listener should have address");
1687
1688        let server = tokio::spawn(async move {
1689            let (stream, _) = listener.accept().await.expect("server should accept");
1690            let mut socket = accept_async(stream)
1691                .await
1692                .expect("server should upgrade websocket");
1693
1694            let first_request = socket
1695                .next()
1696                .await
1697                .expect("request should exist")
1698                .expect("request should be valid");
1699            let payload = first_request
1700                .into_text()
1701                .expect("first request should be text");
1702            assert!(
1703                payload.contains("\"type\":\"response.create\""),
1704                "expected response.create payload, got {payload}"
1705            );
1706            assert!(
1707                !payload.contains("\"previous_response_id\""),
1708                "did not expect previous_response_id in first payload, got {payload}"
1709            );
1710
1711            let failed_response = serde_json::to_value(CompletionResponse {
1712                id: "resp_failed".to_string(),
1713                status: ResponseStatus::Failed,
1714                ..sample_response(ResponseStatus::Completed)
1715            })
1716            .expect("failed response should serialize");
1717
1718            socket
1719                .send(Message::text(
1720                    json!({
1721                        "type": "response.done",
1722                        "response": failed_response,
1723                    })
1724                    .to_string(),
1725                ))
1726                .await
1727                .expect("done event should send");
1728
1729            let second_request = socket
1730                .next()
1731                .await
1732                .expect("request should exist")
1733                .expect("request should be valid");
1734            let payload = second_request
1735                .into_text()
1736                .expect("second request should be text");
1737            assert!(
1738                payload.contains("\"type\":\"response.create\""),
1739                "expected response.create payload, got {payload}"
1740            );
1741            assert!(
1742                !payload.contains("\"previous_response_id\""),
1743                "did not expect chained previous_response_id in payload, got {payload}"
1744            );
1745
1746            let response = serde_json::to_value(CompletionResponse {
1747                id: "resp_2".to_string(),
1748                ..sample_response(ResponseStatus::Completed)
1749            })
1750            .expect("response should serialize");
1751
1752            socket
1753                .send(Message::text(
1754                    json!({
1755                        "type": "response.done",
1756                        "response": response,
1757                    })
1758                    .to_string(),
1759                ))
1760                .await
1761                .expect("done event should send");
1762        });
1763
1764        let base_url = format!("http://{address}/v1");
1765        let client = crate::providers::openai::Client::builder()
1766            .api_key("test-key")
1767            .base_url(&base_url)
1768            .build()
1769            .expect("client should build");
1770        let model = client.completion_model("gpt-4o");
1771        let mut session = client
1772            .responses_websocket("gpt-4o")
1773            .await
1774            .expect("session should connect");
1775
1776        session
1777            .send(model.completion_request("first").build())
1778            .await
1779            .expect("first request should send");
1780        let error = session
1781            .wait_for_completed_response()
1782            .await
1783            .expect_err("failed response should error");
1784        assert!(error.to_string().contains("failed response"));
1785        assert_eq!(session.previous_response_id(), None);
1786
1787        session
1788            .send(model.completion_request("second").build())
1789            .await
1790            .expect("second request should send");
1791        let second = session
1792            .wait_for_completed_response()
1793            .await
1794            .expect("second response should complete");
1795        assert_eq!(second.id, "resp_2");
1796        assert_eq!(session.previous_response_id(), Some("resp_2"));
1797
1798        server.await.expect("server task should finish");
1799    }
1800
1801    #[test]
1802    fn websocket_url_converts_http_to_ws() {
1803        let url = websocket_url("http://localhost:8080/v1").expect("url should convert");
1804        assert_eq!(url, "ws://localhost:8080/v1/responses");
1805    }
1806
1807    #[test]
1808    fn websocket_url_rejects_unsupported_scheme() {
1809        let result = websocket_url("ftp://example.com/v1");
1810        assert!(result.is_err());
1811    }
1812
1813    #[test]
1814    fn websocket_url_trims_trailing_slash() {
1815        let url = websocket_url("https://api.openai.com/v1/").expect("url should convert");
1816        assert_eq!(url, "wss://api.openai.com/v1/responses");
1817    }
1818
1819    #[test]
1820    fn unknown_event_type_is_skipped() {
1821        let payload = json!({
1822            "type": "response.some_future_event",
1823            "data": "hello"
1824        });
1825
1826        let result =
1827            parse_server_event(&payload.to_string()).expect("unknown event should not error");
1828        assert!(result.is_none(), "unknown event should be skipped");
1829    }
1830
1831    #[test]
1832    fn malformed_known_event_returns_error() {
1833        let payload = json!({
1834            "type": "response.completed"
1835        });
1836
1837        let error = parse_server_event(&payload.to_string())
1838            .expect_err("malformed known event should error");
1839        assert!(
1840            error.to_string().contains("StreamingCompletionChunk"),
1841            "expected strict decode failure, got {error}"
1842        );
1843    }
1844
1845    #[tokio::test]
1846    async fn close_is_idempotent() {
1847        let listener = TcpListener::bind("127.0.0.1:0")
1848            .await
1849            .expect("listener should bind");
1850        let address = listener.local_addr().expect("listener should have address");
1851
1852        let server = tokio::spawn(async move {
1853            let (stream, _) = listener.accept().await.expect("server should accept");
1854            let mut socket = accept_async(stream)
1855                .await
1856                .expect("server should upgrade websocket");
1857
1858            let message = socket
1859                .next()
1860                .await
1861                .expect("close frame should arrive")
1862                .expect("close frame should be valid");
1863            assert!(
1864                matches!(message, Message::Close(_)),
1865                "expected close frame, got {message:?}"
1866            );
1867        });
1868
1869        let base_url = format!("http://{address}/v1");
1870        let client = crate::providers::openai::Client::builder()
1871            .api_key("test-key")
1872            .base_url(&base_url)
1873            .build()
1874            .expect("client should build");
1875        let mut session = client
1876            .responses_websocket("gpt-4o")
1877            .await
1878            .expect("session should connect");
1879
1880        session.close().await.expect("first close should succeed");
1881        session.close().await.expect("second close should succeed");
1882
1883        server.await.expect("server task should finish");
1884    }
1885
1886    #[tokio::test]
1887    async fn send_while_in_flight_returns_error() {
1888        let listener = TcpListener::bind("127.0.0.1:0")
1889            .await
1890            .expect("listener should bind");
1891        let address = listener.local_addr().expect("listener should have address");
1892
1893        let server = tokio::spawn(async move {
1894            let (stream, _) = listener.accept().await.expect("server should accept");
1895            let mut socket = accept_async(stream)
1896                .await
1897                .expect("server should upgrade websocket");
1898
1899            // Read the first request but don't respond — keep it in-flight
1900            let _request = socket
1901                .next()
1902                .await
1903                .expect("request should exist")
1904                .expect("request should be valid");
1905
1906            // Wait for client to finish its test
1907            sleep(Duration::from_millis(100)).await;
1908            let _ = socket.close(None).await;
1909        });
1910
1911        let base_url = format!("http://{address}/v1");
1912        let client = crate::providers::openai::Client::builder()
1913            .api_key("test-key")
1914            .base_url(&base_url)
1915            .build()
1916            .expect("client should build");
1917        let model = client.completion_model("gpt-4o");
1918        let mut session = client
1919            .responses_websocket("gpt-4o")
1920            .await
1921            .expect("session should connect");
1922
1923        session
1924            .send(model.completion_request("first").build())
1925            .await
1926            .expect("first request should send");
1927
1928        let error = session
1929            .send(model.completion_request("second").build())
1930            .await
1931            .expect_err("second send while in-flight should error");
1932        assert!(
1933            error.to_string().contains("already in flight"),
1934            "expected in-flight error, got {error}"
1935        );
1936
1937        server.await.expect("server task should finish");
1938    }
1939
1940    #[tokio::test]
1941    async fn send_after_close_returns_error() {
1942        let listener = TcpListener::bind("127.0.0.1:0")
1943            .await
1944            .expect("listener should bind");
1945        let address = listener.local_addr().expect("listener should have address");
1946
1947        let server = tokio::spawn(async move {
1948            let (stream, _) = listener.accept().await.expect("server should accept");
1949            let _socket = accept_async(stream)
1950                .await
1951                .expect("server should upgrade websocket");
1952            sleep(Duration::from_millis(100)).await;
1953        });
1954
1955        let base_url = format!("http://{address}/v1");
1956        let client = crate::providers::openai::Client::builder()
1957            .api_key("test-key")
1958            .base_url(&base_url)
1959            .build()
1960            .expect("client should build");
1961        let model = client.completion_model("gpt-4o");
1962        let mut session = client
1963            .responses_websocket("gpt-4o")
1964            .await
1965            .expect("session should connect");
1966
1967        session.close().await.expect("close should succeed");
1968
1969        let error = session
1970            .send(model.completion_request("after close").build())
1971            .await
1972            .expect_err("send after close should error");
1973        assert!(
1974            error.to_string().contains("session is closed"),
1975            "expected closed-session error, got {error}"
1976        );
1977
1978        server.await.expect("server task should finish");
1979    }
1980
1981    #[tokio::test]
1982    async fn next_event_without_send_returns_error() {
1983        let listener = TcpListener::bind("127.0.0.1:0")
1984            .await
1985            .expect("listener should bind");
1986        let address = listener.local_addr().expect("listener should have address");
1987
1988        let server = tokio::spawn(async move {
1989            let (stream, _) = listener.accept().await.expect("server should accept");
1990            let _socket = accept_async(stream)
1991                .await
1992                .expect("server should upgrade websocket");
1993            sleep(Duration::from_millis(100)).await;
1994        });
1995
1996        let base_url = format!("http://{address}/v1");
1997        let client = crate::providers::openai::Client::builder()
1998            .api_key("test-key")
1999            .base_url(&base_url)
2000            .build()
2001            .expect("client should build");
2002        let mut session = client
2003            .responses_websocket("gpt-4o")
2004            .await
2005            .expect("session should connect");
2006
2007        let error = session
2008            .next_event()
2009            .await
2010            .expect_err("next_event without send should error");
2011        assert!(
2012            error
2013                .to_string()
2014                .contains("No OpenAI websocket response is currently in flight"),
2015            "expected not-in-flight error, got {error}"
2016        );
2017
2018        server.await.expect("server task should finish");
2019    }
2020
2021    #[tokio::test]
2022    async fn unknown_event_is_skipped_during_session() {
2023        let listener = TcpListener::bind("127.0.0.1:0")
2024            .await
2025            .expect("listener should bind");
2026        let address = listener.local_addr().expect("listener should have address");
2027
2028        let server = tokio::spawn(async move {
2029            let (stream, _) = listener.accept().await.expect("server should accept");
2030            let mut socket = accept_async(stream)
2031                .await
2032                .expect("server should upgrade websocket");
2033
2034            let _request = socket
2035                .next()
2036                .await
2037                .expect("request should exist")
2038                .expect("request should be valid");
2039
2040            // Send an unknown event type first
2041            socket
2042                .send(Message::text(
2043                    json!({
2044                        "type": "response.some_future_event",
2045                        "data": "should be skipped"
2046                    })
2047                    .to_string(),
2048                ))
2049                .await
2050                .expect("unknown event should send");
2051
2052            // Then send the real completed response
2053            let response = serde_json::to_value(CompletionResponse {
2054                id: "resp_after_unknown".to_string(),
2055                ..sample_response(ResponseStatus::Completed)
2056            })
2057            .expect("response should serialize");
2058
2059            socket
2060                .send(Message::text(
2061                    json!({
2062                        "type": "response.completed",
2063                        "sequence_number": 1,
2064                        "response": response,
2065                    })
2066                    .to_string(),
2067                ))
2068                .await
2069                .expect("completed event should send");
2070        });
2071
2072        let base_url = format!("http://{address}/v1");
2073        let client = crate::providers::openai::Client::builder()
2074            .api_key("test-key")
2075            .base_url(&base_url)
2076            .build()
2077            .expect("client should build");
2078        let model = client.completion_model("gpt-4o");
2079        let mut session = client
2080            .responses_websocket("gpt-4o")
2081            .await
2082            .expect("session should connect");
2083
2084        session
2085            .send(model.completion_request("hello").build())
2086            .await
2087            .expect("send should succeed");
2088        let response = session
2089            .wait_for_completed_response()
2090            .await
2091            .expect("response should complete despite unknown event");
2092        assert_eq!(response.id, "resp_after_unknown");
2093
2094        server.await.expect("server task should finish");
2095    }
2096}