openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-03 19:53:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use explicit code mode session provider selection

Browse Source
This commit is contained in:
Channing Conger
2026-06-01 15:27:43 -07:00
parent 8034aa2790
commit e19307bb11
13 changed files with 156 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/request_processors/external_agent_config_processor.rs
+1 -3
View File
@@ -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,
codex-rs/app-server/src/request_processors/thread_processor.rs
+1 -3
View File
@@ -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,
codex-rs/code-mode/src/lib.rs
+1
View File
@@ -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;
codex-rs/code-mode/src/service.rs
+46
View File
@@ -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<dyn CodeModeSessionProvider>),
}
impl SessionProviderSelection {
pub async fn create_session(
&self,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> Result<Option<Arc<dyn CodeModeSession>>, 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();
codex-rs/core/src/agent/control.rs
+19 -7
View File
@@ -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<ThreadManagerState>,
state: Arc<AgentRegistry>,
code_mode_session_provider: Option<Arc<dyn codex_code_mode::CodeModeSessionProvider>>,
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<ThreadManagerState>,
code_mode_session_provider: Option<Arc<dyn codex_code_mode::CodeModeSessionProvider>>,
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<Arc<dyn codex_code_mode::CodeModeSessionProvider>> {
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.
codex-rs/core/src/session/mod.rs
+3 -2
View File
@@ -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,
codex-rs/core/src/session/session.rs
+7 -6
View File
@@ -497,7 +497,7 @@ impl Session {
skills_manager: Arc<SkillsManager>,
plugins_manager: Arc<PluginsManager>,
mcp_manager: Arc<McpManager>,
code_mode_session_provider: Option<Arc<dyn codex_code_mode::CodeModeSessionProvider>>,
code_mode_session_provider_selection: codex_code_mode::SessionProviderSelection,
extensions: Arc<codex_extension_api::ExtensionRegistry<crate::config::Config>>,
agent_control: AgentControl,
environment_manager: Arc<EnvironmentManager>,
@@ -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
codex-rs/core/src/session/tests.rs
+3 -3
View File
@@ -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()),
codex-rs/core/src/thread_manager.rs
+12 -20
View File
@@ -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<dyn codex_code_mode::CodeModeSessionProvider>,
> = LazyLock::new(|| Arc::new(codex_code_mode::InProcessCodeModeSessionProvider));
type CapturedOps = Vec<(ThreadId, Op)>;
type SharedCapturedOps = Arc<std::sync::Mutex<CapturedOps>>;
@@ -180,7 +175,7 @@ pub struct StartThreadOptions {
pub initial_history: InitialHistory,
pub session_source: Option<SessionSource>,
pub thread_source: Option<ThreadSource>,
pub code_mode_session_provider: Option<Arc<dyn codex_code_mode::CodeModeSessionProvider>>,
pub code_mode_session_provider_selection: codex_code_mode::SessionProviderSelection,
pub dynamic_tools: Vec<codex_protocol::dynamic_tools::DynamicToolSpec>,
pub persist_extended_history: bool,
pub metrics_service_name: Option<String>,
@@ -409,11 +404,6 @@ impl ThreadManager {
self.state.session_source.clone()
}
pub fn in_process_code_mode_session_provider()
-> Arc<dyn codex_code_mode::CodeModeSessionProvider> {
Arc::clone(&IN_PROCESS_CODE_MODE_SESSION_PROVIDER)
}
pub fn auth_manager(&self) -> Arc<AuthManager> {
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<Arc<dyn codex_code_mode::CodeModeSessionProvider>>,
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)]
codex-rs/core/src/thread_manager_tests.rs
+41 -51
View File
@@ -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 <shell>"));
}
#[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,
codex-rs/core/src/tools/code_mode/mod.rs
+6 -7
View File
@@ -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<Arc<dyn CodeModeSessionProvider>>,
pub(crate) async fn from_provider_selection(
provider_selection: SessionProviderSelection,
) -> Result<Self, String> {
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,
codex-rs/core/tests/suite/subagent_notifications.rs
+15 -18
View File
@@ -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());
codex-rs/memories/write/src/runtime.rs
+1 -3
View File
@@ -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,