From 6d919e65e20dcb1b916be8acd113b394d4dd6cc9 Mon Sep 17 00:00:00 2001 From: Anthony Tafoya Date: Sat, 14 Nov 2026 00:30:43 +0000 Subject: [PATCH] Refresh cloud plugin and skill catalogs during step preparation (#61429) ## Why Installing a plugin without connector tools must still refresh a cached cloud skill catalog so its skills are visible to the next inference. ## What changed - Rely on step preparation for cloud plugin and skill discovery, removing their turn-start callbacks and the `requires_mcp_runtime` lifecycle hook. - Request an MCP runtime refresh after every completed plugin install, including skills-only installs. ## Testing Add regression coverage for skills becoming visible immediately after a connector-free install. Update lifecycle tests to cover discovery after MCP reconciliation, prewarm skipping discovery, cancellation allowing subsequent preparation, and reuse of unchanged connections. Use deterministic DNS lookups in network proxy authorization tests. GitOrigin-RevId: 76036bd29253fb8987e606c3c28411c689590fc4 --- .../core/src/session/turn_start_mcp_tests.rs | 178 ++++++-------- codex-rs/core/src/tasks/lifecycle.rs | 5 - codex-rs/core/src/tasks/regular.rs | 29 +-- .../tools/handlers/request_plugin_install.rs | 6 + .../tests/suite/request_plugin_install.rs | 3 + .../suite/request_plugin_install_skills.rs | 220 ++++++++++++++++++ .../skills_extension/cloud_lifecycle_tests.rs | 4 +- .../ext/extension-api/src/contributors.rs | 6 - codex-rs/ext/mcp/src/cloud_plugin.rs | 32 +-- codex-rs/ext/mcp/src/cloud_plugin_tests.rs | 76 ++---- codex-rs/ext/skills/src/cloud_cache_tests.rs | 3 +- codex-rs/ext/skills/src/cloud_skill.rs | 32 +-- codex-rs/ext/skills/tests/skills_extension.rs | 68 ++---- .../skills_extension/catalog_delta_tests.rs | 2 +- .../skills_extension/merge_catalog_tests.rs | 6 +- .../skills_extension/preparation_tests.rs | 2 +- .../presentation_dedup_tests.rs | 3 +- codex-rs/network-proxy/src/runtime.rs | 24 +- codex-rs/network-proxy/src/socks5.rs | 1 + 19 files changed, 368 insertions(+), 332 deletions(-) create mode 100644 codex-rs/core/tests/suite/request_plugin_install_skills.rs diff --git a/codex-rs/core/src/session/turn_start_mcp_tests.rs b/codex-rs/core/src/session/turn_start_mcp_tests.rs index 3bf66d8955..096d03896f 100644 --- a/codex-rs/core/src/session/turn_start_mcp_tests.rs +++ b/codex-rs/core/src/session/turn_start_mcp_tests.rs @@ -1,4 +1,4 @@ -//! Turn-start MCP reconciliation, discovery cancellation, and prewarm refresh boundaries. +//! Step discovery observes ready MCP connections; prewarm skips discovery and cancellation releases preparation. use super::*; use crate::McpRefreshReason; @@ -8,51 +8,36 @@ use pretty_assertions::assert_eq; struct Recorder { client: McpResourceClient, - requires_mcp: std::sync::atomic::AtomicBool, + pause: std::sync::atomic::AtomicBool, keys: std::sync::Mutex>, discovered: tokio::sync::Notify, } impl codex_extension_api::TurnLifecycleContributor for Recorder { - fn turn_start_phase( - &self, - _thread_store: &codex_extension_api::ExtensionData, - ) -> codex_extension_api::TurnStartPhase { - codex_extension_api::TurnStartPhase::RegularTaskStart - } - - fn requires_mcp_runtime(&self, _thread_store: &codex_extension_api::ExtensionData) -> bool { - self.requires_mcp.load(std::sync::atomic::Ordering::Relaxed) - } - - fn on_turn_start<'a>( + fn refresh_catalog_for_step<'a>( &'a self, - _input: codex_extension_api::TurnStartInput<'a>, - ) -> codex_extension_api::ExtensionFuture<'a, ()> { + _input: codex_extension_api::CatalogRefreshInput<'a>, + ) -> codex_extension_api::ExtensionFuture<'a, bool> { Box::pin(async move { self.keys .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .push(self.client.cache_key()); self.discovered.notify_one(); - std::future::pending::<()>().await; + if self.pause.swap(false, std::sync::atomic::Ordering::Relaxed) { + std::future::pending::<()>().await; + } + false }) } - - fn refresh_catalog_for_step<'a>( - &'a self, - _input: codex_extension_api::CatalogRefreshInput<'a>, - ) -> codex_extension_api::ExtensionFuture<'a, bool> { - panic!("prewarm must not start cloud discovery"); - } } #[tokio::test] -async fn turn_start_refreshes_dirty_mcp_reuses_clean_runtime_and_cancels_discovery() { +async fn step_discovery_reconciles_mcp_skips_prewarm_and_cancels_preparation() { let (mut session, turn_context) = make_session_and_context().await; let recorder = Arc::new(Recorder { client: McpResourceClient::new(Arc::clone(&session.services.mcp_runtime)), - requires_mcp: Default::default(), + pause: std::sync::atomic::AtomicBool::new(true), keys: Default::default(), discovered: Default::default(), }); @@ -67,94 +52,63 @@ async fn turn_start_refreshes_dirty_mcp_reuses_clean_runtime_and_cancels_discove session .mcp_runtime_for_step(&prewarm, &[], &[], &HashSet::new()) .await; + assert!(recorder.keys.lock().expect("recorded keys").is_empty()); let initial_key = recorder.client.cache_key(); - for turn in 0..4 { - recorder - .requires_mcp - .store(turn > 0, std::sync::atomic::Ordering::Relaxed); - if turn == 0 { - session.reproject_mcp_on_next_turn(); - } else if turn == 3 { - session.services.mcp_runtime.reconnect_on_next_refresh(); - session.request_mcp_runtime_refresh(McpRefreshReason::ExplicitInvalidation); - } - session - .spawn_task( - Arc::clone(&turn_context), - Vec::new(), - crate::tasks::RegularTask::new(), - ) - .await; - tokio::time::timeout( - std::time::Duration::from_secs(5), - recorder.discovered.notified(), + session + .spawn_task( + Arc::clone(&turn_context), + Vec::new(), + HeldStepTask { + kind: TaskKind::Regular, + finish: Arc::new(tokio::sync::Notify::new()), + }, ) - .await - .expect("regular-task contributor should run"); - assert_eq!(session.mcp_refresh.is_pending(), turn == 0); - assert_eq!( - recorder - .keys - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .len(), - turn + 1 - ); - if turn > 0 { - session.refresh_mcp_if_dirty().await; - assert!( - recorder.client.cache_key() - == recorder - .keys - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner)[turn], - "a clean refresh must reuse the published MCP runtime" - ); - // Prewarm must also skip discovery while a regular turn is active. - session - .mcp_runtime_for_step(&prewarm, &[], &[], &HashSet::new()) - .await; + .await; + session.mark_mcp_runtime_dirty(McpRefreshReason::ExplicitInvalidation); + let cancellation = CancellationToken::new(); + let capture = session.capture_step_context(Arc::clone(&turn_context), &cancellation); + tokio::pin!(capture); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + tokio::select! { + _ = &mut capture => panic!("step must wait for discovery"), + () = recorder.discovered.notified() => {} } - let (cancellation_token, done) = { - let active_turn = session.active_turn.lock().await; - let task = active_turn - .as_ref() - .and_then(|turn| turn.task.as_ref()) - .expect("discovery should run inside the registered task"); - (task.cancellation_token.clone(), Arc::clone(&task.done)) - }; - let completed = done.notified(); - tokio::pin!(completed); - completed.as_mut().enable(); - cancellation_token.cancel(); - tokio::time::timeout(std::time::Duration::from_secs(5), completed) - .await - .expect("discovery cancellation must finish the task without forced abort"); - assert!( - session.mcp_refresh.is_pending(), - "cancelled discovery must request an MCP refresh before the task completes" - ); - session.abort_all_tasks(TurnAbortReason::Interrupted).await; - } - let keys = recorder - .keys - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - assert!( - keys[0] == initial_key, - "unrelated contributors must not reconcile MCP" - ); - assert!( - keys[1] != initial_key, - "discovery must observe the reconciled generation" - ); - assert!( - keys[2] != keys[1], - "cancelled cloud preparation must request reprojection on the next turn" - ); - assert!( - keys[3] != keys[2], - "the next turn must consume a pending reconnect before discovery" - ); + }) + .await + .expect("regular step should discover catalogs"); + assert!(!session.mcp_refresh.is_pending()); + let prepared_key = recorder.client.cache_key(); + assert!(prepared_key != initial_key); + assert!(recorder.keys.lock().expect("recorded keys")[0] == prepared_key); + + // Prewarm also skips discovery while a regular turn is preparing a step. + session + .mcp_runtime_for_step(&prewarm, &[], &[], &HashSet::new()) + .await; + assert_eq!(recorder.keys.lock().expect("recorded keys").len(), 1); + cancellation.cancel(); + assert!(capture.await.is_err()); + + tokio::time::timeout( + std::time::Duration::from_secs(5), + session.capture_step_context(Arc::clone(&turn_context), &CancellationToken::new()), + ) + .await + .expect("cancelled discovery must allow the next step") + .expect("next step should prepare"); + assert!(recorder.client.cache_key() == prepared_key); + + session.services.mcp_runtime.reconnect_on_next_refresh(); + session.request_mcp_runtime_refresh(McpRefreshReason::ExplicitInvalidation); + session + .capture_step_context(turn_context, &CancellationToken::new()) + .await + .expect("step should consume pending reconnect before discovery"); + let keys = recorder.keys.lock().expect("recorded keys").clone(); + assert_eq!(keys.len(), 3); + assert!(keys[1] == prepared_key); + assert!(keys[2] != prepared_key); + assert!(keys[2] == recorder.client.cache_key()); + session.abort_all_tasks(TurnAbortReason::Interrupted).await; } diff --git a/codex-rs/core/src/tasks/lifecycle.rs b/codex-rs/core/src/tasks/lifecycle.rs index 9fab93cd10..20b8252202 100644 --- a/codex-rs/core/src/tasks/lifecycle.rs +++ b/codex-rs/core/src/tasks/lifecycle.rs @@ -30,11 +30,6 @@ impl Session { if contributor.turn_start_phase(&self.services.thread_extension_data) != phase { continue; } - if phase == TurnStartPhase::RegularTaskStart - && contributor.requires_mcp_runtime(&self.services.thread_extension_data) - { - self.refresh_mcp_if_dirty().await; - } contributor .on_turn_start(codex_extension_api::TurnStartInput { turn_id: turn_context.sub_id.as_str(), diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index dcb07bfd48..6b0d594b97 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -1,4 +1,3 @@ -use crate::McpRefreshReason; use crate::agent::helper_daybreak::turn::SavedProgramResolution; use std::sync::Arc; @@ -75,36 +74,14 @@ impl SessionTask for RegularTask { sess.emit_turn_started(&ctx, TaskKind::Regular).await; sess.start_deferred_mcp_reprojection(); // Regular-start contributors run once, after the task is visible and interruptible. - let prepares_mcp = sess - .services - .extensions - .turn_lifecycle_contributors() - .iter() - .any(|contributor| { - contributor.turn_start_phase(&sess.services.thread_extension_data) - == TurnStartPhase::RegularTaskStart - && contributor.requires_mcp_runtime(&sess.services.thread_extension_data) - }); - let preparation = async { - sess.emit_turn_start_lifecycle( + let preparation = sess + .emit_turn_start_lifecycle( &ctx, /*token_usage_at_turn_start*/ None, TurnStartPhase::RegularTaskStart, ) + .or_cancel(&cancellation_token) .await; - if prepares_mcp { - // Resource discovery can overlap the fetch. Wait before reprojection. - sess.wait_for_apps_tool_refresh().await; - } - } - .or_cancel(&cancellation_token) - .await; - // Even cancelled discovery may have cleared the previous account's catalog. - if prepares_mcp { - sess.request_mcp_runtime_reprojection( - McpRefreshReason::TurnStartContributorPreparedMcp, - ); - } if preparation.is_err() { return SessionStartupPrewarmResolution::Cancelled; } diff --git a/codex-rs/core/src/tools/handlers/request_plugin_install.rs b/codex-rs/core/src/tools/handlers/request_plugin_install.rs index 84f2f875d3..b7c5b70bbb 100644 --- a/codex-rs/core/src/tools/handlers/request_plugin_install.rs +++ b/codex-rs/core/src/tools/handlers/request_plugin_install.rs @@ -275,6 +275,12 @@ impl RequestPluginInstallHandler { false }; + if completed { + // Skills-only installs have no connector tools to refresh. Dirty discovery too, + // so the next step can see the newly installed plugin and its skills. + session.request_mcp_runtime_refresh(crate::McpRefreshReason::PluginInstalled); + } + if completed && let DiscoverableTool::Connector(connector) = &tool { session .merge_connector_selection(HashSet::from([connector.id.clone()])) diff --git a/codex-rs/core/tests/suite/request_plugin_install.rs b/codex-rs/core/tests/suite/request_plugin_install.rs index fe5251fbee..bc28be7c6b 100644 --- a/codex-rs/core/tests/suite/request_plugin_install.rs +++ b/codex-rs/core/tests/suite/request_plugin_install.rs @@ -78,6 +78,9 @@ use super::rmcp_client::remote_aware_stdio_server_bin; #[path = "request_plugin_install_runtime_auth.rs"] mod runtime_auth_tests; +#[path = "request_plugin_install_skills.rs"] +mod skills_tests; + const TOOL_SEARCH_TOOL_NAME: &str = "tool_search"; const LIST_AVAILABLE_PLUGINS_TO_INSTALL_TOOL_NAME: &str = "list_available_plugins_to_install"; const REQUEST_PLUGIN_INSTALL_TOOL_NAME: &str = "request_plugin_install"; diff --git a/codex-rs/core/tests/suite/request_plugin_install_skills.rs b/codex-rs/core/tests/suite/request_plugin_install_skills.rs new file mode 100644 index 0000000000..08efca3338 --- /dev/null +++ b/codex-rs/core/tests/suite/request_plugin_install_skills.rs @@ -0,0 +1,220 @@ +//! Connector-free installs invalidate a warmed cloud skill catalog before the next inference. + +use super::*; +use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_skills_extension::SkillProvider; +use codex_skills_extension::SkillProviders; +use codex_skills_extension::SkillsExtensionConfig; +use codex_skills_extension::catalog::SkillAuthority; +use codex_skills_extension::catalog::SkillCatalog; +use codex_skills_extension::catalog::SkillCatalogEntry; +use codex_skills_extension::catalog::SkillPackageId; +use codex_skills_extension::catalog::SkillProviderError; +use codex_skills_extension::catalog::SkillReadResult; +use codex_skills_extension::catalog::SkillResourceId; +use codex_skills_extension::catalog::SkillSearchResult; +use codex_skills_extension::catalog::SkillSourceKind; +use codex_skills_extension::install_with_providers; +use codex_skills_extension::provider::SkillListQuery; +use codex_skills_extension::provider::SkillProviderFuture; +use codex_skills_extension::provider::SkillReadRequest; +use codex_skills_extension::provider::SkillSearchRequest; +use pretty_assertions::assert_eq; +use std::sync::atomic::AtomicUsize; + +const SKILL_DESCRIPTION: &str = "Plan a week using the newly installed Calendar skill."; + +#[derive(Default)] +struct InstallableCloudSkillProvider { + installed: AtomicBool, + lists: AtomicUsize, +} + +impl SkillProvider for InstallableCloudSkillProvider { + fn list<'a>(&'a self, _query: SkillListQuery<'a>) -> SkillProviderFuture<'a, SkillCatalog> { + Box::pin(async move { + self.lists.fetch_add(1, Ordering::SeqCst); + let entries = if self.installed.load(Ordering::SeqCst) { + vec![ + SkillCatalogEntry::new( + SkillPackageId("skill://calendar/plan-week".to_string()), + SkillAuthority::new(SkillSourceKind::Cloud, CODEX_APPS_MCP_SERVER_NAME), + "calendar:plan-week", + SKILL_DESCRIPTION, + SkillResourceId::new("skill://calendar/plan-week/SKILL.md"), + ) + .with_display_path("skill://calendar/plan-week") + .with_alias_root("skill://calendar"), + ] + } else { + Vec::new() + }; + Ok(SkillCatalog { + is_partial: false, + entries, + warnings: Vec::new(), + }) + }) + } + + fn read<'a>( + &'a self, + _request: SkillReadRequest<'a>, + ) -> SkillProviderFuture<'a, SkillReadResult> { + Box::pin(async { Err(SkillProviderError::new("this test only discovers skills")) }) + } + + fn search(&self, _request: SkillSearchRequest) -> SkillProviderFuture<'_, SkillSearchResult> { + Box::pin(async { Ok(SkillSearchResult::default()) }) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_skills_only_install_refreshes_warmed_cloud_catalog() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let apps_server = AppsTestServer::mount_without_tools(&server).await?; + mount_recommendations( + &server, + ResponseTemplate::new(200).set_body_json(json!({ + "enabled": true, + "plugins": [{ + "id": REMOTE_CALENDAR_PLUGIN_ID, + "name": "calendar", + "display_name": "Calendar" + }] + })), + ) + .await; + Mock::given(method("GET")) + .and(path(format!("/ps/plugins/{REMOTE_CALENDAR_PLUGIN_ID}"))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": REMOTE_CALENDAR_PLUGIN_ID, + "name": "calendar", + "scope": "GLOBAL", + "status": "ENABLED", + "installation_policy": "AVAILABLE", + "authentication_policy": "ON_USE", + "release": { + "display_name": "Calendar", + "description": "Calendar planning skills without connector tools.", + "app_ids": [], + "interface": {} + } + }))) + .expect(1) + .mount(&server) + .await; + let initially_installed = mount_empty_remote_installed_plugins(&server).await; + let provider = Arc::new(InstallableCloudSkillProvider::default()); + let mut extensions = ExtensionRegistryBuilder::new(); + 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 mut builder = test_codex() + .with_exec_server_url("none") + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_extensions(Arc::new(extensions.build())) + .with_config(move |config| { + config + .permissions + .set_permission_profile(PermissionProfile::Disabled) + .expect("test config should allow disabled permissions"); + configure_apps_without_search_tool(config, &apps_server.chatgpt_base_url); + config.include_skill_instructions = true; + config.cloud_skill_enabled = true; + }); + let test = builder.build_with_auto_env(&server).await?; + + let warmup = mount_sse_once( + &server, + sse(vec![ev_response_created("warmup"), ev_completed("warmup")]), + ) + .await; + test.submit_turn("What skills are available?").await?; + assert_eq!(provider.lists.load(Ordering::SeqCst), 1); + assert!( + !warmup + .single_request() + .message_input_texts("developer") + .join("\n") + .contains(SKILL_DESCRIPTION) + ); + + let install_call_id = "install-calendar-skills"; + let mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("install"), + ev_function_call( + install_call_id, + REQUEST_PLUGIN_INSTALL_TOOL_NAME, + &json!({ + "plugin_id": REMOTE_CALENDAR_PLUGIN_CONFIG_ID, + "suggest_reason": "Plan my week" + }) + .to_string(), + ), + ev_completed("install"), + ]), + sse(vec![ + ev_response_created("installed"), + ev_completed("installed"), + ]), + ], + ) + .await; + let elicitation = start_install_turn(&test, "Install Calendar planning skills").await?; + // A second turn must reuse the warmed empty catalog until installation completes. + assert_eq!(provider.lists.load(Ordering::SeqCst), 1); + mount_remote_calendar_installed_plugins(&server).await; + drop(initially_installed); + provider.installed.store(true, Ordering::SeqCst); + resolve_install_elicitation(&test, elicitation, ElicitationAction::Accept).await?; + + let requests = mock.requests(); + assert_eq!(requests.len(), 2); + assert_eq!( + serde_json::from_str::( + &requests[1] + .function_call_output_text(install_call_id) + .expect("install tool output") + )?, + json!({ + "completed": true, + "user_confirmed": true, + "tool_type": "plugin", + "action_type": "install", + "tool_id": REMOTE_CALENDAR_PLUGIN_CONFIG_ID, + "tool_name": "Calendar", + "suggest_reason": "Plan my week" + }) + ); + assert!( + !requests[0] + .message_input_texts("developer") + .join("\n") + .contains(SKILL_DESCRIPTION) + ); + assert!( + requests[1] + .message_input_texts("developer") + .join("\n") + .contains(SKILL_DESCRIPTION), + "the inference immediately after a connector-free install must expose its cloud skill" + ); + assert_eq!(provider.lists.load(Ordering::SeqCst), 2); + Ok(()) +} diff --git a/codex-rs/core/tests/suite/skills_extension/cloud_lifecycle_tests.rs b/codex-rs/core/tests/suite/skills_extension/cloud_lifecycle_tests.rs index 5d3556f8c1..07a01fa035 100644 --- a/codex-rs/core/tests/suite/skills_extension/cloud_lifecycle_tests.rs +++ b/codex-rs/core/tests/suite/skills_extension/cloud_lifecycle_tests.rs @@ -249,7 +249,7 @@ async fn cloud_skills_reuse_cache_and_invalidate_on_connection_or_auth_change() .expect("thread MCP resource client"); let auth = client.auth_cache_key_for_server(CODEX_APPS_MCP_SERVER_NAME); let connection = client.server_cache_key(CODEX_APPS_MCP_SERVER_NAME); - // Turn-start preparation must consume the pending reconnect before discovery. + // Step preparation must consume the pending reconnect before discovery. test.codex.submit(Op::RefreshMcpServers).await?; Some((client, auth, connection)) } else { @@ -325,7 +325,7 @@ async fn cloud_skills_reuse_cache_and_invalidate_on_connection_or_auth_change() assert!( client.server_cache_key(CODEX_APPS_MCP_SERVER_NAME) == *provider.resource_key.lock().expect("resource key lock"), - "turn-start reprojection must not invalidate resource caches: {turn}" + "step reprojection must not invalidate resource caches: {turn}" ); if let Some((client, scope, connection)) = previous_connection_scope { assert!(client.server_cache_key(CODEX_APPS_MCP_SERVER_NAME) != connection); diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 581b5aaab7..0b65d6fb34 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -222,12 +222,6 @@ pub trait TurnLifecycleContributor: Send + Sync { TurnStartPhase::BeforeTaskRegistration } - /// Whether regular-task startup must reconcile MCP before this callback runs. - /// Ignored before task registration; contributors unrelated to MCP keep the default. - fn requires_mcp_runtime(&self, _thread_store: &ExtensionData) -> bool { - false - } - /// Called once in the selected phase, before task-specific work begins. /// Regular-task preparation may be cancelled before this callback completes; /// stop and abort callbacks must tolerate missing or partial initialization. diff --git a/codex-rs/ext/mcp/src/cloud_plugin.rs b/codex-rs/ext/mcp/src/cloud_plugin.rs index 5c7c19d800..4c7429fa5a 100644 --- a/codex-rs/ext/mcp/src/cloud_plugin.rs +++ b/codex-rs/ext/mcp/src/cloud_plugin.rs @@ -1,4 +1,4 @@ -//! Reuses complete, warning-free cloud plugins until their auth or Apps resources change. +//! 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. @@ -17,8 +17,6 @@ use codex_extension_api::McpServerContribution; use codex_extension_api::ThreadLifecycleContributor; use codex_extension_api::ThreadStartInput; use codex_extension_api::TurnLifecycleContributor; -use codex_extension_api::TurnStartInput; -use codex_extension_api::TurnStartPhase; use codex_features::Feature; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::McpResourceClient; @@ -179,34 +177,6 @@ impl ConfigContributor for PluginContributor { } impl TurnLifecycleContributor for PluginContributor { - fn turn_start_phase(&self, thread_store: &ExtensionData) -> TurnStartPhase { - if self.providers.cloud.is_some() - && thread_store - .get::() - .is_some_and(|state| state.contributor_state().plugins_enabled) - { - TurnStartPhase::RegularTaskStart - } else { - TurnStartPhase::BeforeTaskRegistration - } - } - - fn requires_mcp_runtime(&self, thread_store: &ExtensionData) -> bool { - self.turn_start_phase(thread_store) == TurnStartPhase::RegularTaskStart - } - - fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> { - Box::pin(async move { - self.refresh_cloud_catalog(CatalogRefreshInput { - turn_id: input.turn_id, - session_store: input.session_store, - thread_store: input.thread_store, - turn_store: input.turn_store, - }) - .await; - }) - } - fn refresh_catalog_for_step<'a>( &'a self, input: CatalogRefreshInput<'a>, diff --git a/codex-rs/ext/mcp/src/cloud_plugin_tests.rs b/codex-rs/ext/mcp/src/cloud_plugin_tests.rs index 3559a3a955..8a939ffe8d 100644 --- a/codex-rs/ext/mcp/src/cloud_plugin_tests.rs +++ b/codex-rs/ext/mcp/src/cloud_plugin_tests.rs @@ -23,9 +23,6 @@ use codex_extension_api::WorldStateContributionInput; use codex_extension_api::WorldStateSectionContribution; use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; -use codex_protocol::config_types::CollaborationMode; -use codex_protocol::config_types::ModeKind; -use codex_protocol::config_types::Settings; use codex_protocol::protocol::Op; use codex_utils_path_uri::PathUri; use core_test_support::apps_test_server::AppsTestServer; @@ -48,54 +45,6 @@ use std::sync::atomic::Ordering; use std::time::Duration; use tokio::sync::Notify; -#[tokio::test] -async fn cloud_discovery_selects_regular_task_phase() { - struct ExistingContributor; - impl TurnLifecycleContributor for ExistingContributor {} - - let mut builder = ExtensionRegistryBuilder::new(); - builder.turn_lifecycle_contributor(Arc::new(ExistingContributor)); - let executor = Arc::new(ExecutorPluginProvider::new(Arc::new( - EnvironmentManager::default_for_tests(), - ))); - install_plugin_providers(&mut builder, PluginProviders::new(executor.clone())); - install_plugin_providers( - &mut builder, - PluginProviders::new(executor).with_cloud_provider(Arc::new(CloudProvider::default())), - ); - let registry = builder.build(); - let thread = ExtensionData::new("thread"); - for plugins_enabled in [false, true] { - thread - .get_or_init(PluginsThreadState::default) - .contributor_state() - .plugins_enabled = plugins_enabled; - let dependencies = registry - .turn_lifecycle_contributors() - .iter() - .map(|contributor| contributor.requires_mcp_runtime(&thread)) - .collect::>(); - assert_eq!(dependencies, vec![false, false, plugins_enabled]); - let phases = registry - .turn_lifecycle_contributors() - .iter() - .map(|contributor| contributor.turn_start_phase(&thread)) - .collect::>(); - assert_eq!( - phases, - vec![ - TurnStartPhase::BeforeTaskRegistration, - TurnStartPhase::BeforeTaskRegistration, - if plugins_enabled { - TurnStartPhase::RegularTaskStart - } else { - TurnStartPhase::BeforeTaskRegistration - }, - ], - ); - } -} - #[tokio::test] async fn disabling_plugins_clears_cloud_catalog_and_skips_discovery() -> anyhow::Result<()> { let config_home = tempfile::tempdir()?; @@ -104,14 +53,6 @@ async fn disabling_plugins_clears_cloud_catalog_and_skips_discovery() -> anyhow: .fallback_cwd(Some(config_home.path().to_path_buf())) .build() .await?; - let mode = CollaborationMode { - mode: ModeKind::Default, - settings: Settings { - model: "test".to_string(), - reasoning_effort: None, - developer_instructions: None, - }, - }; let provider = Arc::new(CloudProvider::default()); let mut builder = ExtensionRegistryBuilder::new(); let providers = PluginProviders::new(Arc::new(ExecutorPluginProvider::new(Arc::new( @@ -154,10 +95,8 @@ async fn disabling_plugins_clears_cloud_catalog_and_skips_discovery() -> anyhow: } for contributor in registry.turn_lifecycle_contributors() { contributor - .on_turn_start(TurnStartInput { + .refresh_catalog_for_step(CatalogRefreshInput { turn_id, - collaboration_mode: &mode, - token_usage_at_turn_start: None, session_store: &session, thread_store: &thread, turn_store: &ExtensionData::new(turn_id), @@ -921,6 +860,19 @@ async fn run_cloud_catalog_refresh_lifecycle() -> anyhow::Result<()> { ) }) .await??; + // Superseded discovery retries on the next step, not before this request. + assert_eq!( + provider.calls.load(Ordering::SeqCst), + previous_discoveries + 1 + ); + let retry = responses::mount_sse_once( + &server, + sse(vec![ev_response_created("retry"), ev_completed("retry")]), + ) + .await; + test.submit_turn("Retry the catalog after the connection change.") + .await?; + assert_eq!(retry.requests().len(), 1); // Failed retries retain a fallback only for the same auth. let state = test .codex diff --git a/codex-rs/ext/skills/src/cloud_cache_tests.rs b/codex-rs/ext/skills/src/cloud_cache_tests.rs index 710c744a98..0f7764e25c 100644 --- a/codex-rs/ext/skills/src/cloud_cache_tests.rs +++ b/codex-rs/ext/skills/src/cloud_cache_tests.rs @@ -228,8 +228,7 @@ async fn step_cloud_refresh_reuses_accepted_results_and_retries_failures() { let registry = builder.build(); let contributor = registry .turn_lifecycle_contributors() - .iter() - .find(|contributor| contributor.requires_mcp_runtime(&thread_store)) + .last() .expect("cloud skills contributor should be registered"); let refresh_input = CatalogRefreshInput { turn_id: "turn", diff --git a/codex-rs/ext/skills/src/cloud_skill.rs b/codex-rs/ext/skills/src/cloud_skill.rs index e573764f42..e3974282bc 100644 --- a/codex-rs/ext/skills/src/cloud_skill.rs +++ b/codex-rs/ext/skills/src/cloud_skill.rs @@ -1,11 +1,8 @@ -//! Refreshes cloud skills at turn start and through step freshness checks. +//! Refreshes cloud skills before steps when their auth or Apps resources change. use codex_extension_api::CatalogRefreshInput; -use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionFuture; use codex_extension_api::TurnLifecycleContributor; -use codex_extension_api::TurnStartInput; -use codex_extension_api::TurnStartPhase; use super::SkillsExtension; use crate::provider::SkillListQuery; @@ -55,33 +52,6 @@ impl SkillsExtension { } impl TurnLifecycleContributor for SkillsExtension { - fn turn_start_phase(&self, thread_store: &ExtensionData) -> TurnStartPhase { - if thread_store - .get::() - .is_some_and(|state| state.cloud_skill_enabled()) - { - TurnStartPhase::RegularTaskStart - } else { - TurnStartPhase::BeforeTaskRegistration - } - } - - fn requires_mcp_runtime(&self, thread_store: &ExtensionData) -> bool { - self.turn_start_phase(thread_store) == TurnStartPhase::RegularTaskStart - } - - fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> { - Box::pin(async move { - self.refresh_cloud_catalog(CatalogRefreshInput { - turn_id: input.turn_id, - session_store: input.session_store, - thread_store: input.thread_store, - turn_store: input.turn_store, - }) - .await; - }) - } - fn refresh_catalog_for_step<'a>( &'a self, input: CatalogRefreshInput<'a>, diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index 3b87e1ee06..435622d212 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/codex-rs/ext/skills/tests/skills_extension.rs @@ -40,7 +40,6 @@ use codex_extension_api::ToolCall; use codex_extension_api::ToolCallSource; use codex_extension_api::ToolPayload; use codex_extension_api::TurnInputContext; -use codex_extension_api::TurnStartInput; use codex_extension_api::WorldStateContributionInput; use codex_extension_api::WorldStateSectionContribution; use codex_models_manager::model_info::model_info_from_slug; @@ -312,27 +311,17 @@ async fn skill_world_state_fragments( Ok((executor, cloud)) } -async fn start_registered_turn( +async fn refresh_registered_catalogs_for_step( registry: &codex_extension_api::ExtensionRegistry, session_store: &ExtensionData, thread_store: &ExtensionData, turn_id: &str, ) { let turn_store = ExtensionData::new(turn_id); - let mode = codex_protocol::config_types::CollaborationMode { - mode: codex_protocol::config_types::ModeKind::Default, - settings: codex_protocol::config_types::Settings { - model: "test".to_string(), - reasoning_effort: None, - developer_instructions: None, - }, - }; for contributor in registry.turn_lifecycle_contributors() { contributor - .on_turn_start(TurnStartInput { + .refresh_catalog_for_step(codex_extension_api::CatalogRefreshInput { turn_id, - collaboration_mode: &mode, - token_usage_at_turn_start: None, session_store, thread_store, turn_store: &turn_store, @@ -754,7 +743,7 @@ async fn executor_cloud_and_host_share_catalog_world_state_flow() -> TestResult }) .await; // Host rendering stays separate even when its identity matches a Cloud skill. - start_registered_turn(®istry, &session_store, &thread_store, "turn-1").await; + refresh_registered_catalogs_for_step(®istry, &session_store, &thread_store, "turn-1").await; let selected_roots = [SelectedCapabilityRoot { id: "skills".to_string(), @@ -1920,7 +1909,7 @@ async fn default_context_truncates_catalog_descriptions() -> TestResult { thread_store: &thread_store, }) .await; - start_registered_turn(®istry, &session_store, &thread_store, "turn-1").await; + refresh_registered_catalogs_for_step(®istry, &session_store, &thread_store, "turn-1").await; let (executor, cloud) = skill_world_state_fragments(®istry, &session_store, &thread_store, "turn-1").await?; @@ -2086,7 +2075,13 @@ async fn catalog_rebalances_only_to_avoid_omissions_and_retains_the_allocation() ); let turn_id = format!("turn-{step}"); let turn_store = ExtensionData::new(&turn_id); - start_registered_turn(®istry, &session_store, &thread_store, &turn_id).await; + refresh_registered_catalogs_for_step( + ®istry, + &session_store, + &thread_store, + &turn_id, + ) + .await; let sections = prepared_world_state( ®istry, if ready { &selected_roots } else { &[] }, @@ -2270,7 +2265,7 @@ async fn extreme_budget_pressure_removes_descriptions_before_omitting_entries() thread_store: &thread_store, }) .await; - start_registered_turn(®istry, &session_store, &thread_store, "turn-1").await; + refresh_registered_catalogs_for_step(®istry, &session_store, &thread_store, "turn-1").await; let (executor, cloud) = skill_world_state_fragments(®istry, &session_store, &thread_store, "turn-1").await?; @@ -2383,7 +2378,7 @@ async fn skills_list_only_returns_model_visible_bounded_metadata() -> TestResult thread_store: &thread_store, }) .await; - start_registered_turn(®istry, &session_store, &thread_store, "turn-1").await; + refresh_registered_catalogs_for_step(®istry, &session_store, &thread_store, "turn-1").await; let step_store = ExtensionData::new("step"); prepared_world_state( @@ -2644,30 +2639,8 @@ async fn cloud_discovery_is_gated_and_reports_failures() -> TestResult { }) .await; - assert_eq!( - registry - .turn_lifecycle_contributors() - .iter() - .map(|contributor| contributor.turn_start_phase(&thread_store)) - .collect::>(), - vec![ - codex_extension_api::TurnStartPhase::BeforeTaskRegistration, - if cloud_available { - codex_extension_api::TurnStartPhase::RegularTaskStart - } else { - codex_extension_api::TurnStartPhase::BeforeTaskRegistration - }, - ] - ); - assert_eq!( - registry - .turn_lifecycle_contributors() - .iter() - .map(|contributor| contributor.requires_mcp_runtime(&thread_store)) - .collect::>(), - vec![false, cloud_available] - ); - start_registered_turn(®istry, &session_store, &thread_store, "turn-1").await; + refresh_registered_catalogs_for_step(®istry, &session_store, &thread_store, "turn-1") + .await; assert_eq!( list_calls.load(Ordering::Relaxed), usize::from(cloud_available) @@ -2735,10 +2708,13 @@ async fn cloud_discovery_is_gated_and_reports_failures() -> TestResult { 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; + refresh_registered_catalogs_for_step(®istry, &session_store, &thread_store, "turn-2") + .await; + refresh_registered_catalogs_for_step(®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; + refresh_registered_catalogs_for_step(®istry, &session_store, &thread_store, "turn-4") + .await; assert_eq!(list_calls.load(Ordering::Relaxed), 2); } Ok(()) @@ -3120,7 +3096,7 @@ async fn model_context_window_scales_executor_and_cloud_catalogs() -> TestResult thread_store: &thread_store, }) .await; - start_registered_turn(®istry, &session_store, &thread_store, "turn-1").await; + refresh_registered_catalogs_for_step(®istry, &session_store, &thread_store, "turn-1").await; let mut model_info = model_info_from_slug("test-model"); model_info.context_window = Some(10_000); diff --git a/codex-rs/ext/skills/tests/skills_extension/catalog_delta_tests.rs b/codex-rs/ext/skills/tests/skills_extension/catalog_delta_tests.rs index 6f51122c8b..b2c09632a3 100644 --- a/codex-rs/ext/skills/tests/skills_extension/catalog_delta_tests.rs +++ b/codex-rs/ext/skills/tests/skills_extension/catalog_delta_tests.rs @@ -296,7 +296,7 @@ async fn render_catalogs_for_step( thread_store: &thread_store, }) .await; - start_registered_turn(®istry, &session_store, &thread_store, turn_id).await; + refresh_registered_catalogs_for_step(®istry, &session_store, &thread_store, turn_id).await; let roots = [SelectedCapabilityRoot { id: "skills".to_string(), location: CapabilityRootLocation::Environment { diff --git a/codex-rs/ext/skills/tests/skills_extension/merge_catalog_tests.rs b/codex-rs/ext/skills/tests/skills_extension/merge_catalog_tests.rs index 04d2c4aa6b..ca7e0f9e3b 100644 --- a/codex-rs/ext/skills/tests/skills_extension/merge_catalog_tests.rs +++ b/codex-rs/ext/skills/tests/skills_extension/merge_catalog_tests.rs @@ -137,7 +137,7 @@ impl MergeHarness { thread_store: &harness.thread_store, }) .await; - start_registered_turn( + refresh_registered_catalogs_for_step( &harness.registry, &harness.session_store, &harness.thread_store, @@ -471,7 +471,7 @@ async fn executor_aliases_survive_cloud_promotion_without_reassigning_remaining_ ); harness.turn_store = ExtensionData::new("turn-2"); harness.step_store = ExtensionData::new("step-2"); - start_registered_turn( + refresh_registered_catalogs_for_step( &harness.registry, &harness.session_store, &harness.thread_store, @@ -1374,7 +1374,7 @@ async fn v2_shadow_observation_tracks_merged_skills_independently_of_read_route( for turn_id in ["turn-1", "turn-2"] { harness.turn_store = ExtensionData::new(turn_id); harness.step_store = ExtensionData::new(format!("{turn_id}-step")); - start_registered_turn( + refresh_registered_catalogs_for_step( &harness.registry, &harness.session_store, &harness.thread_store, diff --git a/codex-rs/ext/skills/tests/skills_extension/preparation_tests.rs b/codex-rs/ext/skills/tests/skills_extension/preparation_tests.rs index f937ce84e0..301a23c4fe 100644 --- a/codex-rs/ext/skills/tests/skills_extension/preparation_tests.rs +++ b/codex-rs/ext/skills/tests/skills_extension/preparation_tests.rs @@ -85,7 +85,7 @@ async fn config_reload_during_preparation_applies_to_next_step_but_reads_still_c thread_store: &thread, }) .await; - start_registered_turn(®istry, &session, &thread, "turn").await; + refresh_registered_catalogs_for_step(®istry, &session, &thread, "turn").await; let roots = vec![SelectedCapabilityRoot { id: "skills".into(), location: CapabilityRootLocation::Environment { diff --git a/codex-rs/ext/skills/tests/skills_extension/presentation_dedup_tests.rs b/codex-rs/ext/skills/tests/skills_extension/presentation_dedup_tests.rs index ae13ba6a8f..4d01b24bb7 100644 --- a/codex-rs/ext/skills/tests/skills_extension/presentation_dedup_tests.rs +++ b/codex-rs/ext/skills/tests/skills_extension/presentation_dedup_tests.rs @@ -109,7 +109,8 @@ async fn cloud_preference_preserves_aliases_reads_and_executor_fallback() -> Tes ("connected", true), ("disconnected", false), ] { - start_registered_turn(®istry, &session_store, &thread_store, stage).await; + refresh_registered_catalogs_for_step(®istry, &session_store, &thread_store, stage) + .await; let turn_store = ExtensionData::new(stage); let sections = prepared_world_state( ®istry, diff --git a/codex-rs/network-proxy/src/runtime.rs b/codex-rs/network-proxy/src/runtime.rs index 4dbd2cd24a..0288168d6b 100644 --- a/codex-rs/network-proxy/src/runtime.rs +++ b/codex-rs/network-proxy/src/runtime.rs @@ -1897,7 +1897,13 @@ mod tests { assert_eq!( state - .host_blocked("api.openai.com", /*port*/ 80) + .host_blocked_with_lookup( + "api.openai.com", + /*port*/ 80, + /*allow_local_binding*/ None, + HostAuthorization::RequireAllowlist, + |_host, port| async move { Ok(vec![SocketAddr::from(([8, 8, 8, 8], port))]) }, + ) .await .unwrap(), HostBlockDecision::Allowed @@ -1914,14 +1920,26 @@ mod tests { assert_eq!( state - .host_blocked("example.com", /*port*/ 80) + .host_blocked_with_lookup( + "example.com", + /*port*/ 80, + /*allow_local_binding*/ None, + HostAuthorization::RequireAllowlist, + |_host, port| async move { Ok(vec![SocketAddr::from(([8, 8, 8, 8], port))]) }, + ) .await .unwrap(), HostBlockDecision::Allowed ); assert_eq!( state - .host_blocked("api.openai.com", /*port*/ 443) + .host_blocked_with_lookup( + "api.openai.com", + /*port*/ 443, + /*allow_local_binding*/ None, + HostAuthorization::RequireAllowlist, + |_host, port| async move { Ok(vec![SocketAddr::from(([8, 8, 8, 8], port))]) }, + ) .await .unwrap(), HostBlockDecision::Allowed diff --git a/codex-rs/network-proxy/src/socks5.rs b/codex-rs/network-proxy/src/socks5.rs index a5cc5e7df6..2aeb04b5bd 100644 --- a/codex-rs/network-proxy/src/socks5.rs +++ b/codex-rs/network-proxy/src/socks5.rs @@ -1077,6 +1077,7 @@ mod tests { mode: NetworkMode::Full, mitm: true, credential_broker: true, + // This test exercises protocol detection, independently of live DNS resolution. allow_local_binding: Some(true), ..NetworkProxyConfig::default() };