Codex Handbook
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;