diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index ccd49da649..3b2895ad8e 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2529,6 +2529,7 @@ dependencies = [ "codex-app-server-protocol", "codex-apply-patch", "codex-async-utils", + "codex-client", "codex-code-mode", "codex-config", "codex-connectors", @@ -3322,6 +3323,7 @@ dependencies = [ "http 1.4.0", "pretty_assertions", "serde_json", + "tempfile", "tokio", "tracing", "wiremock", diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 4f6ab2e777..88cd61ce37 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -26,6 +26,7 @@ chrono = { workspace = true, features = ["serde"] } clap = { workspace = true, features = ["derive"] } codex-analytics = { workspace = true } codex-api = { workspace = true } +codex-client = { workspace = true } codex-app-server-protocol = { workspace = true } codex-apply-patch = { workspace = true } codex-async-utils = { workspace = true } diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 821611e60e..06c93cbbb0 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 + .native_integrity_surface() + { + new_thread + .thread + .codex + .session + .services + .model_client + .set_native_integrity_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..c2bbbaf10f 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -10,6 +10,7 @@ use crate::context::ContextualUserFragment; use crate::context::SubagentNotification; use crate::init_state_db; use assert_matches::assert_matches; +use codex_client::NativeIntegritySurface; use codex_features::Feature; use codex_login::CodexAuth; use codex_protocol::AgentPath; @@ -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_native_integrity_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 + .native_integrity_surface(), + Some(NativeIntegritySurface::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..2564c323f9 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_client::NativeIntegrityStateContext; +use codex_client::NativeIntegritySurface; 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,6 +183,7 @@ struct ModelClientState { beta_features_header: Option, include_attestation: bool, attestation_provider: Option>, + native_integrity_state: Option, disable_websockets: AtomicBool, cached_websocket_session: StdMutex, } @@ -329,6 +333,12 @@ impl ModelClient { beta_features_header: Option, attestation_provider: Option>, ) -> Self { + let native_integrity_state = auth_manager.as_ref().map(|auth_manager| { + NativeIntegrityStateContext::new( + auth_manager.codex_home().to_path_buf(), + native_integrity_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,6 +363,7 @@ impl ModelClient { beta_features_header, include_attestation, attestation_provider, + native_integrity_state, disable_websockets: AtomicBool::new(false), cached_websocket_session: StdMutex::new(WebsocketSession::default()), }), @@ -390,6 +401,31 @@ 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(NativeIntegritySurface::from_app_server_client_name) + .unwrap_or(NativeIntegritySurface::CodexCli); + self.set_native_integrity_surface(surface); + } + + pub(crate) fn native_integrity_surface(&self) -> Option { + self.state + .native_integrity_state + .as_ref() + .map(NativeIntegrityStateContext::surface) + } + + pub(crate) fn set_native_integrity_surface(&self, surface: NativeIntegritySurface) { + if self + .state + .native_integrity_state + .as_ref() + .is_some_and(|state| state.set_surface(surface)) + { + self.store_cached_websocket_session(WebsocketSession::default()); + } + } + pub(crate) fn set_window_generation(&self, window_generation: u64) { self.state .window_generation @@ -812,7 +848,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.native_integrity_state.clone(), + ); Ok(CurrentClientSetup { auth, api_provider, @@ -1736,6 +1776,21 @@ fn subagent_header_value(session_source: &SessionSource) -> Option { } } +fn native_integrity_surface_for_session_source( + session_source: &SessionSource, +) -> NativeIntegritySurface { + match session_source { + SessionSource::Cli => NativeIntegritySurface::CodexTui, + SessionSource::Exec => NativeIntegritySurface::CodexExec, + SessionSource::VSCode => NativeIntegritySurface::CodexVscode, + SessionSource::SubAgent(_) + | SessionSource::Internal(_) + | SessionSource::Mcp + | SessionSource::Custom(_) + | SessionSource::Unknown => NativeIntegritySurface::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/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 3132420e4c..1f4810b656 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -110,6 +110,17 @@ pub(crate) async fn run_codex_thread_interactive( })) .or_cancel(&cancel_token) .await??; + if let Some(surface) = parent_session + .services + .model_client + .native_integrity_surface() + { + codex + .session + .services + .model_client + .set_native_integrity_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 64e0c16416..607cc20431 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_client::NativeIntegritySurface; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::models::NetworkPermissions; @@ -181,6 +182,40 @@ 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_native_integrity_surface() { + let (parent_session, parent_ctx, _rx_events) = + crate::session::tests::make_session_and_context_with_rx().await; + parent_session + .services + .model_client + .set_native_integrity_surface(NativeIntegritySurface::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 + .native_integrity_surface(), + Some(NativeIntegritySurface::CodexDesktop) + ); + cancel_token.cancel(); +} + #[tokio::test] async fn handle_request_permissions_uses_tool_call_id_for_round_trip() { let (parent_session, parent_ctx, rx_events) = diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index b3dcf504ee..46a0ec5f30 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -781,6 +781,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/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..0bcf17cd97 100644 --- a/codex-rs/model-provider/Cargo.toml +++ b/codex-rs/model-provider/Cargo.toml @@ -32,5 +32,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..c2b01673ef 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_client::NativeIntegrityStateContext; 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,21 @@ 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(_)) = 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..7ccbfd8346 100644 --- a/codex-rs/model-provider/src/bearer_auth_provider.rs +++ b/codex-rs/model-provider/src/bearer_auth_provider.rs @@ -1,4 +1,10 @@ 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::NativeIntegrityStateContext; +use codex_client::NativeIntegritySurface; +use codex_client::is_allowed_chatgpt_request_url; use http::HeaderMap; use http::HeaderValue; @@ -46,10 +52,91 @@ impl AuthProvider for BearerAuthProvider { } } +#[derive(Clone)] +pub(crate) struct NativeIntegrityAuthProvider { + auth: SharedAuthProvider, + state: NativeIntegrityStateContext, + surface: NativeIntegritySurface, +} + +impl NativeIntegrityAuthProvider { + pub(crate) fn new(auth: SharedAuthProvider, state: NativeIntegrityStateContext) -> 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; + } + + match self.state.load_for_surface(self.surface) { + Ok(Some(state_file)) => { + if let Ok(header) = HeaderValue::from_str(&state_file.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(expected_state) = request_headers + .get(INTEGRITY_STATE_HEADER_NAME) + .and_then(|value| value.to_str().ok()) + else { + return; + }; + let Some(next_state) = response_headers + .get(INTEGRITY_STATE_UPDATE_HEADER_NAME) + .and_then(|value| value.to_str().ok()) + else { + return; + }; + + if let Err(error) = self.state.compare_and_store_for_surface( + self.surface, + expected_state, + next_state.to_string(), + ) { + tracing::warn!("failed to rotate native integrity state: {error}"); + } + } +} + #[cfg(test)] mod tests { use super::*; + use codex_client::NativeIntegrityStateFile; + use codex_client::NativeIntegrityStateStore; 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 +194,79 @@ mod tests { Some("true") ); } + + #[test] + fn native_integrity_provider_scopes_and_rotates_surface_state() { + let codex_home = TempDir::new().expect("tempdir"); + let context = NativeIntegrityStateContext::new( + codex_home.path().to_path_buf(), + NativeIntegritySurface::CodexDesktop, + ); + let store = NativeIntegrityStateStore::new(codex_home.path().to_path_buf()); + store + .replace( + NativeIntegritySurface::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(NativeIntegritySurface::CodexCli); + store + .replace( + NativeIntegritySurface::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") + ); + + 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 + .load(NativeIntegritySurface::CodexDesktop) + .expect("state should load") + .expect("state should exist"), + NativeIntegrityStateFile { + state: "ois1.rotated.nonce.ciphertext".to_string(), + } + ); + assert_eq!( + store + .load(NativeIntegritySurface::CodexCli) + .expect("CLI state should load") + .expect("CLI state should exist"), + NativeIntegrityStateFile { + state: "ois1.cli.nonce.ciphertext".to_string(), + } + ); + + 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)); + } } 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;