Hook spawn_agent into tool-use hooks

This commit is contained in:
Abhinav Vedmala
2026-05-19 12:49:10 -07:00
parent b2e51b85ee
commit d6a4eda990
5 changed files with 303 additions and 0 deletions

View File

@@ -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<PreToolUsePayload> {
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<ToolInvocation, FunctionCallError> {
rewrite_collab_hook_input(invocation, updated_input)
}
fn post_tool_use_payload(
&self,
invocation: &ToolInvocation,
result: &Self::Output,
) -> Option<PostToolUsePayload> {
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<JsonValue> {
serde_json::to_value(self).ok()
}
}

View File

@@ -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<String, Functio
}
}
pub(crate) fn collab_hook_tool_input(payload: &ToolPayload) -> Option<JsonValue> {
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<ToolInvocation, FunctionCallError> {
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<T>(value: &T, tool_name: &str) -> String
where
T: Serialize,

View File

@@ -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::<serde_json::Value>(&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::<serde_json::Value>(&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;

View File

@@ -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<PreToolUsePayload> {
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<ToolInvocation, FunctionCallError> {
rewrite_collab_hook_input(invocation, updated_input)
}
fn post_tool_use_payload(
&self,
invocation: &ToolInvocation,
result: &Self::Output,
) -> Option<PostToolUsePayload> {
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<JsonValue> {
serde_json::to_value(self).ok()
}
}

View File

@@ -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