diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 59d3eb730b..5ea1ccab25 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -4098,6 +4098,7 @@ dependencies = [ "arc-swap", "assert_matches", "async-channel", + "base64 0.22.1", "codex-api", "codex-async-utils", "codex-config", @@ -4105,6 +4106,8 @@ dependencies = [ "codex-diagnostics", "codex-exec-server", "codex-exec-server-test-support", + "codex-features", + "codex-http-client", "codex-login", "codex-model-provider", "codex-network-proxy", @@ -4112,12 +4115,14 @@ dependencies = [ "codex-plugin", "codex-protocol", "codex-rmcp-client", + "codex-runtime-agent-instance", "codex-utils-absolute-path", "codex-utils-cargo-bin", "codex-utils-path-uri", "codex-utils-plugins", "codex-utils-process", "futures", + "http 1.4.0", "insta", "lru", "pretty_assertions", @@ -4133,6 +4138,7 @@ dependencies = [ "tracing", "url", "uuid", + "wiremock", ] [[package]] diff --git a/codex-rs/analytics/src/client.rs b/codex-rs/analytics/src/client.rs index 967b45e4ad..6a9dd298ca 100644 --- a/codex-rs/analytics/src/client.rs +++ b/codex-rs/analytics/src/client.rs @@ -406,6 +406,23 @@ impl AnalyticsEventsClient { } } + /// Give a resumed actor its own queue while preserving the caller's analytics setting. + pub fn for_auth_context( + &self, + auth_manager: Arc, + base_url: String, + runtime_auth: Option<(RuntimeAuthSlot, RuntimeAuthMode)>, + api_key_headers: HeaderMap, + ) -> Self { + Self::new_with_runtime( + auth_manager, + base_url, + Some(self.queue.is_some()), + runtime_auth, + api_key_headers, + ) + } + pub fn disabled() -> Self { Self { queue: None } } diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index fc089fb715..6b626c6741 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -484,7 +484,6 @@ impl MessageProcessor { Arc::clone(&model_catalog), ); let apps_processor = AppsRequestProcessor::new( - auth_manager.clone(), Arc::clone(&thread_manager), outgoing.clone(), config_manager.clone(), diff --git a/codex-rs/app-server/src/request_processors/apps_processor.rs b/codex-rs/app-server/src/request_processors/apps_processor.rs index be74a8fb61..3dc9059341 100644 --- a/codex-rs/app-server/src/request_processors/apps_processor.rs +++ b/codex-rs/app-server/src/request_processors/apps_processor.rs @@ -8,7 +8,6 @@ mod read; pub(super) use read::APP_READ_MAX_IDS; pub(crate) struct AppsRequestProcessor { - auth_manager: Arc, thread_manager: Arc, outgoing: Arc, config_manager: ConfigManager, @@ -18,7 +17,6 @@ pub(crate) struct AppsRequestProcessor { impl AppsRequestProcessor { pub(crate) fn new( - auth_manager: Arc, thread_manager: Arc, outgoing: Arc, config_manager: ConfigManager, @@ -26,7 +24,6 @@ impl AppsRequestProcessor { ) -> Self { let shutdown_drop_guard = shutdown_token.clone().drop_guard(); Self { - auth_manager, thread_manager, outgoing, config_manager, @@ -64,13 +61,27 @@ impl AppsRequestProcessor { }; let mut config = self.load_latest_config(fallback_cwd).await?; - if let Some(thread) = thread { + if let Some(thread) = thread.as_ref() { let _ = config .features .set_enabled(Feature::Apps, thread.enabled(Feature::Apps)); } - let auth = self.auth_manager.auth().await; + let mcp_manager = thread + .as_ref() + .map(|thread| thread.mcp_manager()) + .unwrap_or_else(|| self.thread_manager.mcp_manager()); + self.thread_manager + .runtime_auth_for_mcp_manager(&config, &mcp_manager) + .map_err(|error| internal_error(error.to_string()))?; + let auth_manager = mcp_manager.auth_manager(); + if params.force_refetch && !auth_manager.has_external_auth() { + auth_manager + .reload_with_cached_auth() + .await + .map_err(|error| internal_error(error.to_string()))?; + } + let auth = auth_manager.auth().await; if !config .features .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend)) @@ -86,8 +97,7 @@ impl AppsRequestProcessor { let request = request_id.clone(); let outgoing = Arc::clone(&self.outgoing); let environment_manager = self.thread_manager.environment_manager(); - let mcp_manager = self.thread_manager.mcp_manager(); - let plugins_manager = self.thread_manager.plugins_manager(); + let plugins_manager = mcp_manager.plugins_manager(); let shutdown_token = self.shutdown_token.child_token(); tokio::spawn(async move { tokio::select! { @@ -195,10 +205,33 @@ impl AppsRequestProcessor { loaded_plugins.capability_summaries(), ); let plugin_apps = connector_snapshot.connector_ids().to_vec(); - let (mut accessible_connectors, mut all_connectors) = tokio::join!( - connectors::list_cached_accessible_connectors_from_mcp_tools(&config), - connectors::list_cached_all_connectors(&config, &plugin_apps) - ); + let mut config = config; + let auth = match mcp_manager + .auth_manager() + .auth_with_http_client_factory() + .await + { + Some((auth, factory)) => { + config.application_network_policy = factory.network_policy().clone(); + Some(auth) + } + None => None, + }; + let runtime = mcp_manager + .runtime_auth_for_config(&config) + .map(|(slot, mode)| { + slot.get() + .cloned() + .map(|runtime| (runtime, mode)) + .ok_or_else(|| { + internal_error("runtime authentication is not initialized".to_string()) + }) + }) + .transpose()?; + let mut accessible_connectors = + connectors::list_cached_accessible_connectors_from_mcp_tools(&config, auth.as_ref()); + let mut all_connectors = + connectors::list_cached_all_connectors(&config, auth.as_ref(), &plugin_apps); let cached_all_connectors = all_connectors.clone(); let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); @@ -222,6 +255,8 @@ impl AppsRequestProcessor { tokio::spawn(async move { let result = connectors::list_all_connectors_with_options( &all_config, + auth.as_ref(), + runtime.as_ref(), force_refetch, &all_plugin_apps, ) @@ -354,19 +389,32 @@ impl AppsRequestProcessor { .map_err(|err| internal_error(format!("failed to reload config: {err}"))) } - async fn load_apps_config(&self, thread_id: Option<&str>) -> Result { + async fn load_apps_config( + &self, + thread_id: Option<&str>, + ) -> Result<(Config, Arc), JSONRPCErrorError> { let Some(thread_id) = thread_id else { - return self.load_latest_config(/*fallback_cwd*/ None).await; + let config = self.load_latest_config(/*fallback_cwd*/ None).await?; + self.thread_manager + .runtime_auth_for_config(&config) + .map_err(|error| internal_error(error.to_string()))?; + return Ok((config, self.thread_manager.mcp_manager())); }; let (_, thread) = self.load_thread(thread_id).await?; let thread_config = thread.config().await; - self.config_manager + let config = self + .config_manager .load_latest_config_with_session_layers( &thread_config.config_layer_stack, &thread_config.cwd, ) .await - .map_err(|err| internal_error(format!("failed to reload config: {err}"))) + .map_err(|err| internal_error(format!("failed to reload config: {err}")))?; + let mcp_manager = thread.mcp_manager(); + self.thread_manager + .runtime_auth_for_mcp_manager(&config, &mcp_manager) + .map_err(|error| internal_error(error.to_string()))?; + Ok((config, mcp_manager)) } } diff --git a/codex-rs/app-server/src/request_processors/apps_processor/installed.rs b/codex-rs/app-server/src/request_processors/apps_processor/installed.rs index 7631f24266..ce6c105261 100644 --- a/codex-rs/app-server/src/request_processors/apps_processor/installed.rs +++ b/codex-rs/app-server/src/request_processors/apps_processor/installed.rs @@ -58,12 +58,14 @@ impl AppsRequestProcessor { None, ), }; - let auth = self.auth_manager.auth().await; + let mcp_manager = thread.as_ref().map(|thread| thread.mcp_manager()) + .unwrap_or_else(|| self.thread_manager.mcp_manager()); + let auth_manager = mcp_manager.auth_manager(); + let auth = auth_manager.auth().await; let runtime_enabled = config .features .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend)); - let mcp_manager = self.thread_manager.mcp_manager(); let cache_key = connector_runtime_context_key(auth.as_ref()); let previous_snapshot = mcp_manager .codex_apps_tools_cache() @@ -96,9 +98,11 @@ impl AppsRequestProcessor { let cancellation_token = CancellationToken::new(); let codex_apps_auth_manager = host_owned_codex_apps_enabled(&mcp_config, auth.as_ref()) - .then(|| Arc::clone(&self.auth_manager)); + .then(|| Arc::clone(&auth_manager)); + self.thread_manager.runtime_auth_for_mcp_manager(&config, &mcp_manager)?; let runtime = McpRuntime::new(McpRuntimeInput { network_policy_metadata: None, + runtime_auth: mcp_manager.runtime_auth_for_config(&config), startup_policy: McpStartupPolicy::Eager, config: Arc::clone(&mcp_config), plugins_available: false, diff --git a/codex-rs/app-server/src/request_processors/apps_processor/read.rs b/codex-rs/app-server/src/request_processors/apps_processor/read.rs index 21ab3c8084..4348613f35 100644 --- a/codex-rs/app-server/src/request_processors/apps_processor/read.rs +++ b/codex-rs/app-server/src/request_processors/apps_processor/read.rs @@ -26,14 +26,22 @@ impl AppsRequestProcessor { .into_iter() .filter(|app_id| seen_app_ids.insert(app_id.clone())) .collect::>(); - let mut config = self.load_apps_config(thread_id.as_deref()).await?; - let auth = match self.auth_manager.auth_with_http_client_factory().await { + let (mut config, mcp_manager) = self.load_apps_config(thread_id.as_deref()).await?; + let auth = match mcp_manager + .auth_manager() + .auth_with_http_client_factory() + .await + { Some((auth, factory)) => { config.application_network_policy = factory.network_policy().clone(); Some(auth) } None => None, }; + let runtime = self + .thread_manager + .runtime_auth_for_mcp_manager(&config, &mcp_manager) + .map_err(|error| internal_error(error.to_string()))?; if !config .features .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend)) @@ -52,11 +60,16 @@ impl AppsRequestProcessor { let connectors::ConnectorMetadataReadResult { apps, missing_app_ids, - } = connectors::read_connector_metadata(&config, auth, &app_ids, include_tools) - .await - .map_err(|err| internal_error(format!("failed to read app metadata: {err}")))?; - let loaded_plugins = self - .thread_manager + } = connectors::read_connector_metadata( + &config, + auth, + runtime.as_ref(), + &app_ids, + include_tools, + ) + .await + .map_err(|err| internal_error(format!("failed to read app metadata: {err}")))?; + let loaded_plugins = mcp_manager .plugins_manager() .plugins_for_config(&config.plugins_config_input()) .await; 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 aebbab96e4..2d6b065491 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -428,8 +428,15 @@ impl McpRequestProcessor { } None => self.load_global_config().await?, }; - let mcp_manager = self.thread_manager.mcp_manager(); - let auth = self.auth_manager.auth().await; + let mcp_manager = thread + .as_ref() + .map(|thread| thread.mcp_manager()) + .unwrap_or_else(|| self.thread_manager.mcp_manager()); + self.thread_manager + .runtime_auth_for_mcp_manager(&config, &mcp_manager) + .map_err(|error| internal_error(error.to_string()))?; + let auth = mcp_manager.auth_manager().auth().await; + let runtime_auth = mcp_manager.runtime_auth_for_config(&config); let (mcp_config, runtime_context) = match thread.as_ref() { Some(thread) => thread.runtime_mcp_config_and_context(&config).await, None => ( @@ -444,6 +451,7 @@ impl McpRequestProcessor { &mcp_config, auth.as_ref(), Some(&self.auth_manager), + runtime_auth, runtime_context, mcp_manager.codex_apps_tools_cache(), mcp_manager.tool_catalog_cache(), @@ -700,7 +708,11 @@ impl McpRequestProcessor { let mcp_config = mcp_manager.runtime_config(&config).await; let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache(); let tool_catalog_cache = mcp_manager.tool_catalog_cache(); - let auth = self.auth_manager.auth().await; + self.thread_manager + .runtime_auth_for_mcp_manager(&config, &mcp_manager) + .map_err(|error| internal_error(error.to_string()))?; + let runtime_auth = mcp_manager.runtime_auth_for_config(&config); + let auth = mcp_manager.auth_manager().auth().await; let environment_manager = self.thread_manager.environment_manager(); // This threadless resource-read path has no turn cwd or turn-selected // environment. Use config cwd only as the local stdio fallback; named @@ -713,6 +725,7 @@ impl McpRequestProcessor { let result = read_mcp_resource_without_thread( &mcp_config, auth.as_ref(), + runtime_auth, runtime_context, codex_apps_tools_cache, tool_catalog_cache, diff --git a/codex-rs/app-server/src/request_processors/plugins.rs b/codex-rs/app-server/src/request_processors/plugins.rs index 875006ce10..f87e8e5762 100644 --- a/codex-rs/app-server/src/request_processors/plugins.rs +++ b/codex-rs/app-server/src/request_processors/plugins.rs @@ -1113,6 +1113,7 @@ impl PluginRequestProcessor { None => None, }; let app_summaries = load_plugin_app_summaries( + &self.thread_manager, &config, auth.as_ref(), &outcome.plugin.apps, @@ -1219,6 +1220,7 @@ impl PluginRequestProcessor { .map(plugin_app_category_by_id_from_value) .unwrap_or_default(); let app_summaries = load_plugin_app_summaries( + &self.thread_manager, &config, auth.as_ref(), &plugin_apps, @@ -1738,37 +1740,42 @@ impl PluginRequestProcessor { } let is_chatgpt_auth = auth.as_ref().is_some_and(CodexAuth::is_chatgpt_auth); - let apps_needing_auth = if let Some(app_ids_needing_auth) = - installation.app_ids_needing_auth - { - if app_ids_needing_auth.is_empty() - || !config.features.apps_enabled_for_auth(is_chatgpt_auth) - { - Vec::new() - } else { - let plugin_apps = app_ids_needing_auth - .into_iter() - .map(codex_plugin::AppConnectorId) - .collect::>(); - let app_category_by_id = remote_detail - .app_manifest - .as_ref() - .map(plugin_app_category_by_id_from_value) - .unwrap_or_default(); - load_plugin_app_summaries(&config, auth.as_ref(), &plugin_apps, &app_category_by_id) + let apps_needing_auth = + if let Some(app_ids_needing_auth) = installation.app_ids_needing_auth { + if app_ids_needing_auth.is_empty() + || !config.features.apps_enabled_for_auth(is_chatgpt_auth) + { + Vec::new() + } else { + let plugin_apps = app_ids_needing_auth + .into_iter() + .map(codex_plugin::AppConnectorId) + .collect::>(); + let app_category_by_id = remote_detail + .app_manifest + .as_ref() + .map(plugin_app_category_by_id_from_value) + .unwrap_or_default(); + load_plugin_app_summaries( + &self.thread_manager, + &config, + auth.as_ref(), + &plugin_apps, + &app_category_by_id, + ) .await - } - } else { - let plugin_app_declarations = - load_plugin_apps(result.installed_path.as_path(), &result.plugin_id).await; - self.plugin_apps_needing_auth_for_install( - &config, - auth.as_ref(), - &result.plugin_id.as_key(), - &plugin_app_declarations, - ) - .await - }; + } + } else { + let plugin_app_declarations = + load_plugin_apps(result.installed_path.as_path(), &result.plugin_id).await; + self.plugin_apps_needing_auth_for_install( + &config, + auth.as_ref(), + &result.plugin_id.as_key(), + &plugin_app_declarations, + ) + .await + }; Ok(PluginInstallResponse { auth_policy: remote_detail.summary.auth_policy, @@ -1839,7 +1846,13 @@ impl PluginRequestProcessor { .collect(); let environment_manager = self.thread_manager.environment_manager(); let (app_summaries, accessible_connectors_result) = tokio::join!( - load_plugin_app_summaries(config, auth, &plugin_apps, &app_category_by_id), + load_plugin_app_summaries( + &self.thread_manager, + config, + auth, + &plugin_apps, + &app_category_by_id + ), connectors::list_accessible_connectors_from_mcp_tools_with_mcp_manager( config, /*force_refetch*/ true, @@ -1856,8 +1869,7 @@ impl PluginRequestProcessor { "failed to load accessible apps after plugin install: {err:#}" ); ( - connectors::list_cached_accessible_connectors_from_mcp_tools(config) - .await + connectors::list_cached_accessible_connectors_from_mcp_tools(config, auth) .unwrap_or_default(), false, ) @@ -2159,6 +2171,7 @@ impl PluginRequestProcessor { } async fn load_plugin_app_summaries( + thread_manager: &ThreadManager, config: &Config, auth: Option<&CodexAuth>, plugin_apps: &[codex_plugin::AppConnectorId], @@ -2185,9 +2198,17 @@ async fn load_plugin_app_summaries( ) .fresh_records(&app_ids, /*include_tools*/ false), ); + let runtime = thread_manager.runtime_auth_for_config(config); for app_ids in app_ids.chunks(APP_READ_MAX_IDS) { + let runtime = match &runtime { + Ok(runtime) => runtime.as_ref(), + Err(error) => { + warn!("failed to initialize app metadata authentication: {error}"); + break; + } + }; match connectors::read_connector_metadata( - config, auth, app_ids, /*include_tools*/ false, + config, auth, runtime, app_ids, /*include_tools*/ false, ) .await { diff --git a/codex-rs/app-server/tests/suite/v2/app_list.rs b/codex-rs/app-server/tests/suite/v2/app_list.rs index 5805f02ae7..efef823ce0 100644 --- a/codex-rs/app-server/tests/suite/v2/app_list.rs +++ b/codex-rs/app-server/tests/suite/v2/app_list.rs @@ -27,6 +27,8 @@ use codex_app_server_protocol::AppReview; use codex_app_server_protocol::AppScreenshot; use codex_app_server_protocol::AppsListParams; use codex_app_server_protocol::AppsListResponse; +use codex_app_server_protocol::GetAccountParams; +use codex_app_server_protocol::GetAccountResponse; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::LoginAccountResponse; use codex_app_server_protocol::RequestId; @@ -1122,6 +1124,40 @@ async fn list_apps_force_refetch_preserves_previous_cache_on_failure() -> Result assert_eq!(initial_data.len(), 1); assert!(initial_data.iter().all(|app| app.is_accessible)); + let account_request = mcp + .send_get_account_request(GetAccountParams { + refresh_token: false, + }) + .await?; + let initial_account: GetAccountResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(account_request)).await??; + assert!(initial_account.account.is_some()); + + let auth_path = codex_home.path().join("auth.json"); + std::fs::write(&auth_path, "")?; + let refetch_request = mcp + .send_apps_list_request(AppsListParams { + limit: None, + cursor: None, + thread_id: None, + force_refetch: true, + }) + .await?; + let _: JSONRPCError = timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(refetch_request)), + ) + .await??; + + let account_request = mcp + .send_get_account_request(GetAccountParams { + refresh_token: false, + }) + .await?; + let account: GetAccountResponse = + timeout(DEFAULT_TIMEOUT, mcp.read_response(account_request)).await??; + assert_eq!(account, initial_account); + write_chatgpt_auth( codex_home.path(), ChatGptAuthFixture::new("chatgpt-token-invalid") @@ -1586,6 +1622,18 @@ async fn start_apps_server_with_delays_and_control_inner( ); let router = Router::new() + .route( + "/api/codex/accounts/check", + get(|State(state): State>| async move { + Json(json!({ + "accounts": [{ + "id": state.expected_account_id, + "workspace_backend_origin": "https://chatgpt.com", + "account_routing_override": "NO_CONSTRAINT", + }], + })) + }), + ) .route("/connectors/directory/list", get(list_directory_connectors)) .route( "/connectors/directory/list_workspace", diff --git a/codex-rs/app-server/tests/suite/v2/app_read.rs b/codex-rs/app-server/tests/suite/v2/app_read.rs index 058c51dde3..34738b9317 100644 --- a/codex-rs/app-server/tests/suite/v2/app_read.rs +++ b/codex-rs/app-server/tests/suite/v2/app_read.rs @@ -115,14 +115,19 @@ async fn connector_requests_preserve_account_policy() -> Result<()> { let (auth, factory) = manager.auth_with_http_client_factory().await.unwrap(); config.application_network_policy = factory.network_policy().clone(); let app_ids = vec!["alpha".to_string()]; - connectors::read_connector_metadata(&config, &auth, &app_ids, /*include_tools*/ false).await?; + connectors::read_connector_metadata( + &config, &auth, /*runtime*/ None, &app_ids, /*include_tools*/ false, + ) + .await?; assert_eq!(state.requests().len(), 1); policy.invalidate(); controller.publish(policy.revision(), DestinationPolicy::Unrestricted); let error = timeout( DEFAULT_TIMEOUT, - connectors::read_connector_metadata(&config, &auth, &app_ids, /*include_tools*/ false), + connectors::read_connector_metadata( + &config, &auth, /*runtime*/ None, &app_ids, /*include_tools*/ false, + ), ) .await? .err() @@ -139,9 +144,17 @@ async fn connector_requests_preserve_account_policy() -> Result<()> { allowed_hosts: Default::default(), }, ); + let (auth, factory) = manager.auth_with_http_client_factory().await.unwrap(); + config.application_network_policy = factory.network_policy().clone(); let error = timeout( DEFAULT_TIMEOUT, - connectors::list_all_connectors_with_options(&config, /*force_refetch*/ true, &[]), + connectors::list_all_connectors_with_options( + &config, + Some(&auth), + /*runtime*/ None, + /*force_refetch*/ true, + &[], + ), ) .await? .unwrap_err(); diff --git a/codex-rs/chatgpt/src/chatgpt_client.rs b/codex-rs/chatgpt/src/chatgpt_client.rs index fd43eabff0..1b0cfd3d66 100644 --- a/codex-rs/chatgpt/src/chatgpt_client.rs +++ b/codex-rs/chatgpt/src/chatgpt_client.rs @@ -59,14 +59,30 @@ pub(crate) async fn chatgpt_get_request( config: &Config, path: String, ) -> anyhow::Result { - // This helper is used by the standalone apply command. Harness callers pass their owner below. - // The temporary credential manager does not own account changes. - let mut captured_config = config.clone(); - captured_config.application_network_policy = config + let mut config = config.clone(); + let (auth, runtime) = standalone_auth(&mut config).await?; + chatgpt_get_request_with_timeout( + &config, + &auth, + runtime.as_ref(), + path, + /*timeout*/ None, + ) + .await +} + +pub(crate) async fn standalone_auth( + config: &mut Config, +) -> anyhow::Result<( + CodexAuth, + Option<(Arc, RuntimeAuthMode)>, +)> { + // Standalone commands create one owner. Freeze the caller's account policy before loading + // credentials; this temporary manager does not own later account changes. + config.application_network_policy = config .application_network_policy .clone() .for_current_account(); - let config = &captured_config; if config.features.enabled(Feature::RuntimeAgentInstances) { codex_model_provider::validate_runtime_destination(Some(&config.chatgpt_base_url))?; } @@ -85,7 +101,7 @@ pub(crate) async fn chatgpt_get_request( .auth() .await .ok_or_else(|| anyhow::anyhow!("ChatGPT auth not available"))?; - chatgpt_get_request_with_timeout(config, &auth, runtime.as_ref(), path, /*timeout*/ None).await + Ok((auth, runtime)) } pub(crate) async fn chatgpt_get_request_with_timeout( diff --git a/codex-rs/chatgpt/src/chatgpt_client_runtime_auth_tests.rs b/codex-rs/chatgpt/src/chatgpt_client_runtime_auth_tests.rs index d14d53828a..65cfb3c5af 100644 --- a/codex-rs/chatgpt/src/chatgpt_client_runtime_auth_tests.rs +++ b/codex-rs/chatgpt/src/chatgpt_client_runtime_auth_tests.rs @@ -44,6 +44,11 @@ async fn helpers_use_captured_auth_and_strip_cross_origin_proofs() { .expect(1) .mount(&server) .await; + Mock::given(path("/ps/apps/batch")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "apps": [{"id": "calendar", "name": "Calendar", "description": "Calendar metadata"}], + }))) + .expect(1).mount(&server).await; let home = tempfile::TempDir::new().unwrap(); let mut config = codex_core::config::ConfigBuilder::default() .codex_home(home.path().to_path_buf()) @@ -94,6 +99,18 @@ async fn helpers_use_captured_auth_and_strip_cross_origin_proofs() { ) .await .unwrap(); + let metadata = crate::connectors::read_connector_metadata( + &config, + &auth, + Some(&runtime), + &["calendar".into()], + /*include_tools*/ false, + ) + .await + .unwrap(); + assert_eq!(metadata.apps.len(), 1); + assert_eq!(metadata.apps[0].name, "Calendar"); + assert!(metadata.missing_app_ids.is_empty()); assert_eq!(get, json!({"ok": true})); assert_eq!(post, get); let redirected_result = chatgpt_get_request_with_timeout::( @@ -116,9 +133,22 @@ async fn helpers_use_captured_auth_and_strip_cross_origin_proofs() { ); redirected.verify().await; let requests = server.received_requests().await.unwrap(); + let metadata_request = requests + .iter() + .find(|request| request.url.path() == "/ps/apps/batch") + .unwrap(); + assert_eq!( + metadata_request.headers["authorization"], + format!("Bearer {token}") + ); + assert!( + !metadata_request + .headers + .contains_key("x-openai-agent-assertion") + ); let proofs: HashSet<_> = requests .iter() - .filter(|r| r.url.path() != "/agent/register") + .filter(|r| !matches!(r.url.path(), "/agent/register" | "/ps/apps/batch")) .map(|r| { assert_eq!(r.headers["chatgpt-account-id"], "workspace-a"); let authorization = r.headers["authorization"].to_str().unwrap(); @@ -186,6 +216,85 @@ async fn standalone_apply_keeps_custom_backend_auth_when_runtime_feature_is_enab server.verify().await; } +#[tokio::test] +async fn connector_requests_reject_unsupported_openai_destinations() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + for base_url in [ + "http://chatgpt.com", + "https://future.openai.com/backend-api", + ] { + let home = tempfile::TempDir::new().unwrap(); + let mut config = codex_core::config::ConfigBuilder::default() + .codex_home(home.path().to_path_buf()) + .build() + .await + .unwrap(); + config.chatgpt_base_url = base_url.to_string(); + config.features.enable(Feature::Apps).unwrap(); + let auth = CodexAuth::from_external_chatgpt_tokens( + &actor_token(), + "workspace-a", + /*chatgpt_plan_type*/ None, + ) + .unwrap(); + let runtime = ( + Arc::new(RuntimeAuthManager::for_chatgpt( + AuthManager::from_auth_for_testing(auth.clone()), + uuid::Uuid::nil().to_string(), + base_url.to_string(), + RuntimeRegistrationMetadata { + harness: Some("codex".to_string()), + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + )), + mode, + ); + + // Keep a broken validation path from sending the test credential to the network. + let controller = codex_http_client::NetworkPolicyController::default(); + config.application_network_policy = controller.policy(); + assert!(controller.publish( + config.application_network_policy.revision(), + codex_http_client::DestinationPolicy::Restricted { + allowed_hosts: Default::default(), + }, + )); + + let directory_error = crate::connectors::list_all_connectors_with_options( + &config, + Some(&auth), + Some(&runtime), + /*force_refetch*/ true, + &[], + ) + .await + .expect_err("reject unsupported directory destination"); + let metadata_error = crate::connectors::read_connector_metadata( + &config, + &auth, + Some(&runtime), + &["calendar".to_string()], + /*include_tools*/ false, + ) + .await + .err() + .expect("reject unsupported metadata destination"); + for error in [directory_error, metadata_error] { + assert!( + matches!( + error + .downcast_ref::() + .map(codex_protocol::error::CodexErr::details), + Some(codex_protocol::error::CodexErrorDetails::UnsupportedOperation(_)) + ), + "{base_url} ({mode:?}): {error:#}", + ); + } + } + } +} + #[tokio::test] async fn standalone_apply_rejects_unsupported_openai_destination() { let home = tempfile::TempDir::new().unwrap(); diff --git a/codex-rs/chatgpt/src/connectors.rs b/codex-rs/chatgpt/src/connectors.rs index d0c5d82907..f56286c3d7 100644 --- a/codex-rs/chatgpt/src/connectors.rs +++ b/codex-rs/chatgpt/src/connectors.rs @@ -22,11 +22,14 @@ pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_m pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_options; pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_options_and_status; pub use codex_core::connectors::list_cached_accessible_connectors_from_mcp_tools; +use codex_features::RuntimeAuthMode; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_plugin::AppConnectorId; +use codex_runtime_agent_instance::RuntimeAuthManager; use serde::Deserialize; use serde::Serialize; +use std::sync::Arc; const DIRECTORY_CONNECTORS_TIMEOUT: Duration = Duration::from_secs(60); const CONNECTOR_METADATA_TIMEOUT: Duration = Duration::from_secs(60); @@ -41,20 +44,6 @@ async fn apps_enabled(config: &Config) -> anyhow::Result { .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend))) } -async fn connector_auth(config: &Config) -> anyhow::Result { - let auth_manager = - AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await?; - let auth = auth_manager - .auth() - .await - .ok_or_else(|| anyhow::anyhow!("ChatGPT auth not available"))?; - anyhow::ensure!( - auth.uses_codex_backend(), - "ChatGPT connectors require Codex backend auth" - ); - Ok(auth) -} - pub async fn list_connectors(config: &Config) -> anyhow::Result> { if !apps_enabled(config).await? { return Ok(Vec::new()); @@ -75,19 +64,35 @@ pub async fn list_connectors(config: &Config) -> anyhow::Result> { } pub async fn list_all_connectors(config: &Config) -> anyhow::Result> { - list_all_connectors_with_options(config, /*force_refetch*/ false, &[]).await + if !apps_enabled(config).await? { + return Ok(Vec::new()); + } + let mut config = config.clone(); + let (auth, runtime) = crate::chatgpt_client::standalone_auth(&mut config).await?; + list_all_connectors_with_options( + &config, + Some(&auth), + runtime.as_ref(), + /*force_refetch*/ false, + &[], + ) + .await } -pub async fn list_cached_all_connectors( +pub fn list_cached_all_connectors( config: &Config, + auth: Option<&CodexAuth>, plugin_apps: &[AppConnectorId], ) -> Option> { - if !apps_enabled(config).await.ok()? { + if !config + .features + .apps_enabled_for_auth(auth.is_some_and(CodexAuth::uses_codex_backend)) + { return Some(Vec::new()); } - let auth = connector_auth(config).await.ok()?; - let cache_context = connector_directory_cache_context(config, &auth); + let auth = auth?; + let cache_context = connector_directory_cache_context(config, auth); let connectors = codex_connectors::cached_directory_connectors(&cache_context)?; Some(merge_directory_and_plugin_connectors( connectors.to_vec(), @@ -97,21 +102,20 @@ pub async fn list_cached_all_connectors( pub async fn list_all_connectors_with_options( config: &Config, + auth: Option<&CodexAuth>, + runtime: Option<&(Arc, RuntimeAuthMode)>, force_refetch: bool, plugin_apps: &[AppConnectorId], ) -> anyhow::Result> { - if !apps_enabled(config).await? { + if !config + .features + .apps_enabled_for_auth(auth.is_some_and(CodexAuth::uses_codex_backend)) + { return Ok(Vec::new()); } - let mut captured_config = config.clone(); - captured_config.application_network_policy = config - .application_network_policy - .clone() - .for_current_account(); - let config = &captured_config; - let auth = connector_auth(config).await?; - let cache_context = connector_directory_cache_context(config, &auth); - let auth = &auth; + let auth = auth.ok_or_else(|| anyhow::anyhow!("ChatGPT auth not available"))?; + let runtime = codex_model_provider::runtime_auth_for_url(runtime, &config.chatgpt_base_url)?; + let cache_context = connector_directory_cache_context(config, auth); let connectors = codex_connectors::list_all_connectors_with_options( cache_context, auth.is_workspace_account(), @@ -120,7 +124,7 @@ pub async fn list_all_connectors_with_options( chatgpt_get_request_with_timeout::( config, auth, - /*runtime*/ None, + runtime, path, Some(DIRECTORY_CONNECTORS_TIMEOUT), ) @@ -148,6 +152,7 @@ pub struct ConnectorMetadataReadResult { pub async fn read_connector_metadata( config: &Config, auth: &CodexAuth, + runtime: Option<&(Arc, RuntimeAuthMode)>, app_ids: &[String], include_tools: bool, ) -> anyhow::Result { @@ -160,6 +165,7 @@ pub async fn read_connector_metadata( "ChatGPT account ID not available, please re-run codex login" ); + let runtime = codex_model_provider::runtime_auth_for_url(runtime, &config.chatgpt_base_url)?; let store = ConnectorMetadataStore::new( config.chatgpt_base_url.clone(), auth.get_account_id(), @@ -188,7 +194,7 @@ pub async fn read_connector_metadata( }, CONNECTOR_METADATA_TIMEOUT, product_sku, - /*runtime*/ None, + runtime, ) .await?; let mut requested_ids = missing_ids.iter().cloned().collect::>(); diff --git a/codex-rs/codex-mcp/Cargo.toml b/codex-rs/codex-mcp/Cargo.toml index 14ab60fd67..ab46c5fec0 100644 --- a/codex-rs/codex-mcp/Cargo.toml +++ b/codex-rs/codex-mcp/Cargo.toml @@ -22,17 +22,21 @@ codex-config = { workspace = true } codex-connectors = { workspace = true } codex-diagnostics = { workspace = true } codex-exec-server = { workspace = true } +codex-features = { workspace = true } +codex-http-client = { workspace = true } codex-login = { workspace = true } codex-model-provider = { workspace = true } codex-network-proxy = { workspace = true } codex-otel = { workspace = true } codex-protocol = { workspace = true } codex-rmcp-client = { workspace = true } +codex-runtime-agent-instance = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-path-uri = { workspace = true } codex-utils-plugins = { workspace = true } codex-utils-process = { workspace = true } futures = { workspace = true } +http = { workspace = true } lru = { workspace = true } regex-lite = { workspace = true } rmcp = { workspace = true, default-features = false, features = ["base64", "macros", "schemars", "server"] } @@ -47,6 +51,7 @@ url = { workspace = true } uuid = { workspace = true, features = ["v4"] } [dev-dependencies] +base64 = { workspace = true } assert_matches = { workspace = true } codex-exec-server-test-support = { workspace = true } codex-plugin = { workspace = true } @@ -55,3 +60,5 @@ insta = { workspace = true } pretty_assertions = { workspace = true } rmcp = { workspace = true, default-features = false, features = ["base64", "macros", "schemars", "server"] } tempfile = { workspace = true } +uuid = { workspace = true } +wiremock = { workspace = true } diff --git a/codex-rs/codex-mcp/src/auth.rs b/codex-rs/codex-mcp/src/auth.rs new file mode 100644 index 0000000000..d368161cb1 --- /dev/null +++ b/codex-rs/codex-mcp/src/auth.rs @@ -0,0 +1,118 @@ +//! Authentication shared by requests on an MCP connection. + +use std::sync::Arc; + +use codex_api::AuthError; +use codex_api::SharedAuthProvider; +use codex_exec_server::RouteAwareHttpClient; +use codex_features::RuntimeAuthMode; +use codex_http_client::HttpClientFactory; +use codex_login::CodexAuth; +use codex_model_provider::ClientAuth; +use codex_rmcp_client::AuthProviderResolver; +use codex_runtime_agent_instance::RuntimeAuthSlot; +use futures::future::BoxFuture; +use http::HeaderMap; +use tokio::sync::OnceCell; + +use crate::McpRuntimeContext; + +/// A connection keeps its credential source across retries and reconnects. Initial runtime +/// selection happens inside the startup deadline; each HTTP request then selects fresh auth. +#[derive(Clone)] +pub(crate) enum McpAuthProvider { + Credentials(SharedAuthProvider), + Runtime(Arc), +} + +pub(crate) struct RuntimeMcpAuth { + slot: RuntimeAuthSlot, + mode: RuntimeAuthMode, + auth: CodexAuth, + client_auth: OnceCell, +} + +struct RuntimeMcpRequestAuth { + auth: ClientAuth, + http_client_factory: HttpClientFactory, +} + +impl McpAuthProvider { + pub(crate) fn runtime( + (slot, mode): (RuntimeAuthSlot, RuntimeAuthMode), + auth: CodexAuth, + ) -> Self { + Self::Runtime(Arc::new(RuntimeMcpAuth { + slot, + mode, + auth, + client_auth: OnceCell::new(), + })) + } + + pub(crate) async fn resolve( + &self, + mut context: McpRuntimeContext, + url: &str, + ) -> anyhow::Result<(Arc, McpRuntimeContext)> { + match self { + Self::Credentials(provider) => Ok((Arc::new(Arc::clone(provider)), context)), + Self::Runtime(runtime) => { + codex_model_provider::validate_runtime_destination(Some(url))?; + let manager = runtime.slot.get().ok_or_else(|| { + AuthError::Build("runtime authentication is not initialized".to_string()) + })?; + let auth = runtime + .client_auth + .get_or_try_init(|| async { + ClientAuth::for_auth( + &runtime.auth, + Some(&(Arc::clone(manager), runtime.mode)), + context.http_client_factory(), + &HeaderMap::new(), + ) + .await + .map(|(auth, _)| auth) + }) + .await?; + let selected = auth.resolve(&HeaderMap::new()).await?; + let factory = selected.bind_http_client_factory(context.http_client_factory())?; + context.set_local_http_client(Arc::new( + RouteAwareHttpClient::new(factory.clone()).with_tls_backend_fallback(), + )); + Ok(( + Arc::new(RuntimeMcpRequestAuth { + auth: auth.clone(), + http_client_factory: factory, + }), + context, + )) + } + } + } +} + +impl AuthProviderResolver for RuntimeMcpRequestAuth { + fn resolve<'a>( + &'a self, + headers: &'a HeaderMap, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let selected = self.auth.resolve(headers).await?; + let factory = selected.bind_http_client_factory(self.http_client_factory.clone())?; + // Reuse the connection's transport, including its original account-policy binding. + // A newly selected policy must not revive a connection that was already revoked. + if factory != self.http_client_factory { + return Err(AuthError::Build( + "runtime authentication changed the MCP connection's network policy" + .to_string(), + )); + } + Ok(selected.auth) + }) + } +} + +#[cfg(test)] +#[path = "auth_tests.rs"] +mod tests; diff --git a/codex-rs/codex-mcp/src/auth_tests.rs b/codex-rs/codex-mcp/src/auth_tests.rs new file mode 100644 index 0000000000..e8b427d888 --- /dev/null +++ b/codex-rs/codex-mcp/src/auth_tests.rs @@ -0,0 +1,203 @@ +use super::*; + +use base64::Engine; +use codex_config::types::AuthKeyringBackendKind; +use codex_exec_server::EnvironmentManager; +use codex_exec_server::HttpRedirectPolicy; +use codex_exec_server::HttpRequestParams; +use codex_features::RuntimeAuthMode; +use codex_http_client::DestinationPolicy; +use codex_http_client::HttpClientFactory; +use codex_http_client::NetworkPolicyController; +use codex_http_client::OutboundProxyPolicy; +use codex_login::AuthCredentialsStoreMode; +use codex_login::AuthManager; +use codex_login::AuthRouteConfig; +use codex_runtime_agent_instance::AgentHostAuthorization; +use codex_runtime_agent_instance::RuntimeAuthManager; +use codex_runtime_agent_instance::RuntimeRegistrationMetadata; +use codex_runtime_agent_instance::RuntimeRole; +use pretty_assertions::assert_eq; +use serde_json::json; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::path; + +struct RotatingAuth(std::sync::RwLock); + +impl codex_login::ExternalAuth for RotatingAuth { + fn resolve(&self) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(self.0.read().unwrap().clone()) }) + } + + fn refresh( + &self, + _: codex_login::ExternalAuthRefreshContext, + ) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + self.resolve() + } +} + +#[tokio::test] +async fn deferred_connection_keeps_runtime_source_and_captured_account_policy() -> anyhow::Result<()> +{ + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + Mock::given(path("/agent/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "runtime_agent_instance_id": "runtime-mcp", "agent_host_id": "host-a", + "actor": {"kind": "chatgpt", "user_id": "user-a", "account_id": "workspace-a"}, + }))) + .mount(&server) + .await; + Mock::given(path("/mcp")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server) + .await; + let home = tempfile::tempdir()?; + let oauth = |session, token| { + let claims = json!({ + "iss": "https://auth.openai.com/", "session_id": session, "jti": token, + "https://api.openai.com/auth": {"user_id": "user-a", "chatgpt_account_id": "workspace-a"} + }); + let payload = + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string()); + CodexAuth::from_external_chatgpt_tokens( + &format!("header.{payload}.signature"), + "workspace-a", + /*chatgpt_plan_type*/ None, + ) + .unwrap() + }; + let initial = oauth("session-a", "initial"); + let refreshed = oauth("session-a", "refreshed"); + let source = Arc::new(RotatingAuth(std::sync::RwLock::new(initial.clone()))); + let controller = NetworkPolicyController::default(); + let policy = controller.policy(); + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) + .with_network_policy(policy.clone()); + let credentials = Arc::new( + AuthManager::new( + home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::Ephemeral, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::Direct, + AuthRouteConfig::from_http_client_factory(factory.clone()), + ) + .await, + ); + credentials.set_external_auth(source.clone()).await?; + assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); + let runtime = ( + Arc::new(std::sync::OnceLock::from(Arc::new( + RuntimeAuthManager::new( + Arc::clone(&credentials), + AgentHostAuthorization::InstallationId { + installation_id: uuid::Uuid::nil(), + }, + RuntimeRegistrationMetadata { + harness: Some("codex".into()), + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + format!("{}/agent/register", server.uri()).parse()?, + ), + ))), + mode, + ); + let context = McpRuntimeContext::new( + Arc::new(EnvironmentManager::without_environments(factory)), + home.path().to_path_buf(), + ); + let deferred = McpAuthProvider::runtime(runtime.clone(), initial.clone()); + assert!(server.received_requests().await.unwrap().is_empty()); + let error = deferred + .resolve(context.clone(), "https://future.openai.com/mcp") + .await + .err() + .expect("unsupported OpenAI destination must fail before registration"); + assert!( + error + .downcast_ref::() + .is_some_and(|error| matches!( + error.details(), + codex_protocol::error::CodexErrorDetails::UnsupportedOperation(_) + )) + ); + assert!(server.received_requests().await.unwrap().is_empty()); + let (provider, selected_context) = deferred.resolve(context.clone(), &server.uri()).await?; + let first = provider + .resolve(&Default::default()) + .await? + .apply_auth_headers(Default::default()) + .await?; + *source.0.write().unwrap() = refreshed.clone(); + let second = provider + .resolve(&Default::default()) + .await? + .apply_auth_headers(Default::default()) + .await?; + let assertion_header = match mode { + RuntimeAuthMode::DualHeader => "x-openai-agent-assertion", + RuntimeAuthMode::AssertionOnly => "authorization", + }; + assert_ne!(first[assertion_header], second[assertion_header]); + if mode == RuntimeAuthMode::DualHeader { + assert_eq!( + first["authorization"], + format!("Bearer {}", initial.get_token()?) + ); + assert_eq!( + second["authorization"], + format!("Bearer {}", refreshed.get_token()?) + ); + } + let request = || HttpRequestParams { + method: "GET".into(), + url: format!("{}/mcp", server.uri()), + headers: Vec::new(), + body: None, + timeout_ms: Some(1000), + redirect_policy: HttpRedirectPolicy::Stop, + request_id: "mcp-policy-test".into(), + stream_response: false, + }; + let client = selected_context.local_http_client(); + assert_eq!(client.http_request(request()).await?.status, 200); + policy.invalidate(); + assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); + assert!(client.http_request(request()).await.is_err()); + // A new connection captures the new policy and reuses the same registered runtime. + let current = McpAuthProvider::runtime(runtime, refreshed); + let (_, current_context) = current.resolve(context, &server.uri()).await?; + assert_eq!( + current_context + .local_http_client() + .http_request(request()) + .await? + .status, + 200 + ); + let requests = server.received_requests().await.unwrap(); + assert_eq!( + requests + .iter() + .filter(|request| request.url.path() == "/agent/register") + .count(), + 1 + ); + *source.0.write().unwrap() = oauth("another-session", "replacement"); + assert!(provider.resolve(&Default::default()).await.is_err()); + assert_eq!( + requests + .iter() + .filter(|request| request.url.path() == "/mcp") + .count(), + 2 + ); + } + Ok(()) +} diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index edcc1a617f..1bf5651986 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -378,6 +378,7 @@ impl McpConnectionSet { client_mcp_extensions, auth, auth_manager, + runtime_auth, elicitation_reviewer, elicitation_lifecycle, } = input; @@ -546,6 +547,17 @@ impl McpConnectionSet { && should_share_codex_apps_tools_cache(&server_name, uses_env_bearer_token); let shares_codex_apps_tools_cache = eligible_for_apps_cache && codex_apps_tools_cache.shares_snapshot(); + let openai_runtime = runtime_auth.as_ref().filter(|_| { + !uses_env_bearer_token + && configured_config.is_local_environment() + && matches!(configured_config.auth, codex_config::McpServerAuth::ChatGpt) + && auth.is_some_and(codex_login::CodexAuth::uses_codex_backend) + && matches!(&configured_config.transport, + McpServerTransportConfig::StreamableHttp { url, .. } + if codex_http_client::classify_openai_api_url(url) + != codex_http_client::OpenAiDestination::Other) + }); + let runtime_auth_mode = openai_runtime.map(|(_, mode)| *mode); let codex_apps_tools_cache_context = shares_codex_apps_tools_cache.then(|| { // Only equivalent discovery inputs may share executable Apps tools. let mut transport = configured_config.transport.clone(); @@ -561,6 +573,7 @@ impl McpConnectionSet { let mut scope = serde_json::json!([ transport, &configured_config.auth, + runtime_auth_mode, protocol_mode.preferred_protocol_version().as_str(), catalog_item_limit, &client_elicitation_capability, @@ -590,7 +603,14 @@ impl McpConnectionSet { if server_name == CODEX_APPS_MCP_SERVER_NAME && uses_env_bearer_token { None } else { - chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider) + chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider).map( + |provider| match (openai_runtime, auth) { + (Some(runtime), Some(auth)) => { + crate::auth::McpAuthProvider::runtime(runtime.clone(), auth.clone()) + } + _ => crate::auth::McpAuthProvider::Credentials(provider), + }, + ) }; let discover_local_stdio = config.discover_local_stdio && server.can_discover_local_stdio(&server_name, runtime_auth_provider.as_ref()); @@ -624,6 +644,7 @@ impl McpConnectionSet { &resolved_environment, &runtime_context, runtime_auth_provider.as_ref(), + runtime_auth_mode, auth, eligible_for_apps_cache .then(|| (codex_home.clone(), codex_apps_tools_cache_key.clone())), diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 56d19c9e8c..5e4225a417 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -2154,6 +2154,7 @@ async fn shared_apps_snapshot_requires_writable_discovery_and_a_shared_cache() - McpPublicationGate::already_published(), McpRuntimeInput { network_policy_metadata: None, + runtime_auth: None, startup_policy: McpStartupPolicy::Eager, config: Arc::new(config), plugins_available: false, @@ -2234,6 +2235,7 @@ async fn hosted_apps_protocol_mode_is_independent_of_generic_mode() -> anyhow::R McpPublicationGate::already_published(), McpRuntimeInput { network_policy_metadata: None, + runtime_auth: None, startup_policy: McpStartupPolicy::Eager, config: Arc::new(config), plugins_available: false, @@ -2405,6 +2407,7 @@ async fn codex_apps_extension_does_not_share_host_owned_tools_cache() -> anyhow: McpPublicationGate::already_published(), McpRuntimeInput { network_policy_metadata: None, + runtime_auth: None, startup_policy: McpStartupPolicy::Eager, config: Arc::new(config), plugins_available: false, @@ -4973,6 +4976,7 @@ async fn executor_owned_chatgpt_mcp_accepts_only_safe_explicit_authorization() - &resolved_environment, &runtime_context, /*runtime_auth_provider*/ None, + /*runtime_auth_mode*/ None, Some(&hosted_auth), /*codex_apps_cache_identity*/ None, ElicitationCapability::default(), @@ -5020,6 +5024,7 @@ async fn executor_owned_chatgpt_mcp_accepts_only_safe_explicit_authorization() - McpPublicationGate::already_published(), McpRuntimeInput { network_policy_metadata: None, + runtime_auth: None, startup_policy: McpStartupPolicy::Eager, config: Arc::new(runtime_config.clone()), plugins_available: false, @@ -5141,6 +5146,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { McpPublicationGate::already_published(), McpRuntimeInput { network_policy_metadata: None, + runtime_auth: None, startup_policy: McpStartupPolicy::Eager, config: Arc::new(crate::mcp::tests::test_mcp_config( codex_home.path().to_path_buf(), @@ -5521,6 +5527,7 @@ pub(crate) fn reusable_server_identity( &resolved_environment, runtime_context, /*runtime_auth_provider*/ None, + /*runtime_auth_mode*/ None, /*auth*/ None, /*codex_apps_cache_identity*/ None, ElicitationCapability::default(), @@ -5603,6 +5610,7 @@ async fn reconcile_reusable_server_with_mcp_config( McpPublicationGate::already_published(), McpRuntimeInput { network_policy_metadata: None, + runtime_auth: None, startup_policy: McpStartupPolicy::Eager, config: Arc::new(mcp_config), plugins_available: false, @@ -6343,6 +6351,7 @@ fn connection_identity_tracks_only_oauth_credentials_when_oauth_is_active() { &Ok(None), &runtime_context, /*runtime_auth_provider*/ None, + /*runtime_auth_mode*/ None, /*auth*/ None, /*codex_apps_cache_identity*/ None, ElicitationCapability::default(), @@ -6425,6 +6434,7 @@ fn connection_identity_uses_effective_authorization_headers() { &Ok(None), &runtime_context, /*runtime_auth_provider*/ None, + /*runtime_auth_mode*/ None, /*auth*/ None, /*codex_apps_cache_identity*/ None, ElicitationCapability::default(), @@ -6575,6 +6585,7 @@ async fn reconciliation_replaces_connection_when_protocol_mode_changes() { McpPublicationGate::already_published(), McpRuntimeInput { network_policy_metadata: None, + runtime_auth: None, startup_policy: McpStartupPolicy::Eager, config: Arc::new(mcp_config), plugins_available: false, @@ -6639,6 +6650,7 @@ async fn reconciliation_reuses_legacy_stdio_server_when_modern_protocol_is_enabl McpPublicationGate::already_published(), McpRuntimeInput { network_policy_metadata: None, + runtime_auth: None, startup_policy: McpStartupPolicy::Eager, config: Arc::new(mcp_config), plugins_available: false, @@ -6920,6 +6932,7 @@ async fn reconciliation_reconnects_when_host_plugin_root_changes() { &resolved_environment, &runtime_context, /*runtime_auth_provider*/ None, + /*runtime_auth_mode*/ None, /*auth*/ None, /*codex_apps_cache_identity*/ None, ElicitationCapability::default(), @@ -7000,7 +7013,8 @@ async fn connection_identity_distinguishes_accounts_with_the_same_token() -> any McpOAuthRefreshMode::Legacy, &Ok(None), &runtime_context, - Some(&provider), + Some(&crate::auth::McpAuthProvider::Credentials(provider)), + /*runtime_auth_mode*/ None, Some(auth), /*codex_apps_cache_identity*/ None, ElicitationCapability::default(), @@ -7054,7 +7068,8 @@ async fn connection_identity_distinguishes_agent_account_runtime_and_task() -> a McpOAuthRefreshMode::Legacy, &Ok(None), &runtime_context, - Some(&provider), + Some(&crate::auth::McpAuthProvider::Credentials(provider)), + /*runtime_auth_mode*/ None, Some(auth), /*codex_apps_cache_identity*/ None, ElicitationCapability::default(), diff --git a/codex-rs/codex-mcp/src/event_stream.rs b/codex-rs/codex-mcp/src/event_stream.rs index 605e75726e..2075aac942 100644 --- a/codex-rs/codex-mcp/src/event_stream.rs +++ b/codex-rs/codex-mcp/src/event_stream.rs @@ -2,9 +2,9 @@ use std::sync::Arc; +use crate::auth::McpAuthProvider; use anyhow::Result; use anyhow::bail; -use codex_api::SharedAuthProvider; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; @@ -35,7 +35,7 @@ pub(crate) struct EventStreamConnectionSettings { pub oauth_refresh_mode: McpOAuthRefreshMode, pub runtime_context: McpRuntimeContext, pub resolved_environment: std::result::Result>, String>, - pub auth_provider: Option, + pub auth_provider: Option, pub auth_manager: Option>, pub auth: Option, pub protocol_mode: McpProtocolMode, diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 88c1b07555..a6a3596a53 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -118,6 +118,7 @@ pub use mcp::McpPermissionPromptAutoApproveContext; pub use mcp::mcp_permission_prompt_is_auto_approved; pub use mcp::qualified_mcp_tool_name_prefix; +mod auth; mod auth_changes; pub(crate) mod auth_elicitation; mod binding; diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 3f51fa2a1f..492d72a065 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -469,9 +469,14 @@ pub fn tool_plugin_context(config: &McpConfig) -> ToolPluginContext { ToolPluginContext::from_config(config) } +#[allow(clippy::too_many_arguments)] pub async fn read_mcp_resource( config: &McpConfig, auth: Option<&CodexAuth>, + runtime_auth: Option<( + codex_runtime_agent_instance::RuntimeAuthSlot, + codex_features::RuntimeAuthMode, + )>, runtime_context: McpRuntimeContext, codex_apps_tools_cache: ConnectorRuntimeManager, tool_catalog_cache: crate::McpToolCatalogCache, @@ -489,6 +494,7 @@ pub async fn read_mcp_resource( McpPublicationGate::already_published(), McpRuntimeInput { network_policy_metadata: None, + runtime_auth, startup_policy: McpStartupPolicy::Eager, config: Arc::new(runtime_config), plugins_available: false, @@ -534,6 +540,10 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( config: &McpConfig, auth: Option<&CodexAuth>, auth_manager: Option<&Arc>, + runtime_auth: Option<( + codex_runtime_agent_instance::RuntimeAuthSlot, + codex_features::RuntimeAuthMode, + )>, runtime_context: McpRuntimeContext, codex_apps_tools_cache: ConnectorRuntimeManager, tool_catalog_cache: crate::McpToolCatalogCache, @@ -578,6 +588,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( McpPublicationGate::already_published(), McpRuntimeInput { network_policy_metadata: None, + runtime_auth, startup_policy: McpStartupPolicy::Eager, config: Arc::new(runtime_config), plugins_available: false, diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 38acb6060e..bf3451561f 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -32,6 +32,7 @@ use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; +use crate::auth::McpAuthProvider; use crate::client_tool_catalog::ClientToolCatalog; use crate::codex_apps::normalize_codex_apps_callable_name; use crate::codex_apps::normalize_codex_apps_callable_namespace; @@ -59,7 +60,6 @@ use crate::tools::ToolInfo; use anyhow::Result; use anyhow::anyhow; use async_channel::Sender; -use codex_api::SharedAuthProvider; use codex_async_utils::CancelErr; use codex_async_utils::OrCancelExt; use codex_config::McpServerAuth; @@ -285,7 +285,7 @@ struct ManagedClientStartup { tool_catalog_cache_context: Option, runtime_context: McpRuntimeContext, resolved_environment: std::result::Result>, String>, - runtime_auth_provider: Option, + runtime_auth_provider: Option, client_elicitation_capability: ElicitationCapability, client_mcp_extensions: ClientMcpExtensions, auth_changes: Option>, @@ -490,7 +490,7 @@ impl AsyncManagedClient { tool_catalog_cache_context: Option, runtime_context: McpRuntimeContext, resolved_environment: std::result::Result>, String>, - runtime_auth_provider: Option, + runtime_auth_provider: Option, client_elicitation_capability: ElicitationCapability, client_mcp_extensions: ClientMcpExtensions, auth_changes: Option>, @@ -1429,7 +1429,7 @@ pub(crate) async fn make_rmcp_client( oauth_refresh_mode: McpOAuthRefreshMode, runtime_context: McpRuntimeContext, resolved_environment: std::result::Result>, String>, - runtime_auth_provider: Option, + runtime_auth_provider: Option, protocol_mode: McpProtocolMode, ) -> Result { let config = server.config().clone(); @@ -1535,6 +1535,13 @@ pub(crate) async fn make_rmcp_client( bearer_token_env_var, http_headers_helper: _, } => { + let (runtime_auth_provider, runtime_context) = match runtime_auth_provider { + Some(provider) => { + let (provider, context) = provider.resolve(runtime_context, &url).await?; + (Some(provider), context) + } + None => (None, runtime_context), + }; let http_client = runtime_context .http_client_for_server(server.config(), resolved_environment.as_ref()) .map_err(|error| StartupOutcomeError::from(anyhow!(error)))?; diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index e02fd78724..81ce92ec21 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -108,6 +108,10 @@ pub struct McpRuntimeInput { pub client_mcp_extensions: ClientMcpExtensions, pub auth: Option, pub auth_manager: Option>, + pub runtime_auth: Option<( + codex_runtime_agent_instance::RuntimeAuthSlot, + codex_features::RuntimeAuthMode, + )>, pub elicitation_reviewer: Option, pub elicitation_lifecycle: Option, } @@ -1077,6 +1081,14 @@ impl McpRuntimeContext { self.local_process_cwd.clone() } + pub(crate) fn http_client_factory(&self) -> codex_http_client::HttpClientFactory { + self.environment_manager.http_client_factory().clone() + } + + pub(crate) fn set_local_http_client(&mut self, client: Arc) { + self.local_http_client = client; + } + pub(crate) fn local_http_client(&self) -> Arc { Arc::clone(&self.local_http_client) } diff --git a/codex-rs/codex-mcp/src/runtime_shared_connection_tests.rs b/codex-rs/codex-mcp/src/runtime_shared_connection_tests.rs index 812bbcc491..f3009e7d80 100644 --- a/codex-rs/codex-mcp/src/runtime_shared_connection_tests.rs +++ b/codex-rs/codex-mcp/src/runtime_shared_connection_tests.rs @@ -80,6 +80,7 @@ async fn forced_refresh_replaces_pending_connection_without_interrupting_other_c })); let input = McpRuntimeInput { network_policy_metadata: None, + runtime_auth: None, startup_policy: McpStartupPolicy::Eager, config, plugins_available: false, diff --git a/codex-rs/codex-mcp/src/server.rs b/codex-rs/codex-mcp/src/server.rs index 3fa6d853b5..dffab26c55 100644 --- a/codex-rs/codex-mcp/src/server.rs +++ b/codex-rs/codex-mcp/src/server.rs @@ -4,9 +4,9 @@ use std::ffi::OsString; use std::path::PathBuf; use std::sync::Arc; +use crate::auth::McpAuthProvider; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::runtime::McpRuntimeContext; -use codex_api::SharedAuthProvider; use codex_config::AppToolApproval; use codex_config::McpServerAuth; use codex_config::McpServerConfig; @@ -92,7 +92,7 @@ impl EffectiveMcpServer { pub(crate) fn can_discover_local_stdio( &self, server_name: &str, - runtime_auth_provider: Option<&SharedAuthProvider>, + runtime_auth_provider: Option<&McpAuthProvider>, ) -> bool { server_name != CODEX_APPS_MCP_SERVER_NAME && self.config.is_local_environment() @@ -178,6 +178,7 @@ pub(crate) struct McpServerConnectionIdentity { referenced_environment_variables: Vec<(String, Option)>, runtime_auth: Option, runtime_auth_token: Option, + runtime_auth_mode: Option, codex_apps_cache_identity: Option<(PathBuf, ConnectorRuntimeContextKey)>, client_elicitation_capability: ElicitationCapability, client_mcp_extensions: ClientMcpExtensions, @@ -196,7 +197,8 @@ impl McpServerConnectionIdentity { oauth_refresh_mode: McpOAuthRefreshMode, resolved_environment: &Result>, String>, runtime_context: &McpRuntimeContext, - runtime_auth_provider: Option<&SharedAuthProvider>, + runtime_auth_provider: Option<&McpAuthProvider>, + runtime_auth_mode: Option, auth: Option<&CodexAuth>, codex_apps_cache_identity: Option<(PathBuf, ConnectorRuntimeContextKey)>, client_elicitation_capability: ElicitationCapability, @@ -306,6 +308,7 @@ impl McpServerConnectionIdentity { referenced_environment_variables, runtime_auth, runtime_auth_token, + runtime_auth_mode, codex_apps_cache_identity, client_elicitation_capability, client_mcp_extensions, @@ -345,6 +348,7 @@ impl McpServerConnectionIdentity { && self.referenced_environment_variables == other.referenced_environment_variables && same_runtime_auth && self.runtime_auth_token == other.runtime_auth_token + && self.runtime_auth_mode == other.runtime_auth_mode && self.codex_apps_cache_identity == other.codex_apps_cache_identity && self.client_elicitation_capability == other.client_elicitation_capability && self.client_mcp_extensions == other.client_mcp_extensions diff --git a/codex-rs/codex-mcp/src/server_tests.rs b/codex-rs/codex-mcp/src/server_tests.rs index 4f539e6783..3185e384e2 100644 --- a/codex-rs/codex-mcp/src/server_tests.rs +++ b/codex-rs/codex-mcp/src/server_tests.rs @@ -76,6 +76,7 @@ fn connection_identity( &Ok(None), runtime_context, /*runtime_auth_provider*/ None, + /*runtime_auth_mode*/ None, /*auth*/ None, /*codex_apps_cache_identity*/ None, ElicitationCapability::default(), diff --git a/codex-rs/codex-mcp/src/trusted_access.rs b/codex-rs/codex-mcp/src/trusted_access.rs index 4191b127b1..c5c10d39d2 100644 --- a/codex-rs/codex-mcp/src/trusted_access.rs +++ b/codex-rs/codex-mcp/src/trusted_access.rs @@ -116,6 +116,7 @@ pub struct TrustedAccessContext { auth_manager: Arc, chatgpt_base_url: String, http_client: Arc, + runtime_auth: Option<(crate::auth::McpAuthProvider, crate::McpRuntimeContext)>, } impl TrustedAccessContext { @@ -127,12 +128,26 @@ impl TrustedAccessContext { ) { return None; } - Some(Self::new( + let mut context = Self::new( auth.clone(), input.auth_manager.clone()?, input.config.chatgpt_base_url.clone(), input.runtime_context.local_http_client(), - )) + ); + context.runtime_auth = input + .runtime_auth + .as_ref() + .filter(|_| { + codex_http_client::classify_openai_api_url(&input.config.chatgpt_base_url) + != codex_http_client::OpenAiDestination::Other + }) + .map(|runtime| { + ( + crate::auth::McpAuthProvider::runtime(runtime.clone(), auth.clone()), + input.runtime_context.clone(), + ) + }); + Some(context) } pub fn new( @@ -146,6 +161,7 @@ impl TrustedAccessContext { auth_manager, chatgpt_base_url, http_client, + runtime_auth: None, } } @@ -203,8 +219,25 @@ impl TrustedAccessContext { return None; } - let headers = codex_model_provider::auth_provider_from_auth(&auth) - .to_auth_headers() + let (provider, http_client) = match &self.runtime_auth { + Some((resolver, context)) => { + let (resolver, context) = resolver + .resolve(context.clone(), &self.chatgpt_base_url) + .await + .ok()?; + ( + resolver.resolve(&Default::default()).await.ok()?, + context.local_http_client(), + ) + } + None => ( + codex_model_provider::auth_provider_from_auth(&auth), + Arc::clone(&self.http_client), + ), + }; + let headers = provider + .apply_auth_headers(Default::default()) + .await .ok()? .iter() .map(|(name, value)| { @@ -215,8 +248,7 @@ impl TrustedAccessContext { }) }) .collect::>>()?; - let (response, mut response_body) = self - .http_client + let (response, mut response_body) = http_client .http_request_stream(HttpRequestParams { method: "GET".to_string(), url: format!( diff --git a/codex-rs/core-plugins/src/manager.rs b/codex-rs/core-plugins/src/manager.rs index 8558837809..782be99858 100644 --- a/codex-rs/core-plugins/src/manager.rs +++ b/codex-rs/core-plugins/src/manager.rs @@ -730,6 +730,14 @@ impl PluginsManager { ) } + pub fn auth_manager(&self) -> Arc { + Arc::clone(&self.auth_manager) + } + + pub fn runtime_auth_slot(&self) -> RuntimeAuthSlot { + Arc::clone(&self.runtime_auth) + } + /// Reuse host configuration without sharing another actor's remote caches or credentials. pub fn for_auth_context( &self, diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index f3e263d997..a581b78ac1 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -1197,6 +1197,11 @@ impl CodexThread { .collect() } + /// The thread's MCP manager also owns its actor-scoped plugin and runtime auth state. + pub fn mcp_manager(&self) -> Arc { + Arc::clone(&self.session.services.mcp_manager) + } + pub async fn config(&self) -> Arc { self.session.get_config().await } diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 7ec46339cc..04bb4987aa 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -112,21 +112,17 @@ pub(crate) async fn list_tool_suggest_discoverable_tools_with_auth( .collect()) } -pub async fn list_cached_accessible_connectors_from_mcp_tools( +pub fn list_cached_accessible_connectors_from_mcp_tools( config: &Config, + auth: Option<&CodexAuth>, ) -> Option> { - let auth_manager = - AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false) - .await - .ok()?; - let auth = auth_manager.auth().await; if !config .features - .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend)) + .apps_enabled_for_auth(auth.is_some_and(CodexAuth::uses_codex_backend)) { return Some(Vec::new()); } - let cache_key = accessible_connectors_cache_key(config, auth.as_ref()); + let cache_key = accessible_connectors_cache_key(config, auth); read_cached_accessible_connectors(&cache_key) } @@ -205,8 +201,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( environment_manager: Arc, mcp_manager: Arc, ) -> anyhow::Result { - let auth_manager = - AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await?; + let auth_manager = mcp_manager.auth_manager(); let auth = auth_manager.auth().await; if !config .features @@ -248,6 +243,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( .then(|| Arc::clone(&auth_manager)); let mcp_runtime = McpRuntime::new(McpRuntimeInput { network_policy_metadata: None, + runtime_auth: mcp_manager.runtime_auth_for_config(config), startup_policy: McpStartupPolicy::Eager, config: Arc::clone(&mcp_config), plugins_available: false, diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index eba1bd977c..ce3f0b26ad 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -156,6 +156,32 @@ impl McpManager { manager } + pub fn plugins_manager(&self) -> Arc { + Arc::clone(&self.plugins_manager) + } + + pub fn auth_manager(&self) -> Arc { + self.plugins_manager.auth_manager() + } + + pub fn runtime_auth_for_config( + &self, + config: &Config, + ) -> Option<( + codex_runtime_agent_instance::RuntimeAuthSlot, + codex_features::RuntimeAuthMode, + )> { + config + .features + .enabled(Feature::RuntimeAgentInstances) + .then(|| { + ( + self.plugins_manager.runtime_auth_slot(), + config.runtime_auth_mode, + ) + }) + } + pub fn codex_apps_tools_cache(&self) -> ConnectorRuntimeManager { self.codex_apps_tools_cache.clone() } diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index 05e7be9334..99ccfaadc7 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -368,6 +368,16 @@ impl Session { self.mcp_runtime_context(&desired.environments, &desired.local_process_cwd); McpRuntimeInput { network_policy_metadata: None, + runtime_auth: desired + .config + .features + .enabled(Feature::RuntimeAgentInstances) + .then(|| { + ( + Arc::clone(&self.services.runtime_auth), + desired.config.runtime_auth_mode, + ) + }), startup_policy: if matches!(desired.session_source, SessionSource::SubAgent(_)) { McpStartupPolicy::LazyWhenCached } else { diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 1bb4e51d2b..8d8ed5762c 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1200,9 +1200,10 @@ async fn required_mcp_servers_for_input( })); let connector_slug_counts = if turn_context.apps_enabled() && !mentions.plain_names.is_empty() { - let cached_connectors = - connectors::list_cached_accessible_connectors_from_mcp_tools(&turn_context.config) - .await; + let cached_connectors = connectors::list_cached_accessible_connectors_from_mcp_tools( + &turn_context.config, + sess.services.auth_manager.auth().await.as_ref(), + ); let accessible_connectors = match cached_connectors { Some(connectors) => connectors, None => sess diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 5002100b71..19d93d5aca 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -915,6 +915,21 @@ impl ThreadManager { .map(|runtime| (runtime, config.runtime_auth_mode))) } + /// Initialize the selected thread owner's runtime, including after a feature enablement. + pub fn runtime_auth_for_mcp_manager( + &self, + config: &Config, + mcp_manager: &McpManager, + ) -> std::io::Result, codex_features::RuntimeAuthMode)>> { + Ok(crate::runtime_auth::runtime_for_config( + &mcp_manager.plugins_manager().runtime_auth_slot(), + config, + mcp_manager.auth_manager(), + &self.state.installation_id, + ) + .map(|runtime| (runtime, config.runtime_auth_mode))) + } + pub fn skills_service(&self) -> Arc { self.state.skills_service.clone() } @@ -2611,10 +2626,14 @@ impl ThreadManagerState { codex_sandboxing::WindowsSandboxProxySettingsMode::Reconcile }; // A resumed actor and its descendants share an owner, not the harness default's credentials. - let (runtime_auth, plugins_manager) = if Arc::ptr_eq(&auth_manager, &self.auth_manager) { + let (runtime_auth, plugins_manager, analytics_events_client) = if Arc::ptr_eq( + &auth_manager, + &self.auth_manager, + ) { ( inherited_runtime_auth.unwrap_or_else(|| Arc::clone(&self.runtime_auth)), Arc::clone(&self.plugins_manager), + self.analytics_events_client.clone(), ) } else { let threads = self.threads.read().await; @@ -2627,6 +2646,7 @@ impl ThreadManagerState { inherited_runtime_auth .unwrap_or_else(|| Arc::clone(&parent.session.services.runtime_auth)), Arc::clone(&parent.session.services.plugins_manager), + Some(parent.session.services.analytics_events_client.clone()), ) } else { let runtime_auth = inherited_runtime_auth.unwrap_or_default(); @@ -2634,7 +2654,35 @@ impl ThreadManagerState { self.plugins_manager .for_auth_context(Arc::clone(&auth_manager), Arc::clone(&runtime_auth)), ); - (runtime_auth, plugins) + let runtime_provider = config + .features + .enabled(Feature::RuntimeAgentInstances) + .then(|| (Arc::clone(&runtime_auth), config.runtime_auth_mode)); + let analytics = self.analytics_events_client.as_ref().map(|client| { + let api_key_headers = if codex_model_provider::provider_uses_first_party_auth_path( + &config.model_provider, + ) { + match config.model_provider.build_header_map() { + Ok(headers) => headers, + Err(error) => { + tracing::warn!(%error, "failed to resolve analytics resource selectors"); + return codex_analytics::AnalyticsEventsClient::disabled(); + } + } + } else { + Default::default() + }; + client.for_auth_context( + Arc::clone(&auth_manager), + config.chatgpt_base_url.clone(), + runtime_provider, + api_key_headers, + ) + }); + if let Some(client) = &analytics { + plugins.set_analytics_events_client(client.clone()); + } + (runtime_auth, plugins, analytics) } }; let mcp_manager = if Arc::ptr_eq(&plugins_manager, &self.plugins_manager) { @@ -2708,7 +2756,7 @@ impl ThreadManagerState { turn_extension_init, client_mcp_extensions, reserved_thread_id, - analytics_events_client: self.analytics_events_client.clone(), + analytics_events_client, image_store: Arc::clone(&self.image_store), thread_store: Arc::clone(&self.thread_store), attestation_provider: self.attestation_provider.clone(), diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index ffb2463c00..e0a741a263 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -342,6 +342,133 @@ async fn fixed_thread_text_supersedes_the_live_provider_and_isolated_helpers_rej } } +#[test_case::test_case(None, "chatgpt.com", 1; "ordinary_openai")] +#[test_case::test_case(Some(codex_features::RuntimeAuthMode::DualHeader), "chatgpt.com", 0; "dual_header_openai")] +#[test_case::test_case(Some(codex_features::RuntimeAuthMode::AssertionOnly), "chatgpt.com", 0; "assertion_only_openai")] +#[test_case::test_case(Some(codex_features::RuntimeAuthMode::DualHeader), "example.test", 1; "dual_header_custom_backend")] +#[test_case::test_case(Some(codex_features::RuntimeAuthMode::AssertionOnly), "example.test", 1; "assertion_only_custom_backend")] +#[tokio::test] +async fn resumed_analytics_validates_openai_destination( + mode: Option, + host: &str, + expected_requests: usize, +) { + let proxy = MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::path("/codex/analytics-events/events")) + .respond_with(wiremock::ResponseTemplate::new(200)) + .mount(&proxy) + .await; + let home = tempdir().expect("tempdir"); + let mut config = test_config().await; + config.codex_home = home.path().abs(); + config.cwd = home.path().abs(); + config.chatgpt_base_url = format!("http://{host}:{}", proxy.address().port()); + config.model_provider.base_url = Some(proxy.uri()); + config.respect_system_proxy = true; + config.cli_auth_credentials_store_mode = codex_login::AuthCredentialsStoreMode::File; + config.analytics_enabled = Some(true); + if let Some(mode) = mode { + config + .features + .enable(Feature::RuntimeAgentInstances) + .unwrap(); + config.runtime_auth_mode = mode; + } else { + config + .features + .disable(Feature::RuntimeAgentInstances) + .unwrap(); + } + codex_http_client::cache_system_proxy_route_for_test( + &format!("{}/codex/analytics-events/events", config.chatgpt_base_url), + proxy.uri(), + ); + codex_login::login_with_api_key( + &config.codex_home, + "resumed-key", + codex_login::AuthCredentialsStoreMode::File, + codex_login::AuthKeyringBackendKind::default(), + ) + .unwrap(); + let resumed_auth = + AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false) + .await + .unwrap(); + let original_auth = CodexAuth::from_api_key("startup-key"); + let mut manager = ThreadManager::with_models_provider_and_home_for_tests( + original_auth.clone(), + config.model_provider.clone(), + config.codex_home.to_path_buf(), + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + ); + let startup_analytics = AnalyticsEventsClient::new( + AuthManager::from_auth_for_testing(original_auth), + proxy.uri(), + /*analytics_enabled*/ Some(true), + ); + Arc::get_mut(&mut manager.state) + .expect("unshared thread manager state") + .analytics_events_client = Some(startup_analytics); + + let original = manager + .start_thread(StartThreadOptions { + history_mode: Some(ThreadHistoryMode::Legacy), + ..StartThreadOptions::new(config.clone()) + }) + .await + .expect("start original thread"); + original.thread.ensure_rollout_materialized().await; + original.thread.flush_rollout().await.unwrap(); + original.thread.shutdown_and_wait().await.unwrap(); + manager.remove_thread(&original.thread_id).await; + + let resumed = manager + .resume_thread_with_history( + config, + InitialHistory::Resumed(ResumedHistory { + history_revision: None, + conversation_id: original.thread_id, + history: Arc::new(Vec::new()), + rollout_path: original.thread.rollout_path(), + }), + resumed_auth, + /*parent_trace*/ None, + ClientMcpExtensions::default(), + ) + .await + .expect("resume with another auth owner"); + let analytics = &resumed.thread.session.services.analytics_events_client; + analytics.track_plugin_used( + codex_analytics::TrackEventsContext { + turn_metadata: None, + model_slug: "test-model".to_string(), + turn_id: "test-turn".to_string(), + thread_id: resumed.thread_id.to_string(), + product_client_id: "codex_work_cca".to_string(), + }, + codex_plugin::PluginTelemetryMetadata { + plugin_id: Some(codex_plugin::PluginId::parse("test@local").expect("valid plugin ID")), + remote_plugin_id: None, + capability_summary: None, + }, + ); + analytics.flush().await; + resumed.thread.shutdown_and_wait().await.unwrap(); + + let requests: Vec<_> = proxy + .received_requests() + .await + .unwrap() + .into_iter() + .filter(|request| request.url.path() == "/codex/analytics-events/events") + .collect(); + assert_eq!(requests.len(), expected_requests); + for request in requests { + assert_eq!(request.headers["authorization"], "Bearer resumed-key"); + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + } +} + /// A thread opt-out wins over a shared client without disabling its siblings. #[tokio::test] async fn thread_analytics_opt_out_overrides_shared_client() { diff --git a/codex-rs/core/tests/common/apps_test_server.rs b/codex-rs/core/tests/common/apps_test_server.rs index 7cda604de0..3d090cbdea 100644 --- a/codex-rs/core/tests/common/apps_test_server.rs +++ b/codex-rs/core/tests/common/apps_test_server.rs @@ -37,8 +37,8 @@ const DISCOVERABLE_CALENDAR_ID: &str = "connector_2128aebfecb84f64a069897515042a const DISCOVERABLE_GMAIL_ID: &str = "connector_68df038e0ba48191908c8434991bbac2"; const CONNECTOR_DESCRIPTION: &str = "Plan events and manage your calendar."; const CODEX_APPS_META_KEY: &str = "_codex_apps"; -const CODEX_APPS_MCP_PATH_REGEX: &str = "^/api/codex/ps/mcp/?$"; -const HOSTED_PLUGIN_RUNTIME_MCP_PATH_REGEX: &str = "^/api/codex/ps/mcp/?$"; +const CODEX_APPS_MCP_PATH_REGEX: &str = "^/(?:api/codex|backend-api)/ps/mcp/?$"; +const HOSTED_PLUGIN_RUNTIME_MCP_PATH_REGEX: &str = "^/(?:api/codex|backend-api)/ps/mcp/?$"; const PROTOCOL_VERSION: &str = "2025-11-25"; const SERVER_NAME: &str = "codex-apps-test"; const SERVER_VERSION: &str = "1.0.0"; diff --git a/codex-rs/core/tests/suite/mcp_auth_refresh.rs b/codex-rs/core/tests/suite/mcp_auth_refresh.rs index a21b0c5272..baaa61aa28 100644 --- a/codex-rs/core/tests/suite/mcp_auth_refresh.rs +++ b/codex-rs/core/tests/suite/mcp_auth_refresh.rs @@ -92,6 +92,7 @@ async fn hosted_plugin_runtime_ps_mcp_tool_calls_use_current_auth_manager_token( let mcp_config = Arc::new(config.to_mcp_config(&plugins_manager).await); let runtime = McpRuntime::new(McpRuntimeInput { network_policy_metadata: None, + runtime_auth: None, startup_policy: McpStartupPolicy::Eager, config: mcp_config, plugins_available: false, diff --git a/codex-rs/core/tests/suite/runtime_auth.rs b/codex-rs/core/tests/suite/runtime_auth.rs index 51f417aff8..1760eb4d21 100644 --- a/codex-rs/core/tests/suite/runtime_auth.rs +++ b/codex-rs/core/tests/suite/runtime_auth.rs @@ -151,82 +151,7 @@ async fn harness_shares_runtime_and_honors_thread_auth_config() -> Result<()> { .mount(&server) .await; - // The production destination check stays intact. A test-only CA and CONNECT proxy route - // the real host names to wiremock, only inside this isolated test subprocess. - codex_utils_rustls_provider::ensure_rustls_crypto_provider(); - let rcgen::CertifiedKey { cert, signing_key } = rcgen::generate_simple_self_signed(vec![ - "api.openai.com".into(), - "auth.openai.com".into(), - "chatgpt.com".into(), - "us.chatgpt.com".into(), - ])?; - let ca_dir = tempfile::tempdir()?; - let ca_path = ca_dir.path().join("ca.pem"); - std::fs::write(&ca_path, cert.pem())?; - let acceptor = TlsAcceptor::from(Arc::new( - rustls::ServerConfig::builder() - .with_no_client_auth() - .with_single_cert(vec![cert.der().clone()], signing_key.into())?, - )); - let proxy = TcpListener::bind("127.0.0.1:0").await?; - let proxy_url = format!("http://{}", proxy.local_addr()?); - let target = *server.address(); - let proxy_task = tokio::spawn(async move { - let mut connections = JoinSet::new(); - loop { - let (mut stream, _) = proxy.accept().await?; - let acceptor = acceptor.clone(); - connections.spawn(async move { - let mut headers = Vec::new(); - while !headers.ends_with(b"\r\n\r\n") && headers.len() < 8192 { - headers.push(stream.read_u8().await?); - } - let headers = std::str::from_utf8(&headers)?; - anyhow::ensure!( - matches!( - headers.lines().next(), - Some( - "CONNECT api.openai.com:443 HTTP/1.1" - | "CONNECT auth.openai.com:443 HTTP/1.1" - | "CONNECT chatgpt.com:443 HTTP/1.1" - | "CONNECT us.chatgpt.com:443 HTTP/1.1" - ) - ), - "unexpected proxy destination" - ); - stream - .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") - .await?; - let mut stream = acceptor.accept(stream).await?; - let mut target = TcpStream::connect(target).await?; - tokio::io::copy_bidirectional(&mut stream, &mut target).await?; - Ok::<_, anyhow::Error>(()) - }); - } - #[allow(unreachable_code)] - Ok::<_, anyhow::Error>(()) - }); - let mut command = Command::new(std::env::current_exe()?); - command - .args(["--exact", TEST_NAME, "--nocapture"]) - .kill_on_drop(true); - for &key in codex_network_proxy::PROXY_ENV_KEYS { - command.env_remove(key); - } - command - .env(PROBE_SERVER, server.uri()) - .env("HTTPS_PROXY", proxy_url) - .env("NO_PROXY", codex_network_proxy::DEFAULT_NO_PROXY_VALUE) - .env("CODEX_CA_CERTIFICATE", ca_path); - let output = tokio::time::timeout(Duration::from_secs(90), command.output()).await; - proxy_task.abort(); - let output = output??; - assert!( - output.status.success(), - "harness subprocess failed\n{}\n{}", - String::from_utf8_lossy(&output.stdout), - String::from_utf8_lossy(&output.stderr) - ); + run_harness_through_proxy(&server, TEST_NAME).await?; let requests = responses.requests(); let expected = [ @@ -728,3 +653,438 @@ async fn upload_with_harness( assert_eq!(uploaded.uri, "sediment://file-a"); Ok(()) } + +async fn run_harness_through_proxy(server: &MockServer, test_name: &str) -> Result<()> { + // The production destination check stays intact. A test-only CA and CONNECT proxy route + // the real host names to wiremock, only inside this isolated test subprocess. + codex_utils_rustls_provider::ensure_rustls_crypto_provider(); + let rcgen::CertifiedKey { cert, signing_key } = rcgen::generate_simple_self_signed(vec![ + "api.openai.com".into(), + "auth.openai.com".into(), + "chatgpt.com".into(), + "us.chatgpt.com".into(), + ])?; + let ca_dir = tempfile::tempdir()?; + let ca_path = ca_dir.path().join("ca.pem"); + std::fs::write(&ca_path, cert.pem())?; + let acceptor = TlsAcceptor::from(Arc::new( + rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(vec![cert.der().clone()], signing_key.into())?, + )); + let proxy = TcpListener::bind("127.0.0.1:0").await?; + let proxy_url = format!("http://{}", proxy.local_addr()?); + let target = *server.address(); + let proxy_task = tokio::spawn(async move { + let mut connections = JoinSet::new(); + loop { + let (mut stream, _) = proxy.accept().await?; + let acceptor = acceptor.clone(); + connections.spawn(async move { + let mut headers = Vec::new(); + while !headers.ends_with(b"\r\n\r\n") && headers.len() < 8192 { + headers.push(stream.read_u8().await?); + } + let headers = std::str::from_utf8(&headers)?; + anyhow::ensure!( + matches!( + headers.lines().next(), + Some( + "CONNECT api.openai.com:443 HTTP/1.1" + | "CONNECT auth.openai.com:443 HTTP/1.1" + | "CONNECT chatgpt.com:443 HTTP/1.1" + | "CONNECT us.chatgpt.com:443 HTTP/1.1" + ) + ), + "unexpected proxy destination" + ); + stream + .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") + .await?; + let mut stream = acceptor.accept(stream).await?; + let mut target = TcpStream::connect(target).await?; + tokio::io::copy_bidirectional(&mut stream, &mut target).await?; + Ok::<_, anyhow::Error>(()) + }); + } + #[allow(unreachable_code)] + Ok::<_, anyhow::Error>(()) + }); + let mut command = Command::new(std::env::current_exe()?); + command + .args(["--exact", test_name, "--nocapture"]) + .kill_on_drop(true); + for &key in codex_network_proxy::PROXY_ENV_KEYS { + command.env_remove(key); + } + command + .env(PROBE_SERVER, server.uri()) + .env("HTTPS_PROXY", proxy_url) + .env("NO_PROXY", codex_network_proxy::DEFAULT_NO_PROXY_VALUE) + .env("CODEX_CA_CERTIFICATE", ca_path); + let output = tokio::time::timeout(Duration::from_secs(90), command.output()).await; + proxy_task.abort(); + let output = output??; + assert!( + output.status.success(), + "harness subprocess failed\n{}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn mcp_uploads_follow_a_resumed_actors_runtime() -> Result<()> { + use core_test_support::apps_test_server::*; + skip_if_no_network!(Ok(())); + if std::env::var(PROBE_SERVER).is_ok() { + return exercise_mcp_uploads().await; + } + let server = MockServer::start().await; + let other_origin = MockServer::start().await; + let _apps = AppsTestServer::mount(&server).await?; + Mock::given(path("/agent/register")) + .respond_with(|request: &wiremock::Request| { + let body: Value = request.body_json().unwrap(); + let credential = body["actor_auth"]["credential"].as_str().unwrap(); + let claims: Value = serde_json::from_slice( + &base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(credential.split('.').nth(1).unwrap()) + .unwrap(), + ) + .unwrap(); + let user = claims["https://api.openai.com/auth"]["user_id"] + .as_str() + .unwrap(); + ResponseTemplate::new(200).set_body_json(json!({ + "runtime_agent_instance_id": format!("runtime-{user}"), "agent_host_id": "host-a", + "actor": {"kind": "chatgpt", "user_id": user, "account_id": "account_id"}, + })) + }) + .expect(2) + .mount(&server) + .await; + Mock::given(path("/files")) + .respond_with( + ResponseTemplate::new(307).insert_header("location", "/files/create-redirect"), + ) + .expect(3) + .mount(&server) + .await; + Mock::given(path("/files/create-redirect")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "file_id": "file-runtime", "upload_url": format!("{}/blob", server.uri()), + }))) + .expect(3) + .mount(&server) + .await; + Mock::given(path("/files/file-runtime/uploaded")) + .respond_with( + ResponseTemplate::new(307) + .insert_header("location", "/files/file-runtime/uploaded-redirect"), + ) + .expect(3) + .mount(&server) + .await; + Mock::given(path("/files/file-runtime/uploaded-redirect")) + .respond_with( + ResponseTemplate::new(307) + .insert_header("location", format!("{}/finalize", other_origin.uri())), + ) + .expect(3) + .mount(&server) + .await; + Mock::given(path("/finalize")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "status": "success", "download_url": format!("{}/blob", server.uri()), + "file_name": "report.txt", "file_size_bytes": 6, + }))) + .expect(3) + .mount(&other_origin) + .await; + Mock::given(path("/blob")) + .and(method("PUT")) + .respond_with(ResponseTemplate::new(200)) + .expect(3) + .mount(&server) + .await; + let events = (0..3) + .flat_map(|_| { + [ + responses::sse(vec![ + responses::ev_tool_search_call( + "search-file", + &json!({"query": "extract text from uploaded document", "limit": 1}), + ), + responses::ev_completed("search"), + ]), + responses::sse(vec![ + responses::ev_function_call_with_namespace( + "extract-file", + SEARCH_CALENDAR_NAMESPACE, + SEARCH_CALENDAR_EXTRACT_TEXT_TOOL, + &json!({"file": "report.txt"}).to_string(), + ), + responses::ev_completed("extract"), + ]), + responses::sse(vec![ + responses::ev_assistant_message("done", "done"), + responses::ev_completed("done"), + ]), + ] + }) + .collect(); + let responses = responses::mount_sse_sequence(&server, events).await; + run_harness_through_proxy( + &server, + "suite::runtime_auth::mcp_uploads_follow_a_resumed_actors_runtime", + ) + .await?; + assert_eq!(responses.requests().len(), 9); + let requests = server.received_requests().await.unwrap(); + let mut proof_ids = BTreeSet::new(); + let mut runtimes = BTreeSet::new(); + let mut uploaded_calls = 0; + let mut file_requests = 0; + for request in &requests { + let path = request.url.path().trim_end_matches('/'); + if path == "/blob" { + assert!(!request.headers.contains_key("authorization")); + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + assert_eq!(request.body, b"report"); + continue; + } + if !(path.ends_with("/responses") + || path == "/backend-api/ps/mcp" + || path.starts_with("/files")) + { + continue; + } + file_requests += usize::from(path.starts_with("/files")); + let authorization = request.headers["authorization"].to_str()?; + let assertion_only = authorization.starts_with("AgentAssertion "); + let proof = if let Some(proof) = authorization.strip_prefix("AgentAssertion ") { + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + proof + } else { + assert_eq!( + authorization, + format!("Bearer {}", mcp_actor_token("user-a")) + ); + request.headers["x-openai-agent-assertion"].to_str()? + }; + let claims: Value = serde_json::from_slice( + &base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(proof.split('.').nth(2).unwrap())?, + )?; + let runtime = claims["runtime_agent_instance"].as_str().unwrap(); + assert!(matches!(runtime, "runtime-user-a" | "runtime-user-b")); + if !assertion_only { + assert_eq!(runtime, "runtime-user-a"); + } + let proof_id = claims["id"].as_str().unwrap(); + if path.ends_with("-redirect") { + assert!( + proof_ids.contains(proof_id), + "redirect preserves the request proof" + ); + } else { + assert!(proof_ids.insert(proof_id.to_owned())); + } + runtimes.insert(runtime.to_owned()); + if path == "/backend-api/ps/mcp" + && let Ok(body) = request.body_json::() + && body["method"] == "tools/call" + && body["params"]["name"] == CALENDAR_EXTRACT_TEXT_TOOL_NAME + { + assert_eq!( + body["params"]["arguments"]["file"]["file_id"], + "file-runtime" + ); + assert_eq!(assertion_only, uploaded_calls != 0); + assert_eq!( + runtime, + if uploaded_calls < 2 { + "runtime-user-a" + } else { + "runtime-user-b" + } + ); + uploaded_calls += 1; + } + } + assert_eq!(runtimes.len(), 2); + let rpc_calls: Vec<_> = requests + .iter() + .map(|request| { + let body = request.body_json::().unwrap_or_default(); + ( + request.url.path(), + body["method"].clone(), + body["params"]["name"].clone(), + ) + }) + .collect(); + assert_eq!( + uploaded_calls, 3, + "recorded RPC calls: {rpc_calls:?}; file control requests: {file_requests}" + ); + assert_eq!(file_requests, 12); + let redirected = other_origin.received_requests().await.unwrap(); + assert_eq!(redirected.len(), 3); + for request in redirected { + assert!( + !request.headers.contains_key("authorization"), + "cross-origin finalization must not receive authorization", + ); + assert!( + !request.headers.contains_key("x-openai-agent-assertion"), + "cross-origin finalization must not receive the runtime assertion", + ); + } + other_origin.verify().await; + server.verify().await; + Ok(()) +} + +fn mcp_actor_token(user: &str) -> String { + let claims = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode( + json!({ + "iss": "https://auth.openai.com/", "session_id": format!("login-{user}"), + "https://api.openai.com/auth": {"user_id": user, "chatgpt_account_id": "account_id"}, + }) + .to_string(), + ); + format!("header.{claims}.signature") +} + +async fn exercise_mcp_uploads() -> Result<()> { + use core_test_support::apps_test_server::apps_enabled_builder; + let server = MockServer::start().await; + let actor = |user| { + CodexAuth::from_external_chatgpt_tokens( + &mcp_actor_token(user), + "account_id", + /*chatgpt_plan_type*/ None, + ) + .expect("valid test actor") + }; + let mut test = apps_enabled_builder("https://chatgpt.com") + .with_auth(actor("user-a")) + .with_config(|config| { + config.model_provider.base_url = Some("https://api.openai.com/v1".into()); + config.model_provider.env_http_headers = None; + core_test_support::apps_test_server::configure_search_capable_model(config); + config + .features + .enable(Feature::RuntimeAgentInstances) + .expect("enable runtime auth"); + }) + .with_workspace_setup(|cwd, fs| async move { + let report = core_test_support::test_codex::executor_path_uri(cwd.join("report.txt"))?; + fs.write_file( + &report, + b"report".to_vec(), + Default::default(), + /*sandbox*/ None, + ) + .await?; + Ok(()) + }) + .build_with_auto_env(&server) + .await?; + test.submit_turn_with_environments( + "Extract the report text with the app tool.", + Some(vec![test.executor_environment().selection().clone()]), + ) + .await?; + // A second thread uses the same actor and MCP manager but must not inherit dual-header auth. + let mut assertion_only = test.config.clone(); + assertion_only.runtime_auth_mode = RuntimeAuthMode::AssertionOnly; + let assertion_thread = test + .thread_manager + .start_thread(StartThreadOptions::new(assertion_only)) + .await?; + let thread_id = assertion_thread.thread_id; + test.codex = assertion_thread.thread; + test.submit_turn_with_environments( + "Extract the report text with the app tool.", + Some(vec![test.executor_environment().selection().clone()]), + ) + .await?; + let rollout = test.codex.rollout_path().context("saved thread")?; + test.codex.submit(Op::Shutdown).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::ShutdownComplete) + }) + .await; + test.thread_manager.remove_thread(&thread_id).await; + let context = test + .thread_store + .load_latest_model_context(codex_thread_store::LoadThreadHistoryParams { + thread_id, + include_archived: false, + }) + .await?; + let mut config = test.config.clone(); + config.runtime_auth_mode = RuntimeAuthMode::AssertionOnly; + let auth = codex_core::test_support::auth_manager_from_auth_with_home( + actor("user-b"), + test.codex_home_path().to_path_buf(), + ); + let resumed = test + .thread_manager + .resume_thread_with_history( + config, + InitialHistory::Resumed(ResumedHistory { + history_revision: None, + conversation_id: context.thread_id, + history: Arc::new(context.items), + rollout_path: Some(rollout), + }), + auth, + /*parent_trace*/ None, + ClientMcpExtensions::default(), + ) + .await?; + test.codex = resumed.thread; + test.submit_turn_with_environments( + "Extract the report text with the app tool.", + Some(vec![test.executor_environment().selection().clone()]), + ) + .await?; + // Custom MCP hosts use their own authentication, without ambient OpenAI credentials. + let external = MockServer::start().await; + let _apps = core_test_support::apps_test_server::AppsTestServer::mount(&external).await?; + let mut custom = test.config.clone(); + custom.chatgpt_base_url = external.uri(); + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + custom.runtime_auth_mode = mode; + let apps = + codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_mcp_manager( + &custom, + /*force_refetch*/ true, + test.thread_manager.environment_manager(), + test.codex.mcp_manager(), + ) + .await?; + assert!(apps.codex_apps_ready); + assert!(!apps.connectors.is_empty()); + } + let requests = external + .received_requests() + .await + .context("custom MCP destination should record requests")?; + assert!(requests.iter().any(|request| { + request + .body_json::() + .is_ok_and(|body| body["method"] == "tools/list") + })); + for request in requests { + assert!(!request.headers.contains_key("authorization")); + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + } + Ok(()) +} diff --git a/codex-rs/core/tests/suite/scenarios_mcp_sharing.rs b/codex-rs/core/tests/suite/scenarios_mcp_sharing.rs index 54b7aebe95..5b7d2ec1d6 100644 --- a/codex-rs/core/tests/suite/scenarios_mcp_sharing.rs +++ b/codex-rs/core/tests/suite/scenarios_mcp_sharing.rs @@ -1,11 +1,16 @@ //! Shared MCP transports preserve each chat's model-visible tools and input responses. use super::*; +use codex_history::InitialHistory; +use codex_history::ResumedHistory; use codex_protocol::approvals::ElicitationAction; use codex_protocol::approvals::ElicitationRequest; use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::mcp::ClientMcpExtensions; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::InternalSessionSource; +use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadSettingsOverrides; use core_test_support::responses; use pretty_assertions::assert_eq; @@ -16,10 +21,26 @@ async fn shared_mcp_keeps_tools_and_model_requested_input_per_chat() -> Result<( core_test_support::skip_if_remote!(Ok(()), "the shared stdio fixture runs on the host"); skip_if_wine_exec!(Ok(()), "requires a Windows test_stdio_server binary"); + let snapshot = exercise_mcp_sharing(/*resume_with_replacement_auth*/ false).await?; + insta::assert_snapshot!(snapshot); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn resumed_parent_and_child_share_mcp_process() -> Result<()> { + skip_if_no_network!(Ok(())); + core_test_support::skip_if_remote!(Ok(()), "the shared stdio fixture runs on the host"); + skip_if_wine_exec!(Ok(()), "requires a Windows test_stdio_server binary"); + + exercise_mcp_sharing(/*resume_with_replacement_auth*/ true).await?; + Ok(()) +} + +async fn exercise_mcp_sharing(resume_with_replacement_auth: bool) -> Result { let first_server = start_mock_server().await; let second_server = start_mock_server().await; let command = stdio_server_bin()?; - let test = test_codex() + let mut test = test_codex() .with_config(move |config| { super::configure_scenario_catalog(config); config.approvals_reviewer = ApprovalsReviewer::User; @@ -57,6 +78,56 @@ async fn shared_mcp_keeps_tools_and_model_requested_input_per_chat() -> Result<( .await?; wait_for_mcp_server(&test.codex, "shared").await?; + if resume_with_replacement_auth { + let initial = + responses::mount_sse_once(&first_server, responses::sse_completed("before-resume")) + .await; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "Prepare to resume".to_string(), + text_elements: Vec::new(), + }])) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + assert_eq!(initial.requests().len(), 1); + + let parent_id = codex_protocol::ThreadId::from(test.session_configured.session_id); + let rollout = test.codex.rollout_path().expect("parent rollout"); + test.codex.shutdown_and_wait().await?; + test.thread_manager.remove_thread(&parent_id).await; + let context = test + .thread_store + .load_latest_model_context(codex_thread_store::LoadThreadHistoryParams { + thread_id: parent_id, + include_archived: false, + }) + .await?; + let resumed = test + .thread_manager + .resume_thread_with_history( + test.config.clone(), + InitialHistory::Resumed(ResumedHistory { + history_revision: None, + conversation_id: context.thread_id, + history: Arc::new(context.items), + rollout_path: Some(rollout), + }), + codex_core::test_support::auth_manager_from_auth_with_home( + CodexAuth::from_api_key("resumed-key"), + test.codex_home_path().to_path_buf(), + ), + /*parent_trace*/ None, + ClientMcpExtensions::default(), + ) + .await?; + test.codex = resumed.thread; + test.session_configured = resumed.session_configured; + wait_for_mcp_server(&test.codex, "shared").await?; + } + let mut options = test.start_thread_options().await; options.config.model_provider.base_url = Some(format!("{}/v1", second_server.uri())); let mut servers = options.config.mcp_servers.get().clone(); @@ -65,7 +136,17 @@ async fn shared_mcp_keeps_tools_and_model_requested_input_per_chat() -> Result<( .expect("shared server") .enabled_tools = Some(vec!["echo-tool".to_string()]); options.config.mcp_servers.set(servers)?; - let second = test.thread_manager.start_thread(options).await?.thread; + let second = if resume_with_replacement_auth { + options.session_source = Some(SessionSource::Internal( + InternalSessionSource::MemoryConsolidation, + )); + test.thread_manager + .spawn_internal_session(test.session_configured.session_id.into(), options) + .await? + } else { + test.thread_manager.start_thread(options).await? + } + .thread; wait_for_mcp_server(&second, "shared").await?; let chats = [ @@ -168,12 +249,11 @@ async fn shared_mcp_keeps_tools_and_model_requested_input_per_chat() -> Result<( .map(move |request| SnapshotEntry::captured(request).labeled(label)) }) .collect::>(); - insta::assert_snapshot!(context_snapshot::format_context_snapshot( + Ok(context_snapshot::format_context_snapshot( "Two chats share one MCP process, expose different tools, and answer simultaneous model-requested questions in reverse order.", &entries, &ContextSnapshotOptions::default() .rewrite_known_segments() .include_request_settings(), - )); - Ok(()) + )) } diff --git a/codex-rs/ext/mcp/src/event_stream_tests.rs b/codex-rs/ext/mcp/src/event_stream_tests.rs index 48c30e87de..ea920480d6 100644 --- a/codex-rs/ext/mcp/src/event_stream_tests.rs +++ b/codex-rs/ext/mcp/src/event_stream_tests.rs @@ -277,6 +277,7 @@ impl Fixture { let startup_cancellation_token = CancellationToken::new(); let runtime_input = || McpRuntimeInput { network_policy_metadata: None, + runtime_auth: None, startup_policy: McpStartupPolicy::Eager, config: Arc::clone(&config), plugins_available: false, diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index ee67606f3f..14c05028a4 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -1877,6 +1877,13 @@ async fn external_auth_provider_can_install_headers() { headers.insert("x-external-auth", http::HeaderValue::from_static("enabled")); let auth = CodexAuth::Headers(AuthHeaders::new(headers)); let codex_home = tempdir().expect("tempdir"); + login_with_api_key( + codex_home.path(), + "stored-key", + AuthCredentialsStoreMode::Ephemeral, + AuthKeyringBackendKind::default(), + ) + .unwrap(); let manager = AuthManager::new( codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, @@ -1893,6 +1900,7 @@ async fn external_auth_provider_can_install_headers() { .await .expect("external auth should install"); + assert!(!manager.reload_with_cached_auth().await.unwrap()); assert_eq!(manager.auth_cached(), Some(auth)); assert!( manager diff --git a/codex-rs/login/src/auth/change_state_tests.rs b/codex-rs/login/src/auth/change_state_tests.rs index 2dc462752d..873b8e6fe0 100644 --- a/codex-rs/login/src/auth/change_state_tests.rs +++ b/codex-rs/login/src/auth/change_state_tests.rs @@ -117,7 +117,7 @@ fn api_workload_identity_changes_retire_previous_owner() { #[tokio::test] #[serial_test::serial(codex_auth_env)] -async fn auth_refresh_read_failure_is_transient_and_preserves_network_permit() { +async fn auth_reload_and_refresh_read_failure_preserves_network_permit() { let home = tempfile::tempdir().unwrap(); let storage = FileAuthStorage::new(home.path().to_path_buf()); let initial = chatgpt_auth("user-a", "workspace-a", "token-1"); @@ -167,9 +167,15 @@ async fn auth_refresh_read_failure_is_transient_and_preserves_network_permit() { jwt_auth.as_str(), ] { std::fs::write(home.path().join("auth.json"), contents).unwrap(); + let reload_error = manager.reload_with_cached_auth().await.unwrap_err(); let error = manager.refresh_token().await.unwrap_err(); assert!(matches!(error, RefreshTokenError::Transient(_))); - for formatted in [error.to_string(), format!("{error:?}")] { + for formatted in [ + reload_error.to_string(), + format!("{reload_error:?}"), + error.to_string(), + format!("{error:?}"), + ] { assert!( !formatted.contains("private-credential-value"), "{formatted}" diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index a4b61c10d1..79388f6485 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -2869,10 +2869,19 @@ impl AuthManager { /// Reloads stored credentials without hiding storage or hydration failures. /// Failed loads retain the last known owner so a temporary outage does not /// invalidate application network policy as if the user had logged out. - /// Callers must serialize reload and refresh and reject cached headers on error. + /// Leaves host-provided credentials alone and serializes with token refresh. + /// Callers must reject cached headers on error. pub async fn reload_with_cached_auth(&self) -> std::io::Result { + let _refresh_guard = self + .refresh_lock + .acquire() + .await + .map_err(std::io::Error::other)?; loop { let generation = self.auth_change_state_tx.borrow().generation; + if self.has_external_auth() { + return Ok(false); + } let cached = self.auth_cached(); let new_auth = self.try_load_managed_auth(cached.as_ref()).await?; if let Some(publication) = self.publish_reloaded_auth(generation, new_auth) { diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index ba328161e5..2282ca5b8a 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -61,7 +61,7 @@ pub use request_preparation::prepare_provider_request; pub use request_setup::ResolvedProviderRequest; pub use request_setup::resolve_request; pub use runtime_auth::runtime_auth_for_url; -pub use runtime_auth::select_runtime_auth_for_source; +pub(crate) use runtime_auth::select_runtime_auth_for_source; pub use runtime_auth::validate_runtime_destination; pub use runtime_auth::validate_runtime_provider; pub use runtime_auth::validate_runtime_selectors; diff --git a/codex-rs/model-provider/src/runtime_auth.rs b/codex-rs/model-provider/src/runtime_auth.rs index 26a7be196c..e02fb78b4c 100644 --- a/codex-rs/model-provider/src/runtime_auth.rs +++ b/codex-rs/model-provider/src/runtime_auth.rs @@ -31,7 +31,7 @@ pub(crate) async fn select_runtime_auth( } /// Select fresh request auth without allowing a reusable client's source to change. -pub async fn select_runtime_auth_for_source( +pub(crate) async fn select_runtime_auth_for_source( manager: &Arc, source_generation: u64, headers: &HeaderMap, diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index f2e2cf55a3..f6d88975c8 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -15,8 +15,8 @@ use std::sync::PoisonError; use std::time::Duration; use std::time::Instant; +use crate::AuthProviderResolver; use bytes::Bytes; -use codex_api::SharedAuthProvider; use codex_exec_server::ExecServerError; use codex_exec_server::HttpClient; use codex_exec_server::HttpHeader; @@ -79,7 +79,7 @@ pub enum StreamableHttpRedirectMode { pub(crate) struct StreamableHttpClientAdapter { http_client: Arc, default_headers: HeaderMap, - auth_provider: Option, + auth_provider: Option>, event_stream_cancellations: Arc>>>, has_configured_headers: bool, redirect_mode: StreamableHttpRedirectMode, @@ -118,7 +118,7 @@ impl StreamableHttpClientAdapter { pub(crate) fn new( http_client: Arc, default_headers: HeaderMap, - auth_provider: Option, + auth_provider: Option>, has_configured_headers: bool, redirect_mode: StreamableHttpRedirectMode, initialize_deadline: Arc>>, @@ -167,7 +167,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { let has_session_id = session_id.is_some(); let mut headers = self.default_headers.clone(); headers.extend(custom_headers); - self.add_auth_headers(&mut headers)?; + self.add_auth_headers(&mut headers).await?; insert_header( &mut headers, ACCEPT, @@ -464,7 +464,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { ) -> std::result::Result<(), StreamableHttpError> { let mut headers = self.default_headers.clone(); headers.extend(custom_headers); - self.add_auth_headers(&mut headers)?; + self.add_auth_headers(&mut headers).await?; if let Some(auth_token) = auth_token { insert_header( &mut headers, @@ -521,7 +521,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { > { let mut headers = self.default_headers.clone(); headers.extend(custom_headers); - self.add_auth_headers(&mut headers)?; + self.add_auth_headers(&mut headers).await?; insert_header( &mut headers, ACCEPT, @@ -606,17 +606,21 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { } impl StreamableHttpClientAdapter { - fn add_auth_headers( + async fn add_auth_headers( &self, headers: &mut HeaderMap, ) -> Result<(), StreamableHttpError> { - if let Some(auth_provider) = &self.auth_provider { - headers.extend( - auth_provider - .to_auth_headers() - .map_err(StreamableHttpClientAdapterError::from) - .map_err(StreamableHttpError::Client)?, - ); + if let Some(resolver) = &self.auth_provider { + let provider = resolver + .resolve(headers) + .await + .map_err(StreamableHttpClientAdapterError::from) + .map_err(StreamableHttpError::Client)?; + *headers = provider + .apply_auth_headers(std::mem::take(headers)) + .await + .map_err(StreamableHttpClientAdapterError::from) + .map_err(StreamableHttpError::Client)?; } Ok(()) } @@ -792,10 +796,15 @@ fn mcp_redirect_policy( headers: &HeaderMap, has_configured_headers: bool, ) -> HttpRedirectPolicy { - if headers - .get(HEADER_MCP_PROTOCOL_VERSION) - .and_then(|value| value.to_str().ok()) - == Some(ProtocolVersion::V_2026_07_28.as_str()) + if headers.contains_key("x-openai-agent-assertion") + || headers + .get(AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.starts_with("AgentAssertion ")) + || headers + .get(HEADER_MCP_PROTOCOL_VERSION) + .and_then(|value| value.to_str().ok()) + == Some(ProtocolVersion::V_2026_07_28.as_str()) || (mode == StreamableHttpRedirectMode::AgentPluginV1 && (has_configured_headers || headers.contains_key(AUTHORIZATION))) { diff --git a/codex-rs/rmcp-client/src/http_client_adapter_tests.rs b/codex-rs/rmcp-client/src/http_client_adapter_tests.rs index c3effd9983..7d1d06a018 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter_tests.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter_tests.rs @@ -88,6 +88,29 @@ fn authorization_redirects_depend_on_mode() { ); } +#[test] +fn runtime_assertions_stop_redirects_in_both_mcp_modes() { + for mode in [ + StreamableHttpRedirectMode::Legacy, + StreamableHttpRedirectMode::AgentPluginV1, + ] { + for (name, value) in [ + ("x-openai-agent-assertion", "1.protected.payload.signature"), + ( + "authorization", + "AgentAssertion 1.protected.payload.signature", + ), + ] { + let mut headers = HeaderMap::new(); + headers.insert(name, HeaderValue::from_static(value)); + assert_eq!( + mcp_redirect_policy(mode, &headers, /*has_configured_headers*/ false), + HttpRedirectPolicy::Stop, + ); + } + } +} + #[test] fn lf_terminators_reset_the_event_limit() { let mut limit = SseEventSizeLimit::new(Some(8)); @@ -283,3 +306,165 @@ fn size_accounting_saturates_without_overflow() { assert_eq!((error.kind(), limit.failed), (ErrorKind::InvalidData, true)); } + +struct AsyncRequestAuth { + requests: std::sync::atomic::AtomicUsize, +} + +impl crate::AuthProviderResolver for AsyncRequestAuth { + fn resolve<'a>( + &'a self, + _headers: &'a HeaderMap, + ) -> futures::future::BoxFuture<'a, Result> + { + Box::pin(async move { + tokio::task::yield_now().await; + let id = self + .requests + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + Ok(std::sync::Arc::new(SelectedRequestAuth(id)) as codex_api::SharedAuthProvider) + }) + } +} + +struct SelectedRequestAuth(usize); + +impl codex_api::AuthProvider for SelectedRequestAuth { + fn add_auth_headers(&self, headers: &mut HeaderMap) -> Result<(), codex_api::AuthError> { + headers.remove("x-openai-agent-assertion"); + headers.insert( + AUTHORIZATION, + HeaderValue::from_str(&format!("AgentAssertion proof-{}", self.0)).unwrap(), + ); + Ok(()) + } + + fn apply_auth_headers(&self, mut headers: HeaderMap) -> codex_api::AuthHeadersFuture<'_> { + Box::pin(async move { + self.add_auth_headers(&mut headers)?; + Ok(headers) + }) + } +} + +#[tokio::test] +async fn requests_await_auth_and_stop_assertion_redirects() -> anyhow::Result<()> { + use codex_exec_server::RouteAwareHttpClient; + use codex_http_client::HttpClientFactory; + use codex_http_client::OutboundProxyPolicy; + use futures::StreamExt; + use rmcp::transport::streamable_http_client::StreamableHttpClient; + use std::collections::HashMap; + use std::sync::Arc; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::method; + use wiremock::matchers::path; + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(ResponseTemplate::new(202)) + .expect(2) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/mcp")) + .respond_with(ResponseTemplate::new(200).insert_header("content-type", "text/event-stream")) + .expect(2) + .mount(&server) + .await; + Mock::given(method("DELETE")) + .and(path("/mcp")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + Mock::given(path("/redirect")) + .respond_with( + ResponseTemplate::new(307) + .insert_header("location", format!("{}/leaked", server.uri())), + ) + .expect(1) + .mount(&server) + .await; + let mut headers = HeaderMap::new(); + headers.insert(AUTHORIZATION, HeaderValue::from_static("Bearer stale")); + headers.insert( + "x-openai-agent-assertion", + HeaderValue::from_static("stale-proof"), + ); + let adapter = super::StreamableHttpClientAdapter::new( + Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new( + OutboundProxyPolicy::ReqwestDefault, + ))), + headers, + Some(Arc::new(AsyncRequestAuth { + requests: Default::default(), + })), + /*has_configured_headers*/ false, + StreamableHttpRedirectMode::Legacy, + Default::default(), + ); + let uri: Arc = format!("{}/mcp", server.uri()).into(); + let notification: rmcp::model::ClientJsonRpcMessage = + serde_json::from_value(serde_json::json!({ + "jsonrpc": "2.0", "method": "notifications/initialized", + }))?; + for _ in 0..2 { + adapter + .post_message( + uri.clone(), + notification.clone(), + /*session_id*/ None, + /*auth_token*/ None, + HashMap::new(), + ) + .await?; + } + for _ in 0..2 { + let mut stream = adapter + .get_stream( + uri.clone(), + Some("session-a".into()), + /*last_event_id*/ None, + /*auth_token*/ None, + HashMap::new(), + ) + .await?; + assert!(stream.next().await.is_none()); + } + adapter + .delete_session( + uri, + "session-a".into(), + /*auth_token*/ None, + HashMap::new(), + ) + .await?; + assert!( + adapter + .post_message( + format!("{}/redirect", server.uri()).into(), + notification, + /*session_id*/ None, + /*auth_token*/ None, + HashMap::new() + ) + .await + .is_err() + ); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 6); + for (id, request) in requests.iter().enumerate() { + assert_eq!( + request.headers["authorization"], + format!("AgentAssertion proof-{id}") + ); + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + assert_ne!(request.url.path(), "/leaked"); + } + server.verify().await; + Ok(()) +} diff --git a/codex-rs/rmcp-client/src/http_request_auth.rs b/codex-rs/rmcp-client/src/http_request_auth.rs new file mode 100644 index 0000000000..2e955d3109 --- /dev/null +++ b/codex-rs/rmcp-client/src/http_request_auth.rs @@ -0,0 +1,21 @@ +use codex_api::AuthError; +use codex_api::SharedAuthProvider; +use futures::future::BoxFuture; +use http::HeaderMap; + +/// Select request-scoped auth for each HTTP operation on a long-lived MCP connection. +pub trait AuthProviderResolver: Send + Sync { + fn resolve<'a>( + &'a self, + headers: &'a HeaderMap, + ) -> BoxFuture<'a, Result>; +} + +impl AuthProviderResolver for SharedAuthProvider { + fn resolve<'a>( + &'a self, + _headers: &'a HeaderMap, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { Ok(self.clone()) }) + } +} diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 7e459657da..e3fbadac8b 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -14,6 +14,7 @@ mod executor_process_transport; mod http_client_adapter; mod http_client_redirect; mod http_headers; +mod http_request_auth; mod in_process_transport; mod local_stdio_io; mod local_stdio_transport; @@ -76,6 +77,7 @@ pub use enterprise_oauth_login::perform_enterprise_oauth_login_return_url; pub use event_notification_transport::EventNotificationReceiver; pub use http_client_adapter::StreamableHttpRedirectMode; pub use http_headers::with_http_headers_helper; +pub use http_request_auth::AuthProviderResolver; pub use in_process_transport::InProcessTransportFactory; pub use oauth::EmaCredentialLease; pub use oauth::EmaCredentialSnapshot; diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index d2766f0155..221336dfc8 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -14,10 +14,10 @@ use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; +use crate::AuthProviderResolver; use anyhow::Context; use anyhow::Result; use anyhow::anyhow; -use codex_api::SharedAuthProvider; use codex_config::McpServerOAuthConfig; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::McpServerEnvVar; @@ -183,7 +183,7 @@ enum TransportRecipe { pinned_credential_store: Arc>, originator: AuthStorageOriginator, http_client: Arc, - auth_provider: Option, + auth_provider: Option>, redirect_mode: StreamableHttpRedirectMode, oauth_refresh_mode: McpOAuthRefreshMode, initialize_deadline: Arc>>, @@ -573,7 +573,7 @@ impl RmcpClient { store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, http_client: Arc, - auth_provider: Option, + auth_provider: Option>, ) -> Result { Self::new_streamable_http_client_with_protocol_mode( server_name, @@ -601,7 +601,7 @@ impl RmcpClient { store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, http_client: Arc, - auth_provider: Option, + auth_provider: Option>, protocol_mode: McpProtocolMode, ) -> Result { Self::new_streamable_http_client_with_protocol_mode_and_redirect_mode( @@ -633,7 +633,7 @@ impl RmcpClient { store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, http_client: Arc, - auth_provider: Option, + auth_provider: Option>, protocol_mode: McpProtocolMode, redirect_mode: StreamableHttpRedirectMode, oauth_refresh_mode: McpOAuthRefreshMode, diff --git a/codex-rs/rmcp-client/tests/mcp_events.rs b/codex-rs/rmcp-client/tests/mcp_events.rs index 9d3e54e1ee..a76f3c4743 100644 --- a/codex-rs/rmcp-client/tests/mcp_events.rs +++ b/codex-rs/rmcp-client/tests/mcp_events.rs @@ -16,6 +16,7 @@ use axum::response::sse::Sse; use axum::routing::post; use codex_api::AuthError; use codex_api::AuthProvider; +use codex_api::SharedAuthProvider; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; @@ -160,6 +161,7 @@ async fn plugin_runtime_event_streams_are_isolated_and_cancel_locally() -> anyho let base_url = format!("http://{}", listener.local_addr()?); let server = tokio::spawn(async move { axum::serve(listener, router).await }); let auth = Arc::new(RevocableAuth::default()); + let request_auth: SharedAuthProvider = auth.clone(); let client = RmcpClient::new_streamable_http_client( "test-streamable-http", &format!("{base_url}/mcp"), @@ -169,7 +171,7 @@ async fn plugin_runtime_event_streams_are_isolated_and_cancel_locally() -> anyho OAuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), Environment::default_for_tests().get_http_client(), - Some(auth.clone()), + Some(Arc::new(request_auth)), ) .await?; initialize_client(&client).await?;