mirror of
https://github.com/openai/codex.git
synced 2026-11-08 20:40:33 +00:00
471 lines
16 KiB
Rust
471 lines
16 KiB
Rust
use crate::error::CodexErr;
|
|
use crate::error::Result;
|
|
use codex_protocol::ThreadId;
|
|
use codex_protocol::protocol::SessionSource;
|
|
use codex_protocol::protocol::SubAgentSource;
|
|
use rand::prelude::IndexedRandom;
|
|
use std::collections::HashMap;
|
|
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 {
|
|
active_agents: Mutex<ActiveAgents>,
|
|
total_count: AtomicUsize,
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct ActiveAgents {
|
|
threads_set: HashSet<ThreadId>,
|
|
thread_agent_nicknames: HashMap<ThreadId, String>,
|
|
used_agent_nicknames: HashSet<String>,
|
|
nickname_reset_count: usize,
|
|
}
|
|
|
|
fn format_agent_nickname(name: &str, nickname_reset_count: usize) -> String {
|
|
match nickname_reset_count {
|
|
0 => name.to_string(),
|
|
reset_count => {
|
|
let value = reset_count + 1;
|
|
let suffix = match value % 100 {
|
|
11..=13 => "th",
|
|
_ => match value % 10 {
|
|
1 => "st", // codespell:ignore
|
|
2 => "nd", // codespell:ignore
|
|
3 => "rd", // codespell:ignore
|
|
_ => "th", // codespell:ignore
|
|
},
|
|
};
|
|
format!("{name} the {value}{suffix}")
|
|
}
|
|
}
|
|
}
|
|
|
|
fn session_depth(session_source: &SessionSource) -> i32 {
|
|
match session_source {
|
|
SessionSource::SubAgent(SubAgentSource::ThreadSpawn { depth, .. }) => *depth,
|
|
SessionSource::SubAgent(_) => 0,
|
|
_ => 0,
|
|
}
|
|
}
|
|
|
|
pub(crate) fn next_thread_spawn_depth(session_source: &SessionSource) -> i32 {
|
|
session_depth(session_source).saturating_add(1)
|
|
}
|
|
|
|
pub(crate) fn exceeds_thread_spawn_depth_limit(depth: i32, max_depth: i32) -> bool {
|
|
depth > max_depth
|
|
}
|
|
|
|
impl Guards {
|
|
pub(crate) fn reserve_spawn_slot(
|
|
self: &Arc<Self>,
|
|
max_threads: Option<usize>,
|
|
) -> Result<SpawnReservation> {
|
|
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,
|
|
reserved_agent_nickname: None,
|
|
})
|
|
}
|
|
|
|
pub(crate) fn release_spawned_thread(&self, thread_id: ThreadId) {
|
|
let removed = {
|
|
let mut active_agents = self
|
|
.active_agents
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
let removed = active_agents.threads_set.remove(&thread_id);
|
|
active_agents.thread_agent_nicknames.remove(&thread_id);
|
|
removed
|
|
};
|
|
if removed {
|
|
self.total_count.fetch_sub(1, Ordering::AcqRel);
|
|
}
|
|
}
|
|
|
|
fn register_spawned_thread(&self, thread_id: ThreadId, agent_nickname: Option<String>) {
|
|
let mut active_agents = self
|
|
.active_agents
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
active_agents.threads_set.insert(thread_id);
|
|
if let Some(agent_nickname) = agent_nickname {
|
|
active_agents
|
|
.used_agent_nicknames
|
|
.insert(agent_nickname.clone());
|
|
active_agents
|
|
.thread_agent_nicknames
|
|
.insert(thread_id, agent_nickname);
|
|
}
|
|
}
|
|
|
|
fn reserve_agent_nickname(&self, names: &[&str], preferred: Option<&str>) -> Option<String> {
|
|
let mut active_agents = self
|
|
.active_agents
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
let agent_nickname = if let Some(preferred) = preferred {
|
|
preferred.to_string()
|
|
} else {
|
|
if names.is_empty() {
|
|
return None;
|
|
}
|
|
let available_names: Vec<String> = names
|
|
.iter()
|
|
.map(|name| format_agent_nickname(name, active_agents.nickname_reset_count))
|
|
.filter(|name| !active_agents.used_agent_nicknames.contains(name))
|
|
.collect();
|
|
if let Some(name) = available_names.choose(&mut rand::rng()) {
|
|
name.clone()
|
|
} else {
|
|
active_agents.used_agent_nicknames.clear();
|
|
active_agents.nickname_reset_count += 1;
|
|
if let Some(metrics) = codex_otel::metrics::global() {
|
|
let _ = metrics.counter("codex.multi_agent.nickname_pool_reset", 1, &[]);
|
|
}
|
|
format_agent_nickname(
|
|
names.choose(&mut rand::rng())?,
|
|
active_agents.nickname_reset_count,
|
|
)
|
|
}
|
|
};
|
|
active_agents
|
|
.used_agent_nicknames
|
|
.insert(agent_nickname.clone());
|
|
Some(agent_nickname)
|
|
}
|
|
|
|
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<Guards>,
|
|
active: bool,
|
|
reserved_agent_nickname: Option<String>,
|
|
}
|
|
|
|
impl SpawnReservation {
|
|
pub(crate) fn reserve_agent_nickname(&mut self, names: &[&str]) -> Result<String> {
|
|
self.reserve_agent_nickname_with_preference(names, None)
|
|
}
|
|
|
|
pub(crate) fn reserve_agent_nickname_with_preference(
|
|
&mut self,
|
|
names: &[&str],
|
|
preferred: Option<&str>,
|
|
) -> Result<String> {
|
|
let agent_nickname = self
|
|
.state
|
|
.reserve_agent_nickname(names, preferred)
|
|
.ok_or_else(|| {
|
|
CodexErr::UnsupportedOperation("no available agent nicknames".to_string())
|
|
})?;
|
|
self.reserved_agent_nickname = Some(agent_nickname.clone());
|
|
Ok(agent_nickname)
|
|
}
|
|
|
|
pub(crate) fn commit(self, thread_id: ThreadId) {
|
|
self.commit_with_agent_nickname(thread_id, None);
|
|
}
|
|
|
|
pub(crate) fn commit_with_agent_nickname(
|
|
mut self,
|
|
thread_id: ThreadId,
|
|
agent_nickname: Option<String>,
|
|
) {
|
|
let agent_nickname = self.reserved_agent_nickname.take().or(agent_nickname);
|
|
self.state
|
|
.register_spawned_thread(thread_id, agent_nickname);
|
|
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;
|
|
use std::collections::HashSet;
|
|
|
|
#[test]
|
|
fn format_agent_nickname_adds_ordinals_after_reset() {
|
|
assert_eq!(format_agent_nickname("Plato", 0), "Plato");
|
|
assert_eq!(format_agent_nickname("Plato", 1), "Plato the 2nd");
|
|
assert_eq!(format_agent_nickname("Plato", 2), "Plato the 3rd");
|
|
assert_eq!(format_agent_nickname("Plato", 10), "Plato the 11th");
|
|
assert_eq!(format_agent_nickname("Plato", 20), "Plato the 21st");
|
|
}
|
|
|
|
#[test]
|
|
fn session_depth_defaults_to_zero_for_root_sources() {
|
|
assert_eq!(session_depth(&SessionSource::Cli), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn thread_spawn_depth_increments_and_enforces_limit() {
|
|
let session_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
|
|
parent_thread_id: ThreadId::new(),
|
|
depth: 1,
|
|
agent_nickname: None,
|
|
agent_role: None,
|
|
});
|
|
let child_depth = next_thread_spawn_depth(&session_source);
|
|
assert_eq!(child_depth, 2);
|
|
assert!(exceeds_thread_spawn_depth_limit(child_depth, 1));
|
|
}
|
|
|
|
#[test]
|
|
fn non_thread_spawn_subagents_default_to_depth_zero() {
|
|
let session_source = SessionSource::SubAgent(SubAgentSource::Review);
|
|
assert_eq!(session_depth(&session_source), 0);
|
|
assert_eq!(next_thread_spawn_depth(&session_source), 1);
|
|
assert!(!exceeds_thread_spawn_depth_limit(1, 1));
|
|
}
|
|
|
|
#[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);
|
|
}
|
|
|
|
#[test]
|
|
fn failed_spawn_keeps_nickname_marked_used() {
|
|
let guards = Arc::new(Guards::default());
|
|
let mut reservation = guards.reserve_spawn_slot(None).expect("reserve slot");
|
|
let agent_nickname = reservation
|
|
.reserve_agent_nickname(&["alpha"])
|
|
.expect("reserve agent name");
|
|
assert_eq!(agent_nickname, "alpha");
|
|
drop(reservation);
|
|
|
|
let mut reservation = guards.reserve_spawn_slot(None).expect("reserve slot");
|
|
let agent_nickname = reservation
|
|
.reserve_agent_nickname(&["alpha", "beta"])
|
|
.expect("unused name should still be preferred");
|
|
assert_eq!(agent_nickname, "beta");
|
|
}
|
|
|
|
#[test]
|
|
fn agent_nickname_resets_used_pool_when_exhausted() {
|
|
let guards = Arc::new(Guards::default());
|
|
let mut first = guards.reserve_spawn_slot(None).expect("reserve first slot");
|
|
let first_name = first
|
|
.reserve_agent_nickname(&["alpha"])
|
|
.expect("reserve first agent name");
|
|
let first_id = ThreadId::new();
|
|
first.commit(first_id);
|
|
assert_eq!(first_name, "alpha");
|
|
|
|
let mut second = guards
|
|
.reserve_spawn_slot(None)
|
|
.expect("reserve second slot");
|
|
let second_name = second
|
|
.reserve_agent_nickname(&["alpha"])
|
|
.expect("name should be reused after pool reset");
|
|
assert_eq!(second_name, "alpha the 2nd");
|
|
let active_agents = guards
|
|
.active_agents
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
assert_eq!(active_agents.nickname_reset_count, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn released_nickname_stays_used_until_pool_reset() {
|
|
let guards = Arc::new(Guards::default());
|
|
|
|
let mut first = guards.reserve_spawn_slot(None).expect("reserve first slot");
|
|
let first_name = first
|
|
.reserve_agent_nickname(&["alpha"])
|
|
.expect("reserve first agent name");
|
|
let first_id = ThreadId::new();
|
|
first.commit(first_id);
|
|
assert_eq!(first_name, "alpha");
|
|
|
|
guards.release_spawned_thread(first_id);
|
|
|
|
let mut second = guards
|
|
.reserve_spawn_slot(None)
|
|
.expect("reserve second slot");
|
|
let second_name = second
|
|
.reserve_agent_nickname(&["alpha", "beta"])
|
|
.expect("released name should still be marked used");
|
|
assert_eq!(second_name, "beta");
|
|
let second_id = ThreadId::new();
|
|
second.commit(second_id);
|
|
guards.release_spawned_thread(second_id);
|
|
|
|
let mut third = guards.reserve_spawn_slot(None).expect("reserve third slot");
|
|
let third_name = third
|
|
.reserve_agent_nickname(&["alpha", "beta"])
|
|
.expect("pool reset should permit a duplicate");
|
|
let expected_names =
|
|
HashSet::from(["alpha the 2nd".to_string(), "beta the 2nd".to_string()]);
|
|
assert!(expected_names.contains(&third_name));
|
|
let active_agents = guards
|
|
.active_agents
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
assert_eq!(active_agents.nickname_reset_count, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn repeated_resets_advance_the_ordinal_suffix() {
|
|
let guards = Arc::new(Guards::default());
|
|
|
|
let mut first = guards.reserve_spawn_slot(None).expect("reserve first slot");
|
|
let first_name = first
|
|
.reserve_agent_nickname(&["Plato"])
|
|
.expect("reserve first agent name");
|
|
let first_id = ThreadId::new();
|
|
first.commit(first_id);
|
|
assert_eq!(first_name, "Plato");
|
|
guards.release_spawned_thread(first_id);
|
|
|
|
let mut second = guards
|
|
.reserve_spawn_slot(None)
|
|
.expect("reserve second slot");
|
|
let second_name = second
|
|
.reserve_agent_nickname(&["Plato"])
|
|
.expect("reserve second agent name");
|
|
let second_id = ThreadId::new();
|
|
second.commit(second_id);
|
|
assert_eq!(second_name, "Plato the 2nd");
|
|
guards.release_spawned_thread(second_id);
|
|
|
|
let mut third = guards.reserve_spawn_slot(None).expect("reserve third slot");
|
|
let third_name = third
|
|
.reserve_agent_nickname(&["Plato"])
|
|
.expect("reserve third agent name");
|
|
assert_eq!(third_name, "Plato the 3rd");
|
|
let active_agents = guards
|
|
.active_agents
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
assert_eq!(active_agents.nickname_reset_count, 2);
|
|
}
|
|
}
|