Centralize pre-hook tool compatibility

This commit is contained in:
Abhinav Vedmala
2026-05-06 14:16:38 -07:00
parent 173156bf69
commit 1cd3e86056
11 changed files with 330 additions and 413 deletions

View File

@@ -22,12 +22,9 @@ use crate::tools::events::ToolEmitter;
use crate::tools::events::ToolEventCtx;
use crate::tools::handlers::apply_granted_turn_permissions;
use crate::tools::handlers::parse_arguments;
use crate::tools::handlers::rewrite_function_string_argument;
use crate::tools::handlers::updated_hook_command;
use crate::tools::hook_names::HookToolName;
use crate::tools::orchestrator::ToolOrchestrator;
use crate::tools::registry::PostToolUsePayload;
use crate::tools::registry::PreToolUsePayload;
use crate::tools::registry::ToolArgumentDiffConsumer;
use crate::tools::registry::ToolHandler;
use crate::tools::registry::ToolKind;
@@ -243,12 +240,8 @@ fn write_permissions_for_paths(
normalize_additional_permissions(permissions).ok()
}
/// Extracts the raw patch text used as the command-shaped hook input for apply_patch.
///
/// The apply_patch tool can arrive as the older JSON/function shape or as a
/// freeform custom tool call. Both represent the same file edit operation, so
/// hooks see the raw patch body in `tool_input.command` either way.
fn apply_patch_payload_command(payload: &ToolPayload) -> Option<String> {
/// Extracts the raw patch text from either supported apply_patch payload form.
pub(crate) fn apply_patch_payload_command(payload: &ToolPayload) -> Option<String> {
match payload {
ToolPayload::Function { arguments } => parse_arguments::<ApplyPatchToolArgs>(arguments)
.ok()
@@ -318,39 +311,6 @@ impl ToolHandler for ApplyPatchHandler {
Some(Box::<ApplyPatchArgumentDiffConsumer>::default())
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
apply_patch_payload_command(&invocation.payload).map(|command| PreToolUsePayload {
tool_name: HookToolName::apply_patch(),
tool_input: serde_json::json!({ "command": command }),
})
}
// Hooks expose apply_patch through the stable `{ "command": ... }` shape,
// while the underlying tool stores the patch as either the function
// argument `input` or freeform custom-tool input.
fn with_updated_hook_input(
&self,
mut invocation: ToolInvocation,
updated_input: serde_json::Value,
) -> Result<ToolInvocation, FunctionCallError> {
let patch = updated_hook_command(&updated_input)?;
invocation.payload = match invocation.payload {
ToolPayload::Function { arguments } => ToolPayload::Function {
arguments: rewrite_function_string_argument(
&arguments,
"apply_patch",
"input",
patch,
)?,
},
ToolPayload::Custom { .. } => ToolPayload::Custom {
input: patch.to_string(),
},
payload => payload,
};
Ok(invocation)
}
fn post_tool_use_payload(
&self,
invocation: &ToolInvocation,

View File

@@ -49,10 +49,8 @@ async fn pre_tool_use_payload_uses_json_patch_input() {
arguments: json!({ "input": patch }).to_string(),
};
let invocation = invocation_for_payload(payload).await;
let handler = ApplyPatchHandler;
assert_eq!(
handler.pre_tool_use_payload(&invocation),
crate::tools::hook_compat::pre_tool_use_payload(&invocation),
Some(PreToolUsePayload {
tool_name: HookToolName::apply_patch(),
tool_input: json!({ "command": patch }),
@@ -67,10 +65,8 @@ async fn pre_tool_use_payload_uses_freeform_patch_input() {
input: patch.to_string(),
};
let invocation = invocation_for_payload(payload).await;
let handler = ApplyPatchHandler;
assert_eq!(
handler.pre_tool_use_payload(&invocation),
crate::tools::hook_compat::pre_tool_use_payload(&invocation),
Some(PreToolUsePayload {
tool_name: HookToolName::apply_patch(),
tool_input: json!({ "command": patch }),

View File

@@ -10,11 +10,9 @@ use crate::tools::context::ToolOutput;
use crate::tools::context::ToolPayload;
use crate::tools::hook_names::HookToolName;
use crate::tools::registry::PostToolUsePayload;
use crate::tools::registry::PreToolUsePayload;
use crate::tools::registry::ToolHandler;
use crate::tools::registry::ToolKind;
use codex_tools::ToolName;
use serde_json::Value;
pub struct McpHandler {
tool_name: ToolName,
@@ -37,40 +35,6 @@ impl ToolHandler for McpHandler {
ToolKind::Mcp
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
let ToolPayload::Mcp { raw_arguments, .. } = &invocation.payload else {
return None;
};
Some(PreToolUsePayload {
tool_name: HookToolName::new(self.tool_name.display()),
tool_input: mcp_hook_tool_input(raw_arguments),
})
}
// MCP hooks expose the full arguments object, so rewrites replace the
// serialized raw argument payload wholesale rather than patching one
// handler-owned field.
fn with_updated_hook_input(
&self,
mut invocation: ToolInvocation,
updated_input: Value,
) -> Result<ToolInvocation, FunctionCallError> {
invocation.payload = match invocation.payload {
ToolPayload::Mcp { server, tool, .. } => ToolPayload::Mcp {
server,
tool,
raw_arguments: serde_json::to_string(&updated_input).map_err(|err| {
FunctionCallError::RespondToModel(format!(
"failed to serialize rewritten MCP arguments: {err}"
))
})?,
},
payload => payload,
};
Ok(invocation)
}
fn post_tool_use_payload(
&self,
invocation: &ToolInvocation,
@@ -137,14 +101,6 @@ impl ToolHandler for McpHandler {
}
}
fn mcp_hook_tool_input(raw_arguments: &str) -> Value {
if raw_arguments.trim().is_empty() {
return Value::Object(serde_json::Map::new());
}
serde_json::from_str(raw_arguments).unwrap_or_else(|_| Value::String(raw_arguments.to_string()))
}
#[cfg(test)]
mod tests {
use super::*;
@@ -170,13 +126,8 @@ mod tests {
.to_string(),
};
let (session, turn) = make_session_and_context().await;
let handler = McpHandler::new(codex_tools::ToolName::namespaced(
"mcp__memory__",
"create_entities",
));
assert_eq!(
handler.pre_tool_use_payload(&ToolInvocation {
crate::tools::hook_compat::pre_tool_use_payload(&ToolInvocation {
session: session.into(),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
@@ -186,7 +137,7 @@ mod tests {
source: ToolCallSource::Direct,
payload,
}),
Some(PreToolUsePayload {
Some(crate::tools::registry::PreToolUsePayload {
tool_name: HookToolName::new("mcp__memory__create_entities"),
tool_input: json!({
"entities": [{
@@ -262,6 +213,9 @@ mod tests {
#[test]
fn mcp_hook_tool_input_defaults_empty_args_to_object() {
assert_eq!(mcp_hook_tool_input(" "), json!({}));
assert_eq!(
crate::tools::hook_compat::mcp_hook_tool_input(" "),
json!({})
);
}
}

View File

@@ -24,7 +24,6 @@ use codex_sandboxing::policy_transforms::normalize_additional_permissions;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_absolute_path::AbsolutePathBufGuard;
use serde::Deserialize;
use serde_json::Map;
use serde_json::Value;
use std::path::Path;
@@ -54,6 +53,9 @@ pub use shell::ContainerExecHandler;
pub use shell::LocalShellHandler;
pub use shell::ShellCommandHandler;
pub use shell::ShellHandler;
pub(crate) use shell::local_shell_payload_command;
pub(crate) use shell::shell_command_payload_command;
pub(crate) use shell::shell_function_payload_command;
pub use test_sync::TestSyncHandler;
pub use tool_search::ToolSearchHandler;
pub use unavailable_tool::UnavailableToolHandler;
@@ -62,7 +64,7 @@ pub use unified_exec::ExecCommandHandler;
pub use unified_exec::WriteStdinHandler;
pub use view_image::ViewImageHandler;
fn parse_arguments<T>(arguments: &str) -> Result<T, FunctionCallError>
pub(crate) fn parse_arguments<T>(arguments: &str) -> Result<T, FunctionCallError>
where
T: for<'de> Deserialize<'de>,
{
@@ -82,47 +84,6 @@ where
parse_arguments(arguments)
}
fn updated_hook_command(updated_input: &Value) -> Result<&str, FunctionCallError> {
updated_input
.get("command")
.and_then(Value::as_str)
.ok_or_else(|| {
FunctionCallError::RespondToModel(
"hook returned updatedInput without string field `command`".to_string(),
)
})
}
fn rewrite_function_arguments(
arguments: &str,
tool_name: &str,
rewrite: impl FnOnce(&mut Map<String, Value>),
) -> Result<String, FunctionCallError> {
let mut arguments: Value = parse_arguments(arguments)?;
let Value::Object(arguments) = &mut arguments else {
return Err(FunctionCallError::RespondToModel(format!(
"{tool_name} arguments must be an object"
)));
};
rewrite(arguments);
serde_json::to_string(&arguments).map_err(|err| {
FunctionCallError::RespondToModel(format!(
"failed to serialize rewritten {tool_name} arguments: {err}"
))
})
}
fn rewrite_function_string_argument(
arguments: &str,
tool_name: &str,
field_name: &str,
value: &str,
) -> Result<String, FunctionCallError> {
rewrite_function_arguments(arguments, tool_name, |arguments| {
arguments.insert(field_name.to_string(), Value::String(value.to_string()));
})
}
fn resolve_workdir_base_path(
arguments: &str,
default_cwd: &AbsolutePathBuf,

View File

@@ -25,13 +25,9 @@ use crate::tools::handlers::normalize_and_validate_additional_permissions;
use crate::tools::handlers::parse_arguments;
use crate::tools::handlers::parse_arguments_with_base_path;
use crate::tools::handlers::resolve_workdir_base_path;
use crate::tools::handlers::rewrite_function_arguments;
use crate::tools::handlers::rewrite_function_string_argument;
use crate::tools::handlers::updated_hook_command;
use crate::tools::hook_names::HookToolName;
use crate::tools::orchestrator::ToolOrchestrator;
use crate::tools::registry::PostToolUsePayload;
use crate::tools::registry::PreToolUsePayload;
use crate::tools::registry::ToolHandler;
use crate::tools::registry::ToolKind;
use crate::tools::runtimes::shell::ShellRequest;
@@ -59,7 +55,7 @@ pub struct ShellCommandHandler {
backend: ShellCommandBackend,
}
fn shell_function_payload_command(payload: &ToolPayload) -> Option<String> {
pub(crate) fn shell_function_payload_command(payload: &ToolPayload) -> Option<String> {
let ToolPayload::Function { arguments } = payload else {
return None;
};
@@ -69,7 +65,7 @@ fn shell_function_payload_command(payload: &ToolPayload) -> Option<String> {
.map(|params| codex_shell_command::parse_command::shlex_join(&params.command))
}
fn local_shell_payload_command(payload: &ToolPayload) -> Option<String> {
pub(crate) fn local_shell_payload_command(payload: &ToolPayload) -> Option<String> {
let ToolPayload::LocalShell { params } = payload else {
return None;
};
@@ -79,7 +75,7 @@ fn local_shell_payload_command(payload: &ToolPayload) -> Option<String> {
))
}
fn shell_command_payload_command(payload: &ToolPayload) -> Option<String> {
pub(crate) fn shell_command_payload_command(payload: &ToolPayload) -> Option<String> {
let ToolPayload::Function { arguments } = payload else {
return None;
};
@@ -89,35 +85,6 @@ fn shell_command_payload_command(payload: &ToolPayload) -> Option<String> {
.map(|params| params.command)
}
// Hooks expose legacy function shell tools as joined command strings, while
// their function payload stores argv. Split on the way back in to invert the
// hook-facing representation.
fn rewrite_shell_function_updated_hook_input(
mut invocation: ToolInvocation,
updated_input: JsonValue,
tool_name: &str,
) -> Result<ToolInvocation, FunctionCallError> {
let ToolPayload::Function { arguments } = invocation.payload else {
return Err(FunctionCallError::RespondToModel(format!(
"hook input rewrite received unsupported {tool_name} payload"
)));
};
let command = shlex::split(updated_hook_command(&updated_input)?).ok_or_else(|| {
FunctionCallError::RespondToModel(
"hook returned shell input with an invalid command string".to_string(),
)
})?;
invocation.payload = ToolPayload::Function {
arguments: rewrite_function_arguments(&arguments, tool_name, |arguments| {
arguments.insert(
"command".to_string(),
JsonValue::Array(command.into_iter().map(JsonValue::String).collect()),
);
})?,
};
Ok(invocation)
}
struct RunExecLikeArgs {
tool_name: String,
exec_params: ExecParams,
@@ -247,18 +214,6 @@ impl ToolHandler for ShellHandler {
.unwrap_or(true)
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
shell_function_pre_tool_use_payload(invocation)
}
fn with_updated_hook_input(
&self,
invocation: ToolInvocation,
updated_input: JsonValue,
) -> Result<ToolInvocation, FunctionCallError> {
rewrite_shell_function_updated_hook_input(invocation, updated_input, "shell")
}
fn post_tool_use_payload(
&self,
invocation: &ToolInvocation,
@@ -333,18 +288,6 @@ impl ToolHandler for ContainerExecHandler {
.unwrap_or(true)
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
shell_function_pre_tool_use_payload(invocation)
}
fn with_updated_hook_input(
&self,
invocation: ToolInvocation,
updated_input: JsonValue,
) -> Result<ToolInvocation, FunctionCallError> {
rewrite_shell_function_updated_hook_input(invocation, updated_input, "container.exec")
}
fn post_tool_use_payload(
&self,
invocation: &ToolInvocation,
@@ -417,50 +360,6 @@ impl ToolHandler for LocalShellHandler {
!is_known_safe_command(&params.command)
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
local_shell_payload_command(&invocation.payload).map(|command| PreToolUsePayload {
tool_name: HookToolName::bash(),
tool_input: serde_json::json!({ "command": command }),
})
}
// Hooks see a joined shell command string, but local_shell stores argv.
// Split on the way back in to invert the hook-facing representation.
fn with_updated_hook_input(
&self,
mut invocation: ToolInvocation,
updated_input: JsonValue,
) -> Result<ToolInvocation, FunctionCallError> {
let command = updated_hook_command(&updated_input)?;
invocation.payload = match invocation.payload {
ToolPayload::Function { arguments } => {
let command = shlex::split(command).ok_or_else(|| {
FunctionCallError::RespondToModel(
"hook returned shell input with an invalid command string".to_string(),
)
})?;
ToolPayload::Function {
arguments: rewrite_function_arguments(&arguments, "shell", |arguments| {
arguments.insert(
"command".to_string(),
JsonValue::Array(command.into_iter().map(JsonValue::String).collect()),
);
})?,
}
}
ToolPayload::LocalShell { mut params } => {
params.command = shlex::split(command).ok_or_else(|| {
FunctionCallError::RespondToModel(
"hook returned shell input with an invalid command string".to_string(),
)
})?;
ToolPayload::LocalShell { params }
}
payload => payload,
};
Ok(invocation)
}
fn post_tool_use_payload(
&self,
invocation: &ToolInvocation,
@@ -512,13 +411,6 @@ impl ToolHandler for LocalShellHandler {
}
}
fn shell_function_pre_tool_use_payload(invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
shell_function_payload_command(&invocation.payload).map(|command| PreToolUsePayload {
tool_name: HookToolName::bash(),
tool_input: serde_json::json!({ "command": command }),
})
}
fn shell_function_post_tool_use_payload(
invocation: &ToolInvocation,
result: &FunctionToolOutput,
@@ -569,34 +461,6 @@ impl ToolHandler for ShellCommandHandler {
.unwrap_or(true)
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
shell_command_payload_command(&invocation.payload).map(|command| PreToolUsePayload {
tool_name: HookToolName::bash(),
tool_input: serde_json::json!({ "command": command }),
})
}
fn with_updated_hook_input(
&self,
mut invocation: ToolInvocation,
updated_input: JsonValue,
) -> Result<ToolInvocation, FunctionCallError> {
let ToolPayload::Function { arguments } = invocation.payload else {
return Err(FunctionCallError::RespondToModel(
"hook input rewrite received unsupported shell_command payload".to_string(),
));
};
invocation.payload = ToolPayload::Function {
arguments: rewrite_function_string_argument(
&arguments,
"shell_command",
"command",
updated_hook_command(&updated_input)?,
)?,
};
Ok(invocation)
}
fn post_tool_use_payload(
&self,
invocation: &ToolInvocation,

View File

@@ -17,10 +17,7 @@ use crate::tools::context::FunctionToolOutput;
use crate::tools::context::ToolCallSource;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolPayload;
use crate::tools::handlers::ContainerExecHandler;
use crate::tools::handlers::LocalShellHandler;
use crate::tools::handlers::ShellCommandHandler;
use crate::tools::handlers::ShellHandler;
use crate::tools::hook_names::HookToolName;
use crate::tools::registry::ToolHandler;
use crate::turn_diff_tracker::TurnDiffTracker;
@@ -224,10 +221,8 @@ async fn local_shell_pre_tool_use_payload_uses_joined_command() {
},
};
let (session, turn) = make_session_and_context().await;
let handler = LocalShellHandler;
assert_eq!(
handler.pre_tool_use_payload(&ToolInvocation {
crate::tools::hook_compat::pre_tool_use_payload(&ToolInvocation {
session: session.into(),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
@@ -250,12 +245,8 @@ async fn shell_command_pre_tool_use_payload_uses_raw_command() {
arguments: json!({ "command": "printf shell command" }).to_string(),
};
let (session, turn) = make_session_and_context().await;
let handler = ShellCommandHandler {
backend: super::ShellCommandBackend::Classic,
};
assert_eq!(
handler.pre_tool_use_payload(&ToolInvocation {
crate::tools::hook_compat::pre_tool_use_payload(&ToolInvocation {
session: session.into(),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
@@ -282,23 +273,20 @@ async fn shell_handler_rewrites_hook_command_back_to_argv() {
.to_string(),
};
let (session, turn) = make_session_and_context().await;
let handler = ShellHandler;
let invocation = handler
.with_updated_hook_input(
ToolInvocation {
session: session.into(),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
call_id: "call-43".to_string(),
tool_name: codex_tools::ToolName::plain("shell"),
source: ToolCallSource::Direct,
payload,
},
json!({ "command": "bash -lc 'printf new'" }),
)
.expect("shell rewrite should succeed");
let invocation = crate::tools::hook_compat::apply_updated_input(
ToolInvocation {
session: session.into(),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
call_id: "call-43".to_string(),
tool_name: codex_tools::ToolName::plain("shell"),
source: ToolCallSource::Direct,
payload,
},
json!({ "command": "bash -lc 'printf new'" }),
)
.expect("shell rewrite should succeed");
let ToolPayload::Function { arguments } = invocation.payload else {
panic!("shell rewrite should preserve a function payload");
@@ -332,23 +320,20 @@ async fn container_exec_handler_rewrites_hook_command_back_to_argv() {
.to_string(),
};
let (session, turn) = make_session_and_context().await;
let handler = ContainerExecHandler;
let invocation = handler
.with_updated_hook_input(
ToolInvocation {
session: session.into(),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
call_id: "call-44".to_string(),
tool_name: codex_tools::ToolName::plain("container.exec"),
source: ToolCallSource::Direct,
payload,
},
json!({ "command": "bash -lc 'printf new'" }),
)
.expect("container.exec rewrite should succeed");
let invocation = crate::tools::hook_compat::apply_updated_input(
ToolInvocation {
session: session.into(),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),
call_id: "call-44".to_string(),
tool_name: codex_tools::ToolName::plain("container.exec"),
source: ToolCallSource::Direct,
payload,
},
json!({ "command": "bash -lc 'printf new'" }),
)
.expect("container.exec rewrite should succeed");
let ToolPayload::Function { arguments } = invocation.payload else {
panic!("container.exec rewrite should preserve a function payload");

View File

@@ -14,11 +14,8 @@ use crate::tools::handlers::normalize_and_validate_additional_permissions;
use crate::tools::handlers::parse_arguments;
use crate::tools::handlers::parse_arguments_with_base_path;
use crate::tools::handlers::resolve_tool_environment;
use crate::tools::handlers::rewrite_function_string_argument;
use crate::tools::handlers::updated_hook_command;
use crate::tools::hook_names::HookToolName;
use crate::tools::registry::PostToolUsePayload;
use crate::tools::registry::PreToolUsePayload;
use crate::tools::registry::ToolHandler;
use crate::tools::registry::ToolKind;
use crate::unified_exec::ExecCommandRequest;
@@ -48,7 +45,7 @@ pub struct WriteStdinHandler;
#[derive(Debug, Deserialize)]
pub(crate) struct ExecCommandArgs {
cmd: String,
pub(crate) cmd: String,
#[serde(default)]
pub(crate) workdir: Option<String>,
#[serde(default)]
@@ -151,42 +148,6 @@ impl ToolHandler for ExecCommandHandler {
!is_known_safe_command(&command)
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
let ToolPayload::Function { arguments } = &invocation.payload else {
return None;
};
parse_arguments::<ExecCommandArgs>(arguments)
.ok()
.map(|args| PreToolUsePayload {
tool_name: HookToolName::bash(),
tool_input: serde_json::json!({ "command": args.cmd }),
})
}
// Hooks normalize Bash-like tools to `{ "command": ... }`, while the
// exec_command wire schema still names the field `cmd`.
fn with_updated_hook_input(
&self,
mut invocation: ToolInvocation,
updated_input: serde_json::Value,
) -> Result<ToolInvocation, FunctionCallError> {
let ToolPayload::Function { arguments } = invocation.payload else {
return Err(FunctionCallError::RespondToModel(
"hook input rewrite received unsupported exec_command payload".to_string(),
));
};
invocation.payload = ToolPayload::Function {
arguments: rewrite_function_string_argument(
&arguments,
"exec_command",
"cmd",
updated_hook_command(&updated_input)?,
)?,
};
Ok(invocation)
}
fn post_tool_use_payload(
&self,
invocation: &ToolInvocation,

View File

@@ -184,10 +184,8 @@ async fn exec_command_pre_tool_use_payload_uses_raw_command() {
arguments: serde_json::json!({ "cmd": "printf exec command" }).to_string(),
};
let (session, turn) = make_session_and_context().await;
let handler = ExecCommandHandler;
assert_eq!(
handler.pre_tool_use_payload(&ToolInvocation {
crate::tools::hook_compat::pre_tool_use_payload(&ToolInvocation {
session: session.into(),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),
@@ -210,10 +208,8 @@ async fn exec_command_pre_tool_use_payload_skips_write_stdin() {
arguments: serde_json::json!({ "chars": "echo hi" }).to_string(),
};
let (session, turn) = make_session_and_context().await;
let handler = WriteStdinHandler;
assert_eq!(
handler.pre_tool_use_payload(&ToolInvocation {
crate::tools::hook_compat::pre_tool_use_payload(&ToolInvocation {
session: session.into(),
turn: turn.into(),
cancellation_token: tokio_util::sync::CancellationToken::new(),

View File

@@ -0,0 +1,276 @@
use serde_json::Map;
use serde_json::Value;
use crate::function_tool::FunctionCallError;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolPayload;
use crate::tools::handlers::apply_patch::apply_patch_payload_command;
use crate::tools::handlers::local_shell_payload_command;
use crate::tools::handlers::parse_arguments;
use crate::tools::handlers::shell_command_payload_command;
use crate::tools::handlers::shell_function_payload_command;
use crate::tools::handlers::unified_exec::ExecCommandArgs;
use crate::tools::hook_names::HookToolName;
use crate::tools::registry::PreToolUsePayload;
/// Projects native tool payloads into the stable input shape exposed to hooks.
///
/// The hook protocol intentionally hides a few native tool-schema differences:
///
/// - Bash-like tools are exposed as `Bash` with `{ "command": <string> }`,
/// even though their native payloads store commands as argv, `command`, or
/// `cmd` depending on the concrete tool.
/// - `apply_patch` exposes its raw patch body through `{ "command": <patch> }`
/// for compatibility with existing hook consumers.
/// - MCP tools already use arbitrary JSON argument objects, so hooks see those
/// arguments directly rather than through a compatibility shape.
pub(crate) fn pre_tool_use_payload(invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
match invocation.tool_name.name.as_str() {
"shell" | "container.exec" => {
shell_function_payload_command(&invocation.payload).map(bash_payload)
}
"local_shell" => local_shell_payload_command(&invocation.payload).map(bash_payload),
"shell_command" => shell_command_payload_command(&invocation.payload).map(bash_payload),
"exec_command" => exec_command_payload_command(&invocation.payload).map(bash_payload),
"apply_patch" => {
apply_patch_payload_command(&invocation.payload).map(|command| PreToolUsePayload {
tool_name: HookToolName::apply_patch(),
tool_input: serde_json::json!({ "command": command }),
})
}
_ => mcp_payload(invocation),
}
}
/// Rebuilds native tool payloads from hook-facing `updatedInput`.
///
/// This is the inverse of [`pre_tool_use_payload`]: Bash-like and
/// `apply_patch` updates come back through the compatibility `{ "command": ... }`
/// shape and must be written into each tool's native schema, while MCP updates
/// replace the raw argument object wholesale because the hook-facing and native
/// representations are the same.
pub(crate) fn apply_updated_input(
invocation: ToolInvocation,
updated_input: Value,
) -> Result<ToolInvocation, FunctionCallError> {
match invocation.tool_name.name.as_str() {
"shell" => rewrite_shell_function_updated_input(invocation, updated_input, "shell"),
"container.exec" => {
rewrite_shell_function_updated_input(invocation, updated_input, "container.exec")
}
"local_shell" => rewrite_local_shell_updated_input(invocation, updated_input),
"shell_command" => rewrite_shell_command_updated_input(invocation, updated_input),
"exec_command" => rewrite_exec_command_updated_input(invocation, updated_input),
"apply_patch" => rewrite_apply_patch_updated_input(invocation, updated_input),
_ => rewrite_mcp_updated_input(invocation, updated_input),
}
}
fn bash_payload(command: String) -> PreToolUsePayload {
PreToolUsePayload {
tool_name: HookToolName::bash(),
tool_input: serde_json::json!({ "command": command }),
}
}
fn exec_command_payload_command(payload: &ToolPayload) -> Option<String> {
let ToolPayload::Function { arguments } = payload else {
return None;
};
parse_arguments::<ExecCommandArgs>(arguments)
.ok()
.map(|args| args.cmd)
}
fn mcp_payload(invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
let ToolPayload::Mcp { raw_arguments, .. } = &invocation.payload else {
return None;
};
Some(PreToolUsePayload {
tool_name: HookToolName::new(invocation.tool_name.display()),
tool_input: mcp_hook_tool_input(raw_arguments),
})
}
fn updated_hook_command(updated_input: &Value) -> Result<&str, FunctionCallError> {
updated_input
.get("command")
.and_then(Value::as_str)
.ok_or_else(|| {
FunctionCallError::RespondToModel(
"hook returned updatedInput without string field `command`".to_string(),
)
})
}
fn rewrite_function_arguments(
arguments: &str,
tool_name: &str,
rewrite: impl FnOnce(&mut Map<String, Value>),
) -> Result<String, FunctionCallError> {
let mut arguments: Value = parse_arguments(arguments)?;
let Value::Object(arguments) = &mut arguments else {
return Err(FunctionCallError::RespondToModel(format!(
"{tool_name} arguments must be an object"
)));
};
rewrite(arguments);
serde_json::to_string(&arguments).map_err(|err| {
FunctionCallError::RespondToModel(format!(
"failed to serialize rewritten {tool_name} arguments: {err}"
))
})
}
fn rewrite_function_string_argument(
arguments: &str,
tool_name: &str,
field_name: &str,
value: &str,
) -> Result<String, FunctionCallError> {
rewrite_function_arguments(arguments, tool_name, |arguments| {
arguments.insert(field_name.to_string(), Value::String(value.to_string()));
})
}
/// Rehydrates legacy function-style shell tools from hook-facing command text.
fn rewrite_shell_function_updated_input(
mut invocation: ToolInvocation,
updated_input: Value,
tool_name: &str,
) -> Result<ToolInvocation, FunctionCallError> {
let ToolPayload::Function { arguments } = invocation.payload else {
return Err(FunctionCallError::RespondToModel(format!(
"hook input rewrite received unsupported {tool_name} payload"
)));
};
let command = shlex::split(updated_hook_command(&updated_input)?).ok_or_else(|| {
FunctionCallError::RespondToModel(
"hook returned shell input with an invalid command string".to_string(),
)
})?;
invocation.payload = ToolPayload::Function {
arguments: rewrite_function_arguments(&arguments, tool_name, |arguments| {
arguments.insert(
"command".to_string(),
Value::Array(command.into_iter().map(Value::String).collect()),
);
})?,
};
Ok(invocation)
}
/// Rehydrates `local_shell` argv from hook-facing command text.
fn rewrite_local_shell_updated_input(
mut invocation: ToolInvocation,
updated_input: Value,
) -> Result<ToolInvocation, FunctionCallError> {
let command = updated_hook_command(&updated_input)?;
invocation.payload = match invocation.payload {
ToolPayload::LocalShell { mut params } => {
params.command = shlex::split(command).ok_or_else(|| {
FunctionCallError::RespondToModel(
"hook returned shell input with an invalid command string".to_string(),
)
})?;
ToolPayload::LocalShell { params }
}
payload => payload,
};
Ok(invocation)
}
/// Stores hook-facing command text back into the native `shell_command.command`.
fn rewrite_shell_command_updated_input(
mut invocation: ToolInvocation,
updated_input: Value,
) -> Result<ToolInvocation, FunctionCallError> {
let ToolPayload::Function { arguments } = invocation.payload else {
return Err(FunctionCallError::RespondToModel(
"hook input rewrite received unsupported shell_command payload".to_string(),
));
};
invocation.payload = ToolPayload::Function {
arguments: rewrite_function_string_argument(
&arguments,
"shell_command",
"command",
updated_hook_command(&updated_input)?,
)?,
};
Ok(invocation)
}
/// Stores hook-facing command text back into the native `exec_command.cmd`.
fn rewrite_exec_command_updated_input(
mut invocation: ToolInvocation,
updated_input: Value,
) -> Result<ToolInvocation, FunctionCallError> {
let ToolPayload::Function { arguments } = invocation.payload else {
return Err(FunctionCallError::RespondToModel(
"hook input rewrite received unsupported exec_command payload".to_string(),
));
};
invocation.payload = ToolPayload::Function {
arguments: rewrite_function_string_argument(
&arguments,
"exec_command",
"cmd",
updated_hook_command(&updated_input)?,
)?,
};
Ok(invocation)
}
/// Stores hook-facing patch text back into the native apply_patch payload form.
fn rewrite_apply_patch_updated_input(
mut invocation: ToolInvocation,
updated_input: Value,
) -> Result<ToolInvocation, FunctionCallError> {
let patch = updated_hook_command(&updated_input)?;
invocation.payload = match invocation.payload {
ToolPayload::Function { arguments } => ToolPayload::Function {
arguments: rewrite_function_string_argument(&arguments, "apply_patch", "input", patch)?,
},
ToolPayload::Custom { .. } => ToolPayload::Custom {
input: patch.to_string(),
},
payload => payload,
};
Ok(invocation)
}
/// Replaces MCP raw arguments directly because MCP hooks expose that JSON object as-is.
fn rewrite_mcp_updated_input(
mut invocation: ToolInvocation,
updated_input: Value,
) -> Result<ToolInvocation, FunctionCallError> {
invocation.payload = match invocation.payload {
ToolPayload::Mcp { server, tool, .. } => ToolPayload::Mcp {
server,
tool,
raw_arguments: serde_json::to_string(&updated_input).map_err(|err| {
FunctionCallError::RespondToModel(format!(
"failed to serialize rewritten MCP arguments: {err}"
))
})?,
},
payload => {
return Err(FunctionCallError::RespondToModel(format!(
"tool {} does not support hook input rewriting for payload {payload:?}",
invocation.tool_name.display()
)));
}
};
Ok(invocation)
}
pub(crate) fn mcp_hook_tool_input(raw_arguments: &str) -> Value {
if raw_arguments.trim().is_empty() {
return Value::Object(Map::new());
}
serde_json::from_str(raw_arguments).unwrap_or_else(|_| Value::String(raw_arguments.to_string()))
}

View File

@@ -2,6 +2,7 @@ pub(crate) mod code_mode;
pub(crate) mod context;
pub(crate) mod events;
pub(crate) mod handlers;
pub(crate) mod hook_compat;
pub(crate) mod hook_names;
pub(crate) mod network_approval;
pub(crate) mod orchestrator;

View File

@@ -16,6 +16,7 @@ use crate::tools::context::FunctionToolOutput;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolOutput;
use crate::tools::context::ToolPayload;
use crate::tools::hook_compat;
use crate::tools::hook_names::HookToolName;
use crate::tools::tool_dispatch_trace::ToolDispatchTrace;
use codex_hooks::HookEvent;
@@ -69,10 +70,6 @@ pub trait ToolHandler: Send + Sync {
async { false }
}
fn pre_tool_use_payload(&self, _invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
None
}
fn post_tool_use_payload(
&self,
_invocation: &ToolInvocation,
@@ -81,20 +78,6 @@ pub trait ToolHandler: Send + Sync {
None
}
/// Rebuilds a tool invocation from hook-facing `tool_input`.
///
/// Tools that opt into input-rewriting hooks should invert the same stable
/// hook contract they expose from `pre_tool_use_payload`.
fn with_updated_hook_input(
&self,
_invocation: ToolInvocation,
_updated_input: Value,
) -> Result<ToolInvocation, FunctionCallError> {
Err(FunctionCallError::RespondToModel(
"tool does not support hook input rewriting".to_string(),
))
}
/// Creates an optional consumer for streamed tool argument diffs.
fn create_diff_consumer(&self) -> Option<Box<dyn ToolArgumentDiffConsumer>> {
None
@@ -181,14 +164,6 @@ trait AnyToolHandler: Send + Sync {
fn is_mutating<'a>(&'a self, invocation: &'a ToolInvocation) -> BoxFuture<'a, bool>;
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload>;
fn with_updated_hook_input(
&self,
invocation: ToolInvocation,
updated_input: Value,
) -> Result<ToolInvocation, FunctionCallError>;
fn create_diff_consumer(&self) -> Option<Box<dyn ToolArgumentDiffConsumer>>;
fn handle_any<'a>(
&'a self,
@@ -208,18 +183,6 @@ where
Box::pin(ToolHandler::is_mutating(self, invocation))
}
fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {
ToolHandler::pre_tool_use_payload(self, invocation)
}
fn with_updated_hook_input(
&self,
invocation: ToolInvocation,
updated_input: Value,
) -> Result<ToolInvocation, FunctionCallError> {
ToolHandler::with_updated_hook_input(self, invocation, updated_input)
}
fn create_diff_consumer(&self) -> Option<Box<dyn ToolArgumentDiffConsumer>> {
ToolHandler::create_diff_consumer(self)
}
@@ -378,7 +341,7 @@ impl ToolRegistry {
return Err(err);
}
if let Some(pre_tool_use_payload) = handler.pre_tool_use_payload(&invocation) {
if let Some(pre_tool_use_payload) = hook_compat::pre_tool_use_payload(&invocation) {
match run_pre_tool_use_hooks(
&invocation.session,
&invocation.turn,
@@ -396,7 +359,7 @@ impl ToolRegistry {
crate::hook_runtime::PreToolUseHookResult::Continue {
updated_input: Some(updated_input),
} => {
invocation = handler.with_updated_hook_input(invocation, updated_input)?;
invocation = hook_compat::apply_updated_input(invocation, updated_input)?;
}
crate::hook_runtime::PreToolUseHookResult::Continue {
updated_input: None,