mirror of
https://github.com/openai/codex.git
synced 2026-11-14 21:40:39 +00:00
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:
committed by
copyberry
parent
e8f5d2bb55
commit
0b297ec5ee
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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")?]);
|
||||
|
||||
@@ -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>,
|
||||
|
||||
@@ -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",
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user