yapi_ai/api/
codex_websocket.rs

1//! Codex's WebSocket transport: connections, pi's per-session socket cache
2//! and response continuation.
3//!
4//! Port of the WebSocket half of `packages/ai/src/api/openai-codex-responses.ts`
5//! in pi `v1.0.0`. A session keeps one socket per account between requests,
6//! closed after five idle minutes or once it is 55 minutes old. While idle, a
7//! task reads the socket, as Node does: it answers pings, drops stray
8//! messages and notices a close.
9
10use std::collections::{HashMap, HashSet};
11use std::sync::atomic::{AtomicU64, Ordering};
12use std::sync::{LazyLock, Mutex};
13use std::time::Duration;
14
15use futures_util::{SinkExt as _, StreamExt as _};
16use serde::Serialize;
17use serde_json::{Map, Value, json};
18use tokio::sync::oneshot;
19use tokio::task::JoinHandle;
20use tokio::time::Instant;
21use tokio_tungstenite::WebSocketStream;
22use tokio_tungstenite::tungstenite::Message as Frame;
23use tokio_tungstenite::tungstenite::handshake::client::generate_key;
24use tokio_tungstenite::tungstenite::handshake::derive_accept_key;
25use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
26use tokio_tungstenite::tungstenite::protocol::{CloseFrame, Role};
27use tokio_util::sync::CancellationToken;
28use yapi_types::sync::lock;
29
30use crate::http::{self, Headers};
31use crate::stream::{Transport, now_ms};
32
33type Socket = WebSocketStream<reqwest::Upgraded>;
34
35/// A session id and an account id.
36type Key = (String, String);
37
38const IDLE_TTL: Duration = Duration::from_secs(5 * 60);
39const MAX_AGE: Duration = Duration::from_secs(55 * 60);
40const DEFAULT_CONNECT_TIMEOUT_MS: u64 = 15_000;
41/// What pi reports for any failure the runtime's WebSocket gives no detail on.
42const WEBSOCKET_ERROR: &str = "WebSocket error";
43
44enum Slot {
45    /// A request is using the socket.
46    Busy,
47    Idle(Idle),
48}
49
50/// A cached socket between requests, held by its reader task.
51struct Idle {
52    /// Stops the reader: `None` returns the socket, a reason closes it.
53    stop: oneshot::Sender<Option<&'static str>>,
54    /// The socket, unless it closed or was closed.
55    reader: JoinHandle<Option<Socket>>,
56    created: Instant,
57    continuation: Option<Continuation>,
58    generation: u64,
59}
60
61static SOCKETS: LazyLock<Mutex<HashMap<Key, Slot>>> = LazyLock::new(Default::default);
62static GENERATION: AtomicU64 = AtomicU64::new(0);
63/// Sessions whose WebSocket failed, which use server-sent events (SSE) for
64/// the rest of the process.
65static FALLBACK: LazyLock<Mutex<HashSet<String>>> = LazyLock::new(Default::default);
66static STATS: LazyLock<Mutex<HashMap<String, DebugStats>>> = LazyLock::new(Default::default);
67
68/// pi's `OpenAICodexWebSocketDebugStats` for one session.
69#[derive(Clone, Debug, Default, Serialize)]
70#[serde(rename_all = "camelCase")]
71pub struct DebugStats {
72    requests: u64,
73    connections_created: u64,
74    connections_reused: u64,
75    cached_context_requests: u64,
76    store_true_requests: u64,
77    full_context_requests: u64,
78    delta_requests: u64,
79    last_input_items: usize,
80    #[serde(skip_serializing_if = "Option::is_none")]
81    last_delta_input_items: Option<usize>,
82    #[serde(skip_serializing_if = "Option::is_none")]
83    last_previous_response_id: Option<String>,
84    websocket_failures: u64,
85    sse_fallbacks: u64,
86    #[serde(skip_serializing_if = "Option::is_none")]
87    websocket_fallback_active: Option<bool>,
88    #[serde(skip_serializing_if = "Option::is_none")]
89    last_web_socket_error: Option<String>,
90}
91
92fn with_stats(session: &str, update: impl FnOnce(&mut DebugStats)) {
93    update(lock(&STATS).entry(session.to_owned()).or_default());
94}
95
96/// pi's `getOpenAICodexWebSocketDebugStats`: what `session_id`'s WebSocket
97/// requests did in this process.
98pub fn debug_stats(session_id: &str) -> Option<DebugStats> {
99    lock(&STATS).get(session_id).cloned()
100}
101
102/// pi's `resetOpenAICodexWebSocketDebugStats`: forgets the statistics and
103/// the SSE fallback of `session_id`, or of every session.
104pub fn reset_debug_stats(session_id: Option<&str>) {
105    match session_id {
106        Some(session) => {
107            lock(&STATS).remove(session);
108            lock(&FALLBACK).remove(session);
109        }
110        None => {
111            lock(&STATS).clear();
112            lock(&FALLBACK).clear();
113        }
114    }
115}
116
117/// Whether `session`'s WebSocket failed earlier in this process.
118pub(super) fn fallback_active(session: Option<&str>) -> bool {
119    session.is_some_and(|session| lock(&FALLBACK).contains(session))
120}
121
122/// Counts a request of `session` as pi's `processWebSocketStream` does.
123pub(super) fn count_request(session: &str, reused: bool, continues: bool, request: &Value) {
124    with_stats(session, |stats| {
125        stats.requests += 1;
126        if reused {
127            stats.connections_reused += 1;
128        } else {
129            stats.connections_created += 1;
130        }
131        if continues {
132            stats.cached_context_requests += 1;
133        }
134        if request["store"] == true {
135            stats.store_true_requests += 1;
136        }
137        let items = request["input"].as_array().map_or(0, Vec::len);
138        stats.last_input_items = items;
139        match request["previous_response_id"].as_str() {
140            Some(previous) => {
141                stats.delta_requests += 1;
142                stats.last_delta_input_items = Some(items);
143                stats.last_previous_response_id = Some(previous.to_owned());
144            }
145            None => {
146                stats.full_context_requests += 1;
147                stats.last_delta_input_items = None;
148                stats.last_previous_response_id = None;
149            }
150        }
151    });
152}
153
154/// pi's `recordWebSocketSseFallback`: a request of `session` goes over SSE.
155pub(super) fn record_fallback(session: Option<&str>) {
156    if let Some(session) = session {
157        let active = fallback_active(Some(session));
158        with_stats(session, |stats| {
159            stats.sse_fallbacks += 1;
160            stats.websocket_fallback_active = Some(active);
161        });
162    }
163}
164
165/// pi's `recordWebSocketFailure`: `session`'s later requests use SSE.
166pub(super) fn record_failure(session: Option<&str>, message: &str) {
167    if let Some(session) = session {
168        lock(&FALLBACK).insert(session.to_owned());
169        with_stats(session, |stats| {
170            stats.websocket_failures += 1;
171            stats.last_web_socket_error = Some(message.to_owned());
172            stats.websocket_fallback_active = Some(true);
173        });
174    }
175}
176
177/// pi's `closeOpenAICodexWebSocketSessions`: closes the sockets
178/// `session_id` keeps, or every session's. A socket in use closes when its
179/// request ends.
180pub fn close_sessions(session_id: Option<&str>) {
181    let mut sockets = lock(&SOCKETS);
182    let keys: Vec<Key> = sockets
183        .keys()
184        .filter(|(session, _)| session_id.is_none_or(|id| session == id))
185        .cloned()
186        .collect();
187    for key in keys {
188        if let Some(Slot::Idle(idle)) = sockets.remove(&key) {
189            let _ = idle.stop.send(Some("debug_close"));
190        }
191    }
192}
193
194/// Why a WebSocket request failed.
195#[derive(Debug)]
196pub(super) struct Failure {
197    /// The connection failed, as opposed to the backend reporting an error
198    /// or sending a message that is not JSON. Only these fall back to SSE.
199    pub(super) transport: bool,
200    /// pi's error class, for the diagnostic.
201    name: &'static str,
202    pub(super) message: String,
203    /// The backend's error code, or a close code.
204    code: Option<Value>,
205}
206
207impl Failure {
208    /// A failure pi raises as a plain `Error`.
209    pub(super) fn error(message: impl Into<String>) -> Failure {
210        Failure {
211            transport: true,
212            name: "Error",
213            message: message.into(),
214            code: None,
215        }
216    }
217
218    /// An `error` or `response.failed` event, with its code.
219    pub(super) fn api(message: String, code: Option<&str>) -> Failure {
220        Failure {
221            transport: false,
222            name: "CodexApiError",
223            message,
224            code: code.map(Value::from),
225        }
226    }
227
228    /// A message that is not JSON.
229    pub(super) fn protocol(message: String) -> Failure {
230        Failure {
231            transport: false,
232            name: "CodexProtocolError",
233            message,
234            code: None,
235        }
236    }
237
238    fn closed(frame: Option<CloseFrame>) -> Failure {
239        let (code, reason) = frame.map_or((1005, String::new()), |frame| {
240            (u16::from(frame.code), frame.reason.to_string())
241        });
242        let reason = if reason.is_empty() && code == 1009 {
243            "message too big".to_owned()
244        } else {
245            reason
246        };
247        Failure {
248            transport: true,
249            name: "WebSocketCloseError",
250            message: format!("WebSocket closed {code} {reason}")
251                .trim_end()
252                .to_owned(),
253            code: Some(Value::from(code)),
254        }
255    }
256
257    /// The backend's error code.
258    pub(super) fn api_code(&self) -> Option<&str> {
259        self.code
260            .as_ref()
261            .and_then(Value::as_str)
262            .filter(|_| !self.transport)
263    }
264
265    /// pi's `provider_transport_failure` diagnostic: the failure with the
266    /// `transport` setting, whether events had arrived, which decides the SSE
267    /// fallback, and the request's size in bytes.
268    pub(super) fn diagnostic(
269        &self,
270        transport: Transport,
271        events: bool,
272        request_bytes: usize,
273    ) -> Value {
274        let mut error = Map::new();
275        error.insert("name".into(), json!(self.name));
276        error.insert("message".into(), json!(self.message));
277        if let Some(code) = &self.code {
278            error.insert("code".into(), code.clone());
279        }
280        let mut details = Map::new();
281        details.insert("configuredTransport".into(), json!(transport));
282        if !events {
283            details.insert("fallbackTransport".into(), json!("sse"));
284        }
285        details.insert("eventsEmitted".into(), json!(events));
286        let phase = if events {
287            "after_message_stream_start"
288        } else {
289            "before_message_stream_start"
290        };
291        details.insert("phase".into(), json!(phase));
292        details.insert("requestBytes".into(), json!(request_bytes));
293        json!({
294            "type": "provider_transport_failure",
295            "timestamp": now_ms(),
296            "error": error,
297            "details": details,
298        })
299    }
300}
301
302/// The last exchange on a cached socket.
303pub(super) struct Continuation {
304    /// The full request, as sent before any delta was taken.
305    pub(super) body: Value,
306    pub(super) response_id: String,
307    /// The response's output as input items.
308    pub(super) items: Vec<Value>,
309}
310
311/// The request to send on a socket whose last exchange is `continuation`: when
312/// `body` repeats that request with more input after the response's items,
313/// only the new input, referencing the response. Otherwise `body`, and the
314/// continuation is dropped.
315pub(super) fn continued(body: &Value, continuation: &mut Option<Continuation>) -> Value {
316    let Some(last) = continuation else {
317        return body.clone();
318    };
319    match delta(body, last) {
320        Some(input) if !last.response_id.is_empty() => {
321            let mut next = body.clone();
322            next["previous_response_id"] = json!(last.response_id);
323            next["input"] = Value::Array(input);
324            next
325        }
326        _ => {
327            *continuation = None;
328            body.clone()
329        }
330    }
331}
332
333/// pi's `getCachedWebSocketInputDelta`. Requests compare as serialized JSON.
334fn delta(body: &Value, last: &Continuation) -> Option<Vec<Value>> {
335    let rest = |body: &Value| {
336        let mut rest = body.as_object()?.clone();
337        rest.shift_remove("input");
338        rest.shift_remove("previous_response_id");
339        Some(yapi_types::json::stringify(&rest))
340    };
341    if rest(body)? != rest(&last.body)? {
342        return None;
343    }
344    let input = |body: &Value| body["input"].as_array().cloned().unwrap_or_default();
345    let current = input(body);
346    let mut baseline = input(&last.body);
347    baseline.extend(last.items.iter().cloned());
348    let same =
349        |(a, b): (&Value, &Value)| yapi_types::json::stringify(a) == yapi_types::json::stringify(b);
350    if current.len() < baseline.len() || !current.iter().zip(&baseline).all(same) {
351        return None;
352    }
353    Some(current[baseline.len()..].to_vec())
354}
355
356/// The message for a request: `response.create` with the body's fields.
357pub(super) fn frame(body: &Value) -> String {
358    let mut frame = Map::new();
359    frame.insert("type".into(), json!("response.create"));
360    if let Some(fields) = body.as_object() {
361        frame.extend(fields.clone());
362    }
363    yapi_types::json::stringify(&frame)
364}
365
366/// A socket for one request.
367pub(super) struct Lease {
368    socket: Socket,
369    /// The cache slot the socket returns to, or `None` for a one-off socket.
370    key: BusySlot,
371    created: Instant,
372    /// Whether the socket served an earlier request.
373    pub(super) reused: bool,
374    /// The socket's last exchange.
375    pub(super) continuation: Option<Continuation>,
376}
377
378/// A socket to `url` for `key`'s session and account: its idle cached one,
379/// else a new one that is cached once released. While the cached socket is
380/// busy, and without a key, the socket is a one-off.
381pub(super) async fn acquire(
382    url: &str,
383    headers: &Headers,
384    key: Option<Key>,
385    cancel: &CancellationToken,
386    connect_timeout_ms: Option<u64>,
387) -> Result<Lease, Failure> {
388    let open = || connect(url, headers, cancel, connect_timeout_ms);
389    let lease = |socket, key, created| Lease {
390        socket,
391        key: BusySlot(key),
392        created,
393        reused: false,
394        continuation: None,
395    };
396    let Some(key) = key else {
397        return Ok(lease(open().await?, None, Instant::now()));
398    };
399    let found = {
400        let mut sockets = lock(&SOCKETS);
401        let found = sockets.remove(&key);
402        if found.is_some() {
403            sockets.insert(key.clone(), Slot::Busy);
404        }
405        found
406    };
407    match found {
408        Some(Slot::Busy) => return Ok(lease(open().await?, None, Instant::now())),
409        Some(Slot::Idle(idle)) => {
410            let _ = idle.stop.send(None);
411            if let Ok(Some(socket)) = idle.reader.await {
412                if idle.created.elapsed() < MAX_AGE {
413                    return Ok(Lease {
414                        reused: true,
415                        continuation: idle.continuation,
416                        ..lease(socket, Some(key), idle.created)
417                    });
418                }
419                close(socket, "connection_age_limit");
420            }
421        }
422        None => {}
423    }
424    let socket = match open().await {
425        Ok(socket) => socket,
426        Err(failure) => {
427            remove_busy(&key);
428            return Err(failure);
429        }
430    };
431    lock(&SOCKETS).insert(key.clone(), Slot::Busy);
432    Ok(lease(socket, Some(key), Instant::now()))
433}
434
435/// The `Busy` cache slot a request holds. Dropping it frees the slot for the
436/// session's next request, including when the request's task is cancelled
437/// before it releases the slot itself.
438struct BusySlot(Option<Key>);
439
440impl Drop for BusySlot {
441    fn drop(&mut self) {
442        if let Some(key) = &self.0 {
443            remove_busy(key);
444        }
445    }
446}
447
448/// Forgets `key`'s slot when a request holds it.
449fn remove_busy(key: &Key) {
450    let mut sockets = lock(&SOCKETS);
451    if matches!(sockets.get(key), Some(Slot::Busy)) {
452        sockets.remove(key);
453    }
454}
455
456impl Lease {
457    /// Whether the socket returns to the session's cache.
458    pub(super) fn cached(&self) -> bool {
459        self.key.0.is_some()
460    }
461
462    /// Sends `text` as one message.
463    pub(super) async fn send(&mut self, text: String) -> Result<(), Failure> {
464        self.socket
465            .send(Frame::text(text))
466            .await
467            .map_err(|_| Failure::error(WEBSOCKET_ERROR))
468    }
469
470    /// The next text or binary message, as text. Fails when the socket
471    /// closes or errs, when `cancel` fires, and once `idle` passes without a
472    /// message. Pings and pongs do not count.
473    pub(super) async fn next(
474        &mut self,
475        cancel: &CancellationToken,
476        idle: Option<Duration>,
477    ) -> Result<String, Failure> {
478        let deadline = idle.map(|idle| (idle, tokio::time::Instant::now() + idle));
479        loop {
480            let read = self.socket.next();
481            let frame = match deadline {
482                Some((idle, at)) => cancel
483                    .run_until_cancelled(tokio::time::timeout_at(at, read))
484                    .await
485                    .map(|frame| {
486                        frame.map_err(|_| {
487                            Failure::error(format!(
488                                "WebSocket idle timeout after {}ms",
489                                idle.as_millis()
490                            ))
491                        })
492                    }),
493                None => cancel.run_until_cancelled(read).await.map(Ok),
494            };
495            let frame =
496                frame.unwrap_or_else(|| Err(Failure::error(http::ABORTED_DURING_STREAM)))?;
497            // pi ignores empty messages, and they do not restart its idle timer.
498            match frame {
499                Some(Ok(Frame::Text(text))) if !text.is_empty() => return Ok(text.to_string()),
500                Some(Ok(Frame::Binary(bytes))) if !bytes.is_empty() => {
501                    return Ok(String::from_utf8_lossy(&bytes).into_owned());
502                }
503                Some(Ok(Frame::Close(frame))) => return Err(Failure::closed(frame)),
504                Some(Ok(_)) => {}
505                Some(Err(_)) | None => return Err(Failure::error(WEBSOCKET_ERROR)),
506            }
507        }
508    }
509
510    /// Returns a cached socket to its session when `keep` and its slot is
511    /// still the request's, to close after five idle minutes. Otherwise closes
512    /// it.
513    pub(super) fn release(mut self, keep: bool) {
514        let Some(key) = self.key.0.take() else {
515            return close(self.socket, "done");
516        };
517        let mut sockets = lock(&SOCKETS);
518        if !matches!(sockets.get(&key), Some(Slot::Busy)) {
519            return close(self.socket, "done");
520        }
521        if !keep {
522            sockets.remove(&key);
523            return close(self.socket, "done");
524        }
525        let generation = GENERATION.fetch_add(1, Ordering::Relaxed);
526        let (stop, stopped) = oneshot::channel();
527        let reader = tokio::spawn(read_idle(self.socket, stopped, key.clone(), generation));
528        sockets.insert(
529            key,
530            Slot::Idle(Idle {
531                stop,
532                reader,
533                created: self.created,
534                continuation: self.continuation,
535                generation,
536            }),
537        );
538    }
539}
540
541/// Reads an idle socket until a request takes it back, which gets the
542/// socket, or until it closes, errs or idles for five minutes, which ends it
543/// and its slot.
544async fn read_idle(
545    mut socket: Socket,
546    mut stopped: oneshot::Receiver<Option<&'static str>>,
547    key: Key,
548    generation: u64,
549) -> Option<Socket> {
550    let expiry = tokio::time::sleep(IDLE_TTL);
551    tokio::pin!(expiry);
552    let reason = loop {
553        tokio::select! {
554            stop = &mut stopped => match stop {
555                Ok(None) => return Some(socket),
556                Ok(Some(reason)) => break Some(reason),
557                Err(_) => break Some("done"),
558            },
559            () = &mut expiry => break Some("idle_timeout"),
560            frame = socket.next() => match frame {
561                Some(Ok(Frame::Close(_)) | Err(_)) | None => break None,
562                // pi drops what arrives while no request listens.
563                Some(Ok(_)) => {}
564            },
565        }
566    };
567    {
568        let mut sockets = lock(&SOCKETS);
569        if matches!(sockets.get(&key), Some(Slot::Idle(idle)) if idle.generation == generation) {
570            sockets.remove(&key);
571        }
572    }
573    if let Some(reason) = reason {
574        close(socket, reason);
575    }
576    None
577}
578
579/// Closes `socket` with a normal close and `reason`, without waiting.
580fn close(mut socket: Socket, reason: &'static str) {
581    tokio::spawn(async move {
582        let frame = CloseFrame {
583            code: CloseCode::Normal,
584            reason: reason.into(),
585        };
586        let _ = socket.close(Some(frame)).await;
587    });
588}
589
590/// Opens a WebSocket to `url`, an `http` or `https` URL, through an HTTP/1.1
591/// upgrade.
592async fn connect(
593    url: &str,
594    headers: &Headers,
595    cancel: &CancellationToken,
596    timeout_ms: Option<u64>,
597) -> Result<Socket, Failure> {
598    let key = generate_key();
599    let request = headers
600        .apply(http::upgrade_client().get(url))
601        .header("connection", "Upgrade")
602        .header("upgrade", "websocket")
603        .header("sec-websocket-version", "13")
604        .header("sec-websocket-key", &key);
605    let open = async {
606        let failed = |_| Failure::error(WEBSOCKET_ERROR);
607        let response = request.send().await.map_err(failed)?;
608        let accepted = response
609            .headers()
610            .get("sec-websocket-accept")
611            .is_some_and(|accept| {
612                accept.as_bytes() == derive_accept_key(key.as_bytes()).as_bytes()
613            });
614        if response.status() != reqwest::StatusCode::SWITCHING_PROTOCOLS || !accepted {
615            return Err(Failure::error(WEBSOCKET_ERROR));
616        }
617        let upgraded = response.upgrade().await.map_err(failed)?;
618        Ok(WebSocketStream::from_raw_socket(upgraded, Role::Client, None).await)
619    };
620    let timeout_ms = timeout_ms.unwrap_or(DEFAULT_CONNECT_TIMEOUT_MS);
621    let open = async {
622        if timeout_ms == 0 {
623            return open.await;
624        }
625        tokio::time::timeout(Duration::from_millis(timeout_ms), open)
626            .await
627            .unwrap_or_else(|_| {
628                Err(Failure::error(format!(
629                    "WebSocket connect timeout after {timeout_ms}ms"
630                )))
631            })
632    };
633    cancel
634        .run_until_cancelled(open)
635        .await
636        .unwrap_or_else(|| Err(Failure::error(http::ABORTED_DURING_STREAM)))
637}
yapi_ai/api/codex_websocket.rs

yapi_ai/api/
codex_websocket.rs
