diff --git a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs index 91f5990210..db7ea3a44c 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs @@ -8,6 +8,9 @@ use crate::agent::role::DEFAULT_ROLE_NAME; use crate::agent::role::apply_role_to_config; use crate::tools::handlers::multi_agents_spec::SpawnAgentToolOptions; use crate::tools::handlers::multi_agents_spec::create_spawn_agent_tool_v1; +use crate::tools::hook_names::HookToolName; +use crate::tools::registry::PostToolUsePayload; +use crate::tools::registry::PreToolUsePayload; use crate::turn_timing::now_unix_timestamp_ms; use codex_tools::ToolSpec; @@ -202,6 +205,35 @@ impl ToolHandler for Handler { fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } + + fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option { + Some(PreToolUsePayload { + tool_name: HookToolName::spawn_agent(), + tool_input: collab_hook_tool_input(&invocation.payload)?, + }) + } + + fn with_updated_hook_input( + &self, + invocation: ToolInvocation, + updated_input: JsonValue, + ) -> Result { + rewrite_collab_hook_input(invocation, updated_input) + } + + fn post_tool_use_payload( + &self, + invocation: &ToolInvocation, + result: &Self::Output, + ) -> Option { + Some(PostToolUsePayload { + tool_name: HookToolName::spawn_agent(), + tool_use_id: invocation.call_id.clone(), + tool_input: collab_hook_tool_input(&invocation.payload)?, + tool_response: result + .post_tool_use_response(&invocation.call_id, &invocation.payload)?, + }) + } } #[derive(Debug, Deserialize)] @@ -238,4 +270,8 @@ impl ToolOutput for SpawnAgentResult { fn code_mode_result(&self, _payload: &ToolPayload) -> JsonValue { tool_output_code_mode_result(self, "spawn_agent") } + + fn post_tool_use_response(&self, _call_id: &str, _payload: &ToolPayload) -> Option { + serde_json::to_value(self).ok() + } } diff --git a/codex-rs/core/src/tools/handlers/multi_agents_common.rs b/codex-rs/core/src/tools/handlers/multi_agents_common.rs index 968810494d..7a74070bc5 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_common.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_common.rs @@ -6,6 +6,7 @@ use crate::function_tool::FunctionCallError; use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::tools::context::FunctionToolOutput; +use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; use codex_features::Feature; @@ -41,6 +42,41 @@ pub(crate) fn function_arguments(payload: ToolPayload) -> Result Option { + let ToolPayload::Function { arguments } = payload else { + return None; + }; + + if arguments.trim().is_empty() { + return Some(JsonValue::Object(serde_json::Map::new())); + } + + Some( + serde_json::from_str(arguments) + .unwrap_or_else(|_| JsonValue::String(arguments.to_string())), + ) +} + +pub(crate) fn rewrite_collab_hook_input( + mut invocation: ToolInvocation, + updated_input: JsonValue, +) -> Result { + let ToolPayload::Function { .. } = invocation.payload else { + return Err(FunctionCallError::RespondToModel( + "hook input rewrite received unsupported collaboration payload".to_string(), + )); + }; + + invocation.payload = ToolPayload::Function { + arguments: serde_json::to_string(&updated_input).map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to serialize rewritten collaboration arguments: {err}" + )) + })?, + }; + Ok(invocation) +} + pub(crate) fn tool_output_json_text(value: &T, tool_name: &str) -> String where T: Serialize, 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 c6a0007f13..0ec2ab17c6 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -14,6 +14,10 @@ use crate::tools::handlers::multi_agents_v2::ListAgentsHandler as ListAgentsHand use crate::tools::handlers::multi_agents_v2::SendMessageHandler as SendMessageHandlerV2; use crate::tools::handlers::multi_agents_v2::SpawnAgentHandler as SpawnAgentHandlerV2; use crate::tools::handlers::multi_agents_v2::WaitAgentHandler as WaitAgentHandlerV2; +use crate::tools::hook_names::HookToolName; +use crate::tools::registry::PostToolUsePayload; +use crate::tools::registry::PreToolUsePayload; +use crate::tools::registry::ToolHandler; use crate::turn_diff_tracker::TurnDiffTracker; use codex_extension_api::empty_extension_registry; use codex_features::Feature; @@ -193,6 +197,104 @@ async fn handler_rejects_non_function_payloads() { ); } +#[tokio::test] +async fn spawn_agent_hook_payloads_use_agent_alias_and_support_rewrites() { + let (session, turn) = make_session_and_context().await; + let handler = SpawnAgentHandler::default(); + let invocation = invocation( + Arc::new(session), + Arc::new(turn), + "spawn_agent", + function_payload(json!({ + "message": "inspect this repo", + "agent_type": "explorer" + })), + ); + + assert_eq!( + handler.pre_tool_use_payload(&invocation), + Some(PreToolUsePayload { + tool_name: HookToolName::spawn_agent(), + tool_input: json!({ + "message": "inspect this repo", + "agent_type": "explorer" + }), + }) + ); + + let rewritten = handler + .with_updated_hook_input( + invocation, + json!({ + "message": "inspect the tests instead", + "agent_type": "worker" + }), + ) + .expect("spawn hook rewrite should succeed"); + let ToolPayload::Function { arguments } = rewritten.payload else { + panic!("rewritten spawn payload should stay function-shaped"); + }; + assert_eq!( + serde_json::from_str::(&arguments) + .expect("rewritten spawn args should stay json"), + json!({ + "message": "inspect the tests instead", + "agent_type": "worker" + }) + ); +} + +#[tokio::test] +async fn multi_agent_v2_spawn_hook_payloads_use_agent_alias_and_support_rewrites() { + let (session, turn) = make_session_and_context().await; + let handler = SpawnAgentHandlerV2::default(); + let invocation = invocation( + Arc::new(session), + Arc::new(turn), + "spawn_agent", + function_payload(json!({ + "message": "inspect this repo", + "task_name": "scan_repo", + "fork_turns": "none" + })), + ); + + assert_eq!( + handler.pre_tool_use_payload(&invocation), + Some(PreToolUsePayload { + tool_name: HookToolName::spawn_agent(), + tool_input: json!({ + "message": "inspect this repo", + "task_name": "scan_repo", + "fork_turns": "none" + }), + }) + ); + + let rewritten = handler + .with_updated_hook_input( + invocation, + json!({ + "message": "inspect hook tests", + "task_name": "scan_hooks", + "fork_turns": "none" + }), + ) + .expect("v2 spawn hook rewrite should succeed"); + let ToolPayload::Function { arguments } = rewritten.payload else { + panic!("rewritten v2 spawn payload should stay function-shaped"); + }; + assert_eq!( + serde_json::from_str::(&arguments) + .expect("rewritten v2 spawn args should stay json"), + json!({ + "message": "inspect hook tests", + "task_name": "scan_hooks", + "fork_turns": "none" + }) + ); +} + #[tokio::test] async fn spawn_agent_rejects_empty_message() { let (session, turn) = make_session_and_context().await; @@ -211,6 +313,87 @@ async fn spawn_agent_rejects_empty_message() { ); } +#[tokio::test] +async fn spawn_agent_post_tool_use_payload_exposes_spawn_result() { + let (mut session, turn) = make_session_and_context().await; + let manager = thread_manager(); + session.services.agent_control = manager.agent_control(); + let handler = SpawnAgentHandler::default(); + let invocation = invocation( + Arc::new(session), + Arc::new(turn), + "spawn_agent", + function_payload(json!({ + "message": "inspect this repo" + })), + ); + let result = handler + .handle(invocation.clone()) + .await + .expect("spawn_agent should succeed"); + + assert_eq!( + handler.post_tool_use_payload(&invocation, &result), + Some(PostToolUsePayload { + tool_name: HookToolName::spawn_agent(), + tool_use_id: "call-1".to_string(), + tool_input: json!({ + "message": "inspect this repo" + }), + tool_response: serde_json::to_value(&result) + .expect("spawn result should serialize for hooks"), + }) + ); +} + +#[tokio::test] +async fn multi_agent_v2_spawn_post_tool_use_payload_exposes_spawn_result() { + let (mut session, mut turn) = make_session_and_context().await; + let mut config = (*turn.config).clone(); + config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + turn.config = Arc::new(config); + 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.conversation_id = root.thread_id; + let handler = SpawnAgentHandlerV2::default(); + let invocation = invocation( + Arc::new(session), + Arc::new(turn), + "spawn_agent", + function_payload(json!({ + "message": "inspect this repo", + "task_name": "post_hook_worker", + "fork_turns": "none" + })), + ); + let result = handler + .handle(invocation.clone()) + .await + .expect("multi-agent v2 spawn_agent should succeed"); + + assert_eq!( + handler.post_tool_use_payload(&invocation, &result), + Some(PostToolUsePayload { + tool_name: HookToolName::spawn_agent(), + tool_use_id: "call-1".to_string(), + tool_input: json!({ + "message": "inspect this repo", + "task_name": "post_hook_worker", + "fork_turns": "none" + }), + tool_response: serde_json::to_value(&result) + .expect("v2 spawn result should serialize for hooks"), + }) + ); +} + #[tokio::test] async fn spawn_agent_rejects_when_message_and_items_are_both_set() { let (session, turn) = make_session_and_context().await; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs index 3ad7b87145..9704a952e2 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs @@ -7,6 +7,9 @@ use crate::agent::role::DEFAULT_ROLE_NAME; use crate::agent::role::apply_role_to_config; use crate::tools::handlers::multi_agents_spec::SpawnAgentToolOptions; use crate::tools::handlers::multi_agents_spec::create_spawn_agent_tool_v2; +use crate::tools::hook_names::HookToolName; +use crate::tools::registry::PostToolUsePayload; +use crate::tools::registry::PreToolUsePayload; use crate::turn_timing::now_unix_timestamp_ms; use codex_protocol::AgentPath; use codex_protocol::protocol::InterAgentCommunication; @@ -233,6 +236,35 @@ impl ToolHandler for Handler { fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } + + fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option { + Some(PreToolUsePayload { + tool_name: HookToolName::spawn_agent(), + tool_input: collab_hook_tool_input(&invocation.payload)?, + }) + } + + fn with_updated_hook_input( + &self, + invocation: ToolInvocation, + updated_input: JsonValue, + ) -> Result { + rewrite_collab_hook_input(invocation, updated_input) + } + + fn post_tool_use_payload( + &self, + invocation: &ToolInvocation, + result: &Self::Output, + ) -> Option { + Some(PostToolUsePayload { + tool_name: HookToolName::spawn_agent(), + tool_use_id: invocation.call_id.clone(), + tool_input: collab_hook_tool_input(&invocation.payload)?, + tool_response: result + .post_tool_use_response(&invocation.call_id, &invocation.payload)?, + }) + } } #[derive(Debug, Deserialize)] @@ -313,4 +345,8 @@ impl ToolOutput for SpawnAgentResult { fn code_mode_result(&self, _payload: &ToolPayload) -> JsonValue { tool_output_code_mode_result(self, "spawn_agent") } + + fn post_tool_use_response(&self, _call_id: &str, _payload: &ToolPayload) -> Option { + serde_json::to_value(self).ok() + } } diff --git a/codex-rs/core/src/tools/hook_names.rs b/codex-rs/core/src/tools/hook_names.rs index 9d3b6c2409..c2178c8dc9 100644 --- a/codex-rs/core/src/tools/hook_names.rs +++ b/codex-rs/core/src/tools/hook_names.rs @@ -43,6 +43,18 @@ impl HookToolName { Self::new("Bash") } + /// Returns the hook identity for sub-agent spawning. + /// + /// The serialized name remains `spawn_agent`, while `Agent` is accepted as + /// a matcher alias for compatibility with hook configurations that describe + /// agent creation using Claude Code-style names. + pub(crate) fn spawn_agent() -> Self { + Self { + name: "spawn_agent".to_string(), + matcher_aliases: vec!["Agent".to_string()], + } + } + /// Returns the canonical hook name serialized into hook stdin. pub(crate) fn name(&self) -> &str { &self.name