From 89a3b89c4c1d1afaaa93b6669c9e4e03247f8a99 Mon Sep 17 00:00:00 2001 From: Celia Chen Date: Fri, 24 Jul 2026 19:31:52 +0000 Subject: [PATCH] Route MCP auth discovery through runtime HTTP clients (#35239) ## Why MCP authentication checks need to use the same HTTP routing as the MCP transport so servers reached through configured proxies can be discovered reliably. ## What changed - Resolve OAuth discovery and authentication status through each server's runtime HTTP client for both local and managed environments. - Keep local discovery capped at five seconds while allowing explicit login requests to retain their requested timeout. - Resolve refreshed MCP configuration and its runtime context from the same snapshot. ## Testing - Cover OAuth discovery through an environment proxy and macOS system proxy resolution. - Cover proxied MCP startup and runtime refresh with updated authorization headers. - Verify capped and preserved OAuth discovery timeout policies. GitOrigin-RevId: 461fb1d4786e547df8b1e6b2215a8ac40438a3aa --- .../src/request_processors/mcp_processor.rs | 14 +- codex-rs/cli/tests/mcp_list.rs | 196 ++++++++++++++++++ codex-rs/codex-mcp/src/mcp/auth.rs | 55 +++-- codex-rs/core/src/codex_thread.rs | 12 +- .../suite/mcp_startup_refresh_http_proxy.rs | 158 ++++++++++++++ codex-rs/core/tests/suite/mod.rs | 1 + codex-rs/rmcp-client/src/auth_status.rs | 131 +++++++++++- codex-rs/rmcp-client/src/lib.rs | 1 + codex-rs/rmcp-client/src/oauth_http_client.rs | 22 ++ 9 files changed, 541 insertions(+), 49 deletions(-) create mode 100644 codex-rs/core/tests/suite/mcp_startup_refresh_http_proxy.rs diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index 881995bb7f97..a81b2758fce3 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -172,8 +172,12 @@ impl McpRequestProcessor { })?; let discovered_scopes = if scopes.is_none() && server.scopes.is_none() { - discover_supported_scopes_with_http_client(&server.transport, Arc::clone(&http_client)) - .await + discover_supported_scopes_with_http_client( + &server.transport, + Arc::clone(&http_client), + codex_rmcp_client::OAuthDiscoveryTimeout::Requested, + ) + .await } else { None }; @@ -250,11 +254,7 @@ impl McpRequestProcessor { let mcp_manager = self.thread_manager.mcp_manager(); let auth = self.auth_manager.auth().await; let (mcp_config, runtime_context) = match thread { - Some(thread) => { - let mcp_config = thread.runtime_mcp_config(&config).await; - let (_, runtime_context) = thread.current_mcp_config_and_runtime_context().await; - (mcp_config, runtime_context) - } + Some(thread) => thread.runtime_mcp_config_and_context(&config).await, None => { let mcp_config = mcp_manager.runtime_config(&config).await; let runtime_context = McpRuntimeContext::new( diff --git a/codex-rs/cli/tests/mcp_list.rs b/codex-rs/cli/tests/mcp_list.rs index bed3505985f7..6b762baaf0ca 100644 --- a/codex-rs/cli/tests/mcp_list.rs +++ b/codex-rs/cli/tests/mcp_list.rs @@ -1,4 +1,9 @@ +use std::io::Read; +use std::io::Write; +use std::net::TcpListener; use std::path::Path; +use std::time::Duration; +use std::time::Instant; use anyhow::Result; use codex_config::types::McpServerTransportConfig; @@ -10,6 +15,16 @@ use pretty_assertions::assert_eq; use serde_json::Value as JsonValue; use serde_json::json; use tempfile::TempDir; +#[cfg(target_os = "macos")] +use wiremock::Mock; +#[cfg(target_os = "macos")] +use wiremock::MockServer; +#[cfg(target_os = "macos")] +use wiremock::ResponseTemplate; +#[cfg(target_os = "macos")] +use wiremock::matchers::method; +#[cfg(target_os = "macos")] +use wiremock::matchers::path; fn codex_command(codex_home: &Path) -> Result { let mut cmd = assert_cmd::Command::new(codex_utils_cargo_bin::cargo_bin("codex")?); @@ -17,6 +32,18 @@ fn codex_command(codex_home: &Path) -> Result { Ok(cmd) } +async fn configure_http_oauth_server(codex_home: &Path, url: &str) -> Result<()> { + let mut servers = load_global_mcp_servers(codex_home).await?; + servers.insert( + "oauth".to_string(), + toml::from_str(&format!("url = \"{url}\""))?, + ); + ConfigEditsBuilder::new(codex_home) + .replace_mcp_servers(&servers) + .apply_blocking()?; + Ok(()) +} + #[test] fn list_shows_empty_state() -> Result<()> { let codex_home = TempDir::new()?; @@ -30,6 +57,175 @@ fn list_shows_empty_state() -> Result<()> { Ok(()) } +#[tokio::test] +async fn list_discovers_local_oauth_server_through_environment_proxy() -> Result<()> { + let codex_home = TempDir::new()?; + configure_http_oauth_server(codex_home.path(), "http://mcp-proxy.invalid/mcp").await?; + + let listener = TcpListener::bind("127.0.0.1:0")?; + listener.set_nonblocking(true)?; + let proxy_url = format!("http://{}", listener.local_addr()?); + let proxy = std::thread::spawn(move || -> Result> { + let resource_metadata = json!({ + "resource": "http://mcp-proxy.invalid/mcp", + "authorization_servers": ["http://mcp-proxy.invalid"], + }) + .to_string(); + let authorization_metadata = json!({ + "authorization_endpoint": "https://oauth.example/authorize", + "token_endpoint": "https://oauth.example/token", + }) + .to_string(); + let responses = [ + concat!( + "HTTP/1.1 401 Unauthorized\r\n", + "www-authenticate: Bearer resource_metadata=\"http://mcp-proxy.invalid/oauth-resource\"\r\n", + "content-length: 0\r\n", + "connection: close\r\n\r\n" + ) + .to_string(), + format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{resource_metadata}", + resource_metadata.len() + ), + format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{authorization_metadata}", + authorization_metadata.len() + ), + ]; + + let mut requests = Vec::new(); + for response in responses { + let deadline = Instant::now() + Duration::from_secs(30); + let mut stream = loop { + match listener.accept() { + Ok((stream, _)) => break stream, + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { + anyhow::ensure!( + Instant::now() < deadline, + "proxy did not receive OAuth discovery request {}", + requests.len() + 1 + ); + std::thread::sleep(Duration::from_millis(10)); + } + Err(error) => return Err(error.into()), + } + }; + stream.set_read_timeout(Some(Duration::from_secs(5)))?; + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + while !request.windows(4).any(|window| window == b"\r\n\r\n") { + let bytes_read = stream.read(&mut buffer)?; + anyhow::ensure!(bytes_read > 0, "proxy request ended before its headers"); + request.extend_from_slice(&buffer[..bytes_read]); + anyhow::ensure!( + request.len() <= 64 * 1024, + "proxy request headers are too large" + ); + } + let request = String::from_utf8(request)?; + requests.push(request.lines().next().unwrap_or_default().to_string()); + stream.write_all(response.as_bytes())?; + } + + Ok(requests) + }); + + let mut command = codex_command(codex_home.path())?; + command + .env("HTTP_PROXY", &proxy_url) + .env("http_proxy", &proxy_url) + .env_remove("HTTPS_PROXY") + .env_remove("https_proxy") + .env_remove("ALL_PROXY") + .env_remove("all_proxy") + .env_remove("NO_PROXY") + .env_remove("no_proxy") + .args([ + "-c", + "mcp_oauth_credentials_store=\"file\"", + "mcp", + "list", + "--json", + ]); + let output = command.output()?; + assert!( + output.status.success(), + "mcp list failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + let proxy_requests = proxy + .join() + .expect("OAuth discovery proxy thread should finish")?; + + assert_eq!( + proxy_requests, + vec![ + "GET http://mcp-proxy.invalid/mcp HTTP/1.1", + "GET http://mcp-proxy.invalid/oauth-resource HTTP/1.1", + "GET http://mcp-proxy.invalid/.well-known/oauth-authorization-server HTTP/1.1", + ] + ); + let entries: JsonValue = serde_json::from_slice(&output.stdout)?; + assert_eq!(entries[0]["name"], "oauth"); + assert_eq!( + entries[0]["auth_status"], + "not_logged_in", + "OAuth discovery failed after proxy requests {proxy_requests:?}; stderr: {}", + String::from_utf8_lossy(&output.stderr) + ); + Ok(()) +} + +#[cfg(target_os = "macos")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn list_with_macos_proxy_resolution_does_not_panic() -> Result<()> { + let codex_home = TempDir::new()?; + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": "https://oauth.example/authorize", + "token_endpoint": "https://oauth.example/token", + }))) + .expect(2) + .mount(&server) + .await; + configure_http_oauth_server(codex_home.path(), &format!("{}/mcp", server.uri())).await?; + + for respect_system_proxy in [false, true] { + let system_proxy_override = format!("features.respect_system_proxy={respect_system_proxy}"); + let mut command = codex_command(codex_home.path())?; + command + .env_remove("HTTP_PROXY") + .env_remove("http_proxy") + .env_remove("HTTPS_PROXY") + .env_remove("https_proxy") + .env_remove("ALL_PROXY") + .env_remove("all_proxy") + .env_remove("NO_PROXY") + .env_remove("no_proxy") + .args([ + "-c", + &system_proxy_override, + "-c", + "mcp_oauth_credentials_store=\"file\"", + "mcp", + "list", + "--json", + ]); + let output = command.output()?; + assert!( + output.status.success(), + "macOS proxy resolution should not panic with respect_system_proxy={respect_system_proxy}: {}", + String::from_utf8_lossy(&output.stderr) + ); + let entries: JsonValue = serde_json::from_slice(&output.stdout)?; + assert_eq!(entries[0]["auth_status"], "not_logged_in"); + } + Ok(()) +} + #[tokio::test] async fn list_and_get_render_expected_output() -> Result<()> { let codex_home = TempDir::new()?; diff --git a/codex-rs/codex-mcp/src/mcp/auth.rs b/codex-rs/codex-mcp/src/mcp/auth.rs index 8b561edae1f1..f79064a14a33 100644 --- a/codex-rs/codex-mcp/src/mcp/auth.rs +++ b/codex-rs/codex-mcp/src/mcp/auth.rs @@ -10,8 +10,8 @@ use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::HttpClient; use codex_login::CodexAuth; use codex_rmcp_client::McpAuthState; +use codex_rmcp_client::OAuthDiscoveryTimeout; use codex_rmcp_client::OAuthProviderError; -use codex_rmcp_client::determine_streamable_http_auth_status; use codex_rmcp_client::determine_streamable_http_auth_status_with_http_client; use codex_rmcp_client::discover_streamable_http_oauth; use codex_rmcp_client::discover_streamable_http_oauth_with_http_client; @@ -80,6 +80,7 @@ pub async fn oauth_login_support(transport: &McpServerTransportConfig) -> McpOAu pub async fn oauth_login_support_with_http_client( transport: &McpServerTransportConfig, http_client: Arc, + discovery_timeout: OAuthDiscoveryTimeout, ) -> McpOAuthLoginSupport { let Some(mut config) = oauth_login_candidate(transport) else { return McpOAuthLoginSupport::Unsupported; @@ -89,6 +90,7 @@ pub async fn oauth_login_support_with_http_client( config.http_headers.clone(), config.env_http_headers.clone(), http_client, + discovery_timeout, ) .await { @@ -134,8 +136,9 @@ pub async fn discover_supported_scopes( pub async fn discover_supported_scopes_with_http_client( transport: &McpServerTransportConfig, http_client: Arc, + discovery_timeout: OAuthDiscoveryTimeout, ) -> Option> { - match oauth_login_support_with_http_client(transport, http_client).await { + match oauth_login_support_with_http_client(transport, http_client, discovery_timeout).await { McpOAuthLoginSupport::Supported(config) => config.discovered_scopes, McpOAuthLoginSupport::Unsupported | McpOAuthLoginSupport::Unknown(_) => None, } @@ -255,35 +258,27 @@ async fn compute_auth_status( http_headers, env_http_headers, } => { - if config.is_local_environment() { - determine_streamable_http_auth_status( - server_name, - url, - bearer_token_env_var.as_deref(), - http_headers.clone(), - env_http_headers.clone(), - store_mode, - keyring_backend_kind, - ) - .boxed() - .await + let http_client = runtime_context + .resolve_http_client(server_name, config) + .map_err(anyhow::Error::msg)?; + let discovery_timeout = if config.is_local_environment() { + OAuthDiscoveryTimeout::LOCAL } else { - let http_client = runtime_context - .resolve_http_client(server_name, config) - .map_err(anyhow::Error::msg)?; - determine_streamable_http_auth_status_with_http_client( - server_name, - url, - bearer_token_env_var.as_deref(), - http_headers.clone(), - env_http_headers.clone(), - store_mode, - keyring_backend_kind, - http_client, - ) - .boxed() - .await - } + OAuthDiscoveryTimeout::Requested + }; + determine_streamable_http_auth_status_with_http_client( + server_name, + url, + bearer_token_env_var.as_deref(), + http_headers.clone(), + env_http_headers.clone(), + store_mode, + keyring_backend_kind, + http_client, + discovery_timeout, + ) + .boxed() + .await } } } diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index cfd5517f9b0c..c0423c87e318 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -601,9 +601,12 @@ impl CodexThread { self.session.get_config().await } - /// Resolves the MCP runtime configuration using this thread's extension data. - pub async fn runtime_mcp_config(&self, config: &crate::config::Config) -> codex_mcp::McpConfig { - self.session.runtime_mcp_config(config).await + /// Resolves MCP configuration and environment bindings from the same config snapshot. + pub async fn runtime_mcp_config_and_context( + &self, + config: &crate::config::Config, + ) -> (codex_mcp::McpConfig, codex_mcp::McpRuntimeContext) { + self.session.runtime_mcp_config_and_context(config).await } /// Captures the exact MCP config and environment bindings for the current thread state. @@ -611,8 +614,7 @@ impl CodexThread { &self, ) -> (Arc, codex_mcp::McpRuntimeContext) { let config = self.session.get_config().await; - let (mcp_config, runtime_context) = - self.session.runtime_mcp_config_and_context(&config).await; + let (mcp_config, runtime_context) = self.runtime_mcp_config_and_context(&config).await; (Arc::new(mcp_config), runtime_context) } diff --git a/codex-rs/core/tests/suite/mcp_startup_refresh_http_proxy.rs b/codex-rs/core/tests/suite/mcp_startup_refresh_http_proxy.rs new file mode 100644 index 000000000000..70f0dd2ff979 --- /dev/null +++ b/codex-rs/core/tests/suite/mcp_startup_refresh_http_proxy.rs @@ -0,0 +1,158 @@ +use std::collections::HashMap; +use std::time::Duration; + +use anyhow::Result; +use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID; +use codex_config::McpServerConfig; +use codex_config::McpServerTransportConfig; +use codex_features::Feature; +use core_test_support::apps_test_server::AppsTestServer; +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_mcp_server; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use tokio::process::Command; +use wiremock::MockServer; + +const PROXY_TEST_SUBPROCESS_ENV_VAR: &str = "CODEX_MCP_HTTP_PROXY_TEST_SUBPROCESS"; +const TEST_NAME: &str = "suite::mcp_startup_refresh_http_proxy::local_mcp_startup_and_refresh_use_configured_http_client"; +const SERVER_NAME: &str = "proxied_mcp"; +const SERVER_URL: &str = "http://mcp-proxy.invalid/api/codex/ps/mcp"; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn local_mcp_startup_and_refresh_use_configured_http_client() -> Result<()> { + skip_if_no_network!(Ok(())); + + if std::env::var_os(PROXY_TEST_SUBPROCESS_ENV_VAR).is_none() { + let proxy = MockServer::start().await; + let _apps_server = AppsTestServer::mount(&proxy).await?; + let mut command = Command::new(std::env::current_exe()?); + command.arg("--exact").arg(TEST_NAME); + for &key in codex_network_proxy::PROXY_ENV_KEYS { + command.env_remove(key); + } + command + .env(PROXY_TEST_SUBPROCESS_ENV_VAR, "1") + .env("HTTP_PROXY", proxy.uri()) + .env("http_proxy", proxy.uri()) + .env("NO_PROXY", codex_network_proxy::DEFAULT_NO_PROXY_VALUE) + .env("no_proxy", codex_network_proxy::DEFAULT_NO_PROXY_VALUE); + + let output = command.output().await?; + let requests = proxy + .received_requests() + .await + .expect("mock proxy should record MCP requests"); + assert!( + output.status.success(), + "subprocess test `{TEST_NAME}` failed\nstdout:\n{}\nstderr:\n{}\nproxy requests:\n{requests:#?}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr), + ); + let initialize_authorizations = requests + .iter() + .filter_map(|request| { + let body = serde_json::from_slice::(&request.body).ok()?; + (body.get("method").and_then(Value::as_str) == Some("initialize")) + .then(|| { + request + .headers + .get("authorization") + .and_then(|header| header.to_str().ok()) + .map(str::to_string) + }) + .flatten() + }) + .collect::>(); + assert_eq!( + initialize_authorizations, + vec!["Bearer initial", "Bearer refreshed"] + ); + return Ok(()); + } + + let responses_server = responses::start_mock_server().await; + let fixture = test_codex() + .with_config(|config| { + if cfg!(target_os = "linux") { + config + .features + .enable(Feature::RespectSystemProxy) + .expect("test config should allow the system proxy feature"); + config.respect_system_proxy = true; + } + let mut servers = config.mcp_servers.get().clone(); + servers.insert( + SERVER_NAME.to_string(), + McpServerConfig { + auth: Default::default(), + transport: McpServerTransportConfig::StreamableHttp { + url: SERVER_URL.to_string(), + bearer_token_env_var: None, + http_headers: Some(HashMap::from([( + "Authorization".to_string(), + "Bearer initial".to_string(), + )])), + env_http_headers: None, + }, + environment_id: DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + enabled: true, + required: false, + supports_parallel_tool_calls: false, + disabled_reason: None, + startup_timeout_sec: Some(Duration::from_secs(10)), + tool_timeout_sec: None, + default_tools_approval_mode: None, + enabled_tools: None, + disabled_tools: None, + scopes: None, + oauth: None, + oauth_resource: None, + tools: HashMap::new(), + }, + ); + config + .mcp_servers + .set(servers) + .expect("test MCP servers should accept any configuration"); + }) + .build_with_auto_env(&responses_server) + .await?; + wait_for_mcp_server(&fixture.codex, SERVER_NAME).await?; + + let mut refreshed_config = fixture.config.clone(); + let mut servers = refreshed_config.mcp_servers.get().clone(); + let server = servers + .get_mut(SERVER_NAME) + .expect("configured MCP server should exist"); + let McpServerTransportConfig::StreamableHttp { http_headers, .. } = &mut server.transport + else { + unreachable!("test MCP server should use streamable HTTP"); + }; + *http_headers = Some(HashMap::from([( + "Authorization".to_string(), + "Bearer refreshed".to_string(), + )])); + refreshed_config + .mcp_servers + .set(servers) + .expect("test MCP servers should accept the refreshed configuration"); + fixture.codex.refresh_runtime_config(refreshed_config).await; + let result = fixture + .codex + .call_mcp_tool( + SERVER_NAME, + "calendar_create_event", + Some(json!({ + "title": "Proxy refresh", + "starts_at": "2026-07-23T12:00:00Z", + })), + /*meta*/ None, + ) + .await?; + assert_eq!(result.is_error, Some(false)); + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index ca09b01ce98c..ea28172dc9df 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -79,6 +79,7 @@ mod mcp_auth_elicitation; mod mcp_auth_refresh; #[cfg(unix)] mod mcp_refresh_cleanup; +mod mcp_startup_refresh_http_proxy; mod mcp_tool_cache; mod mcp_tool_exposure; mod mcp_turn_metadata; diff --git a/codex-rs/rmcp-client/src/auth_status.rs b/codex-rs/rmcp-client/src/auth_status.rs index b8f0e53a79b1..092de2f056c9 100644 --- a/codex-rs/rmcp-client/src/auth_status.rs +++ b/codex-rs/rmcp-client/src/auth_status.rs @@ -23,6 +23,20 @@ use codex_config::types::OAuthCredentialsStoreMode; const DISCOVERY_TIMEOUT: Duration = Duration::from_secs(5); +/// Timeout policy for OAuth metadata discovery through a supplied HTTP client. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OAuthDiscoveryTimeout { + /// Preserve the timeout requested by the OAuth implementation. + Requested, + /// Cap OAuth discovery requests at the supplied duration. + Capped(Duration), +} + +impl OAuthDiscoveryTimeout { + /// Preserves the existing timeout for local OAuth discovery. + pub const LOCAL: Self = Self::Capped(DISCOVERY_TIMEOUT); +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct StreamableHttpOAuthDiscovery { pub scopes_supported: Option>, @@ -100,6 +114,7 @@ pub async fn determine_streamable_http_auth_status_with_http_client( store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, http_client: Arc, + discovery_timeout: OAuthDiscoveryTimeout, ) -> Result { let default_headers = match auth_status_before_discovery( server_name, @@ -120,6 +135,7 @@ pub async fn determine_streamable_http_auth_status_with_http_client( url, default_headers, http_client, + discovery_timeout, ) .await, ) @@ -224,10 +240,16 @@ pub async fn discover_streamable_http_oauth_with_http_client( http_headers: Option>, env_http_headers: Option>, http_client: Arc, + discovery_timeout: OAuthDiscoveryTimeout, ) -> Result> { let default_headers = build_default_headers(http_headers, env_http_headers)?; - discover_streamable_http_oauth_with_headers_and_http_client(url, default_headers, http_client) - .await + discover_streamable_http_oauth_with_headers_and_http_client( + url, + default_headers, + http_client, + discovery_timeout, + ) + .await } async fn discover_streamable_http_oauth_with_headers( @@ -247,12 +269,18 @@ async fn discover_streamable_http_oauth_with_headers_and_http_client( url: &str, default_headers: HeaderMap, http_client: Arc, + discovery_timeout: OAuthDiscoveryTimeout, ) -> Result> { - let authorization_manager = AuthorizationManager::new_with_oauth_http_client( - url, - Arc::new(OAuthHttpClientAdapter::new(http_client, default_headers)), - ) - .await?; + let oauth_http_client = match discovery_timeout { + OAuthDiscoveryTimeout::Requested => { + OAuthHttpClientAdapter::new(http_client, default_headers) + } + OAuthDiscoveryTimeout::Capped(max_timeout) => { + OAuthHttpClientAdapter::new_with_max_timeout(http_client, default_headers, max_timeout) + } + }; + let authorization_manager = + AuthorizationManager::new_with_oauth_http_client(url, Arc::new(oauth_http_client)).await?; discover_streamable_http_oauth_with_manager(&authorization_manager).await } @@ -298,10 +326,16 @@ mod tests { use axum::http::StatusCode; use axum::http::header::WWW_AUTHENTICATE; use axum::routing::get; + use codex_exec_server::ExecServerError; + use codex_exec_server::HttpRequestParams; + use codex_exec_server::HttpRequestResponse; + use codex_exec_server::HttpResponseBodyStream; + use futures::future::BoxFuture; use pretty_assertions::assert_eq; use serial_test::serial; use std::collections::HashMap; use std::ffi::OsString; + use std::sync::Mutex; use tokio::task::JoinHandle; struct TestServer { @@ -315,6 +349,40 @@ mod tests { } } + #[derive(Default)] + struct RecordingHttpClient { + timeout_ms: Mutex>>, + } + + impl HttpClient for RecordingHttpClient { + fn http_request( + &self, + _params: HttpRequestParams, + ) -> BoxFuture<'_, Result> { + Box::pin(async { + Err(ExecServerError::HttpRequest( + "unexpected buffered request".to_string(), + )) + }) + } + + fn http_request_stream( + &self, + params: HttpRequestParams, + ) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> + { + *self + .timeout_ms + .lock() + .expect("timeout recorder lock should not be poisoned") = Some(params.timeout_ms); + Box::pin(async { + Err(ExecServerError::HttpRequest( + "expected discovery request failure".to_string(), + )) + }) + } + } + async fn spawn_oauth_discovery_server(metadata: serde_json::Value) -> TestServer { let listener = tokio::net::TcpListener::bind("127.0.0.1:0") .await @@ -438,6 +506,55 @@ mod tests { ); } + #[tokio::test] + async fn routed_oauth_discovery_caps_local_discovery_timeout() { + let http_client = Arc::new(RecordingHttpClient::default()); + + let discovery = discover_streamable_http_oauth_with_http_client( + "http://example.com/mcp", + /*http_headers*/ None, + /*env_http_headers*/ None, + http_client.clone(), + OAuthDiscoveryTimeout::LOCAL, + ) + .await; + + assert!(matches!(discovery, Ok(None))); + assert_eq!( + *http_client + .timeout_ms + .lock() + .expect("timeout recorder lock should not be poisoned"), + Some(Some( + u64::try_from(DISCOVERY_TIMEOUT.as_millis()) + .expect("discovery timeout should fit in u64") + )) + ); + } + + #[tokio::test] + async fn routed_oauth_discovery_preserves_requested_timeout() { + let http_client = Arc::new(RecordingHttpClient::default()); + + let discovery = discover_streamable_http_oauth_with_http_client( + "http://example.com/mcp", + /*http_headers*/ None, + /*env_http_headers*/ None, + http_client.clone(), + OAuthDiscoveryTimeout::Requested, + ) + .await; + + assert!(matches!(discovery, Ok(None))); + assert_eq!( + *http_client + .timeout_ms + .lock() + .expect("timeout recorder lock should not be poisoned"), + Some(Some(30_000)) + ); + } + #[tokio::test] async fn discover_streamable_http_oauth_follows_protected_resource_metadata() { let authorization_server = spawn_oauth_discovery_server(serde_json::json!({ diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 8f4bb2953288..f54ca94185c6 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -15,6 +15,7 @@ mod utils; pub use auth_status::McpAuthState; pub use auth_status::McpLoginRequirement; +pub use auth_status::OAuthDiscoveryTimeout; pub use auth_status::StreamableHttpOAuthDiscovery; pub use auth_status::determine_streamable_http_auth_status; pub use auth_status::determine_streamable_http_auth_status_from_credentials; diff --git a/codex-rs/rmcp-client/src/oauth_http_client.rs b/codex-rs/rmcp-client/src/oauth_http_client.rs index 44cf63cb7c77..a3f23d4cf40d 100644 --- a/codex-rs/rmcp-client/src/oauth_http_client.rs +++ b/codex-rs/rmcp-client/src/oauth_http_client.rs @@ -16,6 +16,8 @@ use rmcp::transport::auth::OAuthHttpClientFuture; use rmcp::transport::auth::OAuthHttpRedirectPolicy; use rmcp::transport::auth::OAuthHttpRequest; +use crate::auth_status::OAuthDiscoveryTimeout; + const MAX_OAUTH_HTTP_RESPONSE_BODY_BYTES: usize = 1024 * 1024; static NEXT_OAUTH_REQUEST_ID: AtomicU64 = AtomicU64::new(0); @@ -23,6 +25,7 @@ static NEXT_OAUTH_REQUEST_ID: AtomicU64 = AtomicU64::new(0); pub(crate) struct OAuthHttpClientAdapter { http_client: Arc, default_headers: HeaderMap, + timeout: OAuthDiscoveryTimeout, } impl OAuthHttpClientAdapter { @@ -30,6 +33,19 @@ impl OAuthHttpClientAdapter { Self { http_client, default_headers, + timeout: OAuthDiscoveryTimeout::Requested, + } + } + + pub(crate) fn new_with_max_timeout( + http_client: Arc, + default_headers: HeaderMap, + max_timeout: Duration, + ) -> Self { + Self { + http_client, + default_headers, + timeout: OAuthDiscoveryTimeout::Capped(max_timeout), } } @@ -66,6 +82,12 @@ impl OAuthHttpClientAdapter { }) }) .collect::, OAuthHttpClientError>>()?; + let timeout = match self.timeout { + OAuthDiscoveryTimeout::Requested => timeout, + OAuthDiscoveryTimeout::Capped(max_timeout) => { + Some(timeout.map_or(max_timeout, |timeout| timeout.min(max_timeout))) + } + }; let timeout_ms = timeout.map(|timeout| { u64::try_from(timeout.as_millis()) .unwrap_or(u64::MAX)