elph_ai/api/
codex_transport.rs

1use std::collections::{HashMap, HashSet};
2use std::sync::atomic::{AtomicBool, Ordering};
3use std::sync::{Arc, Mutex};
4use std::time::{Duration, SystemTime, UNIX_EPOCH};
5
6use anyhow::{Result, anyhow};
7use futures_util::{SinkExt, StreamExt};
8use serde_json::Value;
9use tokio_tungstenite::tungstenite::Message as WsMessage;
10
11use crate::api::websocket_connect::{WsStream, connect_websocket_with_proxy};
12use crate::types::ProviderEnv;
13
14use crate::api::common::send_with_abort;
15use crate::api::sse::for_each_sse_json_event;
16use tokio_util::sync::CancellationToken;
17
18const OPENAI_BETA_RESPONSES_WEBSOCKETS: &str = "responses_websockets=2026-02-06";
19const WEBSOCKET_CONNECTION_LIMIT_REACHED: &str = "websocket_connection_limit_reached";
20const SESSION_WEBSOCKET_CACHE_TTL_MS: u64 = 5 * 60 * 1000;
21const SESSION_WEBSOCKET_MAX_AGE_MS: u64 = 55 * 60 * 1000;
22
23#[derive(Clone, Default, PartialEq, Eq)]
24pub enum CodexTransport {
25    #[default]
26    Auto,
27    Sse,
28    WebSocket,
29    WebSocketCached,
30}
31
32pub struct CodexTransportOptions {
33    pub transport: CodexTransport,
34    pub websocket_connect_timeout_ms: Option<u64>,
35    pub session_id: Option<String>,
36    pub signal: Option<CancellationToken>,
37    pub env: Option<ProviderEnv>,
38}
39
40#[derive(Clone, Debug, Default)]
41pub struct CodexWebSocketDebugStats {
42    pub requests: u64,
43    pub connections_created: u64,
44    pub connections_reused: u64,
45    pub cached_context_requests: u64,
46    pub store_true_requests: u64,
47    pub full_context_requests: u64,
48    pub delta_requests: u64,
49    pub last_input_items: u64,
50    pub last_delta_input_items: Option<u64>,
51    pub last_previous_response_id: Option<String>,
52    pub websocket_failures: u64,
53    pub sse_fallbacks: u64,
54    pub websocket_fallback_active: Option<bool>,
55    pub last_websocket_error: Option<String>,
56}
57
58#[derive(Clone)]
59struct CachedWebSocketContinuationState {
60    last_request_body: Value,
61    last_response_id: String,
62    last_response_items: Value,
63}
64
65struct CachedWebSocketConnection {
66    socket: Arc<tokio::sync::Mutex<WsStream>>,
67    busy: Arc<AtomicBool>,
68    created_at: u64,
69    continuation: Arc<Mutex<Option<CachedWebSocketContinuationState>>>,
70    idle_task: Arc<tokio::sync::Mutex<Option<tokio::task::JoinHandle<()>>>>,
71}
72
73struct WebSocketLease {
74    socket: Arc<tokio::sync::Mutex<WsStream>>,
75    entry: Option<Arc<CachedWebSocketConnection>>,
76    session_id: Option<String>,
77    reused: bool,
78    ephemeral: bool,
79}
80
81static SSE_FALLBACK_SESSIONS: once_cell::sync::Lazy<Mutex<HashSet<String>>> =
82    once_cell::sync::Lazy::new(|| Mutex::new(HashSet::new()));
83static WEBSOCKET_SESSION_CACHE: once_cell::sync::Lazy<Mutex<HashMap<String, Arc<CachedWebSocketConnection>>>> =
84    once_cell::sync::Lazy::new(|| Mutex::new(HashMap::new()));
85static WEBSOCKET_DEBUG_STATS: once_cell::sync::Lazy<Mutex<HashMap<String, CodexWebSocketDebugStats>>> =
86    once_cell::sync::Lazy::new(|| Mutex::new(HashMap::new()));
87
88pub fn compress_request_body_zstd(body_json: &str) -> Option<Vec<u8>> {
89    zstd::encode_all(body_json.as_bytes(), 3).ok()
90}
91
92pub fn resolve_codex_websocket_url(base_url: &str) -> String {
93    let normalized = base_url.trim().trim_end_matches('/');
94    let host = normalized.trim_start_matches("https://").trim_start_matches("http://");
95    if normalized.ends_with("/codex/responses") {
96        format!("wss://{host}")
97    } else if normalized.ends_with("/codex") {
98        format!("wss://{host}/responses")
99    } else {
100        format!("wss://{host}/codex/responses")
101    }
102}
103
104pub fn is_websocket_sse_fallback_active(session_id: Option<&str>) -> bool {
105    session_id
106        .map(|id| SSE_FALLBACK_SESSIONS.lock().unwrap().contains(id))
107        .unwrap_or(false)
108}
109
110fn mark_sse_fallback(session_id: Option<&str>) {
111    if let Some(id) = session_id {
112        SSE_FALLBACK_SESSIONS.lock().unwrap().insert(id.to_string());
113        if let Ok(mut stats) = WEBSOCKET_DEBUG_STATS.lock() {
114            let entry = stats.entry(id.to_string()).or_default();
115            entry.sse_fallbacks += 1;
116            entry.websocket_fallback_active = Some(true);
117        }
118    }
119}
120
121pub fn is_connection_limit_error(error: &str) -> bool {
122    error.contains(WEBSOCKET_CONNECTION_LIMIT_REACHED)
123}
124
125fn now_ms() -> u64 {
126    SystemTime::now()
127        .duration_since(UNIX_EPOCH)
128        .unwrap_or_default()
129        .as_millis() as u64
130}
131
132fn update_debug_stats(session_id: &str, update: impl FnOnce(&mut CodexWebSocketDebugStats)) {
133    if let Ok(mut stats) = WEBSOCKET_DEBUG_STATS.lock() {
134        let entry = stats.entry(session_id.to_string()).or_default();
135        update(entry);
136    }
137}
138
139pub fn get_codex_websocket_debug_stats(session_id: &str) -> Option<CodexWebSocketDebugStats> {
140    WEBSOCKET_DEBUG_STATS.lock().ok()?.get(session_id).cloned()
141}
142
143pub fn reset_codex_websocket_debug_stats(session_id: Option<&str>) {
144    if let Ok(mut stats) = WEBSOCKET_DEBUG_STATS.lock() {
145        match session_id {
146            Some(id) => {
147                stats.remove(id);
148                SSE_FALLBACK_SESSIONS.lock().unwrap().remove(id);
149            }
150            None => {
151                stats.clear();
152                SSE_FALLBACK_SESSIONS.lock().unwrap().clear();
153            }
154        }
155    }
156}
157
158pub fn close_codex_websocket_sessions(session_id: Option<&str>) {
159    let close_entry = |entry: &CachedWebSocketConnection| {
160        if let Ok(mut task) = entry.idle_task.try_lock()
161            && let Some(handle) = task.take()
162        {
163            handle.abort();
164        }
165        if let Ok(mut socket) = entry.socket.try_lock() {
166            let _ = futures_util::future::FutureExt::now_or_never(socket.close(None));
167        }
168    };
169
170    if let Ok(mut cache) = WEBSOCKET_SESSION_CACHE.lock() {
171        match session_id {
172            Some(id) => {
173                if let Some(entry) = cache.remove(id) {
174                    close_entry(&entry);
175                }
176                SSE_FALLBACK_SESSIONS.lock().unwrap().remove(id);
177            }
178            None => {
179                for entry in cache.values() {
180                    close_entry(entry);
181                }
182                cache.clear();
183                SSE_FALLBACK_SESSIONS.lock().unwrap().clear();
184            }
185        }
186    }
187}
188
189fn is_session_expired(entry: &CachedWebSocketConnection) -> bool {
190    now_ms().saturating_sub(entry.created_at) >= SESSION_WEBSOCKET_MAX_AGE_MS
191}
192
193async fn close_socket_quietly(socket: &Arc<tokio::sync::Mutex<WsStream>>) {
194    let mut guard = socket.lock().await;
195    let _ = guard.close(None).await;
196}
197
198fn schedule_idle_expiry(session_id: String, entry: Arc<CachedWebSocketConnection>) {
199    let socket = entry.socket.clone();
200    let busy = entry.busy.clone();
201    let task_slot = entry.idle_task.clone();
202    let handle = tokio::spawn(async move {
203        tokio::time::sleep(Duration::from_millis(SESSION_WEBSOCKET_CACHE_TTL_MS)).await;
204        if busy.load(Ordering::SeqCst) {
205            return;
206        }
207        close_socket_quietly(&socket).await;
208        if let Ok(mut cache) = WEBSOCKET_SESSION_CACHE.lock() {
209            cache.remove(&session_id);
210        }
211    });
212    if let Ok(mut slot) = task_slot.try_lock()
213        && let Some(old) = slot.replace(handle)
214    {
215        old.abort();
216    }
217}
218
219async fn connect_websocket(
220    ws_url: &str,
221    headers: &HashMap<String, String>,
222    timeout_ms: u64,
223    env: Option<&ProviderEnv>,
224) -> Result<WsStream> {
225    let mut request_headers = headers.clone();
226    request_headers.insert("OpenAI-Beta".to_string(), OPENAI_BETA_RESPONSES_WEBSOCKETS.to_string());
227    connect_websocket_with_proxy(ws_url, &request_headers, timeout_ms, env).await
228}
229
230async fn acquire_websocket(
231    ws_url: &str,
232    headers: &HashMap<String, String>,
233    session_id: Option<&str>,
234    timeout_ms: u64,
235    env: Option<&ProviderEnv>,
236) -> Result<WebSocketLease> {
237    let Some(session_id) = session_id else {
238        let socket = connect_websocket(ws_url, headers, timeout_ms, env).await?;
239        return Ok(WebSocketLease {
240            socket: Arc::new(tokio::sync::Mutex::new(socket)),
241            entry: None,
242            session_id: None,
243            reused: false,
244            ephemeral: true,
245        });
246    };
247
248    enum CacheAction {
249        Reuse(Arc<CachedWebSocketConnection>),
250        Expired(Arc<CachedWebSocketConnection>),
251        Busy,
252    }
253
254    let cache_action = WEBSOCKET_SESSION_CACHE.lock().ok().and_then(|mut cache| {
255        let entry = cache.get(session_id)?.clone();
256        if let Ok(mut task) = entry.idle_task.try_lock()
257            && let Some(handle) = task.take()
258        {
259            handle.abort();
260        }
261        if entry.busy.load(Ordering::SeqCst) {
262            return Some(CacheAction::Busy);
263        }
264        if is_session_expired(&entry) {
265            cache.remove(session_id);
266            return Some(CacheAction::Expired(entry));
267        }
268        entry.busy.store(true, Ordering::SeqCst);
269        Some(CacheAction::Reuse(entry))
270    });
271
272    if let Some(action) = cache_action {
273        match action {
274            CacheAction::Reuse(entry) => {
275                return Ok(WebSocketLease {
276                    socket: entry.socket.clone(),
277                    entry: Some(entry),
278                    session_id: Some(session_id.to_string()),
279                    reused: true,
280                    ephemeral: false,
281                });
282            }
283            CacheAction::Expired(entry) => {
284                close_socket_quietly(&entry.socket).await;
285            }
286            CacheAction::Busy => {
287                let socket = connect_websocket(ws_url, headers, timeout_ms, env).await?;
288                return Ok(WebSocketLease {
289                    socket: Arc::new(tokio::sync::Mutex::new(socket)),
290                    entry: None,
291                    session_id: Some(session_id.to_string()),
292                    reused: false,
293                    ephemeral: true,
294                });
295            }
296        }
297    }
298
299    let socket = connect_websocket(ws_url, headers, timeout_ms, env).await?;
300    let entry = Arc::new(CachedWebSocketConnection {
301        socket: Arc::new(tokio::sync::Mutex::new(socket)),
302        busy: Arc::new(AtomicBool::new(true)),
303        created_at: now_ms(),
304        continuation: Arc::new(Mutex::new(None)),
305        idle_task: Arc::new(tokio::sync::Mutex::new(None)),
306    });
307    WEBSOCKET_SESSION_CACHE
308        .lock()
309        .unwrap()
310        .insert(session_id.to_string(), entry.clone());
311
312    Ok(WebSocketLease {
313        socket: entry.socket.clone(),
314        entry: Some(entry),
315        session_id: Some(session_id.to_string()),
316        reused: false,
317        ephemeral: false,
318    })
319}
320
321async fn release_websocket(lease: WebSocketLease, keep: bool) {
322    if lease.ephemeral || lease.entry.is_none() {
323        close_socket_quietly(&lease.socket).await;
324        return;
325    }
326
327    let (Some(entry), Some(session_id)) = (lease.entry, lease.session_id) else {
328        close_socket_quietly(&lease.socket).await;
329        return;
330    };
331
332    if !keep {
333        close_socket_quietly(&entry.socket).await;
334        if let Ok(mut cache) = WEBSOCKET_SESSION_CACHE.lock() {
335            cache.remove(&session_id);
336        }
337        return;
338    }
339
340    entry.busy.store(false, Ordering::SeqCst);
341    schedule_idle_expiry(session_id, entry);
342}
343
344fn request_body_without_input(body: &Value) -> Value {
345    let mut copy = body.clone();
346    if let Some(obj) = copy.as_object_mut() {
347        obj.remove("input");
348        obj.remove("previous_response_id");
349    }
350    copy
351}
352
353fn response_inputs_equal(a: &Value, b: &Value) -> bool {
354    serde_json::to_string(a).ok() == serde_json::to_string(b).ok()
355}
356
357fn get_cached_websocket_input_delta(body: &Value, continuation: &CachedWebSocketContinuationState) -> Option<Value> {
358    if !response_inputs_equal(
359        &request_body_without_input(body),
360        &request_body_without_input(&continuation.last_request_body),
361    ) {
362        return None;
363    }
364
365    let current_input = body.get("input").cloned().unwrap_or(Value::Array(vec![]));
366    let baseline_items = continuation
367        .last_request_body
368        .get("input")
369        .cloned()
370        .unwrap_or(Value::Array(vec![]));
371    let response_items = continuation.last_response_items.clone();
372
373    let baseline = match baseline_items {
374        Value::Array(mut items) => {
375            if let Value::Array(extra) = response_items {
376                items.extend(extra);
377            }
378            items
379        }
380        _ => Vec::new(),
381    };
382
383    let current = match current_input {
384        Value::Array(items) => items,
385        _ => return None,
386    };
387
388    if current.len() < baseline.len() {
389        return None;
390    }
391
392    let baseline_len = baseline.len();
393    let prefix = Value::Array(current[..baseline_len].to_vec());
394    let baseline_value = Value::Array(baseline);
395    if !response_inputs_equal(&prefix, &baseline_value) {
396        return None;
397    }
398
399    Some(Value::Array(current[baseline_len..].to_vec()))
400}
401
402fn build_cached_websocket_request_body(entry: &CachedWebSocketConnection, body: &Value) -> Value {
403    let continuation = entry.continuation.lock().ok().and_then(|g| g.clone());
404    let Some(continuation) = continuation else {
405        return body.clone();
406    };
407
408    let delta = get_cached_websocket_input_delta(body, &continuation);
409    let Some(delta) = delta else {
410        if let Ok(mut guard) = entry.continuation.lock() {
411            *guard = None;
412        }
413        return body.clone();
414    };
415
416    if continuation.last_response_id.is_empty() {
417        if let Ok(mut guard) = entry.continuation.lock() {
418            *guard = None;
419        }
420        return body.clone();
421    }
422
423    let mut cached = body.clone();
424    if let Some(obj) = cached.as_object_mut() {
425        obj.insert(
426            "previous_response_id".to_string(),
427            Value::String(continuation.last_response_id.clone()),
428        );
429        obj.insert("input".to_string(), delta);
430    }
431    cached
432}
433
434pub fn update_codex_websocket_continuation(
435    session_id: &str,
436    full_body: &Value,
437    response_id: &str,
438    response_items: Value,
439) {
440    let Ok(cache) = WEBSOCKET_SESSION_CACHE.lock() else {
441        return;
442    };
443    let Some(entry) = cache.get(session_id) else {
444        return;
445    };
446    if let Ok(mut guard) = entry.continuation.lock() {
447        *guard = Some(CachedWebSocketContinuationState {
448            last_request_body: full_body.clone(),
449            last_response_id: response_id.to_string(),
450            last_response_items: response_items,
451        });
452    }
453}
454
455pub fn clear_codex_websocket_continuation(session_id: &str) {
456    if let Ok(cache) = WEBSOCKET_SESSION_CACHE.lock()
457        && let Some(entry) = cache.get(session_id)
458        && let Ok(mut guard) = entry.continuation.lock()
459    {
460        *guard = None;
461    }
462}
463
464pub struct CodexCollectResult {
465    pub events: Vec<Value>,
466    pub websocket_reused: bool,
467    pub used_cached_context: bool,
468}
469
470pub async fn collect_codex_events(
471    base_url: &str,
472    body: Value,
473    headers: HashMap<String, String>,
474    client: &reqwest::Client,
475    sse_url: &str,
476    options: &CodexTransportOptions,
477) -> Result<Vec<Value>> {
478    collect_codex_events_detailed(base_url, body, headers, client, sse_url, options)
479        .await
480        .map(|result| result.events)
481}
482
483pub async fn collect_codex_events_detailed(
484    base_url: &str,
485    body: Value,
486    headers: HashMap<String, String>,
487    client: &reqwest::Client,
488    sse_url: &str,
489    options: &CodexTransportOptions,
490) -> Result<CodexCollectResult> {
491    let transport = options.transport.clone();
492    let session_id = options.session_id.as_deref();
493    let timeout_ms = options.websocket_connect_timeout_ms.unwrap_or(30_000);
494
495    if transport != CodexTransport::Sse && !is_websocket_sse_fallback_active(session_id) {
496        let use_cached_context = transport == CodexTransport::WebSocketCached || transport == CodexTransport::Auto;
497        match try_websocket_stream(base_url, &body, &headers, options, use_cached_context, timeout_ms).await {
498            Ok(result) => return Ok(result),
499            Err(error) => {
500                let msg = error.to_string();
501                if is_connection_limit_error(&msg) && transport == CodexTransport::Auto {
502                    // retry once without session cache entry
503                } else {
504                    mark_sse_fallback(session_id);
505                    if let Some(id) = session_id {
506                        update_debug_stats(id, |stats| {
507                            stats.websocket_failures += 1;
508                            stats.last_websocket_error = Some(msg.clone());
509                            stats.websocket_fallback_active = Some(true);
510                        });
511                    }
512                }
513            }
514        }
515    }
516
517    let events = collect_codex_sse_events(client, sse_url, &body, &headers, &options.signal).await?;
518    Ok(CodexCollectResult {
519        events,
520        websocket_reused: false,
521        used_cached_context: false,
522    })
523}
524
525async fn try_websocket_stream(
526    base_url: &str,
527    body: &Value,
528    headers: &HashMap<String, String>,
529    options: &CodexTransportOptions,
530    use_cached_context: bool,
531    timeout_ms: u64,
532) -> Result<CodexCollectResult> {
533    let ws_url = resolve_codex_websocket_url(base_url);
534    let lease = acquire_websocket(
535        &ws_url,
536        headers,
537        options.session_id.as_deref(),
538        timeout_ms,
539        options.env.as_ref(),
540    )
541    .await?;
542
543    if let (Some(session_id), true) = (options.session_id.as_deref(), lease.reused) {
544        update_debug_stats(session_id, |stats| {
545            stats.connections_reused += 1;
546        });
547    } else if let Some(session_id) = options.session_id.as_deref() {
548        update_debug_stats(session_id, |stats| {
549            stats.connections_created += 1;
550        });
551    }
552
553    let request_body = if use_cached_context {
554        if let Some(entry) = &lease.entry {
555            build_cached_websocket_request_body(entry, body)
556        } else {
557            body.clone()
558        }
559    } else {
560        body.clone()
561    };
562
563    if let Some(session_id) = options.session_id.as_deref() {
564        update_debug_stats(session_id, |stats| {
565            stats.requests += 1;
566            if use_cached_context {
567                stats.cached_context_requests += 1;
568            }
569            if request_body.get("store").and_then(|v| v.as_bool()) == Some(true) {
570                stats.store_true_requests += 1;
571            }
572            stats.last_input_items = request_body
573                .get("input")
574                .and_then(|v| v.as_array())
575                .map(|a| a.len() as u64)
576                .unwrap_or(0);
577            if let Some(prev) = request_body.get("previous_response_id").and_then(|v| v.as_str()) {
578                stats.delta_requests += 1;
579                stats.last_delta_input_items = request_body
580                    .get("input")
581                    .and_then(|v| v.as_array())
582                    .map(|a| a.len() as u64);
583                stats.last_previous_response_id = Some(prev.to_string());
584            } else {
585                stats.full_context_requests += 1;
586                stats.last_delta_input_items = None;
587                stats.last_previous_response_id = None;
588            }
589        });
590    }
591
592    let reused = lease.reused;
593    let collect_result = collect_websocket_events(&lease.socket, &request_body).await;
594    let keep = collect_result.is_ok();
595    if !keep
596        && let Some(entry) = &lease.entry
597        && let Ok(mut guard) = entry.continuation.lock()
598    {
599        *guard = None;
600    }
601    release_websocket(lease, keep).await;
602    let events = collect_result?;
603    Ok(CodexCollectResult {
604        events,
605        websocket_reused: reused,
606        used_cached_context: use_cached_context,
607    })
608}
609
610async fn collect_websocket_events(socket: &Arc<tokio::sync::Mutex<WsStream>>, body: &Value) -> Result<Vec<Value>> {
611    let mut payload = body.clone();
612    if let Some(obj) = payload.as_object_mut() {
613        obj.insert("type".to_string(), Value::String("response.create".to_string()));
614    } else {
615        payload = serde_json::json!({ "type": "response.create", "body": body });
616    }
617
618    {
619        let mut guard = socket.lock().await;
620        guard.send(WsMessage::Text(payload.to_string().into())).await?;
621    }
622
623    let mut events = Vec::new();
624    loop {
625        let msg = {
626            let mut guard = socket.lock().await;
627            guard.next().await
628        };
629        let Some(msg) = msg else {
630            break;
631        };
632        let msg = msg?;
633        if let WsMessage::Text(text) = msg {
634            let event: Value = serde_json::from_str(&text)?;
635            if event.get("type").and_then(|v| v.as_str()) == Some("error") {
636                let code = event.pointer("/error/code").and_then(|v| v.as_str()).unwrap_or("");
637                if code == WEBSOCKET_CONNECTION_LIMIT_REACHED {
638                    return Err(anyhow!(WEBSOCKET_CONNECTION_LIMIT_REACHED));
639                }
640                let message = event
641                    .get("message")
642                    .and_then(|v| v.as_str())
643                    .unwrap_or("Codex WebSocket error");
644                return Err(anyhow!("{message}"));
645            }
646            let event_type = event.get("type").and_then(|v| v.as_str()).map(|s| s.to_string());
647            events.push(event);
648            if matches!(
649                event_type.as_deref(),
650                Some("response.done") | Some("response.completed") | Some("response.incomplete")
651            ) {
652                break;
653            }
654        }
655    }
656    Ok(events)
657}
658
659/// Exposed for integration tests (pi-ai codex websocket cache parity).
660#[doc(hidden)]
661pub fn get_codex_websocket_input_delta(
662    body: &Value,
663    last_request_body: &Value,
664    last_response_items: &Value,
665) -> Option<Value> {
666    let continuation = CachedWebSocketContinuationState {
667        last_request_body: last_request_body.clone(),
668        last_response_id: "resp_test".to_string(),
669        last_response_items: last_response_items.clone(),
670    };
671    get_cached_websocket_input_delta(body, &continuation)
672}
673
674async fn collect_codex_sse_events(
675    client: &reqwest::Client,
676    url: &str,
677    body: &Value,
678    headers: &HashMap<String, String>,
679    signal: &Option<CancellationToken>,
680) -> Result<Vec<Value>> {
681    let body_json = serde_json::to_string(body)?;
682    let mut req = if let Some(compressed) = compress_request_body_zstd(&body_json) {
683        let mut r = client.post(url).body(compressed);
684        r = r.header("Content-Type", "application/json");
685        r = r.header("Content-Encoding", "zstd");
686        r
687    } else {
688        client.post(url).json(body)
689    };
690    for (k, v) in headers {
691        req = req.header(k, v);
692    }
693    let response = send_with_abort(signal, req).await?;
694    let response = crate::api::common::check_response_ok(response).await?;
695    let mut events = Vec::new();
696    for_each_sse_json_event(response, signal, |event| {
697        events.push(event);
698        Ok(())
699    })
700    .await?;
701    Ok(events)
702}
elph_ai/api/codex_transport.rs

elph_ai/api/
codex_transport.rs
