From 19e2f218271a0795ef490909e978d8ff62bbb3de Mon Sep 17 00:00:00 2001 From: Adrian <145513011+adrian-openai@users.noreply.github.com> Date: Mon, 20 Apr 2026 07:24:29 -0700 Subject: [PATCH] [codex] Use background task auth for additional backend calls (#18260) ## Summary Splits the larger PR4.1 background task auth rollout by moving additional backend/control-plane call sites into this downstream PR. This PR keeps callers on the same design as PR4.1: most code asks `AuthManager` for the default ChatGPT backend authorization header, and `AuthManager` decides bearer vs background AgentAssertion internally. Task-pinned inference auth remains separate because it needs the thread's registered task id. ## Stack - PR1: https://github.com/openai/codex/pull/17385 - add `features.use_agent_identity` - PR2: https://github.com/openai/codex/pull/17386 - register agent identities when enabled - PR3: https://github.com/openai/codex/pull/17387 - register agent tasks when enabled - PR3.1: https://github.com/openai/codex/pull/17978 - persist and prewarm registered tasks per thread - PR4: https://github.com/openai/codex/pull/17980 - use task-scoped `AgentAssertion` for downstream calls - PR4.1: https://github.com/openai/codex/pull/18094 - introduce AuthManager-owned background/control-plane `AgentAssertion` auth - PR4.2: this PR - use background task auth for additional backend/control-plane calls ## What Changed - pass full authorization header values through backend-client and cloud-tasks-client call paths where needed - move ChatGPT client, cloud requirements, cloud tasks, thread-manager, and models-manager background auth usage into this downstream slice - make app-server remote control enrollment/websocket auth ask `AuthManager` for the local backend authorization header instead of threading a background auth mode through transport options - keep the same feature-gated bearer fallback behavior from PR4.1 ## Validation - `just fmt` - `cargo check -p codex-core -p codex-login -p codex-analytics -p codex-app-server -p codex-cloud-requirements -p codex-cloud-tasks -p codex-models-manager -p codex-chatgpt -p codex-model-provider -p codex-mcp -p codex-core-skills` - `cargo test -p codex-login agent_identity` - `cargo test -p codex-model-provider bearer_auth_provider` - `cargo test -p codex-core agent_assertion` - `cargo test -p codex-app-server remote_control` - `cargo test -p codex-cloud-requirements fetch_cloud_requirements` - `cargo test -p codex-models-manager manager::tests` - `cargo test -p codex-chatgpt` - `cargo test -p codex-cloud-tasks` - `just fix -p codex-core -p codex-login -p codex-analytics -p codex-app-server -p codex-cloud-requirements -p codex-cloud-tasks -p codex-models-manager -p codex-chatgpt -p codex-model-provider -p codex-mcp -p codex-core-skills` - `just fix -p codex-app-server` - `git diff --check` --- .../app-server/src/codex_message_processor.rs | 18 ++- codex-rs/app-server/src/lib.rs | 24 +-- codex-rs/app-server/src/transport/mod.rs | 3 +- .../src/transport/remote_control/enroll.rs | 18 ++- .../src/transport/remote_control/mod.rs | 41 ++++- .../src/transport/remote_control/websocket.rs | 140 ++++++++++++++++-- codex-rs/backend-client/src/client.rs | 24 +-- codex-rs/chatgpt/src/chatgpt_client.rs | 23 ++- codex-rs/cloud-requirements/src/lib.rs | 33 ++++- codex-rs/cloud-tasks-client/src/http.rs | 10 ++ codex-rs/cloud-tasks/src/lib.rs | 52 +++---- codex-rs/cloud-tasks/src/util.rs | 48 +++--- codex-rs/login/src/auth/manager.rs | 5 + codex-rs/models-manager/src/manager.rs | 23 ++- 14 files changed, 364 insertions(+), 98 deletions(-) diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 33c821bec..229c2a4a6 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -2013,12 +2013,28 @@ impl CodexMessageProcessor { }); } - let client = BackendClient::from_auth(self.config.chatgpt_base_url.clone(), &auth) + let authorization_header_value = self + .auth_manager + .chatgpt_authorization_header_for_auth(&auth) + .await; + let mut client = BackendClient::new(self.config.chatgpt_base_url.clone()) + .map(|client| { + client.with_user_agent(codex_login::default_client::get_codex_user_agent()) + }) .map_err(|err| JSONRPCErrorError { code: INTERNAL_ERROR_CODE, message: format!("failed to construct backend client: {err}"), data: None, })?; + if let Some(authorization_header_value) = authorization_header_value { + client = client.with_authorization_header_value(authorization_header_value); + } + if let Some(account_id) = auth.get_account_id() { + client = client.with_chatgpt_account_id(account_id); + } + if auth.is_fedramp_account() { + client = client.with_fedramp_routing_header(); + } let snapshots = client .get_rate_limits_many() diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 65374a0d9..b7b099960 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -27,10 +27,11 @@ use crate::outgoing_message::QueuedOutgoingMessage; use crate::transport::CHANNEL_CAPACITY; use crate::transport::ConnectionState; use crate::transport::OutboundConnectionState; +use crate::transport::RemoteControlStartOptions; use crate::transport::TransportEvent; use crate::transport::auth::policy_from_settings; use crate::transport::route_outgoing_envelope; -use crate::transport::start_remote_control; +use crate::transport::start_remote_control_with_options; use crate::transport::start_stdio_connection; use crate::transport::start_websocket_acceptor; use codex_analytics::AppServerRpcTransport; @@ -575,16 +576,17 @@ pub async fn run_main_with_transport( )); } - let (remote_control_accept_handle, remote_control_handle) = start_remote_control( - config.chatgpt_base_url.clone(), - state_db.clone(), - auth_manager.clone(), - transport_event_tx.clone(), - transport_shutdown_token.clone(), - app_server_client_name_rx, - remote_control_enabled, - ) - .await?; + let (remote_control_accept_handle, remote_control_handle) = + start_remote_control_with_options(RemoteControlStartOptions { + remote_control_url: config.chatgpt_base_url.clone(), + state_db: state_db.clone(), + auth_manager: auth_manager.clone(), + transport_event_tx: transport_event_tx.clone(), + shutdown_token: transport_shutdown_token.clone(), + app_server_client_name_rx, + initial_enabled: remote_control_enabled, + }) + .await?; transport_accept_handles.push(remote_control_accept_handle); let outbound_handle = tokio::spawn(async move { diff --git a/codex-rs/app-server/src/transport/mod.rs b/codex-rs/app-server/src/transport/mod.rs index 75a497190..78ed01bfa 100644 --- a/codex-rs/app-server/src/transport/mod.rs +++ b/codex-rs/app-server/src/transport/mod.rs @@ -34,7 +34,8 @@ mod stdio; mod websocket; pub(crate) use remote_control::RemoteControlHandle; -pub(crate) use remote_control::start_remote_control; +pub(crate) use remote_control::RemoteControlStartOptions; +pub(crate) use remote_control::start_remote_control_with_options; pub(crate) use stdio::start_stdio_connection; pub(crate) use websocket::start_websocket_acceptor; diff --git a/codex-rs/app-server/src/transport/remote_control/enroll.rs b/codex-rs/app-server/src/transport/remote_control/enroll.rs index dbe18c835..4d9929818 100644 --- a/codex-rs/app-server/src/transport/remote_control/enroll.rs +++ b/codex-rs/app-server/src/transport/remote_control/enroll.rs @@ -17,6 +17,7 @@ const REMOTE_CONTROL_RESPONSE_BODY_MAX_BYTES: usize = 4096; const REQUEST_ID_HEADER: &str = "x-request-id"; const OAI_REQUEST_ID_HEADER: &str = "x-oai-request-id"; const CF_RAY_HEADER: &str = "cf-ray"; +const REMOTE_CONTROL_FEDRAMP_HEADER: &str = "X-OpenAI-Fedramp"; pub(super) const REMOTE_CONTROL_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id"; #[derive(Debug, Clone, PartialEq, Eq)] @@ -29,8 +30,9 @@ pub(super) struct RemoteControlEnrollment { #[derive(Debug, Clone, PartialEq, Eq)] pub(super) struct RemoteControlConnectionAuth { - pub(super) bearer_token: String, + pub(super) authorization_header_value: String, pub(super) account_id: String, + pub(super) is_fedramp_account: bool, } pub(super) async fn load_persisted_remote_control_enrollment( @@ -199,12 +201,15 @@ pub(super) async fn enroll_remote_control_server( app_server_version: env!("CARGO_PKG_VERSION"), }; let client = build_reqwest_client(); - let http_request = client + let mut http_request = client .post(enroll_url) .timeout(REMOTE_CONTROL_ENROLL_TIMEOUT) - .bearer_auth(&auth.bearer_token) - .header(REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id) - .json(&request); + .header("authorization", &auth.authorization_header_value) + .header(REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id); + if auth.is_fedramp_account { + http_request = http_request.header(REMOTE_CONTROL_FEDRAMP_HEADER, "true"); + } + let http_request = http_request.json(&request); let response = http_request.send().await.map_err(|err| { io::Error::other(format!( @@ -445,8 +450,9 @@ mod tests { let err = enroll_remote_control_server( &remote_control_target, &RemoteControlConnectionAuth { - bearer_token: "Access Token".to_string(), + authorization_header_value: "Bearer Access Token".to_string(), account_id: "account_id".to_string(), + is_fedramp_account: false, }, ) .await diff --git a/codex-rs/app-server/src/transport/remote_control/mod.rs b/codex-rs/app-server/src/transport/remote_control/mod.rs index c014c7a2c..8764928ea 100644 --- a/codex-rs/app-server/src/transport/remote_control/mod.rs +++ b/codex-rs/app-server/src/transport/remote_control/mod.rs @@ -4,6 +4,7 @@ mod protocol; mod websocket; use crate::transport::remote_control::websocket::RemoteControlWebsocket; +use crate::transport::remote_control::websocket::RemoteControlWebsocketOptions; pub use self::protocol::ClientId; use self::protocol::ServerEvent; @@ -44,6 +45,17 @@ impl RemoteControlHandle { } } +pub(crate) struct RemoteControlStartOptions { + pub(crate) remote_control_url: String, + pub(crate) state_db: Option>, + pub(crate) auth_manager: Arc, + pub(crate) transport_event_tx: mpsc::Sender, + pub(crate) shutdown_token: CancellationToken, + pub(crate) app_server_client_name_rx: Option>, + pub(crate) initial_enabled: bool, +} + +#[cfg(test)] pub(crate) async fn start_remote_control( remote_control_url: String, state_db: Option>, @@ -53,15 +65,38 @@ pub(crate) async fn start_remote_control( app_server_client_name_rx: Option>, initial_enabled: bool, ) -> io::Result<(JoinHandle<()>, RemoteControlHandle)> { + start_remote_control_with_options(RemoteControlStartOptions { + remote_control_url, + state_db, + auth_manager, + transport_event_tx, + shutdown_token, + app_server_client_name_rx, + initial_enabled, + }) + .await +} + +pub(crate) async fn start_remote_control_with_options( + options: RemoteControlStartOptions, +) -> io::Result<(JoinHandle<()>, RemoteControlHandle)> { + let RemoteControlStartOptions { + remote_control_url, + state_db, + auth_manager, + transport_event_tx, + shutdown_token, + app_server_client_name_rx, + initial_enabled, + } = options; let remote_control_target = if initial_enabled { Some(normalize_remote_control_url(&remote_control_url)?) } else { None }; - let (enabled_tx, enabled_rx) = watch::channel(initial_enabled); let join_handle = tokio::spawn(async move { - RemoteControlWebsocket::new( + RemoteControlWebsocket::from_options(RemoteControlWebsocketOptions { remote_control_url, remote_control_target, state_db, @@ -69,7 +104,7 @@ pub(crate) async fn start_remote_control( transport_event_tx, shutdown_token, enabled_rx, - ) + }) .run(app_server_client_name_rx) .await; }); diff --git a/codex-rs/app-server/src/transport/remote_control/websocket.rs b/codex-rs/app-server/src/transport/remote_control/websocket.rs index 58df3035a..ac91882f5 100644 --- a/codex-rs/app-server/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server/src/transport/remote_control/websocket.rs @@ -49,6 +49,7 @@ use tracing::warn; pub(super) const REMOTE_CONTROL_PROTOCOL_VERSION: &str = "2"; pub(super) const REMOTE_CONTROL_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id"; +const REMOTE_CONTROL_FEDRAMP_HEADER: &str = "X-OpenAI-Fedramp"; const REMOTE_CONTROL_SUBSCRIBE_CURSOR_HEADER: &str = "x-codex-subscribe-cursor"; const REMOTE_CONTROL_WEBSOCKET_PING_INTERVAL: std::time::Duration = std::time::Duration::from_secs(10); @@ -128,6 +129,16 @@ pub(crate) struct RemoteControlWebsocket { enabled_rx: watch::Receiver, } +pub(crate) struct RemoteControlWebsocketOptions { + pub(crate) remote_control_url: String, + pub(crate) remote_control_target: Option, + pub(crate) state_db: Option>, + pub(crate) auth_manager: Arc, + pub(crate) transport_event_tx: mpsc::Sender, + pub(crate) shutdown_token: CancellationToken, + pub(crate) enabled_rx: watch::Receiver, +} + enum ConnectOutcome { Connected(Box>>), Disabled, @@ -135,6 +146,7 @@ enum ConnectOutcome { } impl RemoteControlWebsocket { + #[cfg(test)] pub(crate) fn new( remote_control_url: String, remote_control_target: Option, @@ -144,6 +156,27 @@ impl RemoteControlWebsocket { shutdown_token: CancellationToken, enabled_rx: watch::Receiver, ) -> Self { + Self::from_options(RemoteControlWebsocketOptions { + remote_control_url, + remote_control_target, + state_db, + auth_manager, + transport_event_tx, + shutdown_token, + enabled_rx, + }) + } + + pub(crate) fn from_options(options: RemoteControlWebsocketOptions) -> Self { + let RemoteControlWebsocketOptions { + remote_control_url, + remote_control_target, + state_db, + auth_manager, + transport_event_tx, + shutdown_token, + enabled_rx, + } = options; let shutdown_token = shutdown_token.child_token(); let (server_event_tx, server_event_rx) = mpsc::channel(super::CHANNEL_CAPACITY); let client_tracker = @@ -271,14 +304,16 @@ impl RemoteControlWebsocket { } return ConnectOutcome::Disabled; } - connect_result = connect_remote_control_websocket( - &remote_control_target, - self.state_db.as_deref(), - &self.auth_manager, - &mut self.auth_recovery, - &mut self.enrollment, - subscribe_cursor.as_deref(), - app_server_client_name, + connect_result = connect_remote_control_websocket_with_options( + ConnectRemoteControlWebsocketOptions { + remote_control_target: &remote_control_target, + state_db: self.state_db.as_deref(), + auth_manager: &self.auth_manager, + auth_recovery: &mut self.auth_recovery, + enrollment: &mut self.enrollment, + subscribe_cursor: subscribe_cursor.as_deref(), + app_server_client_name, + }, ) => connect_result, }; @@ -668,12 +703,11 @@ fn build_remote_control_websocket_request( "x-codex-protocol-version", REMOTE_CONTROL_PROTOCOL_VERSION, )?; - set_remote_control_header( - headers, - "authorization", - &format!("Bearer {}", auth.bearer_token), - )?; + set_remote_control_header(headers, "authorization", &auth.authorization_header_value)?; set_remote_control_header(headers, REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id)?; + if auth.is_fedramp_account { + set_remote_control_header(headers, REMOTE_CONTROL_FEDRAMP_HEADER, "true")?; + } if let Some(subscribe_cursor) = subscribe_cursor { set_remote_control_header( headers, @@ -718,8 +752,19 @@ pub(crate) async fn load_remote_control_auth( )); } + let authorization_header_value = auth_manager + .chatgpt_authorization_header_for_auth(&auth) + .await + .ok_or_else(|| { + io::Error::new( + ErrorKind::PermissionDenied, + "remote control requires ChatGPT authentication", + ) + })?; + Ok(RemoteControlConnectionAuth { - bearer_token: auth.get_token().map_err(io::Error::other)?, + authorization_header_value, + is_fedramp_account: auth.is_fedramp_account(), account_id: auth.get_account_id().ok_or_else(|| { io::Error::new( ErrorKind::WouldBlock, @@ -729,6 +774,7 @@ pub(crate) async fn load_remote_control_auth( }) } +#[cfg(test)] pub(super) async fn connect_remote_control_websocket( remote_control_target: &RemoteControlTarget, state_db: Option<&StateRuntime>, @@ -741,6 +787,44 @@ pub(super) async fn connect_remote_control_websocket( WebSocketStream>, tungstenite::http::Response<()>, )> { + connect_remote_control_websocket_with_options(ConnectRemoteControlWebsocketOptions { + remote_control_target, + state_db, + auth_manager, + auth_recovery, + enrollment, + subscribe_cursor, + app_server_client_name, + }) + .await +} + +struct ConnectRemoteControlWebsocketOptions<'a> { + remote_control_target: &'a RemoteControlTarget, + state_db: Option<&'a StateRuntime>, + auth_manager: &'a Arc, + auth_recovery: &'a mut UnauthorizedRecovery, + enrollment: &'a mut Option, + subscribe_cursor: Option<&'a str>, + app_server_client_name: Option<&'a str>, +} + +async fn connect_remote_control_websocket_with_options( + options: ConnectRemoteControlWebsocketOptions<'_>, +) -> io::Result<( + WebSocketStream>, + tungstenite::http::Response<()>, +)> { + let ConnectRemoteControlWebsocketOptions { + remote_control_target, + state_db, + auth_manager, + auth_recovery, + enrollment, + subscribe_cursor, + app_server_client_name, + } = options; + ensure_rustls_crypto_provider(); let auth = load_remote_control_auth(auth_manager).await?; @@ -1003,6 +1087,34 @@ mod tests { } } + #[test] + fn build_remote_control_websocket_request_includes_fedramp_header() { + let request = build_remote_control_websocket_request( + "ws://127.0.0.1/backend-api/wham/remote/control/server", + &RemoteControlEnrollment { + account_id: "account_id".to_string(), + environment_id: "env_test".to_string(), + server_id: "srv_e_test".to_string(), + server_name: "test-server".to_string(), + }, + &RemoteControlConnectionAuth { + authorization_header_value: "AgentAssertion assertion".to_string(), + account_id: "account_id".to_string(), + is_fedramp_account: true, + }, + /*subscribe_cursor*/ None, + ) + .expect("request should build"); + + assert_eq!( + request + .headers() + .get(REMOTE_CONTROL_FEDRAMP_HEADER) + .and_then(|value| value.to_str().ok()), + Some("true") + ); + } + #[tokio::test] async fn connect_remote_control_websocket_includes_http_error_details() { let listener = TcpListener::bind("127.0.0.1:0") diff --git a/codex-rs/backend-client/src/client.rs b/codex-rs/backend-client/src/client.rs index 8f84ef28f..e6ea0253b 100644 --- a/codex-rs/backend-client/src/client.rs +++ b/codex-rs/backend-client/src/client.rs @@ -116,7 +116,7 @@ impl PathStyle { pub struct Client { base_url: String, http: reqwest::Client, - bearer_token: Option, + authorization_header_value: Option, user_agent: Option, chatgpt_account_id: Option, chatgpt_account_is_fedramp: bool, @@ -142,7 +142,7 @@ impl Client { Ok(Self { base_url, http, - bearer_token: None, + authorization_header_value: None, user_agent: None, chatgpt_account_id: None, chatgpt_account_is_fedramp: false, @@ -165,7 +165,12 @@ impl Client { } pub fn with_bearer_token(mut self, token: impl Into) -> Self { - self.bearer_token = Some(token.into()); + self.authorization_header_value = Some(format!("Bearer {}", token.into())); + self + } + + pub fn with_authorization_header_value(mut self, value: impl Into) -> Self { + self.authorization_header_value = Some(value.into()); self } @@ -198,11 +203,10 @@ impl Client { } else { h.insert(USER_AGENT, HeaderValue::from_static("codex-cli")); } - if let Some(token) = &self.bearer_token { - let value = format!("Bearer {token}"); - if let Ok(hv) = HeaderValue::from_str(&value) { - h.insert(AUTHORIZATION, hv); - } + if let Some(value) = &self.authorization_header_value + && let Ok(hv) = HeaderValue::from_str(value) + { + h.insert(AUTHORIZATION, hv); } if let Some(acc) = &self.chatgpt_account_id && let Ok(name) = HeaderName::from_bytes(b"ChatGPT-Account-Id") @@ -816,7 +820,7 @@ mod tests { let codex_client = Client { base_url: "https://example.test".to_string(), http: reqwest::Client::new(), - bearer_token: None, + authorization_header_value: None, user_agent: None, chatgpt_account_id: None, chatgpt_account_is_fedramp: false, @@ -830,7 +834,7 @@ mod tests { let chatgpt_client = Client { base_url: "https://chatgpt.com/backend-api".to_string(), http: reqwest::Client::new(), - bearer_token: None, + authorization_header_value: None, user_agent: None, chatgpt_account_id: None, chatgpt_account_is_fedramp: false, diff --git a/codex-rs/chatgpt/src/chatgpt_client.rs b/codex-rs/chatgpt/src/chatgpt_client.rs index fa3a63dad..6bdb16347 100644 --- a/codex-rs/chatgpt/src/chatgpt_client.rs +++ b/codex-rs/chatgpt/src/chatgpt_client.rs @@ -1,4 +1,5 @@ use codex_core::config::Config; +use codex_login::AuthManager; use codex_login::default_client::create_client; use crate::chatgpt_token::get_chatgpt_token_data; @@ -31,16 +32,32 @@ pub(crate) async fn chatgpt_get_request_with_timeout( let token = get_chatgpt_token_data().ok_or_else(|| anyhow::anyhow!("ChatGPT token not available"))?; + let auth_manager = + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false); + let auth = auth_manager.auth().await; + let is_fedramp_account = auth + .as_ref() + .is_some_and(codex_login::CodexAuth::is_fedramp_account); + let authorization_header_value = match auth.as_ref() { + Some(auth) if auth.is_chatgpt_auth() => auth_manager + .chatgpt_authorization_header_for_auth(auth) + .await + .unwrap_or_else(|| format!("Bearer {}", token.access_token)), + _ => format!("Bearer {}", token.access_token), + }; let account_id = token.account_id.ok_or_else(|| { anyhow::anyhow!("ChatGPT account ID not available, please re-run `codex login`") - }); + })?; let mut request = client .get(&url) - .bearer_auth(&token.access_token) - .header("chatgpt-account-id", account_id?) + .header("authorization", authorization_header_value) + .header("chatgpt-account-id", account_id) .header("Content-Type", "application/json"); + if is_fedramp_account { + request = request.header("X-OpenAI-Fedramp", "true"); + } if let Some(timeout) = timeout { request = request.timeout(timeout); diff --git a/codex-rs/cloud-requirements/src/lib.rs b/codex-rs/cloud-requirements/src/lib.rs index 676d3201d..60a7af20e 100644 --- a/codex-rs/cloud-requirements/src/lib.rs +++ b/codex-rs/cloud-requirements/src/lib.rs @@ -195,11 +195,15 @@ trait RequirementsFetcher: Send + Sync { struct BackendRequirementsFetcher { base_url: String, + auth_manager: Arc, } impl BackendRequirementsFetcher { - fn new(base_url: String) -> Self { - Self { base_url } + fn new(auth_manager: Arc, base_url: String) -> Self { + Self { + base_url, + auth_manager, + } } } @@ -209,7 +213,14 @@ impl RequirementsFetcher for BackendRequirementsFetcher { &self, auth: &CodexAuth, ) -> Result, FetchAttemptError> { - let client = BackendClient::from_auth(self.base_url.clone(), auth) + let authorization_header_value = self + .auth_manager + .chatgpt_authorization_header_for_auth(auth) + .await; + let mut client = BackendClient::new(self.base_url.clone()) + .map(|client| { + client.with_user_agent(codex_login::default_client::get_codex_user_agent()) + }) .inspect_err(|err| { tracing::warn!( error = %err, @@ -217,6 +228,15 @@ impl RequirementsFetcher for BackendRequirementsFetcher { ); }) .map_err(|_| FetchAttemptError::Retryable(RetryableFailureKind::BackendClientInit))?; + if let Some(authorization_header_value) = authorization_header_value { + client = client.with_authorization_header_value(authorization_header_value); + } + if let Some(account_id) = auth.get_account_id() { + client = client.with_chatgpt_account_id(account_id); + } + if auth.is_fedramp_account() { + client = client.with_fedramp_routing_header(); + } let response = client .get_config_requirements_file() @@ -693,8 +713,11 @@ pub fn cloud_requirements_loader( codex_home: PathBuf, ) -> CloudRequirementsLoader { let service = CloudRequirementsService::new( - auth_manager, - Arc::new(BackendRequirementsFetcher::new(chatgpt_base_url)), + auth_manager.clone(), + Arc::new(BackendRequirementsFetcher::new( + auth_manager, + chatgpt_base_url, + )), codex_home, CLOUD_REQUIREMENTS_TIMEOUT, ); diff --git a/codex-rs/cloud-tasks-client/src/http.rs b/codex-rs/cloud-tasks-client/src/http.rs index 4ea098022..3ada70236 100644 --- a/codex-rs/cloud-tasks-client/src/http.rs +++ b/codex-rs/cloud-tasks-client/src/http.rs @@ -37,6 +37,11 @@ impl HttpClient { self } + pub fn with_authorization_header_value(mut self, value: impl Into) -> Self { + self.backend = self.backend.clone().with_authorization_header_value(value); + self + } + pub fn with_user_agent(mut self, ua: impl Into) -> Self { self.backend = self.backend.clone().with_user_agent(ua); self @@ -47,6 +52,11 @@ impl HttpClient { self } + pub fn with_fedramp_routing_header(mut self) -> Self { + self.backend = self.backend.clone().with_fedramp_routing_header(); + self + } + fn tasks_api(&self) -> api::Tasks<'_> { api::Tasks::new(self) } diff --git a/codex-rs/cloud-tasks/src/lib.rs b/codex-rs/cloud-tasks/src/lib.rs index 7006d52b9..155587003 100644 --- a/codex-rs/cloud-tasks/src/lib.rs +++ b/codex-rs/cloud-tasks/src/lib.rs @@ -68,43 +68,45 @@ async fn init_backend(user_agent_suffix: &str) -> anyhow::Result }; append_error_log(format!("startup: base_url={base_url} path_style={style}")); - let auth_manager = util::load_auth_manager().await; - let auth = match auth_manager.as_ref() { - Some(manager) => manager.auth().await, - None => None, + let Some(auth_manager) = util::load_auth_manager(Some(base_url.clone())).await else { + eprintln!( + "Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'." + ); + std::process::exit(1); }; - let auth = match auth { - Some(auth) => auth, - None => { - eprintln!( - "Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'." - ); - std::process::exit(1); - } + let Some(auth) = auth_manager.auth().await else { + eprintln!( + "Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'." + ); + std::process::exit(1); }; if let Some(acc) = auth.get_account_id() { append_error_log(format!("auth: mode=ChatGPT account_id={acc}")); } - let token = match auth.get_token() { - Ok(t) if !t.is_empty() => t, - _ => { - eprintln!( - "Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'." - ); - std::process::exit(1); - } + let authorization_header_value = auth_manager + .chatgpt_authorization_header_for_auth(&auth) + .await; + let Some(authorization_header_value) = authorization_header_value else { + eprintln!( + "Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'." + ); + std::process::exit(1); }; - http = http.with_bearer_token(token.clone()); - if let Some(acc) = auth - .get_account_id() - .or_else(|| util::extract_chatgpt_account_id(&token)) - { + http = http.with_authorization_header_value(authorization_header_value); + if let Some(acc) = auth.get_account_id().or_else(|| { + auth.get_token() + .ok() + .and_then(|token| util::extract_chatgpt_account_id(&token)) + }) { append_error_log(format!("auth: set ChatGPT-Account-Id header: {acc}")); http = http.with_chatgpt_account_id(acc); } + if auth.is_fedramp_account() { + http = http.with_fedramp_routing_header(); + } Ok(BackendContext { backend: Arc::new(http), diff --git a/codex-rs/cloud-tasks/src/util.rs b/codex-rs/cloud-tasks/src/util.rs index cbaed17be..693ff7839 100644 --- a/codex-rs/cloud-tasks/src/util.rs +++ b/codex-rs/cloud-tasks/src/util.rs @@ -3,6 +3,7 @@ use chrono::DateTime; use chrono::Local; use chrono::Utc; use reqwest::header::HeaderMap; +use std::sync::Arc; use codex_core::config::Config; use codex_login::AuthManager; @@ -59,18 +60,18 @@ pub fn extract_chatgpt_account_id(token: &str) -> Option { .map(str::to_string) } -pub async fn load_auth_manager() -> Option { +pub async fn load_auth_manager(chatgpt_base_url: Option) -> Option> { // TODO: pass in cli overrides once cloud tasks properly support them. let config = Config::load_with_cli_overrides(Vec::new()).await.ok()?; - Some(AuthManager::new( - config.codex_home.to_path_buf(), - /*enable_codex_api_key_env*/ false, - config.cli_auth_credentials_store_mode, - )) + let auth_manager = + AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false); + if let Some(chatgpt_base_url) = chatgpt_base_url { + auth_manager.set_chatgpt_backend_base_url(Some(chatgpt_base_url)); + } + Some(auth_manager) } -/// Build headers for ChatGPT-backed requests: `User-Agent`, optional `Authorization`, -/// and optional `ChatGPT-Account-Id`. +/// Build headers for ChatGPT-backed requests. pub async fn build_chatgpt_headers() -> HeaderMap { use reqwest::header::AUTHORIZATION; use reqwest::header::HeaderName; @@ -84,23 +85,34 @@ pub async fn build_chatgpt_headers() -> HeaderMap { USER_AGENT, HeaderValue::from_str(&ua).unwrap_or(HeaderValue::from_static("codex-cli")), ); - if let Some(am) = load_auth_manager().await - && let Some(auth) = am.auth().await - && let Ok(tok) = auth.get_token() - && !tok.is_empty() + let base_url = normalize_base_url( + &std::env::var("CODEX_CLOUD_TASKS_BASE_URL") + .unwrap_or_else(|_| "https://chatgpt.com/backend-api".to_string()), + ); + if let Some(auth_manager) = load_auth_manager(Some(base_url)).await + && let Some(auth) = auth_manager.auth().await { - let v = format!("Bearer {tok}"); - if let Ok(hv) = HeaderValue::from_str(&v) { + if let Some(authorization_header_value) = auth_manager + .chatgpt_authorization_header_for_auth(&auth) + .await + && let Ok(hv) = HeaderValue::from_str(&authorization_header_value) + { headers.insert(AUTHORIZATION, hv); } - if let Some(acc) = auth - .get_account_id() - .or_else(|| extract_chatgpt_account_id(&tok)) - && let Ok(name) = HeaderName::from_bytes(b"ChatGPT-Account-Id") + if let Some(acc) = auth.get_account_id().or_else(|| { + auth.get_token() + .ok() + .and_then(|token| extract_chatgpt_account_id(&token)) + }) && let Ok(name) = HeaderName::from_bytes(b"ChatGPT-Account-Id") && let Ok(hv) = HeaderValue::from_str(&acc) { headers.insert(name, hv); } + if auth.is_fedramp_account() + && let Ok(name) = HeaderName::from_bytes(b"X-OpenAI-Fedramp") + { + headers.insert(name, HeaderValue::from_static("true")); + } } headers } diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 361d3722a..8a5e2fa66 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1581,6 +1581,11 @@ impl AuthManager { } } + pub fn set_chatgpt_backend_base_url(&self, chatgpt_base_url: Option) { + let (_, auth_mode) = self.chatgpt_backend_auth_config(); + self.set_chatgpt_backend_auth_config(chatgpt_base_url, auth_mode); + } + fn chatgpt_backend_auth_config(&self) -> (Option, BackgroundAgentTaskAuthMode) { let chatgpt_base_url = self .chatgpt_base_url diff --git a/codex-rs/models-manager/src/manager.rs b/codex-rs/models-manager/src/manager.rs index c029960a7..6d5bdb9f7 100644 --- a/codex-rs/models-manager/src/manager.rs +++ b/codex-rs/models-manager/src/manager.rs @@ -17,6 +17,7 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::collect_auth_env_telemetry; use codex_login::default_client::build_reqwest_client; +use codex_model_provider::AuthorizationHeaderAuthProvider; use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; @@ -453,7 +454,23 @@ impl ModelsManager { let auth = self.provider.auth().await; let auth_mode = auth.as_ref().map(CodexAuth::auth_mode); let api_provider = self.provider.api_provider().await?; - let api_auth = self.provider.api_auth().await?; + let mut api_auth = self.provider.api_auth().await?; + if let Some(auth_manager) = auth_manager.as_ref() + && let Some(auth) = auth.as_ref().filter(|auth| auth.is_chatgpt_auth()) + && provider_uses_codex_login_auth(self.provider.info()) + && let Some(authorization_header_value) = auth_manager + .chatgpt_authorization_header_for_auth(auth) + .await + { + let mut auth_provider = AuthorizationHeaderAuthProvider::new( + Some(authorization_header_value), + auth.get_account_id(), + ); + if auth.is_fedramp_account() { + auth_provider = auth_provider.with_fedramp_routing_header(); + } + api_auth = Arc::new(auth_provider); + } let auth_env = collect_auth_env_telemetry(self.provider.info(), codex_api_key_env_enabled); let transport = ReqwestTransport::new(build_reqwest_client()); let auth_telemetry = auth_header_telemetry(api_auth.as_ref()); @@ -601,6 +618,10 @@ impl ModelsManager { } } +fn provider_uses_codex_login_auth(provider: &ModelProviderInfo) -> bool { + provider.env_key.is_none() && provider.experimental_bearer_token.is_none() +} + #[cfg(test)] #[path = "manager_tests.rs"] mod tests;