diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 29d3afde00..3a5853ad9c 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2943,6 +2943,7 @@ dependencies = [ "codex-install-context", "codex-login", "codex-mcp", + "codex-mcp-extension", "codex-memories-read", "codex-model-provider", "codex-model-provider-info", @@ -4616,6 +4617,7 @@ dependencies = [ "tokio", "toml 0.9.11+spec-1.1.0", "tracing", + "tracing-subscriber", ] [[package]] diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 661bfdd43a..ff0bbb4f92 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -146,6 +146,7 @@ assert_matches = { workspace = true } codex-exec-server-test-support = { workspace = true } codex-image-generation-extension = { workspace = true } codex-home = { workspace = true } +codex-mcp-extension = { workspace = true } codex-otel = { workspace = true } codex-test-binary-support = { workspace = true } codex-utils-cargo-bin = { workspace = true } diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 9b92194db3..4796147c69 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -183,6 +183,44 @@ impl McpManager { .await } + /// Resolves only plugins from this step's ready roots, independent of shared MCP refresh. + pub(crate) async fn selected_plugins_for_step( + &self, + context: McpServerContributionContext<'_, Config>, + disabled_plugin_ids: &[String], + ) -> SelectedPluginSnapshot { + let roots = context + .ready_selected_capability_roots() + .unwrap_or_default(); + let mut selected = SelectedPluginSnapshot::default(); + if roots.is_empty() { + return selected; + } + for contributor in self.extensions.mcp_server_contributors() { + for SelectedPlugin { + selected_root_id, + plugin_id, + .. + } in contributor.selected_plugins(context).await + { + if !roots.iter().any(|root| root.id == selected_root_id) { + continue; + } + if !context.config().features.enabled(Feature::Plugins) + || disabled_plugin_ids.contains(&plugin_id) + { + selected.disabled_plugin_roots.push(selected_root_id); + } else { + selected.plugins.push(SelectedPluginIdentity { + selected_root_id: Some(selected_root_id), + plugin_id, + }); + } + } + } + selected + } + async fn runtime_config_with_context( &self, context: McpServerContributionContext<'_, Config>, diff --git a/codex-rs/core/src/session/active_plugin_inventory_tests.rs b/codex-rs/core/src/session/active_plugin_inventory_tests.rs index 0d69fd1675..918ad88271 100644 --- a/codex-rs/core/src/session/active_plugin_inventory_tests.rs +++ b/codex-rs/core/src/session/active_plugin_inventory_tests.rs @@ -83,12 +83,12 @@ async fn active_plugin_ids_select_and_validate_by_source() { }) .collect::>(); turn_context.active_host_plugin_identities = Some(identities); - turn_context.extension_data.insert(SelectedPluginSnapshot { + let selected_plugins = SelectedPluginSnapshot { plugins: selected_plugins, ..Default::default() - }); + }; assert_eq!( - turn_context.active_plugin_ids_for_telemetry(), + turn_context.active_plugin_ids_for_telemetry(Some(&selected_plugins)), expected.map(|ids| ids.into_iter().map(str::to_string).collect::>()), "{name}" ); @@ -116,11 +116,8 @@ async fn active_plugin_ids_enforce_bounds_after_selection_and_deduplication() { remote_plugin_id: remote.then_some(id.clone()), }]; turn_context.active_host_plugin_identities = Some(identities); - turn_context - .extension_data - .insert(SelectedPluginSnapshot::default()); assert_eq!( - turn_context.active_plugin_ids_for_telemetry(), + turn_context.active_plugin_ids_for_telemetry(/*selected*/ None), (length == 128).then_some(vec![id]), "length={length}, remote={remote}" ); @@ -131,12 +128,12 @@ async fn active_plugin_ids_enforce_bounds_after_selection_and_deduplication() { plugin_id: id.clone(), }; turn_context.active_host_plugin_identities = Some(Vec::new()); - turn_context.extension_data.insert(SelectedPluginSnapshot { + let selected_plugins = SelectedPluginSnapshot { plugins: vec![selected_plugin], ..Default::default() - }); + }; assert_eq!( - turn_context.active_plugin_ids_for_telemetry(), + turn_context.active_plugin_ids_for_telemetry(Some(&selected_plugins)), (length == 128).then_some(vec![id]), "selected length={length}" ); @@ -161,12 +158,12 @@ async fn active_plugin_ids_enforce_bounds_after_selection_and_deduplication() { }) .collect::>(); turn_context.active_host_plugin_identities = Some(identities); - turn_context.extension_data.insert(SelectedPluginSnapshot { + let selected_plugins = SelectedPluginSnapshot { plugins: selected_plugins, ..Default::default() - }); + }; assert_eq!( - turn_context.active_plugin_ids_for_telemetry(), + turn_context.active_plugin_ids_for_telemetry(Some(&selected_plugins)), (count == 512).then_some(ids), "{count} distinct IDs across host and selected inputs, with one overlap" ); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 9caf797d45..613bdcaf15 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3785,18 +3785,25 @@ impl Session { )), turn::prepare_tool_recommendations(self.as_ref(), turn_context.as_ref()), ); - let mut selected_plugins = self + // A step keeps the plugins from the environments it captured, even if shared MCP + // moves on to another environment. Its skill tools and the model use this same copy. + let selected_plugins = self .services - .thread_extension_data - .get::() - .map(|snapshot| snapshot.as_ref().clone()) - .unwrap_or_default(); - selected_plugins.plugins.retain(|plugin| { - ready_selected_capability_roots - .iter() - .any(|root| plugin.selected_root_id.as_ref() == Some(&root.id)) - }); - extension_data.insert(selected_plugins.clone()); + .mcp_manager + .selected_plugins_for_step( + codex_extension_api::McpServerContributionContext::for_step( + turn_context.config.as_ref(), + &self.services.mcp_thread_init, + &self.services.thread_extension_data, + &turn_context.originator, + &ready_selected_capability_roots, + executor_capability_discovery.as_deref(), + ) + .with_session_source(&turn_context.session_source), + &turn_context.disabled_plugin_ids, + ) + .await; + extension_data.insert(selected_plugins); let tool_router = turn::built_tools( self.as_ref(), turn_context.as_ref(), @@ -3812,7 +3819,7 @@ impl Session { executor_capability_discovery, mcp, tool_router, - selected_plugins, + extension_data, )) }); // Returned warnings must finish delivery even if tools fail or preparation is cancelled. @@ -3827,9 +3834,8 @@ impl Session { executor_capability_discovery, mcp, tool_router, - selected_plugins, + extension_data, ) = prepared_tools??; - turn_context.extension_data.insert(selected_plugins); Ok(Arc::new(StepContext { preempt: turn_context .config @@ -3844,6 +3850,7 @@ impl Session { environments, selected_capability_roots, executor_capability_discovery, + extension_data, mcp, tool_router, loaded_agents_md, diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs index 9f574254b0..2b2428ccf2 100644 --- a/codex-rs/core/src/session/step_context.rs +++ b/codex-rs/core/src/session/step_context.rs @@ -12,6 +12,7 @@ use crate::session::turn_context::TurnEnvironment; use crate::tools::router::ToolRouter; use codex_exec_server::ExecutorCapabilityDiscoverySnapshot; use codex_exec_server::ResolvedSelectedCapabilityRoot; +use codex_extension_api::ExtensionData; use codex_file_system::EnvironmentAccess; use codex_mcp::McpBinding; use codex_otel::SessionTelemetry; @@ -37,6 +38,8 @@ pub(crate) struct StepContext { pub(crate) selected_capability_roots: Vec, /// Executor-materialized capability files shared by MCP and skills in this exact step. pub(crate) executor_capability_discovery: Option>, + /// Keeps the extension inputs used to build this step's tools for its World State as well. + pub(crate) extension_data: ExtensionData, /// The exact MCP connections, configuration, and catalog captured for this step. pub(crate) mcp: Arc, /// The finalized tool plan advertised and executed for this exact sampling request. diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 4496c9b940..c001365742 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -291,6 +291,7 @@ impl StepContext { environments, selected_capability_roots: Vec::new(), executor_capability_discovery: None, + extension_data: codex_extension_api::ExtensionData::new(turn.sub_id.clone()), mcp: Arc::new(codex_mcp::McpBinding::empty(mcp_config_for_test( &turn.config, ))), @@ -9426,12 +9427,6 @@ async fn cancelled_step_capture_finishes_warning_delivery() { ), (INITIAL_SUBMIT_ID.to_owned(), warnings[1].clone()), ); - assert!( - turn.extension_data - .get::() - .is_none() - ); - session .capture_step_context(turn, &CancellationToken::new()) .await diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index f13e98ccb0..03aa0f197e 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -404,7 +404,7 @@ pub(crate) async fn run_turn( .await; } - track_turn_resolved_config_analytics(&sess, &turn_context, &input).await; + track_turn_resolved_config_analytics(&sess, &first_step_context, &input).await; let mut last_agent_message: Option = None; let mut stop_hook_active = false; @@ -1236,9 +1236,10 @@ async fn build_extension_turn_input_items( )] async fn track_turn_resolved_config_analytics( sess: &Session, - turn_context: &TurnContext, + first_step: &StepContext, input: &[TurnInput], ) { + let turn_context = &first_step.turn; let thread_config = sess.thread_config_snapshot().await; let is_first_turn = { let mut state = sess.state.lock().await; @@ -1250,7 +1251,12 @@ async fn track_turn_resolved_config_analytics( turn_id: turn_context.sub_id.clone(), thread_id: sess.thread_id.to_string(), turn_metadata: turn_context.turn_metadata_state.clone(), - active_plugin_ids_at_turn_start: turn_context.active_plugin_ids_for_telemetry(), + active_plugin_ids_at_turn_start: turn_context.active_plugin_ids_for_telemetry( + first_step + .extension_data + .get::() + .as_deref(), + ), num_input_images: input .iter() .filter_map(|item| match item { diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 4db28491b2..e8f9abb256 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -571,15 +571,18 @@ impl TurnContext { self.initial_settings.effective_collaboration_mode() } - /// Combines setup-time host identities with current ready selected packages. + /// Combines setup-time host identities with the first step's ready selected packages. /// Keeps the complete observation or marks it unknown; never truncates membership. - pub(super) fn active_plugin_ids_for_telemetry(&self) -> Option> { + pub(super) fn active_plugin_ids_for_telemetry( + &self, + selected: Option<&SelectedPluginSnapshot>, + ) -> Option> { const MAX_TELEMETRY_PLUGIN_IDS: usize = 512; const MAX_TELEMETRY_PLUGIN_ID_BYTES: usize = 128; let mut identities = self.active_host_plugin_identities.clone()?; // Selected roots provide a package key, not a remote identity for that object. - if let Some(selected) = self.extension_data.get::() { + if let Some(selected) = selected { identities.extend(selected.plugins.iter().map(|plugin| PluginIdentity { plugin_id: plugin.plugin_id.clone(), remote_plugin_id: None, diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index 1dc5b0f219..cb32f51cc5 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/codex-rs/core/src/session/world_state.rs @@ -311,6 +311,7 @@ impl Session { session_store: &self.services.session_extension_data, thread_store: &self.services.thread_extension_data, turn_store: turn_context.extension_data.as_ref(), + step_store: &step_context.extension_data, previous_world_state: previous_world_state.as_ref().map(|state| &state.state), }) .await diff --git a/codex-rs/core/tests/suite/agents_md_refresh.rs b/codex-rs/core/tests/suite/agents_md_refresh.rs index 807ccecaa9..d3eb8efa85 100644 --- a/codex-rs/core/tests/suite/agents_md_refresh.rs +++ b/codex-rs/core/tests/suite/agents_md_refresh.rs @@ -2,14 +2,8 @@ use super::*; use codex_extension_api::ExtensionData; -use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; -use codex_extension_api::SelectedPluginSnapshot; use codex_extension_api::ToolContributor; -use codex_extension_api::TurnErrorInput; -use codex_extension_api::TurnLifecycleContributor; -use codex_extension_api::TurnStartInput; -use codex_extension_api::TurnStopInput; use codex_protocol::dynamic_tools::DynamicToolFunctionSpec; use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::request_user_input::RequestUserInputAnswer; @@ -25,16 +19,6 @@ struct StepPreparationObserver { instructions: RecordingThreadInstructionsProvider, tools_ready: Mutex>>, waiting_for_tools: Mutex>>, - published_plugins: Mutex>, -} - -impl StepPreparationObserver { - fn record_plugins(&self, phase: &'static str, turn_store: &ExtensionData) { - self.published_plugins - .lock() - .expect("plugin observations") - .push((phase, turn_store.get::().is_some())); - } } impl ThreadInstructionsProvider for StepPreparationObserver { @@ -68,20 +52,6 @@ impl ToolContributor for StepPreparationObserver { } } -impl TurnLifecycleContributor for StepPreparationObserver { - fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> { - Box::pin(async move { self.record_plugins("start", input.turn_store) }) - } - - fn on_turn_error<'a>(&'a self, input: TurnErrorInput<'a>) -> ExtensionFuture<'a, ()> { - Box::pin(async move { self.record_plugins("error", input.turn_store) }) - } - - fn on_turn_stop<'a>(&'a self, input: TurnStopInput<'a>) -> ExtensionFuture<'a, ()> { - Box::pin(async move { self.record_plugins("stop", input.turn_store) }) - } -} - #[derive(Clone, Copy)] enum PreparationOutcome { Success, @@ -90,12 +60,12 @@ enum PreparationOutcome { BothErrors, } -#[test_case::test_case(PreparationOutcome::Success; "publish after both branches succeed")] +#[test_case::test_case(PreparationOutcome::Success; "send request after both branches succeed")] #[test_case::test_case(PreparationOutcome::ToolCollision; "warning survives tool failure")] -#[test_case::test_case(PreparationOutcome::InstructionError; "instruction error prevents publication")] +#[test_case::test_case(PreparationOutcome::InstructionError; "instruction error prevents request")] #[test_case::test_case(PreparationOutcome::BothErrors; "instruction error takes precedence")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn concurrent_preparation_preserves_warnings_errors_and_plugin_publication( +async fn concurrent_preparation_preserves_warnings_and_errors( outcome: PreparationOutcome, ) -> Result<()> { let server = start_mock_server().await; @@ -104,11 +74,9 @@ async fn concurrent_preparation_preserves_warnings_errors_and_plugin_publication instructions: RecordingThreadInstructionsProvider::with_text(TASK_USER_INSTRUCTIONS), tools_ready: Mutex::default(), waiting_for_tools: Mutex::default(), - published_plugins: Mutex::default(), }); let mut extensions = ExtensionRegistryBuilder::new(); extensions.tool_contributor(observer.clone()); - extensions.turn_lifecycle_contributor(observer.clone()); let mut builder = test_codex() .with_extensions(Arc::new(extensions.build())) .with_config(|config| { @@ -174,30 +142,19 @@ async fn concurrent_preparation_preserves_warnings_errors_and_plugin_publication }) .await; let mut expected_notices = vec![("warning", PROVIDER_WARNING.to_string())]; - let mut expected_publication = vec![("start", false)]; match outcome { PreparationOutcome::Success => {} PreparationOutcome::ToolCollision => { expected_notices.push(("error", "duplicate tool: functions.update_plan".to_string())); - expected_publication.push(("error", false)); } PreparationOutcome::InstructionError | PreparationOutcome::BothErrors => { expected_notices.push(( "error", "thread instructions exceed the limit of 10000 estimated tokens (10001 estimated tokens provided)".to_string(), )); - expected_publication.push(("error", false)); } } - expected_publication.push(("stop", succeeds)); assert_eq!(notices, expected_notices); - assert_eq!( - *observer - .published_plugins - .lock() - .expect("plugin observations"), - expected_publication - ); assert_eq!(request.requests().len(), usize::from(succeeds)); Ok(()) } diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index 6c143ee226..0ef0ee3013 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/codex-rs/core/tests/suite/scenarios.rs @@ -676,12 +676,42 @@ async fn astra_kickoff_with_skills_plugins_and_remote_compaction() -> Result<()> } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn astra_omits_disabled_executor_skills_from_model_context() -> Result<()> { +async fn astra_omits_disabled_executor_and_plugin_skills_from_model_context() -> Result<()> { + use codex_extension_api::ContextContributor; + use codex_extension_api::ExtensionFuture; + use codex_extension_api::SelectedPluginSnapshot; + use codex_extension_api::WorldStateContributionInput; + use codex_extension_api::WorldStateSectionContribution; + skip_if_no_network!(Ok(())); + struct PauseFirstWorldState { + first: std::sync::Mutex>>, + resume: tokio::sync::Notify, + } + + impl ContextContributor for PauseFirstWorldState { + fn contribute_world_state<'a>( + &'a self, + input: WorldStateContributionInput<'a>, + ) -> ExtensionFuture<'a, Vec> { + Box::pin(async move { + let first = self.first.lock().unwrap().take(); + if let Some(first) = first { + first.send(input.turn_id.to_string()).unwrap(); + self.resume.notified().await; + } + Vec::new() + }) + } + } + let server = start_mock_server().await; - let skill_files = TempDir::new()?; - let skill_root = fs::canonicalize(skill_files.path())?.join("skills"); + let skill_files_dir = TempDir::new()?; + let skill_files = fs::canonicalize(skill_files_dir.path())?; + let skill_root = skill_files.join("skills"); + let other_skill_root = skill_files.join("other-skills"); + let plugin_root = skill_files.join("plugin"); let disabled = write_skill( &skill_root.join("retired-helper"), "retired-helper", @@ -694,8 +724,26 @@ async fn astra_omits_disabled_executor_skills_from_model_context() -> Result<()> "Use the current workflow", "Follow the current workflow.", )?; + write_skill( + &other_skill_root.join("next-helper"), + "next-helper", + "Use the next workflow", + "Follow the next workflow.", + )?; + fs::create_dir_all(plugin_root.join(".codex-plugin"))?; + fs::write( + plugin_root.join(".codex-plugin/plugin.json"), + r#"{"name":"disabled"}"#, + )?; + write_skill( + &plugin_root.join("skills/disabled-plugin-helper"), + "disabled-plugin-helper", + "Use a disabled plugin workflow", + "Follow the disabled plugin workflow.", + )?; + let environment_manager = Arc::new(EnvironmentManager::default_for_tests()); let provider = ExecutorSkillProvider::new_with_restriction_product( - Arc::new(EnvironmentManager::default_for_tests()), + Arc::clone(&environment_manager), /*restriction_product*/ None, ) .with_disabled_skill_paths(HashMap::from([( @@ -703,6 +751,13 @@ async fn astra_omits_disabled_executor_skills_from_model_context() -> Result<()> HashSet::from([PathUri::from_host_native_path(disabled)?]), )])); let mut extensions = ExtensionRegistryBuilder::::new(); + codex_mcp_extension::install_plugins(&mut extensions, environment_manager); + let (started, captured) = oneshot::channel(); + let pause = Arc::new(PauseFirstWorldState { + first: std::sync::Mutex::new(Some(started)), + resume: tokio::sync::Notify::new(), + }); + extensions.prompt_contributor(pause.clone()); install_with_providers( &mut extensions, SkillProviders::new().with_executor_provider(Arc::new(provider)), @@ -716,10 +771,21 @@ async fn astra_omits_disabled_executor_skills_from_model_context() -> Result<()> ); let mock = mount_sse_sequence( &server, - vec![sse(vec![ - ev_assistant_message("skills", "The active-helper skill is available."), - ev_completed("skills-response"), - ])], + vec![ + sse(vec![ + ev_function_call_with_namespace( + "list-skills", + "skills", + "list", + r#"{"authority":{"kind":"executor"}}"#, + ), + ev_completed("first-step"), + ]), + sse(vec![ + ev_assistant_message("skills", "The next-helper skill is available."), + ev_completed("skills-response"), + ]), + ], ) .await; let test = test_codex() @@ -730,25 +796,38 @@ async fn astra_omits_disabled_executor_skills_from_model_context() -> Result<()> .build(&server) .await?; let skill_root = PathUri::from_host_native_path(skill_root)?; - let root_locator = format!( - "skill://workspace-skills/{}", - skill_root - .inferred_native_path_string() - .replace('\\', "/") - .trim_start_matches('/') - ); - let mut thread_extension_init = ExtensionDataInit::new(); - thread_extension_init.insert(vec![SelectedCapabilityRoot { - id: "workspace-skills".to_string(), + let other_skill_root = PathUri::from_host_native_path(other_skill_root)?; + let skill_files = PathUri::from_host_native_path(skill_files)? + .inferred_native_path_string() + .replace('\\', "/"); + let skill_files = format!("{}/", skill_files.trim_start_matches('/')); + let selected_root = |id: &str, path| SelectedCapabilityRoot { + id: id.to_string(), location: CapabilityRootLocation::Environment { environment_id: LOCAL_ENVIRONMENT_ID.to_string(), - path: skill_root, + path, }, - }]); + }; + let mut thread_extension_init = ExtensionDataInit::new(); + thread_extension_init.insert(vec![selected_root( + "plugin-skills", + PathUri::from_host_native_path(plugin_root)?, + )]); + let mut first_environment = test.executor_environment().selection().clone(); + let mut environment_config = environment_config_for_selection(&test.config, &first_environment); + environment_config.selected_capability_roots = + vec![selected_root("workspace-skills", skill_root)]; + first_environment.config = EnvironmentConfigState::Ready(environment_config.clone()); + let mut next_environment = first_environment.clone(); + environment_config.selected_capability_roots = + vec![selected_root("other-skills", other_skill_root)]; + next_environment.config = EnvironmentConfigState::Ready(environment_config); let thread = test .thread_manager .start_thread(StartThreadOptions { + environments: Some(vec![first_environment]), thread_extension_init, + disabled_plugin_ids: Some(vec!["plugin-skills".to_string()]), ..StartThreadOptions::new(test.config.clone()) }) .await? @@ -758,23 +837,74 @@ async fn astra_omits_disabled_executor_skills_from_model_context() -> Result<()> "Which workflow skills are available?", )])) .await?; + let turn_id = tokio::time::timeout(Duration::from_secs(/*secs*/ 10), captured).await??; + let (reply, outcome) = oneshot::channel(); + thread + .submit(Op::TurnSettings { + turn_id, + update: TurnSettingsUpdate { + environments: Some(vec![next_environment]), + ..Default::default() + }, + reply, + }) + .await?; + assert_eq!(outcome.await?, TurnSettingsUpdateOutcome::Applied); + // Refresh shared MCP for the new roots (there is no Apps server in this test), then + // replace its plugin list before the next step captures it or needs another refresh. + let _ = thread.refresh_codex_apps_tools().await; + let shared = thread.thread_extension_data(); + assert_eq!( + shared + .insert(SelectedPluginSnapshot::default()) + .expect("MCP should have published the new plugin list") + .disabled_plugin_roots, + ["plugin-skills"], + ); + pause.resume.notify_one(); wait_for_event(&thread, |event| matches!(event, EventMsg::TurnComplete(_))).await; let requests = mock.requests(); - assert_eq!(requests.len(), 1); - let mut body = requests[0].body_json(); - let input = body["input"].to_string(); - assert!(input.contains("active-helper")); - assert!(!input.contains("retired-helper")); - // Normalize opaque skill locators before snapshot truncation and hashing. - body["input"] = serde_json::from_str( - &input.replace(&root_locator, "skill://workspace-skills/"), - )?; + assert_eq!(requests.len(), 2); + let first_step_tool = requests[1] + .function_call_output_text("list-skills") + .expect("first step's skills.list output"); + assert!(first_step_tool.contains("active-helper")); + assert!(!first_step_tool.contains("next-helper")); + let developer_texts = requests[1].message_input_texts("developer"); + let latest_skills = developer_texts + .iter() + .rev() + .find(|text| text.contains("## Skills")) + .expect("second step's skills section"); + assert!(latest_skills.contains("next-helper")); + assert!(!latest_skills.contains("active-helper")); + let mut bodies = requests + .iter() + .map(core_test_support::responses::ResponsesRequest::body_json) + .collect::>(); + for body in &mut bodies { + let input = body["input"].to_string(); + assert!(!input.contains("retired-helper")); + assert!(!input.contains("disabled-plugin-helper")); + // Normalize opaque skill locators before snapshot truncation and hashing. + body["input"] = serde_json::from_str(&input.replace(&skill_files, "/"))?; + // Cargo and Bazel can serialize the JSON inside the tool output in different key orders. + for item in body["input"].as_array_mut().expect("request input") { + if item["type"] == "function_call_output" && item["call_id"] == "list-skills" { + let mut output: serde_json::Value = + serde_json::from_str(item["output"].as_str().expect("skills.list output"))?; + output.sort_all_objects(); + item["output"] = output.to_string().into(); + } + } + } + let entries = bodies.iter().map(SnapshotEntry::body).collect::>(); insta::assert_snapshot!( "astra_disabled_executor_skills", context_snapshot::format_context_snapshot( - "Astra sees the active executor skill while the caller-disabled skill is omitted.", - &[SnapshotEntry::body(&body)], + "Astra keeps disabled skills hidden and uses each step's selected skills after the environment changes during preparation.", + &entries, &ContextSnapshotOptions::default().include_request_settings(), ) ); diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_disabled_executor_skills.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_disabled_executor_skills.snap index ef38825370..a5546e3f02 100644 --- a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_disabled_executor_skills.snap +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_disabled_executor_skills.snap @@ -1,8 +1,8 @@ --- source: core/tests/suite/scenarios.rs -expression: "context_snapshot::format_context_snapshot(\"Astra sees the active executor skill while the caller-disabled skill is omitted.\",\n&[SnapshotEntry::body(&body)],\n&ContextSnapshotOptions::default().include_request_settings(),)" +expression: "context_snapshot::format_context_snapshot(\"Astra keeps disabled skills hidden and uses each step's selected skills after the environment changes during preparation.\",\n&entries, &ContextSnapshotOptions::default().include_request_settings(),)" --- -Scenario: Astra sees the active executor skill while the caller-disabled skill is omitted. +Scenario: Astra keeps disabled skills hidden and uses each step's selected skills after the environment changes during preparation. ## Window 1 Settings: @@ -87,7 +87,7 @@ Settings: ## Skills A skill is a set of instructions provided through a `SKILL.md` source. Below is the list of skills that ...locators can be expanded using the skill roots table. [hash=f00d6a4d3e192f06] ### Skill roots - - `e0` = `skill://workspace-skills/` + - `e0` = `skill://workspace-skills//skills` Read a skill package directly with `skills.read({"package":""})` to read its `SKILL.md`; root a...ackage is not provided, use `skills.list` to find it. [hash=25bf84ea2123d647] ### Available skills - active-helper: Use the current workflow (executor package: e0/active-helper) @@ -118,7 +118,20 @@ Settings: - [hash=911722cc43b4c3b8] + [hash=e8c1de60307d0a3d] 06:message/user: Which workflow skills are available? +-- request 2 (request) -- +07:function_call/skills.list:{"authority":{"kind":"executor"}} +08:function_call_output:{"next_cursor":null,"skills":[{"authority":{"id":"workspace-skills","kind":"executor"},"description":"Us...//skills/active-helper"}],"warnings":[]} [hash=b72a3cc66e4d3c51] +09:message/developer: + + ## Skills + A skill is a set of instructions provided through a `SKILL.md` source. Below is the list of skills that ...locators can be expanded using the skill roots table. [hash=f00d6a4d3e192f06] + ### Skill roots + - `e0` = `skill://other-skills//other-skills` + Read a skill package directly with `skills.read({"package":""})` to read its `SKILL.md`; root a...ackage is not provided, use `skills.list` to find it. [hash=25bf84ea2123d647] + ### Available skills + - next-helper: Use the next workflow (executor package: e0/next-helper) + diff --git a/codex-rs/ext/extension-api/src/contributors/world_state.rs b/codex-rs/ext/extension-api/src/contributors/world_state.rs index e17873b86c..f5571bb8d8 100644 --- a/codex-rs/ext/extension-api/src/contributors/world_state.rs +++ b/codex-rs/ext/extension-api/src/contributors/world_state.rs @@ -26,6 +26,8 @@ pub struct WorldStateContributionInput<'a> { pub session_store: &'a ExtensionData, pub thread_store: &'a ExtensionData, pub turn_store: &'a ExtensionData, + /// The same captured extension inputs that were used to build this step's tools. + pub step_store: &'a ExtensionData, /// Persisted comparison state from the last recorded step. Its rendered text may have /// left model history; use a retained-fragment matcher before referring back to it. /// After compaction, sections may contain only retained extension metadata. diff --git a/codex-rs/ext/skills/Cargo.toml b/codex-rs/ext/skills/Cargo.toml index cb5cb8c30a..80bcfba8cb 100644 --- a/codex-rs/ext/skills/Cargo.toml +++ b/codex-rs/ext/skills/Cargo.toml @@ -47,3 +47,4 @@ opentelemetry_sdk = { workspace = true } pretty_assertions = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } +tracing-subscriber = { workspace = true } diff --git a/codex-rs/ext/skills/src/skills_extension_state.rs b/codex-rs/ext/skills/src/skills_extension_state.rs index dac4388899..f89246a218 100644 --- a/codex-rs/ext/skills/src/skills_extension_state.rs +++ b/codex-rs/ext/skills/src/skills_extension_state.rs @@ -23,20 +23,11 @@ const MAX_CACHED_CLOUD_CONTENT_BYTES: usize = 8 * 1024 * 1024; /// Stores published cloud and executor catalogs for a thread. #[derive(Clone, Debug, Default)] pub(crate) struct SkillsExtensionState { - pub(crate) executor_catalog_selection: ExecutorCatalogSelection, pub(crate) executor_cache: Vec, pub(crate) executor_discovery_cache: Option, pub(crate) cloud_cache: Option>, } -/// Selects the existing cache entries published by the most recent refresh. -#[derive(Clone, Debug, Default)] -pub(crate) enum ExecutorCatalogSelection { - #[default] - Discovery, - Roots(Vec), -} - /// Keeps cloud metadata and contents within one published auth scope and Apps availability state. /// Warning-free discovery is reused until the Apps resource generation changes. /// A successful refresh replaces the Arc so late reads cannot populate its successor. @@ -79,7 +70,7 @@ pub(crate) struct CachedExecutorCatalog { pub(crate) catalog: SkillCatalog, } -/// Retains the latest projection for snapshots; refresh only reuses successful unchanged inputs. +/// Keeps the latest result; only successful results for unchanged inputs can be reused. #[derive(Clone, Debug)] pub(crate) struct CachedExecutorDiscoveryCatalog { pub(crate) roots: Vec, diff --git a/codex-rs/ext/skills/src/state.rs b/codex-rs/ext/skills/src/state.rs index d7d0ac56c0..e0ba256656 100644 --- a/codex-rs/ext/skills/src/state.rs +++ b/codex-rs/ext/skills/src/state.rs @@ -27,7 +27,6 @@ use crate::skills_extension_state::CachedExecutorCatalog; use crate::skills_extension_state::CachedExecutorDiscoveryCatalog; use crate::skills_extension_state::CloudResourceCache; use crate::skills_extension_state::CloudSkillGeneration; -use crate::skills_extension_state::ExecutorCatalogSelection; use crate::skills_extension_state::SkillReadCacheKey; use crate::sources::SkillProviders; @@ -117,7 +116,7 @@ impl SkillsThreadState { .map(|turn| Arc::clone(&turn.state)) } - /// Refreshes the current step's executor catalog in the existing caches. + /// Returns the catalog for these exact roots, using the existing thread caches. #[tracing::instrument( name = "skills.executor.refresh_executor_catalog", level = "info", @@ -128,57 +127,27 @@ impl SkillsThreadState { &self, providers: &SkillProviders, mut query: SkillListQuery, - ) { - let selection = if query.executor_capability_discovery.is_some() { + ) -> SkillCatalog { + if query.executor_capability_discovery.is_some() { // High-level discovery is enabled: reuse or project its discovery snapshot. self.executor_discovery_catalog_snapshot(providers, query) - .await; - ExecutorCatalogSelection::Discovery + .await } else { // High-level discovery is not enabled: retain the legacy per-root cache path. let roots = std::mem::take(&mut query.executor_roots); - for root in &roots { + let mut catalog = SkillCatalog::default(); + for root in roots { query.executor_roots = vec![root.clone()]; - self.executor_root_catalog(providers, root.clone(), query.clone()) - .await; - } - ExecutorCatalogSelection::Roots(roots) - }; - self.skills_extension_state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .executor_catalog_selection = selection; - } - - /// Reads the most recently refreshed selection without invoking providers. - pub fn executor_catalog_snapshot(&self) -> SkillCatalog { - let state = self - .skills_extension_state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - match &state.executor_catalog_selection { - ExecutorCatalogSelection::Discovery => state - .executor_discovery_cache - .as_ref() - .map(|cached| cached.catalog.clone()) - .unwrap_or_default(), - ExecutorCatalogSelection::Roots(roots) => { - let mut catalog = SkillCatalog::default(); - for root in roots { - if let Some(cached) = state - .executor_cache - .iter() - .find(|cached| &cached.root == root) - { - catalog.extend(cached.catalog.clone()); - } - } - catalog + catalog.extend( + self.executor_root_catalog(providers, root, query.clone()) + .await, + ); } + catalog } } - /// Reuses matching successful discovery or retains a fresh result for snapshot reads. + /// Reuses matching successful discovery; otherwise returns the result of this request. async fn executor_discovery_catalog_snapshot( &self, providers: &SkillProviders, @@ -218,7 +187,7 @@ impl SkillsThreadState { .skills_extension_state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - // Retain failures for read-only snapshots; discovery_failed above prevents reuse. + // Failed discovery is returned to this caller but is never reused by another request. state.executor_discovery_cache = discovery.map(|discovery| CachedExecutorDiscoveryCatalog { roots, diff --git a/codex-rs/ext/skills/src/world_state_catalogs.rs b/codex-rs/ext/skills/src/world_state_catalogs.rs index d0b76d38fa..944f3b4467 100644 --- a/codex-rs/ext/skills/src/world_state_catalogs.rs +++ b/codex-rs/ext/skills/src/world_state_catalogs.rs @@ -184,11 +184,11 @@ impl<'a> CatalogContext<'a> { } async fn discover_executor_catalog(&self, query: SkillListQuery) -> CatalogContribution { - self.thread_state + let mut catalog = self + .thread_state .refresh_executor_catalog(self.providers, query) .await; - let mut catalog = self.thread_state.executor_catalog_snapshot(); - if let Some(selected_plugins) = self.input.turn_store.get::() { + if let Some(selected_plugins) = self.input.step_store.get::() { attribute_executor_plugins(&mut catalog, &selected_plugins); } self.input diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index d4a5be14e6..bf64aba2a3 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/codex-rs/ext/skills/tests/skills_extension.rs @@ -144,6 +144,7 @@ async fn skill_world_state_fragments( session_store, thread_store, turn_store: &turn_store, + step_store: &turn_store, }) .await; @@ -331,6 +332,7 @@ async fn host_world_state_records_catalog_metrics_on_publish_and_change() -> Tes session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &turn_store, }) .await; @@ -357,6 +359,7 @@ async fn host_world_state_records_catalog_metrics_on_publish_and_change() -> Tes session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &turn_store, }) .await; assert!( @@ -396,6 +399,7 @@ async fn host_world_state_records_catalog_metrics_on_publish_and_change() -> Tes session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &turn_store, }) .await; assert!( @@ -469,6 +473,7 @@ async fn persisted_host_snapshot_deduplicates_warning_after_reinitialization() - session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &turn_store, }) .await; let host_section = world_state_section(§ions, "host_skills"); @@ -509,6 +514,7 @@ async fn persisted_host_snapshot_deduplicates_warning_after_reinitialization() - session_store: &session_store, thread_store: &resumed_thread_store, turn_store: &resumed_turn_store, + step_store: &resumed_turn_store, }) .await; assert!( @@ -601,6 +607,7 @@ async fn executor_cloud_and_host_share_catalog_world_state_flow() -> TestResult session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &turn_store, }) .await; @@ -711,6 +718,7 @@ async fn nonempty_executor_empty_host_records_catalog_metrics() -> TestResult { session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &turn_store, }) .await; @@ -792,6 +800,7 @@ async fn host_world_state_uses_provider_catalog_with_core_compatible_rendering() session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &turn_store, }) .await; let host_fragment = world_state_section(§ions, "host_skills") @@ -874,6 +883,7 @@ async fn shadow_selection_uses_host_catalog_when_instructions_are_disabled() -> session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &turn_store, }) .await; let fragments = registry.turn_input_contributors()[0] @@ -1166,6 +1176,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &turn_store, }) .await; assert_eq!(1, available_sections.len()); @@ -1230,6 +1241,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach session_store: &session_store, thread_store: &thread_store, turn_store: &unavailable_turn_store, + step_store: &unavailable_turn_store, }) .await; let (unavailable_snapshot, unavailable_fragment) = @@ -1268,6 +1280,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach session_store: &session_store, thread_store: &thread_store, turn_store: &restored_turn_store, + step_store: &restored_turn_store, }) .await; let (restored_snapshot, restored_fragment) = @@ -1326,6 +1339,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach session_store: &session_store, thread_store: &thread_store, turn_store: &ExtensionData::new(turn_id), + step_store: &ExtensionData::new(turn_id), }) .await; assert_eq!(expected_list_calls, list_calls.load(Ordering::Relaxed)); @@ -1353,6 +1367,7 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach session_store: &session_store, thread_store: &thread_store, turn_store: &listing_disabled_turn_store, + step_store: &listing_disabled_turn_store, }) .await; let (listing_disabled_snapshot, listing_disabled_fragment) = listing_disabled_sections[0] @@ -1380,6 +1395,133 @@ async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cach Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn overlapping_executor_catalogs_render_their_own_roots() -> TestResult { + use tracing::Subscriber; + use tracing::instrument::WithSubscriber; + use tracing_subscriber::Layer; + use tracing_subscriber::layer::Context; + use tracing_subscriber::layer::SubscriberExt; + use tracing_subscriber::registry::LookupSpan; + + // Pause after A's refresh returns, before World State can read the catalog. Pausing the + // provider would be too early: A would still publish its result after B finishes. + struct PauseAfterRefresh { + entered: Mutex>>, + resume: Mutex>, + } + + impl LookupSpan<'a>> Layer for PauseAfterRefresh { + fn on_close(&self, id: tracing::span::Id, context: Context<'_, S>) { + if context + .span(&id) + .is_some_and(|span| span.name() == "skills.executor.refresh_executor_catalog") + && let Some(entered) = self.entered.lock().unwrap().take() + { + entered.send(()).unwrap(); + self.resume + .lock() + .unwrap() + .recv_timeout(std::time::Duration::from_secs(/*secs*/ 10)) + .expect("B should finish before A reads its catalog"); + } + } + } + + let temp = tempfile::tempdir()?; + let root = |name: &str| -> Result> { + let path = temp.path().join(name); + std::fs::create_dir_all(path.join(name))?; + std::fs::write( + path.join(name).join("SKILL.md"), + DEMO_SKILL_CONTENTS.replace("demo", name), + )?; + Ok(SelectedCapabilityRoot { + id: name.to_string(), + location: CapabilityRootLocation::Environment { + environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(), + path: PathUri::from_host_native_path(path)?, + }, + }) + }; + let root_a = root("skill-a")?; + let root_b = root("skill-b")?; + let mut builder = ExtensionRegistryBuilder::new(); + let provider = codex_skills_extension::ExecutorSkillProvider::new_with_restriction_product( + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + /*restriction_product*/ None, + ); + install_with_providers( + &mut builder, + SkillProviders::new().with_executor_provider(Arc::new(provider)), + skills_extension_config, + ); + let registry = Arc::new(builder.build()); + let session_store = Arc::new(ExtensionData::new("session")); + let thread_store = Arc::new(ExtensionData::new("thread")); + registry.thread_lifecycle_contributors()[0] + .on_thread_start(ThreadStartInput { + config: &default_config(), + session_source: &SessionSource::Cli, + persistent_thread_state_available: true, + environments: &[], + mcp_resource_client: None, + extension_metrics: None, + session_store: &session_store, + thread_store: &thread_store, + }) + .await; + let thread_id = codex_protocol::ThreadId::new(); + let render = |root: SelectedCapabilityRoot| { + let registry = Arc::clone(®istry); + let session_store = Arc::clone(&session_store); + let thread_store = Arc::clone(&thread_store); + async move { + let step_store = ExtensionData::new(&root.id); + let sections = registry.context_contributors()[0] + .contribute_world_state(WorldStateContributionInput { + model_info: &catalog_model_info(), + thread_id, + turn_id: "turn", + environments: &[], + ready_selected_capability_roots: &[root], + executor_capability_discovery: None, + extension_metrics: None, + session_store: &session_store, + thread_store: &thread_store, + turn_store: &step_store, + step_store: &step_store, + previous_world_state: None, + }) + .await; + world_state_section(§ions, "skills") + .render_diff(PreviousWorldStateSection::Absent) + .1 + .expect("executor skills should render") + .body() + .to_string() + } + }; + let (entered, paused) = tokio::sync::oneshot::channel(); + let (release, resume) = std::sync::mpsc::channel(); + // Keep tracing aware of uninstrumented threads so a parallel test cannot register + // the refresh callsite as globally disabled before this subscriber reaches it. + let _interest_cache_guard = + tracing::Dispatch::new(tracing::subscriber::NoSubscriber::default()); + let subscriber = tracing_subscriber::registry().with(PauseAfterRefresh { + entered: Mutex::new(Some(entered)), + resume: Mutex::new(resume), + }); + let a = tokio::spawn(render(root_a).with_subscriber(subscriber)); + tokio::time::timeout(std::time::Duration::from_secs(/*secs*/ 10), paused).await??; + let b = render(root_b).await; + release.send(())?; + let a = a.await?; + assert!(a.contains("skill-a") && !a.contains("skill-b"), "{a}"); + assert!(b.contains("skill-b") && !b.contains("skill-a"), "{b}"); + Ok(()) +} + #[tokio::test] async fn default_context_truncates_catalog_descriptions() -> TestResult { let description = "x".repeat(1_025); @@ -1618,6 +1760,7 @@ async fn catalog_rebalances_only_to_avoid_omissions_and_retains_the_allocation() session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &turn_store, }) .await; let rendered = sections @@ -2248,6 +2391,7 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &turn_store, }) .await; let fragments = registry.turn_input_contributors()[0] @@ -2378,6 +2522,7 @@ async fn model_context_window_scales_executor_and_cloud_catalogs() -> TestResult session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &turn_store, }) .await; // Core rebuilds world state before each sampling step. @@ -2394,6 +2539,7 @@ async fn model_context_window_scales_executor_and_cloud_catalogs() -> TestResult session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &turn_store, }) .await; assert!(event_rx.try_recv().is_err()); @@ -2493,6 +2639,7 @@ async fn executor_catalog_emits_at_most_four_warnings() -> TestResult { session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &turn_store, }) .await; registry.turn_input_contributors()[0] 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 9d1747538b..c67d715ff1 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 @@ -123,6 +123,7 @@ async fn cloud_preference_preserves_aliases_reads_and_executor_fallback() -> Tes session_store: &session_store, thread_store: &thread_store, turn_store: &turn_store, + step_store: &turn_store, }) .await; let snapshots: serde_json::Map = sections