core/src/agent/control/execution.rs
110 lines
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<usize>,}pub(crate) struct AgentExecutionGuard { limiter: Arc<AgentExecutionLimiter>,}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<AgentExecutionGuard> { 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<Self>) -> 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;