File size: 3,014 Bytes
52a9af3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | use super::AgentControl;
use crate::codex_thread::CodexThread;
use codex_protocol::error::CodexErr;
use codex_protocol::error::CodexErrorDetails;
use codex_protocol::error::Result as CodexResult;
use codex_protocol::protocol::MultiAgentVersion;
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_turn_start(
&self,
thread: &CodexThread,
) -> CodexResult<()> {
if thread.session.active_turn.lock().await.is_some() {
return Ok(());
}
let config = thread.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::new(CodexErrorDetails::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 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;
|