Add unified exec escalation handling and tests (#6492)

Similar implementation to the shell tool

codex-rs/core/src/tools/handlers/unified_exec.rs

@@ -36,6 +36,10 @@ struct ExecCommandArgs {
     yield_time_ms: Option<u64>,
     #[serde(default)]
     max_output_tokens: Option<usize>,
+    #[serde(default)]
+    with_escalated_permissions: Option<bool>,
+    #[serde(default)]
+    justification: Option<String>,
 }
 
 #[derive(Debug, Deserialize)]
@@ -100,8 +104,30 @@ impl ToolHandler for UnifiedExecHandler {
                         "failed to parse exec_command arguments: {err:?}"
                     ))
                 })?;
-                let workdir = args
-                    .workdir
+                let ExecCommandArgs {
+                    cmd,
+                    workdir,
+                    shell,
+                    login,
+                    yield_time_ms,
+                    max_output_tokens,
+                    with_escalated_permissions,
+                    justification,
+                } = args;
+
+                if with_escalated_permissions.unwrap_or(false)
+                    && !matches!(
+                        context.turn.approval_policy,
+                        codex_protocol::protocol::AskForApproval::OnRequest
+                    )
+                {
+                    return Err(FunctionCallError::RespondToModel(format!(
+                        "approval policy is {policy:?}; reject command — you cannot ask for escalated permissions if the approval policy is {policy:?}",
+                        policy = context.turn.approval_policy
+                    )));
+                }
+
+                let workdir = workdir
                     .as_deref()
                     .filter(|value| !value.is_empty())
                     .map(PathBuf::from);
@@ -113,18 +139,20 @@ impl ToolHandler for UnifiedExecHandler {
                     &context.call_id,
                     None,
                 );
-                let emitter = ToolEmitter::unified_exec(args.cmd.clone(), cwd.clone(), true);
+                let emitter = ToolEmitter::unified_exec(cmd.clone(), cwd.clone(), true);
                 emitter.emit(event_ctx, ToolEventStage::Begin).await;
 
                 manager
                     .exec_command(
                         ExecCommandRequest {
-                            command: &args.cmd,
-                            shell: &args.shell,
-                            login: args.login,
-                            yield_time_ms: args.yield_time_ms,
-                            max_output_tokens: args.max_output_tokens,
+                            command: &cmd,
+                            shell: &shell,
+                            login,
+                            yield_time_ms,
+                            max_output_tokens,
                             workdir,
+                            with_escalated_permissions,
+                            justification,
                         },
                         &context,
                     )

codex-rs/core/src/tools/runtimes/unified_exec.rs

@@ -34,6 +34,8 @@ pub struct UnifiedExecRequest {
     pub command: Vec<String>,
     pub cwd: PathBuf,
     pub env: HashMap<String, String>,
+    pub with_escalated_permissions: Option<bool>,
+    pub justification: Option<String>,
 }
 
 impl ProvidesSandboxRetryData for UnifiedExecRequest {
@@ -49,6 +51,7 @@ impl ProvidesSandboxRetryData for UnifiedExecRequest {
 pub struct UnifiedExecApprovalKey {
     pub command: Vec<String>,
     pub cwd: PathBuf,
+    pub escalated: bool,
 }
 
 pub struct UnifiedExecRuntime<'a> {
@@ -56,8 +59,20 @@ pub struct UnifiedExecRuntime<'a> {
 }
 
 impl UnifiedExecRequest {
-    pub fn new(command: Vec<String>, cwd: PathBuf, env: HashMap<String, String>) -> Self {
-        Self { command, cwd, env }
+    pub fn new(
+        command: Vec<String>,
+        cwd: PathBuf,
+        env: HashMap<String, String>,
+        with_escalated_permissions: Option<bool>,
+        justification: Option<String>,
+    ) -> Self {
+        Self {
+            command,
+            cwd,
+            env,
+            with_escalated_permissions,
+            justification,
+        }
     }
 }
 
@@ -84,6 +99,7 @@ impl Approvable<UnifiedExecRequest> for UnifiedExecRuntime<'_> {
         UnifiedExecApprovalKey {
             command: req.command.clone(),
             cwd: req.cwd.clone(),
+            escalated: req.with_escalated_permissions.unwrap_or(false),
         }
     }
 
@@ -98,7 +114,10 @@ impl Approvable<UnifiedExecRequest> for UnifiedExecRuntime<'_> {
         let call_id = ctx.call_id.to_string();
         let command = req.command.clone();
         let cwd = req.cwd.clone();
-        let reason = ctx.retry_reason.clone();
+        let reason = ctx
+            .retry_reason
+            .clone()
+            .or_else(|| req.justification.clone());
         let risk = ctx.risk.clone();
         Box::pin(async move {
             with_cached_approval(&session.services, key, || async move {
@@ -116,7 +135,16 @@ impl Approvable<UnifiedExecRequest> for UnifiedExecRuntime<'_> {
         policy: AskForApproval,
         sandbox_policy: &SandboxPolicy,
     ) -> bool {
-        requires_initial_appoval(policy, sandbox_policy, &req.command, false)
+        requires_initial_appoval(
+            policy,
+            sandbox_policy,
+            &req.command,
+            req.with_escalated_permissions.unwrap_or(false),
+        )
+    }
+
+    fn wants_escalated_first_attempt(&self, req: &UnifiedExecRequest) -> bool {
+        req.with_escalated_permissions.unwrap_or(false)
     }
 }
 
@@ -127,8 +155,15 @@ impl<'a> ToolRuntime<UnifiedExecRequest, UnifiedExecSession> for UnifiedExecRunt
         attempt: &SandboxAttempt<'_>,
         _ctx: &ToolCtx<'_>,
     ) -> Result<UnifiedExecSession, ToolError> {
-        let spec = build_command_spec(&req.command, &req.cwd, &req.env, None, None, None)
-            .map_err(|_| ToolError::Rejected("missing command line for PTY".to_string()))?;
+        let spec = build_command_spec(
+            &req.command,
+            &req.cwd,
+            &req.env,
+            None,
+            req.with_escalated_permissions,
+            req.justification.clone(),
+        )
+        .map_err(|_| ToolError::Rejected("missing command line for PTY".to_string()))?;
         let exec_env = attempt
             .env_for(&spec)
             .map_err(|err| ToolError::Codex(err.into()))?;

codex-rs/core/src/tools/spec.rs

@@ -177,6 +177,24 @@ fn create_exec_command_tool() -> ToolSpec {
             ),
         },
     );
+    properties.insert(
+        "with_escalated_permissions".to_string(),
+        JsonSchema::Boolean {
+            description: Some(
+                "Whether to request escalated permissions. Set to true if command needs to be run without sandbox restrictions"
+                    .to_string(),
+            ),
+        },
+    );
+    properties.insert(
+        "justification".to_string(),
+        JsonSchema::String {
+            description: Some(
+                "Only set if with_escalated_permissions is true. 1-sentence explanation of why we want to run this command."
+                    .to_string(),
+            ),
+        },
+    );
 
     ToolSpec::Function(ResponsesApiTool {
         name: "exec_command".to_string(),