From a3abb8ba6ee48ec16fa537fd1a44bdb228753af6 Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Mon, 27 Apr 2026 10:03:34 +0000 Subject: [PATCH] Move MCP OAuth orchestration downstream Keep app-server OAuth callsites thin by moving server/runtime-aware OAuth discovery, scope resolution, and runtime HTTP client selection into codex-mcp helpers. Co-authored-by: Codex --- .../app-server/src/codex_message_processor.rs | 45 +--- .../plugin_mcp_oauth.rs | 69 +----- codex-rs/cli/src/mcp_cmd.rs | 5 +- codex-rs/codex-mcp/src/lib.rs | 6 + codex-rs/codex-mcp/src/mcp/auth.rs | 211 +++++++++++++++++- codex-rs/codex-mcp/src/mcp/mod.rs | 6 + codex-rs/core/src/mcp_skill_dependencies.rs | 77 +------ 7 files changed, 249 insertions(+), 170 deletions(-) diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 15c7492a7e..0194c821a0 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -310,11 +310,9 @@ use codex_mcp::McpRuntimeEnvironment; use codex_mcp::McpServerStatusSnapshot; use codex_mcp::McpSnapshotDetail; use codex_mcp::collect_mcp_server_status_snapshot_with_detail; -use codex_mcp::discover_supported_scopes; use codex_mcp::effective_mcp_servers; -use codex_mcp::http_client_for_server; +use codex_mcp::perform_oauth_login_return_url_for_server; use codex_mcp::read_mcp_resource as read_mcp_resource_without_thread; -use codex_mcp::resolve_oauth_scopes; use codex_model_provider::ProviderAccountError; use codex_model_provider::create_model_provider; use codex_models_manager::collaboration_mode_presets::CollaborationModesConfig; @@ -356,7 +354,6 @@ use codex_protocol::protocol::USER_MESSAGE_BEGIN; use codex_protocol::protocol::W3cTraceContext; use codex_protocol::user_input::MAX_USER_INPUT_TEXT_CHARS; use codex_protocol::user_input::UserInput as CoreInputItem; -use codex_rmcp_client::perform_oauth_login_return_url_with_client; use codex_rollout::state_db::StateDbHandle; use codex_rollout::state_db::get_state_db; use codex_rollout::state_db::reconcile_rollout; @@ -5942,13 +5939,8 @@ impl CodexMessageProcessor { return; }; - let (url, http_headers, env_http_headers) = match &server.transport { - McpServerTransportConfig::StreamableHttp { - url, - http_headers, - env_http_headers, - .. - } => (url.clone(), http_headers.clone(), env_http_headers.clone()), + match &server.transport { + McpServerTransportConfig::StreamableHttp { .. } => {} _ => { let error = JSONRPCErrorError { code: INVALID_REQUEST_ERROR_CODE, @@ -5969,39 +5961,16 @@ impl CodexMessageProcessor { config.cwd.to_path_buf(), ), }; - let http_client = match http_client_for_server(server, runtime_environment) { - Ok(http_client) => http_client, - Err(err) => { - let error = JSONRPCErrorError { - code: INVALID_REQUEST_ERROR_CODE, - message: format!("failed to resolve MCP server OAuth environment: {err}"), - data: None, - }; - self.outgoing.send_error(request_id, error).await; - return; - } - }; - let discovered_scopes = if scopes.is_none() && server.scopes.is_none() { - discover_supported_scopes(&server.transport, http_client.clone()).await - } else { - None - }; - let resolved_scopes = - resolve_oauth_scopes(scopes, server.scopes.clone(), discovered_scopes); - - match perform_oauth_login_return_url_with_client( + match perform_oauth_login_return_url_for_server( &name, - &url, + server, config.mcp_oauth_credentials_store_mode, - http_headers, - env_http_headers, - &resolved_scopes.scopes, - server.oauth_resource.as_deref(), + scopes, timeout_secs, config.mcp_oauth_callback_port, config.mcp_oauth_callback_url.as_deref(), - http_client, + runtime_environment, ) .await { diff --git a/codex-rs/app-server/src/codex_message_processor/plugin_mcp_oauth.rs b/codex-rs/app-server/src/codex_message_processor/plugin_mcp_oauth.rs index 30b7c47799..a2afc0bf44 100644 --- a/codex-rs/app-server/src/codex_message_processor/plugin_mcp_oauth.rs +++ b/codex-rs/app-server/src/codex_message_processor/plugin_mcp_oauth.rs @@ -5,14 +5,9 @@ use codex_app_server_protocol::McpServerOauthLoginCompletedNotification; use codex_app_server_protocol::ServerNotification; use codex_config::types::McpServerConfig; use codex_core::config::Config; -use codex_mcp::McpOAuthLoginSupport; +use codex_mcp::McpOAuthLoginOutcome; use codex_mcp::McpRuntimeEnvironment; -use codex_mcp::http_client_for_server; -use codex_mcp::oauth_login_support; -use codex_mcp::resolve_oauth_scopes; -use codex_mcp::should_retry_without_scopes; -use codex_rmcp_client::perform_oauth_login_silent_with_client; -use tracing::warn; +use codex_mcp::perform_oauth_login_silent_for_server; use super::CodexMessageProcessor; @@ -33,33 +28,6 @@ impl CodexMessageProcessor { config.cwd.to_path_buf(), ), }; - let http_client = match http_client_for_server(&server, runtime_environment) { - Ok(http_client) => http_client, - Err(err) => { - warn!( - "failed to resolve MCP OAuth environment for plugin install {name}: {err}" - ); - continue; - } - }; - let oauth_config = match oauth_login_support(&server.transport, http_client.clone()) - .await - { - McpOAuthLoginSupport::Supported(config) => config, - McpOAuthLoginSupport::Unsupported => continue, - McpOAuthLoginSupport::Unknown(err) => { - warn!( - "MCP server may or may not require login for plugin install {name}: {err}" - ); - continue; - } - }; - - let resolved_scopes = resolve_oauth_scopes( - /*explicit_scopes*/ None, - server.scopes.clone(), - oauth_config.discovered_scopes.clone(), - ); let store_mode = config.mcp_oauth_credentials_store_mode; let callback_port = config.mcp_oauth_callback_port; @@ -68,41 +36,20 @@ impl CodexMessageProcessor { let notification_name = name.clone(); tokio::spawn(async move { - let first_attempt = perform_oauth_login_silent_with_client( + let final_result = perform_oauth_login_silent_for_server( &name, - &oauth_config.url, + &server, store_mode, - oauth_config.http_headers.clone(), - oauth_config.env_http_headers.clone(), - &resolved_scopes.scopes, - server.oauth_resource.as_deref(), + /*explicit_scopes*/ None, callback_port, callback_url.as_deref(), - http_client.clone(), + runtime_environment, ) .await; - let final_result = match first_attempt { - Err(err) if should_retry_without_scopes(&resolved_scopes, &err) => { - perform_oauth_login_silent_with_client( - &name, - &oauth_config.url, - store_mode, - oauth_config.http_headers, - oauth_config.env_http_headers, - &[], - server.oauth_resource.as_deref(), - callback_port, - callback_url.as_deref(), - http_client, - ) - .await - } - result => result, - }; - let (success, error) = match final_result { - Ok(()) => (true, None), + Ok(McpOAuthLoginOutcome::Completed) => (true, None), + Ok(McpOAuthLoginOutcome::Unsupported) => return, Err(err) => (false, Some(err.to_string())), }; diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index c795090658..d6d79b3daa 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -16,7 +16,6 @@ use codex_core::config::find_codex_home; use codex_core::config::load_global_mcp_servers; use codex_core::plugins::PluginsManager; use codex_exec_server::Environment; -use codex_exec_server::ReqwestHttpClient; use codex_mcp::McpOAuthLoginSupport; use codex_mcp::McpRuntimeEnvironment; use codex_mcp::ResolvedMcpOAuthScopes; @@ -326,7 +325,7 @@ async fn run_add(config_overrides: &CliConfigOverrides, add_args: AddArgs) -> Re println!("Added global MCP server '{name}'."); - match oauth_login_support(&transport, Arc::new(ReqwestHttpClient)).await { + match oauth_login_support(&transport).await { McpOAuthLoginSupport::Supported(oauth_config) => { println!("Detected OAuth support. Starting OAuth flow…"); let resolved_scopes = resolve_oauth_scopes( @@ -420,7 +419,7 @@ async fn run_login(config_overrides: &CliConfigOverrides, login_args: LoginArgs) let explicit_scopes = (!scopes.is_empty()).then_some(scopes); let discovered_scopes = if explicit_scopes.is_none() && server.scopes.is_none() { - discover_supported_scopes(&server.transport, Arc::new(ReqwestHttpClient)).await + discover_supported_scopes(&server.transport).await } else { None }; diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 9d74ae8451..4e9890c1f9 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -24,13 +24,19 @@ pub use mcp::read_mcp_resource; pub use mcp::McpAuthStatusEntry; pub use mcp::McpOAuthLoginConfig; +pub use mcp::McpOAuthLoginOutcome; pub use mcp::McpOAuthLoginSupport; pub use mcp::McpOAuthScopesSource; pub use mcp::ResolvedMcpOAuthScopes; pub use mcp::compute_auth_statuses; pub use mcp::discover_supported_scopes; +pub use mcp::discover_supported_scopes_for_server; pub use mcp::http_client_for_server; pub use mcp::oauth_login_support; +pub use mcp::oauth_login_support_for_server; +pub use mcp::perform_oauth_login_for_server; +pub use mcp::perform_oauth_login_return_url_for_server; +pub use mcp::perform_oauth_login_silent_for_server; pub use mcp::resolve_oauth_scopes; pub use mcp::should_retry_without_scopes; diff --git a/codex-rs/codex-mcp/src/mcp/auth.rs b/codex-rs/codex-mcp/src/mcp/auth.rs index 042a04d566..b21b185ba5 100644 --- a/codex-rs/codex-mcp/src/mcp/auth.rs +++ b/codex-rs/codex-mcp/src/mcp/auth.rs @@ -9,8 +9,12 @@ use codex_exec_server::ReqwestHttpClient; use codex_login::CodexAuth; use codex_protocol::protocol::McpAuthStatus; use codex_rmcp_client::OAuthProviderError; +use codex_rmcp_client::OauthLoginHandle; use codex_rmcp_client::determine_streamable_http_auth_status_with_client; use codex_rmcp_client::discover_streamable_http_oauth_with_client; +use codex_rmcp_client::perform_oauth_login_return_url_with_client; +use codex_rmcp_client::perform_oauth_login_silent_with_client; +use codex_rmcp_client::perform_oauth_login_with_client; use futures::future::join_all; use std::sync::Arc; use tracing::warn; @@ -54,7 +58,28 @@ pub struct McpAuthStatusEntry { pub auth_status: McpAuthStatus, } -pub async fn oauth_login_support( +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum McpOAuthLoginOutcome { + Completed, + Unsupported, +} + +pub async fn oauth_login_support(transport: &McpServerTransportConfig) -> McpOAuthLoginSupport { + oauth_login_support_with_client(transport, Arc::new(ReqwestHttpClient)).await +} + +pub async fn oauth_login_support_for_server( + config: &McpServerConfig, + runtime_environment: McpRuntimeEnvironment, +) -> McpOAuthLoginSupport { + let http_client = match http_client_for_server(config, runtime_environment) { + Ok(http_client) => http_client, + Err(err) => return McpOAuthLoginSupport::Unknown(err), + }; + oauth_login_support_with_client(&config.transport, http_client).await +} + +async fn oauth_login_support_with_client( transport: &McpServerTransportConfig, http_client: Arc, ) -> McpOAuthLoginSupport { @@ -93,9 +118,25 @@ pub async fn oauth_login_support( pub async fn discover_supported_scopes( transport: &McpServerTransportConfig, +) -> Option> { + discover_supported_scopes_with_client(transport, Arc::new(ReqwestHttpClient)).await +} + +pub async fn discover_supported_scopes_for_server( + config: &McpServerConfig, + runtime_environment: McpRuntimeEnvironment, +) -> Option> { + match oauth_login_support_for_server(config, runtime_environment).await { + McpOAuthLoginSupport::Supported(config) => config.discovered_scopes, + McpOAuthLoginSupport::Unsupported | McpOAuthLoginSupport::Unknown(_) => None, + } +} + +async fn discover_supported_scopes_with_client( + transport: &McpServerTransportConfig, http_client: Arc, ) -> Option> { - match oauth_login_support(transport, http_client).await { + match oauth_login_support_with_client(transport, http_client).await { McpOAuthLoginSupport::Supported(config) => config.discovered_scopes, McpOAuthLoginSupport::Unsupported | McpOAuthLoginSupport::Unknown(_) => None, } @@ -140,6 +181,172 @@ pub fn should_retry_without_scopes(scopes: &ResolvedMcpOAuthScopes, error: &anyh && error.downcast_ref::().is_some() } +#[allow(clippy::too_many_arguments)] +pub async fn perform_oauth_login_return_url_for_server( + server_name: &str, + config: &McpServerConfig, + store_mode: OAuthCredentialsStoreMode, + explicit_scopes: Option>, + timeout_secs: Option, + callback_port: Option, + callback_url: Option<&str>, + runtime_environment: McpRuntimeEnvironment, +) -> Result { + let McpServerTransportConfig::StreamableHttp { + url, + http_headers, + env_http_headers, + .. + } = &config.transport + else { + anyhow::bail!("OAuth login is only supported for streamable HTTP servers."); + }; + + let http_client = http_client_for_server(config, runtime_environment)?; + let discovered_scopes = if explicit_scopes.is_none() && config.scopes.is_none() { + discover_supported_scopes_with_client(&config.transport, http_client.clone()).await + } else { + None + }; + let resolved_scopes = + resolve_oauth_scopes(explicit_scopes, config.scopes.clone(), discovered_scopes); + + perform_oauth_login_return_url_with_client( + server_name, + url, + store_mode, + http_headers.clone(), + env_http_headers.clone(), + &resolved_scopes.scopes, + config.oauth_resource.as_deref(), + timeout_secs, + callback_port, + callback_url, + http_client, + ) + .await +} + +#[allow(clippy::too_many_arguments)] +pub async fn perform_oauth_login_silent_for_server( + server_name: &str, + config: &McpServerConfig, + store_mode: OAuthCredentialsStoreMode, + explicit_scopes: Option>, + callback_port: Option, + callback_url: Option<&str>, + runtime_environment: McpRuntimeEnvironment, +) -> Result { + let http_client = http_client_for_server(config, runtime_environment)?; + let oauth_config = + match oauth_login_support_with_client(&config.transport, http_client.clone()).await { + McpOAuthLoginSupport::Supported(config) => config, + McpOAuthLoginSupport::Unsupported => return Ok(McpOAuthLoginOutcome::Unsupported), + McpOAuthLoginSupport::Unknown(err) => return Err(err), + }; + + let resolved_scopes = resolve_oauth_scopes( + explicit_scopes, + config.scopes.clone(), + oauth_config.discovered_scopes.clone(), + ); + + let first_attempt = perform_oauth_login_silent_with_client( + server_name, + &oauth_config.url, + store_mode, + oauth_config.http_headers.clone(), + oauth_config.env_http_headers.clone(), + &resolved_scopes.scopes, + config.oauth_resource.as_deref(), + callback_port, + callback_url, + http_client.clone(), + ) + .await; + + let final_result = match first_attempt { + Err(err) if should_retry_without_scopes(&resolved_scopes, &err) => { + perform_oauth_login_silent_with_client( + server_name, + &oauth_config.url, + store_mode, + oauth_config.http_headers, + oauth_config.env_http_headers, + &[], + config.oauth_resource.as_deref(), + callback_port, + callback_url, + http_client, + ) + .await + } + result => result, + }; + + final_result.map(|()| McpOAuthLoginOutcome::Completed) +} + +#[allow(clippy::too_many_arguments)] +pub async fn perform_oauth_login_for_server( + server_name: &str, + config: &McpServerConfig, + store_mode: OAuthCredentialsStoreMode, + explicit_scopes: Option>, + callback_port: Option, + callback_url: Option<&str>, + runtime_environment: McpRuntimeEnvironment, +) -> Result { + let http_client = http_client_for_server(config, runtime_environment)?; + let oauth_config = + match oauth_login_support_with_client(&config.transport, http_client.clone()).await { + McpOAuthLoginSupport::Supported(config) => config, + McpOAuthLoginSupport::Unsupported => return Ok(McpOAuthLoginOutcome::Unsupported), + McpOAuthLoginSupport::Unknown(err) => return Err(err), + }; + + let resolved_scopes = resolve_oauth_scopes( + explicit_scopes, + config.scopes.clone(), + oauth_config.discovered_scopes.clone(), + ); + + let first_attempt = perform_oauth_login_with_client( + server_name, + &oauth_config.url, + store_mode, + oauth_config.http_headers.clone(), + oauth_config.env_http_headers.clone(), + &resolved_scopes.scopes, + config.oauth_resource.as_deref(), + callback_port, + callback_url, + http_client.clone(), + ) + .await; + + let final_result = match first_attempt { + Err(err) if should_retry_without_scopes(&resolved_scopes, &err) => { + perform_oauth_login_with_client( + server_name, + &oauth_config.url, + store_mode, + oauth_config.http_headers, + oauth_config.env_http_headers, + &[], + config.oauth_resource.as_deref(), + callback_port, + callback_url, + http_client, + ) + .await + } + result => result, + }; + + final_result.map(|()| McpOAuthLoginOutcome::Completed) +} + pub async fn compute_auth_statuses<'a, I>( servers: I, store_mode: OAuthCredentialsStoreMode, diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 56d2e01e0d..17f3337ce2 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -1,12 +1,18 @@ pub use auth::McpAuthStatusEntry; pub use auth::McpOAuthLoginConfig; +pub use auth::McpOAuthLoginOutcome; pub use auth::McpOAuthLoginSupport; pub use auth::McpOAuthScopesSource; pub use auth::ResolvedMcpOAuthScopes; pub use auth::compute_auth_statuses; pub use auth::discover_supported_scopes; +pub use auth::discover_supported_scopes_for_server; pub use auth::http_client_for_server; pub use auth::oauth_login_support; +pub use auth::oauth_login_support_for_server; +pub use auth::perform_oauth_login_for_server; +pub use auth::perform_oauth_login_return_url_for_server; +pub use auth::perform_oauth_login_silent_for_server; pub use auth::resolve_oauth_scopes; pub use auth::should_retry_without_scopes; diff --git a/codex-rs/core/src/mcp_skill_dependencies.rs b/codex-rs/core/src/mcp_skill_dependencies.rs index 71c3ea4f8a..99dceab936 100644 --- a/codex-rs/core/src/mcp_skill_dependencies.rs +++ b/codex-rs/core/src/mcp_skill_dependencies.rs @@ -11,7 +11,6 @@ use codex_protocol::request_user_input::RequestUserInputArgs; use codex_protocol::request_user_input::RequestUserInputQuestion; use codex_protocol::request_user_input::RequestUserInputQuestionOption; use codex_protocol::request_user_input::RequestUserInputResponse; -use codex_rmcp_client::perform_oauth_login_with_client; use tokio_util::sync::CancellationToken; use tracing::warn; @@ -19,13 +18,10 @@ use crate::SkillMetadata; use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::skills::model::SkillToolDependency; -use codex_mcp::McpOAuthLoginSupport; +use codex_mcp::McpOAuthLoginOutcome; use codex_mcp::McpRuntimeEnvironment; -use codex_mcp::http_client_for_server; use codex_mcp::mcp_permission_prompt_is_auto_approved; -use codex_mcp::oauth_login_support; -use codex_mcp::resolve_oauth_scopes; -use codex_mcp::should_retry_without_scopes; +use codex_mcp::perform_oauth_login_for_server; const SKILL_MCP_DEPENDENCY_PROMPT_ID: &str = "skill_mcp_dependency_install"; const MCP_DEPENDENCY_OPTION_INSTALL: &str = "Install"; @@ -135,23 +131,6 @@ pub(crate) async fn maybe_install_mcp_dependencies( .unwrap_or_else(|| sess.services.environment_manager.local_environment()), turn_context.cwd.to_path_buf(), ); - let http_client = match http_client_for_server(&server_config, runtime_environment) { - Ok(http_client) => http_client, - Err(err) => { - warn!("failed to resolve MCP OAuth environment for dependency {name}: {err}"); - continue; - } - }; - let oauth_config = - match oauth_login_support(&server_config.transport, http_client.clone()).await { - McpOAuthLoginSupport::Supported(config) => config, - McpOAuthLoginSupport::Unsupported => continue, - McpOAuthLoginSupport::Unknown(err) => { - warn!("MCP server may or may not require login for dependency {name}: {err}"); - continue; - } - }; - sess.notify_background_event( turn_context, format!( @@ -160,54 +139,20 @@ pub(crate) async fn maybe_install_mcp_dependencies( ) .await; - let resolved_scopes = resolve_oauth_scopes( - /*explicit_scopes*/ None, - server_config.scopes.clone(), - oauth_config.discovered_scopes.clone(), - ); - let first_attempt = perform_oauth_login_with_client( + match perform_oauth_login_for_server( &name, - &oauth_config.url, + &server_config, config.mcp_oauth_credentials_store_mode, - oauth_config.http_headers.clone(), - oauth_config.env_http_headers.clone(), - &resolved_scopes.scopes, - server_config.oauth_resource.as_deref(), + /*explicit_scopes*/ None, config.mcp_oauth_callback_port, config.mcp_oauth_callback_url.as_deref(), - http_client.clone(), + runtime_environment, ) - .await; - - if let Err(err) = first_attempt { - if should_retry_without_scopes(&resolved_scopes, &err) { - sess.notify_background_event( - turn_context, - format!( - "Retrying MCP {name} authentication without scopes after provider rejection." - ), - ) - .await; - - if let Err(err) = perform_oauth_login_with_client( - &name, - &oauth_config.url, - config.mcp_oauth_credentials_store_mode, - oauth_config.http_headers, - oauth_config.env_http_headers, - &[], - server_config.oauth_resource.as_deref(), - config.mcp_oauth_callback_port, - config.mcp_oauth_callback_url.as_deref(), - http_client, - ) - .await - { - warn!("failed to login to MCP dependency {name}: {err}"); - } - } else { - warn!("failed to login to MCP dependency {name}: {err}"); - } + .await + { + Ok(McpOAuthLoginOutcome::Completed) => {} + Ok(McpOAuthLoginOutcome::Unsupported) => {} + Err(err) => warn!("failed to login to MCP dependency {name}: {err}"), } }