use std::collections::BTreeMap;
use std::path::Path;
use crate::product::agent::apply_patch;
use crate::product::agent::apply_patch::InternalApplyPatchInvocation;
use crate::product::agent::apply_patch::convert_apply_patch_to_protocol;
use crate::product::agent::codex::Session;
use crate::product::agent::codex::TurnContext;
use crate::product::agent::function_tool::FunctionCallError;
use crate::product::agent::tools::context::SharedTurnDiffTracker;
use crate::product::agent::tools::context::ToolInvocation;
use crate::product::agent::tools::context::ToolOutput;
use crate::product::agent::tools::context::ToolPayload;
use crate::product::agent::tools::events::ToolEmitter;
use crate::product::agent::tools::events::ToolEventCtx;
use crate::product::agent::tools::handlers::parse_arguments;
use crate::product::agent::tools::orchestrator::ToolOrchestrator;
use crate::product::agent::tools::registry::ToolHandler;
use crate::product::agent::tools::registry::ToolKind;
use crate::product::agent::tools::runtimes::apply_patch::ApplyPatchRequest;
use crate::product::agent::tools::runtimes::apply_patch::ApplyPatchRuntime;
use crate::product::agent::tools::sandboxing::ToolCtx;
use crate::product::agent::tools::spec::ApplyPatchToolArgs;
use crate::product::agent::tools::spec::JsonSchema;
use crate::product::apply_patch::ApplyPatchAction;
use crate::product::apply_patch::ApplyPatchFileChange;
use crate::product::utils_absolute_path::AbsolutePathBuf;
use async_trait::async_trait;
use lha_llm::FreeformToolDescriptor;
use lha_llm::FreeformToolDescriptorFormat;
use lha_llm::FunctionToolDescriptor;
use lha_llm::ToolDescriptor;
pub struct ApplyPatchHandler;
const APPLY_PATCH_LARK_GRAMMAR: &str = include_str!("tool_apply_patch.lark");
fn file_paths_for_action(action: &ApplyPatchAction) -> Vec<AbsolutePathBuf> {
let mut keys = Vec::new();
let cwd = action.cwd.as_path();
for (path, change) in action.changes() {
if let Some(key) = to_abs_path(cwd, path) {
keys.push(key);
}
if let ApplyPatchFileChange::Update { move_path, .. } = change
&& let Some(dest) = move_path
&& let Some(key) = to_abs_path(cwd, dest)
{
keys.push(key);
}
}
keys
}
fn to_abs_path(cwd: &Path, path: &Path) -> Option<AbsolutePathBuf> {
AbsolutePathBuf::resolve_path_against_base(path, cwd).ok()
}
#[async_trait]
impl ToolHandler for ApplyPatchHandler {
fn kind(&self) -> ToolKind {
ToolKind::Function
}
fn matches_kind(&self, payload: &ToolPayload) -> bool {
matches!(
payload,
ToolPayload::Function { .. } | ToolPayload::Custom { .. }
)
}
async fn is_mutating(&self, _invocation: &ToolInvocation) -> bool {
true
}
async fn handle(&self, invocation: ToolInvocation) -> Result<ToolOutput, FunctionCallError> {
let ToolInvocation {
session,
turn,
tracker,
call_id,
tool_name,
payload,
} = invocation;
let patch_input = match payload {
ToolPayload::Function { arguments } => {
let args: ApplyPatchToolArgs = parse_arguments(&arguments)?;
args.input
}
ToolPayload::Custom { input } => input,
_ => {
return Err(FunctionCallError::RespondToModel(
"apply_patch handler received unsupported payload".to_string(),
));
}
};
let cwd = turn.cwd.clone();
let command = vec!["apply_patch".to_string(), patch_input.clone()];
match crate::product::apply_patch::maybe_parse_apply_patch_verified(&command, &cwd) {
crate::product::apply_patch::MaybeApplyPatchVerified::Body(changes) => {
match apply_patch::apply_patch(turn.as_ref(), changes).await {
InternalApplyPatchInvocation::Output(item) => {
let content = item?;
Ok(ToolOutput::Function {
content,
content_items: None,
success: Some(true),
})
}
InternalApplyPatchInvocation::DelegateToExec(apply) => {
let changes = convert_apply_patch_to_protocol(&apply.action);
let file_paths = file_paths_for_action(&apply.action);
let emitter =
ToolEmitter::apply_patch(changes.clone(), apply.auto_approved);
let event_ctx = ToolEventCtx::new(
session.as_ref(),
turn.as_ref(),
&call_id,
Some(&tracker),
);
emitter.begin(event_ctx).await;
let req = ApplyPatchRequest {
action: apply.action,
file_paths,
changes,
exec_approval_requirement: apply.exec_approval_requirement,
timeout_ms: None,
codex_exe: turn.codex_linux_sandbox_exe.clone(),
};
let mut orchestrator = ToolOrchestrator::new();
let mut runtime = ApplyPatchRuntime::new();
let tool_ctx = ToolCtx {
session: session.as_ref(),
turn: turn.as_ref(),
call_id: call_id.clone(),
tool_name: tool_name.to_string(),
};
let out = orchestrator
.run(&mut runtime, &req, &tool_ctx, &turn, turn.approval_policy)
.await;
let event_ctx = ToolEventCtx::new(
session.as_ref(),
turn.as_ref(),
&call_id,
Some(&tracker),
);
let content = emitter.finish(event_ctx, out).await?;
Ok(ToolOutput::Function {
content,
content_items: None,
success: Some(true),
})
}
}
}
crate::product::apply_patch::MaybeApplyPatchVerified::CorrectnessError(parse_error) => {
Err(FunctionCallError::RespondToModel(format!(
"apply_patch verification failed: {parse_error}"
)))
}
crate::product::apply_patch::MaybeApplyPatchVerified::ShellParseError(error) => {
tracing::trace!("Failed to parse apply_patch input, {error:?}");
Err(FunctionCallError::RespondToModel(
"apply_patch handler received invalid patch input".to_string(),
))
}
crate::product::apply_patch::MaybeApplyPatchVerified::NotApplyPatch => {
Err(FunctionCallError::RespondToModel(
"apply_patch handler received non-apply_patch input".to_string(),
))
}
}
}
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn intercept_apply_patch(
command: &[String],
cwd: &Path,
timeout_ms: Option<u64>,
session: &Session,
turn: &TurnContext,
tracker: Option<&SharedTurnDiffTracker>,
call_id: &str,
tool_name: &str,
) -> Result<Option<ToolOutput>, FunctionCallError> {
match crate::product::apply_patch::maybe_parse_apply_patch_verified(command, cwd) {
crate::product::apply_patch::MaybeApplyPatchVerified::Body(changes) => {
session
.record_model_warning(
format!("apply_patch was requested via {tool_name}. Use the apply_patch tool instead of exec_command."),
turn,
)
.await;
match apply_patch::apply_patch(turn, changes).await {
InternalApplyPatchInvocation::Output(item) => {
let content = item?;
Ok(Some(ToolOutput::Function {
content,
content_items: None,
success: Some(true),
}))
}
InternalApplyPatchInvocation::DelegateToExec(apply) => {
let changes = convert_apply_patch_to_protocol(&apply.action);
let approval_keys = file_paths_for_action(&apply.action);
let emitter = ToolEmitter::apply_patch(changes.clone(), apply.auto_approved);
let event_ctx =
ToolEventCtx::new(session, turn, call_id, tracker.as_ref().copied());
emitter.begin(event_ctx).await;
let req = ApplyPatchRequest {
action: apply.action,
file_paths: approval_keys,
changes,
exec_approval_requirement: apply.exec_approval_requirement,
timeout_ms,
codex_exe: turn.codex_linux_sandbox_exe.clone(),
};
let mut orchestrator = ToolOrchestrator::new();
let mut runtime = ApplyPatchRuntime::new();
let tool_ctx = ToolCtx {
session,
turn,
call_id: call_id.to_string(),
tool_name: tool_name.to_string(),
};
let out = orchestrator
.run(&mut runtime, &req, &tool_ctx, turn, turn.approval_policy)
.await;
let event_ctx =
ToolEventCtx::new(session, turn, call_id, tracker.as_ref().copied());
let content = emitter.finish(event_ctx, out).await?;
Ok(Some(ToolOutput::Function {
content,
content_items: None,
success: Some(true),
}))
}
}
}
crate::product::apply_patch::MaybeApplyPatchVerified::CorrectnessError(parse_error) => {
Err(FunctionCallError::RespondToModel(format!(
"apply_patch verification failed: {parse_error}"
)))
}
crate::product::apply_patch::MaybeApplyPatchVerified::ShellParseError(error) => {
tracing::trace!("Failed to parse apply_patch input, {error:?}");
Ok(None)
}
crate::product::apply_patch::MaybeApplyPatchVerified::NotApplyPatch => Ok(None),
}
}
pub(crate) fn create_apply_patch_freeform_tool() -> ToolDescriptor {
ToolDescriptor::Freeform(FreeformToolDescriptor {
name: "apply_patch".to_string(),
description: "Use the `apply_patch` tool to edit files. This is a FREEFORM tool, so do not wrap the patch in JSON.".to_string(),
format: FreeformToolDescriptorFormat {
r#type: "grammar".to_string(),
syntax: "lark".to_string(),
definition: APPLY_PATCH_LARK_GRAMMAR.to_string(),
},
})
}
pub(crate) fn create_apply_patch_json_tool() -> ToolDescriptor {
let mut properties = BTreeMap::new();
properties.insert(
"input".to_string(),
JsonSchema::String {
description: Some(r#"The entire contents of the apply_patch command"#.to_string()),
enum_values: None,
},
);
ToolDescriptor::Function(FunctionToolDescriptor {
name: "apply_patch".to_string(),
description: r#"Use the `apply_patch` tool to edit files.
Your patch language is a stripped‑down, file‑oriented diff format designed to be easy to parse and safe to apply. You can think of it as a high‑level envelope:
*** Begin Patch
[ one or more file sections ]
*** End Patch
Within that envelope, you get a sequence of file operations.
You MUST include a header to specify the action you are taking.
Each operation starts with one of three headers:
*** Add File: <path> - create a new file. Every following line is a + line (the initial contents).
*** Delete File: <path> - remove an existing file. Nothing follows.
*** Update File: <path> - patch an existing file in place (optionally with a rename).
May be immediately followed by *** Move to: <new path> if you want to rename the file.
Then one or more “hunks”, each introduced by @@ (optionally followed by a hunk header).
Within a hunk each line starts with:
For instructions on [context_before] and [context_after]:
- By default, show 3 lines of code immediately above and 3 lines immediately below each change. If a change is within 3 lines of a previous change, do NOT duplicate the first change’s [context_after] lines in the second change’s [context_before] lines.
- If 3 lines of context is insufficient to uniquely identify the snippet of code within the file, use the @@ operator to indicate the class or function to which the snippet belongs. For instance, we might have:
@@ class BaseClass
[3 lines of pre-context]
- [old_code]
+ [new_code]
[3 lines of post-context]
- If a code block is repeated so many times in a class or function such that even a single `@@` statement and 3 lines of context cannot uniquely identify the snippet of code, you can use multiple `@@` statements to jump to the right context. For instance:
@@ class BaseClass
@@ def method():
[3 lines of pre-context]
- [old_code]
+ [new_code]
[3 lines of post-context]
The full grammar definition is below:
Patch := Begin { FileOp } End
Begin := "*** Begin Patch" NEWLINE
End := "*** End Patch" NEWLINE
FileOp := AddFile | DeleteFile | UpdateFile
AddFile := "*** Add File: " path NEWLINE { "+" line NEWLINE }
DeleteFile := "*** Delete File: " path NEWLINE
UpdateFile := "*** Update File: " path NEWLINE [ MoveTo ] { Hunk }
MoveTo := "*** Move to: " newPath NEWLINE
Hunk := "@@" [ header ] NEWLINE { HunkLine } [ "*** End of File" NEWLINE ]
HunkLine := (" " | "-" | "+") text NEWLINE
A full patch can combine several operations:
*** Begin Patch
*** Add File: hello.txt
+Hello world
*** Update File: src/app.py
*** Move to: src/main.py
@@ def greet():
-print("Hi")
+print("Hello, world!")
*** Delete File: obsolete.txt
*** End Patch
It is important to remember:
- You must include a header with your intended action (Add/Delete/Update)
- You must prefix new lines with `+` even when creating a new file
- File references can only be relative, NEVER ABSOLUTE.
"#
.to_string(),
strict: false,
parameters: JsonSchema::Object {
properties,
required: Some(vec!["input".to_string()]),
additional_properties: Some(false.into()),
},
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::product::apply_patch::MaybeApplyPatchVerified;
use pretty_assertions::assert_eq;
use tempfile::TempDir;
#[test]
fn approval_keys_include_move_destination() {
let tmp = TempDir::new().expect("tmp");
let cwd = tmp.path();
std::fs::create_dir_all(cwd.join("old")).expect("create old dir");
std::fs::create_dir_all(cwd.join("renamed/dir")).expect("create dest dir");
std::fs::write(cwd.join("old/name.txt"), "old content\n").expect("write old file");
let patch = r#"*** Begin Patch
*** Update File: old/name.txt
*** Move to: renamed/dir/name.txt
@@
-old content
+new content
*** End Patch"#;
let argv = vec!["apply_patch".to_string(), patch.to_string()];
let action = match crate::product::apply_patch::maybe_parse_apply_patch_verified(&argv, cwd)
{
MaybeApplyPatchVerified::Body(action) => action,
other => panic!("expected patch body, got: {other:?}"),
};
let keys = file_paths_for_action(&action);
assert_eq!(keys.len(), 2);
}
}