diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index c86d414359..15c7492a7e 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -312,6 +312,7 @@ 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::read_mcp_resource as read_mcp_resource_without_thread; use codex_mcp::resolve_oauth_scopes; use codex_model_provider::ProviderAccountError; @@ -355,7 +356,7 @@ 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; +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; @@ -5960,15 +5961,36 @@ impl CodexMessageProcessor { } }; + let environment_manager = self.thread_manager.environment_manager(); + let runtime_environment = match environment_manager.default_environment() { + Some(environment) => McpRuntimeEnvironment::new(environment, config.cwd.to_path_buf()), + None => McpRuntimeEnvironment::new( + environment_manager.local_environment(), + 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).await + 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( + match perform_oauth_login_return_url_with_client( &name, &url, config.mcp_oauth_credentials_store_mode, @@ -5979,6 +6001,7 @@ impl CodexMessageProcessor { timeout_secs, config.mcp_oauth_callback_port, config.mcp_oauth_callback_url.as_deref(), + http_client, ) .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 b027aef453..30b7c47799 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 @@ -6,10 +6,12 @@ use codex_app_server_protocol::ServerNotification; use codex_config::types::McpServerConfig; use codex_core::config::Config; use codex_mcp::McpOAuthLoginSupport; +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; +use codex_rmcp_client::perform_oauth_login_silent_with_client; use tracing::warn; use super::CodexMessageProcessor; @@ -21,7 +23,28 @@ impl CodexMessageProcessor { plugin_mcp_servers: HashMap, ) { for (name, server) in plugin_mcp_servers { - let oauth_config = match oauth_login_support(&server.transport).await { + let environment_manager = self.thread_manager.environment_manager(); + let runtime_environment = match environment_manager.default_environment() { + Some(environment) => { + McpRuntimeEnvironment::new(environment, config.cwd.to_path_buf()) + } + None => McpRuntimeEnvironment::new( + environment_manager.local_environment(), + 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) => { @@ -45,7 +68,7 @@ impl CodexMessageProcessor { let notification_name = name.clone(); tokio::spawn(async move { - let first_attempt = perform_oauth_login_silent( + let first_attempt = perform_oauth_login_silent_with_client( &name, &oauth_config.url, store_mode, @@ -55,12 +78,13 @@ impl CodexMessageProcessor { server.oauth_resource.as_deref(), callback_port, callback_url.as_deref(), + http_client.clone(), ) .await; let final_result = match first_attempt { Err(err) if should_retry_without_scopes(&resolved_scopes, &err) => { - perform_oauth_login_silent( + perform_oauth_login_silent_with_client( &name, &oauth_config.url, store_mode, @@ -70,6 +94,7 @@ impl CodexMessageProcessor { server.oauth_resource.as_deref(), callback_port, callback_url.as_deref(), + http_client, ) .await } diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index c5b4751322..c795090658 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -15,7 +15,10 @@ use codex_core::config::edit::ConfigEditsBuilder; 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; use codex_mcp::compute_auth_statuses; use codex_mcp::discover_supported_scopes; @@ -323,7 +326,7 @@ async fn run_add(config_overrides: &CliConfigOverrides, add_args: AddArgs) -> Re println!("Added global MCP server '{name}'."); - match oauth_login_support(&transport).await { + match oauth_login_support(&transport, Arc::new(ReqwestHttpClient)).await { McpOAuthLoginSupport::Supported(oauth_config) => { println!("Detected OAuth support. Starting OAuth flow…"); let resolved_scopes = resolve_oauth_scopes( @@ -417,7 +420,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).await + discover_supported_scopes(&server.transport, Arc::new(ReqwestHttpClient)).await } else { None }; @@ -490,6 +493,10 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> mcp_servers.iter(), config.mcp_oauth_credentials_store_mode, /*auth*/ None, + McpRuntimeEnvironment::new( + Arc::new(Environment::default_for_tests()), + config.cwd.to_path_buf(), + ), ) .await; diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 1d3fd17619..9d74ae8451 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -29,6 +29,7 @@ pub use mcp::McpOAuthScopesSource; pub use mcp::ResolvedMcpOAuthScopes; pub use mcp::compute_auth_statuses; pub use mcp::discover_supported_scopes; +pub use mcp::http_client_for_server; pub use mcp::oauth_login_support; 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 6a97b52789..042a04d566 100644 --- a/codex-rs/codex-mcp/src/mcp/auth.rs +++ b/codex-rs/codex-mcp/src/mcp/auth.rs @@ -4,14 +4,19 @@ use anyhow::Result; use codex_config::McpServerConfig; use codex_config::McpServerTransportConfig; use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::HttpClient; +use codex_exec_server::ReqwestHttpClient; use codex_login::CodexAuth; use codex_protocol::protocol::McpAuthStatus; use codex_rmcp_client::OAuthProviderError; -use codex_rmcp_client::determine_streamable_http_auth_status; -use codex_rmcp_client::discover_streamable_http_oauth; +use codex_rmcp_client::determine_streamable_http_auth_status_with_client; +use codex_rmcp_client::discover_streamable_http_oauth_with_client; use futures::future::join_all; +use std::sync::Arc; use tracing::warn; +use crate::runtime::McpRuntimeEnvironment; + use super::CODEX_APPS_MCP_SERVER_NAME; #[derive(Debug, Clone)] @@ -49,7 +54,10 @@ pub struct McpAuthStatusEntry { pub auth_status: McpAuthStatus, } -pub async fn oauth_login_support(transport: &McpServerTransportConfig) -> McpOAuthLoginSupport { +pub async fn oauth_login_support( + transport: &McpServerTransportConfig, + http_client: Arc, +) -> McpOAuthLoginSupport { let McpServerTransportConfig::StreamableHttp { url, bearer_token_env_var, @@ -64,7 +72,13 @@ pub async fn oauth_login_support(transport: &McpServerTransportConfig) -> McpOAu return McpOAuthLoginSupport::Unsupported; } - match discover_streamable_http_oauth(url, http_headers.clone(), env_http_headers.clone()).await + match discover_streamable_http_oauth_with_client( + url, + http_headers.clone(), + env_http_headers.clone(), + http_client, + ) + .await { Ok(Some(discovery)) => McpOAuthLoginSupport::Supported(McpOAuthLoginConfig { url: url.clone(), @@ -79,8 +93,9 @@ pub async fn oauth_login_support(transport: &McpServerTransportConfig) -> McpOAu pub async fn discover_supported_scopes( transport: &McpServerTransportConfig, + http_client: Arc, ) -> Option> { - match oauth_login_support(transport).await { + match oauth_login_support(transport, http_client).await { McpOAuthLoginSupport::Supported(config) => config.discovered_scopes, McpOAuthLoginSupport::Unsupported | McpOAuthLoginSupport::Unknown(_) => None, } @@ -129,6 +144,7 @@ pub async fn compute_auth_statuses<'a, I>( servers: I, store_mode: OAuthCredentialsStoreMode, auth: Option<&CodexAuth>, + runtime_environment: McpRuntimeEnvironment, ) -> HashMap where I: IntoIterator, @@ -136,6 +152,7 @@ where let futures = servers.into_iter().map(|(name, config)| { let name = name.clone(); let config = config.clone(); + let runtime_environment = runtime_environment.clone(); let has_runtime_auth = name == CODEX_APPS_MCP_SERVER_NAME && auth.is_some_and(CodexAuth::uses_codex_backend) && matches!( @@ -146,14 +163,21 @@ where } ); 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 auth_status = match compute_auth_status( + &name, + &config, + store_mode, + has_runtime_auth, + runtime_environment, + ) + .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, @@ -170,6 +194,7 @@ async fn compute_auth_status( config: &McpServerConfig, store_mode: OAuthCredentialsStoreMode, has_runtime_auth: bool, + runtime_environment: McpRuntimeEnvironment, ) -> Result { if !config.enabled { return Ok(McpAuthStatus::Unsupported); @@ -187,27 +212,57 @@ async fn compute_auth_status( http_headers, env_http_headers, } => { - determine_streamable_http_auth_status( + let http_client = http_client_for_server(config, runtime_environment)?; + determine_streamable_http_auth_status_with_client( server_name, url, bearer_token_env_var.as_deref(), http_headers.clone(), env_http_headers.clone(), store_mode, + http_client, ) .await } } } +pub fn http_client_for_server( + config: &McpServerConfig, + runtime_environment: McpRuntimeEnvironment, +) -> Result> { + match config.experimental_environment.as_deref() { + None | Some("local") => Ok(Arc::new(ReqwestHttpClient)), + Some("remote") => { + let environment = runtime_environment.environment(); + if !environment.is_remote() { + anyhow::bail!("remote MCP server requires a remote environment"); + } + Ok(environment.get_http_client()) + } + Some(environment) => anyhow::bail!("unsupported experimental_environment `{environment}`"), + } +} + #[cfg(test)] mod tests { + use std::collections::HashMap; + use std::path::PathBuf; + use std::sync::Arc; + use std::time::Duration; + use anyhow::anyhow; + use codex_config::McpServerConfig; + use codex_config::McpServerTransportConfig; + use codex_exec_server::Environment; use pretty_assertions::assert_eq; + use crate::runtime::McpRuntimeEnvironment; + use super::McpOAuthScopesSource; use super::OAuthProviderError; use super::ResolvedMcpOAuthScopes; + use super::http_client_for_server; use super::resolve_oauth_scopes; use super::should_retry_without_scopes; @@ -318,4 +373,71 @@ mod tests { &anyhow!("timed out waiting for OAuth callback"), )); } + + #[test] + fn local_server_uses_local_http_client_even_with_remote_runtime() { + let config = streamable_config(/*experimental_environment*/ None); + let runtime_environment = remote_runtime_environment(); + + assert!(http_client_for_server(&config, runtime_environment).is_ok()); + } + + #[test] + fn remote_server_uses_runtime_http_client_when_runtime_is_remote() { + let config = streamable_config(Some("remote")); + let runtime_environment = remote_runtime_environment(); + + assert!(http_client_for_server(&config, runtime_environment).is_ok()); + } + + #[tokio::test] + async fn remote_server_without_remote_runtime_returns_clear_error() { + let config = streamable_config(Some("remote")); + let runtime_environment = McpRuntimeEnvironment::new( + Arc::new(Environment::default_for_tests()), + PathBuf::from("/tmp"), + ); + + let Err(error) = http_client_for_server(&config, runtime_environment) else { + panic!("remote server should require remote runtime"); + }; + assert_eq!( + error.to_string(), + "remote MCP server requires a remote environment" + ); + } + + fn streamable_config(experimental_environment: Option<&str>) -> McpServerConfig { + McpServerConfig { + transport: McpServerTransportConfig::StreamableHttp { + url: "http://mcp.example.test/mcp".to_string(), + bearer_token_env_var: None, + http_headers: None, + env_http_headers: None, + }, + experimental_environment: experimental_environment.map(str::to_string), + enabled: true, + required: false, + supports_parallel_tool_calls: false, + disabled_reason: None, + startup_timeout_sec: Some(Duration::from_secs(30)), + tool_timeout_sec: None, + default_tools_approval_mode: None, + enabled_tools: None, + disabled_tools: None, + scopes: None, + oauth_resource: None, + tools: HashMap::new(), + } + } + + fn remote_runtime_environment() -> McpRuntimeEnvironment { + McpRuntimeEnvironment::new( + Arc::new( + Environment::create_for_tests(Some("ws://127.0.0.1:65535".to_string())) + .expect("create remote environment"), + ), + PathBuf::from("/tmp"), + ) + } } diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 080ac889de..56d2e01e0d 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -5,6 +5,7 @@ pub use auth::McpOAuthScopesSource; pub use auth::ResolvedMcpOAuthScopes; pub use auth::compute_auth_statuses; pub use auth::discover_supported_scopes; +pub use auth::http_client_for_server; pub use auth::oauth_login_support; pub use auth::resolve_oauth_scopes; pub use auth::should_retry_without_scopes; @@ -223,6 +224,7 @@ pub async fn read_mcp_resource( mcp_servers.iter(), config.mcp_oauth_credentials_store_mode, auth, + runtime_environment.clone(), ) .await; let (tx_event, rx_event) = unbounded(); @@ -286,6 +288,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( mcp_servers.iter(), config.mcp_oauth_credentials_store_mode, auth, + runtime_environment.clone(), ) .await; diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 456f3a7eac..69e9d0882a 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -253,20 +253,23 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager( }); } + let environment = environment_manager + .default_environment() + .unwrap_or_else(|| environment_manager.local_environment()); + let runtime_environment = + McpRuntimeEnvironment::new(environment.clone(), config.cwd.to_path_buf()); + let auth_status_entries = compute_auth_statuses( mcp_servers.iter(), config.mcp_oauth_credentials_store_mode, auth.as_ref(), + runtime_environment.clone(), ) .await; let (tx_event, rx_event) = unbounded(); drop(rx_event); - let environment = environment_manager - .default_environment() - .unwrap_or_else(|| environment_manager.local_environment()); - let (mcp_connection_manager, cancel_token) = McpConnectionManager::new( &mcp_servers, config.mcp_oauth_credentials_store_mode, @@ -275,7 +278,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager( INITIAL_SUBMIT_ID.to_owned(), tx_event, PermissionProfile::default(), - McpRuntimeEnvironment::new(environment, config.cwd.to_path_buf()), + runtime_environment, config.codex_home.to_path_buf(), codex_apps_tools_cache_key(auth.as_ref()), ToolPluginProvenance::default(), diff --git a/codex-rs/core/src/mcp_skill_dependencies.rs b/codex-rs/core/src/mcp_skill_dependencies.rs index a97424f137..71c3ea4f8a 100644 --- a/codex-rs/core/src/mcp_skill_dependencies.rs +++ b/codex-rs/core/src/mcp_skill_dependencies.rs @@ -11,7 +11,7 @@ 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; +use codex_rmcp_client::perform_oauth_login_with_client; use tokio_util::sync::CancellationToken; use tracing::warn; @@ -20,6 +20,8 @@ use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::skills::model::SkillToolDependency; use codex_mcp::McpOAuthLoginSupport; +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; @@ -126,14 +128,29 @@ pub(crate) async fn maybe_install_mcp_dependencies( } for (name, server_config) in added { - let oauth_config = match oauth_login_support(&server_config.transport).await { - McpOAuthLoginSupport::Supported(config) => config, - McpOAuthLoginSupport::Unsupported => continue, - McpOAuthLoginSupport::Unknown(err) => { - warn!("MCP server may or may not require login for dependency {name}: {err}"); + let runtime_environment = McpRuntimeEnvironment::new( + turn_context + .environment + .clone() + .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, @@ -148,7 +165,7 @@ pub(crate) async fn maybe_install_mcp_dependencies( server_config.scopes.clone(), oauth_config.discovered_scopes.clone(), ); - let first_attempt = perform_oauth_login( + let first_attempt = perform_oauth_login_with_client( &name, &oauth_config.url, config.mcp_oauth_credentials_store_mode, @@ -158,6 +175,7 @@ pub(crate) async fn maybe_install_mcp_dependencies( server_config.oauth_resource.as_deref(), config.mcp_oauth_callback_port, config.mcp_oauth_callback_url.as_deref(), + http_client.clone(), ) .await; @@ -171,7 +189,7 @@ pub(crate) async fn maybe_install_mcp_dependencies( ) .await; - if let Err(err) = perform_oauth_login( + if let Err(err) = perform_oauth_login_with_client( &name, &oauth_config.url, config.mcp_oauth_credentials_store_mode, @@ -181,6 +199,7 @@ pub(crate) async fn maybe_install_mcp_dependencies( server_config.oauth_resource.as_deref(), config.mcp_oauth_callback_port, config.mcp_oauth_callback_url.as_deref(), + http_client, ) .await { diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 86ce79c90f..2aec5a004a 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -32,6 +32,7 @@ use crate::tasks::UndoTask; use crate::tasks::UserShellCommandMode; use crate::tasks::UserShellCommandTask; use crate::tasks::execute_user_shell_command; +use codex_mcp::McpRuntimeEnvironment; use codex_mcp::collect_mcp_snapshot_from_manager; use codex_mcp::compute_auth_statuses; use codex_protocol::models::ContentItem; @@ -538,12 +539,18 @@ pub async fn list_mcp_tools(sess: &Session, config: &Arc, sub_id: String .mcp_manager .effective_servers(config, auth.as_ref()) .await; + let environment = sess + .services + .environment_manager + .default_environment() + .unwrap_or_else(|| sess.services.environment_manager.local_environment()); let snapshot = collect_mcp_snapshot_from_manager( &mcp_connection_manager, compute_auth_statuses( mcp_servers.iter(), config.mcp_oauth_credentials_store_mode, auth.as_ref(), + McpRuntimeEnvironment::new(environment, config.cwd.to_path_buf()), ) .await, ) diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index eae4608c49..a889a13bc0 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -219,8 +219,20 @@ 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, auth.as_ref()).await; + let runtime_environment = McpRuntimeEnvironment::new( + turn_context + .environment + .clone() + .unwrap_or_else(|| self.services.environment_manager.local_environment()), + turn_context.cwd.to_path_buf(), + ); + let auth_statuses = compute_auth_statuses( + mcp_servers.iter(), + store_mode, + auth.as_ref(), + runtime_environment.clone(), + ) + .await; { let mut guard = self.services.mcp_startup_cancellation_token.lock().await; guard.cancel(); @@ -234,13 +246,7 @@ impl Session { turn_context.sub_id.clone(), self.get_tx_event(), turn_context.permission_profile(), - McpRuntimeEnvironment::new( - turn_context - .environment - .clone() - .unwrap_or_else(|| self.services.environment_manager.local_environment()), - turn_context.cwd.to_path_buf(), - ), + runtime_environment, config.codex_home.to_path_buf(), codex_apps_tools_cache_key(auth.as_ref()), tool_plugin_provenance, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index dcadac70a8..e95f1de4bb 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -445,15 +445,20 @@ impl Session { let auth_manager_clone = Arc::clone(&auth_manager); let config_for_mcp = Arc::clone(&config); let mcp_manager_for_mcp = Arc::clone(&mcp_manager); + let environment_manager_for_mcp = Arc::clone(&environment_manager); let auth_and_mcp_fut = async move { let auth = auth_manager_clone.auth().await; let mcp_servers = mcp_manager_for_mcp .effective_servers(&config_for_mcp, auth.as_ref()) .await; + let environment = environment_manager_for_mcp + .default_environment() + .unwrap_or_else(|| environment_manager_for_mcp.local_environment()); let auth_statuses = compute_auth_statuses( mcp_servers.iter(), config_for_mcp.mcp_oauth_credentials_store_mode, auth.as_ref(), + McpRuntimeEnvironment::new(environment, config_for_mcp.cwd.to_path_buf()), ) .await; (auth, mcp_servers, auth_statuses) diff --git a/codex-rs/rmcp-client/src/auth_status.rs b/codex-rs/rmcp-client/src/auth_status.rs index 0b3b3bf6a7..cbb4c5dc5a 100644 --- a/codex-rs/rmcp-client/src/auth_status.rs +++ b/codex-rs/rmcp-client/src/auth_status.rs @@ -1,31 +1,19 @@ use std::collections::HashMap; -use std::time::Duration; +use std::sync::Arc; -use anyhow::Error; use anyhow::Result; +use codex_exec_server::HttpClient; +use codex_exec_server::ReqwestHttpClient; use codex_protocol::protocol::McpAuthStatus; -use reqwest::Client; -use reqwest::StatusCode; -use reqwest::Url; use reqwest::header::AUTHORIZATION; -use reqwest::header::HeaderMap; -use serde::Deserialize; use tracing::debug; +use crate::mcp_oauth_http::OAuthHttpClient; +use crate::mcp_oauth_http::StreamableHttpOAuthDiscovery; use crate::oauth::has_oauth_tokens; -use crate::utils::apply_default_headers; use crate::utils::build_default_headers; use codex_config::types::OAuthCredentialsStoreMode; -const DISCOVERY_TIMEOUT: Duration = Duration::from_secs(5); -const OAUTH_DISCOVERY_HEADER: &str = "MCP-Protocol-Version"; -const OAUTH_DISCOVERY_VERSION: &str = "2024-11-05"; - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct StreamableHttpOAuthDiscovery { - pub scopes_supported: Option>, -} - /// Determine the authentication status for a streamable HTTP MCP server. pub async fn determine_streamable_http_auth_status( server_name: &str, @@ -34,12 +22,36 @@ pub async fn determine_streamable_http_auth_status( http_headers: Option>, env_http_headers: Option>, store_mode: OAuthCredentialsStoreMode, +) -> Result { + determine_streamable_http_auth_status_with_client( + server_name, + url, + bearer_token_env_var, + http_headers, + env_http_headers, + store_mode, + Arc::new(ReqwestHttpClient), + ) + .await +} + +/// Determine the authentication status for a streamable HTTP MCP server using +/// the caller-selected runtime HTTP client. +#[allow(clippy::too_many_arguments)] +pub async fn determine_streamable_http_auth_status_with_client( + server_name: &str, + url: &str, + bearer_token_env_var: Option<&str>, + http_headers: Option>, + env_http_headers: Option>, + store_mode: OAuthCredentialsStoreMode, + http_client: Arc, ) -> Result { if bearer_token_env_var.is_some() { return Ok(McpAuthStatus::BearerToken); } - let default_headers = build_default_headers(http_headers, env_http_headers)?; + let default_headers = build_default_headers(http_headers.clone(), env_http_headers.clone())?; if default_headers.contains_key(AUTHORIZATION) { return Ok(McpAuthStatus::BearerToken); } @@ -48,7 +60,8 @@ pub async fn determine_streamable_http_auth_status( return Ok(McpAuthStatus::OAuth); } - match discover_streamable_http_oauth_with_headers(url, &default_headers).await { + let oauth_http = OAuthHttpClient::from_default_headers(http_client, default_headers); + match oauth_http.discover(url).await { Ok(Some(_)) => Ok(McpAuthStatus::NotLoggedIn), Ok(None) => Ok(McpAuthStatus::Unsupported), Err(error) => { @@ -74,121 +87,23 @@ pub async fn discover_streamable_http_oauth( http_headers: Option>, env_http_headers: Option>, ) -> Result> { - let default_headers = build_default_headers(http_headers, env_http_headers)?; - discover_streamable_http_oauth_with_headers(url, &default_headers).await + discover_streamable_http_oauth_with_client( + url, + http_headers, + env_http_headers, + Arc::new(ReqwestHttpClient), + ) + .await } -async fn discover_streamable_http_oauth_with_headers( +pub async fn discover_streamable_http_oauth_with_client( url: &str, - default_headers: &HeaderMap, + http_headers: Option>, + env_http_headers: Option>, + http_client: Arc, ) -> Result> { - let base_url = Url::parse(url)?; - - // Use no_proxy to avoid a bug in the system-configuration crate that - // can result in a panic. See #8912. - let builder = Client::builder().timeout(DISCOVERY_TIMEOUT).no_proxy(); - let client = apply_default_headers(builder, default_headers).build()?; - - let mut last_error: Option = None; - for candidate_path in discovery_paths(base_url.path()) { - let mut discovery_url = base_url.clone(); - discovery_url.set_path(&candidate_path); - - let response = match client - .get(discovery_url.clone()) - .header(OAUTH_DISCOVERY_HEADER, OAUTH_DISCOVERY_VERSION) - .send() - .await - { - Ok(response) => response, - Err(err) => { - last_error = Some(err.into()); - continue; - } - }; - - if response.status() != StatusCode::OK { - continue; - } - - let metadata = match response.json::().await { - Ok(metadata) => metadata, - Err(err) => { - last_error = Some(err.into()); - continue; - } - }; - - if metadata.authorization_endpoint.is_some() && metadata.token_endpoint.is_some() { - return Ok(Some(StreamableHttpOAuthDiscovery { - scopes_supported: normalize_scopes(metadata.scopes_supported), - })); - } - } - - if let Some(err) = last_error { - debug!("OAuth discovery requests failed for {url}: {err:?}"); - } - - Ok(None) -} - -#[derive(Debug, Deserialize)] -struct OAuthDiscoveryMetadata { - #[serde(default)] - authorization_endpoint: Option, - #[serde(default)] - token_endpoint: Option, - #[serde(default)] - scopes_supported: Option>, -} - -fn normalize_scopes(scopes_supported: Option>) -> Option> { - let scopes_supported = scopes_supported?; - - let mut normalized = Vec::new(); - for scope in scopes_supported { - let scope = scope.trim(); - if scope.is_empty() { - continue; - } - let scope = scope.to_string(); - if !normalized.contains(&scope) { - normalized.push(scope); - } - } - - if normalized.is_empty() { - None - } else { - Some(normalized) - } -} - -/// Implements RFC 8414 section 3.1 for discovering well-known oauth endpoints. -/// This is a requirement for MCP servers to support OAuth. -/// https://datatracker.ietf.org/doc/html/rfc8414#section-3.1 -/// https://github.com/modelcontextprotocol/rust-sdk/blob/main/crates/rmcp/src/transport/auth.rs#L182 -fn discovery_paths(base_path: &str) -> Vec { - let trimmed = base_path.trim_start_matches('/').trim_end_matches('/'); - let canonical = "/.well-known/oauth-authorization-server".to_string(); - - if trimmed.is_empty() { - return vec![canonical]; - } - - let mut candidates = Vec::new(); - let mut push_unique = |candidate: String| { - if !candidates.contains(&candidate) { - candidates.push(candidate); - } - }; - - push_unique(format!("{canonical}/{trimmed}")); - push_unique(format!("/{trimmed}/.well-known/oauth-authorization-server")); - push_unique(canonical); - - candidates + let oauth_http = OAuthHttpClient::new(http_client, http_headers, env_http_headers)?; + oauth_http.discover(url).await } #[cfg(test)] diff --git a/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs b/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs index d1c22f430c..9b90e560c2 100644 --- a/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs @@ -8,6 +8,7 @@ use std::time::Duration; use axum::Router; use axum::body::Body; +use axum::body::Bytes; use axum::extract::Json; use axum::extract::State; use axum::http::HeaderMap; @@ -122,6 +123,7 @@ async fn main() -> Result<(), Box> { SESSION_POST_FAILURE_CONTROL_PATH, post(arm_session_post_failure), ) + .route("/oauth/token", post(oauth_token)) .route( "/.well-known/oauth-authorization-server/mcp", get({ @@ -139,6 +141,7 @@ async fn main() -> Result<(), Box> { serde_json::to_vec(&json!({ "authorization_endpoint": format!("{metadata_base}/oauth/authorize"), "token_endpoint": format!("{metadata_base}/oauth/token"), + "registration_endpoint": format!("{metadata_base}/oauth/register"), "scopes_supported": [""], })).expect("failed to serialize metadata"), )) @@ -386,7 +389,8 @@ async fn require_bearer( request: Request, next: Next, ) -> Result { - if request.uri().path().contains("/.well-known/") { + let path = request.uri().path(); + if path.contains("/.well-known/") || path.starts_with("/oauth/") { return Ok(next.run(request).await); } if request @@ -400,6 +404,33 @@ async fn require_bearer( } } +async fn oauth_token(body: Bytes) -> Result { + let form = String::from_utf8(body.to_vec()).map_err(|_| StatusCode::BAD_REQUEST)?; + if !form.contains("grant_type=refresh_token") { + return Err(StatusCode::BAD_REQUEST); + } + + let access_token = + std::env::var("MCP_OAUTH_ACCESS_TOKEN").unwrap_or_else(|_| "refreshed-oauth-token".into()); + let refresh_token = + std::env::var("MCP_OAUTH_REFRESH_TOKEN").unwrap_or_else(|_| "refresh-token".into()); + + #[expect(clippy::expect_used)] + Response::builder() + .status(StatusCode::OK) + .header(CONTENT_TYPE, "application/json") + .body(Body::from( + serde_json::to_vec(&json!({ + "access_token": access_token, + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": refresh_token, + })) + .expect("failed to serialize token response"), + )) + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR) +} + async fn arm_session_post_failure( State(state): State, Json(request): Json, diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index a1e6680e60..0ab66a117d 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -35,6 +35,8 @@ use rmcp::transport::streamable_http_client::StreamableHttpPostResponse; use sse_stream::Sse; use sse_stream::SseStream; +use crate::oauth::OAuthPersistor; + const EVENT_STREAM_MIME_TYPE: &str = "text/event-stream"; const JSON_MIME_TYPE: &str = "application/json"; const HEADER_SESSION_ID: &str = "Mcp-Session-Id"; @@ -45,6 +47,7 @@ pub(crate) struct StreamableHttpClientAdapter { http_client: Arc, default_headers: HeaderMap, auth_provider: Option, + oauth_persistor: Option, } #[derive(Debug, thiserror::Error)] @@ -55,6 +58,8 @@ pub(crate) enum StreamableHttpClientAdapterError { HttpRequest(#[from] ExecServerError), #[error("invalid HTTP header: {0}")] Header(String), + #[error("OAuth token error: {0}")] + Auth(String), } impl StreamableHttpClientAdapter { @@ -62,11 +67,13 @@ impl StreamableHttpClientAdapter { http_client: Arc, default_headers: HeaderMap, auth_provider: Option, + oauth_persistor: Option, ) -> Self { Self { http_client, default_headers, auth_provider, + oauth_persistor, } } } @@ -82,7 +89,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { auth_token: Option, ) -> std::result::Result> { let mut headers = self.default_headers.clone(); - self.add_auth_headers(&mut headers); + self.add_auth_headers(&mut headers).await?; insert_header( &mut headers, ACCEPT, @@ -179,7 +186,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { auth_token: Option, ) -> std::result::Result<(), StreamableHttpError> { let mut headers = self.default_headers.clone(); - self.add_auth_headers(&mut headers); + self.add_auth_headers(&mut headers).await?; if let Some(auth_token) = auth_token { insert_header( &mut headers, @@ -232,7 +239,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { StreamableHttpError, > { let mut headers = self.default_headers.clone(); - self.add_auth_headers(&mut headers); + self.add_auth_headers(&mut headers).await?; insert_header( &mut headers, ACCEPT, @@ -308,10 +315,27 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { } impl StreamableHttpClientAdapter { - fn add_auth_headers(&self, headers: &mut HeaderMap) { + async fn add_auth_headers( + &self, + headers: &mut HeaderMap, + ) -> std::result::Result<(), StreamableHttpError> { if let Some(auth_provider) = &self.auth_provider { headers.extend(auth_provider.to_auth_headers()); } + if !headers.contains_key(AUTHORIZATION) + && let Some(oauth_persistor) = &self.oauth_persistor + && let Some(access_token) = oauth_persistor.access_token().await.map_err(|err| { + StreamableHttpError::Client(StreamableHttpClientAdapterError::Auth(err.to_string())) + })? + { + insert_header( + headers, + AUTHORIZATION, + format!("Bearer {access_token}"), + StreamableHttpClientAdapterError::Header, + )?; + } + Ok(()) } } diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 57e9f0e800..68a5bf027a 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -3,6 +3,7 @@ mod elicitation_client_service; mod executor_process_transport; mod http_client_adapter; mod logging_client_handler; +mod mcp_oauth_http; mod oauth; mod perform_oauth_login; mod program_resolver; @@ -10,11 +11,13 @@ mod rmcp_client; mod stdio_server_launcher; mod utils; -pub use auth_status::StreamableHttpOAuthDiscovery; pub use auth_status::determine_streamable_http_auth_status; +pub use auth_status::determine_streamable_http_auth_status_with_client; pub use auth_status::discover_streamable_http_oauth; +pub use auth_status::discover_streamable_http_oauth_with_client; pub use auth_status::supports_oauth_login; pub use codex_protocol::protocol::McpAuthStatus; +pub use mcp_oauth_http::StreamableHttpOAuthDiscovery; pub use oauth::StoredOAuthTokens; pub use oauth::WrappedOAuthTokenResponse; pub use oauth::delete_oauth_tokens; @@ -24,7 +27,10 @@ pub use perform_oauth_login::OAuthProviderError; pub use perform_oauth_login::OauthLoginHandle; pub use perform_oauth_login::perform_oauth_login; pub use perform_oauth_login::perform_oauth_login_return_url; +pub use perform_oauth_login::perform_oauth_login_return_url_with_client; pub use perform_oauth_login::perform_oauth_login_silent; +pub use perform_oauth_login::perform_oauth_login_silent_with_client; +pub use perform_oauth_login::perform_oauth_login_with_client; pub use rmcp::model::ElicitationAction; pub use rmcp_client::Elicitation; pub use rmcp_client::ElicitationResponse; diff --git a/codex-rs/rmcp-client/src/mcp_oauth_http.rs b/codex-rs/rmcp-client/src/mcp_oauth_http.rs new file mode 100644 index 0000000000..8c093102fd --- /dev/null +++ b/codex-rs/rmcp-client/src/mcp_oauth_http.rs @@ -0,0 +1,577 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use anyhow::anyhow; +use codex_exec_server::HttpClient; +use codex_exec_server::HttpHeader; +use codex_exec_server::HttpRequestParams; +use oauth2::AuthUrl; +use oauth2::AuthorizationCode; +use oauth2::ClientId; +use oauth2::ClientSecret; +use oauth2::CsrfToken; +use oauth2::EmptyExtraTokenFields; +use oauth2::PkceCodeChallenge; +use oauth2::PkceCodeVerifier; +use oauth2::RedirectUrl; +use oauth2::RefreshToken; +use oauth2::RequestTokenError; +use oauth2::Scope; +use oauth2::StandardErrorResponse; +use oauth2::StandardTokenResponse; +use oauth2::TokenResponse; +use oauth2::TokenUrl; +use oauth2::basic::BasicClient; +use oauth2::basic::BasicErrorResponseType; +use oauth2::basic::BasicTokenType; +use reqwest::StatusCode; +use reqwest::Url; +use reqwest::header::HeaderMap; +use serde::Deserialize; +use serde::Serialize; +use tracing::debug; + +use crate::WrappedOAuthTokenResponse; +use crate::oauth::StoredOAuthTokens; +use crate::oauth::compute_expires_at_millis; +use crate::utils::build_default_headers; + +const DISCOVERY_TIMEOUT: Duration = Duration::from_secs(5); +const OAUTH_HTTP_TIMEOUT: Duration = Duration::from_secs(30); +const OAUTH_DISCOVERY_HEADER: &str = "MCP-Protocol-Version"; +const OAUTH_DISCOVERY_VERSION: &str = "2024-11-05"; + +type OAuthTokenResponse = StandardTokenResponse; +type OAuthErrorResponse = StandardErrorResponse; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StreamableHttpOAuthDiscovery { + pub scopes_supported: Option>, +} + +#[derive(Clone)] +pub(crate) struct OAuthHttpClient { + http_client: Arc, + default_headers: HeaderMap, +} + +#[derive(Debug, thiserror::Error)] +#[error("{0}")] +pub(crate) struct OAuthHttpError(String); + +#[derive(Debug)] +pub(crate) struct OAuthAuthorizationSession { + pub authorization_url: String, + pub csrf_state: CsrfToken, + pkce_verifier: PkceCodeVerifier, + client: BasicClient< + oauth2::EndpointSet, + oauth2::EndpointNotSet, + oauth2::EndpointNotSet, + oauth2::EndpointNotSet, + oauth2::EndpointSet, + >, + client_id: String, +} + +#[derive(Debug, Clone)] +struct OAuthClientConfig { + client_id: String, + client_secret: Option, +} + +impl OAuthHttpClient { + pub(crate) fn new( + http_client: Arc, + http_headers: Option>, + env_http_headers: Option>, + ) -> Result { + let default_headers = build_default_headers(http_headers, env_http_headers)?; + Ok(Self { + http_client, + default_headers, + }) + } + + pub(crate) fn from_default_headers( + http_client: Arc, + default_headers: HeaderMap, + ) -> Self { + Self { + http_client, + default_headers, + } + } + + pub(crate) async fn discover(&self, url: &str) -> Result> { + let metadata = match self.discover_metadata(url).await? { + Some(metadata) => metadata, + None => return Ok(None), + }; + + Ok(Some(StreamableHttpOAuthDiscovery { + scopes_supported: normalize_scopes(metadata.scopes_supported), + })) + } + + async fn discover_metadata(&self, url: &str) -> Result> { + let base_url = Url::parse(url)?; + let mut last_error: Option = None; + + for candidate_path in discovery_paths(base_url.path()) { + let mut discovery_url = base_url.clone(); + discovery_url.set_path(&candidate_path); + + let response = match self + .request( + "GET", + discovery_url.as_str(), + vec![HttpHeader { + name: OAUTH_DISCOVERY_HEADER.to_string(), + value: OAUTH_DISCOVERY_VERSION.to_string(), + }], + None, + Some(DISCOVERY_TIMEOUT), + ) + .await + { + Ok(response) => response, + Err(err) => { + last_error = Some(err); + continue; + } + }; + + if response.status != StatusCode::OK.as_u16() { + continue; + } + + let metadata = match serde_json::from_slice::(&response.body.0) + { + Ok(metadata) => metadata, + Err(err) => { + last_error = Some(err.into()); + continue; + } + }; + + if metadata.authorization_endpoint.is_some() && metadata.token_endpoint.is_some() { + return Ok(Some(metadata)); + } + } + + if let Some(err) = last_error { + debug!("OAuth discovery requests failed for {url}: {err:?}"); + } + + Ok(None) + } + + pub(crate) async fn start_authorization( + &self, + server_url: &str, + scopes: &[String], + redirect_uri: &str, + client_name: &str, + ) -> Result { + let metadata = self + .discover_metadata(server_url) + .await? + .ok_or_else(|| anyhow!("No authorization support detected"))?; + let client_config = self + .register_client(&metadata, client_name, redirect_uri) + .await?; + let client = oauth_client(&metadata, &client_config, redirect_uri)?; + let (pkce_challenge, pkce_verifier) = PkceCodeChallenge::new_random_sha256(); + let mut auth_request = client + .authorize_url(CsrfToken::new_random) + .set_pkce_challenge(pkce_challenge); + for scope in scopes { + auth_request = auth_request.add_scope(Scope::new(scope.clone())); + } + let (authorization_url, csrf_state) = auth_request.url(); + + Ok(OAuthAuthorizationSession { + authorization_url: authorization_url.to_string(), + csrf_state, + pkce_verifier, + client, + client_id: client_config.client_id, + }) + } + + pub(crate) async fn exchange_code( + &self, + session: OAuthAuthorizationSession, + code: &str, + csrf_state: &str, + ) -> Result<(String, OAuthTokenResponse)> { + if session.csrf_state.secret() != csrf_state { + return Err(anyhow!( + "OAuth callback state did not match authorization request" + )); + } + + let http_client = self.clone(); + let token = session + .client + .exchange_code(AuthorizationCode::new(code.to_string())) + .set_pkce_verifier(session.pkce_verifier) + .request_async(&|request| { + let http_client = http_client.clone(); + async move { http_client.oauth_request(request).await } + }) + .await + .or_else(parse_token_from_parse_error)?; + + Ok((session.client_id, token)) + } + + pub(crate) async fn refresh_token( + &self, + tokens: &StoredOAuthTokens, + ) -> Result> { + let refresh_token = match tokens.token_response.0.refresh_token() { + Some(refresh_token) => refresh_token.secret().to_string(), + None => return Ok(None), + }; + + let metadata = self + .discover_metadata(&tokens.url) + .await? + .ok_or_else(|| anyhow!("No authorization support detected"))?; + let client_config = OAuthClientConfig { + client_id: tokens.client_id.clone(), + client_secret: None, + }; + let client = oauth_client(&metadata, &client_config, &tokens.url)?; + let http_client = self.clone(); + let token = client + .exchange_refresh_token(&RefreshToken::new(refresh_token)) + .request_async(&|request| { + let http_client = http_client.clone(); + async move { http_client.oauth_request(request).await } + }) + .await + .or_else(parse_token_from_parse_error)?; + + let expires_at = compute_expires_at_millis(&token); + Ok(Some(StoredOAuthTokens { + server_name: tokens.server_name.clone(), + url: tokens.url.clone(), + client_id: tokens.client_id.clone(), + token_response: WrappedOAuthTokenResponse(token), + expires_at, + })) + } + + async fn register_client( + &self, + metadata: &OAuthDiscoveryMetadata, + client_name: &str, + redirect_uri: &str, + ) -> Result { + let registration_url = metadata + .registration_endpoint + .as_deref() + .ok_or_else(|| anyhow!("Dynamic client registration not supported"))?; + + if let Some(response_types_supported) = metadata.response_types_supported.as_ref() + && !response_types_supported.iter().any(|value| value == "code") + { + return Err(anyhow!( + "OAuth server does not support authorization code flow" + )); + } + + let body = serde_json::to_vec(&ClientRegistrationRequest { + client_name: client_name.to_string(), + redirect_uris: vec![redirect_uri.to_string()], + grant_types: vec![ + "authorization_code".to_string(), + "refresh_token".to_string(), + ], + token_endpoint_auth_method: "none".to_string(), + response_types: vec!["code".to_string()], + })?; + + let response = self + .request( + "POST", + registration_url, + vec![HttpHeader { + name: reqwest::header::CONTENT_TYPE.to_string(), + value: "application/json".to_string(), + }], + Some(body), + Some(OAUTH_HTTP_TIMEOUT), + ) + .await?; + + if !status_is_success(response.status) { + return Err(anyhow!( + "Dynamic registration failed: HTTP {}: {}", + response.status, + String::from_utf8_lossy(&response.body.0) + )); + } + + let registration: ClientRegistrationResponse = serde_json::from_slice(&response.body.0) + .context("failed to parse registration response")?; + Ok(OAuthClientConfig { + client_id: registration.client_id, + client_secret: registration + .client_secret + .filter(|secret| !secret.is_empty()), + }) + } + + async fn oauth_request( + &self, + request: oauth2::HttpRequest, + ) -> std::result::Result { + let (parts, body) = request.into_parts(); + let headers = parts + .headers + .iter() + .map(|(name, value)| { + let value = value + .to_str() + .map_err(|err| OAuthHttpError(format!("invalid OAuth header value: {err}")))?; + Ok(HttpHeader { + name: name.to_string(), + value: value.to_string(), + }) + }) + .collect::, OAuthHttpError>>()?; + let response = self + .request( + parts.method.as_str(), + parts.uri.to_string().as_str(), + headers, + Some(body), + Some(OAUTH_HTTP_TIMEOUT), + ) + .await + .map_err(|err| OAuthHttpError(err.to_string()))?; + + let mut oauth_response = oauth2::HttpResponse::new(response.body.0); + *oauth_response.status_mut() = oauth2::http::StatusCode::from_u16(response.status) + .map_err(|err| OAuthHttpError(format!("invalid OAuth response status: {err}")))?; + for header in response.headers { + let name = oauth2::http::HeaderName::from_bytes(header.name.as_bytes()) + .map_err(|err| OAuthHttpError(format!("invalid OAuth response header: {err}")))?; + let value = oauth2::http::HeaderValue::from_str(&header.value).map_err(|err| { + OAuthHttpError(format!("invalid OAuth response header value: {err}")) + })?; + oauth_response.headers_mut().append(name, value); + } + Ok(oauth_response) + } + + async fn request( + &self, + method: &str, + url: &str, + extra_headers: Vec, + body: Option>, + timeout: Option, + ) -> Result { + let mut headers = protocol_headers(&self.default_headers)?; + headers.extend(extra_headers); + self.http_client + .http_request(HttpRequestParams { + method: method.to_string(), + url: url.to_string(), + headers, + body: body.map(Into::into), + timeout_ms: timeout + .map(|timeout| timeout.as_millis().clamp(1, u64::MAX as u128) as u64), + request_id: "oauth-request".to_string(), + stream_response: false, + }) + .await + .map_err(|err| anyhow!(err)) + } +} + +#[derive(Debug, Deserialize)] +struct OAuthDiscoveryMetadata { + #[serde(default)] + authorization_endpoint: Option, + #[serde(default)] + token_endpoint: Option, + #[serde(default)] + registration_endpoint: Option, + #[serde(default)] + scopes_supported: Option>, + #[serde(default)] + response_types_supported: Option>, +} + +#[derive(Debug, Serialize)] +struct ClientRegistrationRequest { + client_name: String, + redirect_uris: Vec, + grant_types: Vec, + token_endpoint_auth_method: String, + response_types: Vec, +} + +#[derive(Debug, Deserialize)] +struct ClientRegistrationResponse { + client_id: String, + client_secret: Option, +} + +fn oauth_client( + metadata: &OAuthDiscoveryMetadata, + config: &OAuthClientConfig, + redirect_uri: &str, +) -> Result< + BasicClient< + oauth2::EndpointSet, + oauth2::EndpointNotSet, + oauth2::EndpointNotSet, + oauth2::EndpointNotSet, + oauth2::EndpointSet, + >, +> { + let authorization_endpoint = metadata + .authorization_endpoint + .clone() + .ok_or_else(|| anyhow!("OAuth metadata did not include authorization endpoint"))?; + let token_endpoint = metadata + .token_endpoint + .clone() + .ok_or_else(|| anyhow!("OAuth metadata did not include token endpoint"))?; + let mut client = BasicClient::new(ClientId::new(config.client_id.clone())) + .set_auth_uri(AuthUrl::new(authorization_endpoint)?) + .set_token_uri(TokenUrl::new(token_endpoint)?) + .set_redirect_uri(RedirectUrl::new(redirect_uri.to_string())?); + if let Some(secret) = config.client_secret.clone() { + client = client.set_client_secret(ClientSecret::new(secret)); + } + Ok(client) +} + +fn parse_token_from_parse_error( + error: RequestTokenError, +) -> std::result::Result> +{ + match error { + RequestTokenError::Parse(parse_error, body) => { + match serde_json::from_slice::(&body) { + Ok(parsed) => Ok(parsed), + Err(_) => Err(RequestTokenError::Parse(parse_error, body)), + } + } + error => Err(error), + } +} + +fn protocol_headers(headers: &HeaderMap) -> Result> { + headers + .iter() + .map(|(name, value)| { + let value = value + .to_str() + .with_context(|| format!("invalid HTTP header value for `{name}`"))?; + Ok(HttpHeader { + name: name.to_string(), + value: value.to_string(), + }) + }) + .collect() +} + +fn status_is_success(status: u16) -> bool { + (200..300).contains(&status) +} + +pub(crate) fn normalize_scopes(scopes_supported: Option>) -> Option> { + let scopes_supported = scopes_supported?; + + let mut normalized = Vec::new(); + for scope in scopes_supported { + let scope = scope.trim(); + if scope.is_empty() { + continue; + } + let scope = scope.to_string(); + if !normalized.contains(&scope) { + normalized.push(scope); + } + } + + if normalized.is_empty() { + None + } else { + Some(normalized) + } +} + +/// Implements RFC 8414 section 3.1 for discovering well-known oauth endpoints. +pub(crate) fn discovery_paths(base_path: &str) -> Vec { + let trimmed = base_path.trim_start_matches('/').trim_end_matches('/'); + let canonical = "/.well-known/oauth-authorization-server".to_string(); + + if trimmed.is_empty() { + return vec![canonical]; + } + + let mut candidates = Vec::new(); + let mut push_unique = |candidate: String| { + if !candidates.contains(&candidate) { + candidates.push(candidate); + } + }; + + push_unique(format!("{canonical}/{trimmed}")); + push_unique(format!("/{trimmed}/.well-known/oauth-authorization-server")); + push_unique(canonical); + + candidates +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + + #[test] + fn discovery_paths_prefer_rfc_8414_resource_path() { + assert_eq!( + discovery_paths("/mcp"), + vec![ + "/.well-known/oauth-authorization-server/mcp".to_string(), + "/mcp/.well-known/oauth-authorization-server".to_string(), + "/.well-known/oauth-authorization-server".to_string(), + ] + ); + } + + #[test] + fn discovery_paths_deduplicate_root_path() { + assert_eq!( + discovery_paths("/"), + vec!["/.well-known/oauth-authorization-server".to_string()] + ); + } + + #[test] + fn normalize_scopes_trims_empties_and_deduplicates() { + assert_eq!( + normalize_scopes(Some(vec![ + "read".to_string(), + " write ".to_string(), + "".to_string(), + "read".to_string(), + ])), + Some(vec!["read".to_string(), "write".to_string()]) + ); + } +} diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index 5adfae9347..48e6c0526c 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -20,6 +20,7 @@ use anyhow::Context; use anyhow::Error; use anyhow::Result; use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::HttpClient; use oauth2::AccessToken; use oauth2::EmptyExtraTokenFields; use oauth2::RefreshToken; @@ -45,9 +46,9 @@ use tracing::warn; use codex_keyring_store::DefaultKeyringStore; use codex_keyring_store::KeyringStore; -use rmcp::transport::auth::AuthorizationManager; use tokio::sync::Mutex; +use crate::mcp_oauth_http::OAuthHttpClient; use codex_utils_home_dir::find_codex_home; const KEYRING_SERVICE: &str = "Codex MCP Credentials"; @@ -256,8 +257,7 @@ pub(crate) struct OAuthPersistor { struct OAuthPersistorInner { server_name: String, - url: String, - authorization_manager: Arc>, + oauth_http: OAuthHttpClient, store_mode: OAuthCredentialsStoreMode, last_credentials: Mutex>, } @@ -265,107 +265,82 @@ struct OAuthPersistorInner { impl OAuthPersistor { pub(crate) fn new( server_name: String, - url: String, - authorization_manager: Arc>, + oauth_http: OAuthHttpClient, store_mode: OAuthCredentialsStoreMode, initial_credentials: Option, ) -> Self { Self { inner: Arc::new(OAuthPersistorInner { server_name, - url, - authorization_manager, + oauth_http, store_mode, last_credentials: Mutex::new(initial_credentials), }), } } + pub(crate) fn with_http_client( + server_name: String, + http_client: Arc, + default_headers: reqwest::header::HeaderMap, + store_mode: OAuthCredentialsStoreMode, + initial_credentials: Option, + ) -> Self { + Self::new( + server_name, + OAuthHttpClient::from_default_headers(http_client, default_headers), + store_mode, + initial_credentials, + ) + } + /// Persists the latest stored credentials if they have changed. - /// Deletes the credentials if they are no longer present. - #[expect( - clippy::await_holding_invalid_type, - reason = "AuthorizationManager async access must be serialized through its mutex" - )] pub(crate) async fn persist_if_needed(&self) -> Result<()> { - let (client_id, maybe_credentials) = { - let manager = self.inner.authorization_manager.clone(); - let guard = manager.lock().await; - guard.get_credentials().await - }?; - - match maybe_credentials { - Some(credentials) => { - let mut last_credentials = self.inner.last_credentials.lock().await; - let new_token_response = WrappedOAuthTokenResponse(credentials.clone()); - let same_token = last_credentials - .as_ref() - .map(|prev| prev.token_response == new_token_response) - .unwrap_or(false); - let expires_at = if same_token { - last_credentials.as_ref().and_then(|prev| prev.expires_at) - } else { - compute_expires_at_millis(&credentials) - }; - let stored = StoredOAuthTokens { - server_name: self.inner.server_name.clone(), - url: self.inner.url.clone(), - client_id, - token_response: new_token_response, - expires_at, - }; - if last_credentials.as_ref() != Some(&stored) { - save_oauth_tokens(&self.inner.server_name, &stored, self.inner.store_mode)?; - *last_credentials = Some(stored); - } - } - None => { - let mut last_serialized = self.inner.last_credentials.lock().await; - if last_serialized.take().is_some() - && let Err(error) = delete_oauth_tokens( - &self.inner.server_name, - &self.inner.url, - self.inner.store_mode, - ) - { - warn!( - "failed to remove OAuth tokens for server {}: {error}", - self.inner.server_name - ); - } - } + let last_credentials = self.inner.last_credentials.lock().await; + if let Some(stored) = last_credentials.as_ref() { + save_oauth_tokens(&self.inner.server_name, stored, self.inner.store_mode)?; } - Ok(()) } - #[expect( - clippy::await_holding_invalid_type, - reason = "AuthorizationManager async access must be serialized through its mutex" - )] pub(crate) async fn refresh_if_needed(&self) -> Result<()> { - let expires_at = { + let tokens = { let guard = self.inner.last_credentials.lock().await; - guard.as_ref().and_then(|tokens| tokens.expires_at) + let Some(tokens) = guard.as_ref() else { + return Ok(()); + }; + if !token_needs_refresh(tokens.expires_at) { + return Ok(()); + } + tokens.clone() }; - if !token_needs_refresh(expires_at) { - return Ok(()); - } - - { - let manager = self.inner.authorization_manager.clone(); - let guard = manager.lock().await; - guard.refresh_token().await.with_context(|| { + if let Some(refreshed) = self + .inner + .oauth_http + .refresh_token(&tokens) + .await + .with_context(|| { format!( "failed to refresh OAuth tokens for server {}", self.inner.server_name ) - })?; + })? + { + let mut guard = self.inner.last_credentials.lock().await; + *guard = Some(refreshed); } self.persist_if_needed().await } + + pub(crate) async fn access_token(&self) -> Result> { + self.refresh_if_needed().await?; + let guard = self.inner.last_credentials.lock().await; + Ok(guard + .as_ref() + .map(|tokens| tokens.token_response.0.access_token().secret().to_string())) + } } const FALLBACK_FILENAME: &str = ".credentials.json"; diff --git a/codex-rs/rmcp-client/src/perform_oauth_login.rs b/codex-rs/rmcp-client/src/perform_oauth_login.rs index 5bdb315381..f4408a0fc2 100644 --- a/codex-rs/rmcp-client/src/perform_oauth_login.rs +++ b/codex-rs/rmcp-client/src/perform_oauth_login.rs @@ -7,9 +7,9 @@ use anyhow::Context; use anyhow::Result; use anyhow::anyhow; use anyhow::bail; -use reqwest::ClientBuilder; +use codex_exec_server::HttpClient; +use codex_exec_server::ReqwestHttpClient; use reqwest::Url; -use rmcp::transport::auth::OAuthState; use tiny_http::Response; use tiny_http::Server; use tokio::sync::oneshot; @@ -18,10 +18,10 @@ use urlencoding::decode; use crate::StoredOAuthTokens; use crate::WrappedOAuthTokenResponse; +use crate::mcp_oauth_http::OAuthAuthorizationSession; +use crate::mcp_oauth_http::OAuthHttpClient; use crate::oauth::compute_expires_at_millis; use crate::save_oauth_tokens; -use crate::utils::apply_default_headers; -use crate::utils::build_default_headers; use codex_config::types::OAuthCredentialsStoreMode; struct OauthHeaders { @@ -80,6 +80,34 @@ pub async fn perform_oauth_login( oauth_resource: Option<&str>, callback_port: Option, callback_url: Option<&str>, +) -> Result<()> { + perform_oauth_login_with_client( + server_name, + server_url, + store_mode, + http_headers, + env_http_headers, + scopes, + oauth_resource, + callback_port, + callback_url, + Arc::new(ReqwestHttpClient), + ) + .await +} + +#[allow(clippy::too_many_arguments)] +pub async fn perform_oauth_login_with_client( + server_name: &str, + server_url: &str, + store_mode: OAuthCredentialsStoreMode, + http_headers: Option>, + env_http_headers: Option>, + scopes: &[String], + oauth_resource: Option<&str>, + callback_port: Option, + callback_url: Option<&str>, + http_client: Arc, ) -> Result<()> { perform_oauth_login_with_browser_output( server_name, @@ -92,6 +120,7 @@ pub async fn perform_oauth_login( callback_port, callback_url, /*emit_browser_url*/ true, + http_client, ) .await } @@ -107,6 +136,34 @@ pub async fn perform_oauth_login_silent( oauth_resource: Option<&str>, callback_port: Option, callback_url: Option<&str>, +) -> Result<()> { + perform_oauth_login_silent_with_client( + server_name, + server_url, + store_mode, + http_headers, + env_http_headers, + scopes, + oauth_resource, + callback_port, + callback_url, + Arc::new(ReqwestHttpClient), + ) + .await +} + +#[allow(clippy::too_many_arguments)] +pub async fn perform_oauth_login_silent_with_client( + server_name: &str, + server_url: &str, + store_mode: OAuthCredentialsStoreMode, + http_headers: Option>, + env_http_headers: Option>, + scopes: &[String], + oauth_resource: Option<&str>, + callback_port: Option, + callback_url: Option<&str>, + http_client: Arc, ) -> Result<()> { perform_oauth_login_with_browser_output( server_name, @@ -119,6 +176,7 @@ pub async fn perform_oauth_login_silent( callback_port, callback_url, /*emit_browser_url*/ false, + http_client, ) .await } @@ -135,6 +193,7 @@ async fn perform_oauth_login_with_browser_output( callback_port: Option, callback_url: Option<&str>, emit_browser_url: bool, + http_client: Arc, ) -> Result<()> { let headers = OauthHeaders { http_headers, @@ -151,6 +210,7 @@ async fn perform_oauth_login_with_browser_output( callback_port, callback_url, /*timeout_secs*/ None, + http_client, ) .await? .finish(emit_browser_url) @@ -169,6 +229,36 @@ pub async fn perform_oauth_login_return_url( timeout_secs: Option, callback_port: Option, callback_url: Option<&str>, +) -> Result { + perform_oauth_login_return_url_with_client( + server_name, + server_url, + store_mode, + http_headers, + env_http_headers, + scopes, + oauth_resource, + timeout_secs, + callback_port, + callback_url, + Arc::new(ReqwestHttpClient), + ) + .await +} + +#[allow(clippy::too_many_arguments)] +pub async fn perform_oauth_login_return_url_with_client( + server_name: &str, + server_url: &str, + store_mode: OAuthCredentialsStoreMode, + http_headers: Option>, + env_http_headers: Option>, + scopes: &[String], + oauth_resource: Option<&str>, + timeout_secs: Option, + callback_port: Option, + callback_url: Option<&str>, + http_client: Arc, ) -> Result { let headers = OauthHeaders { http_headers, @@ -185,6 +275,7 @@ pub async fn perform_oauth_login_return_url( callback_port, callback_url, timeout_secs, + http_client, ) .await?; @@ -329,7 +420,8 @@ impl OauthLoginHandle { struct OauthLoginFlow { auth_url: String, - oauth_state: OAuthState, + oauth_http: OAuthHttpClient, + oauth_session: Option, rx: oneshot::Receiver, guard: CallbackServerGuard, server_name: String, @@ -412,6 +504,7 @@ impl OauthLoginFlow { callback_port: Option, callback_url: Option<&str>, timeout_secs: Option, + http_client: Arc, ) -> Result { const DEFAULT_OAUTH_TIMEOUT_SECS: i64 = 300; @@ -437,25 +530,19 @@ impl OauthLoginFlow { http_headers, env_http_headers, } = headers; - let default_headers = build_default_headers(http_headers, env_http_headers)?; - let http_client = apply_default_headers(ClientBuilder::new(), &default_headers).build()?; - - let mut oauth_state = OAuthState::new(server_url, Some(http_client)).await?; - let scope_refs: Vec<&str> = scopes.iter().map(String::as_str).collect(); - oauth_state - .start_authorization(&scope_refs, &redirect_uri, Some("Codex")) + let oauth_http = OAuthHttpClient::new(http_client, http_headers, env_http_headers)?; + let oauth_session = oauth_http + .start_authorization(server_url, scopes, &redirect_uri, "Codex") .await?; - let auth_url = append_query_param( - &oauth_state.get_authorization_url().await?, - "resource", - oauth_resource, - ); + let auth_url = + append_query_param(&oauth_session.authorization_url, "resource", oauth_resource); let timeout_secs = timeout_secs.unwrap_or(DEFAULT_OAUTH_TIMEOUT_SECS).max(1); let timeout = Duration::from_secs(timeout_secs as u64); Ok(Self { auth_url, - oauth_state, + oauth_http, + oauth_session: Some(oauth_session), rx, guard, server_name: server_name.to_string(), @@ -503,19 +590,16 @@ impl OauthLoginFlow { CallbackResult::Error(error) => return Err(anyhow!(error)), }; - self.oauth_state - .handle_callback(&code, &csrf_state) + let oauth_session = self + .oauth_session + .take() + .ok_or_else(|| anyhow!("OAuth login flow was already completed"))?; + let (client_id, credentials) = self + .oauth_http + .exchange_code(oauth_session, &code, &csrf_state) .await .context("failed to handle OAuth callback")?; - let (client_id, credentials_opt) = self - .oauth_state - .get_credentials() - .await - .context("failed to retrieve OAuth credentials")?; - let credentials = credentials_opt - .ok_or_else(|| anyhow!("OAuth provider did not return credentials"))?; - let expires_at = compute_expires_at_millis(&credentials); let stored = StoredOAuthTokens { server_name: self.server_name.clone(), diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 5cdb1d441e..7c378b268d 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -12,7 +12,6 @@ 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 codex_exec_server::HttpClient; use futures::FutureExt; @@ -45,9 +44,6 @@ use rmcp::service::RoleClient; use rmcp::service::RunningService; use rmcp::service::{self}; use rmcp::transport::StreamableHttpClientTransport; -use rmcp::transport::auth::AuthClient; -use rmcp::transport::auth::AuthError; -use rmcp::transport::auth::OAuthState; use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig; use rmcp::transport::streamable_http_client::StreamableHttpError; use serde::Deserialize; @@ -68,7 +64,6 @@ use crate::oauth::StoredOAuthTokens; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; use crate::stdio_server_launcher::StdioServerTransport; -use crate::utils::apply_default_headers; use crate::utils::build_default_headers; use codex_config::types::OAuthCredentialsStoreMode; @@ -80,7 +75,7 @@ enum PendingTransport { transport: StreamableHttpClientTransport, }, StreamableHttpWithOAuth { - transport: StreamableHttpClientTransport>, + transport: StreamableHttpClientTransport, oauth_persistor: OAuthPersistor, }, } @@ -708,11 +703,7 @@ impl RmcpClient { oauth_persistor, }) } - Err(err) - if err.downcast_ref::().is_some_and(|auth_err| { - matches!(auth_err, AuthError::NoAuthorizationSupport) - }) => - { + Err(err) if err.to_string().contains("No authorization support") => { let access_token = initial_tokens .token_response .0 @@ -730,6 +721,7 @@ impl RmcpClient { Arc::clone(http_client), default_headers, /*auth_provider*/ None, + /*oauth_persistor*/ None, ), http_config, ); @@ -749,6 +741,7 @@ impl RmcpClient { Arc::clone(http_client), default_headers, auth_provider.clone(), + /*oauth_persistor*/ None, ), http_config, ); @@ -943,52 +936,27 @@ async fn create_oauth_transport_and_runtime( default_headers: HeaderMap, http_client: Arc, ) -> Result<( - StreamableHttpClientTransport>, + StreamableHttpClientTransport, OAuthPersistor, )> { - let builder = apply_default_headers(reqwest::Client::builder(), &default_headers); - let oauth_metadata_client = build_reqwest_client_with_custom_ca(builder)?; - // TODO(aibrahim): teach OAuth bootstrap and refresh to use the same - // shared HTTP client abstraction instead of always creating the local - // reqwest metadata client here. - let mut oauth_state = - OAuthState::new(url.to_string(), Some(oauth_metadata_client.clone())).await?; - - oauth_state - .set_credentials( - &initial_tokens.client_id, - initial_tokens.token_response.0.clone(), - ) - .await?; - - let manager = match oauth_state { - OAuthState::Authorized(manager) => manager, - OAuthState::Unauthorized(manager) => manager, - OAuthState::Session(_) | OAuthState::AuthorizedHttpClient(_) => { - return Err(anyhow!("unexpected OAuth state during client setup")); - } - }; - - let auth_client = AuthClient::new( - StreamableHttpClientAdapter::new(http_client, default_headers, /*auth_provider*/ None), - manager, - ); - let auth_manager = auth_client.auth_manager.clone(); - - let transport = StreamableHttpClientTransport::with_client( - auth_client, - StreamableHttpClientTransportConfig::with_uri(url.to_string()), - ); - - let runtime = OAuthPersistor::new( + let oauth_persistor = OAuthPersistor::with_http_client( server_name.to_string(), - url.to_string(), - auth_manager, + Arc::clone(&http_client), + default_headers.clone(), credentials_store, Some(initial_tokens), ); - Ok((transport, runtime)) + let transport = StreamableHttpClientTransport::with_client( + StreamableHttpClientAdapter::new( + http_client, + default_headers, + /*auth_provider*/ None, + Some(oauth_persistor.clone()), + ), + StreamableHttpClientTransportConfig::with_uri(url.to_string()), + ); + Ok((transport, oauth_persistor)) } #[cfg(test)] diff --git a/codex-rs/rmcp-client/src/utils.rs b/codex-rs/rmcp-client/src/utils.rs index 6176f071ce..c535958544 100644 --- a/codex-rs/rmcp-client/src/utils.rs +++ b/codex-rs/rmcp-client/src/utils.rs @@ -1,7 +1,6 @@ use anyhow::Result; use anyhow::anyhow; use codex_config::types::McpServerEnvVar; -use reqwest::ClientBuilder; use reqwest::header::HeaderMap; use reqwest::header::HeaderName; use reqwest::header::HeaderValue; @@ -115,17 +114,6 @@ pub(crate) fn build_default_headers( Ok(headers) } -pub(crate) fn apply_default_headers( - builder: ClientBuilder, - default_headers: &HeaderMap, -) -> ClientBuilder { - if default_headers.is_empty() { - builder - } else { - builder.default_headers(default_headers.clone()) - } -} - #[cfg(unix)] pub(crate) const DEFAULT_ENV_VARS: &[&str] = &[ "HOME", diff --git a/codex-rs/rmcp-client/tests/streamable_http_remote_oauth.rs b/codex-rs/rmcp-client/tests/streamable_http_remote_oauth.rs new file mode 100644 index 0000000000..b60a05e259 --- /dev/null +++ b/codex-rs/rmcp-client/tests/streamable_http_remote_oauth.rs @@ -0,0 +1,309 @@ +//! Regression coverage for remote-environment Streamable HTTP OAuth. +//! +//! The OAuth issuer in this test uses an unresolvable hostname. If any +//! discovery, registration, or token-exchange request bypasses the injected +//! `HttpClient`, the test fails before the callback can complete. + +mod streamable_http_test_support; + +use std::ffi::OsString; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::MutexGuard; +use std::sync::OnceLock; +use std::sync::PoisonError; + +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::ExecServerError; +use codex_exec_server::HttpClient; +use codex_exec_server::HttpHeader; +use codex_exec_server::HttpRequestParams; +use codex_exec_server::HttpRequestResponse; +use codex_exec_server::HttpResponseBodyStream; +use codex_rmcp_client::StoredOAuthTokens; +use codex_rmcp_client::WrappedOAuthTokenResponse; +use codex_rmcp_client::perform_oauth_login_return_url_with_client; +use codex_rmcp_client::save_oauth_tokens; +use futures::FutureExt as _; +use futures::future::BoxFuture; +use oauth2::AccessToken; +use oauth2::EmptyExtraTokenFields; +use oauth2::RefreshToken; +use oauth2::basic::BasicTokenType; +use pretty_assertions::assert_eq; +use rmcp::transport::auth::OAuthTokenResponse; +use serde_json::Value; +use serde_json::json; +use serial_test::serial; +use tempfile::TempDir; + +use streamable_http_test_support::call_echo_tool; +use streamable_http_test_support::create_remote_oauth_client; +use streamable_http_test_support::expected_echo_result; +use streamable_http_test_support::spawn_exec_server; +use streamable_http_test_support::spawn_streamable_http_server_with_oauth_bearer; + +#[derive(Clone, Default)] +struct RecordingHttpClient { + requests: Arc>>, +} + +impl RecordingHttpClient { + fn recorded_requests(&self) -> Vec { + self.requests + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone() + } +} + +impl HttpClient for RecordingHttpClient { + fn http_request( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result> { + let requests = Arc::clone(&self.requests); + async move { + requests + .lock() + .unwrap_or_else(PoisonError::into_inner) + .push(params.clone()); + + let response = match (params.method.as_str(), params.url.as_str()) { + ("GET", "http://oauth.test/.well-known/oauth-authorization-server/mcp") => { + json_response(json!({ + "authorization_endpoint": "http://oauth.test/authorize", + "token_endpoint": "http://oauth.test/token", + "registration_endpoint": "http://oauth.test/register", + "scopes_supported": ["tools.read", "tools.write"], + "response_types_supported": ["code"], + }))? + } + ("POST", "http://oauth.test/register") => json_response(json!({ + "client_id": "registered-client", + }))?, + ("POST", "http://oauth.test/token") => json_response(json!({ + "access_token": "remote-access-token", + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": "remote-refresh-token", + }))?, + _ => { + return Err(ExecServerError::HttpRequest(format!( + "unexpected HTTP request: {} {}", + params.method, params.url + ))); + } + }; + + Ok(response) + } + .boxed() + } + + fn http_request_stream( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> { + async move { + Err(ExecServerError::HttpRequest(format!( + "unexpected streaming HTTP request: {} {}", + params.method, params.url + ))) + } + .boxed() + } +} + +fn json_response(body: Value) -> Result { + let body = serde_json::to_vec(&body) + .map_err(|err| ExecServerError::HttpRequest(format!("serialize JSON response: {err}")))?; + Ok(HttpRequestResponse { + status: 200, + headers: vec![HttpHeader { + name: "content-type".to_string(), + value: "application/json".to_string(), + }], + body: body.into(), + }) +} + +struct TempCodexHome { + _guard: MutexGuard<'static, ()>, + previous: Option, + _dir: TempDir, +} + +impl TempCodexHome { + fn new() -> anyhow::Result { + static LOCK: OnceLock> = OnceLock::new(); + let guard = LOCK + .get_or_init(Mutex::default) + .lock() + .unwrap_or_else(PoisonError::into_inner); + let previous = std::env::var_os("CODEX_HOME"); + let dir = TempDir::new()?; + unsafe { + std::env::set_var("CODEX_HOME", dir.path()); + } + Ok(Self { + _guard: guard, + previous, + _dir: dir, + }) + } +} + +impl Drop for TempCodexHome { + fn drop(&mut self) { + unsafe { + match self.previous.as_ref() { + Some(value) => std::env::set_var("CODEX_HOME", value), + None => std::env::remove_var("CODEX_HOME"), + } + } + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn browser_callback_flow_uses_injected_http_client_for_oauth_requests() -> anyhow::Result<()> +{ + let _codex_home = TempCodexHome::new()?; + let http_client = RecordingHttpClient::default(); + + let handle = perform_oauth_login_return_url_with_client( + "remote-oauth-test", + "http://oauth.test/mcp", + OAuthCredentialsStoreMode::File, + /*http_headers*/ None, + /*env_http_headers*/ None, + &["tools.read".to_string()], + /*oauth_resource*/ None, + Some(10), + /*callback_port*/ None, + /*callback_url*/ None, + Arc::new(http_client.clone()), + ) + .await?; + + let authorization_url = reqwest::Url::parse(handle.authorization_url())?; + let mut state = None; + let mut redirect_uri = None; + for (name, value) in authorization_url.query_pairs() { + match name.as_ref() { + "state" => state = Some(value.into_owned()), + "redirect_uri" => redirect_uri = Some(value.into_owned()), + _ => {} + } + } + + let state = state.expect("authorization URL includes state"); + let redirect_uri = redirect_uri.expect("authorization URL includes redirect_uri"); + let callback_url = format!("{redirect_uri}?code=provider-code&state={state}"); + let callback_response = reqwest::get(callback_url).await?; + assert_eq!(callback_response.status(), reqwest::StatusCode::OK); + + handle.wait().await?; + + let requests = http_client.recorded_requests(); + assert_eq!( + requests + .iter() + .map(|request| (request.method.as_str(), request.url.as_str())) + .collect::>(), + vec![ + ( + "GET", + "http://oauth.test/.well-known/oauth-authorization-server/mcp" + ), + ("POST", "http://oauth.test/register"), + ("POST", "http://oauth.test/token"), + ] + ); + + let registration_body: Value = serde_json::from_slice( + &requests[1] + .body + .as_ref() + .expect("registration request has a body") + .0, + )?; + assert_eq!( + registration_body, + json!({ + "client_name": "Codex", + "redirect_uris": [redirect_uri], + "grant_types": ["authorization_code", "refresh_token"], + "token_endpoint_auth_method": "none", + "response_types": ["code"], + }) + ); + + let token_body = String::from_utf8( + requests[2] + .body + .as_ref() + .expect("token request has a body") + .0 + .clone(), + )?; + assert!(token_body.contains("grant_type=authorization_code")); + assert!(token_body.contains("code=provider-code")); + assert!(token_body.contains("client_id=registered-client")); + assert!(token_body.contains("code_verifier=")); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn stored_oauth_refreshes_and_authenticates_remote_streamable_http() -> anyhow::Result<()> { + let _codex_home = TempCodexHome::new()?; + let refresh_token = "remote-refresh-token"; + let refreshed_access_token = "refreshed-remote-access-token"; + let (_server, base_url) = + spawn_streamable_http_server_with_oauth_bearer(refreshed_access_token, refresh_token) + .await?; + let exec_server = spawn_exec_server().await?; + + save_oauth_tokens( + "test-streamable-http-remote-oauth", + &expired_oauth_tokens( + "test-streamable-http-remote-oauth", + &format!("{base_url}/mcp"), + "expired-access-token", + refresh_token, + ), + OAuthCredentialsStoreMode::File, + )?; + + let client = create_remote_oauth_client(&base_url, exec_server.client.clone()).await?; + let result = call_echo_tool(&client, "remote-oauth").await?; + + assert_eq!(result, expected_echo_result("remote-oauth")); + + Ok(()) +} + +fn expired_oauth_tokens( + server_name: &str, + url: &str, + access_token: &str, + refresh_token: &str, +) -> StoredOAuthTokens { + let mut token_response = OAuthTokenResponse::new( + AccessToken::new(access_token.to_string()), + BasicTokenType::Bearer, + EmptyExtraTokenFields {}, + ); + token_response.set_refresh_token(Some(RefreshToken::new(refresh_token.to_string()))); + + StoredOAuthTokens { + server_name: server_name.to_string(), + url: url.to_string(), + client_id: "stored-client".to_string(), + token_response: WrappedOAuthTokenResponse(token_response), + expires_at: Some(0), + } +} diff --git a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs index cfff33ab43..215bc76e21 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs @@ -128,10 +128,40 @@ pub(crate) async fn create_remote_client( base_url: &str, http_client: ExecServerClient, ) -> anyhow::Result { - let client = RmcpClient::new_streamable_http_client( + create_remote_client_with_bearer( "test-streamable-http-remote", - &format!("{base_url}/mcp"), + base_url, Some("test-bearer".to_string()), + http_client, + ) + .await +} + +/// Creates a Streamable HTTP RMCP client that authenticates using stored OAuth +/// credentials through the remote runtime HTTP API. +pub(crate) async fn create_remote_oauth_client( + base_url: &str, + http_client: ExecServerClient, +) -> anyhow::Result { + create_remote_client_with_bearer( + "test-streamable-http-remote-oauth", + base_url, + /*bearer_token*/ None, + http_client, + ) + .await +} + +async fn create_remote_client_with_bearer( + server_name: &str, + base_url: &str, + bearer_token: Option, + http_client: ExecServerClient, +) -> anyhow::Result { + let client = RmcpClient::new_streamable_http_client( + server_name, + &format!("{base_url}/mcp"), + bearer_token, /*http_headers*/ None, /*env_http_headers*/ None, OAuthCredentialsStoreMode::File, @@ -193,16 +223,38 @@ pub(crate) async fn arm_session_post_failure( } pub(crate) async fn spawn_streamable_http_server() -> anyhow::Result<(Child, String)> { + spawn_streamable_http_server_with_env(&[]).await +} + +pub(crate) async fn spawn_streamable_http_server_with_oauth_bearer( + bearer_token: &str, + refresh_token: &str, +) -> anyhow::Result<(Child, String)> { + spawn_streamable_http_server_with_env(&[ + ("MCP_EXPECT_BEARER", bearer_token), + ("MCP_OAUTH_ACCESS_TOKEN", bearer_token), + ("MCP_OAUTH_REFRESH_TOKEN", refresh_token), + ]) + .await +} + +async fn spawn_streamable_http_server_with_env( + env: &[(&str, &str)], +) -> anyhow::Result<(Child, String)> { let listener = TcpListener::bind("127.0.0.1:0")?; let port = listener.local_addr()?.port(); drop(listener); let bind_addr = format!("127.0.0.1:{port}"); let base_url = format!("http://{bind_addr}"); - let mut child = Command::new(streamable_http_server_bin()?) + let mut command = Command::new(streamable_http_server_bin()?); + command .kill_on_drop(true) - .env("MCP_STREAMABLE_HTTP_BIND_ADDR", &bind_addr) - .spawn()?; + .env("MCP_STREAMABLE_HTTP_BIND_ADDR", &bind_addr); + for (name, value) in env { + command.env(name, value); + } + let mut child = command.spawn()?; wait_for_streamable_http_server(&mut child, &bind_addr, Duration::from_secs(5)).await?; Ok((child, base_url))