From ed6e5cf919fdb8b388eb5643669dc175f26188cb Mon Sep 17 00:00:00 2001 From: rka-oai Date: Sun, 7 Jun 2026 23:18:23 -0700 Subject: [PATCH 1/6] [codex] Enable standalone web search in code mode (#26719) ## What - Consume plaintext `output` from standalone search while retaining optional `encrypted_output` parsing. - Expose `web.run` to code mode and return search output to nested JavaScript calls. - Cover direct and code-mode standalone search paths with integration tests. ## Why `/v1/alpha/search` now returns plaintext output, which code mode needs to consume standalone search results. ## Test plan - `just test -p codex-api` - `just test -p codex-web-search-extension` - `just test -p codex-core code_mode_can_call_standalone_web_search` - `just test -p codex-app-server standalone_web_search_round_trips_output` --- .../app-server/tests/suite/v2/web_search.rs | 7 +- codex-rs/codex-api/src/endpoint/search.rs | 12 +- codex-rs/codex-api/src/search.rs | 3 +- codex-rs/core/tests/suite/code_mode.rs | 110 ++++++++++++++++++ codex-rs/ext/web-search/src/output.rs | 28 ++--- codex-rs/ext/web-search/src/tool.rs | 8 +- 6 files changed, 141 insertions(+), 27 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/web_search.rs b/codex-rs/app-server/tests/suite/v2/web_search.rs index 13c595d148..626767132b 100644 --- a/codex-rs/app-server/tests/suite/v2/web_search.rs +++ b/codex-rs/app-server/tests/suite/v2/web_search.rs @@ -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) diff --git a/codex-rs/codex-api/src/endpoint/search.rs b/codex-rs/codex-api/src/endpoint/search.rs index d01fbfb780..143048c954 100644 --- a/codex-rs/codex-api/src/endpoint/search.rs +++ b/codex-rs/codex-api/src/endpoint/search.rs @@ -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(), } ); diff --git a/codex-rs/codex-api/src/search.rs b/codex-rs/codex-api/src/search.rs index 061b3ac8c6..bae7c8a7da 100644 --- a/codex-rs/codex-api/src/search.rs +++ b/codex-rs/codex-api/src/search.rs @@ -280,5 +280,6 @@ pub enum AllowedCaller { #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] pub struct SearchResponse { - pub encrypted_output: String, + pub encrypted_output: Option, + pub output: String, } diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index fca4556972..f5f17a09bd 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -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 { 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::::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::() + .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, diff --git a/codex-rs/ext/web-search/src/output.rs b/codex-rs/ext/web-search/src/output.rs index 124271c216..799897b62d 100644 --- a/codex-rs/ext/web-search/src/output.rs +++ b/codex-rs/ext/web-search/src/output.rs @@ -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(), }, ]), } diff --git a/codex-rs/ext/web-search/src/tool.rs b/codex-rs/ext/web-search/src/tool.rs index 9a09b73ebb..35f126dd6b 100644 --- a/codex-rs/ext/web-search/src/tool.rs +++ b/codex-rs/ext/web-search/src/tool.rs @@ -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 for WebSearchTool { } fn exposure(&self) -> ToolExposure { - ToolExposure::DirectModelOnly + ToolExposure::Direct } fn supports_parallel_tool_calls(&self) -> bool { @@ -114,9 +114,7 @@ impl ToolExecutor 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))) } } From 4e803a017c958dd37eb251372ba810232d3e84ba Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 8 Jun 2026 10:24:48 +0200 Subject: [PATCH 2/6] feat: add v2 agent residency lru (#26632) ## Why Multi-agent v2 treats agents as durable logical agents, not just live entries in `ThreadManager`. After the reload-on-delivery change, a v2 agent can be addressed even if its thread is not currently loaded. This PR adds the next layer: loaded v2 subagents can be paged out of `ThreadManager` when the session has too many resident agents. That keeps residency separate from logical identity and prepares the stack for making v2 concurrency count active execution instead of existing agents. ## What Changed - Add an `AgentControl`-scoped LRU for resident v2 subagents. - Reserve residency before spawning or reloading a v2 subagent. - If resident capacity is full, unload the least-recently-used idle v2 subagent from `ThreadManager`. - Keep `ThreadManager` as a primitive loaded-thread store; it does not own the LRU policy. - Keep unloaded agents registered and durable so they can be reloaded by the delivery path. - Preserve the existing v2 cap semantics by using the derived non-root v2 cap for residency. Eviction is intentionally conservative. A thread is unloadable only when it is a v2 subagent, has completed or errored, has no active turn, and has no pending mailbox work. Before removal, the rollout is materialized and flushed. ## Assumptions And Non-Goals - PR #26623 provides the reload-on-delivery path for unloaded v2 agents. - `ThreadManager` membership means loaded/resident, not logical agent existence. - `AgentRegistry` remains the logical identity/metadata source for v2 agents that may be unloaded. - `list_agents` remains a recent/resident view for now. - This does not change active execution concurrency; that is the next PR. - This does not change `close_agent` semantics. - This does not change or remove `resume_agent`. - This does not add a new residency config knob. ## Stack 1. V2 durable lookup and reload on delivery (#26623) - reload unloaded v2 agents before delivering follow-up/input. 2. V2 residency LRU (this PR) - unload idle resident v2 agents from `ThreadManager` when resident capacity is full. 3. V2 active-execution concurrency - count running non-root v2 turns instead of logical agents. 4. V2 close/interrupt semantics - make v2 close interrupt the current turn without deleting durable identity. 5. V2 resume cleanup - remove the manual resume surface for v2 while keeping internal reload support. ## Validation - Added focused coverage for the residency LRU eviction path. - Local clippy/check/tests were not run; CI will cover them. --- codex-rs/core/src/agent/control.rs | 5 + codex-rs/core/src/agent/control/legacy.rs | 1 + codex-rs/core/src/agent/control/residency.rs | 240 ++++++++++++++++++ .../core/src/agent/control/residency_tests.rs | 107 ++++++++ codex-rs/core/src/agent/control/spawn.rs | 30 ++- 5 files changed, 382 insertions(+), 1 deletion(-) create mode 100644 codex-rs/core/src/agent/control/residency.rs create mode 100644 codex-rs/core/src/agent/control/residency_tests.rs diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 6c6650a559..e9994883f5 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -42,9 +42,12 @@ use std::sync::Weak; use tokio::sync::watch; use tracing::warn; +use self::residency::V2Residency; + const ROOT_LAST_TASK_MESSAGE: &str = "Main thread"; mod legacy; +mod residency; mod spawn; #[derive(Clone, Debug, PartialEq, Eq)] @@ -91,6 +94,7 @@ pub(crate) struct AgentControl { /// `ThreadManagerState -> CodexThread -> Session -> SessionServices -> ThreadManagerState`. manager: Weak, state: Arc, + v2_residency: Arc, } impl AgentControl { @@ -183,6 +187,7 @@ impl AgentControl { ) -> CodexResult { 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 diff --git a/codex-rs/core/src/agent/control/legacy.rs b/codex-rs/core/src/agent/control/legacy.rs index e85090b905..aa41472687 100644 --- a/codex-rs/core/src/agent/control/legacy.rs +++ b/codex-rs/core/src/agent/control/legacy.rs @@ -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 } diff --git a/codex-rs/core/src/agent/control/residency.rs b/codex-rs/core/src/agent/control/residency.rs new file mode 100644 index 0000000000..8056fedeba --- /dev/null +++ b/codex-rs/core/src/agent/control/residency.rs @@ -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, +} + +#[derive(Default)] +struct V2ResidencyState { + residents: VecDeque, + pending_slots: usize, +} + +pub(super) struct V2ResidencySlot { + residency: Arc, + 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, + config: &Config, + protected_thread_id: Option, + ) -> CodexResult { + 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, + 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, + manager: &Arc, + capacity: usize, + protected_thread_id: Option, + ) -> CodexResult { + 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, + protected_thread_id: Option, + ) -> 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.flush_rollout().await { + warn!( + "failed to flush 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) -> Option { + 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, 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(_) + ) && 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; diff --git a/codex-rs/core/src/agent/control/residency_tests.rs b/codex-rs/core/src/agent/control/residency_tests.rs new file mode 100644 index 0000000000..a0338eda06 --- /dev/null +++ b/codex-rs/core/src/agent/control/residency_tests.rs @@ -0,0 +1,107 @@ +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::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()); +} + +async fn spawn_v2_subagent( + control: &AgentControl, + state: &Arc, + 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; + // The fixture has no task runner to clear the turn after the terminal event. + *thread.codex.session.active_turn.lock().await = None; +} diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index 8df3b0a07f..e3536e89f9 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -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) @@ -200,7 +208,24 @@ impl AgentControl { ) .await; 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 +289,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 { From 8f1aad58dc3a684afecdd9efd31d074762da3c24 Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 8 Jun 2026 11:39:08 +0200 Subject: [PATCH 3/6] Ignore proc-macro-error2 advisory (#26974) ## Summary - ignore RUSTSEC-2026-0173 in cargo-deny and cargo-audit config - document that proc-macro-error2 is pulled in transitively via i18n-embed-fl/age/codex-secrets - leave the ignore temporary until codex-secrets moves off age or age drops i18n-embed-fl ## Validation - just fmt - cargo deny check --hide-inclusion-graph --- codex-rs/.cargo/audit.toml | 1 + codex-rs/deny.toml | 1 + 2 files changed, 2 insertions(+) diff --git a/codex-rs/.cargo/audit.toml b/codex-rs/.cargo/audit.toml index 6d963b725d..c1d8ff87ee 100644 --- a/codex-rs/.cargo/audit.toml +++ b/codex-rs/.cargo/audit.toml @@ -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 ] diff --git a/codex-rs/deny.toml b/codex-rs/deny.toml index a5c06e862e..983fd5b962 100644 --- a/codex-rs/deny.toml +++ b/codex-rs/deny.toml @@ -81,6 +81,7 @@ ignore = [ { 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. From 743f5aad38accd52da34bf4dcbdd1215a8c3ab9a Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 8 Jun 2026 14:21:28 +0200 Subject: [PATCH 4/6] feat: count V2 concurrency by active execution (#26969) ## Why Multi-Agent V2 concurrency should count active non-root turns, not resident or durable agent threads. The limit is intentionally best effort: admission checks are synchronous, but concurrent successful checks may overshoot slightly. ## What changed - Keep one root-derived execution limit on the shared `AgentControl`. - Count active V2 subagent turns with an RAII guard owned by `RunningTask`. - Check capacity before spawning or starting an idle agent, including direct app-server `turn/start` submissions. - Preserve queued delivery for agents that are already running. - Exempt automatic idle continuations so `/goal` work is not dropped when capacity is temporarily full. - Keep root and V1 turns outside this limiter. ## Test coverage - `execution_guards_count_active_v2_subagent_turns` - `execution_guards_ignore_root_and_v1_turns` - `v2_nested_spawn_checks_shared_active_execution_capacity` --- codex-rs/core/src/agent/control.rs | 33 +++-- codex-rs/core/src/agent/control/execution.rs | 110 ++++++++++++++++ .../core/src/agent/control/execution_tests.rs | 60 +++++++++ codex-rs/core/src/agent/control/spawn.rs | 5 +- codex-rs/core/src/codex_thread.rs | 6 + codex-rs/core/src/session/session.rs | 7 +- codex-rs/core/src/session/tests.rs | 2 +- codex-rs/core/src/state/turn.rs | 2 + codex-rs/core/src/tasks/mod.rs | 5 + codex-rs/core/tests/suite/agent_execution.rs | 122 ++++++++++++++++++ codex-rs/core/tests/suite/mod.rs | 1 + 11 files changed, 340 insertions(+), 13 deletions(-) create mode 100644 codex-rs/core/src/agent/control/execution.rs create mode 100644 codex-rs/core/src/agent/control/execution_tests.rs create mode 100644 codex-rs/core/tests/suite/agent_execution.rs diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index e9994883f5..888f9cdf94 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -42,10 +42,13 @@ 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; @@ -95,6 +98,7 @@ pub(crate) struct AgentControl { manager: Weak, state: Arc, v2_residency: Arc, + agent_execution_limiter: Arc, } impl AgentControl { @@ -106,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 } @@ -120,6 +125,19 @@ impl AgentControl { &self, agent_id: ThreadId, initial_operation: Op, + ) -> CodexResult { + 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, + initial_operation: Op, ) -> CodexResult { let last_task_message = match &initial_operation { Op::InterAgentCommunication { communication } => { @@ -127,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; @@ -153,14 +170,10 @@ impl AgentControl { ) -> CodexResult { 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 { diff --git a/codex-rs/core/src/agent/control/execution.rs b/codex-rs/core/src/agent/control/execution.rs new file mode 100644 index 0000000000..42aaa50e02 --- /dev/null +++ b/codex-rs/core/src/agent/control/execution.rs @@ -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, +} + +pub(crate) struct AgentExecutionGuard { + limiter: Arc, +} + +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 { + 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) -> 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; diff --git a/codex-rs/core/src/agent/control/execution_tests.rs b/codex-rs/core/src/agent/control/execution_tests.rs new file mode 100644 index 0000000000..9cad08f951 --- /dev/null +++ b/codex-rs/core/src/agent/control/execution_tests.rs @@ -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() + ); +} diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index e3536e89f9..a103fbb80d 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -207,6 +207,9 @@ 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 spawn_uses_v2_residency = multi_agent_version == MultiAgentVersion::V2 && session_source @@ -349,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 diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index b3f9cb0e2e..1697f768e2 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -234,6 +234,12 @@ impl CodexThread { trace: Option, client_user_message_id: Option, ) -> CodexResult { + 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, trace, client_user_message_id) .await diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 122ff6aeb9..92f9737310 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -948,7 +948,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 = diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 366cf6f22a..4d914ff985 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5184,7 +5184,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"); diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index eaf26ba853..d1e31b7994 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -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, pub(crate) turn_extension_data: Arc, + pub(crate) _agent_execution_guard: Option, // Timer recorded when the task drops to capture the full turn duration. pub(crate) _timer: Option, } diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index a9a664a505..911f30b319 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -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); diff --git a/codex-rs/core/tests/suite/agent_execution.rs b/codex-rs/core/tests/suite/agent_execution.rs new file mode 100644 index 0000000000..f073d70089 --- /dev/null +++ b/codex-rs/core/tests/suite/agent_execution.rs @@ -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::(&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::(&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(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 11246d7cc4..234fb8a2f0 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -30,6 +30,7 @@ pub static CODEX_ALIASES_TEMP_DIR: Option = { #[cfg(not(target_os = "windows"))] mod abort_tasks; mod additional_context; +mod agent_execution; mod agent_jobs; mod agent_websocket; mod agents_md; From 8d415050fce4b4ebc6da1ba247379844235fa453 Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 8 Jun 2026 14:46:35 +0200 Subject: [PATCH 5/6] Rename multi-agent v2 close_agent to interrupt_agent (#26994) ## Why `close_agent` is the wrong model-facing name for the v2 operation after the residency changes. V2 agents remain reusable by task name, and residency/unloading owns capacity management; the exposed tool should describe the action it actually performs: interrupt the target agent's current turn without making the agent unavailable for future messages or follow-up tasks. ## What changed - Rename the multi-agent v2 tool from `close_agent` to `interrupt_agent`. - Keep the v1 `close_agent` surface unchanged. - Update the v2 handler to send `Op::Interrupt`, keep interrupted agents registered, and reject root/self targets with interrupt-specific errors. - Route interrupt delivery through the existing dead-thread cleanup path so stale resident entries do not keep consuming capacity. - Update tool planning and handler tests for the new v2 surface and semantics. ## Verification Added focused coverage in: - `core/src/tools/spec_plan_tests.rs` - `core/src/tools/handlers/multi_agents_tests.rs` --- codex-rs/core/src/agent/control.rs | 13 +- .../src/tools/handlers/multi_agents_spec.rs | 20 +- .../src/tools/handlers/multi_agents_tests.rs | 231 +++++++++++++----- .../src/tools/handlers/multi_agents_v2.rs | 4 +- .../{close_agent.rs => interrupt_agent.rs} | 77 ++---- .../handlers/multi_agents_v2/message_tool.rs | 8 +- codex-rs/core/src/tools/spec_plan.rs | 4 +- codex-rs/core/src/tools/spec_plan_tests.rs | 12 +- 8 files changed, 239 insertions(+), 130 deletions(-) rename codex-rs/core/src/tools/handlers/multi_agents_v2/{close_agent.rs => interrupt_agent.rs} (56%) diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 888f9cdf94..53a50e8e1f 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -189,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 { 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( @@ -232,6 +237,12 @@ impl AgentControl { self.state.agent_metadata_for_thread(agent_id) } + pub(crate) fn ensure_agent_known(&self, agent_id: ThreadId) -> CodexResult { + 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, diff --git a/codex-rs/core/src/tools/handlers/multi_agents_spec.rs b/codex-rs/core/src/tools/handlers/multi_agents_spec.rs index 1dcb55f436..8455b13036 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_spec.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_spec.rs @@ -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()] } }, diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 67ccf36f91..4f514e23a7 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -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, } +#[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::::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::::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() ) ); diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2.rs index 8bb207e3e1..6b5bd9f241 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2.rs @@ -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; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/close_agent.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/interrupt_agent.rs similarity index 56% rename from codex-rs/core/src/tools/handlers/multi_agents_v2/close_agent.rs rename to codex-rs/core/src/tools/handlers/multi_agents_v2/interrupt_agent.rs index 7b3f8d776e..907367348b 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/close_agent.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/interrupt_agent.rs @@ -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 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, 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 { +) -> Result { 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") } } diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs index 528fa55264..07fd9d70ee 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs @@ -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 diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index b1a01cd7aa..fed62a0ee1 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -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( diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 5fe811b19e..f96f41f0f2 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -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", From 0526cb56ac3501a02968010d03873993c319e290 Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 8 Jun 2026 16:44:50 +0200 Subject: [PATCH 6/6] Avoid reopening v2 descendants on resume (#26997) ## Why Multi-agent v2 residency is intended to keep only the threads that need to be live. The existing rollout resume path still walked persisted open descendants and reopened the entire descendant tree when resuming a v2 root, which turns resume into an eager reload of work that should stay unloaded until it is explicitly needed. The interrupted-agent path has a related residency issue. Interrupted agents remain open by design, so an idle interrupted resident should be eligible for eviction just like an idle completed or errored resident. Otherwise a resident set full of interrupted agents can consume every v2 slot and block later spawns or reloads with `AgentLimitReached`. ## What Changed - Return early from `resume_agent_from_rollout` after resuming a v2 thread so persisted v2 descendants are not reopened eagerly. - Treat idle `Interrupted` v2 residents as unloadable in the LRU residency path. - Add focused coverage for v2 root resume leaving descendants unloaded and for eviction of an idle interrupted v2 resident when a new slot is needed. ## Verification Added targeted `codex-core` tests covering: - v2 root resume with persisted descendants, verifying only the root is loaded after resume. - residency eviction of an idle interrupted v2 agent when the resident set is full. --- codex-rs/core/src/agent/control/residency.rs | 6 +- .../core/src/agent/control/residency_tests.rs | 86 +++++++++++++ codex-rs/core/src/agent/control/spawn.rs | 19 +-- codex-rs/core/src/agent/control_tests.rs | 119 +++++++++++++++--- 4 files changed, 203 insertions(+), 27 deletions(-) diff --git a/codex-rs/core/src/agent/control/residency.rs b/codex-rs/core/src/agent/control/residency.rs index 8056fedeba..901f2c93bd 100644 --- a/codex-rs/core/src/agent/control/residency.rs +++ b/codex-rs/core/src/agent/control/residency.rs @@ -136,9 +136,9 @@ impl V2Residency { continue; } candidate_thread.ensure_rollout_materialized().await; - if let Err(err) = candidate_thread.flush_rollout().await { + if let Err(err) = candidate_thread.shutdown_and_wait().await { warn!( - "failed to flush v2 resident thread before unloading {candidate_thread_id}: {err}" + "failed to shut down v2 resident thread before unloading {candidate_thread_id}: {err}" ); self.touch(candidate_thread_id); continue; @@ -225,7 +225,7 @@ pub(super) fn is_v2_resident_session_source(session_source: &SessionSource) -> b async fn is_unloadable(thread: &CodexThread) -> bool { matches!( thread.agent_status().await, - AgentStatus::Completed(_) | AgentStatus::Errored(_) + AgentStatus::Completed(_) | AgentStatus::Errored(_) | AgentStatus::Interrupted ) && thread.codex.session.active_turn.lock().await.is_none() && !thread .codex diff --git a/codex-rs/core/src/agent/control/residency_tests.rs b/codex-rs/core/src/agent/control/residency_tests.rs index a0338eda06..88d4b004d3 100644 --- a/codex-rs/core/src/agent/control/residency_tests.rs +++ b/codex-rs/core/src/agent/control/residency_tests.rs @@ -12,6 +12,8 @@ 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; @@ -62,6 +64,68 @@ async fn residency_slot_reservation_unloads_oldest_idle_v2_agent() { 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, @@ -102,6 +166,28 @@ async fn mark_thread_completed(thread: &CodexThread) { }), ) .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; } diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index a103fbb80d..a0d52ea4d6 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -527,13 +527,16 @@ impl AgentControl { session_source: SessionSource, ) -> CodexResult { 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); }; @@ -579,7 +582,7 @@ impl AgentControl { )) .await { - Ok(_) => true, + Ok((_, _)) => true, Err(err) => { warn!("failed to resume descendant thread {child_thread_id}: {err}"); false @@ -600,7 +603,7 @@ impl AgentControl { config: Config, thread_id: ThreadId, session_source: SessionSource, - ) -> CodexResult { + ) -> CodexResult<(ThreadId, MultiAgentVersion)> { let state = self.upgrade()?; let state_db_ctx = state.state_db(); let stored_thread = state @@ -705,6 +708,6 @@ impl AgentControl { ) .await; - Ok(resumed_thread.thread_id) + Ok((resumed_thread.thread_id, multi_agent_version)) } } diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index c9cd444925..3ad1360e5a 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -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::::new()); + assert_eq!(report.timed_out, Vec::::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;