From b0017b53330169844ea1641f31505e093f8ea76e Mon Sep 17 00:00:00 2001 From: Edward Frazer Date: Wed, 22 Apr 2026 13:20:52 -0700 Subject: [PATCH] fix: refresh codex auth providers per request --- codex-rs/Cargo.lock | 3 + codex-rs/app-server/Cargo.toml | 1 + .../src/transport/remote_control/enroll.rs | 10 +- .../src/transport/remote_control/websocket.rs | 6 +- codex-rs/backend-client/src/client.rs | 52 ++++----- codex-rs/cli/src/mcp_cmd.rs | 8 +- codex-rs/cloud-tasks-client/src/http.rs | 9 +- codex-rs/cloud-tasks/src/lib.rs | 4 +- codex-rs/cloud-tasks/src/util.rs | 6 +- codex-rs/codex-mcp/Cargo.toml | 1 + codex-rs/codex-mcp/src/mcp/auth.rs | 36 ++++-- codex-rs/codex-mcp/src/mcp/mod.rs | 57 +++++----- .../codex-mcp/src/mcp_connection_manager.rs | 23 ++++ codex-rs/core/src/connectors.rs | 24 ++-- codex-rs/core/src/session/handlers.rs | 7 +- codex-rs/core/src/session/mcp.rs | 4 +- codex-rs/core/src/session/session.rs | 2 + codex-rs/model-provider/src/auth.rs | 34 +++--- codex-rs/model-provider/src/lib.rs | 1 + codex-rs/protocol/src/account.rs | 15 ++- codex-rs/protocol/src/auth.rs | 20 ++-- codex-rs/protocol/src/error.rs | 60 +++++----- codex-rs/rmcp-client/Cargo.toml | 1 + codex-rs/rmcp-client/src/rmcp_client.rs | 103 +++++++++++++++--- .../tests/streamable_http_recovery.rs | 1 + 25 files changed, 307 insertions(+), 181 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index c96db244a2..0aed75c7cb 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1829,6 +1829,7 @@ dependencies = [ "chrono", "clap", "codex-analytics", + "codex-api", "codex-app-server-protocol", "codex-arg0", "codex-backend-client", @@ -2829,6 +2830,7 @@ version = "0.0.0" dependencies = [ "anyhow", "async-channel", + "codex-api", "codex-async-utils", "codex-config", "codex-exec-server", @@ -3135,6 +3137,7 @@ version = "0.0.0" dependencies = [ "anyhow", "axum", + "codex-api", "codex-client", "codex-config", "codex-exec-server", diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index b38b1e2812..69110bafc7 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -30,6 +30,7 @@ axum = { workspace = true, default-features = false, features = [ "ws", ] } codex-analytics = { workspace = true } +codex-api = { workspace = true } codex-arg0 = { workspace = true } codex-cloud-requirements = { workspace = true } codex-config = { workspace = true } 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 6737dc9445..ba69c459e8 100644 --- a/codex-rs/app-server/src/transport/remote_control/enroll.rs +++ b/codex-rs/app-server/src/transport/remote_control/enroll.rs @@ -2,6 +2,7 @@ use super::protocol::EnrollRemoteServerRequest; use super::protocol::EnrollRemoteServerResponse; use super::protocol::RemoteControlTarget; use axum::http::HeaderMap; +use codex_api::SharedAuthProvider; use codex_login::default_client::build_reqwest_client; use codex_state::RemoteControlEnrollmentRecord; use codex_state::StateRuntime; @@ -27,9 +28,8 @@ pub(super) struct RemoteControlEnrollment { pub(super) server_name: String, } -#[derive(Debug, Clone, PartialEq, Eq)] pub(super) struct RemoteControlConnectionAuth { - pub(super) auth_headers: HeaderMap, + pub(super) auth_provider: SharedAuthProvider, pub(super) account_id: String, } @@ -199,10 +199,12 @@ pub(super) async fn enroll_remote_control_server( app_server_version: env!("CARGO_PKG_VERSION"), }; let client = build_reqwest_client(); + let mut auth_headers = HeaderMap::new(); + auth.auth_provider.add_auth_headers(&mut auth_headers); let http_request = client .post(enroll_url) .timeout(REMOTE_CONTROL_ENROLL_TIMEOUT) - .headers(auth.auth_headers.clone()) + .headers(auth_headers) .header(REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id) .json(&request); @@ -445,7 +447,7 @@ mod tests { let err = enroll_remote_control_server( &remote_control_target, &RemoteControlConnectionAuth { - auth_headers: HeaderMap::new(), + auth_provider: codex_model_provider::unauthenticated_auth_provider(), account_id: "account_id".to_string(), }, ) 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 1457eecf34..464832e34a 100644 --- a/codex-rs/app-server/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server/src/transport/remote_control/websocket.rs @@ -680,7 +680,9 @@ fn build_remote_control_websocket_request( "x-codex-protocol-version", REMOTE_CONTROL_PROTOCOL_VERSION, )?; - headers.extend(auth.auth_headers.clone()); + let mut auth_headers = tungstenite::http::HeaderMap::new(); + auth.auth_provider.add_auth_headers(&mut auth_headers); + headers.extend(auth_headers); set_remote_control_header(headers, REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id)?; if let Some(subscribe_cursor) = subscribe_cursor { set_remote_control_header( @@ -727,7 +729,7 @@ pub(crate) async fn load_remote_control_auth( } Ok(RemoteControlConnectionAuth { - auth_headers: codex_model_provider::auth_provider_from_auth(&auth).to_auth_headers(), + auth_provider: codex_model_provider::auth_provider_from_auth(&auth), account_id: auth.get_account_id().ok_or_else(|| { io::Error::new( ErrorKind::WouldBlock, diff --git a/codex-rs/backend-client/src/client.rs b/codex-rs/backend-client/src/client.rs index 8f7fe924a0..6365d527ed 100644 --- a/codex-rs/backend-client/src/client.rs +++ b/codex-rs/backend-client/src/client.rs @@ -5,7 +5,7 @@ use crate::types::RateLimitReachedKind as BackendRateLimitReachedKind; use crate::types::RateLimitStatusPayload; use crate::types::TurnAttemptsSiblingTurnsResponse; use anyhow::Result; -use codex_api::AuthProvider; +use codex_api::SharedAuthProvider; use codex_client::build_reqwest_client_with_custom_ca; use codex_client::with_chatgpt_cloudflare_cookie_store; use codex_login::CodexAuth; @@ -16,7 +16,6 @@ use codex_protocol::protocol::RateLimitReachedType; use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::RateLimitWindow; use reqwest::StatusCode; -use reqwest::header::AUTHORIZATION; use reqwest::header::CONTENT_TYPE; use reqwest::header::HeaderMap; use reqwest::header::HeaderName; @@ -114,18 +113,33 @@ impl PathStyle { } } -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct Client { base_url: String, http: reqwest::Client, - bearer_token: Option, - auth_headers: HeaderMap, + auth_provider: SharedAuthProvider, user_agent: Option, chatgpt_account_id: Option, chatgpt_account_is_fedramp: bool, path_style: PathStyle, } +impl fmt::Debug for Client { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Client") + .field("base_url", &self.base_url) + .field("auth_provider", &"") + .field("user_agent", &self.user_agent) + .field("chatgpt_account_id", &self.chatgpt_account_id) + .field( + "chatgpt_account_is_fedramp", + &self.chatgpt_account_is_fedramp, + ) + .field("path_style", &self.path_style) + .finish_non_exhaustive() + } +} + impl Client { pub fn new(base_url: impl Into) -> Result { let mut base_url = base_url.into(); @@ -147,8 +161,7 @@ impl Client { Ok(Self { base_url, http, - bearer_token: None, - auth_headers: HeaderMap::new(), + auth_provider: codex_model_provider::unauthenticated_auth_provider(), user_agent: None, chatgpt_account_id: None, chatgpt_account_is_fedramp: false, @@ -159,16 +172,11 @@ impl Client { pub fn from_auth(base_url: impl Into, auth: &CodexAuth) -> Result { Ok(Self::new(base_url)? .with_user_agent(get_codex_user_agent()) - .with_auth_provider(codex_model_provider::auth_provider_from_auth(auth).as_ref())) + .with_auth_provider(codex_model_provider::auth_provider_from_auth(auth))) } - pub fn with_auth_provider(mut self, auth: &dyn AuthProvider) -> Self { - auth.add_auth_headers(&mut self.auth_headers); - self - } - - pub fn with_bearer_token(mut self, token: impl Into) -> Self { - self.bearer_token = Some(token.into()); + pub fn with_auth_provider(mut self, auth: SharedAuthProvider) -> Self { + self.auth_provider = auth; self } @@ -201,13 +209,7 @@ 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); - } - } - h.extend(self.auth_headers.clone()); + self.auth_provider.add_auth_headers(&mut h); if let Some(acc) = &self.chatgpt_account_id && let Ok(name) = HeaderName::from_bytes(b"ChatGPT-Account-Id") && let Ok(hv) = HeaderValue::from_str(acc) @@ -820,8 +822,7 @@ mod tests { let codex_client = Client { base_url: "https://example.test".to_string(), http: reqwest::Client::new(), - bearer_token: None, - auth_headers: HeaderMap::new(), + auth_provider: codex_model_provider::unauthenticated_auth_provider(), user_agent: None, chatgpt_account_id: None, chatgpt_account_is_fedramp: false, @@ -835,8 +836,7 @@ mod tests { let chatgpt_client = Client { base_url: "https://chatgpt.com/backend-api".to_string(), http: reqwest::Client::new(), - bearer_token: None, - auth_headers: HeaderMap::new(), + auth_provider: codex_model_provider::unauthenticated_auth_provider(), user_agent: None, chatgpt_account_id: None, chatgpt_account_is_fedramp: false, diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index d413f72ddc..c5b4751322 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -486,8 +486,12 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> let mut entries: Vec<_> = mcp_servers.iter().collect(); entries.sort_by(|(a, _), (b, _)| a.cmp(b)); - let auth_statuses = - compute_auth_statuses(mcp_servers.iter(), config.mcp_oauth_credentials_store_mode).await; + let auth_statuses = compute_auth_statuses( + mcp_servers.iter(), + config.mcp_oauth_credentials_store_mode, + /*auth*/ None, + ) + .await; if list_args.json { let json_entries: Vec<_> = entries diff --git a/codex-rs/cloud-tasks-client/src/http.rs b/codex-rs/cloud-tasks-client/src/http.rs index 8742092435..46fed812ba 100644 --- a/codex-rs/cloud-tasks-client/src/http.rs +++ b/codex-rs/cloud-tasks-client/src/http.rs @@ -14,7 +14,7 @@ use crate::api::TaskText; use chrono::DateTime; use chrono::Utc; -use codex_api::AuthProvider; +use codex_api::SharedAuthProvider; use codex_backend_client as backend; use codex_backend_client::CodeTaskDetailsResponseExt; use codex_git_utils::ApplyGitRequest; @@ -33,17 +33,12 @@ impl HttpClient { Ok(Self { base_url, backend }) } - pub fn with_bearer_token(mut self, token: impl Into) -> Self { - self.backend = self.backend.clone().with_bearer_token(token); - self - } - pub fn with_user_agent(mut self, ua: impl Into) -> Self { self.backend = self.backend.clone().with_user_agent(ua); self } - pub fn with_auth_provider(mut self, auth: &dyn AuthProvider) -> Self { + pub fn with_auth_provider(mut self, auth: SharedAuthProvider) -> Self { self.backend = self.backend.clone().with_auth_provider(auth); self } diff --git a/codex-rs/cloud-tasks/src/lib.rs b/codex-rs/cloud-tasks/src/lib.rs index 9cc45f36c6..e8d6b545b5 100644 --- a/codex-rs/cloud-tasks/src/lib.rs +++ b/codex-rs/cloud-tasks/src/lib.rs @@ -68,7 +68,7 @@ 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_manager = util::load_auth_manager(Some(base_url.clone())).await; let auth = match auth_manager.as_ref() { Some(manager) => manager.auth().await, None => None, @@ -95,7 +95,7 @@ async fn init_backend(user_agent_suffix: &str) -> anyhow::Result } let auth_provider = codex_model_provider::auth_provider_from_auth(&auth); - http = http.with_auth_provider(auth_provider.as_ref()); + http = http.with_auth_provider(auth_provider); if let Some(acc) = auth.get_account_id() { append_error_log(format!("auth: set ChatGPT-Account-Id header: {acc}")); } diff --git a/codex-rs/cloud-tasks/src/util.rs b/codex-rs/cloud-tasks/src/util.rs index fc94b268de..e433b892e5 100644 --- a/codex-rs/cloud-tasks/src/util.rs +++ b/codex-rs/cloud-tasks/src/util.rs @@ -41,14 +41,14 @@ pub fn normalize_base_url(input: &str) -> String { base_url } -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, - Some(config.chatgpt_base_url), + chatgpt_base_url.or(Some(config.chatgpt_base_url)), )) } @@ -65,7 +65,7 @@ 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 + if let Some(am) = load_auth_manager(/*chatgpt_base_url*/ None).await && let Some(auth) = am.auth().await && auth.uses_codex_backend() { diff --git a/codex-rs/codex-mcp/Cargo.toml b/codex-rs/codex-mcp/Cargo.toml index f99f4b7902..a9aacb1925 100644 --- a/codex-rs/codex-mcp/Cargo.toml +++ b/codex-rs/codex-mcp/Cargo.toml @@ -15,6 +15,7 @@ workspace = true anyhow = { workspace = true } async-channel = { workspace = true } codex-async-utils = { workspace = true } +codex-api = { workspace = true } codex-config = { workspace = true } codex-exec-server = { workspace = true } codex-login = { workspace = true } diff --git a/codex-rs/codex-mcp/src/mcp/auth.rs b/codex-rs/codex-mcp/src/mcp/auth.rs index 27d7e13358..9c605c16fb 100644 --- a/codex-rs/codex-mcp/src/mcp/auth.rs +++ b/codex-rs/codex-mcp/src/mcp/auth.rs @@ -1,7 +1,10 @@ use std::collections::HashMap; use anyhow::Result; +use codex_config::McpServerConfig; +use codex_config::McpServerTransportConfig; use codex_config::types::OAuthCredentialsStoreMode; +use codex_login::CodexAuth; use codex_protocol::protocol::McpAuthStatus; use codex_rmcp_client::OAuthProviderError; use codex_rmcp_client::determine_streamable_http_auth_status; @@ -9,8 +12,7 @@ use codex_rmcp_client::discover_streamable_http_oauth; use futures::future::join_all; use tracing::warn; -use codex_config::McpServerConfig; -use codex_config::McpServerTransportConfig; +use super::CODEX_APPS_MCP_SERVER_NAME; #[derive(Debug, Clone)] pub struct McpOAuthLoginConfig { @@ -126,6 +128,7 @@ pub struct McpAuthStatusEntry { pub async fn compute_auth_statuses<'a, I>( servers: I, store_mode: OAuthCredentialsStoreMode, + auth: Option<&CodexAuth>, ) -> HashMap where I: IntoIterator, @@ -133,14 +136,24 @@ where let futures = servers.into_iter().map(|(name, config)| { let name = name.clone(); let config = config.clone(); - async move { - let auth_status = match compute_auth_status(&name, &config, store_mode).await { - Ok(status) => status, - Err(error) => { - warn!("failed to determine auth status for MCP server `{name}`: {error:?}"); - McpAuthStatus::Unsupported + let has_runtime_auth = name == CODEX_APPS_MCP_SERVER_NAME + && auth.is_some_and(CodexAuth::uses_codex_backend) + && matches!( + &config.transport, + McpServerTransportConfig::StreamableHttp { + bearer_token_env_var: None, + .. } - }; + ); + async move { + let auth_status = + match compute_auth_status(&name, &config, store_mode, has_runtime_auth).await { + Ok(status) => status, + Err(error) => { + warn!("failed to determine auth status for MCP server `{name}`: {error:?}"); + McpAuthStatus::Unsupported + } + }; let entry = McpAuthStatusEntry { config, auth_status, @@ -156,11 +169,16 @@ async fn compute_auth_status( server_name: &str, config: &McpServerConfig, store_mode: OAuthCredentialsStoreMode, + has_runtime_auth: bool, ) -> Result { if !config.enabled { return Ok(McpAuthStatus::Unsupported); } + if has_runtime_auth { + return Ok(McpAuthStatus::BearerToken); + } + match &config.transport { McpServerTransportConfig::Stdio { .. } => Ok(McpAuthStatus::Unsupported), McpServerTransportConfig::StreamableHttp { diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 8e6aea91cc..1061a6a542 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -205,22 +205,6 @@ fn codex_apps_mcp_bearer_token_env_var() -> Option { } } -fn codex_apps_mcp_http_headers(auth: Option<&CodexAuth>) -> Option> { - let auth = auth.filter(|auth| auth.uses_codex_backend())?; - let headers = codex_model_provider::auth_provider_from_auth(auth).to_auth_headers(); - let headers: HashMap<_, _> = headers - .iter() - .filter_map(|(name, value)| { - Some((name.as_str().to_string(), value.to_str().ok()?.to_string())) - }) - .collect(); - if headers.is_empty() { - None - } else { - Some(headers) - } -} - fn normalize_codex_apps_base_url(base_url: &str) -> String { let mut base_url = base_url.trim_end_matches('/').to_string(); if (base_url.starts_with("https://chatgpt.com") @@ -247,20 +231,14 @@ pub(crate) fn codex_apps_mcp_url(config: &McpConfig) -> String { codex_apps_mcp_url_for_base_url(&config.chatgpt_base_url) } -fn codex_apps_mcp_server_config(config: &McpConfig, auth: Option<&CodexAuth>) -> McpServerConfig { - let bearer_token_env_var = codex_apps_mcp_bearer_token_env_var(); - let http_headers = if bearer_token_env_var.is_some() { - None - } else { - codex_apps_mcp_http_headers(auth) - }; +fn codex_apps_mcp_server_config(config: &McpConfig) -> McpServerConfig { let url = codex_apps_mcp_url(config); McpServerConfig { transport: McpServerTransportConfig::StreamableHttp { url, - bearer_token_env_var, - http_headers, + bearer_token_env_var: codex_apps_mcp_bearer_token_env_var(), + http_headers: None, env_http_headers: None, }, experimental_environment: None, @@ -287,7 +265,7 @@ pub fn with_codex_apps_mcp( if config.apps_enabled && auth.is_some_and(CodexAuth::uses_codex_backend) { servers.insert( CODEX_APPS_MCP_SERVER_NAME.to_string(), - codex_apps_mcp_server_config(config, auth), + codex_apps_mcp_server_config(config), ); } else { servers.remove(CODEX_APPS_MCP_SERVER_NAME); @@ -320,8 +298,12 @@ pub async fn read_mcp_resource( ) -> anyhow::Result { let mut mcp_servers = effective_mcp_servers(config, auth); mcp_servers.retain(|name, _| name == server); - let auth_statuses = - compute_auth_statuses(mcp_servers.iter(), config.mcp_oauth_credentials_store_mode).await; + let auth_statuses = compute_auth_statuses( + mcp_servers.iter(), + config.mcp_oauth_credentials_store_mode, + auth, + ) + .await; let (tx_event, rx_event) = unbounded(); drop(rx_event); let (manager, cancel_token) = McpConnectionManager::new( @@ -336,6 +318,7 @@ pub async fn read_mcp_resource( config.codex_home.clone(), codex_apps_tools_cache_key(auth), tool_plugin_provenance(config), + auth, ) .await; @@ -386,8 +369,12 @@ pub async fn collect_mcp_snapshot_with_detail( }; } - let auth_status_entries = - compute_auth_statuses(mcp_servers.iter(), config.mcp_oauth_credentials_store_mode).await; + let auth_status_entries = compute_auth_statuses( + mcp_servers.iter(), + config.mcp_oauth_credentials_store_mode, + auth, + ) + .await; let (tx_event, rx_event) = unbounded(); drop(rx_event); @@ -404,6 +391,7 @@ pub async fn collect_mcp_snapshot_with_detail( config.codex_home.clone(), codex_apps_tools_cache_key(auth), tool_plugin_provenance, + auth, ) .await; @@ -461,8 +449,12 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( }; } - let auth_status_entries = - compute_auth_statuses(mcp_servers.iter(), config.mcp_oauth_credentials_store_mode).await; + let auth_status_entries = compute_auth_statuses( + mcp_servers.iter(), + config.mcp_oauth_credentials_store_mode, + auth, + ) + .await; let (tx_event, rx_event) = unbounded(); drop(rx_event); @@ -479,6 +471,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( config.codex_home.clone(), codex_apps_tools_cache_key(auth), tool_plugin_provenance, + auth, ) .await; diff --git a/codex-rs/codex-mcp/src/mcp_connection_manager.rs b/codex-rs/codex-mcp/src/mcp_connection_manager.rs index 8b8d3fe35c..6ada6db8cf 100644 --- a/codex-rs/codex-mcp/src/mcp_connection_manager.rs +++ b/codex-rs/codex-mcp/src/mcp_connection_manager.rs @@ -32,6 +32,7 @@ use anyhow::Context; use anyhow::Result; use anyhow::anyhow; use async_channel::Sender; +use codex_api::SharedAuthProvider; use codex_async_utils::CancelErr; use codex_async_utils::OrCancelExt; use codex_config::Constrained; @@ -485,6 +486,7 @@ impl AsyncManagedClient { codex_apps_tools_cache_context: Option, tool_plugin_provenance: Arc, runtime_environment: McpRuntimeEnvironment, + runtime_auth_provider: Option, ) -> Self { let tool_filter = ToolFilter::from_config(&config); let startup_snapshot = load_startup_cached_codex_apps_tools_snapshot( @@ -507,6 +509,7 @@ impl AsyncManagedClient { config.clone(), store_mode, runtime_environment, + runtime_auth_provider, ) .await?, ); @@ -744,6 +747,7 @@ impl McpConnectionManager { codex_home: PathBuf, codex_apps_tools_cache_key: CodexAppsToolsCacheKey, tool_plugin_provenance: ToolPluginProvenance, + auth: Option<&CodexAuth>, ) -> (Self, CancellationToken) { let cancel_token = CancellationToken::new(); let mut clients = HashMap::new(); @@ -753,6 +757,9 @@ impl McpConnectionManager { ElicitationRequestManager::new(approval_policy.value(), initial_sandbox_policy); let tool_plugin_provenance = Arc::new(tool_plugin_provenance); let startup_submit_id = submit_id.clone(); + let codex_apps_auth_provider = auth + .filter(|auth| auth.uses_codex_backend()) + .map(codex_model_provider::auth_provider_from_auth); let mcp_servers = mcp_servers.clone(); for (server_name, cfg) in mcp_servers.into_iter().filter(|(_, cfg)| cfg.enabled) { if let Some(origin) = transport_origin(&cfg.transport) { @@ -776,6 +783,19 @@ impl McpConnectionManager { } else { None }; + let uses_env_bearer_token = match &cfg.transport { + McpServerTransportConfig::StreamableHttp { + bearer_token_env_var, + .. + } => bearer_token_env_var.is_some(), + McpServerTransportConfig::Stdio { .. } => false, + }; + let runtime_auth_provider = + if server_name == CODEX_APPS_MCP_SERVER_NAME && !uses_env_bearer_token { + codex_apps_auth_provider.clone() + } else { + None + }; let async_managed_client = AsyncManagedClient::new( server_name.clone(), cfg, @@ -786,6 +806,7 @@ impl McpConnectionManager { codex_apps_tools_cache_context, Arc::clone(&tool_plugin_provenance), runtime_environment.clone(), + runtime_auth_provider, ); clients.insert(server_name.clone(), async_managed_client.clone()); let tx_event = tx_event.clone(); @@ -1519,6 +1540,7 @@ async fn make_rmcp_client( config: McpServerConfig, store_mode: OAuthCredentialsStoreMode, runtime_environment: McpRuntimeEnvironment, + runtime_auth_provider: Option, ) -> Result { let McpServerConfig { transport, @@ -1607,6 +1629,7 @@ async fn make_rmcp_client( http_headers, env_http_headers, store_mode, + runtime_auth_provider, ) .await .map_err(StartupOutcomeError::from) diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 803dd592fe..968b93214c 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -8,6 +8,7 @@ use std::time::Instant; use anyhow::Context; use async_channel::unbounded; +use codex_api::SharedAuthProvider; pub use codex_app_server_protocol::AppBranding; pub use codex_app_server_protocol::AppInfo; pub use codex_app_server_protocol::AppMetadata; @@ -252,8 +253,12 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager( }); } - let auth_status_entries = - compute_auth_statuses(mcp_servers.iter(), config.mcp_oauth_credentials_store_mode).await; + let auth_status_entries = compute_auth_statuses( + mcp_servers.iter(), + config.mcp_oauth_credentials_store_mode, + auth.as_ref(), + ) + .await; let (tx_event, rx_event) = unbounded(); drop(rx_event); @@ -274,6 +279,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager( config.codex_home.to_path_buf(), codex_apps_tools_cache_key(auth.as_ref()), ToolPluginProvenance::default(), + auth.as_ref(), ) .await; @@ -440,7 +446,7 @@ async fn list_directory_connectors_for_tool_suggest_with_auth( Some(account_id) if !account_id.is_empty() => account_id, _ => return Ok(Vec::new()), }; - let auth_headers = codex_model_provider::auth_provider_from_auth(auth).to_auth_headers(); + let auth_provider = codex_model_provider::auth_provider_from_auth(auth); let is_workspace_account = auth.is_workspace_account(); let cache_key = AllConnectorsCacheKey::new( config.chatgpt_base_url.clone(), @@ -454,12 +460,12 @@ async fn list_directory_connectors_for_tool_suggest_with_auth( is_workspace_account, /*force_refetch*/ false, |path| { - let auth_headers = auth_headers.clone(); + let auth_provider = auth_provider.clone(); async move { - chatgpt_get_request_with_auth_headers::( + chatgpt_get_request_with_auth_provider::( config, path, - auth_headers, + auth_provider, ) .await } @@ -468,16 +474,16 @@ async fn list_directory_connectors_for_tool_suggest_with_auth( .await } -async fn chatgpt_get_request_with_auth_headers( +async fn chatgpt_get_request_with_auth_provider( config: &Config, path: String, - auth_headers: http::HeaderMap, + auth_provider: SharedAuthProvider, ) -> anyhow::Result { let client = create_client(); let url = format!("{}{}", config.chatgpt_base_url, path); let response = client .get(&url) - .headers(auth_headers) + .headers(auth_provider.to_auth_headers()) .header("Content-Type", "application/json") .timeout(DIRECTORY_CONNECTORS_TIMEOUT) .send() diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 5d262564c4..ec559ee88a 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -489,7 +489,12 @@ pub async fn list_mcp_tools(sess: &Session, config: &Arc, sub_id: String .await; let snapshot = collect_mcp_snapshot_from_manager( &mcp_connection_manager, - compute_auth_statuses(mcp_servers.iter(), config.mcp_oauth_credentials_store_mode).await, + compute_auth_statuses( + mcp_servers.iter(), + config.mcp_oauth_credentials_store_mode, + auth.as_ref(), + ) + .await, ) .await; let event = Event { diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 2e4a3301ea..9d889e743a 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -219,7 +219,8 @@ impl Session { .tool_plugin_provenance(config.as_ref()) .await; let mcp_servers = with_codex_apps_mcp(mcp_servers, auth.as_ref(), &mcp_config); - let auth_statuses = compute_auth_statuses(mcp_servers.iter(), store_mode).await; + let auth_statuses = + compute_auth_statuses(mcp_servers.iter(), store_mode, auth.as_ref()).await; { let mut guard = self.services.mcp_startup_cancellation_token.lock().await; guard.cancel(); @@ -243,6 +244,7 @@ impl Session { config.codex_home.to_path_buf(), codex_apps_tools_cache_key(auth.as_ref()), tool_plugin_provenance, + auth.as_ref(), ) .await; { diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 876008d27a..8fc8d43d7e 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -350,6 +350,7 @@ impl Session { let auth_statuses = compute_auth_statuses( mcp_servers.iter(), config_for_mcp.mcp_oauth_credentials_store_mode, + auth.as_ref(), ) .await; (auth, mcp_servers, auth_statuses) @@ -787,6 +788,7 @@ impl Session { config.codex_home.to_path_buf(), codex_apps_tools_cache_key(auth), tool_plugin_provenance, + auth, ) .instrument(info_span!( "session_init.mcp_manager_init", diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index 7de9ec68c8..d9e31e7827 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -55,31 +55,19 @@ impl AuthProvider for AgentIdentityAuthProvider { } } -#[derive(Clone, Debug)] -struct CodexBearerAuthProvider { - auth: CodexAuth, -} - -impl AuthProvider for CodexBearerAuthProvider { - fn add_auth_headers(&self, headers: &mut HeaderMap) { - BearerAuthProvider { - token: self.auth.get_token().ok(), - account_id: self.auth.get_account_id(), - is_fedramp_account: self.auth.is_fedramp_account(), - } - .add_auth_headers(headers); - } -} - // Some providers are meant to send no auth headers. Examples include local OSS // providers and custom test providers with `requires_openai_auth = false`. #[derive(Clone, Debug)] -struct NoAuthProvider; +struct UnauthenticatedAuthProvider; -impl AuthProvider for NoAuthProvider { +impl AuthProvider for UnauthenticatedAuthProvider { fn add_auth_headers(&self, _headers: &mut HeaderMap) {} } +pub fn unauthenticated_auth_provider() -> SharedAuthProvider { + Arc::new(UnauthenticatedAuthProvider) +} + /// Returns the provider-scoped auth manager when this provider uses command-backed auth. /// /// Providers without custom auth continue using the caller-supplied base manager, when present. @@ -103,7 +91,7 @@ pub(crate) fn resolve_provider_auth( Ok(match auth { Some(auth) => auth_provider_from_auth(auth), - None => Arc::new(NoAuthProvider), + None => unauthenticated_auth_provider(), }) } @@ -128,7 +116,11 @@ pub fn auth_provider_from_auth(auth: &CodexAuth) -> SharedAuthProvider { Arc::new(AgentIdentityAuthProvider { auth: auth.clone() }) } CodexAuth::ApiKey(_) | CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => { - Arc::new(CodexBearerAuthProvider { auth: auth.clone() }) + Arc::new(BearerAuthProvider { + token: auth.get_token().ok(), + account_id: auth.get_account_id(), + is_fedramp_account: auth.is_fedramp_account(), + }) } } } @@ -141,7 +133,7 @@ mod tests { use super::*; #[test] - fn no_auth_provider_adds_no_headers() { + fn unauthenticated_auth_provider_adds_no_headers() { let provider = create_oss_provider_with_base_url("http://localhost:11434/v1", WireApi::Responses); let auth = resolve_provider_auth(/*auth*/ None, &provider).expect("auth should resolve"); diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index 0db7b9c8d6..f5454edd3f 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -4,6 +4,7 @@ mod bearer_auth_provider; mod provider; pub use auth::auth_provider_from_auth; +pub use auth::unauthenticated_auth_provider; pub use bearer_auth_provider::BearerAuthProvider; pub use bearer_auth_provider::BearerAuthProvider as CoreAuthProvider; pub use provider::ModelProvider; diff --git a/codex-rs/protocol/src/account.rs b/codex-rs/protocol/src/account.rs index cc6632a509..5832381cbc 100644 --- a/codex-rs/protocol/src/account.rs +++ b/codex-rs/protocol/src/account.rs @@ -37,9 +37,15 @@ impl PlanType { } pub fn is_workspace_account(self) -> bool { - self.is_team_like() - || self.is_business_like() - || matches!(self, Self::Enterprise | Self::Edu) + matches!( + self, + Self::Team + | Self::SelfServeBusinessUsageBased + | Self::Business + | Self::EnterpriseCbpUsageBased + | Self::Enterprise + | Self::Edu + ) } } @@ -89,7 +95,10 @@ mod tests { assert_eq!(PlanType::Business.is_business_like(), true); assert_eq!(PlanType::EnterpriseCbpUsageBased.is_business_like(), true); assert_eq!(PlanType::Team.is_business_like(), false); + } + #[test] + fn workspace_account_helper_includes_usage_based_workspace_plans() { assert_eq!(PlanType::Team.is_workspace_account(), true); assert_eq!( PlanType::SelfServeBusinessUsageBased.is_workspace_account(), diff --git a/codex-rs/protocol/src/auth.rs b/codex-rs/protocol/src/auth.rs index 604f47855b..99e067bf24 100644 --- a/codex-rs/protocol/src/auth.rs +++ b/codex-rs/protocol/src/auth.rs @@ -50,14 +50,6 @@ pub enum KnownPlan { } impl KnownPlan { - pub fn is_team_like(self) -> bool { - matches!(self, Self::Team | Self::SelfServeBusinessUsageBased) - } - - pub fn is_business_like(self) -> bool { - matches!(self, Self::Business | Self::EnterpriseCbpUsageBased) - } - pub fn display_name(self) -> &'static str { match self { Self::Free => "Free", @@ -91,9 +83,15 @@ impl KnownPlan { } pub fn is_workspace_account(self) -> bool { - self.is_team_like() - || self.is_business_like() - || matches!(self, Self::Enterprise | Self::Edu) + matches!( + self, + Self::Team + | Self::SelfServeBusinessUsageBased + | Self::Business + | Self::EnterpriseCbpUsageBased + | Self::Enterprise + | Self::Edu + ) } } diff --git a/codex-rs/protocol/src/error.rs b/codex-rs/protocol/src/error.rs index d2601f2424..f421db8af5 100644 --- a/codex-rs/protocol/src/error.rs +++ b/codex-rs/protocol/src/error.rs @@ -472,44 +472,40 @@ impl std::fmt::Display for UsageLimitReachedError { } let message = match self.plan_type.as_ref() { - Some(PlanType::Known(plan)) => match plan { - KnownPlan::Plus => format!( - "You've hit your usage limit. Upgrade to Pro (https://chatgpt.com/explore/pro), visit https://chatgpt.com/codex/settings/usage to purchase more credits{}", - retry_suffix_after_or(self.resets_at.as_ref()) - ), - plan if plan.is_team_like() || plan.is_business_like() => { - format!( - "You've hit your usage limit. To get more access now, send a request to your admin{}", - retry_suffix_after_or(self.resets_at.as_ref()) - ) - } - KnownPlan::Free | KnownPlan::Go => { - format!( - "You've hit your usage limit. Upgrade to Plus to continue using Codex (https://chatgpt.com/explore/plus),{}", - retry_suffix_after_or(self.resets_at.as_ref()) - ) - } - KnownPlan::Pro | KnownPlan::ProLite => format!( - "You've hit your usage limit. Visit https://chatgpt.com/codex/settings/usage to purchase more credits{}", - retry_suffix_after_or(self.resets_at.as_ref()) - ), - KnownPlan::Enterprise | KnownPlan::Edu => format!( - "You've hit your usage limit.{}", - retry_suffix(self.resets_at.as_ref()) - ), + Some(PlanType::Known(KnownPlan::Plus)) => format!( + "You've hit your usage limit. Upgrade to Pro (https://chatgpt.com/explore/pro), visit https://chatgpt.com/codex/settings/usage to purchase more credits{}", + retry_suffix_after_or(self.resets_at.as_ref()) + ), + Some(PlanType::Known( KnownPlan::Team | KnownPlan::SelfServeBusinessUsageBased | KnownPlan::Business - | KnownPlan::EnterpriseCbpUsageBased => { - unreachable!("team-like and business-like plans are handled above") - } - }, - Some(PlanType::Unknown(_)) | None => { + | KnownPlan::EnterpriseCbpUsageBased, + )) => { format!( - "You've hit your usage limit.{}", - retry_suffix(self.resets_at.as_ref()) + "You've hit your usage limit. To get more access now, send a request to your admin{}", + retry_suffix_after_or(self.resets_at.as_ref()) ) } + Some(PlanType::Known(KnownPlan::Free)) | Some(PlanType::Known(KnownPlan::Go)) => { + format!( + "You've hit your usage limit. Upgrade to Plus to continue using Codex (https://chatgpt.com/explore/plus),{}", + retry_suffix_after_or(self.resets_at.as_ref()) + ) + } + Some(PlanType::Known(KnownPlan::Pro | KnownPlan::ProLite)) => format!( + "You've hit your usage limit. Visit https://chatgpt.com/codex/settings/usage to purchase more credits{}", + retry_suffix_after_or(self.resets_at.as_ref()) + ), + Some(PlanType::Known(KnownPlan::Enterprise)) + | Some(PlanType::Known(KnownPlan::Edu)) => format!( + "You've hit your usage limit.{}", + retry_suffix(self.resets_at.as_ref()) + ), + Some(PlanType::Unknown(_)) | None => format!( + "You've hit your usage limit.{}", + retry_suffix(self.resets_at.as_ref()) + ), }; write!(f, "{message}") diff --git a/codex-rs/rmcp-client/Cargo.toml b/codex-rs/rmcp-client/Cargo.toml index b81c224a44..f29ade2fd7 100644 --- a/codex-rs/rmcp-client/Cargo.toml +++ b/codex-rs/rmcp-client/Cargo.toml @@ -13,6 +13,7 @@ axum = { workspace = true, default-features = false, features = [ "http1", "tokio", ] } +codex-api = { workspace = true } codex-client = { workspace = true } codex-config = { workspace = true } codex-exec-server = { workspace = true } diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 50891ec8c0..d49169f2fd 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -12,6 +12,7 @@ use std::time::Instant; use anyhow::Result; use anyhow::anyhow; +use codex_api::SharedAuthProvider; use codex_client::build_reqwest_client_with_custom_ca; use codex_config::types::McpServerEnvVar; use futures::FutureExt; @@ -88,11 +89,15 @@ const NON_JSON_RESPONSE_BODY_PREVIEW_BYTES: usize = 8_192; #[derive(Clone)] struct StreamableHttpResponseClient { inner: reqwest::Client, + auth_provider: Option, } impl StreamableHttpResponseClient { - fn new(inner: reqwest::Client) -> Self { - Self { inner } + fn new(inner: reqwest::Client, auth_provider: Option) -> Self { + Self { + inner, + auth_provider, + } } fn reqwest_error( @@ -129,6 +134,9 @@ impl StreamableHttpClient for StreamableHttpResponseClient { .inner .post(uri.as_ref()) .header(ACCEPT, [EVENT_STREAM_MIME_TYPE, JSON_MIME_TYPE].join(", ")); + if let Some(auth_provider) = &self.auth_provider { + request = request.headers(auth_provider.to_auth_headers()); + } if let Some(auth_header) = auth_token { request = request.bearer_auth(auth_header); } @@ -339,6 +347,7 @@ enum TransportRecipe { http_headers: Option>, env_http_headers: Option>, store_mode: OAuthCredentialsStoreMode, + auth_provider: Option, }, } @@ -536,6 +545,7 @@ impl RmcpClient { http_headers: Option>, env_http_headers: Option>, store_mode: OAuthCredentialsStoreMode, + auth_provider: Option, ) -> Result { let transport_recipe = TransportRecipe::StreamableHttp { server_name: server_name.to_string(), @@ -544,6 +554,7 @@ impl RmcpClient { http_headers, env_http_headers, store_mode, + auth_provider, }; let transport = Self::create_pending_transport(&transport_recipe).await?; Ok(Self { @@ -895,22 +906,25 @@ impl RmcpClient { http_headers, env_http_headers, store_mode, + auth_provider, } => { let default_headers = build_default_headers(http_headers.clone(), env_http_headers.clone())?; - let initial_oauth_tokens = - if bearer_token.is_none() && !default_headers.contains_key(AUTHORIZATION) { - match load_oauth_tokens(server_name, url, *store_mode) { - Ok(tokens) => tokens, - Err(err) => { - warn!("failed to read tokens for server `{server_name}`: {err}"); - None - } + let initial_oauth_tokens = if bearer_token.is_none() + && auth_provider.is_none() + && !default_headers.contains_key(AUTHORIZATION) + { + match load_oauth_tokens(server_name, url, *store_mode) { + Ok(tokens) => tokens, + Err(err) => { + warn!("failed to read tokens for server `{server_name}`: {err}"); + None } - } else { - None - }; + } + } else { + None + }; if let Some(initial_tokens) = initial_oauth_tokens.clone() { match create_oauth_transport_and_runtime( @@ -947,7 +961,7 @@ impl RmcpClient { .auth_header(access_token); let http_client = build_http_client(&default_headers)?; let transport = StreamableHttpClientTransport::with_client( - StreamableHttpResponseClient::new(http_client), + StreamableHttpResponseClient::new(http_client, None), http_config, ); Ok(PendingTransport::StreamableHttp { transport }) @@ -964,7 +978,7 @@ impl RmcpClient { let http_client = build_http_client(&default_headers)?; let transport = StreamableHttpClientTransport::with_client( - StreamableHttpResponseClient::new(http_client), + StreamableHttpResponseClient::new(http_client, auth_provider.clone()), http_config, ); Ok(PendingTransport::StreamableHttp { transport }) @@ -1178,7 +1192,10 @@ async fn create_oauth_transport_and_runtime( } }; - let auth_client = AuthClient::new(StreamableHttpResponseClient::new(http_client), manager); + let auth_client = AuthClient::new( + StreamableHttpResponseClient::new(http_client, None), + manager, + ); let auth_manager = auth_client.auth_manager.clone(); let transport = StreamableHttpClientTransport::with_client( @@ -1199,13 +1216,67 @@ async fn create_oauth_transport_and_runtime( #[cfg(test)] mod tests { + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; use std::time::Duration; + use codex_api::AuthProvider; use pretty_assertions::assert_eq; + use reqwest::header::HeaderMap; + use rmcp::model::ClientJsonRpcMessage; + use rmcp::model::ClientNotification; + use rmcp::model::InitializedNotification; use tokio::time; use super::*; + #[derive(Debug)] + struct CountingAuthProvider { + calls: Arc, + } + + impl AuthProvider for CountingAuthProvider { + fn add_auth_headers(&self, _headers: &mut HeaderMap) { + self.calls.fetch_add(1, Ordering::SeqCst); + } + } + + #[tokio::test] + async fn streamable_http_client_applies_runtime_auth_per_request() { + let calls = Arc::new(AtomicUsize::new(0)); + let client = StreamableHttpResponseClient::new( + reqwest::Client::new(), + Some(Arc::new(CountingAuthProvider { + calls: Arc::clone(&calls), + })), + ); + let message = || { + ClientJsonRpcMessage::notification(ClientNotification::InitializedNotification( + InitializedNotification::default(), + )) + }; + + let _ = client + .post_message( + Arc::from("http://127.0.0.1:1/mcp"), + message(), + /*session_id*/ None, + /*auth_token*/ None, + ) + .await; + let _ = client + .post_message( + Arc::from("http://127.0.0.1:1/mcp"), + message(), + /*session_id*/ None, + /*auth_token*/ None, + ) + .await; + + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + #[tokio::test] async fn active_time_timeout_pauses_while_elicitation_is_pending() { let pause_state = ElicitationPauseState::new(); diff --git a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs index c0525aafaf..d1f719918b 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_recovery.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_recovery.rs @@ -77,6 +77,7 @@ async fn create_client(base_url: &str) -> anyhow::Result { /*http_headers*/ None, /*env_http_headers*/ None, OAuthCredentialsStoreMode::File, + /*auth_provider*/ None, ) .await?;