mirror of
https://github.com/openai/codex.git
synced 2026-11-05 20:10:33 +00:00
Merge remote-tracking branch 'origin/main' into dev/python-sdk-goal-turns
# Conflicts: # codex-rs/core/src/codex_thread.rs
This commit is contained in:
Ahmed Ibrahim
Generated
-1
@@ -1420,7 +1420,6 @@
|
||||
"rama-utils_0.3.0-alpha.4": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"ahash\",\"req\":\"^0.8\"},{\"name\":\"const_format\",\"req\":\"^0.2\"},{\"kind\":\"dev\",\"name\":\"loom\",\"req\":\"^0.7\",\"target\":\"cfg(loom)\"},{\"name\":\"parking_lot\",\"req\":\"^0.12\"},{\"name\":\"pin-project-lite\",\"req\":\"^0.2\"},{\"kind\":\"dev\",\"name\":\"quickcheck\",\"req\":\"^1.0\"},{\"name\":\"rama-macros\",\"req\":\"^0.3.0-alpha.4\"},{\"name\":\"regex\",\"req\":\"^1.12\"},{\"features\":[\"derive\"],\"name\":\"serde\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_test\",\"req\":\"^1\"},{\"features\":[\"write\",\"serde\",\"const_generics\",\"const_new\"],\"name\":\"smallvec\",\"req\":\"^1.15\"},{\"name\":\"smol_str\",\"req\":\"^0.3\"},{\"features\":[\"time\",\"macros\"],\"name\":\"tokio\",\"req\":\"^1.48\"},{\"kind\":\"dev\",\"name\":\"tokio-test\",\"req\":\"^0.4\"},{\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0\"},{\"name\":\"wildcard\",\"req\":\"^0.3\"}],\"features\":{}}",
|
||||
"rand_0.10.1": "{\"dependencies\":[{\"default_features\":false,\"features\":[\"rng\"],\"name\":\"chacha20\",\"optional\":true,\"req\":\"^0.10.0\"},{\"name\":\"getrandom\",\"optional\":true,\"req\":\"^0.4.0\"},{\"default_features\":false,\"features\":[\"alloc\"],\"kind\":\"dev\",\"name\":\"postcard\",\"req\":\"^1.1.3\"},{\"default_features\":false,\"name\":\"rand_core\",\"req\":\"^0.10.0\"},{\"kind\":\"dev\",\"name\":\"rand_pcg\",\"req\":\"^0.10\"},{\"kind\":\"dev\",\"name\":\"rayon\",\"req\":\"^1.7\"},{\"features\":[\"derive\"],\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0.103\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0.140\"}],\"features\":{\"alloc\":[],\"chacha\":[\"dep:chacha20\"],\"default\":[\"std\",\"std_rng\",\"sys_rng\",\"thread_rng\"],\"log\":[],\"serde\":[\"dep:serde\"],\"simd_support\":[],\"std\":[\"alloc\",\"getrandom?/std\"],\"std_rng\":[\"dep:chacha20\"],\"sys_rng\":[\"dep:getrandom\",\"getrandom/sys_rng\"],\"thread_rng\":[\"std\",\"std_rng\",\"sys_rng\"],\"unbiased\":[]}}",
|
||||
"rand_0.8.5": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"bincode\",\"req\":\"^1.2.1\"},{\"default_features\":false,\"name\":\"libc\",\"optional\":true,\"req\":\"^0.2.22\",\"target\":\"cfg(unix)\"},{\"name\":\"log\",\"optional\":true,\"req\":\"^0.4.4\"},{\"features\":[\"into_bits\"],\"name\":\"packed_simd\",\"optional\":true,\"package\":\"packed_simd_2\",\"req\":\"^0.3.7\"},{\"default_features\":false,\"name\":\"rand_chacha\",\"optional\":true,\"req\":\"^0.3.0\"},{\"name\":\"rand_core\",\"req\":\"^0.6.0\"},{\"kind\":\"dev\",\"name\":\"rand_pcg\",\"req\":\"^0.3.0\"},{\"features\":[\"derive\"],\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0.103\"}],\"features\":{\"alloc\":[\"rand_core/alloc\"],\"default\":[\"std\",\"std_rng\"],\"getrandom\":[\"rand_core/getrandom\"],\"min_const_gen\":[],\"nightly\":[],\"serde1\":[\"serde\",\"rand_core/serde1\"],\"simd_support\":[\"packed_simd\"],\"small_rng\":[],\"std\":[\"rand_core/std\",\"rand_chacha/std\",\"alloc\",\"getrandom\",\"libc\"],\"std_rng\":[\"rand_chacha\"]}}",
|
||||
"rand_0.8.6": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"bincode\",\"req\":\"^1.2.1\"},{\"default_features\":false,\"name\":\"libc\",\"optional\":true,\"req\":\"^0.2.22\",\"target\":\"cfg(unix)\"},{\"default_features\":false,\"name\":\"rand_chacha\",\"optional\":true,\"req\":\"^0.3.0\"},{\"name\":\"rand_core\",\"req\":\"^0.6.0\"},{\"kind\":\"dev\",\"name\":\"rand_pcg\",\"req\":\"^0.3.0\"},{\"features\":[\"derive\"],\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0.103\"}],\"features\":{\"alloc\":[\"rand_core/alloc\"],\"default\":[\"std\",\"std_rng\"],\"getrandom\":[\"rand_core/getrandom\"],\"log\":[],\"min_const_gen\":[],\"nightly\":[],\"serde1\":[\"serde\",\"rand_core/serde1\"],\"small_rng\":[],\"std\":[\"rand_core/std\",\"rand_chacha/std\",\"alloc\",\"getrandom\",\"libc\"],\"std_rng\":[\"rand_chacha\"]}}",
|
||||
"rand_0.9.3": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"bincode\",\"req\":\"^1.2.1\"},{\"default_features\":false,\"name\":\"rand_chacha\",\"optional\":true,\"req\":\"^0.9.0\"},{\"default_features\":false,\"name\":\"rand_core\",\"req\":\"^0.9.0\"},{\"kind\":\"dev\",\"name\":\"rand_pcg\",\"req\":\"^0.9.0\"},{\"kind\":\"dev\",\"name\":\"rayon\",\"req\":\"^1.7\"},{\"features\":[\"derive\"],\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0.103\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0.140\"}],\"features\":{\"alloc\":[],\"default\":[\"std\",\"std_rng\",\"os_rng\",\"small_rng\",\"thread_rng\"],\"log\":[],\"nightly\":[],\"os_rng\":[\"rand_core/os_rng\"],\"serde\":[\"dep:serde\",\"rand_core/serde\"],\"simd_support\":[],\"small_rng\":[],\"std\":[\"rand_core/std\",\"rand_chacha?/std\",\"alloc\"],\"std_rng\":[\"dep:rand_chacha\"],\"thread_rng\":[\"std\",\"std_rng\",\"os_rng\"],\"unbiased\":[]}}",
|
||||
"rand_chacha_0.3.1": "{\"dependencies\":[{\"default_features\":false,\"features\":[\"simd\"],\"name\":\"ppv-lite86\",\"req\":\"^0.2.8\"},{\"name\":\"rand_core\",\"req\":\"^0.6.0\"},{\"features\":[\"derive\"],\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"}],\"features\":{\"default\":[\"std\"],\"serde1\":[\"serde\"],\"simd\":[],\"std\":[\"ppv-lite86/std\"]}}",
|
||||
"rand_chacha_0.9.0": "{\"dependencies\":[{\"default_features\":false,\"features\":[\"simd\"],\"name\":\"ppv-lite86\",\"req\":\"^0.2.14\"},{\"name\":\"rand_core\",\"req\":\"^0.9.0\"},{\"features\":[\"os_rng\"],\"kind\":\"dev\",\"name\":\"rand_core\",\"req\":\"^0.9.0\"},{\"features\":[\"derive\"],\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"}],\"features\":{\"default\":[\"std\"],\"os_rng\":[\"rand_core/os_rng\"],\"serde\":[\"dep:serde\"],\"std\":[\"ppv-lite86/std\",\"rand_core/std\"]}}",
|
||||
|
||||
@@ -9,4 +9,5 @@ ignore = [
|
||||
"RUSTSEC-2025-0141", # bincode via syntect; remove when syntect drops or updates it
|
||||
"RUSTSEC-2026-0118", # hickory-proto via rama-dns/rama-tcp; remove when rama updates to hickory 0.26.1 or hickory-net
|
||||
"RUSTSEC-2026-0119", # hickory-proto via rama-dns/rama-tcp; remove when rama updates to hickory 0.26.1 or hickory-net
|
||||
"RUSTSEC-2026-0173", # proc-macro-error2 via i18n-embed-fl/age/codex-secrets; remove when local secrets storage migrates off age or age drops i18n-embed-fl
|
||||
]
|
||||
|
||||
Generated
+9
-9
@@ -213,7 +213,7 @@ dependencies = [
|
||||
"lazy_static",
|
||||
"nom 7.1.3",
|
||||
"pin-project",
|
||||
"rand 0.8.6",
|
||||
"rand 0.8.5",
|
||||
"rust-embed",
|
||||
"scrypt",
|
||||
"sha2 0.10.9",
|
||||
@@ -234,7 +234,7 @@ dependencies = [
|
||||
"hkdf 0.12.4",
|
||||
"io_tee",
|
||||
"nom 7.1.3",
|
||||
"rand 0.8.6",
|
||||
"rand 0.8.5",
|
||||
"secrecy",
|
||||
"sha2 0.10.9",
|
||||
]
|
||||
@@ -786,7 +786,7 @@ dependencies = [
|
||||
"hex",
|
||||
"http 1.4.0",
|
||||
"p256",
|
||||
"rand 0.8.6",
|
||||
"rand 0.8.5",
|
||||
"ring",
|
||||
"sha2 0.10.9",
|
||||
"time",
|
||||
@@ -4287,7 +4287,7 @@ dependencies = [
|
||||
"dunce",
|
||||
"glob",
|
||||
"pretty_assertions",
|
||||
"rand 0.8.6",
|
||||
"rand 0.8.5",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tempfile",
|
||||
@@ -9394,7 +9394,7 @@ dependencies = [
|
||||
"chrono",
|
||||
"getrandom 0.2.17",
|
||||
"http 1.4.0",
|
||||
"rand 0.8.6",
|
||||
"rand 0.8.5",
|
||||
"reqwest 0.12.28",
|
||||
"serde",
|
||||
"serde_json",
|
||||
@@ -11033,9 +11033,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "rand"
|
||||
version = "0.8.6"
|
||||
version = "0.8.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a"
|
||||
checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"rand_chacha 0.3.1",
|
||||
@@ -11901,7 +11901,7 @@ dependencies = [
|
||||
"hkdf 0.12.4",
|
||||
"num",
|
||||
"once_cell",
|
||||
"rand 0.8.6",
|
||||
"rand 0.8.5",
|
||||
"serde",
|
||||
"sha2 0.10.9",
|
||||
"zbus",
|
||||
@@ -15566,7 +15566,7 @@ dependencies = [
|
||||
"hex",
|
||||
"nix 0.29.0",
|
||||
"ordered-stream",
|
||||
"rand 0.8.6",
|
||||
"rand 0.8.5",
|
||||
"serde",
|
||||
"serde_repr",
|
||||
"sha1 0.10.6",
|
||||
|
||||
@@ -41,7 +41,7 @@ const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(60);
|
||||
const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
|
||||
#[tokio::test]
|
||||
async fn standalone_web_search_round_trips_encrypted_output() -> Result<()> {
|
||||
async fn standalone_web_search_round_trips_output() -> Result<()> {
|
||||
let call_id = "web-run-1";
|
||||
let server = responses::start_mock_server().await;
|
||||
mount_search_response(&server).await;
|
||||
@@ -170,8 +170,8 @@ async fn standalone_web_search_round_trips_encrypted_output() -> Result<()> {
|
||||
"type": "function_call_output",
|
||||
"call_id": call_id,
|
||||
"output": [{
|
||||
"type": "encrypted_content",
|
||||
"encrypted_content": "ciphertext",
|
||||
"type": "input_text",
|
||||
"text": "Search result",
|
||||
}],
|
||||
})
|
||||
);
|
||||
@@ -259,6 +259,7 @@ async fn mount_search_response(server: &MockServer) {
|
||||
.and(path("/api/codex/alpha/search"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"encrypted_output": "ciphertext",
|
||||
"output": "Search result",
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(server)
|
||||
|
||||
@@ -134,10 +134,13 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn search_posts_typed_request_and_parses_encrypted_output() {
|
||||
async fn search_posts_typed_request_and_parses_output() {
|
||||
let transport = CapturingTransport::new(
|
||||
serde_json::to_vec(&json!({"encrypted_output": "ciphertext"}))
|
||||
.expect("serialize response"),
|
||||
serde_json::to_vec(&json!({
|
||||
"encrypted_output": "ciphertext",
|
||||
"output": "search result",
|
||||
}))
|
||||
.expect("serialize response"),
|
||||
);
|
||||
let client = SearchClient::new(transport.clone(), provider(), Arc::new(DummyAuth));
|
||||
|
||||
@@ -203,7 +206,8 @@ mod tests {
|
||||
assert_eq!(
|
||||
response,
|
||||
SearchResponse {
|
||||
encrypted_output: "ciphertext".to_string(),
|
||||
encrypted_output: Some("ciphertext".to_string()),
|
||||
output: "search result".to_string(),
|
||||
}
|
||||
);
|
||||
|
||||
|
||||
@@ -280,5 +280,6 @@ pub enum AllowedCaller {
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
|
||||
pub struct SearchResponse {
|
||||
pub encrypted_output: String,
|
||||
pub encrypted_output: Option<String>,
|
||||
pub output: String,
|
||||
}
|
||||
|
||||
@@ -42,9 +42,15 @@ 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;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
@@ -91,6 +97,8 @@ pub(crate) struct AgentControl {
|
||||
/// `ThreadManagerState -> CodexThread -> Session -> SessionServices -> ThreadManagerState`.
|
||||
manager: Weak<ThreadManagerState>,
|
||||
state: Arc<AgentRegistry>,
|
||||
v2_residency: Arc<V2Residency>,
|
||||
agent_execution_limiter: Arc<AgentExecutionLimiter>,
|
||||
}
|
||||
|
||||
impl AgentControl {
|
||||
@@ -102,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
|
||||
}
|
||||
|
||||
@@ -116,6 +125,19 @@ impl AgentControl {
|
||||
&self,
|
||||
agent_id: ThreadId,
|
||||
initial_operation: Op,
|
||||
) -> CodexResult<String> {
|
||||
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<ThreadManagerState>,
|
||||
initial_operation: Op,
|
||||
) -> CodexResult<String> {
|
||||
let last_task_message = match &initial_operation {
|
||||
Op::InterAgentCommunication { communication } => {
|
||||
@@ -123,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;
|
||||
@@ -149,14 +170,10 @@ impl AgentControl {
|
||||
) -> CodexResult<String> {
|
||||
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 {
|
||||
@@ -172,7 +189,12 @@ impl AgentControl {
|
||||
/// Interrupt the current task for an existing agent thread.
|
||||
pub(crate) async fn interrupt_agent(&self, agent_id: ThreadId) -> CodexResult<String> {
|
||||
let state = self.upgrade()?;
|
||||
state.send_op(agent_id, Op::Interrupt).await
|
||||
self.handle_thread_request_result(
|
||||
agent_id,
|
||||
&state,
|
||||
state.send_op(agent_id, Op::Interrupt).await,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn handle_thread_request_result(
|
||||
@@ -183,6 +205,7 @@ impl AgentControl {
|
||||
) -> CodexResult<String> {
|
||||
if matches!(result, Err(CodexErr::InternalAgentDied)) {
|
||||
let _ = state.remove_thread(&agent_id).await;
|
||||
self.forget_v2_residency(agent_id);
|
||||
self.state.release_spawned_thread(agent_id);
|
||||
}
|
||||
result
|
||||
@@ -214,6 +237,12 @@ impl AgentControl {
|
||||
self.state.agent_metadata_for_thread(agent_id)
|
||||
}
|
||||
|
||||
pub(crate) fn ensure_agent_known(&self, agent_id: ThreadId) -> CodexResult<AgentMetadata> {
|
||||
self.state
|
||||
.agent_metadata_for_thread(agent_id)
|
||||
.ok_or(CodexErr::ThreadNotFound(agent_id))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_live_agent_subtree_thread_ids(
|
||||
&self,
|
||||
agent_id: ThreadId,
|
||||
|
||||
@@ -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<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;
|
||||
@@ -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()
|
||||
);
|
||||
}
|
||||
@@ -19,6 +19,7 @@ impl AgentControl {
|
||||
state.send_op(agent_id, Op::Shutdown {}).await
|
||||
};
|
||||
let _ = state.remove_thread(&agent_id).await;
|
||||
self.forget_v2_residency(agent_id);
|
||||
self.state.release_spawned_thread(agent_id);
|
||||
result
|
||||
}
|
||||
|
||||
@@ -0,0 +1,240 @@
|
||||
use super::AgentControl;
|
||||
use crate::agent::AgentStatus;
|
||||
use crate::codex_thread::CodexThread;
|
||||
use crate::config::Config;
|
||||
use crate::thread_manager::ThreadManagerState;
|
||||
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::SessionSource;
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use tracing::warn;
|
||||
|
||||
#[derive(Default)]
|
||||
pub(super) struct V2Residency {
|
||||
state: Mutex<V2ResidencyState>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct V2ResidencyState {
|
||||
residents: VecDeque<ThreadId>,
|
||||
pending_slots: usize,
|
||||
}
|
||||
|
||||
pub(super) struct V2ResidencySlot {
|
||||
residency: Arc<V2Residency>,
|
||||
active: bool,
|
||||
}
|
||||
|
||||
impl V2ResidencySlot {
|
||||
pub(super) fn commit(mut self, thread_id: ThreadId) {
|
||||
self.residency.commit_slot(thread_id);
|
||||
self.active = false;
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for V2ResidencySlot {
|
||||
fn drop(&mut self) {
|
||||
if self.active {
|
||||
self.residency.release_pending_slot();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AgentControl {
|
||||
pub(super) async fn reserve_v2_residency_slot(
|
||||
&self,
|
||||
state: &Arc<ThreadManagerState>,
|
||||
config: &Config,
|
||||
protected_thread_id: Option<ThreadId>,
|
||||
) -> CodexResult<V2ResidencySlot> {
|
||||
let capacity = config
|
||||
.effective_agent_max_threads(MultiAgentVersion::V2)
|
||||
.unwrap_or(usize::MAX);
|
||||
Arc::clone(&self.v2_residency)
|
||||
.reserve_slot(state, capacity, protected_thread_id)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(super) async fn touch_loaded_v2_residency(
|
||||
&self,
|
||||
state: &Arc<ThreadManagerState>,
|
||||
thread_id: ThreadId,
|
||||
) {
|
||||
if let Ok(thread) = state.get_thread(thread_id).await
|
||||
&& is_resident_candidate(thread.as_ref())
|
||||
{
|
||||
self.v2_residency.touch(thread_id);
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn forget_v2_residency(&self, thread_id: ThreadId) {
|
||||
self.v2_residency.remove(thread_id);
|
||||
}
|
||||
}
|
||||
|
||||
impl V2Residency {
|
||||
async fn reserve_slot(
|
||||
self: Arc<Self>,
|
||||
manager: &Arc<ThreadManagerState>,
|
||||
capacity: usize,
|
||||
protected_thread_id: Option<ThreadId>,
|
||||
) -> CodexResult<V2ResidencySlot> {
|
||||
loop {
|
||||
if self.try_reserve_pending_slot(capacity) {
|
||||
return Ok(V2ResidencySlot {
|
||||
residency: self,
|
||||
active: true,
|
||||
});
|
||||
}
|
||||
if !self
|
||||
.try_unload_one_resident(manager, protected_thread_id)
|
||||
.await
|
||||
{
|
||||
return Err(CodexErr::AgentLimitReached {
|
||||
max_threads: capacity,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn try_reserve_pending_slot(&self, capacity: usize) -> bool {
|
||||
let mut state = self
|
||||
.state
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if state.residents.len().saturating_add(state.pending_slots) >= capacity {
|
||||
return false;
|
||||
}
|
||||
state.pending_slots += 1;
|
||||
true
|
||||
}
|
||||
|
||||
async fn try_unload_one_resident(
|
||||
&self,
|
||||
manager: &Arc<ThreadManagerState>,
|
||||
protected_thread_id: Option<ThreadId>,
|
||||
) -> bool {
|
||||
let candidates_to_scan = self.resident_count();
|
||||
for _ in 0..candidates_to_scan {
|
||||
let Some(candidate_thread_id) = self.pop_lru_candidate(protected_thread_id) else {
|
||||
return false;
|
||||
};
|
||||
let Some(candidate_thread) = manager
|
||||
.get_thread(candidate_thread_id)
|
||||
.await
|
||||
.ok()
|
||||
.filter(|thread| is_resident_candidate(thread))
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
if !is_unloadable(candidate_thread.as_ref()).await {
|
||||
self.touch(candidate_thread_id);
|
||||
continue;
|
||||
}
|
||||
candidate_thread.ensure_rollout_materialized().await;
|
||||
if let Err(err) = candidate_thread.shutdown_and_wait().await {
|
||||
warn!(
|
||||
"failed to shut down v2 resident thread before unloading {candidate_thread_id}: {err}"
|
||||
);
|
||||
self.touch(candidate_thread_id);
|
||||
continue;
|
||||
}
|
||||
let _ = manager.remove_thread(&candidate_thread_id).await;
|
||||
return true;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
fn resident_count(&self) -> usize {
|
||||
self.state
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.residents
|
||||
.len()
|
||||
}
|
||||
|
||||
fn pop_lru_candidate(&self, protected_thread_id: Option<ThreadId>) -> Option<ThreadId> {
|
||||
let mut state = self
|
||||
.state
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
let candidates_to_scan = state.residents.len();
|
||||
for _ in 0..candidates_to_scan {
|
||||
let candidate_thread_id = state.residents.pop_front()?;
|
||||
if Some(candidate_thread_id) == protected_thread_id {
|
||||
state.residents.push_back(candidate_thread_id);
|
||||
continue;
|
||||
}
|
||||
return Some(candidate_thread_id);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn touch(&self, thread_id: ThreadId) {
|
||||
let mut state = self
|
||||
.state
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
touch_resident(&mut state.residents, thread_id);
|
||||
}
|
||||
|
||||
fn remove(&self, thread_id: ThreadId) {
|
||||
self.state
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.residents
|
||||
.retain(|resident_thread_id| *resident_thread_id != thread_id);
|
||||
}
|
||||
|
||||
fn commit_slot(&self, thread_id: ThreadId) {
|
||||
let mut state = self
|
||||
.state
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
state.pending_slots = state.pending_slots.saturating_sub(1);
|
||||
touch_resident(&mut state.residents, thread_id);
|
||||
}
|
||||
|
||||
fn release_pending_slot(&self) {
|
||||
let mut state = self
|
||||
.state
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
state.pending_slots = state.pending_slots.saturating_sub(1);
|
||||
}
|
||||
}
|
||||
|
||||
fn touch_resident(residents: &mut VecDeque<ThreadId>, thread_id: ThreadId) {
|
||||
residents.retain(|resident_thread_id| *resident_thread_id != thread_id);
|
||||
residents.push_back(thread_id);
|
||||
}
|
||||
|
||||
fn is_resident_candidate(thread: &CodexThread) -> bool {
|
||||
thread.multi_agent_version() == Some(MultiAgentVersion::V2)
|
||||
&& is_v2_resident_session_source(&thread.session_source)
|
||||
}
|
||||
|
||||
pub(super) fn is_v2_resident_session_source(session_source: &SessionSource) -> bool {
|
||||
matches!(session_source, SessionSource::SubAgent(_))
|
||||
}
|
||||
|
||||
async fn is_unloadable(thread: &CodexThread) -> bool {
|
||||
matches!(
|
||||
thread.agent_status().await,
|
||||
AgentStatus::Completed(_) | AgentStatus::Errored(_) | AgentStatus::Interrupted
|
||||
) && thread.codex.session.active_turn.lock().await.is_none()
|
||||
&& !thread
|
||||
.codex
|
||||
.session
|
||||
.input_queue
|
||||
.has_pending_mailbox_items()
|
||||
.await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "residency_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,193 @@
|
||||
use crate::ThreadManager;
|
||||
use crate::agent::AgentControl;
|
||||
use crate::codex_thread::CodexThread;
|
||||
use crate::config::Config;
|
||||
use crate::config::test_config;
|
||||
use crate::thread_manager::ThreadManagerState;
|
||||
use codex_features::Feature;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::error::CodexErr;
|
||||
use codex_protocol::protocol::EventMsg;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use codex_protocol::protocol::SubAgentSource;
|
||||
use codex_protocol::protocol::ThreadSource;
|
||||
use codex_protocol::protocol::TurnAbortReason;
|
||||
use codex_protocol::protocol::TurnAbortedEvent;
|
||||
use codex_protocol::protocol::TurnCompleteEvent;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[tokio::test]
|
||||
async fn residency_slot_reservation_unloads_oldest_idle_v2_agent() {
|
||||
let mut config = test_config().await;
|
||||
let _ = config.features.enable(Feature::MultiAgentV2);
|
||||
config.multi_agent_v2.max_concurrent_threads_per_session = 2;
|
||||
let temp_home = tempfile::tempdir().expect("create temp home");
|
||||
config.codex_home = temp_home.path().to_path_buf().try_into().unwrap();
|
||||
config.cwd = temp_home.path().to_path_buf().try_into().unwrap();
|
||||
let manager = ThreadManager::with_models_provider_and_home_for_tests(
|
||||
CodexAuth::from_api_key("dummy"),
|
||||
config.model_provider.clone(),
|
||||
config.codex_home.to_path_buf(),
|
||||
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
|
||||
);
|
||||
let root = manager
|
||||
.start_thread(config.clone())
|
||||
.await
|
||||
.expect("start root thread");
|
||||
let control = manager.agent_control();
|
||||
let state = control.upgrade().expect("thread manager should be live");
|
||||
|
||||
let first_slot = control
|
||||
.reserve_v2_residency_slot(&state, &config, /*protected_thread_id*/ None)
|
||||
.await
|
||||
.expect("first resident slot");
|
||||
let first =
|
||||
spawn_v2_subagent(&control, &state, config.clone(), root.thread_id, "worker-1").await;
|
||||
first_slot.commit(first.thread_id);
|
||||
mark_thread_completed(first.thread.as_ref()).await;
|
||||
|
||||
let second_slot = control
|
||||
.reserve_v2_residency_slot(&state, &config, /*protected_thread_id*/ None)
|
||||
.await
|
||||
.expect("second resident slot should evict the first idle agent");
|
||||
match manager.get_thread(first.thread_id).await {
|
||||
Err(CodexErr::ThreadNotFound(thread_id)) => assert_eq!(thread_id, first.thread_id),
|
||||
Err(err) => panic!("expected evicted thread to be missing, got {err:?}"),
|
||||
Ok(_) => panic!("expected evicted thread to be missing"),
|
||||
}
|
||||
let second = spawn_v2_subagent(&control, &state, config, root.thread_id, "worker-2").await;
|
||||
second_slot.commit(second.thread_id);
|
||||
|
||||
assert!(manager.get_thread(root.thread_id).await.is_ok());
|
||||
assert!(manager.get_thread(second.thread_id).await.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn interrupted_v2_agent_is_lost_after_residency_eviction() {
|
||||
let mut config = test_config().await;
|
||||
let _ = config.features.enable(Feature::MultiAgentV2);
|
||||
config.multi_agent_v2.max_concurrent_threads_per_session = 2;
|
||||
let temp_home = tempfile::tempdir().expect("create temp home");
|
||||
config.codex_home = temp_home.path().to_path_buf().try_into().unwrap();
|
||||
config.cwd = temp_home.path().to_path_buf().try_into().unwrap();
|
||||
let manager = ThreadManager::with_models_provider_and_home_for_tests(
|
||||
CodexAuth::from_api_key("dummy"),
|
||||
config.model_provider.clone(),
|
||||
config.codex_home.to_path_buf(),
|
||||
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
|
||||
);
|
||||
let root = manager
|
||||
.start_thread(config.clone())
|
||||
.await
|
||||
.expect("start root thread");
|
||||
let control = manager.agent_control();
|
||||
let state = control.upgrade().expect("thread manager should be live");
|
||||
|
||||
let first_slot = control
|
||||
.reserve_v2_residency_slot(&state, &config, /*protected_thread_id*/ None)
|
||||
.await
|
||||
.expect("first resident slot");
|
||||
let first =
|
||||
spawn_v2_subagent(&control, &state, config.clone(), root.thread_id, "worker-1").await;
|
||||
first_slot.commit(first.thread_id);
|
||||
mark_thread_interrupted(first.thread.as_ref()).await;
|
||||
|
||||
let second_slot = control
|
||||
.reserve_v2_residency_slot(&state, &config, /*protected_thread_id*/ None)
|
||||
.await
|
||||
.expect("second resident slot should evict the first interrupted idle agent");
|
||||
match manager.get_thread(first.thread_id).await {
|
||||
Err(CodexErr::ThreadNotFound(thread_id)) => assert_eq!(thread_id, first.thread_id),
|
||||
Err(err) => panic!("expected evicted thread to be missing, got {err:?}"),
|
||||
Ok(_) => panic!("expected evicted thread to be missing"),
|
||||
}
|
||||
let second =
|
||||
spawn_v2_subagent(&control, &state, config.clone(), root.thread_id, "worker-2").await;
|
||||
second_slot.commit(second.thread_id);
|
||||
mark_thread_completed(second.thread.as_ref()).await;
|
||||
|
||||
let err = control
|
||||
.ensure_v2_agent_loaded(config, first.thread_id)
|
||||
.await
|
||||
.expect_err("evicted interrupted agent should stay lost");
|
||||
match err {
|
||||
CodexErr::ThreadNotFound(thread_id) => assert_eq!(thread_id, first.thread_id),
|
||||
err => panic!("expected ThreadNotFound, got {err:?}"),
|
||||
}
|
||||
|
||||
assert!(manager.get_thread(root.thread_id).await.is_ok());
|
||||
assert!(manager.get_thread(second.thread_id).await.is_ok());
|
||||
match manager.get_thread(first.thread_id).await {
|
||||
Err(CodexErr::ThreadNotFound(thread_id)) => assert_eq!(thread_id, first.thread_id),
|
||||
Err(err) => panic!("expected evicted thread to be missing, got {err:?}"),
|
||||
Ok(_) => panic!("expected evicted thread to be missing"),
|
||||
}
|
||||
}
|
||||
|
||||
async fn spawn_v2_subagent(
|
||||
control: &AgentControl,
|
||||
state: &Arc<ThreadManagerState>,
|
||||
config: Config,
|
||||
parent_thread_id: ThreadId,
|
||||
label: &str,
|
||||
) -> crate::thread_manager::NewThread {
|
||||
state
|
||||
.spawn_new_thread_with_source(
|
||||
config,
|
||||
control.clone(),
|
||||
SessionSource::SubAgent(SubAgentSource::Other(label.to_string())),
|
||||
Some(parent_thread_id),
|
||||
/*forked_from_thread_id*/ None,
|
||||
Some(ThreadSource::Subagent),
|
||||
/*metrics_service_name*/ None,
|
||||
/*inherited_shell_snapshot*/ None,
|
||||
/*inherited_exec_policy*/ None,
|
||||
/*environments*/ None,
|
||||
)
|
||||
.await
|
||||
.expect("spawn v2 subagent")
|
||||
}
|
||||
|
||||
async fn mark_thread_completed(thread: &CodexThread) {
|
||||
let turn = thread.codex.session.new_default_turn().await;
|
||||
thread
|
||||
.codex
|
||||
.session
|
||||
.send_event(
|
||||
turn.as_ref(),
|
||||
EventMsg::TurnComplete(TurnCompleteEvent {
|
||||
turn_id: turn.sub_id.clone(),
|
||||
last_agent_message: Some("done".to_string()),
|
||||
completed_at: None,
|
||||
duration_ms: None,
|
||||
time_to_first_token_ms: None,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
clear_active_turn(thread).await;
|
||||
}
|
||||
|
||||
async fn mark_thread_interrupted(thread: &CodexThread) {
|
||||
let turn = thread.codex.session.new_default_turn().await;
|
||||
thread
|
||||
.codex
|
||||
.session
|
||||
.send_event(
|
||||
turn.as_ref(),
|
||||
EventMsg::TurnAborted(TurnAbortedEvent {
|
||||
turn_id: Some(turn.sub_id.clone()),
|
||||
reason: TurnAbortReason::Interrupted,
|
||||
completed_at: None,
|
||||
duration_ms: None,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
clear_active_turn(thread).await;
|
||||
}
|
||||
|
||||
async fn clear_active_turn(thread: &CodexThread) {
|
||||
// The fixture has no task runner to clear the turn after the terminal event.
|
||||
*thread.codex.session.active_turn.lock().await = None;
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
use super::residency::is_v2_resident_session_source;
|
||||
use super::*;
|
||||
|
||||
const AGENT_NAMES: &str = include_str!("../agent_names.txt");
|
||||
@@ -116,6 +117,7 @@ impl AgentControl {
|
||||
) -> CodexResult<()> {
|
||||
let state = self.upgrade()?;
|
||||
if state.get_thread(thread_id).await.is_ok() {
|
||||
self.touch_loaded_v2_residency(&state, thread_id).await;
|
||||
return Ok(());
|
||||
}
|
||||
if self.state.agent_metadata_for_thread(thread_id).is_none() {
|
||||
@@ -143,6 +145,9 @@ impl AgentControl {
|
||||
if initial_history.get_multi_agent_version() != Some(MultiAgentVersion::V2) {
|
||||
return Err(CodexErr::ThreadNotFound(thread_id));
|
||||
}
|
||||
let residency_slot = self
|
||||
.reserve_v2_residency_slot(&state, &config, Some(thread_id))
|
||||
.await?;
|
||||
|
||||
let (session_source, _) = initial_history
|
||||
.get_resumed_session_sources()
|
||||
@@ -170,11 +175,14 @@ impl AgentControl {
|
||||
.await
|
||||
{
|
||||
Ok(reloaded_thread) => {
|
||||
residency_slot.commit(reloaded_thread.thread_id);
|
||||
state.notify_thread_created(reloaded_thread.thread_id);
|
||||
Ok(())
|
||||
}
|
||||
Err(err) => {
|
||||
if state.get_thread(thread_id).await.is_ok() {
|
||||
drop(residency_slot);
|
||||
self.touch_loaded_v2_residency(&state, thread_id).await;
|
||||
return Ok(());
|
||||
}
|
||||
Err(err)
|
||||
@@ -199,8 +207,28 @@ 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 mut reservation = self.state.reserve_spawn_slot(agent_max_threads)?;
|
||||
let spawn_uses_v2_residency = multi_agent_version == MultiAgentVersion::V2
|
||||
&& session_source
|
||||
.as_ref()
|
||||
.is_some_and(is_v2_resident_session_source);
|
||||
let residency_slot = if spawn_uses_v2_residency {
|
||||
Some(
|
||||
self.reserve_v2_residency_slot(&state, &config, /*protected_thread_id*/ None)
|
||||
.await?,
|
||||
)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let reservation_max_threads = if spawn_uses_v2_residency {
|
||||
None
|
||||
} else {
|
||||
agent_max_threads
|
||||
};
|
||||
let mut reservation = self.state.reserve_spawn_slot(reservation_max_threads)?;
|
||||
let inheritance = SpawnAgentThreadInheritance {
|
||||
shell_snapshot: self
|
||||
.inherited_shell_snapshot_for_source(&state, session_source.as_ref())
|
||||
@@ -264,6 +292,9 @@ impl AgentControl {
|
||||
};
|
||||
agent_metadata.agent_id = Some(new_thread.thread_id);
|
||||
reservation.commit(agent_metadata.clone());
|
||||
if let Some(residency_slot) = residency_slot {
|
||||
residency_slot.commit(new_thread.thread_id);
|
||||
}
|
||||
|
||||
if let Some(SessionSource::SubAgent(
|
||||
subagent_source @ SubAgentSource::ThreadSpawn {
|
||||
@@ -321,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
|
||||
@@ -496,13 +527,16 @@ impl AgentControl {
|
||||
session_source: SessionSource,
|
||||
) -> CodexResult<ThreadId> {
|
||||
let root_depth = thread_spawn_depth(&session_source).unwrap_or(0);
|
||||
let resumed_thread_id = Box::pin(self.resume_single_agent_from_rollout(
|
||||
config.clone(),
|
||||
thread_id,
|
||||
session_source,
|
||||
))
|
||||
let (resumed_thread_id, resumed_multi_agent_version) = Box::pin(
|
||||
self.resume_single_agent_from_rollout(config.clone(), thread_id, session_source),
|
||||
)
|
||||
.await?;
|
||||
let state = self.upgrade()?;
|
||||
if config.multi_agent_version_from_features() == MultiAgentVersion::V2
|
||||
|| resumed_multi_agent_version == MultiAgentVersion::V2
|
||||
{
|
||||
return Ok(resumed_thread_id);
|
||||
}
|
||||
let Ok(resumed_thread) = state.get_thread(resumed_thread_id).await else {
|
||||
return Ok(resumed_thread_id);
|
||||
};
|
||||
@@ -548,7 +582,7 @@ impl AgentControl {
|
||||
))
|
||||
.await
|
||||
{
|
||||
Ok(_) => true,
|
||||
Ok((_, _)) => true,
|
||||
Err(err) => {
|
||||
warn!("failed to resume descendant thread {child_thread_id}: {err}");
|
||||
false
|
||||
@@ -569,7 +603,7 @@ impl AgentControl {
|
||||
config: Config,
|
||||
thread_id: ThreadId,
|
||||
session_source: SessionSource,
|
||||
) -> CodexResult<ThreadId> {
|
||||
) -> CodexResult<(ThreadId, MultiAgentVersion)> {
|
||||
let state = self.upgrade()?;
|
||||
let state_db_ctx = state.state_db();
|
||||
let stored_thread = state
|
||||
@@ -674,6 +708,6 @@ impl AgentControl {
|
||||
)
|
||||
.await;
|
||||
|
||||
Ok(resumed_thread.thread_id)
|
||||
Ok((resumed_thread.thread_id, multi_agent_version))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -104,6 +104,10 @@ struct AgentControlHarness {
|
||||
impl AgentControlHarness {
|
||||
async fn new() -> Self {
|
||||
let (home, config) = test_config().await;
|
||||
Self::new_with_config(home, config).await
|
||||
}
|
||||
|
||||
async fn new_with_config(home: TempDir, config: Config) -> Self {
|
||||
let state_db = init_state_db(&config).await;
|
||||
let manager = ThreadManager::with_models_provider_home_and_state_for_tests(
|
||||
CodexAuth::from_api_key("dummy"),
|
||||
@@ -249,6 +253,14 @@ async fn wait_for_live_thread_spawn_children(
|
||||
.expect("expected persisted child tree");
|
||||
}
|
||||
|
||||
async fn assert_thread_not_loaded(manager: &ThreadManager, thread_id: ThreadId) {
|
||||
match manager.get_thread(thread_id).await {
|
||||
Err(CodexErr::ThreadNotFound(id)) => assert_eq!(id, thread_id),
|
||||
Err(err) => panic!("expected ThreadNotFound, got {err:?}"),
|
||||
Ok(_) => panic!("expected thread not to be loaded"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_input_errors_when_manager_dropped() {
|
||||
let control = AgentControl::default();
|
||||
@@ -532,22 +544,7 @@ async fn ensure_v2_agent_loaded_reloads_registered_unloaded_agent() {
|
||||
let (home, mut config) = test_config().await;
|
||||
let _ = config.features.enable(Feature::MultiAgentV2);
|
||||
let _ = config.features.enable(Feature::Sqlite);
|
||||
let state_db = init_state_db(&config).await;
|
||||
let manager = ThreadManager::with_models_provider_home_and_state_for_tests(
|
||||
CodexAuth::from_api_key("dummy"),
|
||||
config.model_provider.clone(),
|
||||
config.codex_home.to_path_buf(),
|
||||
std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
|
||||
state_db.clone(),
|
||||
);
|
||||
let control = manager.agent_control();
|
||||
let harness = AgentControlHarness {
|
||||
_home: home,
|
||||
config,
|
||||
state_db,
|
||||
manager,
|
||||
control,
|
||||
};
|
||||
let harness = AgentControlHarness::new_with_config(home, config).await;
|
||||
let (parent_thread_id, _parent_thread) = harness.start_thread().await;
|
||||
let agent_path = AgentPath::try_from("/root/worker").expect("agent path");
|
||||
let spawned_agent = harness
|
||||
@@ -634,6 +631,96 @@ async fn ensure_v2_agent_loaded_reloads_registered_unloaded_agent() {
|
||||
assert_eq!(captured, Some(expected));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resume_agent_from_rollout_does_not_reopen_v2_descendants() {
|
||||
let (home, mut config) = test_config().await;
|
||||
let _ = config.features.enable(Feature::MultiAgentV2);
|
||||
let _ = config.features.enable(Feature::Sqlite);
|
||||
let harness = AgentControlHarness::new_with_config(home, config).await;
|
||||
let (parent_thread_id, parent_thread) = harness.start_thread().await;
|
||||
let worker_path = AgentPath::root().join("worker").expect("worker path");
|
||||
let worker_thread_id = harness
|
||||
.control
|
||||
.spawn_agent(
|
||||
harness.config.clone(),
|
||||
text_input("hello worker"),
|
||||
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
|
||||
parent_thread_id,
|
||||
depth: 1,
|
||||
agent_path: Some(worker_path.clone()),
|
||||
agent_nickname: None,
|
||||
agent_role: Some("worker".to_string()),
|
||||
})),
|
||||
)
|
||||
.await
|
||||
.expect("worker spawn should succeed");
|
||||
let reviewer_path = worker_path.join("reviewer").expect("reviewer path");
|
||||
let reviewer_thread_id = harness
|
||||
.control
|
||||
.spawn_agent(
|
||||
harness.config.clone(),
|
||||
text_input("hello reviewer"),
|
||||
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
|
||||
parent_thread_id: worker_thread_id,
|
||||
depth: 2,
|
||||
agent_path: Some(reviewer_path.clone()),
|
||||
agent_nickname: None,
|
||||
agent_role: Some("reviewer".to_string()),
|
||||
})),
|
||||
)
|
||||
.await
|
||||
.expect("reviewer spawn should succeed");
|
||||
|
||||
let worker_thread = harness
|
||||
.manager
|
||||
.get_thread(worker_thread_id)
|
||||
.await
|
||||
.expect("worker thread should exist");
|
||||
let reviewer_thread = harness
|
||||
.manager
|
||||
.get_thread(reviewer_thread_id)
|
||||
.await
|
||||
.expect("reviewer thread should exist");
|
||||
persist_thread_for_tree_resume(&parent_thread, "parent persisted").await;
|
||||
persist_thread_for_tree_resume(&worker_thread, "worker persisted").await;
|
||||
persist_thread_for_tree_resume(&reviewer_thread, "reviewer persisted").await;
|
||||
wait_for_live_thread_spawn_children(&harness.control, parent_thread_id, &[worker_thread_id])
|
||||
.await;
|
||||
wait_for_live_thread_spawn_children(&harness.control, worker_thread_id, &[reviewer_thread_id])
|
||||
.await;
|
||||
|
||||
let report = harness
|
||||
.manager
|
||||
.shutdown_all_threads_bounded(Duration::from_secs(5))
|
||||
.await;
|
||||
assert_eq!(report.submit_failed, Vec::<ThreadId>::new());
|
||||
assert_eq!(report.timed_out, Vec::<ThreadId>::new());
|
||||
|
||||
let resumed_manager = ThreadManager::with_models_provider_home_and_state_for_tests(
|
||||
CodexAuth::from_api_key("dummy"),
|
||||
harness.config.model_provider.clone(),
|
||||
harness.config.codex_home.to_path_buf(),
|
||||
std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
|
||||
harness.state_db.clone(),
|
||||
);
|
||||
let resumed_control = resumed_manager.agent_control();
|
||||
let resumed_parent_thread_id = resumed_control
|
||||
.resume_agent_from_rollout(
|
||||
harness.config.clone(),
|
||||
parent_thread_id,
|
||||
SessionSource::Exec,
|
||||
)
|
||||
.await
|
||||
.expect("v2 root resume should succeed");
|
||||
assert_eq!(resumed_parent_thread_id, parent_thread_id);
|
||||
assert_ne!(
|
||||
resumed_control.get_status(parent_thread_id).await,
|
||||
AgentStatus::NotFound
|
||||
);
|
||||
assert_thread_not_loaded(&resumed_manager, worker_thread_id).await;
|
||||
assert_thread_not_loaded(&resumed_manager, reviewer_thread_id).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn encrypted_inter_agent_communication_clears_existing_last_task_message() {
|
||||
let harness = AgentControlHarness::new().await;
|
||||
|
||||
@@ -236,6 +236,12 @@ impl CodexThread {
|
||||
client_user_message_id: Option<String>,
|
||||
initial_goal: Option<InitialGoal>,
|
||||
) -> CodexResult<String> {
|
||||
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,
|
||||
|
||||
@@ -950,7 +950,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 =
|
||||
|
||||
@@ -5185,7 +5185,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");
|
||||
|
||||
@@ -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<TurnContext>,
|
||||
pub(crate) turn_extension_data: Arc<ExtensionData>,
|
||||
pub(crate) _agent_execution_guard: Option<AgentExecutionGuard>,
|
||||
// Timer recorded when the task drops to capture the full turn duration.
|
||||
pub(crate) _timer: Option<codex_otel::Timer>,
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -302,26 +302,30 @@ pub fn create_close_agent_tool_v1() -> ToolSpec {
|
||||
strict: false,
|
||||
defer_loading: None,
|
||||
parameters: JsonSchema::object(properties, Some(vec!["target".to_string()]), Some(false.into())),
|
||||
output_schema: Some(close_agent_output_schema()),
|
||||
output_schema: Some(agent_previous_status_output_schema(
|
||||
"The agent status observed before shutdown was requested.",
|
||||
)),
|
||||
})],
|
||||
})
|
||||
}
|
||||
|
||||
pub fn create_close_agent_tool_v2() -> ToolSpec {
|
||||
pub fn create_interrupt_agent_tool_v2() -> ToolSpec {
|
||||
let properties = BTreeMap::from([(
|
||||
"target".to_string(),
|
||||
JsonSchema::string(Some(
|
||||
"Agent id or canonical task name to close (from spawn_agent).".to_string(),
|
||||
"Agent id or canonical task name to interrupt (from spawn_agent).".to_string(),
|
||||
)),
|
||||
)]);
|
||||
|
||||
ToolSpec::Function(ResponsesApiTool {
|
||||
name: "close_agent".to_string(),
|
||||
description: "Close an agent and any open descendants when they are no longer needed, and return the target agent's previous status before shutdown was requested. Completed agents remain open and count toward the concurrency limit until closed. Don't keep agents open for too long if they are not needed anymore.".to_string(),
|
||||
name: "interrupt_agent".to_string(),
|
||||
description: "Interrupt an agent's current turn, if any, and return its previous status. The agent remains available for messages and follow-up tasks.".to_string(),
|
||||
strict: false,
|
||||
defer_loading: None,
|
||||
parameters: JsonSchema::object(properties, Some(vec!["target".to_string()]), Some(false.into())),
|
||||
output_schema: Some(close_agent_output_schema()),
|
||||
output_schema: Some(agent_previous_status_output_schema(
|
||||
"The agent status observed before the interrupt request was handled.",
|
||||
)),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -501,12 +505,12 @@ fn wait_output_schema_v2() -> Value {
|
||||
})
|
||||
}
|
||||
|
||||
fn close_agent_output_schema() -> Value {
|
||||
fn agent_previous_status_output_schema(previous_status_description: &str) -> Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"previous_status": {
|
||||
"description": "The agent status observed before shutdown was requested.",
|
||||
"description": previous_status_description,
|
||||
"allOf": [agent_status_output_schema()]
|
||||
}
|
||||
},
|
||||
|
||||
@@ -9,8 +9,8 @@ use crate::session::tests::make_session_and_context;
|
||||
use crate::session_prefix::format_subagent_notification_message;
|
||||
use crate::thread_manager::thread_store_from_config;
|
||||
use crate::tools::context::ToolOutput;
|
||||
use crate::tools::handlers::multi_agents_v2::CloseAgentHandler as CloseAgentHandlerV2;
|
||||
use crate::tools::handlers::multi_agents_v2::FollowupTaskHandler as FollowupTaskHandlerV2;
|
||||
use crate::tools::handlers::multi_agents_v2::InterruptAgentHandler;
|
||||
use crate::tools::handlers::multi_agents_v2::ListAgentsHandler as ListAgentsHandlerV2;
|
||||
use crate::tools::handlers::multi_agents_v2::SendMessageHandler as SendMessageHandlerV2;
|
||||
use crate::tools::handlers::multi_agents_v2::SpawnAgentHandler as SpawnAgentHandlerV2;
|
||||
@@ -178,6 +178,11 @@ struct ListedAgentResult {
|
||||
last_task_message: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct InterruptAgentResult {
|
||||
previous_status: AgentStatus,
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn handler_rejects_non_function_payloads() {
|
||||
let (session, turn) = make_session_and_context().await;
|
||||
@@ -1744,6 +1749,82 @@ async fn multi_agent_v2_list_agents_omits_closed_agents() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn multi_agent_v2_list_agents_keeps_interrupted_resident_agents() {
|
||||
let (mut session, mut turn) = make_session_and_context().await;
|
||||
let manager = thread_manager();
|
||||
let root = manager
|
||||
.start_thread((*turn.config).clone())
|
||||
.await
|
||||
.expect("root thread should start");
|
||||
session.services.agent_control = manager.agent_control();
|
||||
session.thread_id = root.thread_id;
|
||||
let mut config = (*turn.config).clone();
|
||||
let _ = config.features.enable(Feature::MultiAgentV2);
|
||||
set_turn_config(&mut turn, config);
|
||||
|
||||
let session = Arc::new(session);
|
||||
let turn = Arc::new(turn);
|
||||
let spawn_output = SpawnAgentHandlerV2::default()
|
||||
.handle(invocation(
|
||||
session.clone(),
|
||||
turn.clone(),
|
||||
"spawn_agent",
|
||||
function_payload(json!({
|
||||
"message": "inspect this repo",
|
||||
"task_name": "worker"
|
||||
})),
|
||||
))
|
||||
.await
|
||||
.expect("spawn_agent should succeed");
|
||||
let _ = expect_text_output(spawn_output);
|
||||
|
||||
let agent_id = session
|
||||
.services
|
||||
.agent_control
|
||||
.resolve_agent_reference(session.thread_id, &turn.session_source, "worker")
|
||||
.await
|
||||
.expect("worker path should resolve");
|
||||
let agent_path = session
|
||||
.services
|
||||
.agent_control
|
||||
.get_agent_metadata(agent_id)
|
||||
.expect("worker metadata should exist")
|
||||
.agent_path
|
||||
.expect("worker path should exist");
|
||||
let interrupt_output = InterruptAgentHandler
|
||||
.handle(invocation(
|
||||
session.clone(),
|
||||
turn.clone(),
|
||||
"interrupt_agent",
|
||||
function_payload(json!({"target": "worker"})),
|
||||
))
|
||||
.await
|
||||
.expect("interrupt_agent should succeed");
|
||||
let _ = expect_text_output(interrupt_output);
|
||||
|
||||
let output = ListAgentsHandlerV2
|
||||
.handle(invocation(
|
||||
session,
|
||||
turn,
|
||||
"list_agents",
|
||||
function_payload(json!({})),
|
||||
))
|
||||
.await
|
||||
.expect("list_agents should succeed");
|
||||
let (content, _) = expect_text_output(output);
|
||||
let result: ListAgentsResult =
|
||||
serde_json::from_str(&content).expect("list_agents result should be json");
|
||||
|
||||
assert_eq!(result.agents.len(), 2);
|
||||
assert_eq!(result.agents[0].agent_name, "/root");
|
||||
assert_eq!(
|
||||
result.agents[0].last_task_message.as_deref(),
|
||||
Some("Main thread")
|
||||
);
|
||||
assert_eq!(result.agents[1].agent_name, agent_path.as_str());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn multi_agent_v2_send_message_rejects_legacy_items_field() {
|
||||
let (mut session, mut turn) = make_session_and_context().await;
|
||||
@@ -3699,7 +3780,7 @@ async fn multi_agent_v2_wait_agent_does_not_return_completed_content() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn multi_agent_v2_close_agent_accepts_task_name_target() {
|
||||
async fn multi_agent_v2_interrupt_agent_accepts_task_name_target() {
|
||||
let (mut session, mut turn) = make_session_and_context().await;
|
||||
let manager = thread_manager();
|
||||
let root = manager
|
||||
@@ -3736,29 +3817,77 @@ async fn multi_agent_v2_close_agent_accepts_task_name_target() {
|
||||
.resolve_agent_reference(session.thread_id, &turn.session_source, "worker")
|
||||
.await
|
||||
.expect("worker path should resolve");
|
||||
|
||||
let output = CloseAgentHandlerV2
|
||||
let worker_thread = manager
|
||||
.get_thread(agent_id)
|
||||
.await
|
||||
.expect("worker thread should be resident");
|
||||
let worker_session = worker_thread.codex.session.clone();
|
||||
SpawnAgentHandlerV2::default()
|
||||
.handle(invocation(
|
||||
session,
|
||||
turn,
|
||||
"close_agent",
|
||||
worker_session.clone(),
|
||||
worker_session.new_default_turn().await,
|
||||
"spawn_agent",
|
||||
function_payload(json!({
|
||||
"message": "inspect a child task",
|
||||
"task_name": "child"
|
||||
})),
|
||||
))
|
||||
.await
|
||||
.expect("child spawn should succeed");
|
||||
let child_id = session
|
||||
.services
|
||||
.agent_control
|
||||
.resolve_agent_reference(session.thread_id, &turn.session_source, "worker/child")
|
||||
.await
|
||||
.expect("child path should resolve");
|
||||
|
||||
let output = InterruptAgentHandler
|
||||
.handle(invocation(
|
||||
session.clone(),
|
||||
turn.clone(),
|
||||
"interrupt_agent",
|
||||
function_payload(json!({"target": "worker"})),
|
||||
))
|
||||
.await
|
||||
.expect("close_agent should succeed for v2 task names");
|
||||
.expect("interrupt_agent should succeed for v2 task names");
|
||||
let (content, success) = expect_text_output(output);
|
||||
let result: close_agent::CloseAgentResult =
|
||||
serde_json::from_str(&content).expect("close_agent result should be json");
|
||||
let result: InterruptAgentResult =
|
||||
serde_json::from_str(&content).expect("interrupt_agent result should be json");
|
||||
assert_ne!(result.previous_status, AgentStatus::NotFound);
|
||||
assert_eq!(success, Some(true));
|
||||
assert_eq!(
|
||||
manager.agent_control().get_status(agent_id).await,
|
||||
AgentStatus::NotFound
|
||||
session
|
||||
.services
|
||||
.agent_control
|
||||
.resolve_agent_reference(session.thread_id, &turn.session_source, "worker")
|
||||
.await
|
||||
.expect("worker path should remain resolvable"),
|
||||
agent_id
|
||||
);
|
||||
manager
|
||||
.get_thread(agent_id)
|
||||
.await
|
||||
.expect("worker should remain resident");
|
||||
manager
|
||||
.get_thread(child_id)
|
||||
.await
|
||||
.expect("child should remain resident");
|
||||
let ops = manager.captured_ops();
|
||||
assert!(
|
||||
ops.iter()
|
||||
.any(|(thread_id, op)| *thread_id == agent_id && matches!(op, Op::Interrupt))
|
||||
);
|
||||
assert!(!ops.iter().any(|(thread_id, op)| {
|
||||
(*thread_id == agent_id || *thread_id == child_id) && matches!(op, Op::Shutdown)
|
||||
}));
|
||||
assert!(
|
||||
!ops.iter()
|
||||
.any(|(thread_id, op)| *thread_id == child_id && matches!(op, Op::Interrupt))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn multi_agent_v2_close_agent_reaps_stale_task_name_target() {
|
||||
async fn multi_agent_v2_interrupt_agent_accepts_unloaded_task_name_target() {
|
||||
let (mut session, mut turn) = make_session_and_context().await;
|
||||
let mut config = (*turn.config).clone();
|
||||
config.multi_agent_v2.max_concurrent_threads_per_session = 2;
|
||||
@@ -3819,18 +3948,18 @@ async fn multi_agent_v2_close_agent_reaps_stale_task_name_target() {
|
||||
.expect("removed worker thread should still accept shutdown");
|
||||
stale_thread.wait_until_terminated().await;
|
||||
|
||||
let output = CloseAgentHandlerV2
|
||||
let output = InterruptAgentHandler
|
||||
.handle(invocation(
|
||||
session.clone(),
|
||||
turn.clone(),
|
||||
"close_agent",
|
||||
"interrupt_agent",
|
||||
function_payload(json!({"target": "worker"})),
|
||||
))
|
||||
.await
|
||||
.expect("close_agent should reap stale v2 task names");
|
||||
.expect("interrupt_agent should accept unloaded v2 task names");
|
||||
let (content, success) = expect_text_output(output);
|
||||
let result: close_agent::CloseAgentResult =
|
||||
serde_json::from_str(&content).expect("close_agent result should be json");
|
||||
let result: InterruptAgentResult =
|
||||
serde_json::from_str(&content).expect("interrupt_agent result should be json");
|
||||
assert_eq!(result.previous_status, AgentStatus::NotFound);
|
||||
assert_eq!(success, Some(true));
|
||||
|
||||
@@ -3841,7 +3970,7 @@ async fn multi_agent_v2_close_agent_reaps_stale_task_name_target() {
|
||||
)
|
||||
.await
|
||||
.expect("open children should load");
|
||||
assert_eq!(open_children, Vec::<ThreadId>::new());
|
||||
assert_eq!(open_children, vec![agent_id]);
|
||||
let closed_children = state_db
|
||||
.list_thread_spawn_children_with_status(
|
||||
root.thread_id,
|
||||
@@ -3849,36 +3978,26 @@ async fn multi_agent_v2_close_agent_reaps_stale_task_name_target() {
|
||||
)
|
||||
.await
|
||||
.expect("closed children should load");
|
||||
assert_eq!(closed_children, vec![agent_id]);
|
||||
assert_eq!(closed_children, Vec::<ThreadId>::new());
|
||||
|
||||
SpawnAgentHandlerV2::default()
|
||||
let output = ListAgentsHandlerV2
|
||||
.handle(invocation(
|
||||
session.clone(),
|
||||
turn.clone(),
|
||||
"spawn_agent",
|
||||
function_payload(json!({
|
||||
"message": "inspect this repo again",
|
||||
"task_name": "replacement"
|
||||
})),
|
||||
"list_agents",
|
||||
function_payload(json!({})),
|
||||
))
|
||||
.await
|
||||
.expect("spawn_agent should succeed after stale close releases the slot");
|
||||
let replacement_id = session
|
||||
.services
|
||||
.agent_control
|
||||
.resolve_agent_reference(session.thread_id, &turn.session_source, "replacement")
|
||||
.await
|
||||
.expect("replacement path should resolve");
|
||||
let _ = session
|
||||
.services
|
||||
.agent_control
|
||||
.shutdown_live_agent(replacement_id)
|
||||
.await
|
||||
.expect("replacement should shut down");
|
||||
.expect("list_agents should succeed");
|
||||
let (content, _) = expect_text_output(output);
|
||||
let result: ListAgentsResult =
|
||||
serde_json::from_str(&content).expect("list_agents result should be json");
|
||||
assert_eq!(result.agents.len(), 1);
|
||||
assert_eq!(result.agents[0].agent_name, "/root");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn multi_agent_v2_close_agent_rejects_root_target_and_id() {
|
||||
async fn multi_agent_v2_interrupt_agent_rejects_root_target_and_id() {
|
||||
let (mut session, mut turn) = make_session_and_context().await;
|
||||
let manager = thread_manager();
|
||||
let root = manager
|
||||
@@ -3896,31 +4015,31 @@ async fn multi_agent_v2_close_agent_rejects_root_target_and_id() {
|
||||
|
||||
let session = Arc::new(session);
|
||||
let turn = Arc::new(turn);
|
||||
let root_path_error = CloseAgentHandlerV2
|
||||
let root_path_error = InterruptAgentHandler
|
||||
.handle(invocation(
|
||||
session.clone(),
|
||||
turn.clone(),
|
||||
"close_agent",
|
||||
"interrupt_agent",
|
||||
function_payload(json!({"target": "/root"})),
|
||||
))
|
||||
.await
|
||||
.err()
|
||||
.expect("close_agent should reject the root path");
|
||||
.expect("interrupt_agent should reject the root path");
|
||||
assert_eq!(
|
||||
root_path_error,
|
||||
FunctionCallError::RespondToModel("root is not a spawned agent".to_string())
|
||||
);
|
||||
|
||||
let root_id_error = CloseAgentHandlerV2
|
||||
let root_id_error = InterruptAgentHandler
|
||||
.handle(invocation(
|
||||
session,
|
||||
turn,
|
||||
"close_agent",
|
||||
"interrupt_agent",
|
||||
function_payload(json!({"target": root.thread_id.to_string()})),
|
||||
))
|
||||
.await
|
||||
.err()
|
||||
.expect("close_agent should reject the root thread id");
|
||||
.expect("interrupt_agent should reject the root thread id");
|
||||
assert_eq!(
|
||||
root_id_error,
|
||||
FunctionCallError::RespondToModel("root is not a spawned agent".to_string())
|
||||
@@ -3928,7 +4047,7 @@ async fn multi_agent_v2_close_agent_rejects_root_target_and_id() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn multi_agent_v2_close_agent_rejects_self_target_by_id() {
|
||||
async fn multi_agent_v2_interrupt_agent_rejects_self_target_by_id() {
|
||||
let (mut session, mut turn) = make_session_and_context().await;
|
||||
let manager = thread_manager();
|
||||
let mut config = (*turn.config).clone();
|
||||
@@ -3976,27 +4095,27 @@ async fn multi_agent_v2_close_agent_rejects_self_target_by_id() {
|
||||
agent_role: None,
|
||||
});
|
||||
|
||||
let err = CloseAgentHandlerV2
|
||||
let err = InterruptAgentHandler
|
||||
.handle(invocation(
|
||||
Arc::new(session),
|
||||
Arc::new(turn),
|
||||
"close_agent",
|
||||
"interrupt_agent",
|
||||
function_payload(json!({"target": child_thread_id.to_string()})),
|
||||
))
|
||||
.await
|
||||
.err()
|
||||
.expect("close_agent should reject self-target by id");
|
||||
.expect("interrupt_agent should reject self-target by id");
|
||||
assert_eq!(
|
||||
err,
|
||||
FunctionCallError::RespondToModel(
|
||||
"an agent cannot close itself; return your result and let the parent close you if needed"
|
||||
"an agent cannot interrupt itself; return your result and let the parent interrupt you if needed"
|
||||
.to_string()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn multi_agent_v2_close_agent_rejects_self_target_by_task_name() {
|
||||
async fn multi_agent_v2_interrupt_agent_rejects_self_target_by_task_name() {
|
||||
let (mut session, mut turn) = make_session_and_context().await;
|
||||
let manager = thread_manager();
|
||||
let mut config = (*turn.config).clone();
|
||||
@@ -4044,20 +4163,20 @@ async fn multi_agent_v2_close_agent_rejects_self_target_by_task_name() {
|
||||
agent_role: None,
|
||||
});
|
||||
|
||||
let err = CloseAgentHandlerV2
|
||||
let err = InterruptAgentHandler
|
||||
.handle(invocation(
|
||||
Arc::new(session),
|
||||
Arc::new(turn),
|
||||
"close_agent",
|
||||
"interrupt_agent",
|
||||
function_payload(json!({"target": child_path.to_string()})),
|
||||
))
|
||||
.await
|
||||
.err()
|
||||
.expect("close_agent should reject self-target by task name");
|
||||
.expect("interrupt_agent should reject self-target by task name");
|
||||
assert_eq!(
|
||||
err,
|
||||
FunctionCallError::RespondToModel(
|
||||
"an agent cannot close itself; return your result and let the parent close you if needed"
|
||||
"an agent cannot interrupt itself; return your result and let the parent interrupt you if needed"
|
||||
.to_string()
|
||||
)
|
||||
);
|
||||
|
||||
@@ -29,15 +29,15 @@ use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
use serde_json::Value as JsonValue;
|
||||
|
||||
pub(crate) use close_agent::Handler as CloseAgentHandler;
|
||||
pub(crate) use followup_task::Handler as FollowupTaskHandler;
|
||||
pub(crate) use interrupt_agent::Handler as InterruptAgentHandler;
|
||||
pub(crate) use list_agents::Handler as ListAgentsHandler;
|
||||
pub(crate) use send_message::Handler as SendMessageHandler;
|
||||
pub(crate) use spawn::Handler as SpawnAgentHandler;
|
||||
pub(crate) use wait::Handler as WaitAgentHandler;
|
||||
|
||||
mod close_agent;
|
||||
mod followup_task;
|
||||
mod interrupt_agent;
|
||||
mod list_agents;
|
||||
mod message_tool;
|
||||
mod send_message;
|
||||
|
||||
+27
-50
@@ -1,5 +1,5 @@
|
||||
use super::*;
|
||||
use crate::tools::handlers::multi_agents_spec::create_close_agent_tool_v2;
|
||||
use crate::tools::handlers::multi_agents_spec::create_interrupt_agent_tool_v2;
|
||||
use crate::turn_timing::now_unix_timestamp_ms;
|
||||
use codex_protocol::error::CodexErr;
|
||||
use codex_tools::ToolSpec;
|
||||
@@ -9,24 +9,26 @@ pub(crate) struct Handler;
|
||||
#[async_trait::async_trait]
|
||||
impl ToolExecutor<ToolInvocation> for Handler {
|
||||
fn tool_name(&self) -> ToolName {
|
||||
ToolName::plain("close_agent")
|
||||
ToolName::plain("interrupt_agent")
|
||||
}
|
||||
|
||||
fn spec(&self) -> ToolSpec {
|
||||
create_close_agent_tool_v2()
|
||||
create_interrupt_agent_tool_v2()
|
||||
}
|
||||
|
||||
async fn handle(
|
||||
&self,
|
||||
invocation: ToolInvocation,
|
||||
) -> Result<Box<dyn crate::tools::context::ToolOutput>, FunctionCallError> {
|
||||
handle_close_agent(invocation).await.map(boxed_tool_output)
|
||||
handle_interrupt_agent(invocation)
|
||||
.await
|
||||
.map(boxed_tool_output)
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_close_agent(
|
||||
async fn handle_interrupt_agent(
|
||||
invocation: ToolInvocation,
|
||||
) -> Result<CloseAgentResult, FunctionCallError> {
|
||||
) -> Result<InterruptAgentResult, FunctionCallError> {
|
||||
let ToolInvocation {
|
||||
session,
|
||||
turn,
|
||||
@@ -35,11 +37,13 @@ async fn handle_close_agent(
|
||||
..
|
||||
} = invocation;
|
||||
let arguments = function_arguments(payload)?;
|
||||
let args: CloseAgentArgs = parse_arguments(&arguments)?;
|
||||
let args: InterruptAgentArgs = parse_arguments(&arguments)?;
|
||||
let agent_id = resolve_agent_target(&session, &turn, &args.target).await?;
|
||||
let receiver_agent = session.services.agent_control.get_agent_metadata(agent_id);
|
||||
let known_agent = receiver_agent.is_some();
|
||||
let receiver_agent = receiver_agent.unwrap_or_default();
|
||||
let receiver_agent = session
|
||||
.services
|
||||
.agent_control
|
||||
.ensure_agent_known(agent_id)
|
||||
.map_err(|err| collab_agent_error(agent_id, err))?;
|
||||
if receiver_agent
|
||||
.agent_path
|
||||
.as_ref()
|
||||
@@ -51,7 +55,7 @@ async fn handle_close_agent(
|
||||
}
|
||||
if agent_id == session.thread_id {
|
||||
return Err(FunctionCallError::RespondToModel(
|
||||
"an agent cannot close itself; return your result and let the parent close you if needed"
|
||||
"an agent cannot interrupt itself; return your result and let the parent interrupt you if needed"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
@@ -67,43 +71,16 @@ async fn handle_close_agent(
|
||||
.into(),
|
||||
)
|
||||
.await;
|
||||
let status = match session
|
||||
let status = session.services.agent_control.get_status(agent_id).await;
|
||||
let result = match session
|
||||
.services
|
||||
.agent_control
|
||||
.subscribe_status(agent_id)
|
||||
.interrupt_agent(agent_id)
|
||||
.await
|
||||
{
|
||||
Ok(mut status_rx) => status_rx.borrow_and_update().clone(),
|
||||
Err(CodexErr::ThreadNotFound(_)) if known_agent => {
|
||||
session.services.agent_control.get_status(agent_id).await
|
||||
}
|
||||
Err(err) => {
|
||||
let status = session.services.agent_control.get_status(agent_id).await;
|
||||
session
|
||||
.send_event(
|
||||
&turn,
|
||||
CollabCloseEndEvent {
|
||||
call_id: call_id.clone(),
|
||||
completed_at_ms: now_unix_timestamp_ms(),
|
||||
sender_thread_id: session.thread_id,
|
||||
receiver_thread_id: agent_id,
|
||||
receiver_agent_nickname: receiver_agent.agent_nickname.clone(),
|
||||
receiver_agent_role: receiver_agent.agent_role.clone(),
|
||||
status,
|
||||
}
|
||||
.into(),
|
||||
)
|
||||
.await;
|
||||
return Err(collab_agent_error(agent_id, err));
|
||||
}
|
||||
Ok(_) | Err(CodexErr::ThreadNotFound(_)) | Err(CodexErr::InternalAgentDied) => Ok(()),
|
||||
Err(err) => Err(collab_agent_error(agent_id, err)),
|
||||
};
|
||||
let result = session
|
||||
.services
|
||||
.agent_control
|
||||
.close_agent(agent_id)
|
||||
.await
|
||||
.map_err(|err| collab_agent_error(agent_id, err))
|
||||
.map(|_| ());
|
||||
session
|
||||
.send_event(
|
||||
&turn,
|
||||
@@ -121,7 +98,7 @@ async fn handle_close_agent(
|
||||
.await;
|
||||
result?;
|
||||
|
||||
Ok(CloseAgentResult {
|
||||
Ok(InterruptAgentResult {
|
||||
previous_status: status,
|
||||
})
|
||||
}
|
||||
@@ -134,18 +111,18 @@ impl CoreToolRuntime for Handler {
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct CloseAgentArgs {
|
||||
struct InterruptAgentArgs {
|
||||
target: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub(crate) struct CloseAgentResult {
|
||||
pub(crate) struct InterruptAgentResult {
|
||||
pub(crate) previous_status: AgentStatus,
|
||||
}
|
||||
|
||||
impl ToolOutput for CloseAgentResult {
|
||||
impl ToolOutput for InterruptAgentResult {
|
||||
fn log_preview(&self) -> String {
|
||||
tool_output_json_text(self, "close_agent")
|
||||
tool_output_json_text(self, "interrupt_agent")
|
||||
}
|
||||
|
||||
fn success_for_logging(&self) -> bool {
|
||||
@@ -153,10 +130,10 @@ impl ToolOutput for CloseAgentResult {
|
||||
}
|
||||
|
||||
fn to_response_item(&self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem {
|
||||
tool_output_response_item(call_id, payload, self, Some(true), "close_agent")
|
||||
tool_output_response_item(call_id, payload, self, Some(true), "interrupt_agent")
|
||||
}
|
||||
|
||||
fn code_mode_result(&self, _payload: &ToolPayload) -> JsonValue {
|
||||
tool_output_code_mode_result(self, "close_agent")
|
||||
tool_output_code_mode_result(self, "interrupt_agent")
|
||||
}
|
||||
}
|
||||
@@ -74,12 +74,8 @@ pub(crate) async fn handle_message_string_tool(
|
||||
let receiver_agent = session
|
||||
.services
|
||||
.agent_control
|
||||
.get_agent_metadata(receiver_thread_id)
|
||||
.ok_or_else(|| {
|
||||
FunctionCallError::RespondToModel(format!(
|
||||
"agent with id {receiver_thread_id} not found"
|
||||
))
|
||||
})?;
|
||||
.ensure_agent_known(receiver_thread_id)
|
||||
.map_err(|err| collab_agent_error(receiver_thread_id, err))?;
|
||||
if mode == MessageDeliveryMode::TriggerTurn
|
||||
&& receiver_agent
|
||||
.agent_path
|
||||
|
||||
@@ -37,8 +37,8 @@ use crate::tools::handlers::multi_agents_common::MAX_WAIT_TIMEOUT_MS;
|
||||
use crate::tools::handlers::multi_agents_common::MIN_WAIT_TIMEOUT_MS;
|
||||
use crate::tools::handlers::multi_agents_spec::SpawnAgentToolOptions;
|
||||
use crate::tools::handlers::multi_agents_spec::WaitAgentTimeoutOptions;
|
||||
use crate::tools::handlers::multi_agents_v2::CloseAgentHandler as CloseAgentHandlerV2;
|
||||
use crate::tools::handlers::multi_agents_v2::FollowupTaskHandler as FollowupTaskHandlerV2;
|
||||
use crate::tools::handlers::multi_agents_v2::InterruptAgentHandler;
|
||||
use crate::tools::handlers::multi_agents_v2::ListAgentsHandler as ListAgentsHandlerV2;
|
||||
use crate::tools::handlers::multi_agents_v2::SendMessageHandler as SendMessageHandlerV2;
|
||||
use crate::tools::handlers::multi_agents_v2::SpawnAgentHandler as SpawnAgentHandlerV2;
|
||||
@@ -742,7 +742,7 @@ fn add_collaboration_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mu
|
||||
exposure,
|
||||
));
|
||||
planned_tools.add_arc(override_tool_exposure(
|
||||
multi_agent_v2_handler(CloseAgentHandlerV2, tool_namespace),
|
||||
multi_agent_v2_handler(InterruptAgentHandler, tool_namespace),
|
||||
exposure,
|
||||
));
|
||||
planned_tools.add_arc(override_tool_exposure(
|
||||
|
||||
@@ -959,6 +959,7 @@ async fn multi_agent_feature_selects_one_agent_tool_family() {
|
||||
"resume_agent",
|
||||
"wait_agent",
|
||||
"close_agent",
|
||||
"interrupt_agent",
|
||||
"send_message",
|
||||
"followup_task",
|
||||
"assign_task",
|
||||
@@ -1011,10 +1012,10 @@ async fn multi_agent_feature_selects_one_agent_tool_family() {
|
||||
"send_message",
|
||||
"followup_task",
|
||||
"wait_agent",
|
||||
"close_agent",
|
||||
"interrupt_agent",
|
||||
"list_agents",
|
||||
]);
|
||||
v2.assert_visible_lacks(&["send_input", "resume_agent", "assign_task"]);
|
||||
v2.assert_visible_lacks(&["send_input", "resume_agent", "assign_task", "close_agent"]);
|
||||
let spawn_agent_description = match v2.visible_spec("spawn_agent") {
|
||||
ToolSpec::Function(tool) => tool.description.as_str(),
|
||||
other => panic!("expected spawn_agent function spec, got {other:?}"),
|
||||
@@ -1096,6 +1097,7 @@ async fn v1_multi_agent_tools_defer_when_tool_search_available() {
|
||||
"resume_agent",
|
||||
"wait_agent",
|
||||
"close_agent",
|
||||
"interrupt_agent",
|
||||
]);
|
||||
for tool_name in [
|
||||
"spawn_agent",
|
||||
@@ -1154,7 +1156,7 @@ async fn multi_agent_v2_can_use_configured_tool_namespace() {
|
||||
"send_message",
|
||||
"followup_task",
|
||||
"wait_agent",
|
||||
"close_agent",
|
||||
"interrupt_agent",
|
||||
"list_agents",
|
||||
] {
|
||||
namespaced.assert_visible_lacks(&[tool_name]);
|
||||
@@ -1244,7 +1246,7 @@ async fn code_mode_only_can_expose_namespaced_multi_agent_v2_as_normal_tools() {
|
||||
"send_message",
|
||||
"followup_task",
|
||||
"wait_agent",
|
||||
"close_agent",
|
||||
"interrupt_agent",
|
||||
"list_agents",
|
||||
] {
|
||||
assert!(
|
||||
@@ -1317,7 +1319,7 @@ async fn hosted_tools_follow_provider_auth_model_and_config_gates() {
|
||||
"send_message",
|
||||
"followup_task",
|
||||
"wait_agent",
|
||||
"close_agent",
|
||||
"interrupt_agent",
|
||||
"list_agents",
|
||||
// Hosted Responses tools.
|
||||
"web_search",
|
||||
|
||||
@@ -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::<serde_json::Value>(&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::<serde_json::Value>(&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(())
|
||||
}
|
||||
@@ -6,9 +6,11 @@ use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
|
||||
use codex_config::types::McpServerConfig;
|
||||
use codex_config::types::McpServerTransportConfig;
|
||||
use codex_core::config::Config;
|
||||
use codex_extension_api::ExtensionRegistryBuilder;
|
||||
use codex_features::Feature;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_models_manager::bundled_models_response;
|
||||
use codex_protocol::config_types::WebSearchMode;
|
||||
use codex_protocol::dynamic_tools::DynamicToolCallOutputContentItem;
|
||||
use codex_protocol::dynamic_tools::DynamicToolResponse;
|
||||
use codex_protocol::dynamic_tools::DynamicToolSpec;
|
||||
@@ -17,6 +19,7 @@ use codex_protocol::protocol::AskForApproval;
|
||||
use codex_protocol::protocol::EventMsg;
|
||||
use codex_protocol::protocol::Op;
|
||||
use codex_protocol::user_input::UserInput;
|
||||
use codex_web_search_extension::install as install_web_search_extension;
|
||||
use core_test_support::apps_test_server::AppsTestServer;
|
||||
use core_test_support::apps_test_server::AppsTestToolLoading;
|
||||
use core_test_support::apps_test_server::DIRECT_CALENDAR_APP_ONLY_TOOL;
|
||||
@@ -45,9 +48,14 @@ use std::collections::HashMap;
|
||||
use std::collections::HashSet;
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use std::time::Instant;
|
||||
use wiremock::Mock;
|
||||
use wiremock::MockServer;
|
||||
use wiremock::ResponseTemplate;
|
||||
use wiremock::matchers::method;
|
||||
use wiremock::matchers::path;
|
||||
|
||||
fn custom_tool_output_items(req: &ResponsesRequest, call_id: &str) -> Vec<Value> {
|
||||
match req.custom_tool_call_output(call_id).get("output") {
|
||||
@@ -191,6 +199,108 @@ async fn run_code_mode_turn_with_config(
|
||||
Ok((test, second_mock))
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn code_mode_can_call_standalone_web_search() -> Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
|
||||
let server = responses::start_mock_server().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/alpha/search"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"output": "Search result",
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
responses::mount_sse_once(
|
||||
&server,
|
||||
sse(vec![
|
||||
ev_response_created("resp-1"),
|
||||
ev_custom_tool_call(
|
||||
"call-1",
|
||||
"exec",
|
||||
r#"
|
||||
const result = await tools.web__run({
|
||||
search_query: [{ q: "standalone web search" }],
|
||||
});
|
||||
text(result);
|
||||
"#,
|
||||
),
|
||||
ev_completed("resp-1"),
|
||||
]),
|
||||
)
|
||||
.await;
|
||||
let follow_up_mock = responses::mount_sse_once(
|
||||
&server,
|
||||
sse(vec![
|
||||
ev_assistant_message("msg-1", "done"),
|
||||
ev_completed("resp-2"),
|
||||
]),
|
||||
)
|
||||
.await;
|
||||
|
||||
let auth = CodexAuth::from_api_key("dummy");
|
||||
let auth_manager = codex_core::test_support::auth_manager_from_auth(auth.clone());
|
||||
let mut extension_builder = ExtensionRegistryBuilder::<Config>::new();
|
||||
install_web_search_extension(&mut extension_builder, auth_manager);
|
||||
let mut builder = test_codex()
|
||||
.with_auth(auth)
|
||||
.with_extensions(Arc::new(extension_builder.build()))
|
||||
.with_model("test-gpt-5.1-codex")
|
||||
.with_config(|config| {
|
||||
config
|
||||
.features
|
||||
.enable(Feature::CodeMode)
|
||||
.expect("code mode should be enabled");
|
||||
config
|
||||
.features
|
||||
.enable(Feature::StandaloneWebSearch)
|
||||
.expect("standalone web search should be enabled");
|
||||
config
|
||||
.web_search_mode
|
||||
.set(WebSearchMode::Live)
|
||||
.expect("web search mode should be accepted");
|
||||
});
|
||||
let test = builder.build(&server).await?;
|
||||
|
||||
test.submit_turn("Search the web from code mode").await?;
|
||||
|
||||
let search_request = server
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("received requests should be available")
|
||||
.into_iter()
|
||||
.find(|request| request.url.path() == "/v1/alpha/search")
|
||||
.expect("standalone search request should be sent");
|
||||
let search_body = search_request
|
||||
.body_json::<Value>()
|
||||
.expect("search request body should be JSON");
|
||||
assert_eq!(
|
||||
search_body["model"],
|
||||
serde_json::json!("test-gpt-5.1-codex")
|
||||
);
|
||||
assert_eq!(
|
||||
search_body["commands"],
|
||||
serde_json::json!({
|
||||
"search_query": [{"q": "standalone web search"}],
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
search_body["settings"],
|
||||
serde_json::json!({
|
||||
"allowed_callers": ["direct"],
|
||||
"external_web_access": true,
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
custom_tool_output_last_non_empty_text(&follow_up_mock.single_request(), "call-1"),
|
||||
Some("Search result".to_string())
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_code_mode_turn_with_rmcp(
|
||||
server: &MockServer,
|
||||
prompt: &str,
|
||||
|
||||
@@ -30,6 +30,7 @@ pub static CODEX_ALIASES_TEMP_DIR: Option<TestBinaryDispatchGuard> = {
|
||||
#[cfg(not(target_os = "windows"))]
|
||||
mod abort_tasks;
|
||||
mod additional_context;
|
||||
mod agent_execution;
|
||||
mod agent_jobs;
|
||||
mod agent_websocket;
|
||||
mod agents_md;
|
||||
|
||||
+1
-1
@@ -75,13 +75,13 @@ ignore = [
|
||||
{ id = "RUSTSEC-2024-0388", reason = "derivative is unmaintained; pulled in via starlark/starlark_syntax v0.14.2 used by execpolicy/cli/core; no fixed starlark release yet" },
|
||||
{ id = "RUSTSEC-2025-0057", reason = "fxhash is unmaintained; pulled in via starlark_map under starlark v0.14.2 used by execpolicy/cli/core and bm25 used by core; remove when upstream dependencies drop fxhash" },
|
||||
{ id = "RUSTSEC-2024-0436", reason = "paste is unmaintained; pulled in via ratatui/rmcp/starlark used by tui/execpolicy; no fixed release yet" },
|
||||
{ id = "RUSTSEC-2026-0173", reason = "proc-macro-error2 is unmaintained; pulled in via i18n-embed-fl v0.9.4 and age v0.11.2 used by codex-secrets; age v0.11.3 still depends on it and no fixed stable release exists" },
|
||||
{ id = "RUSTSEC-2023-0089", reason = "atomic-polyfill is unmaintained; pulled in via postcard/heapless/pagable under starlark v0.14.2 used by execpolicy/cli/core; no fixed starlark release yet" },
|
||||
# TODO(fcoury): remove this exception when syntect drops yaml-rust and bincode, or updates to versions that have fixed the vulnerabilities.
|
||||
{ id = "RUSTSEC-2024-0320", reason = "yaml-rust is unmaintained; pulled in via syntect v5.3.0 used by codex-tui for syntax highlighting; no fixed release yet" },
|
||||
{ id = "RUSTSEC-2025-0141", reason = "bincode is unmaintained; pulled in via syntect v5.3.0 used by codex-tui for syntax highlighting; no fixed release yet" },
|
||||
{ id = "RUSTSEC-2026-0118", reason = "hickory-proto v0.25.2 is pulled in via rama-dns/rama-tcp used by codex-network-proxy; DNSSEC features are not enabled; remove when rama updates to hickory 0.26.1 or hickory-net" },
|
||||
{ id = "RUSTSEC-2026-0119", reason = "hickory-proto v0.25.2 is pulled in via rama-dns/rama-tcp used by codex-network-proxy; no fixed rama release is available yet; remove when rama updates to hickory 0.26.1 or hickory-net" },
|
||||
{ id = "RUSTSEC-2026-0173", reason = "proc-macro-error2 is unmaintained; pulled in via i18n-embed-fl/age used by codex-secrets local storage; remove when codex-secrets migrates off age or age drops i18n-embed-fl" },
|
||||
]
|
||||
# If this is true, then cargo deny will use the git executable to fetch advisory database.
|
||||
# If this is false, then it uses a built-in git library.
|
||||
|
||||
@@ -4,19 +4,19 @@ use codex_protocol::models::FunctionCallOutputContentItem;
|
||||
use codex_protocol::models::FunctionCallOutputPayload;
|
||||
use codex_protocol::models::ResponseInputItem;
|
||||
|
||||
pub(crate) struct EncryptedSearchOutput {
|
||||
encrypted_output: String,
|
||||
pub(crate) struct SearchOutput {
|
||||
output: String,
|
||||
}
|
||||
|
||||
impl EncryptedSearchOutput {
|
||||
pub(crate) fn new(encrypted_output: String) -> Self {
|
||||
Self { encrypted_output }
|
||||
impl SearchOutput {
|
||||
pub(crate) fn new(output: String) -> Self {
|
||||
Self { output }
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolOutput for EncryptedSearchOutput {
|
||||
impl ToolOutput for SearchOutput {
|
||||
fn log_preview(&self) -> String {
|
||||
"[encrypted standalone web search output]".to_string()
|
||||
"[standalone web search output]".to_string()
|
||||
}
|
||||
|
||||
fn success_for_logging(&self) -> bool {
|
||||
@@ -29,8 +29,8 @@ impl ToolOutput for EncryptedSearchOutput {
|
||||
ResponseInputItem::FunctionCallOutput {
|
||||
call_id: call_id.to_string(),
|
||||
output: FunctionCallOutputPayload::from_content_items(vec![
|
||||
FunctionCallOutputContentItem::EncryptedContent {
|
||||
encrypted_content: self.encrypted_output.clone(),
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: self.output.clone(),
|
||||
},
|
||||
]),
|
||||
}
|
||||
@@ -45,12 +45,12 @@ mod tests {
|
||||
use codex_protocol::models::ResponseInputItem;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
use super::EncryptedSearchOutput;
|
||||
use super::SearchOutput;
|
||||
use super::ToolOutput;
|
||||
|
||||
#[test]
|
||||
fn emits_encrypted_function_call_output() {
|
||||
let output = EncryptedSearchOutput::new("encrypted-search-output".to_string());
|
||||
fn emits_plaintext_function_call_output() {
|
||||
let output = SearchOutput::new("search output".to_string());
|
||||
|
||||
assert_eq!(
|
||||
output.to_response_item(
|
||||
@@ -62,8 +62,8 @@ mod tests {
|
||||
ResponseInputItem::FunctionCallOutput {
|
||||
call_id: "call-1".to_string(),
|
||||
output: FunctionCallOutputPayload::from_content_items(vec![
|
||||
FunctionCallOutputContentItem::EncryptedContent {
|
||||
encrypted_content: "encrypted-search-output".to_string(),
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: "search output".to_string(),
|
||||
},
|
||||
]),
|
||||
}
|
||||
|
||||
@@ -26,7 +26,7 @@ use http::HeaderMap;
|
||||
use url::Url;
|
||||
|
||||
use crate::history::recent_input;
|
||||
use crate::output::EncryptedSearchOutput;
|
||||
use crate::output::SearchOutput;
|
||||
use crate::schema::commands_schema;
|
||||
|
||||
pub(crate) const WEB_NAMESPACE: &str = "web";
|
||||
@@ -67,7 +67,7 @@ impl ToolExecutor<ToolCall> for WebSearchTool {
|
||||
}
|
||||
|
||||
fn exposure(&self) -> ToolExposure {
|
||||
ToolExposure::DirectModelOnly
|
||||
ToolExposure::Direct
|
||||
}
|
||||
|
||||
fn supports_parallel_tool_calls(&self) -> bool {
|
||||
@@ -114,9 +114,7 @@ impl ToolExecutor<ToolCall> for WebSearchTool {
|
||||
.emit_completed(web_search_item(&call.call_id, command_action))
|
||||
.await;
|
||||
|
||||
Ok(Box::new(EncryptedSearchOutput::new(
|
||||
response.encrypted_output,
|
||||
)))
|
||||
Ok(Box::new(SearchOutput::new(response.output)))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user