diff --git a/Dockerfile b/Dockerfile index 8829eb3..d532129 100644 --- a/Dockerfile +++ b/Dockerfile @@ -51,7 +51,7 @@ RUN mkdir -p server/src clients/rust/src cxtx/src && \ COPY server/ ./server/ COPY clients/ ./clients/ COPY cxtx/ ./cxtx/ -RUN touch server/src/main.rs && \ +RUN find server/src clients/rust/src cxtx/src -type f -exec touch {} + && \ cargo build --release --manifest-path server/Cargo.toml # ============================================ diff --git a/cxtx/Cargo.toml b/cxtx/Cargo.toml index 135ae65..e45394c 100644 --- a/cxtx/Cargo.toml +++ b/cxtx/Cargo.toml @@ -8,7 +8,7 @@ description = "CLI wrapper that captures claude/codex provider traffic and uploa [dependencies] anyhow = "1.0" async-stream = "0.3" -axum = { version = "0.7", features = ["macros"] } +axum = { version = "0.7", features = ["macros", "ws"] } base64 = "0.22" bytes = "1.10" chrono = { version = "0.4", features = ["clock", "serde"] } @@ -20,6 +20,7 @@ reqwest = { version = "0.12", default-features = false, features = ["json", "rus serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" tokio = { version = "1", features = ["fs", "io-util", "macros", "net", "process", "rt-multi-thread", "signal", "sync", "time"] } +tokio-tungstenite = { version = "0.24", features = ["rustls-tls-webpki-roots"] } url = "2.5" uuid = { version = "1.18", features = ["v4"] } diff --git a/cxtx/README.md b/cxtx/README.md index 57de6bb..a8d4a4a 100644 --- a/cxtx/README.md +++ b/cxtx/README.md @@ -18,7 +18,10 @@ cargo build --release -p cxtx ```bash # Wrap Codex and send captured turns to the local CXDB HTTP endpoint -./target/release/cxtx codex -- --model gpt-5 +./target/release/cxtx --local codex -- --model gpt-5 + +# Wrap Claude and send captured turns to the local CXDB HTTP endpoint +./target/release/cxtx --local claude -- --print stream # Wrap Claude against a specific CXDB server ./target/release/cxtx --url http://127.0.0.1:9010 claude -- --print stream @@ -26,9 +29,13 @@ cargo build --release -p cxtx `cxtx` preserves child stdin, stdout, stderr, and exit status. On successful execution it does not write wrapper-authored stdout. If CXDB is unavailable, it still launches the child, enters queued-delivery mode, and records delivery state in the local ledger until delivery recovers or shutdown drain completes. +For `codex`, `cxtx` now hardens the launch contract so interactive traffic stays on the local proxy path that OSS capture depends on. It injects both `OPENAI_*` and `CXTX_OPENAI_*` proxy base-url env vars, removes inherited upstream base-url overrides before spawning the child, and prepends websocket-disabling flags unless the caller already supplied explicit overrides. + ## Resulting Artifacts - CXDB receives canonical `system`, `user_input`, `assistant_turn`, and tool-related items for the wrapped session. +- Interactive `codex` Responses traffic no longer leaks wrapper/bootstrap scaffolding into uploaded `user_input` turns. +- Websocket-backed provider traffic is captured through the local proxy instead of bypassing ledger and CXDB upload paths. - The first uploaded turn carries `ContextMetadata` and `Provenance`, so the context is queryable in CXDB listings. - `cxtx` publishes the bundled canonical `cxdb.ConversationItem` registry descriptor automatically before the first append when the server does not already have it. - Local evidence is written under `.scratch/cxtx/sessions//`: @@ -42,9 +49,11 @@ cargo build --release -p cxtx - `failed to launch codex` or `failed to launch claude`: - The child binary is missing from `PATH` or is not executable. - `cxtx: CXDB ingest unavailable, entering queued-delivery mode`: - - The wrapper could not reach the configured `--url`. Check the CXDB server, but the child session is still running and the ledger will show queue state. + - The wrapper could not reach the configured default URL, `--local`, or `--url`. Check the CXDB server, but the child session is still running and the ledger will show queue state. - No captured turns appear in CXDB: - - Confirm the child CLI honors the injected provider base URL variables. `cxtx` depends on those environment overrides for transparent capture. + - Confirm the child CLI honors the injected provider base URL variables. `codex` should see both the legacy `OPENAI_*` variables and the `CXTX_OPENAI_*` aliases, and `claude` should see the Anthropic/Claude base-url overrides. +- `codex` still reaches the public OpenAI endpoint directly: + - Check whether the caller explicitly re-enabled websocket features or overrode the proxy base URL in child args. `cxtx` only injects its websocket-disable defaults when the caller did not already make an explicit choice. ## Verification diff --git a/cxtx/src/cli.rs b/cxtx/src/cli.rs index b6189a4..c98273b 100644 --- a/cxtx/src/cli.rs +++ b/cxtx/src/cli.rs @@ -2,20 +2,29 @@ use clap::{Parser, Subcommand}; use crate::provider::ProviderKind; +pub const DEFAULT_LOCAL_CXDB_URL: &str = "http://127.0.0.1:9010"; + #[derive(Debug, Clone, Parser)] #[command( name = "cxtx", about = "Wrap claude or codex, capture provider traffic, and upload canonical conversation context to CXDB", - after_help = "Examples:\n cxtx codex -- --model gpt-5\n cxtx --url http://127.0.0.1:9010 claude -- --print stream" + after_help = "Examples:\n cxtx codex -- --model gpt-5\n cxtx --local claude -- --print stream\n cxtx --url http://127.0.0.1:9010 claude -- --print stream" )] pub struct Cli { #[arg( long, - default_value = "http://127.0.0.1:9010", + default_value = DEFAULT_LOCAL_CXDB_URL, help = "CXDB HTTP base URL used for registry publication, context creation, and turn append" )] pub url: String, + #[arg( + long, + conflicts_with = "url", + help = "Use the local CXDB HTTP endpoint at http://127.0.0.1:9010" + )] + pub local: bool, + #[command(subcommand)] pub command: Command, } @@ -50,6 +59,14 @@ impl Command { } impl Cli { + pub fn effective_url(&self) -> &str { + if self.local { + DEFAULT_LOCAL_CXDB_URL + } else { + &self.url + } + } + pub fn for_tests(provider: ProviderKind, args: Vec, url: &str) -> Self { let command = match provider { ProviderKind::Claude => Command::Claude { args }, @@ -57,7 +74,37 @@ impl Cli { }; Self { url: url.to_string(), + local: false, command, } } } + +#[cfg(test)] +mod tests { + use super::{Cli, Command, DEFAULT_LOCAL_CXDB_URL}; + use clap::Parser; + + #[test] + fn defaults_to_local_cxdb_url() { + let cli = Cli::parse_from(["cxtx", "codex"]); + assert_eq!(cli.effective_url(), DEFAULT_LOCAL_CXDB_URL); + assert!(!cli.local); + assert!(matches!(cli.command, Command::Codex { .. })); + } + + #[test] + fn local_flag_switches_to_local_cxdb_url() { + let cli = Cli::parse_from(["cxtx", "--local", "claude"]); + assert_eq!(cli.effective_url(), DEFAULT_LOCAL_CXDB_URL); + assert!(cli.local); + assert!(matches!(cli.command, Command::Claude { .. })); + } + + #[test] + fn local_flag_conflicts_with_explicit_url() { + let result = + Cli::try_parse_from(["cxtx", "--local", "--url", "http://example.test", "codex"]); + assert!(result.is_err()); + } +} diff --git a/cxtx/src/lib.rs b/cxtx/src/lib.rs index 16d54a4..707afad 100644 --- a/cxtx/src/lib.rs +++ b/cxtx/src/lib.rs @@ -19,22 +19,33 @@ use tokio::process::Command; pub async fn run(cli: Cli) -> Result { let provider = cli.command.provider(); - let args = cli.command.args().to_vec(); + let effective_url = cli.effective_url(); let cxdb_url = cli - .url + .effective_url() .parse() - .with_context(|| format!("invalid CXDB URL: {}", cli.url))?; + .with_context(|| format!("invalid CXDB URL: {effective_url}"))?; let upstream = provider .resolve_upstream_base() .context("failed to resolve provider upstream base URL")?; + let (listener, proxy_base_url) = ProxyServer::bind(provider, &upstream) + .await + .context("failed to reserve local reverse proxy listener")?; + let args = provider.child_args_for_proxy(cli.command.args(), Some(&proxy_base_url)); let allowlisted_env = provider.capture_env_allowlist(); let session = SessionRuntime::new(provider, args.clone(), allowlisted_env)?; let ledger = SessionLedgerWriter::create(&session).await?; - let proxy = ProxyServer::start(provider, upstream, session.clone(), ledger.clone()) - .await - .context("failed to start local reverse proxy")?; + let proxy = ProxyServer::start_with_listener( + provider, + upstream, + session.clone(), + ledger.clone(), + listener, + proxy_base_url.clone(), + ) + .await + .context("failed to start local reverse proxy")?; let delivery = DeliveryHandle::start( cxdb_url, session.clone(), @@ -49,7 +60,10 @@ pub async fn run(cli: Cli) -> Result { command.stdin(Stdio::inherit()); command.stdout(Stdio::inherit()); command.stderr(Stdio::inherit()); - command.envs(provider.injected_env(&proxy.proxy_base_url())); + for name in provider.upstream_base_env_names() { + command.env_remove(name); + } + command.envs(provider.injected_env(&proxy_base_url)); let mut child = match command.spawn() { Ok(child) => child, diff --git a/cxtx/src/provider/anthropic.rs b/cxtx/src/provider/anthropic.rs index af40525..f36824c 100644 --- a/cxtx/src/provider/anthropic.rs +++ b/cxtx/src/provider/anthropic.rs @@ -355,7 +355,12 @@ fn parse_user_blocks(content: &Value) -> Result, String> { Value::Array(blocks) => { let mut history = Vec::new(); let mut text_buffer = String::new(); - for block in blocks { + let content_start = blocks + .iter() + .enumerate() + .find_map(|(index, block)| (!is_bootstrap_user_block(block)).then_some(index)) + .unwrap_or(blocks.len()); + for block in blocks.iter().skip(content_start) { match block.get("type").and_then(Value::as_str) { Some("text") => { let text = block @@ -406,6 +411,23 @@ fn parse_user_blocks(content: &Value) -> Result, String> { } } +fn is_bootstrap_user_block(block: &Value) -> bool { + if block.get("type").and_then(Value::as_str) != Some("text") { + return false; + } + + let text = block + .get("text") + .and_then(Value::as_str) + .unwrap_or_default() + .trim_start(); + + text.starts_with("") + && (text.contains("SessionStart hook additional context") + || text.contains("The following skills are available for use with the Skill tool:") + || text.contains("As you answer the user's questions, you can use the following context:")) +} + fn parse_assistant_content( content: &Value, model: Option<&str>, @@ -540,6 +562,22 @@ mod tests { }), expected_kinds: vec!["user_input", "assistant_turn"], }, + Case { + name: "leading bootstrap reminders are skipped before the real prompt", + payload: json!({ + "model": "claude-sonnet", + "messages": [ + {"role": "user", "content": [ + {"type": "text", "text": "\nSessionStart hook additional context: ..."}, + {"type": "text", "text": "\nThe following skills are available for use with the Skill tool:\n..."}, + {"type": "text", "text": "\nAs you answer the user's questions, you can use the following context:\n..."}, + {"type": "text", "text": "real prompt"} + ]}, + {"role": "assistant", "content": [{"type": "text", "text": "done"}]} + ] + }), + expected_kinds: vec!["user_input", "assistant_turn"], + }, ]; for case in cases { @@ -553,6 +591,12 @@ mod tests { }) .collect::>(); assert_eq!(kinds, case.expected_kinds, "case {}", case.name); + if case.name == "leading bootstrap reminders are skipped before the real prompt" { + match &history[0] { + HistoryItem::UserInput { text, .. } => assert_eq!(text, "real prompt"), + other => panic!("expected user input, got {other:?}"), + } + } } } diff --git a/cxtx/src/provider/mod.rs b/cxtx/src/provider/mod.rs index 1b86b7d..f573615 100644 --- a/cxtx/src/provider/mod.rs +++ b/cxtx/src/provider/mod.rs @@ -53,6 +53,40 @@ impl ProviderKind { } } + pub fn child_args(self, args: &[String]) -> Vec { + self.child_args_for_proxy(args, None) + } + + pub fn child_args_for_proxy( + self, + args: &[String], + _proxy_base_url: Option<&Url>, + ) -> Vec { + match self { + Self::Codex => { + let mut out = Vec::new(); + if !has_codex_config_override(args, "prefer_websockets") { + out.extend(["-c".to_string(), "prefer_websockets=false".to_string()]); + } + if !has_codex_feature_override(args, "responses_websockets") { + out.extend([ + "--disable".to_string(), + "responses_websockets".to_string(), + ]); + } + if !has_codex_feature_override(args, "responses_websockets_v2") { + out.extend([ + "--disable".to_string(), + "responses_websockets_v2".to_string(), + ]); + } + out.extend(args.iter().cloned()); + out + } + Self::Claude => args.to_vec(), + } + } + pub fn labels(self) -> Vec { vec![ "cxtx".to_string(), @@ -115,7 +149,9 @@ impl ProviderKind { match self { Self::Codex => vec![ ("OPENAI_BASE_URL".to_string(), value.clone()), - ("OPENAI_API_BASE".to_string(), value), + ("OPENAI_API_BASE".to_string(), value.clone()), + ("CXTX_OPENAI_BASE_URL".to_string(), value.clone()), + ("CXTX_OPENAI_API_BASE".to_string(), value), ], Self::Claude => { let root = proxy_base_url.origin().unicode_serialization(); @@ -216,7 +252,7 @@ impl ProviderKind { } } - fn upstream_base_env_names(self) -> &'static [&'static str] { + pub fn upstream_base_env_names(self) -> &'static [&'static str] { match self { Self::Codex => &["OPENAI_BASE_URL", "OPENAI_API_BASE"], Self::Claude => &[ @@ -340,6 +376,36 @@ fn find_model_field(value: &Value) -> Option { } } +fn has_codex_config_override(args: &[String], key: &str) -> bool { + args.iter().enumerate().any(|(index, arg)| { + if let Some(value) = arg.strip_prefix("--config=") { + return value.trim_start().starts_with(&format!("{key}=")); + } + if let Some(value) = arg.strip_prefix("-c") { + if !value.is_empty() { + return value.trim_start().starts_with(&format!("{key}=")); + } + } + matches!(arg.as_str(), "-c" | "--config") + && args + .get(index + 1) + .is_some_and(|value| value.trim_start().starts_with(&format!("{key}="))) + }) +} + +fn has_codex_feature_override(args: &[String], feature: &str) -> bool { + args.iter().enumerate().any(|(index, arg)| { + if let Some(value) = arg.strip_prefix("--enable=") { + return value == feature; + } + if let Some(value) = arg.strip_prefix("--disable=") { + return value == feature; + } + matches!(arg.as_str(), "--enable" | "--disable") + && args.get(index + 1).is_some_and(|value| value == feature) + }) || has_codex_config_override(args, &format!("features.{feature}")) +} + #[cfg(test)] mod tests { use super::ProviderKind; @@ -371,4 +437,90 @@ mod tests { "https://example.test/anthropic/v1/messages" ); } + + #[test] + fn codex_child_args_disable_websockets_by_default() { + let args = vec!["exec".to_string(), "say hi".to_string()]; + assert_eq!( + ProviderKind::Codex.child_args(&args), + vec![ + "-c".to_string(), + "prefer_websockets=false".to_string(), + "--disable".to_string(), + "responses_websockets".to_string(), + "--disable".to_string(), + "responses_websockets_v2".to_string(), + "exec".to_string(), + "say hi".to_string(), + ] + ); + } + + #[test] + fn codex_child_args_preserve_explicit_prefer_websockets_override() { + let args = vec![ + "--config".to_string(), + "prefer_websockets=true".to_string(), + "exec".to_string(), + ]; + assert_eq!( + ProviderKind::Codex.child_args(&args), + vec![ + "--disable".to_string(), + "responses_websockets".to_string(), + "--disable".to_string(), + "responses_websockets_v2".to_string(), + "--config".to_string(), + "prefer_websockets=true".to_string(), + "exec".to_string(), + ] + ); + } + + #[test] + fn codex_child_args_preserve_explicit_feature_overrides() { + let args = vec![ + "--disable".to_string(), + "responses_websockets".to_string(), + "--enable".to_string(), + "responses_websockets_v2".to_string(), + "exec".to_string(), + ]; + assert_eq!( + ProviderKind::Codex.child_args(&args), + vec![ + "-c".to_string(), + "prefer_websockets=false".to_string(), + "--disable".to_string(), + "responses_websockets".to_string(), + "--enable".to_string(), + "responses_websockets_v2".to_string(), + "exec".to_string(), + ] + ); + } + + #[test] + fn codex_child_args_include_same_result_when_proxy_is_known() { + let proxy = Url::parse("http://127.0.0.1:48123/v1").unwrap(); + let args = vec!["exec".to_string()]; + assert_eq!( + ProviderKind::Codex.child_args_for_proxy(&args, Some(&proxy)), + vec![ + "-c".to_string(), + "prefer_websockets=false".to_string(), + "--disable".to_string(), + "responses_websockets".to_string(), + "--disable".to_string(), + "responses_websockets_v2".to_string(), + "exec".to_string(), + ] + ); + } + + #[test] + fn claude_child_args_are_unchanged() { + let args = vec!["--print".to_string(), "stream".to_string()]; + assert_eq!(ProviderKind::Claude.child_args(&args), args); + } } diff --git a/cxtx/src/provider/openai.rs b/cxtx/src/provider/openai.rs index 4a4422b..142bf0c 100644 --- a/cxtx/src/provider/openai.rs +++ b/cxtx/src/provider/openai.rs @@ -349,8 +349,19 @@ fn parse_message_history( } fn parse_input_history(input: &[Value], model: Option) -> Result, String> { + let start_index = input + .iter() + .rposition(|item| item.get("role").and_then(Value::as_str) == Some("developer")) + .map(|index| index + 1) + .unwrap_or(0); + let conversation_start = input + .iter() + .enumerate() + .skip(start_index) + .find_map(|(index, item)| (!is_bootstrap_input_item(item)).then_some(index)) + .unwrap_or(input.len()); let mut history = Vec::new(); - for item in input { + for item in input.iter().skip(conversation_start) { if item .get("type") .and_then(Value::as_str) @@ -390,6 +401,23 @@ fn parse_input_history(input: &[Value], model: Option) -> Result bool { + if item.get("type").and_then(Value::as_str) != Some("message") { + return false; + } + if item.get("role").and_then(Value::as_str) != Some("user") { + return false; + } + + let text = content_to_text(item.get("content").unwrap_or(&Value::Null)); + let trimmed = text.trim_start(); + trimmed.starts_with("# AGENTS.md instructions for ") + || trimmed.starts_with("") + || trimmed.starts_with("") + || trimmed.starts_with("") + || trimmed.starts_with("") +} + fn parse_assistant_payload(payload: &Value, fallback_model: Option<&str>) -> Result, String> { if let Some(message) = payload .get("choices") @@ -708,6 +736,32 @@ mod tests { }), expected_kinds: vec!["user_input", "assistant_turn"], }, + Case { + name: "responses api input skips bootstrap conversation before latest developer block", + payload: json!({ + "model": "gpt-5.4", + "input": [ + {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": "sandbox"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "AGENTS bootstrap"}]}, + {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": "current mode"}]}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "previous answer"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "real prompt"}]} + ] + }), + expected_kinds: vec!["assistant_turn", "user_input"], + }, + Case { + name: "responses api input skips leading agents bootstrap user block after developer message", + payload: json!({ + "model": "gpt-5.4", + "input": [ + {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": ""}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "# AGENTS.md instructions for /repo\n\n /repo\n"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "real prompt"}]} + ] + }), + expected_kinds: vec!["user_input"], + }, ]; for case in cases { diff --git a/cxtx/src/proxy.rs b/cxtx/src/proxy.rs index ea251d1..fa429db 100644 --- a/cxtx/src/proxy.rs +++ b/cxtx/src/proxy.rs @@ -1,17 +1,22 @@ use anyhow::{anyhow, Context, Result}; use async_stream::stream; use axum::body::{to_bytes, Body}; -use axum::extract::State; +use axum::extract::ws::{Message as DownstreamWsMessage, WebSocket, WebSocketUpgrade}; +use axum::extract::{FromRequestParts, State}; use axum::http::{HeaderMap, Request, Response, StatusCode}; use axum::response::IntoResponse; use axum::routing::any; use axum::Router; -use futures_util::StreamExt; +use futures_util::{SinkExt, StreamExt}; use reqwest::header::{HeaderName, HeaderValue}; use serde_json::Value; use std::net::SocketAddr; use std::sync::Arc; use tokio::net::TcpListener; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::error::ProtocolError as TungsteniteProtocolError; +use tokio_tungstenite::tungstenite::protocol::Message as UpstreamWsMessage; +use tokio_tungstenite::tungstenite::Error as TungsteniteError; use tokio::sync::{mpsc, oneshot, RwLock}; use url::Url; @@ -39,23 +44,46 @@ pub struct ProxyServer { } impl ProxyServer { + pub async fn bind(provider: ProviderKind, upstream_base: &Url) -> Result<(TcpListener, Url)> { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .context("failed to bind proxy listener")?; + let addr = listener + .local_addr() + .context("missing proxy listener address")?; + let proxy_base_url = proxy_base_url(provider, upstream_base, addr)?; + Ok((listener, proxy_base_url)) + } + pub async fn start( provider: ProviderKind, upstream_base: Url, session: SessionRuntime, ledger: SessionLedgerWriter, + ) -> Result { + let (listener, proxy_base_url) = Self::bind(provider, &upstream_base).await?; + Self::start_with_listener( + provider, + upstream_base, + session, + ledger, + listener, + proxy_base_url, + ) + .await + } + + pub async fn start_with_listener( + provider: ProviderKind, + upstream_base: Url, + session: SessionRuntime, + ledger: SessionLedgerWriter, + listener: TcpListener, + proxy_base_url: Url, ) -> Result { let client = reqwest::Client::builder() .build() .context("failed to construct proxy reqwest client")?; - let listener = TcpListener::bind("127.0.0.1:0") - .await - .context("failed to bind proxy listener")?; - let addr = listener - .local_addr() - .context("missing proxy listener address")?; - let proxy_base_url = proxy_base_url(provider, &upstream_base, addr)?; - let state = ProxyState { provider, upstream_base, @@ -117,6 +145,10 @@ async fn proxy_handler(State(state): State, request: Request) } async fn handle_proxy_request(state: ProxyState, request: Request) -> Result> { + if is_websocket_upgrade_request(&request) { + return handle_websocket_proxy_request(state, request).await; + } + let delivery = state .delivery .read() @@ -226,6 +258,111 @@ async fn handle_proxy_request(state: ProxyState, request: Request) -> Resu } } +async fn handle_websocket_proxy_request( + state: ProxyState, + request: Request, +) -> Result> { + let maybe_delivery = state.delivery.read().await.clone(); + let (mut parts, body) = request.into_parts(); + let _ = body; + let ws = WebSocketUpgrade::from_request_parts(&mut parts, &()) + .await + .map_err(|err| anyhow!("invalid websocket upgrade request: {err}"))?; + let upstream_url = websocket_upstream_url( + &state + .provider + .build_upstream_url(&state.upstream_base, &parts.uri) + .context("failed to derive upstream websocket URL")?, + )?; + let request_headers = websocket_forwardable_headers(&parts.headers); + + let exchange_id = state.session.next_exchange_id(); + let request_artifact = state + .ledger + .record_request(&exchange_id, parts.uri.path(), None, None, &[], None) + .await?; + let artifact_refs = ArtifactRefs::default().with_request_path(Some(request_artifact)); + + let mut upstream_request = upstream_url + .as_str() + .into_client_request() + .context("failed to build upstream websocket request")?; + for (name, value) in request_headers { + upstream_request.headers_mut().insert(name, value); + } + + let (upstream_socket, upstream_response) = + match tokio_tungstenite::connect_async(upstream_request).await { + Ok(result) => result, + Err(err) => { + if let Some(delivery) = maybe_delivery.as_ref() { + delivery + .enqueue_turn(state.session.provider_error_turn( + &exchange_id, + "upstream_websocket_connect_error", + &format!("upstream websocket connect failed: {err}"), + None, + &artifact_refs, + )) + .await + .ok(); + } + return Ok(( + StatusCode::BAD_GATEWAY, + format!("upstream websocket connect failed: {err}"), + ) + .into_response()); + } + }; + + let status = StatusCode::from_u16(upstream_response.status().as_u16()) + .unwrap_or(StatusCode::SWITCHING_PROTOCOLS); + let request_id = state.provider.request_id_from_headers(upstream_response.headers()); + let response_artifact = state + .ledger + .record_response( + &exchange_id, + status.as_u16(), + request_id.as_deref(), + None, + &[], + None, + ) + .await?; + let artifact_refs = artifact_refs.with_response_path(Some(response_artifact)); + let selected_protocol = websocket_selected_protocol(upstream_response.headers()); + let ledger = state.ledger.clone(); + let session = state.session.clone(); + let exchange_id_for_upgrade = exchange_id.clone(); + let request_id_for_upgrade = request_id.clone(); + let artifact_refs_for_upgrade = artifact_refs.clone(); + + let upgrade = if let Some(protocol) = selected_protocol { + ws.protocols([protocol]) + } else { + ws + }; + + Ok(upgrade + .on_upgrade(move |downstream_socket| async move { + if let Err(err) = relay_websocket( + downstream_socket, + upstream_socket, + ledger, + exchange_id_for_upgrade, + maybe_delivery, + session, + request_id_for_upgrade, + artifact_refs_for_upgrade, + ) + .await + { + let _ = err; + } + }) + .into_response()) +} + async fn body_response( state: ProxyState, delivery: DeliveryHandle, @@ -404,6 +541,120 @@ async fn enqueue_turns(delivery: &DeliveryHandle, turns: Vec, + >, + ledger: SessionLedgerWriter, + exchange_id: String, + maybe_delivery: Option, + session: SessionRuntime, + request_id: Option, + artifact_refs: ArtifactRefs, +) -> Result<()> { + let (mut downstream_tx, mut downstream_rx) = downstream_socket.split(); + let (mut upstream_tx, mut upstream_rx) = upstream_socket.split(); + let mut capture = WebsocketCapture::new( + session.provider(), + exchange_id.clone(), + request_id.clone(), + artifact_refs.clone(), + ); + + loop { + tokio::select! { + downstream_message = downstream_rx.next() => { + let Some(downstream_message) = downstream_message else { + let _ = upstream_tx.close().await; + break; + }; + let downstream_message = match downstream_message { + Ok(message) => message, + Err(err) => { + if is_benign_downstream_websocket_read_error(&err) { + let _ = upstream_tx.close().await; + break; + } + return websocket_relay_error( + maybe_delivery, + session, + &exchange_id, + request_id.as_deref(), + &artifact_refs, + "downstream_websocket_error", + &format!("failed to read downstream websocket message: {err}"), + ) + .await; + } + }; + record_websocket_frame(&ledger, &exchange_id, "downstream", &downstream_message).await; + if let Some(delivery) = maybe_delivery.as_ref() { + let turns = capture.observe_downstream_message(&session, &downstream_message); + enqueue_turns(delivery, turns).await; + } + let is_close = matches!(downstream_message, DownstreamWsMessage::Close(_)); + if let Some(upstream_message) = map_downstream_message(downstream_message) { + upstream_tx + .send(upstream_message) + .await + .context("failed to forward websocket message upstream")?; + } + if is_close { + break; + } + } + upstream_message = upstream_rx.next() => { + let Some(upstream_message) = upstream_message else { + let _ = downstream_tx.close().await; + break; + }; + let upstream_message = match upstream_message { + Ok(message) => message, + Err(err) => { + if is_benign_websocket_read_error(&err) { + let _ = downstream_tx.close().await; + break; + } + return websocket_relay_error( + maybe_delivery, + session, + &exchange_id, + request_id.as_deref(), + &artifact_refs, + "upstream_websocket_error", + &format!("failed to read upstream websocket message: {err}"), + ) + .await; + } + }; + record_upstream_websocket_frame(&ledger, &exchange_id, "upstream", &upstream_message).await; + if let Some(delivery) = maybe_delivery.as_ref() { + let turns = capture.observe_upstream_message(&session, &upstream_message); + enqueue_turns(delivery, turns).await; + } + let is_close = matches!(upstream_message, UpstreamWsMessage::Close(_)); + if let Some(downstream_message) = map_upstream_message(upstream_message) { + downstream_tx + .send(downstream_message) + .await + .context("failed to forward websocket message downstream")?; + } + if is_close { + break; + } + } + } + } + + if let Some(delivery) = maybe_delivery.as_ref() { + let turns = capture.finalize_pending(&session); + enqueue_turns(delivery, turns).await; + } + + Ok(()) +} + fn proxy_base_url(provider: ProviderKind, upstream_base: &Url, addr: SocketAddr) -> Result { let mut url = Url::parse(&format!("http://{addr}")).context("failed to build proxy URL")?; let mount_path = provider.proxy_mount_path(upstream_base); @@ -415,7 +666,30 @@ fn forwardable_headers(headers: &HeaderMap) -> Vec<(HeaderName, HeaderValue)> { headers .iter() .filter_map(|(name, value)| { - if is_hop_by_hop(name.as_str()) || name.as_str().eq_ignore_ascii_case("host") { + if is_hop_by_hop(name.as_str()) + || name.as_str().eq_ignore_ascii_case("host") + || name.as_str().eq_ignore_ascii_case("accept-encoding") + { + None + } else { + Some((name.clone(), value.clone())) + } + }) + .collect() +} + +fn websocket_forwardable_headers(headers: &HeaderMap) -> Vec<(HeaderName, HeaderValue)> { + headers + .iter() + .filter_map(|(name, value)| { + let lower = name.as_str().to_ascii_lowercase(); + if is_hop_by_hop(&lower) + || lower == "host" + || matches!( + lower.as_str(), + "sec-websocket-key" | "sec-websocket-version" | "sec-websocket-extensions" + ) + { None } else { Some((name.clone(), value.clone())) @@ -453,6 +727,284 @@ fn header_value_reqwest(headers: &reqwest::header::HeaderMap, name: &str) -> Opt .map(|value| value.to_string()) } +fn header_contains_token(headers: &HeaderMap, name: &str, token: &str) -> bool { + headers + .get(name) + .and_then(|value| value.to_str().ok()) + .map(|value| { + value + .split(',') + .any(|part| part.trim().eq_ignore_ascii_case(token)) + }) + .unwrap_or(false) +} + +fn is_websocket_upgrade_request(request: &Request) -> bool { + request.method() == http::Method::GET + && header_contains_token(request.headers(), "connection", "upgrade") + && request + .headers() + .get("upgrade") + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.eq_ignore_ascii_case("websocket")) +} + +fn websocket_upstream_url(url: &Url) -> Result { + let mut url = url.clone(); + match url.scheme() { + "http" => url + .set_scheme("ws") + .map_err(|_| anyhow!("failed to switch upstream scheme from http to ws"))?, + "https" => url + .set_scheme("wss") + .map_err(|_| anyhow!("failed to switch upstream scheme from https to wss"))?, + "ws" | "wss" => {} + scheme => return Err(anyhow!("unsupported websocket upstream scheme: {scheme}")), + } + Ok(url) +} + +fn websocket_selected_protocol(headers: &reqwest::header::HeaderMap) -> Option { + headers + .get("sec-websocket-protocol") + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.split(',').next()) + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +async fn record_websocket_frame( + ledger: &SessionLedgerWriter, + exchange_id: &str, + direction: &str, + message: &DownstreamWsMessage, +) { + if let Some(frame) = summarize_downstream_websocket_frame(direction, message) { + let _ = ledger.append_stream_frame(exchange_id, &frame).await; + } +} + +async fn record_upstream_websocket_frame( + ledger: &SessionLedgerWriter, + exchange_id: &str, + direction: &str, + message: &UpstreamWsMessage, +) { + if let Some(frame) = summarize_upstream_websocket_frame(direction, message) { + let _ = ledger.append_stream_frame(exchange_id, &frame).await; + } +} + +fn summarize_downstream_websocket_frame( + direction: &str, + message: &DownstreamWsMessage, +) -> Option { + match message { + DownstreamWsMessage::Text(text) => Some(format!("{direction}:text:{text}")), + DownstreamWsMessage::Binary(bytes) => { + Some(format!("{direction}:binary:{} bytes", bytes.len())) + } + DownstreamWsMessage::Ping(bytes) => Some(format!("{direction}:ping:{} bytes", bytes.len())), + DownstreamWsMessage::Pong(bytes) => Some(format!("{direction}:pong:{} bytes", bytes.len())), + DownstreamWsMessage::Close(_) => Some(format!("{direction}:close")), + } +} + +fn summarize_upstream_websocket_frame( + direction: &str, + message: &UpstreamWsMessage, +) -> Option { + match message { + UpstreamWsMessage::Text(text) => Some(format!("{direction}:text:{text}")), + UpstreamWsMessage::Binary(bytes) => { + Some(format!("{direction}:binary:{} bytes", bytes.len())) + } + UpstreamWsMessage::Ping(bytes) => Some(format!("{direction}:ping:{} bytes", bytes.len())), + UpstreamWsMessage::Pong(bytes) => Some(format!("{direction}:pong:{} bytes", bytes.len())), + UpstreamWsMessage::Close(_) => Some(format!("{direction}:close")), + UpstreamWsMessage::Frame(_) => None, + } +} + +fn map_downstream_message(message: DownstreamWsMessage) -> Option { + match message { + DownstreamWsMessage::Text(text) => Some(UpstreamWsMessage::Text(text.to_string())), + DownstreamWsMessage::Binary(bytes) => Some(UpstreamWsMessage::Binary(bytes.to_vec())), + DownstreamWsMessage::Ping(bytes) => Some(UpstreamWsMessage::Ping(bytes.to_vec())), + DownstreamWsMessage::Pong(bytes) => Some(UpstreamWsMessage::Pong(bytes.to_vec())), + DownstreamWsMessage::Close(_) => Some(UpstreamWsMessage::Close(None)), + } +} + +fn map_upstream_message(message: UpstreamWsMessage) -> Option { + match message { + UpstreamWsMessage::Text(text) => Some(DownstreamWsMessage::Text(text.to_string().into())), + UpstreamWsMessage::Binary(bytes) => { + Some(DownstreamWsMessage::Binary(bytes.to_vec().into())) + } + UpstreamWsMessage::Ping(bytes) => Some(DownstreamWsMessage::Ping(bytes.to_vec().into())), + UpstreamWsMessage::Pong(bytes) => Some(DownstreamWsMessage::Pong(bytes.to_vec().into())), + UpstreamWsMessage::Close(_) => Some(DownstreamWsMessage::Close(None)), + UpstreamWsMessage::Frame(_) => None, + } +} + +#[derive(Debug)] +struct WebsocketCapture { + provider: ProviderKind, + exchange_id: String, + request_id: Option, + artifact_refs: ArtifactRefs, + current_state: Option, +} + +impl WebsocketCapture { + fn new( + provider: ProviderKind, + exchange_id: String, + request_id: Option, + artifact_refs: ArtifactRefs, + ) -> Self { + Self { + provider, + exchange_id, + request_id, + artifact_refs, + current_state: None, + } + } + + fn observe_downstream_message( + &mut self, + session: &SessionRuntime, + message: &DownstreamWsMessage, + ) -> Vec { + let DownstreamWsMessage::Text(text) = message else { + return Vec::new(); + }; + self.observe_downstream_text(session, text.as_str()) + } + + fn observe_upstream_message( + &mut self, + session: &SessionRuntime, + message: &UpstreamWsMessage, + ) -> Vec { + let UpstreamWsMessage::Text(text) = message else { + return Vec::new(); + }; + self.observe_upstream_text(session, text.as_str()) + } + + fn observe_downstream_text( + &mut self, + session: &SessionRuntime, + text: &str, + ) -> Vec { + if self.provider != ProviderKind::Codex { + return Vec::new(); + } + let Ok(payload) = serde_json::from_str::(text) else { + return Vec::new(); + }; + if payload.get("type").and_then(Value::as_str) != Some("response.create") { + return Vec::new(); + } + + let mut turns = self.finalize_pending(session); + let prepared = self.provider.prepare_exchange( + session, + self.exchange_id.clone(), + text.as_bytes(), + &self.artifact_refs, + ); + turns.extend(prepared.request_turns); + self.current_state = Some(prepared.state); + turns + } + + fn observe_upstream_text( + &mut self, + session: &SessionRuntime, + text: &str, + ) -> Vec { + let Some(state) = self.current_state.as_mut() else { + return Vec::new(); + }; + let Ok(payload) = serde_json::from_str::(text) else { + return Vec::new(); + }; + let event_type = payload + .get("type") + .and_then(Value::as_str) + .map(|value| value.to_string()); + state.absorb_sse_frame(&openai::SseFrame { + event: event_type.clone(), + data: text.to_string(), + raw: text.to_string(), + }); + if event_type.as_deref() == Some("response.completed") { + return self.finalize_pending(session); + } + Vec::new() + } + + fn finalize_pending(&mut self, session: &SessionRuntime) -> Vec { + let Some(state) = self.current_state.take() else { + return Vec::new(); + }; + state.finalize_stream( + session, + 200, + self.request_id.clone(), + &self.artifact_refs, + None, + ) + } +} + +fn is_benign_websocket_read_error(err: &TungsteniteError) -> bool { + matches!( + err, + TungsteniteError::ConnectionClosed + | TungsteniteError::AlreadyClosed + | TungsteniteError::Protocol(TungsteniteProtocolError::ResetWithoutClosingHandshake) + | TungsteniteError::Protocol(TungsteniteProtocolError::HandshakeIncomplete) + ) +} + +fn is_benign_downstream_websocket_read_error(err: &axum::Error) -> bool { + let message = err.to_string(); + message.contains("Connection closed normally") + || message.contains("Trying to work with closed connection") + || message.contains("Connection reset without closing handshake") + || message.contains("Handshake not finished") +} + +async fn websocket_relay_error( + maybe_delivery: Option, + session: SessionRuntime, + exchange_id: &str, + request_id: Option<&str>, + artifact_refs: &ArtifactRefs, + title: &str, + message: &str, +) -> Result<()> { + if let Some(delivery) = maybe_delivery { + delivery + .enqueue_turn(session.provider_error_turn( + exchange_id, + title, + message, + request_id, + artifact_refs, + )) + .await + .ok(); + } + Err(anyhow!(message.to_string())) +} + fn is_hop_by_hop(name: &str) -> bool { matches!( name.to_ascii_lowercase().as_str(), @@ -469,14 +1021,35 @@ fn is_hop_by_hop(name: &str) -> bool { #[cfg(test)] mod tests { - use super::forwardable_headers; + use super::{ + forwardable_headers, is_benign_websocket_read_error, is_websocket_upgrade_request, + websocket_forwardable_headers, websocket_upstream_url, ProxyServer, WebsocketCapture, + }; + use axum::body::Body; use axum::http::{HeaderMap, HeaderValue}; + use futures_util::{SinkExt, StreamExt}; + use serde_json::Value; + use std::collections::BTreeMap; + use std::sync::{Arc, Mutex}; + use tokio::net::TcpListener; + use tokio_tungstenite::tungstenite::error::ProtocolError as TungsteniteProtocolError; + use tokio_tungstenite::tungstenite::handshake::server::{ + Request as WsRequest, Response as WsResponse, + }; + use tokio_tungstenite::tungstenite::Error as TungsteniteError; + use tokio_tungstenite::tungstenite::Message as WsMessage; + use url::Url; + + use crate::ledger::SessionLedgerWriter; + use crate::provider::ProviderKind; + use crate::session::SessionRuntime; #[test] fn forwardable_headers_drop_host_but_keep_authorization() { let mut headers = HeaderMap::new(); headers.insert("host", HeaderValue::from_static("127.0.0.1:12345")); headers.insert("authorization", HeaderValue::from_static("Bearer test")); + headers.insert("accept-encoding", HeaderValue::from_static("gzip, br")); let forwarded = forwardable_headers(&headers); assert!( @@ -484,10 +1057,206 @@ mod tests { .iter() .any(|(name, _)| name.as_str().eq_ignore_ascii_case("host")) ); + assert!( + !forwarded + .iter() + .any(|(name, _)| name.as_str().eq_ignore_ascii_case("accept-encoding")) + ); + assert!( + forwarded + .iter() + .any(|(name, value)| name == "authorization" && value == "Bearer test") + ); + } + + #[test] + fn websocket_forwardable_headers_drop_client_handshake_headers() { + let mut headers = HeaderMap::new(); + headers.insert("authorization", HeaderValue::from_static("Bearer test")); + headers.insert("sec-websocket-key", HeaderValue::from_static("abc")); + headers.insert("sec-websocket-version", HeaderValue::from_static("13")); + + let forwarded = websocket_forwardable_headers(&headers); assert!( forwarded .iter() .any(|(name, value)| name == "authorization" && value == "Bearer test") ); + assert!( + !forwarded + .iter() + .any(|(name, _)| name.as_str().eq_ignore_ascii_case("sec-websocket-key")) + ); + } + + #[test] + fn websocket_upgrade_detection_matches_standard_headers() { + let request = http::Request::builder() + .method(http::Method::GET) + .uri("/v1/responses") + .header("connection", "keep-alive, Upgrade") + .header("upgrade", "websocket") + .body(Body::empty()) + .unwrap(); + assert!(is_websocket_upgrade_request(&request)); + } + + #[test] + fn websocket_upstream_url_switches_http_scheme() { + let upstream = Url::parse("https://api.openai.com/v1/responses").unwrap(); + assert_eq!( + websocket_upstream_url(&upstream).unwrap().as_str(), + "wss://api.openai.com/v1/responses" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn websocket_upgrade_requests_are_relayed_to_upstream() { + let upstream_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let upstream_addr = upstream_listener.local_addr().unwrap(); + let seen = Arc::new(Mutex::new(None::)); + let seen_for_server = Arc::clone(&seen); + let upstream = tokio::spawn(async move { + let (socket, _) = upstream_listener.accept().await.unwrap(); + let callback = move |request: &WsRequest, response: WsResponse| { + *seen_for_server.lock().unwrap() = Some(request.uri().path().to_string()); + Ok(response) + }; + let mut socket = tokio_tungstenite::accept_hdr_async(socket, callback) + .await + .unwrap(); + let message = socket.next().await.unwrap().unwrap(); + socket.send(message).await.unwrap(); + }); + + let session = + SessionRuntime::new(ProviderKind::Codex, Vec::new(), BTreeMap::new()).unwrap(); + let ledger = SessionLedgerWriter::create(&session).await.unwrap(); + let proxy = ProxyServer::start( + ProviderKind::Codex, + Url::parse(&format!("ws://{upstream_addr}/v1")).unwrap(), + session, + ledger.clone(), + ) + .await + .unwrap(); + + let mut proxy_url = proxy.proxy_base_url(); + proxy_url.set_scheme("ws").unwrap(); + proxy_url.set_path("/v1/responses"); + + let (mut socket, _) = tokio_tungstenite::connect_async(proxy_url.as_str()) + .await + .unwrap(); + socket + .send(WsMessage::Text("hello websocket".to_string())) + .await + .unwrap(); + let echoed = socket.next().await.unwrap().unwrap(); + assert_eq!(echoed.into_text().unwrap(), "hello websocket"); + socket.close(None).await.unwrap(); + + upstream.await.unwrap(); + proxy.shutdown().await.unwrap(); + + assert_eq!( + seen.lock().unwrap().clone(), + Some("/v1/responses".to_string()) + ); + + let ledger_json: Value = + serde_json::from_str(&tokio::fs::read_to_string(ledger.path()).await.unwrap()).unwrap(); + assert_eq!(ledger_json["exchanges"][0]["status_code"], 101); + assert_eq!(ledger_json["exchanges"][0]["endpoint"], "/v1/responses"); + } + + #[test] + fn websocket_capture_turns_real_prompt_into_history_and_answer() { + let session = + SessionRuntime::new(ProviderKind::Codex, Vec::new(), BTreeMap::new()).unwrap(); + let mut capture = WebsocketCapture::new( + ProviderKind::Codex, + "exchange-0001".to_string(), + Some("req_123".to_string()), + crate::turns::ArtifactRefs::default(), + ); + + let bootstrap_turns = capture.observe_downstream_text( + &session, + r#"{"type":"response.create","model":"gpt-5.4","instructions":"bootstrap","input":[]}"#, + ); + assert!(bootstrap_turns.is_empty()); + let bootstrap_answer = capture.observe_upstream_text( + &session, + r#"{"type":"response.completed","response":{"model":"gpt-5.4","status":"completed","output":[]}}"#, + ); + assert!(bootstrap_answer.is_empty()); + + let request_turns = capture.observe_downstream_text( + &session, + r#"{ + "type":"response.create", + "model":"gpt-5.4", + "input":[ + {"type":"message","role":"developer","content":[{"type":"input_text","text":"mode"}]}, + {"type":"message","role":"user","content":[{"type":"input_text","text":"yooo dawg"}]} + ] + }"#, + ); + assert_eq!(request_turns.len(), 1); + assert_eq!(request_turns[0].item.item_type, "user_input"); + assert_eq!( + request_turns[0].item.user_input.as_ref().unwrap().text, + "yooo dawg" + ); + + assert!(capture + .observe_upstream_text( + &session, + r#"{"type":"response.output_text.delta","delta":"yoo"}"#, + ) + .is_empty()); + assert!(capture + .observe_upstream_text( + &session, + r#"{"type":"response.output_text.delta","delta":", what's up?"}"#, + ) + .is_empty()); + let answer_turns = capture.observe_upstream_text( + &session, + r#"{ + "type":"response.completed", + "response":{ + "model":"gpt-5.4", + "status":"completed", + "output":[ + { + "type":"message", + "role":"assistant", + "content":[{"type":"output_text","text":"yoo, what's up?"}] + } + ] + } + }"#, + ); + assert_eq!(answer_turns.len(), 1); + assert_eq!(answer_turns[0].item.item_type, "assistant_turn"); + assert_eq!( + answer_turns[0].item.turn.as_ref().unwrap().text, + "yoo, what's up?" + ); + } + + #[test] + fn websocket_reset_without_close_is_treated_as_benign() { + assert!(is_benign_websocket_read_error(&TungsteniteError::Protocol( + TungsteniteProtocolError::ResetWithoutClosingHandshake + ))); + assert!(is_benign_websocket_read_error(&TungsteniteError::Protocol( + TungsteniteProtocolError::HandshakeIncomplete + ))); + assert!(is_benign_websocket_read_error( + &TungsteniteError::ConnectionClosed + )); } } diff --git a/cxtx/tests/integration.rs b/cxtx/tests/integration.rs index a0b9872..aa56603 100644 --- a/cxtx/tests/integration.rs +++ b/cxtx/tests/integration.rs @@ -19,12 +19,17 @@ use cxdb_server::http::start_http; use cxdb_server::metrics::{Metrics, SessionTracker}; use cxdb_server::registry::Registry; use cxdb_server::store::Store; +use futures_util::{SinkExt, StreamExt}; use predicates::prelude::*; use reqwest::Client; use serde_json::{json, Value}; use tempfile::TempDir; use tokio::net::TcpListener as TokioTcpListener; use tokio::sync::oneshot; +use tokio_tungstenite::tungstenite::handshake::server::{ + Request as WsRequest, Response as WsResponse, +}; +use tokio_tungstenite::tungstenite::Message as WsMessage; use cxtx::cxdb_http::CxdbHttpClient; use cxtx::delivery::DeliveryHandle; @@ -114,8 +119,10 @@ async fn codex_wrapper_preserves_child_io_and_uploads_canonical_turns() { &format!( r#"#!/bin/sh set -eu -printf '%s\n' "$OPENAI_BASE_URL" > "{fixture}/openai_base_url.txt" -printf '%s\n' "$OPENAI_API_BASE" > "{fixture}/openai_api_base.txt" +printf '%s\n' "$OPENAI_BASE_URL" > "{fixture}/openai_base_url_env.txt" +printf '%s\n' "$OPENAI_API_BASE" > "{fixture}/openai_api_base_env.txt" +printf '%s\n' "$CXTX_OPENAI_BASE_URL" > "{fixture}/openai_base_url.txt" +printf '%s\n' "$CXTX_OPENAI_API_BASE" > "{fixture}/openai_api_base.txt" printf '%s\n' "$@" > "{fixture}/args.txt" python3 - <<'PY' import json @@ -123,7 +130,7 @@ import os import sys import urllib.request -base = os.environ["OPENAI_BASE_URL"].rstrip("/") +base = os.environ["CXTX_OPENAI_BASE_URL"].rstrip("/") req = urllib.request.Request( base + "/chat/completions", data=json.dumps({{ @@ -166,10 +173,19 @@ printf 'codex-child-stderr\n' >&2 assert!(fs::read_to_string(fixture_dir.join("openai_base_url.txt")) .unwrap() .contains("/v1")); + assert!(fs::read_to_string(fixture_dir.join("openai_base_url_env.txt")) + .unwrap() + .contains("/v1")); assert_eq!( - fs::read_to_string(fixture_dir.join("args.txt")).unwrap(), - "--model\ngpt-5\n" + fs::read_to_string(fixture_dir.join("openai_api_base_env.txt")).unwrap(), + fs::read_to_string(fixture_dir.join("openai_api_base.txt")).unwrap() ); + let args = fs::read_to_string(fixture_dir.join("args.txt")).unwrap(); + assert!(!args.contains("openai_base_url=")); + assert!(args.starts_with("-c\nprefer_websockets=false\n")); + assert!(args.contains("--disable\nresponses_websockets\n")); + assert!(args.contains("--disable\nresponses_websockets_v2\n")); + assert!(args.ends_with("--model\ngpt-5\n")); let contexts = cxdb.list_contexts().await.unwrap(); let context_id = first_context_id(&contexts); @@ -216,7 +232,7 @@ import json import os import urllib.request -base = os.environ["OPENAI_BASE_URL"].rstrip("/") +base = os.environ["CXTX_OPENAI_BASE_URL"].rstrip("/") headers = { "Content-Type": "application/json", "Authorization": "Bearer test-openai", @@ -292,7 +308,7 @@ import json import os import urllib.request -base = os.environ["OPENAI_BASE_URL"].rstrip("/") +base = os.environ["CXTX_OPENAI_BASE_URL"].rstrip("/") headers = { "Content-Type": "application/json", "Authorization": "Bearer test-openai", @@ -372,7 +388,7 @@ import json import os import urllib.request -base = os.environ["OPENAI_BASE_URL"].rstrip("/") +base = os.environ["CXTX_OPENAI_BASE_URL"].rstrip("/") headers = { "Content-Type": "application/json", "Authorization": "Bearer test-openai", @@ -433,6 +449,164 @@ PY assert_eq!(turns[4]["data"]["turn"]["text"], "tool complete"); } +#[tokio::test(flavor = "multi_thread")] +async fn codex_responses_bootstrap_history_starts_with_real_prompt() { + let scratch = ScratchRoot::new().unwrap(); + let cxdb = TestCxdb::start().await.unwrap(); + let upstream = MockOpenAi::start().await.unwrap(); + let fake_bin_dir = scratch.dir.path().join("bin"); + fs::create_dir_all(&fake_bin_dir).unwrap(); + write_executable( + &fake_bin_dir.join("codex"), + r##"#!/bin/sh +set -eu +python3 - <<'PY' +import json +import os +import sys +import urllib.request + +base = os.environ["CXTX_OPENAI_BASE_URL"].rstrip("/") +req = urllib.request.Request( + base + "/responses", + data=json.dumps({ + "model": "gpt-5.4", + "input": [ + {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": ""}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "# AGENTS.md instructions for /repo\n\n /repo\n"}]}, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "previous answer"}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "real prompt"}]} + ] + }).encode(), + headers={ + "Content-Type": "application/json", + "Authorization": "Bearer test-openai" + }, +) +with urllib.request.urlopen(req) as resp: + sys.stdout.write(resp.read().decode()) +PY +"##, + ) + .unwrap(); + + let mut command = std::process::Command::cargo_bin("cxtx").unwrap(); + command + .current_dir(scratch.dir.path()) + .env("PATH", prepend_path(&fake_bin_dir)) + .env("OPENAI_BASE_URL", upstream.base_url.as_str()) + .arg("--url") + .arg(&cxdb.base_url) + .arg("codex") + .arg("--"); + command.assert().success(); + + let context_id = first_context_id(&cxdb.list_contexts().await.unwrap()); + let turns = cxdb.turns(context_id).await.unwrap(); + let item_types = turn_item_types(&turns); + assert_eq!( + item_types, + vec!["system", "assistant_turn", "user_input", "assistant_turn", "system"] + ); + assert_eq!(turns[1]["data"]["turn"]["text"], "previous answer"); + assert_eq!(turns[2]["data"]["user_input"]["text"], "real prompt"); + assert_eq!(turns[3]["data"]["turn"]["text"], "clean answer"); + + let recorded_requests = upstream.requests.lock().unwrap().clone(); + assert_eq!(recorded_requests.len(), 1); + assert_eq!(recorded_requests[0]["path"], "/v1/responses"); +} + +#[tokio::test(flavor = "multi_thread")] +async fn websocket_proxy_uploads_canonical_turns_into_cxdb() { + let scratch = ScratchRoot::new().unwrap(); + let cxdb = TestCxdb::start().await.unwrap(); + let upstream = MockOpenAiWebsocket::start().await.unwrap(); + let session = SessionRuntime::new( + ProviderKind::Codex, + vec!["interactive".to_string()], + BTreeMap::new(), + ) + .unwrap(); + let ledger = SessionLedgerWriter::create(&session).await.unwrap(); + let proxy = cxtx::proxy::ProxyServer::start( + ProviderKind::Codex, + upstream.base_url.parse().unwrap(), + session.clone(), + ledger.clone(), + ) + .await + .unwrap(); + let delivery = DeliveryHandle::start( + cxdb.base_url.parse().unwrap(), + session.clone(), + ledger.clone(), + "cxtx-tests".to_string(), + ) + .await + .unwrap(); + proxy.set_delivery(delivery.clone()).await; + delivery.enqueue_create_context().await.unwrap(); + delivery.enqueue_turn(session.session_start_turn()).await.unwrap(); + + let mut proxy_url = proxy.proxy_base_url(); + proxy_url.set_scheme("ws").unwrap(); + proxy_url.set_path("/v1/responses"); + + let (mut socket, _) = tokio_tungstenite::connect_async(proxy_url.as_str()) + .await + .unwrap(); + socket + .send(WsMessage::Text( + r#"{ + "type":"response.create", + "model":"gpt-5.4", + "input":[ + {"type":"message","role":"developer","content":[{"type":"input_text","text":"mode"}]}, + {"type":"message","role":"user","content":[{"type":"input_text","text":"websocket hello"}]} + ] + }"# + .to_string(), + )) + .await + .unwrap(); + + while let Some(message) = socket.next().await { + let message = message.unwrap(); + match message { + WsMessage::Text(text) if text.contains("\"response.completed\"") => break, + WsMessage::Close(_) => break, + _ => {} + } + } + let _ = socket.close(None).await; + + delivery + .enqueue_turn(session.session_end_turn(0, true)) + .await + .unwrap(); + proxy.shutdown().await.unwrap(); + delivery.shutdown().await.unwrap(); + ledger.finalize().await.unwrap(); + + let context_id = first_context_id(&cxdb.list_contexts().await.unwrap()); + let turns = cxdb.turns(context_id).await.unwrap(); + let item_types = turn_item_types(&turns); + assert_eq!( + item_types, + vec!["system", "user_input", "assistant_turn", "system"] + ); + assert_eq!(turns[1]["data"]["user_input"]["text"], "websocket hello"); + assert_eq!(turns[2]["data"]["turn"]["text"], "hello from websocket"); + + let ledger_json = find_single_ledger(scratch.dir.path()); + assert_eq!(ledger_json["exchanges"][0]["status_code"], 101); + assert!(ledger_json["exchanges"][0]["stream_path"] + .as_str() + .unwrap() + .ends_with("stream.ndjson")); +} + #[tokio::test(flavor = "multi_thread")] async fn claude_wrapper_streams_to_child_and_uploads_canonical_turns() { let scratch = ScratchRoot::new().unwrap(); @@ -787,7 +961,7 @@ import sys import urllib.error import urllib.request -base = os.environ["OPENAI_BASE_URL"].rstrip("/") +base = os.environ["CXTX_OPENAI_BASE_URL"].rstrip("/") req = urllib.request.Request( base + "/chat/completions", data=json.dumps({ @@ -848,7 +1022,7 @@ import os import sys import urllib.request -base = os.environ["OPENAI_BASE_URL"].rstrip("/") +base = os.environ["CXTX_OPENAI_BASE_URL"].rstrip("/") req = urllib.request.Request( base + "/chat/completions", data=json.dumps({ @@ -1320,6 +1494,7 @@ impl MockOpenAi { let (shutdown_tx, shutdown_rx) = oneshot::channel(); let app = Router::new() .route("/v1/chat/completions", post(mock_openai_handler)) + .route("/v1/responses", post(mock_openai_responses_handler)) .with_state(requests.clone()); tokio::spawn(async move { let _ = axum::serve(listener, app) @@ -1336,6 +1511,64 @@ impl MockOpenAi { } } +struct MockOpenAiWebsocket { + base_url: String, + _shutdown: oneshot::Sender<()>, +} + +impl MockOpenAiWebsocket { + async fn start() -> anyhow::Result { + let listener = TokioTcpListener::bind("127.0.0.1:0").await?; + let addr = listener.local_addr()?; + let (shutdown_tx, mut shutdown_rx) = oneshot::channel(); + tokio::spawn(async move { + tokio::select! { + _ = async { + let (socket, _) = listener.accept().await.unwrap(); + let callback = |_: &WsRequest, response: WsResponse| Ok(response); + let mut socket = tokio_tungstenite::accept_hdr_async(socket, callback) + .await + .unwrap(); + let _request = socket.next().await.unwrap().unwrap(); + socket + .send(WsMessage::Text( + r#"{"type":"response.output_text.delta","delta":"hello from websocket"}"# + .to_string(), + )) + .await + .unwrap(); + socket + .send(WsMessage::Text( + r#"{ + "type":"response.completed", + "response":{ + "model":"gpt-5.4", + "status":"completed", + "output":[ + { + "type":"message", + "role":"assistant", + "content":[{"type":"output_text","text":"hello from websocket"}] + } + ] + } + }"# + .to_string(), + )) + .await + .unwrap(); + let _ = socket.close(None).await; + } => {} + _ = &mut shutdown_rx => {} + } + }); + Ok(Self { + base_url: format!("ws://{addr}/v1"), + _shutdown: shutdown_tx, + }) + } +} + struct MockOpenAiMalformed { base_url: String, _shutdown: oneshot::Sender<()>, @@ -1606,6 +1839,46 @@ async fn mock_openai_handler( ) } +async fn mock_openai_responses_handler( + State(requests): State>>>, + request: axum::http::Request, +) -> impl IntoResponse { + let (parts, body) = request.into_parts(); + let body = axum::body::to_bytes(body, usize::MAX).await.unwrap(); + let json_body: Value = serde_json::from_slice(&body).unwrap(); + requests.lock().unwrap().push(json!({ + "path": parts.uri.path(), + "body": json_body.clone(), + })); + + let response_text = match openai_input_last_user_text(&json_body) { + Some("real prompt") => "clean answer", + Some("websocket hello") => "hello from websocket", + _ => "hi", + }; + + ( + StatusCode::OK, + [ + ("content-type", "application/json"), + ("x-request-id", "req_openai_responses_123"), + ], + Body::from( + json!({ + "id": "resp_123", + "model": "gpt-5.4", + "status": "completed", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": response_text}] + }] + }) + .to_string(), + ), + ) +} + async fn mock_openai_malformed_handler() -> impl IntoResponse { ( StatusCode::OK, @@ -1879,6 +2152,20 @@ fn turn_item_types(turns: &[Value]) -> Vec<&str> { .collect() } +fn openai_input_last_user_text(payload: &Value) -> Option<&str> { + payload["input"] + .as_array() + .and_then(|items| { + items.iter().rev().find_map(|item| { + (item["role"] == "user") + .then_some(item["content"].as_array()) + .flatten() + .and_then(|content| content.last()) + .and_then(|part| part["text"].as_str()) + }) + }) +} + fn first_context_id(contexts: &Value) -> u64 { json_u64(&contexts["contexts"][0]["context_id"]).unwrap() }