simplify hook input rewrites

This commit is contained in:
Abhinav Vedmala
2026-05-04 16:58:54 -07:00
parent bcccc6c208
commit 7932199328
7 changed files with 83 additions and 92 deletions
+2
View File
@@ -48,6 +48,8 @@ pub(crate) enum PreToolUseHookResult {
Blocked(String),
}
pub(crate) const MAX_HOOK_INPUT_REWRITES: usize = 8;
pub(crate) enum PendingInputHookDisposition {
Accepted(Box<PendingInputRecord>),
Blocked { additional_contexts: Vec<String> },
+2 -1
View File
@@ -25,6 +25,7 @@ use crate::guardian::guardian_timeout_message;
use crate::guardian::new_guardian_review_id;
use crate::guardian::review_approval_request;
use crate::guardian::routes_approval_to_guardian;
use crate::hook_runtime::MAX_HOOK_INPUT_REWRITES;
use crate::hook_runtime::run_permission_request_hooks;
use crate::mcp_openai_file::rewrite_mcp_tool_arguments_for_openai_files;
use crate::mcp_tool_approval_templates::RenderedMcpToolApprovalParam;
@@ -215,7 +216,7 @@ pub(crate) async fn handle_mcp_tool_call(
let result = match decision {
McpToolApprovalDecision::AcceptWithUpdatedInput(updated_input) => {
input_rewrites += 1;
if input_rewrites > 8 {
if input_rewrites > MAX_HOOK_INPUT_REWRITES {
notify_mcp_tool_call_skip(
sess.as_ref(),
turn_context.as_ref(),
+11 -27
View File
@@ -22,6 +22,8 @@ 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;
@@ -323,34 +325,16 @@ impl ToolHandler for ApplyPatchHandler {
mut invocation: ToolInvocation,
updated_input: serde_json::Value,
) -> Result<ToolInvocation, FunctionCallError> {
let patch = updated_input
.get("command")
.and_then(serde_json::Value::as_str)
.ok_or_else(|| {
FunctionCallError::RespondToModel(
"hook returned updatedInput without string field `command`".to_string(),
)
})?;
let patch = updated_hook_command(&updated_input)?;
invocation.payload = match invocation.payload {
ToolPayload::Function { arguments } => {
let mut arguments: serde_json::Value = parse_arguments(&arguments)?;
let serde_json::Value::Object(arguments) = &mut arguments else {
return Err(FunctionCallError::RespondToModel(
"apply_patch arguments must be an object".to_string(),
));
};
arguments.insert(
"input".to_string(),
serde_json::Value::String(patch.to_string()),
);
ToolPayload::Function {
arguments: serde_json::to_string(&arguments).map_err(|err| {
FunctionCallError::RespondToModel(format!(
"failed to serialize rewritten apply_patch arguments: {err}"
))
})?,
}
}
ToolPayload::Function { arguments } => ToolPayload::Function {
arguments: rewrite_function_string_argument(
&arguments,
"apply_patch",
"input",
patch,
)?,
},
ToolPayload::Custom { .. } => ToolPayload::Custom {
input: patch.to_string(),
},
+43
View File
@@ -25,6 +25,7 @@ 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;
@@ -74,6 +75,48 @@ 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<(), FunctionCallError>,
) -> 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()));
Ok(())
})
}
fn resolve_workdir_base_path(
arguments: &str,
default_cwd: &AbsolutePathBuf,
+15 -40
View File
@@ -25,6 +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;
@@ -75,17 +78,6 @@ fn shell_command_payload_command(payload: &ToolPayload) -> Option<String> {
.map(|params| params.command)
}
fn updated_hook_command(updated_input: &JsonValue) -> Result<&str, FunctionCallError> {
updated_input
.get("command")
.and_then(JsonValue::as_str)
.ok_or_else(|| {
FunctionCallError::RespondToModel(
"hook returned updatedInput without string field `command`".to_string(),
)
})
}
struct RunExecLikeArgs {
tool_name: String,
exec_params: ExecParams,
@@ -236,21 +228,13 @@ impl ToolHandler for ShellHandler {
"hook returned shell input with an invalid command string".to_string(),
)
})?;
let mut arguments: JsonValue = parse_arguments(&arguments)?;
let JsonValue::Object(arguments) = &mut arguments else {
return Err(FunctionCallError::RespondToModel(
"shell arguments must be an object".to_string(),
));
};
arguments.insert(
"command".to_string(),
JsonValue::Array(command.into_iter().map(JsonValue::String).collect()),
);
ToolPayload::Function {
arguments: serde_json::to_string(&arguments).map_err(|err| {
FunctionCallError::RespondToModel(format!(
"failed to serialize rewritten shell arguments: {err}"
))
arguments: rewrite_function_arguments(&arguments, "shell", |arguments| {
arguments.insert(
"command".to_string(),
JsonValue::Array(command.into_iter().map(JsonValue::String).collect()),
);
Ok(())
})?,
}
}
@@ -391,22 +375,13 @@ impl ToolHandler for ShellCommandHandler {
"hook input rewrite received unsupported shell_command payload".to_string(),
));
};
let mut arguments: JsonValue = parse_arguments(&arguments)?;
let JsonValue::Object(arguments) = &mut arguments else {
return Err(FunctionCallError::RespondToModel(
"shell_command arguments must be an object".to_string(),
));
};
arguments.insert(
"command".to_string(),
JsonValue::String(updated_hook_command(&updated_input)?.to_string()),
);
invocation.payload = ToolPayload::Function {
arguments: serde_json::to_string(&arguments).map_err(|err| {
FunctionCallError::RespondToModel(format!(
"failed to serialize rewritten shell_command arguments: {err}"
))
})?,
arguments: rewrite_function_string_argument(
&arguments,
"shell_command",
"command",
updated_hook_command(&updated_input)?,
)?,
};
Ok(invocation)
}
@@ -14,6 +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_workdir_base_path;
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;
@@ -162,30 +164,13 @@ impl ToolHandler for UnifiedExecHandler {
"hook input rewrite received unsupported unified_exec payload".to_string(),
));
};
let command = updated_input
.get("command")
.and_then(serde_json::Value::as_str)
.ok_or_else(|| {
FunctionCallError::RespondToModel(
"hook returned updatedInput without string field `command`".to_string(),
)
})?;
let mut arguments: serde_json::Value = parse_arguments(&arguments)?;
let serde_json::Value::Object(arguments) = &mut arguments else {
return Err(FunctionCallError::RespondToModel(
"exec_command arguments must be an object".to_string(),
));
};
arguments.insert(
"cmd".to_string(),
serde_json::Value::String(command.to_string()),
);
invocation.payload = ToolPayload::Function {
arguments: serde_json::to_string(&arguments).map_err(|err| {
FunctionCallError::RespondToModel(format!(
"failed to serialize rewritten exec_command arguments: {err}"
))
})?,
arguments: rewrite_function_string_argument(
&arguments,
"exec_command",
"cmd",
updated_hook_command(&updated_input)?,
)?,
};
Ok(invocation)
}
+2 -1
View File
@@ -5,6 +5,7 @@ use std::time::Instant;
use crate::function_tool::FunctionCallError;
use crate::goals::GoalRuntimeEvent;
use crate::hook_runtime::MAX_HOOK_INPUT_REWRITES;
use crate::hook_runtime::record_additional_contexts;
use crate::hook_runtime::run_post_tool_use_hooks;
use crate::hook_runtime::run_pre_tool_use_hooks;
@@ -231,7 +232,7 @@ where
match self.handle(invocation.clone()).await {
Err(FunctionCallError::UpdatedInput(updated_input)) => {
rewrites += 1;
if rewrites > 8 {
if rewrites > MAX_HOOK_INPUT_REWRITES {
return Err(FunctionCallError::RespondToModel(
"hook input rewrite limit exceeded".to_string(),
));