diff --git a/CHANGELOG.md b/CHANGELOG.md index aec5f68d..c3d04692 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,17 @@ title: Changelog description: Release notes for claude-code-proxy. --- +## Unreleased + +- Codex WebSocket connection pacing is adaptive: connections start without + spacing while the origin accepts upgrades, widen after a rejected upgrade, + and relax again after sustained success. A fixed 1s spacing previously + capped each process near one generation per second. Set a floor with + `CCP_CODEX_WS_CONNECT_SPACING_MS` to restore fixed pacing. +- Codex standalone searches keep stable per-Agent sessions without colliding + with sibling Agents that share a Claude Code session, and align all upstream + search identity headers with the request body owner. + ## v0.1.35 (2026-08-19) - Grok web search works reliably with Claude Code, preserves other tools, and diff --git a/docs/src/content/docs/providers/codex.md b/docs/src/content/docs/providers/codex.md index a0061c55..9c76df90 100644 --- a/docs/src/content/docs/providers/codex.md +++ b/docs/src/content/docs/providers/codex.md @@ -49,6 +49,13 @@ Claude Code summary compaction requests are capped at low effort by default beca Structured result DTOs map back to Anthropic `server_tool_use` and `web_search_tool_result` blocks, while standalone text output remains text. The proxy locally estimates input and output tokens and reports search usage. + Standalone search sessions use the same ownership identity as continuation: + Main keeps its Claude Code session ID, while each direct Agent gets a stable, + opaque owner derived from its session and Agent IDs. The parent Agent ID is + validation-only. Missing, malformed, or ambiguous identity headers use a fresh + random search ID instead of sharing state. The search body ID and upstream + `session_id`, `x-client-request-id`, and `x-codex-window-id` headers all carry + that same search owner (with the window header's required `:0` suffix). - Top-level base64 user images map to `input_image`. - Supported base64 images nested in tool results also map to `input_image`. - Remote image URLs, malformed images, and unsupported tool-result image forms remain textual placeholders. @@ -58,6 +65,20 @@ Claude Code summary compaction requests are capped at low effort by default beca WebSocket is the default transport. Set `CCP_CODEX_TRANSPORT=http` for HTTP SSE, or `auto` to use WebSocket with HTTP fallback only when setup fails before a request is sent. +### Connection pacing + +Fresh WebSocket connections are paced adaptively. While the origin accepts upgrades the +proxy opens connections without spacing them; each rejected upgrade widens the spacing +(1s, then doubling up to 8s), and a run of successful connections narrows it back down. + +Because continuation is off by default, every request opens a fresh connection, so a +fixed spacing would cap a single process at roughly one generation per second no matter +how healthy the origin is. Pacing is therefore the price of an observed rejection rather +than a standing tax. + +Set `CCP_CODEX_WS_CONNECT_SPACING_MS` (or `codex.websocketConnectSpacingMs`) to impose a +floor the proxy never relaxes below. The default is `0`. + WebSocket setup honors `HTTP_PROXY` for `ws://`, `HTTPS_PROXY` for the default `wss://` endpoint, `ALL_PROXY` as a fallback, and `NO_PROXY` exclusions. A normal HTTP proxy can therefore carry the default WebSocket connection with CONNECT; TUN mode is not required. Set proxy variables before starting the process and restart after changing them. For example, setting `HTTPS_PROXY` to `http://127.0.0.1:7890` sends HTTPS/WSS destinations through the HTTP proxy at port 7890; it does not require an `https://` proxy URL. `CCP_CODEX_PREVIOUS_RESPONSE_ID=1` enables append-only WebSocket continuation. A valid identity containing only a Claude Code session ID owns the Main continuation for that session. Each valid direct Agent ID owns an independent continuation and reusable WebSocket within the same session. Nested Agents are keyed by their direct child ID; the parent ID is validated but does not become part of the owner key. The proxy sends `previous_response_id` only when the translated request shape and transcript extension are safe, and only on the exact live WebSocket that produced that response. diff --git a/src/config.rs b/src/config.rs index 15c0b072..53ccffd1 100644 --- a/src/config.rs +++ b/src/config.rs @@ -76,6 +76,8 @@ struct CodexConfig { #[serde(rename = "model")] pub model: Option, pub transport: Option, + #[serde(rename = "websocketConnectSpacingMs")] + pub websocket_connect_spacing_ms: Option, } #[derive(Deserialize, Clone)] @@ -875,6 +877,28 @@ fn parse_codex_transport(raw: &str) -> Option { } } +/// Plancher d'espacement des ouvertures de WebSocket Codex, en millisecondes. +/// +/// Zéro (le défaut) laisse le processus ouvrir ses connexions sans attente tant que +/// l'origine n'a refusé aucun upgrade ; l'espacement s'élargit alors tout seul. Un +/// opérateur qui préfère un rythme garanti peut imposer un plancher. +pub fn codex_websocket_connect_spacing() -> std::time::Duration { + let env: HashMap<_, _> = std::env::vars().collect(); + if let Some(raw) = env.get("CCP_CODEX_WS_CONNECT_SPACING_MS") + && let Ok(ms) = raw.trim().parse::() + { + return std::time::Duration::from_millis(ms); + } + let config_dir = paths::config_dir(); + if let Some(file) = read_file_config(&config_dir) + && let Some(codex) = file.codex + && let Some(ms) = codex.websocket_connect_spacing_ms + { + return std::time::Duration::from_millis(ms); + } + std::time::Duration::ZERO +} + pub fn codex_transport() -> CodexTransport { let env: HashMap<_, _> = std::env::vars().collect(); if let Some(raw) = env.get("CCP_CODEX_TRANSPORT") diff --git a/src/providers/codex/client.rs b/src/providers/codex/client.rs index 8f37ab81..c9efae6f 100644 --- a/src/providers/codex/client.rs +++ b/src/providers/codex/client.rs @@ -185,8 +185,19 @@ pub fn build_native_codex_headers( pub fn build_codex_search_headers( auth: &StoredAuth, ctx: &RequestContext, + search_session_id: &str, ) -> Result { let mut headers = build_codex_headers(auth, ctx, false)?; + headers.insert("session_id", header_value("session_id", search_session_id)?); + headers.insert( + "x-client-request-id", + header_value("x-client-request-id", search_session_id)?, + ); + let window_id = format!("{search_session_id}:0"); + headers.insert( + "x-codex-window-id", + header_value("x-codex-window-id", &window_id)?, + ); headers.insert( http::header::ACCEPT, header_value("accept", "application/json")?, @@ -1082,7 +1093,9 @@ impl CodexHttpClient { let mut retries = 0_u32; loop { - let response = self.attempt_post_search(&auth, &body_json, ctx).await?; + let response = self + .attempt_post_search(&auth, &body_json, ctx, &body.id) + .await?; if response.status == 401 && !auth_refresh_attempted { auth_refresh_attempted = true; auth = self @@ -2419,9 +2432,10 @@ impl CodexHttpClient { auth: &StoredAuth, body_json: &str, ctx: &RequestContext, + search_session_id: &str, ) -> Result { let url = search_endpoint(&self.base_url); - let headers = build_codex_search_headers(auth, ctx)?; + let headers = build_codex_search_headers(auth, ctx, search_session_id)?; if let Some(traffic) = ctx.traffic.as_deref() { write_codex_http_request_capture(traffic, &url, &headers, body_json); diff --git a/src/providers/codex/mod.rs b/src/providers/codex/mod.rs index c09f9302..f6e9722d 100644 --- a/src/providers/codex/mod.rs +++ b/src/providers/codex/mod.rs @@ -120,7 +120,7 @@ impl CodexProvider { let (search_request, query) = match search::build_search_request( &body, &resolved.model, - ctx.session_id.as_deref(), + conversation_identity.as_ref(), ) { Ok(request) => request, Err(error) => { diff --git a/src/providers/codex/search.rs b/src/providers/codex/search.rs index 60a846ac..f3f4218e 100644 --- a/src/providers/codex/search.rs +++ b/src/providers/codex/search.rs @@ -5,6 +5,7 @@ use serde_json::{Value, json}; use crate::anthropic::schema::MessagesRequest; use crate::anthropic::sse::encode_sse_event; +use crate::request_identity::ConversationIdentity; use crate::traffic::TrafficCapture; use super::count_tokens::{approx_token_count, truncate_to_token_budget}; @@ -91,14 +92,15 @@ pub fn is_standalone_search_request(req: &MessagesRequest) -> bool { pub fn build_search_request( req: &MessagesRequest, model: &str, - session_id: Option<&str>, + owner: Option<&ConversationIdentity>, ) -> Result<(SearchRequest, String), anyhow::Error> { let query = extract_search_query(req) .ok_or_else(|| anyhow::anyhow!("web_search request does not contain a text query"))?; let input = search_input(req); let filters = search_filters(req); - let id = session_id - .map(str::to_owned) + let id = owner + .map(ConversationIdentity::search_session_id) + .map(Into::into) .unwrap_or_else(|| format!("search-{}", uuid::Uuid::new_v4())); Ok(( @@ -504,6 +506,7 @@ fn emit(out: &mut Vec, traffic: Option<&TrafficCapture>, event: &str, data: mod tests { use super::*; use crate::anthropic::sse::parse_sse_events; + use crate::request_identity::ConversationIdentity; fn request() -> MessagesRequest { serde_json::from_value(json!({ @@ -528,9 +531,10 @@ mod tests { } #[test] - fn request_preserves_luna_and_omits_reasoning() { + fn request_preserves_luna_and_main_search_owner() { + let owner = ConversationIdentity::Main("session-1".into()); let (request, query) = - build_search_request(&request(), "gpt-5.6-luna", Some("session-1")).unwrap(); + build_search_request(&request(), "gpt-5.6-luna", Some(&owner)).unwrap(); assert_eq!(request.model, "gpt-5.6-luna"); assert_eq!(request.reasoning, None); assert_eq!(request.id, "session-1"); @@ -543,6 +547,36 @@ mod tests { ); } + #[test] + fn request_uses_stable_isolated_agent_search_owners() { + let agent = ConversationIdentity::Agent("session-a".into(), "agent-a".into()); + let same_agent = ConversationIdentity::Agent("session-a".into(), "agent-a".into()); + let sibling = ConversationIdentity::Agent("session-a".into(), "agent-b".into()); + let other_session = ConversationIdentity::Agent("session-b".into(), "agent-a".into()); + + let id = build_search_request(&request(), "gpt-5.6-luna", Some(&agent)) + .unwrap() + .0 + .id; + let same_id = build_search_request(&request(), "gpt-5.6-luna", Some(&same_agent)) + .unwrap() + .0 + .id; + let sibling_id = build_search_request(&request(), "gpt-5.6-luna", Some(&sibling)) + .unwrap() + .0 + .id; + let other_session_id = + build_search_request(&request(), "gpt-5.6-luna", Some(&other_session)) + .unwrap() + .0 + .id; + + assert_eq!(id, same_id); + assert_ne!(id, sibling_id); + assert_ne!(id, other_session_id); + } + #[test] fn only_forced_claude_search_uses_standalone_endpoint() { let forced = request(); @@ -555,6 +589,22 @@ mod tests { assert!(!is_standalone_search_request(&automatic)); } + #[test] + fn absent_owner_uses_a_fresh_stateless_search_id() { + let first = build_search_request(&request(), "gpt-5.6-luna", None) + .unwrap() + .0 + .id; + let second = build_search_request(&request(), "gpt-5.6-luna", None) + .unwrap() + .0 + .id; + + assert!(first.starts_with("search-")); + assert!(second.starts_with("search-")); + assert_ne!(first, second); + } + #[test] fn search_input_uses_role_specific_content_and_recent_context() { let mut req = request(); diff --git a/src/providers/codex/websocket.rs b/src/providers/codex/websocket.rs index 9e0d33fd..dd092c98 100644 --- a/src/providers/codex/websocket.rs +++ b/src/providers/codex/websocket.rs @@ -52,7 +52,11 @@ const MAX_POOL_ENTRIES: usize = 10_000; const POOL_CONNECT_CLEANUP_THRESHOLD: usize = 50; const POOL_CONNECT_CLEANUP_TARGET: usize = 40; const MAX_CONNECT_RESPONSE_HEADER_BYTES: usize = 8 * 1024; -const WEBSOCKET_CONNECT_START_SPACING: Duration = Duration::from_secs(1); +// Premier palier appliqué dès qu'une origine refuse un upgrade, puis doublement. +const WEBSOCKET_CONNECT_BACKOFF_STEP: Duration = Duration::from_secs(1); +const WEBSOCKET_CONNECT_MAX_SPACING: Duration = Duration::from_secs(8); +// Connexions réussies consécutives avant de réduire l'espacement de moitié. +const WEBSOCKET_CONNECT_RELAX_AFTER: u32 = 20; const WEBSOCKET_CONNECT_FORBIDDEN_COOLDOWN: Duration = Duration::from_secs(3); const WEBSOCKET_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30); const WEBSOCKET_KEEPALIVE_SEND_TIMEOUT: Duration = Duration::from_secs(10); @@ -334,7 +338,9 @@ static WS_POOL: once_cell::sync::Lazy = AsyncMutex::const_new(()); static WS_CONNECT_GATE: once_cell::sync::Lazy = - once_cell::sync::Lazy::new(|| WebSocketConnectGate::new(WEBSOCKET_CONNECT_START_SPACING)); + once_cell::sync::Lazy::new(|| { + WebSocketConnectGate::new(crate::config::codex_websocket_connect_spacing()) + }); fn next_monotonic_nonzero(sequence: &AtomicU64, label: &str) -> u64 { let previous = sequence @@ -562,26 +568,92 @@ fn cleanup_pool_before_connect() { drop(removed); } +struct GateState { + last_start: Option, + spacing: Duration, + successes: u32, +} + +/// Espacement adaptatif des ouvertures de WebSocket. +/// +/// Un espacement fixe protège l'origine des rafales d'upgrades (voir 529e1d6) mais la +/// facture est permanente : la continuation étant désactivée par défaut, chaque requête +/// ouvre une socket neuve, et un espacement d'une seconde plafonne le PROCESSUS entier +/// à environ une génération par seconde même quand l'origine accepte tout. Mesuré le +/// 2026-09-07 sur un compte au repos : 12 requêtes triviales en 13.27 s (0.90 req/s). +/// +/// L'espacement devient donc le prix d'un refus constaté et non une taxe d'avance : il +/// part du plancher configuré (zéro par défaut), double à chaque upgrade refusé jusqu'à +/// un plafond, puis se relâche après une série de connexions réussies. struct WebSocketConnectGate { - last_start: AsyncMutex>, - start_spacing: Duration, + state: AsyncMutex, + floor: Duration, } impl WebSocketConnectGate { fn new(start_spacing: Duration) -> Self { Self { - last_start: AsyncMutex::new(None), - start_spacing, + state: AsyncMutex::new(GateState { + last_start: None, + spacing: start_spacing, + successes: 0, + }), + floor: start_spacing, } } async fn wait_to_start(&self, before_start: impl std::future::Future) { - let mut last_start = self.last_start.lock().await; - if let Some(previous) = *last_start { - tokio::time::sleep_until(previous + self.start_spacing).await; + let mut state = self.state.lock().await; + if let Some(previous) = state.last_start + && !state.spacing.is_zero() + { + tokio::time::sleep_until(previous + state.spacing).await; } before_start.await; - *last_start = Some(tokio::time::Instant::now()); + state.last_start = Some(tokio::time::Instant::now()); + } + + /// Une origine vient de refuser un upgrade : élargir immédiatement. + fn note_origin_forbidden(&self) { + let mut state = match self.state.try_lock() { + Ok(state) => state, + // Un autre appelant tient l'état : il est en train d'ouvrir une connexion. + // Perdre un signal d'élargissement est sans conséquence — le refus suivant + // le portera — alors que bloquer ici retiendrait le chemin de connexion. + Err(_) => return, + }; + state.successes = 0; + state.spacing = if state.spacing.is_zero() { + WEBSOCKET_CONNECT_BACKOFF_STEP + } else { + state.spacing.saturating_mul(2) + } + .min(WEBSOCKET_CONNECT_MAX_SPACING) + .max(self.floor); + } + + /// Une connexion a abouti : se rapprocher du plancher configuré. + fn note_connect_success(&self) { + let Ok(mut state) = self.state.try_lock() else { + return; + }; + if state.spacing <= self.floor { + return; + } + state.successes += 1; + if state.successes < WEBSOCKET_CONNECT_RELAX_AFTER { + return; + } + state.successes = 0; + state.spacing = (state.spacing / 2).max(self.floor); + if state.spacing < WEBSOCKET_CONNECT_BACKOFF_STEP { + state.spacing = self.floor; + } + } + + #[cfg(test)] + async fn spacing_for_tests(&self) -> Duration { + self.state.lock().await.spacing } } @@ -1903,13 +1975,20 @@ where connect_timeout, ))) }); + // L'espacement suit ce que l'origine répond réellement : un upgrade refusé + // l'élargit, une connexion établie le relâche vers le plancher configuré. if result .as_ref() .is_err_and(ConnectAttemptError::is_origin_forbidden) - && attempt == 0 { - retry_sleep(u64::try_from(forbidden_cooldown.as_millis()).unwrap_or(u64::MAX)).await; - continue; + gate.note_origin_forbidden(); + if attempt == 0 { + retry_sleep(u64::try_from(forbidden_cooldown.as_millis()).unwrap_or(u64::MAX)) + .await; + continue; + } + } else if result.is_ok() { + gate.note_connect_success(); } return result.map_err(ConnectAttemptError::into_error); } @@ -2648,6 +2727,75 @@ mod tests { clear_codex_websocket_pool_for_tests(); } + #[tokio::test(start_paused = true)] + async fn idle_gate_starts_connections_without_spacing_them() { + // Mesuré le 2026-09-07 : avec un espacement fixe d'une seconde, un processus + // plafonne à ~0.90 génération/seconde quelle que soit la santé de l'origine, + // parce que la continuation est désactivée par défaut et que chaque requête + // ouvre donc une socket neuve. L'espacement doit être le prix d'un refus + // constaté, pas une taxe permanente. + let gate = WebSocketConnectGate::new(Duration::ZERO); + let started = tokio::time::Instant::now(); + for _ in 0..8 { + gate.wait_to_start(async {}).await; + } + assert_eq!( + tokio::time::Instant::now().duration_since(started), + Duration::ZERO + ); + } + + #[tokio::test(start_paused = true)] + async fn gate_spaces_starts_after_the_origin_rejects_an_upgrade() { + // L'intention d'origine (529e1d6) est préservée : dès qu'une origine refuse un + // upgrade, on cesse de la marteler. + let gate = WebSocketConnectGate::new(Duration::ZERO); + gate.wait_to_start(async {}).await; + gate.note_origin_forbidden(); + + let before = tokio::time::Instant::now(); + gate.wait_to_start(async {}).await; + assert!( + tokio::time::Instant::now().duration_since(before) >= WEBSOCKET_CONNECT_BACKOFF_STEP, + "a rejected origin must be given room" + ); + } + + #[tokio::test(start_paused = true)] + async fn repeated_rejections_widen_the_spacing_up_to_a_cap() { + let gate = WebSocketConnectGate::new(Duration::ZERO); + for _ in 0..16 { + gate.note_origin_forbidden(); + } + assert_eq!( + gate.spacing_for_tests().await, + WEBSOCKET_CONNECT_MAX_SPACING + ); + } + + #[tokio::test(start_paused = true)] + async fn sustained_success_relaxes_the_spacing_back_to_zero() { + // Sans retour à zéro, un seul 403 taxerait le processus pour toute sa vie. + let gate = WebSocketConnectGate::new(Duration::ZERO); + gate.note_origin_forbidden(); + assert!(gate.spacing_for_tests().await > Duration::ZERO); + for _ in 0..(WEBSOCKET_CONNECT_RELAX_AFTER * 8) { + gate.note_connect_success(); + } + assert_eq!(gate.spacing_for_tests().await, Duration::ZERO); + } + + #[tokio::test(start_paused = true)] + async fn a_configured_floor_is_never_relaxed_away() { + // Un opérateur qui impose un espacement garde son plancher, même après des + // milliers de succès. + let gate = WebSocketConnectGate::new(Duration::from_millis(250)); + for _ in 0..(WEBSOCKET_CONNECT_RELAX_AFTER * 8) { + gate.note_connect_success(); + } + assert_eq!(gate.spacing_for_tests().await, Duration::from_millis(250)); + } + #[tokio::test(start_paused = true)] async fn connect_gate_spaces_starts_without_serializing_handshakes() { let gate = Arc::new(WebSocketConnectGate::new(Duration::from_secs(1))); diff --git a/src/request_identity.rs b/src/request_identity.rs index 0df0a13f..364eeeb7 100644 --- a/src/request_identity.rs +++ b/src/request_identity.rs @@ -1,4 +1,8 @@ +use std::borrow::Cow; +use std::fmt::Write as _; + use http::HeaderMap; +use sha2::{Digest, Sha256}; pub const CLAUDE_SESSION_HEADER: &str = "x-claude-code-session-id"; pub const CLAUDE_AGENT_HEADER: &str = "x-claude-code-agent-id"; @@ -30,6 +34,26 @@ impl ConversationIdentity { _ => None, } } + + pub fn search_session_id(&self) -> Cow<'_, str> { + match self { + Self::Main(session_id) => Cow::Borrowed(session_id), + Self::Agent(session_id, agent_id) => { + let mut digest = Sha256::new(); + digest.update(b"claude-code-proxy:search-agent:v1\0"); + for value in [session_id, agent_id] { + digest.update((value.len() as u64).to_be_bytes()); + digest.update(value.as_bytes()); + } + let digest = digest.finalize(); + let mut encoded = String::with_capacity(digest.len() * 2); + for byte in digest { + write!(&mut encoded, "{byte:02x}").expect("writing to String cannot fail"); + } + Cow::Owned(format!("search-agent-v1-{encoded}")) + } + } + } } #[derive(Debug)] @@ -181,6 +205,42 @@ mod tests { } } + #[test] + fn search_owner_is_stable_and_isolates_identity_tuples() { + let agent = ConversationIdentity::Agent("session-a".into(), "agent-a".into()); + let sibling = ConversationIdentity::Agent("session-a".into(), "agent-b".into()); + let other_session = ConversationIdentity::Agent("session-b".into(), "agent-a".into()); + let shifted_boundary = ConversationIdentity::Agent("session-aa".into(), "gent-a".into()); + + assert_eq!(agent.search_session_id(), agent.search_session_id()); + assert_ne!(agent.search_session_id(), sibling.search_session_id()); + assert_ne!(agent.search_session_id(), other_session.search_session_id()); + assert_ne!( + agent.search_session_id(), + shifted_boundary.search_session_id() + ); + } + + #[test] + fn agent_search_owner_is_bounded_and_does_not_disclose_raw_ids() { + let session = "s".repeat(MAX_IDENTITY_LEN); + let agent = "a".repeat(MAX_IDENTITY_LEN); + let owner = ConversationIdentity::Agent(session.clone(), agent.clone()) + .search_session_id() + .into_owned(); + + assert_eq!(owner.len(), "search-agent-v1-".len() + 64); + assert!(!owner.contains(&session)); + assert!(!owner.contains(&agent)); + } + + #[test] + fn main_search_owner_preserves_the_raw_session_id() { + let main = ConversationIdentity::Main("session-main".into()); + + assert_eq!(main.search_session_id(), "session-main"); + } + #[test] fn parent_is_validation_only_and_never_changes_the_owner() { let direct = ConversationIdentity::from_headers(&headers(&[ diff --git a/tests/smoke_cutover.rs b/tests/smoke_cutover.rs index 12f1622f..7d35fb57 100644 --- a/tests/smoke_cutover.rs +++ b/tests/smoke_cutover.rs @@ -105,6 +105,21 @@ async fn call_messages_body(body: Value) -> Response { .unwrap() } +async fn call_messages_body_with_headers(body: Value, headers: &[(&str, &str)]) -> Response { + let _no_proxy_env = EnvGuard::set("NO_PROXY", "127.0.0.1,localhost"); + let mut request = Request::builder() + .method(Method::POST) + .uri("/v1/messages") + .header("content-type", "application/json"); + for (name, value) in headers { + request = request.header(*name, *value); + } + app(Arc::new(Registry::with_default_alias())) + .oneshot(request.body(Body::from(body.to_string())).unwrap()) + .await + .unwrap() +} + async fn call_responses_body(body: Value) -> Response { let _no_proxy_env = EnvGuard::set("NO_PROXY", "127.0.0.1,localhost"); app_with_options(Arc::new(Registry::with_default_alias()), None, true) @@ -168,6 +183,14 @@ fn traffic_json(files: &[PathBuf], suffix: &str) -> Value { async fn spawn_http_upstream(handler: F) -> String where F: Fn(Value) -> Vec + Send + Sync + 'static, +{ + let handler = Arc::new(handler); + spawn_http_upstream_with_headers(move |_, body| handler(body)).await +} + +async fn spawn_http_upstream_with_headers(handler: F) -> String +where + F: Fn(http::HeaderMap, Value) -> Vec + Send + Sync + 'static, { let handler = Arc::new(handler); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -176,11 +199,11 @@ where let app = axum::Router::new().fallback({ let handler = handler.clone(); - move |body: String| { + move |headers: http::HeaderMap, body: String| { let handler = handler.clone(); async move { let json: Value = serde_json::from_str(&body).unwrap_or_default(); - let response_bytes = handler(json); + let response_bytes = handler(headers, json); http::Response::builder() .status(StatusCode::OK) .header("content-type", "text/event-stream") @@ -883,6 +906,121 @@ async fn spawn_websocket_always_empty_completion_upstream( // Health and routing smoke tests (no env var mutation needed) // --------------------------------------------------------------------------- +#[tokio::test] +#[allow(clippy::await_holding_lock)] +async fn codex_standalone_search_uses_conversation_identity_owner() { + let _guard = env_lock(); + let config = TempDir::new().unwrap(); + write_auth(config.path(), "codex"); + let captured = Arc::new(Mutex::new(Vec::<[String; 4]>::new())); + let upstream = spawn_http_upstream_with_headers({ + let captured = captured.clone(); + move |headers, body| { + captured.lock().unwrap().push([ + body["id"].as_str().unwrap().to_string(), + headers + .get("session_id") + .unwrap() + .to_str() + .unwrap() + .to_string(), + headers + .get("x-client-request-id") + .unwrap() + .to_str() + .unwrap() + .to_string(), + headers + .get("x-codex-window-id") + .unwrap() + .to_str() + .unwrap() + .to_string(), + ]); + serde_json::to_vec(&json!({ + "encrypted_output": null, + "output": "search ok", + "results": [] + })) + .unwrap() + } + }) + .await; + let _config_env = EnvGuard::set("CCP_CONFIG_DIR", config.path()); + let _base_url_env = EnvGuard::set("CCP_CODEX_BASE_URL", &upstream); + let body = json!({ + "model": "gpt-5.6-luna", + "max_tokens": 64, + "messages": [{ + "role": "user", + "content": "Perform a web search for the query: identity owner" + }], + "tools": [{"type": "web_search_20250305", "name": "web_search"}], + "tool_choice": {"type": "tool", "name": "web_search"} + }); + let cases: &[&[(&str, &str)]] = &[ + &[ + ("x-claude-code-session-id", "session-a"), + ("x-claude-code-agent-id", "agent-a"), + ("x-claude-code-parent-agent-id", "parent-one"), + ], + &[ + ("x-claude-code-session-id", "session-a"), + ("x-claude-code-agent-id", "agent-a"), + ("x-claude-code-parent-agent-id", "parent-two"), + ], + &[ + ("x-claude-code-session-id", "session-a"), + ("x-claude-code-agent-id", "agent-b"), + ], + &[ + ("x-claude-code-session-id", "session-b"), + ("x-claude-code-agent-id", "agent-a"), + ], + &[("x-claude-code-session-id", "session-main")], + &[], + &[ + ("x-claude-code-session-id", "session-malformed"), + ("x-claude-code-agent-id", "bad agent"), + ], + ]; + for headers in cases { + let response = call_messages_body_with_headers(body.clone(), headers).await; + assert_eq!(response.status(), StatusCode::OK); + } + + let identities = captured.lock().unwrap(); + assert_eq!(identities.len(), cases.len()); + for carriers in identities.iter() { + assert_eq!(carriers[1], carriers[0], "session_id must match body id"); + assert_eq!( + carriers[2], carriers[0], + "x-client-request-id must match body id" + ); + assert_eq!( + carriers[3], + format!("{}:0", carriers[0]), + "x-codex-window-id must use the body id" + ); + } + assert_eq!( + identities[0][0], identities[1][0], + "same direct agent must keep its owner" + ); + assert_ne!( + identities[0][0], identities[2][0], + "siblings must not share an owner" + ); + assert_ne!( + identities[0][0], identities[3][0], + "sessions must not share an owner" + ); + assert_eq!(identities[4][0], "session-main"); + assert!(identities[5][0].starts_with("search-")); + assert!(identities[6][0].starts_with("search-")); + assert_ne!(identities[5][0], identities[6][0]); +} + #[tokio::test] async fn smoke_healthz_returns_ok() { let app = app(Arc::new(Registry::with_default_alias()));