diff --git a/codex-rs/app-server/src/request_processors/external_agent_session_import.rs b/codex-rs/app-server/src/request_processors/external_agent_session_import.rs index 77b9bfcbf3..bc0a64ef72 100644 --- a/codex-rs/app-server/src/request_processors/external_agent_session_import.rs +++ b/codex-rs/app-server/src/request_processors/external_agent_session_import.rs @@ -215,6 +215,7 @@ impl ExternalAgentSessionImporter { .unwrap_or_else(|| model_info.get_model_instructions(config.personality)), }, dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: Some(MultiAgentVersion::V1), initial_window_id: uuid::Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 4703518f7f..4902471f8f 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -9,6 +9,16 @@ use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_WORKSPACE; const THREAD_LIST_DEFAULT_LIMIT: usize = 25; const THREAD_LIST_MAX_LIMIT: usize = 100; +fn selected_capability_thread_init( + selected_capability_roots: Vec, +) -> ExtensionDataInit { + let mut thread_init = ExtensionDataInit::new(); + if !selected_capability_roots.is_empty() { + thread_init.insert(selected_capability_roots); + } + thread_init +} + struct ThreadListFilters { model_providers: Option>, source_kinds: Option>, @@ -1131,11 +1141,7 @@ impl ThreadRequestProcessor { DynamicToolSpec::Namespace(namespace) => namespace.tools.len(), }) .sum(); - let mut thread_extension_init = ExtensionDataInit::new(); - if !selected_capability_roots.is_empty() { - thread_extension_init.insert(selected_capability_roots); - codex_mcp_extension::initialize_executor_plugin_thread_data(&mut thread_extension_init); - } + let thread_extension_init = selected_capability_thread_init(selected_capability_roots); let create_thread_started_at = std::time::Instant::now(); let NewThread { thread_id, diff --git a/codex-rs/app-server/tests/common/rollout.rs b/codex-rs/app-server/tests/common/rollout.rs index ec01dacc47..8d2567c010 100644 --- a/codex-rs/app-server/tests/common/rollout.rs +++ b/codex-rs/app-server/tests/common/rollout.rs @@ -200,6 +200,7 @@ fn create_fake_rollout_with_source_and_parent_thread_id( model_provider: model_provider.map(str::to_string), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, @@ -288,6 +289,7 @@ pub fn create_fake_rollout_with_text_elements( model_provider: model_provider.map(str::to_string), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/app-server/tests/suite/conversation_summary.rs b/codex-rs/app-server/tests/suite/conversation_summary.rs index 7f964d3147..63ad473261 100644 --- a/codex-rs/app-server/tests/suite/conversation_summary.rs +++ b/codex-rs/app-server/tests/suite/conversation_summary.rs @@ -131,6 +131,7 @@ async fn get_conversation_summary_by_thread_id_reads_pathless_store_thread() -> originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { diff --git a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs index 299ca81aa5..c9467e572b 100644 --- a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs +++ b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs @@ -9,6 +9,10 @@ use codex_app_server_protocol::McpServerToolCallParams; use codex_app_server_protocol::McpServerToolCallResponse; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::SelectedCapabilityRoot; +use codex_app_server_protocol::ThreadForkParams; +use codex_app_server_protocol::ThreadForkResponse; +use codex_app_server_protocol::ThreadResumeParams; +use codex_app_server_protocol::ThreadResumeResponse; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; @@ -32,7 +36,7 @@ const REFRESH_PROBE_SERVER_NAME: &str = "refresh_probe"; const TOOL_CALL_ID: &str = "executor-mcp-call"; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn selected_executor_plugin_exposes_its_stdio_mcp_only_to_that_thread() -> Result<()> { +async fn selected_executor_plugin_mcp_survives_fork_and_resume() -> Result<()> { let responses_server = responses::start_mock_server().await; let codex_home = TempDir::new()?; write_mock_responses_config_toml( @@ -99,27 +103,6 @@ args = ["exec-server", "--listen", "stdio"] ) .await?; - std::fs::write(plugin.path().join(".mcp.json"), r#"{"mcpServers":{}}"#)?; - let config_path = codex_home.path().join("config.toml"); - let mut config = std::fs::read_to_string(&config_path)?; - config.push_str(&format!( - r#" -[mcp_servers.{REFRESH_PROBE_SERVER_NAME}] -command = {} -startup_timeout_sec = 10 -"#, - toml::Value::String(stdio_server_bin()?) - )); - std::fs::write(config_path, config)?; - let request_id = app_server - .send_raw_request("config/mcpServer/reload", /*params*/ None) - .await?; - timeout( - DEFAULT_READ_TIMEOUT, - app_server.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let namespace = format!("mcp__{MCP_SERVER_NAME}"); let response_mock = responses::mount_sse_sequence( &responses_server, @@ -178,6 +161,116 @@ startup_timeout_sec = 10 assert!(output.contains("ECHOING: hello from executor")); assert!(output.contains(EXECUTOR_ENV_VALUE)); + let request_id = app_server + .send_thread_fork_request(ThreadForkParams { + thread_id: selected_thread.clone(), + ..Default::default() + }) + .await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let ThreadForkResponse { + thread: forked_thread, + .. + } = to_response(response)?; + let forked_thread_id = forked_thread.id; + assert!( + mcp_server_names(&mut app_server, forked_thread_id.clone()) + .await? + .iter() + .any(|name| name == MCP_SERVER_NAME) + ); + + let config_path = codex_home.path().join("config.toml"); + let mut config = std::fs::read_to_string(&config_path)?; + config.push_str(&format!( + r#" +[mcp_servers.{REFRESH_PROBE_SERVER_NAME}] +command = {} +startup_timeout_sec = 10 +"#, + toml::Value::String(stdio_server_bin()?) + )); + std::fs::write(config_path, config)?; + + drop(app_server); + let mut app_server = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, app_server.initialize()).await??; + let request_id = app_server + .send_thread_resume_request(ThreadResumeParams { + thread_id: selected_thread.clone(), + ..Default::default() + }) + .await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let ThreadResumeResponse { thread, .. } = to_response(response)?; + assert_eq!(thread.id, selected_thread); + assert!( + mcp_server_names(&mut app_server, thread.id) + .await? + .iter() + .any(|name| name == MCP_SERVER_NAME) + ); + + let request_id = app_server + .send_thread_resume_request(ThreadResumeParams { + thread_id: forked_thread_id.clone(), + ..Default::default() + }) + .await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let ThreadResumeResponse { thread, .. } = to_response(response)?; + assert_eq!(thread.id, forked_thread_id); + assert!( + mcp_server_names(&mut app_server, thread.id) + .await? + .iter() + .any(|name| name == MCP_SERVER_NAME) + ); + + std::fs::write(plugin.path().join(".mcp.json"), r#"{"mcpServers":{}}"#)?; + let request_id = app_server + .send_raw_request("config/mcpServer/reload", /*params*/ None) + .await?; + timeout( + DEFAULT_READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + + let request_id = app_server + .send_mcp_server_tool_call_request(McpServerToolCallParams { + thread_id: selected_thread.clone(), + server: MCP_SERVER_NAME.to_string(), + tool: "echo".to_string(), + arguments: Some(json!({"message": "frozen after reload"})), + meta: None, + }) + .await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let response: McpServerToolCallResponse = to_response(response)?; + assert_eq!( + response + .structured_content + .and_then(|content| content.get("echo").cloned()), + Some(json!("ECHOING: frozen after reload")) + ); + let request_id = app_server .send_mcp_server_tool_call_request(McpServerToolCallParams { thread_id: selected_thread.clone(), diff --git a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs index 8fdb38e5ea..7e0db1b5b3 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs @@ -158,6 +158,7 @@ async fn thread_delete_with_non_local_thread_store_does_not_create_local_persist originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { diff --git a/codex-rs/app-server/tests/suite/v2/thread_read.rs b/codex-rs/app-server/tests/suite/v2/thread_read.rs index 6b427b0d97..9f67761757 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_read.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_read.rs @@ -1369,6 +1369,7 @@ async fn seed_pathless_store_thread( originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index ce95f4810c..533087a11e 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -2070,6 +2070,7 @@ stream_max_retries = 0 model_provider: Some("mock_provider".to_string()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs index f647dc2cfd..231ede47b9 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs @@ -218,6 +218,7 @@ async fn thread_unarchive_preserves_pathless_store_metadata() -> Result<()> { originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { diff --git a/codex-rs/core/src/personality_migration_tests.rs b/codex-rs/core/src/personality_migration_tests.rs index dd0da327c4..396d065644 100644 --- a/codex-rs/core/src/personality_migration_tests.rs +++ b/codex-rs/core/src/personality_migration_tests.rs @@ -59,6 +59,7 @@ async fn write_rollout_with_user_event(dir: &Path, thread_id: ThreadId) -> io::R model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index af64e0bef5..296002414c 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -9,6 +9,7 @@ use crate::skills::SkillError; use crate::state::ActiveTurn; use codex_extension_api::ExtensionDataInit; use codex_protocol::SessionId; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE; use codex_protocol::config_types::ServiceTier; use codex_protocol::permissions::FileSystemPath; @@ -594,6 +595,10 @@ impl Session { text: session_configuration.base_instructions.clone(), }, dynamic_tools: session_configuration.dynamic_tools.clone(), + selected_capability_roots: mcp_thread_init + .get::>() + .map(|roots| roots.as_ref().clone()) + .unwrap_or_default(), multi_agent_version: initial_multi_agent_version, initial_window_id: initial_auto_compact_window_ids .window_id diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 983ed20094..81c0da4b9d 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4037,6 +4037,7 @@ async fn attach_thread_persistence(session: &mut Session) -> PathBuf { originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { @@ -6920,6 +6921,7 @@ async fn shutdown_complete_does_not_append_to_thread_store_after_shutdown() { originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 4846bfd890..907de03705 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -39,6 +39,7 @@ use codex_model_provider_info::OPENAI_PROVIDER_ID; use codex_models_manager::manager::RefreshStrategy; use codex_models_manager::manager::SharedModelsManager; use codex_protocol::ThreadId; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::config_types::CollaborationModeMask; use codex_protocol::config_types::MultiAgentMode; use codex_protocol::error::CodexErr; @@ -1513,7 +1514,7 @@ impl ThreadManagerState { inherited_exec_policy: Option>, parent_trace: Option, environments: Vec, - thread_extension_init: ExtensionDataInit, + mut thread_extension_init: ExtensionDataInit, supports_openai_form_elicitation: bool, user_shell_override: Option, ) -> CodexResult { @@ -1539,6 +1540,19 @@ impl ThreadManagerState { threads.remove(&resumed.conversation_id); } } + if thread_extension_init + .get::>() + .is_none() + { + let selected_capability_roots = initial_history.get_selected_capability_roots(); + if !selected_capability_roots.is_empty() { + thread_extension_init.insert(selected_capability_roots); + } + } + self.extensions + .initialize_thread_data(&mut thread_extension_init) + .await; + let user_instructions = self .user_instructions_for_spawn(&session_source, parent_thread_id, forked_from_thread_id) .await; diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index b1f5c4e24c..75f0bdf08a 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -428,10 +428,29 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { #[tokio::test] async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() { struct InitialDataRecorder { + init_observed: Arc>>, lifecycle_observed: Arc>>, mcp_observed: Arc>>, } + impl codex_extension_api::ThreadExtensionInitContributor for InitialDataRecorder { + fn initialize<'a>( + &'a self, + thread_init: &'a mut codex_extension_api::ExtensionDataInit, + ) -> codex_extension_api::ExtensionFuture<'a, ()> { + Box::pin(async move { + let selected_root = thread_init + .get::>() + .and_then(|roots| roots.first().cloned()) + .expect("selected root should be available"); + self.init_observed + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(selected_root.id); + }) + } + } + impl codex_extension_api::ThreadLifecycleContributor for InitialDataRecorder { fn on_thread_start<'a>( &'a self, @@ -502,13 +521,16 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() config.cwd = config.codex_home.abs(); std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + let init_observed = Arc::new(std::sync::Mutex::new(Vec::new())); let lifecycle_observed = Arc::new(std::sync::Mutex::new(Vec::new())); let mcp_observed = Arc::new(std::sync::Mutex::new(Vec::new())); let recorder = Arc::new(InitialDataRecorder { + init_observed: Arc::clone(&init_observed), lifecycle_observed: Arc::clone(&lifecycle_observed), mcp_observed: Arc::clone(&mcp_observed), }); let mut extensions = codex_extension_api::ExtensionRegistryBuilder::new(); + extensions.thread_extension_init_contributor(recorder.clone()); extensions.thread_lifecycle_contributor(recorder.clone()); extensions.mcp_server_contributor(recorder); let manager = ThreadManager::new( @@ -572,6 +594,12 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() let first_resolved = first_thread.thread.runtime_mcp_config(&config).await; let second_resolved = second_thread.thread.runtime_mcp_config(&config).await; + assert_eq!( + *init_observed + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + vec!["selected-a".to_string(), "selected-b".to_string()] + ); assert_eq!( *lifecycle_observed .lock() diff --git a/codex-rs/core/tests/suite/personality_migration.rs b/codex-rs/core/tests/suite/personality_migration.rs index ca1d148728..3dd0f3c016 100644 --- a/codex-rs/core/tests/suite/personality_migration.rs +++ b/codex-rs/core/tests/suite/personality_migration.rs @@ -75,6 +75,7 @@ async fn write_rollout_with_user_event(dir: &Path, thread_id: ThreadId) -> io::R model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, @@ -127,6 +128,7 @@ async fn write_rollout_with_meta_only(dir: &Path, thread_id: ThreadId) -> io::Re model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/core/tests/suite/sqlite_state.rs b/codex-rs/core/tests/suite/sqlite_state.rs index ddfc8ef210..8828fc97fd 100644 --- a/codex-rs/core/tests/suite/sqlite_state.rs +++ b/codex-rs/core/tests/suite/sqlite_state.rs @@ -372,6 +372,7 @@ async fn backfill_scans_existing_rollouts() -> Result<()> { model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 7d16efd3d4..f23a393deb 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -10,6 +10,7 @@ use codex_tools::ToolCall; use codex_tools::ToolExecutor; use crate::ExtensionData; +use crate::ExtensionDataInit; mod context; mod mcp; @@ -43,6 +44,16 @@ pub use turn_lifecycle::TurnStopInput; /// Boxed, sendable future returned by asynchronous extension contributors. pub type ExtensionFuture<'a, T> = Pin + Send + 'a>>; +/// Extension contribution that seeds thread data before runtime startup. +/// +/// Implementations may derive extension-owned state from host-provided values. +/// Initialization completes before startup-time contributors consume the +/// thread data. +pub trait ThreadExtensionInitContributor: Send + Sync { + /// Adds extension-owned state to the host-provided thread initializer. + fn initialize<'a>(&'a self, thread_init: &'a mut ExtensionDataInit) -> ExtensionFuture<'a, ()>; +} + /// Extension contribution that resolves runtime MCP servers from host config. /// /// Contributors run in registration order. Later contributions for the same diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index de0e0c99f8..9e2aa82133 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -40,6 +40,7 @@ pub use contributors::McpServerContributionContext; pub use contributors::McpServerContributor; pub use contributors::PromptFragment; pub use contributors::PromptSlot; +pub use contributors::ThreadExtensionInitContributor; pub use contributors::ThreadIdleInput; pub use contributors::ThreadLifecycleContributor; pub use contributors::ThreadResumeInput; diff --git a/codex-rs/ext/extension-api/src/registry.rs b/codex-rs/ext/extension-api/src/registry.rs index e8cc989337..1f4575ec2b 100644 --- a/codex-rs/ext/extension-api/src/registry.rs +++ b/codex-rs/ext/extension-api/src/registry.rs @@ -9,6 +9,7 @@ use crate::ExtensionData; use crate::ExtensionEventSink; use crate::McpServerContributor; use crate::NoopExtensionEventSink; +use crate::ThreadExtensionInitContributor; use crate::ThreadLifecycleContributor; use crate::TokenUsageContributor; use crate::ToolContributor; @@ -21,6 +22,7 @@ use crate::TurnLifecycleContributor; pub struct ExtensionRegistryBuilder { event_sink: Arc, thread_lifecycle_contributors: Vec>>, + thread_init_contributors: Vec>, turn_lifecycle_contributors: Vec>, config_contributors: Vec>>, token_usage_contributors: Vec>, @@ -38,6 +40,7 @@ impl Default for ExtensionRegistryBuilder { Self { event_sink: Arc::new(NoopExtensionEventSink), thread_lifecycle_contributors: Vec::new(), + thread_init_contributors: Vec::new(), turn_lifecycle_contributors: Vec::new(), config_contributors: Vec::new(), token_usage_contributors: Vec::new(), @@ -84,6 +87,14 @@ impl ExtensionRegistryBuilder { self.thread_lifecycle_contributors.push(contributor); } + /// Registers one startup-time thread-data contributor. + pub fn thread_extension_init_contributor( + &mut self, + contributor: Arc, + ) { + self.thread_init_contributors.push(contributor); + } + /// Registers one turn-lifecycle contributor. pub fn turn_lifecycle_contributor(&mut self, contributor: Arc) { self.turn_lifecycle_contributors.push(contributor); @@ -134,6 +145,7 @@ impl ExtensionRegistryBuilder { ExtensionRegistry { event_sink: self.event_sink, thread_lifecycle_contributors: self.thread_lifecycle_contributors, + thread_init_contributors: self.thread_init_contributors, turn_lifecycle_contributors: self.turn_lifecycle_contributors, config_contributors: self.config_contributors, token_usage_contributors: self.token_usage_contributors, @@ -152,6 +164,7 @@ impl ExtensionRegistryBuilder { pub struct ExtensionRegistry { event_sink: Arc, thread_lifecycle_contributors: Vec>>, + thread_init_contributors: Vec>, turn_lifecycle_contributors: Vec>, config_contributors: Vec>>, token_usage_contributors: Vec>, @@ -175,6 +188,13 @@ impl ExtensionRegistry { &self.thread_lifecycle_contributors } + /// Applies registered startup-time thread-data contributors in order. + pub async fn initialize_thread_data(&self, thread_init: &mut crate::ExtensionDataInit) { + for contributor in &self.thread_init_contributors { + contributor.initialize(thread_init).await; + } + } + /// Returns the registered turn-lifecycle contributors. pub fn turn_lifecycle_contributors(&self) -> &[Arc] { &self.turn_lifecycle_contributors diff --git a/codex-rs/ext/mcp/src/executor_plugin.rs b/codex-rs/ext/mcp/src/executor_plugin.rs index 69d1a47b31..84e4e83c4f 100644 --- a/codex-rs/ext/mcp/src/executor_plugin.rs +++ b/codex-rs/ext/mcp/src/executor_plugin.rs @@ -6,6 +6,7 @@ use codex_extension_api::ExtensionFuture; use codex_extension_api::McpServerContribution; use codex_extension_api::McpServerContributionContext; use codex_extension_api::McpServerContributor; +use codex_extension_api::ThreadExtensionInitContributor; use codex_protocol::capabilities::SelectedCapabilityRoot; use std::collections::HashMap; use std::sync::Arc; @@ -36,6 +37,22 @@ pub(crate) fn seed_thread_state(thread_init: &mut ExtensionDataInit) { thread_init.insert(SelectedExecutorPluginMcpState::default()); } +pub(crate) struct SelectedExecutorPluginMcpInitializer; + +impl ThreadExtensionInitContributor for SelectedExecutorPluginMcpInitializer { + fn initialize<'a>(&'a self, thread_init: &'a mut ExtensionDataInit) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + if thread_init.get::>().is_some() + && thread_init + .get::() + .is_none() + { + seed_thread_state(thread_init); + } + }) + } +} + pub(crate) struct SelectedExecutorPluginMcpContributor { plugin_provider: ExecutorPluginProvider, mcp_provider: ExecutorPluginMcpProvider, diff --git a/codex-rs/ext/mcp/src/lib.rs b/codex-rs/ext/mcp/src/lib.rs index e32316ac36..d76a83bcfa 100644 --- a/codex-rs/ext/mcp/src/lib.rs +++ b/codex-rs/ext/mcp/src/lib.rs @@ -47,14 +47,10 @@ pub fn install_executor_plugins( builder: &mut ExtensionRegistryBuilder, environment_manager: std::sync::Arc, ) { + builder.thread_extension_init_contributor(std::sync::Arc::new( + executor_plugin::SelectedExecutorPluginMcpInitializer, + )); builder.mcp_server_contributor(std::sync::Arc::new( executor_plugin::SelectedExecutorPluginMcpContributor::new(environment_manager), )); } - -/// Seeds the per-thread snapshot used by selected executor plugin MCP discovery. -pub fn initialize_executor_plugin_thread_data( - thread_init: &mut codex_extension_api::ExtensionDataInit, -) { - executor_plugin::seed_thread_state(thread_init); -} diff --git a/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs b/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs index 7bdbc83d62..423b5785c5 100644 --- a/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs +++ b/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs @@ -109,7 +109,7 @@ async fn selected_plugin_contributions( path: PathUri::from_host_native_path(plugin_root)?, }, }]); - codex_mcp_extension::initialize_executor_plugin_thread_data(&mut thread_init); + registry.initialize_thread_data(&mut thread_init).await; Ok(registry.mcp_server_contributors()[0] .contribute(McpServerContributionContext::for_thread( diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 138b36f10f..640ca7d45d 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -19,6 +19,7 @@ use crate::AgentPath; use crate::SessionId; use crate::ThreadId; use crate::approvals::ElicitationRequestEvent; +use crate::capabilities::SelectedCapabilityRoot; use crate::config_types::ApprovalsReviewer; use crate::config_types::CollaborationMode; use crate::config_types::ModeKind; @@ -2573,6 +2574,12 @@ impl InitialHistory { } } + pub fn get_selected_capability_roots(&self) -> Vec { + self.get_session_meta() + .map(|meta| meta.selected_capability_roots.clone()) + .unwrap_or_default() + } + pub fn get_multi_agent_version(&self) -> Option { match self { InitialHistory::New | InitialHistory::Cleared => None, @@ -3004,6 +3011,9 @@ pub struct SessionMeta { skip_serializing_if = "Option::is_none" )] pub dynamic_tools: Option>, + /// Capability roots selected for this thread by the hosting platform. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub selected_capability_roots: Vec, #[serde(skip_serializing_if = "Option::is_none")] pub memory_mode: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -3033,6 +3043,7 @@ impl Default for SessionMeta { model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/rollout/src/compression_tests.rs b/codex-rs/rollout/src/compression_tests.rs index a4db9a965c..733e6cc5fb 100644 --- a/codex-rs/rollout/src/compression_tests.rs +++ b/codex-rs/rollout/src/compression_tests.rs @@ -472,6 +472,7 @@ fn write_rollout(path: &std::path::Path, thread_id: ThreadId, message: &str) -> model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/rollout/src/metadata_tests.rs b/codex-rs/rollout/src/metadata_tests.rs index 0744e7780e..c0aea248c5 100644 --- a/codex-rs/rollout/src/metadata_tests.rs +++ b/codex-rs/rollout/src/metadata_tests.rs @@ -49,6 +49,7 @@ async fn extract_metadata_from_rollout_uses_session_meta() { model_provider: Some("openai".to_string()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, @@ -106,6 +107,7 @@ async fn extract_metadata_from_rollout_returns_latest_memory_mode() { model_provider: Some("openai".to_string()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, @@ -375,6 +377,7 @@ fn write_rollout_in_sessions_with_cwd( model_provider: Some("test-provider".to_string()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 6459420710..1d5df2d707 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -12,6 +12,7 @@ use std::sync::Mutex; use chrono::SecondsFormat; use codex_protocol::SessionId; use codex_protocol::ThreadId; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::BaseInstructions; use serde_json::Value; @@ -91,6 +92,7 @@ pub enum RolloutRecorderParams { originator: String, base_instructions: BaseInstructions, dynamic_tools: Vec, + selected_capability_roots: Vec, multi_agent_version: Option, initial_window_id: Option, }, @@ -182,6 +184,7 @@ impl RolloutRecorderParams { originator, base_instructions, dynamic_tools, + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: None, } @@ -194,6 +197,20 @@ impl RolloutRecorderParams { self } + pub fn with_selected_capability_roots( + mut self, + selected_capability_roots: Vec, + ) -> Self { + if let Self::Create { + selected_capability_roots: roots, + .. + } = &mut self + { + *roots = selected_capability_roots; + } + self + } + pub fn with_multi_agent_version( mut self, multi_agent_version: Option, @@ -732,6 +749,7 @@ impl RolloutRecorder { originator, base_instructions, dynamic_tools, + selected_capability_roots, multi_agent_version, initial_window_id, } => { @@ -769,6 +787,7 @@ impl RolloutRecorder { } else { Some(dynamic_tools) }, + selected_capability_roots, memory_mode: (!config.generate_memories()).then_some("disabled".to_string()), multi_agent_version, context_window: initial_window_id.map(SessionContextWindow::new), diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index 4b72ff7673..71ec252c23 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -101,6 +101,7 @@ async fn state_db_init_backfills_before_returning() -> anyhow::Result<()> { model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/rollout/src/session_index_tests.rs b/codex-rs/rollout/src/session_index_tests.rs index 7f57c2ba08..be39fb06ab 100644 --- a/codex-rs/rollout/src/session_index_tests.rs +++ b/codex-rs/rollout/src/session_index_tests.rs @@ -41,6 +41,7 @@ fn write_rollout_with_metadata(path: &Path, thread_id: ThreadId) -> std::io::Res model_provider: Some("test-provider".into()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/rollout/src/state_db_tests.rs b/codex-rs/rollout/src/state_db_tests.rs index 5712b65cee..bb341aa85a 100644 --- a/codex-rs/rollout/src/state_db_tests.rs +++ b/codex-rs/rollout/src/state_db_tests.rs @@ -174,6 +174,7 @@ fn write_rollout_with_user_message( model_provider: Some("test-provider".to_string()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/rollout/src/tests.rs b/codex-rs/rollout/src/tests.rs index e6a3c728ef..901461cf56 100644 --- a/codex-rs/rollout/src/tests.rs +++ b/codex-rs/rollout/src/tests.rs @@ -1288,6 +1288,7 @@ async fn test_updated_at_uses_file_mtime() -> Result<()> { model_provider: Some("test-provider".into()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index 5291b679eb..d5d4846a4b 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -341,6 +341,7 @@ mod tests { model_provider: Some("openai".to_string()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, @@ -533,6 +534,7 @@ mod tests { model_provider: Some("openai".to_string()), base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 6f8edc36d2..077a2d7de7 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -2147,6 +2147,7 @@ mod tests { model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: Some("polluted".to_string()), multi_agent_version: None, context_window: None, @@ -2210,6 +2211,7 @@ mod tests { model_provider: None, base_instructions: None, dynamic_tools: None, + selected_capability_roots: Vec::new(), memory_mode: None, multi_agent_version: None, context_window: None, diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index dfed2ffac2..d64510f223 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -126,6 +126,7 @@ mod tests { originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: uuid::Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { @@ -281,6 +282,7 @@ impl InMemoryThreadStore { model_provider: Some(params.metadata.model_provider.clone()), base_instructions: Some(params.base_instructions.clone()), dynamic_tools: (!params.dynamic_tools.is_empty()).then(|| params.dynamic_tools.clone()), + selected_capability_roots: params.selected_capability_roots.clone(), memory_mode: matches!(params.metadata.memory_mode, ThreadMemoryMode::Disabled) .then_some("disabled".to_string()), multi_agent_version: params.multi_agent_version, diff --git a/codex-rs/thread-store/src/local/create_thread.rs b/codex-rs/thread-store/src/local/create_thread.rs index eba6a5a6a9..451931b4dc 100644 --- a/codex-rs/thread-store/src/local/create_thread.rs +++ b/codex-rs/thread-store/src/local/create_thread.rs @@ -38,6 +38,7 @@ pub(super) async fn create_thread( params.dynamic_tools, ) .with_session_id(params.session_id) + .with_selected_capability_roots(params.selected_capability_roots) .with_multi_agent_version(params.multi_agent_version) .with_initial_window_id(params.initial_window_id), ) diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 51e622fc5b..717ae09b69 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -1133,6 +1133,7 @@ mod tests { originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: uuid::Uuid::now_v7().to_string(), metadata: thread_metadata(), diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index 1f761cea93..f96a4df31b 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -5,6 +5,7 @@ use chrono::DateTime; use chrono::Utc; use codex_protocol::SessionId; use codex_protocol::ThreadId; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::BaseInstructions; use codex_protocol::models::PermissionProfile; @@ -85,6 +86,9 @@ pub struct CreateThreadParams { pub base_instructions: BaseInstructions, /// Dynamic tools available to the thread at startup. pub dynamic_tools: Vec, + /// Environment-qualified capability roots selected for this thread. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub selected_capability_roots: Vec, /// Multi-agent runtime selected when the thread was created. pub multi_agent_version: Option, /// Initial context-window identity captured when the thread was created.