From caab5a19ee9d13399c8fb7c40b213ec8872f5b0d Mon Sep 17 00:00:00 2001 From: jimmyfraiture Date: Fri, 26 Sep 2025 14:46:07 +0200 Subject: [PATCH] Move some stuff around --- codex-rs/core/src/apply_patch.rs | 1 + codex-rs/core/src/codex.rs | 103 ++++++-------------------- codex-rs/core/src/sandbox/mod.rs | 3 +- codex-rs/core/src/sandbox/planner.rs | 105 +++++++++++++++++++++++++++ 4 files changed, 131 insertions(+), 81 deletions(-) diff --git a/codex-rs/core/src/apply_patch.rs b/codex-rs/core/src/apply_patch.rs index 1ebbe5d738..8f372fb2d3 100644 --- a/codex-rs/core/src/apply_patch.rs +++ b/codex-rs/core/src/apply_patch.rs @@ -27,6 +27,7 @@ pub(crate) enum InternalApplyPatchInvocation { DelegateToExec(ApplyPatchExec), } +#[derive(Debug)] pub(crate) struct ApplyPatchExec { pub(crate) action: ApplyPatchAction, pub(crate) user_explicitly_approved_this_action: bool, diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 8bdb497906..20127318b2 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -114,12 +114,9 @@ use crate::rollout::RolloutRecorder; use crate::rollout::RolloutRecorderParams; use crate::sandbox::BackendRegistry; use crate::sandbox::ExecPlan; -use crate::sandbox::ExecRequest; use crate::sandbox::ExecRuntimeContext; -use crate::sandbox::PatchExecRequest; -use crate::sandbox::build_exec_params_for_apply_patch; -use crate::sandbox::plan_apply_patch; -use crate::sandbox::plan_exec; +use crate::sandbox::PreparedExec; +use crate::sandbox::prepare_exec_invocation; use crate::sandbox::run_with_plan; use crate::shell; use crate::state::ActiveTurn; @@ -2753,83 +2750,29 @@ async fn handle_container_exec_with_params( MaybeApplyPatchVerified::NotApplyPatch => None, }; - let mut params = params; - let command_for_display; - - let plan = if let Some(apply_patch_exec) = apply_patch_exec.as_ref() { - params = build_exec_params_for_apply_patch(apply_patch_exec, ¶ms)?; - command_for_display = vec![ - "apply_patch".to_string(), - apply_patch_exec.action.patch.clone(), - ]; - - let plan_req = PatchExecRequest { - action: &apply_patch_exec.action, - approval: turn_context.approval_policy, - policy: &turn_context.sandbox_policy, - cwd: &turn_context.cwd, - user_explicitly_approved: apply_patch_exec.user_explicitly_approved_this_action, - }; - - match plan_apply_patch(&plan_req) { - plan @ ExecPlan::Approved { .. } => plan, - ExecPlan::AskUser { .. } => { - return Err(FunctionCallError::RespondToModel( - "patch requires approval but none was recorded".to_string(), - )); - } - ExecPlan::Reject { reason } => { - return Err(FunctionCallError::RespondToModel(format!( - "patch rejected: {reason}" - ))); - } - } - } else { - command_for_display = params.command.clone(); - - let initial_plan = { - let state = sess.state.lock().await; - plan_exec(&ExecRequest { - params: ¶ms, - approval: turn_context.approval_policy, - policy: &turn_context.sandbox_policy, - approved_session_commands: state.approved_commands_ref(), - }) - }; - - match initial_plan { - plan @ ExecPlan::Approved { .. } => plan, - ExecPlan::AskUser { reason } => { - let decision = sess - .request_command_approval( - sub_id.clone(), - call_id.clone(), - params.command.clone(), - params.cwd.clone(), - reason, - ) - .await; - match decision { - ReviewDecision::Approved => ExecPlan::approved(SandboxType::None, false, true), - ReviewDecision::ApprovedForSession => { - sess.add_approved_command(params.command.clone()).await; - ExecPlan::approved(SandboxType::None, false, true) - } - ReviewDecision::Denied | ReviewDecision::Abort => { - return Err(FunctionCallError::RespondToModel( - "exec command rejected by user".to_string(), - )); - } - } - } - ExecPlan::Reject { reason } => { - return Err(FunctionCallError::RespondToModel(format!( - "exec command rejected: {reason:?}" - ))); - } - } + let approved_session_commands = { + let state = sess.state.lock().await; + state.approved_commands_ref().clone() }; + let prepared = prepare_exec_invocation( + sess, + turn_context, + &sub_id, + &call_id, + params, + apply_patch_exec, + approved_session_commands, + ) + .await?; + + let PreparedExec { + params, + plan, + command_for_display, + apply_patch_exec, + } = prepared; + let exec_command_context = ExecCommandContext { sub_id: sub_id.clone(), call_id: call_id.clone(), diff --git a/codex-rs/core/src/sandbox/mod.rs b/codex-rs/core/src/sandbox/mod.rs index abd735992d..d1553f0d17 100644 --- a/codex-rs/core/src/sandbox/mod.rs +++ b/codex-rs/core/src/sandbox/mod.rs @@ -2,7 +2,6 @@ mod apply_patch_adapter; mod backend; mod planner; -pub(crate) use apply_patch_adapter::build_exec_params_for_apply_patch; pub use backend::BackendRegistry; pub use backend::DirectBackend; pub use backend::LinuxBackend; @@ -11,8 +10,10 @@ pub use backend::SpawnBackend; pub use planner::ExecPlan; pub use planner::ExecRequest; pub use planner::PatchExecRequest; +pub(crate) use planner::PreparedExec; pub use planner::plan_apply_patch; pub use planner::plan_exec; +pub(crate) use planner::prepare_exec_invocation; use crate::error::Result; use crate::exec::ExecParams; diff --git a/codex-rs/core/src/sandbox/planner.rs b/codex-rs/core/src/sandbox/planner.rs index fcb1412ab9..2799087871 100644 --- a/codex-rs/core/src/sandbox/planner.rs +++ b/codex-rs/core/src/sandbox/planner.rs @@ -3,9 +3,15 @@ use std::path::Path; use codex_apply_patch::ApplyPatchAction; +use super::apply_patch_adapter::build_exec_params_for_apply_patch; +use crate::apply_patch::ApplyPatchExec; +use crate::codex::Session; +use crate::codex::TurnContext; use crate::exec::ExecParams; use crate::exec::SandboxType; +use crate::function_tool::FunctionCallError; use crate::protocol::AskForApproval; +use crate::protocol::ReviewDecision; use crate::protocol::SandboxPolicy; use crate::safety::SafetyCheck; use crate::safety::assess_command_safety; @@ -101,6 +107,105 @@ pub fn plan_apply_patch(req: &PatchExecRequest<'_>) -> ExecPlan { } } +#[derive(Debug)] +pub(crate) struct PreparedExec { + pub(crate) params: ExecParams, + pub(crate) plan: ExecPlan, + pub(crate) command_for_display: Vec, + pub(crate) apply_patch_exec: Option, +} + +pub(crate) async fn prepare_exec_invocation( + sess: &Session, + turn_context: &TurnContext, + sub_id: &str, + call_id: &str, + params: ExecParams, + apply_patch_exec: Option, + approved_session_commands: HashSet>, +) -> Result { + let mut params = params; + + let (plan, command_for_display) = if let Some(exec) = apply_patch_exec.as_ref() { + params = build_exec_params_for_apply_patch(exec, ¶ms)?; + let command_for_display = vec!["apply_patch".to_string(), exec.action.patch.clone()]; + + let plan_req = PatchExecRequest { + action: &exec.action, + approval: turn_context.approval_policy, + policy: &turn_context.sandbox_policy, + cwd: &turn_context.cwd, + user_explicitly_approved: exec.user_explicitly_approved_this_action, + }; + + let plan = match plan_apply_patch(&plan_req) { + plan @ ExecPlan::Approved { .. } => plan, + ExecPlan::AskUser { .. } => { + return Err(FunctionCallError::RespondToModel( + "patch requires approval but none was recorded".to_string(), + )); + } + ExecPlan::Reject { reason } => { + return Err(FunctionCallError::RespondToModel(format!( + "patch rejected: {reason}" + ))); + } + }; + + (plan, command_for_display) + } else { + let command_for_display = params.command.clone(); + + let initial_plan = plan_exec(&ExecRequest { + params: ¶ms, + approval: turn_context.approval_policy, + policy: &turn_context.sandbox_policy, + approved_session_commands: &approved_session_commands, + }); + + let plan = match initial_plan { + plan @ ExecPlan::Approved { .. } => plan, + ExecPlan::AskUser { reason } => { + let decision = sess + .request_command_approval( + sub_id.to_string(), + call_id.to_string(), + params.command.clone(), + params.cwd.clone(), + reason, + ) + .await; + match decision { + ReviewDecision::Approved => ExecPlan::approved(SandboxType::None, false, true), + ReviewDecision::ApprovedForSession => { + sess.add_approved_command(params.command.clone()).await; + ExecPlan::approved(SandboxType::None, false, true) + } + ReviewDecision::Denied | ReviewDecision::Abort => { + return Err(FunctionCallError::RespondToModel( + "exec command rejected by user".to_string(), + )); + } + } + } + ExecPlan::Reject { reason } => { + return Err(FunctionCallError::RespondToModel(format!( + "exec command rejected: {reason:?}" + ))); + } + }; + + (plan, command_for_display) + }; + + Ok(PreparedExec { + params, + plan, + command_for_display, + apply_patch_exec, + }) +} + fn should_escalate_on_failure(approval: AskForApproval, sandbox: SandboxType) -> bool { matches!( (approval, sandbox),