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() };