From 743f5aad38accd52da34bf4dcbdd1215a8c3ab9a Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 8 Jun 2026 14:21:28 +0200 Subject: [PATCH] feat: count V2 concurrency by active execution (#26969) ## Why Multi-Agent V2 concurrency should count active non-root turns, not resident or durable agent threads. The limit is intentionally best effort: admission checks are synchronous, but concurrent successful checks may overshoot slightly. ## What changed - Keep one root-derived execution limit on the shared `AgentControl`. - Count active V2 subagent turns with an RAII guard owned by `RunningTask`. - Check capacity before spawning or starting an idle agent, including direct app-server `turn/start` submissions. - Preserve queued delivery for agents that are already running. - Exempt automatic idle continuations so `/goal` work is not dropped when capacity is temporarily full. - Keep root and V1 turns outside this limiter. ## Test coverage - `execution_guards_count_active_v2_subagent_turns` - `execution_guards_ignore_root_and_v1_turns` - `v2_nested_spawn_checks_shared_active_execution_capacity` --- codex-rs/core/src/agent/control.rs | 33 +++-- codex-rs/core/src/agent/control/execution.rs | 110 ++++++++++++++++ .../core/src/agent/control/execution_tests.rs | 60 +++++++++ codex-rs/core/src/agent/control/spawn.rs | 5 +- codex-rs/core/src/codex_thread.rs | 6 + codex-rs/core/src/session/session.rs | 7 +- codex-rs/core/src/session/tests.rs | 2 +- codex-rs/core/src/state/turn.rs | 2 + codex-rs/core/src/tasks/mod.rs | 5 + codex-rs/core/tests/suite/agent_execution.rs | 122 ++++++++++++++++++ codex-rs/core/tests/suite/mod.rs | 1 + 11 files changed, 340 insertions(+), 13 deletions(-) create mode 100644 codex-rs/core/src/agent/control/execution.rs create mode 100644 codex-rs/core/src/agent/control/execution_tests.rs create mode 100644 codex-rs/core/tests/suite/agent_execution.rs diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index e9994883f..888f9cdf9 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -42,10 +42,13 @@ use std::sync::Weak; use tokio::sync::watch; use tracing::warn; +pub(crate) use self::execution::AgentExecutionGuard; +use self::execution::AgentExecutionLimiter; use self::residency::V2Residency; const ROOT_LAST_TASK_MESSAGE: &str = "Main thread"; +mod execution; mod legacy; mod residency; mod spawn; @@ -95,6 +98,7 @@ pub(crate) struct AgentControl { manager: Weak, state: Arc, v2_residency: Arc, + agent_execution_limiter: Arc, } impl AgentControl { @@ -106,8 +110,9 @@ impl AgentControl { } } - pub(crate) fn with_session_id(mut self, session_id: SessionId) -> Self { + pub(crate) fn with_session_id(mut self, session_id: SessionId, max_threads: usize) -> Self { self.session_id = session_id; + self.agent_execution_limiter.initialize(max_threads); self } @@ -120,6 +125,19 @@ impl AgentControl { &self, agent_id: ThreadId, initial_operation: Op, + ) -> CodexResult { + let state = self.upgrade()?; + self.ensure_execution_capacity_for_op(agent_id, &initial_operation) + .await?; + self.send_input_after_capacity_check(agent_id, &state, initial_operation) + .await + } + + async fn send_input_after_capacity_check( + &self, + agent_id: ThreadId, + state: &Arc, + initial_operation: Op, ) -> CodexResult { let last_task_message = match &initial_operation { Op::InterAgentCommunication { communication } => { @@ -127,11 +145,10 @@ impl AgentControl { } _ => non_empty_task_message(render_input_preview(&initial_operation)), }; - let state = self.upgrade()?; let result = self .handle_thread_request_result( agent_id, - &state, + state, state.send_op(agent_id, initial_operation).await, ) .await; @@ -153,14 +170,10 @@ impl AgentControl { ) -> CodexResult { let last_task_message = last_task_message_from_communication(&communication); let state = self.upgrade()?; + let op = Op::InterAgentCommunication { communication }; + self.ensure_execution_capacity_for_op(agent_id, &op).await?; let result = self - .handle_thread_request_result( - agent_id, - &state, - state - .send_op(agent_id, Op::InterAgentCommunication { communication }) - .await, - ) + .handle_thread_request_result(agent_id, &state, state.send_op(agent_id, op).await) .await; if result.is_ok() { match last_task_message { diff --git a/codex-rs/core/src/agent/control/execution.rs b/codex-rs/core/src/agent/control/execution.rs new file mode 100644 index 000000000..42aaa50e0 --- /dev/null +++ b/codex-rs/core/src/agent/control/execution.rs @@ -0,0 +1,110 @@ +use super::AgentControl; +use codex_protocol::ThreadId; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; +use codex_protocol::protocol::MultiAgentVersion; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::SessionSource; +use std::sync::Arc; +use std::sync::OnceLock; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +#[derive(Default)] +pub(super) struct AgentExecutionLimiter { + active: AtomicUsize, + max_threads: OnceLock, +} + +pub(crate) struct AgentExecutionGuard { + limiter: Arc, +} + +impl Drop for AgentExecutionGuard { + fn drop(&mut self) { + self.limiter.active.fetch_sub(1, Ordering::AcqRel); + } +} + +impl AgentControl { + pub(crate) async fn ensure_execution_capacity_for_op( + &self, + thread_id: ThreadId, + op: &Op, + ) -> CodexResult<()> { + if !op_starts_turn(op) { + return Ok(()); + } + let state = self.upgrade()?; + let thread = state.get_thread(thread_id).await?; + if thread.codex.session.active_turn.lock().await.is_some() { + return Ok(()); + } + let config = thread.codex.session.get_config().await; + let multi_agent_version = thread + .multi_agent_version() + .unwrap_or_else(|| config.multi_agent_version_from_features()); + self.ensure_execution_capacity(multi_agent_version, &thread.session_source) + } + + pub(crate) fn ensure_execution_capacity( + &self, + multi_agent_version: MultiAgentVersion, + session_source: &SessionSource, + ) -> CodexResult<()> { + if !is_execution_limited(multi_agent_version, session_source) { + return Ok(()); + } + let max_threads = self.agent_execution_limiter.max_threads(); + if self.agent_execution_limiter.has_capacity() { + Ok(()) + } else { + Err(CodexErr::AgentLimitReached { max_threads }) + } + } + + pub(crate) fn execution_guard( + &self, + multi_agent_version: MultiAgentVersion, + session_source: &SessionSource, + ) -> Option { + is_execution_limited(multi_agent_version, session_source) + .then(|| Arc::clone(&self.agent_execution_limiter).guard()) + } +} + +impl AgentExecutionLimiter { + pub(super) fn initialize(&self, max_threads: usize) { + self.max_threads.get_or_init(|| max_threads); + } + + fn max_threads(&self) -> usize { + self.max_threads.get().copied().unwrap_or(usize::MAX) + } + + fn has_capacity(&self) -> bool { + self.active.load(Ordering::Acquire) < self.max_threads() + } + + fn guard(self: Arc) -> AgentExecutionGuard { + self.active.fetch_add(1, Ordering::AcqRel); + AgentExecutionGuard { limiter: self } + } +} + +fn op_starts_turn(op: &Op) -> bool { + matches!(op, Op::UserInput { .. }) + || matches!(op, Op::InterAgentCommunication { communication } if communication.trigger_turn) +} + +fn is_execution_limited( + multi_agent_version: MultiAgentVersion, + session_source: &SessionSource, +) -> bool { + multi_agent_version == MultiAgentVersion::V2 + && matches!(session_source, SessionSource::SubAgent(_)) +} + +#[cfg(test)] +#[path = "execution_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/agent/control/execution_tests.rs b/codex-rs/core/src/agent/control/execution_tests.rs new file mode 100644 index 000000000..9cad08f95 --- /dev/null +++ b/codex-rs/core/src/agent/control/execution_tests.rs @@ -0,0 +1,60 @@ +use crate::agent::AgentControl; +use codex_protocol::error::CodexErr; +use codex_protocol::protocol::MultiAgentVersion; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; +use pretty_assertions::assert_eq; + +fn control_with_limit(max_threads: usize) -> AgentControl { + let control = AgentControl::default(); + control.agent_execution_limiter.initialize(max_threads); + control +} + +#[test] +fn execution_guards_count_active_v2_subagent_turns() { + let control = control_with_limit(/*max_threads*/ 1); + // Child role configs cannot replace the root-derived session limit. + control + .agent_execution_limiter + .initialize(/*max_threads*/ 2); + let source = SessionSource::SubAgent(SubAgentSource::Other("worker".to_string())); + + control + .ensure_execution_capacity(MultiAgentVersion::V2, &source) + .expect("first active turn should fit"); + let first = control + .execution_guard(MultiAgentVersion::V2, &source) + .expect("v2 subagent execution should be counted"); + let Err(err) = control.ensure_execution_capacity(MultiAgentVersion::V2, &source) else { + panic!("second active turn should exceed the derived non-root cap"); + }; + let CodexErr::AgentLimitReached { max_threads } = err else { + panic!("expected AgentLimitReached"); + }; + assert_eq!(max_threads, 1); + + drop(first); + control + .ensure_execution_capacity(MultiAgentVersion::V2, &source) + .expect("capacity should be released when the running task drops"); +} + +#[test] +fn execution_guards_ignore_root_and_v1_turns() { + let control = control_with_limit(/*max_threads*/ 0); + + assert!( + control + .execution_guard(MultiAgentVersion::V2, &SessionSource::Cli) + .is_none() + ); + assert!( + control + .execution_guard( + MultiAgentVersion::V1, + &SessionSource::SubAgent(SubAgentSource::Other("worker".to_string())), + ) + .is_none() + ); +} diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index e3536e89f..a103fbb80 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -207,6 +207,9 @@ impl AgentControl { &config, ) .await; + if let Some(session_source) = session_source.as_ref() { + self.ensure_execution_capacity(multi_agent_version, session_source)?; + } let agent_max_threads = config.effective_agent_max_threads(multi_agent_version); let spawn_uses_v2_residency = multi_agent_version == MultiAgentVersion::V2 && session_source @@ -349,7 +352,7 @@ impl AgentControl { ) .await; - self.send_input(new_thread.thread_id, initial_operation) + self.send_input_after_capacity_check(new_thread.thread_id, &state, initial_operation) .await?; if multi_agent_version != MultiAgentVersion::V2 { let child_reference = agent_metadata diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index b3f9cb0e2..1697f768e 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -234,6 +234,12 @@ impl CodexThread { trace: Option, client_user_message_id: Option, ) -> CodexResult { + self.codex + .session + .services + .agent_control + .ensure_execution_capacity_for_op(self.session_configured.thread_id, &op) + .await?; self.codex .submit_user_input_with_client_user_message_id(op, trace, client_user_message_id) .await diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 122ff6aeb..92f973731 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -948,7 +948,12 @@ impl Session { } else { SessionId::from(thread_id) }; - let agent_control = agent_control.with_session_id(session_id); + let agent_control = agent_control.with_session_id( + session_id, + config + .effective_agent_max_threads(MultiAgentVersion::V2) + .unwrap_or(usize::MAX), + ); let session_extension_data = codex_extension_api::ExtensionData::new(session_id.to_string()); let thread_extension_data = diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 366cf6f22..4d914ff98 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5184,7 +5184,7 @@ async fn resumed_subagent_session_keeps_inherited_session_id() { rollout_path: None, }), session_source, - AgentControl::default().with_session_id(parent_session_id), + AgentControl::default().with_session_id(parent_session_id, /*max_threads*/ usize::MAX), ) .await .expect("resume should succeed"); diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index eaf26ba85..d1e31b799 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -18,6 +18,7 @@ use codex_sandboxing::policy_transforms::merge_permission_profiles; use rmcp::model::RequestId; use tokio::sync::oneshot; +use crate::agent::control::AgentExecutionGuard; use crate::session::TurnInputQueue; use crate::session::turn_context::TurnContext; use crate::tasks::AnySessionTask; @@ -76,6 +77,7 @@ pub(crate) struct RunningTask { pub(crate) handle: AbortOnDropHandle<()>, pub(crate) turn_context: Arc, pub(crate) turn_extension_data: Arc, + pub(crate) _agent_execution_guard: Option, // Timer recorded when the task drops to capture the full turn duration. pub(crate) _timer: Option, } diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index a9a664a50..911f30b31 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -359,6 +359,10 @@ impl Session { let mut active = self.active_turn.lock().await; let turn = active.get_or_insert_with(ActiveTurn::default); debug_assert!(turn.task.is_none()); + let agent_execution_guard = self.services.agent_control.execution_guard( + turn_context.multi_agent_version, + &turn_context.session_source, + ); let done_clone = Arc::clone(&done); let session_ctx = Arc::new(SessionTaskContext::new( Arc::clone(self), @@ -430,6 +434,7 @@ impl Session { cancellation_token, turn_context: Arc::clone(&turn_context), turn_extension_data, + _agent_execution_guard: agent_execution_guard, _timer: timer, }; turn.task = Some(running_task); diff --git a/codex-rs/core/tests/suite/agent_execution.rs b/codex-rs/core/tests/suite/agent_execution.rs new file mode 100644 index 000000000..f073d7008 --- /dev/null +++ b/codex-rs/core/tests/suite/agent_execution.rs @@ -0,0 +1,122 @@ +use anyhow::Result; +use codex_features::Feature; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_function_call; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_once_match; +use core_test_support::responses::sse; +use core_test_support::responses::start_mock_server; +use core_test_support::test_codex::test_codex; +use pretty_assertions::assert_eq; +use serde_json::json; +use std::time::Duration; + +const FIRST_PROMPT: &str = "spawn the first worker"; +const FIRST_TASK: &str = "first worker task"; +const SECOND_TASK: &str = "second worker task"; + +fn body_contains(request: &wiremock::Request, text: &str) -> bool { + serde_json::from_slice::(&request.body) + .is_ok_and(|body| body.to_string().contains(text)) +} + +fn has_function_call_output(request: &wiremock::Request, call_id: &str) -> bool { + serde_json::from_slice::(&request.body).is_ok_and(|body| { + body.get("input") + .and_then(serde_json::Value::as_array) + .is_some_and(|items| { + items.iter().any(|item| { + item.get("type").and_then(serde_json::Value::as_str) + == Some("function_call_output") + && item.get("call_id").and_then(serde_json::Value::as_str) == Some(call_id) + }) + }) + }) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn v2_nested_spawn_checks_shared_active_execution_capacity() -> Result<()> { + let server = start_mock_server().await; + let first_args = serde_json::to_string(&json!({ + "message": FIRST_TASK, + "task_name": "first", + }))?; + mount_sse_once_match( + &server, + |request: &wiremock::Request| body_contains(request, FIRST_PROMPT), + sse(vec![ + ev_response_created("first-response"), + ev_function_call("first-call", "spawn_agent", &first_args), + ev_completed("first-response"), + ]), + ) + .await; + let second_args = serde_json::to_string(&json!({ + "message": SECOND_TASK, + "task_name": "second", + }))?; + mount_sse_once_match( + &server, + |request: &wiremock::Request| { + body_contains(request, FIRST_TASK) && !has_function_call_output(request, "first-call") + }, + sse(vec![ + ev_response_created("first-worker-response"), + ev_function_call("second-call", "spawn_agent", &second_args), + ev_completed("first-worker-response"), + ]), + ) + .await; + let second_followup = mount_sse_once_match( + &server, + |request: &wiremock::Request| has_function_call_output(request, "second-call"), + sse(vec![ + ev_response_created("second-followup-response"), + ev_assistant_message("second-followup-message", "blocked"), + ev_completed("second-followup-response"), + ]), + ) + .await; + mount_sse_once_match( + &server, + |request: &wiremock::Request| has_function_call_output(request, "first-call"), + sse(vec![ + ev_response_created("first-followup-response"), + ev_assistant_message("first-followup-message", "spawned"), + ev_completed("first-followup-response"), + ]), + ) + .await; + + let mut builder = test_codex().with_model("koffing").with_config(|config| { + config + .features + .enable(Feature::Collab) + .expect("test config should allow feature update"); + config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + config.multi_agent_v2.max_concurrent_threads_per_session = 2; + }); + let test = builder.build(&server).await?; + test.submit_turn(FIRST_PROMPT).await?; + + let second_output = tokio::time::timeout(Duration::from_secs(2), async { + loop { + if let Some(output) = second_followup.function_call_output_text("second-call") { + return output; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await?; + assert_eq!( + second_output, + "collab spawn failed: agent thread limit reached" + ); + assert_eq!(test.thread_manager.list_thread_ids().await.len(), 2); + + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 11246d7cc..234fb8a2f 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -30,6 +30,7 @@ pub static CODEX_ALIASES_TEMP_DIR: Option = { #[cfg(not(target_os = "windows"))] mod abort_tasks; mod additional_context; +mod agent_execution; mod agent_jobs; mod agent_websocket; mod agents_md;