diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index d9bcb9335e..b3bb45bc1f 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3935,6 +3935,7 @@ dependencies = [ "codex-protocol", "codex-utils-absolute-path", "codex-utils-path-uri", + "core_test_support", "futures", "pretty_assertions", "serde_json", @@ -3943,6 +3944,7 @@ dependencies = [ "tokio", "tokio-util", "tracing", + "wiremock", ] [[package]] diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 9acb97e52d..7dda112ab1 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -16,6 +16,7 @@ pub use resource_client::McpEventDefinition; pub use resource_client::McpEventNotification; pub use resource_client::McpEventStream; pub use resource_client::McpResourceClient; +pub use resource_client::McpResourceClientAuthKey; pub use resource_client::McpResourceClientCacheKey; pub use resource_client::McpResourcePage; pub use resource_client::McpResourceReadResult; diff --git a/codex-rs/codex-mcp/src/resource_client.rs b/codex-rs/codex-mcp/src/resource_client.rs index 6d2097f79a..0fa40abeb4 100644 --- a/codex-rs/codex-mcp/src/resource_client.rs +++ b/codex-rs/codex-mcp/src/resource_client.rs @@ -219,6 +219,25 @@ impl PartialEq for McpResourceServerCacheKey { impl Eq for McpResourceServerCacheKey {} +/// Opaque auth scope and server availability, independent of connection-set rebuilds. +/// This is for fixed, host-owned sources; it does not fingerprint arbitrary server config. +#[derive(Clone)] +pub struct McpResourceClientAuthKey { + pub(crate) generation: Arc<()>, + pub(crate) server: String, + pub(crate) available: bool, +} + +impl PartialEq for McpResourceClientAuthKey { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.generation, &other.generation) + && self.server == other.server + && self.available == other.available + } +} + +impl Eq for McpResourceClientAuthKey {} + impl std::fmt::Debug for McpResourceClient { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { formatter @@ -243,6 +262,12 @@ impl McpResourceClient { self.runtime.resource_cache_key(server) } + /// Tracks published auth and source availability without invalidating on unrelated + /// environment changes. This key contains no credentials. + pub fn auth_cache_key_for_server(&self, server: &str) -> McpResourceClientAuthKey { + self.runtime.auth_cache_key_for_server(server) + } + /// Returns whether this client can address the named server. /// /// This does not wait for server startup. diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 7cc11b1771..588c9e0edc 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -119,6 +119,7 @@ struct PublishedMcpRuntime { config: Option>, auth: Option, auth_token: Option, + auth_generation: Arc<()>, plugins_available: bool, ready_selected_capability_roots: Vec, environment_selections: Arc<[TurnEnvironmentSelection]>, @@ -145,6 +146,23 @@ fn ensure_host_owned_apps_registration( Ok(()) } +impl PublishedMcpRuntime { + // Shared by dirty detection and cache-key publication so their auth rules agree. + fn auth_matches(&self, auth: Option<&CodexAuth>) -> bool { + match (self.auth.as_ref(), auth) { + (Some(previous), Some(latest)) => { + previous == latest + && previous.get_account_id() == latest.get_account_id() + && previous.get_chatgpt_user_id() == latest.get_chatgpt_user_id() + && previous.is_fedramp_account() == latest.is_fedramp_account() + && self.auth_token == latest.get_token().ok() + } + (None, None) => true, + (Some(_), None) | (None, Some(_)) => false, + } + } +} + struct CachedMcpBinding { catalog_revisions: HashMap, // Reuse a frozen binding while a model step or caller still needs it. @@ -211,6 +229,7 @@ impl McpRuntime { config: None, auth: None, auth_token: None, + auth_generation: Arc::new(()), plugins_available: false, ready_selected_capability_roots: Vec::new(), environment_selections: Arc::default(), @@ -314,6 +333,12 @@ impl McpRuntime { let config = Arc::clone(&input.config); let auth = input.auth.clone(); let auth_token = auth.as_ref().and_then(|auth| auth.get_token().ok()); + let current = self.current.load_full(); + let auth_generation = if current.auth_matches(auth.as_ref()) { + Arc::clone(¤t.auth_generation) + } else { + Arc::new(()) + }; let plugins_available = input.plugins_available; let ready_selected_capability_roots = input.ready_selected_capability_roots.clone(); let environment_selections = Arc::clone(&input.runtime_context.environment_selections); @@ -345,6 +370,7 @@ impl McpRuntime { config: Some(config), auth, auth_token, + auth_generation, plugins_available, ready_selected_capability_roots, environment_selections, @@ -463,17 +489,18 @@ impl McpRuntime { /// Returns whether the published snapshot still belongs to the current credentials. pub fn current_auth_matches(&self, auth: Option<&CodexAuth>) -> bool { + self.current.load().auth_matches(auth) + } + + pub(crate) fn auth_cache_key_for_server( + &self, + server: &str, + ) -> crate::McpResourceClientAuthKey { let current = self.current.load(); - match (current.auth.as_ref(), auth) { - (Some(previous), Some(latest)) => { - previous == latest - && previous.get_account_id() == latest.get_account_id() - && previous.get_chatgpt_user_id() == latest.get_chatgpt_user_id() - && previous.is_fedramp_account() == latest.is_fedramp_account() - && current.auth_token == latest.get_token().ok() - } - (None, None) => true, - (Some(_), None) | (None, Some(_)) => false, + crate::McpResourceClientAuthKey { + generation: Arc::clone(¤t.auth_generation), + server: server.to_string(), + available: current.connections.contains_server(server), } } @@ -982,6 +1009,7 @@ mod tests { config: Some(Arc::new(config)), auth: None, auth_token: None, + auth_generation: Arc::new(()), plugins_available: false, ready_selected_capability_roots: Vec::new(), environment_selections: Arc::default(), @@ -1055,6 +1083,7 @@ mod tests { ))), auth: None, auth_token: None, + auth_generation: Arc::new(()), plugins_available: false, ready_selected_capability_roots: Vec::new(), environment_selections: Arc::default(), diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index f1679f333a..18e14757a6 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -82,6 +82,7 @@ impl<'a> McpEnvironmentScope<'a> { } pub(crate) struct McpThreadIdentity<'a> { + pub(crate) auth_changed: bool, pub(crate) session_source: &'a SessionSource, pub(crate) originator: &'a str, pub(crate) environments: McpEnvironmentScope<'a>, @@ -171,7 +172,8 @@ impl McpManager { ready_selected_capability_roots, executor_capability_discovery, ) - .with_session_source(identity.session_source), + .with_session_source(identity.session_source) + .with_auth_changed(identity.auth_changed), Some(identity.originator), identity.environments, identity.disabled_plugin_ids, diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 0aa971ed8f..baf4cb5777 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -121,23 +121,26 @@ impl Session { &environments, ) .await; - let mcp_projection = self - .services - .mcp_manager - .runtime_config_for_step( - config, - &self.services.mcp_thread_init, - &self.services.thread_extension_data, - McpThreadIdentity { - session_source: &session_source, - originator: &originator, - disabled_plugin_ids: &disabled_plugin_ids, - environments: McpEnvironmentScope::Selected(&environment_selections), - }, - &ready_selected_capability_roots, - executor_capability_discovery.as_deref(), - ) - .await; + let mcp_projection = + self.services + .mcp_manager + .runtime_config_for_step( + config, + &self.services.mcp_thread_init, + &self.services.thread_extension_data, + McpThreadIdentity { + auth_changed: !self.services.mcp_runtime.current_auth_matches( + self.services.auth_manager.auth_cached().as_ref(), + ), + session_source: &session_source, + originator: &originator, + disabled_plugin_ids: &disabled_plugin_ids, + environments: McpEnvironmentScope::Selected(&environment_selections), + }, + &ready_selected_capability_roots, + executor_capability_discovery.as_deref(), + ) + .await; let mcp_config = self .project_selected_environment_mcp_servers(config, &environments, mcp_projection) .await @@ -219,6 +222,10 @@ impl Session { &self.services.mcp_thread_init, &self.services.thread_extension_data, McpThreadIdentity { + auth_changed: !self + .services + .mcp_runtime + .current_auth_matches(desired.auth.as_ref()), session_source: &desired.session_source, originator: &desired.originator, disabled_plugin_ids: &desired.disabled_plugin_ids, @@ -288,6 +295,10 @@ impl Session { &self.services.mcp_thread_init, &self.services.thread_extension_data, McpThreadIdentity { + auth_changed: !self + .services + .mcp_runtime + .current_auth_matches(desired.auth.as_ref()), session_source: &desired.session_source, originator: &desired.originator, disabled_plugin_ids: &desired.disabled_plugin_ids, @@ -688,6 +699,10 @@ impl Session { &self.services.mcp_thread_init, &self.services.thread_extension_data, McpThreadIdentity { + auth_changed: !self + .services + .mcp_runtime + .current_auth_matches(desired.auth.as_ref()), session_source: &turn_context.session_source, originator: &turn_context.originator, disabled_plugin_ids: &disabled_plugin_ids, diff --git a/codex-rs/core/src/session/mcp_prewarm.rs b/codex-rs/core/src/session/mcp_prewarm.rs index b203472367..435699122a 100644 --- a/codex-rs/core/src/session/mcp_prewarm.rs +++ b/codex-rs/core/src/session/mcp_prewarm.rs @@ -9,6 +9,11 @@ impl Session { pub(crate) fn request_mcp_runtime_refresh(&self) { // Plugin changes can reuse connections but still change their skill resources. self.services.mcp_runtime.invalidate_resource_caches(); + self.request_mcp_runtime_reprojection(); + } + + /// Reproject contributor state without invalidating cached MCP resources. + pub(crate) fn request_mcp_runtime_reprojection(&self) { self.mark_mcp_runtime_dirty(); self.schedule_mcp_prewarm(); } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index f27eae416f..cf806a804f 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1125,6 +1125,7 @@ impl Session { mcp_thread_init_for_startup, thread_extension_data_for_mcp, McpThreadIdentity { + auth_changed: false, session_source: &mcp_session_source, originator: &mcp_originator, disabled_plugin_ids: &mcp_disabled_plugin_ids, @@ -1818,6 +1819,10 @@ impl Session { &sess.services.mcp_thread_init, &sess.services.thread_extension_data, McpThreadIdentity { + auth_changed: !sess + .services + .mcp_runtime + .current_auth_matches(latest_auth.as_ref()), session_source: &session_configuration.session_source, originator: &session_configuration.originator, disabled_plugin_ids: &session_configuration.disabled_plugin_ids, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 3bb0bfb433..8417461e78 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3,6 +3,9 @@ mod notification_tests; use super::mcp_refresh::McpRefresh; use super::step_context::StepInputs; + +#[path = "turn_start_mcp_tests.rs"] +mod turn_start_mcp_tests; use super::step_settings::ResolvedStepSettings; use super::step_settings::StepSettings; use super::step_settings::StepSettingsUpdate; @@ -3285,7 +3288,7 @@ async fn turn_start_lifecycle_exposes_turn_metadata_and_token_baseline() { turn_level_id: String, turn_id: String, collaboration_mode: CollaborationMode, - token_usage_at_turn_start: TokenUsage, + token_usage_at_turn_start: Option, saw_session_store: bool, saw_thread_store: bool, } @@ -3309,7 +3312,7 @@ async fn turn_start_lifecycle_exposes_turn_metadata_and_token_baseline() { turn_level_id: input.turn_store.level_id().to_string(), turn_id: input.turn_id.to_string(), collaboration_mode: input.collaboration_mode.clone(), - token_usage_at_turn_start: input.token_usage_at_turn_start.clone(), + token_usage_at_turn_start: input.token_usage_at_turn_start.cloned(), saw_session_store: input .session_store .get::() @@ -3356,7 +3359,7 @@ async fn turn_start_lifecycle_exposes_turn_metadata_and_token_baseline() { turn_level_id: turn_context.sub_id.clone(), turn_id: turn_context.sub_id.clone(), collaboration_mode: turn_context.collaboration_mode(), - token_usage_at_turn_start, + token_usage_at_turn_start: Some(token_usage_at_turn_start), saw_session_store: true, saw_thread_store: true, }; diff --git a/codex-rs/core/src/session/turn_start_mcp_tests.rs b/codex-rs/core/src/session/turn_start_mcp_tests.rs new file mode 100644 index 0000000000..c325d53026 --- /dev/null +++ b/codex-rs/core/src/session/turn_start_mcp_tests.rs @@ -0,0 +1,142 @@ +//! Turn-start MCP reconciliation and cooperative discovery cancellation regressions. + +use super::*; +use codex_mcp::McpResourceClient; +use codex_mcp::McpResourceClientCacheKey; +use pretty_assertions::assert_eq; + +struct Recorder { + client: McpResourceClient, + requires_mcp: std::sync::atomic::AtomicBool, + keys: std::sync::Mutex>, + discovered: tokio::sync::Notify, +} + +impl codex_extension_api::TurnLifecycleContributor for Recorder { + fn turn_start_phase( + &self, + _thread_store: &codex_extension_api::ExtensionData, + ) -> codex_extension_api::TurnStartPhase { + codex_extension_api::TurnStartPhase::RegularTaskStart + } + + fn requires_mcp_runtime(&self, _thread_store: &codex_extension_api::ExtensionData) -> bool { + self.requires_mcp.load(std::sync::atomic::Ordering::Relaxed) + } + + fn on_turn_start<'a>( + &'a self, + _input: codex_extension_api::TurnStartInput<'a>, + ) -> codex_extension_api::ExtensionFuture<'a, ()> { + Box::pin(async move { + self.keys + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(self.client.cache_key()); + self.discovered.notify_one(); + std::future::pending::<()>().await; + }) + } +} + +#[tokio::test] +async fn turn_start_refreshes_dirty_mcp_reuses_clean_runtime_and_cancels_discovery() { + let (mut session, turn_context) = make_session_and_context().await; + let recorder = Arc::new(Recorder { + client: McpResourceClient::new(Arc::clone(&session.services.mcp_runtime)), + requires_mcp: Default::default(), + keys: Default::default(), + discovered: Default::default(), + }); + let initial_key = recorder.client.cache_key(); + let mut builder = codex_extension_api::ExtensionRegistryBuilder::new(); + builder.turn_lifecycle_contributor(recorder.clone()); + session.services.extensions = Arc::new(builder.build()); + let session = Arc::new(session); + let turn_context = Arc::new(turn_context); + + for turn in 0..4 { + recorder + .requires_mcp + .store(turn > 0, std::sync::atomic::Ordering::Relaxed); + if turn == 0 { + session.mark_mcp_runtime_dirty(); + } else if turn == 3 { + session.services.mcp_runtime.reconnect_on_next_refresh(); + session.request_mcp_runtime_refresh(); + } + session + .spawn_task( + Arc::clone(&turn_context), + Vec::new(), + crate::tasks::RegularTask::new(), + ) + .await; + tokio::time::timeout( + std::time::Duration::from_secs(5), + recorder.discovered.notified(), + ) + .await + .expect("regular-task contributor should run"); + assert_eq!(session.mcp_refresh.is_pending(), turn == 0); + assert_eq!( + recorder + .keys + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .len(), + turn + 1 + ); + if turn > 0 { + session.refresh_mcp_if_dirty().await; + assert!( + recorder.client.cache_key() + == recorder + .keys + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner)[turn], + "a clean refresh must reuse the published MCP runtime" + ); + } + let (cancellation_token, done) = { + let active_turn = session.active_turn.lock().await; + let task = active_turn + .as_ref() + .and_then(|turn| turn.task.as_ref()) + .expect("discovery should run inside the registered task"); + (task.cancellation_token.clone(), Arc::clone(&task.done)) + }; + let completed = done.notified(); + tokio::pin!(completed); + completed.as_mut().enable(); + cancellation_token.cancel(); + tokio::time::timeout(std::time::Duration::from_secs(5), completed) + .await + .expect("discovery cancellation must finish the task without forced abort"); + assert!( + session.mcp_refresh.is_pending(), + "cancelled discovery must request an MCP refresh before the task completes" + ); + session.abort_all_tasks(TurnAbortReason::Interrupted).await; + } + let keys = recorder + .keys + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + assert!( + keys[0] == initial_key, + "unrelated contributors must not reconcile MCP" + ); + assert!( + keys[1] != initial_key, + "discovery must observe the reconciled generation" + ); + assert!( + keys[2] != keys[1], + "cancelled cloud preparation must request reprojection on the next turn" + ); + assert!( + keys[3] != keys[2], + "the next turn must consume a pending reconnect before discovery" + ); +} diff --git a/codex-rs/core/src/tasks/lifecycle.rs b/codex-rs/core/src/tasks/lifecycle.rs index cde3fc7d72..10057aaefa 100644 --- a/codex-rs/core/src/tasks/lifecycle.rs +++ b/codex-rs/core/src/tasks/lifecycle.rs @@ -1,5 +1,8 @@ +use std::sync::Arc; + use codex_extension_api::ExtensionData; use codex_extension_api::ThreadIdleCause; +use codex_extension_api::TurnStartPhase; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TurnAbortReason; @@ -9,12 +12,21 @@ use crate::session::turn_context::TurnContext; impl Session { pub(super) async fn emit_turn_start_lifecycle( - &self, + self: &Arc, turn_context: &TurnContext, - token_usage_at_turn_start: &TokenUsage, + token_usage_at_turn_start: Option<&TokenUsage>, + phase: TurnStartPhase, ) { let collaboration_mode = turn_context.collaboration_mode(); for contributor in self.services.extensions.turn_lifecycle_contributors() { + if contributor.turn_start_phase(&self.services.thread_extension_data) != phase { + continue; + } + if phase == TurnStartPhase::RegularTaskStart + && contributor.requires_mcp_runtime(&self.services.thread_extension_data) + { + self.refresh_mcp_if_dirty().await; + } contributor .on_turn_start(codex_extension_api::TurnStartInput { turn_id: turn_context.sub_id.as_str(), diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 3caff155a6..751279328b 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -322,8 +322,12 @@ impl Session { self.input_queue .extend_pending_input_for_turn_state(turn_state.as_ref(), pending_items) .await; - self.emit_turn_start_lifecycle(turn_context.as_ref(), &token_usage_at_turn_start) - .await; + self.emit_turn_start_lifecycle( + turn_context.as_ref(), + Some(&token_usage_at_turn_start), + codex_extension_api::TurnStartPhase::BeforeTaskRegistration, + ) + .await; let mut active = self.active_turn.lock().await; let turn = active.get_or_insert_with(ActiveTurn::default); diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index 1dc274437f..5068ebf5a8 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -1,5 +1,7 @@ use std::sync::Arc; +use codex_async_utils::OrCancelExt; +use codex_extension_api::TurnStartPhase; use tokio_util::sync::CancellationToken; use crate::session::TurnInput; @@ -47,6 +49,32 @@ impl SessionTask for RegularTask { // not wait on startup prewarm resolution. let prewarmed_client_session = async { sess.emit_turn_started(&ctx).await; + // Regular-start contributors run once, after the task is visible and interruptible. + let prepares_mcp = sess + .services + .extensions + .turn_lifecycle_contributors() + .iter() + .any(|contributor| { + contributor.turn_start_phase(&sess.services.thread_extension_data) + == TurnStartPhase::RegularTaskStart + && contributor.requires_mcp_runtime(&sess.services.thread_extension_data) + }); + let preparation = sess + .emit_turn_start_lifecycle( + &ctx, + /*token_usage_at_turn_start*/ None, + TurnStartPhase::RegularTaskStart, + ) + .or_cancel(&cancellation_token) + .await; + // Even cancelled discovery may have cleared the previous account's catalog. + if prepares_mcp { + sess.request_mcp_runtime_reprojection(); + } + if preparation.is_err() { + return SessionStartupPrewarmResolution::Cancelled; + } sess.set_server_reasoning_included(/*included*/ false).await; sess.consume_startup_prewarm_for_regular_turn(&cancellation_token) .await diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 8c2df4e169..89c39d5e44 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -1704,6 +1704,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() &first_session.services.mcp_thread_init, &first_session.services.thread_extension_data, McpThreadIdentity { + auth_changed: false, session_source: &SessionSource::Exec, originator: &first_originator, disabled_plugin_ids: &[], @@ -1725,6 +1726,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() &second_session.services.mcp_thread_init, &second_session.services.thread_extension_data, McpThreadIdentity { + auth_changed: false, session_source: &second_session_source, originator: &second_originator, disabled_plugin_ids: &[], @@ -1799,6 +1801,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() &first_session.services.mcp_thread_init, &first_session.services.thread_extension_data, McpThreadIdentity { + auth_changed: false, session_source: &SessionSource::Exec, originator: &first_originator, disabled_plugin_ids: &disabled_plugin_ids, diff --git a/codex-rs/core/tests/common/apps_test_server.rs b/codex-rs/core/tests/common/apps_test_server.rs index 92c9b38cae..eebd4ea848 100644 --- a/codex-rs/core/tests/common/apps_test_server.rs +++ b/codex-rs/core/tests/common/apps_test_server.rs @@ -102,6 +102,7 @@ enum AppsTestToolsListBehavior { AlwaysAvailable, AvailableWhen(Arc), AlwaysUnavailable, + Custom(Arc>>), } impl AppsTestServer { @@ -226,6 +227,13 @@ impl AppsTestServer { .await } + pub async fn mount_with_tools( + server: &MockServer, + tools: Arc>>, + ) -> Result { + Self::mount_with_tools_list_behavior(server, AppsTestToolsListBehavior::Custom(tools)).await + } + pub async fn mount_without_tools(server: &MockServer) -> Result { Self::mount_with_tools_list_behavior(server, AppsTestToolsListBehavior::AlwaysUnavailable) .await @@ -551,7 +559,8 @@ impl Respond for CodexAppsJsonRpcResponder { "notifications/initialized" => ResponseTemplate::new(202), "tools/list" => { let tools_available = match &self.tools_list_behavior { - AppsTestToolsListBehavior::AlwaysAvailable => true, + AppsTestToolsListBehavior::AlwaysAvailable + | AppsTestToolsListBehavior::Custom(_) => true, AppsTestToolsListBehavior::AvailableWhen(tools_available) => { tools_available.load(Ordering::SeqCst) } @@ -660,6 +669,14 @@ impl Respond for CodexAppsJsonRpcResponder { "nextCursor": null } }); + if let AppsTestToolsListBehavior::Custom(tools) = &self.tools_list_behavior { + response["result"]["tools"] = json!( + tools + .lock() + .expect("Apps test tools lock should not be poisoned") + .clone() + ); + } if !tools_available && let Some(tools) = response .pointer_mut("/result/tools") diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 2880e21b6d..1afb93a360 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -1426,6 +1426,31 @@ pub fn test_codex() -> TestCodexBuilder { } } +pub fn run_test_with_large_stack(name: &str, test: F) -> Result<()> +where + F: FnOnce() -> Fut + Send + 'static, + Fut: Future> + Send + 'static, +{ + const WORKER_THREADS: usize = 2; + const TEST_STACK_SIZE_BYTES: usize = 8 * 1024 * 1024; + + let handle = std::thread::Builder::new() + .name(name.to_string()) + .stack_size(TEST_STACK_SIZE_BYTES) + .spawn(move || -> Result<()> { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(WORKER_THREADS) + .thread_stack_size(TEST_STACK_SIZE_BYTES) + .enable_all() + .build()?; + runtime.block_on(Box::pin(test())) + })?; + + handle + .join() + .map_err(|_| anyhow!("{name} thread panicked"))? +} + #[cfg(test)] mod tests { use super::*; diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 58a9909ce1..6b5acb267f 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -55,6 +55,7 @@ pub use turn_input::TurnInputEnvironment; pub use turn_lifecycle::TurnAbortInput; pub use turn_lifecycle::TurnErrorInput; pub use turn_lifecycle::TurnStartInput; +pub use turn_lifecycle::TurnStartPhase; pub use turn_lifecycle::TurnStopInput; pub use world_state::PreviousWorldStateSection; pub use world_state::RenderedWorldStateFragment; @@ -185,8 +186,21 @@ pub trait ThreadLifecycleContributor: Send + Sync { /// extension-private turn state. The host exposes stable identifiers and /// extension stores instead of core runtime objects. pub trait TurnLifecycleContributor: Send + Sync { - /// Called after turn-scoped extension stores are created, before the task - /// for the turn starts running. + /// Selects the start phase using the current thread policy. Disabled callbacks + /// may retain the default phase and return without doing any work. + fn turn_start_phase(&self, _thread_store: &ExtensionData) -> TurnStartPhase { + TurnStartPhase::BeforeTaskRegistration + } + + /// Whether regular-task startup must reconcile MCP before this callback runs. + /// Ignored before task registration; contributors unrelated to MCP keep the default. + fn requires_mcp_runtime(&self, _thread_store: &ExtensionData) -> bool { + false + } + + /// Called once in the selected phase, before task-specific work begins. + /// Regular-task preparation may be cancelled before this callback completes; + /// stop and abort callbacks must tolerate missing or partial initialization. fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> { Box::pin(async move { let _self = self; diff --git a/codex-rs/ext/extension-api/src/contributors/mcp.rs b/codex-rs/ext/extension-api/src/contributors/mcp.rs index 374dd77f8f..69f8c6cdbb 100644 --- a/codex-rs/ext/extension-api/src/contributors/mcp.rs +++ b/codex-rs/ext/extension-api/src/contributors/mcp.rs @@ -14,6 +14,8 @@ use crate::ExtensionDataInit; pub struct McpServerContributionContext<'a, C> { /// Host configuration visible during MCP resolution. config: &'a C, + /// Whether pending auth differs from published auth; invalidate catalogs before projection. + auth_changed: bool, /// Extension-owned data for the active thread, when resolution is thread-scoped. thread_store: Option<&'a ExtensionData>, /// Stable host inputs for the active thread, when resolution is thread-scoped. @@ -41,6 +43,7 @@ impl<'a, C> McpServerContributionContext<'a, C> { pub fn global(config: &'a C) -> Self { Self { config, + auth_changed: false, thread_store: None, thread_init: None, session_source: None, @@ -61,6 +64,7 @@ impl<'a, C> McpServerContributionContext<'a, C> { ) -> Self { Self { config, + auth_changed: false, thread_store: Some(thread_store), thread_init: Some(thread_init), session_source: None, @@ -70,6 +74,17 @@ impl<'a, C> McpServerContributionContext<'a, C> { } } + /// Marks whether this projection replaces the published MCP authentication. + pub fn with_auth_changed(mut self, auth_changed: bool) -> Self { + self.auth_changed = auth_changed; + self + } + + /// Returns whether auth-scoped catalogs must be invalidated before projection. + pub fn auth_changed(&self) -> bool { + self.auth_changed + } + /// Attaches the stable source of the active thread to this contribution. pub fn with_session_source(mut self, session_source: &'a SessionSource) -> Self { self.session_source = Some(session_source); diff --git a/codex-rs/ext/extension-api/src/contributors/turn_lifecycle.rs b/codex-rs/ext/extension-api/src/contributors/turn_lifecycle.rs index e2c9e4bf7b..ce9e854930 100644 --- a/codex-rs/ext/extension-api/src/contributors/turn_lifecycle.rs +++ b/codex-rs/ext/extension-api/src/contributors/turn_lifecycle.rs @@ -1,3 +1,5 @@ +//! Turn lifecycle inputs and scheduling phases for host-owned contributors. + use codex_protocol::config_types::CollaborationMode; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::TokenUsage; @@ -5,6 +7,13 @@ use codex_protocol::protocol::TurnAbortReason; use crate::ExtensionData; +/// Runs before task registration or during cancellable regular-task startup. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum TurnStartPhase { + BeforeTaskRegistration, + RegularTaskStart, +} + /// Input supplied when the host starts a turn. pub struct TurnStartInput<'a> { /// Stable host-owned turn identifier. @@ -12,7 +21,9 @@ pub struct TurnStartInput<'a> { /// Effective collaboration mode for this turn. pub collaboration_mode: &'a CollaborationMode, /// Total token usage snapshot captured when the turn started. - pub token_usage_at_turn_start: &'a TokenUsage, + /// Present before task registration; absent during regular-task preparation. + /// Token-accounting contributors must remain in `BeforeTaskRegistration`. + pub token_usage_at_turn_start: Option<&'a TokenUsage>, /// Store scoped to the host session runtime. pub session_store: &'a ExtensionData, /// Store scoped to this thread runtime. diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index ea2d0eb046..9cda9f62b2 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -91,6 +91,7 @@ pub use contributors::TurnInputEnvironment; pub use contributors::TurnItemContributor; pub use contributors::TurnLifecycleContributor; pub use contributors::TurnStartInput; +pub use contributors::TurnStartPhase; pub use contributors::TurnStopInput; pub use contributors::WorldStateContributionInput; pub use contributors::WorldStateSectionContribution; diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index f5dc85138a..1ec9ee6378 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -234,6 +234,11 @@ where return; } + let Some(token_usage_at_turn_start) = input.token_usage_at_turn_start else { + tracing::warn!("skipping goal turn accounting: token baseline unavailable"); + return; + }; + if let Err(err) = self .state_dbs .thread_goals() @@ -247,7 +252,7 @@ where accounting.start_turn( input.turn_id, input.collaboration_mode.mode, - input.token_usage_at_turn_start, + token_usage_at_turn_start, ); if matches!( input.collaboration_mode.mode, diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index 124bfbd80f..4b41740f89 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -263,6 +263,19 @@ async fn create_goal_resets_baseline_before_turn_stop_accounting() -> anyhow::Re let thread_id = test_thread_id()?; seed_thread_metadata(runtime.as_ref(), thread_id).await?; let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; + // A missing baseline must not panic or seed accounting with invented token usage. + for contributor in harness.registry.turn_lifecycle_contributors() { + contributor + .on_turn_start(TurnStartInput { + turn_id: "missing-baseline", + collaboration_mode: &default_collaboration_mode(), + token_usage_at_turn_start: None, + session_store: &harness.session_store, + thread_store: &harness.thread_store, + turn_store: &ExtensionData::new("missing-baseline"), + }) + .await; + } harness .start_turn( "turn-1", @@ -1766,7 +1779,7 @@ impl GoalExtensionHarness { .on_turn_start(TurnStartInput { turn_id, collaboration_mode: &collaboration_mode, - token_usage_at_turn_start: usage, + token_usage_at_turn_start: Some(usage), session_store: &self.session_store, thread_store: &self.thread_store, turn_store: &turn_store, diff --git a/codex-rs/ext/mcp/Cargo.toml b/codex-rs/ext/mcp/Cargo.toml index bfc8a40f6a..5185f6cab1 100644 --- a/codex-rs/ext/mcp/Cargo.toml +++ b/codex-rs/ext/mcp/Cargo.toml @@ -33,6 +33,7 @@ tracing = { workspace = true } tokio = { workspace = true, features = ["sync"] } [dev-dependencies] +core_test_support = { workspace = true } axum = { workspace = true } codex-exec-server-test-support = { workspace = true } codex-login = { workspace = true } @@ -41,3 +42,4 @@ pretty_assertions = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } tokio-util = { workspace = true, features = ["rt"] } +wiremock = { workspace = true } diff --git a/codex-rs/ext/mcp/src/cloud_plugin.rs b/codex-rs/ext/mcp/src/cloud_plugin.rs new file mode 100644 index 0000000000..55cd2b1bbc --- /dev/null +++ b/codex-rs/ext/mcp/src/cloud_plugin.rs @@ -0,0 +1,164 @@ +//! Refreshes the cloud catalog each regular turn and projects Apps-only plugin packages. + +use crate::PluginListQuery; +use crate::PluginsThreadState; +use crate::plugin_contributor::PluginContributor; +use crate::plugin_contributor_state::SelectedPluginMetadata; +use codex_core::config::Config; +use codex_core_plugins::PluginCatalog; +use codex_core_plugins::PluginIdentity; +use codex_extension_api::ConfigContributor; +use codex_extension_api::ExtensionData; +use codex_extension_api::ExtensionFuture; +use codex_extension_api::ThreadLifecycleContributor; +use codex_extension_api::ThreadStartInput; +use codex_extension_api::TurnLifecycleContributor; +use codex_extension_api::TurnStartInput; +use codex_extension_api::TurnStartPhase; +use codex_features::Feature; +use codex_mcp::McpResourceClient; +use std::collections::HashSet; +use std::sync::Arc; + +const MAX_PROJECTED_CLOUD_PLUGINS: usize = 100; +const MAX_PLUGIN_DISPLAY_NAME_CHARS: usize = 64; +const MAX_CONNECTOR_IDS_PER_PLUGIN: usize = 64; + +impl ThreadLifecycleContributor for PluginContributor { + fn on_thread_start<'a>( + &'a self, + input: ThreadStartInput<'a, Config>, + ) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let state = input.thread_store.get_or_init(PluginsThreadState::default); + state.contributor_state().plugins_enabled = + input.config.features.enabled(Feature::Plugins); + if let Some(resources) = &input.mcp_resource_client { + input.session_store.insert(resources.as_ref().clone()); + } + }) + } +} + +impl ConfigContributor for PluginContributor { + fn on_config_changed( + &self, + _session_store: &ExtensionData, + thread_store: &ExtensionData, + _previous_config: &Config, + new_config: &Config, + ) { + let plugins_enabled = new_config.features.enabled(Feature::Plugins); + let state = thread_store.get_or_init(PluginsThreadState::default); + let mut state = state.contributor_state(); + state.plugins_enabled = plugins_enabled; + if self.providers.cloud.is_none() || !plugins_enabled { + state.cloud_generation = None; + } + } +} + +impl TurnLifecycleContributor for PluginContributor { + fn turn_start_phase(&self, thread_store: &ExtensionData) -> TurnStartPhase { + if self.providers.cloud.is_some() + && thread_store + .get::() + .is_some_and(|state| state.contributor_state().plugins_enabled) + { + TurnStartPhase::RegularTaskStart + } else { + TurnStartPhase::BeforeTaskRegistration + } + } + + fn requires_mcp_runtime(&self, thread_store: &ExtensionData) -> bool { + self.turn_start_phase(thread_store) == TurnStartPhase::RegularTaskStart + } + + fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let Some(provider) = &self.providers.cloud else { + return; + }; + if !input + .thread_store + .get::() + .is_some_and(|state| state.contributor_state().plugins_enabled) + { + return; + } + let resources = input + .session_store + .get::() + .map(|resources| resources.as_ref().clone()); + let state = input.thread_store.get_or_init(PluginsThreadState::default); + // Change generations before the fallible lookup so only failures within the + // same auth scope may retain a previously published catalog. + let auth_cache_key = state + .contributor_state() + .prepare_cloud_generation(resources.clone()); + match provider + .list(PluginListQuery { + thread_id: input.thread_store.level_id().to_string(), + turn_id: input.turn_id.to_string(), + mcp_resources: resources.map(Arc::new), + }) + .await + { + Ok(catalog) => { + let mut state = state.contributor_state(); + if !state.publish_cloud_catalog(auth_cache_key.as_ref(), catalog) { + tracing::warn!("cloud plugin discovery finished after auth changed"); + } + } + Err(_) => tracing::warn!("cloud plugin discovery failed"), + } + }) + } +} + +/// Cloud packages contribute Apps attribution, never executor MCP declarations. +pub(crate) struct CloudPluginMetadata { + pub(crate) selection_order: usize, + pub(crate) selected_root_id: String, + pub(crate) metadata: SelectedPluginMetadata, +} + +pub(crate) fn catalog_to_metadata(catalog: &PluginCatalog) -> Vec { + catalog + .entries + .iter() + .enumerate() + .filter_map(|(selection_order, entry)| { + let PluginIdentity::Remote { remote_plugin_id } = &entry.id else { + return None; + }; + let mut seen_connector_ids = HashSet::new(); + Some(CloudPluginMetadata { + selection_order, + selected_root_id: format!("cloud:{remote_plugin_id}"), + metadata: SelectedPluginMetadata { + plugin_id: remote_plugin_id.clone(), + plugin_display_name: entry + .display_name + .chars() + .take(MAX_PLUGIN_DISPLAY_NAME_CHARS) + .collect(), + servers: Vec::new(), + connector_ids: entry + .connector_ids + .iter() + .filter(|connector_id| seen_connector_ids.insert((*connector_id).clone())) + .take(MAX_CONNECTOR_IDS_PER_PLUGIN) + .cloned() + .collect(), + }, + }) + }) + .take(MAX_PROJECTED_CLOUD_PLUGINS) + .collect() +} + +#[cfg(test)] +#[path = "cloud_plugin_tests.rs"] +mod tests; diff --git a/codex-rs/ext/mcp/src/cloud_plugin_tests.rs b/codex-rs/ext/mcp/src/cloud_plugin_tests.rs new file mode 100644 index 0000000000..8ac0f36b27 --- /dev/null +++ b/codex-rs/ext/mcp/src/cloud_plugin_tests.rs @@ -0,0 +1,583 @@ +//! Tests cloud plugin lifecycle, tool calls, scheduling, and plugin policy. + +use super::*; +use crate::PluginProvider; +use crate::PluginProviderError; +use crate::PluginProviderFuture; +use crate::PluginProviders; +use crate::PluginsThreadState; +use crate::install_plugin_providers; +use anyhow::Context; +use codex_core_plugins::ExecutorPluginProvider; +use codex_core_plugins::PluginCatalog; +use codex_core_plugins::PluginCatalogEntry; +use codex_core_plugins::PluginIdentity; +use codex_exec_server::EnvironmentManager; +use codex_extension_api::ContextContributor; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::SelectedPluginSnapshot; +use codex_extension_api::WorldStateContributionInput; +use codex_extension_api::WorldStateSectionContribution; +use codex_protocol::config_types::CollaborationMode; +use codex_protocol::config_types::ModeKind; +use codex_protocol::config_types::Settings; +use core_test_support::apps_test_server::AppsTestServer; +use core_test_support::apps_test_server::recorded_apps_tool_calls; +use core_test_support::responses; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_function_call_with_namespace; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_sequence; +use core_test_support::responses::sse; +use core_test_support::test_codex::run_test_with_large_stack; +use core_test_support::test_codex::test_codex; +use pretty_assertions::assert_eq; +use serde_json::json; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; +use tokio::sync::Notify; + +#[tokio::test] +async fn cloud_discovery_selects_regular_task_phase() { + struct ExistingContributor; + impl TurnLifecycleContributor for ExistingContributor {} + + let mut builder = ExtensionRegistryBuilder::new(); + builder.turn_lifecycle_contributor(Arc::new(ExistingContributor)); + let executor = Arc::new(ExecutorPluginProvider::new(Arc::new( + EnvironmentManager::default_for_tests(), + ))); + install_plugin_providers(&mut builder, PluginProviders::new(executor.clone())); + install_plugin_providers( + &mut builder, + PluginProviders::new(executor).with_cloud_provider(Arc::new(CloudProvider::default())), + ); + let registry = builder.build(); + let thread = ExtensionData::new("thread"); + for plugins_enabled in [false, true] { + thread + .get_or_init(PluginsThreadState::default) + .contributor_state() + .plugins_enabled = plugins_enabled; + let dependencies = registry + .turn_lifecycle_contributors() + .iter() + .map(|contributor| contributor.requires_mcp_runtime(&thread)) + .collect::>(); + assert_eq!(dependencies, vec![false, false, plugins_enabled]); + let phases = registry + .turn_lifecycle_contributors() + .iter() + .map(|contributor| contributor.turn_start_phase(&thread)) + .collect::>(); + assert_eq!( + phases, + vec![ + TurnStartPhase::BeforeTaskRegistration, + TurnStartPhase::BeforeTaskRegistration, + if plugins_enabled { + TurnStartPhase::RegularTaskStart + } else { + TurnStartPhase::BeforeTaskRegistration + }, + ], + ); + } +} + +#[tokio::test] +async fn disabling_plugins_clears_cloud_catalog_and_skips_discovery() -> anyhow::Result<()> { + let config_home = tempfile::tempdir()?; + let mut config = codex_core::config::ConfigBuilder::default() + .codex_home(config_home.path().to_path_buf()) + .fallback_cwd(Some(config_home.path().to_path_buf())) + .build() + .await?; + let mode = CollaborationMode { + mode: ModeKind::Default, + settings: Settings { + model: "test".to_string(), + reasoning_effort: None, + developer_instructions: None, + }, + }; + let provider = Arc::new(CloudProvider::default()); + let mut builder = ExtensionRegistryBuilder::new(); + let providers = PluginProviders::new(Arc::new(ExecutorPluginProvider::new(Arc::new( + EnvironmentManager::default_for_tests(), + )))) + .with_cloud_provider(provider.clone()); + install_plugin_providers(&mut builder, providers); + let registry = builder.build(); + let session = ExtensionData::new("session"); + let thread = ExtensionData::new("thread"); + let thread_init = codex_extension_api::ExtensionDataInit::default(); + for (turn_id, plugins_enabled, reply) in [ + ("off", false, Ok(cloud_catalog())), + ("on", true, Ok(cloud_catalog())), + ("plugins-off", false, Ok(cloud_catalog())), + ( + "re-enabled-failed", + true, + Err("discovery unavailable".to_string()), + ), + ("re-enabled", true, Ok(cloud_catalog())), + ] { + let expected_catalog = if plugins_enabled { + reply.clone().ok() + } else { + None + }; + let previous_discoveries = provider.calls.load(Ordering::SeqCst); + provider.set_reply(reply); + let previous = config.clone(); + config + .features + .set_enabled(Feature::Plugins, plugins_enabled)?; + for contributor in registry.config_contributors() { + contributor.on_config_changed(&session, &thread, &previous, &config); + } + if !plugins_enabled && let Some(state) = thread.get::() { + assert_eq!(state.cloud_catalog(), None); + } + for contributor in registry.turn_lifecycle_contributors() { + contributor + .on_turn_start(TurnStartInput { + turn_id, + collaboration_mode: &mode, + token_usage_at_turn_start: None, + session_store: &session, + thread_store: &thread, + turn_store: &ExtensionData::new(turn_id), + }) + .await; + } + registry.mcp_server_contributors()[0] + .contribute(codex_extension_api::McpServerContributionContext::for_step( + &config, + &thread_init, + &thread, + "test", + &[], + /*executor_capability_discovery*/ None, + )) + .await; + if let Some(state) = thread.get::() { + assert_eq!(state.cloud_catalog(), expected_catalog, "turn {turn_id}"); + } + assert_eq!( + provider.calls.load(Ordering::SeqCst), + previous_discoveries + usize::from(plugins_enabled), + "turn {turn_id}", + ); + } + Ok(()) +} + +struct CloudProvider { + reply: Mutex>, + calls: AtomicUsize, + resources: Mutex>>, + pause_discovery: AtomicBool, + discovery_started: Notify, + resume_discovery: Notify, +} + +#[derive(Default)] +struct StepPackages(Mutex>>); + +impl ContextContributor for StepPackages { + fn contribute_world_state<'a>( + &'a self, + input: WorldStateContributionInput<'a>, + ) -> ExtensionFuture<'a, Vec> { + Box::pin(async move { + *self + .0 + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = input + .thread_store + .get::() + .map(|snapshot| { + snapshot + .plugins + .iter() + .map(|plugin| plugin.selected_root_id.clone()) + .collect() + }); + Vec::new() + }) + } +} + +impl Default for CloudProvider { + fn default() -> Self { + Self { + reply: Mutex::new(Ok(PluginCatalog::default())), + calls: AtomicUsize::new(0), + resources: Mutex::new(None), + pause_discovery: AtomicBool::new(false), + discovery_started: Notify::new(), + resume_discovery: Notify::new(), + } + } +} + +impl CloudProvider { + fn set_reply(&self, reply: Result) { + *self + .reply + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = reply; + } +} + +impl PluginProvider for CloudProvider { + fn list(&self, query: PluginListQuery) -> PluginProviderFuture<'_, PluginCatalog> { + *self + .resources + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = query.mcp_resources; + self.calls.fetch_add(1, Ordering::SeqCst); + let reply = self + .reply + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + let pause = self.pause_discovery.swap(false, Ordering::SeqCst); + Box::pin(async move { + if pause { + self.discovery_started.notify_one(); + self.resume_discovery.notified().await; + } + reply.map_err(PluginProviderError::Message) + }) + } +} + +fn cloud_catalog() -> PluginCatalog { + PluginCatalog { + entries: vec![PluginCatalogEntry { + id: PluginIdentity::Remote { + remote_plugin_id: "plugin_remote".to_string(), + }, + display_name: "Cloud Demo".to_string(), + version: Some("1.0.0".to_string()), + mcp_servers: Default::default(), + connector_ids: vec!["calendar".to_string()], + locations: vec![codex_core_plugins::PluginSourceLocation::Cloud { + resource_uri: "plugin://plugin_remote".to_string(), + bundle_uri: None, + }], + }], + warnings: Vec::new(), + } +} + +fn app_tool(connector: &str, action: &str) -> serde_json::Value { + json!({ + "name": format!("{connector}_{action}"), + "description": format!("Run {action} in {connector}."), + "annotations": {"readOnlyHint": true}, + "inputSchema": {"type": "object", "properties": {}, "additionalProperties": false}, + "_meta": { + "connector_id": connector, + "connector_name": connector, + "_codex_apps": { + "resource_uri": format!("connector://{connector}/tools/{connector}_{action}"), + "contains_mcp_source": true, + "connector_id": connector, + }, + }, + }) +} + +#[test] +fn cloud_catalog_refreshes_per_turn_and_invalidates_on_auth_change() -> anyhow::Result<()> { + run_test_with_large_stack( + "cloud-plugin-lifecycle", + run_cloud_catalog_refresh_lifecycle, + ) +} + +async fn run_cloud_catalog_refresh_lifecycle() -> anyhow::Result<()> { + let server = responses::start_mock_server().await; + // Each account has a fixed inventory; only the auth switch changes the server's tools. + let apps_tools = Arc::new(Mutex::new(vec![app_tool("calendar", "list_events")])); + let apps_server = AppsTestServer::mount_with_tools(&server, apps_tools.clone()).await?; + let mut gmail_catalog = cloud_catalog(); + gmail_catalog.entries[0].id = PluginIdentity::Remote { + remote_plugin_id: "plugin_gmail".to_string(), + }; + gmail_catalog.entries[0].display_name = "Gmail Demo".to_string(); + gmail_catalog.entries[0].connector_ids = vec!["gmail".to_string()]; + let provider = Arc::new(CloudProvider::default()); + let step_packages = Arc::new(StepPackages::default()); + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.prompt_contributor(step_packages.clone()); + install_plugin_providers( + &mut extensions, + PluginProviders::new(Arc::new(ExecutorPluginProvider::new(Arc::new( + EnvironmentManager::default_for_tests(), + )))) + .with_cloud_provider(provider.clone()), + ); + let mut builder = test_codex() + .with_exec_server_url("none") + .with_extensions(Arc::new(extensions.build())) + .with_auth(codex_login::CodexAuth::from_external_chatgpt_tokens( + "header.e30.initial", + "account-a", + /*chatgpt_plan_type*/ None, + )?) + .with_model_info_override("gpt-5.5", |model| { + model.supports_search_tool = false; + }) + .with_config(move |config| { + config + .features + .enable(Feature::Plugins) + .expect("enable plugins in fixture"); + config + .features + .enable(Feature::Apps) + .expect("enable apps in fixture"); + config.chatgpt_base_url = apps_server.chatgpt_base_url; + }); + let test = builder.build_with_auto_env(&server).await?; + let calendar = ("calendar", "list_events", "Cloud Demo"); + let gmail = ("gmail", "search_email", "Gmail Demo"); + for (turn, reply, expected_plugin, new_account, tool_call) in [ + ( + "empty", + Ok(PluginCatalog::default()), + None, + None, + Some(calendar), + ), + ( + "discovered", + Ok(cloud_catalog()), + Some("cloud:plugin_remote"), + None, + Some(calendar), + ), + ( + "same-account-failure", + Err("temporarily unavailable".to_string()), + Some("cloud:plugin_remote"), + None, + None, + ), + ("removed", Ok(PluginCatalog::default()), None, None, None), + ( + "restored", + Ok(cloud_catalog()), + Some("cloud:plugin_remote"), + None, + None, + ), + ( + "auth-changed-failure", + Err("unavailable after auth changed".to_string()), + None, + Some("account-b"), + None, + ), + ( + "new-account-discovered", + Ok(gmail_catalog), + Some("cloud:plugin_gmail"), + None, + Some(gmail), + ), + ( + "in-flight-auth-change", + Ok(cloud_catalog()), + None, + None, + None, + ), + ] { + let previous_discoveries = provider.calls.load(Ordering::SeqCst); + if let Some(account) = new_account { + // Do not manually refresh MCP: the auth transition must pick up this new inventory. + *apps_tools + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + vec![app_tool("gmail", "search_email")]; + let resources = provider + .resources + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + .context("cloud provider resource client missing")?; + let previous_auth = + resources.auth_cache_key_for_server(codex_mcp::CODEX_APPS_MCP_SERVER_NAME); + codex_login::auth::login_with_chatgpt_auth_tokens( + test.codex_home_path(), + "header.e30.changed", + account, + /*chatgpt_plan_type*/ None, + )?; + test.thread_manager.auth_manager().reload().await; + tokio::time::timeout(std::time::Duration::from_secs(10), async { + loop { + if resources.auth_cache_key_for_server(codex_mcp::CODEX_APPS_MCP_SERVER_NAME) + != previous_auth + && test + .codex + .thread_extension_data() + .get::() + .is_some_and(|snapshot| snapshot.plugins.is_empty()) + { + break; + } + tokio::task::yield_now().await; + } + }) + .await?; + assert_eq!(provider.calls.load(Ordering::SeqCst), previous_discoveries); + } + provider.set_reply(reply); + let previous_calls = recorded_apps_tool_calls(&server).await.len(); + let call_id = format!("{turn}-call"); + let mut events = vec![ev_response_created("resp")]; + if let Some((connector, action, _)) = tool_call { + events.push(ev_function_call_with_namespace( + &call_id, + &format!("mcp__codex_apps__{connector}"), + &format!("_{action}"), + "{}", + )); + } + events.push(ev_completed("resp")); + let mut responses = vec![sse(events)]; + if tool_call.is_some() { + responses.push(sse(vec![ + ev_response_created("after-tool"), + ev_completed("after-tool"), + ])); + } + let response = mount_sse_sequence(&server, responses).await; + *step_packages + .0 + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = None; + if turn == "in-flight-auth-change" { + provider.pause_discovery.store(true, Ordering::SeqCst); + tokio::time::timeout(Duration::from_secs(/*secs*/ 10), async { + tokio::try_join!( + test.submit_turn("What capabilities are available?"), + async { + provider.discovery_started.notified().await; + let resources = provider + .resources + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + .context("cloud provider resource client missing")?; + let previous_auth = resources + .auth_cache_key_for_server(codex_mcp::CODEX_APPS_MCP_SERVER_NAME); + // Rotate credentials without replacing the captured runtime/client. + codex_login::auth::login_with_chatgpt_auth_tokens( + test.codex_home_path(), + "header.e30.rotated", + "account-b", + /*chatgpt_plan_type*/ None, + )?; + test.thread_manager.auth_manager().reload().await; + while resources + .auth_cache_key_for_server(codex_mcp::CODEX_APPS_MCP_SERVER_NAME) + == previous_auth + { + tokio::task::yield_now().await; + } + provider.resume_discovery.notify_one(); + anyhow::Ok(()) + } + ) + }) + .await??; + // The old request must not republish its catalog after credentials change. + let state = test + .codex + .thread_extension_data() + .get::() + .context("plugin state missing")?; + assert_eq!(state.cloud_catalog(), None); + assert_eq!( + state + .contributor_state() + .cloud_generation + .as_ref() + .and_then(|generation| generation.catalog.clone()), + None, + ); + } else { + test.submit_turn("What capabilities are available?").await?; + } + let requests = response.requests(); + let request = &requests[0]; + if let Some((connector, action, plugin_name)) = tool_call { + assert_eq!(requests.len(), 2); + let calls = recorded_apps_tool_calls(&server).await; + let output = requests[1].function_call_output(&call_id); + assert_eq!(calls.len(), previous_calls + 1, "turn {turn}: {output}"); + let call = calls.last().context("expected an Apps tool call")?; + assert_eq!( + call["params"]["name"], + json!(format!("{connector}_{action}")) + ); + assert_eq!(call["params"]["arguments"], json!({})); + let (other_namespace, other_tool) = if connector == "calendar" { + ("mcp__codex_apps__gmail", "_search_email") + } else { + ("mcp__codex_apps__calendar", "_list_events") + }; + assert!(request.tool_by_name(other_namespace, other_tool).is_none()); + let tool = request + .tool_by_name( + &format!("mcp__codex_apps__{connector}"), + &format!("_{action}"), + ) + .expect("the current account's tool should be advertised"); + assert_eq!( + tool["description"] + .as_str() + .context("tool description missing")? + .contains(plugin_name), + expected_plugin.is_some(), + "cloud discovery should attribute tools to the current account's plugin" + ); + // The server echoes this URI in its result; it must reach the next model step. + assert!( + output["output"] + .as_str() + .context("tool output missing")? + .contains(&format!( + "connector://{connector}/tools/{connector}_{action}" + )), + "turn {turn}: {output}" + ); + } + assert_eq!( + *step_packages + .0 + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + Some(expected_plugin.into_iter().map(str::to_string).collect()), + "turn {turn}" + ); + assert_eq!( + provider.calls.load(Ordering::SeqCst), + previous_discoveries + 1, + "turn {turn}" + ); + } + Ok(()) +} diff --git a/codex-rs/ext/mcp/src/lib.rs b/codex-rs/ext/mcp/src/lib.rs index 62689947e9..2ca9e12a8a 100644 --- a/codex-rs/ext/mcp/src/lib.rs +++ b/codex-rs/ext/mcp/src/lib.rs @@ -7,10 +7,14 @@ use codex_extension_api::McpServerContributor; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::hosted_plugin_runtime_mcp_server_config; +mod cloud_plugin; #[cfg(test)] #[path = "event_stream_tests.rs"] mod event_stream_tests; mod plugin; +mod plugin_contributor; +mod plugin_contributor_state; +mod plugin_providers; mod stream_manager; pub use codex_core_plugins::PluginListQuery; @@ -18,6 +22,10 @@ pub use codex_core_plugins::PluginProvider; pub use codex_core_plugins::PluginProviderError; pub use codex_core_plugins::PluginProviderFuture; pub use codex_core_plugins::PluginProviderResult; +pub use plugin_contributor::install_plugin_providers; +pub use plugin_contributor::install_plugins; +pub use plugin_contributor_state::PluginsThreadState; +pub use plugin_providers::PluginProviders; pub use stream_manager::McpEventStreamManager; pub use stream_manager::McpEventStreamUpdate; @@ -59,16 +67,6 @@ pub fn install(builder: &mut ExtensionRegistryBuilder) { builder.mcp_server_contributor(std::sync::Arc::new(HostedPluginRuntimeExtension)); } -/// Installs discovery for MCP servers and apps declared by thread-selected plugins. -pub fn install_plugins( - builder: &mut ExtensionRegistryBuilder, - environment_manager: std::sync::Arc, -) { - builder.mcp_server_contributor(std::sync::Arc::new(plugin::PluginContributor::new( - environment_manager, - ))); -} - #[cfg(test)] #[path = "lib_tests.rs"] mod tests; diff --git a/codex-rs/ext/mcp/src/plugin.rs b/codex-rs/ext/mcp/src/plugin.rs index 4908e2aedd..21aac3743d 100644 --- a/codex-rs/ext/mcp/src/plugin.rs +++ b/codex-rs/ext/mcp/src/plugin.rs @@ -1,10 +1,8 @@ use codex_config::types::PluginMcpServerConfig; use codex_connectors_extension::PluginAppProvider; use codex_core::config::Config; -use codex_core_plugins::ExecutorPluginProvider; use codex_core_plugins::loader::apply_configured_plugin_mcp_server_policies; use codex_core_plugins::loader::configured_plugin_mcp_server_policies; -use codex_exec_server::EnvironmentManager; use codex_extension_api::ExtensionFuture; use codex_extension_api::McpServerContribution; use codex_extension_api::McpServerContributionContext; @@ -12,51 +10,19 @@ use codex_extension_api::McpServerContributor; use codex_features::Feature; use codex_protocol::capabilities::SelectedCapabilityRoot; use std::collections::HashMap; -use std::sync::Arc; -use std::sync::Mutex; use self::provider::PluginMcpProvider; +use crate::PluginsThreadState; +use crate::cloud_plugin::CloudPluginMetadata; +use crate::cloud_plugin::catalog_to_metadata; +use crate::plugin_contributor::PluginContributor; +use crate::plugin_contributor_state::CachedSelectedRoot; +use crate::plugin_contributor_state::SelectedPluginMetadata; mod discovery; mod provider; -/// Frozen MCP and app declarations for one selected package. -/// -/// Each server config retains the stable logical environment ID. Reconnection may replace the -/// concrete environment instance without changing that authority. -#[derive(Clone)] -struct SelectedPluginMetadata { - plugin_id: String, - plugin_display_name: String, - servers: Vec<(String, codex_config::McpServerConfig)>, - connector_ids: Vec, -} - -#[derive(Default)] -pub(crate) struct PluginContributorState { - cache: Mutex>, -} - -struct CachedSelectedRoot { - root: SelectedCapabilityRoot, - metadata: Option, -} - -pub(crate) struct PluginContributor { - plugin_provider: ExecutorPluginProvider, - mcp_provider: PluginMcpProvider, - app_provider: PluginAppProvider, -} - impl PluginContributor { - pub(crate) fn new(environment_manager: Arc) -> Self { - Self { - plugin_provider: ExecutorPluginProvider::new(Arc::clone(&environment_manager)), - mcp_provider: PluginMcpProvider, - app_provider: PluginAppProvider, - } - } - /// Returns metadata for one stable selected root. /// /// Successful resolution, including a root that is not a plugin or declares no capabilities, @@ -65,20 +31,19 @@ impl PluginContributor { #[tracing::instrument(name = "mcp.plugin.metadata.load", skip_all)] async fn metadata_for_root( &self, - state: &PluginContributorState, + state: &PluginsThreadState, selected_root: &SelectedCapabilityRoot, ) -> Option { if let Some(cached) = state - .cache - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) + .contributor_state() + .executor_cache .iter() .find(|cached| cached.root == *selected_root) { return cached.metadata.clone(); } - let plugin = match self.plugin_provider.resolve_bound(selected_root).await { + let plugin = match self.providers.executor.resolve_bound(selected_root).await { Ok(plugin) => plugin, Err(err) => { tracing::warn!( @@ -95,8 +60,8 @@ impl PluginContributor { // executor-owned files. Read them together so a remote environment only // pays for the slower read instead of both reads back-to-back. let (servers, app_declarations) = tokio::join!( - self.mcp_provider.load(&plugin), - self.app_provider.load(&plugin) + PluginMcpProvider.load(&plugin), + PluginAppProvider.load(&plugin) ); let servers = servers.unwrap_or_else(|err| { tracing::warn!( @@ -127,10 +92,8 @@ impl PluginContributor { } None => None, }; - let mut cache = state - .cache - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut state = state.contributor_state(); + let cache = &mut state.executor_cache; if let Some(cached) = cache.iter().find(|cached| cached.root == *selected_root) { return cached.metadata.clone(); } @@ -155,9 +118,17 @@ impl McpServerContributor for PluginContributor { let Some(thread_store) = context.thread_store() else { return Vec::new(); }; - let Some(selected_roots) = context.ready_selected_capability_roots() else { - return Vec::new(); - }; + // Cloud projection is gated independently; executor projection remains unchanged. + let cloud_plugins_enabled = self.providers.cloud.is_some() + && context.config().features.enabled(Feature::Plugins); + let state = thread_store.get_or_init(PluginsThreadState::default); + if !cloud_plugins_enabled || context.auth_changed() { + // Clear stale cloud metadata before projecting a replacement runtime. + state.contributor_state().cloud_generation = None; + } + let selected_roots = context + .ready_selected_capability_roots() + .unwrap_or_default(); let plugin_policies = configured_plugin_mcp_server_policies(&context.config().config_layer_stack); let mut contributions = Vec::new(); @@ -189,7 +160,6 @@ impl McpServerContributor for PluginContributor { )); } } else { - let state = thread_store.get_or_init(PluginContributorState::default); for (selection_order, selected_root) in selected_roots.iter().enumerate() { let Some(plugin) = self.metadata_for_root(&state, selected_root).await else { continue; @@ -204,6 +174,24 @@ impl McpServerContributor for PluginContributor { } } + // V1 has no remote plugin identities, so cloud packages are additive for now. + // Clone the current snapshot after executor loading and release the lock before projection. + if cloud_plugins_enabled && let Some(catalog) = state.cloud_catalog() { + for CloudPluginMetadata { + selection_order, + selected_root_id, + metadata, + } in catalog_to_metadata(&catalog) + { + contributions.extend(project_metadata( + context.config(), + /*plugin_policy*/ None, + selection_order, + &selected_root_id, + metadata, + )); + } + } contributions }) } diff --git a/codex-rs/ext/mcp/src/plugin_contributor.rs b/codex-rs/ext/mcp/src/plugin_contributor.rs new file mode 100644 index 0000000000..ad5a60f95f --- /dev/null +++ b/codex-rs/ext/mcp/src/plugin_contributor.rs @@ -0,0 +1,40 @@ +//! Coordinates cloud turn refreshes and executor MCP contributions with thread-owned catalogs. +//! +//! One contributor is registered for both lifecycles. Provider instances are shared, while +//! catalogs and executor metadata live together in thread extension data. + +use crate::PluginProviders; +use codex_core::config::Config; +use codex_core_plugins::ExecutorPluginProvider; +use codex_exec_server::EnvironmentManager; +use codex_extension_api::ExtensionRegistryBuilder; +use std::sync::Arc; + +pub(crate) struct PluginContributor { + pub(crate) providers: PluginProviders, +} + +/// Installs one coordinator for cloud turn-start and executor MCP callbacks. +/// The optional cloud provider never disables executor plugins, and Core's plugin policy applies +/// to both paths. +pub fn install_plugin_providers( + builder: &mut ExtensionRegistryBuilder, + providers: PluginProviders, +) { + let contributor = Arc::new(PluginContributor { providers }); + builder.thread_lifecycle_contributor(contributor.clone()); + builder.config_contributor(contributor.clone()); + builder.turn_lifecycle_contributor(contributor.clone()); + builder.mcp_server_contributor(contributor); +} + +/// Installs plugin support backed by selected environment roots. +pub fn install_plugins( + builder: &mut ExtensionRegistryBuilder, + environment_manager: Arc, +) { + install_plugin_providers( + builder, + PluginProviders::new(Arc::new(ExecutorPluginProvider::new(environment_manager))), + ); +} diff --git a/codex-rs/ext/mcp/src/plugin_contributor_state.rs b/codex-rs/ext/mcp/src/plugin_contributor_state.rs new file mode 100644 index 0000000000..ab088f37a2 --- /dev/null +++ b/codex-rs/ext/mcp/src/plugin_contributor_state.rs @@ -0,0 +1,122 @@ +//! Thread-owned cloud catalog and legacy selected-root cache, guarded by a short-lived mutex. + +use codex_core_plugins::PluginCatalog; +use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_mcp::McpResourceClient; +use codex_mcp::McpResourceClientAuthKey; +use codex_protocol::capabilities::SelectedCapabilityRoot; +use std::sync::Mutex; +use std::sync::MutexGuard; + +/// Plugin configuration and caches owned by one thread; mutable state stays extension-private. +#[derive(Default)] +pub struct PluginsThreadState { + state: Mutex, +} + +impl PluginsThreadState { + /// Returns an owned cloud snapshot only while its auth generation is current. + pub fn cloud_catalog(&self) -> Option { + self.contributor_state().cloud_catalog().cloned() + } + + pub(crate) fn contributor_state(&self) -> MutexGuard<'_, PluginContributorState> { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } +} + +/// `None` is unpublished; an empty catalog is a successful discovery result. +#[derive(Default)] +pub(crate) struct PluginContributorState { + pub(crate) plugins_enabled: bool, + pub(crate) cloud_generation: Option, + pub(crate) executor_cache: Vec, +} + +/// Keeps cloud plugin metadata within one published auth scope and Apps availability state. +pub(crate) struct CloudPluginGeneration { + pub(crate) auth_cache_key: Option, + pub(crate) mcp_resources: Option, + pub(crate) catalog: Option, +} + +impl CloudPluginGeneration { + pub(crate) fn new(mcp_resources: Option) -> Self { + let auth_cache_key = mcp_resources + .as_ref() + .map(|client| client.auth_cache_key_for_server(CODEX_APPS_MCP_SERVER_NAME)); + Self { + auth_cache_key, + mcp_resources, + catalog: None, + } + } + + pub(crate) fn is_current(&self) -> bool { + self.auth_cache_key + == self + .mcp_resources + .as_ref() + .map(|client| client.auth_cache_key_for_server(CODEX_APPS_MCP_SERVER_NAME)) + } +} + +impl PluginContributorState { + /// Returns the cloud catalog only while it belongs to the live auth generation. + fn cloud_catalog(&self) -> Option<&PluginCatalog> { + self.cloud_generation + .as_ref() + .filter(|generation| generation.is_current()) + .and_then(|generation| generation.catalog.as_ref()) + } + + pub(crate) fn prepare_cloud_generation( + &mut self, + mcp_resources: Option, + ) -> Option { + let auth_cache_key = mcp_resources + .as_ref() + .map(|client| client.auth_cache_key_for_server(CODEX_APPS_MCP_SERVER_NAME)); + if self + .cloud_generation + .as_ref() + .is_none_or(|generation| generation.auth_cache_key != auth_cache_key) + { + self.cloud_generation = Some(CloudPluginGeneration::new(mcp_resources)); + } + auth_cache_key + } + + pub(crate) fn publish_cloud_catalog( + &mut self, + auth_cache_key: Option<&McpResourceClientAuthKey>, + catalog: PluginCatalog, + ) -> bool { + let Some(generation) = self + .cloud_generation + .as_mut() + .filter(|generation| generation.auth_cache_key.as_ref() == auth_cache_key) + .filter(|generation| generation.is_current()) + else { + return false; + }; + generation.catalog = Some(catalog); + true + } +} + +pub(crate) struct CachedSelectedRoot { + pub(crate) root: SelectedCapabilityRoot, + pub(crate) metadata: Option, +} + +/// Frozen declarations retain the logical environment ID across executor reconnections. +#[derive(Clone)] +pub(crate) struct SelectedPluginMetadata { + pub(crate) plugin_id: String, + pub(crate) plugin_display_name: String, + pub(crate) servers: Vec<(String, codex_config::McpServerConfig)>, + pub(crate) connector_ids: Vec, +} diff --git a/codex-rs/ext/mcp/src/plugin_providers.rs b/codex-rs/ext/mcp/src/plugin_providers.rs new file mode 100644 index 0000000000..8255107556 --- /dev/null +++ b/codex-rs/ext/mcp/src/plugin_providers.rs @@ -0,0 +1,27 @@ +//! Host-assembled plugin sources; discovery results belong to the active thread, not this registry. + +use codex_core_plugins::ExecutorPluginProvider; +use codex_core_plugins::PluginProvider; +use std::sync::Arc; + +/// Cloud discovery is optional and never replaces executor plugin support. +pub struct PluginProviders { + // CCA owns and supplies the cloud provider implementation. + pub(crate) cloud: Option>, + // Keep concrete access to resolve_bound so loaders retain the exact executor filesystem. + pub(crate) executor: Arc, +} + +impl PluginProviders { + pub fn new(executor: Arc) -> Self { + Self { + cloud: None, + executor, + } + } + + pub fn with_cloud_provider(mut self, provider: Arc) -> Self { + self.cloud = Some(provider); + self + } +} diff --git a/codex-rs/ext/mcp/tests/plugin_mcp.rs b/codex-rs/ext/mcp/tests/plugin_mcp.rs index 49b184e429..7b05bea104 100644 --- a/codex-rs/ext/mcp/tests/plugin_mcp.rs +++ b/codex-rs/ext/mcp/tests/plugin_mcp.rs @@ -1,8 +1,20 @@ +use anyhow::Context; use codex_config::AppToolApproval; use codex_config::McpServerToolConfig; use codex_config::test_support::CloudConfigBundleFixture; +use codex_core::EnvironmentConfig; use codex_core::config::Config; use codex_core::config::ConfigBuilder; +use codex_core::windows_sandbox::WindowsSandboxLevelExt; +use codex_core_plugins::ExecutorPluginProvider; +use codex_core_plugins::PluginCatalog; +use codex_core_plugins::PluginCatalogEntry; +use codex_core_plugins::PluginIdentity; +use codex_core_plugins::PluginListQuery; +use codex_core_plugins::PluginProvider; +use codex_core_plugins::PluginProviderError; +use codex_core_plugins::PluginProviderFuture; +use codex_core_plugins::PluginSourceLocation; use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecutorCapabilityDiscoveryCache; use codex_exec_server::LOCAL_ENVIRONMENT_ID; @@ -12,12 +24,34 @@ use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::McpServerContribution; use codex_extension_api::McpServerContributionContext; use codex_features::Feature; +use codex_login::CodexAuth; +use codex_mcp_extension::PluginProviders; +use codex_mcp_extension::PluginsThreadState; +use codex_mcp_extension::install_plugin_providers; use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; +use codex_protocol::config_types::WindowsSandboxLevel; +use codex_protocol::models::PermissionProfileSnapshot; use codex_utils_path_uri::PathUri; +use core_test_support::apps_test_server::AppsTestServer; +use core_test_support::apps_test_server::recorded_apps_tool_calls; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_function_call_with_namespace; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_sequence; +use core_test_support::responses::sse; +use core_test_support::responses::start_mock_server; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::run_test_with_large_stack; +use core_test_support::test_codex::test_codex; use pretty_assertions::assert_eq; +use serde_json::json; use std::collections::HashMap; +use std::fs; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; type TestResult = Result<(), Box>; @@ -405,7 +439,7 @@ async fn raw_selected_plugin_contributions( None }; - Ok(registry.mcp_server_contributors()[0] + let contributions = registry.mcp_server_contributors()[0] .contribute(McpServerContributionContext::for_step( config, &thread_init, @@ -414,5 +448,252 @@ async fn raw_selected_plugin_contributions( &selected_capability_roots, executor_capability_discovery.as_ref(), )) - .await) + .await; + Ok(contributions) +} + +struct CloudCatalogFixture { + reply: Result, + calls: AtomicUsize, +} + +impl PluginProvider for CloudCatalogFixture { + fn list(&self, _query: PluginListQuery) -> PluginProviderFuture<'_, PluginCatalog> { + Box::pin(async { + self.calls.fetch_add(1, Ordering::SeqCst); + self.reply.clone().map_err(PluginProviderError::Message) + }) + } +} + +#[test] +fn cloud_plugins_append_apps_without_changing_executor_tools() -> anyhow::Result<()> { + run_test_with_large_stack( + "cloud-plugin-projection", + run_cloud_plugin_projection_scenarios, + ) +} + +async fn run_cloud_plugin_projection_scenarios() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + core_test_support::skip_if_remote!(Ok(()), "V1 plugin fixtures use host-local files"); + // Fresh sessions keep unavailable-cloud fallback independent of a previously cached catalog. + for (case, installed, cloud_enabled, cloud_available, cloud_empty) in [ + ("cloud_only", false, true, true, false), + ("cloud_and_executor", true, true, true, false), + ("unavailable_cloud", true, true, false, false), + ("rollout_disabled", true, false, true, false), + ("empty_cloud", true, true, true, true), + ] { + let server = start_mock_server().await; + let apps = AppsTestServer::mount(&server).await?; + let plugin_server = start_mock_server().await; + let plugin_mcp = AppsTestServer::mount_with_tools( + &plugin_server, + Arc::new(std::sync::Mutex::new(vec![json!({ + "name": "echo", "description": "Echo a note.", + "annotations": {"readOnlyHint": true}, + "inputSchema": {"type": "object", "properties": {"title": {"type": "string"}}} + })])), + ) + .await?; + wiremock::Mock::given(wiremock::matchers::method("POST")) + .and(wiremock::matchers::body_partial_json(json!({"method": "tools/call"}))) + .respond_with(|request: &wiremock::Request| { + let body: serde_json::Value = serde_json::from_slice(&request.body) + .unwrap_or_else(|error| panic!("invalid MCP request JSON: {error}")); + wiremock::ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", "id": body["id"], + "result": {"content": [], "structuredContent": {"echo": body["params"]["arguments"]["title"]}, "isError": false} + })) + }) + .with_priority(1) + .mount(&plugin_server).await; + let declaration = + json!({"url": format!("{}/api/codex/ps/mcp", plugin_mcp.chatgpt_base_url)}); + let cloud = Arc::new(CloudCatalogFixture { + reply: if cloud_empty { + Ok(PluginCatalog::default()) + } else if cloud_available { + Ok(PluginCatalog { + entries: vec![PluginCatalogEntry { + id: PluginIdentity::Remote { + remote_plugin_id: "plugin_notes".into(), + }, + display_name: "Cloud Notes".into(), + version: Some("1".into()), + mcp_servers: [("notes".into(), declaration.to_string())].into(), + connector_ids: vec!["calendar".into()], + locations: vec![PluginSourceLocation::Cloud { + resource_uri: "plugin://plugin_notes".into(), + bundle_uri: None, + }], + }], + warnings: Vec::new(), + }) + } else { + Err("plugin-service unavailable".into()) + }, + calls: AtomicUsize::new(0), + }); + let mut extensions = ExtensionRegistryBuilder::new(); + let mut providers = PluginProviders::new(Arc::new(ExecutorPluginProvider::new(Arc::new( + EnvironmentManager::default_for_tests(), + )))); + if cloud_enabled { + providers = providers.with_cloud_provider(cloud.clone()); + } + install_plugin_providers(&mut extensions, providers); + let mut builder = test_codex() + .with_extensions(Arc::new(extensions.build())) + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_model_info_override("gpt-5.5", |model| model.supports_search_tool = false) + .with_config(move |config| { + assert!(config.features.enable(Feature::Plugins).is_ok()); + assert!(config.features.enable(Feature::Apps).is_ok()); + assert!( + config + .features + .enable(Feature::ExecutorCapabilityDiscovery) + .is_ok() + ); + config.chatgpt_base_url = apps.chatgpt_base_url; + }); + let test = builder.build(&server).await?; + if installed { + let selection = test + .codex + .environment_selections() + .await + .into_iter() + .next() + .context("selected environment missing")?; + let root = selection.cwd.join("notes")?; + let root_path = root.to_abs_path()?; + fs::create_dir_all(root_path.join(".codex-plugin"))?; + fs::write(root_path.join(".codex-plugin/plugin.json"), json!({ + "name": "notes", "version": "1", "interface": {"displayName": "Installed Notes"}, + "mcpServers": "./.mcp.json" + }).to_string())?; + fs::write( + root_path.join(".mcp.json"), + json!({"mcpServers": {"notes": declaration}}).to_string(), + )?; + test.codex + .environment_ready( + &selection, + EnvironmentConfig { + allow_login_shell: false, + workspace_roots: selection.workspace_roots.clone(), + permission_profile: PermissionProfileSnapshot::legacy( + test.config.permissions.permission_profile().clone(), + ), + shell_environment_policy: Default::default(), + windows_sandbox_level: WindowsSandboxLevel::from_config(&test.config), + windows_sandbox_type: test.config.permissions.windows_sandbox_type, + use_legacy_landlock: test.config.features.use_legacy_landlock(), + exec_policy: None, + mcp_policy: None, + network_policy: None, + selected_capability_roots: vec![SelectedCapabilityRoot { + id: "notes@marketplace".into(), + location: CapabilityRootLocation::Environment { + environment_id: selection.environment_id.clone(), + path: root, + }, + }], + }, + ) + .await?; + } + let response = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("call"), + ev_function_call_with_namespace( + "notes-call", + "mcp__notes", + "echo", + r#"{"title":"release check"}"#, + ), + ev_completed("call"), + ]), + sse(vec![ + ev_assistant_message("done", "Finished checking Notes."), + ev_completed("done"), + ]), + ], + ) + .await; + test.submit_turn("Use the Notes MCP to check the release.") + .await?; + let requests = response.requests(); + assert_eq!(requests.len(), 2, "{case}"); + let tool = requests[0].tool_by_name("mcp__notes", "echo"); + assert_eq!(tool.is_some(), installed, "{case}"); + let output = requests[1].function_call_output("notes-call"); + if installed { + assert!( + tool.context("Notes tool missing")?["description"] + .as_str() + .context("tool description missing")? + .contains("Installed Notes") + ); + assert!( + output.to_string().contains("release check"), + "{case}: {output}" + ); + } + let calls = recorded_apps_tool_calls(&plugin_server).await; + assert_eq!(calls.len(), usize::from(installed), "{case}"); + if installed { + assert_eq!(calls[0]["params"]["name"], "echo"); + assert_eq!( + calls[0]["params"]["arguments"], + json!({"title": "release check"}) + ); + } + let initialized = plugin_server + .received_requests() + .await + .context("MCP requests missing")? + .iter() + .any(|request| { + serde_json::from_slice::(&request.body) + .is_ok_and(|body| body["method"] == "initialize") + }); + assert_eq!( + initialized, installed, + "{case}: cloud-only declarations must not start MCPs" + ); + assert_eq!( + cloud.calls.load(Ordering::SeqCst), + usize::from(cloud_enabled), + "{case}" + ); + let state = test + .codex + .thread_extension_data() + .get::() + .context("plugin state missing")?; + assert_eq!( + state.cloud_catalog().is_some(), + cloud_enabled && cloud_available, + "{case}" + ); + let calendar = requests[0] + .tool_by_name("mcp__codex_apps__calendar", "_list_events") + .context("Calendar tool missing")?; + assert_eq!( + calendar["description"] + .as_str() + .context("Apps description missing")? + .contains("Cloud Notes"), + cloud_enabled && cloud_available && !cloud_empty, + "{case}" + ); + test.codex.shutdown_and_wait().await?; + } + Ok(()) }