From 350ffbe96d78fbee2ca3442a92907768b94673db Mon Sep 17 00:00:00 2001 From: Zyy Date: Fri, 12 Jun 2026 11:52:44 +0800 Subject: [PATCH] feat: add codex websocket responses transport Co-authored-by: Codex --- Cargo.lock | 54 ++++++- Cargo.toml | 1 + docs/implementation-status.md | 6 +- .../2026-06-12-openai-gpt-codex-parity.md | 13 +- ...26-06-12-openai-gpt-codex-parity-design.md | 1 + src/codex/client.rs | 23 ++- src/codex/websocket.rs | 150 +++++++++++++++++- src/http/v1.rs | 106 +++++++++++++ tests/codex_websocket_test.rs | 111 +++++++++++++ tests/v1_upstream_route_test.rs | 72 +++++++++ 10 files changed, 526 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 03a487371..4628fe3c6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -393,6 +393,7 @@ dependencies = [ "tempfile", "thiserror", "tokio", + "tokio-tungstenite", "tower", "tower-http", "tracing", @@ -653,6 +654,12 @@ dependencies = [ "cmov", ] +[[package]] +name = "data-encoding" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" + [[package]] name = "deadpool" version = "0.12.3" @@ -2314,6 +2321,17 @@ dependencies = [ "unsafe-libyaml", ] +[[package]] +name = "sha1" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +dependencies = [ + "cfg-if", + "cpufeatures 0.2.17", + "digest 0.10.7", +] + [[package]] name = "sha1" version = "0.11.0" @@ -2528,7 +2546,7 @@ dependencies = [ "log", "percent-encoding", "serde", - "sha1", + "sha1 0.11.0", "sha2 0.11.0", "sqlx-core", "thiserror", @@ -2814,6 +2832,22 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-tungstenite" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f72a05e828585856dacd553fba484c242c46e391fb0e58917c942ee9202915c" +dependencies = [ + "futures-util", + "log", + "rustls", + "rustls-native-certs", + "rustls-pki-types", + "tokio", + "tokio-rustls", + "tungstenite", +] + [[package]] name = "tokio-util" version = "0.7.18" @@ -3005,6 +3039,24 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "tungstenite" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c01152af293afb9c7c2a57e4b559c5620b421f6d133261c60dd2d0cdb38e6b8" +dependencies = [ + "bytes", + "data-encoding", + "http", + "httparse", + "log", + "rand 0.9.4", + "rustls", + "rustls-pki-types", + "sha1 0.10.6", + "thiserror", +] + [[package]] name = "typeid" version = "1.0.3" diff --git a/Cargo.toml b/Cargo.toml index c64386b07..c6f72f44b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,6 +31,7 @@ sha2 = "0.11.0" sqlx = { version = "0.9.0", features = ["runtime-tokio", "tls-rustls", "sqlite", "chrono", "uuid", "migrate"] } thiserror = "2.0.18" tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "signal", "fs", "time", "sync", "net"] } +tokio-tungstenite = { version = "0.29.0", default-features = false, features = ["connect", "rustls-tls-native-roots"] } tower = { version = "0.5.3", features = ["util"] } tower-http = { version = "0.6.11", features = ["trace", "cors", "request-id", "timeout"] } tracing = "0.1.44" diff --git a/docs/implementation-status.md b/docs/implementation-status.md index 63b7faefa..2f1be8373 100644 --- a/docs/implementation-status.md +++ b/docs/implementation-status.md @@ -16,11 +16,11 @@ Update this file before each feature commit. | Account pool and refresh | Partial | `0f97c5c` + `2f25db8` + `4dc0c9a` + `c2cf7a6` + `636fa79` + `fdb9374` | `cargo test account_pool_skips_expired_disabled_banned_and_quota_exhausted_accounts`; `cargo test --test account_pool_scheduling_test`; `cargo test --test account_refresh_test`; `cargo test --test oauth_refresh_test`; `cargo test --test v1_upstream_route_test`; `cargo test --test account_pool_startup_test`; full `fmt/test/clippy` before commit | Implements base account status filtering, refresh scheduler semantics, real OpenAI OAuth refresh-token exchange, 401-triggered refresh and single retry for `/v1/responses`, max concurrent slots, stale slot cleanup, least-used/round-robin/sticky strategies, tier priority, quota-cache filtering, Cloudflare cooldown skip, request staggering metadata, preferred account selection, exclude IDs, capacity summary, settings-backed pool construction, same-process import-to-pool sync, and startup pool snapshots from SQLite with usage-based last-used metadata. Still missing durable release-time slot state, broader upstream retry/fallback wiring, and admin account CRUD. | | Account repository | Completed | `14fe706` | `cargo test --test account_repository_test`; `cargo fmt --check`; `cargo clippy --all-targets --all-features --locked -- -D warnings`; `cargo test` | Adds encrypted SQLite account token storage, decrypting reads, cursor pagination, status/label updates, release-time usage counter persistence, and refresh-token-preserving token updates. | | Translation | Completed | `eff182e` + `b9c30ba` | `cargo test --test routes_chat_test`; `cargo test --test chat_completions_route_test`; full `cargo fmt --check && cargo test && cargo clippy --all-targets --all-features --locked -- -D warnings` before `b9c30ba` | Adds OpenAI Chat request types, expanded Codex Responses request type, OpenAI Chat to Codex translation for system/developer instructions, multimodal content, tool/function calls, response formats, reasoning effort, service tier, and parallel tool calls. Codex SSE now converts back to OpenAI chat completions and chat chunks with text deltas, reasoning content, tool calls, usage details, and OpenAI-compatible error bodies. | -| HTTP routes | Partial | `ab0ad8f` + `3830b21` + `6179d44` + `cfec8ba` + `6341336` + `228febd` + `c2cf7a6` + `e9fdea0` + `ee0ed90` + `fdb9374` + `9ca5c7d` + `b9c30ba` | `cargo test v1_requires_client_api_key_not_admin_cookie`; `cargo test models_route_returns_openai_compatible_codex_model_list`; `cargo test --test admin_settings_route_test`; `cargo test --test model_catalog_test`; `cargo test --test admin_accounts_route_test`; `cargo test --test admin_api_keys_route_test`; `cargo test --test admin_usage_stats_route_test`; `cargo test --test v1_upstream_route_test`; `cargo test --test chat_completions_route_test`; full `cargo fmt --check && cargo test && cargo clippy --all-targets --all-features --locked -- -D warnings` before `b9c30ba` | Adds Axum router, health route, `/v1/models`, `/v1/models/catalog`, `/v1/models/{id}`, `/v1/models/{id}/info`, `/debug/models`, `/admin/accounts`, `/admin/accounts/import`, `/admin/api-keys`, `/admin/usage-stats`, `/admin/usage-stats/summary`, `/v1/responses` HTTP SSE collection with output item/delta reconstruction plus `output_text` sync for non-streaming clients, default streaming Responses passthrough when `stream` is omitted, Responses field parity for service tier, tool choice, parallel tool calls, text format, prompt cache key, include, client metadata, reasoning include default, and Codex context headers, plus a dedicated `/v1/chat/completions` route backed by imported Codex accounts. `/v1` unauthorized and model errors remain OpenAI-compatible, global request-id response headers are preserved, `gpt-5.5` remains the default/test model baseline, and `/admin/settings` is read-only with admin-session auth and lower camelCase envelope. `/admin/accounts` reads metadata only and does not decrypt tokens for list views. Still missing OAuth/device/CLI account login flows, admin account mutation/health/quota/Cookie routes, full admin API-key mutation/export/import routes, settings mutation/persistence, diagnostics, backend model refresh, model cache persistence, live Chat streaming transform, and real WebSocket Responses transport. | +| HTTP routes | Partial | `ab0ad8f` + `3830b21` + `6179d44` + `cfec8ba` + `6341336` + `228febd` + `c2cf7a6` + `e9fdea0` + `ee0ed90` + `fdb9374` + `9ca5c7d` + `b9c30ba` + `pending` | `cargo test v1_requires_client_api_key_not_admin_cookie`; `cargo test models_route_returns_openai_compatible_codex_model_list`; `cargo test --test admin_settings_route_test`; `cargo test --test model_catalog_test`; `cargo test --test admin_accounts_route_test`; `cargo test --test admin_api_keys_route_test`; `cargo test --test admin_usage_stats_route_test`; `cargo test --test v1_upstream_route_test`; `cargo test --test chat_completions_route_test`; `cargo test --test codex_websocket_test`; full `cargo fmt --check && cargo test && cargo clippy --all-targets --all-features --locked -- -D warnings` before `b9c30ba`; Task 4 targeted verification before `pending` | Adds Axum router, health route, `/v1/models`, `/v1/models/catalog`, `/v1/models/{id}`, `/v1/models/{id}/info`, `/debug/models`, `/admin/accounts`, `/admin/accounts/import`, `/admin/api-keys`, `/admin/usage-stats`, `/admin/usage-stats/summary`, `/v1/responses` HTTP SSE collection with output item/delta reconstruction plus `output_text` sync for non-streaming clients, default streaming Responses passthrough when `stream` is omitted, WebSocket-backed Responses streaming for `previous_response_id`, Responses field parity for service tier, tool choice, parallel tool calls, text format, prompt cache key, include, client metadata, reasoning include default, and Codex context headers, plus a dedicated `/v1/chat/completions` route backed by imported Codex accounts. `/v1` unauthorized and model errors remain OpenAI-compatible, global request-id response headers are preserved, `gpt-5.5` remains the default/test model baseline, and `/admin/settings` is read-only with admin-session auth and lower camelCase envelope. `/admin/accounts` reads metadata only and does not decrypt tokens for list views. Still missing OAuth/device/CLI account login flows, admin account mutation/health/quota/Cookie routes, full admin API-key mutation/export/import routes, settings mutation/persistence, diagnostics, backend model refresh, model cache persistence, and live Chat streaming transform. | | Admin/frontend response contract | Completed | `1c34ea7` + `b0c10f2` | `cargo test admin_envelope_serializes_lower_camel_case_request_id`; `cargo test admin_logs_reject_missing_admin_session_cookie`; `cargo test v1_response_has_request_id_header_without_admin_body_field`; full `fmt/test/clippy` before commit | Adds reusable admin envelopes with body codes, lower camelCase JSON, paginated list shape, requestId bodies, admin session cookie enforcement, and keeps `/v1/*` OpenAI-compatible. | | HTTP auth boundary | Completed | `cec4607` + `e9fdea0` | `cargo test --test http_auth_test`; `cargo test --test routes_responses_test`; `cargo test --test admin_api_keys_route_test`; `cargo test --test admin_logs_route_test`; `cargo fmt --check`; `cargo clippy --all-targets --all-features --locked -- -D warnings`; `cargo test` | Adds a real `src/http/auth.rs` boundary for client API key and admin session cookie parsing. `/v1/*` accepts only enabled local `Bearer cpr_...` client API keys verified from SQLite, admin session parsing ignores client API keys, and route handlers no longer own ad hoc credential parsing. | | Model catalog | Partial | `6341336` | `cargo test --test model_catalog_test`; `cargo test model_catalog_route_returns_codex_metadata_without_alias_entries`; full `fmt/test/clippy` before commit | Adds a native static Codex model catalog with `gpt-5.5` default, configured aliases, reasoning/service-tier suffix parsing, OpenAI-compatible `/v1/models`, extended catalog/detail/info routes, and debug summary. Still missing backend Codex model fetching per plan, cache persistence, `/admin/refresh-models`, plan allowlist integration from fetched models, and full model metadata parity. | -| Upstream lifecycle | Partial | `88cad87` + `225fbff` + `c2cf7a6` + `636fa79` + `ee0ed90` + `fdb9374` + `5bf012c` + `9ca5c7d` + `b9c30ba` | `cargo test responses_route_rejects_non_codex_provider_models`; `cargo test v1_requires_client_api_key_not_admin_cookie`; `cargo test --test codex_sse_usage_test`; `cargo test --test codex_client_test`; `cargo test --test oauth_refresh_test`; `cargo test --test v1_upstream_route_test`; `cargo test --test account_pool_startup_test`; full `cargo fmt --check && cargo test && cargo clippy --all-targets --all-features --locked -- -D warnings` before `b9c30ba` | `/v1` enforces client API key shape, rejects non-Codex provider models through the model catalog, returns OpenAI-compatible errors, includes reusable SSE/usage helpers, sends Codex Desktop HTTP SSE requests through imported accounts, replays/captures encrypted Cookies, forwards Codex request context headers, normalizes `service_tier: fast` to upstream `priority`, collects `response.completed`, reconstructs missing assistant output from `response.output_item.done` or text deltas, synchronizes `output_text` for non-streaming clients, passes through Responses streaming SSE by default, adapts Chat Completions responses/chunks, refreshes and retries once after upstream 401 when a refresh token is stored, records token usage, writes v1 lifecycle event logs, records refresh failure classification for upstream 401, and restores the account pool from SQLite on startup. Still missing SSE error/failed mapping, rate-limit header capture, retry/fallback across accounts beyond refresh retry, 429 quota cache, 402/403 classification, previous-response/session affinity, and live Chat streaming transform. | -| Codex WebSocket boundary | Completed | `cc0516d` + `b9c30ba` | `cargo test --test codex_websocket_test`; full `cargo fmt --check && cargo test && cargo clippy --all-targets --all-features --locked -- -D warnings` before `b9c30ba` | Adds `src/codex/websocket.rs` with tested transport classification for HTTP SSE versus WebSocket-only requests. `previous_response_id` and explicit WebSocket requests now fail through a typed boundary instead of being silently downgraded, and the local-only `use_websocket` transport flag is not serialized into upstream JSON. Real WebSocket transport remains an upstream lifecycle follow-up. | +| Upstream lifecycle | Partial | `88cad87` + `225fbff` + `c2cf7a6` + `636fa79` + `ee0ed90` + `fdb9374` + `5bf012c` + `9ca5c7d` + `b9c30ba` + `pending` | `cargo test responses_route_rejects_non_codex_provider_models`; `cargo test v1_requires_client_api_key_not_admin_cookie`; `cargo test --test codex_sse_usage_test`; `cargo test --test codex_client_test`; `cargo test --test oauth_refresh_test`; `cargo test --test v1_upstream_route_test`; `cargo test --test account_pool_startup_test`; `cargo test --test codex_websocket_test`; full `cargo fmt --check && cargo test && cargo clippy --all-targets --all-features --locked -- -D warnings` before `b9c30ba`; Task 4 targeted verification before `pending` | `/v1` enforces client API key shape, rejects non-Codex provider models through the model catalog, returns OpenAI-compatible errors, includes reusable SSE/usage helpers, sends Codex Desktop HTTP SSE requests through imported accounts, sends WebSocket `response.create` requests when `previous_response_id` requires upstream history, replays/captures encrypted Cookies, forwards Codex request context headers, normalizes `service_tier: fast` to upstream `priority`, collects `response.completed`, reconstructs missing assistant output from `response.output_item.done` or text deltas, synchronizes `output_text` for non-streaming clients, passes through Responses streaming SSE by default, adapts Chat Completions responses/chunks, refreshes and retries once after upstream 401 when a refresh token is stored, records token usage, writes v1 lifecycle event logs, records refresh failure classification for upstream 401, and restores the account pool from SQLite on startup. Still missing SSE error/failed mapping, rate-limit header capture, retry/fallback across accounts beyond refresh retry, 429 quota cache, 402/403 classification, previous-response account affinity, and live Chat streaming transform. | +| Codex WebSocket boundary | Completed | `cc0516d` + `b9c30ba` + `pending` | `cargo test --test codex_websocket_test`; `cargo test --test v1_upstream_route_test v1_responses_should_use_websocket_for_previous_response_id_streaming`; full `cargo fmt --check && cargo test && cargo clippy --all-targets --all-features --locked -- -D warnings` before `b9c30ba`; Task 4 targeted verification before `pending` | Adds `src/codex/websocket.rs` with tested transport classification for HTTP SSE versus WebSocket-only requests. `previous_response_id` and explicit WebSocket requests use typed transport selection, local-only `use_websocket` is not serialized into upstream JSON, WebSocket `response.create` carries Codex headers and parity fields, incoming WebSocket messages are emitted as SSE-compatible events, and `/v1/responses` returns WebSocket-backed `text/event-stream` output for `previous_response_id` streaming requests. | | Fingerprint updates | Partial | `dd26d31` + `e32220a` | `cargo test update_manifest_updates_app_version_and_build_number`; `cargo test fingerprint_updater_fetches_manifest_and_persists_history`; full `fmt/test/clippy` before commit | Parses a JSON update manifest and persists app version/build number history to SQLite. Still missing original appcast XML polling, update-state persistence, runtime header fingerprint selection from the latest stored fingerprint, and optional local Codex.app Chromium extraction. | | Runtime docs and packaging | Completed | `33b63f3` | `cargo fmt --check`; `cargo clippy --all-targets --all-features --locked -- -D warnings`; `cargo test`; `cargo run`; `curl -s -i http://127.0.0.1:8080/health` | README scope, `.env.example`, Dockerfile, compose file, and health smoke test are complete. | diff --git a/docs/superpowers/plans/2026-06-12-openai-gpt-codex-parity.md b/docs/superpowers/plans/2026-06-12-openai-gpt-codex-parity.md index 9a2850be0..3e80835a6 100644 --- a/docs/superpowers/plans/2026-06-12-openai-gpt-codex-parity.md +++ b/docs/superpowers/plans/2026-06-12-openai-gpt-codex-parity.md @@ -8,7 +8,7 @@ **Tech Stack:** Rust, Axum, Reqwest/rustls, serde, sqlx/SQLite, wiremock, tokio, futures, optional WebSocket transport dependency. -**Checkpoint 2026-06-12:** Commit `b9c30ba` completed Tasks 1-3 for the in-scope OpenAI GPT/Codex path: Chat Completions route/translation/output, Responses request field parity, default Responses streaming, Codex context headers, and local-only WebSocket flag serialization cleanup. Remaining major work starts at Task 4: real WebSocket transport for `previous_response_id`, then account fallback/rate limits, model refresh/cache, and scoped admin operations. +**Checkpoint 2026-06-12:** Commit `b9c30ba` completed Tasks 1-3 for the in-scope OpenAI GPT/Codex path: Chat Completions route/translation/output, Responses request field parity, default Responses streaming, Codex context headers, and local-only WebSocket flag serialization cleanup. Task 4 is implemented in the current pending stage: `previous_response_id` requests use WebSocket `response.create` and return SSE-compatible output. Remaining major work starts at Task 5: account fallback/rate limits, model refresh/cache, and scoped admin operations. --- @@ -56,12 +56,15 @@ - Modify: `Cargo.toml` - Modify: `src/codex/websocket.rs` - Modify: `src/codex/client.rs` +- Modify: `src/http/v1.rs` - Test: `tests/codex_websocket_test.rs` +- Test: `tests/v1_upstream_route_test.rs` -- [ ] Write failing tests proving `previous_response_id` uses WebSocket and HTTP fallback is not used when server-side history is required. -- [ ] Run: `cargo test --test codex_websocket_test` -- [ ] Add WebSocket response.create transport that emits SSE-compatible output. -- [ ] Run: `cargo test --test codex_websocket_test` +- [x] Write failing tests proving `previous_response_id` uses WebSocket and HTTP fallback is not used when server-side history is required. +- [x] Run: `cargo test --test codex_websocket_test` +- [x] Add WebSocket response.create transport that emits SSE-compatible output. +- [x] Run: `cargo test --test codex_websocket_test` +- [x] Run: `cargo test --test v1_upstream_route_test v1_responses_should_use_websocket_for_previous_response_id_streaming` ### Task 5: Upstream Retry, Fallback, And Rate Limits diff --git a/docs/superpowers/specs/2026-06-12-openai-gpt-codex-parity-design.md b/docs/superpowers/specs/2026-06-12-openai-gpt-codex-parity-design.md index d7fa79cca..60dbe71ef 100644 --- a/docs/superpowers/specs/2026-06-12-openai-gpt-codex-parity-design.md +++ b/docs/superpowers/specs/2026-06-12-openai-gpt-codex-parity-design.md @@ -23,6 +23,7 @@ Responses compatibility keeps native Responses passthrough semantics but adds th - `src/translation/openai_to_codex.rs`: OpenAI Chat request schema and Chat-to-Codex translation, including messages, tools/functions, multimodal input, response formats, reasoning effort, and service tier. - `src/translation/codex_to_openai.rs`: Codex SSE to OpenAI Chat JSON/SSE conversion, including text, reasoning, tool calls, usage, and errors. - `src/codex/types.rs`: Expanded Codex request/response types for Responses and WebSocket. +- `src/codex/websocket.rs`: WebSocket transport selection and `response.create` conversion into SSE-compatible events for upstream-history requests such as `previous_response_id`. - `src/codex/client.rs`: HTTP SSE and WebSocket Codex transport, context headers, cookies, and rate-limit header capture. - `src/http/v1.rs`: Separate `/v1/responses` and `/v1/chat/completions` handlers sharing upstream lifecycle helpers. - `src/accounts/pool.rs` and repositories: fallback selection, quota/rate-limit updates, durable usage release, and account state mutation hooks. diff --git a/src/codex/client.rs b/src/codex/client.rs index d5b66b4a7..d1622b092 100644 --- a/src/codex/client.rs +++ b/src/codex/client.rs @@ -10,7 +10,10 @@ use crate::{ sse::SseError, types::CodexResponsesRequest, usage::{extract_sse_usage, TokenUsage}, - websocket::{ensure_http_sse_supported, WebSocketSupportError}, + websocket::{ + create_response_via_websocket, ensure_http_sse_supported, transport_for_request, + CodexTransport, CodexWebSocketError, WebSocketSupportError, + }, }, fingerprint::model::Fingerprint, }; @@ -25,6 +28,8 @@ pub enum CodexClientError { InvalidHeaderValue(#[from] reqwest::header::InvalidHeaderValue), #[error("unsupported transport: {0}")] UnsupportedTransport(#[from] WebSocketSupportError), + #[error("websocket transport error: {0}")] + WebSocket(#[from] CodexWebSocketError), #[error("invalid upstream SSE response: {0}")] InvalidSse(#[from] SseError), #[error("upstream returned status {status}: {body}")] @@ -83,6 +88,22 @@ impl CodexBackendClient { request: &CodexResponsesRequest, context: CodexRequestContext<'_>, ) -> CodexClientResult { + if transport_for_request(request) == CodexTransport::WebSocketRequired { + let body = create_response_via_websocket( + &self.base_url, + request, + self.request_headers(context)?, + ) + .await?; + let usage = extract_sse_usage(&body)?; + return Ok(CodexBackendResponse { + body, + usage, + turn_state: None, + set_cookie_headers: Vec::new(), + }); + } + let response = self.send_response_request(request, context).await?; let status = response.status(); let turn_state = turn_state(&response); diff --git a/src/codex/websocket.rs b/src/codex/websocket.rs index 957edd193..e52d9ec3d 100644 --- a/src/codex/websocket.rs +++ b/src/codex/websocket.rs @@ -1,6 +1,15 @@ use thiserror::Error; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::{ + connect_async, + tungstenite::{http::Request as WsRequest, Message}, +}; -use crate::codex::types::CodexResponsesRequest; +use futures::{SinkExt, StreamExt}; +use reqwest::header::HeaderMap; +use serde_json::{json, Value}; + +use crate::codex::{sse::encode_sse_event, types::CodexResponsesRequest}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum CodexTransport { @@ -16,6 +25,16 @@ pub enum WebSocketSupportError { ExplicitWebSocketRequired, } +#[derive(Debug, Error)] +pub enum CodexWebSocketError { + #[error("invalid WebSocket request: {0}")] + InvalidRequest(#[from] tokio_tungstenite::tungstenite::http::Error), + #[error("websocket transport error: {0}")] + Transport(#[from] tokio_tungstenite::tungstenite::Error), + #[error("websocket response ended before any events")] + EmptyResponse, +} + pub fn transport_for_request(request: &CodexResponsesRequest) -> CodexTransport { if request.previous_response_id.is_some() || request.use_websocket { CodexTransport::WebSocketRequired @@ -35,3 +54,132 @@ pub fn ensure_http_sse_supported( } Ok(()) } + +pub async fn create_response_via_websocket( + base_url: &str, + request: &CodexResponsesRequest, + headers: HeaderMap, +) -> Result { + let ws_request = build_ws_request(base_url, headers)?; + let (mut websocket, _) = connect_async(ws_request).await?; + websocket + .send(Message::Text( + websocket_request_body(request).to_string().into(), + )) + .await?; + + let mut body = String::new(); + while let Some(message) = websocket.next().await { + let message = message?; + let Some(raw) = websocket_message_text(message) else { + continue; + }; + let event = websocket_event_type(&raw); + body.push_str(&encode_sse_event( + event.as_deref().unwrap_or_default(), + &raw, + )); + if event.as_deref().is_some_and(is_terminal_websocket_event) { + break; + } + } + if body.is_empty() { + return Err(CodexWebSocketError::EmptyResponse); + } + Ok(body) +} + +fn build_ws_request( + base_url: &str, + headers: HeaderMap, +) -> Result, CodexWebSocketError> { + let mut request = websocket_url(base_url).into_client_request()?; + for (name, value) in &headers { + let Ok(name) = + tokio_tungstenite::tungstenite::http::HeaderName::from_bytes(name.as_str().as_bytes()) + else { + continue; + }; + let Ok(value) = + tokio_tungstenite::tungstenite::http::HeaderValue::from_bytes(value.as_bytes()) + else { + continue; + }; + request.headers_mut().insert(name, value); + } + Ok(request) +} + +fn websocket_url(base_url: &str) -> String { + let url = format!("{}/codex/responses", base_url.trim_end_matches('/')); + if let Some(rest) = url.strip_prefix("https://") { + format!("wss://{rest}") + } else if let Some(rest) = url.strip_prefix("http://") { + format!("ws://{rest}") + } else { + url + } +} + +fn websocket_request_body(request: &CodexResponsesRequest) -> Value { + let mut body = json!({ + "type": "response.create", + "model": request.model, + "instructions": request.instructions, + "input": request.input, + "stream": true, + "store": false, + "tool_choice": request.tool_choice.clone().unwrap_or_else(|| json!("auto")), + "parallel_tool_calls": request.parallel_tool_calls.unwrap_or(true), + }); + if let Some(previous_response_id) = &request.previous_response_id { + body["previous_response_id"] = Value::String(previous_response_id.clone()); + } + if let Some(reasoning) = &request.reasoning { + body["reasoning"] = reasoning.clone(); + } + if let Some(tools) = request.tools.as_ref().filter(|tools| !tools.is_empty()) { + body["tools"] = Value::Array(tools.clone()); + } + if let Some(text) = &request.text { + body["text"] = text.clone(); + } + if let Some(service_tier) = &request.service_tier { + body["service_tier"] = Value::String(service_tier.clone()); + } + if let Some(prompt_cache_key) = &request.prompt_cache_key { + body["prompt_cache_key"] = Value::String(prompt_cache_key.clone()); + } + if let Some(include) = request + .include + .as_ref() + .filter(|include| !include.is_empty()) + { + body["include"] = Value::Array(include.iter().cloned().map(Value::String).collect()); + } + if let Some(client_metadata) = &request.client_metadata { + body["client_metadata"] = client_metadata.clone(); + } + body +} + +fn websocket_message_text(message: Message) -> Option { + match message { + Message::Text(text) => Some(text.to_string()), + Message::Binary(bytes) => String::from_utf8(bytes.to_vec()).ok(), + _ => None, + } +} + +fn websocket_event_type(raw: &str) -> Option { + serde_json::from_str::(raw).ok().and_then(|value| { + value + .get("type") + .and_then(Value::as_str) + .map(str::to_string) + }) +} + +fn is_terminal_websocket_event(event: &str) -> bool { + event == "response.completed" || event == "response.failed" || event == "error" +} diff --git a/src/http/v1.rs b/src/http/v1.rs index 02dd3a5db..523ad2e38 100644 --- a/src/http/v1.rs +++ b/src/http/v1.rs @@ -33,6 +33,7 @@ use crate::{ sse::parse_sse_events, types::CodexResponsesRequest, usage::{extract_sse_usage, TokenUsage}, + websocket::{transport_for_request, CodexTransport}, }, fingerprint::model::Fingerprint, http::{auth::client_api_key, middleware::RequestId}, @@ -695,6 +696,10 @@ async fn responses_stream( account: Account, log_context: V1LogContext, ) -> Response { + if transport_for_request(&request) == CodexTransport::WebSocketRequired { + return responses_websocket_stream(state, request, account, log_context).await; + } + let stream_response = send_codex_stream_request_with_refresh_retry( &state, &request, @@ -781,6 +786,107 @@ async fn responses_stream( }) } +async fn responses_websocket_stream( + state: AppState, + request: CodexResponsesRequest, + account: Account, + log_context: V1LogContext, +) -> Response { + let response = send_codex_request_with_refresh_retry( + &state, + &request, + &account, + log_context.request_id.as_str(), + ) + .await; + state.account_pool().lock().await.release(&account.id); + + let response = match response { + Ok(response) => response, + Err(error) => { + let error_response = codex_client_error_response(error); + log_v1_response( + &state, + &log_context, + error_response.0, + EventLevel::Error, + "v1 responses websocket stream upstream request failed", + json!({"stream": true, "transport": "websocket"}), + ) + .await; + return error_response.into_response(); + } + }; + + if persist_upstream_cookies(&state, &account.id, &response.set_cookie_headers) + .await + .is_err() + { + log_v1_response( + &state, + &log_context, + StatusCode::INTERNAL_SERVER_ERROR, + EventLevel::Error, + "v1 responses websocket stream cookie persistence failed", + json!({"stream": true, "transport": "websocket", "cookieStoreError": true}), + ) + .await; + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(openai_error( + "Failed to persist upstream cookies", + "cookie_store_error", + )), + ) + .into_response(); + } + + let mut level = EventLevel::Info; + let mut message = "v1 responses websocket stream completed"; + let mut metadata = json!({ + "stream": true, + "transport": "websocket", + "usage": response.usage, + }); + if let Some(usage) = response.usage { + if record_usage(&state, &account.id, usage).await.is_err() { + level = EventLevel::Warn; + message = "v1 responses websocket stream completed with usage store error"; + metadata = json!({ + "stream": true, + "transport": "websocket", + "usage": usage, + "usageStoreError": true, + }); + } + } + log_v1_response( + &state, + &log_context, + StatusCode::OK, + level, + message, + metadata, + ) + .await; + + Response::builder() + .status(StatusCode::OK) + .header(CONTENT_TYPE, "text/event-stream") + .header(CACHE_CONTROL, "no-cache") + .body(Body::from(response.body)) + .unwrap_or_else(|_| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(openai_error( + "Failed to build stream response", + "stream_response_error", + )), + ) + .into_response() + }) +} + async fn refresh_account_after_unauthorized( state: &AppState, request: &CodexResponsesRequest, diff --git a/tests/codex_websocket_test.rs b/tests/codex_websocket_test.rs index 9f888fbf6..a11a1ff45 100644 --- a/tests/codex_websocket_test.rs +++ b/tests/codex_websocket_test.rs @@ -1,4 +1,18 @@ +use std::sync::Arc; + +use futures::{SinkExt, StreamExt}; +use serde_json::{json, Value}; +use tokio::{net::TcpListener, sync::Mutex}; +use tokio_tungstenite::{ + accept_hdr_async, + tungstenite::{ + handshake::server::{Request as WsRequest, Response as WsResponse}, + Message, + }, +}; + use codex_proxy_rs::codex::{ + client::{build_reqwest_client, CodexBackendClient, CodexRequestContext}, types::CodexResponsesRequest, websocket::{ensure_http_sse_supported, transport_for_request, CodexTransport}, }; @@ -37,6 +51,103 @@ fn use_websocket_should_not_serialize_to_upstream_json() { assert!(body.get("useWebSocket").is_none()); } +#[tokio::test] +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn previous_response_id_should_use_websocket_transport() { + let received_headers = Arc::new(Mutex::new(None)); + let received_request = Arc::new(Mutex::new(None)); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let headers_for_task = Arc::clone(&received_headers); + let request_for_task = Arc::clone(&received_request); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = + accept_hdr_async(stream, |request: &WsRequest, response: WsResponse| { + let mut headers = Vec::new(); + for (name, value) in request.headers() { + let value = value.to_str().unwrap_or_default().to_string(); + headers.push((name.as_str().to_string(), value)); + } + let headers_for_callback = Arc::clone(&headers_for_task); + tokio::spawn(async move { + *headers_for_callback.lock().await = Some(headers); + }); + Ok(response) + }) + .await + .unwrap(); + let message = websocket.next().await.unwrap().unwrap(); + let text = message.into_text().unwrap(); + *request_for_task.lock().await = Some(serde_json::from_str::(&text).unwrap()); + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_ws", + "object": "response", + "usage": { + "input_tokens": 2, + "output_tokens": 1 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + + let mut request = CodexResponsesRequest::new_http_sse("gpt-5.5", "be brief", Vec::new()); + request.previous_response_id = Some("resp_prev".to_string()); + let client = CodexBackendClient::new( + build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_rs::fingerprint::model::Fingerprint::default_for_tests(), + ); + + let response = client + .create_response( + &request, + CodexRequestContext { + access_token: "access-token", + account_id: Some("chatgpt-account"), + request_id: "req_ws", + turn_state: None, + turn_metadata: None, + beta_features: None, + include_timing_metrics: None, + version: None, + codex_window_id: None, + parent_thread_id: None, + cookie_header: None, + }, + ) + .await + .unwrap(); + + server.await.unwrap(); + let request = received_request.lock().await.clone().unwrap(); + assert_eq!(request["type"], "response.create"); + assert_eq!(request["model"], "gpt-5.5"); + assert_eq!(request["instructions"], "be brief"); + assert_eq!(request["previous_response_id"], "resp_prev"); + assert_eq!(request["stream"], true); + assert!(response.body.contains("event: response.completed")); + assert!(response.body.contains("\"id\":\"resp_ws\"")); + assert_eq!(response.usage.unwrap().input_tokens, 2); + let headers = received_headers.lock().await.clone().unwrap(); + assert!(headers + .iter() + .any(|(name, value)| { name == "authorization" && value == "Bearer access-token" })); +} + fn base_request() -> CodexResponsesRequest { CodexResponsesRequest::new_http_sse("gpt-5.5", "", Vec::new()) } diff --git a/tests/v1_upstream_route_test.rs b/tests/v1_upstream_route_test.rs index 19bd10d9c..7822ebf1e 100644 --- a/tests/v1_upstream_route_test.rs +++ b/tests/v1_upstream_route_test.rs @@ -6,8 +6,11 @@ use axum::{ http::{Request, StatusCode}, }; use chrono::Utc; +use futures::{SinkExt, StreamExt}; use secrecy::ExposeSecret; use serde_json::{json, Value}; +use tokio::{net::TcpListener, sync::oneshot}; +use tokio_tungstenite::{accept_async, tungstenite::Message}; use tower::ServiceExt; use wiremock::{ matchers::{body_json, header, method, path}, @@ -743,6 +746,75 @@ async fn v1_responses_should_default_to_streaming_when_stream_is_omitted() { assert!(body.contains("event: response.completed")); } +#[tokio::test] +async fn v1_responses_should_use_websocket_for_previous_response_id_streaming() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (request_tx, request_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let message = websocket.next().await.unwrap().unwrap(); + let request = serde_json::from_str::(&message.into_text().unwrap()).unwrap(); + request_tx.send(request).unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_route_ws", + "object": "response", + "usage": { + "input_tokens": 8, + "output_tokens": 5 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let imported = build_imported_app(format!("http://{addr}")).await; + + let response = imported + .app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header( + "authorization", + format!("Bearer {}", imported.client_api_key), + ) + .header("content-type", "application/json") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"previous_response_id":"resp_prev"}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + assert!(content_type.starts_with("text/event-stream")); + let body = response_text(response).await; + assert!(body.contains("event: response.completed")); + assert!(body.contains("\"id\":\"resp_route_ws\"")); + let request = request_rx.await.unwrap(); + assert_eq!(request["type"], "response.create"); + assert_eq!(request["previous_response_id"], "resp_prev"); + server.await.unwrap(); +} + #[tokio::test] async fn v1_responses_should_refresh_after_401_and_retry_non_stream() { let server = MockServer::start().await;