diff --git a/codex-rs/app-server/src/request_processors/external_agent_config_processor.rs b/codex-rs/app-server/src/request_processors/external_agent_config_processor.rs index b3a2a5e06c..b2899f6ba7 100644 --- a/codex-rs/app-server/src/request_processors/external_agent_config_processor.rs +++ b/codex-rs/app-server/src/request_processors/external_agent_config_processor.rs @@ -308,9 +308,7 @@ impl ExternalAgentConfigRequestProcessor { initial_history: InitialHistory::Forked(rollout_items), session_source: None, thread_source: None, - code_mode_session_provider: Some( - ThreadManager::in_process_code_mode_session_provider(), - ), + code_mode_session_provider_selection: Default::default(), dynamic_tools: Vec::new(), persist_extended_history: false, metrics_service_name: None, 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 99ddb34130..6cd4ac2996 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -1105,9 +1105,7 @@ impl ThreadRequestProcessor { }, session_source: None, thread_source, - code_mode_session_provider: Some( - ThreadManager::in_process_code_mode_session_provider(), - ), + code_mode_session_provider_selection: Default::default(), dynamic_tools: core_dynamic_tools, persist_extended_history: false, metrics_service_name: service_name, diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index 37ba1d2a91..95a571bc90 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -39,6 +39,7 @@ pub use service::CodeModeSessionResultFuture; pub use service::InProcessCodeModeSessionProvider; pub use service::NoopCodeModeSessionDelegate; pub use service::NotificationFuture; +pub use service::SessionProviderSelection; pub use service::StartedCell; pub use service::ToolInvocationFuture; diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index 6589e6ec5a..cb881ae894 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -156,6 +156,31 @@ pub trait CodeModeSessionProvider: Send + Sync { ) -> CodeModeSessionProviderFuture<'a>; } +/// Selects how a code-mode session is provisioned for a Codex session. +#[derive(Clone, Default)] +pub enum SessionProviderSelection { + #[default] + InProcess, + Disabled, + Custom(Arc), +} + +impl SessionProviderSelection { + pub async fn create_session( + &self, + delegate: Arc, + ) -> Result>, String> { + match self { + Self::InProcess => InProcessCodeModeSessionProvider + .create_session(delegate) + .await + .map(Some), + Self::Disabled => Ok(None), + Self::Custom(provider) => provider.create_session(delegate).await.map(Some), + } + } +} + #[derive(Default)] pub struct InProcessCodeModeSessionProvider; @@ -870,6 +895,7 @@ mod tests { use super::PendingRuntimeMode; use super::RuntimeCommand; use super::RuntimeResponse; + use super::SessionProviderSelection; use super::WaitOutcome; use super::WaitRequest; use super::WaitToPendingOutcome; @@ -943,6 +969,26 @@ mod tests { ); } + #[tokio::test] + async fn provider_selection_defaults_to_in_process_and_can_disable() { + let delegate = Arc::new(NoopCodeModeSessionDelegate); + + assert!( + SessionProviderSelection::default() + .create_session(delegate.clone()) + .await + .unwrap() + .is_some() + ); + assert!( + SessionProviderSelection::Disabled + .create_session(delegate) + .await + .unwrap() + .is_none() + ); + } + #[tokio::test] async fn stored_values_are_shared_between_cells_but_not_sessions() { let first_session = CodeModeService::new(); diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index f39d38e5ab..026e1840d3 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -151,7 +151,7 @@ fn is_multi_agent_v2_usage_hint_message(item: &ResponseItem, usage_hint_texts: & /// An `AgentControl` instance is intended to be created at most once per root thread/session /// tree. That same `AgentControl` is then shared with every sub-agent spawned from that root, /// which keeps the registry scoped to that root thread rather than the entire `ThreadManager`. -#[derive(Clone, Default)] +#[derive(Clone)] pub(crate) struct AgentControl { /// ID shared by the whole agent control session. This means every sub-agents from a common /// root share the same session ID. @@ -161,18 +161,30 @@ pub(crate) struct AgentControl { /// `ThreadManagerState -> CodexThread -> Session -> SessionServices -> ThreadManagerState`. manager: Weak, state: Arc, - code_mode_session_provider: Option>, + code_mode_session_provider_selection: codex_code_mode::SessionProviderSelection, +} + +impl Default for AgentControl { + fn default() -> Self { + Self { + session_id: SessionId::default(), + manager: Weak::new(), + state: Arc::new(AgentRegistry::default()), + code_mode_session_provider_selection: + codex_code_mode::SessionProviderSelection::Disabled, + } + } } impl AgentControl { /// Construct a new `AgentControl` that can spawn/message agents via the given manager state. pub(crate) fn new( manager: Weak, - code_mode_session_provider: Option>, + code_mode_session_provider_selection: codex_code_mode::SessionProviderSelection, ) -> Self { Self { manager, - code_mode_session_provider, + code_mode_session_provider_selection, ..Default::default() } } @@ -186,10 +198,10 @@ impl AgentControl { self.session_id } - pub(crate) fn code_mode_session_provider( + pub(crate) fn code_mode_session_provider_selection( &self, - ) -> Option> { - self.code_mode_session_provider.clone() + ) -> codex_code_mode::SessionProviderSelection { + self.code_mode_session_provider_selection.clone() } /// Spawn a new agent thread and submit the initial prompt. diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 6e04be7f09..32606b2f05 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -608,7 +608,8 @@ impl Codex { // Generate a unique ID for the lifetime of this Codex session. let session_source_clone = session_configuration.session_source.clone(); let (agent_status_tx, agent_status_rx) = watch::channel(AgentStatus::PendingInit); - let code_mode_session_provider = agent_control.code_mode_session_provider(); + let code_mode_session_provider_selection = + agent_control.code_mode_session_provider_selection(); let session = Session::new( session_configuration, @@ -624,7 +625,7 @@ impl Codex { skills_manager, plugins_manager, mcp_manager.clone(), - code_mode_session_provider, + code_mode_session_provider_selection, extensions, agent_control, environment_manager, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 860295c345..546b840a05 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -497,7 +497,7 @@ impl Session { skills_manager: Arc, plugins_manager: Arc, mcp_manager: Arc, - code_mode_session_provider: Option>, + code_mode_session_provider_selection: codex_code_mode::SessionProviderSelection, extensions: Arc>, agent_control: AgentControl, environment_manager: Arc, @@ -1053,11 +1053,12 @@ impl Session { session_configuration.parent_thread_id, ), ), - code_mode_service: crate::tools::code_mode::CodeModeService::from_provider( - code_mode_session_provider, - ) - .await - .map_err(anyhow::Error::msg)?, + code_mode_service: + crate::tools::code_mode::CodeModeService::from_provider_selection( + code_mode_session_provider_selection, + ) + .await + .map_err(anyhow::Error::msg)?, environment_manager, }; services diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index d5036dd8d3..344aa21337 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4508,7 +4508,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { skills_manager, plugins_manager, mcp_manager, - /*code_mode_session_provider*/ None, + codex_code_mode::SessionProviderSelection::Disabled, Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), AgentControl::default(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -4858,7 +4858,7 @@ async fn make_session_with_config_and_rx( skills_manager, plugins_manager, mcp_manager, - /*code_mode_session_provider*/ None, + codex_code_mode::SessionProviderSelection::Disabled, Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), AgentControl::default(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -4964,7 +4964,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( skills_manager, plugins_manager, mcp_manager, - /*code_mode_session_provider*/ None, + codex_code_mode::SessionProviderSelection::Disabled, Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), agent_control, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 2644c8d682..8d9f4c52ef 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -68,7 +68,6 @@ use std::collections::HashMap; use std::collections::HashSet; use std::path::PathBuf; use std::sync::Arc; -use std::sync::LazyLock; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Duration; @@ -83,10 +82,6 @@ const THREAD_CREATED_CHANNEL_CAPACITY: usize = 1024; /// In production builds this value should remain at its default (`false`) and /// must not be toggled. static FORCE_TEST_THREAD_MANAGER_BEHAVIOR: AtomicBool = AtomicBool::new(false); -static IN_PROCESS_CODE_MODE_SESSION_PROVIDER: LazyLock< - Arc, -> = LazyLock::new(|| Arc::new(codex_code_mode::InProcessCodeModeSessionProvider)); - type CapturedOps = Vec<(ThreadId, Op)>; type SharedCapturedOps = Arc>; @@ -180,7 +175,7 @@ pub struct StartThreadOptions { pub initial_history: InitialHistory, pub session_source: Option, pub thread_source: Option, - pub code_mode_session_provider: Option>, + pub code_mode_session_provider_selection: codex_code_mode::SessionProviderSelection, pub dynamic_tools: Vec, pub persist_extended_history: bool, pub metrics_service_name: Option, @@ -409,11 +404,6 @@ impl ThreadManager { self.state.session_source.clone() } - pub fn in_process_code_mode_session_provider() - -> Arc { - Arc::clone(&IN_PROCESS_CODE_MODE_SESSION_PROVIDER) - } - pub fn auth_manager(&self) -> Arc { self.state.auth_manager.clone() } @@ -587,7 +577,7 @@ impl ThreadManager { initial_history: InitialHistory::New, session_source: None, thread_source: None, - code_mode_session_provider: Some(Self::in_process_code_mode_session_provider()), + code_mode_session_provider_selection: Default::default(), dynamic_tools, persist_extended_history, metrics_service_name: None, @@ -616,8 +606,9 @@ impl ThreadManager { .unwrap_or_else(|| (self.state.session_source.clone(), None)); let session_source = options.session_source.unwrap_or(resumed_session_source); let thread_source = options.thread_source.or(resumed_thread_source); - let agent_control = - self.agent_control_with_code_mode_session_provider(options.code_mode_session_provider); + let agent_control = self.agent_control_with_code_mode_session_provider_selection( + options.code_mode_session_provider_selection, + ); Box::pin(self.state.spawn_thread_with_source( options.config, options.initial_history, @@ -957,16 +948,17 @@ impl ThreadManager { } pub(crate) fn agent_control(&self) -> AgentControl { - self.agent_control_with_code_mode_session_provider(Some( - Self::in_process_code_mode_session_provider(), - )) + self.agent_control_with_code_mode_session_provider_selection(Default::default()) } - fn agent_control_with_code_mode_session_provider( + fn agent_control_with_code_mode_session_provider_selection( &self, - code_mode_session_provider: Option>, + code_mode_session_provider_selection: codex_code_mode::SessionProviderSelection, ) -> AgentControl { - AgentControl::new(Arc::downgrade(&self.state), code_mode_session_provider) + AgentControl::new( + Arc::downgrade(&self.state), + code_mode_session_provider_selection, + ) } #[cfg(test)] diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index ffbfdb8933..7d0f4a09df 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -342,26 +342,23 @@ async fn start_thread_rejects_explicit_local_environment_when_default_provider_i environment_manager, ); - let result = - manager - .start_thread_with_options(StartThreadOptions { - config: config.clone(), - initial_history: InitialHistory::New, - session_source: None, - thread_source: None, - code_mode_session_provider: Some( - ThreadManager::in_process_code_mode_session_provider(), - ), - dynamic_tools: Vec::new(), - persist_extended_history: false, - metrics_service_name: None, - parent_trace: None, - environments: vec![TurnEnvironmentSelection { - environment_id: "local".to_string(), - cwd: config.cwd.clone(), - }], - }) - .await; + let result = manager + .start_thread_with_options(StartThreadOptions { + config: config.clone(), + initial_history: InitialHistory::New, + session_source: None, + thread_source: None, + code_mode_session_provider_selection: Default::default(), + dynamic_tools: Vec::new(), + persist_extended_history: false, + metrics_service_name: None, + parent_trace: None, + environments: vec![TurnEnvironmentSelection { + environment_id: "local".to_string(), + cwd: config.cwd.clone(), + }], + }) + .await; let err = match result { Ok(_) => panic!("explicit local environment should not resolve when provider is disabled"), Err(err) => err, @@ -464,14 +461,6 @@ args = ["dev", "cd /tmp && true"] assert!(!environment_context.contains("\n ")); } -#[test] -fn in_process_code_mode_session_provider_is_a_singleton() { - assert!(Arc::ptr_eq( - &ThreadManager::in_process_code_mode_session_provider(), - &ThreadManager::in_process_code_mode_session_provider(), - )); -} - #[tokio::test] async fn start_thread_options_code_mode_provider_is_propagated_to_spawned_agents() { let temp_dir = tempdir().expect("tempdir"); @@ -497,7 +486,9 @@ async fn start_thread_options_code_mode_provider_is_propagated_to_spawned_agents initial_history: InitialHistory::New, session_source: None, thread_source: None, - code_mode_session_provider: Some(provider), + code_mode_session_provider_selection: codex_code_mode::SessionProviderSelection::Custom( + provider, + ), dynamic_tools: Vec::new(), persist_extended_history: false, metrics_service_name: None, @@ -533,26 +524,23 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { config.codex_home.to_path_buf(), Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), ); - let thread = - manager - .start_thread_with_options(StartThreadOptions { - config, - initial_history: InitialHistory::New, - session_source: Some(SessionSource::Internal( - InternalSessionSource::MemoryConsolidation, - )), - thread_source: None, - code_mode_session_provider: Some( - ThreadManager::in_process_code_mode_session_provider(), - ), - dynamic_tools: Vec::new(), - persist_extended_history: false, - metrics_service_name: None, - parent_trace: None, - environments: Vec::new(), - }) - .await - .expect("internal thread should start"); + let thread = manager + .start_thread_with_options(StartThreadOptions { + config, + initial_history: InitialHistory::New, + session_source: Some(SessionSource::Internal( + InternalSessionSource::MemoryConsolidation, + )), + thread_source: None, + code_mode_session_provider_selection: Default::default(), + dynamic_tools: Vec::new(), + persist_extended_history: false, + metrics_service_name: None, + parent_trace: None, + environments: Vec::new(), + }) + .await + .expect("internal thread should start"); assert_eq!(manager.list_thread_ids().await, Vec::new()); assert!(manager.get_thread(thread.thread_id).await.is_err()); @@ -601,7 +589,8 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { initial_history: InitialHistory::New, session_source: None, thread_source: None, - code_mode_session_provider: None, + code_mode_session_provider_selection: + codex_code_mode::SessionProviderSelection::Disabled, dynamic_tools: Vec::new(), persist_extended_history: false, metrics_service_name: None, @@ -874,7 +863,8 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { initial_history: InitialHistory::New, session_source: None, thread_source: Some(ThreadSource::User), - code_mode_session_provider: None, + code_mode_session_provider_selection: + codex_code_mode::SessionProviderSelection::Disabled, dynamic_tools: Vec::new(), persist_extended_history: false, metrics_service_name: None, diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index a1d90b5d77..588182a331 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -11,9 +11,9 @@ use std::time::Duration; use codex_code_mode::CellId; use codex_code_mode::CodeModeNestedToolCall; use codex_code_mode::CodeModeSession; -use codex_code_mode::CodeModeSessionProvider; use codex_code_mode::CodeModeToolKind; use codex_code_mode::RuntimeResponse; +use codex_code_mode::SessionProviderSelection; use codex_protocol::models::FunctionCallOutputContentItem; use serde_json::Value as JsonValue; use tokio_util::sync::CancellationToken; @@ -75,14 +75,13 @@ impl CodeModeService { } } - pub(crate) async fn from_provider( - provider: Option>, + pub(crate) async fn from_provider_selection( + provider_selection: SessionProviderSelection, ) -> Result { let dispatch_broker = Arc::new(CodeModeDispatchBroker::new()); - let session = match provider { - Some(provider) => Some(provider.create_session(dispatch_broker.clone()).await?), - None => None, - }; + let session = provider_selection + .create_session(dispatch_broker.clone()) + .await?; Ok(Self { session, dispatch_broker, diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index b4f3da0e89..3c387767df 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -1,7 +1,6 @@ use anyhow::Result; use codex_core::StartThreadOptions; use codex_core::ThreadConfigSnapshot; -use codex_core::ThreadManager; use codex_core::config::AgentRoleConfig; use codex_features::Feature; use codex_protocol::ThreadId; @@ -745,23 +744,21 @@ async fn subagent_stop_replaces_stop_and_skips_internal_subagents() -> Result<() // This matcher would catch the old synthetic "review" SubagentStop target // because the SubagentStop hook above intentionally matches all agent types. - let internal_thread = - test.thread_manager - .start_thread_with_options(StartThreadOptions { - config: test.config.clone(), - initial_history: InitialHistory::New, - session_source: Some(SessionSource::SubAgent(SubAgentSource::Review)), - thread_source: None, - code_mode_session_provider: Some( - ThreadManager::in_process_code_mode_session_provider(), - ), - dynamic_tools: Vec::new(), - persist_extended_history: false, - metrics_service_name: None, - parent_trace: None, - environments: Vec::new(), - }) - .await?; + let internal_thread = test + .thread_manager + .start_thread_with_options(StartThreadOptions { + config: test.config.clone(), + initial_history: InitialHistory::New, + session_source: Some(SessionSource::SubAgent(SubAgentSource::Review)), + thread_source: None, + code_mode_session_provider_selection: Default::default(), + dynamic_tools: Vec::new(), + persist_extended_history: false, + metrics_service_name: None, + parent_trace: None, + environments: Vec::new(), + }) + .await?; let (sandbox_policy, permission_profile) = turn_permission_fields(PermissionProfile::Disabled, test.cwd_path()); diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index 94b3d46c20..dbfac7d8a4 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -247,9 +247,7 @@ impl MemoryStartupContext { InternalSessionSource::MemoryConsolidation, )), thread_source: Some(ThreadSource::MemoryConsolidation), - code_mode_session_provider: Some( - ThreadManager::in_process_code_mode_session_provider(), - ), + code_mode_session_provider_selection: Default::default(), dynamic_tools: Vec::new(), persist_extended_history: false, metrics_service_name: None,