From f1c961d5f7a00033c5c668a8a6ddbe96d2004762 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 21 Jan 2026 09:39:11 +0000 Subject: [PATCH] feat: max threads config (#9483) # External (non-OpenAI) Pull Request Requirements Before opening this Pull Request, please read the dedicated "Contributing" markdown file or your PR may be closed: https://github.com/openai/codex/blob/main/docs/contributing.md If your PR conforms to our contribution guidelines, replace this text with a detailed and high quality description of your changes. Include a link to a bug report or enhancement request. --- codex-rs/core/config.schema.json | 20 +++ codex-rs/core/src/agent/control.rs | 139 +++++++++++++++++++- codex-rs/core/src/agent/guards.rs | 193 ++++++++++++++++++++++++++++ codex-rs/core/src/agent/mod.rs | 2 + codex-rs/core/src/config/mod.rs | 33 +++++ codex-rs/core/src/error.rs | 10 +- codex-rs/core/src/thread_manager.rs | 3 +- 7 files changed, 394 insertions(+), 6 deletions(-) create mode 100644 codex-rs/core/src/agent/guards.rs diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index aea80a91d..e9ec0d7e6 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -4,6 +4,14 @@ "description": "Base config deserialized from ~/.codex/config.toml.", "type": "object", "properties": { + "agents": { + "description": "Agent-related settings (thread limits, etc.).", + "allOf": [ + { + "$ref": "#/definitions/AgentsToml" + } + ] + }, "analytics": { "description": "When `false`, disables analytics across Codex product surfaces in this machine. Defaults to `true`.", "allOf": [ @@ -436,6 +444,18 @@ "description": "A path that is guaranteed to be absolute and normalized (though it is not guaranteed to be canonicalized or exist on the filesystem).\n\nIMPORTANT: When deserializing an `AbsolutePathBuf`, a base path must be set using [AbsolutePathBufGuard::new]. If no base path is set, the deserialization will fail unless the path being deserialized is already absolute.", "type": "string" }, + "AgentsToml": { + "type": "object", + "properties": { + "max_threads": { + "description": "Maximum number of agent threads that can be open concurrently. When unset, no limit is enforced.", + "type": "integer", + "format": "uint", + "minimum": 1.0 + } + }, + "additionalProperties": false + }, "AltScreenMode": { "description": "Controls whether the TUI uses the terminal's alternate screen buffer.\n\n**Background:** The alternate screen buffer provides a cleaner fullscreen experience without polluting the terminal's scrollback history. However, it conflicts with terminal multiplexers like Zellij that strictly follow the xterm specification, which defines that alternate screen buffers should not have scrollback.\n\n**Zellij's behavior:** Zellij intentionally disables scrollback in alternate screen mode (see https://github.com/zellij-org/zellij/pull/1032) to comply with the xterm spec. This is by design and not configurable in Zellij—there is no option to enable scrollback in alternate screen mode.\n\n**Solution:** This setting provides a pragmatic workaround: - `auto` (default): Automatically detect the terminal multiplexer. If running in Zellij, disable alternate screen to preserve scrollback. Enable it everywhere else. - `always`: Always use alternate screen mode (original behavior before this fix). - `never`: Never use alternate screen mode. Runs in inline mode, preserving scrollback in all multiplexers.\n\nThe CLI flag `--no-alt-screen` can override this setting at runtime.", "oneOf": [ diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 4467494fc..6c5ebf9f0 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -1,4 +1,5 @@ use crate::agent::AgentStatus; +use crate::agent::guards::Guards; use crate::error::CodexErr; use crate::error::Result as CodexResult; use crate::thread_manager::ThreadManagerState; @@ -12,18 +13,25 @@ use tokio::sync::watch; /// Control-plane handle for multi-agent operations. /// `AgentControl` is held by each session (via `SessionServices`). It provides capability to /// spawn new agents and the inter-agent communication layer. +/// An `AgentControl` instance is shared per "user session" which means the same `AgentControl` +/// is used for every sub-agent spawned by Codex. By doing so, we make sure the guards are +/// scoped to a user session. #[derive(Clone, Default)] pub(crate) struct AgentControl { /// Weak handle back to the global thread registry/state. /// This is `Weak` to avoid reference cycles and shadow persistence of the form /// `ThreadManagerState -> CodexThread -> Session -> SessionServices -> ThreadManagerState`. manager: Weak, + state: Arc, } impl AgentControl { /// Construct a new `AgentControl` that can spawn/message agents via the given manager state. pub(crate) fn new(manager: Weak) -> Self { - Self { manager } + Self { + manager, + ..Default::default() + } } /// Spawn a new agent thread and submit the initial prompt. @@ -33,7 +41,11 @@ impl AgentControl { prompt: String, ) -> CodexResult { let state = self.upgrade()?; + let reservation = self.state.reserve_spawn_slot(config.agent_max_threads)?; + + // The same `AgentControl` is sent to spawn the thread. let new_thread = state.spawn_new_thread(config, self.clone()).await?; + reservation.commit(new_thread.thread_id); // Notify a new thread has been created. This notification will be processed by clients // to subscribe or drain this newly created thread. @@ -67,6 +79,7 @@ impl AgentControl { .await; if matches!(result, Err(CodexErr::InternalAgentDied)) { let _ = state.remove_thread(&agent_id).await; + self.state.release_spawned_thread(agent_id); } result } @@ -82,6 +95,7 @@ impl AgentControl { let state = self.upgrade()?; let result = state.send_op(agent_id, Op::Shutdown {}).await; let _ = state.remove_thread(&agent_id).await; + self.state.release_spawned_thread(agent_id); result } @@ -132,17 +146,25 @@ mod tests { use codex_protocol::protocol::TurnStartedEvent; use pretty_assertions::assert_eq; use tempfile::TempDir; + use toml::Value as TomlValue; - async fn test_config() -> (TempDir, Config) { + async fn test_config_with_cli_overrides( + cli_overrides: Vec<(String, TomlValue)>, + ) -> (TempDir, Config) { let home = TempDir::new().expect("create temp dir"); let config = ConfigBuilder::default() .codex_home(home.path().to_path_buf()) + .cli_overrides(cli_overrides) .build() .await .expect("load default test config"); (home, config) } + async fn test_config() -> (TempDir, Config) { + test_config_with_cli_overrides(Vec::new()).await + } + struct AgentControlHarness { _home: TempDir, config: Config, @@ -373,4 +395,117 @@ mod tests { .find(|entry| *entry == expected); assert_eq!(captured, Some(expected)); } + + #[tokio::test] + async fn spawn_agent_respects_max_threads_limit() { + let max_threads = 1usize; + let (_home, config) = test_config_with_cli_overrides(vec![( + "agents.max_threads".to_string(), + TomlValue::Integer(max_threads as i64), + )]) + .await; + let manager = ThreadManager::with_models_provider_and_home( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.clone(), + ); + let control = manager.agent_control(); + + let _ = manager + .start_thread(config.clone()) + .await + .expect("start thread"); + + let first_agent_id = control + .spawn_agent(config.clone(), "hello".to_string()) + .await + .expect("spawn_agent should succeed"); + + let err = control + .spawn_agent(config, "hello again".to_string()) + .await + .expect_err("spawn_agent should respect max threads"); + let CodexErr::AgentLimitReached { + max_threads: seen_max_threads, + } = err + else { + panic!("expected CodexErr::AgentLimitReached"); + }; + assert_eq!(seen_max_threads, max_threads); + + let _ = control + .shutdown_agent(first_agent_id) + .await + .expect("shutdown agent"); + } + + #[tokio::test] + async fn spawn_agent_releases_slot_after_shutdown() { + let max_threads = 1usize; + let (_home, config) = test_config_with_cli_overrides(vec![( + "agents.max_threads".to_string(), + TomlValue::Integer(max_threads as i64), + )]) + .await; + let manager = ThreadManager::with_models_provider_and_home( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.clone(), + ); + let control = manager.agent_control(); + + let first_agent_id = control + .spawn_agent(config.clone(), "hello".to_string()) + .await + .expect("spawn_agent should succeed"); + let _ = control + .shutdown_agent(first_agent_id) + .await + .expect("shutdown agent"); + + let second_agent_id = control + .spawn_agent(config.clone(), "hello again".to_string()) + .await + .expect("spawn_agent should succeed after shutdown"); + let _ = control + .shutdown_agent(second_agent_id) + .await + .expect("shutdown agent"); + } + + #[tokio::test] + async fn spawn_agent_limit_shared_across_clones() { + let max_threads = 1usize; + let (_home, config) = test_config_with_cli_overrides(vec![( + "agents.max_threads".to_string(), + TomlValue::Integer(max_threads as i64), + )]) + .await; + let manager = ThreadManager::with_models_provider_and_home( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.clone(), + ); + let control = manager.agent_control(); + let cloned = control.clone(); + + let first_agent_id = cloned + .spawn_agent(config.clone(), "hello".to_string()) + .await + .expect("spawn_agent should succeed"); + + let err = control + .spawn_agent(config, "hello again".to_string()) + .await + .expect_err("spawn_agent should respect shared guard"); + let CodexErr::AgentLimitReached { max_threads } = err else { + panic!("expected CodexErr::AgentLimitReached"); + }; + assert_eq!(max_threads, 1); + + let _ = control + .shutdown_agent(first_agent_id) + .await + .expect("shutdown agent"); + } } diff --git a/codex-rs/core/src/agent/guards.rs b/codex-rs/core/src/agent/guards.rs new file mode 100644 index 000000000..c384ed7cd --- /dev/null +++ b/codex-rs/core/src/agent/guards.rs @@ -0,0 +1,193 @@ +use crate::error::CodexErr; +use crate::error::Result; +use codex_protocol::ThreadId; +use std::collections::HashSet; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +/// This structure is used to add some limits on the multi-agent capabilities for Codex. In +/// the current implementation, it limits: +/// * Total number of sub-agents (i.e. threads) per user session +/// +/// This structure is shared by all agents in the same user session (because the `AgentControl` +/// is). +#[derive(Default)] +pub(crate) struct Guards { + threads_set: Mutex>, + total_count: AtomicUsize, +} + +impl Guards { + pub(crate) fn reserve_spawn_slot( + self: &Arc, + max_threads: Option, + ) -> Result { + if let Some(max_threads) = max_threads { + if !self.try_increment_spawned(max_threads) { + return Err(CodexErr::AgentLimitReached { max_threads }); + } + } else { + self.total_count.fetch_add(1, Ordering::AcqRel); + } + Ok(SpawnReservation { + state: Arc::clone(self), + active: true, + }) + } + + pub(crate) fn release_spawned_thread(&self, thread_id: ThreadId) { + let removed = { + let mut threads = self + .threads_set + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + threads.remove(&thread_id) + }; + if removed { + self.total_count.fetch_sub(1, Ordering::AcqRel); + } + } + + fn register_spawned_thread(&self, thread_id: ThreadId) { + let mut threads = self + .threads_set + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + threads.insert(thread_id); + } + + fn try_increment_spawned(&self, max_threads: usize) -> bool { + let mut current = self.total_count.load(Ordering::Acquire); + loop { + if current >= max_threads { + return false; + } + match self.total_count.compare_exchange_weak( + current, + current + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return true, + Err(updated) => current = updated, + } + } + } +} + +pub(crate) struct SpawnReservation { + state: Arc, + active: bool, +} + +impl SpawnReservation { + pub(crate) fn commit(mut self, thread_id: ThreadId) { + self.state.register_spawned_thread(thread_id); + self.active = false; + } +} + +impl Drop for SpawnReservation { + fn drop(&mut self) { + if self.active { + self.state.total_count.fetch_sub(1, Ordering::AcqRel); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + + #[test] + fn reservation_drop_releases_slot() { + let guards = Arc::new(Guards::default()); + let reservation = guards.reserve_spawn_slot(Some(1)).expect("reserve slot"); + drop(reservation); + + let reservation = guards.reserve_spawn_slot(Some(1)).expect("slot released"); + drop(reservation); + } + + #[test] + fn commit_holds_slot_until_release() { + let guards = Arc::new(Guards::default()); + let reservation = guards.reserve_spawn_slot(Some(1)).expect("reserve slot"); + let thread_id = ThreadId::new(); + reservation.commit(thread_id); + + let err = match guards.reserve_spawn_slot(Some(1)) { + Ok(_) => panic!("limit should be enforced"), + Err(err) => err, + }; + let CodexErr::AgentLimitReached { max_threads } = err else { + panic!("expected CodexErr::AgentLimitReached"); + }; + assert_eq!(max_threads, 1); + + guards.release_spawned_thread(thread_id); + let reservation = guards + .reserve_spawn_slot(Some(1)) + .expect("slot released after thread removal"); + drop(reservation); + } + + #[test] + fn release_ignores_unknown_thread_id() { + let guards = Arc::new(Guards::default()); + let reservation = guards.reserve_spawn_slot(Some(1)).expect("reserve slot"); + let thread_id = ThreadId::new(); + reservation.commit(thread_id); + + guards.release_spawned_thread(ThreadId::new()); + + let err = match guards.reserve_spawn_slot(Some(1)) { + Ok(_) => panic!("limit should still be enforced"), + Err(err) => err, + }; + let CodexErr::AgentLimitReached { max_threads } = err else { + panic!("expected CodexErr::AgentLimitReached"); + }; + assert_eq!(max_threads, 1); + + guards.release_spawned_thread(thread_id); + let reservation = guards + .reserve_spawn_slot(Some(1)) + .expect("slot released after real thread removal"); + drop(reservation); + } + + #[test] + fn release_is_idempotent_for_registered_threads() { + let guards = Arc::new(Guards::default()); + let reservation = guards.reserve_spawn_slot(Some(1)).expect("reserve slot"); + let first_id = ThreadId::new(); + reservation.commit(first_id); + + guards.release_spawned_thread(first_id); + + let reservation = guards.reserve_spawn_slot(Some(1)).expect("slot reused"); + let second_id = ThreadId::new(); + reservation.commit(second_id); + + guards.release_spawned_thread(first_id); + + let err = match guards.reserve_spawn_slot(Some(1)) { + Ok(_) => panic!("limit should still be enforced"), + Err(err) => err, + }; + let CodexErr::AgentLimitReached { max_threads } = err else { + panic!("expected CodexErr::AgentLimitReached"); + }; + assert_eq!(max_threads, 1); + + guards.release_spawned_thread(second_id); + let reservation = guards + .reserve_spawn_slot(Some(1)) + .expect("slot released after second thread removal"); + drop(reservation); + } +} diff --git a/codex-rs/core/src/agent/mod.rs b/codex-rs/core/src/agent/mod.rs index 0387eda49..180f70dbe 100644 --- a/codex-rs/core/src/agent/mod.rs +++ b/codex-rs/core/src/agent/mod.rs @@ -1,4 +1,6 @@ pub(crate) mod control; +// Do not put in `pub` or `pub(crate)`. This code should not be used somewhere else. +mod guards; pub(crate) mod role; pub(crate) mod status; diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index b65a20b79..1e409559e 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -89,6 +89,7 @@ pub use codex_git::GhostSnapshotConfig; /// files are *silently truncated* to this size so we do not take up too much of /// the context window. pub(crate) const PROJECT_DOC_MAX_BYTES: usize = 32 * 1024; // 32 KiB +pub(crate) const DEFAULT_AGENT_MAX_THREADS: Option = None; pub const CONFIG_TOML_FILE: &str = "config.toml"; @@ -299,6 +300,9 @@ pub struct Config { /// Token budget applied when storing tool/function outputs in the context manager. pub tool_output_token_limit: Option, + /// Maximum number of agent threads that can be open concurrently. + pub agent_max_threads: Option, + /// Directory containing all Codex state (defaults to `~/.codex` but can be /// overridden by the `CODEX_HOME` environment variable). pub codex_home: PathBuf, @@ -924,6 +928,9 @@ pub struct ConfigToml { /// Nested tools section for feature toggles pub tools: Option, + /// Agent-related settings (thread limits, etc.). + pub agents: Option, + /// User-level skill config entries keyed by SKILL.md path. pub skills: Option, @@ -1033,6 +1040,15 @@ pub struct ToolsToml { pub view_image: Option, } +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] +#[schemars(deny_unknown_fields)] +pub struct AgentsToml { + /// Maximum number of agent threads that can be open concurrently. + /// When unset, no limit is enforced. + #[schemars(range(min = 1))] + pub max_threads: Option, +} + impl From for Tools { fn from(tools_toml: ToolsToml) -> Self { Self { @@ -1392,6 +1408,18 @@ impl Config { let history = cfg.history.unwrap_or_default(); + let agent_max_threads = cfg + .agents + .as_ref() + .and_then(|agents| agents.max_threads) + .or(DEFAULT_AGENT_MAX_THREADS); + if agent_max_threads == Some(0) { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "agents.max_threads must be at least 1", + )); + } + let ghost_snapshot = { let mut config = GhostSnapshotConfig::default(); if let Some(ghost_snapshot) = cfg.ghost_snapshot.as_ref() @@ -1530,6 +1558,7 @@ impl Config { }) .collect(), tool_output_token_limit: cfg.tool_output_token_limit, + agent_max_threads, codex_home, config_layer_stack, history, @@ -3718,6 +3747,7 @@ model_verbosity = "high" project_doc_max_bytes: PROJECT_DOC_MAX_BYTES, project_doc_fallback_filenames: Vec::new(), tool_output_token_limit: None, + agent_max_threads: None, codex_home: fixture.codex_home(), config_layer_stack: Default::default(), history: History::default(), @@ -3806,6 +3836,7 @@ model_verbosity = "high" project_doc_max_bytes: PROJECT_DOC_MAX_BYTES, project_doc_fallback_filenames: Vec::new(), tool_output_token_limit: None, + agent_max_threads: None, codex_home: fixture.codex_home(), config_layer_stack: Default::default(), history: History::default(), @@ -3909,6 +3940,7 @@ model_verbosity = "high" project_doc_max_bytes: PROJECT_DOC_MAX_BYTES, project_doc_fallback_filenames: Vec::new(), tool_output_token_limit: None, + agent_max_threads: None, codex_home: fixture.codex_home(), config_layer_stack: Default::default(), history: History::default(), @@ -3998,6 +4030,7 @@ model_verbosity = "high" project_doc_max_bytes: PROJECT_DOC_MAX_BYTES, project_doc_fallback_filenames: Vec::new(), tool_output_token_limit: None, + agent_max_threads: None, codex_home: fixture.codex_home(), config_layer_stack: Default::default(), history: History::default(), diff --git a/codex-rs/core/src/error.rs b/codex-rs/core/src/error.rs index 3b490436e..e9830f518 100644 --- a/codex-rs/core/src/error.rs +++ b/codex-rs/core/src/error.rs @@ -78,6 +78,9 @@ pub enum CodexErr { #[error("no thread with id: {0}")] ThreadNotFound(ThreadId), + #[error("agent thread limit reached (max {max_threads})")] + AgentLimitReached { max_threads: usize }, + #[error("session configured event was not the first event in the stream")] SessionConfiguredNotFirstEvent, @@ -199,6 +202,7 @@ impl CodexErr { | CodexErr::RetryLimit(_) | CodexErr::ContextWindowExceeded | CodexErr::ThreadNotFound(_) + | CodexErr::AgentLimitReached { .. } | CodexErr::Spawn | CodexErr::SessionConfiguredNotFirstEvent | CodexErr::UsageLimitReached(_) => false, @@ -497,9 +501,9 @@ impl CodexErr { CodexErr::SessionConfiguredNotFirstEvent | CodexErr::InternalServerError | CodexErr::InternalAgentDied => CodexErrorInfo::InternalServerError, - CodexErr::UnsupportedOperation(_) | CodexErr::ThreadNotFound(_) => { - CodexErrorInfo::BadRequest - } + CodexErr::UnsupportedOperation(_) + | CodexErr::ThreadNotFound(_) + | CodexErr::AgentLimitReached { .. } => CodexErrorInfo::BadRequest, CodexErr::Sandbox(_) => CodexErrorInfo::SandboxError, _ => CodexErrorInfo::Other, } diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 01dc99d91..6aa4d378d 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -366,7 +366,8 @@ impl ThreadManagerState { codex, session_configured.rollout_path.clone(), )); - self.threads.write().await.insert(thread_id, thread.clone()); + let mut threads = self.threads.write().await; + threads.insert(thread_id, thread.clone()); Ok(NewThread { thread_id,