Retry incomplete cloud catalogs with exponential backoff (#61400)

## What changed

- Keep partial cloud skill and plugin catalogs, and catalogs with warnings, usable while retrying discovery with exponential backoff. Start retries after 5 seconds and cap the delay at 5 minutes; clear backoff after a complete, warning-free result.
- Replace previous catalog entries with each accepted result, including partial results, and clear cached cloud skill resources when publishing a replacement.
- Check catalog freshness across turns and steps using auth and resource identity. Auth or resource changes bypass retry cooldowns, including within the same turn.

## Testing

Add and update tests for retry delays and recovery, concurrent refreshes, auth and resource invalidation, partial catalog replacement, and cloud skill read-cache replacement.

GitOrigin-RevId: 0ce5a9d5596e8a73d7084abbbaeb4fb8cda84aa0
This commit is contained in:
Anthony Tafoya
2026-11-13 23:14:47 +00:00
committed by copyberry
parent e8f5d2bb55
commit 0b297ec5ee
12 changed files with 471 additions and 399 deletions
@@ -10,6 +10,7 @@ use std::sync::atomic::AtomicUsize;
struct CloudTurnProvider {
reply: Mutex<Result<Option<(&'static str, u32)>, ()>>,
lists: AtomicUsize,
partial: AtomicBool,
reads: AtomicUsize,
resource_key: Mutex<Option<codex_mcp::McpResourceServerCacheKey>>,
}
@@ -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"
@@ -932,14 +932,5 @@ async fn cloud_skill_can_read_referenced_resource_without_an_executor() -> Resul
assert!(skill_fragments[0].contains(&format!("<name>{SKILL_NAME}</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(())
}
@@ -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<McpResourceClientAuthKey>,
@@ -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 {
+27 -46
View File
@@ -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<McpResourceClientAuthKey>);
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::<LatestCloudPluginAuth>()
.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::<LatestCloudPluginAuth>();
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::<LatestCloudPluginAuth>();
self.refresh_cloud_catalog(CatalogRefreshInput {
turn_id: input.turn_id,
session_store: input.session_store,
+202 -73
View File
@@ -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::<LatestCloudPluginAuth>();
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::<PluginsThreadState>()
.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(())
}
@@ -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<bool>,
/// Serializes discovery before checking freshness and publishing the next catalog.
pub(crate) cloud_refresh_lock: AsyncMutex<()>,
state: Mutex<PluginContributorState>,
}
@@ -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<McpResourceServerCacheKey>,
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;
@@ -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")?]);
+1 -1
View File
@@ -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<McpResourceClientAuthKey>,
+166 -186
View File
@@ -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::<LatestCloudSkillsAuth>()
.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::<LatestCloudSkillsAuth>().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::<LatestCloudSkillsAuth>().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",
);
}
}
}
+2 -6
View File
@@ -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<C: Send + Sync + 'static> SkillsExtension<C> {
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<C: Send + Sync + 'static> TurnLifecycleContributor for SkillsExtension<C> {
fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> {
Box::pin(async move {
input.turn_store.remove::<LatestCloudSkillsAuth>();
self.refresh_cloud_catalog(CatalogRefreshInput {
turn_id: input.turn_id,
session_store: input.session_store,
+20 -46
View File
@@ -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<Arc<dyn ExtensionMetrics>>,
}
/// Auth scope of the latest catalog accepted for this turn.
pub(crate) struct LatestCloudSkillsAuth(Option<McpResourceClientAuthKey>);
/// 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<bool>,
/// Serializes discovery before checking freshness and publishing the next catalog.
cloud_refresh_lock: AsyncMutex<()>,
config: Mutex<SkillsExtensionConfig>,
cloud_skills_available: bool,
catalog_state: Mutex<SkillCatalogState>,
@@ -74,7 +69,7 @@ impl SkillsThreadState {
extension_metrics: Option<Arc<dyn ExtensionMetrics>>,
) -> 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::<LatestCloudSkillsAuth>()
.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::<LatestCloudSkillsAuth>();
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(())
}
+13 -2
View File
@@ -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(&registry, &session_store, &thread_store, "turn-2").await;
start_registered_turn(&registry, &session_store, &thread_store, "turn-3").await;
assert_eq!(list_calls.load(Ordering::Relaxed), 2);
start_registered_turn(&registry, &session_store, &thread_store, "turn-4").await;
assert_eq!(list_calls.load(Ordering::Relaxed), 2);
}
Ok(())
}