Extract pre-tool-use hook handling into a helper (#58357)

## What changed

Move hook execution and input updates into `apply_pre_tool_use`, with distinct errors for blocked calls, failures, and cancellation. Preserve dispatch error reporting and clear the serving assessment when a hook changes the invocation payload.

GitOrigin-RevId: d3c5b37692970b739393cb0c998f6cff0ff39af2
This commit is contained in:
swastik-oai
2026-10-30 15:02:16 +00:00
committed by copyberry
parent 9b8c2bb50b
commit 85d042c71a
+68 -46
View File
@@ -852,54 +852,27 @@ impl ToolRegistry {
Ok(environments) => environments,
Err(err) => return Err(fail_before_execution(&invocation, err).await),
};
if let Some(pre_tool_use_payload) = tool.pre_tool_use_payload(&invocation) {
let tool_cwd = match tool.pre_tool_use_cwd(&invocation) {
Ok(cwd) => cwd,
Err(err) => return Err(fail_before_execution(&invocation, err).await),
};
match run_pre_tool_use_hooks(
&invocation.session,
invocation.step_context.as_ref(),
invocation.call_id.clone(),
&pre_tool_use_payload.tool_name,
&pre_tool_use_payload.tool_input,
tool_cwd,
&invocation.environments.hook_ids(),
)
.or_cancel_tool(&cancellation_token)
.await?
{
PreToolUseHookResult::Blocked(message) => {
if tool.is_builtin_control_tool() {
let mut analytics = ControlToolCallGuard::new(&invocation);
analytics.finish(ControlToolCallStatus::Rejected);
}
let err = FunctionCallError::RespondToModel(message);
dispatch_trace.record_failed(&err);
notify_tool_finish_if_unclaimed(
&invocation,
call_state.as_deref(),
ToolCallOutcome::Blocked,
)
.await;
return Err(err);
match apply_pre_tool_use(tool.as_ref(), &mut invocation).await {
Ok(()) => {}
Err(PreToolUseError::Blocked(message)) => {
if tool.is_builtin_control_tool() {
let mut analytics = ControlToolCallGuard::new(&invocation);
analytics.finish(ControlToolCallStatus::Rejected);
}
PreToolUseHookResult::Continue {
updated_input: Some(updated_input),
} => match tool.with_updated_hook_input(invocation.clone(), updated_input) {
Ok(mut updated_invocation) => {
if updated_invocation.payload != invocation.payload {
Arc::make_mut(&mut updated_invocation.step_context)
.serving_assessment = None;
}
invocation = updated_invocation;
}
Err(err) => return Err(fail_before_execution(&invocation, err).await),
},
PreToolUseHookResult::Continue {
updated_input: None,
} => {}
let err = FunctionCallError::RespondToModel(message);
dispatch_trace.record_failed(&err);
notify_tool_finish_if_unclaimed(
&invocation,
call_state.as_deref(),
ToolCallOutcome::Blocked,
)
.await;
return Err(err);
}
Err(PreToolUseError::Failed(err)) => {
return Err(fail_before_execution(&invocation, err).await);
}
Err(PreToolUseError::Cancelled(err)) => return Err(err),
}
if let Some(assessment) = &invocation.step_context.serving_assessment {
@@ -1104,6 +1077,55 @@ async fn handle_any_tool(
Ok(result)
}
enum PreToolUseError {
Blocked(String),
Failed(FunctionCallError),
Cancelled(FunctionCallError),
}
async fn apply_pre_tool_use(
tool: &dyn CoreToolRuntime,
invocation: &mut ToolInvocation,
) -> Result<(), PreToolUseError> {
let Some(payload) = tool.pre_tool_use_payload(invocation) else {
return Ok(());
};
let cwd = tool
.pre_tool_use_cwd(invocation)
.map_err(PreToolUseError::Failed)?;
let outcome = run_pre_tool_use_hooks(
&invocation.session,
&invocation.step_context,
invocation.call_id.clone(),
&payload.tool_name,
&payload.tool_input,
cwd,
&invocation.environments.hook_ids(),
)
.or_cancel_tool(&invocation.cancellation_token)
.await
.map_err(PreToolUseError::Cancelled)?;
match outcome {
PreToolUseHookResult::Blocked(reason) => Err(PreToolUseError::Blocked(reason)),
PreToolUseHookResult::Continue {
updated_input: None,
} => Ok(()),
PreToolUseHookResult::Continue {
updated_input: Some(input),
} => {
let mut updated = tool
.with_updated_hook_input(invocation.clone(), input)
.map_err(PreToolUseError::Failed)?;
if updated.payload != invocation.payload {
Arc::make_mut(&mut updated.step_context).serving_assessment = None;
}
*invocation = updated;
Ok(())
}
}
}
fn function_hook_tool_name(invocation: &ToolInvocation) -> HookToolName {
if invocation.tool_name.name == "spawn_agent"
&& (invocation.tool_name.is_default_namespace()