diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index d2a0d0bde5..4f0af0d04f 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2547,6 +2547,7 @@ dependencies = [ "codex-feedback", "codex-git-utils", "codex-hooks", + "codex-http-state", "codex-install-context", "codex-login", "codex-mcp", @@ -3321,6 +3322,7 @@ dependencies = [ "codex-aws-auth", "codex-client", "codex-feedback", + "codex-http-state", "codex-login", "codex-model-provider-info", "codex-models-manager", @@ -3330,6 +3332,7 @@ dependencies = [ "http 1.4.0", "pretty_assertions", "serde_json", + "tempfile", "tokio", "tracing", "wiremock", diff --git a/codex-rs/codex-api/src/auth.rs b/codex-rs/codex-api/src/auth.rs index 41394a2258..a6fd6622f6 100644 --- a/codex-rs/codex-api/src/auth.rs +++ b/codex-rs/codex-api/src/auth.rs @@ -41,6 +41,27 @@ pub trait AuthProvider: Send + Sync { headers } + /// Adds auth headers for a concrete HTTP request URL. + /// + /// URL-sensitive providers can override this to scope request-specific + /// headers without exposing them through generic/non-HTTP auth helpers. + fn add_auth_headers_for_url(&self, _request_url: &str, headers: &mut HeaderMap) { + self.add_auth_headers(headers); + } + + /// Observes response headers for auth state that may need to rotate after a request. + /// + /// Most providers do not need this. Providers with server-minted, + /// request-scoped state may validate the URL and selectively persist + /// response headers here. + fn observe_response_headers( + &self, + _request_url: &str, + _request_headers: &HeaderMap, + _response_headers: &HeaderMap, + ) { + } + /// Applies auth to a complete outbound request and returns the request to send. /// /// The input `request` is moved into this method. Implementations may mutate @@ -54,7 +75,7 @@ pub trait AuthProvider: Send + Sync { /// If this returns [`AuthError`], the request should not be sent. async fn apply_auth(&self, request: Request) -> Result { let mut request = request; - self.add_auth_headers(&mut request.headers); + self.add_auth_headers_for_url(&request.url, &mut request.headers); Ok(request) } } diff --git a/codex-rs/codex-api/src/endpoint/mod.rs b/codex-rs/codex-api/src/endpoint/mod.rs index 106c5d73ff..adbf5629cb 100644 --- a/codex-rs/codex-api/src/endpoint/mod.rs +++ b/codex-rs/codex-api/src/endpoint/mod.rs @@ -6,6 +6,7 @@ pub(crate) mod realtime_call; pub(crate) mod realtime_websocket; pub(crate) mod responses; pub(crate) mod responses_websocket; +mod responses_websocket_auth; pub(crate) mod search; mod session; diff --git a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs index ad26549a26..5df495786e 100644 --- a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs +++ b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs @@ -1,3 +1,4 @@ +use crate::auth::SharedAuthProvider; use crate::endpoint::realtime_websocket::methods_common::conversation_function_call_output_message; use crate::endpoint::realtime_websocket::methods_common::conversation_item_create_message; use crate::endpoint::realtime_websocket::methods_common::normalized_session_mode; @@ -550,11 +551,21 @@ fn contains_transcript_entry(entries: &[RealtimeTranscriptEntry], role: &str, te pub struct RealtimeWebsocketClient { provider: Provider, + auth: Option, } impl RealtimeWebsocketClient { pub fn new(provider: Provider) -> Self { - Self { provider } + Self { + provider, + auth: None, + } + } + + /// Applies URL-aware auth during handshakes and observes handshake responses. + pub fn with_auth(mut self, auth: SharedAuthProvider) -> Self { + self.auth = Some(auth); + self } pub async fn connect( @@ -653,6 +664,10 @@ impl RealtimeWebsocketClient { default_headers, ); request.headers_mut().extend(headers); + if let Some(auth) = self.auth.as_ref() { + auth.add_auth_headers_for_url(ws_url.as_str(), request.headers_mut()); + } + let request_headers = request.headers().clone(); info!("connecting realtime websocket: {ws_url}"); // Realtime websocket TLS should honor the same custom-CA env vars as the rest of Codex's @@ -660,14 +675,39 @@ impl RealtimeWebsocketClient { let connector = maybe_build_rustls_client_config_with_custom_ca() .map_err(|err| ApiError::Stream(format!("failed to configure websocket TLS: {err}")))? .map(tokio_tungstenite::Connector::Rustls); - let (stream, response) = tokio_tungstenite::connect_async_tls_with_config( + let connect_result = tokio_tungstenite::connect_async_tls_with_config( request, Some(websocket_config()), false, connector, ) - .await - .map_err(|err| ApiError::Stream(format!("failed to connect realtime websocket: {err}")))?; + .await; + let (stream, response) = match connect_result { + Ok((stream, response)) => { + if let Some(auth) = self.auth.as_ref() { + auth.observe_response_headers( + ws_url.as_str(), + &request_headers, + response.headers(), + ); + } + (stream, response) + } + Err(err) => { + if let Some(auth) = self.auth.as_ref() + && let WsError::Http(response) = &err + { + auth.observe_response_headers( + ws_url.as_str(), + &request_headers, + response.headers(), + ); + } + return Err(ApiError::Stream(format!( + "failed to connect realtime websocket: {err}" + ))); + } + }; info!( ws_url = %ws_url, status = %response.status(), diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index f0a0019817..70a0365b77 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -1,3 +1,4 @@ +use super::responses_websocket_auth::WebsocketAuthContext; use crate::auth::SharedAuthProvider; use crate::common::ResponseEvent; use crate::common::ResponseProcessedWsRequest; @@ -167,6 +168,7 @@ pub struct ResponsesWebsocketConnection { server_reasoning_included: bool, models_etag: Option, server_model: Option, + auth_context: WebsocketAuthContext, telemetry: Option>, } @@ -178,6 +180,7 @@ impl std::fmt::Debug for ResponsesWebsocketConnection { .field("server_reasoning_included", &self.server_reasoning_included) .field("models_etag", &self.models_etag) .field("server_model", &self.server_model) + .field("auth_context", &"") .field("telemetry", &self.telemetry.as_ref().map(|_| "")) .finish() } @@ -190,6 +193,7 @@ impl ResponsesWebsocketConnection { server_reasoning_included: bool, models_etag: Option, server_model: Option, + auth_context: WebsocketAuthContext, telemetry: Option>, ) -> Self { Self { @@ -198,6 +202,7 @@ impl ResponsesWebsocketConnection { server_reasoning_included, models_etag, server_model, + auth_context, telemetry, } } @@ -258,6 +263,7 @@ impl ResponsesWebsocketConnection { let server_reasoning_included = self.server_reasoning_included; let models_etag = self.models_etag.clone(); let server_model = self.server_model.clone(); + let auth_context = self.auth_context.clone(); let telemetry = self.telemetry.clone(); let request_body = serde_json::to_value(&request).map_err(|err| { ApiError::Stream(format!("failed to encode websocket request: {err}")) @@ -297,6 +303,7 @@ impl ResponsesWebsocketConnection { tx_event.clone(), request_body, idle_timeout, + &auth_context, telemetry, connection_reused, ) @@ -377,19 +384,35 @@ impl ResponsesWebsocketClient { .provider .websocket_url_for_path("responses") .map_err(|err| ApiError::Stream(format!("failed to build websocket URL: {err}")))?; + let request_url = ws_url.to_string(); let mut headers = merge_request_headers(&self.provider.headers, extra_headers, default_headers); - self.auth.add_auth_headers(&mut headers); + self.auth + .add_auth_headers_for_url(&request_url, &mut headers); + let auth_context = + WebsocketAuthContext::new(self.auth.clone(), request_url, headers.clone()); - let (stream, _status, server_reasoning_included, models_etag, server_model) = - connect_websocket(ws_url, headers, turn_state.clone()).await?; + let ( + stream, + _status, + server_reasoning_included, + models_etag, + server_model, + response_headers, + ) = connect_websocket(ws_url, headers, turn_state.clone()) + .await + .inspect_err(|error| { + auth_context.observe_error_headers(error); + })?; + auth_context.observe_response_headers(&response_headers); Ok(ResponsesWebsocketConnection::new( stream, self.provider.stream_idle_timeout, server_reasoning_included, models_etag, server_model, + auth_context, telemetry, )) } @@ -411,13 +434,22 @@ impl ResponsesWebsocketClient { .provider .websocket_url_for_path("responses") .map_err(|err| ApiError::Stream(format!("failed to build websocket URL: {err}")))?; + let request_url = ws_url.to_string(); let mut headers = merge_request_headers(&self.provider.headers, extra_headers, default_headers); - self.auth.add_auth_headers(&mut headers); + self.auth + .add_auth_headers_for_url(&request_url, &mut headers); + let auth_context = + WebsocketAuthContext::new(self.auth.clone(), request_url, headers.clone()); - let (mut stream, status, reasoning_included, models_etag, server_model) = - connect_websocket(ws_url.clone(), headers, /*turn_state*/ None).await?; + let (mut stream, status, reasoning_included, models_etag, server_model, response_headers) = + connect_websocket(ws_url.clone(), headers, /*turn_state*/ None) + .await + .inspect_err(|error| { + auth_context.observe_error_headers(error); + })?; + auth_context.observe_response_headers(&response_headers); let immediate_close = tokio::time::timeout(immediate_close_timeout, stream.next()) .await .ok() @@ -472,7 +504,17 @@ async fn connect_websocket( url: Url, headers: HeaderMap, turn_state: Option>>, -) -> Result<(WsStream, StatusCode, bool, Option, Option), ApiError> { +) -> Result< + ( + WsStream, + StatusCode, + bool, + Option, + Option, + HeaderMap, + ), + ApiError, +> { ensure_rustls_crypto_provider(); info!("connecting to websocket: {url}"); @@ -499,10 +541,7 @@ async fn connect_websocket( let (stream, response) = match response { Ok((stream, response)) => { - info!( - "successfully connected to websocket: {url}, headers: {:?}", - response.headers() - ); + info!("successfully connected to websocket: {url}"); (stream, response) } Err(err) => { @@ -536,6 +575,7 @@ async fn connect_websocket( reasoning_included, models_etag, server_model, + response.headers().clone(), )) } @@ -630,12 +670,12 @@ fn map_wrapped_websocket_error_event( Some(ApiError::Transport(TransportError::Http { status, url: None, - headers: headers.map(json_headers_to_http_headers), + headers: headers.as_ref().map(json_headers_to_http_headers), body: Some(original_payload), })) } -fn json_headers_to_http_headers(headers: JsonMap) -> HeaderMap { +fn json_headers_to_http_headers(headers: &JsonMap) -> HeaderMap { let mut mapped = HeaderMap::new(); for (name, value) in headers { let Ok(header_name) = HeaderName::from_bytes(name.as_bytes()) else { @@ -649,9 +689,9 @@ fn json_headers_to_http_headers(headers: JsonMap) -> HeaderMap { mapped } -fn json_header_value(value: Value) -> Option { +fn json_header_value(value: &Value) -> Option { let value = match value { - Value::String(value) => value, + Value::String(value) => value.clone(), Value::Number(value) => value.to_string(), Value::Bool(value) => value.to_string(), _ => return None, @@ -664,6 +704,7 @@ async fn run_websocket_response_stream( tx_event: mpsc::Sender>, request_body: Value, idle_timeout: Duration, + auth_context: &WebsocketAuthContext, telemetry: Option>, connection_reused: bool, ) -> Result<(), ApiError> { @@ -702,12 +743,25 @@ async fn run_websocket_response_stream( match message { Message::Text(text) => { - trace!("websocket event: {text}"); - if let Some(wrapped_error) = parse_wrapped_websocket_error_event(&text) - && let Some(error) = + let wrapped_error = parse_wrapped_websocket_error_event(&text); + if wrapped_error.is_some() { + trace!("websocket error event"); + } else { + trace!("websocket event: {text}"); + } + if let Some(wrapped_error) = wrapped_error { + if let Some(response_headers) = wrapped_error.headers.as_ref() { + auth_context.observe_response_headers(&json_headers_to_http_headers( + response_headers, + )); + } + if let Some(error) = map_wrapped_websocket_error_event(wrapped_error, text.to_string()) - { - return Err(error); + { + return Err(error); + } + debug!("ignoring unmapped websocket error event"); + continue; } let event = match serde_json::from_str::(&text) { diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket_auth.rs b/codex-rs/codex-api/src/endpoint/responses_websocket_auth.rs new file mode 100644 index 0000000000..83d17963ba --- /dev/null +++ b/codex-rs/codex-api/src/endpoint/responses_websocket_auth.rs @@ -0,0 +1,52 @@ +use crate::auth::SharedAuthProvider; +use crate::error::ApiError; +use codex_client::TransportError; +use http::HeaderMap; +use std::sync::Arc; +use std::sync::Mutex; + +#[derive(Clone)] +pub(super) struct WebsocketAuthContext { + auth: SharedAuthProvider, + request_url: String, + request_headers: Arc>, +} + +impl WebsocketAuthContext { + pub(super) fn new( + auth: SharedAuthProvider, + request_url: String, + request_headers: HeaderMap, + ) -> Self { + Self { + auth, + request_url, + request_headers: Arc::new(Mutex::new(request_headers)), + } + } + + pub(super) fn observe_response_headers(&self, response_headers: &HeaderMap) { + let mut request_headers = self + .request_headers + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + self.auth + .observe_response_headers(&self.request_url, &request_headers, response_headers); + self.auth + .add_auth_headers_for_url(&self.request_url, &mut request_headers); + } + + pub(super) fn observe_error_headers(&self, error: &ApiError) { + if let ApiError::Transport(TransportError::Http { + headers: Some(response_headers), + .. + }) = error + { + self.observe_response_headers(response_headers); + } + } +} + +#[cfg(test)] +#[path = "responses_websocket_auth_tests.rs"] +mod tests; diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket_auth_tests.rs b/codex-rs/codex-api/src/endpoint/responses_websocket_auth_tests.rs new file mode 100644 index 0000000000..5c38e2dbc8 --- /dev/null +++ b/codex-rs/codex-api/src/endpoint/responses_websocket_auth_tests.rs @@ -0,0 +1,124 @@ +use super::*; +use crate::auth::AuthProvider; +use http::HeaderValue; +use http::StatusCode; +use pretty_assertions::assert_eq; +use std::sync::Mutex as StdMutex; + +const REQUEST_URL: &str = "wss://chatgpt.com/backend-api/codex/responses"; +const SENT_STATE: &str = "sent-state"; +const HANDSHAKE_STATE: &str = "handshake-state"; +const WRAPPED_ERROR_STATE: &str = "wrapped-error-state"; + +struct RecordingAuthProvider { + state: StdMutex, + observed_updates: StdMutex>, +} + +impl Default for RecordingAuthProvider { + fn default() -> Self { + Self { + state: StdMutex::new(SENT_STATE.to_string()), + observed_updates: StdMutex::new(Vec::new()), + } + } +} + +impl AuthProvider for RecordingAuthProvider { + fn add_auth_headers(&self, _headers: &mut HeaderMap) {} + + fn add_auth_headers_for_url(&self, _request_url: &str, headers: &mut HeaderMap) { + let state = self + .state + .lock() + .expect("state lock should not be poisoned"); + headers.insert( + "x-test-state", + HeaderValue::from_str(&state).expect("state should be a valid header value"), + ); + } + + fn observe_response_headers( + &self, + request_url: &str, + request_headers: &HeaderMap, + response_headers: &HeaderMap, + ) { + let Some(sent_state) = request_headers + .get("x-test-state") + .and_then(|value| value.to_str().ok()) + else { + return; + }; + let Some(update_state) = response_headers + .get("x-test-state-update") + .and_then(|value| value.to_str().ok()) + else { + return; + }; + self.observed_updates + .lock() + .expect("observed updates lock should not be poisoned") + .push(( + request_url.to_string(), + sent_state.to_string(), + update_state.to_string(), + )); + *self + .state + .lock() + .expect("state lock should not be poisoned") = update_state.to_string(); + } +} + +#[test] +fn observes_handshake_and_wrapped_error_headers() { + let auth = Arc::new(RecordingAuthProvider::default()); + let mut request_headers = HeaderMap::new(); + auth.add_auth_headers_for_url(REQUEST_URL, &mut request_headers); + assert_eq!( + request_headers + .get("x-test-state") + .and_then(|value| value.to_str().ok()), + Some(SENT_STATE) + ); + let context = WebsocketAuthContext::new(auth.clone(), REQUEST_URL.to_string(), request_headers); + + let mut handshake_headers = HeaderMap::new(); + handshake_headers.insert( + "x-test-state-update", + HeaderValue::from_static(HANDSHAKE_STATE), + ); + context.observe_response_headers(&handshake_headers); + + let mut wrapped_error_headers = HeaderMap::new(); + wrapped_error_headers.insert( + "x-test-state-update", + HeaderValue::from_static(WRAPPED_ERROR_STATE), + ); + context.observe_error_headers(&ApiError::Transport(TransportError::Http { + status: StatusCode::BAD_REQUEST, + url: Some(REQUEST_URL.to_string()), + headers: Some(wrapped_error_headers), + body: None, + })); + + assert_eq!( + *auth + .observed_updates + .lock() + .expect("observed updates lock should not be poisoned"), + vec![ + ( + REQUEST_URL.to_string(), + SENT_STATE.to_string(), + HANDSHAKE_STATE.to_string(), + ), + ( + REQUEST_URL.to_string(), + HANDSHAKE_STATE.to_string(), + WRAPPED_ERROR_STATE.to_string(), + ), + ] + ); +} diff --git a/codex-rs/codex-api/src/endpoint/session.rs b/codex-rs/codex-api/src/endpoint/session.rs index 132c3abd90..8bf5ce2e4b 100644 --- a/codex-rs/codex-api/src/endpoint/session.rs +++ b/codex-rs/codex-api/src/endpoint/session.rs @@ -102,7 +102,16 @@ impl EndpointSession { let transport = &self.transport; async move { let req = auth.apply_auth(req).await.map_err(TransportError::from)?; - transport.execute(req).await + let request_url = req.url.clone(); + let request_headers = req.headers.clone(); + let response = transport.execute(req).await; + observe_auth_response_headers( + auth.as_ref(), + &request_url, + &request_headers, + &response, + ); + response } }, ) @@ -143,7 +152,16 @@ impl EndpointSession { let transport = &self.transport; async move { let req = auth.apply_auth(req).await.map_err(TransportError::from)?; - transport.stream(req).await + let request_url = req.url.clone(); + let request_headers = req.headers.clone(); + let response = transport.stream(req).await; + observe_auth_response_headers( + auth.as_ref(), + &request_url, + &request_headers, + &response, + ); + response } }, ) @@ -152,3 +170,44 @@ impl EndpointSession { Ok(stream) } } + +fn observe_auth_response_headers( + auth: &dyn crate::auth::AuthProvider, + request_url: &str, + request_headers: &HeaderMap, + response: &Result, +) where + T: ResponseHeaders, +{ + match response { + Ok(response) => { + auth.observe_response_headers(request_url, request_headers, response.headers()) + } + Err(TransportError::Http { + headers: Some(headers), + .. + }) => auth.observe_response_headers(request_url, request_headers, headers), + Err(_) => {} + } +} + +/// Exposes transport response headers so auth providers can observe request-scoped updates. +trait ResponseHeaders { + fn headers(&self) -> &HeaderMap; +} + +impl ResponseHeaders for Response { + fn headers(&self) -> &HeaderMap { + &self.headers + } +} + +impl ResponseHeaders for StreamResponse { + fn headers(&self) -> &HeaderMap { + &self.headers + } +} + +#[cfg(test)] +#[path = "session_tests.rs"] +mod tests; diff --git a/codex-rs/codex-api/src/endpoint/session_tests.rs b/codex-rs/codex-api/src/endpoint/session_tests.rs new file mode 100644 index 0000000000..884b704113 --- /dev/null +++ b/codex-rs/codex-api/src/endpoint/session_tests.rs @@ -0,0 +1,174 @@ +use super::*; +use crate::provider::RetryConfig; +use async_trait::async_trait; +use codex_client::Request; +use http::HeaderValue; +use http::StatusCode; +use pretty_assertions::assert_eq; +use std::sync::Mutex as StdMutex; +use std::time::Duration; + +const RESPONSES_URL: &str = "https://chatgpt.com/backend-api/codex/responses"; + +#[derive(Default)] +struct RecordingAuthProvider { + observed_headers: StdMutex>, +} + +impl crate::auth::AuthProvider for RecordingAuthProvider { + fn add_auth_headers(&self, _headers: &mut HeaderMap) {} + + fn add_auth_headers_for_url(&self, _request_url: &str, headers: &mut HeaderMap) { + headers.insert("x-test-state", HeaderValue::from_static("sent-state")); + } + + fn observe_response_headers( + &self, + request_url: &str, + request_headers: &HeaderMap, + response_headers: &HeaderMap, + ) { + self.observed_headers + .lock() + .expect("recording auth lock should not be poisoned") + .push(( + request_url.to_string(), + request_headers.clone(), + response_headers.clone(), + )); + } +} + +#[derive(Clone)] +struct RejectingTransport { + requests: Arc>>, + response_headers: HeaderMap, +} + +impl RejectingTransport { + fn new(response_headers: HeaderMap) -> Self { + Self { + requests: Arc::new(StdMutex::new(Vec::new())), + response_headers, + } + } + + fn http_error(&self, request: Request) -> TransportError { + self.requests + .lock() + .expect("recording transport lock should not be poisoned") + .push((request.url.clone(), request.headers)); + TransportError::Http { + status: StatusCode::UNAUTHORIZED, + url: Some(request.url), + headers: Some(self.response_headers.clone()), + body: None, + } + } +} + +#[async_trait] +impl HttpTransport for RejectingTransport { + async fn execute(&self, request: Request) -> Result { + Err(self.http_error(request)) + } + + async fn stream(&self, request: Request) -> Result { + Err(self.http_error(request)) + } +} + +#[tokio::test] +async fn execute_attaches_url_scoped_auth_and_observes_http_error_headers() { + let (session, transport, auth, request_headers, response_headers) = test_session(); + + session + .execute( + Method::POST, + "responses", + HeaderMap::new(), + /*body*/ None, + ) + .await + .expect_err("request should fail"); + + assert_request_and_response_headers(&transport, &auth, request_headers, response_headers); +} + +#[tokio::test] +async fn stream_attaches_url_scoped_auth_and_observes_http_error_headers() { + let (session, transport, auth, request_headers, response_headers) = test_session(); + + let result = session + .stream_with( + Method::POST, + "responses", + HeaderMap::new(), + /*body*/ None, + |_| {}, + ) + .await; + assert!(result.is_err(), "request should fail"); + + assert_request_and_response_headers(&transport, &auth, request_headers, response_headers); +} + +fn test_session() -> ( + EndpointSession, + RejectingTransport, + Arc, + HeaderMap, + HeaderMap, +) { + let mut request_headers = HeaderMap::new(); + request_headers.insert("x-test-state", HeaderValue::from_static("sent-state")); + let mut response_headers = HeaderMap::new(); + response_headers.insert( + "x-test-state-update", + HeaderValue::from_static("rotated-state"), + ); + let transport = RejectingTransport::new(response_headers.clone()); + let auth = Arc::new(RecordingAuthProvider::default()); + let session = EndpointSession::new(transport.clone(), provider(), auth.clone()); + (session, transport, auth, request_headers, response_headers) +} + +fn provider() -> Provider { + Provider { + name: "test".to_string(), + base_url: "https://chatgpt.com/backend-api/codex".to_string(), + query_params: None, + headers: HeaderMap::new(), + retry: RetryConfig { + max_attempts: 1, + base_delay: Duration::from_millis(1), + retry_429: false, + retry_5xx: false, + retry_transport: false, + }, + stream_idle_timeout: Duration::from_secs(1), + } +} + +fn assert_request_and_response_headers( + transport: &RejectingTransport, + auth: &RecordingAuthProvider, + request_headers: HeaderMap, + response_headers: HeaderMap, +) { + assert_eq!( + transport + .requests + .lock() + .expect("recording transport lock should not be poisoned") + .as_slice(), + &[(RESPONSES_URL.to_string(), request_headers.clone())] + ); + assert_eq!( + auth.observed_headers + .lock() + .expect("recording auth lock should not be poisoned") + .as_slice(), + &[(RESPONSES_URL.to_string(), request_headers, response_headers,)] + ); +} diff --git a/codex-rs/codex-api/src/files.rs b/codex-rs/codex-api/src/files.rs index d1e2840066..ddabfc4357 100644 --- a/codex-rs/codex-api/src/files.rs +++ b/codex-rs/codex-api/src/files.rs @@ -129,7 +129,9 @@ pub async fn upload_local_file( .unwrap_or("file") .to_string(); let create_url = format!("{}/files", base_url.trim_end_matches('/')); - let create_response = authorized_request(auth, reqwest::Method::POST, &create_url) + let (create_request, create_request_headers) = + authorized_request(auth, reqwest::Method::POST, &create_url); + let create_response = create_request .json(&serde_json::json!({ "file_name": file_name, "file_size": metadata.len(), @@ -141,6 +143,11 @@ pub async fn upload_local_file( url: create_url.clone(), source, })?; + auth.observe_response_headers( + &create_url, + &create_request_headers, + create_response.headers(), + ); let create_status = create_response.status(); let create_body = create_response.text().await.unwrap_or_default(); if !create_status.is_success() { @@ -191,7 +198,9 @@ pub async fn upload_local_file( ); let finalize_started_at = Instant::now(); loop { - let finalize_response = authorized_request(auth, reqwest::Method::POST, &finalize_url) + let (finalize_request, finalize_request_headers) = + authorized_request(auth, reqwest::Method::POST, &finalize_url); + let finalize_response = finalize_request .json(&serde_json::json!({})) .send() .await @@ -199,6 +208,11 @@ pub async fn upload_local_file( url: finalize_url.clone(), source, })?; + auth.observe_response_headers( + &finalize_url, + &finalize_request_headers, + finalize_response.headers(), + ); let finalize_status = finalize_response.status(); let finalize_body = finalize_response.text().await.unwrap_or_default(); if !finalize_status.is_success() { @@ -255,15 +269,18 @@ fn authorized_request( auth: &dyn AuthProvider, method: reqwest::Method, url: &str, -) -> reqwest::RequestBuilder { +) -> (reqwest::RequestBuilder, http::HeaderMap) { let mut headers = http::HeaderMap::new(); - auth.add_auth_headers(&mut headers); + auth.add_auth_headers_for_url(url, &mut headers); let client = build_reqwest_client(); - client - .request(method, url) - .timeout(OPENAI_FILE_REQUEST_TIMEOUT) - .headers(headers) + ( + client + .request(method, url) + .timeout(OPENAI_FILE_REQUEST_TIMEOUT) + .headers(headers.clone()), + headers, + ) } fn build_reqwest_client() -> reqwest::Client { @@ -279,6 +296,7 @@ mod tests { use pretty_assertions::assert_eq; use reqwest::header::HeaderValue; use std::sync::Arc; + use std::sync::Mutex; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use tempfile::TempDir; @@ -291,8 +309,16 @@ mod tests { use wiremock::matchers::method; use wiremock::matchers::path; - #[derive(Clone, Copy)] - struct ChatGptTestAuth; + const SENT_STATE: &str = "sent-state"; + const CREATED_STATE: &str = "created-state"; + const RETRY_STATE: &str = "retry-state"; + const FINALIZED_STATE: &str = "finalized-state"; + type ObservedUpdate = (String, Option); + + #[derive(Clone, Default)] + struct ChatGptTestAuth { + observed_updates: Arc>>, + } impl AuthProvider for ChatGptTestAuth { fn add_auth_headers(&self, headers: &mut reqwest::header::HeaderMap) { @@ -302,10 +328,33 @@ mod tests { ); headers.insert("ChatGPT-Account-ID", HeaderValue::from_static("account_id")); } + + fn add_auth_headers_for_url(&self, _request_url: &str, headers: &mut http::HeaderMap) { + self.add_auth_headers(headers); + headers.insert("x-test-state", HeaderValue::from_static(SENT_STATE)); + } + + fn observe_response_headers( + &self, + request_url: &str, + _request_headers: &http::HeaderMap, + response_headers: &http::HeaderMap, + ) { + self.observed_updates + .lock() + .expect("observed updates lock should not be poisoned") + .push(( + request_url.to_string(), + response_headers + .get("x-test-state-update") + .and_then(|value| value.to_str().ok()) + .map(ToString::to_string), + )); + } } fn chatgpt_auth() -> ChatGptTestAuth { - ChatGptTestAuth + ChatGptTestAuth::default() } fn base_url_for(server: &MockServer) -> String { @@ -318,6 +367,7 @@ mod tests { Mock::given(method("POST")) .and(path("/backend-api/files")) .and(header("chatgpt-account-id", "account_id")) + .and(header("x-test-state", SENT_STATE)) .and(body_json(serde_json::json!({ "file_name": "hello.txt", "file_size": 5, @@ -325,6 +375,7 @@ mod tests { }))) .respond_with( ResponseTemplate::new(200) + .insert_header("x-test-state-update", CREATED_STATE) .set_body_json(serde_json::json!({"file_id": "file_123", "upload_url": format!("{}/upload/file_123", server.uri())})), ) .mount(&server) @@ -340,20 +391,25 @@ mod tests { let download_url = format!("{}/download/file_123", server.uri()); Mock::given(method("POST")) .and(path("/backend-api/files/file_123/uploaded")) + .and(header("x-test-state", SENT_STATE)) .respond_with(move |_request: &Request| { if finalize_attempts_responder.fetch_add(1, Ordering::SeqCst) == 0 { - return ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "status": "retry" - })); + return ResponseTemplate::new(200) + .insert_header("x-test-state-update", RETRY_STATE) + .set_body_json(serde_json::json!({ + "status": "retry" + })); } - ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "status": "success", - "download_url": download_url, - "file_name": "hello.txt", - "mime_type": "text/plain", - "file_size_bytes": 5 - })) + ResponseTemplate::new(200) + .insert_header("x-test-state-update", FINALIZED_STATE) + .set_body_json(serde_json::json!({ + "status": "success", + "download_url": download_url, + "file_name": "hello.txt", + "mime_type": "text/plain", + "file_size_bytes": 5 + })) }) .mount(&server) .await; @@ -363,7 +419,8 @@ mod tests { let path = dir.path().join("hello.txt"); tokio::fs::write(&path, b"hello").await.expect("write file"); - let uploaded = upload_local_file(&base_url, &chatgpt_auth(), &path) + let auth = chatgpt_auth(); + let uploaded = upload_local_file(&base_url, &auth, &path) .await .expect("upload succeeds"); @@ -376,5 +433,25 @@ mod tests { assert_eq!(uploaded.file_name, "hello.txt"); assert_eq!(uploaded.mime_type, Some("text/plain".to_string())); assert_eq!(finalize_attempts.load(Ordering::SeqCst), 2); + assert_eq!( + *auth + .observed_updates + .lock() + .expect("observed updates lock should not be poisoned"), + vec![ + ( + format!("{}/files", base_url.trim_end_matches('/')), + Some(CREATED_STATE.to_string()), + ), + ( + format!("{}/files/file_123/uploaded", base_url.trim_end_matches('/')), + Some(RETRY_STATE.to_string()), + ), + ( + format!("{}/files/file_123/uploaded", base_url.trim_end_matches('/')), + Some(FINALIZED_STATE.to_string()), + ), + ] + ); } } diff --git a/codex-rs/codex-api/tests/realtime_websocket_e2e.rs b/codex-rs/codex-api/tests/realtime_websocket_e2e.rs index cb9d7122f4..5d29a7bfb7 100644 --- a/codex-rs/codex-api/tests/realtime_websocket_e2e.rs +++ b/codex-rs/codex-api/tests/realtime_websocket_e2e.rs @@ -1,7 +1,11 @@ use std::collections::HashMap; use std::future::Future; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::time::Duration; +use codex_api::AuthProvider; use codex_api::Provider; use codex_api::RealtimeAudioFrame; use codex_api::RealtimeEvent; @@ -19,11 +23,15 @@ use codex_protocol::protocol::RealtimeVoice; use futures::SinkExt; use futures::StreamExt; use http::HeaderMap; +use http::HeaderValue; use serde_json::Value; use serde_json::json; use tokio::net::TcpListener; use tokio_tungstenite::accept_async; +use tokio_tungstenite::accept_hdr_async; use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::tungstenite::handshake::server::Request; +use tokio_tungstenite::tungstenite::handshake::server::Response; type RealtimeWsStream = tokio_tungstenite::WebSocketStream; @@ -75,6 +83,37 @@ fn test_provider(base_url: String) -> Provider { } } +struct HandshakeAuthProvider { + observed_response: Arc, +} + +impl AuthProvider for HandshakeAuthProvider { + fn add_auth_headers(&self, _headers: &mut HeaderMap) {} + + fn add_auth_headers_for_url(&self, request_url: &str, headers: &mut HeaderMap) { + assert!(request_url.starts_with("ws://127.0.0.1:")); + headers.insert("x-test-auth-state", HeaderValue::from_static("sent-state")); + } + + fn observe_response_headers( + &self, + request_url: &str, + request_headers: &HeaderMap, + response_headers: &HeaderMap, + ) { + assert!(request_url.starts_with("ws://127.0.0.1:")); + assert_eq!( + request_headers.get("x-test-auth-state"), + Some(&HeaderValue::from_static("sent-state")), + ); + if response_headers.get("x-test-auth-state-update") + == Some(&HeaderValue::from_static("rotated-state")) + { + self.observed_response.store(true, Ordering::Relaxed); + } + } +} + #[tokio::test] async fn realtime_ws_e2e_session_create_and_event_flow() { let (addr, server) = spawn_realtime_ws_server(|mut ws: RealtimeWsStream| async move { @@ -202,6 +241,66 @@ async fn realtime_ws_e2e_session_create_and_event_flow() { server.await.expect("server task"); } +#[tokio::test] +async fn realtime_ws_e2e_applies_and_observes_url_aware_auth() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let addr = listener.local_addr().expect("local addr"); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept"); + let mut ws = accept_hdr_async(stream, |request: &Request, mut response: Response| { + assert_eq!( + request.headers().get("x-test-auth-state"), + Some(&HeaderValue::from_static("sent-state")), + ); + response.headers_mut().insert( + "x-test-auth-state-update", + HeaderValue::from_static("rotated-state"), + ); + Ok(response) + }) + .await + .expect("accept ws"); + let first = ws + .next() + .await + .expect("first msg") + .expect("first msg ok") + .into_text() + .expect("text"); + let first_json: Value = serde_json::from_str(&first).expect("json"); + assert_eq!(first_json["type"], "session.update"); + }); + + let observed_response = Arc::new(AtomicBool::new(false)); + let client = RealtimeWebsocketClient::new(test_provider(format!("http://{addr}"))).with_auth( + Arc::new(HandshakeAuthProvider { + observed_response: Arc::clone(&observed_response), + }), + ); + let connection = client + .connect( + RealtimeSessionConfig { + instructions: "backend prompt".to_string(), + model: Some("realtime-test-model".to_string()), + session_id: Some("conv_123".to_string()), + event_parser: RealtimeEventParser::V1, + session_mode: RealtimeSessionMode::Conversational, + output_modality: RealtimeOutputModality::Audio, + voice: RealtimeVoice::Cove, + }, + HeaderMap::new(), + HeaderMap::new(), + ) + .await + .expect("connect"); + + assert!(observed_response.load(Ordering::Relaxed)); + connection.close().await.expect("close"); + server.await.expect("server task"); +} + #[tokio::test] async fn realtime_ws_connect_webrtc_sideband_retries_join_until_server_is_available() { let reserving_listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); diff --git a/codex-rs/codex-client/src/chatgpt_hosts.rs b/codex-rs/codex-client/src/chatgpt_hosts.rs index dd0b99589c..7ab6819e79 100644 --- a/codex-rs/codex-client/src/chatgpt_hosts.rs +++ b/codex-rs/codex-client/src/chatgpt_hosts.rs @@ -10,6 +10,18 @@ pub fn is_allowed_chatgpt_host(host: &str) -> bool { .any(|suffix| host.ends_with(suffix)) } +/// Returns whether `url` is an HTTPS or secure WebSocket ChatGPT URL that Codex +/// may treat as first-party traffic. +pub fn is_allowed_chatgpt_request_url(url: &str) -> bool { + let Ok(url) = reqwest::Url::parse(url) else { + return false; + }; + if !matches!(url.scheme(), "https" | "wss") { + return false; + } + url.host_str().is_some_and(is_allowed_chatgpt_host) +} + #[cfg(test)] mod tests { use super::*; @@ -36,4 +48,25 @@ mod tests { assert!(!is_allowed_chatgpt_host(host)); } } + + #[test] + fn recognizes_secure_chatgpt_request_urls() { + for url in [ + "https://chatgpt.com/backend-api/codex/responses", + "https://preview.chatgpt.com/backend-api/codex/models", + "wss://chatgpt-staging.com/backend-api/codex/responses", + ] { + assert!(is_allowed_chatgpt_request_url(url)); + } + + for url in [ + "http://chatgpt.com/backend-api/codex/responses", + "ws://chatgpt.com/backend-api/codex/responses", + "https://api.openai.com/v1/responses", + "https://chatgpt.com.evil.example/backend-api", + "not a url", + ] { + assert!(!is_allowed_chatgpt_request_url(url)); + } + } } diff --git a/codex-rs/codex-client/src/default_client.rs b/codex-rs/codex-client/src/default_client.rs index 56b3ce4b16..a57b1dc338 100644 --- a/codex-rs/codex-client/src/default_client.rs +++ b/codex-rs/codex-client/src/default_client.rs @@ -1,3 +1,5 @@ +use crate::INTEGRITY_STATE_HEADER_NAME; +use crate::INTEGRITY_STATE_UPDATE_HEADER_NAME; use http::Error as HttpError; use http::HeaderMap; use http::HeaderName; @@ -115,11 +117,12 @@ impl CodexRequestBuilder { match self.builder.headers(headers).send().await { Ok(response) => { + let headers = redacted_headers_for_logging(response.headers()); tracing::debug!( method = %self.method, url = %self.url, status = %response.status(), - headers = ?response.headers(), + headers = ?headers, version = ?response.version(), "Request completed" ); @@ -141,6 +144,19 @@ impl CodexRequestBuilder { } } +fn redacted_headers_for_logging(headers: &HeaderMap) -> HeaderMap { + let mut headers = headers.clone(); + for name in [ + INTEGRITY_STATE_HEADER_NAME, + INTEGRITY_STATE_UPDATE_HEADER_NAME, + ] { + if headers.contains_key(name) { + headers.insert(name, HeaderValue::from_static("")); + } + } + headers +} + struct HeaderMapInjector<'a>(&'a mut HeaderMap); impl<'a> Injector for HeaderMapInjector<'a> { @@ -174,10 +190,39 @@ mod tests { use opentelemetry::trace::TracerProvider; use opentelemetry_sdk::propagation::TraceContextPropagator; use opentelemetry_sdk::trace::SdkTracerProvider; + use pretty_assertions::assert_eq; use tracing::trace_span; use tracing_subscriber::layer::SubscriberExt; use tracing_subscriber::util::SubscriberInitExt; + #[test] + fn redacted_headers_for_logging_hides_integrity_state() { + let mut headers = HeaderMap::new(); + headers.insert( + INTEGRITY_STATE_HEADER_NAME, + HeaderValue::from_static("ois1.request.nonce.ciphertext"), + ); + headers.insert( + INTEGRITY_STATE_UPDATE_HEADER_NAME, + HeaderValue::from_static("ois1.response.nonce.ciphertext"), + ); + headers.insert("x-request-id", HeaderValue::from_static("request-id")); + + let redacted = redacted_headers_for_logging(&headers); + + let mut expected = HeaderMap::new(); + expected.insert( + INTEGRITY_STATE_HEADER_NAME, + HeaderValue::from_static(""), + ); + expected.insert( + INTEGRITY_STATE_UPDATE_HEADER_NAME, + HeaderValue::from_static(""), + ); + expected.insert("x-request-id", HeaderValue::from_static("request-id")); + assert_eq!(redacted, expected); + } + #[test] fn inject_trace_headers_uses_current_span_context() { global::set_text_map_propagator(TraceContextPropagator::new()); diff --git a/codex-rs/codex-client/src/lib.rs b/codex-rs/codex-client/src/lib.rs index 0f503fb3e2..a8eea88922 100644 --- a/codex-rs/codex-client/src/lib.rs +++ b/codex-rs/codex-client/src/lib.rs @@ -11,6 +11,7 @@ mod transport; pub use crate::chatgpt_cloudflare_cookies::with_chatgpt_cloudflare_cookie_store; pub use crate::chatgpt_hosts::is_allowed_chatgpt_host; +pub use crate::chatgpt_hosts::is_allowed_chatgpt_request_url; pub use crate::custom_ca::BuildCustomCaTransportError; /// Test-only subprocess hook for custom CA coverage. /// @@ -40,3 +41,6 @@ pub use crate::transport::ByteStream; pub use crate::transport::HttpTransport; pub use crate::transport::ReqwestTransport; pub use crate::transport::StreamResponse; + +pub const INTEGRITY_STATE_HEADER_NAME: &str = "x-oai-is"; +pub const INTEGRITY_STATE_UPDATE_HEADER_NAME: &str = "x-oai-is-update"; diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 97e104464f..ff8b2306d2 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -48,6 +48,7 @@ codex-shell-command = { workspace = true } codex-execpolicy = { workspace = true } codex-git-utils = { workspace = true } codex-hooks = { workspace = true } +codex-http-state = { workspace = true } codex-install-context = { workspace = true } codex-network-proxy = { workspace = true } codex-otel = { workspace = true } diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 821611e60e..36fdb91a86 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -307,6 +307,21 @@ impl AgentControl { { let client_metadata = match state.get_thread(*parent_thread_id).await { Ok(parent_thread) => { + if let Some(surface) = parent_thread + .codex + .session + .services + .model_client + .http_state_surface() + { + new_thread + .thread + .codex + .session + .services + .model_client + .set_http_state_surface(surface); + } parent_thread .codex .session diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index f9462e94f6..3971434c08 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -11,6 +11,7 @@ use crate::context::SubagentNotification; use crate::init_state_db; use assert_matches::assert_matches; use codex_features::Feature; +use codex_http_state::HttpStateSurface; use codex_login::CodexAuth; use codex_protocol::AgentPath; use codex_protocol::config_types::ModeKind; @@ -1584,6 +1585,51 @@ async fn spawn_child_completion_notifies_parent_history() { assert_eq!(wait_for_subagent_notification(&parent_thread).await, true); } +#[tokio::test] +async fn spawn_thread_subagent_inherits_parent_http_state_surface() { + let harness = AgentControlHarness::new().await; + let (parent_thread_id, parent_thread) = harness.start_thread().await; + parent_thread + .set_app_server_client_info( + Some("codex_desktop".to_string()), + /*app_server_client_version*/ None, + /*mcp_elicitations_auto_deny*/ false, + ) + .await + .expect("parent client info should update"); + + let child_thread_id = harness + .control + .spawn_agent( + harness.config.clone(), + text_input("hello child"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: Some("explorer".to_string()), + })), + ) + .await + .expect("child spawn should succeed"); + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should exist"); + + assert_eq!( + child_thread + .codex + .session + .services + .model_client + .http_state_surface(), + Some(HttpStateSurface::CodexDesktop) + ); +} + #[tokio::test] async fn multi_agent_v2_completion_ignores_dead_direct_parent() { let harness = AgentControlHarness::new().await; diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 124dfb18c0..dc112747ff 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -62,6 +62,8 @@ use codex_api::build_session_headers; use codex_api::create_text_param_for_request; use codex_api::response_create_client_metadata; use codex_app_server_protocol::AuthMode; +use codex_http_state::HttpStateContext; +use codex_http_state::HttpStateSurface; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::RefreshTokenError; @@ -120,6 +122,7 @@ use codex_login::auth_env_telemetry::AuthEnvTelemetry; use codex_login::auth_env_telemetry::collect_auth_env_telemetry; use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; +use codex_model_provider::with_native_integrity_state; #[cfg(test)] use codex_model_provider_info::DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS; use codex_model_provider_info::ModelProviderInfo; @@ -180,7 +183,9 @@ struct ModelClientState { beta_features_header: Option, include_attestation: bool, attestation_provider: Option>, + http_state: Option, disable_websockets: AtomicBool, + websocket_session_generation: AtomicU64, cached_websocket_session: StdMutex, } @@ -237,6 +242,7 @@ pub struct ModelClient { /// contract and can cause routing bugs. pub struct ModelClientSession { client: ModelClient, + websocket_session_generation: u64, websocket_session: WebsocketSession, /// Turn state for sticky routing. /// @@ -296,17 +302,7 @@ pub(crate) struct RealtimeWebrtcCallStart { pub(crate) sdp: String, pub(crate) call_id: String, pub(crate) sideband_headers: ApiHeaderMap, -} - -/// Reuses the API-auth material that created the WebRTC call for the sideband WebSocket join. -/// -/// API-key sessions send that API bearer. ChatGPT-auth sessions send their bearer plus account id; -/// transceiver is responsible for accepting that same call-create identity on the direct -/// `api.openai.com` sideband path. -fn sideband_websocket_auth_headers(api_auth: &dyn AuthProvider) -> ApiHeaderMap { - let mut headers = ApiHeaderMap::new(); - api_auth.add_auth_headers(&mut headers); - headers + pub(crate) sideband_auth: SharedAuthProvider, } impl ModelClient { @@ -329,6 +325,12 @@ impl ModelClient { beta_features_header: Option, attestation_provider: Option>, ) -> Self { + let http_state = auth_manager.as_ref().map(|auth_manager| { + HttpStateContext::new( + auth_manager.codex_home().to_path_buf(), + http_state_surface_for_session_source(&session_source), + ) + }); let model_provider = create_model_provider(provider_info, auth_manager); let codex_api_key_env_enabled = model_provider .auth_manager() @@ -353,7 +355,9 @@ impl ModelClient { beta_features_header, include_attestation, attestation_provider, + http_state, disable_websockets: AtomicBool::new(false), + websocket_session_generation: AtomicU64::new(0), cached_websocket_session: StdMutex::new(WebsocketSession::default()), }), prompt_cache_key_override: None, @@ -379,9 +383,12 @@ impl ModelClient { /// This constructor does not perform network I/O itself; the session opens a websocket lazily /// when the first stream request is issued. pub fn new_session(&self) -> ModelClientSession { + let (websocket_session_generation, websocket_session) = + self.take_cached_websocket_session(); ModelClientSession { client: self.clone(), - websocket_session: self.take_cached_websocket_session(), + websocket_session_generation, + websocket_session, turn_state: Arc::new(OnceLock::new()), } } @@ -390,16 +397,45 @@ impl ModelClient { self.state.provider.auth_manager() } + pub(crate) fn set_app_server_client_name(&self, client_name: Option<&str>) { + let surface = client_name + .map(HttpStateSurface::from_app_server_client_name) + .unwrap_or(HttpStateSurface::CodexCli); + self.set_http_state_surface(surface); + } + + pub(crate) fn http_state_surface(&self) -> Option { + self.state + .http_state + .as_ref() + .map(HttpStateContext::surface) + } + + pub(crate) fn http_state_context(&self) -> Option { + self.state.http_state.clone() + } + + pub(crate) fn set_http_state_surface(&self, surface: HttpStateSurface) { + if self + .state + .http_state + .as_ref() + .is_some_and(|state| state.set_surface(surface)) + { + self.invalidate_cached_websocket_session(); + } + } + pub(crate) fn set_window_generation(&self, window_generation: u64) { self.state .window_generation .store(window_generation, Ordering::Relaxed); - self.store_cached_websocket_session(WebsocketSession::default()); + self.invalidate_cached_websocket_session(); } pub(crate) fn advance_window_generation(&self) { self.state.window_generation.fetch_add(1, Ordering::Relaxed); - self.store_cached_websocket_session(WebsocketSession::default()); + self.invalidate_cached_websocket_session(); } pub(crate) fn current_window_id(&self) -> String { @@ -408,21 +444,45 @@ impl ModelClient { format!("{thread_id}:{window_generation}") } - fn take_cached_websocket_session(&self) -> WebsocketSession { + fn take_cached_websocket_session(&self) -> (u64, WebsocketSession) { let mut cached_websocket_session = self .state .cached_websocket_session .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - std::mem::take(&mut *cached_websocket_session) + let generation = self + .state + .websocket_session_generation + .load(Ordering::Relaxed); + (generation, std::mem::take(&mut *cached_websocket_session)) } - fn store_cached_websocket_session(&self, websocket_session: WebsocketSession) { - *self + fn invalidate_cached_websocket_session(&self) { + let mut cached_websocket_session = self .state .cached_websocket_session .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) = websocket_session; + .unwrap_or_else(std::sync::PoisonError::into_inner); + self.state + .websocket_session_generation + .fetch_add(1, Ordering::Relaxed); + *cached_websocket_session = WebsocketSession::default(); + } + + fn store_cached_websocket_session(&self, generation: u64, websocket_session: WebsocketSession) { + let mut cached_websocket_session = self + .state + .cached_websocket_session + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if self + .state + .websocket_session_generation + .load(Ordering::Relaxed) + == generation + { + *cached_websocket_session = websocket_session; + } } pub(crate) fn force_http_fallback( @@ -442,7 +502,7 @@ impl ModelClient { ); } - self.store_cached_websocket_session(WebsocketSession::default()); + self.invalidate_cached_websocket_session(); activated } @@ -554,10 +614,8 @@ impl ModelClient { if let Some(header_value) = self.generate_attestation_header_for().await { extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value); } - let mut sideband_headers = extra_headers.clone(); - sideband_headers.extend(sideband_websocket_auth_headers( - client_setup.api_auth.as_ref(), - )); + let sideband_headers = extra_headers.clone(); + let sideband_auth = Arc::clone(&client_setup.api_auth); let transport = ReqwestTransport::new(build_reqwest_client()); let response = ApiRealtimeCallClient::new(transport, client_setup.api_provider, client_setup.api_auth) @@ -568,6 +626,7 @@ impl ModelClient { sdp: response.sdp, call_id: response.call_id, sideband_headers, + sideband_auth, }) } @@ -812,7 +871,11 @@ impl ModelClient { async fn current_client_setup(&self) -> Result { let auth = self.state.provider.auth().await; let api_provider = self.state.provider.api_provider().await?; - let api_auth = self.state.provider.api_auth().await?; + let api_auth = with_native_integrity_state( + self.state.provider.api_auth().await?, + auth.as_ref(), + self.state.http_state.clone(), + ); Ok(CurrentClientSetup { auth, api_provider, @@ -953,7 +1016,7 @@ impl Drop for ModelClientSession { fn drop(&mut self) { let websocket_session = std::mem::take(&mut self.websocket_session); self.client - .store_cached_websocket_session(websocket_session); + .store_cached_websocket_session(self.websocket_session_generation, websocket_session); } } @@ -1736,6 +1799,19 @@ fn subagent_header_value(session_source: &SessionSource) -> Option { } } +fn http_state_surface_for_session_source(session_source: &SessionSource) -> HttpStateSurface { + match session_source { + SessionSource::Cli => HttpStateSurface::CodexTui, + SessionSource::Exec => HttpStateSurface::CodexExec, + SessionSource::VSCode => HttpStateSurface::CodexVscode, + SessionSource::SubAgent(_) + | SessionSource::Internal(_) + | SessionSource::Mcp + | SessionSource::Custom(_) + | SessionSource::Unknown => HttpStateSurface::CodexCli, + } +} + fn parent_thread_id_header_value(parent_thread_id: Option) -> Option { parent_thread_id.map(|parent_thread_id| parent_thread_id.to_string()) } diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index f8d904036a..419b28f2d2 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -13,12 +13,15 @@ use crate::GenerateAttestationFuture; use codex_api::ApiError; use codex_api::ResponseEvent; use codex_app_server_protocol::AuthMode; +use codex_http_state::HttpStateStore; +use codex_http_state::HttpStateSurface; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider::BearerAuthProvider; use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::WireApi; +use codex_model_provider_info::built_in_model_providers; use codex_model_provider_info::create_oss_provider_with_base_url; use codex_otel::SessionTelemetry; use codex_protocol::SessionId; @@ -67,11 +70,32 @@ fn test_model_client(session_source: SessionSource) -> ModelClient { fn test_model_client_with_parent( session_source: SessionSource, parent_thread_id: Option, +) -> ModelClient { + test_model_client_with_parent_and_auth( + session_source, + parent_thread_id, + /*auth_manager*/ None, + ) +} + +fn test_model_client_with_parent_and_auth( + session_source: SessionSource, + parent_thread_id: Option, + auth_manager: Option>, ) -> ModelClient { let provider = create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses); + test_model_client_with_provider(session_source, parent_thread_id, auth_manager, provider) +} + +fn test_model_client_with_provider( + session_source: SessionSource, + parent_thread_id: Option, + auth_manager: Option>, + provider: ModelProviderInfo, +) -> ModelClient { let thread_id = ThreadId::new(); ModelClient::new( - /*auth_manager*/ None, + auth_manager, thread_id.into(), thread_id, /*installation_id*/ "11111111-1111-4111-8111-111111111111".to_string(), @@ -322,6 +346,77 @@ fn build_ws_client_metadata_includes_window_lineage_and_turn_metadata() { ); } +#[tokio::test] +async fn current_client_setup_attaches_and_rotates_native_http_state() { + let codex_home = TempDir::new().expect("tempdir"); + let store = HttpStateStore::new(codex_home.path().to_path_buf()); + store + .set(HttpStateSurface::CodexTui, "old-state".to_string()) + .expect("state should store"); + let auth_manager = AuthManager::from_auth_for_testing_with_home( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + codex_home.path().to_path_buf(), + ); + let client = test_model_client_with_provider( + SessionSource::Cli, + /*parent_thread_id*/ None, + Some(auth_manager), + built_in_model_providers(/*openai_base_url*/ None)["openai"].clone(), + ); + let setup = client + .current_client_setup() + .await + .expect("client setup should resolve"); + let request_url = "https://chatgpt.com/backend-api/codex/responses"; + let mut request_headers = http::HeaderMap::new(); + setup + .api_auth + .add_auth_headers_for_url(request_url, &mut request_headers); + assert_eq!( + request_headers + .get("x-oai-is") + .and_then(|value| value.to_str().ok()), + Some("old-state"), + ); + + let mut response_headers = http::HeaderMap::new(); + response_headers.insert( + "x-oai-is-update", + http::HeaderValue::from_static("new-state"), + ); + setup + .api_auth + .observe_response_headers(request_url, &request_headers, &response_headers); + assert_eq!( + store + .get(HttpStateSurface::CodexTui) + .expect("state should load"), + Some("new-state".to_string()), + ); +} + +#[test] +fn surface_change_discards_checked_out_websocket_session() { + let codex_home = TempDir::new().expect("tempdir"); + let auth_manager = AuthManager::from_auth_for_testing_with_home( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + codex_home.path().to_path_buf(), + ); + let client = test_model_client_with_parent_and_auth( + SessionSource::Cli, + /*parent_thread_id*/ None, + Some(auth_manager), + ); + let mut session = client.new_session(); + session.websocket_session.last_response_from_untraced_warmup = true; + + client.set_http_state_surface(HttpStateSurface::CodexDesktop); + drop(session); + + let (_, cached_session) = client.take_cached_websocket_session(); + assert!(!cached_session.last_response_from_untraced_warmup); +} + #[tokio::test] async fn summarize_memories_returns_empty_for_empty_input() { let client = test_model_client(SessionSource::Cli); diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 678a6e5d10..88f85b0b11 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -110,6 +110,13 @@ pub(crate) async fn run_codex_thread_interactive( })) .or_cancel(&cancel_token) .await??; + if let Some(surface) = parent_session.services.model_client.http_state_surface() { + codex + .session + .services + .model_client + .set_http_state_surface(surface); + } let thread_config = codex.thread_config_snapshot().await; let client_metadata = parent_session.app_server_client_metadata().await; emit_subagent_session_started( diff --git a/codex-rs/core/src/codex_delegate_tests.rs b/codex-rs/core/src/codex_delegate_tests.rs index 7dabb09d33..0cdb91a06b 100644 --- a/codex-rs/core/src/codex_delegate_tests.rs +++ b/codex-rs/core/src/codex_delegate_tests.rs @@ -2,6 +2,7 @@ use super::*; use crate::mcp_tool_call::MCP_TOOL_APPROVAL_DECLINE_SYNTHETIC; use crate::mcp_tool_call::MCP_TOOL_APPROVAL_QUESTION_ID_PREFIX; use async_channel::bounded; +use codex_http_state::HttpStateSurface; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::models::NetworkPermissions; @@ -181,6 +182,36 @@ async fn run_codex_thread_interactive_respects_pre_cancelled_spawn() { assert!(matches!(result, Err(CodexErr::TurnAborted))); } +#[tokio::test] +async fn run_codex_thread_interactive_inherits_parent_http_state_surface() { + let (parent_session, parent_ctx, _rx_events) = + crate::session::tests::make_session_and_context_with_rx().await; + parent_session + .services + .model_client + .set_http_state_surface(HttpStateSurface::CodexDesktop); + let cancel_token = CancellationToken::new(); + + let delegated = run_codex_thread_interactive( + parent_ctx.config.as_ref().clone(), + Arc::clone(&parent_session.services.auth_manager), + Arc::clone(&parent_session.services.models_manager), + parent_session, + parent_ctx, + cancel_token.clone(), + SubAgentSource::Review, + /*initial_history*/ None, + ) + .await + .expect("delegate should spawn"); + + assert_eq!( + delegated.session.services.model_client.http_state_surface(), + Some(HttpStateSurface::CodexDesktop) + ); + cancel_token.cancel(); +} + #[tokio::test] async fn handle_request_permissions_uses_tool_call_id_for_round_trip() { let (parent_session, mut parent_ctx, rx_events) = diff --git a/codex-rs/core/src/mcp_openai_file.rs b/codex-rs/core/src/mcp_openai_file.rs index ae44515c62..821cab7de9 100644 --- a/codex-rs/core/src/mcp_openai_file.rs +++ b/codex-rs/core/src/mcp_openai_file.rs @@ -13,6 +13,7 @@ use crate::session::session::Session; use crate::session::turn_context::TurnContext; use codex_api::upload_local_file; +use codex_http_state::HttpStateContext; use codex_login::CodexAuth; use serde_json::Value as JsonValue; @@ -33,15 +34,21 @@ pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files( return Ok(Some(arguments_value)); }; let auth = sess.services.auth_manager.auth().await; + let http_state = sess.services.model_client.http_state_context(); let mut rewritten_arguments = arguments.clone(); for field_name in openai_file_input_params { let Some(value) = arguments.get(field_name) else { continue; }; - let Some(uploaded_value) = - rewrite_argument_value_for_openai_files(turn_context, auth.as_ref(), field_name, value) - .await? + let Some(uploaded_value) = rewrite_argument_value_for_openai_files( + turn_context, + auth.as_ref(), + http_state.clone(), + field_name, + value, + ) + .await? else { continue; }; @@ -58,6 +65,7 @@ pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files( async fn rewrite_argument_value_for_openai_files( turn_context: &TurnContext, auth: Option<&CodexAuth>, + http_state: Option, field_name: &str, value: &JsonValue, ) -> Result, String> { @@ -66,6 +74,7 @@ async fn rewrite_argument_value_for_openai_files( let rewritten = build_uploaded_local_argument_value( turn_context, auth, + http_state, field_name, /*index*/ None, path_or_file_ref, @@ -82,6 +91,7 @@ async fn rewrite_argument_value_for_openai_files( let rewritten = build_uploaded_local_argument_value( turn_context, auth, + http_state.clone(), field_name, Some(index), path_or_file_ref, @@ -98,6 +108,7 @@ async fn rewrite_argument_value_for_openai_files( async fn build_uploaded_local_argument_value( turn_context: &TurnContext, auth: Option<&CodexAuth>, + http_state: Option, field_name: &str, index: Option, file_path: &str, @@ -114,7 +125,11 @@ async fn build_uploaded_local_argument_value( "ChatGPT auth is required to upload local files for Codex Apps tools".to_string(), ); } - let upload_auth = codex_model_provider::auth_provider_from_auth(auth); + let upload_auth = codex_model_provider::with_native_integrity_state( + codex_model_provider::auth_provider_from_auth(auth), + Some(auth), + http_state, + ); let uploaded = upload_local_file( turn_context.config.chatgpt_base_url.trim_end_matches('/'), upload_auth.as_ref(), @@ -229,6 +244,7 @@ mod tests { let rewritten = build_uploaded_local_argument_value( &turn_context, Some(&auth), + /*http_state*/ None, "file", /*index*/ None, "file_report.csv", @@ -312,6 +328,7 @@ mod tests { let rewritten = rewrite_argument_value_for_openai_files( &turn_context, Some(&auth), + /*http_state*/ None, "file", &serde_json::json!("file_report.csv"), ) @@ -429,6 +446,7 @@ mod tests { let rewritten = rewrite_argument_value_for_openai_files( &turn_context, Some(&auth), + /*http_state*/ None, "files", &serde_json::json!(["one.csv", "two.csv"]), ) diff --git a/codex-rs/core/src/realtime_conversation.rs b/codex-rs/core/src/realtime_conversation.rs index 7f71142e19..1040d9e2d7 100644 --- a/codex-rs/core/src/realtime_conversation.rs +++ b/codex-rs/core/src/realtime_conversation.rs @@ -19,6 +19,7 @@ use codex_api::RealtimeSessionMode; use codex_api::RealtimeWebsocketClient; use codex_api::RealtimeWebsocketEvents; use codex_api::RealtimeWebsocketWriter; +use codex_api::SharedAuthProvider; use codex_api::map_api_error; use codex_app_server_protocol::AuthMode; use codex_config::config_toml::RealtimeWsMode; @@ -320,6 +321,7 @@ impl RealtimeConversationManager { session_config, call_id: call.call_id, sideband_headers: call.sideband_headers, + sideband_auth: call.sideband_auth, input_channels, events_tx, handoff_state: handoff.clone(), @@ -1029,6 +1031,7 @@ struct RealtimeWebrtcSidebandInputTask { session_config: RealtimeSessionConfig, call_id: String, sideband_headers: HeaderMap, + sideband_auth: SharedAuthProvider, input_channels: RealtimeInputChannels, events_tx: Sender, handoff_state: RealtimeHandoffState, @@ -1043,6 +1046,7 @@ fn spawn_webrtc_sideband_input_task(input: RealtimeWebrtcSidebandInputTask) -> J session_config, call_id, sideband_headers, + sideband_auth, input_channels, events_tx, handoff_state, @@ -1057,6 +1061,7 @@ fn spawn_webrtc_sideband_input_task(input: RealtimeWebrtcSidebandInputTask) -> J } let connection = match client + .with_auth(sideband_auth) .connect_webrtc_sideband( session_config, &call_id, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 5a92c9c825..d536669943 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -782,6 +782,10 @@ impl Codex { app_server_client_version: Option, mcp_elicitations_auto_deny: bool, ) -> ConstraintResult<()> { + self.session + .services + .model_client + .set_app_server_client_name(app_server_client_name.as_deref()); self.session .update_settings(SessionSettingsUpdate { app_server_client_name, diff --git a/codex-rs/http-state/src/lib.rs b/codex-rs/http-state/src/lib.rs index 61b69fbbc0..c1d6ee3be8 100644 --- a/codex-rs/http-state/src/lib.rs +++ b/codex-rs/http-state/src/lib.rs @@ -241,6 +241,10 @@ impl HttpStateContext { self.store.get(surface) } + pub fn set_for_surface(&self, surface: HttpStateSurface, state: String) -> io::Result<()> { + self.store.set(surface, state) + } + pub fn compare_and_set_for_surface( &self, surface: HttpStateSurface, diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 22e2c9f3b9..e6a402626d 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1401,6 +1401,10 @@ impl AuthManager { }) } + pub fn codex_home(&self) -> &Path { + &self.codex_home + } + /// Current cached auth (clone) without attempting a refresh. pub fn auth_cached(&self) -> Option { self.inner.read().ok().and_then(|c| c.auth.clone()) diff --git a/codex-rs/model-provider/Cargo.toml b/codex-rs/model-provider/Cargo.toml index 58235ab24d..965b118431 100644 --- a/codex-rs/model-provider/Cargo.toml +++ b/codex-rs/model-provider/Cargo.toml @@ -19,6 +19,7 @@ codex-agent-identity = { workspace = true } codex-aws-auth = { workspace = true } codex-client = { workspace = true } codex-feedback = { workspace = true } +codex-http-state = { workspace = true } codex-login = { workspace = true } codex-model-provider-info = { workspace = true } codex-models-manager = { workspace = true } @@ -32,5 +33,6 @@ tracing = { workspace = true, features = ["log"] } [dev-dependencies] pretty_assertions = { workspace = true } serde_json = { workspace = true } +tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt"] } wiremock = { workspace = true } diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index 3c7f4dbd0f..ddf1c0bbcc 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -5,6 +5,7 @@ use codex_agent_identity::AgentTaskAuthorizationTarget; use codex_agent_identity::authorization_header_for_agent_task; use codex_api::AuthProvider; use codex_api::SharedAuthProvider; +use codex_http_state::HttpStateContext; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderInfo; @@ -12,6 +13,7 @@ use http::HeaderMap; use http::HeaderValue; use crate::bearer_auth_provider::BearerAuthProvider; +use crate::bearer_auth_provider::NativeIntegrityAuthProvider; #[derive(Clone, Debug)] struct AgentIdentityAuthProvider { @@ -119,6 +121,23 @@ pub fn auth_provider_from_auth(auth: &CodexAuth) -> SharedAuthProvider { } } +/// Decorates first-party ChatGPT auth with native integrity-state transport. +pub fn with_native_integrity_state( + auth_provider: SharedAuthProvider, + auth: Option<&CodexAuth>, + state: Option, +) -> SharedAuthProvider { + let Some(state) = state else { + return auth_provider; + }; + let Some(CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) | CodexAuth::AgentIdentity(_)) = + auth + else { + return auth_provider; + }; + Arc::new(NativeIntegrityAuthProvider::new(auth_provider, state)) +} + #[cfg(test)] mod tests { use codex_model_provider_info::WireApi; diff --git a/codex-rs/model-provider/src/bearer_auth_provider.rs b/codex-rs/model-provider/src/bearer_auth_provider.rs index a28e069227..83ccb22e78 100644 --- a/codex-rs/model-provider/src/bearer_auth_provider.rs +++ b/codex-rs/model-provider/src/bearer_auth_provider.rs @@ -1,7 +1,15 @@ use codex_api::AuthProvider; +use codex_api::SharedAuthProvider; +use codex_client::INTEGRITY_STATE_HEADER_NAME; +use codex_client::INTEGRITY_STATE_UPDATE_HEADER_NAME; +use codex_client::is_allowed_chatgpt_request_url; +use codex_http_state::HttpStateContext; +use codex_http_state::HttpStateSurface; use http::HeaderMap; use http::HeaderValue; +const OAI_CLIENT_TYPE_HEADER_NAME: &str = "oai-client-type"; + /// Bearer-token auth provider for OpenAI-compatible model-provider requests. #[derive(Clone, Default)] pub struct BearerAuthProvider { @@ -46,10 +54,97 @@ impl AuthProvider for BearerAuthProvider { } } +#[derive(Clone)] +pub(crate) struct NativeIntegrityAuthProvider { + auth: SharedAuthProvider, + state: HttpStateContext, + surface: HttpStateSurface, +} + +impl NativeIntegrityAuthProvider { + pub(crate) fn new(auth: SharedAuthProvider, state: HttpStateContext) -> Self { + let surface = state.surface(); + Self { + auth, + state, + surface, + } + } +} + +impl AuthProvider for NativeIntegrityAuthProvider { + fn add_auth_headers(&self, headers: &mut HeaderMap) { + self.auth.add_auth_headers(headers); + } + + fn add_auth_headers_for_url(&self, request_url: &str, headers: &mut HeaderMap) { + self.auth.add_auth_headers_for_url(request_url, headers); + if !is_allowed_chatgpt_request_url(request_url) { + return; + } + + if let Ok(header) = HeaderValue::from_str(self.surface.as_str()) { + let _ = headers.insert(OAI_CLIENT_TYPE_HEADER_NAME, header); + } + match self.state.get_for_surface(self.surface) { + Ok(Some(state)) => { + if let Ok(header) = HeaderValue::from_str(&state) { + let _ = headers.insert(INTEGRITY_STATE_HEADER_NAME, header); + } + } + Ok(None) => {} + Err(error) => { + tracing::warn!("failed to load native integrity state: {error}"); + } + } + } + + fn observe_response_headers( + &self, + request_url: &str, + request_headers: &HeaderMap, + response_headers: &HeaderMap, + ) { + self.auth + .observe_response_headers(request_url, request_headers, response_headers); + if !is_allowed_chatgpt_request_url(request_url) { + return; + } + + let Some(next_state) = response_headers + .get(INTEGRITY_STATE_UPDATE_HEADER_NAME) + .and_then(|value| value.to_str().ok()) + else { + return; + }; + + let expected_state = request_headers + .get(INTEGRITY_STATE_HEADER_NAME) + .and_then(|value| value.to_str().ok()); + let result = match expected_state { + Some(expected_state) => self.state.compare_and_set_for_surface( + self.surface, + expected_state, + next_state.to_string(), + ), + None => self + .state + .set_for_surface(self.surface, next_state.to_string()) + .map(|()| true), + }; + if let Err(error) = result { + tracing::warn!("failed to rotate native integrity state: {error}"); + } + } +} + #[cfg(test)] mod tests { use super::*; + use codex_http_state::HttpStateStore; use pretty_assertions::assert_eq; + use std::sync::Arc; + use tempfile::TempDir; #[test] fn bearer_auth_provider_reports_when_auth_header_will_attach() { @@ -107,4 +202,115 @@ mod tests { Some("true") ); } + + #[test] + fn native_integrity_provider_scopes_and_rotates_surface_state() { + let codex_home = TempDir::new().expect("tempdir"); + let context = HttpStateContext::new( + codex_home.path().to_path_buf(), + HttpStateSurface::CodexDesktop, + ); + let store = HttpStateStore::new(codex_home.path().to_path_buf()); + store + .set( + HttpStateSurface::CodexDesktop, + "ois1.initial.nonce.ciphertext".to_string(), + ) + .expect("state should store"); + let provider = NativeIntegrityAuthProvider::new( + Arc::new(BearerAuthProvider::new("access-token".to_string())), + context.clone(), + ); + context.set_surface(HttpStateSurface::CodexCli); + store + .set( + HttpStateSurface::CodexCli, + "ois1.cli.nonce.ciphertext".to_string(), + ) + .expect("CLI state should store"); + + let mut request_headers = HeaderMap::new(); + provider.add_auth_headers_for_url( + "https://chatgpt.com/backend-api/codex/responses", + &mut request_headers, + ); + assert_eq!( + request_headers + .get(INTEGRITY_STATE_HEADER_NAME) + .and_then(|value| value.to_str().ok()), + Some("ois1.initial.nonce.ciphertext") + ); + assert_eq!( + request_headers + .get(OAI_CLIENT_TYPE_HEADER_NAME) + .and_then(|value| value.to_str().ok()), + Some("codex_desktop") + ); + + let mut response_headers = HeaderMap::new(); + response_headers.insert( + INTEGRITY_STATE_UPDATE_HEADER_NAME, + HeaderValue::from_static("ois1.rotated.nonce.ciphertext"), + ); + provider.observe_response_headers( + "https://chatgpt.com/backend-api/codex/responses", + &request_headers, + &response_headers, + ); + assert_eq!( + store + .get(HttpStateSurface::CodexDesktop) + .expect("state should load") + .expect("state should exist"), + "ois1.rotated.nonce.ciphertext" + ); + assert_eq!( + store + .get(HttpStateSurface::CodexCli) + .expect("CLI state should load") + .expect("CLI state should exist"), + "ois1.cli.nonce.ciphertext" + ); + + let mut external_headers = HeaderMap::new(); + provider.add_auth_headers_for_url( + "https://example.com/backend-api/codex/responses", + &mut external_headers, + ); + assert!(!external_headers.contains_key(INTEGRITY_STATE_HEADER_NAME)); + assert!(!external_headers.contains_key(OAI_CLIENT_TYPE_HEADER_NAME)); + } + + #[test] + fn native_integrity_provider_seeds_empty_surface_from_response() { + let codex_home = TempDir::new().expect("tempdir"); + let context = HttpStateContext::new( + codex_home.path().to_path_buf(), + HttpStateSurface::CodexDesktop, + ); + let store = HttpStateStore::new(codex_home.path().to_path_buf()); + let provider = NativeIntegrityAuthProvider::new( + Arc::new(BearerAuthProvider::new("access-token".to_string())), + context, + ); + let mut response_headers = HeaderMap::new(); + response_headers.insert( + INTEGRITY_STATE_UPDATE_HEADER_NAME, + HeaderValue::from_static("ois1.seeded.nonce.ciphertext"), + ); + + provider.observe_response_headers( + "https://chatgpt.com/backend-api/codex/responses", + &HeaderMap::new(), + &response_headers, + ); + + assert_eq!( + store + .get(HttpStateSurface::CodexDesktop) + .expect("state should load") + .expect("state should exist"), + "ois1.seeded.nonce.ciphertext" + ); + } } diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index 4e4660812b..70be6fb2b3 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -6,6 +6,7 @@ mod provider; pub use auth::auth_provider_from_auth; pub use auth::unauthenticated_auth_provider; +pub use auth::with_native_integrity_state; pub use bearer_auth_provider::BearerAuthProvider; pub use bearer_auth_provider::BearerAuthProvider as CoreAuthProvider; pub use codex_protocol::account::ProviderAccount;