Add explicit cloud catalog refresh before turn admission (#61449)

## Why

Hosts need to refresh cloud plugin and skill catalogs without starting a turn or waiting for step preparation.

## What changed

- Add `CodexThread::refresh_cloud_plugins()` and a turn-independent extension callback that bypass cached catalog reuse and failure cooldowns.
- Run catalog contributors concurrently while serializing publication with step snapshot capture. Invalidate MCP resource caches before explicit discovery and schedule MCP reconciliation afterward.
- Accept complete and partial catalogs, report discovery errors, and keep failed or cancelled refreshes retryable.
- Allow hosts to gate cloud plugin discovery through effective configuration.

## Testing

Add coverage for discovery before turn admission, catalog replacement and reuse, immediate retry after failure, concurrent contributors, and cancellation-safe publication and invalidation.

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