agent/codex
1
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: custom watcher for multi-agent v2 (#15576)

Browse Source 
The new wait tool just returns `Wait timed out.` or `Wait completed.`.
The actual content is done through the notification watcher
This commit is contained in:
jif-oai
2026-03-23 23:27:55 +00:00
committed by GitHub
Unverified
parent 0b5ba25b46
commit 527244910f
4 changed files with 123 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/tools/handlers/multi_agents_tests.rs
+60 -11
View File
@@ -9,12 +9,14 @@ use crate::config::types::ShellEnvironmentPolicy;
use crate::function_tool::FunctionCallError;
use crate::protocol::AgentStatus;
use crate::protocol::AskForApproval;
use crate::protocol::EventMsg;
use crate::protocol::FileSystemSandboxPolicy;
use crate::protocol::NetworkSandboxPolicy;
use crate::protocol::Op;
use crate::protocol::SandboxPolicy;
use crate::protocol::SessionSource;
use crate::protocol::SubAgentSource;
use crate::protocol::TurnCompleteEvent;
use crate::state::TaskKind;
use crate::tasks::SessionTask;
use crate::tasks::SessionTaskContext;
@@ -1414,7 +1416,7 @@ async fn multi_agent_v2_wait_agent_accepts_targets_argument() {
assert_eq!(
result,
crate::tools::handlers::multi_agents_v2::wait::WaitAgentResult {
status: HashMap::from([(target, AgentStatus::NotFound)]),
message: "Wait completed.".to_string(),
timed_out: false,
}
);
@@ -1582,12 +1584,7 @@ async fn wait_agent_returns_final_status_without_timeout() {
}
#[tokio::test]
async fn multi_agent_v2_wait_agent_returns_statuses_keyed_by_path() {
#[derive(Debug, Deserialize)]
struct SpawnAgentResult {
task_name: String,
}
async fn multi_agent_v2_wait_agent_returns_summary_for_named_targets() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let root = manager
@@ -1617,9 +1614,7 @@ async fn multi_agent_v2_wait_agent_returns_statuses_keyed_by_path() {
))
.await
.expect("spawn_agent should succeed");
let (content, _) = expect_text_output(spawn_output);
let spawn_result: SpawnAgentResult =
serde_json::from_str(&content).expect("spawn result should parse");
let _ = expect_text_output(spawn_output);
let agent_id = session
.services
@@ -1667,13 +1662,67 @@ async fn multi_agent_v2_wait_agent_returns_statuses_keyed_by_path() {
assert_eq!(
result,
crate::tools::handlers::multi_agents_v2::wait::WaitAgentResult {
status: HashMap::from([(spawn_result.task_name, AgentStatus::Shutdown)]),
message: "Wait completed.".to_string(),
timed_out: false,
}
);
assert_eq!(success, None);
}
#[tokio::test]
async fn multi_agent_v2_wait_agent_does_not_return_completed_content() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
session.services.agent_control = manager.agent_control();
let mut config = (*turn.config).clone();
config
.features
.enable(Feature::MultiAgentV2)
.expect("test config should allow feature update");
turn.config = Arc::new(config.clone());
let thread = manager.start_thread(config).await.expect("start thread");
let agent_id = thread.thread_id;
let child_turn = thread.thread.codex.session.new_default_turn().await;
thread
.thread
.codex
.session
.send_event(
child_turn.as_ref(),
EventMsg::TurnComplete(TurnCompleteEvent {
turn_id: child_turn.sub_id.clone(),
last_agent_message: Some("sensitive child output".to_string()),
}),
)
.await;
let output = WaitAgentHandlerV2
.handle(invocation(
Arc::new(session),
Arc::new(turn),
"wait_agent",
function_payload(json!({
"targets": [agent_id.to_string()],
"timeout_ms": 1000
})),
))
.await
.expect("wait_agent should succeed");
let (content, success) = expect_text_output(output);
let result: crate::tools::handlers::multi_agents_v2::wait::WaitAgentResult =
serde_json::from_str(&content).expect("wait_agent result should be json");
assert_eq!(
result,
crate::tools::handlers::multi_agents_v2::wait::WaitAgentResult {
message: "Wait completed.".to_string(),
timed_out: false,
}
);
assert!(!content.contains("sensitive child output"));
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;
codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs
+16 -22
View File
@@ -35,21 +35,12 @@ impl ToolHandler for Handler {
let args: WaitArgs = parse_arguments(&arguments)?;
let receiver_thread_ids = resolve_agent_targets(&session, &turn, args.targets).await?;
let mut receiver_agents = Vec::with_capacity(receiver_thread_ids.len());
let mut target_by_thread_id = HashMap::with_capacity(receiver_thread_ids.len());
for receiver_thread_id in &receiver_thread_ids {
let agent_metadata = session
.services
.agent_control
.get_agent_metadata(*receiver_thread_id)
.unwrap_or_default();
target_by_thread_id.insert(
*receiver_thread_id,
agent_metadata
.agent_path
.as_ref()
.map(ToString::to_string)
.unwrap_or_else(|| receiver_thread_id.to_string()),
);
receiver_agents.push(CollabAgentRef {
thread_id: *receiver_thread_id,
agent_nickname: agent_metadata.agent_nickname,
@@ -152,18 +143,7 @@ impl ToolHandler for Handler {
let timed_out = statuses.is_empty();
let statuses_by_id = statuses.clone().into_iter().collect::<HashMap<_, _>>();
let agent_statuses = build_wait_agent_statuses(&statuses_by_id, &receiver_agents);
let result = WaitAgentResult {
status: statuses
.into_iter()
.filter_map(|(thread_id, status)| {
target_by_thread_id
.get(&thread_id)
.cloned()
.map(|target| (target, status))
})
.collect(),
timed_out,
};
let result = WaitAgentResult::from_timed_out(timed_out);
session
.send_event(
@@ -191,10 +171,24 @@ struct WaitArgs {
#[derive(Debug, Deserialize, Serialize, PartialEq, Eq)]
pub(crate) struct WaitAgentResult {
pub(crate) status: HashMap<String, AgentStatus>,
pub(crate) message: String,
pub(crate) timed_out: bool,
}
impl WaitAgentResult {
fn from_timed_out(timed_out: bool) -> Self {
let message = if timed_out {
"Wait timed out."
} else {
"Wait completed."
};
Self {
message: message.to_string(),
timed_out,
}
}
}
impl ToolOutput for WaitAgentResult {
fn log_preview(&self) -> String {
tool_output_json_text(self, "wait_agent")
codex-rs/core/src/tools/spec.rs
+44 -21
View File
@@ -178,23 +178,41 @@ fn resume_agent_output_schema() -> JsonValue {
})
}
fn wait_output_schema() -> JsonValue {
json!({
"type": "object",
"properties": {
"status": {
"type": "object",
"description": "Final statuses keyed by canonical task name when available, otherwise by agent id.",
"additionalProperties": agent_status_output_schema()
fn wait_output_schema(multi_agent_v2: bool) -> JsonValue {
if multi_agent_v2 {
json!({
"type": "object",
"properties": {
"message": {
"type": "string",
"description": "Brief wait summary without the agent's final content."
},
"timed_out": {
"type": "boolean",
"description": "Whether the wait call returned due to timeout before any agent reached a final status."
}
},
"timed_out": {
"type": "boolean",
"description": "Whether the wait call returned due to timeout before any agent reached a final status."
}
},
"required": ["status", "timed_out"],
"additionalProperties": false
})
"required": ["message", "timed_out"],
"additionalProperties": false
})
} else {
json!({
"type": "object",
"properties": {
"status": {
"type": "object",
"description": "Final statuses keyed by canonical task name when available, otherwise by agent id.",
"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."
}
},
"required": ["status", "timed_out"],
"additionalProperties": false
})
}
}
fn close_agent_output_schema() -> JsonValue {
@@ -1422,7 +1440,7 @@ fn create_resume_agent_tool() -> ToolSpec {
})
}
fn create_wait_agent_tool() -> ToolSpec {
fn create_wait_agent_tool(multi_agent_v2: bool) -> ToolSpec {
let mut properties = BTreeMap::new();
properties.insert(
"targets".to_string(),
@@ -1445,8 +1463,13 @@ 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."
.to_string(),
description: if multi_agent_v2 {
"Wait for agents to reach a final status. Returns a brief wait summary instead of the agent's final content. Returns a timeout summary when no agent reaches a final status before the deadline."
.to_string()
} else {
"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."
.to_string()
},
strict: false,
defer_loading: None,
parameters: JsonSchema::Object {
@@ -1454,7 +1477,7 @@ fn create_wait_agent_tool() -> ToolSpec {
required: Some(vec!["targets".to_string()]),
additional_properties: Some(false.into()),
},
output_schema: Some(wait_output_schema()),
output_schema: Some(wait_output_schema(multi_agent_v2)),
})
}
@@ -3006,7 +3029,7 @@ pub(crate) fn build_specs_with_discoverable_tools(
}
push_tool_spec(
&mut builder,
create_wait_agent_tool(),
create_wait_agent_tool(config.multi_agent_v2),
/*supports_parallel_tool_calls*/ false,
config.code_mode_enabled,
);
codex-rs/core/src/tools/spec_tests.rs
+3 -3
View File
@@ -469,7 +469,7 @@ fn test_full_toolset_specs_for_gpt5_codex_unified_exec_web_search() {
create_view_image_tool(config.can_request_original_image_detail),
create_spawn_agent_tool(&config),
create_send_input_tool(),
create_wait_agent_tool(),
create_wait_agent_tool(config.multi_agent_v2),
create_close_agent_tool(),
] {
expected.insert(tool_name(&spec).to_string(), spec);
@@ -607,8 +607,8 @@ fn test_build_specs_multi_agent_v2_uses_task_names_and_hides_resume() {
.as_ref()
.expect("wait_agent should define output schema");
assert_eq!(
output_schema["properties"]["status"]["description"],
json!("Final statuses keyed by canonical task name when available, otherwise by agent id.")
output_schema["properties"]["message"]["description"],
json!("Brief wait summary without the agent's final content.")
);
assert_lacks_tool_name(&tools, "resume_agent");
}