From c2fef72f12129aabe91affa4302e6b1973ef228a Mon Sep 17 00:00:00 2001 From: ychhabria Date: Tue, 24 Mar 2026 00:05:32 -0700 Subject: [PATCH] codex-exp-2: tune fanout and batch waiting --- codex-rs/cli/src/main.rs | 3 +- .../src/tools/handlers/multi_agents/wait.rs | 122 +++++++++++++----- .../src/tools/handlers/multi_agents_tests.rs | 100 ++++++++++++++ codex-rs/core/src/tools/spec.rs | 18 ++- .../core/templates/agents/orchestrator.md | 3 +- .../templates/collab/experimental_prompt.md | 1 + 6 files changed, 207 insertions(+), 40 deletions(-) diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 8e00c6bb79..0470e5b304 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -86,11 +86,12 @@ struct MultitoolCli { } const CODEX_EXP_2_BIN_NAME: &str = "codex-exp-2"; +// Benchmarked on a 16-core M4 Max with a 20-task cross-repo workload. const CODEX_EXP_2_DEFAULT_OVERRIDES: [&str; 5] = [ r#"model="gpt-5.4""#, r#"model_reasoning_effort="xhigh""#, r#"service_tier="fast""#, - "agents.max_threads=16", + "agents.max_threads=8", "agents.max_depth=3", ]; diff --git a/codex-rs/core/src/tools/handlers/multi_agents/wait.rs b/codex-rs/core/src/tools/handlers/multi_agents/wait.rs index 2d655ce86d..2f50fb4f18 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/wait.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/wait.rs @@ -119,43 +119,21 @@ impl ToolHandler for Handler { } } - let statuses = if !initial_final_statuses.is_empty() { - initial_final_statuses - } else { - let mut futures = FuturesUnordered::new(); - for (id, rx) in status_rxs.into_iter() { - let session = session.clone(); - futures.push(wait_for_final_status(session, id, rx)); - } - let mut results = Vec::new(); - let deadline = Instant::now() + Duration::from_millis(timeout_ms as u64); - loop { - match timeout_at(deadline, futures.next()).await { - Ok(Some(Some(result))) => { - results.push(result); - break; - } - Ok(Some(None)) => continue, - Ok(None) | Err(_) => break, - } - } - if !results.is_empty() { - loop { - match futures.next().now_or_never() { - Some(Some(Some(result))) => results.push(result), - Some(Some(None)) => continue, - Some(None) | None => break, - } - } - } - results - }; - - let statuses_map = statuses.clone().into_iter().collect::>(); + let expected_status_count = receiver_thread_ids.len(); + let deadline = Instant::now() + Duration::from_millis(timeout_ms as u64); + let statuses_map = wait_for_requested_statuses( + session.clone(), + status_rxs, + initial_final_statuses, + expected_status_count, + args.wait_for_all, + deadline, + ) + .await; let agent_statuses = build_wait_agent_statuses(&statuses_map, &receiver_agents); let result = WaitAgentResult { status: statuses_map.clone(), - timed_out: statuses.is_empty(), + timed_out: wait_timed_out(statuses_map.len(), expected_status_count, args.wait_for_all), }; session @@ -179,6 +157,8 @@ impl ToolHandler for Handler { struct WaitArgs { ids: Vec, timeout_ms: Option, + #[serde(default)] + wait_for_all: bool, } #[derive(Debug, Deserialize, Serialize, PartialEq, Eq)] @@ -226,3 +206,77 @@ async fn wait_for_final_status( } } } + +async fn wait_for_requested_statuses( + session: Arc, + status_rxs: Vec<(ThreadId, Receiver)>, + initial_statuses: Vec<(ThreadId, AgentStatus)>, + expected_count: usize, + wait_for_all: bool, + deadline: Instant, +) -> HashMap { + let mut statuses = initial_statuses.into_iter().collect::>(); + if wait_condition_satisfied(statuses.len(), expected_count, wait_for_all) { + return statuses; + } + + let mut futures = FuturesUnordered::new(); + for (id, rx) in status_rxs { + if statuses.contains_key(&id) { + continue; + } + let session = session.clone(); + futures.push(wait_for_final_status(session, id, rx)); + } + + loop { + if wait_condition_satisfied(statuses.len(), expected_count, wait_for_all) { + break; + } + + match timeout_at(deadline, futures.next()).await { + Ok(Some(Some((id, status)))) => { + statuses.insert(id, status); + if !wait_for_all { + break; + } + } + Ok(Some(None)) => continue, + Ok(None) | Err(_) => break, + } + } + + drain_ready_final_statuses(&mut futures, &mut statuses); + statuses +} + +fn drain_ready_final_statuses( + futures: &mut FuturesUnordered>>, + statuses: &mut HashMap, +) { + loop { + match futures.next().now_or_never() { + Some(Some(Some((id, status)))) => { + statuses.insert(id, status); + } + Some(Some(None)) => continue, + Some(None) | None => break, + } + } +} + +fn wait_condition_satisfied( + observed_count: usize, + expected_count: usize, + wait_for_all: bool, +) -> bool { + if wait_for_all { + observed_count >= expected_count + } else { + observed_count > 0 + } +} + +fn wait_timed_out(observed_count: usize, expected_count: usize, wait_for_all: bool) -> bool { + !wait_condition_satisfied(observed_count, expected_count, wait_for_all) +} 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 99afe8ac25..fea1ae2db6 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -31,6 +31,7 @@ use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; use tokio::sync::Mutex; +use tokio::time::sleep; use tokio::time::timeout; fn invocation( @@ -970,6 +971,105 @@ async fn wait_agent_returns_final_status_without_timeout() { assert_eq!(success, None); } +#[tokio::test] +async fn wait_agent_wait_for_all_does_not_return_after_first_final_status() { + let (mut session, turn) = make_session_and_context().await; + let manager = thread_manager(); + session.services.agent_control = manager.agent_control(); + let config = turn.config.as_ref().clone(); + let thread_a = manager + .start_thread(config.clone()) + .await + .expect("start thread"); + let thread_b = manager.start_thread(config).await.expect("start thread"); + let id_a = thread_a.thread_id; + let thread_a_handle = Arc::clone(&thread_a.thread); + + tokio::spawn(async move { + sleep(Duration::from_millis(20)).await; + let _ = thread_a_handle.submit(Op::Shutdown {}).await; + }); + + let invocation = invocation( + Arc::new(session), + Arc::new(turn), + "wait_agent", + function_payload(json!({ + "ids": [id_a.to_string(), thread_b.thread_id.to_string()], + "wait_for_all": true, + "timeout_ms": 1000 + })), + ); + + let early = timeout( + Duration::from_millis(80), + WaitAgentHandler.handle(invocation), + ) + .await; + assert!( + early.is_err(), + "wait_agent(wait_for_all=true) should keep waiting after the first agent finishes" + ); + + let _ = thread_b + .thread + .submit(Op::Shutdown {}) + .await + .expect("shutdown should submit"); +} + +#[tokio::test] +async fn wait_agent_wait_for_all_returns_all_final_statuses() { + let (mut session, turn) = make_session_and_context().await; + let manager = thread_manager(); + session.services.agent_control = manager.agent_control(); + let config = turn.config.as_ref().clone(); + let thread_a = manager + .start_thread(config.clone()) + .await + .expect("start thread"); + let thread_b = manager.start_thread(config).await.expect("start thread"); + let id_a = thread_a.thread_id; + let id_b = thread_b.thread_id; + let thread_a_handle = Arc::clone(&thread_a.thread); + let thread_b_handle = Arc::clone(&thread_b.thread); + + tokio::spawn(async move { + sleep(Duration::from_millis(20)).await; + let _ = thread_a_handle.submit(Op::Shutdown {}).await; + }); + tokio::spawn(async move { + sleep(Duration::from_millis(40)).await; + let _ = thread_b_handle.submit(Op::Shutdown {}).await; + }); + + let invocation = invocation( + Arc::new(session), + Arc::new(turn), + "wait_agent", + function_payload(json!({ + "ids": [id_a.to_string(), id_b.to_string()], + "wait_for_all": true, + "timeout_ms": 1000 + })), + ); + let output = WaitAgentHandler + .handle(invocation) + .await + .expect("wait_agent should succeed"); + let (content, success) = expect_text_output(output); + let result: wait::WaitAgentResult = + serde_json::from_str(&content).expect("wait_agent result should be json"); + assert_eq!( + result, + wait::WaitAgentResult { + status: HashMap::from([(id_a, AgentStatus::Shutdown), (id_b, AgentStatus::Shutdown)]), + timed_out: false + } + ); + assert_eq!(success, None); +} + #[tokio::test] async fn close_agent_submits_shutdown_and_returns_previous_status() { let (mut session, turn) = make_session_and_context().await; diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 45fadae7a5..0ea794bf78 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -178,12 +178,12 @@ fn wait_output_schema() -> JsonValue { "properties": { "status": { "type": "object", - "description": "Final statuses keyed by agent id for agents that finished before the timeout.", + "description": "Final statuses keyed by agent id for agents that finished before the timeout or before the requested wait condition was satisfied.", "additionalProperties": agent_status_output_schema() }, "timed_out": { "type": "boolean", - "description": "Whether the wait call returned due to timeout before any agent reached a final status." + "description": "Whether the wait call returned due to timeout before the requested wait condition was satisfied. With wait_for_all=true, partial statuses may still be returned." } }, "required": ["status", "timed_out"], @@ -1114,6 +1114,7 @@ fn create_spawn_agent_tool(config: &ToolsConfig) -> ToolSpec { ### After you delegate - Call wait_agent very sparingly. Only call wait_agent when you need the result immediately for the next critical-path step and you are blocked until it returns. +- If you launch a fixed batch and need every result before the next step, spawn the whole batch first and use one wait_agent call with wait_for_all=true and a sufficiently long timeout instead of a short wait loop. - Do not redo delegated subagent tasks yourself; focus on integrating results or tackling non-overlapping work. - While the subagent is running in the background, do meaningful non-overlapping work immediately. - Do not repeatedly wait by reflex. @@ -1370,7 +1371,16 @@ fn create_wait_agent_tool() -> ToolSpec { JsonSchema::Array { items: Box::new(JsonSchema::String { description: None }), description: Some( - "Agent ids to wait on. Pass multiple ids to wait for whichever finishes first." + "Agent ids to wait on. Pass multiple ids to wait for whichever finishes first unless wait_for_all is true." + .to_string(), + ), + }, + ); + properties.insert( + "wait_for_all".to_string(), + JsonSchema::Boolean { + description: Some( + "When true, wait until every requested agent reaches a final status or the timeout expires. When false (default), return after the first requested agent reaches a final status." .to_string(), ), }, @@ -1386,7 +1396,7 @@ fn create_wait_agent_tool() -> ToolSpec { ToolSpec::Function(ResponsesApiTool { name: "wait_agent".to_string(), - description: "Wait for agents to reach a final status. Completed statuses may include the agent's final message. Returns empty status when timed out. Once the agent reaches a final status, a notification message will be received containing the same completed status." + description: "Wait for agents to reach a final status. By default this returns when any requested agent finishes; set wait_for_all=true to wait for the whole batch. Completed statuses may include the agent's final message. Returns empty status when timed out before any requested agent finishes; with wait_for_all=true, partial statuses may still be returned on timeout." .to_string(), strict: false, defer_loading: None, diff --git a/codex-rs/core/templates/agents/orchestrator.md b/codex-rs/core/templates/agents/orchestrator.md index 39d86c2c2c..8bf94bf23f 100644 --- a/codex-rs/core/templates/agents/orchestrator.md +++ b/codex-rs/core/templates/agents/orchestrator.md @@ -100,7 +100,8 @@ Sub-agents are their to make you go fast and time is a big constraint so leverag ## Flow 1. Understand the task. -2. Spawn the optimal necessary sub-agents. +2. Spawn the optimal necessary sub-agents. When multiple independent agents are needed, launch the whole batch up front before waiting. 3. Coordinate them via wait_agent / send_input. + If you need every agent result before the next step, prefer one wait_agent call with `wait_for_all=true` and a long timeout instead of repeatedly waiting for whichever agent finishes first. 4. Iterate on this. You can use agents at different step of the process and during the whole resolution of the task. Never forget to use them. 5. Ask the user before shutting sub-agents down unless you need to because you reached the agent limit. diff --git a/codex-rs/core/templates/collab/experimental_prompt.md b/codex-rs/core/templates/collab/experimental_prompt.md index 1c390adec9..190cf4aace 100644 --- a/codex-rs/core/templates/collab/experimental_prompt.md +++ b/codex-rs/core/templates/collab/experimental_prompt.md @@ -12,4 +12,5 @@ This feature must be used wisely. For simple or straightforward tasks, you don't * Running tests or some config commands can output a large amount of logs. In order to optimize your own context, you can spawn an agent and ask it to do it for you. In such cases, you must tell this agent that it can't spawn another agent himself (to prevent infinite recursion) * When you're done with a sub-agent, don't forget to close it using `close_agent`. * Be careful on the `timeout_ms` parameter you choose for `wait_agent`. It should be wisely scaled. +* If you launch a fixed batch and need every result before the next step, spawn the whole batch first and use one `wait_agent(..., wait_for_all=true)` call with a long timeout instead of a short wait loop. * Sub-agents have access to the same set of tools as you do so you must tell them if they are allowed to spawn sub-agents themselves or not.