use std::collections::BTreeSet;use std::collections::HashMap;use std::path::Path;use std::path::PathBuf;use std::sync::Arc;use std::time::Duration;use std::time::Instant;use crate::apply_patch;use crate::apply_patch::InternalApplyPatchInvocation;use crate::apply_patch::convert_apply_patch_to_protocol;use crate::function_tool::FunctionCallError;use crate::session::session::Session;use crate::session::turn_context::TurnContext;use crate::session::turn_context::TurnEnvironment;use crate::tools::context::ApplyPatchToolOutput;use crate::tools::context::FunctionToolOutput;use crate::tools::context::SharedTurnDiffTracker;use crate::tools::context::ToolInvocation;use crate::tools::context::ToolPayload;use crate::tools::context::boxed_tool_output;use crate::tools::events::ToolEmitter;use crate::tools::events::ToolEventCtx;use crate::tools::handlers::apply_granted_turn_permissions;use crate::tools::handlers::apply_patch_spec::create_apply_patch_freeform_tool;use crate::tools::handlers::resolve_tool_environment;use crate::tools::handlers::updated_hook_command;use crate::tools::hook_names::HookToolName;use crate::tools::orchestrator::ToolOrchestrator;use crate::tools::registry::CoreToolRuntime;use crate::tools::registry::PostToolUsePayload;use crate::tools::registry::PreToolUsePayload;use crate::tools::registry::ToolArgumentDiffConsumer;use crate::tools::registry::ToolExecutor;use crate::tools::runtimes::apply_patch::ApplyPatchRequest;use crate::tools::runtimes::apply_patch::ApplyPatchRuntime;use crate::tools::sandboxing::ToolCtx;use codex_apply_patch::ApplyPatchAction;use codex_apply_patch::ApplyPatchFileChange;use codex_apply_patch::Hunk;use codex_apply_patch::StreamingPatchParser;use codex_exec_server::ExecutorFileSystem;use codex_features::Feature;use codex_protocol::models::AdditionalPermissionProfile;use codex_protocol::models::FileSystemPermissions;use codex_protocol::protocol::EventMsg;use codex_protocol::protocol::FileChange;use codex_protocol::protocol::PatchApplyUpdatedEvent;use codex_sandboxing::policy_transforms::effective_file_system_sandbox_policy;use codex_sandboxing::policy_transforms::merge_permission_profiles;use codex_sandboxing::policy_transforms::normalize_additional_permissions;use codex_tools::ToolName;use codex_tools::ToolSpec;use codex_utils_absolute_path::AbsolutePathBuf;use codex_utils_path_uri::PathUri;const APPLY_PATCH_ARGUMENT_DIFF_BUFFER_INTERVAL: Duration = Duration::from_millis(500);/// Handles freeform `apply_patch` requests and routes verified patches to the/// selected environment filesystem.#[derive(Default)]pub struct ApplyPatchHandler {    multi_environment: bool,}impl ApplyPatchHandler {    pub(crate) fn new(multi_environment: bool) -> Self {        Self { multi_environment }    }}#[derive(Default)]struct ApplyPatchArgumentDiffConsumer {    parser: StreamingPatchParser,    last_sent_at: Option<Instant>,    pending: Option<PatchApplyUpdatedEvent>,}impl ToolArgumentDiffConsumer for ApplyPatchArgumentDiffConsumer {    fn consume_diff(        &mut self,        turn: &TurnContext,        call_id: String,        diff: &str,    ) -> Option<EventMsg> {        if !turn.features.enabled(Feature::ApplyPatchStreamingEvents) {            return None;        }        self.push_delta(call_id, diff)            .map(EventMsg::PatchApplyUpdated)    }    fn finish(&mut self) -> Result<Option<EventMsg>, FunctionCallError> {        self.finish_update_on_complete()            .map(|event| event.map(EventMsg::PatchApplyUpdated))    }}impl ApplyPatchArgumentDiffConsumer {    fn push_delta(&mut self, call_id: String, delta: &str) -> Option<PatchApplyUpdatedEvent> {        let hunks = self.parser.push_delta(delta).ok()?;        if hunks.is_empty() {            return None;        }        let changes = convert_apply_patch_hunks_to_protocol(&hunks);        let event = PatchApplyUpdatedEvent { call_id, changes };        let now = Instant::now();        match self.last_sent_at {            Some(last_sent_at)                if now.duration_since(last_sent_at) < APPLY_PATCH_ARGUMENT_DIFF_BUFFER_INTERVAL =>            {                self.pending = Some(event);                None            }            Some(_) | None => {                self.pending = None;                self.last_sent_at = Some(now);                Some(event)            }        }    }    fn finish_update_on_complete(        &mut self,    ) -> Result<Option<PatchApplyUpdatedEvent>, FunctionCallError> {        self.parser.finish().map_err(|err| {            FunctionCallError::RespondToModel(format!("failed to parse apply_patch: {err}"))        })?;        let event = self.pending.take();        if event.is_some() {            self.last_sent_at = Some(Instant::now());        }        Ok(event)    }}fn convert_apply_patch_hunks_to_protocol(hunks: &[Hunk]) -> HashMap<PathBuf, FileChange> {    hunks        .iter()        .map(|hunk| {            let path = hunk_source_path(hunk).to_path_buf();            let change = match hunk {                Hunk::AddFile { contents, .. } => FileChange::Add {                    content: contents.clone(),                },                Hunk::DeleteFile { .. } => FileChange::Delete {                    content: String::new(),                },                Hunk::UpdateFile {                    chunks, move_path, ..                } => FileChange::Update {                    unified_diff: format_update_chunks_for_progress(chunks),                    move_path: move_path.clone(),                },            };            (path, change)        })        .collect()}fn hunk_source_path(hunk: &Hunk) -> &Path {    match hunk {        Hunk::AddFile { path, .. } | Hunk::DeleteFile { path } | Hunk::UpdateFile { path, .. } => {            path        }    }}fn format_update_chunks_for_progress(chunks: &[codex_apply_patch::UpdateFileChunk]) -> String {    let mut unified_diff = String::new();    for chunk in chunks {        match &chunk.change_context {            Some(context) => {                unified_diff.push_str("@@ ");                unified_diff.push_str(context);                unified_diff.push('\n');            }            None => {                unified_diff.push_str("@@");                unified_diff.push('\n');            }        }        for line in &chunk.old_lines {            unified_diff.push('-');            unified_diff.push_str(line);            unified_diff.push('\n');        }        for line in &chunk.new_lines {            unified_diff.push('+');            unified_diff.push_str(line);            unified_diff.push('\n');        }        if chunk.is_end_of_file {            unified_diff.push_str("*** End of File");            unified_diff.push('\n');        }    }    unified_diff}fn file_paths_for_action(action: &ApplyPatchAction) -> Vec<AbsolutePathBuf> {    let mut keys = Vec::new();    let cwd = &action.cwd;    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: &AbsolutePathBuf, path: &Path) -> Option<AbsolutePathBuf> {    Some(AbsolutePathBuf::resolve_path_against_base(path, cwd))}fn write_permissions_for_paths(    file_paths: &[AbsolutePathBuf],    file_system_sandbox_policy: &codex_protocol::permissions::FileSystemSandboxPolicy,    cwd: &AbsolutePathBuf,) -> Option<AdditionalPermissionProfile> {    let write_paths = file_paths        .iter()        .map(|path| {            path.parent()                .unwrap_or_else(|| path.clone())                .into_path_buf()        })        .filter(|path| {            !file_system_sandbox_policy.can_write_path_with_cwd(path.as_path(), cwd.as_path())        })        .collect::<BTreeSet<_>>()        .into_iter()        .map(AbsolutePathBuf::from_absolute_path)        .collect::<Result<Vec<_>, _>>()        .ok()?;    let permissions = (!write_paths.is_empty()).then_some(AdditionalPermissionProfile {        file_system: Some(FileSystemPermissions::from_read_write_roots(            Some(vec![]),            Some(write_paths),        )),        ..Default::default()    })?;    normalize_additional_permissions(permissions).ok()}/// Extracts the raw patch text used as the command-shaped hook input for apply_patch.fn apply_patch_payload_command(payload: &ToolPayload) -> Option<String> {    match payload {        ToolPayload::Custom { input } => Some(input.clone()),        _ => None,    }}async fn effective_patch_permissions(    session: &Session,    turn: &TurnContext,    environment_id: &str,    action: &ApplyPatchAction,    cwd: &AbsolutePathBuf,) -> (    Vec<AbsolutePathBuf>,    crate::tools::handlers::EffectiveAdditionalPermissions,    codex_protocol::permissions::FileSystemSandboxPolicy,) {    let file_paths = file_paths_for_action(action);    let granted_permissions = merge_permission_profiles(        session            .granted_session_permissions(environment_id)            .await            .as_ref(),        session            .granted_turn_permissions(environment_id)            .await            .as_ref(),    );    let base_file_system_sandbox_policy = turn.file_system_sandbox_policy();    let file_system_sandbox_policy = effective_file_system_sandbox_policy(        &base_file_system_sandbox_policy,        granted_permissions.as_ref(),    );    let effective_additional_permissions = apply_granted_turn_permissions(        session,        environment_id,        cwd.as_path(),        crate::sandboxing::SandboxPermissions::UseDefault,        write_permissions_for_paths(&file_paths, &file_system_sandbox_policy, cwd),    )    .await;    (        file_paths,        effective_additional_permissions,        file_system_sandbox_policy,    )}impl ToolExecutor<ToolInvocation> for ApplyPatchHandler {    fn tool_name(&self) -> ToolName {        ToolName::plain("apply_patch")    }    fn spec(&self) -> ToolSpec {        create_apply_patch_freeform_tool(self.multi_environment)    }    fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> {        Box::pin(self.handle_call(invocation))    }}impl ApplyPatchHandler {    async fn handle_call(        &self,        invocation: ToolInvocation,    ) -> Result<Box<dyn crate::tools::context::ToolOutput>, FunctionCallError> {        let ToolInvocation {            session,            turn,            tracker,            call_id,            tool_name,            payload,            ..        } = invocation;        let ToolPayload::Custom { input: patch_input } = payload else {            return Err(FunctionCallError::RespondToModel(                "apply_patch handler received unsupported payload".to_string(),            ));        };        let args = match codex_apply_patch::parse_patch(&patch_input) {            Ok(args) => args,            Err(parse_error) => {                return Err(FunctionCallError::RespondToModel(format!(                    "apply_patch verification failed: {parse_error}"                )));            }        };        let selected_environment_id =            require_environment_id(args.environment_id.as_deref(), self.multi_environment)?;        // Verify the parsed patch against the selected environment filesystem.        let Some(turn_environment) =            resolve_tool_environment(turn.as_ref(), selected_environment_id.as_deref())?        else {            return Err(FunctionCallError::RespondToModel(                "apply_patch is unavailable in this session".to_string(),            ));        };        // TODO(anp): Migrate apply-patch verification and permission accounting to PathUri so        // patches can target environment-native foreign paths without host projection.        let cwd = turn_environment.cwd().to_abs_path().map_err(|err| {            FunctionCallError::RespondToModel(format!(                "apply_patch cwd `{}` is not native to the Codex host: {err}",                turn_environment.cwd()            ))        })?;        let fs = turn_environment.environment.get_filesystem();        let sandbox = turn.file_system_sandbox_context(            /*additional_permissions*/ None,            turn_environment.cwd(),        );        match codex_apply_patch::verify_apply_patch_args(args, &cwd, fs.as_ref(), Some(&sandbox))            .await        {            codex_apply_patch::MaybeApplyPatchVerified::Body(changes) => {                let (file_paths, effective_additional_permissions, file_system_sandbox_policy) =                    effective_patch_permissions(                        session.as_ref(),                        turn.as_ref(),                        &turn_environment.environment_id,                        &changes,                        &cwd,                    )                    .await;                match apply_patch::apply_patch(turn.as_ref(), &file_system_sandbox_policy, changes)                    .await                {                    InternalApplyPatchInvocation::Output(item) => {                        let content = item?;                        Ok(boxed_tool_output(ApplyPatchToolOutput::from_text(content)))                    }                    InternalApplyPatchInvocation::DelegateToRuntime(apply) => {                        let changes = convert_apply_patch_to_protocol(&apply.action);                        let emitter = ToolEmitter::apply_patch_for_environment(                            changes.clone(),                            apply.auto_approved,                            turn_environment.environment_id.clone(),                        );                        let event_ctx = ToolEventCtx::new(                            session.as_ref(),                            turn.as_ref(),                            &call_id,                            Some(&tracker),                        );                        emitter.begin(event_ctx).await;                        let req = ApplyPatchRequest {                            turn_environment: turn_environment.clone(),                            action: apply.action,                            file_paths,                            changes,                            exec_approval_requirement: apply.exec_approval_requirement,                            additional_permissions: effective_additional_permissions                                .additional_permissions,                            permissions_preapproved: effective_additional_permissions                                .permissions_preapproved,                        };                        let mut orchestrator = ToolOrchestrator::new();                        let mut runtime = ApplyPatchRuntime::new();                        let tool_ctx = ToolCtx {                            session: session.clone(),                            turn: turn.clone(),                            call_id: call_id.clone(),                            tool_name: tool_name.clone(),                        };                        let out = orchestrator                            .run(                                &mut runtime,                                &req,                                &tool_ctx,                                turn.as_ref(),                                turn.approval_policy.value(),                            )                            .await                            .map(|result| result.output);                        let (out, delta) = match out {                            Ok(output) => (Ok(output.exec_output), Some(output.delta)),                            Err(error) => (Err(error), Some(runtime.committed_delta().clone())),                        };                        let event_ctx = ToolEventCtx::new(                            session.as_ref(),                            turn.as_ref(),                            &call_id,                            Some(&tracker),                        );                        let content = emitter.finish(event_ctx, out, delta.as_ref()).await?;                        Ok(boxed_tool_output(ApplyPatchToolOutput::from_text(content)))                    }                }            }            codex_apply_patch::MaybeApplyPatchVerified::CorrectnessError(parse_error) => {                Err(FunctionCallError::RespondToModel(format!(                    "apply_patch verification failed: {parse_error}"                )))            }            codex_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(),                ))            }            codex_apply_patch::MaybeApplyPatchVerified::NotApplyPatch => {                Err(FunctionCallError::RespondToModel(                    "apply_patch handler received non-apply_patch input".to_string(),                ))            }        }    }}impl CoreToolRuntime for ApplyPatchHandler {    fn matches_kind(&self, payload: &ToolPayload) -> bool {        matches!(payload, ToolPayload::Custom { .. })    }    fn create_diff_consumer(&self) -> Option<Box<dyn ToolArgumentDiffConsumer>> {        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 }),        })    }    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::Custom { .. } => ToolPayload::Custom {                input: patch.to_string(),            },            payload => payload,        };        Ok(invocation)    }    fn post_tool_use_payload(        &self,        invocation: &ToolInvocation,        result: &dyn crate::tools::context::ToolOutput,    ) -> Option<PostToolUsePayload> {        let tool_response =            result.post_tool_use_response(&invocation.call_id, &invocation.payload)?;        Some(PostToolUsePayload {            tool_name: HookToolName::apply_patch(),            tool_use_id: invocation.call_id.clone(),            tool_input: serde_json::json!({                "command": apply_patch_payload_command(&invocation.payload)?,            }),            tool_response,        })    }}#[allow(clippy::too_many_arguments)]pub(crate) async fn intercept_apply_patch(    command: &[String],    cwd: &AbsolutePathBuf,    fs: &dyn ExecutorFileSystem,    turn_environment: TurnEnvironment,    session: Arc<Session>,    turn: Arc<TurnContext>,    tracker: Option<&SharedTurnDiffTracker>,    call_id: &str,    tool_name: &str,) -> Result<Option<FunctionToolOutput>, FunctionCallError> {    let sandbox_cwd = PathUri::from_abs_path(cwd);    let sandbox =        turn.file_system_sandbox_context(/*additional_permissions*/ None, &sandbox_cwd);    match codex_apply_patch::maybe_parse_apply_patch_verified(command, cwd, fs, Some(&sandbox))        .await    {        codex_apply_patch::MaybeApplyPatchVerified::Body(changes) => {            let (approval_keys, effective_additional_permissions, file_system_sandbox_policy) =                effective_patch_permissions(                    session.as_ref(),                    turn.as_ref(),                    &turn_environment.environment_id,                    &changes,                    cwd,                )                .await;            match apply_patch::apply_patch(turn.as_ref(), &file_system_sandbox_policy, changes)                .await            {                InternalApplyPatchInvocation::Output(item) => {                    let content = item?;                    Ok(Some(FunctionToolOutput::from_text(content, Some(true))))                }                InternalApplyPatchInvocation::DelegateToRuntime(apply) => {                    let changes = convert_apply_patch_to_protocol(&apply.action);                    let emitter = ToolEmitter::apply_patch_for_environment(                        changes.clone(),                        apply.auto_approved,                        turn_environment.environment_id.clone(),                    );                    let event_ctx = ToolEventCtx::new(                        session.as_ref(),                        turn.as_ref(),                        call_id,                        tracker.as_ref().copied(),                    );                    emitter.begin(event_ctx).await;                    let req = ApplyPatchRequest {                        turn_environment,                        action: apply.action,                        file_paths: approval_keys,                        changes,                        exec_approval_requirement: apply.exec_approval_requirement,                        additional_permissions: effective_additional_permissions                            .additional_permissions,                        permissions_preapproved: effective_additional_permissions                            .permissions_preapproved,                    };                    let mut orchestrator = ToolOrchestrator::new();                    let mut runtime = ApplyPatchRuntime::new();                    let tool_ctx = ToolCtx {                        session: session.clone(),                        turn: turn.clone(),                        call_id: call_id.to_string(),                        tool_name: ToolName::plain(tool_name),                    };                    let out = orchestrator                        .run(                            &mut runtime,                            &req,                            &tool_ctx,                            turn.as_ref(),                            turn.approval_policy.value(),                        )                        .await                        .map(|result| result.output);                    let (out, delta) = match out {                        Ok(output) => (Ok(output.exec_output), Some(output.delta)),                        Err(error) => (Err(error), Some(runtime.committed_delta().clone())),                    };                    let event_ctx = ToolEventCtx::new(                        session.as_ref(),                        turn.as_ref(),                        call_id,                        tracker.as_ref().copied(),                    );                    let content = emitter.finish(event_ctx, out, delta.as_ref()).await?;                    Ok(Some(FunctionToolOutput::from_text(content, Some(true))))                }            }        }        codex_apply_patch::MaybeApplyPatchVerified::CorrectnessError(parse_error) => {            Err(FunctionCallError::RespondToModel(format!(                "apply_patch verification failed: {parse_error}"            )))        }        codex_apply_patch::MaybeApplyPatchVerified::ShellParseError(error) => {            tracing::trace!("Failed to parse apply_patch input, {error:?}");            Ok(None)        }        codex_apply_patch::MaybeApplyPatchVerified::NotApplyPatch => Ok(None),    }}fn require_environment_id(    parsed_environment_id: Option<&str>,    allow_environment_id: bool,) -> Result<Option<String>, FunctionCallError> {    match parsed_environment_id {        Some(_) if !allow_environment_id => Err(FunctionCallError::RespondToModel(            "apply_patch environment selection is unavailable for this turn".to_string(),        )),        Some(environment_id) => Ok(Some(environment_id.to_string())),        None => Ok(None),    }}#[cfg(test)]#[path = "apply_patch_tests.rs"]mod tests;