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 c31a8debbf..5d3556f8c1 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 @@ -10,6 +10,7 @@ use std::sync::atomic::AtomicUsize; struct CloudTurnProvider { reply: Mutex, ()>>, lists: AtomicUsize, + partial: AtomicBool, reads: AtomicUsize, resource_key: Mutex>, } @@ -41,7 +42,7 @@ impl SkillProvider for CloudTurnProvider { }) .collect(); Ok(SkillCatalog { - is_partial: false, + is_partial: self.partial.load(Ordering::SeqCst), entries, warnings: Vec::new(), }) @@ -101,6 +102,7 @@ async fn cloud_skills_reuse_cache_and_invalidate_on_connection_or_auth_change() let provider = Arc::new(CloudTurnProvider { reply: Mutex::new(Ok(Some(calendar))), lists: AtomicUsize::new(0), + partial: AtomicBool::new(false), reads: AtomicUsize::new(0), resource_key: Mutex::new(None), }); @@ -153,10 +155,26 @@ async fn cloud_skills_reuse_cache_and_invalidate_on_connection_or_auth_change() 1, 1, ), + ( + "partial-replaces-complete", + Ok(Some(("calendar", 3))), + Some(("calendar", 3)), + None, + 1, + 1, + ), + ( + "partial-cached", + Ok(Some(("calendar", 4))), + Some(("calendar", 3)), + None, + 0, + 0, + ), ( "same-account-failure", Err(()), - Some(updated_calendar), + Some(("calendar", 3)), None, 1, 0, @@ -208,9 +226,14 @@ async fn cloud_skills_reuse_cache_and_invalidate_on_connection_or_auth_change() let previous_lists = provider.lists.load(Ordering::SeqCst); let previous_reads = provider.reads.load(Ordering::SeqCst); *provider.reply.lock().expect("cloud reply lock") = next_reply; + provider.partial.store( + matches!(turn, "partial-replaces-complete" | "partial-cached"), + Ordering::SeqCst, + ); let previous_connection_scope = if matches!( turn, "refreshed" + | "partial-replaces-complete" | "same-account-failure" | "removed" | "restored" diff --git a/codex-rs/core/tests/suite/skills_extension/cloud_skill_tests.rs b/codex-rs/core/tests/suite/skills_extension/cloud_skill_tests.rs index 33e2e95e5e..263de50b9f 100644 --- a/codex-rs/core/tests/suite/skills_extension/cloud_skill_tests.rs +++ b/codex-rs/core/tests/suite/skills_extension/cloud_skill_tests.rs @@ -932,14 +932,5 @@ async fn cloud_skill_can_read_referenced_resource_without_an_executor() -> Resul assert!(skill_fragments[0].contains(&format!("{SKILL_NAME}"))); assert!(skill_fragments[0].contains(SKILL_MARKER)); assert!(skill_fragments[0].contains(SKILL_REFERENCE_URI)); - // Partial discovery retries on the next turn without an MCP or auth invalidation. - assert_eq!( - *provider.reads.lock().await, - vec![ - SKILL_MAIN_PROMPT_URI, - SKILL_REFERENCE_URI, - SKILL_MAIN_PROMPT_URI - ] - ); Ok(()) } diff --git a/codex-rs/ext/extension-api/src/catalog_refresh.rs b/codex-rs/ext/extension-api/src/catalog_refresh.rs index b5bb07ccd8..15964437ab 100644 --- a/codex-rs/ext/extension-api/src/catalog_refresh.rs +++ b/codex-rs/ext/extension-api/src/catalog_refresh.rs @@ -5,7 +5,7 @@ use codex_mcp::McpResourceServerCacheKey; use std::time::Duration; use tokio::time::Instant; -/// Store per catalog and thread; clear on success or configuration reset. +/// Store per catalog and thread; clear on a complete, warning-free result or configuration reset. #[derive(Clone)] pub struct CatalogRefreshBackoff { auth: Option, @@ -42,7 +42,7 @@ impl CatalogRefreshBackoff { ) -> Self { let delay = previous .filter(|previous| previous.auth == auth && previous.resources == resources) - .map_or(Duration::from_secs(30), |previous| { + .map_or(Duration::from_secs(5), |previous| { (previous.delay * 2).min(Duration::from_secs(300)) }); Self { diff --git a/codex-rs/ext/mcp/src/cloud_plugin.rs b/codex-rs/ext/mcp/src/cloud_plugin.rs index d2ee3755a1..5c7c19d800 100644 --- a/codex-rs/ext/mcp/src/cloud_plugin.rs +++ b/codex-rs/ext/mcp/src/cloud_plugin.rs @@ -1,6 +1,6 @@ -//! Reuses complete, warning-free cloud plugins until auth or Apps resources change. -//! Hosts must invalidate Apps resources when the remote plugin catalog changes. -//! Partial catalogs retry next turn; failures use thread-owned exponential backoff. +//! Reuses complete, warning-free cloud plugins until 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. use crate::PluginListQuery; use crate::PluginsThreadState; @@ -22,16 +22,12 @@ use codex_extension_api::TurnStartPhase; use codex_features::Feature; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::McpResourceClient; -use codex_mcp::McpResourceClientAuthKey; use std::collections::HashSet; 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; -/// Auth scope of the latest catalog accepted for this turn. -struct LatestCloudPluginAuth(Option); - impl PluginContributor { #[expect( clippy::await_holding_invalid_type, @@ -45,10 +41,7 @@ impl PluginContributor { return false; }; // Wait without holding catalog state or the MCP runtime permit, then recheck the cache. - let (mut refresh_succeeded, waited) = match state.cloud_refresh_lock.try_lock() { - Ok(guard) => (guard, false), - Err(_) => (state.cloud_refresh_lock.lock().await, true), - }; + 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(), @@ -69,44 +62,28 @@ impl PluginContributor { .mcp_resources .as_ref() .and_then(|client| client.server_cache_key(CODEX_APPS_MCP_SERVER_NAME)); - if input - .turn_store - .get::() - .is_some_and(|previous| previous.0 == auth) - && state - .cloud_generation - .as_ref() - .is_some_and(|generation| generation.auth_cache_key == auth) - && state.cloud_catalog().is_some() - { - return false; - } if state.cloud_retry.as_ref().is_some_and(|retry| { retry.is_pending(auth.as_ref(), resource_cache_key.as_ref()) }) { return false; } - input.turn_store.remove::(); - if let Some(generation) = &state.cloud_generation - && generation.auth_cache_key == auth - && generation.is_auth_current() - && generation.resource_cache_key == resource_cache_key - && generation.catalog.as_ref().is_some_and(|catalog| { - (waited && *refresh_succeeded) - || (!catalog.is_partial && catalog.warnings.is_empty()) - }) - { - state.cloud_retry = None; - input.turn_store.insert(LatestCloudPluginAuth(auth)); - return false; - } - // Waiting callers must not reuse a failed or cancelled refresh. - *refresh_succeeded = 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 = @@ -117,17 +94,22 @@ impl PluginContributor { let mut state = state.contributor_state(); let catalog_published = match result { Ok(catalog) => { + let needs_retry = catalog.is_partial || !catalog.warnings.is_empty(); let published = state.publish_cloud_catalog( auth_cache_key.as_ref(), - resource_cache_key, + resource_cache_key.clone(), catalog, ); if published { - state.cloud_retry = None; - input - .turn_store - .insert(LatestCloudPluginAuth(auth_cache_key)); - *refresh_succeeded = true; + state.cloud_retry = if needs_retry { + Some(CatalogRefreshBackoff::after_failure( + auth_cache_key, + resource_cache_key, + state.cloud_retry.as_ref(), + )) + } else { + None + }; } else { tracing::debug!( "cloud plugin discovery finished after auth or resources changed" @@ -215,7 +197,6 @@ impl TurnLifecycleContributor for PluginContributor { fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> { Box::pin(async move { - input.turn_store.remove::(); self.refresh_cloud_catalog(CatalogRefreshInput { turn_id: input.turn_id, session_store: input.session_store, diff --git a/codex-rs/ext/mcp/src/cloud_plugin_tests.rs b/codex-rs/ext/mcp/src/cloud_plugin_tests.rs index f06cb95298..3559a3a955 100644 --- a/codex-rs/ext/mcp/src/cloud_plugin_tests.rs +++ b/codex-rs/ext/mcp/src/cloud_plugin_tests.rs @@ -249,7 +249,92 @@ async fn account_change_clears_hosted_plugins_before_waiting_for_executor() -> a } #[tokio::test(start_paused = true)] -async fn cloud_refresh_serializes_waiters_and_retries() { +async fn partial_or_warning_cloud_catalogs_retry_with_backoff_and_recover() { + for (is_partial, warnings) in [ + (true, Vec::new()), + (false, vec!["catalog diagnostic".to_string()]), + (true, vec!["catalog diagnostic".to_string()]), + ] { + for initial_client in [ + None, + Some(McpResourceClient::new(Arc::new( + codex_mcp::McpRuntime::empty(/*prefix_mcp_tool_names*/ false), + ))), + ] { + let provider = Arc::new(CloudProvider::default()); + let retryable = PluginCatalog { + is_partial, + warnings: warnings.clone(), + ..cloud_catalog() + }; + provider.set_reply(Ok(retryable.clone())); + 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"); + if let Some(client) = initial_client { + session.insert(client); + } + let thread = ExtensionData::new("thread"); + let turn = ExtensionData::new("turn"); + let input = CatalogRefreshInput { + turn_id: "turn", + session_store: &session, + thread_store: &thread, + turn_store: &turn, + }; + let next_turn = ExtensionData::new("next turn"); + let next_input = CatalogRefreshInput { + turn_id: "next turn", + turn_store: &next_turn, + ..input + }; + let state = thread.get_or_init(PluginsThreadState::default); + state.contributor_state().plugins_enabled = true; + for (index, delay) in [5, 10, 20, 40, 80, 160, 300, 300].into_iter().enumerate() { + assert!(contributor.refresh_catalog_for_step(input).await); + assert_eq!(state.cloud_catalog(), Some(retryable.clone())); + assert_eq!(provider.calls.load(Ordering::SeqCst), index + 1); + tokio::time::advance(Duration::from_secs(delay - 1)).await; + assert!(!contributor.refresh_catalog_for_step(next_input).await); + assert_eq!(provider.calls.load(Ordering::SeqCst), index + 1); + tokio::time::advance(Duration::from_secs(1)).await; + } + let complete = PluginCatalog { + warnings: Vec::new(), + ..cloud_catalog() + }; + provider.set_reply(Ok(complete.clone())); + assert!(contributor.refresh_catalog_for_step(input).await); + assert_eq!(state.cloud_catalog(), Some(complete.clone())); + assert!(state.contributor_state().cloud_retry.is_none()); + tokio::time::advance(Duration::from_secs(600)).await; + assert!(!contributor.refresh_catalog_for_step(input).await); + assert_eq!(provider.calls.load(Ordering::SeqCst), 9); + + // A new auth scope bypasses a retry cooldown immediately. + provider.set_reply(Ok(retryable)); + session.insert(McpResourceClient::new(Arc::new( + codex_mcp::McpRuntime::empty(/*prefix_mcp_tool_names*/ false), + ))); + assert!(contributor.refresh_catalog_for_step(input).await); + assert!(!contributor.refresh_catalog_for_step(input).await); + session.insert(McpResourceClient::new(Arc::new( + codex_mcp::McpRuntime::empty(/*prefix_mcp_tool_names*/ false), + ))); + provider.set_reply(Ok(complete.clone())); + assert!(contributor.refresh_catalog_for_step(input).await); + assert_eq!(state.cloud_catalog(), Some(complete)); + assert_eq!(provider.calls.load(Ordering::SeqCst), 11); + } + } +} + +#[tokio::test(start_paused = true)] +async fn cloud_refresh_reuses_catalogs_and_retries() { let new_catalog = PluginCatalog::default(); for (reply, cancelled) in [ (Ok(cloud_catalog()), false), @@ -257,7 +342,7 @@ async fn cloud_refresh_serializes_waiters_and_retries() { (Ok(cloud_catalog()), true), ] { let succeeded = reply.is_ok() && !cancelled; - let mut expected_catalog = if succeeded { + let expected_catalog = if succeeded { cloud_catalog() } else { new_catalog.clone() @@ -273,7 +358,7 @@ async fn cloud_refresh_serializes_waiters_and_retries() { let session = ExtensionData::new("session"); let thread = ExtensionData::new("thread"); let turn = ExtensionData::new("turn"); - let prewarm_turn = ExtensionData::new("prewarm"); + let next_turn = ExtensionData::new("next turn"); let input = CatalogRefreshInput { turn_id: "turn", session_store: &session, @@ -283,76 +368,53 @@ async fn cloud_refresh_serializes_waiters_and_retries() { let state = thread.get_or_init(PluginsThreadState::default); state.contributor_state().plugins_enabled = true; - // A delayed prewarm waits without superseding the real turn's request. + // Concurrent preparation waits for discovery before checking freshness. provider.pause_discovery.store(true, Ordering::SeqCst); - let mut turn_start = contributor.refresh_catalog_for_step(input); - assert!(futures::poll!(&mut turn_start).is_pending()); - let prewarm_input = CatalogRefreshInput { - turn_id: "prewarm", - turn_store: &prewarm_turn, + let mut discovery = contributor.refresh_catalog_for_step(input); + assert!(futures::poll!(&mut discovery).is_pending()); + let next_input = CatalogRefreshInput { + turn_id: "next turn", + turn_store: &next_turn, ..input }; - let mut prewarm = contributor.refresh_catalog_for_step(prewarm_input); - assert!(futures::poll!(&mut prewarm).is_pending()); - let mut cancelled_waiter = contributor.refresh_catalog_for_step(prewarm_input); - assert!(futures::poll!(&mut cancelled_waiter).is_pending()); - drop(cancelled_waiter); + 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(turn_start); + drop(discovery); } else { provider.resume_discovery.notify_one(); - assert_eq!(turn_start.await, succeeded); + assert_eq!(discovery.await, succeeded); } - // Waiting turns share failure cooldowns; cancellation still allows immediate retry. - assert_eq!( - tokio::time::timeout(Duration::from_secs(5), prewarm) - .await - .expect("waiting discovery should complete after its predecessor stops"), - cancelled, - ); + // Later turns share failure cooldowns; cancellation still allows immediate retry. + assert_eq!(waiting.await, cancelled); if !succeeded && !cancelled { assert_eq!(state.cloud_catalog(), None); assert_eq!(provider.calls.load(Ordering::SeqCst), 1); - tokio::time::advance(Duration::from_secs(30)).await; - assert!(contributor.refresh_catalog_for_step(prewarm_input).await); + tokio::time::advance(Duration::from_secs(5)).await; + assert!(contributor.refresh_catalog_for_step(next_input).await); } assert_eq!(state.cloud_catalog(), Some(expected_catalog.clone())); - assert!(!contributor.refresh_catalog_for_step(prewarm_input).await); + assert!(!contributor.refresh_catalog_for_step(next_input).await); assert_eq!( provider.calls.load(Ordering::SeqCst), 1 + usize::from(!succeeded), ); - // Retained partial catalogs do not need MCP rebuilding after failed retries. - prewarm_turn.remove::(); - expected_catalog.is_partial = true; - state - .contributor_state() - .cloud_generation - .as_mut() - .unwrap() - .catalog = Some(expected_catalog.clone()); provider.set_reply(Err("discovery failed".to_string())); - let previous_calls = provider.calls.load(Ordering::SeqCst); - for _ in 0..2 { - assert!(!contributor.refresh_catalog_for_step(prewarm_input).await); - assert_eq!(state.cloud_catalog(), Some(expected_catalog.clone())); - } - assert_eq!(provider.calls.load(Ordering::SeqCst), previous_calls + 1); // Clearing an old auth scope still requires rebuilding, even on failure. session.insert(McpResourceClient::new(Arc::new( codex_mcp::McpRuntime::empty(/*prefix_mcp_tool_names*/ false), ))); - assert!(contributor.refresh_catalog_for_step(prewarm_input).await); + assert!(contributor.refresh_catalog_for_step(next_input).await); assert_eq!(state.cloud_catalog(), None); - assert!(!contributor.refresh_catalog_for_step(prewarm_input).await); + assert!(!contributor.refresh_catalog_for_step(next_input).await); provider.set_reply(Ok(new_catalog.clone())); - tokio::time::advance(Duration::from_secs(30)).await; - assert!(contributor.refresh_catalog_for_step(prewarm_input).await); + tokio::time::advance(Duration::from_secs(5)).await; + assert!(contributor.refresh_catalog_for_step(next_input).await); } } @@ -566,6 +628,13 @@ async fn run_cloud_catalog_refresh_lifecycle() -> anyhow::Result<()> { None, None, ), + ( + "partial-cached", + Ok(cloud_catalog()), + Some("plugin_gmail"), + None, + None, + ), ( "partial-replaces-partial", Ok(PluginCatalog { @@ -583,6 +652,13 @@ async fn run_cloud_catalog_refresh_lifecycle() -> anyhow::Result<()> { None, None, ), + ( + "partial-empty-cached", + Ok(cloud_catalog()), + None, + None, + None, + ), ( "empty", Ok(PluginCatalog::default()), @@ -601,7 +677,7 @@ async fn run_cloud_catalog_refresh_lifecycle() -> anyhow::Result<()> { ( "partial-after-complete-empty", Ok(partial_catalog.clone()), - None, + Some("plugin_gmail"), None, None, ), @@ -622,9 +698,9 @@ async fn run_cloud_catalog_refresh_lifecycle() -> anyhow::Result<()> { ( "partial-refresh", Ok(partial_catalog.clone()), - Some("plugin_remote"), + Some("plugin_gmail"), + None, None, - Some(calendar), ), ( "partial-empty-refresh", @@ -650,10 +726,11 @@ async fn run_cloud_catalog_refresh_lifecycle() -> anyhow::Result<()> { None, None, ), + ("warning-cached", Ok(cloud_catalog()), None, None, None), ( "partial-after-disable", Ok(partial_catalog.clone()), - None, + Some("plugin_gmail"), None, None, ), @@ -686,6 +763,13 @@ async fn run_cloud_catalog_refresh_lifecycle() -> anyhow::Result<()> { None, Some(gmail), ), + ( + "new-account-partial-cached", + Ok(PluginCatalog::default()), + Some("plugin_gmail"), + None, + Some(gmail), + ), ( "new-account-discovered", Ok(gmail_catalog), @@ -702,20 +786,19 @@ async fn run_cloud_catalog_refresh_lifecycle() -> anyhow::Result<()> { ), ] { let previous_discoveries = provider.calls.load(Ordering::SeqCst); - let cached = matches!(turn, "cached" | "empty-cached"); + let cached = matches!( + turn, + "cached" + | "empty-cached" + | "partial-cached" + | "partial-empty-cached" + | "warning-cached" + | "new-account-partial-cached" + ); let in_flight_change = matches!(turn, "in-flight-auth-change" | "in-flight-resource-change"); - let retry_discovery = in_flight_change; - // Complete catalogs refresh only after invalidation, just like cloud skills. - if matches!( - turn, - "in-flight-resource-change" - | "partial-after-complete-empty" - | "partial-refresh" - | "removed" - | "restored" - | "in-flight-auth-change" - ) { + // Invalidation bypasses retry cooldowns; clean catalogs await invalidation. + if !cached && turn != "initial-partial" && new_account.is_none() { test.codex.submit(Op::RefreshMcpServers).await?; } if let Some(account) = new_account { @@ -757,15 +840,11 @@ async fn run_cloud_catalog_refresh_lifecycle() -> anyhow::Result<()> { .await?; assert_eq!(provider.calls.load(Ordering::SeqCst), previous_discoveries); } - if matches!( - turn, - "complete-with-warning-disables-plugin" | "new-account-partial" - ) { - // Reconnecting permits recovery before the thread's cooldown expires. - test.codex - .submit(codex_protocol::protocol::Op::RefreshMcpServers) - .await?; - } + let expected_catalog = if cached { + None + } else { + reply.as_ref().ok().cloned() + }; provider.set_reply(reply); let previous_calls = recorded_apps_tool_calls(&server).await.len(); let call_id = format!("{turn}-call"); @@ -852,6 +931,17 @@ async fn run_cloud_catalog_refresh_lifecycle() -> anyhow::Result<()> { assert_eq!(state.cloud_catalog(), expected); } else { test.submit_turn("What capabilities are available?").await?; + if let Some(expected_catalog) = expected_catalog { + assert_eq!( + test.codex + .thread_extension_data() + .get::() + .context("plugin state missing")? + .cloud_catalog(), + Some(expected_catalog), + "turn {turn}" + ); + } } let requests = response.requests(); let request = &requests[0]; @@ -912,11 +1002,50 @@ async fn run_cloud_catalog_refresh_lifecycle() -> anyhow::Result<()> { ); assert_eq!( provider.calls.load(Ordering::SeqCst), - previous_discoveries - + usize::from(!cached) - + if retry_discovery { requests.len() } else { 0 }, + previous_discoveries + usize::from(!cached) + usize::from(in_flight_change), "turn {turn}" ); } + // Re-enabling a server-filtered plugin can change Apps resources within one turn. + let resources = provider.resources.lock().unwrap().clone().unwrap(); + let session = ExtensionData::new("same-turn-session"); + session.insert(resources.as_ref().clone()); + let thread = ExtensionData::new("same-turn-thread"); + let turn = ExtensionData::new("same-turn"); + let state = thread.get_or_init(PluginsThreadState::default); + state.contributor_state().plugins_enabled = true; + let contributor = PluginContributor { + providers: PluginProviders::new(Arc::new(ExecutorPluginProvider::new(Arc::new( + EnvironmentManager::default_for_tests(), + )))) + .with_cloud_provider(provider.clone()), + }; + let input = CatalogRefreshInput { + turn_id: "same-turn", + session_store: &session, + thread_store: &thread, + turn_store: &turn, + }; + provider.set_reply(Ok(PluginCatalog::default())); + assert!(contributor.refresh_catalog_for_step(input).await); + provider.set_reply(Ok(cloud_catalog())); + assert!(!contributor.refresh_catalog_for_step(input).await); + assert_eq!(state.cloud_catalog(), Some(PluginCatalog::default())); + let key = resources.server_cache_key(CODEX_APPS_MCP_SERVER_NAME); + let auth = resources.auth_cache_key_for_server(CODEX_APPS_MCP_SERVER_NAME); + test.codex.submit(Op::RefreshMcpServers).await?; + tokio::time::timeout(Duration::from_secs(10), async { + while resources.server_cache_key(CODEX_APPS_MCP_SERVER_NAME) == key { + tokio::task::yield_now().await; + } + }) + .await?; + assert!(resources.server_cache_key(CODEX_APPS_MCP_SERVER_NAME) != key); + assert!(resources.auth_cache_key_for_server(CODEX_APPS_MCP_SERVER_NAME) == auth); + assert!(contributor.refresh_catalog_for_step(input).await); + assert_eq!(state.cloud_catalog(), Some(cloud_catalog())); + let calls = provider.calls.load(Ordering::SeqCst); + assert!(!contributor.refresh_catalog_for_step(input).await); + assert_eq!(provider.calls.load(Ordering::SeqCst), calls); Ok(()) } diff --git a/codex-rs/ext/mcp/src/plugin_contributor_state.rs b/codex-rs/ext/mcp/src/plugin_contributor_state.rs index 5f6415ed89..592072959f 100644 --- a/codex-rs/ext/mcp/src/plugin_contributor_state.rs +++ b/codex-rs/ext/mcp/src/plugin_contributor_state.rs @@ -1,5 +1,5 @@ //! Thread-owned catalogs, reused within their auth and resource scope, and selected-root metadata. -//! A short-lived mutex guards publication; incomplete refreshes preserve same-auth fallback data. +//! A short-lived mutex guards publication; failed refreshes preserve same-auth fallback data. use crate::catalog_observability::record_catalog; use codex_core_plugins::PluginCatalog; @@ -24,8 +24,8 @@ 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; true only when the last refresh published successfully. - pub(crate) cloud_refresh_lock: AsyncMutex, + /// Serializes discovery before checking freshness and publishing the next catalog. + pub(crate) cloud_refresh_lock: AsyncMutex<()>, state: Mutex, } @@ -144,13 +144,13 @@ impl PluginContributorState { auth_cache_key } - /// Accepts current-auth results, retaining prior complete entries as partial after an incomplete refresh. + /// Replaces the catalog with each accepted complete or partial result. /// Returns false when auth or resource identity changed during discovery. pub(crate) fn publish_cloud_catalog( &mut self, auth_cache_key: Option<&McpResourceClientAuthKey>, resource_cache_key: Option, - mut catalog: PluginCatalog, + catalog: PluginCatalog, ) -> bool { let Some(generation) = self .cloud_generation @@ -161,18 +161,6 @@ impl PluginContributorState { else { return false; }; - if catalog.is_partial - && let Some(previous) = generation - .catalog - .as_ref() - .filter(|previous| !previous.is_partial) - { - catalog = previous.clone(); - catalog.is_partial = true; - tracing::warn!( - "cloud plugin discovery was partial; retaining prior auth-scoped entries as partial" - ); - } record_catalog(self.extension_metrics.as_deref(), "cloud", &catalog); generation.catalog = Some(catalog); generation.resource_cache_key = resource_cache_key; diff --git a/codex-rs/ext/mcp/src/plugin_contributor_state_tests.rs b/codex-rs/ext/mcp/src/plugin_contributor_state_tests.rs index b232892b99..49175e2663 100644 --- a/codex-rs/ext/mcp/src/plugin_contributor_state_tests.rs +++ b/codex-rs/ext/mcp/src/plugin_contributor_state_tests.rs @@ -109,7 +109,7 @@ fn executor_snapshot( } #[test] -fn partial_refresh_retains_entries_without_excluding_new_executor_plugins() -> anyhow::Result<()> { +fn partial_refresh_replaces_entries_and_merges_executor_plugins() -> anyhow::Result<()> { let mut state = PluginContributorState::default(); let auth_key = state.prepare_cloud_generation(/*mcp_resources*/ None); let original = PluginCatalog { @@ -119,19 +119,18 @@ fn partial_refresh_retains_entries_without_excluding_new_executor_plugins() -> a assert!(state.publish_cloud_catalog( auth_key.as_ref(), /*resource_cache_key*/ None, - original.clone() + original )); + let expected = PluginCatalog { + entries: vec![cloud_plugin("plugin_a")], + is_partial: true, + ..Default::default() + }; assert!(state.publish_cloud_catalog( auth_key.as_ref(), /*resource_cache_key*/ None, - PluginCatalog { - entries: vec![cloud_plugin("plugin_a")], - is_partial: true, - ..Default::default() - }, + expected.clone(), )); - let mut expected = original; - expected.is_partial = true; assert_eq!(state.cloud_catalog(), Some(&expected)); state.update_executor_catalog(&[executor_snapshot("executor-a")?]); diff --git a/codex-rs/ext/skills/src/catalog_state.rs b/codex-rs/ext/skills/src/catalog_state.rs index e3fc3fedbf..b2090b3871 100644 --- a/codex-rs/ext/skills/src/catalog_state.rs +++ b/codex-rs/ext/skills/src/catalog_state.rs @@ -67,7 +67,7 @@ impl SkillCatalogState { } /// Keeps cloud metadata within one published auth scope and Apps availability state. -/// Warning-free discovery is reused until the Apps resource generation changes. +/// Accepted discovery, including partial catalogs, is reused until auth or resources change. /// A published refresh replaces the Arc so late reads cannot populate its successor. pub(crate) struct CloudSkillGeneration { pub(crate) auth_cache_key: Option, diff --git a/codex-rs/ext/skills/src/cloud_cache_tests.rs b/codex-rs/ext/skills/src/cloud_cache_tests.rs index a9c0ccdf06..710c744a98 100644 --- a/codex-rs/ext/skills/src/cloud_cache_tests.rs +++ b/codex-rs/ext/skills/src/cloud_cache_tests.rs @@ -151,18 +151,10 @@ 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)); - let turn_store = ExtensionData::new("turn"); state - .refresh_cloud_catalog(&providers, query(Some(client())), &turn_store) + .refresh_cloud_catalog(&providers, query(Some(client()))) .await .expect("cloud catalog refresh should succeed"); - assert!( - turn_store - .get::() - .expect("successful discovery should record its auth scope") - .0 - .is_none() - ); let catalog = state .available_cloud_catalog() .expect("cloud catalog should be available after refresh"); @@ -190,7 +182,7 @@ async fn frozen_cloud_catalog_and_reads_ignore_live_mcp_changes() { } ); state - .refresh_cloud_catalog(&providers, query(mcp_resources), &turn_store) + .refresh_cloud_catalog(&providers, query(mcp_resources)) .await .expect("cloud catalog refresh should succeed"); assert_eq!(state.available_cloud_catalog(), Some(catalog.clone())); @@ -208,7 +200,7 @@ async fn frozen_cloud_catalog_and_reads_ignore_live_mcp_changes() { } #[tokio::test(start_paused = true)] -async fn step_cloud_refresh_serializes_waiters_and_retries() { +async fn step_cloud_refresh_reuses_accepted_results_and_retries_failures() { #[derive(Default)] struct WarningCounter(AtomicUsize); @@ -246,7 +238,7 @@ async fn step_cloud_refresh_serializes_waiters_and_retries() { turn_store: &turn_store, }; - for (index, delay) in [30, 60, 120, 240, 300, 300].into_iter().enumerate() { + for (index, delay) in [5, 10, 20, 40, 80, 160, 300, 300].into_iter().enumerate() { cloud.fail_list.store(true, Ordering::SeqCst); assert!(!contributor.refresh_catalog_for_step(refresh_input).await); assert_eq!(thread_state.available_cloud_catalog(), None); @@ -268,16 +260,15 @@ async fn step_cloud_refresh_serializes_waiters_and_retries() { Some(SkillCatalog::default()); assert!(!contributor.refresh_catalog_for_step(refresh_input).await); assert!(!contributor.refresh_catalog_for_step(refresh_input).await); - assert_eq!(cloud.lists.load(Ordering::SeqCst), 8); + assert_eq!(cloud.lists.load(Ordering::SeqCst), 10); - // Waiting callers reuse success and share failures' cooldowns; cancellation allows retry. + // Later steps reuse the accepted result or failure cooldown; cancellation allows retry. for (fails, cancelled) in [(false, false), (true, false), (false, true)] { session_store.insert(SkillsSessionState { mcp_resources: Some(client()), extension_metrics: None, }); let previous_calls = cloud.lists.load(Ordering::SeqCst); - let prewarm_store = ExtensionData::new("prewarm"); let first_catalog = SkillCatalog { warnings: vec!["catalog warning".into()], ..Default::default() @@ -291,16 +282,6 @@ async fn step_cloud_refresh_serializes_waiters_and_retries() { cloud.pause_list.store(true, Ordering::SeqCst); let mut turn_start = contributor.refresh_catalog_for_step(refresh_input); assert!(futures::poll!(&mut turn_start).is_pending()); - let prewarm_input = CatalogRefreshInput { - turn_id: "prewarm", - turn_store: &prewarm_store, - ..refresh_input - }; - let mut prewarm = contributor.refresh_catalog_for_step(prewarm_input); - assert!(futures::poll!(&mut prewarm).is_pending()); - let mut cancelled_waiter = contributor.refresh_catalog_for_step(prewarm_input); - assert!(futures::poll!(&mut cancelled_waiter).is_pending()); - drop(cancelled_waiter); assert_eq!(cloud.lists.load(Ordering::SeqCst), previous_calls + 1); *cloud .next_catalog @@ -313,16 +294,12 @@ async fn step_cloud_refresh_serializes_waiters_and_retries() { cloud.resume.notify_one(); assert!(!turn_start.await); } - assert!( - !tokio::time::timeout(std::time::Duration::from_secs(5), prewarm) - .await - .expect("waiting discovery should complete after its predecessor stops") - ); + assert!(!contributor.refresh_catalog_for_step(refresh_input).await); if fails { assert_eq!(thread_state.available_cloud_catalog(), None); assert_eq!(cloud.lists.load(Ordering::SeqCst), previous_calls + 1); - tokio::time::advance(Duration::from_secs(30)).await; - assert!(!contributor.refresh_catalog_for_step(prewarm_input).await); + tokio::time::advance(Duration::from_secs(5)).await; + assert!(!contributor.refresh_catalog_for_step(refresh_input).await); } assert_eq!( thread_state.available_cloud_catalog(), @@ -332,7 +309,6 @@ async fn step_cloud_refresh_serializes_waiters_and_retries() { first_catalog }) ); - assert!(!contributor.refresh_catalog_for_step(prewarm_input).await); assert!(!contributor.refresh_catalog_for_step(refresh_input).await); assert_eq!( cloud.lists.load(Ordering::SeqCst), @@ -366,10 +342,10 @@ async fn step_cloud_refresh_serializes_waiters_and_retries() { turn_store: &next_turn, ..refresh_input }; - tokio::time::advance(Duration::from_secs(30)).await; + tokio::time::advance(Duration::from_secs(5)).await; assert!(!contributor.refresh_catalog_for_step(next_input).await); assert_eq!(cloud.lists.load(Ordering::SeqCst), previous_calls + 2); - tokio::time::advance(Duration::from_secs(59)).await; + tokio::time::advance(Duration::from_secs(9)).await; assert!(!contributor.refresh_catalog_for_step(refresh_input).await); assert_eq!(cloud.lists.load(Ordering::SeqCst), previous_calls + 2); tokio::time::advance(Duration::from_secs(1)).await; @@ -385,17 +361,13 @@ async fn late_read_cannot_return_or_populate_a_replaced_cache() { let provider = Arc::new(Provider::default()); let providers = SkillProviders::new().with_cloud_provider(provider.clone()); let original = Some(client()); - // None here exercises a successful per-turn refresh on the SAME generation. + // None here exercises a refresh within the same auth generation. let replacement = replacement.or_else(|| original.clone()); provider.pause_read.store(true, Ordering::SeqCst); let (late, ()) = tokio::join!(state.read_skill(&providers, request(original)), async { provider.entered.notified().await; state - .refresh_cloud_catalog( - &providers, - query(replacement.clone()), - &ExtensionData::new("turn"), - ) + .refresh_cloud_catalog(&providers, query(replacement.clone())) .await .expect("cloud catalog refresh should succeed"); provider.resume.notify_one(); @@ -418,9 +390,8 @@ async fn late_discovery_cannot_publish_into_a_new_generation() { let provider = Arc::new(Provider::default()); let providers = SkillProviders::new().with_cloud_provider(provider.clone()); provider.pause_list.store(true, Ordering::SeqCst); - let turn_store = ExtensionData::new("turn"); let (late, ()) = tokio::join!( - state.refresh_cloud_catalog(&providers, query(Some(client())), &turn_store), + state.refresh_cloud_catalog(&providers, query(Some(client()))), async { provider.entered.notified().await; state @@ -432,7 +403,6 @@ async fn late_discovery_cannot_publish_into_a_new_generation() { } ); late.expect("superseded discovery should be discarded without an error"); - assert!(turn_store.get::().is_none()); assert!( state .catalog_state @@ -447,71 +417,110 @@ async fn late_discovery_cannot_publish_into_a_new_generation() { } #[tokio::test(start_paused = true)] -async fn partial_cloud_catalog_without_warnings_is_refreshed_next_turn() { - let state = state(); - let provider = Arc::new(Provider::default()); - provider.partial_list.store(true, Ordering::SeqCst); - let providers = SkillProviders::new().with_cloud_provider(provider.clone()); - let client = Some(client()); - let turn_store = ExtensionData::new("turn"); - state - .refresh_cloud_catalog(&providers, query(client.clone()), &turn_store) - .await - .expect("cloud catalog refresh should succeed"); - // Reuse partial catalogs for the rest of the turn. - state - .refresh_cloud_catalog(&providers, query(client.clone()), &turn_store) - .await - .expect("cloud catalog refresh should succeed"); - assert_eq!(provider.lists.load(Ordering::SeqCst), 1); - assert_eq!( - state - .available_cloud_catalog() - .expect("cloud catalog should be available after refresh"), - SkillCatalog { - is_partial: true, - ..Default::default() +async fn cloud_catalogs_are_reused_while_fresh_or_backing_off_until_auth_changes() { + for is_partial in [false, true] { + for warnings in [Vec::new(), vec!["catalog diagnostic".to_string()]] { + let state = state(); + let provider = Arc::new(Provider::default()); + let accepted = SkillCatalog { + is_partial, + warnings, + ..Default::default() + }; + *provider.next_catalog.lock().unwrap() = Some(accepted.clone()); + let providers = SkillProviders::new().with_cloud_provider(provider.clone()); + let original_client = client(); + state + .refresh_cloud_catalog(&providers, query(Some(original_client.clone()))) + .await + .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(); + state + .refresh_cloud_catalog(&providers, request) + .await + .expect("accepted cloud catalog should be reused"); + assert_eq!(state.available_cloud_catalog(), Some(accepted.clone())); + assert_eq!(provider.lists.load(Ordering::SeqCst), 1); + } + *provider.next_catalog.lock().unwrap() = Some(SkillCatalog::default()); + state + .refresh_cloud_catalog(&providers, query(Some(client()))) + .await + .expect("auth changes should refresh the accepted catalog"); + assert_eq!( + state.available_cloud_catalog(), + Some(SkillCatalog::default()) + ); + assert_eq!(provider.lists.load(Ordering::SeqCst), 2); } - ); - state - .refresh_cloud_catalog( - &providers, - query(client.clone()), - &ExtensionData::new("next turn"), - ) - .await - .expect("cloud catalog refresh should succeed"); - assert_eq!( - state - .available_cloud_catalog() - .expect("cloud catalog should be available after refresh"), - SkillCatalog { - is_partial: false, - warnings: vec!["account catalog".into()], - ..Default::default() + } +} + +#[tokio::test(start_paused = true)] +async fn partial_or_warning_cloud_catalogs_retry_with_backoff_and_recover() { + for (is_partial, warnings) in [ + (true, Vec::new()), + (false, vec!["catalog diagnostic".to_string()]), + (true, vec!["catalog diagnostic".to_string()]), + ] { + for original_client in [None, Some(client())] { + let state = state(); + let provider = Arc::new(Provider::default()); + let providers = SkillProviders::new().with_cloud_provider(provider.clone()); + let retryable = SkillCatalog { + entries: vec![SkillCatalogEntry::new( + SkillPackageId("skill://demo".into()), + SkillAuthority::new(SkillSourceKind::Cloud, "codex_apps"), + "demo", + "retryable catalog", + SkillResourceId::new("skill://demo/SKILL.md"), + )], + is_partial, + warnings: warnings.clone(), + }; + 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())) + .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())) + .await + .expect("retry cooldown should suppress discovery"); + assert_eq!(provider.lists.load(Ordering::SeqCst), index + 1); + tokio::time::advance(Duration::from_secs(1)).await; + } + let complete = SkillCatalog { + is_partial: false, + warnings: Vec::new(), + ..retryable + }; + *provider.next_catalog.lock().unwrap() = Some(complete.clone()); + state + .refresh_cloud_catalog(&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)) + .await + .expect("clean catalogs should remain cached"); + assert_eq!(state.available_cloud_catalog(), Some(complete)); + assert_eq!(provider.lists.load(Ordering::SeqCst), 9); } - ); - // Retry failures even when the previous catalog remains available. - let retry_turn = ExtensionData::new("retry turn"); - provider.fail_list.store(true, Ordering::SeqCst); - assert!( - state - .refresh_cloud_catalog(&providers, query(client.clone()), &retry_turn) - .await - .is_err() - ); - assert!(state.available_cloud_catalog().is_some()); - assert!(retry_turn.get::().is_none()); - tokio::time::advance(Duration::from_secs(30)).await; - state - .refresh_cloud_catalog(&providers, query(client), &retry_turn) - .await - .expect("cloud catalog refresh should succeed"); - assert_eq!(provider.lists.load(Ordering::SeqCst), 4); + } } #[tokio::test] -async fn cloud_refresh_preserves_complete_catalogs_only_within_the_same_auth_scope() { +async fn cloud_refresh_replaces_catalogs_and_read_cache_for_every_accepted_result() { let catalog = |name: &str, is_partial| SkillCatalog { entries: vec![SkillCatalogEntry::new( SkillPackageId("skill://demo".into()), @@ -520,101 +529,72 @@ async fn cloud_refresh_preserves_complete_catalogs_only_within_the_same_auth_sco "test catalog", SkillResourceId::new("skill://demo/SKILL.md"), )], - // Complete listings can have diagnostics; these also keep discovery eligible for refresh. warnings: vec![format!("{name} catalog diagnostic")], is_partial, }; for previous_partial in [false, true] { for incoming_partial in [false, true] { - for changed_auth in [false, true] { - let state = state(); - let provider = Arc::new(Provider::default()); - let providers = SkillProviders::new().with_cloud_provider(provider.clone()); - let original_client = client(); - let mut previous = catalog("previous", previous_partial); - *provider - .next_catalog - .lock() - .expect("test provider catalog mutex should not be poisoned") = - Some(previous.clone()); + let state = state(); + let provider = Arc::new(Provider::default()); + let providers = SkillProviders::new().with_cloud_provider(provider.clone()); + let original_client = client(); + let previous = catalog("previous", previous_partial); + *provider + .next_catalog + .lock() + .expect("test provider catalog mutex should not be poisoned") = + Some(previous.clone()); + state + .refresh_cloud_catalog(&providers, query(Some(original_client.clone()))) + .await + .expect("cloud catalog refresh should succeed"); + let generation = state + .catalog_state + .lock() + .expect("catalog state mutex should not be poisoned") + .cloud_generation + .clone() + .expect("initial refresh should publish a cloud generation"); + assert_eq!( state - .refresh_cloud_catalog( - &providers, - query(Some(original_client.clone())), - &ExtensionData::new("turn"), - ) + .read_skill(&providers, request(Some(original_client.clone()))) .await - .expect("cloud catalog refresh should succeed"); - let generation = state + .expect("cloud skill read should succeed") + .contents, + "contents 0" + ); + + let incoming = catalog("incoming", incoming_partial); + *provider + .next_catalog + .lock() + .expect("test provider catalog mutex should not be poisoned") = + Some(incoming.clone()); + let current_client = client(); + state + .refresh_cloud_catalog(&providers, query(Some(current_client.clone()))) + .await + .expect("cloud catalog refresh should succeed"); + assert_eq!(generation.catalog, Some(previous)); + assert_eq!(state.available_cloud_catalog(), Some(incoming)); + assert!(!Arc::ptr_eq( + &generation, + state .catalog_state .lock() .expect("catalog state mutex should not be poisoned") .cloud_generation - .clone() - .expect("initial refresh should publish a cloud generation"); - let contents = state - .read_skill(&providers, request(Some(original_client.clone()))) - .await - .expect("cloud skill read should succeed"); - - let incoming = catalog("incoming", incoming_partial); - *provider - .next_catalog - .lock() - .expect("test provider catalog mutex should not be poisoned") = - Some(incoming.clone()); - let current_client = if changed_auth { - client() - } else { - original_client - }; + .as_ref() + .expect("refresh should publish a cloud generation") + )); + assert_eq!( state - .refresh_cloud_catalog( - &providers, - query(Some(current_client.clone())), - &ExtensionData::new("next turn"), - ) + .read_skill(&providers, request(Some(current_client))) .await - .expect("cloud catalog refresh should succeed"); - let retained = !changed_auth && !previous_partial && incoming_partial; - assert_eq!( - generation - .catalog - .as_ref() - .expect("published generation should contain a catalog") - .is_partial, - previous_partial - ); - if retained { - previous.is_partial = true; - } - assert_eq!( - state.available_cloud_catalog(), - Some(if retained { previous } else { incoming }) - ); - assert!(!Arc::ptr_eq( - &generation, - state - .catalog_state - .lock() - .expect("catalog state mutex should not be poisoned") - .cloud_generation - .as_ref() - .expect("refresh should publish a cloud generation") - )); - assert_eq!( - state - .read_skill(&providers, request(Some(current_client))) - .await - .expect("cloud skill read should succeed") - .contents, - if retained { - contents.contents - } else { - "contents 1".to_string() - }, - ); - } + .expect("cloud skill read should succeed") + .contents, + "contents 1", + ); } } } diff --git a/codex-rs/ext/skills/src/cloud_skill.rs b/codex-rs/ext/skills/src/cloud_skill.rs index 12a363fd15..e573764f42 100644 --- a/codex-rs/ext/skills/src/cloud_skill.rs +++ b/codex-rs/ext/skills/src/cloud_skill.rs @@ -1,4 +1,4 @@ -//! Refreshes cloud skills at turn start and after auth changes. +//! Refreshes cloud skills at turn start and through step freshness checks. use codex_extension_api::CatalogRefreshInput; use codex_extension_api::ExtensionData; @@ -9,7 +9,6 @@ use codex_extension_api::TurnStartPhase; use super::SkillsExtension; use crate::provider::SkillListQuery; -use crate::state::LatestCloudSkillsAuth; use crate::state::SkillsSessionState; use crate::state::SkillsThreadState; @@ -38,9 +37,7 @@ impl SkillsExtension { executor_capability_discovery: None, executor_capability_discovery_v2: None, }; - if let Err(error) = state - .refresh_cloud_catalog(&self.providers, query, input.turn_store) - .await + if let Err(error) = state.refresh_cloud_catalog(&self.providers, query).await && input .turn_store .insert(CloudDiscoveryWarningEmitted) @@ -75,7 +72,6 @@ impl TurnLifecycleContributor for SkillsExtension { fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> { Box::pin(async move { - input.turn_store.remove::(); self.refresh_cloud_catalog(CatalogRefreshInput { turn_id: input.turn_id, session_store: input.session_store, diff --git a/codex-rs/ext/skills/src/state.rs b/codex-rs/ext/skills/src/state.rs index 4105e7c815..389db2cdaf 100644 --- a/codex-rs/ext/skills/src/state.rs +++ b/codex-rs/ext/skills/src/state.rs @@ -8,10 +8,8 @@ use tokio::sync::Mutex as AsyncMutex; use codex_exec_server::EnvironmentAccessKey; use codex_extension_api::CatalogRefreshBackoff; -use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionMetrics; use codex_mcp::McpResourceClient; -use codex_mcp::McpResourceClientAuthKey; use crate::SkillsExtensionConfig; use crate::aliases::SkillAliasState; @@ -49,13 +47,10 @@ pub(crate) struct SkillsSessionState { pub(crate) extension_metrics: Option>, } -/// Auth scope of the latest catalog accepted for this turn. -pub(crate) struct LatestCloudSkillsAuth(Option); - /// Thread-owned skill configuration and caches; consumers can only read catalog snapshots. pub struct SkillsThreadState { - /// Serializes discovery; true only when the last refresh published successfully. - cloud_refresh_lock: AsyncMutex, + /// Serializes discovery before checking freshness and publishing the next catalog. + cloud_refresh_lock: AsyncMutex<()>, config: Mutex, cloud_skills_available: bool, catalog_state: Mutex, @@ -74,7 +69,7 @@ impl SkillsThreadState { extension_metrics: Option>, ) -> Self { Self { - cloud_refresh_lock: AsyncMutex::new(false), + cloud_refresh_lock: AsyncMutex::new(()), config: Mutex::new(config), cloud_skills_available, catalog_state: Mutex::new(SkillCatalogState::default()), @@ -289,7 +284,7 @@ impl SkillsThreadState { .and_then(|cache| cache.catalog.clone()) } - /// Reuses discovery within a turn; backs off failures. + /// 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" @@ -299,13 +294,9 @@ impl SkillsThreadState { &self, providers: &SkillProviders, query: SkillListQuery<'_>, - turn_store: &ExtensionData, ) -> SkillProviderResult<()> { // Wait without holding config/catalog state or the MCP runtime permit. - let (mut refresh_succeeded, waited) = match self.cloud_refresh_lock.try_lock() { - Ok(guard) => (guard, false), - Err(_) => (self.cloud_refresh_lock.lock().await, true), - }; + let _refresh = self.cloud_refresh_lock.lock().await; let (cache, resource_cache_key) = { let config = self .config @@ -318,38 +309,25 @@ impl SkillsThreadState { .catalog_state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - // Check this turn's discovery using the provider's auth scope. let cache = catalogs.prepare_cloud_generation( providers.cloud_mcp_resources(query.mcp_resources.as_deref()), ); - if cache.is_auth_current() - && cache.catalog.is_some() - && turn_store - .get::() - .is_some_and(|previous| previous.0 == cache.auth_cache_key) - { - return Ok(()); - } 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(()); } - turn_store.remove::(); if cache.is_auth_current() - && cache.catalog.as_ref().is_some_and(|catalog| { - (waited && *refresh_succeeded) - || (!catalog.is_partial && catalog.warnings.is_empty()) - }) + && 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; - turn_store.insert(LatestCloudSkillsAuth(cache.auth_cache_key.clone())); return Ok(()); } - // Waiting callers must not reuse a failed or cancelled refresh. - *refresh_succeeded = false; (cache, resource_cache_key) }; let mut result = providers.list_cloud_for_turn(query).await; @@ -370,7 +348,7 @@ impl SkillsThreadState { // Ignore superseded results, including errors. return Ok(()); } - let mut catalog = match result { + let catalog = match result { Ok(catalog) => catalog, Err(error) => { catalogs.cloud_retry = Some(CatalogRefreshBackoff::after_failure( @@ -381,18 +359,16 @@ impl SkillsThreadState { return Err(error); } }; - catalogs.cloud_retry = None; - if catalog.is_partial - && let Some(previous) = cache - .catalog - .as_ref() - .filter(|previous| !previous.is_partial) - { - catalog = previous.clone(); - catalog.is_partial = true; + catalogs.cloud_retry = if catalog.is_partial || !catalog.warnings.is_empty() { + Some(CatalogRefreshBackoff::after_failure( + cache.auth_cache_key.clone(), + resource_cache_key.clone(), + catalogs.cloud_retry.as_ref(), + )) } else { - catalogs.cloud_resources = CloudResourceCache::default(); - } + None + }; + catalogs.cloud_resources = CloudResourceCache::default(); record_catalog(self.extension_metrics.as_deref(), "cloud", &catalog); let generation = Arc::new(CloudSkillGeneration { auth_cache_key: cache.auth_cache_key.clone(), @@ -400,9 +376,7 @@ impl SkillsThreadState { mcp_resources: cache.mcp_resources.clone(), catalog: Some(catalog), }); - catalogs.cloud_generation = Some(Arc::clone(&generation)); - turn_store.insert(LatestCloudSkillsAuth(generation.auth_cache_key.clone())); - *refresh_succeeded = true; + catalogs.cloud_generation = Some(generation); Ok(()) } diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index 7c7dbbd9e4..3b87e1ee06 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/codex-rs/ext/skills/tests/skills_extension.rs @@ -2601,14 +2601,18 @@ async fn skills_list_only_returns_model_visible_bounded_metadata() -> TestResult Ok(()) } -#[tokio::test] +#[tokio::test(start_paused = true)] async fn cloud_discovery_is_gated_and_reports_failures() -> TestResult { for (cloud_skill_enabled, has_cloud_provider) in [(false, true), (true, true), (true, false)] { let list_calls = Arc::new(AtomicUsize::new(0)); let mut providers = SkillProviders::new(); if has_cloud_provider { providers = providers.with_cloud_provider(Arc::new(StaticSkillProvider { - catalog: SkillCatalog::default(), + catalog: SkillCatalog { + is_partial: true, + warnings: vec!["catalog diagnostic".to_string()], + ..SkillCatalog::default() + }, read_requests: Arc::new(Mutex::new(Vec::new())), list_calls: Some(Arc::clone(&list_calls)), fail_first_list: true, @@ -2729,6 +2733,13 @@ async fn cloud_discovery_is_gated_and_reports_failures() -> TestResult { assert_eq!(warning.thread_id, thread_store.level_id()); assert_eq!(warning.turn_id.as_deref(), Some("turn-1")); assert!(warning.message.contains("temporary cloud failure")); + + tokio::time::advance(std::time::Duration::from_secs(30)).await; + start_registered_turn(®istry, &session_store, &thread_store, "turn-2").await; + start_registered_turn(®istry, &session_store, &thread_store, "turn-3").await; + assert_eq!(list_calls.load(Ordering::Relaxed), 2); + start_registered_turn(®istry, &session_store, &thread_store, "turn-4").await; + assert_eq!(list_calls.load(Ordering::Relaxed), 2); } Ok(()) }