diff --git a/codex-rs/core-plugins/src/marketplace_staging_tests.rs b/codex-rs/core-plugins/src/marketplace_staging_tests.rs index a4b9b6ee04..487347f485 100644 --- a/codex-rs/core-plugins/src/marketplace_staging_tests.rs +++ b/codex-rs/core-plugins/src/marketplace_staging_tests.rs @@ -60,6 +60,13 @@ fn cleanup_preserves_live_process_and_reclaims_its_abandoned_checkout() -> Resul child.stdin.take().unwrap().write_all(b"x")?; assert!(child.wait()?.success()); assert!(payload.exists(), "process exit must bypass TempDir cleanup"); + // Other tests can fork while the earlier cleanup or checkout creation holds + // a lock. Wait for inherited locks and release them before cleanup's try-locks. + for name in [CLEANUP_LOCK, COORDINATION_LOCK] { + let lock = open_lock(&tmp_root.join(name))?; + lock.lock()?; + lock.unlock()?; + } drop( crate::configure_trusted_git_repository( &mut std::process::Command::new("git"), diff --git a/codex-rs/core-plugins/src/provider.rs b/codex-rs/core-plugins/src/provider.rs index 39106e8d4c..e2cc0aa8bc 100644 --- a/codex-rs/core-plugins/src/provider.rs +++ b/codex-rs/core-plugins/src/provider.rs @@ -13,7 +13,8 @@ use thiserror::Error; #[derive(Clone, Debug)] pub struct PluginListQuery { pub thread_id: String, - pub turn_id: String, + /// Absent for explicit discovery before turn admission. + pub turn_id: Option, pub mcp_resources: Option>, /// Already-acquired V2 inventories; empty for cloud listing. pub snapshots: Vec, diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 9f1fb465d8..801670b023 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -1392,6 +1392,13 @@ impl CodexThread { self.session.refresh_codex_apps_tools().await } + /// Refreshes cloud plugin and skill catalogs before input admission, without starting a turn. + /// Serializes publication with step preparation and marks MCP dirty for later reconciliation. + /// Accepts complete or partial catalogs; failed discovery remains retryable. + pub async fn refresh_cloud_plugins(&self) -> anyhow::Result<()> { + self.session.refresh_cloud_plugins().await + } + /// Returns the environments configured for future turns. pub async fn environment_selections(&self) -> Vec { self.session.configured_environment_selections().await diff --git a/codex-rs/core/src/session/cloud_plugins.rs b/codex-rs/core/src/session/cloud_plugins.rs new file mode 100644 index 0000000000..5fa2e4f8f0 --- /dev/null +++ b/codex-rs/core/src/session/cloud_plugins.rs @@ -0,0 +1,67 @@ +//! Refreshes independent cloud catalogs concurrently, serialized with step preparation. + +use super::Session; +use crate::McpRefreshReason; +use codex_extension_api::CloudCatalogRefreshInput; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use tokio::sync::Mutex; + +#[derive(Default)] +pub(super) struct CloudCatalogPreparation { + /// Serializes catalog refreshes and snapshot capture. + refresh_lock: Mutex<()>, + /// Requests resource invalidation on the next lock acquisition. + /// Survives cancellation while waiting. + refresh_pending: AtomicBool, +} + +impl Session { + pub(crate) async fn lock_cloud_catalogs(&self) -> tokio::sync::MutexGuard<'_, ()> { + let guard = self.cloud_catalog_preparation.refresh_lock.lock().await; + if self + .cloud_catalog_preparation + .refresh_pending + .swap(false, Ordering::AcqRel) + { + self.services.mcp_runtime.invalidate_resource_caches(); + } + guard + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "catalog publication must remain serialized with step preparation" + )] + pub(crate) async fn refresh_cloud_plugins(self: &Arc) -> anyhow::Result<()> { + self.cloud_catalog_preparation + .refresh_pending + .store(true, Ordering::Release); + let _catalogs = self.lock_cloud_catalogs().await; + self.refresh_mcp_if_dirty().await; + let results = futures::future::join_all( + self.services + .extensions + .turn_lifecycle_contributors() + .iter() + .map(|contributor| async move { + let result = contributor + .refresh_cloud_catalog(CloudCatalogRefreshInput { + session_store: &self.services.session_extension_data, + thread_store: &self.services.thread_extension_data, + }) + .await; + // Retain publication even if another provider is still pending at cancellation. + self.mark_mcp_runtime_dirty(McpRefreshReason::ExtensionCatalogChanged); + result + }), + ) + .await; + drop(_catalogs); + self.schedule_mcp_prewarm(); + let errors: Vec<_> = results.into_iter().filter_map(Result::err).collect(); + anyhow::ensure!(errors.is_empty(), "{}", errors.join("; ")); + Ok(()) + } +} diff --git a/codex-rs/core/src/session/config_refresh.rs b/codex-rs/core/src/session/config_refresh.rs index ff7a7be638..2e417271ff 100644 --- a/codex-rs/core/src/session/config_refresh.rs +++ b/codex-rs/core/src/session/config_refresh.rs @@ -194,9 +194,6 @@ impl Session { } state.session_configuration.original_config_do_not_use = Arc::clone(&config); drop(policy_guard); - if matches!(scope, RuntimeConfigRefresh::Mcp) { - self.services.mcp_runtime.invalidate_resource_caches(); - } self.mark_mcp_runtime_dirty(match scope { RuntimeConfigRefresh::User => McpRefreshReason::UserConfigReloaded, RuntimeConfigRefresh::Mcp => McpRefreshReason::McpConfigReloaded, diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 69ce341eb8..928a721215 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -482,20 +482,27 @@ impl Session { }; if is_active_regular_turn { // Refresh cloud catalogs without holding the MCP permit or active-turn lock. - for contributor in self.services.extensions.turn_lifecycle_contributors() { - if contributor - .refresh_catalog_for_step(codex_extension_api::CatalogRefreshInput { - turn_id: &turn_context.sub_id, - session_store: &self.services.session_extension_data, - thread_store: &self.services.thread_extension_data, - turn_store: turn_context.extension_data.as_ref(), - }) - .await - { - // Preserve MCP updates if a later refresh is cancelled. - self.mark_mcp_runtime_dirty(McpRefreshReason::ExtensionCatalogChanged); - } - } + futures::future::join_all( + self.services + .extensions + .turn_lifecycle_contributors() + .iter() + .map(|contributor| async move { + if contributor + .refresh_catalog_for_step(codex_extension_api::CatalogRefreshInput { + turn_id: &turn_context.sub_id, + session_store: &self.services.session_extension_data, + thread_store: &self.services.thread_extension_data, + turn_store: turn_context.extension_data.as_ref(), + }) + .await + { + // Preserve each publication if the refresh batch is cancelled. + self.mark_mcp_runtime_dirty(McpRefreshReason::ExtensionCatalogChanged); + } + }), + ) + .await; self.refresh_mcp_if_dirty().await; } self.wait_for_apps_tool_refresh().await; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 7c11e6d4ee..f323f7c850 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -266,6 +266,7 @@ mod reasoning_effort; mod submission; pub(crate) use reasoning_effort::RequestEffortUsage; pub(crate) use submission::Submission; +mod cloud_plugins; mod input_queue; mod mcp; mod mcp_prewarm; @@ -4325,7 +4326,7 @@ impl Session { #[expect( clippy::await_holding_invalid_type, - reason = "step settings and both environment snapshots must be captured under active_turn then state" + reason = "step settings and environments require active_turn then state; cloud snapshots require the catalog publication lock" )] #[tracing::instrument(name = "step_context.capture", level = "info", skip_all)] async fn capture_step_context_inner( @@ -4495,6 +4496,7 @@ impl Session { { required_servers.push(codex_mcp::CODEX_APPS_MCP_SERVER_NAME.to_owned()); } + let catalogs = self.lock_cloud_catalogs().await; let (mcp, prepared_recommendations) = tokio::join!( // MCP refresh can be large; keep it off the sampling request's stack. Box::pin(self.mcp_runtime_for_step( @@ -4615,6 +4617,7 @@ impl Session { .or_cancel(cancellation_token) .await?; } + drop(catalogs); extension_data.insert(selected_plugins); let tool_router = turn::built_tools( self.as_ref(), diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 9d3687700f..71dfc8f711 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -99,6 +99,8 @@ pub(crate) struct Session { pub(super) multi_agent_version: OnceLock, /// Owns invalidation and serializes refreshes without blocking captured calls. pub(super) mcp_refresh: McpRefresh, + /// Serializes cloud publication with capture of a step's plugin and skill snapshots. + pub(super) cloud_catalog_preparation: super::cloud_plugins::CloudCatalogPreparation, /// Non-owning lookup for approval data retained by running MCP invocations. pub(crate) mcp_tool_approval_metadata: std::sync::Mutex, pub(super) mcp_elicitation_reviewer_handle: OnceLock, @@ -1981,6 +1983,7 @@ impl Session { windows_sandbox_proxy_settings_mode, multi_agent_version, mcp_refresh: McpRefresh::new(), + cloud_catalog_preparation: Default::default(), mcp_tool_approval_metadata: Default::default(), mcp_elicitation_reviewer_handle: OnceLock::new(), mcp_elicitation_lifecycle_handle: OnceLock::new(), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 48707528cd..8731bfe3ec 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -7044,6 +7044,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { codex_sandboxing::WindowsSandboxProxySettingsMode::Reconcile, multi_agent_version: OnceLock::from(config.multi_agent_version_from_features()), mcp_refresh: McpRefresh::new(), + cloud_catalog_preparation: Default::default(), mcp_tool_approval_metadata: Default::default(), mcp_elicitation_reviewer_handle: OnceLock::new(), mcp_elicitation_lifecycle_handle: OnceLock::new(), @@ -9391,6 +9392,7 @@ where codex_sandboxing::WindowsSandboxProxySettingsMode::Reconcile, multi_agent_version: OnceLock::from(config.multi_agent_version_from_features()), mcp_refresh: McpRefresh::new(), + cloud_catalog_preparation: Default::default(), mcp_tool_approval_metadata: Default::default(), mcp_elicitation_reviewer_handle: OnceLock::new(), mcp_elicitation_lifecycle_handle: OnceLock::new(), diff --git a/codex-rs/core/src/session/turn_start_mcp_tests.rs b/codex-rs/core/src/session/turn_start_mcp_tests.rs index 096d03896f..21c806cd16 100644 --- a/codex-rs/core/src/session/turn_start_mcp_tests.rs +++ b/codex-rs/core/src/session/turn_start_mcp_tests.rs @@ -1,4 +1,4 @@ -//! Step discovery observes ready MCP connections; prewarm skips discovery and cancellation releases preparation. +//! Concurrent catalog discovery retains publication and releases preparation on cancellation. use super::*; use crate::McpRefreshReason; @@ -8,6 +8,7 @@ use pretty_assertions::assert_eq; struct Recorder { client: McpResourceClient, + ready: Arc, pause: std::sync::atomic::AtomicBool, keys: std::sync::Mutex>, discovered: tokio::sync::Notify, @@ -23,6 +24,7 @@ impl codex_extension_api::TurnLifecycleContributor for Recorder { .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .push(self.client.cache_key()); + self.ready.wait().await; self.discovered.notify_one(); if self.pause.swap(false, std::sync::atomic::Ordering::Relaxed) { std::future::pending::<()>().await; @@ -32,17 +34,27 @@ impl codex_extension_api::TurnLifecycleContributor for Recorder { } } +/// Step discovery overlaps providers on ready connections; cancellation releases the catalog lock. #[tokio::test] -async fn step_discovery_reconciles_mcp_skips_prewarm_and_cancels_preparation() { +async fn step_discovery_reconciles_mcp_skips_prewarm_and_releases_lock_on_cancellation() { let (mut session, turn_context) = make_session_and_context().await; + let ready = Arc::new(tokio::sync::Barrier::new(2)); let recorder = Arc::new(Recorder { client: McpResourceClient::new(Arc::clone(&session.services.mcp_runtime)), + ready: Arc::clone(&ready), pause: std::sync::atomic::AtomicBool::new(true), keys: Default::default(), discovered: Default::default(), }); let mut builder = codex_extension_api::ExtensionRegistryBuilder::new(); builder.turn_lifecycle_contributor(recorder.clone()); + builder.turn_lifecycle_contributor(Arc::new(Recorder { + client: recorder.client.clone(), + ready, + pause: std::sync::atomic::AtomicBool::new(false), + keys: Default::default(), + discovered: Default::default(), + })); session.services.extensions = Arc::new(builder.build()); let session = Arc::new(session); let turn_context = Arc::new(turn_context); @@ -95,7 +107,7 @@ async fn step_discovery_reconciles_mcp_skips_prewarm_and_cancels_preparation() { session.capture_step_context(Arc::clone(&turn_context), &CancellationToken::new()), ) .await - .expect("cancelled discovery must allow the next step") + .expect("cancelled discovery must release the catalog lock") .expect("next step should prepare"); assert!(recorder.client.cache_key() == prepared_key); @@ -112,3 +124,75 @@ async fn step_discovery_reconciles_mcp_skips_prewarm_and_cancels_preparation() { assert!(keys[2] == recorder.client.cache_key()); session.abort_all_tasks(TurnAbortReason::Interrupted).await; } + +/// Explicit refresh overlaps providers under one lock and retains publication while a sibling waits. +#[tokio::test] +async fn explicit_cloud_refresh_runs_concurrently_and_awaits_all_results() { + struct CatalogGate { + entered: Arc, + finish: tokio::sync::Notify, + result: Result<(), &'static str>, + } + + impl codex_extension_api::TurnLifecycleContributor for CatalogGate { + fn refresh_cloud_catalog<'a>( + &'a self, + _input: codex_extension_api::CloudCatalogRefreshInput<'a>, + ) -> codex_extension_api::ExtensionFuture<'a, Result<(), String>> { + Box::pin(async move { + self.entered.wait().await; + self.finish.notified().await; + self.result.map_err(str::to_string) + }) + } + } + + let (mut session, _) = make_session_and_context().await; + let entered = Arc::new(tokio::sync::Barrier::new(3)); + let failing = Arc::new(CatalogGate { + entered: Arc::clone(&entered), + finish: Default::default(), + result: Err("plugin discovery failed"), + }); + let succeeding = Arc::new(CatalogGate { + entered: Arc::clone(&entered), + finish: Default::default(), + result: Ok(()), + }); + let mut builder = codex_extension_api::ExtensionRegistryBuilder::new(); + builder.turn_lifecycle_contributor(failing.clone()); + builder.turn_lifecycle_contributor(succeeding.clone()); + session.services.extensions = Arc::new(builder.build()); + let session = Arc::new(session); + let mut refresh = Box::pin(session.refresh_cloud_plugins()); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + tokio::select! { + result = &mut refresh => panic!("refresh completed before both providers started: {result:?}"), + _ = entered.wait() => {} + } + }) + .await + .expect("both providers must start before either finishes"); + assert!(!session.mcp_refresh.is_pending()); + + let mut lock = Box::pin(session.lock_cloud_catalogs()); + assert!(futures::poll!(&mut lock).is_pending()); + drop(lock); + + failing.finish.notify_one(); + assert!(futures::poll!(&mut refresh).is_pending()); + assert!(session.mcp_refresh.is_pending()); + succeeding.finish.notify_one(); + assert_eq!( + refresh.await.expect_err("one provider failed").to_string(), + "plugin discovery failed" + ); + drop( + tokio::time::timeout( + std::time::Duration::from_secs(5), + session.lock_cloud_catalogs(), + ) + .await + .expect("completed refresh must release the catalog lock"), + ); +} diff --git a/codex-rs/core/tests/suite/scenarios_frozen_cloud_skills.rs b/codex-rs/core/tests/suite/scenarios_frozen_cloud_skills.rs index 7a758f60b8..1cd9b4ef3c 100644 --- a/codex-rs/core/tests/suite/scenarios_frozen_cloud_skills.rs +++ b/codex-rs/core/tests/suite/scenarios_frozen_cloud_skills.rs @@ -87,8 +87,13 @@ async fn frozen_cloud_skills_with_named_executor_request_history() -> Result<()> }) .build(&server) .await?; - let executor = - core_test_support::exec_server::ExecServerProcess::start_with_version("0.164.0").await?; + test.codex + .refresh_cloud_plugins() + .await + .expect("explicit startup discovery"); + let executor = core_test_support::exec_server::ExecServerProcess::start_with_version("0.164.0") + .await + .expect("executor fixture"); test.thread_manager .environment_manager() .upsert_environment( diff --git a/codex-rs/core/tests/suite/skills_extension.rs b/codex-rs/core/tests/suite/skills_extension.rs index 5a61c4ccbe..bb68c4be1d 100644 --- a/codex-rs/core/tests/suite/skills_extension.rs +++ b/codex-rs/core/tests/suite/skills_extension.rs @@ -161,7 +161,7 @@ impl SkillProvider for PausedCatalogSkillProvider { } if self.first_discovery.swap(/*val*/ false, Ordering::SeqCst) { self.started - .send(query.turn_id.clone()) + .send(query.turn_id.clone().expect("turn discovery")) .expect("discovery receiver"); self.resume .acquire() diff --git a/codex-rs/core/tests/suite/skills_extension/cloud_lifecycle_tests.rs b/codex-rs/core/tests/suite/skills_extension/cloud_lifecycle_tests.rs index 07a01fa035..a9cff19118 100644 --- a/codex-rs/core/tests/suite/skills_extension/cloud_lifecycle_tests.rs +++ b/codex-rs/core/tests/suite/skills_extension/cloud_lifecycle_tests.rs @@ -90,6 +90,197 @@ impl ThreadLifecycleContributor for SessionResources { } } +async fn mount_cloud_skill_list( + server: &wiremock::MockServer, + call_id: &str, +) -> responses::ResponseMock { + responses::mount_sse_sequence( + server, + vec![ + sse(vec![ + ev_response_created(call_id), + responses::ev_function_call_with_namespace( + call_id, + "skills", + "list", + &json!({"authority":{"kind":"cloud"}}).to_string(), + ), + ev_completed(call_id), + ]), + sse(vec![ev_response_created("done"), ev_completed("done")]), + ], + ) + .await +} + +/// Hosts refresh before admission; publication is serialized with step preparation. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn explicit_cloud_refresh_publishes_before_admission_and_reuses_catalogs() -> Result<()> { + #[derive(Default)] + struct CatalogGate { + pause: AtomicBool, + entered: tokio::sync::Notify, + resume: tokio::sync::Notify, + } + + impl codex_extension_api::TurnLifecycleContributor for CatalogGate { + fn refresh_catalog_for_step<'a>( + &'a self, + _: codex_extension_api::CatalogRefreshInput<'a>, + ) -> ExtensionFuture<'a, bool> { + Box::pin(async move { + if self.pause.swap(false, Ordering::SeqCst) { + self.entered.notify_one(); + self.resume.notified().await; + } + false + }) + } + } + + core_test_support::skip_if_no_network!(Ok(())); + let server = responses::start_mock_server().await; + let apps = AppsTestServer::mount(&server).await?; + let provider = Arc::new(CloudTurnProvider { + reply: Mutex::new(Ok(Some(("calendar", 1)))), + lists: AtomicUsize::new(0), + partial: AtomicBool::new(false), + reads: AtomicUsize::new(0), + resource_key: Mutex::new(None), + }); + let mut extensions = ExtensionRegistryBuilder::new(); + let resources = Arc::new(SessionResources::default()); + extensions.thread_lifecycle_contributor(resources.clone()); + install_with_providers( + &mut extensions, + SkillProviders::new().with_cloud_provider(provider.clone()), + |_config: &Config| SkillsExtensionConfig { + include_instructions: true, + catalog_deltas_enabled: false, + omit_descriptions: false, + max_context_tokens: None, + bundled_skills_enabled: false, + cloud_skill_enabled: true, + shadow_selection_enabled: false, + }, + ); + let gate = Arc::new(CatalogGate::default()); + extensions.turn_lifecycle_contributor(gate.clone()); + let mut builder = apps_enabled_builder(apps.chatgpt_base_url) + .with_exec_server_url("none") + .with_extensions(Arc::new(extensions.build())) + .with_auth(CodexAuth::from_external_chatgpt_tokens( + "header.e30.initial", + "account-a", + /*chatgpt_plan_type*/ None, + )?) + .with_config(|config| { + config.include_skill_instructions = true; + config.cloud_skill_enabled = true; + }); + let test = builder.build_with_auto_env(&server).await?; + for (index, plugin) in ["calendar", "gmail"].into_iter().enumerate() { + *provider.reply.lock().expect("reply") = Ok(Some((plugin, 1))); + provider.partial.store(plugin == "gmail", Ordering::SeqCst); + test.codex.refresh_cloud_plugins().await?; + assert_eq!(provider.lists.load(Ordering::SeqCst), index + 1); + test.codex.refresh_codex_apps_tools().await?; + let current = test.codex.config().await; + let mut next = (*current).clone(); + next.config_layer_stack = next + .config_layer_stack + .with_session_plugins(toml::toml! { _default = { enabled = false } }.into())?; + assert_eq!( + test.codex.refresh_mcp_config(current, next).await, + codex_core::ConfigRefreshOutcome::Published, + ); + let list_id = format!("list-{index}"); + let response = mount_cloud_skill_list(&server, &list_id).await; + test.submit_turn("List available cloud skills.").await?; + assert_eq!( + provider.lists.load(Ordering::SeqCst), + index + 1, + "tool/config refresh and steps must reuse explicitly published catalogs" + ); + let requests = response.requests(); + let listed = requests[1] + .function_call_output_text(&list_id) + .expect("skill list"); + assert!(listed.contains(plugin), "expected {plugin}: {listed}"); + } + provider.partial.store(false, Ordering::SeqCst); + let accepted_key = provider + .resource_key + .lock() + .expect("resource key lock") + .clone() + .expect("accepted catalog should have an Apps resource key"); + *provider.reply.lock().expect("reply") = Err(()); + assert!(test.codex.refresh_cloud_plugins().await.is_err()); + assert!( + provider + .resource_key + .lock() + .expect("resource key lock") + .as_ref() + .is_some_and(|key| key != &accepted_key), + "explicit refresh must invalidate resources before discovery" + ); + *provider.reply.lock().expect("reply") = Ok(Some(("gmail", 2))); + let response = mount_cloud_skill_list(&server, "retried-list").await; + test.submit_turn("List skills after retrying discovery.") + .await?; + assert_eq!(provider.lists.load(Ordering::SeqCst), 4); + assert!( + response.requests()[1] + .function_call_output_text("retried-list") + .expect("skill list") + .contains("Cloud gmail revision 2") + ); + + // Cancelling an explicit refresh while it waits must preserve pending invalidation. + let previous_calls = provider.lists.load(Ordering::SeqCst); + *provider.reply.lock().expect("reply") = Ok(Some(("gmail", 3))); + gate.pause.store(true, Ordering::SeqCst); + let response = mount_cloud_skill_list(&server, "cancelled-refresh-list").await; + tokio::time::timeout(Duration::from_secs(10), async { + tokio::try_join!( + test.submit_turn("List skills after a cancelled refresh request."), + async { + gate.entered.notified().await; + let client = resources.0.lock().expect("resources").clone().unwrap(); + let key = client.server_cache_key(CODEX_APPS_MCP_SERVER_NAME); + let mut refresh = Box::pin(test.codex.refresh_cloud_plugins()); + assert!(futures::poll!(&mut refresh).is_pending()); + drop(refresh); + assert!(client.server_cache_key(CODEX_APPS_MCP_SERVER_NAME) == key); + assert_eq!(provider.lists.load(Ordering::SeqCst), previous_calls); + gate.resume.notify_one(); + anyhow::Ok(()) + } + ) + }) + .await??; + assert_eq!(provider.lists.load(Ordering::SeqCst), previous_calls + 1); + assert!( + response.requests()[1] + .function_call_output_text("cancelled-refresh-list") + .expect("skill list") + .contains("Cloud gmail revision 2") + ); + let response = mount_cloud_skill_list(&server, "after-cancellation-list").await; + test.submit_turn("List the catalog recovered after cancellation.") + .await?; + assert_eq!(provider.lists.load(Ordering::SeqCst), previous_calls + 1); + assert!( + response.requests()[1] + .function_call_output_text("after-cancellation-list") + .expect("skill list") + .contains("Cloud gmail revision 3") + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn cloud_skills_reuse_cache_and_invalidate_on_connection_or_auth_change() -> Result<()> { core_test_support::skip_if_no_network!(Ok(())); diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 0b65d6fb34..e6ed59bc88 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -63,6 +63,7 @@ pub use tool_lifecycle::ToolTimingInput; pub use turn_input::TurnInputContext; pub use turn_input::TurnInputEnvironment; pub use turn_lifecycle::CatalogRefreshInput; +pub use turn_lifecycle::CloudCatalogRefreshInput; pub use turn_lifecycle::TurnAbortInput; pub use turn_lifecycle::TurnErrorInput; pub use turn_lifecycle::TurnStartInput; @@ -215,6 +216,12 @@ pub trait ThreadLifecycleContributor: Send + Sync { /// Implementations should use these callbacks to seed, observe, or clear /// extension-private turn state. The host exposes stable identifiers and /// extension stores instead of core runtime objects. +/// +/// Core serializes explicit catalog refresh and step preparation per thread. +/// Within a refresh, contributors run concurrently and must not depend on each +/// other's publication order. +/// Cloud discovery runs through those refresh callbacks; hosts requesting discovery +/// should use Core's refresh API rather than invoke hooks directly. pub trait TurnLifecycleContributor: Send + Sync { /// Selects the start phase using the current thread policy. Disabled callbacks /// may retain the default phase and return without doing any work. @@ -233,6 +240,7 @@ pub trait TurnLifecycleContributor: Send + Sync { } /// Refreshes catalogs before a step; returns true when MCP needs rebuilding. + /// Called under the same per-thread catalog lock as explicit refresh and step preparation. fn refresh_catalog_for_step<'a>( &'a self, _input: CatalogRefreshInput<'a>, @@ -240,6 +248,19 @@ pub trait TurnLifecycleContributor: Send + Sync { Box::pin(async { false }) } + /// Publishes cloud catalogs after explicit invalidation, without requiring a turn. + /// Core marks MCP dirty after publication. Accepted complete or partial catalogs + /// replace the previous result; failed or superseded discovery returns an error. + /// Core invalidates resource keys before calling. MCP-backed catalogs record + /// the new key only on acceptance, keeping failed or cancelled refreshes stale. + /// Called under the same per-thread catalog lock as step preparation. + fn refresh_cloud_catalog<'a>( + &'a self, + _input: CloudCatalogRefreshInput<'a>, + ) -> ExtensionFuture<'a, Result<(), String>> { + Box::pin(async { Ok(()) }) + } + /// Observes a completed item without changing it or delaying streamed deltas. fn on_item_completed<'a>( &'a self, 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 c6c7571de5..a9f4622e4f 100644 --- a/codex-rs/ext/extension-api/src/contributors/turn_lifecycle.rs +++ b/codex-rs/ext/extension-api/src/contributors/turn_lifecycle.rs @@ -42,6 +42,22 @@ pub struct CatalogRefreshInput<'a> { pub turn_store: &'a ExtensionData, } +/// Inputs for explicit cloud discovery, including before a thread has any turns. +#[derive(Clone, Copy)] +pub struct CloudCatalogRefreshInput<'a> { + pub session_store: &'a ExtensionData, + pub thread_store: &'a ExtensionData, +} + +impl<'a> From> for CloudCatalogRefreshInput<'a> { + fn from(input: CatalogRefreshInput<'a>) -> Self { + Self { + session_store: input.session_store, + thread_store: input.thread_store, + } + } +} + /// Input supplied when the host completes a turn. pub struct TurnStopInput<'a> { /// Store scoped to the host session runtime. diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 7295cd03b8..5d87b363e9 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -58,6 +58,7 @@ pub use contributors::ApprovalDecision; pub use contributors::ApprovalDecisionInput; pub use contributors::ApprovalReviewContributor; pub use contributors::CatalogRefreshInput; +pub use contributors::CloudCatalogRefreshInput; pub use contributors::CodeModeCellTiming; pub use contributors::CommandStartInput; pub use contributors::ConfigContributor; diff --git a/codex-rs/ext/mcp/src/cloud_plugin.rs b/codex-rs/ext/mcp/src/cloud_plugin.rs index 4c7429fa5a..3b12c8df28 100644 --- a/codex-rs/ext/mcp/src/cloud_plugin.rs +++ b/codex-rs/ext/mcp/src/cloud_plugin.rs @@ -1,6 +1,7 @@ -//! Refreshes cloud plugins before a step when their auth or Apps resources change. -//! Partial catalogs and catalogs with warnings remain usable and retry with thread-owned exponential backoff. -//! Accepted complete and partial catalogs replace prior entries; failures retain same-auth fallback. +//! Refreshes cloud plugins explicitly or when their accepted auth/resource generation is stale. +//! Partial catalogs and warnings remain usable and retry with exponential backoff. +//! Complete and partial results replace the catalog; failed attempts retain same-auth fallback. +//! Core serializes discovery and step snapshots with its per-thread catalog lock. use crate::PluginListQuery; use crate::PluginsThreadState; @@ -9,7 +10,8 @@ use codex_core::config::Config; use codex_core_plugins::PluginCatalog; use codex_core_plugins::PluginIdentity; use codex_extension_api::CatalogRefreshBackoff; -use codex_extension_api::CatalogRefreshInput; +use codex_extension_api::CatalogRefreshInput as StepCatalogRefreshInput; +use codex_extension_api::CloudCatalogRefreshInput; use codex_extension_api::ConfigContributor; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionFuture; @@ -17,7 +19,6 @@ use codex_extension_api::McpServerContribution; use codex_extension_api::ThreadLifecycleContributor; use codex_extension_api::ThreadStartInput; use codex_extension_api::TurnLifecycleContributor; -use codex_features::Feature; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::McpResourceClient; use std::collections::HashSet; @@ -26,71 +27,59 @@ const MAX_PROJECTED_CLOUD_PLUGINS: usize = 100; pub(crate) const MAX_PLUGIN_DISPLAY_NAME_CHARS: usize = 64; pub(crate) const MAX_CONNECTOR_IDS_PER_PLUGIN: usize = 64; +struct CloudPluginRefreshError { + error: String, + catalog_cleared: bool, + /// Cooldown for the step caller to store after a current-generation failure. + /// Superseded attempts return `None`, allowing the next step to retry immediately. + retry_backoff: Option, +} + +fn cloud_catalog_query( + input: CloudCatalogRefreshInput<'_>, + turn_id: Option<&str>, +) -> PluginListQuery { + PluginListQuery { + thread_id: input.thread_store.level_id().to_string(), + turn_id: turn_id.map(str::to_owned), + mcp_resources: input.session_store.get::(), + snapshots: Vec::new(), + } +} + impl PluginContributor { - #[expect( - clippy::await_holding_invalid_type, - reason = "cloud discovery must remain serialized through fetching and catalog publication" - )] - async fn refresh_cloud_catalog(&self, input: CatalogRefreshInput<'_>) -> bool { + /// Returns whether publishing or clearing a catalog requires MCP reprojection. + async fn discover_cloud_catalog( + &self, + state: &PluginsThreadState, + query: PluginListQuery, + ) -> Result { let Some(provider) = &self.providers.cloud else { - return false; + return Ok(false); }; - let Some(state) = input.thread_store.get::() else { - return false; + let (auth_cache_key, resource_cache_key, catalog_cleared) = { + let mut state = state.contributor_state(); + if !state.plugins_enabled { + return Ok(false); + } + let resource_cache_key = query + .mcp_resources + .as_ref() + .and_then(|client| client.server_cache_key(CODEX_APPS_MCP_SERVER_NAME)); + let previous_catalog_auth = state + .cloud_generation + .as_ref() + .filter(|generation| generation.catalog.is_some()) + .map(|generation| generation.auth_cache_key.clone()); + let auth_cache_key = + state.prepare_cloud_generation(query.mcp_resources.as_deref().cloned()); + let catalog_cleared = + previous_catalog_auth.is_some_and(|previous| previous != auth_cache_key); + (auth_cache_key, resource_cache_key, catalog_cleared) }; - // Wait without holding catalog state or the MCP runtime permit, then recheck the cache. - let _refresh = state.cloud_refresh_lock.lock().await; - let query = PluginListQuery { - thread_id: input.thread_store.level_id().to_string(), - turn_id: input.turn_id.to_string(), - mcp_resources: input.session_store.get::(), - snapshots: Vec::new(), - }; - let (auth_cache_key, resource_cache_key, catalog_cleared) = - { - let mut state = state.contributor_state(); - if !state.plugins_enabled { - return false; - } - let auth = query - .mcp_resources - .as_ref() - .map(|client| client.auth_cache_key_for_server(CODEX_APPS_MCP_SERVER_NAME)); - let resource_cache_key = query - .mcp_resources - .as_ref() - .and_then(|client| client.server_cache_key(CODEX_APPS_MCP_SERVER_NAME)); - if state.cloud_retry.as_ref().is_some_and(|retry| { - retry.is_pending(auth.as_ref(), resource_cache_key.as_ref()) - }) { - return false; - } - let previous_catalog_auth = state - .cloud_generation - .as_ref() - .filter(|generation| generation.catalog.is_some()) - .map(|generation| generation.auth_cache_key.clone()); - if let Some(generation) = &state.cloud_generation - && previous_catalog_auth.as_ref() == Some(&auth) - && generation.is_auth_current() - && generation.resource_cache_key == resource_cache_key - && generation - .catalog - .as_ref() - .is_some_and(|catalog| !catalog.is_partial && catalog.warnings.is_empty()) - { - state.cloud_retry = None; - return false; - } - let auth_cache_key = - state.prepare_cloud_generation(query.mcp_resources.as_deref().cloned()); - let catalog_cleared = - previous_catalog_auth.is_some_and(|previous| previous != auth_cache_key); - (auth_cache_key, resource_cache_key, catalog_cleared) - }; let result = provider.list(query).await; let mut state = state.contributor_state(); - let catalog_published = match result { + match result { Ok(catalog) => { let needs_retry = catalog.is_partial || !catalog.warnings.is_empty(); let published = state.publish_cloud_catalog( @@ -112,26 +101,38 @@ impl PluginContributor { tracing::debug!( "cloud plugin discovery finished after auth or resources changed" ); + return Err(CloudPluginRefreshError { + error: "cloud plugin discovery was superseded".to_string(), + catalog_cleared, + retry_backoff: None, + }); } - published } - Err(_) => { - if state.cloud_generation.as_ref().is_some_and(|generation| { - generation.auth_cache_key == auth_cache_key - && generation.is_auth_current() - && generation.current_resource_cache_key() == resource_cache_key - }) { - state.cloud_retry = Some(CatalogRefreshBackoff::after_failure( - auth_cache_key, - resource_cache_key, - state.cloud_retry.as_ref(), - )); - } + Err(error) => { + let retry_backoff = state + .cloud_generation + .as_ref() + .filter(|generation| { + generation.auth_cache_key == auth_cache_key + && generation.is_auth_current() + && generation.current_resource_cache_key() == resource_cache_key + }) + .map(|_| { + CatalogRefreshBackoff::after_failure( + auth_cache_key, + resource_cache_key, + state.cloud_retry.as_ref(), + ) + }); tracing::warn!("cloud plugin discovery failed"); - false + return Err(CloudPluginRefreshError { + error: error.to_string(), + catalog_cleared, + retry_backoff, + }); } - }; - catalog_cleared || catalog_published + } + Ok(true) } } @@ -144,7 +145,7 @@ impl ThreadLifecycleContributor for PluginContributor { let state = input.thread_store.get_or_init(PluginsThreadState::default); { let mut state = state.contributor_state(); - state.plugins_enabled = input.config.features.enabled(Feature::Plugins); + state.plugins_enabled = self.providers.cloud_enabled(input.config); state.extension_metrics = input.extension_metrics; } if let Some(resources) = &input.mcp_resource_client { @@ -162,14 +163,14 @@ impl ConfigContributor for PluginContributor { _previous_config: &Config, new_config: &Config, ) { - let plugins_enabled = new_config.features.enabled(Feature::Plugins); + let plugins_enabled = self.providers.cloud_enabled(new_config); let state = thread_store.get_or_init(PluginsThreadState::default); let mut state = state.contributor_state(); if state.plugins_enabled != plugins_enabled { state.cloud_retry = None; } state.plugins_enabled = plugins_enabled; - if self.providers.cloud.is_none() || !plugins_enabled { + if !plugins_enabled { state.cloud_generation = None; state.cloud_retry = None; } @@ -179,9 +180,70 @@ impl ConfigContributor for PluginContributor { impl TurnLifecycleContributor for PluginContributor { fn refresh_catalog_for_step<'a>( &'a self, - input: CatalogRefreshInput<'a>, + input: StepCatalogRefreshInput<'a>, ) -> ExtensionFuture<'a, bool> { - Box::pin(self.refresh_cloud_catalog(input)) + Box::pin(async move { + let Some(state) = input.thread_store.get::() else { + return false; + }; + let query = cloud_catalog_query(input.into(), Some(input.turn_id)); + { + let mut state = state.contributor_state(); + let auth = query + .mcp_resources + .as_ref() + .map(|client| client.auth_cache_key_for_server(CODEX_APPS_MCP_SERVER_NAME)); + let resources = query + .mcp_resources + .as_ref() + .and_then(|client| client.server_cache_key(CODEX_APPS_MCP_SERVER_NAME)); + if let Some(generation) = &state.cloud_generation + && generation + .catalog + .as_ref() + .is_some_and(|catalog| !catalog.is_partial && catalog.warnings.is_empty()) + && generation.auth_cache_key == auth + && generation.is_auth_current() + && generation.resource_cache_key == resources + { + state.cloud_retry = None; + return false; + } + if state + .cloud_retry + .as_ref() + .is_some_and(|retry| retry.is_pending(auth.as_ref(), resources.as_ref())) + { + return false; + } + } + match self.discover_cloud_catalog(&state, query).await { + Ok(changed) => changed, + Err(failure) => { + if let Some(retry_backoff) = failure.retry_backoff { + state.contributor_state().cloud_retry = Some(retry_backoff); + } + failure.catalog_cleared + } + } + }) + } + + fn refresh_cloud_catalog<'a>( + &'a self, + input: CloudCatalogRefreshInput<'a>, + ) -> ExtensionFuture<'a, Result<(), String>> { + Box::pin(async move { + let Some(state) = input.thread_store.get::() else { + return Ok(()); + }; + // A failed explicit attempt gets an immediate retry on the next step. + state.contributor_state().cloud_retry = None; + self.discover_cloud_catalog(&state, cloud_catalog_query(input, /*turn_id*/ None)) + .await + .map(|_| ()) + .map_err(|failure| failure.error) + }) } } diff --git a/codex-rs/ext/mcp/src/cloud_plugin_tests.rs b/codex-rs/ext/mcp/src/cloud_plugin_tests.rs index 8a939ffe8d..722faf8e2c 100644 --- a/codex-rs/ext/mcp/src/cloud_plugin_tests.rs +++ b/codex-rs/ext/mcp/src/cloud_plugin_tests.rs @@ -21,6 +21,7 @@ use codex_extension_api::McpServerContributor; use codex_extension_api::SelectedPluginSnapshot; use codex_extension_api::WorldStateContributionInput; use codex_extension_api::WorldStateSectionContribution; +use codex_features::Feature; use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::protocol::Op; @@ -45,6 +46,50 @@ use std::sync::atomic::Ordering; use std::time::Duration; use tokio::sync::Notify; +/// The host can suppress the first cloud refresh even when plugins are enabled. +#[tokio::test] +async fn cloud_source_policy_skips_startup_discovery() -> anyhow::Result<()> { + let home = tempfile::tempdir()?; + let mut config = codex_core::config::ConfigBuilder::default() + .codex_home(home.path().to_path_buf()) + .fallback_cwd(Some(home.path().to_path_buf())) + .build() + .await?; + config.features.enable(Feature::Plugins)?; + config.features.disable(Feature::Apps)?; + let provider = Arc::new(CloudProvider::default()); + let contributor = PluginContributor { + providers: PluginProviders::new(Arc::new(ExecutorPluginProvider::new(Arc::new( + EnvironmentManager::default_for_tests(), + )))) + .with_cloud_provider(provider.clone()) + .with_cloud_enabled(|config| config.features.enabled(Feature::Apps)), + }; + let session_store = ExtensionData::new("session"); + let thread_store = ExtensionData::new("thread"); + contributor + .on_thread_start(ThreadStartInput { + config: &config, + session_source: &codex_protocol::protocol::SessionSource::Exec, + persistent_thread_state_available: false, + environments: &[], + mcp_resource_client: None, + extension_metrics: None, + session_store: &session_store, + thread_store: &thread_store, + }) + .await; + contributor + .refresh_cloud_catalog(CloudCatalogRefreshInput { + session_store: &session_store, + thread_store: &thread_store, + }) + .await + .expect("disabled discovery succeeds"); + assert_eq!(provider.calls.load(Ordering::SeqCst), 0); + Ok(()) +} + #[tokio::test] async fn disabling_plugins_clears_cloud_catalog_and_skips_discovery() -> anyhow::Result<()> { let config_home = tempfile::tempdir()?; @@ -95,7 +140,7 @@ async fn disabling_plugins_clears_cloud_catalog_and_skips_discovery() -> anyhow: } for contributor in registry.turn_lifecycle_contributors() { contributor - .refresh_catalog_for_step(CatalogRefreshInput { + .refresh_catalog_for_step(StepCatalogRefreshInput { turn_id, session_store: &session, thread_store: &thread, @@ -219,14 +264,14 @@ async fn partial_or_warning_cloud_catalogs_retry_with_backoff_and_recover() { } let thread = ExtensionData::new("thread"); let turn = ExtensionData::new("turn"); - let input = CatalogRefreshInput { + let input = StepCatalogRefreshInput { turn_id: "turn", session_store: &session, thread_store: &thread, turn_store: &turn, }; let next_turn = ExtensionData::new("next turn"); - let next_input = CatalogRefreshInput { + let next_input = StepCatalogRefreshInput { turn_id: "next turn", turn_store: &next_turn, ..input @@ -298,7 +343,7 @@ async fn cloud_refresh_reuses_catalogs_and_retries() { let thread = ExtensionData::new("thread"); let turn = ExtensionData::new("turn"); let next_turn = ExtensionData::new("next turn"); - let input = CatalogRefreshInput { + let input = StepCatalogRefreshInput { turn_id: "turn", session_store: &session, thread_store: &thread, @@ -307,28 +352,29 @@ async fn cloud_refresh_reuses_catalogs_and_retries() { let state = thread.get_or_init(PluginsThreadState::default); state.contributor_state().plugins_enabled = true; - // Concurrent preparation waits for discovery before checking freshness. + // Core serializes callers; exercise the next preparation after completion or cancellation. provider.pause_discovery.store(true, Ordering::SeqCst); - let mut discovery = contributor.refresh_catalog_for_step(input); - assert!(futures::poll!(&mut discovery).is_pending()); - let next_input = CatalogRefreshInput { + let mut turn_start = contributor.refresh_catalog_for_step(input); + assert!(futures::poll!(&mut turn_start).is_pending()); + let next_input = StepCatalogRefreshInput { turn_id: "next turn", turn_store: &next_turn, ..input }; - let mut waiting = contributor.refresh_catalog_for_step(next_input); - assert!(futures::poll!(&mut waiting).is_pending()); assert_eq!(provider.calls.load(Ordering::SeqCst), 1); provider.set_reply(Ok(new_catalog.clone())); if cancelled { - drop(discovery); + drop(turn_start); } else { provider.resume_discovery.notify_one(); - assert_eq!(discovery.await, succeeded); + assert_eq!(turn_start.await, succeeded); } // Later turns share failure cooldowns; cancellation still allows immediate retry. - assert_eq!(waiting.await, cancelled); + assert_eq!( + contributor.refresh_catalog_for_step(next_input).await, + cancelled + ); if !succeeded && !cancelled { assert_eq!(state.cloud_catalog(), None); assert_eq!(provider.calls.load(Ordering::SeqCst), 1); @@ -357,6 +403,43 @@ async fn cloud_refresh_reuses_catalogs_and_retries() { } } +#[tokio::test(start_paused = true)] +async fn explicit_cloud_refresh_replaces_cached_catalog() { + let provider = Arc::new(CloudProvider::default()); + provider.set_reply(Ok(cloud_catalog())); + let contributor = PluginContributor { + providers: PluginProviders::new(Arc::new(ExecutorPluginProvider::new(Arc::new( + EnvironmentManager::default_for_tests(), + )))) + .with_cloud_provider(provider.clone()), + }; + let session = ExtensionData::new("session"); + let thread = ExtensionData::new("thread"); + let turn = ExtensionData::new("turn"); + let state = thread.get_or_init(PluginsThreadState::default); + state.contributor_state().plugins_enabled = true; + let input = StepCatalogRefreshInput { + turn_id: "turn", + session_store: &session, + thread_store: &thread, + turn_store: &turn, + }; + assert!(contributor.refresh_catalog_for_step(input).await); + let refreshed_catalog = PluginCatalog { + is_partial: true, + ..Default::default() + }; + provider.set_reply(Ok(refreshed_catalog.clone())); + assert!(!contributor.refresh_catalog_for_step(input).await); + contributor + .refresh_cloud_catalog(input.into()) + .await + .unwrap(); + assert_eq!(state.cloud_catalog(), Some(refreshed_catalog)); + assert!(!contributor.refresh_catalog_for_step(input).await); + assert_eq!(provider.calls.load(Ordering::SeqCst), 2); +} + struct CloudProvider { reply: Mutex>, calls: AtomicUsize, @@ -972,7 +1055,7 @@ async fn run_cloud_catalog_refresh_lifecycle() -> anyhow::Result<()> { )))) .with_cloud_provider(provider.clone()), }; - let input = CatalogRefreshInput { + let input = StepCatalogRefreshInput { turn_id: "same-turn", session_store: &session, thread_store: &thread, diff --git a/codex-rs/ext/mcp/src/plugin_contributor_state.rs b/codex-rs/ext/mcp/src/plugin_contributor_state.rs index 592072959f..9d95ac871e 100644 --- a/codex-rs/ext/mcp/src/plugin_contributor_state.rs +++ b/codex-rs/ext/mcp/src/plugin_contributor_state.rs @@ -19,13 +19,10 @@ use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex; use std::sync::MutexGuard; -use tokio::sync::Mutex as AsyncMutex; /// Plugin configuration and caches owned by one thread; mutable state stays extension-private. #[derive(Default)] pub struct PluginsThreadState { - /// Serializes discovery before checking freshness and publishing the next catalog. - pub(crate) cloud_refresh_lock: AsyncMutex<()>, state: Mutex, } diff --git a/codex-rs/ext/mcp/src/plugin_providers.rs b/codex-rs/ext/mcp/src/plugin_providers.rs index 8255107556..e2307f061e 100644 --- a/codex-rs/ext/mcp/src/plugin_providers.rs +++ b/codex-rs/ext/mcp/src/plugin_providers.rs @@ -1,15 +1,20 @@ //! Host-assembled plugin sources; discovery results belong to the active thread, not this registry. +use codex_core::config::Config; use codex_core_plugins::ExecutorPluginProvider; use codex_core_plugins::PluginProvider; +use codex_features::Feature; use std::sync::Arc; +type CloudEnabled = dyn Fn(&Config) -> bool + Send + Sync; + /// 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, + cloud_enabled: Option>, } impl PluginProviders { @@ -17,6 +22,7 @@ impl PluginProviders { Self { cloud: None, executor, + cloud_enabled: None, } } @@ -24,4 +30,23 @@ impl PluginProviders { self.cloud = Some(provider); self } + + /// Restricts cloud discovery using effective host configuration, independently of executor + /// plugins. Configured sources with unavailable connections must remain enabled for recovery. + pub fn with_cloud_enabled( + mut self, + enabled: impl Fn(&Config) -> bool + Send + Sync + 'static, + ) -> Self { + self.cloud_enabled = Some(Arc::new(enabled)); + self + } + + pub(crate) fn cloud_enabled(&self, config: &Config) -> bool { + self.cloud.is_some() + && config.features.enabled(Feature::Plugins) + && self + .cloud_enabled + .as_ref() + .is_none_or(|enabled| enabled(config)) + } } diff --git a/codex-rs/ext/skills/src/catalog_state.rs b/codex-rs/ext/skills/src/catalog_state.rs index b2090b3871..647dfb653c 100644 --- a/codex-rs/ext/skills/src/catalog_state.rs +++ b/codex-rs/ext/skills/src/catalog_state.rs @@ -123,7 +123,7 @@ impl CachedExecutorV2Catalog { #[derive(Clone, Debug)] pub(crate) struct ExecutorCatalogCacheKey { - pub(crate) turn_id: String, + pub(crate) turn_id: Option, pub(crate) roots: Vec<(SelectedCapabilityRoot, EnvironmentAccessKey)>, pub(crate) include_bundled_skills: bool, } diff --git a/codex-rs/ext/skills/src/cloud_cache_tests.rs b/codex-rs/ext/skills/src/cloud_cache_tests.rs index 0f7764e25c..8f12e2f223 100644 --- a/codex-rs/ext/skills/src/cloud_cache_tests.rs +++ b/codex-rs/ext/skills/src/cloud_cache_tests.rs @@ -114,7 +114,7 @@ fn client() -> Arc { fn query(client: Option>) -> SkillListQuery<'static> { SkillListQuery { - turn_id: "turn".into(), + turn_id: Some("turn".into()), executor_roots: Vec::new(), host_snapshot: None, include_host_skills: false, @@ -152,7 +152,7 @@ async fn frozen_cloud_catalog_and_reads_ignore_live_mcp_changes() { .expect("frozen cloud skill fixture should be valid"); let providers = SkillProviders::new().with_cloud_provider(Arc::new(provider)); state - .refresh_cloud_catalog(&providers, query(Some(client()))) + .refresh_cloud_catalog_for_step(&providers, query(Some(client()))) .await .expect("cloud catalog refresh should succeed"); let catalog = state @@ -182,7 +182,7 @@ async fn frozen_cloud_catalog_and_reads_ignore_live_mcp_changes() { } ); state - .refresh_cloud_catalog(&providers, query(mcp_resources)) + .refresh_cloud_catalog_for_step(&providers, query(mcp_resources)) .await .expect("cloud catalog refresh should succeed"); assert_eq!(state.available_cloud_catalog(), Some(catalog.clone())); @@ -366,7 +366,7 @@ async fn late_read_cannot_return_or_populate_a_replaced_cache() { let (late, ()) = tokio::join!(state.read_skill(&providers, request(original)), async { provider.entered.notified().await; state - .refresh_cloud_catalog(&providers, query(replacement.clone())) + .refresh_cloud_catalog_for_step(&providers, query(replacement.clone())) .await .expect("cloud catalog refresh should succeed"); provider.resume.notify_one(); @@ -390,7 +390,7 @@ async fn late_discovery_cannot_publish_into_a_new_generation() { let providers = SkillProviders::new().with_cloud_provider(provider.clone()); provider.pause_list.store(true, Ordering::SeqCst); let (late, ()) = tokio::join!( - state.refresh_cloud_catalog(&providers, query(Some(client()))), + state.refresh_cloud_catalog_for_step(&providers, query(Some(client()))), async { provider.entered.notified().await; state @@ -435,9 +435,9 @@ async fn cloud_catalogs_are_reused_while_fresh_or_backing_off_until_auth_changes .expect("every successful provider response should be accepted"); for turn_id in ["turn", "next turn"] { let mut request = query(Some(original_client.clone())); - request.turn_id = turn_id.to_string(); + request.turn_id = Some(turn_id.to_string()); state - .refresh_cloud_catalog(&providers, request) + .refresh_cloud_catalog_for_step(&providers, request) .await .expect("accepted cloud catalog should be reused"); assert_eq!(state.available_cloud_catalog(), Some(accepted.clone())); @@ -445,7 +445,7 @@ async fn cloud_catalogs_are_reused_while_fresh_or_backing_off_until_auth_changes } *provider.next_catalog.lock().unwrap() = Some(SkillCatalog::default()); state - .refresh_cloud_catalog(&providers, query(Some(client()))) + .refresh_cloud_catalog_for_step(&providers, query(Some(client()))) .await .expect("auth changes should refresh the accepted catalog"); assert_eq!( @@ -482,14 +482,14 @@ async fn partial_or_warning_cloud_catalogs_retry_with_backoff_and_recover() { for (index, delay) in [5, 10, 20, 40, 80, 160, 300, 300].into_iter().enumerate() { *provider.next_catalog.lock().unwrap() = Some(retryable.clone()); state - .refresh_cloud_catalog(&providers, query(original_client.clone())) + .refresh_cloud_catalog_for_step(&providers, query(original_client.clone())) .await .expect("partial or warning catalogs should remain usable"); assert_eq!(state.available_cloud_catalog(), Some(retryable.clone())); assert_eq!(provider.lists.load(Ordering::SeqCst), index + 1); tokio::time::advance(Duration::from_secs(delay - 1)).await; state - .refresh_cloud_catalog(&providers, query(original_client.clone())) + .refresh_cloud_catalog_for_step(&providers, query(original_client.clone())) .await .expect("retry cooldown should suppress discovery"); assert_eq!(provider.lists.load(Ordering::SeqCst), index + 1); @@ -502,14 +502,14 @@ async fn partial_or_warning_cloud_catalogs_retry_with_backoff_and_recover() { }; *provider.next_catalog.lock().unwrap() = Some(complete.clone()); state - .refresh_cloud_catalog(&providers, query(original_client.clone())) + .refresh_cloud_catalog_for_step(&providers, query(original_client.clone())) .await .expect("retry should recover a complete catalog"); assert_eq!(state.available_cloud_catalog(), Some(complete.clone())); assert!(state.catalog_state.lock().unwrap().cloud_retry.is_none()); tokio::time::advance(Duration::from_secs(600)).await; state - .refresh_cloud_catalog(&providers, query(original_client)) + .refresh_cloud_catalog_for_step(&providers, query(original_client)) .await .expect("clean catalogs should remain cached"); assert_eq!(state.available_cloud_catalog(), Some(complete)); @@ -544,7 +544,7 @@ async fn cloud_refresh_replaces_catalogs_and_read_cache_for_every_accepted_resul .expect("test provider catalog mutex should not be poisoned") = Some(previous.clone()); state - .refresh_cloud_catalog(&providers, query(Some(original_client.clone()))) + .refresh_cloud_catalog_for_step(&providers, query(Some(original_client.clone()))) .await .expect("cloud catalog refresh should succeed"); let generation = state @@ -569,7 +569,7 @@ async fn cloud_refresh_replaces_catalogs_and_read_cache_for_every_accepted_resul .lock() .expect("test provider catalog mutex should not be poisoned") = Some(incoming.clone()); - let current_client = client(); + let current_client = original_client; state .refresh_cloud_catalog(&providers, query(Some(current_client.clone()))) .await diff --git a/codex-rs/ext/skills/src/cloud_skill.rs b/codex-rs/ext/skills/src/cloud_skill.rs index e3974282bc..563c5c13e2 100644 --- a/codex-rs/ext/skills/src/cloud_skill.rs +++ b/codex-rs/ext/skills/src/cloud_skill.rs @@ -1,8 +1,10 @@ -//! Refreshes cloud skills before steps when their auth or Apps resources change. +//! Refreshes cloud skills explicitly or through step freshness checks. -use codex_extension_api::CatalogRefreshInput; +use codex_extension_api::CatalogRefreshInput as StepCatalogRefreshInput; +use codex_extension_api::CloudCatalogRefreshInput; use codex_extension_api::ExtensionFuture; use codex_extension_api::TurnLifecycleContributor; +use std::sync::Arc; use super::SkillsExtension; use crate::provider::SkillListQuery; @@ -11,54 +13,70 @@ use crate::state::SkillsThreadState; struct CloudDiscoveryWarningEmitted; -impl SkillsExtension { - async fn refresh_cloud_catalog(&self, input: CatalogRefreshInput<'_>) { - let Some(state) = input.thread_store.get::() else { - return; - }; - if !state.cloud_skill_enabled() { - return; - } - let config = state.config(); - let query = SkillListQuery { - turn_id: input.turn_id.to_string(), - executor_roots: Vec::new(), - host_snapshot: None, - include_host_skills: false, - include_bundled_skills: config.bundled_skills_enabled, - include_cloud_skills: true, - mcp_resources: input - .session_store - .get::() - .and_then(|state| state.mcp_resources.clone()), - executor_capability_discovery: None, - executor_capability_discovery_v2: None, - }; - if let Err(error) = state.refresh_cloud_catalog(&self.providers, query).await - && input - .turn_store - .insert(CloudDiscoveryWarningEmitted) - .is_none() - { - self.emit_warning( - input.thread_store.level_id(), - Some(input.turn_id), - format!( - "Cloud skill discovery failed; retaining any catalog from the same cloud auth scope: {error}" - ), - ); - } +fn cloud_catalog_query( + input: CloudCatalogRefreshInput<'_>, + turn_id: Option<&str>, +) -> Option<(Arc, SkillListQuery<'static>)> { + let state = input.thread_store.get::()?; + if !state.cloud_skill_enabled() { + return None; } + let query = SkillListQuery { + turn_id: turn_id.map(str::to_owned), + executor_roots: Vec::new(), + host_snapshot: None, + include_host_skills: false, + include_bundled_skills: state.config().bundled_skills_enabled, + include_cloud_skills: true, + mcp_resources: input + .session_store + .get::() + .and_then(|state| state.mcp_resources.clone()), + executor_capability_discovery: None, + executor_capability_discovery_v2: None, + }; + Some((state, query)) } impl TurnLifecycleContributor for SkillsExtension { fn refresh_catalog_for_step<'a>( &'a self, - input: CatalogRefreshInput<'a>, + input: StepCatalogRefreshInput<'a>, ) -> ExtensionFuture<'a, bool> { Box::pin(async move { - self.refresh_cloud_catalog(input).await; + let Some((state, query)) = cloud_catalog_query(input.into(), Some(input.turn_id)) + else { + return false; + }; + if let Err(error) = state + .refresh_cloud_catalog_for_step(&self.providers, query) + .await + && input + .turn_store + .insert(CloudDiscoveryWarningEmitted) + .is_none() + { + let message = format!( + "Cloud skill discovery failed; retaining any catalog from the same cloud auth scope: {error}" + ); + self.emit_warning(input.thread_store.level_id(), Some(input.turn_id), message); + } false }) } + + fn refresh_cloud_catalog<'a>( + &'a self, + input: CloudCatalogRefreshInput<'a>, + ) -> ExtensionFuture<'a, Result<(), String>> { + Box::pin(async move { + let Some((state, query)) = cloud_catalog_query(input, /*turn_id*/ None) else { + return Ok(()); + }; + state + .refresh_cloud_catalog(&self.providers, query) + .await + .map_err(|error| error.to_string()) + }) + } } diff --git a/codex-rs/ext/skills/src/extension.rs b/codex-rs/ext/skills/src/extension.rs index 4c1c7191e0..7ec6aa9989 100644 --- a/codex-rs/ext/skills/src/extension.rs +++ b/codex-rs/ext/skills/src/extension.rs @@ -211,7 +211,7 @@ where let catalog = self .providers .list_for_turn(SkillListQuery { - turn_id: thread_store.level_id().to_string(), + turn_id: Some(thread_store.level_id().to_string()), executor_roots: Vec::new(), host_snapshot: None, include_host_skills: false, @@ -355,7 +355,7 @@ where let config = &step_state.config; let host_skills = step_state.host_catalog.as_ref(); let query = SkillListQuery { - turn_id: input.turn_id.clone(), + turn_id: Some(input.turn_id.clone()), executor_roots: Vec::new(), host_snapshot: host_snapshot.clone(), include_host_skills: host_skills.is_none() && !host_catalog_in_world_state, diff --git a/codex-rs/ext/skills/src/preparation.rs b/codex-rs/ext/skills/src/preparation.rs index 81e911fc0a..d7abeae18a 100644 --- a/codex-rs/ext/skills/src/preparation.rs +++ b/codex-rs/ext/skills/src/preparation.rs @@ -64,7 +64,7 @@ async fn discover_executor_catalog( cloud_enabled: bool, ) -> SkillCatalog { let query = SkillListQuery { - turn_id: input.turn_id.to_string(), + turn_id: Some(input.turn_id.to_string()), executor_roots: input.ready_selected_capability_roots.to_vec(), host_snapshot: None, include_host_skills: false, @@ -95,7 +95,7 @@ async fn discover_host_catalog( Some( providers .list_host_for_turn(SkillListQuery { - turn_id: input.turn_id.to_string(), + turn_id: Some(input.turn_id.to_string()), executor_roots: Vec::new(), host_snapshot: Some(host_snapshot), include_host_skills: true, diff --git a/codex-rs/ext/skills/src/provider.rs b/codex-rs/ext/skills/src/provider.rs index 361f1f228a..00a7a4f542 100644 --- a/codex-rs/ext/skills/src/provider.rs +++ b/codex-rs/ext/skills/src/provider.rs @@ -31,7 +31,8 @@ pub(crate) const MAX_SKILL_RESOURCE_CONTENT_BYTES: usize = 1024 * 1024; /// Catalog inputs with executor roots bound to filesystem access for this callback. #[derive(Clone)] pub struct SkillListQuery<'a> { - pub turn_id: String, + /// Absent for explicit cloud discovery before turn admission. + pub turn_id: Option, pub executor_roots: Vec>, pub host_snapshot: Option>, pub include_host_skills: bool, diff --git a/codex-rs/ext/skills/src/render_tests.rs b/codex-rs/ext/skills/src/render_tests.rs index 22e4c17541..3d9123f3e7 100644 --- a/codex-rs/ext/skills/src/render_tests.rs +++ b/codex-rs/ext/skills/src/render_tests.rs @@ -638,7 +638,7 @@ async fn host_alias_roots_follow_core_discovery_order() -> Result<(), Box, config: Mutex, cloud_skills_available: bool, catalog_state: Mutex, @@ -62,6 +59,13 @@ pub struct SkillsThreadState { pub(crate) shadow_task_context: Arc, } +struct CloudSkillRefreshError { + error: SkillProviderError, + /// Cooldown for the step caller to store after a current-generation failure. + /// Superseded attempts return `None`, allowing the next step to retry immediately. + retry_backoff: Option, +} + impl SkillsThreadState { pub(crate) fn new( config: SkillsExtensionConfig, @@ -69,7 +73,6 @@ impl SkillsThreadState { extension_metrics: Option>, ) -> Self { Self { - cloud_refresh_lock: AsyncMutex::new(()), config: Mutex::new(config), cloud_skills_available, catalog_state: Mutex::new(SkillCatalogState::default()), @@ -284,19 +287,82 @@ impl SkillsThreadState { .and_then(|cache| cache.catalog.clone()) } - /// Reuses complete, warning-free catalogs; backs off partials, warnings, and failures. - #[expect( - clippy::await_holding_invalid_type, - reason = "cloud discovery must remain serialized through fetching and catalog publication" - )] - #[tracing::instrument(name = "skills.cloud.refresh_cloud_catalog", level = "info", skip_all)] + /// Checks freshness and backoff before using the shared cloud refresh path. + pub(crate) async fn refresh_cloud_catalog_for_step( + &self, + providers: &SkillProviders, + query: SkillListQuery<'_>, + ) -> SkillProviderResult<()> { + { + let config = self + .config + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if !config.cloud_skill_enabled || !self.cloud_skills_available { + return Ok(()); + } + let mut catalogs = self + .catalog_state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let cache = catalogs.prepare_cloud_generation( + providers.cloud_mcp_resources(query.mcp_resources.as_deref()), + ); + let resource_cache_key = cache.current_resource_cache_key(); + if cache.is_auth_current() + && cache + .catalog + .as_ref() + .is_some_and(|catalog| !catalog.is_partial && catalog.warnings.is_empty()) + && cache.resource_cache_key == resource_cache_key + { + catalogs.cloud_retry = None; + return Ok(()); + } + if catalogs.cloud_retry.as_ref().is_some_and(|retry| { + retry.is_pending(cache.auth_cache_key.as_ref(), resource_cache_key.as_ref()) + }) { + return Ok(()); + } + } + match self.discover_cloud_catalog(providers, query).await { + Ok(()) => Ok(()), + Err(failure) => { + let Some(retry_backoff) = failure.retry_backoff else { + // A superseded attempt can be retried on the next step. + return Ok(()); + }; + self.catalog_state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .cloud_retry = Some(retry_backoff); + Err(failure.error) + } + } + } + + /// Explicit requests bypass cache reuse and leave failures retryable on the next step. pub(crate) async fn refresh_cloud_catalog( &self, providers: &SkillProviders, query: SkillListQuery<'_>, ) -> SkillProviderResult<()> { - // Wait without holding config/catalog state or the MCP runtime permit. - let _refresh = self.cloud_refresh_lock.lock().await; + self.catalog_state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .cloud_retry = None; + self.discover_cloud_catalog(providers, query) + .await + .map_err(|failure| failure.error) + } + + /// Fetches and publishes one catalog; Core serializes discovery and step snapshots. + #[tracing::instrument(name = "skills.cloud.refresh_cloud_catalog", level = "info", skip_all)] + async fn discover_cloud_catalog( + &self, + providers: &SkillProviders, + query: SkillListQuery<'_>, + ) -> Result<(), CloudSkillRefreshError> { let (cache, resource_cache_key) = { let config = self .config @@ -313,21 +379,6 @@ impl SkillsThreadState { providers.cloud_mcp_resources(query.mcp_resources.as_deref()), ); let resource_cache_key = cache.current_resource_cache_key(); - if catalogs.cloud_retry.as_ref().is_some_and(|retry| { - retry.is_pending(cache.auth_cache_key.as_ref(), resource_cache_key.as_ref()) - }) { - return Ok(()); - } - if cache.is_auth_current() - && cache - .catalog - .as_ref() - .is_some_and(|catalog| !catalog.is_partial && catalog.warnings.is_empty()) - && cache.resource_cache_key == resource_cache_key - { - catalogs.cloud_retry = None; - return Ok(()); - } (cache, resource_cache_key) }; let mut result = providers.list_cloud_for_turn(query).await; @@ -345,18 +396,22 @@ impl SkillsThreadState { .as_ref() .is_some_and(|current| Arc::ptr_eq(current, &cache)) { - // Ignore superseded results, including errors. - return Ok(()); + return Err(CloudSkillRefreshError { + error: SkillProviderError::new("cloud skill discovery was superseded"), + retry_backoff: None, + }); } let catalog = match result { Ok(catalog) => catalog, Err(error) => { - catalogs.cloud_retry = Some(CatalogRefreshBackoff::after_failure( - cache.auth_cache_key.clone(), - resource_cache_key, - catalogs.cloud_retry.as_ref(), - )); - return Err(error); + return Err(CloudSkillRefreshError { + error, + retry_backoff: Some(CatalogRefreshBackoff::after_failure( + cache.auth_cache_key.clone(), + resource_cache_key, + catalogs.cloud_retry.as_ref(), + )), + }); } }; catalogs.cloud_retry = if catalog.is_partial || !catalog.warnings.is_empty() { diff --git a/codex-rs/ext/skills/tests/executor_file_system_authority.rs b/codex-rs/ext/skills/tests/executor_file_system_authority.rs index 73a93dfd2b..71c4978e95 100644 --- a/codex-rs/ext/skills/tests/executor_file_system_authority.rs +++ b/codex-rs/ext/skills/tests/executor_file_system_authority.rs @@ -523,7 +523,7 @@ async fn selected_root_id_distinguishes_identical_executor_paths() { }); let catalog = provider .list(SkillListQuery { - turn_id: "turn-1".to_string(), + turn_id: Some("turn-1".to_string()), executor_roots: selected_roots .iter() .map(|root| AccessedSelectedCapabilityRoot { root, fs: &access }) @@ -614,7 +614,7 @@ async fn executor_discovery_preserves_posix_and_windows_locator_alias_roots() { ); let catalog = provider .list(SkillListQuery { - turn_id: "turn-1".to_string(), + turn_id: Some("turn-1".to_string()), executor_roots: Vec::new(), host_snapshot: None, include_host_skills: false, @@ -736,7 +736,7 @@ async fn executor_discovery_routes_produce_equivalent_catalog_metadata() { ), ); let query = |executor_capability_discovery| SkillListQuery { - turn_id: "turn-1".to_string(), + turn_id: Some("turn-1".to_string()), executor_roots: executor_roots .iter() .map(|root| AccessedSelectedCapabilityRoot { root, fs: &access }) @@ -927,7 +927,7 @@ async fn pre_discovered_executor_catalog_snapshot() { .await; let catalog = provider .list(SkillListQuery { - turn_id: "turn-1".to_string(), + turn_id: Some("turn-1".to_string()), executor_roots: Vec::new(), host_snapshot: None, include_host_skills: false, @@ -1046,7 +1046,7 @@ async fn direct_executor_discovery_preserves_hidden_nested_and_probed_metadata() ); let direct = provider .list(SkillListQuery { - turn_id: "turn-1".to_string(), + turn_id: Some("turn-1".to_string()), executor_roots: vec![AccessedSelectedCapabilityRoot { root: &selected_root, fs: &access, @@ -1113,7 +1113,7 @@ async fn high_level_discovery_reuses_materialized_skill_contents_for_reads() { .await; let catalog = provider .list(SkillListQuery { - turn_id: "turn-1".to_string(), + turn_id: Some("turn-1".to_string()), executor_roots: Vec::new(), host_snapshot: None, include_host_skills: false, diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index 435622d212..af7b07d59c 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/codex-rs/ext/skills/tests/skills_extension.rs @@ -2713,9 +2713,18 @@ async fn cloud_discovery_is_gated_and_reports_failures() -> TestResult { refresh_registered_catalogs_for_step(®istry, &session_store, &thread_store, "turn-3") .await; assert_eq!(list_calls.load(Ordering::Relaxed), 2); + for contributor in registry.turn_lifecycle_contributors() { + contributor + .refresh_cloud_catalog(codex_extension_api::CloudCatalogRefreshInput { + session_store: &session_store, + thread_store: &thread_store, + }) + .await + .map_err(|error| format!("partial catalogs should be accepted: {error}"))?; + } refresh_registered_catalogs_for_step(®istry, &session_store, &thread_store, "turn-4") .await; - assert_eq!(list_calls.load(Ordering::Relaxed), 2); + assert_eq!(list_calls.load(Ordering::Relaxed), 3); } Ok(()) } diff --git a/codex-rs/ext/skills/tests/skills_extension/discovery_v2_tests.rs b/codex-rs/ext/skills/tests/skills_extension/discovery_v2_tests.rs index 6a9ff3619a..278b1138fc 100644 --- a/codex-rs/ext/skills/tests/skills_extension/discovery_v2_tests.rs +++ b/codex-rs/ext/skills/tests/skills_extension/discovery_v2_tests.rs @@ -90,7 +90,7 @@ async fn v2_steps_refresh_metadata_and_preserve_separate_environments() -> TestR }); } let query = SkillListQuery { - turn_id: "turn-1".to_string(), + turn_id: Some("turn-1".to_string()), executor_roots: Vec::new(), host_snapshot: None, include_host_skills: false, @@ -297,7 +297,7 @@ async fn v2_empty_failed_and_partial_results_never_load_legacy_roots() -> TestRe }; let access = FileSystemEnvironmentAccessor::unrestricted(&LOCAL_FS); let query = |snapshots| SkillListQuery { - turn_id: "turn-1".to_string(), + turn_id: Some("turn-1".to_string()), // V2 must ignore this unavailable legacy root. executor_roots: vec![AccessedSelectedCapabilityRoot { root: &root, diff --git a/codex-rs/file-search/src/lib_tests.rs b/codex-rs/file-search/src/lib_tests.rs index ddb369072f..4893ec24a6 100644 --- a/codex-rs/file-search/src/lib_tests.rs +++ b/codex-rs/file-search/src/lib_tests.rs @@ -259,22 +259,66 @@ fn session_scanned_file_count_is_monotonic_across_queries() { #[test] fn session_streams_updates_before_walk_complete() { - let dir = create_temp_tree(/*file_count*/ 600); + let dir = tempfile::tempdir().unwrap(); let reporter = Arc::new(RecordingReporter::default()); - let session = create_session( - vec![dir.path().to_path_buf()], - FileSearchOptions::default(), - reporter.clone(), - /*cancel_flag*/ None, - ) - .expect("session"); + let (work_tx, work_rx) = unbounded(); + let notify_tx = work_tx.clone(); + let nucleo = Nucleo::new( + Config::DEFAULT.match_paths(), + Arc::new(move || { + let _ = notify_tx.send(WorkSignal::NucleoNotify); + }), + /*num_threads*/ Some(1), + /*columns*/ 1, + ); + let injector = nucleo.injector(); + let inner = Arc::new(SessionInner { + search_directories: vec![dir.path().to_path_buf()], + limit: 50, + threads: 1, + compute_indices: false, + respect_gitignore: true, + cancelled: Arc::new(AtomicBool::new(false)), + shutdown: Arc::new(AtomicBool::new(false)), + reporter: reporter.clone(), + work_tx, + }); + let session = FileSearchSession { + inner: inner.clone(), + }; + let worker = thread::spawn(move || matcher_worker(inner, work_rx, nucleo)); + // Feed a walk result directly so a fast filesystem cannot finish the walk + // before the matcher has a chance to publish a streaming update. + injector.push( + IndexedEntry::Ascii { + root_idx: 0, + match_type: MatchType::File, + }, + |_, columns| columns[0] = Utf32String::from("file-0000.txt"), + ); session.update_query("file-0"); - let completed = reporter.wait_for_complete(Duration::from_secs(5)); - - assert!(completed); - let updates = reporter.updates(); - assert!(updates.iter().any(|snapshot| !snapshot.walk_complete)); + assert!(reporter.wait_until( + &reporter.updates, + &reporter.update_cv, + Duration::from_secs(5), + |updates| updates.iter().any(|snapshot| { + !snapshot.walk_complete + && snapshot.query == "file-0" + && snapshot + .matches + .iter() + .any(|matched| matched.path == Path::new("file-0000.txt")) + }), + )); + session + .inner + .work_tx + .send(WorkSignal::WalkComplete) + .unwrap(); + assert!(reporter.wait_for_complete(Duration::from_secs(5))); + drop(session); + worker.join().unwrap().unwrap(); } #[test]