From a27251eb5722ed8f0ac1ddda661e3225e883f698 Mon Sep 17 00:00:00 2001 From: Dennis Henry Date: Fri, 29 May 2026 09:59:19 -0400 Subject: [PATCH] feat: add headersHelper support for codex --- .../src/request_processors/mcp_processor.rs | 11 +- .../src/request_processors/plugins.rs | 4 +- codex-rs/cli/src/mcp_cmd.rs | 133 ++++- codex-rs/cli/tests/mcp_add_remove.rs | 80 +++ .../codex-mcp/src/connection_manager_tests.rs | 3 + codex-rs/codex-mcp/src/mcp/auth.rs | 20 +- codex-rs/codex-mcp/src/mcp/mod.rs | 1 + codex-rs/codex-mcp/src/mcp/mod_tests.rs | 2 + codex-rs/codex-mcp/src/rmcp_client.rs | 2 + codex-rs/codex-mcp/src/runtime.rs | 1 + codex-rs/config/src/lib.rs | 2 + codex-rs/config/src/mcp_edit.rs | 24 + codex-rs/config/src/mcp_edit_tests.rs | 65 +++ codex-rs/config/src/mcp_types.rs | 94 +++- codex-rs/config/src/mcp_types_tests.rs | 83 ++++ codex-rs/config/src/types.rs | 2 + codex-rs/core-plugins/src/manager_tests.rs | 4 + codex-rs/core/config.schema.json | 45 ++ codex-rs/core/src/config/config_tests.rs | 11 + codex-rs/core/src/config/edit.rs | 24 + codex-rs/core/src/config/edit_tests.rs | 1 + codex-rs/core/src/mcp_skill_dependencies.rs | 5 +- codex-rs/core/tests/suite/rmcp_client.rs | 2 + codex-rs/external-agent-migration/src/lib.rs | 167 +++++++ codex-rs/rmcp-client/src/auth_status.rs | 159 +++++- .../rmcp-client/src/perform_oauth_login.rs | 14 +- codex-rs/rmcp-client/src/rmcp_client.rs | 13 +- codex-rs/rmcp-client/src/utils.rs | 458 +++++++++++++++++- .../tests/streamable_http_test_support.rs | 2 + 29 files changed, 1395 insertions(+), 37 deletions(-) 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 ae62e2e7855..21583b72ece 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -131,13 +131,19 @@ impl McpRequestProcessor { ))); }; - let (url, http_headers, env_http_headers) = match &server.transport { + let (url, http_headers, env_http_headers, http_headers_helper) = match &server.transport { McpServerTransportConfig::StreamableHttp { url, http_headers, env_http_headers, + http_headers_helper, .. - } => (url.clone(), http_headers.clone(), env_http_headers.clone()), + } => ( + url.clone(), + http_headers.clone(), + env_http_headers.clone(), + http_headers_helper.clone(), + ), _ => { return Err(invalid_request( "OAuth login is only supported for streamable HTTP servers.", @@ -159,6 +165,7 @@ impl McpRequestProcessor { config.mcp_oauth_credentials_store_mode, http_headers, env_http_headers, + http_headers_helper, &resolved_scopes.scopes, server.oauth_client_id(), server.oauth_resource.as_deref(), diff --git a/codex-rs/app-server/src/request_processors/plugins.rs b/codex-rs/app-server/src/request_processors/plugins.rs index fd45a8a732c..0aaa433c41b 100644 --- a/codex-rs/app-server/src/request_processors/plugins.rs +++ b/codex-rs/app-server/src/request_processors/plugins.rs @@ -1567,7 +1567,7 @@ impl PluginRequestProcessor { ) { for (name, server) in plugin_mcp_servers { let oauth_config = match oauth_login_support(&server.transport).await { - McpOAuthLoginSupport::Supported(config) => config, + McpOAuthLoginSupport::Supported(config) => *config, McpOAuthLoginSupport::Unsupported => continue, McpOAuthLoginSupport::Unknown(err) => { warn!( @@ -1597,6 +1597,7 @@ impl PluginRequestProcessor { store_mode, oauth_config.http_headers.clone(), oauth_config.env_http_headers.clone(), + oauth_config.http_headers_helper.clone(), &resolved_scopes.scopes, oauth_client_id, server.oauth_resource.as_deref(), @@ -1613,6 +1614,7 @@ impl PluginRequestProcessor { store_mode, oauth_config.http_headers, oauth_config.env_http_headers, + oauth_config.http_headers_helper, &[], oauth_client_id, server.oauth_resource.as_deref(), diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index 0103782653b..0b8d662d2a8 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::path::PathBuf; use std::sync::Arc; use anyhow::Context; @@ -6,8 +7,10 @@ use anyhow::Result; use anyhow::anyhow; use anyhow::bail; use clap::ArgGroup; +use codex_config::AbsolutePathBuf; use codex_config::types::AppToolApproval; use codex_config::types::McpServerConfig; +use codex_config::types::McpServerHttpHeadersHelperConfig; use codex_config::types::McpServerOAuthConfig; use codex_config::types::McpServerTransportConfig; use codex_core::McpManager; @@ -143,6 +146,27 @@ pub struct AddMcpStreamableHttpArgs { /// Optional OAuth resource parameter to include during MCP login. #[arg(long = "oauth-resource", value_name = "RESOURCE", requires = "url")] pub oauth_resource: Option, + + /// Command that prints a JSON object of HTTP headers for this MCP server. + #[arg(long = "http-headers-helper", value_name = "COMMAND", requires = "url")] + pub http_headers_helper: Option, + + /// Argument to pass to --http-headers-helper. May be specified multiple times. + #[arg( + long = "http-headers-helper-arg", + value_name = "ARG", + requires = "http_headers_helper", + allow_hyphen_values = true + )] + pub http_headers_helper_args: Vec, + + /// Working directory for --http-headers-helper. Relative paths are resolved from the current directory. + #[arg( + long = "http-headers-helper-cwd", + value_name = "CWD", + requires = "http_headers_helper" + )] + pub http_headers_helper_cwd: Option, } #[derive(Debug, clap::Parser)] @@ -213,6 +237,7 @@ async fn perform_oauth_login_retry_without_scopes( store_mode: codex_config::types::OAuthCredentialsStoreMode, http_headers: Option>, env_http_headers: Option>, + http_headers_helper: Option, resolved_scopes: &ResolvedMcpOAuthScopes, oauth_client_id: Option<&str>, oauth_resource: Option<&str>, @@ -225,6 +250,7 @@ async fn perform_oauth_login_retry_without_scopes( store_mode, http_headers.clone(), env_http_headers.clone(), + http_headers_helper.clone(), &resolved_scopes.scopes, oauth_client_id, oauth_resource, @@ -242,6 +268,7 @@ async fn perform_oauth_login_retry_without_scopes( store_mode, http_headers, env_http_headers, + http_headers_helper, &[], oauth_client_id, oauth_resource, @@ -325,18 +352,38 @@ async fn run_add(config_overrides: &CliConfigOverrides, add_args: AddArgs) -> Re bearer_token_env_var, oauth_client_id, oauth_resource, + http_headers_helper, + http_headers_helper_args, + http_headers_helper_cwd, }), .. - } => ( - McpServerTransportConfig::StreamableHttp { - url, - bearer_token_env_var, - http_headers: None, - env_http_headers: None, - }, - oauth_client_id, - oauth_resource, - ), + } => { + let http_headers_helper = http_headers_helper + .map(|command| { + if command.trim().is_empty() { + bail!("--http-headers-helper command must not be empty"); + } + let mut helper = + McpServerHttpHeadersHelperConfig::new(command, http_headers_helper_args); + if let Some(cwd) = http_headers_helper_cwd { + helper.cwd = AbsolutePathBuf::relative_to_current_dir(cwd) + .context("failed to resolve --http-headers-helper-cwd")?; + } + Ok(helper) + }) + .transpose()?; + ( + McpServerTransportConfig::StreamableHttp { + url, + bearer_token_env_var, + http_headers: None, + env_http_headers: None, + http_headers_helper, + }, + oauth_client_id, + oauth_resource, + ) + } AddMcpTransportArgs { .. } => bail!("exactly one of --command or --url must be provided"), }; @@ -372,8 +419,22 @@ async fn run_add(config_overrides: &CliConfigOverrides, add_args: AddArgs) -> Re println!("Added global MCP server '{name}'."); + if matches!( + &transport, + McpServerTransportConfig::StreamableHttp { + http_headers_helper: Some(_), + .. + } + ) { + println!( + "Stored HTTP headers helper. Skipping automatic OAuth discovery so the helper is not run during add. Run `codex mcp login {name}` to run the helper and complete OAuth if required." + ); + return Ok(()); + } + match oauth_login_support(&transport).await { McpOAuthLoginSupport::Supported(oauth_config) => { + let oauth_config = *oauth_config; println!("Detected OAuth support. Starting OAuth flow…"); let resolved_scopes = resolve_oauth_scopes( /*explicit_scopes*/ None, @@ -386,6 +447,7 @@ async fn run_add(config_overrides: &CliConfigOverrides, add_args: AddArgs) -> Re config.mcp_oauth_credentials_store_mode, oauth_config.http_headers, oauth_config.env_http_headers, + oauth_config.http_headers_helper, &resolved_scopes, oauth_client_id.as_deref(), oauth_resource.as_deref(), @@ -455,13 +517,19 @@ async fn run_login(config_overrides: &CliConfigOverrides, login_args: LoginArgs) bail!("No MCP server named '{name}' found."); }; - let (url, http_headers, env_http_headers) = match &server.transport { + let (url, http_headers, env_http_headers, http_headers_helper) = match &server.transport { McpServerTransportConfig::StreamableHttp { url, http_headers, env_http_headers, + http_headers_helper, .. - } => (url.clone(), http_headers.clone(), env_http_headers.clone()), + } => ( + url.clone(), + http_headers.clone(), + env_http_headers.clone(), + http_headers_helper.clone(), + ), _ => bail!("OAuth login is only supported for streamable HTTP servers."), }; @@ -480,6 +548,7 @@ async fn run_login(config_overrides: &CliConfigOverrides, login_args: LoginArgs) config.mcp_oauth_credentials_store_mode, http_headers, env_http_headers, + http_headers_helper, &resolved_scopes, server.oauth_client_id(), server.oauth_resource.as_deref(), @@ -573,6 +642,7 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> bearer_token_env_var, http_headers, env_http_headers, + http_headers_helper, } => { serde_json::json!({ "type": "streamable_http", @@ -580,6 +650,7 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> "bearer_token_env_var": bearer_token_env_var, "http_headers": http_headers, "env_http_headers": env_http_headers, + "http_headers_helper": http_headers_helper, }) } }; @@ -610,7 +681,7 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> } let mut stdio_rows: Vec<[String; 7]> = Vec::new(); - let mut http_rows: Vec<[String; 5]> = Vec::new(); + let mut http_rows: Vec<[String; 6]> = Vec::new(); for (name, cfg) in entries { match &cfg.transport { @@ -651,6 +722,7 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> McpServerTransportConfig::StreamableHttp { url, bearer_token_env_var, + http_headers_helper, .. } => { let status = format_mcp_status(cfg); @@ -661,10 +733,15 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> .to_string(); let bearer_token_display = bearer_token_env_var.as_deref().unwrap_or("-").to_string(); + let headers_helper_display = http_headers_helper + .as_ref() + .map(|helper| helper.command.clone()) + .unwrap_or_else(|| "-".to_string()); http_rows.push([ name.clone(), url.clone(), bearer_token_display, + headers_helper_display, status, auth_status, ]); @@ -736,6 +813,7 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> "Name".len(), "Url".len(), "Bearer Token Env Var".len(), + "Headers Helper".len(), "Status".len(), "Auth".len(), ]; @@ -746,32 +824,36 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> } println!( - "{name: Re bearer_token_env_var, http_headers, env_http_headers, + http_headers_helper, } => serde_json::json!({ "type": "streamable_http", "url": url, "bearer_token_env_var": bearer_token_env_var, "http_headers": http_headers, "env_http_headers": env_http_headers, + "http_headers_helper": http_headers_helper, }), }; let output = serde_json::to_string_pretty(&serde_json::json!({ @@ -898,6 +982,7 @@ async fn run_get(config_overrides: &CliConfigOverrides, get_args: GetArgs) -> Re bearer_token_env_var, http_headers, env_http_headers, + http_headers_helper, } => { println!(" transport: streamable_http"); println!(" url: {url}"); @@ -929,6 +1014,12 @@ async fn run_get(config_overrides: &CliConfigOverrides, get_args: GetArgs) -> Re _ => "-".to_string(), }; println!(" env_http_headers: {env_headers_display}"); + let helper_display = match http_headers_helper { + Some(helper) if helper.args.is_empty() => helper.command.clone(), + Some(helper) => format!("{} {}", helper.command, helper.args.join(" ")), + None => "-".to_string(), + }; + println!(" http_headers_helper: {helper_display}"); } } if let Some(timeout) = server.startup_timeout_sec { diff --git a/codex-rs/cli/tests/mcp_add_remove.rs b/codex-rs/cli/tests/mcp_add_remove.rs index d0fc5f327db..5ba7a52fb19 100644 --- a/codex-rs/cli/tests/mcp_add_remove.rs +++ b/codex-rs/cli/tests/mcp_add_remove.rs @@ -144,11 +144,13 @@ async fn add_streamable_http_without_manual_token() -> Result<()> { bearer_token_env_var, http_headers, env_http_headers, + http_headers_helper, } => { assert_eq!(url, "https://example.com/mcp"); assert!(bearer_token_env_var.is_none()); assert!(http_headers.is_none()); assert!(env_http_headers.is_none()); + assert!(http_headers_helper.is_none()); } other => panic!("unexpected transport: {other:?}"), } @@ -160,6 +162,82 @@ async fn add_streamable_http_without_manual_token() -> Result<()> { Ok(()) } +#[tokio::test] +async fn add_streamable_http_with_headers_helper() -> Result<()> { + let codex_home = TempDir::new()?; + + let mut add_cmd = codex_command(codex_home.path())?; + add_cmd + .args([ + "mcp", + "add", + "litellm", + "--url", + "https://llm.atko.ai/mcp", + "--http-headers-helper", + "ocm", + "--http-headers-helper-arg", + "auth", + "--http-headers-helper-arg", + "litellm", + "--http-headers-helper-arg", + "--site", + "--http-headers-helper-arg", + "llm.atko.ai", + ]) + .assert() + .success() + .stdout(contains( + "Skipping automatic OAuth discovery so the helper is not run during add.", + )); + + let servers = load_global_mcp_servers(codex_home.path()).await?; + let litellm = servers.get("litellm").expect("litellm server should exist"); + match &litellm.transport { + McpServerTransportConfig::StreamableHttp { + url, + bearer_token_env_var, + http_headers, + env_http_headers, + http_headers_helper, + } => { + assert_eq!(url, "https://llm.atko.ai/mcp"); + assert!(bearer_token_env_var.is_none()); + assert!(http_headers.is_none()); + assert!(env_http_headers.is_none()); + assert_eq!( + http_headers_helper + .as_ref() + .map(|helper| (helper.command.as_str(), helper.args.as_slice())), + Some(( + "ocm", + &[ + "auth".to_string(), + "litellm".to_string(), + "--site".to_string(), + "llm.atko.ai".to_string() + ][..], + )) + ); + } + other => panic!("unexpected transport: {other:?}"), + } + + let serialized = std::fs::read_to_string(codex_home.path().join("config.toml"))?; + assert_eq!( + serialized, + r#"[mcp_servers.litellm] +url = "https://llm.atko.ai/mcp" + +[mcp_servers.litellm.http_headers_helper] +command = "ocm" +args = ["auth", "litellm", "--site", "llm.atko.ai"] +"# + ); + + Ok(()) +} + #[tokio::test] async fn add_streamable_http_with_custom_env_var() -> Result<()> { let codex_home = TempDir::new()?; @@ -186,11 +264,13 @@ async fn add_streamable_http_with_custom_env_var() -> Result<()> { bearer_token_env_var, http_headers, env_http_headers, + http_headers_helper, } => { assert_eq!(url, "https://example.com/issues"); assert_eq!(bearer_token_env_var.as_deref(), Some("GITHUB_TOKEN")); assert!(http_headers.is_none()); assert!(env_http_headers.is_none()); + assert!(http_headers_helper.is_none()); } other => panic!("unexpected transport: {other:?}"), } diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 60a0026eeb2..6fa5a902aa4 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1148,6 +1148,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, }, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, @@ -1247,6 +1248,7 @@ fn mcp_init_error_display_prompts_for_github_pat() { bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, }, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, @@ -1300,6 +1302,7 @@ fn mcp_init_error_display_reports_generic_errors() { bearer_token_env_var: Some("TOKEN".to_string()), http_headers: None, env_http_headers: None, + http_headers_helper: None, }, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, diff --git a/codex-rs/codex-mcp/src/mcp/auth.rs b/codex-rs/codex-mcp/src/mcp/auth.rs index 12f832f9e99..4e4ded38d71 100644 --- a/codex-rs/codex-mcp/src/mcp/auth.rs +++ b/codex-rs/codex-mcp/src/mcp/auth.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use anyhow::Result; use codex_config::McpServerConfig; +use codex_config::McpServerHttpHeadersHelperConfig; use codex_config::McpServerTransportConfig; use codex_config::types::OAuthCredentialsStoreMode; use codex_login::CodexAuth; @@ -21,12 +22,13 @@ pub struct McpOAuthLoginConfig { pub url: String, pub http_headers: Option>, pub env_http_headers: Option>, + pub http_headers_helper: Option, pub discovered_scopes: Option>, } #[derive(Debug)] pub enum McpOAuthLoginSupport { - Supported(McpOAuthLoginConfig), + Supported(Box), Unsupported, Unknown(anyhow::Error), } @@ -57,6 +59,7 @@ pub async fn oauth_login_support(transport: &McpServerTransportConfig) -> McpOAu bearer_token_env_var, http_headers, env_http_headers, + http_headers_helper, } = transport else { return McpOAuthLoginSupport::Unsupported; @@ -66,14 +69,21 @@ 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( + url, + http_headers.clone(), + env_http_headers.clone(), + http_headers_helper.clone(), + ) + .await { - Ok(Some(discovery)) => McpOAuthLoginSupport::Supported(McpOAuthLoginConfig { + Ok(Some(discovery)) => McpOAuthLoginSupport::Supported(Box::new(McpOAuthLoginConfig { url: url.clone(), http_headers: http_headers.clone(), env_http_headers: env_http_headers.clone(), + http_headers_helper: http_headers_helper.clone(), discovered_scopes: discovery.scopes_supported, - }), + })), Ok(None) => McpOAuthLoginSupport::Unsupported, Err(err) => McpOAuthLoginSupport::Unknown(err), } @@ -196,6 +206,7 @@ async fn compute_auth_status( bearer_token_env_var, http_headers, env_http_headers, + http_headers_helper, } => { determine_streamable_http_auth_status( server_name, @@ -203,6 +214,7 @@ async fn compute_auth_status( bearer_token_env_var.as_deref(), http_headers.clone(), env_http_headers.clone(), + http_headers_helper.clone(), store_mode, ) .await diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 51a6f186860..75666fc8172 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -457,6 +457,7 @@ fn codex_apps_mcp_server_config(config: &McpConfig) -> McpServerConfig { bearer_token_env_var: codex_apps_mcp_bearer_token_env_var(), http_headers, env_http_headers: None, + http_headers_helper: None, }, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, diff --git a/codex-rs/codex-mcp/src/mcp/mod_tests.rs b/codex-rs/codex-mcp/src/mcp/mod_tests.rs index d4d888f2022..c8bab198e43 100644 --- a/codex-rs/codex-mcp/src/mcp/mod_tests.rs +++ b/codex-rs/codex-mcp/src/mcp/mod_tests.rs @@ -316,6 +316,7 @@ async fn effective_mcp_servers_preserve_user_servers_and_add_codex_apps() { bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, }, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, @@ -341,6 +342,7 @@ async fn effective_mcp_servers_preserve_user_servers_and_add_codex_apps() { bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, }, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 78078720e7e..eca4aabf9da 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -630,6 +630,7 @@ async fn make_rmcp_client( url, http_headers, env_http_headers, + http_headers_helper, bearer_token_env_var, } => { let http_client = resolved_environment.as_ref().map_or_else( @@ -647,6 +648,7 @@ async fn make_rmcp_client( resolved_bearer_token, http_headers, env_http_headers, + http_headers_helper, store_mode, http_client, runtime_auth_provider, diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 5404bf1eb53..479b82f643c 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -161,6 +161,7 @@ mod tests { bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, }, environment_id: environment_id.to_string(), ..stdio_server(environment_id) diff --git a/codex-rs/config/src/lib.rs b/codex-rs/config/src/lib.rs index 5ae2688036a..78c1e6436e6 100644 --- a/codex-rs/config/src/lib.rs +++ b/codex-rs/config/src/lib.rs @@ -94,10 +94,12 @@ pub use marketplace_edit::remove_user_marketplace_config; pub use mcp_edit::ConfigEditsBuilder; pub use mcp_edit::load_global_mcp_servers; pub use mcp_types::AppToolApproval; +pub use mcp_types::DEFAULT_MCP_HTTP_HEADERS_HELPER_TIMEOUT_MS; pub use mcp_types::DEFAULT_MCP_SERVER_ENVIRONMENT_ID; pub use mcp_types::McpServerConfig; pub use mcp_types::McpServerDisabledReason; pub use mcp_types::McpServerEnvVar; +pub use mcp_types::McpServerHttpHeadersHelperConfig; pub use mcp_types::McpServerOAuthConfig; pub use mcp_types::McpServerToolConfig; pub use mcp_types::McpServerTransportConfig; diff --git a/codex-rs/config/src/mcp_edit.rs b/codex-rs/config/src/mcp_edit.rs index 9f007de1130..2b04763e066 100644 --- a/codex-rs/config/src/mcp_edit.rs +++ b/codex-rs/config/src/mcp_edit.rs @@ -13,8 +13,10 @@ use toml_edit::value; use crate::AppToolApproval; use crate::CONFIG_TOML_FILE; +use crate::DEFAULT_MCP_HTTP_HEADERS_HELPER_TIMEOUT_MS; use crate::McpServerConfig; use crate::McpServerEnvVar; +use crate::McpServerHttpHeadersHelperConfig; use crate::McpServerTransportConfig; pub async fn load_global_mcp_servers( @@ -154,6 +156,7 @@ fn serialize_mcp_server(config: &McpServerConfig) -> TomlItem { bearer_token_env_var, http_headers, env_http_headers, + http_headers_helper, } => { entry["url"] = value(url.clone()); if let Some(env_var) = bearer_token_env_var { @@ -169,6 +172,9 @@ fn serialize_mcp_server(config: &McpServerConfig) -> TomlItem { { entry["env_http_headers"] = table_from_pairs(headers.iter()); } + if let Some(helper) = http_headers_helper { + entry["http_headers_helper"] = http_headers_helper_table(helper); + } } } @@ -257,6 +263,24 @@ fn array_from_strings(values: &[String]) -> TomlItem { TomlItem::Value(array.into()) } +fn http_headers_helper_table(helper: &McpServerHttpHeadersHelperConfig) -> TomlItem { + let mut table = TomlTable::new(); + table.set_implicit(false); + table["command"] = value(helper.command.clone()); + if !helper.args.is_empty() { + table["args"] = array_from_strings(&helper.args); + } + if helper.timeout_ms.get() != DEFAULT_MCP_HTTP_HEADERS_HELPER_TIMEOUT_MS + && let Ok(timeout_ms) = i64::try_from(helper.timeout_ms.get()) + { + table["timeout_ms"] = value(timeout_ms); + } + if !helper.is_default_cwd() { + table["cwd"] = value(helper.cwd.to_string_lossy().to_string()); + } + TomlItem::Table(table) +} + fn array_from_env_vars(env_vars: &[McpServerEnvVar]) -> TomlItem { let mut array = toml_edit::Array::new(); for env_var in env_vars { diff --git a/codex-rs/config/src/mcp_edit_tests.rs b/codex-rs/config/src/mcp_edit_tests.rs index 030b0a02335..4ad645cd81f 100644 --- a/codex-rs/config/src/mcp_edit_tests.rs +++ b/codex-rs/config/src/mcp_edit_tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::McpServerHttpHeadersHelperConfig; use crate::McpServerOAuthConfig; use crate::McpServerToolConfig; use pretty_assertions::assert_eq; @@ -100,6 +101,7 @@ async fn replace_mcp_servers_serializes_oauth_client_id() -> anyhow::Result<()> bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, }, environment_id: crate::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, @@ -144,3 +146,66 @@ client_id = "eci-prd-pub-codex-123" Ok(()) } + +#[tokio::test] +async fn replace_mcp_servers_serializes_http_headers_helper() -> anyhow::Result<()> { + let unique_suffix = SystemTime::now().duration_since(UNIX_EPOCH)?.as_nanos(); + let codex_home = std::env::temp_dir().join(format!( + "codex-config-mcp-headers-helper-edit-test-{}-{unique_suffix}", + std::process::id() + )); + let servers = BTreeMap::from([( + "auth_mcp".to_string(), + McpServerConfig { + transport: McpServerTransportConfig::StreamableHttp { + url: "https://example.com/mcp".to_string(), + bearer_token_env_var: None, + http_headers: None, + env_http_headers: None, + http_headers_helper: Some(McpServerHttpHeadersHelperConfig::new( + "ocm".to_string(), + vec!["auth".to_string(), "mcp".to_string()], + )), + }, + environment_id: crate::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + enabled: true, + required: false, + supports_parallel_tool_calls: false, + disabled_reason: None, + startup_timeout_sec: None, + 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(), + }, + )]); + + ConfigEditsBuilder::new(&codex_home) + .replace_mcp_servers(&servers) + .apply() + .await?; + + let config_path = codex_home.join(CONFIG_TOML_FILE); + let serialized = std::fs::read_to_string(&config_path)?; + assert_eq!( + serialized, + r#"[mcp_servers.auth_mcp] +url = "https://example.com/mcp" + +[mcp_servers.auth_mcp.http_headers_helper] +command = "ocm" +args = ["auth", "mcp"] +"# + ); + + let loaded = load_global_mcp_servers(&codex_home).await?; + assert_eq!(loaded, servers); + + std::fs::remove_dir_all(&codex_home)?; + + Ok(()) +} diff --git a/codex-rs/config/src/mcp_types.rs b/codex-rs/config/src/mcp_types.rs index 0d30ba8b23f..3d8e4563ee0 100644 --- a/codex-rs/config/src/mcp_types.rs +++ b/codex-rs/config/src/mcp_types.rs @@ -2,9 +2,11 @@ use std::collections::HashMap; use std::fmt; +use std::num::NonZeroU64; use std::path::PathBuf; use std::time::Duration; +use crate::AbsolutePathBuf; use schemars::JsonSchema; use serde::Deserialize; use serde::Deserializer; @@ -15,6 +17,7 @@ use crate::RequirementSource; /// Effective MCP environment id when config omits `environment_id`. pub const DEFAULT_MCP_SERVER_ENVIRONMENT_ID: &str = "local"; +pub const DEFAULT_MCP_HTTP_HEADERS_HELPER_TIMEOUT_MS: u64 = 10_000; #[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, Default, JsonSchema)] #[serde(rename_all = "snake_case")] @@ -126,6 +129,49 @@ pub struct McpServerOAuthConfig { pub client_id: Option, } +/// Command that prints a JSON object of HTTP headers for a streamable HTTP MCP server. +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema)] +#[schemars(deny_unknown_fields)] +pub struct McpServerHttpHeadersHelperConfig { + /// Command to execute. Bare names are resolved via `PATH`; paths are resolved against `cwd`. + pub command: String, + + /// Command arguments. + #[serde(default)] + pub args: Vec, + + /// Maximum time to wait for the helper command to exit successfully. + #[serde(default = "default_mcp_http_headers_helper_timeout_ms")] + pub timeout_ms: NonZeroU64, + + /// Working directory used when running the helper command. + #[serde( + default = "default_mcp_http_headers_helper_cwd", + skip_serializing_if = "is_default_mcp_http_headers_helper_cwd" + )] + #[schemars(skip_serializing_if = "is_default_mcp_http_headers_helper_cwd")] + pub cwd: AbsolutePathBuf, +} + +impl McpServerHttpHeadersHelperConfig { + pub fn new(command: String, args: Vec) -> Self { + Self { + command, + args, + timeout_ms: default_mcp_http_headers_helper_timeout_ms(), + cwd: default_mcp_http_headers_helper_cwd(), + } + } + + pub fn timeout(&self) -> Duration { + Duration::from_millis(self.timeout_ms.get()) + } + + pub fn is_default_cwd(&self) -> bool { + is_default_mcp_http_headers_helper_cwd(&self.cwd) + } +} + #[derive(Serialize, Debug, Clone, PartialEq)] pub struct McpServerConfig { #[serde(flatten)] @@ -228,6 +274,8 @@ pub struct RawMcpServerConfig { pub http_headers: Option>, #[serde(default)] pub env_http_headers: Option>, + #[serde(default)] + pub http_headers_helper: Option, // streamable_http pub url: Option, @@ -282,6 +330,7 @@ impl TryFrom for McpServerConfig { cwd, http_headers, env_http_headers, + http_headers_helper, url, bearer_token, bearer_token_env_var, @@ -317,6 +366,9 @@ impl TryFrom for McpServerConfig { Err(format!("{field} is not supported for {transport}")) } + let environment_id = + environment_id.unwrap_or_else(|| DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string()); + let transport = if let Some(command) = command { throw_if_set("stdio", "url", url.as_ref())?; throw_if_set( @@ -327,6 +379,7 @@ impl TryFrom for McpServerConfig { throw_if_set("stdio", "bearer_token", bearer_token.as_ref())?; throw_if_set("stdio", "http_headers", http_headers.as_ref())?; throw_if_set("stdio", "env_http_headers", env_http_headers.as_ref())?; + throw_if_set("stdio", "http_headers_helper", http_headers_helper.as_ref())?; throw_if_set("stdio", "oauth", oauth.as_ref())?; throw_if_set("stdio", "oauth_resource", oauth_resource.as_ref())?; let env_vars = env_vars.unwrap_or_default(); @@ -346,18 +399,29 @@ impl TryFrom for McpServerConfig { throw_if_set("streamable_http", "env_vars", env_vars.as_ref())?; throw_if_set("streamable_http", "cwd", cwd.as_ref())?; throw_if_set("streamable_http", "bearer_token", bearer_token.as_ref())?; + if let Some(helper) = http_headers_helper.as_ref() + && helper.command.trim().is_empty() + { + return Err("http_headers_helper.command must not be empty".to_string()); + } + if http_headers_helper.is_some() && environment_id != DEFAULT_MCP_SERVER_ENVIRONMENT_ID + { + return Err( + "http_headers_helper is only supported for local streamable_http MCP servers" + .to_string(), + ); + } McpServerTransportConfig::StreamableHttp { url, bearer_token_env_var, http_headers, env_http_headers, + http_headers_helper, } } else { return Err("invalid transport".to_string()); }; - let environment_id = - environment_id.unwrap_or_else(|| DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string()); validate_remote_stdio_cwd(&transport, &environment_id)?; Ok(Self { @@ -448,9 +512,35 @@ pub enum McpServerTransportConfig { /// HTTP headers where the value is sourced from an environment variable. #[serde(default, skip_serializing_if = "Option::is_none")] env_http_headers: Option>, + /// Command that prints a JSON object of HTTP headers to include in requests. + #[serde(default, skip_serializing_if = "Option::is_none")] + http_headers_helper: Option, }, } +fn default_mcp_http_headers_helper_timeout_ms() -> NonZeroU64 { + match NonZeroU64::new(DEFAULT_MCP_HTTP_HEADERS_HELPER_TIMEOUT_MS) { + Some(value) => value, + None => panic!("DEFAULT_MCP_HTTP_HEADERS_HELPER_TIMEOUT_MS must be non-zero"), + } +} + +fn default_mcp_http_headers_helper_cwd() -> AbsolutePathBuf { + let deserializer = serde::de::value::StrDeserializer::::new("."); + if let Ok(cwd) = AbsolutePathBuf::deserialize(deserializer) { + return cwd; + } + + match AbsolutePathBuf::current_dir() { + Ok(cwd) => cwd, + Err(err) => panic!("http_headers_helper cwd must resolve: {err}"), + } +} + +fn is_default_mcp_http_headers_helper_cwd(path: &AbsolutePathBuf) -> bool { + path == &default_mcp_http_headers_helper_cwd() +} + mod option_duration_secs { use serde::Deserialize; use serde::Deserializer; diff --git a/codex-rs/config/src/mcp_types_tests.rs b/codex-rs/config/src/mcp_types_tests.rs index 5f3933ce43e..0f297090e37 100644 --- a/codex-rs/config/src/mcp_types_tests.rs +++ b/codex-rs/config/src/mcp_types_tests.rs @@ -271,6 +271,7 @@ fn deserialize_streamable_http_server_config() { bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, } ); assert!(cfg.enabled); @@ -293,6 +294,7 @@ fn deserialize_streamable_http_server_config_with_env_var() { bearer_token_env_var: Some("GITHUB_TOKEN".to_string()), http_headers: None, env_http_headers: None, + http_headers_helper: None, } ); assert!(cfg.enabled); @@ -319,10 +321,60 @@ fn deserialize_streamable_http_server_config_with_headers() { "X-Token".to_string(), "TOKEN_ENV".to_string() )])), + http_headers_helper: None, } ); } +#[test] +fn deserialize_streamable_http_server_config_with_headers_helper() { + let cfg: McpServerConfig = toml::from_str( + r#" + url = "https://example.com/mcp" + + [http_headers_helper] + command = "ocm" + args = ["auth", "mcp"] + timeout_ms = 5000 + "#, + ) + .expect("should deserialize http config with headers helper"); + + assert_eq!( + cfg.transport, + McpServerTransportConfig::StreamableHttp { + url: "https://example.com/mcp".to_string(), + bearer_token_env_var: None, + http_headers: None, + env_http_headers: None, + http_headers_helper: Some(McpServerHttpHeadersHelperConfig { + command: "ocm".to_string(), + args: vec!["auth".to_string(), "mcp".to_string()], + timeout_ms: NonZeroU64::new(5_000).expect("timeout should be non-zero"), + cwd: default_mcp_http_headers_helper_cwd(), + }), + } + ); +} + +#[test] +fn deserialize_rejects_headers_helper_for_remote_streamable_http_server() { + let err = toml::from_str::( + r#" + url = "https://example.com/mcp" + environment_id = "remote" + http_headers_helper = { command = "helper" } + "#, + ) + .expect_err("should reject headers helper for remote http transport"); + + assert!( + err.to_string() + .contains("http_headers_helper is only supported for local streamable_http"), + "unexpected error: {err}" + ); +} + #[test] fn deserialize_streamable_http_server_config_with_oauth_resource() { let cfg: McpServerConfig = toml::from_str( @@ -516,6 +568,20 @@ fn deserialize_rejects_headers_for_stdio() { ) .expect_err("should reject env_http_headers for stdio transport"); + let err = toml::from_str::( + r#" + command = "echo" + http_headers_helper = { command = "helper" } + "#, + ) + .expect_err("should reject http_headers_helper for stdio transport"); + + assert!( + err.to_string() + .contains("http_headers_helper is not supported for stdio"), + "unexpected error: {err}" + ); + let err = toml::from_str::( r#" command = "echo" @@ -544,6 +610,23 @@ fn deserialize_rejects_headers_for_stdio() { ); } +#[test] +fn deserialize_rejects_empty_headers_helper_command() { + let err = toml::from_str::( + r#" + url = "https://example.com" + http_headers_helper = { command = " " } + "#, + ) + .expect_err("should reject empty http_headers_helper command"); + + assert!( + err.to_string() + .contains("http_headers_helper.command must not be empty"), + "unexpected error: {err}" + ); +} + #[test] fn deserialize_rejects_inline_bearer_token_field() { let err = toml::from_str::( diff --git a/codex-rs/config/src/types.rs b/codex-rs/config/src/types.rs index 30d297030ac..51aed8914dd 100644 --- a/codex-rs/config/src/types.rs +++ b/codex-rs/config/src/types.rs @@ -4,9 +4,11 @@ // definitions that do not contain business logic. pub use crate::mcp_types::AppToolApproval; +pub use crate::mcp_types::DEFAULT_MCP_HTTP_HEADERS_HELPER_TIMEOUT_MS; pub use crate::mcp_types::McpServerConfig; pub use crate::mcp_types::McpServerDisabledReason; pub use crate::mcp_types::McpServerEnvVar; +pub use crate::mcp_types::McpServerHttpHeadersHelperConfig; pub use crate::mcp_types::McpServerOAuthConfig; pub use crate::mcp_types::McpServerToolConfig; pub use crate::mcp_types::McpServerTransportConfig; diff --git a/codex-rs/core-plugins/src/manager_tests.rs b/codex-rs/core-plugins/src/manager_tests.rs index 0131ef22ed1..5d05c45df29 100644 --- a/codex-rs/core-plugins/src/manager_tests.rs +++ b/codex-rs/core-plugins/src/manager_tests.rs @@ -218,6 +218,7 @@ async fn load_plugins_loads_default_skills_and_mcp_servers() { bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, }, environment_id: "local".to_string(), enabled: true, @@ -721,6 +722,7 @@ async fn load_plugins_uses_manifest_configured_component_paths() { bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, }, environment_id: "local".to_string(), enabled: true, @@ -833,6 +835,7 @@ async fn load_plugins_ignores_manifest_component_paths_without_dot_slash() { bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, }, environment_id: "local".to_string(), enabled: true, @@ -995,6 +998,7 @@ fn capability_index_filters_inactive_and_zero_capability_plugins() { bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, }, environment_id: "local".to_string(), enabled: true, diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index b7bdff813bb..6e318fc613a 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -1255,6 +1255,43 @@ } ] }, + "McpServerHttpHeadersHelperConfig": { + "additionalProperties": false, + "description": "Command that prints a JSON object of HTTP headers for a streamable HTTP MCP server.", + "properties": { + "args": { + "default": [], + "description": "Command arguments.", + "items": { + "type": "string" + }, + "type": "array" + }, + "command": { + "description": "Command to execute. Bare names are resolved via `PATH`; paths are resolved against `cwd`.", + "type": "string" + }, + "cwd": { + "allOf": [ + { + "$ref": "#/definitions/AbsolutePathBuf" + } + ], + "description": "Working directory used when running the helper command." + }, + "timeout_ms": { + "default": 10000, + "description": "Maximum time to wait for the helper command to exit successfully.", + "format": "uint64", + "minimum": 1.0, + "type": "integer" + } + }, + "required": [ + "command" + ], + "type": "object" + }, "McpServerOAuthConfig": { "additionalProperties": false, "description": "OAuth client settings used when Codex launches an MCP OAuth flow.", @@ -2319,6 +2356,14 @@ }, "type": "object" }, + "http_headers_helper": { + "allOf": [ + { + "$ref": "#/definitions/McpServerHttpHeadersHelperConfig" + } + ], + "default": null + }, "name": { "default": null, "description": "Legacy display-name field accepted for backward compatibility.", diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 95010356251..a40222b6368 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -141,6 +141,7 @@ fn http_mcp(url: &str) -> McpServerConfig { bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, }, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, @@ -5693,6 +5694,7 @@ async fn replace_mcp_servers_streamable_http_serializes_bearer_token() -> anyhow bearer_token_env_var: Some("MCP_TOKEN".to_string()), http_headers: None, env_http_headers: None, + http_headers_helper: None, }, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, @@ -5735,11 +5737,13 @@ startup_timeout_sec = 2.0 bearer_token_env_var, http_headers, env_http_headers, + http_headers_helper, } => { assert_eq!(url, "https://example.com/mcp"); assert_eq!(bearer_token_env_var.as_deref(), Some("MCP_TOKEN")); assert!(http_headers.is_none()); assert!(env_http_headers.is_none()); + assert!(http_headers_helper.is_none()); } other => panic!("unexpected transport {other:?}"), } @@ -5763,6 +5767,7 @@ async fn replace_mcp_servers_streamable_http_serializes_custom_headers() -> anyh "X-Auth".to_string(), "DOCS_AUTH".to_string(), )])), + http_headers_helper: None, }, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, @@ -5845,6 +5850,7 @@ async fn replace_mcp_servers_streamable_http_removes_optional_sections() -> anyh "X-Auth".to_string(), "DOCS_AUTH".to_string(), )])), + http_headers_helper: None, }, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, @@ -5880,6 +5886,7 @@ async fn replace_mcp_servers_streamable_http_removes_optional_sections() -> anyh bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, }, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, @@ -5918,11 +5925,13 @@ url = "https://example.com/mcp" bearer_token_env_var, http_headers, env_http_headers, + http_headers_helper, } => { assert_eq!(url, "https://example.com/mcp"); assert!(bearer_token_env_var.is_none()); assert!(http_headers.is_none()); assert!(env_http_headers.is_none()); + assert!(http_headers_helper.is_none()); } other => panic!("unexpected transport {other:?}"), } @@ -5950,6 +5959,7 @@ async fn replace_mcp_servers_streamable_http_isolates_headers_between_servers() "X-Auth".to_string(), "DOCS_AUTH".to_string(), )])), + http_headers_helper: None, }, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, @@ -6218,6 +6228,7 @@ async fn replace_mcp_servers_streamable_http_serializes_oauth_resource() -> anyh bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, }, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, diff --git a/codex-rs/core/src/config/edit.rs b/codex-rs/core/src/config/edit.rs index 290edf6fb81..a1073671149 100644 --- a/codex-rs/core/src/config/edit.rs +++ b/codex-rs/core/src/config/edit.rs @@ -201,8 +201,10 @@ pub fn model_availability_nux_count_edits(shown_count: &HashMap) -> // TODO(jif) move to a dedicated file mod document_helpers { use codex_config::types::AppToolApproval; + use codex_config::types::DEFAULT_MCP_HTTP_HEADERS_HELPER_TIMEOUT_MS; use codex_config::types::McpServerConfig; use codex_config::types::McpServerEnvVar; + use codex_config::types::McpServerHttpHeadersHelperConfig; use codex_config::types::McpServerToolConfig; use codex_config::types::McpServerTransportConfig; use codex_config::types::ToolSuggestDisabledTool; @@ -279,6 +281,7 @@ mod document_helpers { bearer_token_env_var, http_headers, env_http_headers, + http_headers_helper, } => { entry["url"] = value(url.clone()); if let Some(env_var) = bearer_token_env_var { @@ -294,6 +297,9 @@ mod document_helpers { { entry["env_http_headers"] = table_from_pairs(headers.iter()); } + if let Some(helper) = http_headers_helper { + entry["http_headers_helper"] = http_headers_helper_table(helper); + } } } @@ -496,6 +502,24 @@ mod document_helpers { TomlItem::Value(array.into()) } + fn http_headers_helper_table(helper: &McpServerHttpHeadersHelperConfig) -> TomlItem { + let mut table = TomlTable::new(); + table.set_implicit(false); + table["command"] = value(helper.command.clone()); + if !helper.args.is_empty() { + table["args"] = array_from_iter(helper.args.iter().cloned()); + } + if helper.timeout_ms.get() != DEFAULT_MCP_HTTP_HEADERS_HELPER_TIMEOUT_MS + && let Ok(timeout_ms) = i64::try_from(helper.timeout_ms.get()) + { + table["timeout_ms"] = value(timeout_ms); + } + if !helper.is_default_cwd() { + table["cwd"] = value(helper.cwd.to_string_lossy().to_string()); + } + TomlItem::Table(table) + } + fn table_from_pairs<'a, I>(pairs: I) -> TomlItem where I: IntoIterator, diff --git a/codex-rs/core/src/config/edit_tests.rs b/codex-rs/core/src/config/edit_tests.rs index dce192831b6..bc58af863ff 100644 --- a/codex-rs/core/src/config/edit_tests.rs +++ b/codex-rs/core/src/config/edit_tests.rs @@ -916,6 +916,7 @@ fn blocking_replace_mcp_servers_round_trips() { .collect(), ), env_http_headers: None, + http_headers_helper: None, }, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: false, diff --git a/codex-rs/core/src/mcp_skill_dependencies.rs b/codex-rs/core/src/mcp_skill_dependencies.rs index 936f1b5fb5c..d0fe48e0992 100644 --- a/codex-rs/core/src/mcp_skill_dependencies.rs +++ b/codex-rs/core/src/mcp_skill_dependencies.rs @@ -138,7 +138,7 @@ 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::Supported(config) => *config, McpOAuthLoginSupport::Unsupported => continue, McpOAuthLoginSupport::Unknown(err) => { warn!("MCP server may or may not require login for dependency {name}: {err}"); @@ -158,6 +158,7 @@ pub(crate) async fn maybe_install_mcp_dependencies( config.mcp_oauth_credentials_store_mode, oauth_config.http_headers.clone(), oauth_config.env_http_headers.clone(), + oauth_config.http_headers_helper.clone(), &resolved_scopes.scopes, oauth_client_id, server_config.oauth_resource.as_deref(), @@ -174,6 +175,7 @@ pub(crate) async fn maybe_install_mcp_dependencies( config.mcp_oauth_credentials_store_mode, oauth_config.http_headers, oauth_config.env_http_headers, + oauth_config.http_headers_helper, &[], oauth_client_id, server_config.oauth_resource.as_deref(), @@ -360,6 +362,7 @@ fn mcp_dependency_to_server_config( bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, }, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 2c86cb393ce..a90407ba56a 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -1912,6 +1912,7 @@ async fn streamable_http_tool_call_round_trip() -> anyhow::Result<()> { bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, }, TestMcpServerOptions { environment_id: remote_aware_environment_id(), @@ -2098,6 +2099,7 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { bearer_token_env_var: None, http_headers: None, env_http_headers: None, + http_headers_helper: None, }, TestMcpServerOptions { environment_id: remote_aware_environment_id(), diff --git a/codex-rs/external-agent-migration/src/lib.rs b/codex-rs/external-agent-migration/src/lib.rs index d7b8801796b..0679f7fdcbf 100644 --- a/codex-rs/external-agent-migration/src/lib.rs +++ b/codex-rs/external-agent-migration/src/lib.rs @@ -388,6 +388,9 @@ fn mcp_server_toml_table( if let Some(headers) = server_config.get("headers").and_then(JsonValue::as_object) { append_header_config(&mut table, headers)?; } + if let Some(headers_helper) = server_config.get("headersHelper") { + append_headers_helper_config(&mut table, headers_helper)?; + } } else { return None; } @@ -395,6 +398,65 @@ fn mcp_server_toml_table( Some(table) } +fn append_headers_helper_config( + table: &mut toml::map::Map, + headers_helper: &JsonValue, +) -> Option<()> { + let (command, args, timeout_ms) = match headers_helper { + JsonValue::String(command) => (command.clone(), Vec::new(), None), + JsonValue::Object(helper) => { + let command = helper.get("command").and_then(json_string)?; + let args = helper.get("args").map(json_string_vec).unwrap_or_default(); + let timeout_ms = helper + .get("timeout_ms") + .or_else(|| helper.get("timeoutMs")) + .and_then(json_u64); + (command, args, timeout_ms) + } + JsonValue::Null | JsonValue::Bool(_) | JsonValue::Number(_) | JsonValue::Array(_) => { + return None; + } + }; + + let command = command.trim(); + if command.is_empty() + || contains_env_placeholder(command) + || is_relative_path_command(command) + || args.iter().any(|arg| contains_env_placeholder(arg)) + { + return None; + } + + let mut helper_table = toml::map::Map::new(); + helper_table.insert( + "command".to_string(), + TomlValue::String(command.to_string()), + ); + if !args.is_empty() { + helper_table.insert( + "args".to_string(), + TomlValue::Array(args.into_iter().map(TomlValue::String).collect()), + ); + } + if let Some(timeout_ms) = timeout_ms { + let timeout_ms = i64::try_from(timeout_ms).ok()?; + if timeout_ms == 0 { + return None; + } + helper_table.insert("timeout_ms".to_string(), TomlValue::Integer(timeout_ms)); + } + + table.insert( + "http_headers_helper".to_string(), + TomlValue::Table(helper_table), + ); + Some(()) +} + +fn is_relative_path_command(command: &str) -> bool { + (command.contains('/') || command.contains('\\')) && !Path::new(command).is_absolute() +} + fn mcp_server_is_disabled( server_name: &str, server_config: &serde_json::Map, @@ -1486,6 +1548,111 @@ bearer_token_env_var = "VAULT_TOKEN" ); } + #[test] + fn mcp_migration_converts_headers_helper() { + let root = tempfile::TempDir::new().expect("tempdir"); + fs::write( + root.path().join(".mcp.json"), + r#"{ + "mcpServers": { + "litellm": { + "url": "https://llm.atko.ai/mcp", + "headersHelper": { + "command": "ocm", + "args": ["auth", "litellm", "--site", "llm.atko.ai"], + "timeoutMs": 5000 + } + } + } + }"#, + ) + .expect("write mcp"); + + assert_eq!( + build_mcp_config_from_external( + root.path(), + /*external_agent_home*/ None, + /*settings*/ None, + ) + .unwrap(), + toml::from_str( + r#" +[mcp_servers.litellm] +url = "https://llm.atko.ai/mcp" + +[mcp_servers.litellm.http_headers_helper] +command = "ocm" +args = ["auth", "litellm", "--site", "llm.atko.ai"] +timeout_ms = 5000 +"# + ) + .unwrap() + ); + } + + #[test] + fn mcp_migration_converts_string_headers_helper() { + let root = tempfile::TempDir::new().expect("tempdir"); + fs::write( + root.path().join(".mcp.json"), + r#"{ + "mcpServers": { + "authed": { + "url": "https://example.invalid/mcp", + "headersHelper": "/usr/local/bin/mcp-headers" + } + } + }"#, + ) + .expect("write mcp"); + + assert_eq!( + build_mcp_config_from_external( + root.path(), + /*external_agent_home*/ None, + /*settings*/ None, + ) + .unwrap(), + toml::from_str( + r#" +[mcp_servers.authed] +url = "https://example.invalid/mcp" + +[mcp_servers.authed.http_headers_helper] +command = "/usr/local/bin/mcp-headers" +"# + ) + .unwrap() + ); + } + + #[test] + fn mcp_migration_skips_relative_path_headers_helper() { + let root = tempfile::TempDir::new().expect("tempdir"); + fs::write( + root.path().join(".mcp.json"), + r#"{ + "mcpServers": { + "authed": { + "url": "https://example.invalid/mcp", + "headersHelper": "./scripts/mcp-headers" + } + } + }"#, + ) + .expect("write mcp"); + + assert_eq!( + build_mcp_config_from_external( + root.path(), + /*external_agent_home*/ None, + /*settings*/ None, + ) + .unwrap(), + toml::Value::Table(Default::default()) + ); + } + #[test] fn mcp_migration_reads_matching_project_entries_from_repo_external_project_config() { let root = tempfile::TempDir::new().expect("tempdir"); diff --git a/codex-rs/rmcp-client/src/auth_status.rs b/codex-rs/rmcp-client/src/auth_status.rs index 0b3b3bf6a76..3964073fd34 100644 --- a/codex-rs/rmcp-client/src/auth_status.rs +++ b/codex-rs/rmcp-client/src/auth_status.rs @@ -3,6 +3,7 @@ use std::time::Duration; use anyhow::Error; use anyhow::Result; +use codex_config::McpServerHttpHeadersHelperConfig; use codex_protocol::protocol::McpAuthStatus; use reqwest::Client; use reqwest::StatusCode; @@ -33,13 +34,19 @@ pub async fn determine_streamable_http_auth_status( bearer_token_env_var: Option<&str>, http_headers: Option>, env_http_headers: Option>, + _http_headers_helper: Option, store_mode: OAuthCredentialsStoreMode, ) -> 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, + env_http_headers, + /*http_headers_helper*/ None, + ) + .await?; if default_headers.contains_key(AUTHORIZATION) { return Ok(McpAuthStatus::BearerToken); } @@ -64,6 +71,7 @@ pub async fn determine_streamable_http_auth_status( pub async fn supports_oauth_login(url: &str) -> Result { Ok(discover_streamable_http_oauth( url, /*http_headers*/ None, /*env_http_headers*/ None, + /*http_headers_helper*/ None, ) .await? .is_some()) @@ -73,8 +81,10 @@ pub async fn discover_streamable_http_oauth( url: &str, http_headers: Option>, env_http_headers: Option>, + http_headers_helper: Option, ) -> Result> { - let default_headers = build_default_headers(http_headers, env_http_headers)?; + let default_headers = + build_default_headers(http_headers, env_http_headers, http_headers_helper).await?; discover_streamable_http_oauth_with_headers(url, &default_headers).await } @@ -196,11 +206,15 @@ mod tests { use super::*; use axum::Json; use axum::Router; + use axum::http::StatusCode; + use axum::response::IntoResponse; use axum::routing::get; use pretty_assertions::assert_eq; use serial_test::serial; use std::collections::HashMap; use std::ffi::OsString; + use std::path::Path; + use std::path::PathBuf; use tokio::task::JoinHandle; struct TestServer { @@ -239,6 +253,45 @@ mod tests { } } + async fn spawn_oauth_discovery_server_requiring_header( + metadata: serde_json::Value, + header_name: &'static str, + header_value: &'static str, + ) -> TestServer { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let address = listener.local_addr().expect("listener should have address"); + let app = Router::new().route( + "/.well-known/oauth-authorization-server/mcp", + get({ + let metadata = metadata.clone(); + move |headers: axum::http::HeaderMap| { + let metadata = metadata.clone(); + async move { + if headers + .get(header_name) + .and_then(|value| value.to_str().ok()) + == Some(header_value) + { + Json(metadata).into_response() + } else { + StatusCode::UNAUTHORIZED.into_response() + } + } + } + }), + ); + let handle = tokio::spawn(async move { + axum::serve(listener, app).await.expect("server should run"); + }); + + TestServer { + url: format!("http://{address}/mcp"), + handle, + } + } + struct EnvVarGuard { key: String, original: Option, @@ -271,6 +324,49 @@ mod tests { } } + fn write_headers_helper_script(dir: &Path, output: &str) -> PathBuf { + let script_path = dir.join(helper_script_name()); + std::fs::write(&script_path, helper_script_contents(output)) + .expect("headers helper script should be written"); + make_executable(&script_path); + script_path + } + + #[cfg(unix)] + fn helper_script_name() -> &'static str { + "headers-helper" + } + + #[cfg(windows)] + fn helper_script_name() -> &'static str { + "headers-helper.cmd" + } + + #[cfg(unix)] + fn helper_script_contents(output: &str) -> String { + format!("#!/bin/sh\nprintf '%s\\n' '{output}'\n") + } + + #[cfg(windows)] + fn helper_script_contents(output: &str) -> String { + format!("@echo off\r\necho {output}\r\n") + } + + #[cfg(unix)] + fn make_executable(path: &Path) { + use std::os::unix::fs::PermissionsExt; + + let mut permissions = std::fs::metadata(path) + .expect("headers helper script should have metadata") + .permissions(); + permissions.set_mode(0o755); + std::fs::set_permissions(path, permissions) + .expect("headers helper script should be executable"); + } + + #[cfg(windows)] + fn make_executable(_path: &Path) {} + #[tokio::test] async fn determine_auth_status_uses_bearer_token_when_authorization_header_present() { let status = determine_streamable_http_auth_status( @@ -282,6 +378,7 @@ mod tests { "Bearer token".to_string(), )])), /*env_http_headers*/ None, + /*http_headers_helper*/ None, OAuthCredentialsStoreMode::Keyring, ) .await @@ -303,6 +400,7 @@ mod tests { "Authorization".to_string(), "CODEX_RMCP_CLIENT_AUTH_STATUS_TEST_TOKEN".to_string(), )])), + /*http_headers_helper*/ None, OAuthCredentialsStoreMode::Keyring, ) .await @@ -311,6 +409,26 @@ mod tests { assert_eq!(status, McpAuthStatus::BearerToken); } + #[tokio::test] + async fn determine_auth_status_does_not_treat_headers_helper_as_bearer_token() { + let status = determine_streamable_http_auth_status( + "server", + "not-a-url", + /*bearer_token_env_var*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + Some(McpServerHttpHeadersHelperConfig::new( + "does-not-need-to-exist".to_string(), + Vec::new(), + )), + OAuthCredentialsStoreMode::Keyring, + ) + .await + .expect("status should compute without running helper"); + + assert_eq!(status, McpAuthStatus::Unsupported); + } + #[tokio::test] async fn discover_streamable_http_oauth_returns_normalized_scopes() { let server = spawn_oauth_discovery_server(serde_json::json!({ @@ -324,6 +442,7 @@ mod tests { &server.url, /*http_headers*/ None, /*env_http_headers*/ None, + /*http_headers_helper*/ None, ) .await .expect("discovery should succeed") @@ -335,6 +454,41 @@ mod tests { ); } + #[tokio::test] + async fn discover_streamable_http_oauth_uses_headers_helper() { + let server = spawn_oauth_discovery_server_requiring_header( + serde_json::json!({ + "authorization_endpoint": "https://example.com/authorize", + "token_endpoint": "https://example.com/token", + "scopes_supported": ["mcp.read"], + }), + "authorization", + "Bearer helper", + ) + .await; + let temp_dir = tempfile::TempDir::new().expect("tempdir"); + let script_path = + write_headers_helper_script(temp_dir.path(), r#"{"Authorization":"Bearer helper"}"#); + + let discovery = discover_streamable_http_oauth( + &server.url, + /*http_headers*/ None, + /*env_http_headers*/ None, + Some(McpServerHttpHeadersHelperConfig::new( + script_path.display().to_string(), + Vec::new(), + )), + ) + .await + .expect("discovery should succeed") + .expect("oauth support should be detected"); + + assert_eq!( + discovery.scopes_supported, + Some(vec!["mcp.read".to_string()]) + ); + } + #[tokio::test] async fn discover_streamable_http_oauth_ignores_empty_scopes() { let server = spawn_oauth_discovery_server(serde_json::json!({ @@ -348,6 +502,7 @@ mod tests { &server.url, /*http_headers*/ None, /*env_http_headers*/ None, + /*http_headers_helper*/ None, ) .await .expect("discovery should succeed") diff --git a/codex-rs/rmcp-client/src/perform_oauth_login.rs b/codex-rs/rmcp-client/src/perform_oauth_login.rs index a416b3d9b01..c562d033a1e 100644 --- a/codex-rs/rmcp-client/src/perform_oauth_login.rs +++ b/codex-rs/rmcp-client/src/perform_oauth_login.rs @@ -29,11 +29,13 @@ 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::McpServerHttpHeadersHelperConfig; use codex_config::types::OAuthCredentialsStoreMode; struct OauthHeaders { http_headers: Option>, env_http_headers: Option>, + http_headers_helper: Option, } struct CallbackServerGuard { @@ -83,6 +85,7 @@ pub async fn perform_oauth_login( store_mode: OAuthCredentialsStoreMode, http_headers: Option>, env_http_headers: Option>, + http_headers_helper: Option, scopes: &[String], oauth_client_id: Option<&str>, oauth_resource: Option<&str>, @@ -95,6 +98,7 @@ pub async fn perform_oauth_login( store_mode, http_headers, env_http_headers, + http_headers_helper, scopes, oauth_client_id, oauth_resource, @@ -112,6 +116,7 @@ pub async fn perform_oauth_login_silent( store_mode: OAuthCredentialsStoreMode, http_headers: Option>, env_http_headers: Option>, + http_headers_helper: Option, scopes: &[String], oauth_client_id: Option<&str>, oauth_resource: Option<&str>, @@ -124,6 +129,7 @@ pub async fn perform_oauth_login_silent( store_mode, http_headers, env_http_headers, + http_headers_helper, scopes, oauth_client_id, oauth_resource, @@ -141,6 +147,7 @@ async fn perform_oauth_login_with_browser_output( store_mode: OAuthCredentialsStoreMode, http_headers: Option>, env_http_headers: Option>, + http_headers_helper: Option, scopes: &[String], oauth_client_id: Option<&str>, oauth_resource: Option<&str>, @@ -151,6 +158,7 @@ async fn perform_oauth_login_with_browser_output( let headers = OauthHeaders { http_headers, env_http_headers, + http_headers_helper, }; OauthLoginFlow::new( server_name, @@ -177,6 +185,7 @@ pub async fn perform_oauth_login_return_url( store_mode: OAuthCredentialsStoreMode, http_headers: Option>, env_http_headers: Option>, + http_headers_helper: Option, scopes: &[String], oauth_client_id: Option<&str>, oauth_resource: Option<&str>, @@ -187,6 +196,7 @@ pub async fn perform_oauth_login_return_url( let headers = OauthHeaders { http_headers, env_http_headers, + http_headers_helper, }; let flow = OauthLoginFlow::new( server_name, @@ -479,8 +489,10 @@ impl OauthLoginFlow { let OauthHeaders { http_headers, env_http_headers, + http_headers_helper, } = headers; - let default_headers = build_default_headers(http_headers, env_http_headers)?; + let default_headers = + build_default_headers(http_headers, env_http_headers, http_headers_helper).await?; let http_client = apply_default_headers(ClientBuilder::new(), &default_headers).build()?; let scope_refs: Vec<&str> = scopes.iter().map(String::as_str).collect(); diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 90b09d724c3..d0384224379 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -13,6 +13,7 @@ use anyhow::Result; use anyhow::anyhow; use codex_api::SharedAuthProvider; use codex_client::maybe_build_rustls_client_config_with_custom_ca; +use codex_config::McpServerHttpHeadersHelperConfig; use codex_config::types::McpServerEnvVar; use codex_exec_server::HttpClient; use futures::FutureExt; @@ -116,6 +117,7 @@ enum TransportRecipe { bearer_token: Option, http_headers: Option>, env_http_headers: Option>, + http_headers_helper: Option, store_mode: OAuthCredentialsStoreMode, http_client: Arc, auth_provider: Option, @@ -344,6 +346,7 @@ impl RmcpClient { bearer_token: Option, http_headers: Option>, env_http_headers: Option>, + http_headers_helper: Option, store_mode: OAuthCredentialsStoreMode, http_client: Arc, auth_provider: Option, @@ -354,6 +357,7 @@ impl RmcpClient { bearer_token, http_headers, env_http_headers, + http_headers_helper, store_mode, http_client, auth_provider, @@ -726,12 +730,17 @@ impl RmcpClient { bearer_token, http_headers, env_http_headers, + http_headers_helper, store_mode, http_client, auth_provider, } => { - let default_headers = - build_default_headers(http_headers.clone(), env_http_headers.clone())?; + let default_headers = build_default_headers( + http_headers.clone(), + env_http_headers.clone(), + http_headers_helper.clone(), + ) + .await?; let initial_oauth_tokens = if bearer_token.is_none() && auth_provider.is_none() diff --git a/codex-rs/rmcp-client/src/utils.rs b/codex-rs/rmcp-client/src/utils.rs index 892df544cba..eede8e85f2d 100644 --- a/codex-rs/rmcp-client/src/utils.rs +++ b/codex-rs/rmcp-client/src/utils.rs @@ -1,5 +1,8 @@ +use crate::program_resolver; + use anyhow::Result; use anyhow::anyhow; +use codex_config::McpServerHttpHeadersHelperConfig; use codex_config::types::McpServerEnvVar; use reqwest::ClientBuilder; use reqwest::header::HeaderMap; @@ -7,7 +10,21 @@ use reqwest::header::HeaderName; use reqwest::header::HeaderValue; use std::collections::HashMap; use std::env; +use std::ffi::OsStr; use std::ffi::OsString; +use std::io; +use std::path::Path; +use std::path::PathBuf; +use std::process::Stdio; +use std::time::Duration; +use tokio::io::AsyncRead; +use tokio::io::AsyncReadExt; +use tokio::process::Command; + +const MAX_HTTP_HEADERS_HELPER_STDOUT_BYTES: usize = 64 * 1024; +const MAX_HTTP_HEADERS_HELPER_STDERR_BYTES: usize = 8 * 1024; +const MAX_HTTP_HEADERS_HELPER_HEADERS: usize = 64; +const MAX_HTTP_HEADERS_HELPER_HEADER_VALUE_BYTES: usize = 16 * 1024; pub(crate) fn create_env_for_mcp_server( extra_env: Option>, @@ -57,9 +74,10 @@ fn local_stdio_env_var_names(env_vars: &[McpServerEnvVar]) -> Result>, env_http_headers: Option>, + http_headers_helper: Option, ) -> Result { let mut headers = HeaderMap::new(); @@ -112,9 +130,221 @@ pub(crate) fn build_default_headers( } } + if let Some(helper) = http_headers_helper { + let helper_headers = run_http_headers_helper(&helper).await?; + if helper_headers.len() > MAX_HTTP_HEADERS_HELPER_HEADERS { + return Err(anyhow!( + "MCP HTTP headers helper `{}` produced {} headers, exceeding the limit of {}", + helper.command, + helper_headers.len(), + MAX_HTTP_HEADERS_HELPER_HEADERS + )); + } + for (name, value) in helper_headers { + let header_name = HeaderName::from_bytes(name.as_bytes()) + .map_err(|err| anyhow!("invalid HTTP header name `{name}` from helper: {err}"))?; + if value.len() > MAX_HTTP_HEADERS_HELPER_HEADER_VALUE_BYTES { + return Err(anyhow!( + "HTTP header value from helper for `{name}` exceeds the limit of {MAX_HTTP_HEADERS_HELPER_HEADER_VALUE_BYTES} bytes" + )); + } + let header_value = HeaderValue::from_str(value.as_str()).map_err(|err| { + anyhow!("invalid HTTP header value from helper for `{name}`: {err}") + })?; + headers.insert(header_name, header_value); + } + } + Ok(headers) } +async fn run_http_headers_helper( + helper: &McpServerHttpHeadersHelperConfig, +) -> Result> { + let program = resolve_http_headers_helper_program(helper)?; + let mut command = Command::new(program); + command + .args(&helper.args) + .current_dir(helper.cwd.as_path()) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true); + + let mut child = command.spawn().map_err(|err| { + anyhow!( + "MCP HTTP headers helper `{}` failed to start: {err}", + helper.command + ) + })?; + let stdout = child.stdout.take().ok_or_else(|| { + anyhow!( + "MCP HTTP headers helper `{}` failed to capture stdout", + helper.command + ) + })?; + let stderr = child.stderr.take().ok_or_else(|| { + anyhow!( + "MCP HTTP headers helper `{}` failed to capture stderr", + helper.command + ) + })?; + let stdout_task = tokio::spawn(read_capped(stdout, MAX_HTTP_HEADERS_HELPER_STDOUT_BYTES)); + let stderr_task = tokio::spawn(read_capped(stderr, MAX_HTTP_HEADERS_HELPER_STDERR_BYTES)); + + let status = match tokio::time::timeout(helper.timeout(), child.wait()).await { + Ok(result) => result.map_err(|err| { + anyhow!( + "MCP HTTP headers helper `{}` failed while waiting for exit: {err}", + helper.command + ) + })?, + Err(_) => { + let _ = child.kill().await; + let _ = stdout_task.await; + let _ = stderr_task.await; + return Err(anyhow!( + "MCP HTTP headers helper `{}` timed out after {} ms", + helper.command, + helper.timeout_ms.get() + )); + } + }; + + let stdout = + join_capped_output(&helper.command, stdout_task, "stdout", helper.timeout()).await?; + let stderr = + join_capped_output(&helper.command, stderr_task, "stderr", helper.timeout()).await?; + + if !status.success() { + let stderr_suffix = if stderr.bytes.is_empty() { + String::new() + } else { + format!( + "; stderr omitted ({} byte{})", + stderr.bytes.len(), + if stderr.bytes.len() == 1 { "" } else { "s" } + ) + }; + return Err(anyhow!( + "MCP HTTP headers helper `{}` exited with status {status}{stderr_suffix}", + helper.command + )); + } + + if stdout.truncated { + return Err(anyhow!( + "MCP HTTP headers helper `{}` wrote more than {} bytes to stdout", + helper.command, + MAX_HTTP_HEADERS_HELPER_STDOUT_BYTES + )); + } + + let stdout = String::from_utf8(stdout.bytes).map_err(|_| { + anyhow!( + "MCP HTTP headers helper `{}` wrote non-UTF-8 data to stdout", + helper.command + ) + })?; + let output = stdout.trim(); + if output.is_empty() { + return Err(anyhow!( + "MCP HTTP headers helper `{}` produced empty output", + helper.command + )); + } + + serde_json::from_str(output).map_err(|err| { + anyhow!( + "MCP HTTP headers helper `{}` must output a JSON object with string values: {err}", + helper.command + ) + }) +} + +fn resolve_http_headers_helper_program( + helper: &McpServerHttpHeadersHelperConfig, +) -> Result { + let command = Path::new(&helper.command); + if command.is_absolute() { + return Ok(command.as_os_str().to_os_string()); + } + if has_path_separator(command.as_os_str()) { + return Ok(helper.cwd.as_path().join(command).into_os_string()); + } + + program_resolver::resolve( + command.as_os_str().to_os_string(), + &env::vars_os().collect(), + helper.cwd.as_path(), + ) + .map_err(|err| { + anyhow!( + "MCP HTTP headers helper `{}` could not be resolved: {err}", + helper.command + ) + }) +} + +fn has_path_separator(value: &OsStr) -> bool { + let path = PathBuf::from(value); + path.components().count() > 1 +} + +struct CappedOutput { + bytes: Vec, + truncated: bool, +} + +async fn read_capped(mut reader: R, max_bytes: usize) -> io::Result +where + R: AsyncRead + Unpin, +{ + let mut bytes = Vec::new(); + let mut truncated = false; + let mut chunk = [0_u8; 8192]; + + loop { + let read = reader.read(&mut chunk).await?; + if read == 0 { + break; + } + + let remaining = max_bytes.saturating_sub(bytes.len()); + if remaining == 0 { + truncated = true; + continue; + } + + let to_copy = remaining.min(read); + bytes.extend_from_slice(&chunk[..to_copy]); + if to_copy < read { + truncated = true; + } + } + + Ok(CappedOutput { bytes, truncated }) +} + +async fn join_capped_output( + command: &str, + mut task: tokio::task::JoinHandle>, + stream_name: &str, + timeout: Duration, +) -> Result { + tokio::select! { + result = &mut task => { + result + .map_err(|err| anyhow!("MCP HTTP headers helper `{command}` {stream_name} task failed: {err}"))? + .map_err(|err| anyhow!("MCP HTTP headers helper `{command}` failed to read {stream_name}: {err}")) + } + _ = tokio::time::sleep(timeout) => { + task.abort(); + Err(anyhow!("MCP HTTP headers helper `{command}` timed out while reading {stream_name}")) + } + } +} + pub(crate) fn apply_default_headers( builder: ClientBuilder, default_headers: &HeaderMap, @@ -148,10 +378,13 @@ pub(crate) const DEFAULT_ENV_VARS: &[&str] = #[cfg(test)] mod tests { use super::*; + use codex_config::AbsolutePathBuf; use pretty_assertions::assert_eq; use serial_test::serial; use std::ffi::OsStr; + use std::path::Path; + use std::path::PathBuf; struct EnvVarGuard { key: String, @@ -290,6 +523,229 @@ mod tests { ); } + fn write_headers_helper_script(dir: &Path, name: &str, output: &str) -> PathBuf { + let script_path = dir.join(helper_script_name(name)); + std::fs::write(&script_path, helper_script_contents(output)) + .expect("headers helper script should be written"); + make_executable(&script_path); + script_path + } + + fn write_failing_headers_helper_script(dir: &Path, name: &str, stderr_output: &str) -> PathBuf { + let script_path = dir.join(helper_script_name(name)); + std::fs::write(&script_path, failing_helper_script_contents(stderr_output)) + .expect("headers helper script should be written"); + make_executable(&script_path); + script_path + } + + #[cfg(unix)] + fn helper_script_name(name: &str) -> String { + name.to_string() + } + + #[cfg(windows)] + fn helper_script_name(name: &str) -> String { + format!("{name}.cmd") + } + + #[cfg(unix)] + fn helper_script_contents(output: &str) -> String { + format!("#!/bin/sh\nprintf '%s\\n' '{output}'\n") + } + + #[cfg(unix)] + fn failing_helper_script_contents(stderr_output: &str) -> String { + format!("#!/bin/sh\nprintf '%s\\n' '{stderr_output}' >&2\nexit 1\n") + } + + #[cfg(windows)] + fn helper_script_contents(output: &str) -> String { + format!("@echo off\r\necho {output}\r\n") + } + + #[cfg(windows)] + fn failing_helper_script_contents(stderr_output: &str) -> String { + format!("@echo off\r\necho {stderr_output} 1>&2\r\nexit /b 1\r\n") + } + + #[cfg(unix)] + fn make_executable(path: &Path) { + use std::os::unix::fs::PermissionsExt; + + let mut permissions = std::fs::metadata(path) + .expect("headers helper script should have metadata") + .permissions(); + permissions.set_mode(0o755); + std::fs::set_permissions(path, permissions) + .expect("headers helper script should be executable"); + } + + #[cfg(windows)] + fn make_executable(_path: &Path) {} + + #[tokio::test] + #[serial(extra_rmcp_env)] + async fn build_default_headers_runs_helper_and_allows_helper_to_override() { + let temp_dir = tempfile::TempDir::new().expect("tempdir"); + let script_path = write_headers_helper_script( + temp_dir.path(), + "headers-helper", + r#"{"Authorization":"Bearer helper","X-Static":"helper","X-Arg":"from-helper"}"#, + ); + let _guard = EnvVarGuard::set("CODEX_RMCP_CLIENT_HEADER_TEST", "from-env"); + + let headers = build_default_headers( + Some(HashMap::from([ + ("Authorization".to_string(), "Bearer static".to_string()), + ("X-Static".to_string(), "static".to_string()), + ])), + Some(HashMap::from([( + "X-Env".to_string(), + "CODEX_RMCP_CLIENT_HEADER_TEST".to_string(), + )])), + Some(McpServerHttpHeadersHelperConfig::new( + script_path.display().to_string(), + Vec::new(), + )), + ) + .await + .expect("headers should build"); + + assert_eq!( + headers + .get("authorization") + .and_then(|value| value.to_str().ok()), + Some("Bearer helper") + ); + assert_eq!( + headers + .get("x-static") + .and_then(|value| value.to_str().ok()), + Some("helper") + ); + assert_eq!( + headers.get("x-env").and_then(|value| value.to_str().ok()), + Some("from-env") + ); + assert_eq!( + headers.get("x-arg").and_then(|value| value.to_str().ok()), + Some("from-helper") + ); + } + + #[tokio::test] + async fn build_default_headers_resolves_relative_helper_against_cwd() { + let temp_dir = tempfile::TempDir::new().expect("tempdir"); + write_headers_helper_script( + temp_dir.path(), + "headers-helper", + r#"{"Authorization":"Bearer helper"}"#, + ); + let helper_program = format!( + ".{}{}", + std::path::MAIN_SEPARATOR, + helper_script_name("headers-helper") + ); + let mut helper = McpServerHttpHeadersHelperConfig::new(helper_program, Vec::new()); + helper.cwd = AbsolutePathBuf::from_absolute_path(temp_dir.path()) + .expect("tempdir path should be absolute"); + + let headers = build_default_headers( + /*http_headers*/ None, + /*env_http_headers*/ None, + Some(helper), + ) + .await + .expect("headers should build"); + + assert_eq!( + headers + .get("authorization") + .and_then(|value| value.to_str().ok()), + Some("Bearer helper") + ); + } + + #[tokio::test] + async fn build_default_headers_rejects_invalid_helper_output() { + let temp_dir = tempfile::TempDir::new().expect("tempdir"); + let script_path = + write_headers_helper_script(temp_dir.path(), "bad-headers-helper", "not-json"); + + let err = build_default_headers( + /*http_headers*/ None, + /*env_http_headers*/ None, + Some(McpServerHttpHeadersHelperConfig::new( + script_path.display().to_string(), + Vec::new(), + )), + ) + .await + .expect_err("invalid helper output should fail"); + + assert!( + err.to_string() + .contains("must output a JSON object with string values"), + "unexpected error: {err}" + ); + } + + #[tokio::test] + async fn build_default_headers_rejects_oversized_helper_output() { + let temp_dir = tempfile::TempDir::new().expect("tempdir"); + let output = "x".repeat(MAX_HTTP_HEADERS_HELPER_STDOUT_BYTES + 1); + let script_path = + write_headers_helper_script(temp_dir.path(), "oversized-headers-helper", &output); + + let err = build_default_headers( + /*http_headers*/ None, + /*env_http_headers*/ None, + Some(McpServerHttpHeadersHelperConfig::new( + script_path.display().to_string(), + Vec::new(), + )), + ) + .await + .expect_err("oversized helper output should fail"); + + assert!( + err.to_string().contains("wrote more than"), + "unexpected error: {err}" + ); + } + + #[tokio::test] + async fn build_default_headers_omits_helper_stderr_from_error() { + let temp_dir = tempfile::TempDir::new().expect("tempdir"); + let script_path = write_failing_headers_helper_script( + temp_dir.path(), + "failing-headers-helper", + "secret-token-value", + ); + + let err = build_default_headers( + /*http_headers*/ None, + /*env_http_headers*/ None, + Some(McpServerHttpHeadersHelperConfig::new( + script_path.display().to_string(), + Vec::new(), + )), + ) + .await + .expect_err("failing helper should fail"); + let message = err.to_string(); + + assert!( + message.contains("stderr omitted"), + "unexpected error: {err}" + ); + assert!( + !message.contains("secret-token-value"), + "stderr leaked in error: {err}" + ); + } + #[cfg(unix)] #[test] #[serial(extra_rmcp_env)] 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 822acef1a26..184cd6d5e7a 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_test_support.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_test_support.rs @@ -80,6 +80,7 @@ pub(crate) async fn create_client(base_url: &str) -> anyhow::Result Some("test-bearer".to_string()), /*http_headers*/ None, /*env_http_headers*/ None, + /*http_headers_helper*/ None, OAuthCredentialsStoreMode::File, Environment::default_for_tests().get_http_client(), /*auth_provider*/ None, @@ -118,6 +119,7 @@ pub(crate) async fn create_remote_client( Some("test-bearer".to_string()), /*http_headers*/ None, /*env_http_headers*/ None, + /*http_headers_helper*/ None, OAuthCredentialsStoreMode::File, Arc::new(http_client), /*auth_provider*/ None,